[WIP] Control where additional nprop steps are added in the lambda protocol

blues/integrators.py

@@ -69,6 +69,10 @@ class AlchemicalExternalLangevinIntegrator(AlchemicalNonequilibriumLangevinInteg
     nprop : int (Default: 1)
         Controls the number of propagation steps to add in the lambda
         region defined by `prop_lambda`.
+    propRegion: {'sterics', 'electrostatics'} (Default: 'sterics')
+        Specifies the where additional propogation steps are focused. If 'sterics',
+        then additional propogation steps are placed around the middle depending on
+        prop_lambda. If 'electrostatics', then the propogation steps are placed outisde prop_lambda.
 
     Attributes
     ----------
@@ -86,6 +90,26 @@ class AlchemicalExternalLangevinIntegrator(AlchemicalNonequilibriumLangevinInteg
     - An NCMC algorithm with Metropolized integrator:
         splitting="O { V R H R V } O"
 
+    Changing how the sterics and electrostatics are distributed
+
+    40% electrostatics, 60% sterics (Default)
+    >>> alchemical_functions = {lambda_sterics:'min(1, (1/0.3)abs(lambda-0.5))',
+    >>> lambda_electrostatics:'step(0.2-lambda) - 1/0.2lambdastep(0.2-lambda) + 1/0.2(lambda-0.8)*step(lambda-0.8)'}
+    >>> AlchemicalExternalLangevinIntegrator(alchemical_functions=alchemical_functions)
+
+    65% electrostatics, 35% sterics
+    >>> alchemical_functions = {lambda_sterics:'min(1, (1/0.175)abs(lambda-0.5)',
+    >>> lambda_electrostatics:'step(0.325-lambda) - 1/0.325lambda step(0.325-lambda) + 1/0.325(lambda-0.675)*step(lambda-0.675)'}
+    >>> AlchemicalExternalLangevinIntegrator(alchemical_functions=alchemical_functions)
+
+    90% electrostatics, 10% sterics, with 5 propogation steps per lambda increase/decrease in the electrostatic region
+    >>> alchemical_functions = {lambda_sterics:'min(1, (1/0.05)abs(lambda-0.5)',
+    >>> lambda_electrostatics:'step(0.45-lambda) - 1/0.45lambdastep(0.45-lambda) + 1/0.45(lambda-0.55)*step(lambda-0.55)'}
+    >>> AlchemicalExternalLangevinIntegrator(alchemical_functions=alchemical_functions, nprop=5, propRegion='electrostatics')
+
+
+
+
 
     References
     ----------
@@ -107,6 +131,7 @@ def __init__(self,
                  nsteps_neq=100,
                  nprop=1,
                  prop_lambda=0.3,
+                 propRegion='sterics',
                  *args,
                  **kwargs):
         # call the base class constructor
@@ -133,6 +158,10 @@ def __init__(self,
         self.addGlobalVariable("prop", 1)
         self.addGlobalVariable("prop_lambda_min", self._prop_lambda[0])
         self.addGlobalVariable("prop_lambda_max", self._prop_lambda[1])
+        if propRegion=='sterics':
+            self.addGlobalVariable("propRegion", 1)
+        elif propRegion=='electrostatics':
+            self.addGlobalVariable("propRegion", -1)
         # Behavior changed in https://github.com/choderalab/openmmtools/commit/7c2630050631e126d61b67f56e941de429b2d643#diff-5ce4bc8893e544833c827299a5d48b0d
         self._step_dispatch_table['H'] = (self._add_alchemical_perturbation_step, False)
         #$self._registered_step_types['H'] = (
@@ -174,39 +203,55 @@ def _add_integrator_steps(self):
         # Main body
         if self._n_steps_neq == 0:
             # If nsteps = 0, we need to force execution on the first step only.
-            self.beginIfBlock('step = 0')
+            self.beginIfBlock('step = 0')#
             super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
             self.addComputeGlobal("step", "step + 1")
-            self.endBlock()
+            self.endBlock()#
         else:
             #call the superclass function to insert the appropriate steps, provided the step number is less than n_steps
-            self.beginIfBlock("step < nsteps")
+            self.beginIfBlock("step < nsteps")#
             self.addComputeGlobal("perturbed_pe", "energy")
-            self.beginIfBlock("first_step < 1")
+            self.beginIfBlock("first_step < 1")##
             #TODO write better test that checks that the initial work isn't gigantic
             self.addComputeGlobal("first_step", "1")
             self.addComputeGlobal("unperturbed_pe", "energy")
-            self.endBlock()
+            self.endBlock()##
             #initial iteration
             self.addComputeGlobal("protocol_work", "protocol_work + (perturbed_pe - unperturbed_pe)")
             super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
-            #if more propogation steps are requested
-            self.beginIfBlock("lambda > prop_lambda_min")
-            self.beginIfBlock("lambda <= prop_lambda_max")
+            #if more propogation steps are requested, focus at middle of region (sterics)
+            self.beginIfBlock("propRegion > 0")######
+            self.beginIfBlock("lambda > prop_lambda_min")###
+            self.beginIfBlock("lambda <= prop_lambda_max")####
 
-            self.beginWhileBlock("prop < nprop")
+            self.beginWhileBlock("prop < nprop")#####
             self.addComputeGlobal("prop", "prop + 1")
 
             super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
-            self.endBlock()
-            self.endBlock()
-            self.endBlock()
+            self.endBlock()#####
+            self.endBlock()####
+            self.endBlock()###
+            self.endBlock()######
+            #else we want to focus the propogation steps at the beginning (electrostatics)
+            self.beginIfBlock("propRegion < 0")######
+            self.beginIfBlock("lambda < prop_lambda_min")###
+            self.beginIfBlock("lambda >= prop_lambda_max")####
+
+            self.beginWhileBlock("prop < nprop")#####
+            self.addComputeGlobal("prop", "prop + 1")
+
+            super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
+            self.endBlock()#####
+            self.endBlock()####
+            self.endBlock()###
+            self.endBlock()######
+
             #ending variables to reset
             self.addComputeGlobal("unperturbed_pe", "energy")
             self.addComputeGlobal("step", "step + 1")
             self.addComputeGlobal("prop", "1")
 
-            self.endBlock()
+            self.endBlock()#
 
     def _add_alchemical_perturbation_step(self):
         """

blues/reporters.py

@@ -63,11 +63,11 @@ def addLoggingLevel(levelName, levelNum, methodName=None):
         methodName = levelName.lower()
 
     if hasattr(logging, levelName):
-        logging.warn('{} already defined in logging module'.format(levelName))
+        logging.warning('{} already defined in logging module'.format(levelName))
     if hasattr(logging, methodName):
-        logging.warn('{} already defined in logging module'.format(methodName))
+        logging.warning('{} already defined in logging module'.format(methodName))
     if hasattr(logging.getLoggerClass(), methodName):
-        logging.warn('{} already defined in logger class'.format(methodName))
+        logging.warning('{} already defined in logger class'.format(methodName))
 
     # This method was inspired by the answers to Stack Overflow post
     # http://stackoverflow.com/q/2183233/2988730, especially

blues/tests/test_simulation.py

@@ -444,6 +444,8 @@ def test_blues_simulationRunYAML(self, tmpdir, structure, tol_atom_indices, syst
               nIter: 1
               nstepsMD: 2
               nstepsNC: 2
+              propRegion: sterics
+              nprop: 2
               platform: CPU
 
             md_reporters:
@@ -533,3 +535,135 @@ def test_blues_simulationRunPython(self, systems, simulations, engine, tmpdir, s
         #Check that our system has run dynamics
         pos_compare = np.not_equal(before_iter, after_iter).all()
         assert pos_compare
+
+    def test_blues_NCMCSteps_simulation(self, tmpdir, structure, tol_atom_indices, system_cfg, engine):
+        yaml_cfg = """
+            output_dir: .
+            outfname: tol-test
+            logger:
+              level: info
+              stream: True
+
+            system:
+              nonbondedMethod: PME
+              nonbondedCutoff: 8.0 * angstroms
+              constraints: HBonds
+
+            simulation:
+              dt: 0.002 * picoseconds
+              friction: 1 * 1/picoseconds
+              temperature: 300 * kelvin
+              nIter: 1
+              nstepsMD: 2
+              nstepsNC: 20
+              propRegion: sterics
+              propLambda: 0.3
+
+              nprop: 2
+              platform: CPU
+
+            md_reporters:
+              stream:
+                title: md
+                reportInterval: 1
+                totalSteps: 2 # nIter * nstepsMD
+                step: True
+                speed: True
+                progress: True
+                remainingTime: True
+                currentIter : True
+            ncmc_reporters:
+              stream:
+                title: ncmc
+                reportInterval: 1
+                totalSteps: 2 # Use nstepsNC
+                step: True
+                speed: True
+                progress: True
+                remainingTime: True
+                protocolWork : True
+                alchemicalLambda : True
+                currentIter : True
+        """
+        print('Testing Simulation.run() from YAML')
+        yaml_cfg = Settings(yaml_cfg)
+        cfg = yaml_cfg.asDict()
+        cfg['output_dir'] = tmpdir
+        # os.getenv is equivalent, and can also give a default value instead of `None`
+        PLATFORM = os.getenv('OMM_PLATFORM', 'CPU')
+        cfg['simulation']['platform'] = PLATFORM
+        systems = SystemFactory(structure, tol_atom_indices, cfg['system'])
+        simulations = SimulationFactory(systems, engine, cfg['simulation'], None, None)
+
+        blues = BLUESSimulation(simulations)
+        blues._md_sim.minimizeEnergy()
+        blues._alch_sim.minimizeEnergy()
+        blues._ncmc_sim.minimizeEnergy()
+        blues.run()
+        debug_var1 = blues._ncmc_sim.integrator.getGlobalVariableByName('debug')
+
+        yaml_cfg = """
+            output_dir: .
+            outfname: tol-test
+            logger:
+              level: info
+              stream: True
+
+            system:
+              nonbondedMethod: PME
+              nonbondedCutoff: 8.0 * angstroms
+              constraints: HBonds
+
+            simulation:
+              dt: 0.002 * picoseconds
+              friction: 1 * 1/picoseconds
+              temperature: 300 * kelvin
+              nIter: 1
+              nstepsMD: 2
+              nstepsNC: 20
+              propRegion: electrostatics
+              nprop: 2
+              propLambda: 0.2
+              platform: CPU
+
+            md_reporters:
+              stream:
+                title: md
+                reportInterval: 1
+                totalSteps: 2 # nIter * nstepsMD
+                step: True
+                speed: True
+                progress: True
+                remainingTime: True
+                currentIter : True
+            ncmc_reporters:
+              stream:
+                title: ncmc
+                reportInterval: 1
+                totalSteps: 2 # Use nstepsNC
+                step: True
+                speed: True
+                progress: True
+                remainingTime: True
+                protocolWork : True
+                alchemicalLambda : True
+                currentIter : True
+        """
+        print('Testing Simulation.run() from YAML')
+        yaml_cfg = Settings(yaml_cfg)
+        cfg = yaml_cfg.asDict()
+        cfg['output_dir'] = tmpdir
+        # os.getenv is equivalent, and can also give a default value instead of `None`
+        PLATFORM = os.getenv('OMM_PLATFORM', 'CPU')
+        cfg['simulation']['platform'] = PLATFORM
+        systems = SystemFactory(structure, tol_atom_indices, cfg['system'])
+        simulations = SimulationFactory(systems, engine, cfg['simulation'], None, None)
+
+        blues = BLUESSimulation(simulations)
+        blues._md_sim.minimizeEnergy()
+        blues._alch_sim.minimizeEnergy()
+        blues._ncmc_sim.minimizeEnergy()
+        blues.run()
+        debug_var2 = blues._ncmc_sim.integrator.getGlobalVariableByName('debug')
+        assert debug_var2 == debug_var1
+

blues/utils.py

@@ -86,7 +86,7 @@ def print_host_info(simulation):
     logger.info(msg)
 
 
-def calculateNCMCSteps(nstepsNC=0, nprop=1, propLambda=0.3, **kwargs):
+def calculateNCMCSteps(nstepsNC=0, nprop=1, propLambda=0.3, propRegion='sterics', **kwargs):
     """
     Calculates the number of NCMC switching steps.
 
@@ -114,24 +114,38 @@ def calculateNCMCSteps(nstepsNC=0, nprop=1, propLambda=0.3, **kwargs):
             logger.error(msg)
             sys.exit(1)
     # Calculate the total number of lambda switching steps
-    lambdaSteps = nstepsNC / (2 * (nprop * propLambda + 0.5 - propLambda))
+    if propRegion == 'sterics':
+        lambdaSteps = nstepsNC / (2 * (nprop * propLambda + 0.5 - propLambda))
+    elif propRegion == 'electrostatics':
+        #TODO update this
+        lambdaSteps = nstepsNC / (2*(propLambda+nprop*(0.5 - propLambda)))
+    else:
+        msg = 'propRegion must be either `sterics`, or `electrostatics`'
+        logger.error(msg)
+        sys.exit(1)
     if int(lambdaSteps) % 2 == 0:
         lambdaSteps = int(lambdaSteps)
     else:
         lambdaSteps = int(lambdaSteps) + 1
-
-    # Calculate number of lambda steps inside/outside region with extra propgation steps
+        # Calculate number of lambda steps inside/outside region with extra propgation steps
     in_portion = (propLambda) * lambdaSteps
     out_portion = (0.5 - propLambda) * lambdaSteps
-    in_prop = int(nprop * (2 * floor(in_portion)))
-    out_prop = int((2 * ceil(out_portion)))
+    if propRegion == 'sterics':
+        in_prop = int(nprop * (2 * floor(in_portion)))
+        out_prop = int((2 * ceil(out_portion)))
+
+    elif propRegion == 'electrostatics':
+        in_prop = int((2 * floor(in_portion)))
+        out_prop = int(nprop *(2 * ceil(out_portion)))
     propSteps = int(in_prop + out_prop)
 
+
+
     if propSteps != nstepsNC:
-        logger.warn("nstepsNC=%s is incompatible with prop_lambda=%s and nprop=%s." % (nstepsNC, propLambda, nprop))
-        logger.warn("Changing NCMC protocol to %s lambda switching within %s total propagation steps." % (lambdaSteps,
+        logger.warning("nstepsNC=%s is incompatible with prop_lambda=%s and nprop=%s." % (nstepsNC, propLambda, nprop))
+        logger.warning("Changing NCMC protocol to %s lambda switching within %s total propagation steps." % (lambdaSteps,
                                                                                                           propSteps))
-        nstepsNC = lambdaSteps
+    nstepsNC = lambdaSteps
 
     moveStep = int(nstepsNC / 2)
     ncmc_parameters = {