Research translation via KdTree

moyo/src/search/primitive_cell.rs

@@ -3,7 +3,8 @@ use std::collections::BTreeMap;
 use nalgebra::{Dyn, Matrix3, OMatrix, Vector3, U3};
 
 use super::solve::{
-    pivot_site_indices, solve_correspondence_naive, symmetrize_translation_from_permutation,
+    pivot_site_indices, solve_correspondence, symmetrize_translation_from_permutation,
+    PeriodicKdTree,
 };
 use crate::base::{
     orbits_from_permutations, Cell, Linear, MoyoError, Permutation, Position, Rotation,
@@ -47,6 +48,7 @@ impl PrimitiveCell {
         }
 
         // Try possible translations: overlap the `src`the site to the `dst`th site
+        let pkdtree = PeriodicKdTree::new(&reduced_cell, rough_symprec);
         let pivot_site_indices = pivot_site_indices(&reduced_cell.numbers);
         let mut permutations_translations_tmp = vec![];
         let src = pivot_site_indices[0];
@@ -60,8 +62,7 @@ impl PrimitiveCell {
 
             // Because the translation may not be optimal to minimize distance between input and acted positions,
             // use a larger symprec (diameter of a Ball) for finding correspondence
-            if let Some(permutation) =
-                solve_correspondence_naive(&reduced_cell, &new_positions, rough_symprec)
+            if let Some(permutation) = solve_correspondence(&pkdtree, &reduced_cell, &new_positions)
             {
                 permutations_translations_tmp.push((permutation, translation));
             }

moyo/src/search/solve.rs

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
 
 use itertools::iproduct;
 use kiddo::{KdTree, SquaredEuclidean};
-use nalgebra::{matrix, Vector3};
+use nalgebra::{Rotation3, Vector3};
 
 use crate::base::{AtomicSpecie, Cell, Lattice, Permutation, Position, Rotation, Translation};
 
@@ -23,18 +23,13 @@ pub struct PeriodicNeighbor {
 impl PeriodicKdTree {
     /// Construct a periodic kd-tree from the given **Minkowski-reduced** cell.
     pub fn new(reduced_cell: &Cell, symprec: f64) -> Self {
+        // Randomly rotate lattice to avoid align along x,y,z axes
+        let random_rotation_matrix = Rotation3::new(Vector3::new(0.01, 0.02, 0.03));
+        let new_lattice = reduced_cell.lattice.rotate(&random_rotation_matrix.into());
+
         // Twice the padding for safety
         let padding = 2.0 * symprec
-            / (3.0 * (reduced_cell.lattice.basis * reduced_cell.lattice.basis.transpose()).trace())
-                .sqrt();
-
-        // Randomly rotate lattice to avoid align along x,y,z axes
-        let random_rotation_matrix = matrix![
-            1.54407599, 1.85427383, 1.28175466;
-            1.79032253, 1.59655   , 0.67848383;
-            1.39057248, 1.77092511, 0.13151589;
-        ];
-        let new_lattice = reduced_cell.lattice.rotate(&random_rotation_matrix);
+            / (3.0 * (new_lattice.basis * new_lattice.basis.transpose()).trace()).sqrt();
 
         let mut entries = vec![];
         let mut indices = vec![];
@@ -84,9 +79,14 @@ impl PeriodicKdTree {
         }
 
         let item = within[0].item as usize;
+        let distance = within[0].distance.sqrt();
+        if distance > self.symprec {
+            return None;
+        }
+
         Some(PeriodicNeighbor {
             index: self.indices[item],
-            distance: within[0].distance.sqrt(),
+            distance,
         })
     }
 }
@@ -116,7 +116,6 @@ pub fn solve_correspondence(
     pkdtree: &PeriodicKdTree,
     reduced_cell: &Cell,
     new_positions: &[Position],
-    symprec: f64,
 ) -> Option<Permutation> {
     let num_atoms = pkdtree.num_sites;
     let mut mapping = vec![0; num_atoms];
@@ -128,9 +127,6 @@ pub fn solve_correspondence(
         if visited[j] || reduced_cell.numbers[i] != reduced_cell.numbers[j] {
             return None;
         }
-        if neighbor.distance > symprec {
-            return None;
-        }
 
         mapping[i] = j;
         visited[j] = true;
@@ -308,7 +304,7 @@ mod tests {
             assert_eq!(actual_naive, expect);
 
             let actual_kdtree =
-                solve_correspondence(&pkdtree, &reduced_cell, &new_positions, symprec).unwrap();
+                solve_correspondence(&pkdtree, &reduced_cell, &new_positions).unwrap();
             assert_eq!(actual_kdtree, expect);
         }
         {
@@ -323,8 +319,7 @@ mod tests {
             let actual_naive = solve_correspondence_naive(&reduced_cell, &new_positions, symprec);
             assert_eq!(actual_naive, None);
 
-            let actual_kdtree =
-                solve_correspondence(&pkdtree, &reduced_cell, &new_positions, symprec);
+            let actual_kdtree = solve_correspondence(&pkdtree, &reduced_cell, &new_positions);
             assert_eq!(actual_kdtree, None);
         }
     }

moyo/src/search/symmetry_search.rs

@@ -4,7 +4,8 @@ use std::collections::HashSet;
 use nalgebra::{Matrix3, Vector3};
 
 use super::solve::{
-    pivot_site_indices, solve_correspondence_naive, symmetrize_translation_from_permutation,
+    pivot_site_indices, solve_correspondence, symmetrize_translation_from_permutation,
+    PeriodicKdTree,
 };
 use crate::base::{
     AngleTolerance, Cell, Lattice, MoyoError, Operations, Permutation, Position, Rotation, EPS,
@@ -35,6 +36,7 @@ impl PrimitiveSymmetrySearch {
         }
 
         // Search symmetry operations
+        let pkdtree = PeriodicKdTree::new(primitive_cell, rough_symprec);
         let bravais_group =
             search_bravais_group(&primitive_cell.lattice, symprec, angle_tolerance)?;
         let pivot_site_indices = pivot_site_indices(&primitive_cell.numbers);
@@ -52,7 +54,7 @@ impl PrimitiveSymmetrySearch {
                     .collect();
 
                 if let Some(permutation) =
-                    solve_correspondence_naive(primitive_cell, &new_positions, rough_symprec)
+                    solve_correspondence(&pkdtree, primitive_cell, &new_positions)
                 {
                     symmetries_tmp.push((*rotation, translation, permutation));
                     // If a translation part is found, it should be unique (up to lattice translations)