createMixedSystem() for switching a whole waterbox

[JohannesKarwou]

I would like to use the OpenMM-ML package to switch a complete waterbox from MM to ML description using the createMixedSystem() functionality. I know that technically it is not made for that, but I wonder if that might - in theory - be possible.
For the simulation, I’m using the CutoffPeriodic option and set up the ML system the following way (I can provide a more detailed minimal example if needed):

potential = MLPotential("ani2x")
system = potential.createMixedSystem(
    psf.topology,
    mm_system,
    ml_atoms,
    interpolate=True,
    removeConstraints=True,
    implementation="torchani",
)
…
simulation.context.setParameter("lambda_interpolate", lamb)

I believe that when using a cutoff, the CustomBondForce within the CustomCVForces may cause issues in reproducing the Nonbonded Interaction for particles described only as ML atoms (https://github.com/openmm/openmm-ml/blob/c3d8c28eb92bf5c4b16efb81ad7a44b707fc5907/openmmml/mlpotential.py#LL298C14-L298C14). It seems that the PBC is not replicated for the ML portion of the system, which may not be problematic if only one small molecule is being treated by ML. However, when dealing with a system where all (or multiple) particles are described using ML (like the waterbox I would be interested in), this is not describing PBC correctly and leads to strange behavior of the box during the simulation (box blowing up). Therefore, I have created a Nonbonded Force using the openmm.CustomNonbondedForce according to the steps outlined in this discussion (openmm/openmm#3281). In the context of the CustomCVForce this might look like this:

for force in system.getForces():
  if isinstance(force, openmm.NonbondedForce):
    if topology.getNumAtoms() == len(atomList):
      eps_solvent = force.getReactionFieldDielectric()
      cutoff = force.getCutoffDistance()
      krf = (1/ (cutoff**3)) * (eps_solvent - 1) / (2*eps_solvent + 1)
      crf = (1/ cutoff) * (3* eps_solvent) / (2*eps_solvent + 1)
      ONE_4PI_EPS0 = 138.935456
      energy_expression = "4*epsilon*((sigma/r)^12 - (sigma/r)^6) + ONE_4PI_EPS0*chargeprod*(1/r + krf*r*r - crf);"
      energy_expression += "krf = {:f};".format(krf.value_in_unit(unit.nanometer**-3))
      energy_expression += "crf = {:f};".format(crf.value_in_unit(unit.nanometer**-1))
      energy_expression += "epsilon = sqrt(epsilon1*epsilon2);"
      energy_expression += "sigma = 0.5*(sigma1+sigma2);"
      energy_expression += "chargeprod = charge1*charge2;"
      energy_expression += "ONE_4PI_EPS0 = {:f};".format(ONE_4PI_EPS0)
      custom_nonbonded_force = openmm.CustomNonbondedForce(energy_expression)
      custom_nonbonded_force.addPerParticleParameter('charge')
      custom_nonbonded_force.addPerParticleParameter('sigma')
      custom_nonbonded_force.addPerParticleParameter('epsilon')
      custom_nonbonded_force.setNonbondedMethod(openmm.CustomNonbondedForce.CutoffPeriodic)
      custom_nonbonded_force.setCutoffDistance(cutoff)

      for index in range(force.getNumParticles()):
        charge, sigma, epsilon = force.getParticleParameters(index)
        custom_nonbonded_force.addParticle([charge, sigma, epsilon])
      for index in range(force.getNumExceptions()):
        j, k, chargeprod, sigma, epsilon = force.getExceptionParameters(index)
        custom_nonbonded_force.addExclusion(j, k)

      if custom_nonbonded_force.getNumParticles() > 0:
        name = f'mmForce{len(mmVarNames)+1}'
        cv.addCollectiveVariable(name, custom_nonbonded_force)
        mmVarNames.append(name)

Using this approach, I have been able to reproduce the expected energies that would result from simulating the box using only MM (ff.createSystem) or only ML (potential.createSystem).

force	MM (createSystem)	MM/ML (createMixedSystem)	MM/ML (createMixedSystem)	ML (createSystem)
lambda		0 (=MM)	1 (=ML)
HarmonicBondForce	4.4e-06	4.4e-06	4.4e-06	–
HarmonicAngleForce	3.8e-07	3.8e-07	3.8e-07	–
NonbondedForce	-1142.246	-1142.2456	-1142.246	–
TorchForce	–	-5993629	-5993629	-5993629
CustomCVForce		-1142.245	-5993629	–

In your opinion, do you believe that this is a valid approach for transitioning a box from MM to ML? Or is there anything I am missing?

[peastman]

It's an interesting question. createMixedSystem() was created with the assumption you'd be simulating a ligand with ML and everything else with MM. In that context, it makes sense to create a bonded force with no periodic boundary conditions to handle interactions within the ligand. And I think you're right that we should call setUsesPeriodicBoundaryConditions(True) on the CustomBondForce. A PR would be welcome!

For your case, though, this is all unnecessary complexity. You don't need to separate out internal from external interactions. You really just need to create a CustomCVForce, add all the ML and MM forces to it, and set the energy expression to interpolate between them. You can copy everything you need from the existing code.

openmm-ml/openmmml/mlpotential.py

Lines 275 to 279 in c3d8c28

	mlVarNames = []
	for i, force in enumerate(tempSystem.getForces()):
	name = f'mlForce{i+1}'
	cv.addCollectiveVariable(name, deepcopy(force))
	mlVarNames.append(name)

openmm-ml/openmmml/mlpotential.py

Lines 288 to 292 in c3d8c28

	mmVarNames = []
	for i, force in enumerate(bondedForces):
	name = f'mmForce{i+1}'
	cv.addCollectiveVariable(name, deepcopy(force))
	mmVarNames.append(name)

openmm-ml/openmmml/mlpotential.py

Lines 336 to 338 in c3d8c28

	mlSum = '+'.join(mlVarNames) if len(mlVarNames) > 0 else '0'
	mmSum = '+'.join(mmVarNames) if len(mmVarNames) > 0 else '0'
	cv.setEnergyFunction(f'lambda_interpolate*({mlSum}) + (1-lambda_interpolate)*({mmSum})')

[jchodera]

@peastman : Perhaps we can add functionality for easily simulating periodic ML systems into the next feature release?

[peastman]

createSystem() works fine for periodic systems. createMixedSystem() does too as long as the ML part is just a ligand. It also should work fine if interpolate=False.

[JohannesKarwou]

Thank you very much for these insights and the helpful comment!
I would have one additional (maybe stupid) question:
Is there a reason, why it is not recommended to have constraints and interpolate=True (#12 (comment))? Because in my context, it would be nice to have at least cons=HBonds for the MM and for the ML part using the option removeConstraints =False.

[peastman]

Interpolating between MM and ML potentials is complicated, and I'm not sure all the possible combinations have been fully explored or thought out yet. You need to make sure both of them are valid. In your case you're simulating a water box. Most classical water models are designed to be fully rigid. That's how they were parametrized. You can make them flexible, but they may not do a good job of reproducing experimental quantities if you do. ML models are generally designed to be completely flexible with no constraints. It's possible there's a way you could add some constraints and still get realistic behavior, but that hasn't been studied much yet. At the very least, the constrained distances would likely need to be different from the ones in the classical water model.

The iterpolate option was really created to support one particular use case: a protein-ligand binding simulation that's done mainly with a classical potential, and then at the end you run a much shorter calculation to correct the endpoints to ML accuracy for the ligand only. See https://www.biorxiv.org/content/10.1101/2020.07.29.227959v1.full.pdf. Other use cases would have other issues that may not be handled properly yet.