Created Boresch restraint integrator

Author: Steven Ayoub

blues/moves.py

@@ -195,24 +195,51 @@ class RandomLigandRotationMove(Move):
         'LIG'
     """
 
-    def __init__(self, structure, resname='LIG', ligand_indices=None, random_state=None):
+    def __init__(self, 
+        structure, 
+        resname='LIG', 
+        ligand_indices=None, 
+        random_state=None, ligand_pdb_file=None,
+        restraints=None, K_r=None, K_angle=None,
+        restrained_receptor_atoms=None, restrained_ligand_atoms=None,
+        lambda_restraints = "max(0, 1-(1/0.10)*abs(lambda-0.5))",
+        no_move=False,
+        ):
         self.structure = structure
         self.resname = resname
         self.random_state = random_state
         self.atom_indices = self.getAtomIndices(structure, self.resname)
         self.ligand_indices = ligand_indices
         self.topology = structure[self.atom_indices].topology
         self.totalmass = 0
+        self.restraints = restraints
         self.masses = []
         self.center_of_mass = None
+        self.K_r = K_r
+        self.K_angle = K_angle
+        self.lambda_restraints = lambda_restraints
+        self.restrained_receptor_atoms = restrained_receptor_atoms
+        self.restrained_ligand_atoms = restrained_ligand_atoms
         self.positions = structure[self.atom_indices].positions
+        self.no_move = no_move
         if self.ligand_indices:
             self.topology = structure[self.ligand_indices].topology
             self.positions = structure[self.ligand_indices].positions
-
+        
+        if ligand_pdb_file is not None:
+            self._loadBindingModeTraj(ligand_pdb_file)
         self._calculateProperties()
 
 
+    def _loadBindingModeTraj(self, ligand_pdb_file):
+        """
+        Load the binding mode trajectory from the file.
+        """
+        import mdtraj as md
+        binding_mode_traj = []
+        binding_mode_traj.append(md.load(ligand_pdb_file))
+        self.binding_mode_traj = binding_mode_traj
+    
     def initializeSystem(self, system, integrator, config):
         """
         Changes the system by adding forces corresponding to restraints (if specified)
@@ -299,30 +326,17 @@ def initializeRestraints(self, system: openmm.System, integrator: openmm.Integra
         integrator_kwargs = config or {}
 
         # Get integrator kwargs if available, otherwise use defaults
-        # Create new integrator with restraints
-
-        if self.old_restraint:
-            new_int = AlchemicalExternalRestrainedLangevinIntegrator(
-                restraint_group=set(self.restraint_groups.values()),
-                lambda_restraints=self.lambda_restraints, 
-                alchemical_functions = old_int._alchemical_functions,
-                nsteps_neq=integrator_kwargs['nstepsNC'],
-                nprop=integrator_kwargs['nprop'],
-                prop_lambda=integrator_kwargs['propLambda'],
-                splitting=integrator_kwargs['splitting'],)
-        
-        else:
-            new_int = AlchemicalExternalLangevinIntegrator(
-                restraint_group=set(self.restraint_groups.values()),
-                lambda_restraints=self.lambda_restraints, 
-                alchemical_functions = old_int._alchemical_functions,
-                nsteps_neq=integrator_kwargs['nstepsNC'],
-                nprop=integrator_kwargs['nprop'],
-                prop_lambda=integrator_kwargs['propLambda'],
-                splitting=integrator_kwargs['splitting'])
-                #**old_int.int_kwargs)
-
-        new_int.reset()
+        new_int = AlchemicalExternalLangevinIntegrator(
+            restraint_group=set(self.restraint_groups.values()),
+            lambda_restraints=self.lambda_restraints, 
+            alchemical_functions = old_int._alchemical_functions,
+            nsteps_neq=integrator_kwargs['nstepsNC'],
+            nprop=integrator_kwargs['nprop'],
+            prop_lambda=integrator_kwargs['propLambda'],
+            splitting=integrator_kwargs['splitting'])
+            #**old_int.int_kwargs)
+
+        #new_int.reset()
 
         # Verify we have the required trajectory data
         if not hasattr(self, 'binding_mode_traj') or len(self.binding_mode_traj) == 0:
@@ -350,8 +364,7 @@ def initializeRestraints(self, system: openmm.System, integrator: openmm.Integra
                 new_sys = add_boresch_restraints(sys=new_sys, struct=self.structure, pos=pose_allpos, ligand_atoms=self.atom_indices, 
                                                  pose_num=index, force_group=self.restraint_groups['boresch'],
                                             restrained_receptor_atoms=self.restrained_receptor_atoms, restrained_ligand_atoms=self.restrained_ligand_atoms,
-                                            K_r=self.K_r, K_angle=self.K_angle, K_RMSD=self.K_RMSD, RMSD0=self.RMSD0,
-                                            K_com=self.K_com)
+                                            K_r=self.K_r, K_angle=self.K_angle)
 
             if 'boresch' not in self.restraints and 'rmsd' not in self.restraints:
                 raise ValueError(f'Invalid restraint type: {self.restraints}')
@@ -469,6 +482,8 @@ def move(self, context):
         context: openmm.openmm.Context object
             The same input context, but whose positions were changed by this function.
         """
+        if self.no_move: 
+            return context
         positions = context.getState(getPositions=True).getPositions(asNumpy=True)
 
         self.positions = positions[self.atom_indices]

blues/simulation.py

@@ -805,7 +805,10 @@ def generateSimulationSet(self, config=None):
 
         #Initialize the Move Engine with the Alchemical System and NCMC Integrator
         for move in self._move_engine.moves:
-            self._alch_system, self.ncmc_integrator = move.initializeSystem(self._alch_system, self.ncmc_integrator)
+            try:
+                self._alch_system, self.ncmc_integrator = move.initializeSystem(self._alch_system, self.ncmc_integrator)
+            except TypeError:
+                self._alch_system, self.ncmc_integrator = move.initializeSystem(self._alch_system, self.ncmc_integrator, config=config)
         self.ncmc = self.generateSimFromStruct(self._structure, self._alch_system, self.ncmc_integrator, **config)
         utils.print_host_info(self.ncmc)