Add methods to retrieve all bonds, internal bonds, and external bonds to Topology.Residue

wrappers/python/simtk/openmm/app/topology.py

@@ -371,6 +371,21 @@ class Residue(object):
         """Iterate over all Atoms in the Residue."""
         return iter(self._atoms)
 
+    def bonds(self):
+        """Iterate over all Bonds involving any atom in this residue."""
+        bonds = [ bond for bond in residue.chain.topology.bonds() if ((bond[0] in self._atoms) or (bond[1] in self._atoms)) ]
+        return iter(bonds)
+
+    def internal_bonds(self):
+        """Iterate over all internal Bonds."""
+        bonds = [ bond for bond in residue.chain.topology.bonds() if ((bond[0] in self._atoms) and (bond[1] in self._atoms)) ]
+        return iter(bonds)
+
+    def external_bonds(self):
+        """Iterate over all Bonds to external atoms."""
+        bonds = [ bond for bond in residue.chain.topology.bonds() if ((bond[0] in self._atoms) != (bond[1] in self._atoms)) ]
+        return iter(bonds)
+
     def __len__(self):
         return len(self._atoms)
 

wrappers/python/tests/TestTopology.py

@@ -26,5 +26,26 @@ class TestTopology(unittest.TestCase):
         """Test getters for number of atoms, residues, chains."""
         self.check_pdbfile('systems/1T2Y.pdb', 271, 25, 1)
 
+    def test_residue_bonds(self):
+        """Test retrieving bonds for a residue produces expected results."""
+        # Create a test topology
+        # atom connectivity = A1-|-B1-B2-|-C1
+        topology = Topology()
+        chain = topology.addChain(id='A')
+        residue1 = topology.addResidue('AAA', chain)
+        residue2 = topology.addResidue('BBB', chain)
+        residue3 = topology.addResidue('CCC', chain)
+        atom_A1 = topology.addAtom('A1', element.carbon, residue1)
+        atom_B1 = topology.addAtom('B1', element.carbon, residue2)
+        atom_B2 = topology.addAtom('B2', element.carbon, residue2)
+        atom_C1 = topology.addAtom('C1', element.carbon, residue3)
+        # Check bonds
+        all_bonds = [ bond for bond in residue.bonds() ]
+        internal_bonds = [ bond for bond in residue.internal_bonds() ]
+        external_bonds = [ bond for bond in residue.external_bonds() ]
+        self.assertEqual(all_bonds, [ (atom_A1, atom_B1), (atom_B1, atom_B2), (atom_B2, atom_C1) ])
+        self.assertEqual(internal_bonds, [ (atom_B1, atom_B2) ])
+        self.assertEqual(external_bonds, [ (atom_A1, atom_B1), (atom_B2, atom_C1) ])
+
 if __name__ == '__main__':
     unittest.main()