multipledihedrals for impropers (#387)

* change topology object

* fix top.to_dict() error for impropers

* black

* fix coordinates part

---------

Co-authored-by: Eric Hartmann <hartmaec@rh05659.villa-bosch.de>

src/kimmdy/coordinates.py

@@ -19,6 +19,7 @@
     Dihedral,
     DihedralType,
     ProperDihedralId,
+    ImproperDihedralId,
     MultipleDihedrals,
     Interaction,
     InteractionType,
@@ -157,8 +158,12 @@ def merge_dihedrals(
     dihedral_key: tuple[str, str, str, str],
     dihedral_a: Optional[Dihedral],
     dihedral_b: Optional[Dihedral],
-    dihedral_types_a: dict[ProperDihedralId, DihedralType],
-    dihedral_types_b: dict[ProperDihedralId, DihedralType],
+    dihedral_types_a: Union[
+        dict[ProperDihedralId, DihedralType], dict[ImproperDihedralId, DihedralType]
+    ],
+    dihedral_types_b: Union[
+        dict[ProperDihedralId, DihedralType], dict[ImproperDihedralId, DihedralType]
+    ],
     molA: MoleculeType,
     molB: MoleculeType,
     funct: str,
@@ -447,28 +452,58 @@ def merge_top_moleculetypes_slow_growth(
                 )
 
     # improper dihedrals
-    # all impropers in amber99SB ffbonded.itp have a periodicity of 2
-    # but not the ones defined in aminoacids.rtp. For now, I am assuming
-    # a periodicity of 2 in this section
+    # TODO: duplicate of proper dihedrals, could refactor
 
     keys = set(molA.improper_dihedrals.keys()) | set(molB.improper_dihedrals.keys())
     for key in keys:
-        interactionA = molA.improper_dihedrals.get(key)
-        interactionB = molB.improper_dihedrals.get(key)
+        multiple_dihedralsA = molA.improper_dihedrals.get(key)
+        multiple_dihedralsB = molB.improper_dihedrals.get(key)
 
-        if interactionA != interactionB:
-            molB.improper_dihedrals[key] = merge_dihedrals(
-                key,
-                interactionA,
-                interactionB,
-                ff.improper_dihedraltypes,
-                ff.improper_dihedraltypes,
-                molA,
-                molB,
-                "4",
-                "2",
+        if multiple_dihedralsA != multiple_dihedralsB:
+            multiple_dihedralsA = get_explicit_MultipleDihedrals(
+                key, molA, multiple_dihedralsA, ff
+            )
+            multiple_dihedralsB = get_explicit_MultipleDihedrals(
+                key, molB, multiple_dihedralsB, ff
+            )
+            keysA = (
+                set(multiple_dihedralsA.dihedrals.keys())
+                if multiple_dihedralsA
+                else set()
+            )
+            keysB = (
+                set(multiple_dihedralsB.dihedrals.keys())
+                if multiple_dihedralsB
+                else set()
             )
 
+            molB.improper_dihedrals[key] = MultipleDihedrals(*key, "9", {})
+            periodicities = keysA | keysB
+            for periodicity in periodicities:
+                assert isinstance(periodicity, str)
+                interactionA = (
+                    multiple_dihedralsA.dihedrals.get(periodicity)
+                    if multiple_dihedralsA
+                    else None
+                )
+                interactionB = (
+                    multiple_dihedralsB.dihedrals.get(periodicity)
+                    if multiple_dihedralsB
+                    else None
+                )
+
+                molB.improper_dihedrals[key].dihedrals[periodicity] = merge_dihedrals(
+                    key,
+                    interactionA,
+                    interactionB,
+                    ff.improper_dihedraltypes,
+                    ff.improper_dihedraltypes,
+                    molA,
+                    molB,
+                    "9",
+                    periodicity,
+                )
+
     # amber fix for breaking/binding atom types without LJ potential
     # breakpoint()
 

src/kimmdy/topology/atomic.py

@@ -559,17 +559,18 @@ def from_top_line(cls, l: list[str]):
 
 @dataclass()
 class ResidueImproperSpec:
-    """Information about one imroper dihedral in a residue
+    """Information about one improper dihedral in a residue
 
-    ; atom1 atom2 atom3 atom4 q0 cq
+    ; atom1 atom2 atom3 atom4 c0(q0) c1(cp) c2(mult)
     """
 
     atom1: str
     atom2: str
     atom3: str
     atom4: str
-    q0: Optional[str]
-    cq: Optional[str]
+    c0: Optional[str]
+    c1: Optional[str]
+    c2: Optional[str]
 
     @classmethod
     def from_top_line(cls, l: list[str]):
@@ -578,23 +579,26 @@ def from_top_line(cls, l: list[str]):
             atom2=l[1],
             atom3=l[2],
             atom4=l[3],
-            q0=field_or_none(l, 4),
-            cq=field_or_none(l, 5),
+            c0=field_or_none(l, 4),
+            c1=field_or_none(l, 5),
+            c2=field_or_none(l, 6),
         )
 
 
 @dataclass()
 class ResidueProperSpec:
-    """Information about one imroper dihedral in a residue
+    """Information about one proper dihedral in a residue
 
-    ; atom1 atom2 atom3 atom4 q0 cq
+    ; atom1 atom2 atom3 atom4 c0(q0) c1(cq) c2
     """
 
     atom1: str
     atom2: str
     atom3: str
     atom4: str
-    q0: Optional[str]
+    c0: Optional[str]
+    c1: Optional[str]
+    c2: Optional[str]
 
     @classmethod
     def from_top_line(cls, l: list[str]):
@@ -603,7 +607,9 @@ def from_top_line(cls, l: list[str]):
             atom2=l[1],
             atom3=l[2],
             atom4=l[3],
-            q0=field_or_none(l, 4),
+            c0=field_or_none(l, 4),
+            c1=field_or_none(l, 5),
+            c2=field_or_none(l, 6),
         )
 
 

src/kimmdy/topology/topology.py

@@ -72,7 +72,7 @@ def __init__(
         self.pairs: dict[tuple[str, str], Pair] = {}
         self.angles: dict[tuple[str, str, str], Angle] = {}
         self.proper_dihedrals: dict[tuple[str, str, str, str], MultipleDihedrals] = {}
-        self.improper_dihedrals: dict[tuple[str, str, str, str], Dihedral] = {}
+        self.improper_dihedrals: dict[tuple[str, str, str, str], MultipleDihedrals] = {}
         self.position_restraints: dict[str, PositionRestraint] = {}
         self.dihedral_restraints: dict[tuple[str, str, str, str], DihedralRestraint] = (
             {}
@@ -183,7 +183,11 @@ def _parse_dihedrals(self):
             dihedral = Dihedral.from_top_line(l)
             key = (dihedral.ai, dihedral.aj, dihedral.ak, dihedral.al)
             if dihedral.funct == "4":
-                self.improper_dihedrals[key] = dihedral
+                if self.improper_dihedrals.get(key) is None:
+                    self.improper_dihedrals[key] = MultipleDihedrals(
+                        *key, dihedral.funct, dihedrals={}
+                    )
+                self.improper_dihedrals[key].dihedrals[dihedral.periodicity] = dihedral
             else:
                 if self.proper_dihedrals.get(key) is None:
                     self.proper_dihedrals[key] = MultipleDihedrals(
@@ -263,7 +267,11 @@ def _update_atomics_dict(self):
             attributes_to_list(x)
             for dihedrals in self.proper_dihedrals.values()
             for x in dihedrals.dihedrals.values()
-        ] + [attributes_to_list(x) for x in self.improper_dihedrals.values()]
+        ] + [
+            attributes_to_list(x)
+            for dihedrals in self.improper_dihedrals.values()
+            for x in dihedrals.dihedrals.values()
+        ]
         self.atomics["position_restraints"] = [
             attributes_to_list(x) for x in self.position_restraints.values()
         ]
@@ -464,13 +472,18 @@ def _regenerate_topology_from_bound_to(self, ff):
         for atom in self.atoms.values():
             impropers = self._get_atom_improper_dihedrals(atom.nr, ff)
             for key, improper in impropers:
-                self.improper_dihedrals[key] = Dihedral(
-                    improper.atom1,
-                    improper.atom2,
-                    improper.atom3,
-                    improper.atom4,
+                self.improper_dihedrals[key] = MultipleDihedrals(
+                    *key,
                     "4",
-                    improper.cq,
+                    dihedrals={
+                        "": Dihedral(
+                            *key,
+                            "4",
+                            c0=improper.c0,
+                            c1=improper.c1,
+                            periodicity=improper.c2,
+                        )
+                    },
                 )
 
     def reindex_atomnrs(self) -> dict[str, str]:
@@ -1256,11 +1269,12 @@ def bind_bond(
             atompair_nrs[0], self.ff
         ) + reactive_moleculetype._get_atom_improper_dihedrals(atompair_nrs[1], self.ff)
         for key, value in dihedral_k_v:
+            if value.c2 is None:
+                value.c2 = ""
             if reactive_moleculetype.improper_dihedrals.get(key) is None:
-                # TODO: fix this
-                c2 = ""
-                if value.q0 is not None:
-                    c2 = "1"
-                reactive_moleculetype.improper_dihedrals[key] = Dihedral(
-                    key[0], key[1], key[2], key[3], "4", value.q0, value.cq, c2
+                reactive_moleculetype.improper_dihedrals[key] = MultipleDihedrals(
+                    *key, "4", dihedrals={}
                 )
+            reactive_moleculetype.improper_dihedrals[key].dihedrals[value.c2] = (
+                Dihedral(*key, "4", c0=value.c0, c1=value.c1, periodicity=value.c2)
+            )

tests/test_coordinates.py

@@ -145,7 +145,9 @@ def test_merge_prm_top(arranged_tmp_path):
     assert top_merge.bonds[("26", "27")].funct == "3"
     assert top_merge.angles[("17", "19", "20")].c3 is not None
     assert top_merge.proper_dihedrals[("15", "17", "19", "24")].dihedrals["3"].c5 == "3"
-    assert top_merge.improper_dihedrals[("17", "20", "19", "24")].c5 == "2"
+    assert (
+        top_merge.improper_dihedrals[("17", "20", "19", "24")].dihedrals["2"].c5 == "2"
+    )
     # assert one dihedral merge improper/proper
 
 

tests/test_integration.py

@@ -47,6 +47,7 @@ def test_integration_valid_input_files(arranged_tmp_path):
 def test_grompp_with_kimmdy_topology(arranged_tmp_path):
     raw_top = read_top(Path("minimal.top"))
     top = Topology(raw_top)
+    top_dict = top.to_dict()
     write_top(top.to_dict(), Path("output.top"))
     assert sp.run(
         [