search

openmm / openmm-ml

High level API for using machine learning models in OpenMM simulations
Other
80 stars 25 forks source link

Add MACE model support #61

Closed sef43 closed 7 months ago

sef43 commented 12 months ago

This adds support for MACE models: https://github.com/ACEsuit/mace.

The MACE model framework is well defined. With this PR a user can run an OpenMM simulation using a MACE model trained with the scripts in https://github.com/ACEsuit/mace.

cc @jharrymoore @jmichel80

JMorado commented 7 months ago

I have been working on this PR and would like to discuss the best approach for incorporating the computation of the neighbour list. There are 3 possible options:

  1. Utilise the getNeighborPairs function from the NNPOps package.
  2. Utilise the simple_nl function introduced in the utils module, which is included with this PR.
  3. Include both and leave the choice to the user.

What do you think would be preferable? The first approach has the advantage of not needing additional modules to be added to openmm-ml, but it does make NNPOps a mandatory dependency of openmm-ml. Additionally, it necessitates recalculating the vectors between neighbors during the forward pass for shift computation. Perhaps integrating the computation of shifts into the NNPOps implementations could be beneficial, as it may also prove useful for other MLPs.

peastman commented 7 months ago

but it does make NNPOps a mandatory dependency of openmm-ml.

It's not mandatory: it's only needed if you use MACE. I think that's ok. I expect this package to eventually have interfaces to a lot of other libraries, each of which might or might not be installed. So for example if MACE is installed, you can create MACE models. And if it isn't installed, that doesn't stop you from creating other types of models.

Additionally, it necessitates recalculating the vectors between neighbors during the forward pass for shift computation.

Compared to the cost of computing the model, computing the vector between two atoms should be totally negligible!

JMorado commented 7 months ago

Sounds good. And I agree that the calculation of the vector should be negligible.

JMorado commented 7 months ago

This is ready for review. I've addressed most of the comments raised by @raimis. In addition to general refactoring, I have removed the utils.py module as now getNeighborPairs from NNPops is used.

JMorado commented 7 months ago

Thank you for the review, @peastman. All fixed now.

JMorado commented 7 months ago

I've also added a fix to ensure tensors are placed on the correct devices. I was only testing the implementation on the CPU and Reference platforms and was not having any issues. I've now conducted testing across all platforms, and everything should be functioning correctly.

peastman commented 7 months ago

Looks good. Is this ready to merge?

JMorado commented 7 months ago

Ready to merge.

peastman commented 7 months ago

Thanks!