Merge pull request #95 from MobleyLab/propagation

Features to boost move acceptance and general code cleanup.

blues/example.py

@@ -18,20 +18,48 @@
 import parmed
 from simtk import openmm
 from optparse import OptionParser
+import sys
+import logging
 
-def runNCMC(platform_name):
-    #Define some options
-    opt = { 'temperature' : 300.0, 'friction' : 1, 'dt' : 0.002,
-            'nIter' : 10, 'nstepsNC' : 10, 'nstepsMD' : 5000,
-            'nonbondedMethod' : 'PME', 'nonbondedCutoff': 10, 'constraints': 'HBonds',
-            'trajectory_interval' : 1000, 'reporter_interval' : 1000,
-            'platform' : platform_name,
-            'verbose' : True }
+def init_logger():
+    logger = logging.getLogger()
+    logger.setLevel(logging.DEBUG)
+    formatter = logging.Formatter('%(asctime)s-%(levelname)s-%(message)s')
+    # Write to File
+    fh = logging.FileHandler('blues-example.log')
+    fh.setLevel(logging.INFO)
+    fh.setFormatter(formatter)
+    logger.addHandler(fh)
+
+    # Stream to terminal
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.INFO)
+    ch.setFormatter(formatter)
+    logger.addHandler(ch)
+    return logger
+
+def runNCMC(platform_name, nstepsNC, nprop, outfname):
+
+    logger = init_logger()
 
     #Generate the ParmEd Structure
     prmtop = utils.get_data_filename('blues', 'tests/data/eqToluene.prmtop')#
     inpcrd = utils.get_data_filename('blues', 'tests/data/eqToluene.inpcrd')
     struct = parmed.load_file(prmtop, xyz=inpcrd)
+    logger.info('Structure: %s' % struct.topology)
+
+    #Define some options
+    opt = { 'temperature' : 300.0, 'friction' : 1, 'dt' : 0.002,
+            'nIter' : 100, 'nstepsNC' : nstepsNC, 'nstepsMD' : 5000, 'nprop' : nprop,
+            'nonbondedMethod' : 'PME', 'nonbondedCutoff': 10,
+            'constraints': 'HBonds', 'freeze_distance' : 5.0,
+            'trajectory_interval' : 1000, 'reporter_interval' : 1000,
+            'ncmc_traj' : None, 'write_move' : True,
+            'platform' : platform_name,
+            'verbose' : False}
+
+    for k,v in opt.items():
+        logger.debug('Options: {} = {}'.format(k,v))
 
     #Define the 'model' object we are perturbing here.
     # Calculate particle masses of object to be moved
@@ -45,23 +73,41 @@ def runNCMC(platform_name):
     simulations = SimulationFactory(struct, ligand_mover, **opt)
     simulations.createSimulationSet()
 
+    # Add reporters to MD simulation.
+    traj_reporter = openmm.app.DCDReporter(outfname+'-nc{}.dcd'.format(nstepsNC), opt['trajectory_interval'])
+    progress_reporter = openmm.app.StateDataReporter(sys.stdout, separator="\t",
+                                reportInterval=opt['reporter_interval'],
+                                step=True, totalSteps=opt['nIter']*opt['nstepsMD'],
+                                time=True, speed=True, progress=True,
+                                elapsedTime=True, remainingTime=True)
+    simulations.md.reporters.append(traj_reporter)
+    simulations.md.reporters.append(progress_reporter)
+
+    # Run BLUES Simulation
     blues = Simulation(simulations, ligand_mover, **opt)
-    blues.runNCMC()
+    blues.run(opt['nIter'])
 
 parser = OptionParser()
 parser.add_option('-f', '--force', action='store_true', default=False,
                   help='run BLUES example without GPU platform')
+parser.add_option('-n','--ncmc', dest='nstepsNC', type='int', default=5000,
+                  help='number of NCMC steps')
+parser.add_option('-p','--nprop', dest='nprop', type='int', default=5,
+                  help='number of propgation steps')
+parser.add_option('-o','--output', dest='outfname', type='str', default="blues",
+                  help='Filename for output DCD')
 (options, args) = parser.parse_args()
 
-platformNames = [openmm.Platform.getPlatform(i).getName() for i in range(openmm.Platform.getNumPlatforms())]
 
-if 'OpenCL' in platformNames:
-    runNCMC('OpenCL')
-elif 'CUDA' in platformNames:
-    runNCMC('CUDA')
+
+platformNames = [openmm.Platform.getPlatform(i).getName() for i in range(openmm.Platform.getNumPlatforms())]
+if 'CUDA' in platformNames:
+    runNCMC('CUDA', options.nstepsNC, options.nprop, options.outfname)
+elif 'OpenCL' in platformNames:
+    runNCMC('OpenCL',options.nstepsNC, options.nprop, options.outfname)
 else:
     if options.force:
-        runNCMC('CPU')
+        runNCMC('CPU', options.nstepsNC, options.outfname)
     else:
         print('WARNING: Could not find a valid CUDA/OpenCL platform. BLUES is not recommended on CPUs.')
         print("To run on CPU: 'python blues/example.py -f'")

blues/integrators.py

@@ -68,6 +68,9 @@ def __init__(self,
                  measure_shadow_work=False,
                  measure_heat=True,
                  nsteps_neq=100,
+                 nprop=1,
+                 prop_lambda_min=0.2,
+                 prop_lambda_max=0.8,
                  *args, **kwargs):
         """
         Parameters
@@ -113,6 +116,15 @@ def __init__(self,
         self.addGlobalVariable("perturbed_pe", 0)
         self.addGlobalVariable("unperturbed_pe", 0)
         self.addGlobalVariable("first_step", 0)
+        self.addGlobalVariable("nprop", nprop)
+        self.addGlobalVariable("prop", 1)
+        self.addGlobalVariable("prop_lambda_min", prop_lambda_min)
+        self.addGlobalVariable("prop_lambda_max", prop_lambda_max)
+        self._registered_step_types['H'] = (self._add_alchemical_perturbation_step, False)
+        self.addGlobalVariable("debug", 0)
+
+
+
         try:
             self.getGlobalVariableByName("shadow_work")
         except:
@@ -149,16 +161,48 @@ def _add_integrator_steps(self):
             self.addComputeGlobal("first_step", "1")
             self.addComputeGlobal("unperturbed_pe", "energy")
             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")
 
+            self.beginWhileBlock("prop < nprop")
+            self.addComputeGlobal("prop", "prop + 1")
+
+            super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
+            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()
 
+    def _add_alchemical_perturbation_step(self):
+        """
+        Add alchemical perturbation step, accumulating protocol work.
+        TODO: Extend this to be able to handle force groups?
+        """
+        # Store initial potential energy
+        self.beginIfBlock("prop = 1")
+        self.addComputeGlobal("debug", "debug + 1")
+        self.addComputeGlobal("Eold", "energy")
+
+        # Update lambda and increment that tracks updates.
+        self.addComputeGlobal('lambda', '(lambda_step+1)/n_lambda_steps')
+        self.addComputeGlobal('lambda_step', 'lambda_step + 1')
+
+        # Update all slaved alchemical parameters
+        self._add_update_alchemical_parameters_step()
+
+        # Accumulate protocol work
+        self.addComputeGlobal("Enew", "energy")
+        self.addComputeGlobal("protocol_work", "protocol_work + (Enew-Eold)")
+        self.endBlock()
 
     def getLogAcceptanceProbability(self, context):
         #TODO remove context from arguments if/once ncmc_switching is changed
@@ -170,10 +214,10 @@ def getLogAcceptanceProbability(self, context):
     def reset(self):
         self.setGlobalVariableByName("step", 0)
         self.setGlobalVariableByName("lambda", 0.0)
-#        self.setGlobalVariableByName("total_work", 0.0)
         self.setGlobalVariableByName("protocol_work", 0.0)
         self.setGlobalVariableByName("shadow_work", 0.0)
         self.setGlobalVariableByName("first_step", 0)
         self.setGlobalVariableByName("perturbed_pe", 0.0)
         self.setGlobalVariableByName("unperturbed_pe", 0.0)
-
+        self.setGlobalVariableByName("prop", 1)
+        super(AlchemicalExternalLangevinIntegrator, self).reset()