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 a PyTorch model, which can then be integrated into OpenMM using a PythonForce. Note that you must provide a model, as no general purpose model is available.

To use a model in OpenMM-ML, you must provide the path to the model checkpoint file, 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.ckpt',
                            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 model checkpoint file.

  • 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.

createMixedSystem(topology, system, atoms, ...)

Creates a mixed system using a potential-specific embedding method.

getMLLongRange()

Returns whether ML interactions are local (False), or long-ranged (True) when periodic boundary conditions are present.

getSupportedEmbeddings()

Retrieves a list of names of supported embedding methods (for the creation of mixed ML/MM systems) specific to this potential.
addForces(topology: Topology, system: System, atoms: Iterable[int] | None, forceGroup: int, precision: str | None = None, atomTypes: List[int] | None = None, **args)#

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.

createMixedSystem(topology: Topology, system: System, atoms: list[int], forceGroup: int, interpolate: bool, embedding: str, **args) → System#

Creates a mixed system using a potential-specific embedding method.

This is invoked by MLPotential.createMixedSystem(). It will only be called with one of the names returned by getSupportedEmbeddings(). If subclasses support any potential-specific embedding methods, they must also provide implementations of those methods by overriding this method. If not, it does not need to be implemented.

Each embedding method is responsible for implementing interpolation; if interpolate is True, a global parameter “lambda_interpolate” should be present in the returned system, with the behavior as described by MLPotential.createMixedSystem().

Parameters:

  • topology (Topology) – the Topology for which to create a System

  • system (System) – a System that models the Topology with a conventional force field

  • atoms (Iterable[int]) – the indices of all atoms whose interactions should be computed with this potential

  • forceGroup (int) – the force group the ML potential’s Forces should be placed in

  • interpolate (bool) – if True, create a System that can smoothly interpolate between the conventional and ML potentials

  • embedding (str) – the name of the embedding method (will always be one in the list returned by the getSupportedEmbeddings() method)

  • args – any additional arguments for the potential or embedding method

Returns:

  • a newly created System object that uses this potential function and the

  • requested embedding method to model the Topology

getMLLongRange() → bool | None#

Returns whether ML interactions are local (False), or long-ranged (True) when periodic boundary conditions are present. This controls which interactions are included by some ML/MM embedding methods. Consult the documentation for each embedding method for more details. None can be returned if the nature of the interactions is not known by the MLPotentialImpl.

The default implementation of this method, used if a subclass does not override it, always returns None.

getSupportedEmbeddings() → list[str]#

Retrieves a list of names of supported embedding methods (for the creation of mixed ML/MM systems) specific to this potential. If one of these names is provided instead of providing the name of a generic embedding plugin, embedding will be performed by the potential itself.

This is invoked by MLPotential.createMixedSystem(). If a subclass wishes to define one or more potential-specific embedding methods, it should implement this method; otherwise, it does not need to.