Merge pull request #6 from ADicksonLab/attr_force

Merge attr_force branch

examples/brd_ghost_distribution/build_minimize_heat_noMLF.py

@@ -18,6 +18,8 @@
 
 from flexibletopology.utils.reporters import H5Reporter
 from flexibletopology.utils.openmmutils import read_params, getParameters, setParameters
+from flexibletopology.forces.static import build_protein_restraint_force
+from flexibletopology.attr_minimizer import AttrMinimizer
 import sys
 
 from contforceplugin import ContForce
@@ -85,7 +87,7 @@
 pos_arr = np.array(crd.positions.value_in_unit(unit.nanometers))
 
 # MD simulations settings
-TEMPERATURES = [10, 20, 50, 100, 150, 200, 250, 300]
+TEMPERATURES = [10, 10, 20, 50, 100, 150, 200, 250, 300]
 FRICTION_COEFFICIENT = 1/unit.picosecond
 TIMESTEP = 0.002*unit.picoseconds
 NUM_STEPS = [10000 for t in TEMPERATURES]
@@ -97,7 +99,7 @@
 # system building values
 WIDTH = 0.3 # nm
 # MIN_DIST = 0.04 # nm
-MIN_DIST = 0.05 # nm
+MIN_DIST = 0.2 # nm
 
 # force group values
 ghostghost_group = 29
@@ -136,8 +138,10 @@
     #_______________BUILD SYSTEM & SIMULATION OBJECT_______________#
 
     bs_idxs = pdb_file.topology.select(BS_SELECTION_STRING)
+    prot_idxs = pdb_file.topology.select('protein')
+    bb_idxs = pdb_file.topology.select('protein and backbone')
 
-    print("n_ghosts is ",n_ghosts)
+    print(f"Building a system with {n_ghosts} ghost particles..")
     n_system = pdb_file.n_atoms
     gst_idxs = list(range(n_system, n_system+n_ghosts))
 
@@ -148,23 +152,19 @@
         cont_force_idxs = gst_idxs + SYSTEM_CONT_FORCE_IDXS
         con_force.addBond(cont_force_idxs, len(cont_force_idxs), 0.25, 10000)
 
-    BUILD_UTILS = SystemBuild(psf=psf, crd=crd, pdb=pdb_file, n_ghosts=n_ghosts,
+    BUILD_UTILS = SystemBuild(psf=psf, pos=pos_arr, 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, 
+                              min_dist=MIN_DIST, gg_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()
+    system, initial_signals, n_ghosts, psf_top, 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')
+    barostat = omm.MonteCarloBarostat(pressure, temp)     # add this to the system after heating step 0
 
     coeffs = {'lambda': rest_coeff,
               'charge': rest_coeff,
@@ -178,7 +178,7 @@
 
     simulation = omma.Simulation(psf_top, system, integrator, platform, prop)
 
-    simulation.context.setPositions(crd_pos)
+    simulation.context.setPositions(pos)
 
     # add reporters                                                                     
     if not osp.exists(OUTPUTS_PATH):
@@ -188,34 +188,49 @@
     omma.PDBFile.writeFile(psf.topology, pre_min_positions, open(osp.join(OUTPUTS_PATH,'struct_before_min.pdb'), 'w'))
 
     #_______________MINIMIZE_______________#
-    print('Running minimization')
+    print('Running minimization..')
     print('Before min: E=', simulation.context.getState(getEnergy=True).getPotentialEnergy())
+
+    attr_min = AttrMinimizer(simulation, n_ghosts, BOUNDS)
+    print('old attr:',attr_min.get_attributes())
+    print('new_attr:',attr_min.attr_minimize())
+    print('After attr min: E=', simulation.context.getState(getEnergy=True).getPotentialEnergy())
+
     begin = time.time()
     simulation.minimizeEnergy(tolerance=TOL ,maxIterations=MAXITR)
     end = time.time()
     print('After min: E=', simulation.context.getState(getEnergy=True).getPotentialEnergy())
 
     # save a PDB of the minimized positions
-    latest_state = simulation.context.getState(getPositions=True)
+    latest_state = simulation.context.getState(getPositions=True,enforcePeriodicBox=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'))
 
-    print("Minimization Ends")
     print(f"Minimization run time = {np.round(end - begin, 3)}s")
 
-    print('Heating system')
-    begin = time.time()
-
+    print('Heating system..')
+    heat_begin = time.time()    
     for temp_idx, TEMP in enumerate(TEMPERATURES):
+        heat_step_begin = time.time()    
         H5REPORTER_FILE = osp.join(OUTPUTS_PATH,f'traj{temp_idx}.h5')
 
         integrator = CustomHybridIntegratorRestrictedChargeVariance(n_ghosts, TEMP*unit.kelvin, FRICTION_COEFFICIENT,
                                                        TIMESTEP, attr_fric_coeffs=coeffs, attr_bounds=BOUNDS)
 
 
         # _______________HEAT SYSTEM_______________ #
-
+
+        if temp_idx == 0: # restrain all system atoms
+            box_sizes = np.zeros((3))
+            box_sizes[0] = system.getDefaultPeriodicBoxVectors()[0][0].value_in_unit(unit.nanometers)
+            box_sizes[1] = system.getDefaultPeriodicBoxVectors()[1][1].value_in_unit(unit.nanometers)
+            box_sizes[2] = system.getDefaultPeriodicBoxVectors()[2][2].value_in_unit(unit.nanometers)
+
+            positions = simulation.context.getState(getPositions=True,enforcePeriodicBox=True).getPositions()
+            prot_restr_force = build_protein_restraint_force(positions,prot_idxs,bb_idxs,box_sizes)
+            prot_restr_force_idx = system.addForce(prot_restr_force)
+
         simulation = omma.Simulation(psf_top, system, integrator, platform, prop)
 
         simulation.context.setState(latest_state)
@@ -245,18 +260,21 @@
         latest_par = getParameters(simulation, n_ghosts)    
         latest_state = simulation.context.getState(getPositions=True)
 
-    end = time.time()
-    print(f"Heating run time = {np.round(end - begin, 3)}s")
-    print('Done heating system; dcd saved')
-
+        if temp_idx == 0:
+            system.removeForce(prot_restr_force_idx)
+            system.addForce(barostat)
 
+        heat_step_end = time.time()
+        print(f"Heating step {temp_idx} completed. Run time = {np.round(heat_step_end - heat_step_begin, 3)}s")
+
+    heat_end = time.time()
+    print(f"Total heating run time = {np.round(heat_end - heat_begin, 3)}s")
 
 
     #_______________SAVE SIM INPUTS_______________#
 
     # save the system, topology, simulation object, positions, and parameters
-    print(" n_ghosts",  n_ghosts)
-    print('Saving simulation input files')
+    print('Saving simulation input files..')
 
     with open(osp.join(OUTPUTS_PATH, 'system.pkl'), 'wb') as new_file:
         pkl.dump(system, new_file)

examples/brd_ghost_distribution/system_build_util.py

@@ -22,13 +22,13 @@ class SystemBuild(object):
     A class that generates a ready-to-equilibrate OpenMM system object.
     """
 
-    def __init__(self, psf=None, crd=None, pdb=None, target_pkl=None, n_ghosts=None, toppar_str=None, inputs_path=None,
+    def __init__(self, psf=None, pos=None, pdb=None, target_pkl=None, n_ghosts=None, toppar_str=None, inputs_path=None,
                  ani_model=None, width=None, binding_site_idxs=None, min_dist=0.15, gg_min_dist=0.05,
                  sf_weights=None, gg_group=None, gg_nb_group=None, mlforce_group=None, sg_group=None, mlforce_scale=None,
                  ghost_mass=None,attr_bounds=None,assignFreq=None,rmax_delta=None, rest_k=None, contForce=None):
 
         self.psf = psf
-        self.crd = crd
+        self.pos = pos
         self.pdb = pdb
         self.target_pkl = target_pkl
         self.toppar_str = toppar_str
@@ -151,21 +151,19 @@ def build_system_forces(self):
         init_attr = gen_init_attr(self.n_ghosts, self.attr_bounds,total_charge=0,init_lambda=1.0)
 
         com_bs = self.generate_COM(self.binding_site_idxs, self.pdb)
-        pos_arr = np.array(self.crd.positions.value_in_unit(unit.nanometers))
-        pdb_pos = np.array([pos_arr])
-        init_positions = gen_init_pos(self.n_ghosts, com_bs, self.width, pdb_pos, self.min_dist, self.gg_min_dist)
+        init_positions = gen_init_pos(self.n_ghosts, com_bs, self.width, self.pos, self.min_dist, self.gg_min_dist)
 
         # calculating box length
-        box_lengths = pos_arr.max(axis=0) - pos_arr.min(axis=0)
+        box_lengths = self.pos.max(axis=0) - self.pos.min(axis=0)
         self.psf.setBox(box_lengths[0] * unit.nanometers,
                    box_lengths[1] * unit.nanometers,
                    box_lengths[2] * unit.nanometers)
 
         params = read_params(self.toppar_str, self.inputs_path)
 
-        n_part_system = len(self.crd.positions)
-        self.crd.positions.extend(unit.quantity.Quantity(init_positions,
-                                            unit.nanometers))
+        n_part_system = len(self.pos)
+        self.pos = np.append(self.pos, init_positions, axis=0)
+
         system = self.psf.createSystem(params,
                                   nonbondedMethod=omma.forcefield.CutoffPeriodic,
                                   nonbondedCutoff=1*unit.nanometers,
@@ -201,9 +199,10 @@ def build_system_forces(self):
                                  n_part_system=n_part_system,
                                  group_num=self.gg_nb_group,
                                  initial_sigmas=init_attr['sigma'],
+                                 initial_charges=init_attr['charge'],
                                  nb_exclusion_list=exclusion_list)
 
-        return system, init_attr, self.n_ghosts, new_psf.topology, self.crd.positions, target_features
+        return system, init_attr, self.n_ghosts, new_psf.topology, self.pos, target_features
 
 
 

src/flexibletopology/attr_minimizer.py

@@ -0,0 +1,102 @@
+from scipy.optimize import Bounds, LinearConstraint, NonlinearConstraint, minimize
+from openmm import unit
+import numpy as np
+
+class AttrMinimizer(object):
+    """
+    A class to minimize ghost particle simulations according to their attributes,
+    which are saved within an OpenMM context as global variables.
+
+    This is valuable as the standard simulation.minimize function does not
+    minimize the attributes, only the positions.
+
+    Internally, the dependent variables are represented by x, which stores the 
+    attributes in the following order:
+
+    x[0] - charge[0]
+    x[1] - sigma[0]
+    x[2] - epsilon[0]
+    x[3] - lambda[0]
+    x[4] - charge[1]
+    ...
+
+    """
+    def __init__(self, simulation, n_ghosts, bounds, charge_sum=0.0, max_charge_var=0.08):
+        self.simulation = simulation
+
+        self.n_ghosts = n_ghosts
+
+        lower_bounds = [bounds['charge'][0],bounds['sigma'][0],bounds['epsilon'][0],bounds['lambda'][0]] * n_ghosts
+        upper_bounds = [bounds['charge'][1],bounds['sigma'][1],bounds['epsilon'][1],bounds['lambda'][1]] * n_ghosts
+        self.bounds = Bounds(lower_bounds,upper_bounds)
+
+        self.constraints = []
+        if charge_sum is not None:
+            sum_charge_vector = [1,0,0,0] * n_ghosts
+            self.constraints.append(LinearConstraint([sum_charge_vector],[charge_sum],[charge_sum]))
+
+        if max_charge_var is not None:
+            def cons_f(x):
+                sum_sq = 0.0
+                for i in range(n_ghosts):
+                    sum_sq += x[4*i]*x[4*i]
+                sum_sq /= n_ghosts
+                return [sum_sq]
+            def cons_J(x):
+                jvec = []
+                for i in range(n_ghosts):
+                    for j in range(4):
+                        if j == 0:
+                            jvec.append(2*x[4*i]/n_ghosts)
+                        else:
+                            jvec.append(0)
+                return [jvec]
+
+            self.constraints.append(NonlinearConstraint(cons_f, 0, max_charge_var*max_charge_var, jac=cons_J, hess='2-point'))
+
+    def get_attributes(self):
+        pars = self.simulation.context.getParameters()
+
+        x = np.zeros((4*self.n_ghosts))
+        for i in range(self.n_ghosts):
+            x[4*i] = pars[f'charge_g{i}']
+            x[4*i + 1] = pars[f'sigma_g{i}']
+            x[4*i + 2] = pars[f'epsilon_g{i}']
+            x[4*i + 3] = pars[f'lambda_g{i}']
+
+        return x
+
+    def set_attributes(self, x):
+        for i in range(self.n_ghosts):
+            self.simulation.context.setParameter(f'charge_g{i}',x[4*i])
+            self.simulation.context.setParameter(f'sigma_g{i}',x[4*i + 1])
+            self.simulation.context.setParameter(f'epsilon_g{i}',x[4*i + 2])
+            self.simulation.context.setParameter(f'lambda_g{i}',x[4*i + 3])
+
+        return
+
+    def _energy(self, x):
+        self.set_attributes(x)
+        return self.simulation.context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
+
+    def _energy_derivs(self, x):
+        self.set_attributes(x)
+        par_derivs = self.simulation.context.getState(getParameterDerivatives=True).getEnergyParameterDerivatives()
+
+        v = np.zeros((4*self.n_ghosts))
+        for i in range(self.n_ghosts):
+            v[4*i] = par_derivs[f'charge_g{i}']
+            v[4*i + 1] = par_derivs[f'sigma_g{i}']
+            v[4*i + 2] = par_derivs[f'epsilon_g{i}']
+            v[4*i + 3] = par_derivs[f'lambda_g{i}']
+
+        return v
+
+    def attr_minimize(self,method='trust-constr', options={'verbose':1}):
+        res = minimize(self._energy, self.get_attributes(), method=method, jac=self._energy_derivs, hess='2-point',
+                       constraints=self.constraints, options=options, bounds=self.bounds)
+
+        self.set_attributes(res.x)
+
+        return res.x
+

src/flexibletopology/forces/attributes.py

@@ -0,0 +1,44 @@
+import openmm as omm
+
+def add_attr_force(system,
+                   n_ghosts=None,
+                   group_num=None,
+                   initial_attr=None):
+
+    eps_terms = ''
+    charge_terms = ''
+    fsig_eqns = ''
+    for gh_idx in range(n_ghosts):
+        eps_terms += f'A*fsig_{gh_idx}*(epsilon_g{gh_idx}-0.25)^2' 
+        charge_terms += f'B*fsig_{gh_idx}*(charge_g{gh_idx}-0.25)^2'
+        fsig_eqns += f'fsig_{gh_idx} = 1/(1+exp(100*(sigma_g{gh_idx}-0.3))); '
+        if gh_idx < n_ghosts-1:
+            eps_terms += ' + '
+            charge_terms += ' + '
+    energy_function = eps_terms + ' + ' + charge_terms + '; ' + fsig_eqns
+
+    attr_force = omm.CustomCVForce(energy_function)
+
+    attr_force.addGlobalParameter('A', 941.5)
+    attr_force.addGlobalParameter('B', 275.3)
+
+    for gh_idx in range(n_ghosts):
+        # set to initial values
+        attr_force.addGlobalParameter(f'charge_g{gh_idx}', initial_attr['charge'][gh_idx])
+        attr_force.addGlobalParameter(f'sigma_g{gh_idx}', initial_attr['sigma'][gh_idx])
+        attr_force.addGlobalParameter(f'epsilon_g{gh_idx}', initial_attr['epsilon'][gh_idx])
+        attr_force.addGlobalParameter(f'lambda_g{gh_idx}', initial_attr['lambda'][gh_idx])
+
+        # adding the del(signal)s [needed in the integrator]
+        attr_force.addEnergyParameterDerivative(f'charge_g{gh_idx}')
+        attr_force.addEnergyParameterDerivative(f'sigma_g{gh_idx}')
+        attr_force.addEnergyParameterDerivative(f'epsilon_g{gh_idx}')
+        attr_force.addEnergyParameterDerivative(f'lambda_g{gh_idx}')
+
+    # set force parameters
+    attr_force.setForceGroup(group_num)
+    system.addForce(attr_force)
+
+    return system
+
+