Add OpenMM-torch interface #664

frostedoyster · 2024-06-21T02:48:56Z

Adds a MD interface to OpenMM via OpenMM-torch

Contributor (creator of pull-request) checklist

Tests updated (for new features and bugfixes)?
Documentation updated (for new features)?
Issue referenced (for PRs that solve an issue)?

Reviewer checklist

CHANGELOG updated with public API or any other important changes?

📚 Download documentation preview for this pull-request

⚙️ Download Python wheels for this pull-request (you can install these with pip)

Luthaf

Why is this using openmm-ml instead of openmm-torch directly? My understanding of the former was that it is supposed to provide specific models, not a whole new set of models?

This will also need tests, documentation and examples. And I guess this required #402 to be done first as well, since OpenMM relies heavily on conda.

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

Luthaf

I could not find how this code handle dtype/device? The model might be on a different device than the positions/cell.

For the dtype, the model might require a specific one and the system should be converted to this dtype

Luthaf · 2024-06-28T12:00:30Z

docs/src/atomistic/reference/openmm.rst

+.. autoclass:: metatensor.torch.atomistic.openmm_interface.get_metatensor_force
+ :show-inheritance:
+ :members:


Suggested change

.. autoclass:: metatensor.torch.atomistic.openmm_interface.get_metatensor_force

:show-inheritance:

:members:

.. autofunction:: metatensor.torch.atomistic.openmm_interface.get_metatensor_force

Luthaf · 2024-06-28T12:01:28Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+ forceGroup: int = 0,
+ atoms: Optional[Iterable[int]] = None,
+ check_consistency: bool = False,
+) -> openmm.System:


this needs a docstring & doc for all parameters

Luthaf · 2024-06-28T12:02:11Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+ path: str,
+ extensions_directory: Optional[str] = None,
+ forceGroup: int = 0,
+ atoms: Optional[Iterable[int]] = None,


Suggested change

atoms: Optional[Iterable[int]] = None,

selected_atoms: Optional[Iterable[int]] = None,

Luthaf · 2024-06-28T12:03:08Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+ def __init__(
+ self,
+ model: torch.jit._script.RecursiveScriptModule,
+ atomic_numbers: List[int],


Suggested change

atomic_numbers: List[int],

atomic_types: List[int],

Luthaf · 2024-06-28T12:06:58Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+ cell = system.cell
+
+ # Get the neighbor pairs, shifts and edge indices.
+ neighbors, interatomic_vectors, _, _ = getNeighborPairs(


Suggested change

neighbors, interatomic_vectors, _, _ = getNeighborPairs(

neighbors, interatomic_vectors, _, _ = NNPOps.neighbors.getNeighborPairs(

for clarity

Luthaf · 2024-06-28T12:07:34Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+ neighbors, interatomic_vectors, _, _ = getNeighborPairs(
+ system.positions,
+ self.requested_neighbor_list.engine_cutoff("nm"),
+ -1,


could you give the name of this argument? (and maybe others as well)

Luthaf · 2024-06-28T12:08:16Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+ -1,
+ cell,
+ )
+ mask = neighbors[0] >= 0


will there negative values in neighbors[0] if n_pairs=-1 above?

Yes, I think that NL returns a huge Tensor with all pairs, where only a certain number are not -1, i.e., masked

I understand this behavior if we are setting max_n_pairs=300: if there are fewer pairs, the other ones are filled with -1. You are saying this is also the case when giving max_n_pairs=-1?

Yes, I tried a few tests (with max_n_pairs=-1) and the mask is still necessary. I think they're always returning a number of pairs that is a multiple of 4. E.g. if you have 27 pairs, 28 will be returned and one of them is masked with center=neighbor=-1

Luthaf · 2024-06-28T12:10:28Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+ energy = (
+ self.model(
+ [system],
+ self.evaluation_options,
+ check_consistency=self.check_consistency,
+ )["energy"]
+ .block()
+ .values.reshape(())
+ )


Suggested change

energy = (

self.model(

[system],

self.evaluation_options,

check_consistency=self.check_consistency,

)["energy"]

.block()

.values.reshape(())

)

outputs = self.model(

[system],

self.evaluation_options,

check_consistency=self.check_consistency,

)

energy = outputs["energy"].block().values.reshape(())

Luthaf · 2024-06-28T12:13:16Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+ )
+
+ model = load_atomistic_model(path, extensions_directory=extensions_directory)
+


Could you print model.metadata() to the log somewhere? This prints information about the model.

Also, is there a mechanism to register papers to cite in OpenMM? If so, we should register the model references.

I think they have some MD loggers, but I don't think they're accessible from our function (they're attached to a Simulation object that is independent). Literally zero clue regarding the citation mechanism

Perhaps the best solution is to print() and whatever

printing everything is a good starting point, and enough for this PR!

Luthaf · 2024-06-28T12:13:49Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_interface.py

+
+ def __init__(
+ self,
+ model: torch.jit._script.RecursiveScriptModule,


Suggested change

model: torch.jit._script.RecursiveScriptModule,

model,

We don't really need the types in __init__, and this one is private

frostedoyster · 2024-06-28T14:06:22Z

Typical code to start a simulation would look like this:

platform = openmm.Platform.getPlatformByName('CUDA')
properties = {'Precision': 'double'}
simulation = openmm.app.Simulation(topology, system, integrator, platform, properties)

The device is handled by OpenMM torch (it moves the model to the desired device).
The dtype is more tricky, as it is not handled by OpenMM (I think they're trying to allow model developers to do things in mixed precision, so they won't call to(dtype)). At the moment, if the precision is wrong, the model will crash at runtime, and if the user matches the dtypes correctly, the model will run.

Let me know what you think the best solution is (leave as is, check at each step, etc)

Luthaf · 2024-07-02T12:16:46Z

Let me know what you think the best solution is (leave as is, check at each step, etc)

The way we are doing this in other engines is to always convert from the engine dtype to the model dtype for inputs, and the other way around for outputs. We could add a warning about mismatched dtype if we do the conversion, but I would rather allow the code to run instead of crashing.

frostedoyster · 2024-07-05T06:34:31Z

Printing metadata doesn't work because ::print() is not registered it seems, I'll see if I can fix it

This reverts commit 8c403c1.

Luthaf

Looks good! We should find a way to run the tests for this. I guess the main issue is that openmm requires conda for installation, right?

Either we keep the code in this repo, and we find a way to use conda from tox; or we move the code to another repo, with its own CI setup.

Maybe adding a section in tox.ini for "tests that require conda", and in there creating a conda env by calling conda explicitly could work.

Luthaf · 2024-07-08T09:32:52Z

docs/src/atomistic/reference/openmm.rst

+``conda install -c conda-forge openmm-torch nnpops``. Subsequently,
+metatensor can be installed with ``pip install metatensor[torch]``, and a minimal


Since there is #665, should we recommend building metatensor-torch from sources for now? Something like pip install --no-binary=metatensor-torch metatensor[torch]

Sounds good

Luthaf · 2024-07-08T09:33:53Z

python/metatensor-operations/CHANGELOG.md

@@ -17,10 +17,6 @@ follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 ### Removed
 -->

-### Changed


this removal looks like a merge issue

Luthaf · 2024-07-08T09:34:38Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_force.py

+ :param topology: The OpenMM Topology object.
+ :param path: The path to the exported metatensor model.
+ :param extensions_directory: The path to the extensions for the model.
+ :param forceGroup: The force group to which the force should be added.


should we link to OpenMM docs for force group here? It might not be clear to everyone.

Luthaf · 2024-07-08T09:35:25Z

docs/src/atomistic/reference/openmm.rst

+.. autofunction:: metatensor.torch.atomistic.openmm_force.get_metatensor_force
+
+In order to run simulations with ``metatensor.torch.atomistic`` and ``OpenMM``,
+we recommend installing ``OpenMM`` from conda, using


you should mention than NNPOps is a required dependency as well in here

Luthaf · 2024-07-08T09:38:00Z

python/metatensor-torch/metatensor/torch/atomistic/openmm_force.py

+ module = torch.jit.script(metatensor_force)
+
+ # create the OpenMM force
+ force = openmmtorch.TorchForce(module)


so model.to(device) is called by TorchForce? Could you leave a comment if this is the case, so future readers don't wonder what's happening?

Luthaf · 2024-07-08T09:39:56Z

python/metatensor-torch/tests/atomistic/openmm_force.py

+ )
+
+
+def model_different_units():


this might not be so interesting to run here, I did it in ASE to ensure our code did unit conversions right, but we don't have to do it for all interfaces.

Luthaf · 2024-07-08T09:41:20Z

python/metatensor-torch/tests/atomistic/openmm_force.py

+ force = get_metatensor_force(system, topology, model_path, check_consistency=True)
+ system.addForce(force)
+ integrator = openmm.VerletIntegrator(0.001 * openmm.unit.picoseconds)
+ platform = openmm.Platform.getPlatformByName("CPU")


should we try to run this test on CUDA if it is available?

Luthaf · 2024-07-08T09:41:45Z

python/metatensor-torch/tests/atomistic/openmm_force.py

+ state.getForces(asNumpy=True).value_in_unit(
+ openmm.unit.ev / (openmm.unit.angstrom * openmm.unit.mole)
+ )
+ / 6.0221367e23


maybe put this in a constant on top?

also, why are you doing * openmm.unit.mole above, just to divide again here?

Luthaf · 2024-07-08T09:45:07Z

python/metatensor-torch/tests/atomistic/openmm_force.py

+try:
+ import NNPOps # noqa: F401
+ import openmm
+ import openmmtorch # noqa: F401
+
+ HAS_OPENMM = True
+except ImportError:
+ HAS_OPENMM = False


Since we want to skip the whole file if one of the dependency is missing, we should use importskip:

Suggested change

try:

import NNPOps # noqa: F401

import openmm

import openmmtorch # noqa: F401

HAS_OPENMM = True

except ImportError:

HAS_OPENMM = False

openmm = pytest.importskip("openmm")

pytest.importskip("NNPOps")

pytest.importskip("openmmtorch")

frostedoyster requested a review from Luthaf June 21, 2024 02:49

frostedoyster force-pushed the openmm branch 2 times, most recently from 2818ef7 to 8fc6b55 Compare June 21, 2024 08:06

Luthaf reviewed Jun 21, 2024

View reviewed changes

frostedoyster requested a review from Luthaf June 21, 2024 15:03

frostedoyster changed the title ~~Add OpenMM-ML interface~~ Add OpenMM-torch interface Jun 24, 2024

frostedoyster marked this pull request as ready for review June 27, 2024 18:37

Luthaf reviewed Jun 28, 2024

View reviewed changes

frostedoyster requested a review from Luthaf June 28, 2024 14:08

frostedoyster and others added 13 commits July 5, 2024 08:41

Register ModelMetadataHolder::print()

8c403c1

Add OpenMM interface

1e3a0a6

Linter

1792629

Change interface to openmm-torch

a673e36

Let the user add the force to the system

cea817b

Add basic documentation

f2e3477

Add neighbor list

21b2629

Fix bugs

7e2f643

Extract neighbor list

5e94523

Add tests

e74d0fd

Implement suggestions

c660e07

Update fake classes for type hints

94e0973

Print metadata

3ea577a

frostedoyster force-pushed the openmm branch from e0ada11 to 3ea577a Compare July 5, 2024 06:43

frostedoyster changed the base branch from master to metadata-print July 5, 2024 06:44

frostedoyster added 3 commits July 5, 2024 11:38

Add warning for dtype conversion

aed0dcf

Revert "Register ModelMetadataHolder::print()"

4624483

This reverts commit 8c403c1.

Use print(metadata) to print the metadata

fc19b41

frostedoyster changed the base branch from metadata-print to master July 5, 2024 18:46

Luthaf reviewed Jul 8, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add OpenMM-torch interface #664

Add OpenMM-torch interface #664

frostedoyster commented Jun 21, 2024 •

edited by github-actions bot

Loading

Luthaf left a comment •

edited

Loading

Luthaf left a comment

Luthaf Jun 28, 2024

Luthaf Jun 28, 2024

Luthaf Jun 28, 2024

Luthaf Jun 28, 2024

Luthaf Jun 28, 2024

Luthaf Jun 28, 2024

Luthaf Jun 28, 2024

frostedoyster Jun 28, 2024

Luthaf Jul 2, 2024

frostedoyster Jul 5, 2024

Luthaf Jun 28, 2024

Luthaf Jun 28, 2024

frostedoyster Jun 28, 2024

frostedoyster Jun 28, 2024

Luthaf Jul 2, 2024

Luthaf Jun 28, 2024

frostedoyster commented Jun 28, 2024

Luthaf commented Jul 2, 2024

frostedoyster commented Jul 5, 2024

Luthaf left a comment •

edited

Loading

Luthaf Jul 8, 2024

frostedoyster Jul 8, 2024

Luthaf Jul 8, 2024

Luthaf Jul 8, 2024

Luthaf Jul 8, 2024

Luthaf Jul 8, 2024

Luthaf Jul 8, 2024

Luthaf Jul 8, 2024

Luthaf Jul 8, 2024

Luthaf Jul 8, 2024

Luthaf Jul 8, 2024

	atoms: Optional[Iterable[int]] = None,
	selected_atoms: Optional[Iterable[int]] = None,

	neighbors, interatomic_vectors, _, _ = getNeighborPairs(
	neighbors, interatomic_vectors, _, _ = NNPOps.neighbors.getNeighborPairs(

		)

		model = load_atomistic_model(path, extensions_directory=extensions_directory)

		``conda install -c conda-forge openmm-torch nnpops``. Subsequently,
		metatensor can be installed with ``pip install metatensor[torch]``, and a minimal

Add OpenMM-torch interface #664

Are you sure you want to change the base?

Add OpenMM-torch interface #664

Conversation

frostedoyster commented Jun 21, 2024 • edited by github-actions bot Loading

Contributor (creator of pull-request) checklist

Reviewer checklist

Luthaf left a comment • edited Loading

Choose a reason for hiding this comment

Luthaf 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

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

Choose a reason for hiding this comment

frostedoyster commented Jun 28, 2024

Luthaf commented Jul 2, 2024

frostedoyster commented Jul 5, 2024

Luthaf left a comment • edited Loading

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

frostedoyster commented Jun 21, 2024 •

edited by github-actions bot

Loading

Luthaf left a comment •

edited

Loading

Luthaf left a comment •

edited

Loading