diff --git a/src/DiffAtomComp.py b/src/DiffAtomComp.py index 14cbef0..c9cb60a 100644 --- a/src/DiffAtomComp.py +++ b/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) diff --git a/src/DiffFit_bins.py b/src/DiffFit_bins.py index 8dcdbfc..cb0d4c4 100644 --- a/src/DiffFit_bins.py +++ b/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. #