Add method to Compound that checks for overlapping particles

environment-dev.yml

@@ -10,6 +10,7 @@ dependencies:
   - boltons
   - lark>=1.2
   - lxml
+  - freud
   - intermol
   - mdtraj
   - pydantic>=2

mbuild/compound.py

@@ -969,9 +969,17 @@ def available_ports(self):
 
         return [p for p in self.labels.values() if isinstance(p, Port) and not p.used]
 
-    def direct_bonds(self):
+    def direct_bonds(self, graph_depth=1):
         """Return a list of particles that this particle bonds to.
 
+        Parameters
+        ----------
+        graph_depth : int, default=1
+            Determines how many subsequent bonded neighbors to count.
+            A value of 1 returns only paricles this particle is directly bonded to.
+            A value of 2 returns direct bonded neighbors,
+            plus their direct bonded neighbors.
+
         Returns
         -------
         List of mb.Compound
@@ -981,13 +989,23 @@ def direct_bonds(self):
         bond_graph.edges_iter : Iterations over all edges in a BondGraph
         Compound.n_direct_bonds : Returns the number of bonds a particle contains
         """
+        if graph_depth <= 0 or not isinstance(graph_depth, int):
+            raise ValueError("`graph_depth` must be an integer >= 1.")
         if list(self.particles()) != [self]:
             raise MBuildError(
                 "The direct_bonds method can only "
                 "be used on compounds at the bottom of their hierarchy."
             )
-        for b1, b2 in self.root.bond_graph.edges(self):
-            yield b2
+        if not self.parent:
+            return None
+        # Get all nodes within n edges (graph_depth=n) using BFS
+        top_ancestor = [i for i in self.ancestors()][-1]
+        neighbor_dict = nx.single_source_shortest_path_length(
+            top_ancestor.bond_graph, self, cutoff=graph_depth
+        )
+        # Exclude the source node itself
+        all_neighbors = {n for n, depth in neighbor_dict.items() if depth > 0}
+        return all_neighbors
 
     def bonds(self, return_bond_order=False):
         """Return all bonds in the Compound and sub-Compounds.
@@ -1036,7 +1054,10 @@ def n_direct_bonds(self):
                 "The direct_bonds method can only "
                 "be used on compounds at the bottom of their hierarchy."
             )
-        return sum(1 for _ in self.direct_bonds())
+        if self.direct_bonds(graph_depth=1):
+            return len(self.direct_bonds(graph_depth=1))
+        else:
+            return 0
 
     @property
     def n_bonds(self):
@@ -1130,13 +1151,7 @@ def generate_bonds(self, name_a, name_b, dmin, dmax):
                     added_bonds.append(bond_tuple)
 
     @experimental_feature()
-    def freud_generate_bonds(
-        self,
-        name_a,
-        name_b,
-        dmin,
-        dmax,
-    ):
+    def freud_generate_bonds(self, name_a, name_b, dmin, dmax):
         """Add Bonds between all pairs of types a/b within [dmin, dmax].
 
         Parameters
@@ -1156,33 +1171,7 @@ def freud_generate_bonds(
 
         """
         freud = import_("freud")
-        if self.box is None:
-            box = self.get_boundingbox()
-        else:
-            box = self.box
-        moved_positions = self.xyz - np.array([box.Lx / 2, box.Ly / 2, box.Lz / 2])
-
-        # quadruple box lengths for non-periodic dimensions
-        # since freud boxes are centered at the origin, extend box
-        #   lengths 2x in the positive and negative direction
-
-        # we are periodic in all directions, no need to change anything
-        if all(self.periodicity):
-            freud_box = freud.box.Box.from_matrix(box.vectors.T)
-        # not periodic in some dimensions, lets make them pseudo-periodic
-        else:
-            tmp_lengths = [length for length in box.lengths]
-            max_tmp_length = max(tmp_lengths)
-            for i, is_periodic in enumerate(self.periodicity):
-                if is_periodic:
-                    continue
-                else:
-                    tmp_lengths[i] = tmp_lengths[i] + 4 * max_tmp_length
-            tmp_box = Box.from_lengths_angles(lengths=tmp_lengths, angles=box.angles)
-            freud_box = freud.box.Box.from_matrix(tmp_box.vectors.T)
-
-        freud_box.periodic = (True, True, True)
-
+        moved_positions, freud_box = self._to_freud()
         a_indices = []
         b_indices = []
         for i, part in enumerate(self.particles()):
@@ -1209,11 +1198,7 @@ def freud_generate_bonds(
 
         nlist = aq.query(
             moved_positions[a_indices],
-            dict(
-                r_min=dmin,
-                r_max=dmax,
-                exclude_ii=exclude_ii,
-            ),
+            dict(r_min=dmin, r_max=dmax, exclude_ii=exclude_ii),
         ).toNeighborList()
 
         part_list = [part for part in self.particles(include_ports=False)]
@@ -1457,6 +1442,42 @@ def is_independent(self):
                         return False
             return True
 
+    def check_for_overlap(self, excluded_bond_depth, minimum_distance=0.10):
+        """Check if a compound contains overlapping particles.
+
+        Parameters:
+        -----------
+        excluded_bond_depth : int, required
+            The depth of bonded neighbors to exclude from overlap check.
+            see Compound.direct_bonds()
+        minimum_distance : float, default=0.10
+            Distance in nanometers used as the distance threshold in
+            determining if a pair of particles overlap.
+        """
+        if excluded_bond_depth < 0 or not isinstance(excluded_bond_depth, int):
+            raise ValueError("`excluded_bond_depth must be an integer >= 0.")
+        freud = import_("freud")
+        moved_positions, freud_box = self._to_freud()
+        aq = freud.locality.AABBQuery(freud_box, moved_positions)
+        aq_query = aq.query(
+            query_points=moved_positions,
+            query_args=dict(r_min=0.0, r_max=minimum_distance, exclude_ii=True),
+        )
+        nlist = aq_query.toNeighborList()
+        # nlist contains each pair twice, get the set
+        pairs_set = set([tuple(sorted((i, j))) for i, j in nlist])
+        all_particles = [p for p in self.particles()]
+        overlapping_particles = []
+        for i, j in pairs_set:
+            # Exclude bonded neighbors that are within min distance
+            if excluded_bond_depth > 0:
+                i_bonds = all_particles[i].direct_bonds(graph_depth=excluded_bond_depth)
+                if all_particles[j] not in i_bonds:
+                    overlapping_particles.append([i, j])
+            else:  # Don't exclude bonded neighbors
+                overlapping_particles.append([i, j])
+        return overlapping_particles
+
     def get_boundingbox(self, pad_box=None):
         """Compute the bounding box of the compound.
 
@@ -3246,6 +3267,37 @@ def get_smiles(self):
         smiles = pybel_cmp.write().split()[0]
         return smiles
 
+    def _to_freud(self):
+        """Convert a compound to a freud system (freud box and shifted coordinates)."""
+        freud = import_("freud")
+        if self.box is None:
+            box = self.get_boundingbox()
+        else:
+            box = self.box
+        moved_positions = self.xyz - np.array([box.Lx / 2, box.Ly / 2, box.Lz / 2])
+
+        # quadruple box lengths for non-periodic dimensions
+        # since freud boxes are centered at the origin, extend box
+        #   lengths 2x in the positive and negative direction
+
+        # we are periodic in all directions, no need to change anything
+        if all(self.periodicity):
+            freud_box = freud.box.Box.from_matrix(box.vectors.T)
+        # not periodic in some dimensions, lets make them pseudo-periodic
+        else:
+            tmp_lengths = [length for length in box.lengths]
+            max_tmp_length = max(tmp_lengths)
+            for i, is_periodic in enumerate(self.periodicity):
+                if is_periodic:
+                    continue
+                else:
+                    tmp_lengths[i] = tmp_lengths[i] + 4 * max_tmp_length
+            tmp_box = Box.from_lengths_angles(lengths=tmp_lengths, angles=box.angles)
+            freud_box = freud.box.Box.from_matrix(tmp_box.vectors.T)
+
+        freud_box.periodic = (True, True, True)
+        return moved_positions, freud_box
+
     def __getitem__(self, selection):
         """Get item from Compound."""
         if isinstance(selection, int):

mbuild/tests/test_compound.py

@@ -94,11 +94,29 @@ def test_direct_bonds_parent(self, methane):
         with pytest.raises(MBuildError):
             [i for i in methane.direct_bonds()]
 
+    def test_direct_bonds_no_bonds(self):
+        compound = mb.Compound(name="A")
+        assert compound.direct_bonds() is None
+
+    def test_direct_bonds_depth(self, ethane):
+        bonded_particles = ethane[0].direct_bonds(graph_depth=1)
+        assert len(bonded_particles) == 4
+        bonded_particles = ethane[0].direct_bonds(graph_depth=2)
+        assert len(bonded_particles) == ethane.n_particles - 1
+
     def test_direct_bonds(self, methane):
-        bond_particles = [i for i in methane[0].direct_bonds()]
+        bond_particles = methane[0].direct_bonds()
         for H in methane.particles_by_name("H"):
             assert H in bond_particles
 
+    def test_direct_bonds_bad_inputs(self, methane):
+        with pytest.raises(ValueError):
+            methane[0].direct_bonds(graph_depth=0)
+        with pytest.raises(ValueError):
+            methane[0].direct_bonds(graph_depth=-2)
+        with pytest.raises(ValueError):
+            methane[0].direct_bonds(graph_depth=2.4)
+
     def test_bond_order(self, methane):
         # test the default behavior
         bonds = [bond for bond in methane.bonds(return_bond_order=True)]
@@ -436,6 +454,64 @@ def test_save_resnames_single(self, c3, n4):
         assert struct.residues[0].number == 1
         assert struct.residues[1].number == 2
 
+    def test_check_for_overlap(self, ethane):
+        assert not ethane.check_for_overlap(excluded_bond_depth=0)
+        overlap = ethane.check_for_overlap(excluded_bond_depth=0, minimum_distance=0.5)
+        assert len(overlap) == ethane.n_particles * (ethane.n_particles - 1) / 2
+        assert not ethane.check_for_overlap(
+            excluded_bond_depth=0, minimum_distance=0.001
+        )
+        compound = mb.Compound([ethane, mb.clone(ethane)])
+        overlap = compound.check_for_overlap(
+            excluded_bond_depth=0, minimum_distance=0.5
+        )
+        assert len(overlap) == compound.n_particles * (compound.n_particles - 1) / 2
+
+    def test_check_overlap_bond_depth(self):
+        methane = mb.load("C", smiles=True)
+        overlap = methane.check_for_overlap(
+            excluded_bond_depth=0, minimum_distance=0.13
+        )
+        assert len(overlap) == 4
+        assert not methane.check_for_overlap(
+            excluded_bond_depth=1, minimum_distance=0.13
+        )
+        methane2 = mb.clone(methane)
+        methane2.translate(np.array([5, 0, 0]))
+        compound = mb.Compound([methane, methane2])
+        assert not compound.check_for_overlap(
+            excluded_bond_depth=2, minimum_distance=0.5
+        )
+
+    def test_check_overlap_chain(self):
+        positions = [
+            [0.10, 0.0, 0.0],
+            [0.20, 0.0, 0.0],
+            [0.30, 0.0, 0.0],
+            [0.30, 0.10, 0.0],
+            [0.20, 0.10, 0.0],
+            [0.10, 0.10, 0.0],
+        ]
+        comp = mb.Compound()
+        last_bead = None
+        for pos in positions:
+            bead = mb.Compound(pos=pos, name="A")
+            comp.add(bead)
+            if last_bead:
+                comp.add_bond([bead, last_bead])
+            last_bead = bead
+
+        overlap = comp.check_for_overlap(excluded_bond_depth=4, minimum_distance=0.11)
+        assert len(overlap) == 1
+        assert overlap[0] == [0, 5]
+        assert not comp.check_for_overlap(excluded_bond_depth=5, minimum_distance=0.11)
+
+    def test_check_for_overlap_bad_inputs(self, ethane):
+        with pytest.raises(ValueError):
+            ethane.check_for_overlap(excluded_bond_depth=-2, minimum_distance=0.5)
+        with pytest.raises(ValueError):
+            ethane.check_for_overlap(excluded_bond_depth=2.4, minimum_distance=0.5)
+
     def test_clone_with_box(self, ethane):
         ethane.box = ethane.get_boundingbox()
         ethane.periodicity = (True, True, False)