moyo 0.10.0

/*!
# moyo

**moyo** is a fast and robust crystal symmetry finder.

## Using **moyo**

Simply add the following to your `Cargo.toml` file:

```ignore
[dependencies]
// TODO: replace the * with the latest version
moyo = "*"
```

## Examples

The basic usage of **moyo** is to create a [`moyo::base::Cell`](Cell) representing a crystal structure, and then create a [`moyo::MoyoDataset`](MoyoDataset) from the [`moyo::base::Cell`](Cell).
The [`moyo::MoyoDataset`](MoyoDataset) contains symmetry information of the input crystal structure: for example, the space group number, symmetry operations, and standardized cell.

Magnetic symmetry is also supported in **moyo**.
Magnetic moments are represented by a struct implementing the [`moyo::base::MagneticMoment`](MagneticMoment) trait: for example, [`moyo::base::Collinear`](base::Collinear) or [`moyo::base::NonCollinear`](base::NonCollinear).
Magnetic cell is represented by a [`moyo::base::MagneticCell`](MagneticCell) struct.
The [`moyo::MoyoMagneticDataset`](MoyoMagneticDataset) contains magnetic symmetry information of the input magnetic cell: for example, the magnetic space-group type, magnetic symmetry operations, and standardized magnetic cell.

```
use nalgebra::{matrix, vector, Matrix3, Vector3};
use moyo::{MoyoDataset, MoyoMagneticDataset};
use moyo::base::{Cell, MagneticCell, AngleTolerance, Lattice, NonCollinear, RotationMagneticMomentAction};
use moyo::data::Setting;

let lattice = Lattice::new(matrix![
    4.603, 0.0, 0.0;
    0.0, 4.603, 0.0;
    0.0, 0.0, 2.969;
]);
let x_4f = 0.3046;
let positions = vec![
    Vector3::new(0.0, 0.0, 0.0),                // Ti(2a)
    Vector3::new(0.5, 0.5, 0.5),                // Ti(2a)
    Vector3::new(x_4f, x_4f, 0.0),              // O(4f)
    Vector3::new(-x_4f, -x_4f, 0.0),            // O(4f)
    Vector3::new(-x_4f + 0.5, x_4f + 0.5, 0.5), // O(4f)
    Vector3::new(x_4f + 0.5, -x_4f + 0.5, 0.5), // O(4f)
];
let numbers = vec![0, 0, 1, 1, 1, 1];
let cell = Cell::new(lattice.clone(), positions.clone(), numbers.clone());

let symprec = 1e-4;

let dataset = MoyoDataset::with_default(&cell, symprec).unwrap();
assert_eq!(dataset.number, 136);  // P4_2/mnm

let magnetic_moments = vec![
    NonCollinear(vector![0.0, 0.0, 0.7]),
    NonCollinear(vector![0.0, 0.0, -0.7]),
    NonCollinear(vector![0.0, 0.0, 0.0]),
    NonCollinear(vector![0.0, 0.0, 0.0]),
    NonCollinear(vector![0.0, 0.0, 0.0]),
    NonCollinear(vector![0.0, 0.0, 0.0]),
];
let magnetic_cell = MagneticCell::new(lattice, positions, numbers, magnetic_moments);

let action = RotationMagneticMomentAction::Axial;

let magnetic_dataset = MoyoMagneticDataset::with_default(&magnetic_cell, symprec, action).unwrap();
assert_eq!(magnetic_dataset.uni_number, 1159);  // BNS 136.499
```

## Features
- Support most of the symmetry search functionality in Spglib
- Primitive cell search
- Symmetry operation search
- Space-group type identification
- Wyckoff position assignment
- Crystal structure symmetrization
- Magnetic space group support

*/
#[allow(unused_imports)]
#[macro_use]
extern crate approx;

pub mod base;
pub mod data;
pub mod identify;
pub mod math;
pub mod search; // Public for benchmarking
pub mod utils;

mod symmetrize;

pub use base::MoyoError as Error;

use crate::base::{
    AngleTolerance, Cell, LayerCell, Linear, MagneticCell, MagneticMoment, MagneticOperations,
    MoyoError, Operation, Operations, OriginShift, RotationMagneticMomentAction, Transformation,
    UnimodularTransformation,
};
use crate::data::{
    ArithmeticCrystalClassEntry, HallNumber, HallSymbol, HallSymbolEntry,
    LayerArithmeticCrystalClassEntry, LayerHallNumber, LayerHallSymbolEntry, LayerNumber,
    LayerSetting, Number, Setting, UNINumber, arithmetic_crystal_class_entry, hall_symbol_entry,
    layer_arithmetic_crystal_class_entry, layer_hall_symbol_entry,
};
use crate::identify::{LayerGroup, MagneticSpaceGroup, Normalizer, SpaceGroup};
use crate::search::{
    LayerPrimitiveCell, iterative_layer_symmetry_search, iterative_magnetic_symmetry_search,
    iterative_symmetry_search, magnetic_operations_in_magnetic_cell, operations_in_cell,
};
use crate::symmetrize::{
    StandardizedCell, StandardizedLayerCell, StandardizedMagneticCell, orbits_in_cell,
};
use crate::utils::{to_3_slice, to_3x3_slice};

use nalgebra::Matrix3;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Deserialize, Serialize)]
/// A dataset containing symmetry information of the input crystal structure.
pub struct MoyoDataset {
    // ------------------------------------------------------------------------
    // Identification
    // ------------------------------------------------------------------------
    /// Space group number.
    pub number: Number,
    /// Hall symbol number.
    pub hall_number: HallNumber,
    /// Hermann-Mauguin symbol in short notation (e.g., "Fd-3m" for space group 227).
    pub hm_symbol: String,
    // ------------------------------------------------------------------------
    // Symmetry operations in the input cell
    // ------------------------------------------------------------------------
    /// Symmetry operations in the input cell.
    pub operations: Operations,
    // ------------------------------------------------------------------------
    // Site symmetry
    // ------------------------------------------------------------------------
    /// Spglib's `crystallographic_orbits` not `equivalent_atoms`
    /// The `i`th atom in the input cell is equivalent to the `orbits[i]`th atom in the **input** cell.
    /// For example, orbits=[0, 0, 2, 2, 2, 2] means the first two atoms are equivalent and the last four atoms are equivalent to each other.
    pub orbits: Vec<usize>,
    /// Wyckoff letters for each site in the input cell.
    pub wyckoffs: Vec<char>,
    /// Site symmetry symbols for each site in the input cell.
    /// The orientation of the site symmetry is w.r.t. the standardized cell.
    pub site_symmetry_symbols: Vec<String>,
    // ------------------------------------------------------------------------
    // Standardized cell
    // ------------------------------------------------------------------------
    /// Standardized cell.
    ///
    /// The input cell is related to the standardized cell by
    /// `(std_linear, std_origin_shift)` and `std_rotation_matrix`:
    ///
    ///   Lattice (column-vector convention):
    ///     std_cell.lattice.basis = std_rotation_matrix * cell.lattice.basis * std_linear
    ///
    ///   Fractional positions:
    ///     x_std = std_linear^-1 * (x_input - std_origin_shift)
    ///
    /// `std_rotation_matrix` is a rigid rotation (orthogonal matrix) applied
    /// only to the Cartesian lattice basis. It does not affect fractional
    /// coordinates.
    pub std_cell: Cell,
    /// Linear part of transformation from the input cell to the standardized cell.
    pub std_linear: Matrix3<f64>,
    /// Origin shift of transformation from the input cell to the standardized cell.
    pub std_origin_shift: OriginShift,
    /// Rigid rotation (orthogonal matrix) applied to the lattice basis.
    pub std_rotation_matrix: Matrix3<f64>,
    /// Pearson symbol for standardized cell.
    pub pearson_symbol: String,
    // ------------------------------------------------------------------------
    // Primitive standardized cell
    // ------------------------------------------------------------------------
    /// Primitive standardized cell.
    ///
    /// Same transformation convention as the standardized cell above:
    ///
    ///   Lattice:
    ///     prim_std_cell.lattice.basis = std_rotation_matrix * cell.lattice.basis * prim_std_linear
    ///
    ///   Fractional positions:
    ///     x_prim_std = prim_std_linear^-1 * (x_input - prim_std_origin_shift)
    pub prim_std_cell: Cell,
    /// Linear part of transformation from the input cell to the primitive standardized cell.
    pub prim_std_linear: Matrix3<f64>,
    /// Origin shift of transformation from the input cell to the primitive standardized cell.
    pub prim_std_origin_shift: OriginShift,
    /// Mapping sites in the input cell to those in the primitive standardized cell.
    /// The `i`th atom in the input cell is mapped to the `mapping_to_std_prim[i]`th atom in the primitive standardized cell.
    pub mapping_std_prim: Vec<usize>,
    // ------------------------------------------------------------------------
    // Final parameters
    // ------------------------------------------------------------------------
    /// Actually used `symprec` in iterative symmetry search.
    pub symprec: f64,
    /// Actually used `angle_tolerance` in iterative symmetry search.
    pub angle_tolerance: AngleTolerance,
}

impl MoyoDataset {
    /// Create a new [`MoyoDataset`] from the input cell, `cell`.
    /// `symprec` controls the tolerance for searching symmetry operations in the unit of `cell.lattice.basis`.
    /// `setting` determines the preference for the "standardized" setting of a detected space-group type.
    /// `rotate_basis` specifies whether to rotate the basis vectors of the input cell to those of the standardized cell.
    /// If the search fails, [`MoyoError`] is returned.
    pub fn new(
        cell: &Cell,
        symprec: f64,
        angle_tolerance: AngleTolerance,
        setting: Setting,
        rotate_basis: bool,
    ) -> Result<Self, MoyoError> {
        let (prim_cell, symmetry_search, symprec, angle_tolerance) =
            iterative_symmetry_search(cell, symprec, angle_tolerance)?;
        let operations = operations_in_cell(&prim_cell, &symmetry_search.operations);

        // Space-group type identification
        let epsilon = symprec / prim_cell.cell.lattice.volume().powf(1.0 / 3.0);
        let space_group = SpaceGroup::new(&symmetry_search.operations, setting, epsilon)?;

        // Standardized cell
        let std_cell = StandardizedCell::new(
            &prim_cell.cell,
            &symmetry_search.operations,
            &symmetry_search.permutations,
            &space_group,
            symprec,
            epsilon,
            rotate_basis,
        )?;

        // site symmetry
        let orbits = orbits_in_cell(
            prim_cell.cell.num_atoms(),
            &symmetry_search.permutations,
            &prim_cell.site_mapping,
        );
        // StandardizedCell.prim_cell and prim_cell have the same site order
        let mapping_std_prim = prim_cell.site_mapping.clone();
        let wyckoffs = lift_wyckoffs_to_input_cell(
            prim_cell.cell.num_atoms(),
            &std_cell.wyckoffs,
            &std_cell.site_mapping,
            &mapping_std_prim,
        )?;

        let (std_linear, std_origin_shift, prim_std_linear, prim_std_origin_shift) =
            compose_std_transformations(
                prim_cell.linear,
                &std_cell.transformation,
                &std_cell.prim_transformation,
            );

        // Pearson symbol
        let hall_symbol = hall_symbol_entry(space_group.hall_number).unwrap();
        let arithmetic_entry =
            arithmetic_crystal_class_entry(hall_symbol.arithmetic_number).unwrap();
        let bravais_class = arithmetic_entry.bravais_class;
        let pearson_symbol = format!("{}{}", bravais_class.to_string(), std_cell.cell.num_atoms());

        Ok(Self {
            // Space-group type
            number: space_group.number,
            hall_number: space_group.hall_number,
            hm_symbol: hall_symbol.hm_short.to_string(),
            // Symmetry operations in the input cell
            operations,
            // Standardized cell
            std_cell: std_cell.cell,
            std_linear,
            std_origin_shift,
            std_rotation_matrix: std_cell.rotation_matrix,
            pearson_symbol,
            // Primitive standardized cell
            prim_std_cell: std_cell.prim_cell,
            prim_std_linear,
            prim_std_origin_shift,
            mapping_std_prim,
            // Site symmetry
            orbits,
            wyckoffs: wyckoffs.iter().map(|w| w.letter).collect(),
            site_symmetry_symbols: wyckoffs
                .iter()
                .map(|w| w.site_symmetry.to_string())
                .collect(),
            // Final parameters
            symprec,
            angle_tolerance,
        })
    }

    /// Create a new [`MoyoDataset`] from the input cell, `cell`, with default parameters.
    /// `symprec` controls the tolerance for searching symmetry operations in the unit of `cell.lattice.basis`.
    pub fn with_default(cell: &Cell, symprec: f64) -> Result<Self, MoyoError> {
        Self::new(
            cell,
            symprec,
            AngleTolerance::default(),
            Setting::default(),
            true,
        )
    }

    /// Return the number of symmetry operations in the input cell.
    pub fn num_operations(&self) -> usize {
        self.operations.len()
    }

    /// Returns `std_linear` as a 3x3 array. See [`std_linear`](Self::std_linear) field docs.
    pub fn std_linear_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.std_linear)
    }

    /// Returns `std_origin_shift` as a `[f64; 3]` array.
    pub fn std_origin_shift_as_array(&self) -> [f64; 3] {
        to_3_slice(&self.std_origin_shift)
    }

    /// Returns `std_rotation_matrix` as a 3x3 array.
    pub fn std_rotation_matrix_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.std_rotation_matrix)
    }

    /// Returns `prim_std_linear` as a 3x3 array.
    pub fn prim_std_linear_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.prim_std_linear)
    }

    /// Returns `prim_std_origin_shift` as a `[f64; 3]` array.
    pub fn prim_std_origin_shift_as_array(&self) -> [f64; 3] {
        to_3_slice(&self.prim_std_origin_shift)
    }

    /// Returns the Hall symbol entry for the space group.
    pub fn hall_symbol(&self) -> HallSymbolEntry {
        hall_symbol_entry(self.hall_number).unwrap()
    }

    /// Returns the arithmetic crystal class entry for the space group.
    pub fn arithmetic_crystal_class(&self) -> ArithmeticCrystalClassEntry {
        arithmetic_crystal_class_entry(self.hall_symbol().arithmetic_number).unwrap()
    }

    /// Compute the Euclidean normalizer of the identified space group in the
    /// primitive standardized basis (`prim_std_cell.lattice`).
    ///
    /// `preserve_chirality = true` restricts to the chirality-preserving
    /// subgroup N_E^+(G).
    pub fn euclidean_normalizer(&self, preserve_chirality: bool) -> Result<Normalizer, MoyoError> {
        // `self.hall_number` was produced by a successful identification, so
        // looking it up in the Hall-symbol database always succeeds.
        let hall_symbol = HallSymbol::from_hall_number(self.hall_number).unwrap();
        let prim_operations = hall_symbol.primitive_traverse();
        let prim_generators = hall_symbol.primitive_generators();
        Normalizer::from_lattice(
            &self.prim_std_cell.lattice,
            &prim_operations,
            &prim_generators,
            self.symprec,
            self.angle_tolerance,
            preserve_chirality,
        )
    }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
/// A dataset containing layer-group symmetry information of the input crystal
/// structure (a 2D-periodic system whose third basis vector is the aperiodic
/// stacking direction; paper Fu et al. 2024 §2).
///
/// Mirrors [`MoyoDataset`] for the bulk space-group case, but the input must
/// satisfy the layer-group periodicity contract (`c` perpendicular to `a, b`)
/// enforced by [`LayerCell::new`].
pub struct MoyoLayerDataset {
    // ------------------------------------------------------------------------
    // Identification
    // ------------------------------------------------------------------------
    /// Layer group number (1 - 80, paper Fu et al. 2024 Table 5).
    pub number: LayerNumber,
    /// Layer Hall symbol number (1 - 116).
    pub hall_number: LayerHallNumber,
    /// Hermann-Mauguin symbol in short notation.
    pub hm_symbol: String,
    // ------------------------------------------------------------------------
    // Symmetry operations in the input cell
    // ------------------------------------------------------------------------
    /// Layer-group operations in the input cell.
    pub operations: Operations,
    // ------------------------------------------------------------------------
    // Site symmetry
    // ------------------------------------------------------------------------
    /// The `i`th atom in the input cell is equivalent to the `orbits[i]`th atom
    /// in the input cell. Convention matches [`MoyoDataset::orbits`].
    pub orbits: Vec<usize>,
    /// Wyckoff letters for each site in the input cell.
    pub wyckoffs: Vec<char>,
    /// Site symmetry symbols for each site in the input cell, oriented w.r.t.
    /// the standardized cell.
    pub site_symmetry_symbols: Vec<String>,
    // ------------------------------------------------------------------------
    // Standardized layer cell
    // ------------------------------------------------------------------------
    /// Conventional standardized layer cell. See
    /// `moyopy/docs/layer_standardization.md` for the output convention.
    pub std_cell: LayerCell,
    /// Linear part of the transformation from the input cell to `std_cell`.
    pub std_linear: Matrix3<f64>,
    /// Origin shift of the transformation from the input cell to `std_cell`.
    pub std_origin_shift: OriginShift,
    /// Rigid rotation applied to the lattice basis when `rotate_basis = true`,
    /// identity otherwise.
    pub std_rotation_matrix: Matrix3<f64>,
    /// Pearson symbol for the standardized layer cell. The first two characters
    /// are the 2D Bravais type (`mp`, `op`, `oc`, `tp`, or `hp`).
    pub pearson_symbol: String,
    // ------------------------------------------------------------------------
    // Primitive standardized layer cell
    // ------------------------------------------------------------------------
    /// Primitive standardized layer cell.
    pub prim_std_cell: LayerCell,
    /// Linear part of the transformation from the input cell to `prim_std_cell`.
    pub prim_std_linear: Matrix3<f64>,
    /// Origin shift of the transformation from the input cell to `prim_std_cell`.
    pub prim_std_origin_shift: OriginShift,
    /// Mapping sites in the input cell to those in `prim_std_cell`.
    pub mapping_std_prim: Vec<usize>,
    // ------------------------------------------------------------------------
    // Final parameters
    // ------------------------------------------------------------------------
    /// `symprec` used for the symmetry search.
    pub symprec: f64,
    /// `angle_tolerance` used for the symmetry search.
    pub angle_tolerance: AngleTolerance,
}

impl MoyoLayerDataset {
    /// Identify the layer group of `cell` and produce a [`MoyoLayerDataset`].
    ///
    /// `cell` must satisfy the layer-group periodicity contract (`c`
    /// perpendicular to `a, b`); inputs that violate it are rejected with
    /// [`MoyoError::AperiodicAxisNotOrthogonal`].
    pub fn new(
        cell: &Cell,
        symprec: f64,
        angle_tolerance: AngleTolerance,
        setting: LayerSetting,
        rotate_basis: bool,
    ) -> Result<Self, MoyoError> {
        let (prim_layer, symmetry_search, symprec, angle_tolerance) =
            iterative_layer_symmetry_search(cell, symprec, angle_tolerance)?;

        let prim_volume = prim_layer.layer_cell.lattice().basis().determinant().abs();
        let epsilon = symprec / prim_volume.powf(1.0 / 3.0);

        let layer_group = LayerGroup::new(&symmetry_search.operations, setting, epsilon)?;

        let std_layer = StandardizedLayerCell::new(
            &prim_layer.layer_cell,
            &symmetry_search.operations,
            &symmetry_search.permutations,
            &layer_group,
            symprec,
            epsilon,
            rotate_basis,
        )?;

        let operations = layer_operations_in_cell(&prim_layer, &symmetry_search.operations);

        // `StandardizedLayerCell.prim_layer_cell` and `prim_layer.layer_cell`
        // share the same site order (mirrors the bulk pipeline), so the
        // prim-cell Wyckoffs are looked up via `site_mapping`.
        let mapping_std_prim = prim_layer.site_mapping.clone();
        let orbits = orbits_in_cell(
            prim_layer.layer_cell.num_atoms(),
            &symmetry_search.permutations,
            &prim_layer.site_mapping,
        );
        let wyckoffs = lift_wyckoffs_to_input_cell(
            prim_layer.layer_cell.num_atoms(),
            &std_layer.wyckoffs,
            &std_layer.site_mapping,
            &mapping_std_prim,
        )?;

        let (std_linear, std_origin_shift, prim_std_linear, prim_std_origin_shift) =
            compose_std_transformations(
                prim_layer.linear,
                &std_layer.transformation,
                &std_layer.prim_transformation,
            );

        let entry = layer_hall_symbol_entry(layer_group.hall_number).unwrap();
        let layer_arith = layer_arithmetic_crystal_class_entry(entry.arithmetic_number).unwrap();
        let pearson_symbol = format!(
            "{}{}",
            layer_arith.layer_bravais_class.to_string(),
            std_layer.layer_cell.num_atoms()
        );

        Ok(Self {
            number: layer_group.number,
            hall_number: layer_group.hall_number,
            hm_symbol: entry.hm_short.to_string(),
            operations,
            orbits,
            wyckoffs: wyckoffs.iter().map(|w| w.letter).collect(),
            site_symmetry_symbols: wyckoffs
                .iter()
                .map(|w| w.site_symmetry.to_string())
                .collect(),
            std_cell: std_layer.layer_cell,
            std_linear,
            std_origin_shift,
            std_rotation_matrix: std_layer.rotation_matrix,
            pearson_symbol,
            prim_std_cell: std_layer.prim_layer_cell,
            prim_std_linear,
            prim_std_origin_shift,
            mapping_std_prim,
            symprec,
            angle_tolerance,
        })
    }

    /// Identify the layer group of `cell` with default parameters
    /// ([`AngleTolerance::Default`], [`LayerSetting::Standard`],
    /// `rotate_basis = true`).
    pub fn with_default(cell: &Cell, symprec: f64) -> Result<Self, MoyoError> {
        Self::new(
            cell,
            symprec,
            AngleTolerance::default(),
            LayerSetting::default(),
            true,
        )
    }

    /// Number of layer-group operations in the input cell.
    pub fn num_operations(&self) -> usize {
        self.operations.len()
    }

    /// `std_linear` as a 3x3 array.
    pub fn std_linear_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.std_linear)
    }

    /// `std_origin_shift` as a 3-array.
    pub fn std_origin_shift_as_array(&self) -> [f64; 3] {
        to_3_slice(&self.std_origin_shift)
    }

    /// `std_rotation_matrix` as a 3x3 array.
    pub fn std_rotation_matrix_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.std_rotation_matrix)
    }

    /// `prim_std_linear` as a 3x3 array.
    pub fn prim_std_linear_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.prim_std_linear)
    }

    /// `prim_std_origin_shift` as a 3-array.
    pub fn prim_std_origin_shift_as_array(&self) -> [f64; 3] {
        to_3_slice(&self.prim_std_origin_shift)
    }

    /// Hall symbol entry for the identified layer group.
    pub fn hall_symbol(&self) -> LayerHallSymbolEntry {
        layer_hall_symbol_entry(self.hall_number).unwrap().clone()
    }

    /// Arithmetic crystal class entry for the identified layer group.
    pub fn arithmetic_crystal_class(&self) -> LayerArithmeticCrystalClassEntry {
        layer_arithmetic_crystal_class_entry(self.hall_symbol().arithmetic_number).unwrap()
    }
}

/// Lift primitive-layer-cell operations to operations in the input cell.
/// Mirrors [`crate::search::operations_in_cell`] but consumes
/// [`LayerPrimitiveCell`] (whose pure-translation set is in-plane only).
fn layer_operations_in_cell(prim: &LayerPrimitiveCell, prim_operations: &Operations) -> Operations {
    let input_operations =
        Transformation::from_linear(prim.linear).transform_operations(prim_operations);
    let mut operations = vec![];
    for t1 in prim.translations.iter() {
        for op2 in input_operations.iter() {
            // (E, t1) (rotation, t2) = (rotation, t1 + t2)
            let t12 = (t1 + op2.translation).map(|e| e % 1.);
            operations.push(Operation::new(op2.rotation, t12));
        }
    }
    operations
}

/// Lift Wyckoff assignments from the standardized cell back to the input cell.
///
/// The standardized cell and the input primitive cell share the same site order
/// (the standardize step preserves it), so per-orbit Wyckoffs collected in
/// `std_site_mapping` order can be re-indexed by `input_to_prim_mapping` to
/// produce one Wyckoff per input-cell site. Returns
/// [`MoyoError::WyckoffPositionAssignmentError`] if any primitive site is left
/// without a Wyckoff (which would mean the standardize step missed an orbit).
fn lift_wyckoffs_to_input_cell<W: Clone>(
    prim_num_atoms: usize,
    std_wyckoffs: &[W],
    std_site_mapping: &[usize],
    input_to_prim_mapping: &[usize],
) -> Result<Vec<W>, MoyoError> {
    let mut prim_wyckoffs: Vec<Option<W>> = vec![None; prim_num_atoms];
    for (i, wyckoff) in std_wyckoffs.iter().enumerate() {
        let j = std_site_mapping[i];
        if prim_wyckoffs[j].is_none() {
            prim_wyckoffs[j] = Some(wyckoff.clone());
        }
    }
    let lifted: Option<Vec<W>> = input_to_prim_mapping
        .iter()
        .map(|&i| prim_wyckoffs[i].clone())
        .collect();
    lifted.ok_or(MoyoError::WyckoffPositionAssignmentError)
}

/// Compose the (input -> primitive) and (primitive -> standardized) maps into
/// the (input -> conventional std) and (input -> primitive std) transformations
/// reported on the dataset. Returns
/// `(std_linear, std_origin_shift, prim_std_linear, prim_std_origin_shift)`.
///
/// `prim_linear` is the primitive-to-input matrix (`PrimitiveCell::linear` for
/// bulk, `LayerPrimitiveCell::linear` for layer); inverting it yields the
/// input-to-primitive map that the standardized transformations are then
/// composed with.
fn compose_std_transformations(
    prim_linear: Linear,
    std_transformation: &Transformation,
    prim_std_transformation: &UnimodularTransformation,
) -> (Matrix3<f64>, OriginShift, Matrix3<f64>, OriginShift) {
    let prim_inv = prim_linear.map(|e| e as f64).try_inverse().unwrap();
    (
        prim_inv * std_transformation.linear_as_f64(),
        prim_inv * std_transformation.origin_shift,
        prim_inv * prim_std_transformation.linear_as_f64(),
        prim_inv * prim_std_transformation.origin_shift,
    )
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct MoyoMagneticDataset<M: MagneticMoment> {
    // ------------------------------------------------------------------------
    // Magnetic space-group type
    // ------------------------------------------------------------------------
    pub uni_number: UNINumber,
    // ------------------------------------------------------------------------
    // Magnetic symmetry operations in the input cell
    // ------------------------------------------------------------------------
    /// Magnetic symmetry operations in the input cell.
    pub magnetic_operations: MagneticOperations,
    // ------------------------------------------------------------------------
    // Site symmetry
    // ------------------------------------------------------------------------
    /// The `i`th atom in the input magnetic cell is equivalent to the `orbits[i]`th atom in the **input** magnetic cell.
    /// For example, orbits=[0, 0, 2, 2, 2, 2] means the first two atoms are equivalent and the last four atoms are equivalent to each other.
    pub orbits: Vec<usize>,
    // ------------------------------------------------------------------------
    // Standardized magnetic cell
    // ------------------------------------------------------------------------
    /// Standardized magnetic cell.
    ///
    /// The input magnetic cell is related to the standardized magnetic cell by
    /// `(std_linear, std_origin_shift)` and `std_rotation_matrix`:
    ///
    ///   Lattice (column-vector convention):
    ///     std_mag_cell.cell.lattice.basis = std_rotation_matrix * mag_cell.cell.lattice.basis * std_linear
    ///
    ///   Fractional positions:
    ///     x_std = std_linear^-1 * (x_input - std_origin_shift)
    ///
    /// `std_rotation_matrix` is a rigid rotation (orthogonal matrix) applied
    /// only to the Cartesian lattice basis. It does not affect fractional
    /// coordinates.
    pub std_mag_cell: MagneticCell<M>,
    /// Linear part of transformation from the input magnetic cell to the standardized one.
    pub std_linear: Matrix3<f64>,
    /// Origin shift of transformation from the input magnetic cell to the standardized one.
    pub std_origin_shift: OriginShift,
    /// Rigid rotation (orthogonal matrix) applied to the lattice basis.
    pub std_rotation_matrix: Matrix3<f64>,
    // ------------------------------------------------------------------------
    // Primitive standardized magnetic cell
    // ------------------------------------------------------------------------
    /// Primitive standardized magnetic cell.
    ///
    /// Same transformation convention as the standardized magnetic cell above:
    ///
    ///   Lattice:
    ///     prim_std_mag_cell.cell.lattice.basis = std_rotation_matrix * mag_cell.cell.lattice.basis * prim_std_linear
    ///
    ///   Fractional positions:
    ///     x_prim_std = prim_std_linear^-1 * (x_input - prim_std_origin_shift)
    pub prim_std_mag_cell: MagneticCell<M>,
    /// Linear part of transformation from the input magnetic cell to the primitive standardized magnetic cell.
    pub prim_std_linear: Matrix3<f64>,
    /// Origin shift of transformation from the input magnetic cell to the primitive standardized magnetic cell.
    pub prim_std_origin_shift: OriginShift,
    /// Mapping sites in the input magnetic cell to those in the primitive standardized magnetic cell.
    /// The `i`th atom in the input magnetic cell is mapped to the `mapping_to_std_prim[i]`th atom in the primitive standardized magnetic cell.
    pub mapping_std_prim: Vec<usize>,
    // ------------------------------------------------------------------------
    // Final parameters
    // ------------------------------------------------------------------------
    /// Actually used `symprec` in iterative symmetry search.
    pub symprec: f64,
    /// Actually used `angle_tolerance` in iterative symmetry search.
    pub angle_tolerance: AngleTolerance,
    /// Actually used `mag_symprec` in iterative symmetry search.
    pub mag_symprec: f64,
}

impl<M: MagneticMoment> MoyoMagneticDataset<M> {
    /// Create a new [`MoyoMagneticDataset`] from the input magnetic cell, `magnetic_cell`.
    /// `symprec` controls the tolerance for searching symmetry operations in the unit of `magnetic_cell.cell.lattice.basis`.
    /// `mag_symprec` controls the tolerance for searching magnetic symmetry operations in the unit of `magnetic_cell.magnetic_moments`.
    /// `action` specifies how a magnetic symmetry operation acts on magnetic moments.
    /// `rotate_basis` specifies whether to rotate the basis vectors of the input cell to those of the standardized cell.
    /// If the search fails, [`MoyoError`] is returned.
    pub fn new(
        magnetic_cell: &MagneticCell<M>,
        symprec: f64,
        angle_tolerance: AngleTolerance,
        mag_symprec: Option<f64>,
        action: RotationMagneticMomentAction,
        rotate_basis: bool,
    ) -> Result<Self, MoyoError> {
        let (prim_mag_cell, magnetic_symmetry_search, symprec, angle_tolerance, mag_symprec) =
            iterative_magnetic_symmetry_search(
                magnetic_cell,
                symprec,
                angle_tolerance,
                mag_symprec,
                action,
            )?;
        let magnetic_operations = magnetic_operations_in_magnetic_cell(
            &prim_mag_cell,
            &magnetic_symmetry_search.magnetic_operations,
        );

        // Magnetic space-group type identification
        let epsilon = symprec
            / prim_mag_cell
                .magnetic_cell
                .cell
                .lattice
                .volume()
                .powf(1.0 / 3.0);
        let magnetic_space_group =
            MagneticSpaceGroup::new(&magnetic_symmetry_search.magnetic_operations, epsilon)?;

        // Standardized magnetic cell
        let std_mag_cell = StandardizedMagneticCell::new(
            &prim_mag_cell,
            &magnetic_symmetry_search,
            &magnetic_space_group,
            symprec,
            mag_symprec,
            epsilon,
            action,
            rotate_basis,
        )?;

        // Site symmetry
        // StandardizedMagneticCell.prim_mag_cell and prim_mag_cell have the same site order
        let mapping_std_prim = prim_mag_cell.site_mapping.clone();
        let orbits = orbits_in_cell(
            prim_mag_cell.magnetic_cell.num_atoms(),
            &magnetic_symmetry_search.permutations,
            &mapping_std_prim,
        );

        // magnetic_cell <-(prim_mag_cell.linear, 0)- prim_mag_cell.magnetic_cell -(std_mag_cell.transformation)-> std_mag_cell.mag_cell
        // (std_linear, std_origin_shift) = (prim_mag_cell.linear^-1, 0) * std_mag_cell.transformation
        let prim_mag_cell_linear_inv = prim_mag_cell
            .linear
            .map(|e| e as f64)
            .try_inverse()
            .unwrap();
        let std_linear = prim_mag_cell_linear_inv * std_mag_cell.transformation.linear_as_f64();
        let std_origin_shift = prim_mag_cell_linear_inv * std_mag_cell.transformation.origin_shift;

        // (prim_std_linear, prim_std_origin_shift) = (prim_mag_cell.linear^-1, 0) * std_mag_cell.prim_transformation
        let prim_std_linear =
            prim_mag_cell_linear_inv * std_mag_cell.prim_transformation.linear_as_f64();
        let prim_std_origin_shift =
            prim_mag_cell_linear_inv * std_mag_cell.prim_transformation.origin_shift;

        Ok(Self {
            // Magnetic space-group type
            uni_number: magnetic_space_group.uni_number,
            // Magnetic symmetry operations in the input cell
            magnetic_operations,
            // Site symmetry
            orbits,
            // Standardized magnetic cell
            std_mag_cell: std_mag_cell.mag_cell,
            std_linear,
            std_origin_shift,
            std_rotation_matrix: std_mag_cell.rotation_matrix,
            // Primitive standardized magnetic cell
            prim_std_mag_cell: std_mag_cell.prim_mag_cell,
            prim_std_linear,
            prim_std_origin_shift,
            mapping_std_prim,
            // Final parameters
            symprec,
            angle_tolerance,
            mag_symprec,
        })
    }

    /// Create a new [`MoyoMagneticDataset`] from the input magnetic cell, `magnetic_cell`, with default parameters.
    /// `symprec` controls the tolerance for searching symmetry operations in the unit of `magnetic_cell.cell.lattice.basis`.
    /// `action` specifies how a magnetic symmetry operation acts on magnetic moments.
    pub fn with_default(
        magnetic_cell: &MagneticCell<M>,
        symprec: f64,
        action: RotationMagneticMomentAction,
    ) -> Result<Self, MoyoError> {
        Self::new(
            magnetic_cell,
            symprec,
            AngleTolerance::default(),
            None,
            action,
            true,
        )
    }

    /// Return the number of magnetic symmetry operations in the input magnetic cell.
    pub fn num_magnetic_operations(&self) -> usize {
        self.magnetic_operations.len()
    }

    /// Returns `std_linear` as a 3x3 array. See [`std_linear`](Self::std_linear) field docs.
    pub fn std_linear_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.std_linear)
    }

    /// Returns `std_origin_shift` as a `[f64; 3]` array.
    pub fn std_origin_shift_as_array(&self) -> [f64; 3] {
        to_3_slice(&self.std_origin_shift)
    }

    /// Returns `std_rotation_matrix` as a 3x3 array.
    pub fn std_rotation_matrix_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.std_rotation_matrix)
    }

    /// Returns `prim_std_linear` as a 3x3 array.
    pub fn prim_std_linear_as_array(&self) -> [[f64; 3]; 3] {
        to_3x3_slice(&self.prim_std_linear)
    }

    /// Returns `prim_std_origin_shift` as a `[f64; 3]` array.
    pub fn prim_std_origin_shift_as_array(&self) -> [f64; 3] {
        to_3_slice(&self.prim_std_origin_shift)
    }
}