move clustering solution from test playground to the main code base; #44

src/DiffAtomComp.py

@@ -214,7 +214,7 @@ def cluster_and_sort_sqd_fast(e_sqd_log, shift_tolerance: float = 3.0, angle_tol
         T_ID_dict = {}
         for i in range(len(mol_shift)):
             ptf = T[i]
-            in_cluster = b.one_in_cluster_transform(ptf)
+            in_cluster = b.any_close_transform(ptf)
             if in_cluster is None:
                 b.add_transform(ptf)
                 mol_transform_label.append(unique_id)

src/DiffFit_bins.py

@@ -99,6 +99,21 @@ def one_in_cluster_transform(self, tf):
 
         return None
 
+    # -------------------------------------------------------------------------
+    #
+    def any_close_transform(self, tf):
+        '''Check the center bin first for a close transform to improve speed when most queries have a close transform.'''
+        bc = tuple(int(x / bs) for x, bs in zip(self.bin_point(tf), self.bins.bin_size))
+        if bc in self.bins.bins:
+            itf = tf.inverse()
+            for c, btf in self.bins.bins[bc]:
+                dx, dy, dz = btf.translation() - tf.translation()
+                if (dx * dx + dy * dy + dz * dz <= self.d2max and
+                        (btf * itf).rotation_angle() < self.angle):
+                    return btf
+
+        return self.one_in_cluster_transform(tf)
+
 # -----------------------------------------------------------------------------
 # Bin objects in a grid for fast lookup of objects close to a given object.
 #