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    /// Wraps a position vector back into the primary cell unit.
87    pub fn wrap_vector(&self, mut p: DVec3) -> DVec3 {
88        match self.cell_type {
89            CellType::None => p,
90            CellType::Orthorhombic { size } => {
91                p.x = p.x.rem_euclid(size.x);
92                p.y = p.y.rem_euclid(size.y);
93                p.z = p.z.rem_euclid(size.z);
94                p
95            }
96            CellType::Triclinic { .. } => {
97                let mut f = self.inv_matrix * p;
98                f.x = f.x.rem_euclid(1.0);
99                f.y = f.y.rem_euclid(1.0);
100                f.z = f.z.rem_euclid(1.0);
101                self.matrix * f
102            }
103        }
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn test_orthorhombic_distance() {
113        let size = DVec3::new(10.0, 10.0, 10.0);
114        let cell = UnitCell::new_orthorhombic(size);
115        
116        let p1 = DVec3::new(1.0, 1.0, 1.0);
117        let p2 = DVec3::new(9.0, 9.0, 9.0);
118        
119        let dist_vec = cell.distance_vector(p1, p2);
120        // Minimum image convention: dist should be (-2, -2, -2) or (2, 2, 2) length-wise
121        assert!((dist_vec.length() - (3.0 * 2.0f64.powi(2)).sqrt()).abs() < 1e-9);
122    }
123
124    #[test]
125    fn test_none_distance() {
126        let cell = UnitCell::new_none();
127        let p1 = DVec3::new(1.0, 1.0, 1.0);
128        let p2 = DVec3::new(9.0, 9.0, 9.0);
129        let dist_vec = cell.distance_vector(p1, p2);
130        assert!((dist_vec.length() - (3.0 * 8.0f64.powi(2)).sqrt()).abs() < 1e-9);
131    }
132}