Skip to main content

uff_relax/
cell.rs

1use glam::DMat3;
2use glam::DVec3;
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub enum CellType {
7    None,
8    Orthorhombic { size: DVec3 },
9    Triclinic { matrix: DMat3 },
10}
11
12#[derive(Debug, Clone)]
13pub struct UnitCell {
14    pub cell_type: CellType,
15    // Matrix where columns are cell vectors a, b, c
16    matrix: DMat3,
17    // Inverse matrix for fractional coordinate conversion
18    inv_matrix: DMat3,
19}
20
21impl UnitCell {
22    pub fn new_none() -> Self {
23        Self {
24            cell_type: CellType::None,
25            matrix: DMat3::IDENTITY,
26            inv_matrix: DMat3::IDENTITY,
27        }
28    }
29
30    pub fn new_orthorhombic(size: DVec3) -> Self {
31        let matrix = DMat3::from_cols(
32            DVec3::new(size.x, 0.0, 0.0),
33            DVec3::new(0.0, size.y, 0.0),
34            DVec3::new(0.0, 0.0, size.z),
35        );
36        Self {
37            cell_type: CellType::Orthorhombic { size },
38            matrix,
39            inv_matrix: matrix.inverse(),
40        }
41    }
42
43    pub fn new_triclinic(matrix: DMat3) -> Self {
44        Self {
45            cell_type: CellType::Triclinic { matrix },
46            matrix,
47            inv_matrix: matrix.inverse(),
48        }
49    }
50
51    /// Returns the shortest displacement vector from p2 to p1 considering PBC.
52    pub fn distance_vector(&self, p1: DVec3, p2: DVec3) -> DVec3 {
53        let mut diff = p1 - p2;
54        match self.cell_type {
55            CellType::None => diff,
56            CellType::Orthorhombic { size } => {
57                if diff.x > size.x * 0.5 { diff.x -= size.x; }
58                else if diff.x < -size.x * 0.5 { diff.x += size.x; }
59                
60                if diff.y > size.y * 0.5 { diff.y -= size.y; }
61                else if diff.y < -size.y * 0.5 { diff.y += size.y; }
62                
63                if diff.z > size.z * 0.5 { diff.z -= size.z; }
64                else if diff.z < -size.z * 0.5 { diff.z += size.z; }
65                diff
66            }
67            CellType::Triclinic { .. } => {
68                // Convert to fractional coordinates
69                let f_diff = self.inv_matrix * diff;
70                // Apply PBC in fractional space [-0.5, 0.5]
71                let f_diff_pbc = DVec3::new(
72                    f_diff.x - f_diff.x.round(),
73                    f_diff.y - f_diff.y.round(),
74                    f_diff.z - f_diff.z.round(),
75                );
76                // Convert back to Cartesian
77                self.matrix * f_diff_pbc
78            }
79        }
80    }
81
82    pub fn matrix(&self) -> DMat3 {
83        self.matrix
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn test_orthorhombic_distance() {
93        let size = DVec3::new(10.0, 10.0, 10.0);
94        let cell = UnitCell::new_orthorhombic(size);
95        
96        let p1 = DVec3::new(1.0, 1.0, 1.0);
97        let p2 = DVec3::new(9.0, 9.0, 9.0);
98        
99        let dist_vec = cell.distance_vector(p1, p2);
100        // Minimum image convention: dist should be (-2, -2, -2) or (2, 2, 2) length-wise
101        assert!((dist_vec.length() - (3.0 * 2.0f64.powi(2)).sqrt()).abs() < 1e-9);
102    }
103
104    #[test]
105    fn test_none_distance() {
106        let cell = UnitCell::new_none();
107        let p1 = DVec3::new(1.0, 1.0, 1.0);
108        let p2 = DVec3::new(9.0, 9.0, 9.0);
109        let dist_vec = cell.distance_vector(p1, p2);
110        assert!((dist_vec.length() - (3.0 * 8.0f64.powi(2)).sqrt()).abs() < 1e-9);
111    }
112}