diff --git a/MDRestraintsGenerator/search.py b/MDRestraintsGenerator/search.py index 13bc1dc..64df90d 100644 --- a/MDRestraintsGenerator/search.py +++ b/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, diff --git a/MDRestraintsGenerator/tests/test_search.py b/MDRestraintsGenerator/tests/test_search.py index 9c336aa..8583b41 100644 --- a/MDRestraintsGenerator/tests/test_search.py +++ b/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"""