sci_form/forcefield/
dg_terms.rs

1use super::traits::ForceFieldContribution;
2
3/// Distance Violation contribution for Distance Geometry refinement.
4/// E = weight * (d/ub - 1)^2 if d > ub, or weight * (2*lb/(lb+d) - 1)^2 if d < lb.
5pub struct DistanceViolation {
6    pub atom_i_idx: usize,
7    pub atom_j_idx: usize,
8    pub lower_bound: f64,
9    pub upper_bound: f64,
10    pub weight: f64,
11}
12
13impl ForceFieldContribution for DistanceViolation {
14    fn evaluate_energy_and_inject_gradient(&self, coords: &[f64], grad: &mut [f64]) -> f64 {
15        let root_i = self.atom_i_idx * 3;
16        let root_j = self.atom_j_idx * 3;
17
18        let dx = coords[root_i] - coords[root_j];
19        let dy = coords[root_i + 1] - coords[root_j + 1];
20        let dz = coords[root_i + 2] - coords[root_j + 2];
21
22        let d2 = dx * dx + dy * dy + dz * dz;
23        let d = d2.sqrt();
24
25        let ub = self.upper_bound;
26        let lb = self.lower_bound;
27        let ub2 = ub * ub;
28        let lb2 = lb * lb;
29
30        if d2 > ub2 {
31            let val = (d2 / ub2) - 1.0;
32            let energy = self.weight * val * val;
33            let pre_factor = 4.0 * self.weight * val / ub2;
34
35            grad[root_i] += pre_factor * dx;
36            grad[root_i + 1] += pre_factor * dy;
37            grad[root_i + 2] += pre_factor * dz;
38
39            grad[root_j] -= pre_factor * dx;
40            grad[root_j + 1] -= pre_factor * dy;
41            grad[root_j + 2] -= pre_factor * dz;
42
43            energy
44        } else if d2 < lb2 && d > 1e-8 {
45            let l2d2 = lb2 + d2;
46            let val = (2.0 * lb2 / l2d2) - 1.0;
47            let energy = self.weight * val * val;
48
49            // dE/dd = 2 * weight * (2*lb^2/(lb^2+d^2) - 1) * (-4*lb^2*d / (lb^2+d^2)^2)
50            // pre_factor = dE/dd / d = -8 * weight * lb2 * val / (l2d2 * l2d2)
51            // Note: RDKit uses a slightly different pre-factor for d < lb, but this is the analytical one.
52            let pre_factor = -8.0 * self.weight * lb2 * val / (l2d2 * l2d2);
53
54            grad[root_i] += pre_factor * dx;
55            grad[root_i + 1] += pre_factor * dy;
56            grad[root_i + 2] += pre_factor * dz;
57
58            grad[root_j] -= pre_factor * dx;
59            grad[root_j + 1] -= pre_factor * dy;
60            grad[root_j + 2] -= pre_factor * dz;
61
62            energy
63        } else {
64            0.0
65        }
66    }
67}
68
69/// Chiral Violation contribution for Distance Geometry refinement.
70pub struct ChiralViolation {
71    pub atom_idx: [usize; 4], // [i, j, k, l] -> l is center? No, usually i, j, k are neighbors, l is center
72    pub lower_vol: f64,
73    pub upper_vol: f64,
74    pub weight: f64,
75}
76
77impl ForceFieldContribution for ChiralViolation {
78    fn evaluate_energy_and_inject_gradient(&self, coords: &[f64], grad: &mut [f64]) -> f64 {
79        let i = self.atom_idx[0] * 3;
80        let j = self.atom_idx[1] * 3;
81        let k = self.atom_idx[2] * 3;
82        let l = self.atom_idx[3] * 3;
83
84        let v1 = [
85            coords[i] - coords[l],
86            coords[i + 1] - coords[l + 1],
87            coords[i + 2] - coords[l + 2],
88        ];
89        let v2 = [
90            coords[j] - coords[l],
91            coords[j + 1] - coords[l + 1],
92            coords[j + 2] - coords[l + 2],
93        ];
94        let v3 = [
95            coords[k] - coords[l],
96            coords[k + 1] - coords[l + 1],
97            coords[k + 2] - coords[l + 2],
98        ];
99
100        let v2xv3 = [
101            v2[1] * v3[2] - v2[2] * v3[1],
102            v2[2] * v3[0] - v2[0] * v3[2],
103            v2[0] * v3[1] - v2[1] * v3[0],
104        ];
105        let vol = v1[0] * v2xv3[0] + v1[1] * v2xv3[1] + v1[2] * v2xv3[2];
106
107        let pre_factor;
108        let energy;
109        if vol < self.lower_vol {
110            let diff = vol - self.lower_vol;
111            energy = self.weight * diff * diff;
112            pre_factor = 2.0 * self.weight * diff;
113        } else if vol > self.upper_vol {
114            let diff = vol - self.upper_vol;
115            energy = self.weight * diff * diff;
116            pre_factor = 2.0 * self.weight * diff;
117        } else {
118            return 0.0;
119        }
120
121        // dV/dpos1 = v2 x v3
122        grad[i] += pre_factor * v2xv3[0];
123        grad[i + 1] += pre_factor * v2xv3[1];
124        grad[i + 2] += pre_factor * v2xv3[2];
125
126        // dV/dpos2 = v3 x v1
127        let v3xv1 = [
128            v3[1] * v1[2] - v3[2] * v1[1],
129            v3[2] * v1[0] - v3[0] * v1[2],
130            v3[0] * v1[1] - v3[1] * v1[0],
131        ];
132        grad[j] += pre_factor * v3xv1[0];
133        grad[j + 1] += pre_factor * v3xv1[1];
134        grad[j + 2] += pre_factor * v3xv1[2];
135
136        // dV/dpos3 = v1 x v2
137        let v1xv2 = [
138            v1[1] * v2[2] - v1[2] * v2[1],
139            v1[2] * v2[0] - v1[0] * v2[2],
140            v1[0] * v2[1] - v1[1] * v2[0],
141        ];
142        grad[k] += pre_factor * v1xv2[0];
143        grad[k + 1] += pre_factor * v1xv2[1];
144        grad[k + 2] += pre_factor * v1xv2[2];
145
146        // dV/dpos4 = -(dV/dpos1 + dV/dpos2 + dV/dpos3)
147        grad[l] -= pre_factor * (v2xv3[0] + v3xv1[0] + v1xv2[0]);
148        grad[l + 1] -= pre_factor * (v2xv3[1] + v3xv1[1] + v1xv2[1]);
149        grad[l + 2] -= pre_factor * (v2xv3[2] + v3xv1[2] + v1xv2[2]);
150
151        energy
152    }
153}
sci_form/forcefield/dg_terms.rs

sci_form/forcefield/
dg_terms.rs
