Eliminated MDAnalysis' FastNS, which had a catastrophic bug

requirements.txt

@@ -1 +1 @@
-MDAnalysis>=0.19.0
+

setup.py

@@ -13,7 +13,6 @@
         "scipy",
         "ase",
         "sklearn",
-        "sitator",
-        "MDAnalysis>=0.19.0"
+        "sitator"
       ],
       zip_safe = True)

surfator/StructureGroupAnalysis.py

@@ -4,9 +4,7 @@
 import numbers
 import itertools
 
-#from scipy.spatial import cKDTree
-from MDAnalysis.lib.nsgrid import FastNS
-from MDAnalysis.lib.mdamath import triclinic_box
+from scipy.spatial import cKDTree
 
 import ase
 import ase.geometry
@@ -63,7 +61,7 @@ class StructureGroupAnalysis(object):
             to the winning structure group. All compatible structure groups to the
             winner are available for assignment, but sites belonging to the
             winning structure group can be made more appealing. A value of 0
-            indicates no bias; a value of 1 ensures the choice of a winning
+            indicates no bias; a value of nearly 1 ensures the choice of a winning
             structure group site (assuming one exists within the cutoff distnace).
         - error_on_no_majority (bool): If True, a majority below `min_winner_percentage`
             will result in an error. If False, that agreement group will simply
@@ -95,12 +93,13 @@ def run(self,
             traj,
             cutoff,
             k_neighbor = 6,
+            kdtree_eps = 0.05,
             agreement_group_function = surfator.agreement_groups.all_atoms_agree,
             structure_group_compatability = None,
             return_assignments = False):
         """
         Args:
-            ref_sn (SiteNetwork): A `SiteNetwork` containing the sites to which
+            ref_sn (SiteNetwork): A ``SiteNetwork`` containing the sites to which
                 the mobile atoms will be assigned. Can contain multiple, possibly
                 exclusive, sets of sites called "structure groups." The site
                 attribute `STRUCTURE_GROUP_ATTRIBUTE` indicates, for each site,
@@ -137,6 +136,8 @@ def run(self,
                 labels for agreement groups might change from
                 frame to frame and be meaningless -- please check your
                 `agreement_group_function` to see if this makes sense.
+            kdtree_eps (float): The epsilon parameter for the ``scipy.spatial.cKDTree``
+                used to find nearest sites.
         Returns:
             a SiteTrajectory[, agreegrp_assignments, structgrp_assignments]
         """
@@ -184,29 +185,20 @@ def run(self,
         pbcc.wrap_points(wrapped_traj)
         wrapped_traj.shape = traj.shape
 
-        # Build set of reference sites with periodic copies
-        # images = list(itertools.product(range(-1, 2), repeat = 3))
-        # image_000 = images.index((0, 0, 0))
-        # ref_site_ids = np.tile(np.arange(ref_sn.n_sites), len(images))
-        # all_centers = np.tile(ref_sn.centers, (len(images), 1, 1))
-        # for image_i, image in enumerate(images):
-        #     all_centers[image_i] += np.dot(ref_sn.structure.cell.T, image)
-        # all_centers.shape = (-1, 3)
-
-        # kdtree = cKDTree(
-        #     data = all_centers,
-        # )
-
-        our_triclinic_cell = triclinic_box(
-            cell[0],
-            cell[1],
-            cell[2]
+        ref_full = ase.Atoms(
+            positions = ref_sn.centers,
+            cell = ref_sn.structure.cell,
+            pbc = ref_sn.structure.pbc,
         )
-        ref_fastns = FastNS(
-            cutoff = cutoff,
-            coords = ref_sn.centers,
-            box = our_triclinic_cell,
-            pbc = True
+        ref_full.set_tags(np.arange(len(ref_full)))
+        ref_full *= tuple(3 if p else 1 for p in ref_full.get_pbc())
+        # Recenter image 000:
+        ref_full.translate(-sum(ref_sn.structure.cell[i] * p for i, p in enumerate(ref_full.get_pbc())))
+        # KDTree gives index one too high so we need an extra element
+        ref_full_to_real = np.concatenate((ref_full.get_tags(), [-1]))
+
+        kdtree = cKDTree(
+            data = ref_full.positions,
         )
 
         # Will be passed to the agreement group function
@@ -262,10 +254,17 @@ def run(self,
             site_distance_to_atom.fill(np.inf)
 
             # Distances don't change between rounds or agreegroups:
-            neighbor_res = ref_fastns.search(mobile_struct.positions)
-            all_neighbor_dists = neighbor_res.get_distances()
+            all_neighbor_dists, all_neighbor_idexs = kdtree.query(
+                mobile_struct.positions,
+                k = k_neighbor,
+                eps = kdtree_eps,
+                distance_upper_bound = cutoff,
+            )
             all_neighbor_dists_copy = all_neighbor_dists.copy()
-            all_neighbor_idexs = neighbor_res.get_indices()
+            ishape_orig = all_neighbor_idexs.shape
+            all_neighbor_idexs.shape = (-1,)
+            all_neighbor_idexs = ref_full_to_real[all_neighbor_idexs]
+            all_neighbor_idexs.shape = ishape_orig
 
             # - (2) - In order, assign the agreegrps
             for agreegrp_i, agreegrp_mask in enumerate(agreegrp_masks):
@@ -321,11 +320,9 @@ def run(self,
                             neighbor_dist *= site_weights[neighbor_idex]
                             # Apply can_assign_to
                             neighbor_dist[~can_assign_to[neighbor_idex]] = np.inf
-                            # Strict distance cutoff:
-                            # Unneeded cause KDTree is strict
-                            # neighbor_dist[neighbor_dist > cutoff] = np.inf
                             nearest_neighbor = np.argmin(neighbor_dist)
                             dist_to_nn = neighbor_dist[nearest_neighbor]
+
                         if dist_to_nn < np.inf:
                             assign_to_site = neighbor_idex[nearest_neighbor]
                             nearest_neighbors[mob_i] = assign_to_site
@@ -342,6 +339,8 @@ def run(self,
                                     # This atom has already been displaced
                                     displaced_by_closer.append(mob_i)
                         else:
+                            # It could be NaN but shouldn't be
+                            assert not np.isnan(dist_to_nn)
                             # Can't assign this one
                             logger.warning("At frame %i couldn't assign mobile atom %i" % (frame_idex, mob_i))
                             nearest_neighbors[mob_i] = -1

surfator/agreement_groups/layers.py

@@ -2,9 +2,6 @@
 import numpy as np
 import math
 
-from MDAnalysis.lib.nsgrid import FastNS
-from MDAnalysis.lib.mdamath import triclinic_box
-
 from scipy.sparse.csgraph import connected_components
 
 from sitator.util import PBCCalculator
@@ -133,17 +130,13 @@ def agree_within_components_of_groups(groupfunc,
             for them to be considered part of the same deposit.
     """
 
-    pbcc, our_triclinic_cell, connmat, newtags, layer_mask = None, None, None, None, None
+    pbcc, dmat, connmat, newtags, layer_mask = None, None, None, None, None
     def func(atoms, **kwargs):
-        nonlocal pbcc, our_triclinic_cell, connmat, newtags, layer_mask
+        nonlocal pbcc, dmat, connmat, newtags, layer_mask
         # preallocate buffers
         if pbcc is None:
             pbcc = PBCCalculator(atoms.cell)
-            our_triclinic_cell = triclinic_box(
-                atoms.cell[0],
-                atoms.cell[1],
-                atoms.cell[2]
-            )
+            dmat = np.empty(shape = (len(atoms), len(atoms)))
             connmat = np.empty(shape = (len(atoms), len(atoms)), dtype = np.bool)
             newtags = np.empty(shape = len(atoms), dtype = np.int)
             layer_mask = np.empty(shape = len(atoms), dtype = np.bool)
@@ -153,18 +146,8 @@ def func(atoms, **kwargs):
         layers.sort()
         newtags.fill(-1)
 
-        ns = FastNS(
-            cutoff = cutoff,
-            coords = atoms.positions,
-            box = our_triclinic_cell,
-            pbc = True
-        )
-        nsres = ns.self_search()
-        connmat.fill(False)
-        neighbor_idexes = nsres.get_indices()
-        for i in range(len(neighbor_idexes)):
-            # Don't need symmetry cause use directed = False later
-            connmat[i, neighbor_idexes[i]] = True
+        pbcc.pairwise_distances(atoms.positions, out = dmat)
+        np.less_equal(dmat, cutoff, out = connmat)
 
         agreegrp_conns = []
         nexttag = 0

surfator/analysis/coord.py

@@ -1,9 +1,6 @@
 
 import numpy as np
 
-from MDAnalysis.lib.nsgrid import FastNS
-from MDAnalysis.lib.mdamath import triclinic_box
-
 from sitator.util.progress import tqdm
 from sitator.util import PBCCalculator
 
@@ -42,6 +39,8 @@ def calculate_coord_numbers(traj, atoms, cutoff):
 
 # This algorithm is linear time, but is broken until this bug is fixed in
 # MDAnalysis: https://github.com/MDAnalysis/mdanalysis/issues/2345
+# from MDAnalysis.lib.nsgrid import FastNS
+# from MDAnalysis.lib.mdamath import triclinic_box
 # def calculate_coord_numbers(traj, atoms, cutoff, skin = 0, mask = None):
 #     """Compute the coordination numbers for `mask` atoms at all times in `traj`.
 #