Merge pull request #470 from choderalab/bryce-exp

SAMS sampler and other improvements

analyse_sams_convergence.py

@@ -0,0 +1,39 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
+import os, sys
+from netCDF4 import Dataset
+from glob import glob
+
+
+if __name__ == '__main__':
+    directory = sys.argv[1]
+    files = sorted(glob(os.path.join(os.getcwd(), directory, '*.nc')))
+    files = [x for x in files if 'checkpoint' not in x]
+    f, axarr = plt.subplots(2, 3, sharex=False, sharey=False, figsize=(16, 8))
+    for i, filename in enumerate(files):
+        phase = filename.split('-')[1].rstrip('.nc')
+        ncfile = Dataset(filename, 'r')
+        logZ = ncfile.groups['online_analysis'].variables['logZ_history']
+        n_iterations, n_states = logZ.shape
+        print(n_iterations, n_states)
+        axarr[i, 0].plot(logZ, '.')
+        axarr[i, 0].set_xlabel('iteration')
+        axarr[i, 0].set_ylabel('logZ / kT')
+        axarr[i, 0].set_title('%s_%s' % (phase, directory))
+        #axarr[i, 1].plot(np.transpose(np.transpose(logZ[:-2]) - logZ[:-2, 0]), '.')
+        #axarr[i, 1].set_xlabel('iteration')
+        #axarr[i, 1].set_ylabel('logZ / kT')
+        states = ncfile.variables['states']
+        n_iterations, n_replicas = states.shape
+        axarr[i, 1].plot(states, '.')
+        axarr[i, 1].set_xlabel('iteration')
+        axarr[i, 1].set_ylabel('thermodynamic state')
+        axarr[i, 1].axis([0, n_iterations, 0, n_states])
+        gamma = ncfile.groups['online_analysis'].variables['gamma_history']
+        axarr[i, 2].plot(gamma, '.')
+        axarr[i, 2].set_xlabel('iteration')
+        axarr[i, 2].set_ylabel('gamma')
+
+    f.tight_layout()
+    f.savefig('%s.png' % directory, dpi=300)

devtools/conda-recipe/meta.yaml

@@ -33,6 +33,7 @@ requirements:
     - pdbfixer
     - lxml
     - networkx >=2.0
+    - yank
 
   run:
     - python
@@ -61,6 +62,7 @@ requirements:
     - dask
     - distributed
     - progressbar2
+    - yank
 
 test:
   requires:

examples/cdk2-example/cdk2_setup.yaml

@@ -9,7 +9,7 @@ ligand_file: CDK2_ligands.mol2
 old_ligand_index: 14
 new_ligand_index: 15
 
-#Provide the list of forcefield files. Non-standard (such as gaff.xml) should
+#Provide the list of forcefield files. Non-standard files should
 #be provided with a full path
 forcefield_files:
     - gaff.xml
@@ -27,41 +27,61 @@ solvent_padding: 9.0
 
 
 #The name of the pickle file where we will save the setup object
-save_setup_pickle_as: fesetup.pkl
+save_setup_pickle_as: fesetup_hbonds.pkl
 
-#whether to compute the solvent or complex phase
-phase: solvent
+#the type of free energy calculation.
+#currently, this could be either nonequilibrium or sams
+fe_type: nonequilibrium
 
-#the forward switching functions. The reverse ones will be computed from this
+#the forward switching functions. The reverse ones will be computed from this. Only valid for nonequilibrium switching
 forward_functions:
     lambda_sterics: lambda
     lambda_electrostatics: lambda
     lambda_bonds: lambda
     lambda_angles: lambda
     lambda_torsions: lambda
 
-#The number of equilibrium steps to take between nonequilibrium switching events
+#the number of equilibration iterations:
+n_equilibration_iterations: 100
+
+#The number of equilibrium steps to take between nonequilibrium switching events (only valid for nonequiibrium switching)
 n_equilibrium_steps_per_iteration: 100
 
-#The length of the ncmc protocol
+#The length of the ncmc protocol (only valid for nonequilibrium switching)
 n_steps_ncmc_protocol: 50
 
 #The number of NCMC steps per move application. This controls the output frequency
-#1 step/move application means writing out the work at every step.
-n_steps_per_move_application: 1
+#1 step/move application means writing out the work at every step. (only valid for nonequilibrium switching)
+n_steps_per_move_application: 10
 
 #where to put the trajectories
-trajectory_directory: solvent_test
+trajectory_directory: cdk2_neq_hbonds
 
 #how to prefix the trajectory files (project-specific name)
-trajectory_prefix: cdk02
+trajectory_prefix: cdk2
 
 #which atoms to save (MDTraj selection syntax)
-atom_selection: all
+atom_selection: not water
+
+#which phases do we want to run (solvent, complex, or both solvent and complex in the list are acceptable):
+phases:
+    - solvent
+
+#timestep in fs
+timestep: 1.0
+
+#splitting string for equilibrium integration (only valid for nonequilibrium switching):
+eq_splitting: V R O R V
+
+#splitting string for nonequilibrium switching (only valid for nonequilibrium switching):
+neq_splitting: V R O H R V
+
+#whether to measure shadow work (only valid for nonequilibrium switching):
+measure_shadow_work: True
 
-#The location of the schduler. If it's null, a localhost scheduler is created
-scheduler_address: localhost
+#The location of the schduler. If it's null, a localhost scheduler is created (only valid for nonequilibrium switching)
+scheduler_address: null
 
-#how many iterations to run (n_cycles*n_iterations_per_cycle)
+#how many iterations to run (n_cycles*n_iterations_per_cycle) (only valid for nonequilibrium switching)
 n_cycles: 5
 n_iterations_per_cycle: 1

hybrid_sams_sampler_test.py

@@ -0,0 +1,23 @@
+import numpy as np
+from simtk import unit
+from openmmtools import mcmc
+from perses.dispersed.relative_setup import HybridTopologyFactory, HybridSAMSSampler
+import logging
+logging.basicConfig(level=logging.DEBUG)
+
+if __name__ == '__main__':
+    mcl1tp = np.load('/home/ballen/PycharmProjects/perses/examples/cdk2-example/cdk2_sams_hbonds/cdk2topology_proposals.npy').item()
+    factory = HybridTopologyFactory(mcl1tp['complex_topology_proposal'], mcl1tp['complex_old_positions'],
+                                    mcl1tp['complex_new_positions'])
+    chss = HybridSAMSSampler(mcmc_moves=mcmc.LangevinDynamicsMove(timestep=2.0 * unit.femtosecond,
+                                                                  collision_rate=5.0 / unit.picosecond,
+                                                                  n_steps=1000,
+                                                                  reassign_velocities=False,
+                                                                  n_restart_attempts=6),
+                             hybrid_factory=factory)
+    chss.setup(n_states=100, temperature=300.0 * unit.kelvin,
+               storage_file='/media/ballen/overflow/perses/complex_test_100.nc', checkpoint_interval=1)
+    chss.minimize()
+    chss.equilibrate(10)
+    chss.extend(10000)
+    print("DONE FINALLY!!!")

perses/annihilation/new_relative.py

@@ -98,6 +98,20 @@ def __init__(self, topology_proposal, current_positions, new_positions, use_disp
         if softcore_method not in self._known_softcore_methods:
             raise ValueError("Softcore method {} is not a valid method. Acceptable options are default, amber, and classic".format(softcore_method))
 
+        if softcore_alpha is None:
+            self.softcore_alpha = 0.5
+        else:
+            self.softcore_alpha = softcore_alpha
+
+        if softcore_beta is None:
+            self.softcore_beta = 12*unit.angstrom**2
+        else:
+            self.softcore_beta = softcore_beta
+
+        if softcore_method not in self._known_softcore_methods:
+            raise ValueError("Softcore method {} is not a valid method. Acceptable options are default, amber, and classic".format(softcore_method))
+
+
         self._softcore_method = softcore_method
 
         if functions:
@@ -945,8 +959,8 @@ def _nonbonded_custom_bond_force(self, sterics_energy_expression, electrostatics
             sterics_energy_expression += 'lambda_sterics = ' + self._functions['lambda_sterics']
             electrostatics_energy_expression += 'lambda_electrostatics = ' + self._functions['lambda_electrostatics']
         custom_bond_force = openmm.CustomBondForce("U_sterics + U_electrostatics;" + sterics_energy_expression + electrostatics_energy_expression)
-        custom_bond_force.addGlobalParameter("lambda_electrostatics", 0.0)
-        custom_bond_force.addGlobalParameter("lambda_sterics", 0.0)
+        #custom_bond_force.addGlobalParameter("lambda_electrostatics", 0.0)
+        #custom_bond_force.addGlobalParameter("lambda_sterics", 0.0)
         custom_bond_force.addGlobalParameter("softcore_alpha", self.softcore_alpha)
         custom_bond_force.addGlobalParameter("softcore_beta", self.softcore_beta)
         custom_bond_force.addPerBondParameter("chargeprodA")

perses/data/cdk2-example/cdk2_setup.yaml

@@ -9,7 +9,7 @@ ligand_file: CDK2_ligands.mol2
 old_ligand_index: 14
 new_ligand_index: 15
 
-#Provide the list of forcefield files. Non-standard (such as gaff.xml) should
+#Provide the list of forcefield files. Non-standard files should
 #be provided with a full path
 forcefield_files:
     - gaff.xml
@@ -27,41 +27,61 @@ solvent_padding: 9.0
 
 
 #The name of the pickle file where we will save the setup object
-save_setup_pickle_as: fesetup.pkl
+save_setup_pickle_as: fesetup_hbonds.pkl
 
-#whether to compute the solvent or complex phase
-phase: solvent
+#the type of free energy calculation.
+#currently, this could be either nonequilibrium or sams
+fe_type: nonequilibrium
 
-#the forward switching functions. The reverse ones will be computed from this
+#the forward switching functions. The reverse ones will be computed from this. Only valid for nonequilibrium switching
 forward_functions:
     lambda_sterics: lambda
     lambda_electrostatics: lambda
     lambda_bonds: lambda
     lambda_angles: lambda
     lambda_torsions: lambda
 
-#The number of equilibrium steps to take between nonequilibrium switching events
+#the number of equilibration iterations:
+n_equilibration_iterations: 100
+
+#The number of equilibrium steps to take between nonequilibrium switching events (only valid for nonequiibrium switching)
 n_equilibrium_steps_per_iteration: 100
 
-#The length of the ncmc protocol
+#The length of the ncmc protocol (only valid for nonequilibrium switching)
 n_steps_ncmc_protocol: 50
 
 #The number of NCMC steps per move application. This controls the output frequency
-#1 step/move application means writing out the work at every step.
-n_steps_per_move_application: 1
+#1 step/move application means writing out the work at every step. (only valid for nonequilibrium switching)
+n_steps_per_move_application: 10
 
 #where to put the trajectories
-trajectory_directory: solvent_test
+trajectory_directory: cdk2_neq_hbonds
 
 #how to prefix the trajectory files (project-specific name)
-trajectory_prefix: cdk02
+trajectory_prefix: cdk2
 
 #which atoms to save (MDTraj selection syntax)
-atom_selection: all
+atom_selection: not water
+
+#which phases do we want to run (solvent, complex, or both solvent and complex in the list are acceptable):
+phases:
+    - solvent
+
+#timestep in fs
+timestep: 1.0
+
+#splitting string for equilibrium integration (only valid for nonequilibrium switching):
+eq_splitting: V R O R V
+
+#splitting string for nonequilibrium switching (only valid for nonequilibrium switching):
+neq_splitting: V R O H R V
+
+#whether to measure shadow work (only valid for nonequilibrium switching):
+measure_shadow_work: True
 
-#The location of the schduler. If it's null, a localhost scheduler is created
-scheduler_address: localhost
+#The location of the schduler. If it's null, a localhost scheduler is created (only valid for nonequilibrium switching)
+scheduler_address: null
 
-#how many iterations to run (n_cycles*n_iterations_per_cycle)
+#how many iterations to run (n_cycles*n_iterations_per_cycle) (only valid for nonequilibrium switching)
 n_cycles: 5
 n_iterations_per_cycle: 1