User Guide

Introduction

OpenMM-ML is a high level API for using machine learning interatomic potentials (MLIPs) in OpenMM simulations. With just a few lines of code, you can set up a simulation that uses a standard, pretrained model to represent some or all of the interactions in a system.

OpenMM-ML does not provide the models itself. It relies on other packages such as MACE, NequIP, TorchANI, and others for that. It provides a common interface for working with them, and handles all the details of connecting them to OpenMM.

Installation

The stable version of OpenMM-ML can be installed with pip.

pip install openmmml

This installs only OpenMM-ML itself, not the packages that provide specific models. Those must be installed separately by following the instructions from the package developers.

OpenMM-ML can also be installed with conda or mamba.

mamba install -c conda-forge openmm-ml

However, we generally recommend using pip because most of the packages that provide models are only available with pip, not conda.

Usage

The central class in OpenMM-ML is MLPotential. It represents an MLIP that can be used for simulations. In the simplest cases you just provide it the name of the pretrained potential function to use. You can then call createSystem() to create a System object for a simulation. For example,

from openmmml import MLPotential
potential = MLPotential('ani2x')
system = potential.createSystem(topology)

To use a model you provide yourself rather than a standard pretrained one, include the modelPath option to point to the file containing your model. For example,

potential = MLPotential('mace', modelPath='MACE.model')

Other options that are specific to particular types of models are described below.

Rather than simulating the entire model with an MLIP, you can use createMixedSystem() to create a System where part is modeled with the MLIP and the rest is modeled with a conventional force field. To do this, first create a System that is entirely modeled with the force field. Then call createMixedSystem(), providing it the list of atoms to replace with the MLIP.

As an example, suppose the Topology contains three chains. Chain 0 is a protein, chain 1 is a ligand, and chain 2 is solvent. The following code creates a System in which the internal energy of the ligand is computed with ANI2x, while everything else (including interactions between the ligand and the rest of the system) is computed with Amber14.

forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')
mm_system = forcefield.createSystem(topology)
chains = list(topology.chains())
ml_atoms = [atom.index for atom in chains[1].atoms()]
potential = MLPotential('ani2x')
ml_system = potential.createMixedSystem(topology, mm_system, ml_atoms)

See the API documentation for additional information about the arguments to createSystem() and createMixedSystem().

Supported Models

OpenMM-ML supports models created with a number of packages. For each one we list the supported model names (the first argument to the MLPotential constructor), as well as additional arguments that can be passed to the constructor and to createSystem().

MACE

The MACE package can be used to create models based on the MACE architecture. This includes both pretrained models and custom models you create yourself. The following model names are supported.

Name	Model
mace-off23-small mace-off23-medium mace-off23-large mace-off24-medium	Pretrained MACE-OFF models
mace-mpa-0-medium	Pretrained MACE-MPA-0 model
mace-omat-0-small mace-omat-0-medium	Pretrained MACE-OMAT-0 models
mace-omol-0-extra-large	Pretrained MACE-OMOL-0 model
mace	Custom MACE models specified with the modelPath argument

When creating MACE models, the following keyword arguments to the MLPotential constructor are supported.

Argument	Description
modelPath	For custom models, the path to the file containing the model

When using MACE models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
precision	The numerical precision of the model. Supported options are 'single' and 'double'. If None, the default precision of the model is used.
returnEnergyType	Whether to return the interaction energy or the energy including the self-energy. The default is 'interaction_energy'. Supported options are 'interaction_energy' and 'energy'.
device	The PyTorch device to perform calculations on, either a torch.device object or a string (such as 'cuda' or 'cpu'.) If omitted, a device is chosen automatically.
charge	The total charge of the system. If omitted, it is assumed to be 0. This is only used by MACE-OMOL-0. For other models it is ignored.
multiplicity	The spin multiplicity of the system. If omitted, it is assumed to be 1. This is only used by MACE-OMOL-0. For other models it is ignored.

AIMNet2

The aimnet package can be used to create models using the pretrained AIMNet2 potential. The following model names are supported.

Name	Model
aimnet2	Pretrained AIMNet2 models

When using AIMNet2 models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
charge	The total charge of the system. If omitted, it is assumed to be 0.
multiplicity	The spin multiplicity of the system. If omitted, it is assumed to be 1.

NequIP

The NequIP package can be used to create models based on the NequIP architecture. That includes models created with Allegro, which is an extension package for NequIP. The following model names are supported.

Name	Model
nequip	Custom NequIP or Allegro models specified with the modelPath argument

When creating NequIP models, the following keyword arguments to the MLPotential constructor are supported.

Argument	Description
modelPath	The path to the checkpoint file containing the model
lengthScale	The conversion factor from the model’s length units to nanometers (e.g. 0.1 if the model uses Angstroms)
energyScale	The conversion factor from the model’s energy units to kJ/mol (e.g. 4.184 if the model uses kcal/mol)

When using NequIP models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
precision	The numerical precision of the model. Supported options are 'single' and 'double'. If None, the default precision of the model is used.
atomTypes	A list of integers corresponding to the NequiIP 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.
device	The PyTorch device to perform calculations on, either a torch.device object or a string (such as 'cuda' or 'cpu'.) If omitted, a device is chosen automatically.

TorchANI

The TorchANI package can be used to create models using the pretrained ANI-1ccx and ANI-2x potentials.

Both ANI-1ccx and ANI-2x are ensembles of eight models. Averaging over all eight models leads to slightly more accurate results than any one individually. You can optionally use only a single model, which leads to a large improvement in speed at the cost of a small decrease in accuracy.

The following model names are supported.

Name	Model
ani1ccx	Pretrained ANI-1ccx model
ani2x	Pretrained ANI-2x model

When using TorchANI models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
modelIndex	The index of the model within the ensemble to use. If it is None, the full ensemble of eight models is used.
device	The PyTorch device to perform calculations on, either a torch.device object or a string (such as 'cuda' or 'cpu'.) If omitted, a device is chosen automatically.

DeePMD-kit

The DeePMD-kit package can be used to create models based on the Deep Potential architecture. The following model names are supported.

Name	Model
deepmd	Custom models specified with the model argument

When creating DeePMD-kit models, the following keyword arguments to the MLPotential constructor are supported.

Argument	Description
model	The path to the file containing the model
coordinatesCoefficient	The conversion factor between the model’s length units and nanometers. The default value is 10, corresponding to DeePMD-kit’s default units of Angstroms.
forceCoefficient	The conversion factor between the model’s force units and kJ/mol/nm. The default value is 964.8792534459, corresponding to DeePMD-kit’s default units of eV/Å.
energyCoefficient	The conversion factor between the model’s energy units and kJ/mol. The default value is 96.48792534459, corresponding to DeePMD-kit’s default units of eV.

When using DeePMD-kit models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
lambdaName	The name of a lambda parameter to use for alchemical calculations. The default value is None, which does not create a lambda parameter.
lambdaValue	The initial value of the lambda parameter

TorchMD-Net

The TorchMD-Net package can be used to create models.

Pretrained and your own local models are supported.

Name	Model
aceff-1.0	Pretrained AceFF-1.0 TensorNet model.
aceff-1.1	Pretrained AceFF-1.1 TensorNet model.
aceff-2.0	Pretrained AceFF-2.0 TensorNet2 model.
torchmdnet	Custom TorchMD-Net models specified with the modelPath argument.

When creating TorchMD-Net models, the following keyword arguments to the MLPotential constructor are supported.

Argument	Description
modelPath	The path to the file containing the model.
lengthScale	The conversion factor from the model’s length units to nanometers (e.g. 0.1 if the model uses Angstroms). If omitted the default is 0.1 (model in Angstroms).
energyScale	The conversion factor from the model’s energy units to kJ/mol (e.g. 4.184 if the model uses kcal/mol). If omitted the default is 96.4916 (model in eV).

When using TorchMD-Net models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
charge	The total charge of the system. If omitted, it is assumed to be 0.
coulomb_cutoff	The cutoff distance to apply to Coulomb interactions, in nanometers. If omitted, it defaults to 1.2 nm. For models without an explicit Coulomb term, this is ignored.
remove_ref_energy	Argument passed to the TorchMD-Net model, please see here. Default is True.
max_num_neighbors	Argument passed to the TorchMD-Net model, please see here. Default is the minimum of 64 or the number of atoms in the molecule.
batch	Argument passed to the forward call of the TorchMD-Net model, please see here. With this argument you can denote different atoms to be in different batches. The format should be a 1d list containing the batch index of each atom. e.g. for two molecules each with three atoms to be treated as seperate batches you would pass batch = [0, 0, 0, 1, 1, 1].

FeNNix

OpenMM-ML has an interface to the FeNNol package for simulating FeNNix models. The following model names are supported.

Name	Model
fennix-bio1-small	Pretrained FeNNix-Bio1(S) model.
fennix-bio1-medium	Pretrained FeNNix-Bio1(M) model.
fennix-bio1-small-finetune-ions	Pretrained FeNNix-Bio1(S) model finetuned to AMOEBA03 ions.
fennix-bio1-medium-finetune-ions	Pretrained FeNNix-Bio1(M) model finetuned to AMOEBA03 ions.
fennix	Use a custom FeNNix model loaded from a local file.

When creating FeNNix models, the following keyword arguments to the MLPotential constructor are supported.

Argument	Description
modelPath	The path to the file containing the model: required for fennix, ignored for all other models.

When using FeNNix models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
charge	The total charge of the system. If omitted, it is assumed to be 0.
precision	'single' for single precision (the default if not specified) or 'double' for double precision.

Orb

OpenMM-ML supports the Orb models from Orbital Materials. The following model names are supported.

Name	Model
orb-v3-conservative-inf-omat	Orb-v3 model pretrained on OMat24.
orb-v3-conservative-omol	Orb-v3 model pretrained on OMol25.

When using Orb models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
charge	The total charge of the system. If omitted, it is assumed to be 0.
multiplicity	The spin multiplicity of the system. If omitted, it is assumed to be 1.
device	The PyTorch device to perform calculations on, either a torch.device object or a string (such as 'cuda' or 'cpu'.) If omitted, a device is chosen automatically.

ASE

OpenMM-ML can use an arbitrary ASE Calculator to perform calculations. This allows it to use any model or code for which a Calculator is available, including a wide variety of MLIPs and quantum chemistry programs. Simply pass the Calculator to createSystem():

potential = MLPotential('ase')
system = potential.createSystem(topology, calculator=calculator)

An ASE Atoms object is created automatically based on the OpenMM Topology. You can optionally provide values to add to its info dict. Some Calculators use this as a mechanism to specify parameters like total charge and spin multiplicity:

system = potential.createSystem(topology, calculator=calculator, info={'charge':2})

In cases where you need more direct control, you can instead create an Atoms object yourself and provide it directly:

system = potential.createSystem(topology, aseAtoms=atoms)

In this case, the Atoms should already be fully configured, including having a Calculator set. It is up to you to make sure the Atoms and Topology are fully consistent with each other.

The following model names are supported.

Name	Model
ase	Custom models specified with either the calculator or aseAtoms argument

When using ASE models, the following extra keyword arguments to createSystem() and createMixedSystem() are supported.

Argument	Description
calculator	The Calculator to use for computations. An Atoms object is created automatically.
aseAtoms	An Atoms object to use for computations.
info	Values that should be added to the info dict of the Atoms object.

Other Packages

OpenMM-ML is based on a plugin architecture, allowing other packages to provide their own interfaces to it. The packages listed above are the ones for which OpenMM-ML has built in support. Other packages can interface to it by defining two classes that subclass MLPotentialImpl and MLPotentialImplFactory, then registering them by specifying an entry point in the group openmmml.potentials. Consult the documentation for other packages to see whether they provide interfaces for OpenMM-ML.