Created Boresch restraint integrator within RandomLigandRotation

Author: Steven Ayoub

blues/integrators.py

@@ -310,140 +310,120 @@ def __init__(self,
                  nsteps_neq=0,
                  nprop=1,
                  prop_lambda=0.3,
-                 lambda_restraints='max(0, 1 - 2*abs(lambda - 0.5))',  # tent fn (0→1 at 0.5→0)
+                 lambda_restraints = 'max(0, 1-(1/0.10)*abs(lambda-0.5))',
+                #relax_steps=500, #'max(0, 1-(1/0.10)*abs(lambda-0.5))', #"3*lambda^2 - 2*lambda^3", # old: 'max(0, 1-(1/0.10)*abs(lambda-0.5))'
                  relax_steps=50,
                  *args, **kwargs):
-
-        # Store the expression string separately to avoid name collisions
-        self.lambda_restraints_expr = lambda_restraints
-
-        # Which energy group holds the restraint energy (so we can subtract it cleanly)
-        self.restraint_energy_group = str(restraint_group)
-        self.restraint_energy_ref = f"energy{self.restraint_energy_group}"
-
-        super().__init__(
-            alchemical_functions=alchemical_functions,
-            splitting=splitting,
-            temperature=temperature,
-            collision_rate=collision_rate,
-            timestep=timestep,
-            constraint_tolerance=constraint_tolerance,
-            measure_shadow_work=measure_shadow_work,
-            measure_heat=measure_heat,
-            nsteps_neq=nsteps_neq,
-            nprop=nprop,
-            prop_lambda=prop_lambda,
-            *args, **kwargs
-        )
-
+        
+        self.lambda_restraints = lambda_restraints
+        self.restraint_energy = "energy"+str(restraint_group)
+
+        super(AlchemicalExternalRestrainedLangevinIntegrator, self).__init__(
+                     alchemical_functions,
+                     splitting,
+                     temperature,
+                     collision_rate,
+                     timestep,
+                     constraint_tolerance,
+                     measure_shadow_work,
+                     measure_heat,
+                     nsteps_neq,
+                     nprop,
+                     prop_lambda,
+                     *args, **kwargs)
+        
+        try:
+            self.addGlobalVariable("restraint_energy", 0)
+        except:
+            pass
         logger.info(f'[LAMBDA STEPS] N_lambda_steps: {self._n_lambda_steps}')
-
-        # === Declare NEW integrator globals (names MUST match force globals if you want propagation) ===
+        # Only declare NEW variables
         self.addGlobalVariable("debug_lambda", 0.0)
-        self.addGlobalVariable("restraint_energy", 0.0)      # running value read via energy group
+        #self.addGlobalVariable("restraint_energy", 0.0)
 
+        # Set existing globals from parent
+        # self.setGlobalVariableByName("lambda_step", 0.0)
+        # self.setGlobalVariableByName("lambda", 0.0)
 
         # Optional debug
         self.addComputeGlobal("debug_lambda", "lambda")
 
-        logger.info(f"Current nsteps_neq: {nsteps_neq}")
-        logger.info(f'lambda_restraints expression: {self.lambda_restraints_expr}')
-
+        # Now safe to use lambda_restraints in update
+        #self.updateRestraints()
 
+        logger.info(f"Current nsteps_neq: {nsteps_neq}")
+        logger.info(f'lambda_restraints selected: {self.lambda_restraints}')
+        # compute the mid‐point slice index once
+        mid = int(self._n_lambda_steps/2)
+        self.addGlobalVariable("mid_step", float(mid))
+        self.addGlobalVariable("relax_counter", 0.0)
+        self.addGlobalVariable("relax_steps", float(relax_steps))
 
-    # ---- Helper: recompute + propagate lambda_restraints each time lambda changes ----
-    def _recompute_and_push_lambda_restraints(self):
-        # 1) Recompute integrator global from current 'lambda'
-        self.addComputeGlobal('lambda_restraints', self.lambda_restraints_expr)
-        # 2) Propagate integrator globals -> Context parameters (names must match in forces)
-        self.addUpdateContextState()
+    def updateRestraints(self):
+        logger.info(f"UPDATE RESTAINTS: {self.lambda_restraints}")
+        self.addComputeGlobal('lambda_restraints', self.lambda_restraints)
 
-    # ---- Helper: refresh restraint energy and split PE bookkeeping ----
-    def _refresh_restraint_and_pes(self):
-        self.addComputeGlobal("restraint_energy", self.restraint_energy_ref)
-        # Exclude restraint energy so protocol work uses only alchemical parts
-        self.addComputeGlobal("perturbed_pe",   "energy - restraint_energy")
-        self.addComputeGlobal("unperturbed_pe", "energy - restraint_energy")
 
     def _add_integrator_steps(self):
         """
-        Override to insert (a) consistent initialization and (b) frequent
-        lambda_restraints recompute + propagation around each propagation call.
+        Override the base class to insert reset steps around the integrator.
         """
-
-        # =========================
-        # Initialization (step == 0)
-        # =========================
+        
+        # First step: Constrain positions and velocities and reset work accumulators and alchemical integrators
+        logger.info("sams protocol")
         self.beginIfBlock('step = 0')
-
-        # Make sure restraints start consistent with lambda=0
-        self._recompute_and_push_lambda_restraints()
-        self._refresh_restraint_and_pes()
-
-        # Usual initialization: constrain and reset accumulators
+        self.addComputeGlobal("restraint_energy", self.restraint_energy)
+        self.addComputeGlobal("perturbed_pe", "energy - restraint_energy")
+        self.addComputeGlobal("unperturbed_pe", "energy - restraint_energy")
         self.addConstrainPositions()
         self.addConstrainVelocities()
         self._add_reset_protocol_work_step()
         self._add_alchemical_reset_step()
+        self.endBlock()
 
-        self.endBlock()  # end init
-
-        # =========================
         # Main body
-        # =========================
+
         if self._n_steps_neq == 0:
-            # Force a single execution on first step
+            # If nsteps = 0, we need to force execution on the first step only.
             self.beginIfBlock('step = 0')
-            # Track lambda_restraints before the parent propagation
-            self._recompute_and_push_lambda_restraints()
-            self._refresh_restraint_and_pes()
-            super()._add_integrator_steps()   # call parent’s propagation
+            super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
             self.addComputeGlobal("step", "step + 1")
             self.endBlock()
-            return
-
-        # NEQ loop
-        self.beginIfBlock("step < n_lambda_steps")
-
-        # --- First substep at this NEQ step ---
-        # Update lambda_restraints to track the *new* lambda for this substep
-        self._recompute_and_push_lambda_restraints()
-        self._refresh_restraint_and_pes()
-
-        # First-step bookkeeping (only once)
-        self.beginIfBlock("first_step < 1")
-        self.addComputeGlobal("first_step", "1")
-        self._refresh_restraint_and_pes()
-        self.endBlock()
-
-        # Protocol work increment (before propagation) using split PEs
-        self.addComputeGlobal("protocol_work", "protocol_work + (perturbed_pe - unperturbed_pe)")
-
-        # === Primary propagation for this lambda value ===
-        self._recompute_and_push_lambda_restraints()
-        super()._add_integrator_steps()   # parent does the actual R/V/O/H... steps
-
-        # ============ Optional extra propagation windows at same lambda ============
-        self.beginIfBlock("lambda > prop_lambda_min")
-        self.beginIfBlock("lambda <= prop_lambda_max")
-
-        self.beginWhileBlock("prop < nprop")
-        self.addComputeGlobal("prop", "prop + 1")
-
-        self._recompute_and_push_lambda_restraints()
-        super()._add_integrator_steps()
-
-        self.endBlock()  # while
-        self.endBlock()  # if <= prop_lambda_max
-        self.endBlock()  # if > prop_lambda_min
+        else:
+            #call the superclass function to insert the appropriate steps, provided the step number is less than n_steps
+            self.beginIfBlock("step < n_lambda_steps")#
+            self.addComputeGlobal("restraint_energy", self.restraint_energy)
+            self.addComputeGlobal("perturbed_pe", "energy - restraint_energy")
+            self.beginIfBlock("first_step < 1")##
+            #TODO write better test that checks that the initial work isn't gigantic
+            self.addComputeGlobal("first_step", "1")
+            self.addComputeGlobal("restraint_energy", self.restraint_energy)
+            self.addComputeGlobal("unperturbed_pe", "energy-restraint_energy")
+            self.endBlock()##
+            #initial iteration
+            # Gill put this first: 
+            self.addComputeGlobal("protocol_work", 
+                                  "protocol_work + (perturbed_pe - unperturbed_pe)"
+            )
+            super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
+            logger.info("COMPUTE WORKSSSS")
+            #if more propogation steps are requested
+            self.beginIfBlock("lambda > prop_lambda_min")###
+            self.beginIfBlock("lambda <= prop_lambda_max")####
 
-        # ---- End-of-step bookkeeping ----
-        # Recompute to ensure end-of-step energies are coherent
-        self._recompute_and_push_lambda_restraints()
-        self._refresh_restraint_and_pes()
+            self.beginWhileBlock("prop < nprop")#####
+            self.addComputeGlobal("prop", "prop + 1")
 
-        self.addComputeGlobal("step", "step + 1")
-        self.addComputeGlobal("prop", "1")
+            super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
+            self.endBlock()#####
+            self.endBlock()####
+            self.endBlock()###
+            #ending variables to reset
+            self.updateRestraints()
+            self.addComputeGlobal("restraint_energy", self.restraint_energy)
+            self.addComputeGlobal("unperturbed_pe", "energy-restraint_energy")
+            self.addComputeGlobal("step", "step + 1")
+            self.addComputeGlobal("prop", "1")
 
-        self.endBlock()  # if step < n_lambda_steps
+            self.endBlock()#
 

blues/moves.py

@@ -483,7 +483,7 @@ def beforeMove(self, context):
         if not self.restraints:
             return context
         for i in range(len(self.binding_mode_traj)):
-            context.setParameter(f'restraint_pose_{i}', 0.0)
+            context.setParameter(f'restraint_pose_{i}', 1.0)
 
         return context
 

blues/simulation.py

@@ -1043,7 +1043,19 @@ def _syncStatesMDtoNCMC(self):
 
     def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
         """Advance the NCMC simulation."""
+        #print("Running _stepNCMC...")
+        #print(f"nsteps: {nstepsNC}, moveStep: {moveStep}")
 
+        self.worker_tracker[self.currentIter] = []
+        self.lambda_tracker[self.currentIter] = []
+        self.steric_tracker[self.currentIter] = []
+        self.restraint_tracker[self.currentIter] = []
+
+        logger.info('Advancing %i NCMC switching steps...' % (nstepsNC))
+
+
+        # Retrieve NCMC state before proposed move
+        logger.info('running getStateFromContext')
         ncmc_state0 = self.getStateFromContext(self._ncmc_sim.context, self._state_keys)
         #print(ncmc_state0['positions'])
         self._setStateTable('ncmc', 'state0', ncmc_state0)
@@ -1062,12 +1074,8 @@ def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
             try:
                 if not step:
                     logger.info("Calling beforeMove()")
-                    self._ncmc_sim.context = move_engine.selected_move.beforeMove(self._ncmc_sim.context) 
-                    try:            
-                        lambda_rest = self._ncmc_sim.context.getParameter("lambda_restraints")
-                        logger.info(f"[Step {step}] lambda_restraints = {lambda_rest}")               
-                    except Exception as e:
-                        logger.warning(f"[Step {step}] Could not retrieve lambda_restraints: {e}")
+                    self._ncmc_sim.context = move_engine.selected_move.beforeMove(self._ncmc_sim.context)             
+
                 if step == moveStep:
                     if hasattr(logger, 'report'):
                         logger.info = logger.report
@@ -1082,9 +1090,7 @@ def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
                     # Perform the NCMC move (lambda 0 → 0.5 and apply move)
                     self._ncmc_sim.context = move_engine.runEngine(self._ncmc_sim.context)
 
-                # state = self._ncmc_sim.context.getState(getPositions=True, getEnergy=True)
-                # steric_state = self._ncmc_sim.context.getState(getEnergy=True, groups={move_engine.selected_move.steric_group})
-                # steric_energy = steric_state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
+
                 lambda_val = self._ncmc_sim.context._integrator.getGlobalVariableByName("lambda")
                 if abs(lambda_val - 0.0) < 1e-4 or abs(lambda_val - 0.5) < 1e-4 or abs(lambda_val - 0.9) < 1e-4:
                     steric_state = self._ncmc_sim.context.getState(getEnergy=True, groups={move_engine.selected_move.steric_group})
@@ -1095,17 +1101,19 @@ def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
                     lambda_s = self._ncmc_sim.context.getParameter("lambda_sterics")
                     lambda_e = self._ncmc_sim.context.getParameter("lambda_electrostatics")
                     logger.info(f"[Step {step}] λ_sterics = {lambda_s}, λ_electrostatics = {lambda_e}")
-                    try:
-                        lambda_rest = self._ncmc_sim.context.getParameter("lambda_restraints")
-                        logger.info(f"[Step {step}] lambda_restraints = {lambda_rest}")               
-                    except Exception as e:
-                        logger.warning(f"[Step {step}] Could not retrieve lambda_restraints: {e}")
+                    self._log_restraint_energies(step=step, move_engine = move_engine)
+               
                 self._ncmc_sim.step(1)
 
                 integrator = self._ncmc_sim.context._integrator
                 lambda_val = integrator.getGlobalVariableByName("lambda")
+                protocol_work = integrator.getGlobalVariableByName("protocol_work")
                 #logger.info(f"[NCMC step {step}] lambda = {lambda_val:.6f} | protocol_work = {protocol_work:.4f}")
-                #self._log_restraint_energies(step=step, move_engine = move_engine, print_output=False)
+                self.worker_tracker[self.currentIter].append(protocol_work)
+                self.lambda_tracker[self.currentIter].append(lambda_val)
+                steric_state = self._ncmc_sim.context.getState(getEnergy=True, groups={move_engine.selected_move.steric_group})
+                steric_energy = steric_state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
+                self.steric_tracker[self.currentIter].append(steric_energy)
 
                 if step == lastStep:
                     logger.info("AFTER MOVE WILL BE CALLED")
@@ -1116,20 +1124,14 @@ def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
                     lambda_s = self._ncmc_sim.context.getParameter("lambda_sterics")
                     lambda_e = self._ncmc_sim.context.getParameter("lambda_electrostatics")
                     logger.info(f"[Step {step}] λ_sterics = {lambda_s}, λ_electrostatics = {lambda_e}")
-                    #self._log_restraint_energies(step=step, move_engine = move_engine)
+                    self._log_restraint_energies(step=step, move_engine = move_engine)
                     # Log sterics after move
                     try:
                         steric_state = self._ncmc_sim.context.getState(getEnergy=True, groups={move_engine.selected_move.steric_group})
                         steric_energy = steric_state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
                         logger.info(f"[Step {step}] Steric Energy AFTER MOVE (group {move_engine.selected_move.steric_group}): {steric_energy:.4f} kJ/mol")
                     except Exception as e:
                         logger.warning(f"[Step {step}] Could not retrieve steric energy after move: {e}")
-
-                    try:
-                        lambda_rest = self._ncmc_sim.context.getParameter("lambda_restraints")
-                        logger.info(f"[Step {step}] lambda_restraints = {lambda_rest}")               
-                    except Exception as e:
-                        logger.warning(f"[Step {step}] Could not retrieve lambda_restraints: {e}")
 
                     self._ncmc_sim.context = move_engine.selected_move.afterMove(self._ncmc_sim.context)
                     # Debug: print positions after afterMove                    
@@ -1145,6 +1147,26 @@ def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
         ncmc_state1 = self.getStateFromContext(self._ncmc_sim.context, self._state_keys)
         self._setStateTable('ncmc', 'state1', ncmc_state1)
 
+
+    def _log_restraint_energies(self, step, move_engine, print_output=True):
+
+        try:
+            for name, group in move_engine.selected_move.restraint_groups.items():
+                state = self._ncmc_sim.context.getState(getEnergy=True, groups={group})
+                energy = state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
+                if print_output:
+                    logger.info(f"[Step {step}] {name} restraint energy (group {group}): {energy:.4f} kJ/mol")
+                else:
+                    self.restraint_tracker[self.currentIter].append(energy)
+            
+            lambda_rest = self._ncmc_sim.context.getParameter("lambda_restraints")
+            logger.info(f"[Step {step}] lambda_restraints = {lambda_rest}")
+        
+        except Exception as e:
+            pass
+            #  logger.warning(f"[Step {step}] Could not retrieve lambda_restraints: {e}")
+
+        
     # def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
     #     """Advance the NCMC simulation."""
     #     #print("Running _stepNCMC...")
@@ -1355,6 +1377,11 @@ def run(self, nIter=0, nstepsNC=0, moveStep=0, nstepsMD=0, temperature=300, writ
         if not nstepsMD: nstepsMD = self._config['nstepsMD']
         if not moveStep: moveStep = self._config['moveStep']
 
+        self.worker_tracker = {}
+        self.lambda_tracker = {}
+        self.steric_tracker = {}
+        self.restraint_tracker = {}
+
         logger.info('Running %i BLUES iterations...' % (nIter))
         for N in range(int(nIter)):
             self.currentIter = N