Proposed API for embedding potentials

[JMorado]

Hi @peastman,

Following our discussion yesterday, here are some ideas on an API for supporting embedding potentials.

While OpenMM-ML currently assumes that the ML potential replaces only the intramolecular interactions of the ML region, embedding potentials differ in that they replace (either fully or partially) the intermolecular interactions between the ML and MM regions. This means that the Coulomb and/or Lennard-Jones interactions between these regions need to be zeroed out to avoid double-counting. Because this will require specific operations, I propose we create an EmbeddingPotential class that handles this logic, something along the lines of:

class EmbeddingPotential:
    def __init__(self, name: str, **args):
        self._impl = EmbeddingPotential._implFactories[name].createImpl(name, **args)

    def _removeCoulomb(self, …):
        pass

    def _removeLJ(self, ...):
        pass

This class would be used just like MLPotential to create a specific embedding potential:

emle = EmbeddingPotential('emle')

Embedding potentials should be combined with an "intramolecular" MLP, and thus the EmbeddingPotential class does not need any of the logic for creating OpenMM systems, which is already implemented in MLPotential. Furthermore, because embedding potentials only make sense in mixed systems, the main modifications required are in the createMixedSystem method. Therefore, I propose we add an additional parameter to this method, called embeddingPotential, which is by default set to None (mechanical embedding ML/MM) but can also accept an embedding potential. This would preserve the philosophy that these embedding schemes are to be applied on top of ML potentials, e.g.:

potential = MLPotential('ani2x')
emle = EmbeddingPotential('emle')
ml_system = potential.createMixedSystem(topology, mm_system, ml_atoms, embeddingPotential=emle)

Now, the question is: which ML–MM interactions should be zeroed out inside createMixedSystem? Only the electrostatics? Only the LJ? Or both? To minimise the risk of users shooting themselves in the foot, I think the best option is for the embedding potential implementation to determine, during addForces and based on the user choices, which interactions should be removed and to set instance-level flags indicating these, e.g. REMOVE_LJ and REMOVE_COULOMB. Depending on the values of these flags, createMixedSystem would then remove (or not) the corresponding interactions from the system, just as it already does for the bonded terms.

A more tricky situation is when interpolation=True is passed to createMixedSystem. In this case, the approach I propose is to create a CustomNonbondedForce for the intermolecular interactions that need to be removed, restricted only to the MM-ML interaction groups. The question here is whether this force should vary with the same lambda_interpolate used for the underlying MLP, or whether it should be controlled independently. In practice, I have found it useful to couple an indepenent lambda parameter. This makes it possible to individually account for the corrections introduced when going from MM to ML/MM(mechanical embedding) and from ML/MM(mechanical embedding) to ML/MM (electrostatic embedding). Of course, the main disadvantage of using a CustomNonbondedForce is that the PME treatment of the ML-MM nonbonded interactions is lost. I'm not sure whether there is an obvious way around this, or if it's simply a limitation we have to live with. We can also assume that, in general, embedding potentials will have their own cutoff schemes, long-range treatments, switching functions, etc., which are built in and therefore not easily controlled.

Regarding the implementation of the embedding potentials themselves, here there are no majorchanges required to the current functionality of OpenMM-ML, so a standard structure like this would work:

class EMLEPotentialImplFactory(EmbeddingPotentialImplFactory):
    def createImpl(self, name: str, **args) -> EmbeddingPotentialImpl:
        return EMLEPotentialImpl(name)

class EMLEPotentialImpl(EmbeddingPotentialImpl):

where EMLEPotentialImpl would have and addForces method and the OpenMM-Torch force implementation (similar to what I did here).

I'm looking forward to hearing any thoughts on how to improve these ideas. I'm also cc'ing @jmichel80 and @lohedges as they may have additional suggestions.

Many thanks.

[peastman]

You're assuming that an embedding is a separate object, independent of either the MM or ML potential, and can be freely mixed and matched. That isn't always the case. For example, AIMNet2 has built in support for electrostatic embedding as part of the ML model. You can provide a list of external charges. It polarizes the ML region based on them and computes the interaction energy.

Another approach that has been tried is to have the ML model predict charges and possibly dipoles for the atoms in the ML region. The interaction energy then gets computed with the classical force field, but using the multipoles from the ML model.

Let's think about an API that can uniformly handle all of these possibilities. For example, what about having an embedding argument to createMixedSystem() that takes a string? The default value would be 'mechanical', which corresponds to the current behavior. A particular ML potential could support other values. For example, AIMNet2 might accept the value 'electrostatic', which would tell it to use its built in electrostatic embedding. And then we could have a plugin mechanism for embeddings that aren't tied to a particular ML potential. If you specified embedding='emle', it would invoke the embedding plugin registered with that name.

[JMorado]

You're assuming that an embedding is a separate object, independent of either the MM or ML potential, and can be freely mixed and matched. That isn't always the case. For example, AIMNet2 has built in support for electrostatic embedding as part of the ML model. You can provide a list of external charges. It polarizes the ML region based on them and computes the interaction energy.

Right. This use case seems to be the one requiring less changes to the current code but it's definitely important to support.

Another approach that has been tried is to have the ML model predict charges and possibly dipoles for the atoms in the ML region. The interaction energy then gets computed with the classical force field, but using the multipoles from the ML model.

EMLE also predicts multipoles using an ML model, but directly outputs the interaction energy. Are you referring here to dynamic charge models? If the intention is to reuse the classical FF forces implemented by OpenMM wouldn't this require updating the context parameters at every iteration? Or would these interactions be handled differently, e.g. by a TorchForce that predicts the charges using an ML model and computes the interactions?

Let's think about an API that can uniformly handle all of these possibilities. For example, what about having an embedding argument to createMixedSystem() that takes a string? The default value would be 'mechanical', which corresponds to the current behavior. A particular ML potential could support other values. For example, AIMNet2 might accept the value 'electrostatic', which would tell it to use its built in electrostatic embedding. And then we could have a plugin mechanism for embeddings that aren't tied to a particular ML potential. If you specified embedding='emle', it would invoke the embedding plugin registered with that name.

Seems like a viable approach to me. The only part I don't understand is how the user would pass extra parameters to the plugin-based embeddings. Some of these plugins allow user defined settings.

[peastman]

The only part I don't understand is how the user would pass extra parameters to the plugin-based embeddings.

A few possibilities.

1. Pass them as arguments to createMixedSystem(). It already allows arbitrary arguments that get passed to the MLPotential. It could pass them to the embedding as well.
2. Allow the value of embedding to be an EmbeddingPotential that you've pre-configured.

[JMorado]

I was trying to avoid option 1 because I tend not to like the embedding_kwargs approach, as it isn’t very explicit. In my opinion option 2 would be the best way to go.

[peastman]

I wasn't suggesting an embedding_kwargs argument. You would just pass options to createMixedSystem(), the same way you can already pass arguments related to the potential.

[JMorado]

Thanks for the input. It sounds like we have a general plan for this.

Another approach that has been tried is to have the ML model predict charges and possibly dipoles for the atoms in the ML region. The interaction energy then gets computed with the classical force field, but using the multipoles from the ML model.

Is there anything special about these models that would change what we have proposed above? If not, shall I go ahead and make a draft PR to start working on this?

[peastman]

Let's wait a little. I want to create the infrastructure for it myself. But before doing that, I want to get out another release that adds some more models. Support for other embeddings will be part of the release after that, which will go with OpenMM 8.5.

[peastman]

I'm getting closer to being ready to start on this, so let's pin down the design. How does the following sound?

As far as the user is concerned, there are just a few minor changes:

• createMixedSystem() will get a new embedding argument with the default value 'mechanical'. For some embeddings, there might be additional new arguments you can pass.
• MLPotential will get a new getSupportedEmbeddings() method returning a list of values you can pass for that argument.

All potentials will support mechanical embedding. Some may define additional ones. For example, AIMNet2 has built in support for electrostatic embedding, so it would allow embedding='electrostatic'.

Embeddings could also be defined by plugins. They would extend a particular superclass and get registered as entry points, the same way potentials work. The interface would be very simple: after createMixedSystem() created the System it would pass it to the embedding, which could modify it in any way, or even create a new System to replace it. Given how open ended the field of embeddings is, I don't want to restrict what they can do. This interface will hopefully be able to support just about anything someone invents.

[JMorado]

Thanks! That sounds good to me.

Just to clarify the use case where interpolate=True is passed to createMixedSystem(): as I understand it, no major changes will be made to the current implementation. This means that the ML potential (and any plugin extending it) will be coupled to the lambda parameter, enabling turning it (them) on and off.

For MLPs with built-in support for embedding schemes, do these affect the ML-MM interactions or they are only environment corrections to the intramolecular energy? If the former is the case, then there must be a way to determine which ML-MM interactions are being affected, as these need to also be coupled to the lamba parameter. If they only afffect the intramolecular energy, then the current implementation should work out of the box.

For plugins such as EMLE and other frameworks, it is likely necessary to couple the lambda parameter to the electrostatic and/or Lennard-Jones terms. As I understand it, we leave this responsibility to the plugin itself, since it can modify the system as needed. Is that correct? If yes, this seems like a quite flexible approach.

[peastman]

That's the idea, right.

Modifying the system to add the embedding will require a lot of care. We'll need to look at what nonbonded method the MM system uses. And then we need to understand exactly what interactions are and aren't included in the embedding model. For example, if it's a periodic system, does it include the interactions of the ML region with other periodic copies of itself? What about interactions of the ML region with other periodic copies of the MM region? We need to make sure we're computing exactly the right interactions for the nonbonded method, none being missed, none being double counted, and no extra ones being included.

[epretti]

I've been looking at this to understand what changes will need to be made under the hood to provide support for embedding schemes other than mechanical. This has raised a few questions.

First, it seems to me like createMixedSystem() might be double-counting some interactions even with the existing mechanical embedding implementation. Specifically, the ML-ML nonbonded interactions are removed by adding exceptions:

openmm-ml/openmmml/mlpotential.py

Lines 308 to 310 in cf4efac

	for i in range(len(atomList)):
	for j in range(i):
	force.addException(atomList[i], atomList[j], 0, 1, 0, True)

This seems like it will double-count interactions between atoms in the ML region and any periodic copies, since these will be computed by the ML potential when PBCs are in use and will not be removed from the MM part here. Am I missing something?

Second, I'm not yet sure about the best way for embeddings to handle interpolation. Different cases will need different approaches, and there is also discussion above about allowing different interpolation parameters for ML-ML and ML-MM. Each embedding could be responsible for interpolation on its own, but this could lead to a lot of code duplication especially if we eventually have a few potentials that offer their own electrostatic embedding as well as standalone mechanical and EMLE embedding classes, for instance. Alternatively, MLPotential could try to handle this in a unified way, but this could be complicated and potentially limiting as to how embedding plugins could work.

[JMorado]

I think your concern about double-counting interactions is valid but subtle. The issue centers on what happens to ML-ML interactions beyond the MLP's cutoff (<=6 Å for most short-range models).

In typical biomolecular simulation boxes, ML atoms rarely see their periodic copies because they're beyond this cutoff range. The current implementation doesn't use setExceptionsUsePeriodicBoundaryConditions(True), which means the MM FF (e.g., using PME) still computes long-range electrostatic interactions between ML atoms and their distant periodic copies. This seems reasonable since it provides the missing long-range ML-ML electrostatics which the short-range MLP doesn't capture.

The "correct" approach depends on whether you believe the MLP should be the sole source of truth for ML-ML interactions (even if incomplete beyond its cutoff), or whether using the MM FF as a long-range approximation is acceptable.

You're right, however, that double-counting can occur in specific cases: for example, with MLPs that include long-range corrections, or in the infinite chain scenario described in the documentation. Given this, I think the option to enable PBC for exceptions should at least be exposed to users and properly documented, allowing them to choose the behavior that best matches their use case.

I’ll think a bit more about the interpolation points you raised and come back with a more concrete answer next week.

[epretti]

the MM FF (e.g., using PME) still computes long-range electrostatic interactions between ML atoms and their distant periodic copies. This seems reasonable since it provides the missing long-range ML-ML electrostatics which the short-range MLP doesn't capture.

The "correct" approach depends on whether you believe the MLP should be the sole source of truth for ML-ML interactions (even if incomplete beyond its cutoff), or whether using the MM FF as a long-range approximation is acceptable.

I would argue that if we pass box vectors to an MLIP for a periodic system, we shouldn't assume that it's going to provide inaccurate long-range behavior; this is an issue for the design of the MLIP itself. In addition, the current approach won't provide an MM approximation to long-range interactions between distant parts of the same periodic image, and it would seem to be inconsistent with the behavior you'd get just using createSystem() to simulate an entire system with the MLIP.

double-counting can occur in specific cases: for example, with MLPs that include long-range corrections

For instance, I want to add support for MACE-POLAR which does its own long-range calculations internally.

Given this, I think the option to enable PBC for exceptions should at least be exposed to users and properly documented, allowing them to choose the behavior that best matches their use case.

I think this wouldn't do what we'd want in general. This would subtract regular $1/r$ Coulomb terms with PBCs, rather than properly cancel out the PME ML-ML interaction energy, when enabled.

[peastman]

I think the difference is in what you expect the ML model to calculate. Is it the interaction of a single molecule with itself, or the interaction of that molecule with all possible periodic copies of itself? Both options are possible. Some models do one and some do the other.

[epretti]

Is it documented anywhere which models behave in which manner? Right now the mechanical embedding seems to assume the ML model won't include interactions with periodic copies, and so these are approximated with the MM force field. Do we still want that behavior then, or should it vary per model, or change altogether?

When I wrote up the proposal for electrostatic embedding, I assumed and specified that ML models would include the interactions with periodic copies for the terms they compute, and nobody we shared it with suggested that this would be a problem; I'd assume this makes sense for models that are doing their own electrostatics. If we can't assume this about some other models, then this would seem to be a problem if a user wants to use, e.g., EMLE with a model that doesn't include interactions with periodic copies.