From 617079bd403d1a0157a1783ca7f4c60557f2e0a3 Mon Sep 17 00:00:00 2001 From: Justin Gilmer Date: Thu, 30 Sep 2021 16:30:12 -0500 Subject: [PATCH] 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 * Update KDTree name * Fix import * Update mbuild/lib/recipes/tiled_compound.py Co-authored-by: Co Quach Co-authored-by: Umesh Timalsina --- mbuild/compound.py | 18 ++-- mbuild/lib/recipes/tiled_compound.py | 23 ++--- mbuild/periodic_kdtree.py | 131 ++++++++++++++++++--------- 3 files changed, 108 insertions(+), 64 deletions(-) diff --git a/mbuild/compound.py b/mbuild/compound.py index fc8543aaf..37adb260a 100755 --- a/mbuild/compound.py +++ b/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 ) diff --git a/mbuild/lib/recipes/tiled_compound.py b/mbuild/lib/recipes/tiled_compound.py index d9d2748e0..0d8729f41 100755 --- a/mbuild/lib/recipes/tiled_compound.py +++ b/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] diff --git a/mbuild/periodic_kdtree.py b/mbuild/periodic_kdtree.py index 0060f48e4..94281e218 100755 --- a/mbuild/periodic_kdtree.py +++ b/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,33 +107,32 @@ 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( @@ -145,7 +140,57 @@ def __init__(self, data, leafsize=10, box=None, bounds=None): ) # 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: