started working on distance based probabilities

protex/residue.py

@@ -81,6 +81,12 @@ def __init__(
     def __str__(self) -> str:
         return f"Residue {self.current_name}, {self.residue}"
 
+    def __eq__(self, other) -> bool:
+        return self.current_name == other.current_name and self.residue.index == other.residue.index
+
+    def __hash__(self):
+        return hash((self.current_name, self.residue.index))
+
     @property
     def has_equivalent_atom(self) -> bool:
         """Determines if the current residue has an equivalent atom defined.
@@ -500,7 +506,7 @@ def _get_offset(self, name, force_name=None):
     def _get_offset_thole(self, name):
         # get offset for atom idx for thole parameters
         force_name = "DrudeForceThole"
-        print(self.parameters[name])
+        #print(self.parameters[name])
         return min(
             itertools.chain(
                 *[query_parm[0:2] for query_parm in self.parameters[name][force_name]]

protex/testsystems.py

@@ -63,6 +63,7 @@ def setup_system(
     dummy_atom_type: str = "DUMH",
     cutoff: float = 11.0,
     switch: float = 10.0,
+    barostat=None
 ):
     if dummy_atom_type is not None:
         # print(params.atom_types_str["DUM"].epsilon)
@@ -103,6 +104,10 @@ def setup_system(
         print(
             "Only contraints=None or constraints=HBonds (given as string in function call) implemented"
         )
+
+
+    if barostat is not None:
+        system.addForce(barostat)
 
     for force in system.getForces():
         if type(force).__name__ == "CMMotionRemover":
@@ -630,6 +635,7 @@ def generate_im1h_oac_dummy_system(
     dummy_atom_type: str = "DUMH",
     dummies: list[tuple[str, str]] = [("IM1", "H7"), ("OAC", "H")],
     use_plugin: bool = True,
+    barostat=None
 ):
     """Set up a solvated and parametrized system for IM1H/OAC."""
     base = f"{protex.__path__[0]}/forcefield"
@@ -654,7 +660,7 @@ def generate_im1h_oac_dummy_system(
         boxl=boxl,
     )
     system = setup_system(
-        psf, params, constraints=constraints, dummy_atom_type=dummy_atom_type
+        psf, params, constraints=constraints, dummy_atom_type=dummy_atom_type, barostat=barostat
     )
 
     simulation = setup_simulation(

protex/update.py

@@ -55,12 +55,14 @@ def __init__(
         all_forces: bool,
         include_equivalent_atom: bool,
         reorient: bool,
+        K:int = 300,
     ) -> None:
         self.ionic_liquid: ProtexSystem = ionic_liquid
         self.to_adapt: list[tuple[str, int, frozenset[str]]] = to_adapt
         self.include_equivalent_atom: bool = include_equivalent_atom
         self.reorient: bool = reorient
         self.all_forces: bool = all_forces
+        self.K = K
         allowed_forces: list[str] = [  # change charges only
             "NonbondedForce",  # BUG: Charge stored in the DrudeForce does NOT get updated, probably you want to allow DrudeForce as well!
             "CustomNonbondedForce",  # NEW
@@ -102,7 +104,7 @@ def _update(self, candidates: list[tuple], nr_of_steps: int) -> None:
         pass
 
     def _adapt_probabilities(
-        self, to_adapt=list[tuple[str, int, frozenset[str]]]
+        self, to_adapt=list[tuple[str, int, frozenset[str]]], K=int
     ) -> None:
         """Adapt the probability for certain events depending on the current equilibrium, in order to stay close to a given reference
         i.e. prob_neu = prob_orig + K*( x(t) - x(eq) )^3 where x(t) is the current percentage in the system of one species.
@@ -119,7 +121,7 @@ def _adapt_probabilities(
             len([e for e, c in counts.items() if c > 1]) == 0
         ), "No duplicates for the transfer reactions allowed!"
 
-        K = 300
+
         current_numbers: dict[
             str, int
         ] = self.ionic_liquid.get_current_number_of_each_residue_type()
@@ -182,9 +184,10 @@ def __init__(
         to_adapt: list[tuple[str, int, frozenset[str]]] or None = None,
         include_equivalent_atom: bool = False,
         reorient: bool = False,
+        K: int = 300
     ) -> None:
         super().__init__(
-            ionic_liquid, to_adapt, all_forces, include_equivalent_atom, reorient
+            ionic_liquid, to_adapt, all_forces, include_equivalent_atom, reorient, K
         )
 
     def dump(self, fname: str) -> None:
@@ -216,9 +219,9 @@ def _reorient_atoms(self, candidate, positions, positions_copy):
                 )
 
                 pos_atom = positions_copy[atom_idx]
-                print(f"{pos_atom=}")
+                #print(f"{pos_atom=}")
                 pos_equivalent = positions_copy[equivalent_idx]
-                print(f"{pos_equivalent=}")
+                #print(f"{pos_equivalent=}")
 
                 positions[atom_idx] = pos_equivalent
                 positions[equivalent_idx] = pos_atom
@@ -237,9 +240,9 @@ def _reorient_atoms(self, candidate, positions, positions_copy):
                 #     positions[atom_idx+6] = pos_equivalent_lp12
                 #     positions[equivalent_idx+5] = pos_atom_lp21
                 #     positions[equivalent_idx+6] = pos_atom_lp22
-                print(
-                    f"setting position of {self.ionic_liquid.templates.get_atom_name_for(resi.current_name)} to {positions[atom_idx]} and {self.ionic_liquid.templates.get_equivalent_atom_for(resi.current_name)} to {positions[equivalent_idx]}"
-                )
+                # print(
+                #     f"setting position of {self.ionic_liquid.templates.get_atom_name_for(resi.current_name)} to {positions[atom_idx]} and {self.ionic_liquid.templates.get_equivalent_atom_for(resi.current_name)} to {positions[equivalent_idx]}"
+                # )
 
         # set new H position:
         if "H" in self.ionic_liquid.templates.get_atom_name_for(
@@ -312,9 +315,9 @@ def _reorient_atoms(self, candidate, positions, positions_copy):
         # update position of the once-dummy H on the acceptor - original H line
         positions[idx_accepted_H] = pos_accepted_H
 
-        print(
-            f"donated H: {pos_donated_H}, acceptor atom: {pos_acceptor_atom}, H set to: {pos_accepted_H}"
-        )
+        # print(
+        #     f"donated H: {pos_donated_H}, acceptor atom: {pos_acceptor_atom}, H set to: {pos_accepted_H}"
+        # )
 
         return positions
 
@@ -373,10 +376,10 @@ def _update(self, candidates: list[tuple], nr_of_steps: int) -> None:
 
         for candidate in candidates:
             candidate1_residue, candidate2_residue = candidate
-            print(f"candidate pair {candidates.index(candidate)}")
-            print(
-                f"candidate1 used equivalent atom: {candidate1_residue.used_equivalent_atom}, candidate2 used equivalent atom: {candidate2_residue.used_equivalent_atom}"
-            )
+            #print(f"candidate pair {candidates.index(candidate)}")
+            # print(
+            #     f"candidate1 used equivalent atom: {candidate1_residue.used_equivalent_atom}, candidate2 used equivalent atom: {candidate2_residue.used_equivalent_atom}"
+            # )
 
             positions = self._reorient_atoms(candidate, positions, positions_copy)
 
@@ -571,13 +574,15 @@ def load(fname: str, updateMethod: Update) -> StateUpdate:
             from_pickle = pickle.load(inp)  # ensure correct order of arguments
         state_update.history = from_pickle[0]
         state_update.update_trial = from_pickle[1]
+        state_update.prob_function = from_pickle[2]
         return state_update
 
-    def __init__(self, updateMethod: Update) -> None:
+    def __init__(self, updateMethod: Update, prob_function: str = None) -> None:
         self.updateMethod: Update = updateMethod
         self.ionic_liquid: ProtexSystem = self.updateMethod.ionic_liquid
         self.history: deque = deque(maxlen=10)
         self.update_trial: int = 0
+        self.prob_function = prob_function
 
     def dump(self, fname: str) -> None:
         """Pickle the StateUpdate instance.
@@ -590,6 +595,7 @@ def dump(self, fname: str) -> None:
         to_pickle = [
             self.history,
             self.update_trial,
+            self.prob_function
         ]  # enusre correct order of arguments
         with open(fname, "wb") as outp:
             pickle.dump(to_pickle, outp, pickle.HIGHEST_PROTOCOL)
@@ -658,6 +664,7 @@ def _print_start(self):
             ##############################
             ##############################
             --- Update trial: {self.update_trial} ---
+            {self.updateMethod.K=}
             ##############################
             ##############################
             """
@@ -706,7 +713,7 @@ def update(self, nr_of_steps: int = 2) -> list[tuple[Residue, Residue]]:
         self.update_trial += 1
 
         if self.updateMethod.to_adapt is not None:
-            self.updateMethod._adapt_probabilities(self.updateMethod.to_adapt)
+            self.updateMethod._adapt_probabilities(self.updateMethod.to_adapt, self.updateMethod.K)
 
         self._print_stop()
 
@@ -736,6 +743,12 @@ def _rPBC(
                 dz = boxl - dz
             return np.sqrt(dx * dx + dy * dy + dz * dz)
 
+        def distance_based_probability(r, r_min, r_max, prob):
+            if self.prob_function == "linear":
+                return -(prob/(r_max-r_min))*r+prob*r_max/(r_max-r_min)
+            elif self.prob_function == "cosine":
+                return (prob/2)*np.cos(np.pi/(r_max-r_min)*r-np.pi*r_min/(r_max-r_min))+prob/2
+
         # calculate distance matrix between the two molecules
         if use_pbc:
             logger.debug("Using PBC correction for distance calculation")
@@ -774,12 +787,27 @@ def _rPBC(
                 prob = self.ionic_liquid.templates.allowed_updates[
                     frozenset([residue1.current_name, residue2.current_name])
                 ]["prob"]
-                logger.debug(f"{r_max=}, {prob=}")
                 r = distance[candidate_idx1, candidate_idx2]
+                if self.prob_function is None:
+                    dist_prob = prob
+                    logger.debug(f"{r=}, {r_max=}, {prob=}")
+                else:
+                    r_min = self.ionic_liquid.templates.allowed_updates[
+                    frozenset([residue1.current_name, residue2.current_name])
+                    ]["r_min"]
+                    if r < r_min:
+                        dist_prob = prob
+                    elif r <= r_max:
+                        dist_prob = distance_based_probability(r, r_min, r_max, prob)
+                    else:
+                        dist_prob = 0
+                    logger.debug(f"{r=}, {r_min=}, {r_max=}, {prob=}, {dist_prob=}, {self.prob_function=}")
+                #print(f"{r=}, {r_min=}, {r_max=}, {prob=}, {dist_prob=}, {self.prob_function=}")
+
                 # break for loop if no pair can fulfill distance condition
                 if r > self.ionic_liquid.templates.overall_max_distance:
                     break
-                elif r <= r_max and random.random() <= prob:  # random enough?
+                elif r <= r_max and random.random() <= dist_prob:  # random enough?
                     charge_candidate_idx1 = residue1.endstate_charge
                     charge_candidate_idx2 = residue2.endstate_charge