Improved performance of bond look up

mbuild/compound.py

@@ -12,6 +12,7 @@
 
 import ele
 import numpy as np
+from boltons.setutils import IndexedSet
 from ele.element import Element, element_from_name, element_from_symbol
 from ele.exceptions import ElementError
 
@@ -972,7 +973,7 @@ def direct_bonds(self):
         for i in self.root.bond_graph._adj[self]:
             yield i
 
-    def bonds(self):
+    def _bonds(self):
         """Return all bonds in the Compound and sub-Compounds.
 
         Yields
@@ -995,6 +996,78 @@ def bonds(self):
         else:
             return iter(())
 
+    def bonds(self, use_connected=True):
+        """Return all bonds in the Compound and sub-Compounds. By default,
+        this uses the connected components to improve performance when possible.
+
+        Yields
+        ------
+        tuple of mb.Compound
+            The next bond in the Compound
+
+        Parameters
+        ----------
+        use_connected : bool, optional, default=True
+            Uses the connected_components information in the bond_graph to
+            retrieve bonded information for sub-Compounds. In most cases,
+            this will be substantially faster.
+
+        See Also
+        --------
+        bond_graph.edges_iter : Iterates over all edges in a BondGraph
+        Compound.n_bonds : Returns the total number of bonds in the Compound and sub-Compounds
+        Compound._bonds : Return all bonds in the Compound and sub-Compounds
+        """
+
+        if use_connected:
+            if self.root.bond_graph:
+                connected_subgraph = self.root.bond_graph.connected_components()
+                bonds = []
+                molecule_tags = {}
+                for molecule in connected_subgraph:
+                    # if molecule contains a single particle it will not contain any bonds and can be skipped
+                    if len(molecule) > 1:
+                        # if the Compound of interest is not an ancestor of molecule, move to the next molecule
+                        ancestors = [ann for ann in molecule[0].ancestors()]
+                        if self in ancestors:
+                            if len(ancestors) > 1:
+                                # this will set ancestors[0] to the lowest level Compound that
+                                # contains all the connected components in molecule
+                                ancestors = IndexedSet(molecule[0].ancestors())
+                                for particle in molecule[1:]:
+                                    ancestors = ancestors.intersection(
+                                        IndexedSet(particle.ancestors())
+                                    )
+                            # if the compound we are interested in is the lowest level ancestor or root, find the bonds
+                            if ancestors[0] == self or self == self.root:
+                                for b in self.root.bond_graph.subgraph(
+                                    molecule
+                                ).edges_iter():
+                                    bonds.append(b)
+                            # if the Compound of interest is not at a lower level than the lowest level container (i.e. ancestors[0])
+                            # then ancestors[0] is a sub-Compound of the Compound of interest
+                            elif self not in [
+                                part for part in ancestors[0].successors()
+                            ]:
+                                for b in self.root.bond_graph.subgraph(
+                                    molecule
+                                ).edges_iter():
+                                    bonds.append(b)
+
+                # If we did not add any bonds to the list that does not mean there are no bonds in the Compound of interest.
+                # It may mean that we are looking at a Compound with bonds to other Compounds
+                # (e.g., a CH3 in ethane; ethane would be the lowest level connected component stored in ancestors[0])
+                # As such, we will just return the original bond routine to retrieve the relevant info
+                if len(bonds) == 0:
+                    return self._bonds()
+                else:
+                    return iter(bonds)
+            else:
+                return self._bonds()
+
+        else:
+            return self._bonds()
+
     @property
     def n_direct_bonds(self):
         """Return the number of bonds a particle is directly involved in.

mbuild/tests/test_compound.py

@@ -138,6 +138,32 @@ def test_direct_bonds_cloning(self, ethane):
         for p1, p2 in zip(ethane.particles(), ethane_clone.particles()):
             assert p1.n_direct_bonds == p2.n_direct_bonds
 
+    def test_bond_speedup(self, ethane):
+        # does the efficient code give the same bonds as the slower code?
+        ethane_clone = mb.clone(ethane)
+        assert sum([1 for _ in ethane_clone.bonds()]) == sum(
+            [1 for _ in ethane_clone._bonds()]
+        )
+
+        # with added hierarchy, do we still get the same answers
+        ethane_clone2 = mb.clone(ethane)
+        system = mb.Compound()
+        system.add(ethane_clone)
+        system.add(ethane_clone2)
+
+        assert sum([1 for _ in system.bonds()]) == sum(
+            [1 for _ in system._bonds()]
+        )
+        assert sum([1 for _ in system.bonds()]) == sum(
+            [1 for _ in system.bonds(use_connected=False)]
+        )
+
+        # check we can traverse the hierarchy and get the correct number of bonds
+        ethane_clone3 = mb.clone(ethane)
+        assert sum([1 for _ in ethane_clone3.bonds()]) == sum(
+            [1 for _ in ethane_clone2.bonds()]
+        )
+
     def test_load_protein(self):
         # Testing the loading function with complicated protein,
         # The protein file is taken from RCSB protein data bank
@@ -2154,6 +2180,6 @@ def test_xyz_setter(self):
     def test_ordered_bonds(self):
         ethane = mb.load("CC", smiles=True)
         ethane2 = mb.load("CC", smiles=True)
-        for bond2, bond in zip(ethane2.bonds(), ethane.bonds()):
+        for bond2, bond in zip(ethane2._bonds(), ethane._bonds()):
             assert bond2[0].name == bond[0].name
             assert all(bond2[0].pos == bond[0].pos)