Feat: ramping and continue_until_unbound options

openmm_ramd/logger.py

@@ -27,7 +27,7 @@ def log(self, string, prefix=RAMD_PREFIX, print_also=False):
         full_string = prefix + string + "\n"
         self.file.write(full_string)
         if print_also:
-            print("full_string")
+            print(full_string)
         return
 
     def exit_log(self, string, print_also=False):

openmm_ramd/openmm_ramd.py

@@ -94,7 +94,7 @@ def __init__(self, topology, system, integrator, ramd_force_magnitude,
                 topology, system, integrator, platform, properties)
 
     def recompute_RAMD_force(self):
-        print("recomputing force!")
+        #print("recomputing force!")
         self.force_handler.set_new_RAMD_force_vector()
         self.force_handler.force_object.updateParametersInContext(self.context)
 
@@ -271,18 +271,99 @@ def RAMD_step(self, numSteps=50):
         self.old_lig_com = lig_com
         return lig_com, rec_com
 
-    def run_RAMD_sim(self, max_num_steps=1e8):
-        self.counter = 0
-        start_lig_com = self.RAMD_start()
-
-        # Do the simulation steps and loop here. These are done every step
+    def main_ramd_loop(self, max_num_steps):
         while self.counter < max_num_steps:
             lig_com, rec_com = self.RAMD_step(self.ramdSteps)
-            lig_prot_com_distance = np.linalg.norm(lig_com.value_in_unit(unit.angstroms) - rec_com.value_in_unit(unit.angstroms))
+            lig_prot_com_distance = np.linalg.norm(lig_com.value_in_unit(unit.angstroms) - 
+                                                   rec_com.value_in_unit(unit.angstroms))
             if lig_prot_com_distance > self.maxDist:
                 self.max_distance_exceeded(self.counter)
-                break
+                return self.counter
+
+    def run_RAMD_sim(self, max_num_steps=1e8, ramping=False, continue_until_unbound=False,
+                     extra_force=2.0, extra_steps=500_000):
+        """
+        Run the RAMD simulation until the maximum distance is exceeded.
+
+        Parameters
+        ----------
+        max_num_steps : int
+            Maximum number of simulation steps to perform.
+        ramping : bool
+            Whether to ramp up the RAMD force magnitude over time.
+            Starts with no force and increases linearly to the target.
+        continue_until_unbound : bool
+            Whether to continue the simulation until the ligand is unbound,
+            even if the maximum number of steps is reached.
+            Will continue for extra_steps additional steps at an extra extra_force kcal/mol*Angstrom.
+            This repeats with increasing force until unbound.
+        extra_force : float
+            The extra force magnitude to add when continuing after max steps.
+            Default is 2.0 kcal/mol*Angstrom.
+        extra_steps : int
+            The number of extra steps to run when continuing after max steps.
+            Default is 500,000 steps (2 ns at 4 fs timestep).
+        """
+        self.counter = 0
+        start_lig_com = self.RAMD_start()
+
+        # ramping option
+        if ramping:
+            # create linear schedule from 0 to max force in steps of 2 kcal/mol*Angstrom
+            max_force_magnitude = self.force_handler.random_force_magnitude.value_in_unit(
+                kcal_per_mole_per_angstrom)
+            # start schedule at 0 force and then go up in steps of 2 from 1/2 max force
+            force_schedule = np.concatenate(([0], np.arange(max_force_magnitude/2, max_force_magnitude, 2.0)))
+            # steps_between_force_updates = 50000  # e.g., 100/200 ps with 2/4 fs timestep
+            # start with 25000 steps (100 ps at 4 fs timestep) to randomize initial state
+            # then run ramd of the other forces at increasingly larger step counts
+            # go up from 0.5 ns to 4 ns (125_000 to 1_000_000 steps at 4 fs timestep)
+            # use same n elements as in force schedule
+            num_forces = len(force_schedule)
+            steps_between_force_updates = np.concatenate(([25_000], 
+                                            np.linspace(125_000, 1_000_000, num_forces - 1, dtype=int)))
+            self.logger.log(f"Ramping RAMD force magnitude: {force_schedule} kcal/mol*Angstrom " +
+                            f"with {steps_between_force_updates} steps between updates.")
+
+            # loop over force schedule
+            for force, n_steps in dict(zip(force_schedule, steps_between_force_updates)).items():
+                # set updated force magnitude
+                self.force_handler.random_force_magnitude = force * kcal_per_mole_per_angstrom
+                self.recompute_RAMD_force()
+
+                # set number of steps to run until next force update
+                # starts at 0 steps
+                next_n_steps = self.counter + n_steps
+                # run regular RAMD simulation until next force update
+                self.main_ramd_loop(next_n_steps)
+
+                self.logger.log(f"Ramped RAMD force magnitude to {force} kcal/mol*Angstrom at step {self.counter}.")
+            self.logger.log("Completed RAMD force ramping.")
+
+        # Do the standard RAMD simulation steps and loop here. These are done every step
+        # by default 0.1/0.2 ps or 50 steps with 2/4 fs timestep before potential force update
+        self.main_ramd_loop(max_num_steps)
+
+        # if continue_until_unbound is set, continue with extra force if max steps reached
+        if continue_until_unbound:
+            while True:
+                # increase max_num_steps by extra_steps
+                max_num_steps += extra_steps
+                # set new increased force magnitude
+                ramd_force_magnitude = self.force_handler.random_force_magnitude.value_in_unit(
+                    kcal_per_mole_per_angstrom) + extra_force
+                self.logger.log(f"WARNING: Maximum number of steps increased to {max_num_steps}. " +
+                                f"Continuing RAMD simulation with extra force " +
+                                f"of {extra_force} kcal/mol*Angstrom for {extra_steps} extra steps. " +
+                                f"Total force magnitude of {ramd_force_magnitude} kcal/mol*Angstrom.",
+                                print_also=True)
+                self.force_handler.random_force_magnitude = ramd_force_magnitude * kcal_per_mole_per_angstrom
+                self.recompute_RAMD_force()
+
+                # run additional steps with increased force
+                self.main_ramd_loop(max_num_steps)
 
+        # return the total number of steps performed
         return self.counter
 
 if __name__ == "__main__":