🌟 [FEATURE] OpenMM

[asedova]

Is your feature request related to a problem? Please describe.
Using LAMMPS can be a bit of a pain.

Describe the solution you'd like
OpenMM is a really easy to use, Pythonic MD code that also allows a user-defined set of forces and runs ridiculously fast on one GPU. It seems they are also staring to incorporate some ML methods. It would be amazing if one could run NEQUIP MD with OpenMM.

Describe alternatives you've considered
ASE is an option, but wouldn't get the speed of a GPU-accelerated MD code. OpenMM will JIT user-defined forces to the GPU.

[asedova]

Left an question on OpenMM's github about how this would work, and here is the reply: openmm/openmm-ml#48 (comment) I am about to see how hard this would be for a high-level user to do themselves. Any help would be appreciated!

[Linux-cpp-lisp]

Hi @asedova ,

Thanks very much for your interest in our methods and working on this integration!

We previously had some discussion about this (#76) and the great news is that the main thing that was missing---a PyTorch neighborlist--- has now been implemented very nicely by Felix Musil with good performance and what appear to be pretty rigorous correctness tests: https://github.com/felixmusil/torch_nl. So it's now the perfect time for a "high-level user" to get this done finally! 🎉

Reading @peastman 's comment in the linked thread, I think the best way forward would be a 1-to-1 translation of the OpenMM-ML ANI interface https://github.com/openmm/openmm-ml/blob/main/openmmml/models/anipotential.py.

It would be something like:

• Load a deployed NequIP model using nequip.scripts.deploy.load_deployed_model (nequip.ase.NequIPCalculator is a good example of this); the result is a TorchScript model.
• Similar to ANIForce, define a wrapper around the deployed NequIP model that:
  • Uses torch_nl to construct edge_index for NequIP from the positions and, optionally, cell + PBC. It is important to set edge_cell_shift correctly as well, see here for how we do this with the ase neighborlist. For complete PBC correctness, make sure to replicate the self-interaction filtering logic as well.
  • Builds a NequIP input dictionary (Dict[str, torch.Tensor], keys can be found here). Necessary keys are pos, edge_index, atom_types and, if PBC are used, cell and edge_cell_shift.
  • Somehow map OpenMM atom types / species into the NequIP nequip.data.AtomicDataDict.ATOM_TYPES_KEY 0-based index. See again (nequip.ase.NequIPCalculator for an example of how to do this. I'm not super familiar with how the OpenMM type machinery works, but basically there should be a user-defined mapping of OpenMM types into NequIP type names, which the code then internally uses to create a mapping of OpenMM types to NequIP type numbers.
  • Call the deployed NequIP model with this dict as input, and get the total energy and forces from the result: model(in_dict)[nequip.data.AtomicDataDict.TOTAL_ENERGY_KEY / FORCE_KEY]
  • Do any unit conversions according to the user's provided conversions
  • Return this total energy and forces, see https://github.com/openmm/openmm-torch#computing-forces-in-the-model
• torch.jit.script this whole thing, then torch.jit.save it to a tempfile and pass it along to OpenMM-Torch.

It would be very nice to see this PRed into OpenMM-ML; please let me know if you have more questions, and feel free to ask/open PRs (off develop) here if there are any additions to nequip needed to make it work.

[Linux-cpp-lisp]

OK @asedova I've edited the above to reflect more recent versions of OpenMM-Torch 🙂

[asedova]

Thanks so much! I am going to try to see what I can do. I just started looking at NEQUIP after some time with DeePMDkit, so it may take me some time to work through all this. So happy there is momentum in this direction!

[Linux-cpp-lisp]

Great! Let me know if you have any questions. It looks like a lot but I think it should hopefully be straightforward, mostly "copy-paste" style coding.

[peastman]

there should be a user-defined mapping of OpenMM types into NequIP type names, which the code then internally uses to create a mapping of OpenMM types to NequIP type numbers.

OpenMM doesn't have any fixed set of types. It treats atom types as a property of a particular force field. addForces() receives a Topology object which contains chemical information: elements, bonds, residue names, etc. No set of atom types have been applied to them.

How does NequIP define atom types?

[Linux-cpp-lisp]

It treats atom types as a property of a particular force field.

That sounds quite compatible.

In nequip, atom types are just 0-based integer indexes of which there can be arbitrarily many. Each is also associated with a string name for usability. What they mean is specific to a given trained model and how the user did the training, but typically the names just correspond to element symbols. So clearly there should be some kind of user-defined mapping from OpenMM "groups" of atoms (however that is mostly naturally defined) to nequip type names; the user is then responsible for making sure this mapping makes sense for the model they are using.

Similarly, nequip places no restriction on system of units, as long as they are self-consistent; a trained model has the units of the data used to train it. So it's also necessary to allow the user to supply the unit (conversion) from the units of the particular model they are using into kJ/mol for OpenMM.

[asedova]

@peastman what troubles me is that with NEQUIP (and other similar codes like DeePMDkit) there is no topology, no bonds, no residues. And for good reason-- since it's trained on DFT, these all can actually change during a simulation, just like in AIMD. I'm not super familiar with ANI, but with these codes, input file would just be atom type (element) + features in, then the model gets energy + forces on each atom. The most compatible traditional input file to use as an example would be an xyz file, like CP2K uses.

[Linux-cpp-lisp]

receives a Topology object which contains chemical information: elements, bonds, residue names, etc.

I think as a first solution, you can just ignore all of that, take the elements, convert them to chemical symbols, and use those as NequIP type names? For most models this will be sufficient... what is missing would be something like how do you specify a model that treats, say, the oxygen atoms in one molecule as different from the oxygen atoms in a second molecule. That is clearly some kind of a function of the topology, I'm just not sure in OpenMM language how you let the user specify that.

[asedova]

That would be ok-- for now, we can let the model decide how the oxygen will act in different contexts, which is what I planned to start with anyway. I also think OpenMM would be ok if the topology object is not coming from any of the normal topology files from packages like AMBER and GROMACS. OpenMM in general allows a good amount of user-defined things, it's just not how I normally use it so I need to do some reading up. Here's an example: https://gpantel.github.io/computational-method/LJsimulation/

[peastman]

what is missing would be something like how do you specify a model that treats, say, the oxygen atoms in one molecule as different from the oxygen atoms in a second molecule.

In any model that allows forming and breaking bonds, I think the only option is to equate atom types with elements. Everything else can change with time. If the molecules are required to stay fixed over a simulation, that does allow incorporating more information into the atom types. I've experimented with including formal charges and/or hybridization states into the definitions of types, and it does seem to help accuracy.

Most force fields in OpenMM identify atom types by matching templates to residues. That isn't required though. For example, OpenFF uses a completely different mechanism of assigning parameters based on matching SMIRKS patterns (and not involving atom types at all). For ML models, you can build your model however you want. The PyTorch model and the Python code to build it are both completely in your control.

[Linux-cpp-lisp]

👍 I vote then to start with just taking OpenMM element -> chemical symbol as NequIP type name, possibly with an option of letting the user provide their own function that maps OpenMM topology -> NequIP type names, since every trained NequIP model will be different in this way.