Skip to main content

sci_form/forcefield/
traits.rs

1/// Core interface for a molecular force-field contribution.
2pub trait ForceFieldContribution: Send + Sync {
3    /// Evaluate the energy term and accumulate its Cartesian gradient contribution.
4    fn evaluate_energy_and_inject_gradient(&self, coords: &[f64], grad: &mut [f64]) -> f64;
5}
6
7pub struct MolecularForceField {
8    pub iter_terms: Vec<Box<dyn ForceFieldContribution>>,
9}
10
11impl Default for MolecularForceField {
12    fn default() -> Self {
13        Self::new()
14    }
15}
16
17impl MolecularForceField {
18    pub fn new() -> Self {
19        MolecularForceField {
20            iter_terms: Vec::new(),
21        }
22    }
23
24    pub fn insert_dynamic_term(&mut self, term: Box<dyn ForceFieldContribution>) {
25        self.iter_terms.push(term);
26    }
27
28    /// Evaluate all registered terms for the current geometry.
29    pub fn compute_system_energy_and_gradients(&self, coords: &[f64], grad: &mut [f64]) -> f64 {
30        grad.fill(0.0);
31
32        #[cfg(feature = "parallel")]
33        {
34            use rayon::prelude::*;
35
36            let (total_energy, total_grad) = self
37                .iter_terms
38                .par_iter()
39                .map(|term| {
40                    let mut local_grad = vec![0.0; grad.len()];
41                    let energy = term.evaluate_energy_and_inject_gradient(coords, &mut local_grad);
42                    (energy, local_grad)
43                })
44                .reduce(
45                    || (0.0, vec![0.0; grad.len()]),
46                    |mut left, right| {
47                        left.0 += right.0;
48                        for (dst, src) in left.1.iter_mut().zip(right.1.into_iter()) {
49                            *dst += src;
50                        }
51                        left
52                    },
53                );
54
55            grad.copy_from_slice(&total_grad);
56            total_energy
57        }
58
59        #[cfg(not(feature = "parallel"))]
60        {
61            let mut total_energy = 0.0;
62            for term in &self.iter_terms {
63                total_energy += term.evaluate_energy_and_inject_gradient(coords, grad);
64            }
65
66            total_energy
67        }
68    }
69}