xtalforge-core 0.0.2

//! P1 expansion with provenance back to the asymmetric site and operation.

use serde::{Deserialize, Serialize};

use crate::Error;
use crate::Lattice;
use crate::Structure;
use crate::spatial::{FractionalGrid, perpendicular_widths};
use crate::symop::SymOp;

/// One atom of an expanded P1 cell.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct P1Atom {
    /// Element symbol.
    pub element: String,
    /// Site label inherited from the asymmetric unit.
    pub label: String,
    /// Fractional coordinates in the cell.
    pub frac: [f64; 3],
    /// Site occupancy.
    pub occupancy: f64,
    /// Index of the parent site.
    pub site_idx: usize,
    /// Index of the operation that generated this atom.
    pub op_idx: usize,
    /// Integer lattice shift applied after the operation.
    pub tau: [i32; 3],
}

/// The expanded P1 cell.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct P1Cell {
    /// Cell lattice.
    pub lattice: Lattice,
    /// Expanded atoms.
    pub atoms: Vec<P1Atom>,
}

impl P1Cell {
    /// Expands a structure with its explicit operations, or the identity when
    /// it has none.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidInput`] for a negative or non-finite tolerance.
    pub fn expand(structure: &Structure, tolerance: f64) -> Result<Self, Error> {
        if !tolerance.is_finite() || tolerance < 0.0 {
            return Err(Error::InvalidInput(
                "P1 expansion tolerance must be non-negative and finite".to_owned(),
            ));
        }
        let identity = [SymOp::identity()];
        let operations = if structure.operations().is_empty() {
            identity.as_slice()
        } else {
            structure.operations()
        };
        let widths = perpendicular_widths(structure.lattice())?;
        let mut grid = FractionalGrid::new(widths, tolerance);
        let mut atoms: Vec<P1Atom> = Vec::new();
        for (site_idx, site) in structure.sites().iter().enumerate() {
            for (op_idx, operation) in operations.iter().enumerate() {
                let base = operation.apply(site.frac());
                let [bx, by, bz] = base;
                let floor_x = bx.floor();
                let floor_y = by.floor();
                let floor_z = bz.floor();
                let frac = [bx - floor_x, by - floor_y, bz - floor_z];
                let tau = [-floor_x as i32, -floor_y as i32, -floor_z as i32];
                let mut duplicate = false;
                grid.for_each_candidate(site.element(), frac, |existing_index| {
                    let Some(existing) = atoms.get(existing_index) else {
                        return true;
                    };
                    if structure.lattice().distance(existing.frac, frac) <= tolerance {
                        duplicate = true;
                        return false;
                    }
                    true
                });
                if duplicate {
                    continue;
                }
                let atom_index = atoms.len();
                atoms.push(P1Atom {
                    element: site.element().to_owned(),
                    label: site.label().to_owned(),
                    frac,
                    occupancy: site.occupancy(),
                    site_idx,
                    op_idx,
                    tau,
                });
                grid.insert(site.element(), frac, atom_index);
            }
        }
        Ok(Self {
            lattice: *structure.lattice(),
            atoms,
        })
    }

    /// Number of atoms in the P1 cell.
    #[must_use]
    pub fn n_atoms(&self) -> usize {
        self.atoms.len()
    }

    /// Element symbols in atom order.
    #[must_use]
    pub fn elements(&self) -> Vec<String> {
        self.atoms.iter().map(|atom| atom.element.clone()).collect()
    }

    /// Fractional coordinates in atom order.
    #[must_use]
    pub fn frac(&self) -> Vec<[f64; 3]> {
        self.atoms.iter().map(|atom| atom.frac).collect()
    }

    /// Occupancies in atom order.
    #[must_use]
    pub fn occupancy(&self) -> Vec<f64> {
        self.atoms.iter().map(|atom| atom.occupancy).collect()
    }

    /// Atoms generated by one asymmetric site.
    #[must_use]
    pub fn images_of(&self, site_idx: usize) -> Vec<&P1Atom> {
        self.atoms
            .iter()
            .filter(|atom| atom.site_idx == site_idx)
            .collect()
    }
}

impl Structure {
    /// Expands this structure to P1 with the default tolerance.
    ///
    /// # Errors
    ///
    /// Returns the errors of [`P1Cell::expand`].
    pub fn p1(&self) -> Result<P1Cell, Error> {
        P1Cell::expand(self, 0.05)
    }
}