🌟 [FEATURE] OpenMM #288
Comments
|
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! |
|
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:
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 |
|
OK @asedova I've edited the above to reflect more recent versions of OpenMM-Torch 🙂 |
|
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! |
|
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. |
OpenMM doesn't have any fixed set of types. It treats atom types as a property of a particular force field. How does NequIP define atom types? |
That sounds quite compatible. In Similarly, |
|
@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. |
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. |
|
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/ |
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. |
|
👍 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. |
|
Hi @asedova , Just curious if you'd had any luck with this. Thanks! |
|
Hi @asedova @Linux-cpp-lisp how is progress on this going? I am an RSE working with OpenMM, I have time available to help write/write a wrapper in openmm-ml for NequIP. Let me know if you have any questions and how far you have got with this! |
|
Hey all- Progress is slow, since we have never used NEQUIP that much. We are much more familiar with DeePMD. I should be able to get back to this soon; we are finishing up some DeePMD work. @sef43 I was wondering if there would be a way to pass the virial/stress tensor to OpenMM, based on this: #291 (comment) |
|
@sef43 hi, thank you for reaching out! 👋 We'd love to see this happen, so your help and hopefully leadership on it as an RSE who knows OpenMM well would be fantastic. I think all my relevant thoughts are collected in the post above (#288 (comment)), but please feel free to get in touch either here or by email (see my GitHub profile) if there is more we should discuss. @asedova no worries; do you have any code we ought work off, or would it be better for @sef43 to start a fresh repo / fork to work on? (I'm assuming @sef43 that, as laid out above, all of this wrapper should live in the fork of openmm-ml and eventually be merged there, but if there is code that we think is more appropriate in our repos we can make that happen easily too.) |
|
I've just added a few more notes to the post above with information on how we can go about this, FYI. Regarding stress/virial @asedova, my understanding is that OpenMM doesn't support or need it: openmm/openmm-torch#91 (comment) |
|
I would suggest just starting with a basic box of water example first, as done in the nature comms paper. Our systems are really difficult and we are still working out some of the science. Has getting a condensed-phase NNP to work with OpenMM happened before? I am confused because it seems there is a lot of gas phase work within ANI/TorchANI |
|
Sure--- I think, however, that the science questions that you bring up are fairly decoupled from the software-level interfacing questions, which are fairly mechanical and binary (either they correctly transfer the predictions of the neural network potential or not, regardless of how useful / accurate that potential model is) At the software level, I think the biggest difficulty moving from gas to condensed phase is PBC, but |
|
I will start doing this then!
Yes I will make a fork with the aim that it ends up merged into openmm-ml main, It shouldn't require any changes to the nequip repo. You steps in the comment above are very clear thank you, should be most of what I need, I will let you know when I have questions. As @Linux-cpp-lisp has said there should be no major reasons why condensed phase simulations wouldn't work provided an appropriately trained neural network is used. |
|
Fantastic! Thanks again @sef43 , looking forward to having this!
Sounds good, feel free to reach out any time here or by email as is convenient for you. |
|
@Linux-cpp-lisp I have a WIP fork branch here: https://github.com/sef43/openmm-ml/tree/nequip I have done a first implementation. It works for the Toluene example, I have put an example here, which can be run in colab It uses a trained deployed model I ran using your example Things I still need to do:
Some issues:
pytorch error from using nequip.scripts.deploy.load_deployed_model: |
|
@sef43 fantastic! That was quick 😆 , thanks! Regarding the issues, sort of in order:
Two easy options would be:
Both should be in units of eV and Å. If it's easier, I'm happy to train a model quickly and email it to you as well if you prefer.
Am I right in understanding that in OpenMM atom type == atomic number, and that OpenMM only reasons about atoms in terms of their atomic numbers? Is that why you have it set up to possibly allow mapping multiple
👍 yep, NequIP/Allegro is entirely agnostic to units and preserves those of the training data (as you've seen), so a user-facing option in the OpenMM wrapper seems appropriate. A finally, some random things: https://github.com/sef43/openmm-ml/blob/nequip/openmmml/models/nequippotential.py#L128
|
I wouldn't read anything into what I have currently done, I need to understand a bit more about the nequip types! Each particle in OpenMM is an As I understand nequip types are zero indexed and user defined. So there can be multiple types for the same chemical element. OpenMM is not going to know about this from reading in files (it's own concept of atom types and atom classes is tied to which ForceField you are using). The most flexible way would be to explicitly specify an array of length Natoms containing the nequip types of each atom and pass this to the MLP constructor, but not very user friendly
Ah ok I will avoid PyTorch 1.12 then |
This is correct, yes.
I see--- that sounds like a useful, optional choice to have, but I agree that it is an unpleasant interface for the common cases, which are best handled with a mapping. Probably best to restrict the mapping to one-to-one atomic numbers to types and let more complicated cases be handled explicitly, though, to avoid easy but hard to debug issues. (Incidentally, the common case for which then sorts them by atomic number (https://github.com/mir-group/nequip/blob/main/nequip/data/transforms.py#L30-L36) and assigns consecutive atom types starting at zero. The other common case is |
|
We can generate atom types in pretty much any way we want. The only requirement is that it has to be possible to determine them from the arguments to |
|
I've now added PBC support, It seems to work for the I have set the NequIP types default behaviour to be a 1 to 1 mapping from the OpenMM atomic element/symbol, with the option of passing a list that explicitly defines the zero indexed NequIP types of each atom. And conversion factors from the model units to OpenMM units need to be given when constructing the potential. It should work on both GPU and CPU I have put an example of running a deployed toy EMT model with PBC and NPT: |
not quite, this line: needs to check the value rather than the key, the key always gets set to exist here, the value gets set to Also |
|
Absolutely right on both counts, thanks for catching this error @sef43 --- I'll get those fixed. |
|
Yep looks good, works for me! I think this is ready for people to try using, I have put a README here: https://github.com/sef43/openmm-ml/blob/nequip/examples/nequip/README.md Once we know it works for people we can think about PRing it into OpenMM-ML main, tagging @peastman for insight on how to proceede with that |
|
That looks really nice! If we want to get it into the main OpenMM-ML distribution, there are a couple of things we need to consider. First, it would be very helpful if all dependencies can be installed with conda. Otherwise we can't build and test the conda packages against them. The other issue is more of a conceptual question. OpenMM-ML is intended to provide general purpose, ready to use models, not just model architectures. For example, you can write potential = MLPotential('ani2x')and immediately start simulating your own system with that potential. What you've implemented just provides support for an architecture. It requires the user to provide their own model. That's also useful, for example if you want to create a special purpose model for your particular system, but it requires a lot more work from the user. The goal of OpenMM-ML is to make ML simulations just as easy to run as ones with conventional force fields. Are there any plans to create general purpose, ready to use models based on NequIP? |
|
@sef43 saw this issue of yours (felixmusil/torch_nl#1 (comment)) on |
Yes it will be! |
|
👍 Is there a reason not to just wrap the coordinates in PyTorch in the wrapper and keep using |
|
yes could just wrap, it is a cheap operation, but actually when running on GPU the naive calculation is significantly faster: I guess the simpler operations map much better to vectorised GPU code by the jit compiler/optimizer To reproduce both start getting out of memory errors for more than 20,000 atoms. |
|
This code shifts = torch.round(triu_deltas/cell_size)
wrapped_triu_deltas = triu_deltas - shifts*cell_sizeis only accurate for rectangular boxes. If you want to support triclinic boxes its a little more complicated. Still not very complicated though. |
Yes that is true, it adds some compute time. I have updated it to support OpenMM triclinic boxes. For completeness here is the benchmarks for the 3 cases:
Updated code here: https://colab.research.google.com/drive/1Vz4QN8Y_cEUkCL2YxY86QhSruYPxfwtr?usp=sharing @peastman this nicely demonstrates the speed of the NNPOps implementation! once openmm/NNPOps#92 is fixed and released to the conda-forge build we can default to using NNPOps if available |
|
Wow, very nice! What a plot... |
@peastman the types of systems that people using NEQUIP (or DeePMD or other such tools) are studying vary dramatically. For example our team is studying molten salts, organic solvents, ionic liquids, etc. I doubt there will be a totally generalizable general ML potential for any systems from the periodic table any time soon. Therefore, the fact that people have to train their own models is not a deal breaker by any means. We already do all that. Having a fast, Pythonic, efficient way to run the MD inference is going to be really helpful. Especially since we often need many rounds of active-learning iterative training (with potentially some enhanced sampling) before we get stuff right, like with reactive molecular systems. Here are some examples: https://www.sciencedirect.com/science/article/abs/pii/S092058612100136X |
|
Hi, I am taking over this task from @sef43. Could you please clarify if there are any changes to the NequIP code that impact this feature? I have been working on this implementation, and it is working fine for me. The major change is that the NNPops neighbour list is now being used. Perhaps worth discussing is the PyTorch version that is installed when following a standard installation like this: The output of From this, we can clearly see that the PyTorch version is greater than 1.11, and therefore, I see the following warning when running this example: Unfortunately, I don't think we can install the latest versions of NNPOps with PyTorch 1.11. As far as I understand, we would have to revert to NNPOps 0.2 for compatibility, which does not include the neighbor list. What's your recommendation in this scenario? Is there any progress in making NequIP compatible with PyTorch >= 2.1.0? Many thanks in advance. |
|
Hi @JMorado , (CC @cw-tan re neighborlists) Great to hear, thanks for your efforts on this! The main issue that remains unclear for newer PyTorch versions is this: #311. I think it is reasonable to try latest stable 2.* versions and see if this issue is present. You can see the potential workaround (as yet not entirely proved) here (https://github.com/mir-group/nequip/blob/develop/nequip/utils/_global_options.py#L85-L107), which may be necessary to implement in the OpenMM interface as well? I see that you use If you'd like to discuss in more detail, please feel free to email me and we can have a call. |
|
Hi @Linux-cpp-lisp, Thank you for the clarifications! Since the workaround is already implemented in I'm going to run a long MD simulation with one of the test systems and see how it behaves. I'll report the results later. |
|
@JMorado I agree, as long as you are using Looking forward to hear whether it works! |
|
@Linux-cpp-lisp, I have made tests using toluene and Si. The performance looks pretty stable both in terms of GPU utilization and simulation speed:
I'm using PyTorch version 2.1.2: $ micromamba list | grep torch
libtorch 2.1.2 cuda120_h2aa5df7_303 conda-forge
openmm-torch 1.4 cuda120py311h478a873_4 conda-forge
pytorch 2.1.2 cuda120_py311h25b6552_303 conda-forge
torchani 2.2.4 cuda120py311he2766f7_3 conda-forgeDo you have any idea on why the Si model yields ca. four times more efficient simulations, even though the system is larger? |
Hi @sef43 ! Could you please upload this example again. It is not in the repo at the moment. |
|
@JMorado Hi! I have tried to create a mixed system for toluene in explicit water using NequIP for the toluene and TIP3P for the water. However, only the single point energy is correct, there are some weird value for forces on water molecules and the simulation quickly failed. While ANI works well. I guess there are some problems for NequIP implementaion when you setoutputforces=True. I'm using the example from here: https://github.com/sef43/openmm-ml/blob/nequip/test/test_nequip.py Could you check that? |
|
@CalvinGe, thank you for reporting. There is indeed an issue with the forces for periodic systems, which appears to be resolved (at least for pure ML systems) if the positions are not downcasted to float32. I am still trying to determine if it's the same type of issue causing hybrid simulations to fail, or if it's simply a matter of initial conditions. I will report back once I have reached a conclusion. |
|
What was causing your simulations to fail was actually a different issue. Since The fix to this issue is here. Would you mind giving it a try? It works smoothly for me using the following script to run a simulation: import openmm as mm
import openmm.app as app
import openmm.unit as unit
from openmmml import MLPotential
from sys import stdout
# Load the prmtop and inpcrd files
prmtop = app.AmberPrmtopFile("toluene-explicit.prm7")
inpcrd = app.AmberInpcrdFile("toluene-explicit.rst7")
# Define the atoms to be treated with the ML model
mlAtoms = list(range(15))
# Create the MM and ML/MM systems
mmSystem = prmtop.createSystem(nonbondedMethod=app.PME, nonbondedCutoff=0.9 * unit.nanometers)
potential = MLPotential("nequip", modelPath="toluene-nequip.pth", lengthScale=0.1, energyScale=4.184)
mixedSystem = potential.createMixedSystem(prmtop.topology, mmSystem, mlAtoms, interpolate=False)
# Add a barostat
mixedSystem.addForce(mm.MonteCarloBarostat(1 * unit.atmosphere, 300 * unit.kelvin, 25))
# Define the integrator, platform, and simulation
integrator = mm.LangevinMiddleIntegrator(300 * unit.kelvin, 1.0 / unit.picoseconds, 1 * unit.femtoseconds)
platform = mm.Platform.getPlatformByName("CUDA")
simulation = app.Simulation(prmtop.topology, mixedSystem, integrator, platform)
# Set the positions and velocities
simulation.context.setPositions(inpcrd.positions)
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(300 * unit.kelvin)
# Run the simulation
simulation.reporters.append(
app.StateDataReporter(
stdout, 100, step=True, potentialEnergy=True, temperature=True, speed=True
)
)
simulation.step(1000) |
|
@JMorado It works well for my test! The forces for MM region in ML/MM now also matches with the pure MM forces. Thank you very much! |
|
Just to clarify, this means that the NequIP OpenMM integration supports ML/MM? Are there specific pieces of setup necessary to run in that mode? |
According to your plot this is a NequIP and an Allegro model, so I would not expect them to be directly comparable. |
Yes, it does support hybrid ML/MM simulations. There aren't really any specific pieces of setup required, just the standard way of first creating the MM and MLPotential, defining the ML atoms, and then setting up the mixed system. The procedure is described here in more detail. |
How do you deal with the interaction between ML and MM atoms? |
Mechanical embedding is used, meaning that the interaction between ML and MM atoms is treated at the MM level using standard Coulomb and Lennard-Jones 12-6 potentials if you're using class I FFs. |
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.
The text was updated successfully, but these errors were encountered: