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: DMat3,
17 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 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 let f_diff = self.inv_matrix * diff;
70 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 self.matrix * f_diff_pbc
78 }
79 }
80 }
81
82 pub fn matrix(&self) -> DMat3 {
83 self.matrix
84 }
85
86 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 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}