Remove duplicated key search

bench/mp.py

@@ -39,6 +39,7 @@ def main():
                 pbar.set_postfix_str(f"{material_id=} number={moyopy_dataset.number}")
             except:  # noqa: E722
                 print(f"Failed: {material_id=}")
+                print(f"Elapsed time: {time_moyopy=}, {time_spglib=}")
 
                 with open(f"{material_id}.json", "w") as f:
                     json.dump(structure.as_dict(), f, indent=4)

justfile

@@ -16,6 +16,9 @@ install-python:
 test-python:
     pytest -v moyopy/python/tests
 
+profile:
+    cargo flamegraph --test test_moyo_dataset --root
+
 pre-commit:
     pre-commit run --all-files
 

moyo/src/data/hall_symbol.rs

@@ -1,3 +1,4 @@
+use std::collections::hash_map::Entry;
 use std::collections::{HashMap, VecDeque};
 
 use nalgebra::{matrix, Matrix3, Vector3};
@@ -172,10 +173,11 @@ impl HallSymbol {
 
         while !queue.is_empty() {
             let (rotation_lhs, translation_lhs) = queue.pop_front().unwrap();
-            if hm_translations.contains_key(&rotation_lhs) {
+            let entry = hm_translations.entry(rotation_lhs);
+            if let Entry::Occupied(_) = entry {
                 continue;
             }
-            hm_translations.insert(rotation_lhs, translation_lhs);
+            entry.or_insert(translation_lhs);
             rotations.push(rotation_lhs);
             translations.push(translation_lhs);
 
@@ -185,17 +187,17 @@ impl HallSymbol {
                 .iter()
                 .zip(self.generators.translations.iter())
             {
-                let rotation = rotation_lhs * rotation_rhs;
-                let translation =
+                let new_rotation = rotation_lhs * rotation_rhs;
+                let new_translation =
                     rotation_lhs.map(|e| e as f64) * translation_rhs + translation_lhs;
-                let translation_mod1 = translation.map(|e| {
+                let new_translation_mod1 = new_translation.map(|e| {
                     let mut eint = (e * (MAX_DENOMINATOR as f64)).round() as i32;
                     eint = eint.rem_euclid(MAX_DENOMINATOR);
                     (eint as f64) / (MAX_DENOMINATOR as f64)
                 });
 
-                if !hm_translations.contains_key(&rotation) {
-                    queue.push_back((rotation, translation_mod1));
+                if !hm_translations.contains_key(&new_rotation) {
+                    queue.push_back((new_rotation, new_translation_mod1));
                 }
             }
         }

moyo/src/identify/space_group.rs

@@ -168,11 +168,8 @@ fn match_origin_shift(
         .enumerate()
     {
         // Correction transformation matrix may not be normalizer of the point group. For example, mm2 -> 2mm
-        if !hm_translations.contains_key(rotation) {
-            return None;
-        }
+        let target_translation = hm_translations.get(rotation)?;
 
-        let target_translation = hm_translations.get(rotation).unwrap();
         let ak = rotation - Matrix3::<i32>::identity();
         let bk = other_translation - target_translation;
         for i in 0..3 {

moyo/src/search/primitive_cell.rs

@@ -196,13 +196,15 @@ fn primitive_cell_from_transformation(
     Some((primitive_cell, site_mapping))
 }
 
-#[allow(clippy::map_entry)]
 fn site_mapping_from_orbits(orbits: &[usize]) -> Vec<usize> {
     let mut mapping = BTreeMap::new();
+    let mut count = 0;
     for ri in orbits.iter() {
-        if !mapping.contains_key(&ri) {
-            mapping.insert(ri, mapping.len());
-        }
+        mapping.entry(ri).or_insert_with(|| {
+            let value = count;
+            count += 1;
+            value
+        });
     }
 
     orbits.iter().map(|ri| *mapping.get(&ri).unwrap()).collect()

moyo/tests/test_moyo_dataset.rs

@@ -403,17 +403,17 @@ fn test_with_trigonal_Sc() {
     let numbers = vec![0];
     let cell = Cell::new(lattice, positions, numbers);
 
-    let symprec = 1e-4;
+    let symprec = 1e-1; // This structure is distorted
     let angle_tolerance = AngleTolerance::Default;
     let setting = Setting::Standard;
 
     let dataset = assert_dataset(&cell, symprec, angle_tolerance, setting);
     assert_dataset(&dataset.std_cell, symprec, angle_tolerance, setting);
     assert_dataset(&dataset.prim_std_cell, symprec, angle_tolerance, setting);
 
-    // assert_eq!(dataset.number, 178);
-    // assert_eq!(dataset.hall_number, 472);
-    // assert_eq!(dataset.num_operations(), 12);
-    // assert_eq!(dataset.orbits, vec![0, 0, 0, 0, 0, 0]);
-    // assert_eq!(dataset.wyckoffs, vec!['a', 'a', 'a', 'a', 'a', 'a']);
+    assert_eq!(dataset.number, 166);
+    assert_eq!(dataset.hall_number, 458);
+    assert_eq!(dataset.num_operations(), 12); // Rhombohedral setting
+    assert_eq!(dataset.orbits, vec![0]);
+    assert_eq!(dataset.wyckoffs, vec!['a']);
 }