1use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct ConnectionPoint {
12 pub position: [f64; 3],
14 pub direction: [f64; 3],
16 pub kind: String,
18}
19
20#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
22pub enum CoordinationGeometry {
23 Linear,
24 Trigonal,
25 Tetrahedral,
26 SquarePlanar,
27 Octahedral,
28}
29
30impl CoordinationGeometry {
31 pub fn num_connections(self) -> usize {
33 match self {
34 CoordinationGeometry::Linear => 2,
35 CoordinationGeometry::Trigonal => 3,
36 CoordinationGeometry::Tetrahedral => 4,
37 CoordinationGeometry::SquarePlanar => 4,
38 CoordinationGeometry::Octahedral => 6,
39 }
40 }
41
42 pub fn ideal_directions(self) -> Vec<[f64; 3]> {
44 match self {
45 CoordinationGeometry::Linear => vec![[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]],
46 CoordinationGeometry::Trigonal => {
47 let s3 = (3.0f64).sqrt() / 2.0;
48 vec![[1.0, 0.0, 0.0], [-0.5, s3, 0.0], [-0.5, -s3, 0.0]]
49 }
50 CoordinationGeometry::Tetrahedral => {
51 let c = 1.0 / (3.0f64).sqrt();
52 vec![[c, c, c], [c, -c, -c], [-c, c, -c], [-c, -c, c]]
53 }
54 CoordinationGeometry::SquarePlanar => {
55 vec![
56 [1.0, 0.0, 0.0],
57 [0.0, 1.0, 0.0],
58 [-1.0, 0.0, 0.0],
59 [0.0, -1.0, 0.0],
60 ]
61 }
62 CoordinationGeometry::Octahedral => {
63 vec![
64 [1.0, 0.0, 0.0],
65 [-1.0, 0.0, 0.0],
66 [0.0, 1.0, 0.0],
67 [0.0, -1.0, 0.0],
68 [0.0, 0.0, 1.0],
69 [0.0, 0.0, -1.0],
70 ]
71 }
72 }
73 }
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize)]
78pub struct Sbu {
79 pub label: String,
81 pub elements: Vec<u8>,
83 pub positions: Vec<[f64; 3]>,
85 pub connections: Vec<ConnectionPoint>,
87 pub geometry: Option<CoordinationGeometry>,
89}
90
91impl Sbu {
92 pub fn new(label: &str) -> Self {
94 Self {
95 label: label.to_string(),
96 elements: Vec::new(),
97 positions: Vec::new(),
98 connections: Vec::new(),
99 geometry: None,
100 }
101 }
102
103 pub fn add_atom(&mut self, element: u8, position: [f64; 3]) {
105 self.elements.push(element);
106 self.positions.push(position);
107 }
108
109 pub fn add_connection(&mut self, position: [f64; 3], direction: [f64; 3], kind: &str) {
111 let mag = (direction[0].powi(2) + direction[1].powi(2) + direction[2].powi(2)).sqrt();
112 let normalized = if mag > 1e-12 {
113 [direction[0] / mag, direction[1] / mag, direction[2] / mag]
114 } else {
115 [0.0, 0.0, 1.0]
116 };
117 self.connections.push(ConnectionPoint {
118 position,
119 direction: normalized,
120 kind: kind.to_string(),
121 });
122 }
123
124 pub fn num_atoms(&self) -> usize {
126 self.elements.len()
127 }
128
129 pub fn num_connections(&self) -> usize {
131 self.connections.len()
132 }
133
134 pub fn centroid(&self) -> [f64; 3] {
136 let n = self.positions.len() as f64;
137 if n < 1.0 {
138 return [0.0; 3];
139 }
140 let mut c = [0.0; 3];
141 for p in &self.positions {
142 c[0] += p[0];
143 c[1] += p[1];
144 c[2] += p[2];
145 }
146 [c[0] / n, c[1] / n, c[2] / n]
147 }
148
149 pub fn metal_node(element: u8, bond_length: f64, geometry: CoordinationGeometry) -> Self {
155 let mut sbu = Self::new("metal_node");
156 sbu.add_atom(element, [0.0, 0.0, 0.0]);
157 sbu.geometry = Some(geometry);
158
159 for dir in geometry.ideal_directions() {
160 let pos = [
161 dir[0] * bond_length,
162 dir[1] * bond_length,
163 dir[2] * bond_length,
164 ];
165 sbu.add_connection(pos, dir, "metal");
166 }
167
168 sbu
169 }
170
171 pub fn linear_linker(elements: &[u8], spacing: f64, kind: &str) -> Self {
175 let mut sbu = Self::new("linear_linker");
176 let n = elements.len();
177 let total_len = (n - 1) as f64 * spacing;
178
179 for (i, &e) in elements.iter().enumerate() {
180 let x = i as f64 * spacing - total_len / 2.0;
181 sbu.add_atom(e, [x, 0.0, 0.0]);
182 }
183
184 let half = total_len / 2.0 + spacing / 2.0;
186 sbu.add_connection([-half, 0.0, 0.0], [-1.0, 0.0, 0.0], kind);
187 sbu.add_connection([half, 0.0, 0.0], [1.0, 0.0, 0.0], kind);
188
189 sbu
190 }
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196
197 #[test]
198 fn test_metal_node_tetrahedral() {
199 let sbu = Sbu::metal_node(30, 2.0, CoordinationGeometry::Tetrahedral); assert_eq!(sbu.num_atoms(), 1);
201 assert_eq!(sbu.num_connections(), 4);
202 assert_eq!(sbu.elements[0], 30);
203 }
204
205 #[test]
206 fn test_metal_node_octahedral() {
207 let sbu = Sbu::metal_node(26, 2.1, CoordinationGeometry::Octahedral); assert_eq!(sbu.num_connections(), 6);
209 }
210
211 #[test]
212 fn test_linear_linker() {
213 let linker = Sbu::linear_linker(&[6, 6, 6, 6, 6, 6], 1.4, "carboxylate"); assert_eq!(linker.num_atoms(), 6);
215 assert_eq!(linker.num_connections(), 2);
216 }
217
218 #[test]
219 fn test_connection_direction_normalized() {
220 let mut sbu = Sbu::new("test");
221 sbu.add_connection([0.0, 0.0, 0.0], [3.0, 4.0, 0.0], "test");
222 let dir = sbu.connections[0].direction;
223 let mag = (dir[0] * dir[0] + dir[1] * dir[1] + dir[2] * dir[2]).sqrt();
224 assert!(
225 (mag - 1.0).abs() < 1e-10,
226 "Direction should be unit vector, got mag={mag:.6}"
227 );
228 }
229
230 #[test]
231 fn test_sbu_centroid() {
232 let mut sbu = Sbu::new("test");
233 sbu.add_atom(6, [0.0, 0.0, 0.0]);
234 sbu.add_atom(6, [2.0, 0.0, 0.0]);
235 let c = sbu.centroid();
236 assert!((c[0] - 1.0).abs() < 1e-10);
237 assert!((c[1]).abs() < 1e-10);
238 }
239
240 #[test]
241 fn test_coordination_num_connections() {
242 assert_eq!(CoordinationGeometry::Linear.num_connections(), 2);
243 assert_eq!(CoordinationGeometry::Tetrahedral.num_connections(), 4);
244 assert_eq!(CoordinationGeometry::Octahedral.num_connections(), 6);
245 }
246}