1use nalgebra::DMatrix;
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
12pub struct EeqParams {
13 pub chi: f64,
15 pub eta: f64,
17 pub r_eeq: f64,
19 pub r_cov: f64,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct EeqConfig {
26 pub total_charge: f64,
28 pub regularization: f64,
30}
31
32impl Default for EeqConfig {
33 fn default() -> Self {
34 Self {
35 total_charge: 0.0,
36 regularization: 1e-10,
37 }
38 }
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct EeqChargeResult {
44 pub charges: Vec<f64>,
46 pub coordination_numbers: Vec<f64>,
48 pub total_charge: f64,
50}
51
52pub fn get_eeq_params(z: u8) -> EeqParams {
54 match z {
55 1 => EeqParams {
56 chi: 2.20,
57 eta: 13.6,
58 r_eeq: 0.80,
59 r_cov: 0.32,
60 },
61 5 => EeqParams {
62 chi: 2.04,
63 eta: 8.30,
64 r_eeq: 1.40,
65 r_cov: 0.85,
66 },
67 6 => EeqParams {
68 chi: 2.55,
69 eta: 10.0,
70 r_eeq: 1.30,
71 r_cov: 0.77,
72 },
73 7 => EeqParams {
74 chi: 3.04,
75 eta: 14.5,
76 r_eeq: 1.20,
77 r_cov: 0.75,
78 },
79 8 => EeqParams {
80 chi: 3.44,
81 eta: 13.4,
82 r_eeq: 1.10,
83 r_cov: 0.73,
84 },
85 9 => EeqParams {
86 chi: 3.98,
87 eta: 17.4,
88 r_eeq: 1.00,
89 r_cov: 0.71,
90 },
91 14 => EeqParams {
92 chi: 1.90,
93 eta: 8.15,
94 r_eeq: 1.75,
95 r_cov: 1.17,
96 },
97 15 => EeqParams {
98 chi: 2.19,
99 eta: 10.5,
100 r_eeq: 1.60,
101 r_cov: 1.10,
102 },
103 16 => EeqParams {
104 chi: 2.58,
105 eta: 10.4,
106 r_eeq: 1.50,
107 r_cov: 1.04,
108 },
109 17 => EeqParams {
110 chi: 3.16,
111 eta: 13.0,
112 r_eeq: 1.40,
113 r_cov: 0.99,
114 },
115 35 => EeqParams {
116 chi: 2.96,
117 eta: 11.8,
118 r_eeq: 1.55,
119 r_cov: 1.14,
120 },
121 53 => EeqParams {
122 chi: 2.66,
123 eta: 10.5,
124 r_eeq: 1.70,
125 r_cov: 1.33,
126 },
127 26 => EeqParams {
128 chi: 1.83,
129 eta: 7.90,
130 r_eeq: 1.70,
131 r_cov: 1.24,
132 },
133 29 => EeqParams {
134 chi: 1.90,
135 eta: 7.73,
136 r_eeq: 1.60,
137 r_cov: 1.32,
138 },
139 30 => EeqParams {
140 chi: 1.65,
141 eta: 9.39,
142 r_eeq: 1.65,
143 r_cov: 1.22,
144 },
145 22 => EeqParams {
147 chi: 1.54,
148 eta: 6.83,
149 r_eeq: 1.80,
150 r_cov: 1.47,
151 }, 24 => EeqParams {
153 chi: 1.66,
154 eta: 6.77,
155 r_eeq: 1.75,
156 r_cov: 1.39,
157 }, 25 => EeqParams {
159 chi: 1.55,
160 eta: 7.43,
161 r_eeq: 1.73,
162 r_cov: 1.39,
163 }, 27 => EeqParams {
165 chi: 1.88,
166 eta: 7.86,
167 r_eeq: 1.67,
168 r_cov: 1.26,
169 }, 28 => EeqParams {
171 chi: 1.91,
172 eta: 7.64,
173 r_eeq: 1.65,
174 r_cov: 1.21,
175 }, 44 => EeqParams {
177 chi: 2.20,
178 eta: 7.36,
179 r_eeq: 1.75,
180 r_cov: 1.46,
181 }, 46 => EeqParams {
183 chi: 2.20,
184 eta: 8.34,
185 r_eeq: 1.70,
186 r_cov: 1.39,
187 }, 47 => EeqParams {
189 chi: 1.93,
190 eta: 7.58,
191 r_eeq: 1.72,
192 r_cov: 1.45,
193 }, 78 => EeqParams {
195 chi: 2.28,
196 eta: 8.96,
197 r_eeq: 1.72,
198 r_cov: 1.36,
199 }, 79 => EeqParams {
201 chi: 2.54,
202 eta: 9.23,
203 r_eeq: 1.70,
204 r_cov: 1.36,
205 }, 3 => EeqParams {
208 chi: 0.98,
209 eta: 5.39,
210 r_eeq: 2.10,
211 r_cov: 1.28,
212 }, 11 => EeqParams {
214 chi: 0.93,
215 eta: 5.14,
216 r_eeq: 2.40,
217 r_cov: 1.66,
218 }, 12 => EeqParams {
220 chi: 1.31,
221 eta: 7.65,
222 r_eeq: 2.00,
223 r_cov: 1.41,
224 }, 13 => EeqParams {
226 chi: 1.61,
227 eta: 5.99,
228 r_eeq: 1.90,
229 r_cov: 1.21,
230 }, 19 => EeqParams {
232 chi: 0.82,
233 eta: 4.34,
234 r_eeq: 2.75,
235 r_cov: 2.03,
236 }, 20 => EeqParams {
238 chi: 1.00,
239 eta: 6.11,
240 r_eeq: 2.31,
241 r_cov: 1.76,
242 }, _ => EeqParams {
244 chi: 2.20,
245 eta: 10.0,
246 r_eeq: 1.50,
247 r_cov: 1.00,
248 },
249 }
250}
251
252pub fn fractional_coordination(elements: &[u8], positions: &[[f64; 3]]) -> Vec<f64> {
257 let n = elements.len();
258 let mut cn = vec![0.0; n];
259
260 for i in 0..n {
261 let pi = get_eeq_params(elements[i]);
262 for j in (i + 1)..n {
263 let pj = get_eeq_params(elements[j]);
264 let r_cov_ij = pi.r_cov + pj.r_cov;
265 let dx = positions[i][0] - positions[j][0];
266 let dy = positions[i][1] - positions[j][1];
267 let dz = positions[i][2] - positions[j][2];
268 let r_ij = (dx * dx + dy * dy + dz * dz).sqrt();
269
270 if r_ij < 1e-10 {
271 continue;
272 }
273
274 let f = 1.0 / (1.0 + (-16.0 * (r_cov_ij / r_ij - 1.0)).exp());
275 cn[i] += f;
276 cn[j] += f;
277 }
278 }
279
280 cn
281}
282
283fn gamma_damped(r_ij: f64, r_eeq_i: f64, r_eeq_j: f64) -> f64 {
285 if r_ij < 1e-10 {
286 return 0.0;
287 }
288 let sigma_ij = (r_eeq_i * r_eeq_i + r_eeq_j * r_eeq_j).sqrt();
289 let arg = std::f64::consts::SQRT_2 / sigma_ij * r_ij;
290 erf_approx(arg) / r_ij
291}
292
293pub(crate) fn erf_approx(x: f64) -> f64 {
295 let a1 = 0.254829592;
296 let a2 = -0.284496736;
297 let a3 = 1.421413741;
298 let a4 = -1.453152027;
299 let a5 = 1.061405429;
300 let p = 0.3275911;
301
302 let sign = if x < 0.0 { -1.0 } else { 1.0 };
303 let x = x.abs();
304 let t = 1.0 / (1.0 + p * x);
305 let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp();
306 sign * y
307}
308
309pub fn compute_eeq_charges(
314 elements: &[u8],
315 positions: &[[f64; 3]],
316 config: &EeqConfig,
317) -> EeqChargeResult {
318 let n = elements.len();
319
320 if n == 0 {
321 return EeqChargeResult {
322 charges: vec![],
323 coordination_numbers: vec![],
324 total_charge: config.total_charge,
325 };
326 }
327
328 if n == 1 {
329 return EeqChargeResult {
330 charges: vec![config.total_charge],
331 coordination_numbers: vec![0.0],
332 total_charge: config.total_charge,
333 };
334 }
335
336 let cn = fractional_coordination(elements, positions);
337
338 let dim = n + 1;
339 let mut a = DMatrix::zeros(dim, dim);
340 let mut b_vec = vec![0.0; dim];
341
342 let params: Vec<EeqParams> = elements.iter().map(|&z| get_eeq_params(z)).collect();
343
344 for i in 0..n {
345 a[(i, i)] = params[i].eta + config.regularization;
346
347 for j in (i + 1)..n {
348 let dx = positions[i][0] - positions[j][0];
349 let dy = positions[i][1] - positions[j][1];
350 let dz = positions[i][2] - positions[j][2];
351 let r_ij = (dx * dx + dy * dy + dz * dz).sqrt();
352 let gij = gamma_damped(r_ij, params[i].r_eeq, params[j].r_eeq);
353 a[(i, j)] = gij;
354 a[(j, i)] = gij;
355 }
356
357 a[(i, n)] = 1.0;
358 a[(n, i)] = 1.0;
359
360 let cn_correction = -0.1 * (cn[i] - 2.0);
361 b_vec[i] = -(params[i].chi + cn_correction);
362 }
363
364 b_vec[n] = config.total_charge;
365
366 let b_nalg = nalgebra::DVector::from_vec(b_vec);
367 let solution = a.lu().solve(&b_nalg);
368
369 let charges = match solution {
370 Some(sol) => (0..n).map(|i| sol[i]).collect(),
371 None => vec![0.0; n],
372 };
373
374 let total: f64 = charges.iter().sum();
375
376 EeqChargeResult {
377 charges,
378 coordination_numbers: cn,
379 total_charge: total,
380 }
381}
382
383#[cfg(test)]
384mod tests {
385 use super::*;
386
387 fn water_positions() -> Vec<[f64; 3]> {
388 vec![[0.0, 0.0, 0.0], [0.757, 0.586, 0.0], [-0.757, 0.586, 0.0]]
389 }
390
391 #[test]
392 fn test_coordination_number_water() {
393 let elements = [8, 1, 1];
394 let pos = water_positions();
395 let cn = fractional_coordination(&elements, &pos);
396 assert!(cn[0] > 1.5, "O CN = {}", cn[0]);
397 assert!(cn[1] > 0.5 && cn[1] < 1.5, "H CN = {}", cn[1]);
398 }
399
400 #[test]
401 fn test_eeq_charge_neutrality() {
402 let elements = [6, 6, 8, 1, 1, 1, 1, 1];
403 let pos = [
404 [0.0, 0.0, 0.0],
405 [1.54, 0.0, 0.0],
406 [2.57, 1.03, 0.0],
407 [-0.63, 0.89, 0.0],
408 [-0.63, -0.89, 0.0],
409 [1.54, -0.63, 0.89],
410 [1.54, -0.63, -0.89],
411 [3.52, 0.93, 0.0],
412 ];
413 let config = EeqConfig::default();
414 let result = compute_eeq_charges(&elements, &pos, &config);
415 assert!(
416 result.total_charge.abs() < 0.01,
417 "Charge not neutral: {}",
418 result.total_charge
419 );
420 }
421
422 #[test]
423 fn test_eeq_oxygen_negative() {
424 let elements = [8, 1, 1];
425 let pos = water_positions();
426 let config = EeqConfig::default();
427 let result = compute_eeq_charges(&elements, &pos, &config);
428 assert!(result.charges[0] < 0.0, "O charge = {}", result.charges[0]);
429 assert!(result.charges[1] > 0.0, "H charge = {}", result.charges[1]);
430 }
431
432 #[test]
433 fn test_gamma_damped_asymptotic() {
434 let g = gamma_damped(10.0, 1.0, 1.0);
435 assert!((g - 0.1).abs() < 0.01, "gamma(10) = {}, expected ~0.1", g);
436 }
437}