Adding models¶

ML-PEG gets its model list from ml_peg/models/models.yml by default. Each top-level key is the model name used in calculation outputs, analysis, and the app. Calculation scripts load these entries with load_models and then call model.get_calculator(...) to obtain an ASE calculator.

You can either edit the default registry or keep local/private entries in a separate YAML file and pass it on the command line:

ml_peg calc --models-file my_models.yml --models my-mace-model
ml_peg analyse --models-file my_models.yml --models my-mace-model
ml_peg app --models-file my_models.yml --models my-mace-model

To check which model names ML-PEG can see, use:

ml_peg list models
ml_peg list models --models-file my_models.yml

Model entry format¶

A typical entry has this shape:

model-name:
  module: package.module
  class_name: CalculatorFactoryOrClass
  device: cuda
  trained_on_dispersion: false
  level_of_theory: PBE
  kwargs:
    option_name: option_value
  dispersion_kwargs:
    option_name: option_value

The common fields are:

module: Python module containing the calculator factory or class.
class_name: Name imported from module.
device: Device passed to calculators that support it. Options: cuda, cpu or auto.
trained_on_dispersion: Whether the model’s training data already included dispersion corrections. Some benchmark scripts call add_d3_calculator. If this field is false, ML-PEG may add a separate D3 correction for those benchmarks; if it is true, the base calculator is used unchanged.
level_of_theory: Functional or method represented by the model training data. The app compares this string against benchmark metric metadata and displays warnings when they differ. See Levels of theory for naming conventions.
kwargs: Keyword arguments forwarded to the calculator constructor or factory. Here you can input kwargs you would usually use for the calculator.
dispersion_kwargs: Optional settings used by add_d3_calculator when a benchmark adds a separate dispersion correction. These are option/value pairs passed through to the dispersion wrapper, so the accepted keys depend on that wrapper.
overwrite_dtype: Optional precision override for model wrappers that support it. Most benchmarks request either precision="high" or precision="low"; this field forces a specific dtype regardless of that request.

Examples¶

MACE-MP foundation model:

mace-mp-0a:
  module: mace.calculators
  class_name: mace_mp
  device: cuda
  trained_on_dispersion: false
  level_of_theory: PBE
  kwargs:
    model: medium

MACE checkpoint from a local path:

my-mace-model:
  module: mace.calculators
  class_name: mace_mp
  device: cuda
  trained_on_dispersion: false
  level_of_theory: PBE
  kwargs:
    model: /absolute/path/to/my_mace_checkpoint.model

Models with a specific head or task can pass that head in kwargs:

mace-mh-1-omol:
module: mace.calculators
class_name: mace_mp
device: "cuda"
trained_on_dispersion: true
level_of_theory: ωB97M-V/def2-TZVPD
kwargs:
   model: "mh-1"
   head: omol

uma-s-1p1-omol:
module: fairchem.core
class_name: FAIRChemCalculator
device: "cuda"
trained_on_dispersion: true
level_of_theory: ωB97M-V/def2-TZVPD
kwargs:
   model_name: "uma-s-1p1"
   task_name: "omol"

Supported model families¶

MACE entries use the generic ASE calculator wrapper. Common class names include mace_mp, mace_off, mace_omol and mace_polar.

MACE-Polar foundation model:

mace-polar-1-m:
  module: mace.calculators
  class_name: mace_polar
  device: cuda
  trained_on_dispersion: true
  level_of_theory: ωB97M-V
  kwargs:
    model: polar-1-m

ORB entries use the dedicated OrbCalc wrapper:

orb-v3-consv-inf-omat:
  module: orb_models.inference.calculator
  class_name: OrbCalc
  device: cuda
  trained_on_dispersion: false
  level_of_theory: PBE
  kwargs:
    name: orb_v3_conservative_inf_omat

FairChem/UMA entries use the dedicated FAIRChemCalculator branch:

uma-s-1p1-omat:
  module: fairchem.core
  class_name: FAIRChemCalculator
  device: cuda
  trained_on_dispersion: false
  level_of_theory: PBE
  kwargs:
    model_name: uma-s-1p1
    task_name: omat

PET-MAD entries use the dedicated PETMADCalculator branch:

pet-mad:
  module: pet_mad.calculator
  class_name: PETMADCalculator
  device: cuda
  trained_on_dispersion: false
  level_of_theory: PBEsol
  kwargs:
    version: v1.0.2
  dispersion_kwargs:
    xc: pbesol

Other ASE-compatible MLIP calculators can usually be added by specifying their module, class_name and constructor kwargs. That is enough when the calculator accepts the same common arguments as the generic ML-PEG model wrapper.

For a completely new model family, or for a calculator that needs special setup logic, add a small wrapper class in ml_peg/models/models.py and route to it from load_models in ml_peg/models/get_models.py. In those cases the YAML entry records the model configuration, while the Python wrapper defines how ML-PEG constructs the calculator.

Checking a new entry¶

After adding a model, run a small calculation first:

ml_peg list models --models-file my_models.yml
ml_peg calc --category molecular_crystal --test X23 --models-file my_models.yml --models my-mace-model

If loading fails, check that the optional dependency is installed, the module/class_name pair can be imported, local checkpoint paths are valid, and the model’s kwargs match the calculator API.