Add MACE model support #61

sef43 · 2023-10-05T16:07:34Z

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

raimis · 2023-10-09T07:27:58Z

openmmml/models/macepotential.py

+
+        import torch
+        import openmmtorch
+        from e3nn.util import jit


This should give a better error message than import failed. An explanation of missing MACE and pip command should be shown.

raimis · 2023-10-09T07:28:46Z

openmmml/models/macepotential.py

+                super(MACEForce, self).__init__()
+
+
+                if not torch.cuda.is_available():


OpenMM-Torch takes care of loading to the right devices. Here you just need to load on CPU.

raimis · 2023-10-09T07:32:02Z

openmmml/models/macepotential.py

+                # get the dtype of the saved model
+                model_dtype = [p.dtype for p in self.model.parameters()][0]
+
+                # check is user has request a different dtype    


The model should run on the precision with which it was trained. If OpenMM runs at a different precision (double rather float), one just casts inputs/outputs accordingly rather than change the model.

raimis · 2023-10-09T07:33:44Z

openmmml/models/macepotential.py

+                else:
+                    self.default_dtype = model_dtype
+
+                torch.set_default_dtype(self.default_dtype)


This a global PyTorch setting. It is better to set tensor precisions explicitly.

raimis · 2023-10-09T07:35:09Z

openmmml/models/macepotential.py

+
+                # setup input
+                N = len(atomic_numbers)
+                self.ptr = torch.nn.Parameter(torch.tensor([0,N], dtype=torch.long), requires_grad=False)


Why these are parameters rather than buffers?

raimis · 2023-10-09T07:40:44Z

openmmml/models/macepotential.py

+        if atoms is not None:
+            includedAtoms = [includedAtoms[i] for i in atoms]
+
+        class MACEForce(torch.nn.Module):


For better readability, don't nest the classes.

raimis · 2023-10-09T07:42:14Z

openmmml/models/macepotential.py

+                              "shifts": shifts}
+
+                # predict
+                out = self.model(input_dict,compute_force=False)


Why do the forces are computed here?

raimis · 2023-10-09T07:43:31Z

openmmml/models/macepotential.py

+                out = self.model(input_dict,compute_force=False)
+
+                energy = out[self.return_energy_type]
+                if energy is None:


How can you get energy None? The simulation should fail rather than silently assume it is 0.

raimis · 2023-10-09T07:44:59Z

openmmml/models/utils.py

+from typing import Tuple
+
+@torch.jit.script
+def simple_nl(positions: torch.Tensor, cell: torch.Tensor, pbc: bool, cutoff: float, sorti: bool=False) -> Tuple[torch.Tensor, torch.Tensor]:


Do we really need to invent a new neighbor list rather than use the one in NNPOps?

JMorado · 2024-02-27T19:21:58Z

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:

Utilise the getNeighborPairs function from the NNPOps package.
Utilise the simple_nl function introduced in the utils module, which is included with this PR.
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 · 2024-02-27T19:56:36Z

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 · 2024-02-27T20:07:35Z

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

JMorado · 2024-02-28T12:34:48Z

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.

peastman

This looks great! Just a few very minor comments.

peastman · 2024-02-28T19:48:41Z

openmmml/models/macepotential.py

+    self-energy, instead of the default interaction energy, by setting `returnEnergyType` to
+    `energy`. For example:
+    >>> system = potential.createSystem(topology, returnEnergyType=`energy`)


Here and in similar places, you probably meant to use single quotes rather than back quotes? Make it clear that 'energy' is a string literal

peastman · 2024-02-28T19:52:48Z

openmmml/models/macepotential.py

+        # Convert it to TorchScript and save it.
+        module = torch.jit.script(maceForce)
+        module.save(filename)
+
+        # Create the TorchForce and add it to the System.
+        force = openmmtorch.TorchForce(filename)


Saving to a file isn't needed anymore. TorchForce can now take a model directly. You can get rid of the filename argument and just write

force = openmmtorch.TorchForce(module)

Should this be changed here too?

openmm-ml/openmmml/models/anipotential.py

Lines 142 to 149 in 09be77f

# Convert it to TorchScript and save it.

module = torch.jit.script(aniForce)

module.save(filename)

# Create the TorchForce and add it to the System.

force = openmmtorch.TorchForce(filename)

Yes, it would probably be good to clean that up eventually.

Ok, I'll take care of it.

peastman · 2024-02-28T19:57:42Z

openmmml/models/macepotential.py

+        atoms: Optional[Iterable[int]],
+        forceGroup: int,
+        precision: Optional[str] = None,
+        returnEnergyType: str = "energy",


Here you set the default to 'energy', but in the docstring you say the default is 'interaction_energy'. Probably interaction energy is a better default.

…ry model writing, and corrected default value of returnEnergyType parameter

JMorado · 2024-02-29T10:28:40Z

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

JMorado · 2024-02-29T12:39:43Z

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.

openmmml/models/macepotential.py

peastman · 2024-02-29T16:59:20Z

Looks good. Is this ready to merge?

JMorado · 2024-03-04T18:07:45Z

Ready to merge.

peastman · 2024-03-04T18:09:58Z

Thanks!

sef43 added 11 commits February 23, 2023 16:13

implement MACE

fb5eb67

change NL to fix PBCs

217f330

improve simple_nl code

f2f1143

triclinic simple_nl

c27e3d3

triclinic pbc

26f01d1

update to openmm-torch 1.1 syntax

7f27816

Merge branch 'openmm:main' into mace

94fbd96

update

b3f048c

cleanup file

0e22700

update notebook

aa6ee0f

fix colab install

d5ff44e

raimis suggested changes Oct 9, 2023

View reviewed changes

raimis requested a review from peastman October 9, 2023 07:45

danielhollas mentioned this pull request Oct 30, 2023

Adding OpenMM as an optional MD engine duartegroup/mlp-train#55

Merged

JMorado added 3 commits February 28, 2024 12:27

Overall refactor of macepotential.py

1b844d5

Removed deprecated utils.py module

ecabaf5

Fixed formatting

3756b44

JMorado added 2 commits February 28, 2024 12:37

Removed examples as new ones will be created for the new implementation

8870825

Removed import of utils

5041b17

peastman reviewed Feb 28, 2024

View reviewed changes

Fixed improper use of backquotes and single quotes, removed unnecessa…

1482cba

…ry model writing, and corrected default value of returnEnergyType parameter

Fix to ensure correct device allocation

8dfc4fc

peastman reviewed Feb 29, 2024

View reviewed changes

openmmml/models/macepotential.py Outdated Show resolved Hide resolved

Fixed typo

1bd4dd7

JMorado added a commit to JMorado/openmm-ml that referenced this pull request Mar 1, 2024

openmm#61 Removed need for file saving in ANIPotentialImpl

5069496

JMorado mentioned this pull request Mar 1, 2024

Remove file saving in ANIPotentialImp #71

Merged

peastman pushed a commit that referenced this pull request Mar 1, 2024

#61 Removed need for file saving in ANIPotentialImpl (#71)

90984d5

peastman merged commit 4cc85f4 into openmm:main Mar 4, 2024

This pull request was closed.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add MACE model support #61

Add MACE model support #61

sef43 commented Oct 5, 2023

raimis Oct 9, 2023

raimis Oct 9, 2023

raimis Oct 9, 2023

raimis Oct 9, 2023

raimis Oct 9, 2023

raimis Oct 9, 2023

raimis Oct 9, 2023

raimis Oct 9, 2023

raimis Oct 9, 2023

JMorado commented Feb 27, 2024 •

edited

Loading

peastman commented Feb 27, 2024

JMorado commented Feb 27, 2024

JMorado commented Feb 28, 2024

peastman left a comment

peastman Feb 28, 2024

peastman Feb 28, 2024

JMorado Feb 29, 2024

peastman Feb 29, 2024

JMorado Feb 29, 2024

peastman Feb 28, 2024

JMorado commented Feb 29, 2024

JMorado commented Feb 29, 2024

peastman commented Feb 29, 2024

JMorado commented Mar 4, 2024

peastman commented Mar 4, 2024

		super(MACEForce, self).__init__()


		if not torch.cuda.is_available():

	# Convert it to TorchScript and save it.

	module = torch.jit.script(aniForce)
	module.save(filename)

	# Create the TorchForce and add it to the System.

	force = openmmtorch.TorchForce(filename)

Add MACE model support #61

Add MACE model support #61

Conversation

sef43 commented Oct 5, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JMorado commented Feb 27, 2024 • edited Loading

peastman commented Feb 27, 2024

JMorado commented Feb 27, 2024

JMorado commented Feb 28, 2024

peastman left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JMorado commented Feb 29, 2024

JMorado commented Feb 29, 2024

peastman commented Feb 29, 2024

JMorado commented Mar 4, 2024

peastman commented Mar 4, 2024

JMorado commented Feb 27, 2024 •

edited

Loading