openmmml.models.nequippotential.NequIPPotentialImpl

class openmmml.models.nequippotential.NequIPPotentialImpl(name: str, modelPath: str, lengthScale: float, energyScale: float)

This is the MLPotentialImpl implementing support for E(3)-equivariant interatomic potentials generated by NequIP or Allegro.

The potential must be constructed using the NequIP/Allegro code to build and deploy a PyTorch model, which can then be integrated into OpenMM using a TorchForce. The TorchForce is used to add it to the OpenMM System. Note that you must provide a deployed model, as no general purpose model is available.

To use a deployed model in openmm-ml, you must provide the path to the model, and the conversion factors between the model length and energy units and OpenMM units (nm and kJ/mol, respectively). For example:

>>> potential = MLPotential('nequip',
                            modelPath='example_model_deployed.pth',
                            lengthScale=0.1, # Angstrom to nm
                            energyScale=4.184 # kcal/mol to kJ/mol
                            )

During system creation, if the model was trained with custom atom types, you can specify this by passing to the atomTypes parameter a list of of integers corresponding to the nequip atom type for each particle that will be modeled using this potential. This argument should, therefore, have the same length as the number of ML atoms in the system. Note that by default the model uses the atomic number to map the atom type. This will work if you trained your model using the standard chemical_symbols option.

Additionally, you can specify the precision of the model using the precision keyword argument. Supported options are ‘single’ and ‘double’. For example:

>>> system = potential.createSystem(topology, precision='single')

By default, the implementation uses the precision of the loaded model. Note that models deployed before NequIP v0.6.0 don’t contain information about their precision, so precision='double' should only be used if the model was explicitly trained with default_dtype=float64, as by default the model is trained with default_dtype=float32.

__init__(name: str, modelPath: str, lengthScale: float, energyScale: float) → None

Initialize the NequIPPotentialImpl.

Parameters:

• name (str) – The name specified by the MLPotential constructor, viz. ‘nequip’.

• modelPath (str, optional) – The path to the deployed model.

• lengthScale (float) – The length conversion factor from the model units to nanometers.

• energyScale (float) – The energy conversion factor from the model units to kJ/mol.

Methods

__init__(name, modelPath, lengthScale, ...)	Initialize the NequIPPotentialImpl.
addForces(topology, system, atoms, forceGroup)	Add the NequIPForce to the OpenMM System.

addForces(topology: Topology, system: System, atoms: Iterable[int] | None, forceGroup: int, precision: str | None = None, atomTypes: List[int] | None = None, **kwargs)

Add the NequIPForce to the OpenMM System.

Parameters:

• topology (openmm.app.Topology) – The topology of the system.

• system (openmm.System) – The system to which the force will be added.

• atoms (iterable of int) – The indices of the atoms to include in the model. If None, all atoms are included.

• forceGroup (int) – The force group to which the force should be assigned.

• precision (str, optional) – The precision of the model. Supported options are ‘single’ and ‘double’. If None, the default precision of the model is used. This is the recommended option. Models deployed before NequIP v0.6.0 don’t contain information about their precision, so precision='double' should only be used if the model was explicitly trained with default_dtype=float64, as by default the model is trained with default_dtype=float32.

• atomTypes (List[int], optional) – A list of integers corresponding to the nequip atom type for each ML atom in the system. This is only required if the model was trained with custom atom types. If None, the atomic number is used to determine the atom type. This list should have the same length as the number of ML atoms in the system.