Support terminal AA CN expansion (#31)

MDRestraintsGenerator/search.py

@@ -105,12 +105,24 @@ def _get_bonded_host_cn_atoms(atomgroup, anchor_ix):
     anchor_ag = atomgroup.select_atoms(f'index {anchor_ix}')
     resid = anchor_ag.atoms[0].resid
     c_atom_str = f"(backbone) and (name C) and (resid {resid})"
-    n_atom_str = f"(backbone) and (name N) and (resid {resid+1})"
-
     c_atom = atomgroup.select_atoms(c_atom_str)
+    c_atom_ix = c_atom.atoms[0].ix
+    n_atom_str = (f"(backbone) and (name N) and (resid {resid+1}) and "
+                  f"(bonded index {c_atom_ix})")
     n_atom = atomgroup.select_atoms(n_atom_str)
 
-    return c_atom.atoms[0].ix, n_atom.atoms[0].ix
+    if n_atom:
+        return c_atom.atoms[0].ix, n_atom.atoms[0].ix
+    else:
+        # Issue 30 - terminal residue
+        # reverse selection to go to residue-1 rather than residue+1
+        n_atom_str = f"(backbone) and (name N) and (resid {resid})"
+        n_atom = atomgroup.select_atoms(n_atom_str)
+        n_atom_ix = n_atom.atoms[0].ix
+        c_atom_str = (f"(backbone) and (name C) and (resid {resid-1}) and "
+                      f"(bonded index {n_atom_ix})")
+        c_atom = atomgroup.select_atoms(c_atom_str)
+        return n_atom.atoms[0].ix, c_atom.atoms[0].ix
 
 
 def _get_bonded_host_atoms(atomgroup, anchor_ix,

MDRestraintsGenerator/tests/test_search.py

@@ -90,6 +90,16 @@ def test_findhostatoms_names(u):
         assert u.atoms[atoms[2]].name == "N"
 
 
+@pytest.mark.parametrize('ix1, ix2, ix3', [
+    [2560, 2577, 2579], [2581, 2579, 2577]
+])
+def test_bonded_cn(u, ix1, ix2, ix3):
+    """Tests for getting bonded CN atoms"""
+    second_atom, third_atom = search._get_bonded_host_cn_atoms(u, ix1)
+    assert second_atom == ix2
+    assert third_atom == ix3
+
+
 def test_basic_bonded(u):
     """Basic test for getting a bonded atom"""