add missing bonds to topology

SimonBoothroyd · Nov 22, 2024 · a8899d3 · a8899d3
1 parent abffa10
commit a8899d3
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 7 deletions.
diff --git a/absolv/runner.py b/absolv/runner.py
@@ -98,6 +98,7 @@ def _rebuild_topology(
     atom_counts_per_residue = collections.defaultdict(
         lambda: collections.defaultdict(int)
     )
+    atoms = []
 
     last_residue_idx = -1
     residue = None
@@ -113,14 +114,27 @@ def _rebuild_topology(
         symbol = "X" if element is None else element.symbol
 
         atom_counts_per_residue[residue_idx][atomic_num] += 1
-        topology.addAtom(
+        atom = topology.addAtom(
             f"{symbol}{atom_counts_per_residue[residue_idx][atomic_num]}".ljust(3, "x"),
             element,
             residue,
         )
 
+        atoms.append(atom)
+
     _rename_residues(topology)
 
+    atom_idx_to_particle_idx = {j: i for i, j in particle_idx_to_atom_idx.items()}
+
+    for bond in orig_top.bonds:
+        if atoms[atom_idx_to_particle_idx[bond.atom1_index]].residue.name == "HOH":
+            continue
+
+        topology.addBond(
+            atoms[atom_idx_to_particle_idx[bond.atom1_index]],
+            atoms[atom_idx_to_particle_idx[bond.atom2_index]],
+        )
+
     coords_with_v_sites = []
 
     for particle_idx in range(system.getNumParticles()):
@@ -158,6 +172,9 @@ def _rename_residues(topology: openmm.app.Topology):
         if symbols == ["H", "H", "O"]:
             residue.name = "HOH"
 
+            for i, atom in enumerate(residue.atoms()):
+                atom.name = "OW" if atom.element.symbol == "O" else f"HW{i}"
+
 
 def _setup_solvent(
     solvent_idx: typing.Literal["solvent-a", "solvent-b"],

diff --git a/absolv/tests/test_runner.py b/absolv/tests/test_runner.py
@@ -132,12 +132,12 @@ def test_rebuild_topology():
         ("H3x", "H", 0, "UNK"),
         ("H4x", "H", 0, "UNK"),
         ("H5x", "H", 0, "UNK"),
-        ("O1x", "O", 1, "HOH"),
-        ("H1x", "H", 1, "HOH"),
-        ("H2x", "H", 1, "HOH"),
-        ("O1x", "O", 2, "HOH"),
-        ("H1x", "H", 2, "HOH"),
-        ("H2x", "H", 2, "HOH"),
+        ("OW", "O", 1, "HOH"),
+        ("HW1", "H", 1, "HOH"),
+        ("HW2", "H", 1, "HOH"),
+        ("OW", "O", 2, "HOH"),
+        ("HW1", "H", 2, "HOH"),
+        ("HW2", "H", 2, "HOH"),
         ("X1x", None, 3, "UNK"),
         ("X1x", None, 4, "UNK"),
         ("X1x", None, 5, "UNK"),
@@ -178,6 +178,27 @@ def test_rebuild_topology():
 
     assert numpy.allclose(box_vectors, expected_box_vectors)
 
+    expected_bonds = [
+        ("C1x", "C2x"),
+        ("C2x", "C3x"),
+        ("C3x", "N1x"),
+        ("N1x", "C4x"),
+        ("C4x", "C5x"),
+        ("C5x", "C1x"),
+        ("C1x", "H1x"),
+        ("C2x", "H2x"),
+        ("C3x", "H3x"),
+        ("C4x", "H4x"),
+        ("C5x", "H5x"),
+        ("C1x", "C2x"),
+        ("C1x", "C3x"),
+        ("C1x", "C2x"),
+        ("C1x", "C3x"),
+    ]
+
+    actual_bonds = [(bond.atom1.name, bond.atom2.name) for bond in top.bonds()]
+    assert actual_bonds == expected_bonds
+
 
 def test_setup_fn():
     system = absolv.config.System(