started on idea for acceptor and donor modes for each residue, settin…

…g H and D nonbonded parameters extra after transfer

protex/residue.py

@@ -45,6 +45,13 @@ class Residue:
     equivalent_atoms: dict[str, bool]
         if orignal_name and alternative name have equivalent atoms
     force_idxs:
+
+    modes: tuple(str)
+        whether the residue can be a donor, acceptor or both
+    donors: tuple(str)
+        atom names (NOTE: maybe use indices within molecule?) that are currently real Hs (NOTE: at the moment residues with only acceptor mode may also have donor Hs, e.g. OH)
+    acceptors: tuple(str)
+        atom names (NOTE: maybe use indices within molecule?) that are currently dummies
     """
 
     def __init__(
@@ -55,6 +62,9 @@ def __init__(
         inital_parameters,
         alternativ_parameters,
         has_equivalent_atoms,
+        modes,
+        donors,
+        acceptors,
         force_idxs=dict(),
     ) -> None:
         self.residue = residue
@@ -76,6 +86,9 @@ def __init__(
             }
         self.equivalent_atom_pos_in_list: int = None
         self.used_equivalent_atom: bool = False
+        self.modes = modes
+        self.donors = donors
+        self.acceptors = acceptors
         self.force_idxs = force_idxs
 
     def __str__(self) -> str:
@@ -97,11 +110,12 @@ def has_equivalent_atom(self) -> bool:
     @property
     def alternativ_name(self) -> str:
         """Alternative name for the residue, e.g. the corresponding name for the protonated/deprotonated form.
-
+            
         Returns
         -------
         str
             The alternative name
+        TODO: handle ampholytes (OH_H2O_H3O)
         """
         for name in self.parameters.keys():
             if name != self.current_name:

protex/system.py

@@ -104,6 +104,9 @@ def __init__(
         # store names in variables, in case syntax for states dict changes
         self._atom_name: str = "atom_name"
         self._equivalent_atom: str = "equivalent_atom"
+        self._donors: str = "donors"
+        self._acceptors: str = "acceptors"
+        self._modes: str = "modes"
 
     def dump(self, fname: str) -> None:
         """Pickle the current ProtexTemplates object.
@@ -130,6 +133,51 @@ def get_atom_name_for(self, resname: str) -> str:
             The atom name
         """
         return self.states[resname][self._atom_name]
+
+    def get_donors_for(self, resname: str) -> tuple:
+        """Get the atom names of donors for a specific residue.
+
+        Parameters
+        ----------
+        resname : str
+            The residue name
+
+        Returns
+        -------
+        tuple
+            The atom names
+        """
+        return self.states[resname][self._donors]
+
+    def get_acceptors_for(self, resname: str) -> tuple:
+        """Get the atom names of acceptors for a specific residue.
+
+        Parameters
+        ----------
+        resname : str
+            The residue name
+
+        Returns
+        -------
+        tuple
+            The atom names
+        """
+        return self.states[resname][self._acceptors]
+
+    def get_modes_for(self, resname: str) -> tuple:
+        """Get the possible modes for a specific residue.
+
+        Parameters
+        ----------
+        resname : str
+            The residue name
+
+        Returns
+        -------
+        tuple
+            tuple of mode(s)
+        """
+        return self.states[resname][self._modes]
 
     def has_equivalent_atom(self, resname: str) -> bool:
         """Check if a given residue has an equivalent atom defined.
@@ -386,13 +434,15 @@ def __init__(
         simulation: openmm.app.simulation.Simulation,
         templates: ProtexTemplates,
         simulation_for_parameters: openmm.app.simulation.Simulation = None,
+        real_Hs: list[tuple[str,str]] = [("IM1H", "H7"), ("HOAC", "H")],
         fast: bool = True,
     ) -> None:
         self.system: openmm.openmm.System = simulation.system
         self.topology: openmm.app.topology.Topology = simulation.topology
         self.simulation: openmm.app.simulation.Simulation = simulation
         self.templates: ProtexTemplates = templates
         self.simulation_for_parameters = simulation_for_parameters
+        self.real_Hs = real_Hs
         self._check_forces()
         self.detected_forces: set[str] = self._detect_forces()
         self.fast: bool = fast
@@ -591,6 +641,54 @@ def _extract_templates(self, query_name: str) -> defaultdict:
         else:
             raise RuntimeError("residue not found")
         return forces_dict
+
+
+    def _extract_templates_Hs(self, query_name: str) -> defaultdict:
+        # returns the nonbonded parameters of real Hs for the residue name
+        forces_dict = defaultdict(list)
+
+        # if there is an additional parameter file with all possible residues,
+        # use this for getting the templates
+        if self.simulation_for_parameters is not None:
+            sim = self.simulation_for_parameters
+        else:
+            sim = self.simulation
+
+        for residue in sim.topology.residues():
+            if query_name == residue.name:
+                logger.debug(residue.name)
+                logger.debug(self.real_Hs)
+                atom_names = [self.real_Hs[i][1] for i in range(len(self.real_Hs)) if self.real_Hs[i][0] == residue.name]
+                atom_idxs_all = [atom.index for atom in residue.atoms()]
+                atom_names_all = [atom.name for atom in residue.atoms()]
+                atom_idxs = [atom_idxs_all[i] for i in range(len(atom_idxs_all)) if atom_names_all[i] in atom_names]
+                logger.debug(atom_idxs)
+                logger.debug(atom_names)
+
+                for force in sim.system.getForces():
+                    forcename = type(force).__name__
+                    if forcename == "NonbondedForce":
+                        forces_dict[forcename] = [
+                            force.getParticleParameters(idx) for idx in atom_idxs
+                        ]
+                        # Also add exceptions TODO: what do we do with these? they need atom idxes and can be applied to e.g H1 and H2. is there a difference with dummies?
+                        for exc_id in range(force.getNumExceptions()):
+                            f = force.getExceptionParameters(exc_id)
+                            idx1 = f[0]
+                            idx2 = f[1]
+                            if (idx1 in atom_idxs and idx2 in atom_idxs_all) or (idx2 in atom_idxs and idx1 in atom_idxs_all):
+                                forces_dict[forcename + "Exceptions"].append(f)
+
+                    # double-checking for molecules with multiple equivalent atoms, then use one set of parameters only (we assume here that all acidic Hs in the residue are the same, e.g. MeOH2, H2O, H2OAc)
+                    if len(atom_names) > 1:
+                        for i in range(len(atom_names)):
+                            assert (forces_dict[forcename][0][1], forces_dict[forcename][0][2], forces_dict[forcename][0][3]) == (forces_dict[forcename][i][1], forces_dict[forcename][i][2], forces_dict[forcename][i][3])
+                        forces_dict = forces_dict[0]
+                break  # do this only for the relevant residue once
+        else:
+            raise RuntimeError("residue not found")
+        logger.debug(forces_dict)
+        return forces_dict
 
     @staticmethod
     def _check_nr_of_forces(
@@ -742,6 +840,16 @@ def _fill_residue_templates(self, name):
                 return
             self.residue_templates[name] = self._extract_templates(name)
 
+    def _fill_H_templates(self, name):
+        if name in self.templates.names:
+            if name in self.H_templates:
+                return
+            self.H_templates[name] = self._extract_templates_Hs(name)
+        else:
+            if name in self.H_templates:  # or name_of_paired_ion in templates:
+                return
+            self.H_templates[name] = self._extract_templates_Hs(name)
+
     def _set_initial_states(self) -> list:
         """set_initial_states.
 
@@ -750,6 +858,7 @@ def _set_initial_states(self) -> list:
         """
         residues = []
         self.residue_templates = dict()
+        self.H_templates = dict()
         # this will become a dict of the form:
         # self.per_residue_forces[(min,max tuple of the atom idxs for each residue)][forcegroup][forcename]: list of the idxs of this force for this residue
         self.per_residue_forces = {}
@@ -760,6 +869,8 @@ def _set_initial_states(self) -> list:
             self.per_residue_forces[(mini, maxi)] = {}
             name = r.name
             self._fill_residue_templates(name)
+            self._fill_H_templates(name)
+
 
         if self.fast:
             # this takes some time, but the update calls on the residues are then much faster
@@ -793,6 +904,9 @@ def _set_initial_states(self) -> list:
                             self.templates.has_equivalent_atom(name),
                             self.templates.has_equivalent_atom(name_of_paired_ion),
                         ),
+                        self.templates.get_modes_for(name),
+                        self.templates.get_donors_for(name),
+                        self.templates.get_acceptors_for(name),
                         force_idxs=self.per_residue_forces[minmax],
                     )
                 else:
@@ -806,6 +920,9 @@ def _set_initial_states(self) -> list:
                             self.templates.has_equivalent_atom(name),
                             self.templates.has_equivalent_atom(name_of_paired_ion),
                         ),
+                        self.templates.get_modes_for(name),
+                        self.templates.get_donors_for(name),
+                        self.templates.get_acceptors_for(name),
                     )
                 residues.append(new_res)
                 residues[
@@ -815,7 +932,7 @@ def _set_initial_states(self) -> list:
                 )
             else:
                 parameters_state1 = self.residue_templates[name]
-                r = Residue(r, None, self.system, parameters_state1, None, None)
+                r = Residue(r, None, self.system, parameters_state1, None, None, None, None, None)
                 # the residue is not part of any proton transfers,
                 # we still need it in the residue list for the parmed hack...
                 # there we need the current_name attribute, hence give it to the residue
@@ -825,6 +942,7 @@ def _set_initial_states(self) -> list:
             #    raise RuntimeError(
             #        f"Found resiude not present in Templates: {r.name}"
             #    )  # we want to ignore meoh, doesn't work the way it actually is
+
         return residues
 
     # def save_current_names(self, file: str) -> None:

protex/tests/test_h2o.py

@@ -186,8 +186,10 @@ def test_run_simulation(tmp_path):
 
 
 def test_scratch():
-    a = "hello"
-    print(f"{'o' in a}")
+    a = [('IM1H', 'H7'), ('HOAC', 'H')]
+    for j in range(len(a)):
+        print(a[j][1])
+
 
 def test_create_ProtexTemplate():
     allowed_updates = {}

protex/tests/test_system.py

@@ -642,8 +642,8 @@ def test_create_IonicLiquidTemplate():
 
 
 def test_create_IonicLiquid():
-    # simulation = generate_small_box(use_plugin=False)
-    simulation = generate_im1h_oac_system(use_plugin=False)
+    simulation = generate_small_box(use_plugin=False)
+    #simulation = generate_im1h_oac_system(use_plugin=False)
     allowed_updates = {}
     allowed_updates[frozenset(["IM1H", "OAC"])] = {"r_max": 0.16, "prob": 2.33}
     allowed_updates[frozenset(["IM1", "HOAC"])] = {"r_max": 0.16, "prob": -2.33}

protex/testsystems.py

@@ -988,20 +988,19 @@ def generate_single_hpts_meoh_system(
 
 IM1H_IM1 = {
     "IM1H": {
-        "atom_name": "H7",
+        "donors": ("H7"), "acceptors": (), "modes": ("donor")
     },
     "IM1": {
-        "atom_name": "N2",
+        "donors": (), "acceptors": ("H7"), "modes": ("acceptor")
     },
 }
 
 OAC_HOAC = {
     "OAC": {
-        "atom_name": "O2",
-        "equivalent_atom": "O1",
+        "donors": (), "acceptors": ("H"), "modes": ("acceptor")
     },
     "HOAC": {
-        "atom_name": "H",
+        "donors": ("H"), "acceptors": (), "modes": ("donor")
     },
 }
 
@@ -1027,39 +1026,33 @@ def generate_single_hpts_meoh_system(
 }
 
 # started working on new structure with possible donor / acceptor H sites in templates
-# idea: specify atom names of all Hs and which is a real H ("donor") in the supplied topology -> automatically define acceptors
-# find all possible donor / acceptor combinations based on number of donors in molecule
-def get_all_states(possible_atoms, num_donors):
-        states = []
-        donors = list(combinations(possible_atoms, num_donors))
-        for i in range(0, len(donors)):
-            donor = donors[i]
-            acceptor = (tuple(set(possible_atoms).symmetric_difference(donor)))
-            states.append({"donors": donor, "acceptors": acceptor})
-        return states
-
-def OH_H2O_H3O():
-
-    # NOTE: take care whether we want to use H2O or SWM4, SPCE etc. for residue name
-    OH_H2O_H3O =  {
-        "OH": {"possible_atoms" : ("H1", "H2", "H3", "H4"), "num_donors" : 1, "starting_donors" : ("H1",), "modes" : ("acceptor")},
-        "H2O": {"possible_atoms" : ("H1", "H2", "H3", "H4"), "num_donors" : 2, "starting_donors" : ("H1", "H2"), "modes" : ("acceptor", "donor")},
-        "H3O": {"possible_atoms" : ("H1", "H2", "H3", "H4"), "num_donors" : 3, "starting_donors" : ("H1", "H2", "H3"), "modes" : ("donor")},
-    }
-
-    for i in OH_H2O_H3O:
-        OH_H2O_H3O[i]["starting_acceptors"] = tuple(set(OH_H2O_H3O[i]["possible_atoms"]).symmetric_difference(OH_H2O_H3O[i]["starting_donors"]))
-        OH_H2O_H3O[i]["possible_states"] = get_all_states(OH_H2O_H3O[i]["possible_atoms"], OH_H2O_H3O[i]["num_donors"])
-
-    return OH_H2O_H3O
+# def get_all_states(possible_atoms, num_donors): # probably not needed with more generic donor-acceptor structure
+#         states = []
+#         donors = list(combinations(possible_atoms, num_donors))
+#         for i in range(0, len(donors)):
+#             donor = donors[i]
+#             acceptor = (tuple(set(possible_atoms).symmetric_difference(donor)))
+#             states.append({"donors": donor, "acceptors": acceptor})
+#         return states
+
+
+# NOTE: take care whether we want to use H2O or SWM4, SPCE etc. for residue name
+# TODO: at the moment fixed atom names, will revert to this at the beginning of each run -> reformulate so that donors and acceptors are filled based on atom type
+
+OH_H2O_H3O =  {
+    "OH":  {"donors" : ("H1",), "acceptors" : ("H2", "H3", "H4"), "modes" : ("acceptor",)},
+    "H2O": {"donors" : ("H1", "H2"), "acceptors" : ("H3", "H4"), "modes" : ("acceptor", "donor")},
+    "H3O": {"donors" : ("H1", "H2", "H3"), "acceptors" : ("H4",), "modes" : ("donor",)},
+}
 
 # TODO:
-# find a way to switch parameters around to get the parameter sets for each possible state (problem: how to get original state for each new run)
-    # maybe something like: by going from OHHDD to OHDHD, we switch atoms 3 and 4,
-    # so the original bond forces between O and H or D were [1, 3, k13, r13] and [1, 4, k14, r14],
-    # now we need [1, 3, k14, r14] and [1, 4, k13, r13]
-# keep track of actual state of each residue
-    #  maybe like "actual_state" : {"donors" : (), "acceptors" : ()}
+# keep track of what H is real at the moment: like "donors" : (), "acceptors" : ()
+    # caution: will have to change how we write the psf as well
 # get atom index or something like that in the update step and swap H to D and vice versa
     # something similar to the way we check whether the equivalent atom was used
-    # then e.g. OHHDD -> OHHHD: H3 was used -> we need parameters from H3O, where H1, H2, H3 are donors, H4 is acceptor
+# find a way to switch parameters around to get the parameter sets for each possible state (problem: how to get original state for each new run)
+    # at the moment only nonbonded parameters change, all possible donors and acceptors in a single molecule are equivalent (i.e. no two different acidic side chains)
+        # update like now
+        # set parameters for H and D that were used extra
+            # D: force.setParticleParameters(atomidx, 0, 0, 0] 
+            # H: extract nonbonded parameters from template first