Fix double work accumulation in AlchemicalExternalLangevinIntegrator

- Remove redundant work accumulation from _add_integrator_steps()
- Fix lambda_step reset in reset() method to prevent lambda drift
- Add debug variables for work tracking in different phases
- Remove super call from propagation loop to prevent double counting
- Fix freeze_radius usage in example_ncmc.py

Author: stevenayoub

blues/integrators.py

@@ -135,6 +135,9 @@ def __init__(self,
         self.addGlobalVariable("prop", 1)
         self.addGlobalVariable("prop_lambda_min", self._prop_lambda[0])
         self.addGlobalVariable("prop_lambda_max", self._prop_lambda[1])
+        self.addGlobalVariable("debug_work", 0.0)
+        self.addGlobalVariable("work_0_to_05", 0.0)
+        self.addGlobalVariable("work_05_to_1", 0.0)
         # Behavior changed in https://github.com/choderalab/openmmtools/commit/7c2630050631e126d61b67f56e941de429b2d643#diff-5ce4bc8893e544833c827299a5d48b0d
         self._step_dispatch_table['H'] = (self._add_alchemical_perturbation_step, False)
         #$self._registered_step_types['H'] = (
@@ -192,16 +195,18 @@ def _add_integrator_steps(self):
             self.addComputeGlobal("unperturbed_pe", "energy")
             self.endBlock()
             #initial iteration
-            self.addComputeGlobal("protocol_work", "protocol_work + (perturbed_pe - unperturbed_pe)")
+            # Work accumulation is handled in _add_alchemical_perturbation_step()
             super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
             #if more propogation steps are requested
             self.beginIfBlock("lambda > prop_lambda_min")
             self.beginIfBlock("lambda <= prop_lambda_max")
 
             self.beginWhileBlock("prop < nprop")
             self.addComputeGlobal("prop", "prop + 1")
-
             super(AlchemicalNonequilibriumLangevinIntegrator, self)._add_integrator_steps()
+            # Propagation steps - just do additional integration without lambda changes
+            # The parent's integrator steps are already called above, so we don't need to call them again
+            # This prevents double work accumulation
             self.endBlock()
             self.endBlock()
             self.endBlock()
@@ -234,6 +239,20 @@ def _add_alchemical_perturbation_step(self):
         # Accumulate protocol work
         self.addComputeGlobal("Enew", "energy")
         self.addComputeGlobal("protocol_work", "protocol_work + (Enew-Eold)")
+        
+        # Track work in different phases
+        self.beginIfBlock("lambda <= 0.5")
+        self.addComputeGlobal("work_0_to_05", "work_0_to_05 + (Enew-Eold)")
+        self.endBlock()
+        
+        self.beginIfBlock("lambda > 0.5")
+        self.addComputeGlobal("work_05_to_1", "work_05_to_1 + (Enew-Eold)")
+        self.endBlock()
+        
+        # Debug: Print work at move step (λ=0.5)
+        self.beginIfBlock("abs(lambda - 0.5) < 0.001")
+        self.addComputeGlobal("debug_work", "Enew - Eold")
+        self.endBlock()
         self.endBlock()
 
     def getLogAcceptanceProbability(self, context):
@@ -243,17 +262,33 @@ def getLogAcceptanceProbability(self, context):
         logp_accept = -1.0 * (protocol + shadow) * _OPENMM_ENERGY_UNIT / self.kT
         logger.info(f'[WORK] protocol work: {protocol}')
         logger.info(f'[shadow] shadow: {shadow}')
+        
+        # Debug: Print work in different phases
+        try:
+            work_0_to_05 = self.getGlobalVariableByName("work_0_to_05")
+            work_05_to_1 = self.getGlobalVariableByName("work_05_to_1")
+            debug_work = self.getGlobalVariableByName("debug_work")
+            logger.info(f'[DEBUG] Work 0→0.5: {work_0_to_05}')
+            logger.info(f'[DEBUG] Work 0.5→1.0: {work_05_to_1}')
+            logger.info(f'[DEBUG] Work at λ=0.5: {debug_work}')
+        except:
+            pass
+        
         return logp_accept
 
     def reset(self):
         self.setGlobalVariableByName("step", 0)
         self.setGlobalVariableByName("lambda", 0.0)
+        self.setGlobalVariableByName("lambda_step", 0.0)
         self.setGlobalVariableByName("protocol_work", 0.0)
         self.setGlobalVariableByName("shadow_work", 0.0)
         self.setGlobalVariableByName("first_step", 0)
         self.setGlobalVariableByName("perturbed_pe", 0.0)
         self.setGlobalVariableByName("unperturbed_pe", 0.0)
         self.setGlobalVariableByName("prop", 1)
+        self.setGlobalVariableByName("debug_work", 0.0)
+        self.setGlobalVariableByName("work_0_to_05", 0.0)
+        self.setGlobalVariableByName("work_05_to_1", 0.0)
         super(AlchemicalExternalLangevinIntegrator, self).reset()
 
 

blues/simulation.py

@@ -1072,7 +1072,14 @@ def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
                     #print("Running move_engine.runEngine() at moveStep")
                     self._ncmc_sim.context = move_engine.runEngine(self._ncmc_sim.context)
 
-
+                lambda_val = self._ncmc_sim.context._integrator.getGlobalVariableByName("lambda")
+                # if lambda val is near 0.0
+                if lambda_val < 0.000100 or  step == 0 or abs(lambda_val - 0.5) < 1e-4 or abs(lambda_val - 1.0) < 1e-4:
+                    state = self._ncmc_sim.context.getState(getPositions=True, getEnergy=True)
+                    logger.info(f"[λ={lambda_val:.6f} and Step={step}] Total potential energy: {state.getPotentialEnergy()}")
+                    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._ncmc_sim.step(1)
 
                 if step == lastStep:
@@ -1090,6 +1097,7 @@ 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)
 
+        
         # # Optional: check difference
         # import numpy as np
         # delta = np.abs(ncmc_state1['positions'] - ncmc_state0['positions'])