[WIP] Solvent-solute splitting

.gitignore

@@ -98,5 +98,8 @@ ENV/
 # mkdocs documentation
 /site
 
+# pycharm
+.idea
+
 # mypy
 .mypy_cache/

examples/mts/gbaoab_with_solvent_solute_splitting.py

@@ -0,0 +1,39 @@
+from openmmtools import testsystems
+from openmmtools.forcefactories import split_nb_using_exceptions
+
+testsystem = testsystems.AlanineDipeptideExplicit()  # using PME! should be starting with reaction field
+
+new_system = split_nb_using_exceptions(testsystem.system, testsystem.mdtraj_topology)
+
+from simtk.openmm.app import Simulation
+from openmmtools.integrators import LangevinIntegrator
+from simtk import unit
+
+collision_rate = 1/unit.picoseconds
+#langevin = LangevinIntegrator(splitting='V R O R V', timestep=6 * unit.femtosecond)
+mts = LangevinIntegrator(splitting="V0 V1 R R O R R V1 V1 R R O R R V1 V0", collision_rate=collision_rate, timestep=8 * unit.femtosecond)
+sim = Simulation(testsystem.topology, new_system, mts)
+sim.context.setPositions(testsystem.positions)
+
+print('minimizing energy')
+sim.minimizeEnergy(maxIterations=5)
+
+print('first step')
+sim.step(1)
+pos = sim.context.getState(getPositions=True).getPositions(asNumpy=True)
+print(pos)
+
+print('second step')
+sim.step(2)
+pos = sim.context.getState(getPositions=True).getPositions(asNumpy=True)
+print(pos)
+
+print('next 10 steps')
+sim.step(10)
+pos = sim.context.getState(getPositions=True).getPositions(asNumpy=True)
+print(pos)
+
+print('next 1000 steps')
+sim.step(1000)
+pos = sim.context.getState(getPositions=True).getPositions(asNumpy=True)
+print(pos)

openmmtools/forcefactories.py

@@ -174,6 +174,240 @@ def restrain_atoms(thermodynamic_state, sampler_state, restrained_atoms, sigma=3
     thermodynamic_state.system = system
 
 
+from openmmtools.constants import ONE_4PI_EPS0
+default_energy_expression = "(4*epsilon*((sigma/r)^12-(sigma/r)^6) + k * chargeprod/r); sigma=0.5*(sigma1+sigma2); epsilon=sqrt(epsilon1*epsilon2); chargeprod=charge1*charge2; k={};".format(ONE_4PI_EPS0)
+
+
+def clone_nonbonded_parameters(nonbonded_force, 
+                               energy_expression=default_energy_expression,
+                               energy_prefactor=''):
+    """Creates a new CustomNonbonded force with the same global parameters,
+    per-particle parameters, and exception parameters as """
+
+    allowable_nb_methods = {openmm.NonbondedForce.CutoffPeriodic, openmm.NonbondedForce.CutoffNonPeriodic}
+    assert(nonbonded_force.getNonbondedMethod() in allowable_nb_methods)
+
+
+    # call constructor
+    # The 'energy_prefactor' allows us to easily change sign, or e.g. halve the energy if needed
+    new_force = openmm.CustomNonbondedForce(energy_prefactor+energy_expression)
+    new_force.addPerParticleParameter('charge')
+    new_force.addPerParticleParameter('sigma')
+    new_force.addPerParticleParameter('epsilon')
+
+
+    # go through all of the setter and getter methods
+    new_force.setCutoffDistance(nonbonded_force.getCutoffDistance())
+    #new_force.setEwaldErrorTolerance(nonbonded_force.getEwaldErrorTolerance())
+    #new_force.setForceGroup(nonbonded_force.getForceGroup())
+    new_force.setNonbondedMethod(nonbonded_force.getNonbondedMethod())
+    #new_force.setPMEParameters(*nonbonded_force.getPMEParameters())
+    #new_force.setReactionFieldDielectric(nonbonded_force.getReactionFieldDielectric())
+    #new_force.setReciprocalSpaceForceGroup(nonbonded_force.getReciprocalSpaceForceGroup())
+    new_force.setUseSwitchingFunction(nonbonded_force.getUseSwitchingFunction())
+    new_force.setSwitchingDistance(nonbonded_force.getSwitchingDistance())
+    new_force.setUseLongRangeCorrection(nonbonded_force.getUseDispersionCorrection())
+
+
+    # now add all the particle parameters
+    num_particles = nonbonded_force.getNumParticles()
+    for i in range(num_particles):
+        new_force.addParticle(nonbonded_force.getParticleParameters(i))
+
+    # now add all the exceptions? # TODO: check if we want to do this...
+    #    for exception_index in range(nonbonded_force.getNumExceptions()):
+    #        new_force.addException(nonbonded_force.getExceptionParameters(exception_index))
+
+    # TODO: There's probably a cleaner, more Pythonic way to do this
+    # (this was the most obvious way, but may not work in the future if the OpenMM API changes)
+
+    return new_force
+
+
+def split_nb_using_interaction_groups(system, md_topology):
+    """Construct a copy of system where its nonbonded force has been replaced
+    with two nonbonded forces, one using an interaction group restricted to
+    water-water interactions, and the other using interaction groups
+    restricted to water-solute and solute-solute interactions. Water-water
+    interactions are in force group 0, and all other interactions are in force
+    group 1.
+    """
+
+    # create a copy of the original system: only touch this
+    new_system = copy.deepcopy(system)
+
+    # find the default nonbonded force
+    force_index, nb_force = forces.find_forces(new_system, openmm.NonbondedForce, only_one=True)
+    # create copies for each interaction. Only half in solvent/solvent and solute/solute as we double-count.
+    nb_only_solvent_solvent = clone_nonbonded_parameters(nb_force, energy_prefactor='0.5*')
+    nb_solute = clone_nonbonded_parameters(nb_force)
+
+    # NOTE: these need to be python ints -- not np.int64s -- when passing to addInteractionGroup later!
+    solvent_indices = list(map(int, md_topology.select('water')))
+    solute_indices = list(map(int, md_topology.select('not water')))
+
+    #all_indices = list(range(md_topology.n_atoms))
+
+    nb_only_solvent_solvent.addInteractionGroup(set1=solvent_indices, set2=solvent_indices)
+
+    nb_solute.addInteractionGroup(set1=solute_indices, set2=solute_indices)
+    nb_solute.addInteractionGroup(set1=solute_indices, set2=solvent_indices)
+
+    #nb_solute.addInteractionGroup(set1=solute_indices, set2=all_indices)
+
+    # remove original force, add new forces
+    new_system.removeForce(force_index)
+    new_system.addForce(nb_solute)
+    new_system.addForce(nb_only_solvent_solvent)
+
+    # Set solvent-solvent to fg 0, everything else to fg1
+    nb_only_solvent_solvent.setForceGroup(0)
+    nb_solute.setForceGroup(1)
+
+    # handle non-NonbondedForce's
+    for force in new_system.getForces():
+        if 'Nonbonded' not in force.__class__.__name__:
+            force.setForceGroup(1)
+
+    return new_system
+
+
+def split_nb_using_exceptions(system, md_topology):
+    """Construct a new system where force group 0 contains slow, expensive solvent-solvent
+    interactions, and force group 1 contains fast, cheaper solute-solute and solute-solvent
+    interactions.
+
+    TODO: correct this:
+    Force group 0: default nonbonded force with exceptions for solute-{solute,solvent} pairs
+    Force group 1:
+    * CustomNonbonded force with interaction group for solute-{solute,solvent} pairs
+    * non-Nonbonded forces
+
+    TODO: more informative docstring
+    """
+
+    # create a copy of the original system: only touch this
+    new_system = copy.deepcopy(system)
+
+    # find the default nonbonded force
+    force_index, nb_force = forces.find_forces(new_system, openmm.NonbondedForce, only_one=True)
+    # assert('Cutoff' in nb_force.getNonbondedMethod()) # TODO: this, less jankily
+
+    # create copies for each interaction. Only half in solute/solute as we double count.
+    nb_only_solvent_solvent = nb_force
+    nb_only_solvent_solute = clone_nonbonded_parameters(nb_force)
+    nb_only_solute_solute = clone_nonbonded_parameters(nb_force,energy_prefactor='0.5*')
+
+    # identify solvent-solute pairs
+    # find the pairs (i,j) where i in solute, j in solvent
+
+    # NOTE: these need to be python ints -- not np.int64s -- when passing to addInteractionGroup later!
+    solvent_indices = list(map(int, md_topology.select('water')))
+    solute_indices = list(map(int, md_topology.select('not water')))
+
+    # for each (i,j) set chargeprod, sigma, epsilon to 0, causing these pairs to be omitted completely from force calculation
+    # Remove solute-solute interactions
+    for i in solute_indices:
+        for j in solute_indices:
+            nb_only_solvent_solvent.addException(i, j, 0, 0, 0, replace=True)
+
+    # Remove solvent-solute interactions
+    for i in solvent_indices:
+        for j in solute_indices:
+            nb_only_solvent_solvent.addException(i, j, 0, 0, 0,replace=True)
+
+
+    # Add appropriate interaction groups
+    nb_only_solvent_solute.addInteractionGroup(set1=solute_indices, set2=solvent_indices)
+    nb_only_solute_solute.addInteractionGroup(set1=solute_indices, set2=solute_indices)
+
+    # add new forces
+    new_system.addForce(nb_only_solute_solute)
+    new_system.addForce(nb_only_solvent_solute)
+
+    # Set solvent-solvent to fg 0, everything else to fg1
+    nb_only_solvent_solvent.setForceGroup(0) 
+    nb_only_solvent_solute.setForceGroup(1)
+    nb_only_solute_solute.setForceGroup(1)
+
+    # handle non-NonbondedForce's
+    for force in new_system.getForces():
+        if 'Nonbonded' not in force.__class__.__name__:
+            force.setForceGroup(1)
+
+    return new_system
+
+
+
+# surrogate could deviate in functional form from the NonbondedForce defaults (e.g. using soft-core)
+surrogate_energy_expression = default_energy_expression
+half_surrogate_energy_expression = '0.5*' + surrogate_energy_expression
+minus_surrogate_energy_expression = '-' + surrogate_energy_expression
+minus_half_surrogate_energy_expression = '-' + half_surrogate_energy_expression
+
+# TODO: different energy expressions for protein-protein and protein-solvent interactions?
+
+def split_nb_using_subtraction(system, md_topology,
+                               switching_distance=4.0 * unit.angstrom,
+                               cutoff_distance=5.0 * unit.angstrom,
+                               ):
+    """
+    Force group 0: default NonbondedForce minus surrogate
+    Force group 1: surrogate + non-Nonbonded"""
+
+    new_system = copy.deepcopy(system)
+
+
+    # find the default nonbonded force
+    force_index, force = forces.find_forces(new_system, openmm.NonbondedForce, only_one=True)
+    force.setForceGroup(0)
+
+    # find atom indices for solvent, atom indices for solute
+    # NOTE: these need to be python ints -- not np.int64s -- when passing to addInteractionGroup later!
+    solvent_indices = list(map(int, md_topology.select('water')))
+    solute_indices = list(map(int, md_topology.select('not water')))
+
+    # TODO: handle counterions
+
+    # TODO: smooth cutoff, and check how bad this is with very small cutoff (setUseSwitchingFunction(True)
+    # TODO: soft-core LJ or other variants (maybe think about playing with effective vdW radii?
+
+    # define surrogate forces that need to be added
+    protein_protein_force = clone_nonbonded_parameters(force, half_surrogate_energy_expression)
+    protein_solvent_force = clone_nonbonded_parameters(force, surrogate_energy_expression)
+
+    # define surrogate forces that need to be subtracted
+    minus_protein_protein_force = clone_nonbonded_parameters(force, minus_half_surrogate_energy_expression)
+    minus_protein_solvent_force = clone_nonbonded_parameters(force, minus_surrogate_energy_expression)
+
+    # add forces to new_system
+    for new_force in [protein_protein_force, protein_solvent_force, minus_protein_protein_force, minus_protein_solvent_force]:
+        new_force.setSwitchingDistance(switching_distance)
+        new_force.setCutoffDistance(cutoff_distance)
+        new_system.addForce(new_force)
+
+    # set interaction groups
+    protein_protein_force.addInteractionGroup(set1=solute_indices, set2=solute_indices)
+    protein_solvent_force.addInteractionGroup(set1=solute_indices, set2=solvent_indices)
+
+    minus_protein_protein_force.addInteractionGroup(set1=solute_indices, set2=solute_indices)
+    minus_protein_solvent_force.addInteractionGroup(set1=solute_indices, set2=solvent_indices)
+
+    # set force groups
+    minus_protein_protein_force.setForceGroup(0)
+    minus_protein_solvent_force.setForceGroup(0)
+
+    protein_protein_force.setForceGroup(1)
+    protein_solvent_force.setForceGroup(1)
+
+    # handle non-NonbondedForce's
+    for force in new_system.getForces():
+        if 'Nonbonded' not in force.__class__.__name__:
+            force.setForceGroup(1)
+
+    return new_system
+
+
 if __name__ == '__main__':
     import doctest
     doctest.testmod()

openmmtools/tests/test_forcefactories.py

@@ -196,3 +196,31 @@ def test_replace_reaction_field():
                     name=test_name, platform=platform)
         f.description = "Testing replace_reaction_field on system {} with shifted=True".format(test_name)
         yield f
+
+
+def check_force_group_decomposition_valid(testsystem, force_group_splitter):
+    """force_group_splitter(system, md_topology) returns a copy of system that
+    splits system's forces in some convenient way.
+
+    Check that the copy describes the same system overall, by checking whether it
+    produces identical forces as the original on a slightly randomized configuration.
+    """
+    new_system = force_group_splitter(testsystem.system, testsystem.mdtraj_topology)
+    positions = generate_new_positions(testsystem.system, testsystem.positions, nsteps=50)
+    compare_system_forces(testsystem.system, new_system, positions)
+
+
+from simtk.openmm import app
+def test_split_nb_using_interaction_groups():
+    testsystem = testsystems.AlanineDipeptideExplicit(nonbondedMethod=app.CutoffPeriodic, use_dispersion_correction=False)
+    check_force_group_decomposition_valid(testsystem, split_nb_using_interaction_groups)
+
+
+def test_split_nb_using_exceptions():
+    testsystem = testsystems.AlanineDipeptideExplicit(nonbondedMethod=app.CutoffPeriodic, use_dispersion_correction=False)
+    check_force_group_decomposition_valid(testsystem, split_nb_using_exceptions)
+
+
+def test_split_nb_using_subtraction():
+    testsystem = testsystems.AlanineDipeptideExplicit(nonbondedMethod=app.CutoffPeriodic, use_dispersion_correction=False)
+    check_force_group_decomposition_valid(testsystem, split_nb_using_subtraction)

openmmtools/utils.py

@@ -961,6 +961,37 @@ def _compute_class_hash(openmm_class):
         return float(zlib.adler32(openmm_class.__name__.encode()))
 
 
+def generate_solvent_solute_splitting_string(base_integrator="VRORV", K_p=1, K_r=3):
+    """Generate string representing sequence of V0, V1, R, O steps, where:
+    * force group 1 is assumed to contain fast-changing, cheap-to-evaluate forces, and
+    * force group 0 is assumed to contain slow-changing, expensive-to-evaluate forces.
+
+    Currently only supports solvent-solute splittings of the VRORV (BAOAB / g-BAOAB)
+    integrator.
+
+    Parameters
+    -----------
+    base_integrator: string
+        Currently only supports VRORV
+    K_p: int
+        Number of times to evaluate force group 1 per timestep.
+    K_r: int
+        Number of inner-loop iterations
+
+    Returns
+    -------
+    splitting_string: string
+        Sequence of V0, V1, R, O steps, to be passed to LangevinSplittingIntegrator
+
+    TODO: add doctests
+    """
+    assert (base_integrator == "VRORV" or base_integrator == "BAOAB")
+    Rs = "R " * K_r
+    inner_loop = "V1 " + Rs + "O " + Rs + "V1 "
+    s = "V0 " + inner_loop * K_p + "V0"
+    return s
+
+
 if __name__ == '__main__':
     import doctest
     doctest.testmod()