Remaking forces in OpenMM

Repo

OpenMM is a great tool for molecular dynamics, in part because it is fully customizable via the Python API. On the other hand, with great flexibility comes great complexity. A single incorrect choice of settings can lead to simulations that explode for seemingly no reason, or (worse) run fine but produce incorrect results. One common way to sanity check anything involving customized behaviour is to calculate the potential energy of the customized system and compare to the native version.

To that end, I re-wrote the energy functions (as Custom*Force objects) for bonds, angles, torsions, and non-bonded (electrostatic and lennard-jones) forces, which covers everything for common additive force fields. While this just involves re-writing the energy function and re-using the parameters of a native system, it becomes more useful when adding terms that distort the potential energy for use in enhanced sampling or free energy calculation.

For instance, these custom force functions are the starting point for implementing free energy estimation techniques such as generalized tempering or alchemical annihilation, which I'll come back to!

The replacement functions are copied here for convenience:

Replacing bonds:

def replaceBonds(system):
    forces = system.getForces()
    for c, f in enumerate(forces):
        if isinstance(f, HarmonicBondForce):  
            energy_expression = "0.5*k*(r-r0)^2;"      
            new_bond_force = CustomBondForce(energy_expression)
            new_bond_force.addPerBondParameter("k")
            new_bond_force.addPerBondParameter("r0")   
            for bond in range(f.getNumBonds()):
                
                pars = f.getBondParameters(bond)
                new_bond_force.addBond(pars[0], pars[1], [pars[3], pars[2]])
            system.addForce(new_bond_force)
        
    for c, f in enumerate(forces):
        if isinstance(f, HarmonicBondForce):
            system.removeForce(c)
      

Replacing angles:

def replaceAngles(system):
    forces = system.getForces()
    for c, f in enumerate(forces):
        if isinstance(f, HarmonicAngleForce):
            energy_expression = "0.5*k*(theta-theta0)^2;"
            new_angle_force = CustomAngleForce(energy_expression)
            new_angle_force.addPerAngleParameter("k")
            new_angle_force.addPerAngleParameter("theta0")
            for angle in range(f.getNumAngles()):

                a1,a2,a3, theta0, k = f.getAngleParameters(angle)
                new_angle_force.addAngle(a1, a2, a3, [k,theta0])
            system.addForce(new_angle_force)

    for c, f in enumerate(forces):
        if isinstance(f, HarmonicAngleForce):
            system.removeForce(c)
      

Replacing torsions:

def replaceTorsions(system):
    forces = system.getForces()
    for c, f in enumerate(forces):
        if isinstance(f, PeriodicTorsionForce):
            energy_expression = "k*(1+cos(periodicity*theta-theta0))"
            new_torsion_force = CustomTorsionForce(energy_expression)
            new_torsion_force.addPerTorsionParameter("k");
            new_torsion_force.addPerTorsionParameter("periodicity")
            new_torsion_force.addPerTorsionParameter("theta0");

            for torsion_index in range(f.getNumTorsions()):
                a0, a1, a2, a3, periodicity, phase, k = f.getTorsionParameters(torsion_index)
                new_torsion_force.addTorsion(a0, a1, a2, a3, [k, periodicity,phase])
            system.addForce(new_torsion_force)

    for c, f in enumerate(forces):
        if isinstance(f, PeriodicTorsionForce):
            system.removeForce(c)
      

Replacing nonbonded-interactions (note: this uses reaction field and, at present, requires turning the long-range dispersion correction off):

def replaceNonbonded(system):
    forces = system.getForces()
    for c, f in enumerate(forces):
        if isinstance(f, NonbondedForce):
            original_nbforce = f
            ONE_4PI_EPS0 = 138.935456
            epsilon_solvent = original_nbforce.getReactionFieldDielectric()
            r_cutoff = original_nbforce.getCutoffDistance()
            k_rf = r_cutoff**(-3) * ((epsilon_solvent - 1) / (2*epsilon_solvent + 1))
            c_rf = r_cutoff**(-1) * ((3*epsilon_solvent) / (2*epsilon_solvent + 1))

            energy_expression = "electrostatics+sterics;"

            energy_expression += "electrostatics=ONE_4PI_EPS0*chargeprod*(r^(-1) + k_rf*r^2-c_rf);"
            energy_expression += "chargeprod = charge1*charge2;"
            energy_expression += "k_rf = %f;" % (k_rf.value_in_unit_system(md_unit_system))
            energy_expression += "c_rf = %f;" % (c_rf.value_in_unit_system(md_unit_system))
            energy_expression += "ONE_4PI_EPS0 = %f;" % ONE_4PI_EPS0

            energy_expression += "sterics=4*epsilon*((sigma/r)^12 - (sigma/r)^6);"
            energy_expression += "epsilon = sqrt(epsilon1*epsilon2);"
            energy_expression += "sigma = 0.5*(sigma1+sigma2);"

            new_custom_nonbonded_force = openmm.CustomNonbondedForce(energy_expression)
            new_custom_nonbonded_force.addPerParticleParameter('charge')
            new_custom_nonbonded_force.addPerParticleParameter('sigma')
            new_custom_nonbonded_force.addPerParticleParameter('epsilon')

            new_custom_nonbonded_force.setNonbondedMethod(original_nbforce.getNonbondedMethod())

            new_custom_nonbonded_force.setCutoffDistance(original_nbforce.getCutoffDistance())
            new_custom_nonbonded_force.setUseLongRangeCorrection(False)

            energy_expression = "4*epsilon*((sigma/r)^12 - (sigma/r)^6) + ONE_4PI_EPS0*chargeprod/r;"
            energy_expression += "ONE_4PI_EPS0 = {:f};".format(ONE_4PI_EPS0)  # already in OpenMM units
            new_custom_bond_force = openmm.CustomBondForce(energy_expression)
            new_custom_bond_force.addPerBondParameter('chargeprod')
            new_custom_bond_force.addPerBondParameter('sigma')
            new_custom_bond_force.addPerBondParameter('epsilon')

            for index in range(system.getNumParticles()):
                [charge, sigma, epsilon] = original_nbforce.getParticleParameters(index)
                new_custom_nonbonded_force.addParticle([charge, sigma, epsilon])

            for index in range(original_nbforce.getNumExceptions()):
                idx, jdx, c, s, eps = original_nbforce.getExceptionParameters(index)
                new_custom_nonbonded_force.addExclusion(idx, jdx)
                c_value = c/elementary_charge**2
                eps_value = eps/(kilojoule/mole)
                if c_value != 0 or eps_value!=0:
                    new_custom_bond_force.addBond(idx, jdx, [c, s, eps])

            system.addForce(new_custom_nonbonded_force)
            system.addForce(new_custom_bond_force)


    for c, f in enumerate(forces):
        if isinstance(f, NonbondedForce):
            system.removeForce(c)