Implement support for decoupled electrostatics with exact PME and GROMACS-type direct-space PME

[JIMonroe]

I'm looking to perform some alchemical free energy calculations and recently started using OpenMM to do this. You've all done some nice work here, so I don't want to re-invent the wheel, but I initially set out to try and make these calculations work without openmmtools.alchemy, mainly to deepen my own understanding of how this all works.

Just to clearly define the problem, I want to take a particular molecule and decouple it from the rest of the system, leaving all direct-space intramolecular interactions on. From what I understand, this is what GROMACS achieves (someone please correct me here if necessary), and so I've been checking potential energies at different states against calculations from GROMACS.

This is relatively straight-forward in terms of sterics, as there is no reciprocal space term to deal with. However, electrostatics is a different story. Ideally, the coupled state includes all direct and reciprocal space interactions while the uncoupled state includes direct and reciprocal space interactions for the system minus the alchemical molecule and the direct-space interactions for the molecule with itself. I think I figured out a way to do this with OpenMM, and it mostly matches with GROMACS, but it doesn't seem to match any of the options for handling electrostatics provided by AbsoluteAlchemicalFactory. In my approach, I use the following scheme:

• NonbondedForce using PME with all electrostatics (except 1-4 terms and with all alchemical/alchemical terms set to zero with exceptions), and all non-alchemical/non-alchemical sterics (LJ)

• CustomNonbondedForce using CutoffPeriodic implementing a soft-core potential for alchemical/non-alchemical sterics

• NonbondedForce using PME with ONLY exceptions (zero for all particle charges and epsilons), including 1-4 terms and other alchemical/alchemical non-bonded interactions

Is this doing what I think it is? And for the options you have in AbsoluteAlchemicalFactory, what are the exact end states (coupled molecule with lambda of say 0, to uncoupled molecule with lambda of 1)? In terms of the potential energy function?

In my tests, I have set up the alchemical system with the following code:

alchFactory = alchemy.AbsoluteAlchemicalFactory(consistent_exceptions=False, 
                                                alchemical_pme_treatment='coulomb',
                                                disable_alchemical_dispersion_correction=False,
                                                split_alchemical_forces=True
                                               )

alchemicalRegion = alchemy.AlchemicalRegion(alchemical_atoms=soluteIndices, 
                                            annihilate_electrostatics=False,
                                            softcore_alpha=0.5,
                                            softcore_a=1, softcore_b=1, softcore_c=1,
                                            softcore_beta=0.0,
                                            softcore_d=1, softcore_e=1, softcore_f=1
                                           )

Obviously consistent_exceptions plays a role here too, but based on other discussions I've seen around this point I'm pretty sure I want to keep it as False.

I'm sorry if this is documented and I missed it. I also know there has been a lot of tangentially-related discussion on github, but I'm having difficulty synthesizing it into something cohesive.

[andrrizzi]

I may be missing something, of course, but you may need an extra CustomNonbondedForce to model the nonalchemical-alchemical electrostatics as well.

Direct-space and Coulomb electrostatics in alchemy is modeled with a CustomNonbondedForce currently. In general, we have two actually, one controlling the alchemical-alchemical electrostatics, and another controlling nonalchemical-alchemical electrostatics, but since you only want to decouple (i.e., not annihilate) the electrostatics, I think a single non-alchemical/alchemical custom force will suffice in your case. We use interaction groups to restrict those custom forces to the group of atoms of interest.

Exact treatment is implemented using a single NonbondedForce and using NonbondedForce.updateParametersInContext() to modify directly the charges, but we'll soon change this implementation when OpenMM 7.3 will come out to use the new parameter offset feature.

alchemy also models the 1-4 exceptions using CustomBondForces instead of an extra NonbondedForce, but in principle, your approach could work too.

Hope this helps!

[jchodera]

One thing you'll need to decide is how you want to represent intra-alchemical electrostatics, especially for the fully coupled state. If you want this to be identical to the PME energy, if you simply use a Coulomb term, you are going to miss the interaction with periodic images of the alchemical system, which can be problematic for charged or even highly polar ligands. If you want to keep this as PME interactions, you will run into trouble with charged ligands. Check out the new "offset" feature in the development build of openmm 7.3, which allows you to add a charge offset to particles and exceptions you can control with a context parameter.

Let's keep this issue open to remind ourselves to more clearly explain our alchemical modifications in the Sphinx theory docs.

[JIMonroe]

Thanks! The offset feature sounds promising, but I may wait for the official release.

If you want this to be identical to the PME energy, if you simply use a Coulomb term, you are going to miss the interaction with periodic images of the alchemical

From this forum post I was lead to believe that even if a particle-particle interaction is excluded, the reciprocal space interaction is still present as long as the charges are not zero. Is this right? If it works this way, you can use exceptions to turn off only direct-space intra-alchemical electrostatics without getting rid of the reciprocal space term. I then migrate these direct-space intra-alchemical terms to exceptions in a new NonbondedForce.

From my tests, the potential energy of the alchemically-modified system with the alchemical molecule fully coupled is different than before making the alchemical modifications if I use 'direct-space' or 'coulomb'. It is the same for 'exact', but then I have to annihilate the charges, which isn't the worst thing in the world, it just means I have to do one more calculation. I realize this is expected behavior, but in my mind it should be very clear that when you use 'direct-space' or 'coulomb' you are NOT simulating the same coupled state anymore, and are NOT using PME in the sense that one might imagine. You do say this in the code, but thank you for clarifying here.

[andrrizzi]

From this forum post I was lead to believe that even if a particle-particle interaction is excluded, the reciprocal space interaction is still present as long as the charges are not zero. Is this right?

I don't think this is the case. I believe OpenMM subtracts the self-term using the original charges even for pairs of atoms that are an exclusion so effectively they should not be included in the reciprocal space.

It may be possible to implement in Python a direct-space PME preserving the reciprocal space by keeping all the charges on in the NonbondedForce and adding a CustomNonbondedForce that subtracts the alchemical direct-space electrostatics energy. It will be probably less efficient since you end up computing the same terms twice, but it may be more useful than what we currently have implemented.

[JIMonroe]

I believe OpenMM subtracts the self-term using the original charges even for pairs of atoms that are an exclusion so effectively they should not be included in the reciprocal space.

Ah, thanks a ton! That got me on the right track and I found this discussion that really made things clear.

I don't want to spend too much more time on this since it sounds like the "offset" feature will make this discussion more or less obsolete, but I do want to sort of sum up what I've learned (please correct me if I'm still off somewhere). The approach I outlined is not exactly what I described I want because, although it does still include non-alchemical/alchemical electrostatics in the original non-bonded force, including the reciprocal part, it does not include the reciprocal part of the alchemical-alchemical interaction. So as @jchodera mentioned, there will still be issues with charged or highly-polar ligands. In 'direct-space' or 'coulomb' you're more honest about what you're doing and switch off all reciprocal space interactions involving the alchemical molecule, leaving only the direct-space interactions which are encoded in the custom forces you create.

This has really helped my understanding, so thank you! Hopefully it also helps write the documentation clearly (assuming the "offset" stuff doesn't completely supersede all of this).

[JIMonroe]

Not trying to re-open this, really, but as an update, I've figured out how to get what I'm looking for and also checked out the new parameter offset features (No use cases, but a description is here).

In the coupled state, I want full PME treatment of everything, and in the decoupled state, I want PME of everything except the alchemical molecule plus regular intramolecular Coulombic interactions in the alchemical molecule (with appropriate exclusions and 1-4's). To achieve this, I scale the charges of the alchemical molecule to zero in a NonbondedForce while I simply turn on a CustomNonbonded force that defines regular Coulomb interactions but applies an interaction group for just alchemical-alchemical. The expression for the CustomNonbondedForce is as follows so that the sum of this and the NonbondedForce produces the full intramolecular Coulombic interactions (lambdaQ below directly scales each charge of the alchemical molecule in the NonbondedForce).

soluteCoulFunction = '(1.0-(lambdaQ^2))*ONE_4PI_EPS0*charge/r;'
soluteCoulFunction += 'ONE_4PI_EPS0 = %.16e;' % (ONE_4PI_EPS0)
soluteCoulFunction += 'charge = charge1*charge2'

I can get it to work very efficiently as long as I keep a different context object in memory for every alchemical state (otherwise it's slow if you want dispersion corrections). I tested this against GROMACS and get identical potential energies in the coupled and de-coupled states for all the systems I've tested from the FreeSolv database (of course, I also add a CustomNonbondedForce for alchemical-chemical soft-core interactions and another CustomNonbondedForce to handle intramolecular LJ interactions).

The offset seems like a nicer way to switch off the charges in the NonbondedForce. I'm not sure this fully gets you around having to do something similar for keeping on the alchemical molecule intramolecular interactions - you might gain efficiency and be able to get rid of the CustomNonbondedForce in favor of some exceptions. It seems, however, that you would need to be careful and clever with how you then modify exceptions in order to remove alchemical-chemical interactions and leave on intramolecular interactions of the alchemical molecule. Have I interpreted this right?

[andrrizzi]

turn on a CustomNonbonded force that defines regular Coulomb interactions but applies an interaction group for just alchemical-alchemical

This is something that we can implement in AbsoluteAlchemicalFactory to enable decoupling of electrostatics with exact PME treatment! In the gas phase, we're not interested in the interactions between periodic copies of the alchemical molecule anyway. I hope I'll be able to wrap up the offset enhancement in the next few days, but if you want this feature in openmmtools feel free to open a PR with the changes, and I can help you out.

I can get it to work very efficiently as long as I keep a different context object in memory for every alchemical state (otherwise it's slow if you want dispersion corrections).

The Coulomb forces should have dispersion interactions set to False. For the sterics, our strategy is to disable them (with AbsoluteAlchemicalFactory(disable_alchemical_dispersion_correction=False)) and reweight the end points to account for them again. This allows us to use a single Context for everything, which means we can simulate as many states as we want without limits imposed by the GPU memory, and the loss of overlap is usually very small since the number of alchemical atoms is usually small.

[andrrizzi]

I've changed the title of this to remind us to implement decoupled electrostatics with exact PME and direct-space PME preserving the reciprocal space during the alchemical transformation.

[JIMonroe]

Thank you! Sorry for the silence for awhile, but I will fork and take a stab at modifying alchemy.py to include this option when I get the chance. For now it's been easier for me to just stick with my own code, but you guys have a lot of nice features I eventually want to take advantage of.

[jchodera]

@JIMonroe : Just wanted to check in to see if you had a chance to try to make these changes, or if we should give it a try!

[JIMonroe]

I unfortunately have not gotten the chance to make these changes, and I'm not sure I will in the near future. I ended up developing working code and am running with it for now. However, I'm happy to share what I've learned, as having this as a workable feature in a package like openmmtools would have made my life quite a bit easier. (So hopefully it will help someone else in the future!)

My code for making the system alchemical is shown below:

#The below sets up an alchemical system

#First get the solute indices
soluteIndices = []
for res in top.residues:
    if res.name == 'MOL':
        for atom in res.atoms:
            soluteIndices.append(atom.idx)

#We need to add a custom non-bonded force for the solute being alchemically changed
#Will be helpful to have handle on non-bonded force handling LJ and coulombic interactions
NBForce = None
for frc in system.getForces():
    if (isinstance(frc, mm.NonbondedForce)):
        NBForce = frc

#Separate out alchemical and regular particles using set objects
alchemicalParticles = set(soluteIndices)
chemicalParticles = set(range(system.getNumParticles())) - alchemicalParticles

#Define the soft-core function for turning on/off LJ interactions
#In energy expressions for CustomNonbondedForce, r is a special variable and refers to the distance between particles
#All other variables must be defined somewhere in the function.
#The exception are variables like sigma1 and sigma2.
#It is understood that a parameter will be added called 'sigma' and that the '1' and '2' are to specify the combining rule.
softCoreFunction = '4.0*lambdaLJ*epsilon*x*(x-1.0); x = (1.0/reff_sterics);'
softCoreFunction += 'reff_sterics = (0.5*(1.0-lambdaLJ) + ((r/sigma)^6));'
softCoreFunction += 'sigma=0.5*(sigma1+sigma2); epsilon = sqrt(epsilon1*epsilon2)'
#Define the system force for this function and its parameters
SoftCoreForce = mm.CustomNonbondedForce(softCoreFunction)
SoftCoreForce.addGlobalParameter('lambdaLJ', 1.0) #Throughout, should follow convention that lambdaLJ=1.0 is fully-interacting state
SoftCoreForce.addPerParticleParameter('sigma')
SoftCoreForce.addPerParticleParameter('epsilon')

#Will turn off electrostatics completely in the original non-bonded force
#In the end-state, only want electrostatics inside the alchemical molecule
#To do this, just turn ON a custom force as we turn OFF electrostatics in the original force
ONE_4PI_EPS0 = 138.935456 #in kJ/mol nm/e^2
soluteCoulFunction = '(1.0-(lambdaQ^2))*ONE_4PI_EPS0*charge/r;'
soluteCoulFunction += 'ONE_4PI_EPS0 = %.16e;' % (ONE_4PI_EPS0)
soluteCoulFunction += 'charge = charge1*charge2'
SoluteCoulForce = mm.CustomNonbondedForce(soluteCoulFunction)
#Note this lambdaQ will be different than for soft core (it's also named differently, which is CRITICAL)
#This lambdaQ corresponds to the lambda that scales the charges to zero
#To turn on this custom force at the same rate, need to multiply by (1.0-lambdaQ**2), which we do
SoluteCoulForce.addGlobalParameter('lambdaQ', 1.0) 
SoluteCoulForce.addPerParticleParameter('charge')

#Also create custom force for intramolecular alchemical LJ interactions
#Could include with electrostatics, but nice to break up
#We could also do this with a separate NonbondedForce object, but it would be a little more work, actually
soluteLJFunction = '4.0*epsilon*x*(x-1.0); x = (sigma/r)^6;'
soluteLJFunction += 'sigma=0.5*(sigma1+sigma2); epsilon=sqrt(epsilon1*epsilon2)'
SoluteLJForce = mm.CustomNonbondedForce(soluteLJFunction)
SoluteLJForce.addPerParticleParameter('sigma')
SoluteLJForce.addPerParticleParameter('epsilon')

#Loop over all particles and add to custom forces
#As we go, will also collect full charges on the solute particles
#AND we will set up the solute-solute interaction forces
alchemicalCharges = [[0]]*len(soluteIndices)
for ind in range(system.getNumParticles()):
    #Get current parameters in non-bonded force
    [charge, sigma, epsilon] = NBForce.getParticleParameters(ind)
    #Make sure that sigma is not set to zero! Fine for some ways of writing LJ energy, but NOT OK for soft-core!
    if sigma/u.nanometer == 0.0:
        newsigma = 0.3*u.nanometer #This 0.3 is what's used by GROMACS as a default value for sc-sigma
    else:
        newsigma = sigma
    #Add the particle to the soft-core force (do for ALL particles)
    SoftCoreForce.addParticle([newsigma, epsilon])
    #Also add the particle to the solute only forces
    SoluteCoulForce.addParticle([charge])
    SoluteLJForce.addParticle([sigma, epsilon])
    #If the particle is in the alchemical molecule, need to set it's LJ interactions to zero in original force
    if ind in soluteIndices:
        NBForce.setParticleParameters(ind, charge, sigma, epsilon*0.0)
        #And keep track of full charge so we can scale it right by lambda
        alchemicalCharges[soluteIndices.index(ind)] = charge

#Now we need to handle exceptions carefully
for ind in range(NBForce.getNumExceptions()):
    [p1, p2, excCharge, excSig, excEps] = NBForce.getExceptionParameters(ind)
    #For consistency, must add exclusions where we have exceptions for custom forces
    SoftCoreForce.addExclusion(p1, p2)
    SoluteCoulForce.addExclusion(p1, p2)
    SoluteLJForce.addExclusion(p1, p2)

#Only compute interactions between the alchemical and other particles for the soft-core force
SoftCoreForce.addInteractionGroup(alchemicalParticles, chemicalParticles)

#And only compute alchemical/alchemical interactions for other custom forces
SoluteCoulForce.addInteractionGroup(alchemicalParticles, alchemicalParticles)
SoluteLJForce.addInteractionGroup(alchemicalParticles, alchemicalParticles)

#Set other soft-core parameters as needed
SoftCoreForce.setCutoffDistance(10.0*u.angstroms)
SoftCoreForce.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
SoftCoreForce.setUseSwitchingFunction(True)
SoftCoreForce.setSwitchingDistance(9.0*u.angstroms)
SoftCoreForce.setUseLongRangeCorrection(True) 
system.addForce(SoftCoreForce)

#Set other parameters as needed - note that for the solute force would like to set no cutoff
#However, OpenMM won't allow a bunch of potentials with cutoffs then one without...
#So as long as the solute is smaller than the cut-off, won't have any problems!
SoluteCoulForce.setCutoffDistance(10.0*u.angstroms)
SoluteCoulForce.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
SoluteCoulForce.setUseSwitchingFunction(True)
SoluteCoulForce.setSwitchingDistance(9.0*u.angstroms)
SoluteCoulForce.setUseLongRangeCorrection(False) #DON'T want long-range correction here!
system.addForce(SoluteCoulForce)

SoluteLJForce.setCutoffDistance(10.0*u.angstroms)
SoluteLJForce.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
SoluteLJForce.setUseSwitchingFunction(True)
SoluteLJForce.setSwitchingDistance(9.0*u.angstroms)
SoluteLJForce.setUseLongRangeCorrection(False) 
system.addForce(SoluteLJForce)

So again, sorry to just sort of dump this on you guys. In the above, all you need to know is that I have an OpenMM system object called system and a related topology object I can pull solute indices from. The idea is to use your exact option for alchemical_pme_treatment, but then slowly turn on a custom non-bonded force that defines a direct-space coulombic interaction only between intramolecular alchemical atoms. In all honesty, I sort of chickened out a bit because I didn't see a really clean, quick way to do this in alchemy.py, but hopefully you do.

I've made some comments in the code that are hopefully helpful, but let me know if something is unclear. I'm not sure the resulting alchemical path is identical to that used in GROMACS, so I can't say anything precisely about it's efficiency or behavior, which I haven't well-characterized. However, in all of my tests, I get the same potential energies as GROMACS for the coupled and decoupled states, which is what's important to me as a validation (in addition to that I can also reproduce free energies of solvation).