Support for chemical bonds crossing the ML-MM boundary in ML/MM simulations using link atoms

[epretti]

We want to support the link-atom method for ML/MM simulations in which the ML-MM region boundary cuts across one or more bonds. I already have a tentative implementation for mechanical embedding but want to open this issue to solicit feedback about API details and other related matters.

For mechanical embedding, support can be largely automatic, and bonds from the ML region that have been cut and require capping with additional atoms can be identified from the topology and ML region atom indices.

• Is there any reason to support capping atoms other than hydrogen?
• Our tests suggest the best approach for positioning the capping atoms is to place them at a fixed distance along the cut bonds. We can guess an appropriate distance using a built-in table of covalent radii. However, we'll want to allow users to override this; would something like linkAtomDistances=[(mlAtomIndex, mmAtomIndex, distance), ...] as a keyword argument for the embedding make sense?

A more general question relates to support for other generic embedding plugins and potential-specific embedding methods. To set up the link-atom method requires adding the appropriate virtual sites, as well as finding the appropriate bonded terms to remove. We don't want each plugin to have to implement this over and over. What kind of customization, if any, will embeddings need to be able to make to this process? (An embedding should also be able to reject ML regions cutting across bonds if it can't support them.)

[epretti]

An initial draft of an implementation can be found in #161. The changes to the mechanical embedding plugin should illustrate what needs to be done to an embedding method to support it.

One issue with this method in the current OpenMM-ML API is that createMixedSystem() only returns a System. This makes it hard to work with when link atoms are present, since extra virtual sites get added that won't be in the Topology. The current code generates a Topology with the virtual sites, but currently doesn't return it. Doing so would break compatibility for existing users of createMixedSystem(). Modifying the Topology in place is possible but probably a bad idea.

[mattiabondanza]

I started to work on this too starting from what you have already done https://github.com/mattiabondanza/openmm-ml .

I noticed a few problems that I am trying to address but I would like to know your opinion on that (sorry habit, I'm using QM to indicate ML part actually (: ):

1. As link atoms are something that is needed just by the QM part, I think it's not worth (and actually causes some problems) to have them in the OpenMM system. We can just "add" them on the fly to the QM system in the correct position before running the calculations and that should be enough.
2. QM forces on LA should be projected on MM and QM atoms defining the LA bond: this is needed because of the constraints on LA position and without doing that it's not clear to me where this force should go (if you just remove it, Newton's third low is broken).

I addressed those two problems, and now it feels a more natural implementation to me. Next I will check that all bonded and non-bonded terms are in the correct place, after that it should work.

Note that I implemented that only for MACE.

If you can give me a feedback it would be precious.

[epretti]

2. QM forces on LA should be projected on MM and QM atoms defining the LA bond: this is needed because of the constraints on LA position and without doing that it's not clear to me where this force should go (if you just remove it, Newton's third low is broken).

This is why I chose to use virtual sites for the link atoms; although they are a bit of trouble to have in the topology, their positions will be calculated automatically and any forces on them should be correctly propagated to the parent atoms by OpenMM.

The implementation in #161 should presently support interpolation and work with mechanical embedding for any ML model without requiring any model-specific modifications.

[mattiabondanza]

Tks for the quick clarification.
Now I see how your code works and why you chose to use VirtualSites for link atoms. This choice is clever as it simplifies the force propagation and position calculations, but it still feels a bit awkward to me; Those virtual sites are altering the topology without being really involved in the MM calculations and once you add them you have to take them in account during any subsequent manipulation of the System from the calling script (eg. coordinates updates, particle count etc), which should not be necessary.
Moreover the math needed to perform the placement and projection is pretty easy and can be implemented in utilities functions and called when performed operations that are common to all models (extracting indices from positions and put forces back in place). In my last commit I wrote to utilities function to isolate the extraction of coordinates and the projection/expansion of forces in the full system array. This should be super easy to be replaced in every method (actually in the current implementation there are blocks that are just the same or very similar in every model):

    positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
    numAtoms = positions.shape[0]
    if indices is not None:
        positions = positions[indices]

    if indices is not None:
        f = np.zeros((numAtoms, 3), dtype=forces.dtype)
        f[indices] = forces
        forces = f

Those blocks can be replaced by something like:

    _positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
    numAtoms = _positions.shape[0]
    positions = utilities.system_positions_to_ml(_positions, indices, linkBondsData)

and

    forces = utilities.ml_forces_to_system(_positions, forces, indices, linkBondsData)

In my opinion this has the advantage to be slightly more complicated on the developer perspective but really much more transparent on the user interface. Let me know what to you think about it; and if you like it feel free to use my code as you wish.