Merge pull request #5 from ADicksonLab/nbforces_test

The Philadelphia Updates
ADicksonLab · Oct 10, 2023 · 9c18738 · 9c18738
2 parents 125b143 + a942c5b
commit 9c18738
Show file tree

Hide file tree

Showing 20 changed files with 1,894 additions and 223 deletions.
diff --git a/examples/brd_ghost_distribution/build_minimize_heat_noMLF.py b/examples/brd_ghost_distribution/build_minimize_heat_noMLF.py
@@ -17,7 +17,7 @@
 from flexibletopology.utils.integrators import CustomHybridIntegratorRestrictedChargeVariance
 
 from flexibletopology.utils.reporters import H5Reporter
-from flexibletopology.utils.openmmutils import read_params
+from flexibletopology.utils.openmmutils import read_params, getParameters, setParameters
 import sys
 
 from contforceplugin import ContForce
@@ -59,8 +59,9 @@
 SYSTEM_PSF = osp.join(INPUTS_PATH, 'brd2_water_trim.psf') # name of psf file 
 SYSTEM_PDB = osp.join(INPUTS_PATH, 'brd_water_trim.pdb')  # name of pdb file
 BS_SELECTION_STRING = "resid 88 29 81 24" # selection of residues where the centroid is the middle of the binding site
-SYSTEM_CONT_FORCE_IDXS = [1380,1381,1382] # These are indices of atoms in the target that will be made to be continuous with the FT atoms
-                                           # (set to empty list if not using)
+#SYSTEM_CONT_FORCE_IDXS = [1380,1381,1382] # These are indices of atoms in the target that will be made to be continuous with the FT atoms
+#                                           # (set to empty list if not using)
+SYSTEM_CONT_FORCE_IDXS = []
 
 sigma_LowerBound = 0.2
 sigma_coeff = 100000.0
@@ -84,8 +85,6 @@
 pos_arr = np.array(crd.positions.value_in_unit(unit.nanometers))
 
 # MD simulations settings
-CONVERT_FAC = -0.2390057
-
 TEMPERATURES = [10, 20, 50, 100, 150, 200, 250, 300]
 FRICTION_COEFFICIENT = 1/unit.picosecond
 TIMESTEP = 0.002*unit.picoseconds
@@ -102,6 +101,7 @@
 
 # force group values
 ghostghost_group = 29
+ghostghost_nb_group = 30
 systemghost_group = 31
 
 # minimization values
@@ -142,7 +142,7 @@
     gst_idxs = list(range(n_system, n_system+n_ghosts))
 
     con_force = ContForce()
-    con_force.addBond(gst_idxs, len(gst_idxs), 0.18, 10000)
+    con_force.addBond(gst_idxs, len(gst_idxs), 0.25, 10000)
 
     if len(SYSTEM_CONT_FORCE_IDXS) > 0:
         cont_force_idxs = gst_idxs + SYSTEM_CONT_FORCE_IDXS
@@ -151,14 +151,18 @@
     BUILD_UTILS = SystemBuild(psf=psf, crd=crd, pdb=pdb_file, n_ghosts=n_ghosts,
                               toppar_str=TOPPAR_STR, inputs_path=INPUTS_PATH,
                               width=WIDTH, binding_site_idxs=bs_idxs,
-                              min_dist=MIN_DIST, ep_convert=CONVERT_FAC,
-                              gg_group=ghostghost_group, sg_group=systemghost_group,
+                              min_dist=MIN_DIST, 
+                              gg_group=ghostghost_group, gg_nb_group=ghostghost_nb_group, sg_group=systemghost_group,
                               ghost_mass=GHOST_MASS, attr_bounds=BOUNDS,
                               contForce=con_force)
 
     print('Building the system..')
     system, initial_signals, n_ghosts, psf_top, crd_pos, _ = BUILD_UTILS.build_system_forces()
 
+    pressure = 1.0*unit.atmospheres
+    temp = 300*unit.kelvin
+    barostat = omm.MonteCarloBarostat(pressure, temp)
+    system.addForce(barostat)
 
     print('System built')
 
@@ -171,7 +175,7 @@
 
     integrator = CustomHybridIntegratorRestrictedChargeVariance(n_ghosts, TEMPERATURES[0], FRICTION_COEFFICIENT,
                                                 TIMESTEP, attr_fric_coeffs=coeffs, attr_bounds=BOUNDS)
-                                                   
+
     simulation = omma.Simulation(psf_top, system, integrator, platform, prop)
 
     simulation.context.setPositions(crd_pos)
@@ -181,7 +185,7 @@
         os.makedirs(OUTPUTS_PATH)
 
     pre_min_positions = simulation.context.getState(getPositions=True).getPositions()
-    omma.PDBFile.writeFile(psf_top, pre_min_positions, open(osp.join(OUTPUTS_PATH,'struct_before_min.pdb'), 'w'))
+    omma.PDBFile.writeFile(psf.topology, pre_min_positions, open(osp.join(OUTPUTS_PATH,'struct_before_min.pdb'), 'w'))
 
     #_______________MINIMIZE_______________#
     print('Running minimization')
@@ -193,6 +197,7 @@
 
     # save a PDB of the minimized positions
     latest_state = simulation.context.getState(getPositions=True)
+    latest_par = getParameters(simulation, n_ghosts)
 
     omma.PDBFile.writeFile(psf_top, latest_state.getPositions(), open(osp.join(OUTPUTS_PATH,f'minimized_pos.pdb'),'w'))
 
@@ -204,27 +209,26 @@
 
     for temp_idx, TEMP in enumerate(TEMPERATURES):
         H5REPORTER_FILE = osp.join(OUTPUTS_PATH,f'traj{temp_idx}.h5')
-        # integrator = CustomHybridIntegratorConstCharge(n_ghosts, TEMP*unit.kelvin, FRICTION_COEFFICIENT,
-        #                                                TIMESTEP, attr_fric_coeffs=coeffs, attr_bounds=BOUNDS)
 
         integrator = CustomHybridIntegratorRestrictedChargeVariance(n_ghosts, TEMP*unit.kelvin, FRICTION_COEFFICIENT,
                                                        TIMESTEP, attr_fric_coeffs=coeffs, attr_bounds=BOUNDS)
 
 
         # _______________HEAT SYSTEM_______________ #
 
-
-
         simulation = omma.Simulation(psf_top, system, integrator, platform, prop)
 
         simulation.context.setState(latest_state)
-
+        setParameters(simulation, latest_par)
+
         simulation.reporters.append(H5Reporter(H5REPORTER_FILE,
                                                reportInterval=REPORT_STEPS,
                                                coordinates=False,
+                                               velocities=False,
                                                assignments=False,
+                                               global_variables=True,
+                                               global_variable_forces=False,
                                                groups=systemghost_group, num_ghosts=n_ghosts))
-
 
         simulation.reporters.append(mdj.reporters.DCDReporter(osp.join(OUTPUTS_PATH,
                                                                         f'heating{temp_idx}.dcd'),
@@ -237,7 +241,8 @@
                                                            temperature=True))
 
         simulation.step(NUM_STEPS[temp_idx])
-
+
+        latest_par = getParameters(simulation, n_ghosts)    
         latest_state = simulation.context.getState(getPositions=True)
 
     end = time.time()
@@ -265,7 +270,7 @@
         pkl.dump(final_pos, new_file)
     print(" n_ghosts",  n_ghosts)
 
-    par = BUILD_UTILS.getParameters(simulation, n_ghosts)
+    par = getParameters(simulation, n_ghosts)
     print("parameters: ", par)
 
     with open(osp.join(OUTPUTS_PATH, 'parameters.pkl'), 'wb') as new_file: