Add OpenMM-torch interface

[frostedoyster]

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.

[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]

Suggested change

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

[Luthaf]

this needs a docstring & doc for all parameters

[Luthaf]

Suggested change

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

[Luthaf]

Suggested change

	atomic_numbers: List[int],
	atomic_types: List[int],

[Luthaf]

Suggested change

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

for clarity

[Luthaf]

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

[Luthaf]

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

[frostedoyster]

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

[Luthaf]

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]

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.

[frostedoyster]

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

[frostedoyster]

Perhaps the best solution is to print() and whatever

[Luthaf]

Suggested change

	model: torch.jit._script.RecursiveScriptModule,
	model,

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

[frostedoyster]

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)