1use chemx_core::Molecule;
2use chemx_integrals::{GradientError, GradientProvider, IntegralProvider};
3use chemx_linalg::{mat_from_row_major, mat_to_row_major};
4use chemx_scf::{RangeSeparation, XcContributor};
5
6pub type Gradient = Vec<[f64; 3]>;
7
8pub fn hf_gradient<P>(
9 provider: &P,
10 molecule: &Molecule,
11 density_alpha: &[f64],
12 density_beta: &[f64],
13) -> Result<Gradient, GradientError>
14where
15 P: IntegralProvider + GradientProvider,
16{
17 gradient_core(
18 provider,
19 molecule,
20 density_alpha,
21 density_beta,
22 1.0,
23 None,
24 None,
25 )
26}
27
28pub fn ks_gradient<P>(
29 provider: &P,
30 molecule: &Molecule,
31 xc: &dyn XcContributor,
32 density_alpha: &[f64],
33 density_beta: &[f64],
34 restricted: bool,
35) -> Result<Gradient, GradientError>
36where
37 P: IntegralProvider + GradientProvider,
38{
39 let n = provider.n_basis();
40 let xc_grad = xc
41 .gradient(density_alpha, density_beta, restricted)
42 .ok_or_else(|| {
43 GradientError::Backend(
44 "XC contributor exposes no analytic gradient; use finite differences".to_string(),
45 )
46 })?;
47 let contrib = xc.eval(density_alpha, density_beta, n, restricted);
48 let mut grad = gradient_core(
49 provider,
50 molecule,
51 density_alpha,
52 density_beta,
53 xc.exx_fraction(),
54 Some((&contrib.vxc_alpha, &contrib.vxc_beta)),
55 xc.range_separation(),
56 )?;
57 assert_eq!(xc_grad.len(), grad.len(), "XC gradient atom count");
58 for (g, x) in grad.iter_mut().zip(&xc_grad) {
59 for k in 0..3 {
60 g[k] += x[k];
61 }
62 }
63 Ok(grad)
64}
65
66fn gradient_core<P>(
67 provider: &P,
68 molecule: &Molecule,
69 density_alpha: &[f64],
70 density_beta: &[f64],
71 c_x: f64,
72 vxc: Option<(&[f64], &[f64])>,
73 rs: Option<RangeSeparation>,
74) -> Result<Gradient, GradientError>
75where
76 P: IntegralProvider + GradientProvider,
77{
78 let n = provider.n_basis();
79 let nn = n * n;
80 let natom = molecule.len();
81 assert_eq!(density_alpha.len(), nn, "density_alpha must be n²");
82 assert_eq!(density_beta.len(), nn, "density_beta must be n²");
83
84 let mut p = vec![0.0; nn];
85 for i in 0..nn {
86 p[i] = density_alpha[i] + density_beta[i];
87 }
88
89 let hcore = mat_to_row_major(&provider.core_hamiltonian());
90 let jk = provider.build_jk(&[
91 mat_from_row_major(n, density_alpha),
92 mat_from_row_major(n, density_beta),
93 ]);
94 let ja = mat_to_row_major(&jk.coulomb[0]);
95 let jb = mat_to_row_major(&jk.coulomb[1]);
96 let mut ka = mat_to_row_major(&jk.exchange[0]);
97 let mut kb = mat_to_row_major(&jk.exchange[1]);
98 let c_x_fock = match &rs {
99 Some(rs) => {
100 debug_assert!(
101 (c_x - rs.alpha).abs() < 1e-12,
102 "RS hybrid: exx_fraction must equal α"
103 );
104 let klr = provider
105 .build_k_erf(
106 &[
107 mat_from_row_major(n, density_alpha),
108 mat_from_row_major(n, density_beta),
109 ],
110 rs.omega,
111 )
112 .ok_or_else(|| {
113 GradientError::Backend(
114 "range-separated gradient: the integral backend has no erf-attenuated \
115 exchange (build_k_erf declined)"
116 .to_string(),
117 )
118 })?;
119 let klr_a = mat_to_row_major(&klr[0]);
120 let klr_b = mat_to_row_major(&klr[1]);
121 for i in 0..nn {
122 ka[i] = rs.alpha * ka[i] + rs.beta * klr_a[i];
123 kb[i] = rs.alpha * kb[i] + rs.beta * klr_b[i];
124 }
125 1.0
126 }
127 None => c_x,
128 };
129 let mut fa = vec![0.0; nn];
130 let mut fb = vec![0.0; nn];
131 for i in 0..nn {
132 let j_total = ja[i] + jb[i];
133 fa[i] = hcore[i] + j_total - c_x_fock * ka[i];
134 fb[i] = hcore[i] + j_total - c_x_fock * kb[i];
135 }
136 if let Some((va, vb)) = vxc {
137 assert_eq!(va.len(), nn, "vxc_alpha must be n²");
138 assert_eq!(vb.len(), nn, "vxc_beta must be n²");
139 for i in 0..nn {
140 fa[i] += va[i];
141 fb[i] += vb[i];
142 }
143 }
144
145 let wa = triple_product(density_alpha, &fa, density_alpha, n);
146 let wb = triple_product(density_beta, &fb, density_beta, n);
147 let mut w = vec![0.0; nn];
148 for i in 0..nn {
149 w[i] = wa[i] + wb[i];
150 }
151
152 let dt = provider.kinetic_gradient()?;
153 let dv = provider.nuclear_gradient()?;
154 let ds = provider.overlap_gradient()?;
155 assert_eq!(
156 dt.natom(),
157 natom,
158 "integral-derivative atom count {} disagrees with molecule {natom}",
159 dt.natom()
160 );
161
162 let mut grad: Gradient = vec![[0.0; 3]; natom];
163 for (atom, g_atom) in grad.iter_mut().enumerate() {
164 for (axis, g) in g_atom.iter_mut().enumerate() {
165 let bt = dt.block(atom, axis);
166 let bv = dv.block(atom, axis);
167 let bs = ds.block(atom, axis);
168 let mut acc = 0.0;
169 for i in 0..nn {
170 acc += p[i] * (bt[i] + bv[i]) - w[i] * bs[i];
171 }
172 *g = acc;
173 }
174 }
175
176 if let Some(g_ecp) = provider.ecp_gradient_contract(&p)? {
177 assert_eq!(
178 g_ecp.len(),
179 natom,
180 "ECP gradient atom count {} disagrees with molecule {natom}",
181 g_ecp.len()
182 );
183 for (atom, g_atom) in grad.iter_mut().enumerate() {
184 for axis in 0..3 {
185 g_atom[axis] += g_ecp[atom][axis];
186 }
187 }
188 }
189
190 let gamma = two_particle_density(&p, density_alpha, density_beta, n, c_x);
191 let g2e = provider.eri_gradient_contract(&gamma)?;
192 for (atom, g_atom) in grad.iter_mut().enumerate() {
193 for axis in 0..3 {
194 g_atom[axis] += g2e[atom][axis];
195 }
196 }
197
198 if let Some(rs) = &rs {
199 let gamma_lr = exchange_density(density_alpha, density_beta, n, rs.beta);
200 let g_lr = provider
201 .eri_gradient_contract_erf(&gamma_lr, rs.omega)?
202 .ok_or_else(|| {
203 GradientError::Backend(
204 "range-separated gradient: the integral backend has no erf-attenuated \
205 ERI derivative (eri_gradient_contract_erf declined)"
206 .to_string(),
207 )
208 })?;
209 for (atom, g_atom) in grad.iter_mut().enumerate() {
210 for axis in 0..3 {
211 g_atom[axis] += g_lr[atom][axis];
212 }
213 }
214 }
215
216 let zeff = provider.effective_nuclear_charges().unwrap_or_else(|| {
217 molecule
218 .atoms
219 .iter()
220 .map(|a| a.element.z() as f64)
221 .collect()
222 });
223 assert_eq!(zeff.len(), natom, "effective-charge count");
224 let vnn = nuclear_repulsion_gradient(molecule, &zeff);
225 for (atom, g_atom) in grad.iter_mut().enumerate() {
226 for axis in 0..3 {
227 g_atom[axis] += vnn[atom][axis];
228 }
229 }
230
231 Ok(grad)
232}
233
234fn two_particle_density(p: &[f64], da: &[f64], db: &[f64], n: usize, c_x: f64) -> Vec<f64> {
235 let mut gamma = vec![0.0; n * n * n * n];
236 for pi in 0..n {
237 for q in 0..n {
238 let pq = p[pi * n + q];
239 for r in 0..n {
240 let da_pr = da[pi * n + r];
241 let db_pr = db[pi * n + r];
242 let base = ((pi * n + q) * n + r) * n;
243 for s in 0..n {
244 gamma[base + s] = 0.5 * pq * p[r * n + s]
245 - c_x * (0.5 * da_pr * da[q * n + s] + 0.5 * db_pr * db[q * n + s]);
246 }
247 }
248 }
249 }
250 gamma
251}
252
253fn exchange_density(da: &[f64], db: &[f64], n: usize, c: f64) -> Vec<f64> {
254 let mut gamma = vec![0.0; n * n * n * n];
255 for pi in 0..n {
256 for q in 0..n {
257 for r in 0..n {
258 let da_pr = da[pi * n + r];
259 let db_pr = db[pi * n + r];
260 let base = ((pi * n + q) * n + r) * n;
261 for s in 0..n {
262 gamma[base + s] =
263 -c * (0.5 * da_pr * da[q * n + s] + 0.5 * db_pr * db[q * n + s]);
264 }
265 }
266 }
267 }
268 gamma
269}
270
271fn nuclear_repulsion_gradient(mol: &Molecule, charges: &[f64]) -> Gradient {
272 let natom = mol.len();
273 let mut g: Gradient = vec![[0.0; 3]; natom];
274 for (a, g_a) in g.iter_mut().enumerate() {
275 let za = charges[a];
276 let ra = mol.atoms[a].position;
277 for (b, other) in mol.atoms.iter().enumerate() {
278 if a == b {
279 continue;
280 }
281 let zb = charges[b];
282 let rb = other.position;
283 let d = [ra[0] - rb[0], ra[1] - rb[1], ra[2] - rb[2]];
284 let r2 = d[0] * d[0] + d[1] * d[1] + d[2] * d[2];
285 let r3 = r2 * r2.sqrt();
286 let f = za * zb / r3;
287 for k in 0..3 {
288 g_a[k] -= f * d[k];
289 }
290 }
291 }
292 g
293}
294
295fn triple_product(a: &[f64], m: &[f64], b: &[f64], n: usize) -> Vec<f64> {
296 let am = matmul(a, m, n);
297 matmul(&am, b, n)
298}
299
300fn matmul(a: &[f64], b: &[f64], n: usize) -> Vec<f64> {
301 let mut c = vec![0.0; n * n];
302 for i in 0..n {
303 for k in 0..n {
304 let a_ik = a[i * n + k];
305 if a_ik == 0.0 {
306 continue;
307 }
308 let brow = k * n;
309 let crow = i * n;
310 for j in 0..n {
311 c[crow + j] += a_ik * b[brow + j];
312 }
313 }
314 }
315 c
316}
317
318#[cfg(test)]
319mod tests {
320 use super::*;
321 use chemx_core::{Atom, Element};
322
323 #[test]
324 fn vnn_gradient_diatomic() {
325 let mol = Molecule::new(
326 vec![
327 Atom::new(Element::from_z(1).unwrap(), [0.0, 0.0, 0.0]),
328 Atom::new(Element::from_z(1).unwrap(), [0.0, 0.0, 1.4]),
329 ],
330 0,
331 1,
332 );
333 let g = nuclear_repulsion_gradient(&mol, &[1.0, 1.0]);
334 let expected = 1.0 / (1.4 * 1.4);
335 assert!((g[0][2] - expected).abs() < 1e-12, "g0z = {}", g[0][2]);
336 assert!((g[1][2] + expected).abs() < 1e-12, "g1z = {}", g[1][2]);
337 for (axis, (a, b)) in g[0].iter().zip(g[1].iter()).enumerate() {
338 let s = a + b;
339 assert!(s.abs() < 1e-12, "Σg axis {axis} = {s}");
340 }
341 }
342}