Hack in bonded interactions

src/relentless/simulate/hoomd.py

@@ -64,6 +64,8 @@ def _call_v3(self, sim):
         # parse masses by type
         snap = sim["engine"]["_hoomd"].state.get_snapshot()
         sim.masses = self._get_masses_from_snapshot(sim, snap)
+        sim[self]["_bonds"] = self._get_bonds_from_snapshot(sim, snap)
+        sim[self]["_angles"] = self._get_angles_from_snapshot(sim, snap)
         self._assert_dimension_safe(sim, snap)
         # create the potentials, defer attaching until later
         neighbor_list = hoomd.md.nlist.Tree(buffer=sim.potentials.pair.neighbor_buffer)
@@ -152,6 +154,12 @@ def _get_masses_from_snapshot(self, sim, snap):
         masses.update(masses_)
         return masses
 
+    def _get_bonds_from_snapshot(self, sim, snap):
+        return snap.bonds.group
+
+    def _get_angles_from_snapshot(self, sim, snap):
+        return snap.angles.group
+
     def _assert_dimension_safe(self, sim, snap):
         if sim.dimension == 3:
             dim_safe = True
@@ -987,6 +995,8 @@ def _pre_run_v3(self, sim, sim_op):
             system=sim["engine"]["_hoomd"].state,
             rdf_params=self._get_rdf_params(sim),
             constraints=self._get_constrained_quantities(sim, sim_op),
+            bonds=sim[sim.initializer]["_bonds"],
+            angles=sim[sim.initializer]["_angles"],
         )
         sim[self]["_hoomd_thermo_callback"] = hoomd.write.CustomWriter(
             trigger=self.every,
@@ -1178,6 +1188,8 @@ def __init__(
             system,
             rdf_params=None,
             constraints=None,
+            bonds=None,
+            angles=None,
         ):
             if dimension not in (2, 3):
                 raise ValueError("Only 2 or 3 dimensions supported")
@@ -1188,6 +1200,8 @@ def __init__(
             self.system = system
             self.rdf_params = rdf_params
             self.constraints = constraints if constraints is not None else {}
+            self.bonds = bonds
+            self.angles = angles
 
             # this method handles all the initialization
             self.reset()
@@ -1282,6 +1296,40 @@ def act(self, timestep):
                             ),
                         ).toNeighborList()
 
+                        # filter bonds from the neighbor list if they are present
+                        # bond exclusions apply regardless of order, so
+                        # consider both (i,j) and (j,i) permutations
+                        if snap.bonds is not None and len(neighbors[:]) > 0:
+                            bonds = numpy.vstack(
+                                [self.bonds, numpy.flip(self.bonds, axis=1)],
+                            )
+                            # list intersect using Cantor Pairing Function (pi):
+                            # https://en.wikipedia.org/wiki/Pairing_function
+                            pi_bond = (bonds[:, 0] + bonds[:, 1]) * (
+                                bonds[:, 0] + bonds[:, 1] + 1
+                            ) / 2 + bonds[:, 1]
+                            pi_neighbor = (neighbors[:, 0] + neighbors[:, 1]) * (
+                                neighbors[:, 0] + neighbors[:, 1] + 1
+                            ) / 2 + neighbors[:, 1]
+                            bond_exclusion_filter = ~numpy.isin(pi_neighbor, pi_bond)
+
+                            neighbors.filter(bond_exclusion_filter)
+
+                        if snap.angles is not None and len(neighbors[:]) > 0:
+                            angles = numpy.vstack(
+                                [self.angles, numpy.flip(self.angles, axis=1)],
+                            )
+                            # list intersect using Cantor Pairing Function (pi):
+                            # https://en.wikipedia.org/wiki/Pairing_function
+                            pi_angle = (angles[:, 0] + angles[:, 2]) * (
+                                angles[:, 0] + angles[:, 2] + 1
+                            ) / 2 + angles[:, 2]
+                            pi_neighbor = (neighbors[:, 0] + neighbors[:, 1]) * (
+                                neighbors[:, 0] + neighbors[:, 1] + 1
+                            ) / 2 + neighbors[:, 1]
+                            angle_exclusion_filter = ~numpy.isin(pi_neighbor, pi_angle)
+
+                            neighbors.filter(angle_exclusion_filter)
                         for i in self.types:
                             for j in self.types:
                                 filter_ij = numpy.logical_and(
@@ -1294,6 +1342,7 @@ def act(self, timestep):
                                     range=(0, self.rdf_params["stop"]),
                                 )
                                 self._rdf_counts[i, j] += counts
+
             self.num_samples += 1
             if compute_rdf:
                 self.num_rdf_samples += 1