Allow for long-range dispersion correction to be (efficiently?) computed when parameter offsets are used?

[msuruzhon]

Hello,

Consider this basic code which creates a simple system of two Lennard-Jones particles with a periodic NonbondedForce in OpenMM 7.4.1/7.4.2:

from simtk import openmm, unit

def generateSystem():
    system = openmm.System()
    system.addParticle(1)
    system.addParticle(1)
    return system

def generateContext(system):
    context = openmm.Context(system, openmm.LangevinIntegrator(298 * unit.kelvin, 1 / unit.picosecond, 2 * unit.femtosecond))
    context.setPeriodicBoxVectors(openmm.Vec3(1, 0, 0), openmm.Vec3(0, 1, 0), openmm.Vec3(0, 0, 1))
    context.setPositions([openmm.Vec3(0, 0, 0), openmm.Vec3(0.4, 0, 0)])
    return context

# generate a simple system of two Lennard-Jones particles
system = generateSystem()
force = openmm.NonbondedForce()
system.addForce(force)
force.addParticle(0 * unit.elementary_charge, 0.2 * unit.nanometer, 10 * unit.kilocalories_per_mole)
force.addParticle(0 * unit.elementary_charge, 0.2 * unit.nanometer, 5 * unit.kilocalories_per_mole)
force.addGlobalParameter("switch", 0)
force.setNonbondedMethod(openmm.NonbondedForce.CutoffPeriodic)
force.setCutoffDistance(0.5 * unit.nanometer)
context = generateContext(system)

print("Unperturbed:", context.getState(getEnergy=True).getPotentialEnergy())

# perturb the LJ parameters
force.setParticleParameters(1, 0 * unit.elementary_charge, 0.2 * unit.nanometer, 1 * unit.kilocalories_per_mole)
force.addParticleParameterOffset("switch", 1, 0, 0, 4 * unit.kilocalories_per_mole)
context = generateContext(system)

print("Perturbed before switch:", context.getState(getEnergy=True).getPotentialEnergy())
context.setParameter("switch", 1)
print("Perturbed after switch:", context.getState(getEnergy=True).getPotentialEnergy())

When we perturb the epsilon parameter back to the initial value of the unperturbed system we don't recover the original energy:

Unperturbed: -2.3476030676357142 kJ/mol
Perturbed before switch: -1.1524368965005585 kJ/mol
Perturbed after switch: -2.1586144410943695 kJ/mol

However, had we instantiated the switch parameter to 1 before defining the parameter offset, we would have got the expected result:

Unperturbed: -2.3476030676357142 kJ/mol
Perturbed before switch: -2.3476030676357142 kJ/mol
Perturbed after switch: -2.3476030676357142 kJ/mol

This also seems to only happen when periodic boundary conditions are used, because when we swap CutoffPeriodic for CutoffNonPeriodic we get the expected behaviour:

Unperturbed: -1.8201923370361328 kJ/mol
Perturbed before switch: -0.8140147924423218 kJ/mol
Perturbed after switch: -1.8201923370361328 kJ/mol

It therefore seems that there is a bug specifically related to parameter offsets in periodic systems? Any advice would be appreciated.

EDIT: It seems that this problem is related to the dispersion correction, which I assume doesn't get updated with the parameter. Is there any way to trigger this update manually without calling the integrator?

[peastman]

I just finished tracking down the problem, then found you'd already come to the same conclusion. It's best to add a new comment rather than editing the previous one. Editing a comment doesn't cause emails to get sent out, so no one realizes you've done it.

Anyway, the difference does in fact come from the long range dispersion correction. The magnitude of that correction is just a constant divided by the box volume. So it adds negligible cost at each time step, but computing the constant at the start of the simulation can be expensive. That's why it gets computed using the base parameters for each particle, not taking offsets into account. If you disable the correction, you'll see the difference goes away:

force.setUseDispersionCorrection(False)

[msuruzhon]

@peastman Many thanks for the reply. Even if it is expensive, is there a way to force the constant recalculation from Python (ideally without having to reinitialise the Context)? I can't find anything in the documentation.

[peastman]

Reinitializing the context is the only way. And even then, it will do it based on the default value of the parameter, not the current value.

[jchodera]

We've run into this issue before as well. Would it be possible to write a GPU kernel to recompute the dispersion correction on the fly when the parameters change?

[peastman]

Of course it's possible, but it would be very expensive.

[msuruzhon]

I think in many practical situations only a few parameters are changed using offsets, so maybe the update could only recalculate the contributions involving the perturbed parameters? I am not sure how straightforward it would be to implement this though.

[jchodera]

I think in many practical situations only a few parameters are changed using offsets, so maybe the update could only recalculate the contributions involving the perturbed parameters? I am not sure how straightforward it would be to implement this though.

Good point! We could cluster the interactions of the non-offset particles into specific classes of (ncounts, epsilon, sigma), and then only integrate all (offset particle, non-offset particle class) combinations whenever the offset parameters are changed. This could be done in parallel for each combination and in a kernel that runs independently of others?