Models

For each recipe, the models to evaluate are defined in the models.py file. In most cases, you can use mlipx.GenericASECalculator to access models. However, in certain scenarios, you may need to provide a custom calculator.

Below, we demonstrate how to write a custom calculator node for SevenCalc. While this is an example, note that SevenCalc could also be used with mlipx.GenericASECalculator.

Defining Models

Here is the content of a typical models.py file:

Content of models.py

import mlipx
from src import SevenCalc

mace_medium = mlipx.GenericASECalculator(
   module="mace.calculators",
   class_name="MACECalculator",
   device='auto',
   kwargs={
      "model_paths": "mace_models/y7uhwpje-medium.model",
   },
)

mace_agnesi = mlipx.GenericASECalculator(
   module="mace.calculators",
   class_name="MACECalculator",
   device='auto',
   kwargs={
      "model_paths": "mace_models/mace_mp_agnesi_medium.model",
   },
)

sevennet = SevenCalc(model='7net-0')

MODELS = {
   "mace_medm": mace_medium,
   "mace_agne": mace_agnesi,
   "7net": sevennet,
}

Custom Calculator Example

The SevenCalc class, used in the example above, is defined in src/__init__.py as follows:

Content of src/__init__.py

import dataclasses
from ase.calculators.calculator import Calculator

@dataclasses.dataclass
class SevenCalc:
   model: str

   def get_calculator(self, **kwargs) -> Calculator:
      from sevenn.sevennet_calculator import SevenNetCalculator
      sevennet = SevenNetCalculator(self.model, device='cpu')

      return sevennet

For more details, refer to the Build your own Graph section.

Updating Dataset Keys

In some cases, models may need to be defined to convert existing dataset keys into the format mlipx expects. For example, you may need to provide isolated atom energies or convert data where energies are stored as atoms.info['DFT_ENERGY'] and forces as atoms.arrays['DFT_FORCES'].

Here’s how to define a model for such a scenario:

import mlipx

REFERENCE = mlipx.UpdateFramesCalc(
    results_mapping={"energy": "DFT_ENERGY", "forces": "DFT_FORCES"},
    info_mapping={mlipx.abc.ASEKeys.isolated_energies.value: "isol_ene"},
)