Update PeriodicKDTree to support box overhaul (#952)

* test

* Update PeriodicKDTree to support box overhaul

There were certain situations where the periodic KDTree was not updated
properly to support the updates to the periodicity and the box overhaul.

There is still one test case failing for `TiledCompound`, but this PR
helps a lot with generating bonds. Once the `TiledCompound` error is
fixed, this should be good to merge.

* Fix periodic kd tree, add in kwargs for k value in query

* Update KDTree with a Compound creation method

Based on discussions with @umesh-timalsina and @daico007, the
PeriodicCKDTree has been updated with a class method to generate a
KDTree by passing in a compound.

This makes creation of kdtrees much easier and simpler for use in
mBuild.

* Update mbuild/periodic_kdtree.py

* Apply suggestions from code review

Co-authored-by: Umesh Timalsina <[email protected]>

* Update KDTree name

* Fix import

* Update mbuild/lib/recipes/tiled_compound.py

Co-authored-by: Co Quach <[email protected]>
Co-authored-by: Umesh Timalsina <[email protected]>

mbuild/compound.py

@@ -19,7 +19,7 @@
 from mbuild.box import Box
 from mbuild.coordinate_transform import _rotate, _translate
 from mbuild.exceptions import MBuildError
-from mbuild.periodic_kdtree import PeriodicCKDTree
+from mbuild.periodic_kdtree import PeriodicKDTree
 from mbuild.utils.exceptions import RemovedFuncError
 from mbuild.utils.io import import_, run_from_ipython
 from mbuild.utils.jsutils import overwrite_nglview_default
@@ -983,7 +983,11 @@ def generate_bonds(self, name_a, name_b, dmin, dmax):
         dmax : float
             The maximum distance between Particles for considering a bond
         """
-        particle_kdtree = PeriodicCKDTree(data=self.xyz, box=self.box)
+        if self.box is None:
+            self.box = self.get_boundingbox()
+        particle_kdtree = PeriodicKDTree.from_compound(
+            compound=self, leafsize=10
+        )
         particle_array = np.array(list(self.particles()))
         added_bonds = list()
         for p1 in self.particles_by_name(name_a):
@@ -1351,7 +1355,7 @@ def particles_in_range(
             Maximum distance from 'compound' to look for Particles
         max_particles : int, optional, default=20
             Maximum number of Particles to return
-        particle_kdtree : mb.PeriodicCKDTree, optional
+        particle_kdtree : mb.PeriodicKDTree, optional
             KD-tree for looking up nearest neighbors. If not provided, a KD-
             tree will be generated from all Particles in self
         particle_array : np.ndarray, shape=(n,), dtype=mb.Compound, optional
@@ -1365,11 +1369,13 @@ def particles_in_range(
 
         See Also
         --------
-        periodic_kdtree.PerioidicCKDTree : mBuild implementation of kd-trees
-        scipy.spatial.ckdtree : Further details on kd-trees
+        periodic_kdtree.PerioidicKDTree : mBuild implementation of kd-trees
+        scipy.spatial.kdtree : Further details on kd-trees
         """
+        if self.box is None:
+            self.box = self.get_boundingbox()
         if particle_kdtree is None:
-            particle_kdtree = PeriodicCKDTree(data=self.xyz, box=self.box)
+            particle_kdtree = PeriodicKDTree.from_compound(self, leafsize=10)
         _, idxs = particle_kdtree.query(
             compound.pos, k=max_particles, distance_upper_bound=dmax
         )

mbuild/lib/recipes/tiled_compound.py

@@ -7,7 +7,7 @@
 
 from mbuild import Box, Compound, Port, clone
 from mbuild.exceptions import MBuildError
-from mbuild.periodic_kdtree import PeriodicCKDTree
+from mbuild.periodic_kdtree import PeriodicKDTree
 
 
 class TiledCompound(Compound):
@@ -26,7 +26,7 @@ class TiledCompound(Compound):
         Descriptive string for the compound.
     """
 
-    def __init__(self, tile, n_tiles, name=None):
+    def __init__(self, tile, n_tiles, name=None, **kwargs):
         super(TiledCompound, self).__init__()
 
         n_tiles = np.asarray(n_tiles)
@@ -83,14 +83,6 @@ def __init__(self, tile, n_tiles, name=None):
                 continue
         dist_thresh = np.min(threshold_calc) / 2
 
-        # Create the bounds for the periodicKDtree, non-periodic dimensions are 0
-        bounds = [0, 0, 0]
-        length_array = np.asarray(tile.box.lengths)
-        for i, dim in enumerate(bounds):
-            if tile.periodicity[i]:
-                bounds[i] = self.box.lengths[i]
-            else:
-                continue
         # Bonds that were periodic in the original tile.
         periodic_bonds = set()
         for particle1, particle2 in tile.bonds():
@@ -99,7 +91,7 @@ def __init__(self, tile, n_tiles, name=None):
                 periodic_bonds.add((particle1.index, particle2.index))
 
         # Build a periodic kdtree of all particle positions.
-        self.particle_kdtree = PeriodicCKDTree(data=self.xyz, bounds=bounds)
+        self.particle_kdtree = PeriodicKDTree.from_compound(self, leafsize=10)
         all_particles = np.asarray(list(self.particles(include_ports=False)))
 
         # Store bonds to remove/add since we'll be iterating over all bonds.
@@ -116,10 +108,10 @@ def __init__(self, tile, n_tiles, name=None):
             if dist > dist_thresh:
                 bonds_to_remove.add((particle1, particle2))
                 image2 = self._find_particle_image(
-                    particle1, particle2, all_particles
+                    particle1, particle2, all_particles, **kwargs
                 )
                 image1 = self._find_particle_image(
-                    particle2, particle1, all_particles
+                    particle2, particle1, all_particles, **kwargs
                 )
 
                 if (image2, particle1) not in bonds_to_add:
@@ -150,9 +142,10 @@ def _hoist_ports(self, new_tile):
             if isinstance(port, Port):
                 self.add(port, containment=False)
 
-    def _find_particle_image(self, query, match, all_particles):
+    def _find_particle_image(self, query, match, all_particles, **kwargs):
         """Find particle with the same index as match in a neighboring tile."""
-        _, idxs = self.particle_kdtree.query(query.pos, k=10)
+        k = kwargs.get("k", 10)
+        _, idxs = self.particle_kdtree.query(query.pos, k=k)
 
         neighbors = all_particles[idxs]
 

mbuild/periodic_kdtree.py

@@ -38,6 +38,8 @@
 import numpy as np
 from scipy.spatial import KDTree
 
+import mbuild as mb
+
 
 def _gen_relevant_images(x, bounds, distance_upper_bound):
     """Map x onto canonical unit cell and produce mirror images."""
@@ -74,10 +76,10 @@ def _gen_relevant_images(x, bounds, distance_upper_bound):
     return xs
 
 
-class PeriodicCKDTree(KDTree):
+class PeriodicKDTree(KDTree):
     """Cython kd-tree for nearest-neighbor lookup with periodic boundaries.
 
-    See scipy.spatial.ckdtree for details on kd-trees.
+    See scipy.spatial.kdtree for details on kd-trees.
 
     Searches with periodic boundaries are implemented by mapping all initial
     data points to one canonical periodic image, building an ordinary kd-tree
@@ -86,21 +88,15 @@ class PeriodicCKDTree(KDTree):
 
     Parameters
     ----------
-    box : mbuild.Box
-        The box object containing the periodicity of the system.
-        A zero value for one of the box lengths mean that the
-        space is not periodic in that dimension.
-    bounds : array_like, shape (k,)
-        Deprecated. Will be removed in version 0.11. Please use
-        the box argument instead.
+    data : array-like, shape (n,m), required
+        The n data points of dimension m to be indexed. This array is not copied
+        unless this is necessary to produce a contiguous array of doubles, and
+        so modifying this data will result in bogus results.
+    bounds : array_like, shape (m,), required
         Size of the periodic box along each spatial dimension.  A
         negative or zero size for dimension k means that space is not
         periodic along k.
-    data : array-like, shape (n,m)
-        The n data points of dimension mto be indexed. This array is not copied
-        unless this is necessary to produce a contiguous array of doubles, and
-        so modifying this data will result in bogus results.
-    leafsize : positive integer
+    leafsize : positive integer, optional, default=10
         The number of points at which the algorithm switches over to
         brute-force.
 
@@ -111,41 +107,90 @@ class PeriodicCKDTree(KDTree):
     point and a data point to half the smallest box dimension.
     """
 
-    def __init__(self, data, leafsize=10, box=None, bounds=None):
-        if bounds is not None and box is not None:
-            raise ValueError(
-                "Only one of 'bounds' and 'box' may be specified. "
-                "Since 'bounds' is deprecated, please use 'box'."
-            )
-        # Map all points to canonical periodic image.
-        if box is None:
-            if bounds is None:
-                bounds = np.array([0.0, 0.0, 0.0])
-        elif np.allclose(box.angles, 90.0):
-            bounds = box.lengths
-        else:
-            raise NotImplementedError(
-                "Periodic KCDTree search only implemented"
-                "for orthorhombic periodic boundaries"
-            )
+    def __init__(self, bounds, data, leafsize=10):
+        """Construct Cython kd-tree for nearest-neighbor lookup with periodic boundaries.
+
+        Parameters
+        ----------
+        bounds : array_like, shape (m,)
+            Size of the periodic box along each spatial dimension.  A
+            negative or zero size for dimension k means that space is not
+            periodic along k.
+        data : array-like, shape (n,m)
+            The n data points of dimension m to be indexed. This array is
+            not copied unless this is necessary to produce a contiguous
+            array of doubles, and so modifying this data will result in
+            bogus results.
+        leafsize : positive integer
+            The number of points at which the algorithm switches over to
+            brute-force.
+        """
+        # Map all points to canonical periodic image
         self.bounds = np.array(bounds)
         self.real_data = np.asarray(data)
-
-        wrapped_data = self.real_data
-
-        for i, row in enumerate(self.real_data):
-            for j, coord in enumerate(row):
-                if bounds[j] > 0.0:
-                    wrap = np.floor(self.real_data[i, j] / bounds[j])
-                    wrapped_data[i, j] = self.real_data[i, j] - wrap * bounds[j]
+        wrapped_data = self.real_data - np.where(
+            self.bounds > 0.0,
+            (np.floor(self.real_data / self.bounds) * self.bounds),
+            0.0,
+        )
 
         # Calculate maximum distance_upper_bound
         self.max_distance_upper_bound = np.min(
             np.where(self.bounds > 0, 0.5 * self.bounds, np.inf)
         )
 
         # Set up underlying kd-tree
-        super(PeriodicCKDTree, self).__init__(wrapped_data, leafsize)
+        super(PeriodicKDTree, self).__init__(wrapped_data, leafsize)
+
+    @classmethod
+    def from_compound(cls, compound, leafsize=10):
+        """Create a PeriodicKDTree from a compound.
+
+        See scipy.spatial.kdtree for details on kd-trees.
+
+        Searches with periodic boundaries are implemented by mapping all initial
+        data points to one canonical periodic image, building an ordinary kd-tree
+        with these points, then querying this kd-tree multiple times, if necessary,
+        with all the relevant periodic images of the query point.
+
+        Parameters
+        ----------
+        compound : mb.Compound, required
+            The mbuild Compound to gather periodicity and box information for the
+            PeriodicKDTree.
+        leafsize : positive integer
+            The number of points at which the algorithm switches over to
+            brute-force.
+
+        Note
+        ----
+        To ensure that no two distinct images of the same point appear in the
+        results, it is essential to restrict the maximum distance between a query
+        point and a data point to half the smallest box dimension.
+        """
+        if not isinstance(compound, mb.Compound):
+            raise TypeError(
+                f"Incorrect type of compound. Provided compound of type {type(compound)}. Expected mbuild.Compound"
+            )
+        if not isinstance(compound.box, mb.Box):
+            raise TypeError(
+                f"Incorrect type of box. Was provided box of type {type(compound.box)}. Expected mbuild.Box"
+            )
+        if not np.allclose(compound.box.angles, 90.0):
+            raise NotImplementedError(
+                "Periodic KDTree search only implemented"
+                "for orthorhombic periodic boundaries"
+            )
+        # Map all points to canonical periodic image.
+        compound_bounds = []
+        for box_len, period in zip(compound.box.lengths, compound.periodicity):
+            if period:
+                compound_bounds.append(box_len)
+            else:
+                compound_bounds.append(0)
+
+        # Set up underlying kd-tree
+        return cls(bounds=compound_bounds, data=compound.xyz, leafsize=leafsize)
 
     # Ideally, KDTree and cKDTree would expose identical query and __query
     # interfaces.  But they don't, and cKDTree.__query is also inaccessible
@@ -167,7 +212,7 @@ def __query(self, x, k=1, eps=0, p=2, distance_upper_bound=np.inf):
         for real_x in _gen_relevant_images(
             x, self.bounds, distance_upper_bound
         ):
-            d, i = super(PeriodicCKDTree, self).query(
+            d, i = super(PeriodicKDTree, self).query(
                 real_x, k, eps, p, distance_upper_bound
             )
             if k > 1:
@@ -298,7 +343,7 @@ def __query_ball_point(self, x, r, p=2.0, eps=0):
         results = []
         for real_x in _gen_relevant_images(x, self.bounds, r):
             results.extend(
-                super(PeriodicCKDTree, self).query_ball_point(real_x, r, p, eps)
+                super(PeriodicKDTree, self).query_ball_point(real_x, r, p, eps)
             )
         return results
 
@@ -330,7 +375,7 @@ def query_ball_point(self, x, r, p=2.0, eps=0):
         -----
         If you have many points whose neighbors you want to find, you may
         save substantial amounts of time by putting them in a
-        PeriodicCKDTree and using query_ball_tree.
+        PeriodicKDTree and using query_ball_tree.
         """
         x = np.asarray(x).astype(float)
         if x.shape[-1] != self.m: