nequip.integrations.ase

class nequip.integrations.ase.NequIPCalculator(model: Module, device: str | device, energy_units_to_eV: float = 1.0, length_units_to_A: float = 1.0, transforms: List[Callable] = [], **kwargs)[source]

NequIP framework ASE Calculator.

This ASE Calculator is compatible with models from the NequIP framework, including NequIP and Allegro models.

The recommended way to use this Calculator is with a compiled model, i.e. nequip-compile the model and load it into the Calculator with NequIPCalculator.from_compiled_model(...). If one uses --mode aotinductor during nequip-compile, it is important to use the flag --target ase for the compiled model file to work with this ASE Calculator.

Warning

If you are running MD with custom species, please make sure to set the correct masses for ASE.

Parameters:

  • model – a model in the NequIP framework

  • device (str/torch.device) – device for model to evaluate on, e.g. cpu or cuda

  • energy_units_to_eV (float) – energy conversion factor (default 1.0)

  • length_units_to_A (float) – length units conversion factor (default 1.0)

  • transforms (List[Callable]) – list of data transforms

implemented_properties: list[str] = ['energy', 'energies', 'forces', 'stress', 'free_energy']

Properties calculator can handle (energy, forces, …)

classmethod NequIPCalculator.from_compiled_model(compile_path: str | Path, device: str | device = 'cpu', chemical_species_to_atom_type_map: Dict[str, str] | bool | None = None, neighborlist_backend: str = 'matscipy', **kwargs)

Build an integration calculator from a compiled model artifact.

Parameters:

  • compile_path – path to the compiled model artifact.

  • device – device where the model is loaded and evaluated.

  • chemical_species_to_atom_type_map – optional chemical species mapping override.

  • neighborlist_backend – neighbor list backend used by neighbor transforms.

  • **kwargs – forwarded to the integration class constructor.