diff --git a/README.md b/README.md index ed09b2b16ec8e1446c990c5312a175eec23629c2..bbe05046dd7c4c5ab4ea740d130fa7268fc73cbe 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ - `outputs`: 取值Vec\。用于计算结束后输出结果。可输出的信息包括： - `dipole` 偶极 - `fchk`　 Gaussian程序的fchk文件 + - 双杂化 (PT2 族) 的解析力计算 (`job_type = "force"`) 额外在 fchk 中写入弛豫总密度 `Total MP2 Density` 段；开壳层 (UHF 参考) 另行写入 `Spin MP2 Density` 段 (alpha − beta)。 - `cube_orb` 格点化的轨道文件信息 - `molden`　　结果输出为molden程序的格式 - `geometry` 输出分子结构文件 diff --git a/src/analdrv/multipole/interface.rs b/src/analdrv/multipole/interface.rs index a3b80079e12c5de1ecc473bf87a000b36291dbee..6eca3ec1b215e5f07b3cd6e0613ecd1969698903 100644 --- a/src/analdrv/multipole/interface.rs +++ b/src/analdrv/multipole/interface.rs @@ -174,8 +174,10 @@ fn multipole_interface_r<'a>( // the fchk (head + MO coefficients) is regenerated first, so the appended section // always sits after the librest2fch-written MO coefficients (the content is - // identical to the earlier `post_scf_output` write) - scf_data.save_fchk_of_gaussian(); + // identical to the earlier `post_scf_output` write); the primitives are called + // directly, since `save_fchk_of_gaussian` would also write the DH density dump + scf_data.create_fchk_head(); + scf_data.fchk_write_mo(); // pack the density into the fchk layout: Gaussian AO order (consistent with the // librest2fch-written MO coefficients), lower triangle by column diff --git a/src/analdrv/response/rgfock_interface.rs b/src/analdrv/response/rgfock_interface.rs index 8708768a931c69268751ea5b78480e766e54bd99..8ab997b0ddcfd00d93a13616550eb1ce473edecf 100644 --- a/src/analdrv/response/rgfock_interface.rs +++ b/src/analdrv/response/rgfock_interface.rs @@ -458,6 +458,22 @@ pub fn rgfock_dh_interface<'a>(scf_data: &'a SCF) -> RGFockDH<'a> { // --- RI-PT2 (correlation contribution; an ordinary element of the list) --- // + // the PT2 kernel streams its amplitudes over occupied windows (storage bounded by the + // `nvir * nocc * naux` class): size the windows so the per-window transients stay within a + // few times that class and inside `max_memory` + let (nmo_, naux_) = (mo_coeff.shape()[1], rimatr.to_rstsr_view(&device).shape()[1]); + let nvir = nmo_ - nocc; + let mem_avail = scf_data.mol.ctrl.max_memory.map(|m| m - crate::utilities::memory_batch::detect_used_memory_mb("proc")); + let elems_per_step = 2 * nvir * nvir * nocc + nvir * naux_; + crate::utilities::memory_batch::handle_memory_exceed( + elems_per_step as f64 * 8.0 / 1048576.0, + mem_avail, + scf_data.mol.ctrl.abort_on_mem_exceed, + ); + let index_occ_outer = crate::ri_pt2::occ_batch_index( + nocc, + crate::ri_pt2::occ_batch_step(nocc, elems_per_step, nvir * nocc * naux_, mem_avail, 4.0), + ); // the RI-PT2 working precision follows the `[ri_pt2] fp_mode` control keyword (default FP32; // the pre-transformed f32 `cderi_vox` is not supplied here, so the kernel transforms in f64 // and casts on the fly); all outputs (`gfock_part`, `rdm1_corr`, `e_corr`) are f64 regardless @@ -468,7 +484,7 @@ pub fn rgfock_dh_interface<'a>(scf_data: &'a SCF) -> RGFockDH<'a> { mo_energy.to_owned(), rimatr.to_rstsr_view(&device).into_cow(), None, - vec![0, nocc], + index_occ_outer.clone(), c_os, c_ss, ))), @@ -480,7 +496,7 @@ pub fn rgfock_dh_interface<'a>(scf_data: &'a SCF) -> RGFockDH<'a> { mo_energy.to_owned(), rimatr.to_rstsr_view(&device).into_cow(), None, - vec![0, nocc], + index_occ_outer.clone(), c_os, c_ss, ))), diff --git a/src/analdrv/response/uresp_interface.rs b/src/analdrv/response/uresp_interface.rs index 0c02ee01663c6a8f2dbbe0172c83f653e2df870e..1c22adf734977d49afd438f83ab2278b942a690a 100644 --- a/src/analdrv/response/uresp_interface.rs +++ b/src/analdrv/response/uresp_interface.rs @@ -35,6 +35,28 @@ pub fn scf_xc_func_list_uks(scf_data: &SCF) -> Vec<(f64, LibXCFunctional)> { crate::analdrv::response::rresp_interface::scf_xc_func_list_with_spin(scf_data, LibXCSpin::Polarized) } +/// List of `(scale, functional)` pairs of the final-energy functional's DFT (XC numint) part of a +/// DH (spin-polarized); `None` when the final functional has no DFT part. +/// +/// Spin-polarized sibling of [`dh_xc_func_list`](crate::analdrv::response::rgfock_interface:: +/// dh_xc_func_list). +pub fn dh_xc_func_list_uks(scf_data: &SCF) -> Option> { + let xc_data = &scf_data.mol.xc_data; + let (xc_code, xc_params) = match (&xc_data.dfa_compnt_pos, &xc_data.dfa_paramr_pos) { + (Some(code), Some(_)) if code.is_empty() => return None, + (Some(code), Some(param)) => (code, param), + (None, None) => return None, + _ => panic!("dfa_compnt_pos and dfa_paramr_pos must be present or absent together."), + }; + Some( + xc_code + .iter() + .zip(xc_params.iter()) + .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, LibXCSpin::Polarized))) + .collect_vec(), + ) +} + /// The response (fock/response) objects of all electron-interaction contributions of an /// unrestricted SCF, as an owned list of `URespAPI` trait objects. /// diff --git a/src/grad/mod.rs b/src/grad/mod.rs index a966d82cd63ddb0605dfc1f762893372d4143ca7..e792fb044035655a972e3655522db856e80eab2f 100644 --- a/src/grad/mod.rs +++ b/src/grad/mod.rs @@ -1,3 +1,5 @@ +pub mod rdh; +pub mod udh; pub mod uhf; pub mod rks; pub mod uks; diff --git a/src/grad/rdh.rs b/src/grad/rdh.rs new file mode 100644 index 0000000000000000000000000000000000000000..fdd74a989a7382c9f53003f6af1123712436ee0f --- /dev/null +++ b/src/grad/rdh.rs @@ -0,0 +1,916 @@ +//! Restricted doubly-hybrid (XYG3-family and pure-MP2) analytic gradient. +//! +//! The SCF-iteration functional's gradient is provided by the existing `"SCF"` entry of +//! [`crate::main_driver::eval_force`]; this entry covers the D_r-carrying (correlation) +//! remainder plus the final-functional difference terms, following pyscf-forge +//! `dh/grad/dfdh.py`: +//! +//! - `de_h`: hcore-derivative term `` (forge `_get_gradient_pt2`, first term); +//! - `de_s1`: overlap-Pulay term `` with `W = W_I + W_II + W_III` assembled from the +//! full four-block PT2 generalized Fock, the orbital-energy terms and the SCF response upon +//! the relaxed correlation density (forge `_get_gradient_pt2`, second term); +//! - `de_jk`: `t1 + t2 + t4` of forge `get_gradient_jk`, restricted to the D_r-carrying weights +//! (`C1 = cx C D_r^S` with `cx = dfa_hybrid_scf`; the `0.5 cx_n C D_mo` exchange part belongs +//! to the `"SCF"` entry and the difference term below); +//! - `de_rint`: `4 ` (forge `_get_gradient_pt2`, third term); +//! - `de_k_dh`: exchange-derivative difference upon the SCF density with `factor_k = +//! dfa_hybrid_pos - dfa_hybrid_scf` (the D_scf part of forge `get_gradient_jk`'s `C D_mo` +//! exchange piece; the J part is unchanged); +//! - `de_xc_n`: skeleton derivative of the final functional's XC part upon the SCF density +//! minus the SCF functional's one (forge `_get_gradient_gga` explicit part; the latter is +//! already carried by the `"SCF"` entry's `de_xc`); +//! - `de_resp`: ``, the explicit AO-derivative of the SCF functional's XC +//! potential contracted with the relaxed density (forge `_get_gradient_gga` response part; +//! the fxc-upon-D_r half enters the CP-SCF kernel of the Z-vector instead); +//! - `de_ovlp_dh`: occupied-Pulay difference, `Tr(S1, W_n - dme0)` with `W_n = C_oo F_n(D_scf)_oo +//! C_oo^T` the final-functional energy-weighted density (forge `_get_gradient_enfunc`). +//! +//! The relaxed density itself is the unrelaxed PT2 rdm1 with its vir-occupied block replaced by +//! the DH-combined Z-vector: the Lagrangian carries the formal-SCF (Brillouin-violation) term +//! `4 Cv^T F_n(D_scf) Co` of the final functional on top of the PT2 part. All the +//! final-functional terms vanish for the pure-MP2 family (no `dfa_compnt_pos`, equal hybrid +//! coefficients), which leaves the previous RI-MP2-only behaviour untouched. +//! +//! Every in-loop metric solve of the old per-atom scheme is replaced by an exact reordering: +//! the persistent MO partners are pre-contracted once as `Ȳ = partner % S` (`S` = the solve +//! operator that generated `rimatr`, obtained from `get_j2c_decomp(ctrl.j2c_decomp)`), and each +//! shell batch reduces directly (`rt::vecdot`) against `Ȳ`, never materializing +//! `[naux, 3, x, y]` intermediates. Each partner is stored in a single layout `[i, Q, x]` (the +//! metric-solve output permuted once); both the `u` and the `v` fold slots and the ip2-side +//! reductions consume it in place, so the persistent cost stays at the two `Ȳ` tensors +//! themselves. The metric-derivative (`L_1`) term follows forge's Cholesky-generator identity +//! under the `Cd` convention; under `Eig` the same term is evaluated through the (self-adjoint) +//! Fréchet derivative of `J^{-1/2}`, which folds the weight matrices once so that the raw +//! metric-derivative integrals contract against them directly. + +use crate::analdrv::response::rresp_interface::{rscf_resp_interface, scf_xc_func_list, RRespSCF}; +use crate::analdrv::response::rgfock_interface::{dh_jk_factors, dh_xc_func_list, solve_z_vector}; +use crate::analdrv::response::trait_rgfock::{GFockFlags, RGFockAPI}; +use crate::analdrv::response::trait_rresp::RRespAPI; +use crate::dft::num_int::{NumInt, XCData}; +use crate::dft::numint_matmul::nimatmul::{regroup_grids_by_atom, NIMatmul}; +use crate::dft::numint_matmul::resp_rks::RRespKSNIMatmul; +use crate::dft::xc_deriv::XCType; +use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; +use crate::grad::rhf::{get_grad_dao_ovlp, RIRHFGradient, RIHFGradientFlagsBuilder}; +use crate::grad::rks::get_vxc_rayon_new; +use crate::grad::traits::GradAPI; +use crate::mpi_io::MPIOperator; +use crate::ri_jk::resp_r::RRespRIJK; +use crate::ri_jk::util; +use crate::ri_jk::{get_j2c_decomp, J2CDecompose, J2C_THRESH}; +use crate::ri_pt2::rgfock_pt2::RGFockPT2; +use crate::ri_pt2::{occ_batch_index, occ_batch_step}; +use crate::scf_io::SCF; +use crate::utilities::memory_batch::{blocksize_partition, calc_batch_size, detect_used_memory_mb, handle_memory_exceed}; +use crate::utilities::rstsr_util::{RestTensorToRstsrTsrAPI, RestTensorToRstsrViewAPI, Tsr, TsrView}; +use rest_libcint::prelude::*; +use rstsr::prelude::*; +use std::collections::HashMap; +use tensors::MatrixFull; + +/// Analytic gradient of the PT2-correlation (RI-MP2) energy, restricted case. +pub struct RDHGradient<'a> { + scf_data: &'a SCF, + mpi_operator: &'a Option, + /// Gradient parts, keyed by `de_h`/`de_s1`/`de_jk`/`de_rint`/`de`, each `[3, natm]`. + pub result: HashMap>, + /// PT2 correlation energy (side product of the generalized-Fock evaluation). + pub energy: f64, + /// Total relaxed density (SCF density plus the relaxed correlation density) prepared for the + /// fchk density dump, in fchk section order: `[Total MP2 Density]`. `Some` only when + /// `[ctrl] outputs` requests `fchk`; the same density the multipole task dumps. + pub rdm1_dump: Option>>, +} + +impl<'a> RDHGradient<'a> { + pub fn new(scf_data: &'a SCF, mpi_operator: &'a Option) -> Self { + Self { scf_data, mpi_operator, result: HashMap::new(), energy: 0.0, rdm1_dump: None } + } + + pub fn calc(&mut self) -> &mut Self { + let scf_data = self.scf_data; + let mol_obj = &scf_data.mol; + let device = DeviceBLAS::default(); + assert!(!mol_obj.ctrl.spin_polarization, "RDHGradient supports only restricted (spin-polarization = false) references."); + assert!(mol_obj.start_mo == 0, "RDHGradient does not support frozen core (start_mo)."); + assert!(self.mpi_operator.is_none(), "RDHGradient is not MPI-parallelized yet."); + assert!( + scf_data.mol.xc_data.omega().is_none(), + "RSH doubly-hybrid gradients are not supported." + ); + assert!( + scf_data.mol.geom.ghost_pc_chrg.is_empty() && !scf_data.mol.ctrl.solvent_enabled, + "The DH analytic gradient does not support QMMM point charges or solvent models yet." + ); + assert!( + !matches!(scf_data.mol.ctrl.j2c_decomp.uplo, Lower), + "RDHGradient does not support lower-triangular Cholesky factors (developer option)." + ); + // this driver reads `rimatr` through the full-space pair tables; a storage-level pruned + // tensor stores its retained rows only, so refuse it rather than contract wrong rows + crate::ri_jk::require_unpruned_rimatr(&scf_data.rimatr_pair_map, "the DH analytic gradient (RDHGradient)"); + let rimatr = scf_data.rimatr.as_ref().expect( + "Decomposed ERI (rimatr) not found; the DH analytic gradient requires the streaming/new-driver RI-PT2 engine which keeps it alive.", + ); + let ederi_utp = rt::asarray((&rimatr.0.data, rimatr.0.size, &device)); + let mol = util::get_cint_mol(mol_obj); + let aux = util::get_cint_aux(mol_obj); + let natm = mol_obj.geom.elem.len(); + let naux = aux.nao(); + + let mo_coeff = (&scf_data.eigenvectors[0]).to_rstsr(&device); + let mo_occ = (&scf_data.occupation[0]).to_rstsr(&device); + let mo_energy = (&scf_data.eigenvalues[0]).to_rstsr(&device); + let nmo = mo_coeff.shape()[1]; + let nao = mo_coeff.shape()[0]; + let nocc = mo_occ.view().greater(0).sum(); + let nvir = nmo - nocc; + let nao_tp = nao * (nao + 1) / 2; + let so = rt::slice!(0, nocc); + let sv = rt::slice!(nocc, nmo); + let c0 = mo_coeff.i((.., so.clone())).into_contig(ColMajor); + let cv = mo_coeff.i((.., sv.clone())).into_contig(ColMajor); + let d_hf = (&scf_data.density_matrix[0]).to_rstsr(&device); + + // ── response objects, relaxed correlation density, full four-block gfock ── + let config = mol_obj.ctrl.analdrv.clone().unwrap_or_default(); + let mut resp_objs: RRespSCF = rscf_resp_interface(scf_data, &config); + let [c_os, c_ss]: [f64; 2] = mol_obj + .xc_data + .dfa_paramr_adv + .clone() + .expect("Doubly-hybrid gradient requires the PT2 spin factors (dfa_paramr_adv).") + .try_into() + .expect("dfa_paramr_adv must have exactly two entries."); + + // gates of the final-functional (doubly-hybrid) terms; the pure-MP2 family has none, and + // a bDH whose final functional equals the SCF one is covered by the "SCF" entry entirely + let xc_data = &mol_obj.xc_data; + let delta_hyb = xc_data.dfa_hybrid_pos.unwrap_or(xc_data.dfa_hybrid_scf) - xc_data.dfa_hybrid_scf; + let xc_func_pos = dh_xc_func_list(scf_data); + let xc_n_is_scf = match (&xc_data.dfa_compnt_pos, &xc_data.dfa_paramr_pos) { + (Some(code), Some(param)) => { + code.iter().eq(xc_data.dfa_compnt_scf.iter()) + && param.len() == xc_data.dfa_paramr_scf.len() + && param.iter().zip(xc_data.dfa_paramr_scf.iter()).all(|(a, b)| (a - b).abs() < 1e-12) + }, + _ => true, + }; + let add_delta_k = delta_hyb.abs() > 1.0e-10; + let add_xc_n = xc_func_pos.is_some() && !xc_n_is_scf; + let has_final_functional = add_delta_k || add_xc_n; + let add_resp = has_scf_xc(scf_data); + + // atom-grouped copy of the SCF grid, shared by the final-Fock and response terms + let grids_regrouped = if has_final_functional || add_resp { + let grids = scf_data.grids.as_ref().expect( + "Doubly-hybrid gradient requires the DFT grids (`eval_force` regenerates them; `xdh_calculations` frees them).", + ); + Some(regroup_grids_by_atom( + grids.coordinates.clone(), + grids.weights.clone(), + grids.atm_idx.clone(), + grids.quadrature_weights.clone(), + natm, + )) + } else { + None + }; + + // F_n(D_scf) = h + J - (hyb_pos / 2) K + v_xc[xc_n]; it enters the DH-combined Lagrangian + // (the formal-SCF/Brillouin term `4 Cv^T F_n Co`) and the occupied-Pulay difference + let fock_n = if has_final_functional { + let (coords, weights, atm_idx_grids, quadrature_weights) = grids_regrouped.as_ref().unwrap(); + let (factor_j_dh, factor_k_dh) = dh_jk_factors(scf_data); + let cderi_dh = rimatr.0.to_rstsr_view(&device).into_cow(); + let mut resp_rijk = RRespRIJK::new_with_cderi(factor_j_dh, factor_k_dh, cderi_dh); + let h_core = scf_data.h_core.to_matrixfull().unwrap().to_rstsr(&device); + let mut fock_n = &h_core + resp_rijk.get_fock_rdm(d_hf.view(), true); + if xc_func_pos.is_some() { + let ni_dh = NIMatmul::new(&mol, coords, weights, atm_idx_grids, quadrature_weights); + let mut resp_ks = RRespKSNIMatmul::new(dh_xc_func_list(scf_data).unwrap(), ni_dh, mol_obj.ctrl.print_level > 2); + *&mut fock_n += &resp_ks.get_fock_rdm(d_hf.view(), true); + } + Some(fock_n.into_contig(ColMajor)) + } else { + None + }; + // the PT2 kernel streams its amplitudes over occupied windows (storage bounded by the + // `nvir * nocc * naux` class): size the windows so the per-window transients stay within + // a few times that class and inside `max_memory` + let mem_avail = mol_obj.ctrl.max_memory.map(|m| m - detect_used_memory_mb("proc")); + let abort = mol_obj.ctrl.abort_on_mem_exceed; + let elems_per_step = 2 * nvir * nvir * nocc + nvir * naux; + handle_memory_exceed(elems_per_step as f64 * 8.0 / 1048576.0, mem_avail, abort); + let step_pt2 = occ_batch_step(nocc, elems_per_step, nvir * nocc * naux, mem_avail, 4.0); + if mol_obj.ctrl.print_level > 1 { + println!("DH gradient amplitude windows: PT2 occ step {step_pt2} (nocc {nocc}, nvir {nvir}, naux {naux})."); + } + let index_occ_outer = occ_batch_index(nocc, step_pt2); + let j3c = rimatr.0.to_rstsr_view(&device); + let mut pt2 = RGFockPT2::::new( + mo_coeff.clone(), + mo_occ.clone(), + mo_energy.clone(), + j3c.into_cow(), + None, + index_occ_outer, + c_os, + c_ss, + ); + pt2.dump_g_vix = true; + resp_objs.make_cpscf_preparation(mo_coeff.view(), pt2.mo_occ.view(), pt2.mo_energy.view()); + // full four-block request (also caches `g_vix` and the SCF response upon rdm1_corr), + // then the relaxed density through the DH-level Z-vector + let _ = pt2.make_gfock(Some(&mut resp_objs), GFockFlags::OO | GFockFlags::VV); + let rdm1 = pt2.result["rdm1"].to_owned(); + let rdm1_resp = match &fock_n { + Some(fock_n) => { + // DH-combined Lagrangian: the PT2 part plus the formal-SCF term `4 Cv^T F_n Co` + let lag_vo = pt2.make_lagrangian_vo(&mut resp_objs); + let lag_xc_n: Tsr = 4.0 * (cv.view().t() % fock_n.view() % c0.view()); + let lag_dh = (&lag_vo + &lag_xc_n).into_contig(ColMajor); + let z = solve_z_vector(lag_dh.view(), &mut resp_objs); + let mut r = rdm1.clone(); + r.i_mut((sv.clone(), so.clone())).assign(&z.i((sv.clone(), so.clone()))); + r + }, + None => { + let _ = pt2.make_rdm1_resp(&mut resp_objs); + pt2.result["rdm1_resp"].to_owned() + }, + }; + let gfock_part = pt2.result["gfock_part_full"].to_owned(); + let g_vix = pt2.result["g_vix"].to_owned(); // [nvir, nocc, naux] + self.energy = pt2.e_corr.unwrap_or(0.0); + drop(pt2); + + let d_r_sym = (&rdm1_resp + &rdm1_resp.t()).into_contig(ColMajor); + let d_r_s: Tsr = d_r_sym.mapv(|x| 0.5 * x); + // the exchange (K) part of the SCF functional linearized upon the relaxed density carries + // the SCF hybrid coefficient (forge `C1 = cx C D_r_symm`; unity for the MP2/HF reference) + let cx_scf = mol_obj.xc_data.dfa_hybrid_scf; + let c1_dr: Tsr = (mo_coeff.view() % d_r_s.view()).mapv(|x| cx_scf * x).into_contig(ColMajor); + let d_r_ao = (mo_coeff.view() % d_r_s.view() % mo_coeff.view().t()).into_contig(ColMajor); + // total relaxed density (SCF + relaxed correlation) for the fchk density dump of force + // jobs; only materialized when the fchk output is requested (`SCF::dh_rdm1_resp`) + if mol_obj.ctrl.outputs.iter().any(|output| output.eq("fchk")) { + self.rdm1_dump = Some(vec![to_matrix_full(&d_hf + &d_r_ao)]); + } + + // SCF response upon the *relaxed* correlation density (forge W_III); the cached `axd` + // of the generalized-Fock evaluation acts on the unrelaxed rdm1 instead. + let axd_dr = { + let resp_ao = resp_objs.get_response_rdm(d_r_ao.view(), true); + (mo_coeff.view().t() % resp_ao % mo_coeff.view()).into_contig(ColMajor) + }; + + // ── W = W_I + W_II + W_III (MO), from the axd-free `gfock_part_full` blocks ── + let eps2: Tsr = 2.0 * mo_energy.view(); + let mut w_mo: Tsr = rt::zeros(([nmo, nmo], &device)); + { + let oo = (&gfock_part.i((so.clone(), so.clone())) - &eps2.i((None, so.clone())) * rdm1.i((so.clone(), so.clone()))) + .mapv(|x: f64| -0.5 * x); + let vv = (&gfock_part.i((sv.clone(), sv.clone())) - &eps2.i((None, sv.clone())) * rdm1.i((sv.clone(), sv.clone()))) + .mapv(|x: f64| -0.5 * x); + w_mo.i_mut((so.clone(), so.clone())).assign(&oo); + w_mo.i_mut((sv.clone(), sv.clone())).assign(&vv); + w_mo.i_mut((sv.clone(), so.clone())).assign(&gfock_part.i((so.clone(), sv.clone())).t().mapv(|x: f64| -x)); + } + // forge W_II = -D_r (.) eps with the UNSYMMETRISED relaxed density (its [ov] half is + // zero); the raw (non-doubled) orbital energies enter here + *&mut w_mo -= &rdm1_resp * &mo_energy.view().i((None, ..)); + *&mut w_mo.i_mut((so.clone(), so.clone())) -= 0.5 * axd_dr.i((so.clone(), so.clone())); + let w_ao: Tsr = (mo_coeff.view() % w_mo.view() % mo_coeff.view().t()).into_contig(ColMajor); + + // ── de_h and de_s1 ── + let mut de_h: Tsr = rt::zeros(([3, natm], &device)); + let mut de_s1: Tsr = rt::zeros(([3, natm], &device)); + let tsr_ipovlp = { + let (out, shape) = mol.integrate("int1e_ipovlp", "s1", None).into(); + rt::asarray((out, shape, &device)) + }; + { + let mut hcore_gen = crate::grad::rhf::generator_deriv_hcore(scf_data); + // two half-row contractions of the skeleton overlap derivative (forge S_1_ao) + let s1_mu = (&tsr_ipovlp * &w_ao.i((.., .., None))).sum_axes(1); + let s1_nu = (&tsr_ipovlp * &w_ao.view().t().i((.., .., None))).sum_axes(1); + let ao_slice = mol_obj.aoslice_by_atom(); + for atm in 0..natm { + let h1 = hcore_gen(atm); + de_h.i_mut((.., atm)).assign(&(&h1 * &d_r_ao.i((.., .., None))).sum_axes([0, 1])); + let [_, _, p0, p1] = ao_slice[atm]; + let part = (s1_mu.i(p0..p1).sum_axes(0) + s1_nu.i(p0..p1).sum_axes(0)).mapv(|x: f64| -x); + de_s1.i_mut((.., atm)).assign(&part); + } + } + + // ── metric objects, following the rimatr convention (ctrl) ── + let j2c_decomp = get_j2c_decomp(&aux, &device, mol_obj.ctrl.j2c_decomp); + // `S` = the solve operator that generated the rimatr: `L_inv` (Cd) or `J^-1/2` (Eig) + let s_op: Tsr = match &j2c_decomp { + J2CDecompose::Cd { j2c_l, uplo: Upper, .. } => { + let eye = rt::eye((naux, &device)); + let u_inv = rt::linalg::solve_triangular((j2c_l.view(), eye.view(), Upper)); + u_inv.t().into_contig(ColMajor) + }, + J2CDecompose::Eig { j2c_l_inv, .. } => j2c_l_inv.clone(), + _ => panic!("RDHGradient does not support lower-triangular Cholesky factors (developer option)."), + }; + let tsr_int2c2e_ip1 = { + let (out, shape) = aux.integrate("int2c2e_ip1", "s1", None).into(); + rt::asarray((out, shape, &device)) + }; + // strictly-lower mask with 1/2 diagonal of forge's `generator_L_1`; only the Cholesky + // convention uses it per atom (`Eig` folds the weight matrices once instead, below) + let l1_cd = match &j2c_decomp { + J2CDecompose::Cd { j2c_l, .. } => { + let l_lower = j2c_l.t().into_contig(ColMajor); + let mut l_mask: Tsr = rt::zeros(([naux, naux], &device)); + for i in 0..naux { + for j in 0..i { + l_mask[[i, j]] = 1.0; + } + l_mask[[i, i]] = 0.5; + } + Some((l_lower, l_mask)) + }, + J2CDecompose::Eig { .. } => None, + }; + + // ── persistent MO intermediates; density dots of the rimatr by two GEMVs ── + let mem_avail = mol_obj.ctrl.max_memory.map(|m| m - detect_used_memory_mb("proc")); + let abort = mol_obj.ctrl.abort_on_mem_exceed; + // true peak of the partner intermediates: each source coexists with its metric-solve + // transient, which in turn coexists with the permuted `[i, Q, x]` layout replacing it + // (the partner built first stays resident throughout the second one's construction) + let (n_y, n_g) = ((nocc * nmo * naux) as f64, (nvir * nocc * naux) as f64); + handle_memory_exceed((2.0 * n_y).max(n_y + 2.0 * n_g) * 8.0 / 1048576.0, mem_avail, abort); + let d_hf_tp = crate::grad::rhf::pack_triu_tilde(d_hf.view()); + let d_r_ao_tp = crate::grad::rhf::pack_triu_tilde(d_r_ao.view()); + let y_dot_d: Tsr = (ederi_utp.view().t() % d_hf_tp.view().reshape([nao_tp, 1])).into_shape([naux]); + let y_dot_dr: Tsr = (ederi_utp.view().t() % d_r_ao_tp.view().reshape([nao_tp, 1])).into_shape([naux]); + let mut y_ip: Tsr = rt::zeros(([nocc, nmo, naux], &device)); + let mut r_jk: Tsr = rt::zeros(([naux, naux], &device)); + let mut r_ri: Tsr = rt::zeros(([naux, naux], &device)); + let g2d = g_vix.view().into_shape([nvir * nocc, naux]); // (a i)-packed rows, P columns + for p in 0..naux { + let col = ederi_utp.i((.., p)).unpack_tri(Upper, FlagSymm::Sy); + y_ip.i_mut((.., .., p)).assign(&(c0.view().t() % &col % mo_coeff.view())); + // forge RI-term metric coupling: R2[P, p] = sum_{i a} G[a, i, P] (i a|p) + let m_q = (c0.view().t() % &col % cv.view()).into_contig(ColMajor); // [nocc(i), nvir(a)] + let mt = m_q.view().t().into_contig(ColMajor); // [nvir, nocc], (a i)-packed on flatten + r_ri.i_mut((.., p)).assign(&(g2d.t() % mt.view().reshape(nvir * nocc))); + } + // forge t4 metric coupling: R[P, Q] = sum_{i p q} Y_ip[P, i, p] D_r_s[p, q] Y_ip[Q, i, q] + for i in 0..nocc { + let y2t = y_ip.i((i, .., ..)).into_contig(ColMajor); // [nmo(p), naux(P)] + let a_i = y2t.view().t() % d_r_s.view(); // [naux(P), nmo(q)] + *&mut r_jk += &a_i.view() % y2t.view(); + } + *&mut r_jk *= cx_scf; + // `Eig` metric-derivative weights: the y-dot cross terms and the exchange weight form + // `M_jk`, the RI weight `M_ri`; each is folded once by the Fréchet derivative of + // `J^-1/2`, so the atom loop contracts the raw metric derivatives directly + let l1_eig = match &j2c_decomp { + J2CDecompose::Eig { j2c_e: Some(j2c_e), j2c_v: Some(j2c_v), .. } => { + let thresh = mol_obj.ctrl.j2c_decomp.threshold.unwrap_or(J2C_THRESH); + let outer: Tsr = y_dot_dr.view().i((.., None)) % y_dot_d.view().i((None, ..)); + let m_jk: Tsr = (&outer + &outer.t()).mapv(|x: f64| -x) + 2.0 * r_jk.view(); + let m_ri: Tsr = -4.0 * r_ri.view(); + Some(( + l1_fold_eig(m_jk.view(), j2c_e.view(), j2c_v.view(), thresh), + l1_fold_eig(m_ri.view(), j2c_e.view(), j2c_v.view(), thresh), + )) + }, + J2CDecompose::Eig { .. } => panic!( + "The eigen-decomposed 2c-2e ERI carries no stored eigenpairs; cannot build the DH metric-derivative kernels." + ), + J2CDecompose::Cd { .. } => None, + }; + + // ── pre-contractions `Ȳ = partner % S` in the single kept layout `[i, Q, x]` ── + // one stored layout per partner: the solve result `[(i m), Q]` / `[(a i), Q]` is permuted + // once into `[i, Q, m]` / `[i, Q, a]`, and each source is dropped before the permutation + let ybar_jk: Tsr = { + let solved: Tsr = y_ip.view().into_shape([nocc * nmo, naux]) % s_op.view(); // [(i m), Q] + drop(y_ip); + let ybar = solved.view().into_shape([nocc, nmo, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [i, Q, m] + ybar + }; + let ybar_ri: Tsr = { + let solved: Tsr = g_vix.view().into_shape([nvir * nocc, naux]) % s_op.view(); // [(a i), Q] + drop(g_vix); + let ybar = solved.view().into_shape([nvir, nocc, naux]).transpose([1, 2, 0]).into_contig(ColMajor); // [i, Q, a] + ybar + }; + + // ── derivative-integral passes: shell batches over all atoms, atom scatter in-batch ── + let mol_loc = mol.ao_loc(); + let aux_loc = aux.ao_loc(); + let mol_ao_slice = mol_obj.aoslice_by_atom(); + let aux_slice = mol_obj.make_auxmol_fake().aoslice_by_atom(); + let nbas = mol.nbas(); + let nauxbas = aux.nbas(); + let ncomp = 3usize; + // transient per AO function: raw ip1 + transposed copy + fold scratch + partner fold + let ip1_unit = 6 * nao * naux * ncomp + 2 * naux * nmo; + let shell_batch = calc_batch_size::(ip1_unit, mem_avail, None, None).max(1); + // transient per aux function: raw ip2 + unpacked + folded (g_o, f_jk, f_ri) + let ip2_unit = 3 * nao_tp + 3 * nao * nao + 3 * nocc * (nao + nmo + nvir); + let aux_batch = calc_batch_size::(ip2_unit, mem_avail, None, None).max(1); + handle_memory_exceed(ip1_unit as f64 * 8.0 / 1048576.0, mem_avail, abort); + handle_memory_exceed(ip2_unit as f64 * 8.0 / 1048576.0, mem_avail, abort); + + let mut de_jk: Tsr = rt::zeros(([3, natm], &device)); + let mut de_rint: Tsr = rt::zeros(([3, natm], &device)); + let mut de_jk_acc = vec![0.0f64; 3 * natm]; + let mut de_rint_acc = vec![0.0f64; 3 * natm]; + let mut ydd_sol_atoms: Vec> = Vec::with_capacity(natm); + let mut ydr_sol_atoms: Vec> = Vec::with_capacity(natm); + + // ip1 side: u-slot / v-slot of the derivative integrals, per atom + for atm in 0..natm { + let [shl_a0, shl_a1, _, _] = mol_ao_slice[atm]; + let mut ydd_raw: Tsr = rt::zeros(([naux, ncomp], &device)); + let mut ydr_raw: Tsr = rt::zeros(([naux, ncomp], &device)); + let mut t4_3c = [0.0f64; 3]; + let mut ri_3c = [0.0f64; 3]; + let shl_batches = blocksize_partition(&mol_loc[shl_a0..=shl_a1], shell_batch) + .into_iter() + .map(|[i0, i1]| [shl_a0 + i0, shl_a0 + i1]); + for [s0, s1] in shl_batches { + let ip1 = { + let shl_slices = [[s0, s1], [0, nbas], [0, nauxbas]]; + let (out, shape) = CInt::integrate_cross("int3c2e_ip1", [&mol, &mol, &aux], "s1", shl_slices).into(); + rt::asarray((out, shape, &device)) + }; // [nub, naov, naux, ncomp] + let nub = ip1.shape()[0]; + let naov = ip1.shape()[1]; + // forge t1/t2 density vectors of this batch (`u` rows of the densities) + let ip1_mat = ip1.view().into_shape([nub * naov, naux * ncomp]); + { + let d_rows = d_hf.i((mol_loc[s0]..mol_loc[s1], ..)).into_contig(ColMajor); + let ydd_seg = ip1_mat.view().t() % d_rows.view().reshape([nub * naov, 1]); + *&mut ydd_raw += &ydd_seg.into_shape([naux, ncomp]); + } + { + let d_rows = d_r_ao.i((mol_loc[s0]..mol_loc[s1], ..)).into_contig(ColMajor); + let ydr_seg = ip1_mat.view().t() % d_rows.view().reshape([nub * naov, 1]); + *&mut ydr_raw += &ydr_seg.into_shape([naux, ncomp]); + } + // transposed copy `[naov, (nub naux ncomp)]`: the full-AO axis becomes a GEMM index + let ip1_t = ip1.view().transpose([1, 0, 2, 3]).into_contig(ColMajor); + let ip1_t2 = ip1_t.view().into_shape([naov, nub * naux * ncomp]); + let r = mol_loc[s0]..mol_loc[s1]; + // u-slot (C0 on the atom rows) and v-slot (C1_dr / C0 swapped), t4 then RI term + t4_3c = add_fold_ip1_batched(ip1_t2.view(), nub, c0.i((r.clone(), ..)), c1_dr.view(), ybar_jk.view(), 2.0, t4_3c); + t4_3c = add_fold_ip1_batched_v(ip1_t2.view(), nub, c1_dr.i((r.clone(), ..)), c0.view(), ybar_jk.view(), 2.0, t4_3c); + ri_3c = add_fold_ip1_batched(ip1_t2.view(), nub, c0.i((r.clone(), ..)), cv.view(), ybar_ri.view(), -4.0, ri_3c); + ri_3c = add_fold_ip1_batched_v(ip1_t2.view(), nub, cv.i((r, ..)), c0.view(), ybar_ri.view(), -4.0, ri_3c); + } + ydd_sol_atoms.push(s_op.view() % ydd_raw.view()); + ydr_sol_atoms.push(s_op.view() % ydr_raw.view()); + for t in 0..ncomp { + de_jk_acc[t * natm + atm] += t4_3c[t]; + de_rint_acc[t * natm + atm] += ri_3c[t]; + } + } + + // ip2 side: Q-slot batches spanning all aux shells, scattered by atom afterwards + let mut wdd_all: Tsr = rt::zeros(([naux, ncomp], &device)); + let mut wdr_all: Tsr = rt::zeros(([naux, ncomp], &device)); + let q_batches = blocksize_partition(&aux_loc[0..=nauxbas], aux_batch); + for [q0, q1] in q_batches { + let ip2 = { + let shl_slices = [[0, nbas], [0, nbas], [q0, q1]]; + let (out, shape) = CInt::integrate_cross("int3c2e_ip2", [&mol, &mol, &aux], "s2ij", shl_slices).into(); + rt::asarray((out, shape, &device)) + }; // [nao_tp, nb_q, ncomp] + let nb_q = ip2.shape()[1]; + let ip2_2d = ip2.view().into_shape([nao_tp, nb_q * ncomp]); + // density vectors of this aux batch, scattered into the atom-local rows + let wdd_batch = (ip2_2d.view().t() % d_hf_tp.view().reshape([nao_tp, 1])).into_shape([nb_q, ncomp]); + let wdr_batch = (ip2_2d.view().t() % d_r_ao_tp.view().reshape([nao_tp, 1])).into_shape([nb_q, ncomp]); + let (qa, qb) = (aux_loc[q0], aux_loc[q1]); + wdd_all.i_mut((qa..qb, ..)).assign(&wdd_batch); + wdr_all.i_mut((qa..qb, ..)).assign(&wdr_batch); + // MO folds with C_o first, then the nmo/nvir side (never forming [naux, 3, x, y]) + // two batched half-transforms (the per-(q,t) `c0^T % col % w_right`, batched over the + // aux batch), never forming the `[naux, 3, x, y]` intermediate of the metric solve + let ip2_unp = ip2_2d.unpack_tri(Upper, FlagSymm::Sy); // [mu, nu, (q t)] + let w1 = c0.view().t() % &ip2_unp; // [i, nu, (q t)] + // the batched-matmul output stores the batch axis first; materialize standard + // col-major so the `(i, m)` / `(i, a)` axes pair with the partner's `[i, q, x]` slots + let f_jk = (w1.view() % c1_dr.view()).into_contig(ColMajor); // [i, m, (q t)] + let f_ri = (w1.view() % cv.view()).into_contig(ColMajor); // [i, a, (q t)] + for t in 0..ncomp { + // `r_jk_q[q] = sum_{i m} F_jk[i, m, (q t)] Ȳ[i, q, m]` — the partner is stored once, + // so the two contraction axes are paired individually (`(i m)` is not merged) + let r_jk_q = rt::vecdot( + &ybar_jk.i((.., qa..qb, ..)), + &f_jk.i((.., .., t * nb_q..(t + 1) * nb_q)).transpose([0, 2, 1]), + ([0, 2], [0, 2]), + ); + let r_ri_q = rt::vecdot( + &ybar_ri.i((.., qa..qb, ..)), + &f_ri.i((.., .., t * nb_q..(t + 1) * nb_q)).transpose([0, 2, 1]), + ([0, 2], [0, 2]), + ); + for atm in 0..natm { + let [_, _, auq0, auq1] = aux_slice[atm]; + let (lo, hi) = (qa.max(auq0), qb.min(auq1)); + if lo >= hi { + continue; + } + let f0 = lo - qa; + let f1 = hi - qa; + de_jk_acc[t * natm + atm] += 2.0 * r_jk_q.i(f0..f1).sum_all(); + de_rint_acc[t * natm + atm] -= 4.0 * r_ri_q.i(f0..f1).sum_all(); + } + } + } + + // ── per-atom metric couplings, L_1 terms and assembly ── + for atm in 0..natm { + let [_, _, auq0, auq1] = aux_slice[atm]; + let s_op_cols = s_op.i((.., auq0..auq1)); + let ydd2 = &s_op_cols % wdd_all.i((auq0..auq1, ..)); // [naux, 3] + let ydr2 = &s_op_cols % wdr_all.i((auq0..auq1, ..)); + let mut ydd_l1: Tsr = rt::zeros(([naux, ncomp], &device)); + let mut ydr_l1: Tsr = rt::zeros(([naux, ncomp], &device)); + let mut t4_l1 = [0.0f64; 3]; + let mut ri_l1 = [0.0f64; 3]; + // `Cd`: forge's L_1 generator of this atom, contracted with the solve operator; + // `Eig`: the raw metric derivative contracts against the folded weights directly + for t in 0..ncomp { + match &l1_cd { + Some((l_lower, l_mask)) => { + let m0 = &s_op_cols % tsr_int2c2e_ip1.i((auq0..auq1, .., t)); + let mut m: Tsr = &m0 % s_op.view().t(); + let m_t = m.view().t().into_owned(); + *&mut m += &m_t; + let l1_t = (l_lower.view() % &(l_mask * &m)).mapv(|x: f64| -x); + let l1di = s_op.view() % l1_t.view(); + ydd_l1.i_mut((.., t)).assign(&(l1di.view() % y_dot_d.view())); + ydr_l1.i_mut((.., t)).assign(&(l1di.view() % y_dot_dr.view())); + t4_l1[t] = 2.0 * rt::vecdot(&l1di, r_jk.view(), -1).sum_all(); + ri_l1[t] = -4.0 * rt::vecdot(&l1di, r_ri.view(), -1).sum_all(); + }, + None => { + let (k_jk, k_ri) = + l1_eig.as_ref().expect("the metric-derivative kernels must exist for a non-Cholesky policy."); + let dj = tsr_int2c2e_ip1.i((auq0..auq1, .., t)); + t4_l1[t] = rt::vecdot(&dj, &k_jk.i((auq0..auq1, ..)), -1).sum_all(); + ri_l1[t] = rt::vecdot(&dj, &k_ri.i((auq0..auq1, ..)), -1).sum_all(); + }, + } + } + + // assemble forge t1 + t2 + t4 (D_r-carrying) and the RI term + let ydd_sol = &ydd_sol_atoms[atm]; + let ydr_sol = &ydr_sol_atoms[atm]; + let mut grad = [0.0f64; 3]; + let mut grad_ri = [0.0f64; 3]; + for t in 0..ncomp { + let v_ydd = ydd_sol.i((.., t)); + let v_ydr = ydr_sol.i((.., t)); + let c_ydd: Tsr = ydd2.i((.., t)) + ydd_l1.i((.., t)); + let c_ydr: Tsr = 2.0 * &v_ydr + &ydr2.i((.., t)) + &ydr_l1.i((.., t)); + let s1: f64 = rt::vecdot(&y_dot_d, c_ydr.view(), -1).sum_all(); + let s2: f64 = rt::vecdot(&y_dot_dr, (&c_ydd + &(2.0 * &v_ydd)).view(), -1).sum_all(); + grad[t] = -s1 - s2 + t4_l1[t] + de_jk_acc[t * natm + atm]; + grad_ri[t] = ri_l1[t] + de_rint_acc[t * natm + atm]; + } + de_jk.i_mut((.., atm)).assign(&rt::asarray((grad.as_slice(), [3], &device))); + de_rint.i_mut((.., atm)).assign(&rt::asarray((grad_ri.as_slice(), [3], &device))); + } + + // ── doubly-hybrid (final-functional) terms, all absent for the pure-MP2 family ── + // `de_k_dh`: exchange-derivative difference upon the SCF density (`factor_k = + // dfa_hybrid_pos - dfa_hybrid_scf`; the J part is unchanged and covered by the "SCF" + // entry), i.e. the D_scf part of forge `get_gradient_jk`'s `0.5 cx_n C D_mo` piece; + // `de_xc_n`: XC skeleton difference `xc_n - xc` upon D_scf (forge `_get_gradient_gga` + // explicit part; the SCF-functional one is the "SCF" entry's `de_xc`); + // `de_resp`: ``, the explicit AO-derivative of the SCF functional's XC + // potential (forge `_get_gradient_gga` response part; its fxc-on-D_r half enters the + // CP-SCF kernel of the Z-vector instead); + // `de_ovlp_dh`: occupied-Pulay difference `-2 Tr(S1, W_n) - (-Tr(S1, dme0))` with + // `W_n = C_oo F_n(D_scf)_oo C_oo^T` (forge `_get_gradient_enfunc`), against the + // "SCF" entry's `dme0` term. + let mut de_k_dh: Tsr = rt::zeros(([3, natm], &device)); + let mut de_xc_n: Tsr = rt::zeros(([3, natm], &device)); + let mut de_resp: Tsr = rt::zeros(([3, natm], &device)); + let mut de_ovlp_dh: Tsr = rt::zeros(([3, natm], &device)); + if has_final_functional || add_resp { + if add_delta_k { + let mut grad_helper = dh_grad_helper(scf_data, self.mpi_operator); + grad_helper.flags.factor_k = Some(delta_hyb); + grad_helper.calc_de_jk(); + let de_k = &grad_helper.result["de_k"]; + de_k_dh.assign(&rt::asarray((&de_k.data, de_k.size, &device))); + // the auxiliary-basis-response part of the exchange gradient + if let Some(de_kaux) = grad_helper.result.get("de_kaux") { + *&mut de_k_dh += &rt::asarray((&de_kaux.data, de_kaux.size, &device)); + } + } + if add_xc_n { + // the DH total energy carries the *final* functional's XC skeleton explicitly + // (forge `grad_gga`), while the "SCF" entry contributes the SCF functional's one; + // add their difference + let dao_xc_n = xc_skeleton_dao(scf_data, self.mpi_operator, false, scf_data.grids.as_ref().unwrap()); + let dao_xc_scf = xc_skeleton_dao(scf_data, self.mpi_operator, true, scf_data.grids.as_ref().unwrap()); + let dao_xc: Tsr = get_grad_dao_ovlp((&dao_xc_n - &dao_xc_scf).view(), d_hf.view()); + for atm in 0..natm { + let [_, _, p0, p1] = mol_ao_slice[atm]; + *&mut de_xc_n.i_mut((.., atm)) += &dao_xc.i(p0..p1).sum_axes(0); + } + } + if add_resp { + let (coords, weights, atm_idx_grids, quadrature_weights) = grids_regrouped.as_ref().unwrap(); + let dm0 = util::get_dm0_restricted(mo_coeff.view(), mo_occ.view()); + let mut ni = NIMatmul::new(&mol, coords, weights, atm_idx_grids, quadrature_weights); + // the explicit AO-derivative of the SCF functional's XC potential (forge + // `hessian.rks._get_vxc_deriv1`), contracted with the relaxed density on the fly; + // the grid-shift terms are absent, following the forge reference + de_resp = de_xc_response_term(&mol, &scf_xc_func_list(scf_data), &mut ni, dm0.view(), d_r_ao.view()); + } + if has_final_functional { + let fock_n = fock_n.as_ref().unwrap(); + let f_n_oo = (c0.view().t() % fock_n.view() % c0.view()).into_contig(ColMajor); + let w_n = (c0.view() % &f_n_oo % c0.view().t()).into_contig(ColMajor); + // the "SCF" entry's `de_ovlp` carries the `dme0 = C (n eps) C^T` term, while the + // final-functional Pulay reads `-2 Tr(S1, W_n)` (no occupancy factor); add the + // difference `f(2 W_n - dme0)` only (`f` = the two-half contraction pattern) + let dme0 = crate::grad::rhf::get_dme0(mo_coeff.view(), mo_occ.view(), mo_energy.view()); + let dao_ovlp: Tsr = + get_grad_dao_ovlp(tsr_ipovlp.view(), (&(&w_n * 2.0) - &dme0).view()); + for atm in 0..natm { + let [_, _, p0, p1] = mol_ao_slice[atm]; + *&mut de_ovlp_dh.i_mut((.., atm)) += &dao_ovlp.i(p0..p1).sum_axes(0); + } + } + } + + // ── assemble the parts ── + let de = de_h.view() + de_s1.view() + de_jk.view() + de_rint.view() + de_k_dh.view() + de_xc_n.view() + + de_resp.view() + de_ovlp_dh.view(); + for (key, val) in [ + ("de_h", de_h), + ("de_s1", de_s1), + ("de_jk", de_jk), + ("de_rint", de_rint), + ("de_k_dh", de_k_dh), + ("de_xc_n", de_xc_n), + ("de_resp", de_resp), + ("de_ovlp_dh", de_ovlp_dh), + ("de", de.into_owned()), + ] { + self.result.insert(key.to_string(), to_matrix_full(val)); + } + self + } +} + +/// Response-density XC term `` of the DH gradient (forge `_get_gradient_gga` +/// response part), with `F_1[A, t, mu, nu]` the explicit AO-derivative of the SCF functional's XC +/// potential (the gradient-level analogue of forge `hessian.rks._get_vxc_deriv1`, including its +/// fxc-upon-skeleton-density part; the grid-shift terms are absent, following the forge +/// reference). +/// +/// Chunk-parallel over the (atom-grouped) grid on first-and-second-order AO channels, with the +/// relaxed density contracted immediately, so only the Hessian-level kernels of +/// [`crate::dft::numint_matmul::hess_rks`] are reused and no full-grid +/// `[nao, nao, 3, natm]` accumulator is held. +fn de_xc_response_term( + mol: &CInt, + xc_func_list: &[(f64, libxc::functional::LibXCFunctional)], + ni: &mut NIMatmul, + dm0: TsrView, + d_r_ao: TsrView, +) -> Tsr { + use crate::dft::numint_matmul::hess_rks::{get_drho, get_rho_exc_vxc_fxc, get_vmat_fxc, get_vmat_ip, get_vmat_vxc}; + use crate::dft::xceff::prelude::XCDenType::*; + + let natm = mol.natm(); + let nao = mol.nao(); + let ngrids = ni.weights.len(); + let device = dm0.device().clone(); + let aoslices = mol.aoslice_by_atom(); + let func_refs: Vec<&libxc::functional::LibXCFunctional> = xc_func_list.iter().map(|(_, f)| f).collect(); + let xc_type = determine_den_type_from_list(&func_refs); + // gradient-level AO derivative order (pyscf `grad.rks` convention; the Hessian needs one more) + let ao_deriv = match xc_type { + RHO => 1, + _ => 2, + }; + let ncomp_ao_dm0 = xc_type.num_ao_comp(); + + let de_resp: Tsr = rt::zeros(([3, natm].f(), &device)); + let nchunk_target = (rayon::current_num_threads() * 8).max(1); + let nsplit = ((ngrids + nchunk_target - 1) / nchunk_target).max(1); + let chunks = (0..ngrids).step_by(nsplit).map(|s| (s, (s + nsplit).min(ngrids))).collect::>(); + let guard = std::sync::Mutex::new(()); + use rayon::prelude::*; + let weights_full = ni.weights.clone(); + let d_r_ao_ref = &d_r_ao; + let dm0_owned: Tsr = dm0.into_contig(ColMajor).into_owned(); + let ni_ref = &*ni; + let device_ref = &device; + let aoslices_ref = &aoslices; + chunks.into_par_iter().for_each(|(start, end)| { + let mut ni_chunk = ni_ref.split_batch(start, end); + let weights = rt::asarray((&weights_full[start..end], device_ref)); + let ao = ni_chunk.get_cached_ao(ao_deriv); + let ao_dm0 = ao.i((Ellipsis, ..ncomp_ao_dm0)) % &dm0_owned; + let (_rho, _exc, vxc, fxc) = get_rho_exc_vxc_fxc(xc_func_list, ao.view(), ao_dm0.view()); + let wv = &weights * &vxc; + let wf = &weights * &fxc; + let drho = get_drho(xc_type, ao.view(), ao_dm0.view(), aoslices_ref); + let vmat_ip = get_vmat_ip(xc_type, ao.view(), wv.view()); + let vmat_fxc = get_vmat_fxc(xc_type, ao.view(), drho.view(), wf.view()); + let vmat_vxc = get_vmat_vxc(vmat_ip.view(), aoslices_ref); + let vmat = &vmat_fxc + &vmat_vxc; + let mut de_chunk: Tsr = rt::zeros(([3, natm].f(), &device)); + for atm in 0..natm { + for t in 0..3 { + let v: f64 = rt::vecdot( + &d_r_ao_ref.view().into_shape([nao * nao]), + &vmat.i((.., .., t, atm)).into_shape([nao * nao]), + -1, + ) + .sum_all(); + de_chunk[[t, atm]] = v; + } + } + let _lock = guard.lock().unwrap(); + unsafe { *&mut de_resp.force_mut() += &de_chunk }; + }); + de_resp +} + +/// Device tensor (e.g. a `[3, natm]` gradient part or a `[nao, nao]` density) into a +/// `MatrixFull` of the same shape. +fn to_matrix_full(tsr: Tsr) -> MatrixFull { + let shape: [usize; 2] = tsr.shape().to_vec().try_into().unwrap(); + MatrixFull::from_vec(shape, tsr.into_shape(-1).into_raw()).unwrap() +} + +/// Fold a metric-derivative weight matrix for the eigen convention: the `L_1` term +/// `sum_weights . raw d(J^-1/2)` equals `` with `F` the (self-adjoint) Fréchet +/// derivative of `J^-1/2`, i.e. `V (D ∘ (V^T M^T V)) V^T` with +/// `D_ij = -1 / (sqrt(e_i) (sqrt(e_i) + sqrt(e_j)))`; the returned `G + G^T` pairs with one +/// atom block of the raw `int2c2e_ip1` (its transposed half included). Eigenpairs below +/// `threshold` are discarded, exactly as in the decomposition itself. +fn l1_fold_eig(weight: TsrView, j2c_e: TsrView, j2c_v: TsrView, threshold: f64) -> Tsr { + let n = j2c_e.view().less(threshold).sum(); + let sr = j2c_e.i(n..).pow(0.5); + let v = j2c_v.i((.., n..)).into_contig(ColMajor); + let ssum = &sr.i((.., None)) + &sr.i((None, ..)); + let d: Tsr = (&sr.i((.., None)) * &ssum).mapv(|x: f64| -1.0 / x); + let p = v.view().t() % (weight.t() % v.view()); // V^T M^T V + let g: Tsr = v.view() % &(&d * &p) % v.view().t(); + (&g + &g.t()).into_owned() +} + +/// One ip1 derivative-integral pass on an AO-shell batch, reduced against the pre-contracted +/// partner `Ȳ` (the metric solve is hoisted out of the atom loop). This is the `u` slot, where +/// the atom-side weight contracts the partner's leading axis (`v` slot: +/// [`add_fold_ip1_batched_v`]). +/// +/// Computes `acc[t] += factor sum_{u, Q, y} G[(u Q), y](t) K[(u Q), y]`, where the derivative +/// fold `G[(u Q), y](t) = sum_v ip1_t2[v, (u Q t)] c_col[v, y]` and the partner fold +/// `K[(u Q), y] = sum_x c_row[u, x] ybar[x, Q, y]` share the same contiguous `(u Q)` row order, +/// so the reduction is a last-axis [`rt::vecdot`]. +fn add_fold_ip1_batched( + ip1_t2: TsrView, // [naov, nub naux ncomp] + nub: usize, + c_row: TsrView, // [nub, x], the atom-side MO weight + c_col: TsrView, // [naov, y], the full-AO-side MO weight + ybar: TsrView, // [x, naux, y], pre-contracted partner + factor: f64, + mut acc: [f64; 3], +) -> [f64; 3] { + let device = ip1_t2.device().clone(); + let naov = ip1_t2.shape()[0]; + let naux = ip1_t2.shape()[1] / (nub * 3); + let x = c_row.shape()[1]; + let y = c_col.shape()[1]; + let mut k_buf: Tsr = rt::zeros(([nub, naux * y], &device)); + k_buf.view_mut().matmul_from(&c_row, &ybar.reshape([x, naux * y]), 1.0, 0.0); + let k_buf_view = k_buf.view(); + let k2 = k_buf_view.reshape([nub * naux, y]); + let mut g_buf: Tsr = rt::zeros(([nub * naux, y], &device)); + for t in 0..3 { + let blk = ip1_t2.i((.., t * nub * naux..(t + 1) * nub * naux)); + g_buf.view_mut().matmul_from(&blk.t(), &c_col, 1.0, 0.0); + let s: f64 = rt::vecdot(&g_buf, &k2, -1).sum_all(); + acc[t] += factor * s; + } + acc +} + +/// [`add_fold_ip1_batched`] counterpart for the `v` slot, where the partner `Ȳ[i, Q, m]` is +/// contracted over its slow `m` axis instead: the same two GEMMs are kept, but the partner +/// enters on the left of a single `[nocc naux, nmo] % [nmo, nub]` fold while the full-AO-side +/// fold is assembled from one `[y, naov] % [naov, naux]` GEMM per batch row `u` (each landing +/// contiguously in the paired buffer). Both stay within the `u` slot's GEMM structure and the +/// reduction remains a `rt::vecdot` over matching layouts, so no second partner layout is +/// stored. +fn add_fold_ip1_batched_v( + ip1_t2: TsrView, // [naov, nub naux ncomp] + nub: usize, + c_row: TsrView, // [nub, x], the atom-side MO weight (contracted partner axis) + c_col: TsrView, // [naov, y], the full-AO-side MO weight (partner's leading axis) + ybar: TsrView, // [y, naux, x], pre-contracted partner in the single kept layout + factor: f64, + mut acc: [f64; 3], +) -> [f64; 3] { + let device = ip1_t2.device().clone(); + let naov = ip1_t2.shape()[0]; + let naux = ip1_t2.shape()[1] / (nub * 3); + let x = c_row.shape()[1]; + let y = c_col.shape()[1]; + let mut k_buf: Tsr = rt::zeros(([y * naux, nub], &device)); + k_buf.view_mut().matmul_from(&ybar.reshape([y * naux, x]), &c_row.t(), 1.0, 0.0); + let mut g_buf: Tsr = rt::zeros(([y, naux, nub], &device)); + let c_col_t = c_col.t(); + let ip1_t3 = ip1_t2.reshape([naov, nub, naux * 3]); // [naov, u, (Q t)] + for t in 0..3 { + for u in 0..nub { + let blk = ip1_t3.i((.., u, t * naux..(t + 1) * naux)); + g_buf.i_mut((.., .., u)).matmul_from(&c_col_t, &blk, 1.0, 0.0); + } + let s: f64 = rt::vecdot(&k_buf, &g_buf.view().reshape([y * naux, nub]), 0).sum_all(); + acc[t] += factor * s; + } + acc +} + +impl GradAPI for RDHGradient<'_> { + fn get_gradient(&self) -> MatrixFull { + self.result["de"].clone() + } + fn get_energy(&self) -> f64 { + self.energy + } +} + +/// Whether the SCF-iteration functional carries a non-HF XC part (the gate of the +/// response-density XC term; the kernels handle LDA/GGA/MGGA alike, unlike forge's GGA-only gate). +fn has_scf_xc(scf_data: &SCF) -> bool { + !scf_data.mol.xc_data.dfa_compnt_scf.is_empty() +} + +/// A bare gradient helper carrying the SCF control flags, for reuse of the gradient kernels +/// (`calc_de_jk` with custom J/K factors, `get_vxc_rayon_new`). +fn dh_grad_helper<'a>(scf_data: &'a SCF, mpi_operator: &'a Option) -> RIRHFGradient<'a> { + let mol_obj = &scf_data.mol; + let flags = RIHFGradientFlagsBuilder::default() + .print_level(mol_obj.ctrl.print_level) + .max_memory(mol_obj.ctrl.max_memory) + .auxbasis_response(mol_obj.ctrl.auxbasis_response) + .factor_j(None) + .factor_k(None) + .build() + .unwrap(); + RIRHFGradient { scf_data, flags, mpi_operator, result: HashMap::new() } +} + +/// AO-derivative XC potential contraction `dao_vxc[nao, 3]` of one functional (the row-scatter, +/// contracted with the SCF density and the two-half factor, is its skeleton-gradient +/// contribution); `scf` selects the SCF-iteration functional, else the final (`pos`) one. +fn xc_skeleton_dao(scf_data: &SCF, mpi_operator: &Option, scf: bool, grids: &crate::dft::Grids) -> Tsr { + let mol_obj = &scf_data.mol; + let xc_data = &mol_obj.xc_data; + let (code, param) = if scf { + (&xc_data.dfa_compnt_scf, &xc_data.dfa_paramr_scf) + } else { + ( + xc_data.dfa_compnt_pos.as_ref().expect("the final functional has no XC component list."), + xc_data.dfa_paramr_pos.as_ref().expect("the final functional has no XC parameter list."), + ) + }; + let xc_type = if code.iter().any(|&c| xc_data.init_libxc(&c).needs_tau()) { + XCType::MGGA + } else if code + .iter() + .any(|&c| !matches!(xc_data.init_libxc(&c).family(), libxc::enums::LibXCFamily::LDA | libxc::enums::LibXCFamily::HybLDA)) + { + XCType::GGA + } else { + XCType::LDA + }; + let xc_data_sel = XCData { + xc_type, + xc_code: code, + xc_params: param, + device: DeviceOpenBLAS::default(), + dm: &scf_data.density_matrix, + mo_coeffs: None, + occ: None, + }; + let mut grids_clone = grids.clone(); + get_vxc_rayon_new(&dh_grad_helper(scf_data, mpi_operator), &xc_data_sel, &mut grids_clone, mol_obj, 16) +} diff --git a/src/grad/rks.rs b/src/grad/rks.rs index df5b45a7d76fec8fd01b33b82beb99d69d4fdf5a..342f8a83297f46451cbd8953c7991f7ebf89c1bb 100644 --- a/src/grad/rks.rs +++ b/src/grad/rks.rs @@ -640,7 +640,14 @@ fn get_vxc_rayon(gradient_method: &RIRHFGradient, xc_data: &XCData, grids: &mut } -fn get_vxc_rayon_new(gradient_method: &RIRHFGradient, xc_data: &XCData, grids: &mut Grids, mol: &Molecule, max_memory: usize) -> Tsr { +/// Grid contribution of the XC skeleton gradient, `dao_vxc[nao, nao, 3]` (the AO-derivative +/// potential contraction whose row-scatter, times `2`, is the XC skeleton gradient), for the +/// XC functional carried by `xc_data`. +/// +/// Public so that post-SCF (doubly-hybrid) gradients can evaluate the skeleton derivative of +/// the *final* functional, whose component list is not the SCF one (`xc_data` is built by the +/// caller, e.g. from `dfa_compnt_pos` / `dfa_paramr_pos`). +pub fn get_vxc_rayon_new(gradient_method: &RIRHFGradient, xc_data: &XCData, grids: &mut Grids, mol: &Molecule, max_memory: usize) -> Tsr { let default_omp_num_threads = omp_get_num_threads_wrapper(); let num_grids = grids.weights.len(); diff --git a/src/grad/udh.rs b/src/grad/udh.rs new file mode 100644 index 0000000000000000000000000000000000000000..8b755380a69b94c91038b51bc3461fa6262e9db1 --- /dev/null +++ b/src/grad/udh.rs @@ -0,0 +1,1172 @@ +//! Unrestricted doubly-hybrid (pure-MP2 family) analytic gradient; the unrestricted twin of +//! [`crate::grad::rdh`], following pyscf-forge `dh/grad/udfdh.py`. +//! +//! The SCF-iteration functional's gradient is provided by the existing `"SCF"` entry of +//! [`crate::main_driver::eval_force`]; this entry covers the D_r-carrying (correlation) +//! remainder, per spin $\sigma$ (`de_*` keys mirror the restricted module): +//! +//! - `de_h`: hcore-derivative term $\Sigma_\sigma \langle H_1(A), D_r^\sigma\rangle$; +//! - `de_s1`: overlap-Pulay term $\Sigma_\sigma \langle W^\sigma, S_1(A)\rangle$, with $W$ the +//! unrestricted `W_I + W_II + W_III` assembly (forge `prepare_lagrangian`/`_get_gradient_pt2`); +//! - `de_jk`: `t1 + t2 + t4` of forge `get_gradient_jk`, restricted to the D_r-carrying weights: +//! the RI-J bilinears `t1`/`t2` run over all four spin pairs (the total SCF density meets the +//! per-spin relaxed densities), while the exchange `t4` is per spin with +//! `C1^s = cx C^s D_r^{S,s}` (`cx = dfa_hybrid_scf`); +//! - `de_rint`: $\Sigma_\sigma \langle G^\sigma, Y_1^\sigma(A)\rangle$ with the unrestricted +//! coefficient `1` (restricted: `4`), `G` carrying forge's `4`/`2` spin-block normalization. +//! +//! The relaxed density is solved per spin through the coupled CP-SCF of +//! [`uscf_resp_interface`](crate::analdrv::response::uresp_interface::uscf_resp_interface) (no +//! $\alpha\beta$ exchange in the kernel), driven by the DH-combined Lagrangian +//! $L^\sigma = \mathscr{F}^{\sigma}_{vo} - \mathscr{F}^{\sigma}_{ov\top} + A^\sigma(D^-)_{vo} +//! + 2 C_v^{\sigma\top} F_n^\sigma(D^\sigma) C_o^\sigma$. +//! +//! The final-functional (XYG3-family) terms mirror the restricted module +//! [`crate::grad::rdh`], per spin $\sigma$ and all absent for the pure-MP2 family: +//! `de_k_dh`: exchange-derivative difference upon the SCF density (`factor_k = +//! dfa_hybrid_pos - dfa_hybrid_scf`, with its aux-basis-response part); `de_xc_n`: XC skeleton +//! difference `xc_n - xc` upon D_scf; `de_resp`: $\langle D_{r,\mathrm{ao}}^\sigma, F_1 +//! [\mathrm{xc}]\rangle$ through the UKS Hessian `vmat_deriv1` kernels (coefficient 1 per spin); +//! `de_ovlp_dh`: occupied-Pulay difference `Tr(S1, W_n^\sigma - dme0^\sigma)` with +//! $W_n^\sigma = C_o^\sigma F_n^\sigma(D_{scf})_{oo} C_o^{\sigma\top}$ (no occupancy factor, +//! against the restricted 2). + +use crate::analdrv::response::rgfock_interface::dh_jk_factors; +use crate::analdrv::response::trait_uresp::URespAPI; +use crate::analdrv::response::uresp_interface::{ + dh_xc_func_list_uks, scf_xc_func_list_uks, uscf_resp_interface, URespSCF, +}; +use crate::dft::num_int::XCData; +use crate::dft::numint_matmul::nimatmul::{regroup_grids_by_atom, NIMatmul}; +use crate::dft::numint_matmul::resp_uks::URespKSNIMatmul; +use crate::dft::xc_deriv::XCType; +use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; +use crate::grad::rhf::{ + generator_deriv_hcore, get_dme0, get_grad_dao_ovlp, pack_triu_tilde, RIHFGradientFlagsBuilder, +}; +use crate::grad::traits::GradAPI; +use crate::grad::uhf::RIUHFGradient; +use crate::mpi_io::MPIOperator; +use crate::ri_jk::resp_u::URespRIJK; +use crate::ri_jk::util; +use crate::ri_jk::util::get_dm0_restricted; +use crate::ri_jk::{get_j2c_decomp, J2CDecompose, J2C_THRESH}; +use crate::ri_pt2::pure_pt2_u_elecderiv::{ + get_rupt2_elec_deriv_incore, UPT2ElecDerivIncoreArg, UPT2ElecDerivIncoreInp, +}; +use crate::ri_pt2::{occ_batch_index, occ_batch_step}; +use crate::scf_io::{SCF, SCFType}; +use crate::utilities::memory_batch::{ + blocksize_partition, calc_batch_size, detect_used_memory_mb, handle_memory_exceed, +}; +use crate::utilities::rstsr_util::{RestTensorToRstsrTsrAPI, RestTensorToRstsrViewAPI, Tsr, TsrView}; +use rest_libcint::prelude::*; +use rstsr::prelude::*; +use std::collections::HashMap; +use std::ops::Range; +use tensors::MatrixFull; + +/// Analytic gradient of the PT2-correlation (RI-MP2) energy, unrestricted case. +pub struct UDHGradient<'a> { + scf_data: &'a SCF, + mpi_operator: &'a Option, + /// Gradient parts, keyed by `de_h`/`de_s1`/`de_jk`/`de_rint`/`de`, each `[3, natm]`. + pub result: HashMap>, + /// PT2 correlation energy (side product of the electronic-derivative evaluation). + pub energy: f64, + /// Total relaxed density (SCF density plus the relaxed correlation density) prepared for the + /// fchk density dump, in fchk section order: `[Total MP2 Density, Spin MP2 Density]` + /// (alpha + beta, alpha - beta). `Some` only when `[ctrl] outputs` requests `fchk`; the same + /// density the multipole task dumps. + pub rdm1_dump: Option>>, +} + +impl<'a> UDHGradient<'a> { + pub fn new(scf_data: &'a SCF, mpi_operator: &'a Option) -> Self { + Self { scf_data, mpi_operator, result: HashMap::new(), energy: 0.0, rdm1_dump: None } + } + + pub fn calc(&mut self) -> &mut Self { + let scf_data = self.scf_data; + let mol_obj = &scf_data.mol; + let device = DeviceBLAS::default(); + assert!(mol_obj.ctrl.spin_polarization, "UDHGradient supports only unrestricted (spin-polarization = true) references."); + assert!(matches!(scf_data.scftype, SCFType::UHF), "UDHGradient supports only UHF references."); + assert!(mol_obj.start_mo == 0, "UDHGradient does not support frozen core (start_mo)."); + assert!(self.mpi_operator.is_none(), "UDHGradient is not MPI-parallelized yet."); + assert!( + scf_data.mol.xc_data.omega().is_none(), + "RSH doubly-hybrid gradients are not supported." + ); + assert!( + scf_data.mol.geom.ghost_pc_chrg.is_empty() && !scf_data.mol.ctrl.solvent_enabled, + "The DH analytic gradient does not support QMMM point charges or solvent models yet." + ); + assert!( + !matches!(scf_data.mol.ctrl.j2c_decomp.uplo, Lower), + "UDHGradient does not support lower-triangular Cholesky factors (developer option)." + ); + // this driver reads `rimatr` through the full-space pair tables; a storage-level pruned + // tensor stores its retained rows only, so refuse it rather than contract wrong rows + crate::ri_jk::require_unpruned_rimatr(&scf_data.rimatr_pair_map, "the DH analytic gradient (UDHGradient)"); + let rimatr = scf_data.rimatr.as_ref().expect( + "Decomposed ERI (rimatr) not found; the DH analytic gradient requires the streaming/new-driver RI-PT2 engine which keeps it alive.", + ); + let ederi_utp = rt::asarray((&rimatr.0.data, rimatr.0.size, &device)); + let mol = util::get_cint_mol(mol_obj); + let aux = util::get_cint_aux(mol_obj); + let natm = mol_obj.geom.elem.len(); + let naux = aux.nao(); + + let [c_os, c_ss]: [f64; 2] = mol_obj + .xc_data + .dfa_paramr_adv + .clone() + .expect("Doubly-hybrid gradient requires the PT2 spin factors (dfa_paramr_adv).") + .try_into() + .expect("dfa_paramr_adv must have exactly two entries."); + + // gates of the final-functional (doubly-hybrid) terms; the pure-MP2 family has none, and + // a bDH whose final functional equals the SCF one is covered by the "SCF" entry entirely + let xc_data = &mol_obj.xc_data; + let delta_hyb = xc_data.dfa_hybrid_pos.unwrap_or(xc_data.dfa_hybrid_scf) - xc_data.dfa_hybrid_scf; + let xc_func_pos = dh_xc_func_list_uks(scf_data); + let xc_n_is_scf = match (&xc_data.dfa_compnt_pos, &xc_data.dfa_paramr_pos) { + (Some(code), Some(param)) => { + code.iter().eq(xc_data.dfa_compnt_scf.iter()) + && param.len() == xc_data.dfa_paramr_scf.len() + && param.iter().zip(xc_data.dfa_paramr_scf.iter()).all(|(a, b)| (a - b).abs() < 1e-12) + }, + _ => true, + }; + let add_delta_k = delta_hyb.abs() > 1.0e-10; + let add_xc_n = xc_func_pos.is_some() && !xc_n_is_scf; + let has_final_functional = add_delta_k || add_xc_n; + let add_resp = has_scf_xc(scf_data); + + // atom-grouped copy of the SCF grid, shared by the final-Fock and response terms + let grids_regrouped = if has_final_functional || add_resp { + let grids = scf_data.grids.as_ref().expect( + "Doubly-hybrid gradient requires the DFT grids (`eval_force` regenerates them; `xdh_calculations` frees them).", + ); + Some(regroup_grids_by_atom( + grids.coordinates.clone(), + grids.weights.clone(), + grids.atm_idx.clone(), + grids.quadrature_weights.clone(), + natm, + )) + } else { + None + }; + + // ── per-spin orbital data ── // + let mo_coeff: [Tsr; 2] = + [(&scf_data.eigenvectors[0]).to_rstsr(&device), (&scf_data.eigenvectors[1]).to_rstsr(&device)]; + let mo_occ: [Tsr; 2] = + [(&scf_data.occupation[0]).to_rstsr(&device), (&scf_data.occupation[1]).to_rstsr(&device)]; + let mo_energy: [Tsr; 2] = + [(&scf_data.eigenvalues[0]).to_rstsr(&device), (&scf_data.eigenvalues[1]).to_rstsr(&device)]; + let nao = mo_coeff[0].shape()[0]; + let nmo = mo_coeff[0].shape()[1]; + assert_eq!(mo_coeff[1].shape()[1], nmo, "UDHGradient expects equal MO counts per spin."); + let nocc = [ + mo_occ[0].view().greater(0).sum(), + mo_occ[1].view().greater(0).sum(), + ]; + let nvir = [nmo - nocc[0], nmo - nocc[1]]; + assert!(nocc[0] > 0 && nocc[1] > 0, "UDHGradient expects both spin channels occupied."); + let nao_tp = nao * (nao + 1) / 2; + let so: [Range; 2] = [0..nocc[0], 0..nocc[1]]; + let sv: [Range; 2] = [nocc[0]..nmo, nocc[1]..nmo]; + let c0: [Tsr; 2] = [ + mo_coeff[0].i((.., so[0].clone())).into_contig(ColMajor), + mo_coeff[1].i((.., so[1].clone())).into_contig(ColMajor), + ]; + let cv: [Tsr; 2] = [ + mo_coeff[0].i((.., sv[0].clone())).into_contig(ColMajor), + mo_coeff[1].i((.., sv[1].clone())).into_contig(ColMajor), + ]; + let eo: [Tsr; 2] = [ + mo_energy[0].i(so[0].clone()).into_contig(ColMajor), + mo_energy[1].i(so[1].clone()).into_contig(ColMajor), + ]; + let ev: [Tsr; 2] = [ + mo_energy[0].i(sv[0].clone()).into_contig(ColMajor), + mo_energy[1].i(sv[1].clone()).into_contig(ColMajor), + ]; + let d_hf: [Tsr; 2] = [ + (&scf_data.density_matrix[0]).to_rstsr(&device), + (&scf_data.density_matrix[1]).to_rstsr(&device), + ]; + + // F_n^s(D_scf) = h + J[D_a + D_b] - hyb_pos K^s[D^s] + v_xc[xc_n]^s; it enters the + // DH-combined Lagrangian (the formal-SCF/Brillouin term) and the occupied-Pulay difference + let fock_n: Option<[Tsr; 2]> = if has_final_functional { + let (coords, weights, atm_idx_grids, quadrature_weights) = grids_regrouped.as_ref().unwrap(); + let (factor_j_dh, factor_k_dh) = dh_jk_factors(scf_data); + let cderi_dh = rimatr.0.to_rstsr_view(&device).into_cow(); + let h_core = scf_data.h_core.to_matrixfull().unwrap().to_rstsr(&device); + let mut resp_rijk = URespRIJK::new_with_cderi(factor_j_dh, factor_k_dh, cderi_dh); + let f_jk = resp_rijk.get_fock_rdm(&[d_hf[0].view(), d_hf[1].view()], true); + let mut fock_n = [(&h_core + &f_jk[0]).into_contig(ColMajor), (&h_core + &f_jk[1]).into_contig(ColMajor)]; + if let Some(xc_func_list) = dh_xc_func_list_uks(scf_data) { + let ni_dh = NIMatmul::new(&mol, coords, weights, atm_idx_grids, quadrature_weights); + let mut resp_ks = URespKSNIMatmul::new(xc_func_list, ni_dh, mol_obj.ctrl.print_level > 2); + let f_ks = resp_ks.get_fock_rdm(&[d_hf[0].view(), d_hf[1].view()], true); + *&mut fock_n[0] += &f_ks[0]; + *&mut fock_n[1] += &f_ks[1]; + } + Some(fock_n) + } else { + None + }; + + // ── response objects, PT2 electronic derivative, relaxed correlation density ── // + let config = mol_obj.ctrl.analdrv.clone().unwrap_or_default(); + let mut resp_objs: URespSCF = uscf_resp_interface(scf_data, &config); + // the in-core electronic-derivative evaluation streams its amplitudes (storage bounded + // by the `nvir * nocc * naux` class): the same-spin blocks window their outer + // occupied index, the αβ block its α virtual index (full occupied range per window, cf. + // the kernel docs). Size the windows so the per-window transients stay within a few + // times the `nvir * nocc * naux` class and inside `max_memory`. + let mem_avail = mol_obj.ctrl.max_memory.map(|m| m - detect_used_memory_mb("proc")); + let abort = mol_obj.ctrl.abort_on_mem_exceed; + let nocc_ = [c0[0].shape()[1], c0[1].shape()[1]]; + let nvir_ = [cv[0].shape()[1], cv[1].shape()[1]]; + let headroom = 4.0; + let mut step_ss = [1usize; 2]; + for s in 0..2 { + if nocc_[s] > 0 && nvir_[s] > 0 { + let elems_per_step = 2 * nvir_[s] * nvir_[s] * nocc_[s] + nvir_[s] * naux; + step_ss[s] = occ_batch_step(nocc_[s], elems_per_step, nvir_[s] * nocc_[s] * naux, mem_avail, headroom); + } + } + // the αβ block holds its amplitude window plus ~3 layout copies of it, in each of its + // two passes (α-side respectively β-side rdm1 Grams) + let mut step_ab = [nvir_[0].max(1), nvir_[1].max(1)]; + let mut ab_elems_per_step = [0usize; 2]; + if nocc_[0] > 0 && nocc_[1] > 0 && nvir_[0] > 0 && nvir_[1] > 0 { + ab_elems_per_step[0] = 4 * nocc_[0] * nvir_[1] * nocc_[1]; + step_ab[0] = occ_batch_step(nvir_[0], ab_elems_per_step[0], nvir_[0] * nocc_[0] * naux, mem_avail, headroom); + ab_elems_per_step[1] = 4 * nvir_[0] * nocc_[0] * nocc_[1]; + step_ab[1] = occ_batch_step(nvir_[1], ab_elems_per_step[1], nvir_[1] * nocc_[1] * naux, mem_avail, headroom); + } + let index_occ_outer = [occ_batch_index(nocc_[0], step_ss[0]), occ_batch_index(nocc_[1], step_ss[1])]; + let index_vir_outer = [occ_batch_index(nvir_[0], step_ab[0]), occ_batch_index(nvir_[1], step_ab[1])]; + if mol_obj.ctrl.print_level > 1 { + println!( + "DH gradient amplitude windows: σσ occ steps {:?}, αβ virtual steps {:?} (nocc {nocc_:?}, nvir {nvir_:?}, naux {naux}).", + step_ss, step_ab + ); + } + // guard: even the smallest possible window must fit when batching is floored at one + let peak_elems = [ + if nocc_[0] > 0 && nvir_[0] > 0 { step_ss[0] * (2 * nvir_[0] * nvir_[0] * nocc_[0] + nvir_[0] * naux) } else { 0 }, + if nocc_[1] > 0 && nvir_[1] > 0 { step_ss[1] * (2 * nvir_[1] * nvir_[1] * nocc_[1] + nvir_[1] * naux) } else { 0 }, + step_ab[0] * ab_elems_per_step[0], + step_ab[1] * ab_elems_per_step[1], + ] + .into_iter() + .max() + .unwrap_or(0); + handle_memory_exceed(peak_elems as f64 * 8.0 / 1048576.0, mem_avail, abort); + let j3c = rimatr.0.to_rstsr_view(&device); + let out = get_rupt2_elec_deriv_incore( + &UPT2ElecDerivIncoreInp { + cderi: j3c.view(), + cderi_vox: [None, None], + occ_coeff: [c0[0].view(), c0[1].view()], + vir_coeff: [cv[0].view(), cv[1].view()], + occ_energy: [eo[0].view(), eo[1].view()], + vir_energy: [ev[0].view(), ev[1].view()], + index_occ_outer_vec: index_occ_outer, + index_vir_outer_vec: index_vir_outer, + }, + &UPT2ElecDerivIncoreArg { c_os, c_ss }, + ); + let rdm1 = out.rdm1_corr; + let gfock_part = out.gfock_part; + let g_vix = out.g_vix; + self.energy = out.e_corr; + + resp_objs.make_cpscf_preparation( + &[mo_coeff[0].view(), mo_coeff[1].view()], + &[mo_occ[0].view(), mo_occ[1].view()], + &[mo_energy[0].view(), mo_energy[1].view()], + ); + // SCF response upon the unrelaxed correlation density (the Lagrangian's `Ax0` term), and + // upon the *relaxed* one (the W_III term) below; both are rdm-form A-contractions + let d_corr_ao: [Tsr; 2] = [ + (mo_coeff[0].view() % rdm1[0].view() % mo_coeff[0].view().t()).into_contig(ColMajor), + (mo_coeff[1].view() % rdm1[1].view() % mo_coeff[1].view().t()).into_contig(ColMajor), + ]; + let axd = { + let resp_ao = resp_objs.get_response_rdm(&[d_corr_ao[0].view(), d_corr_ao[1].view()], true); + [ + (mo_coeff[0].view().t() % resp_ao[0].view() % mo_coeff[0].view()).into_contig(ColMajor), + (mo_coeff[1].view().t() % resp_ao[1].view() % mo_coeff[1].view()).into_contig(ColMajor), + ] + }; + + // L^s = gfock[vo] - gfock[ov]^T + axd[vo]; Z-vector through the coupled CP-SCF + let rdm1_resp: [Tsr; 2] = { + let e_ai = [ + resp_objs.cpscf_state().e_ai_shift[0].to_owned(), + resp_objs.cpscf_state().e_ai_shift[1].to_owned(), + ]; + let mut rhs = [ + rt::zeros(([nmo, nocc[0]].f(), &device)), + rt::zeros(([nmo, nocc[1]].f(), &device)), + ]; + for s in 0..2 { + let mut lag_vo = (&gfock_part[s].i((sv[s].clone(), so[s].clone())) + - &gfock_part[s].i((so[s].clone(), sv[s].clone())).t() + + &axd[s].i((sv[s].clone(), so[s].clone()))) + .into_contig(ColMajor); + // the DH-combined Lagrangian carries the formal-SCF (Brillouin-violation) term + // of the final functional: `2 Cv^T F_n^s(D^s) Co` per spin + if let Some(fock_n) = &fock_n { + *&mut lag_vo += &(2.0 * (cv[s].view().t() % fock_n[s].view() % c0[s].view())); + } + rhs[s].i_mut(sv[s].clone()).assign(&(-lag_vo / &e_ai[s])); + } + let z = resp_objs.solve_dimless_cpscf(&[rhs[0].view(), rhs[1].view()]); + [ + { + let mut r = rdm1[0].to_owned(); + r.i_mut((sv[0].clone(), so[0].clone())).assign(&z[0].i((sv[0].clone(), so[0].clone()))); + r + }, + { + let mut r = rdm1[1].to_owned(); + r.i_mut((sv[1].clone(), so[1].clone())).assign(&z[1].i((sv[1].clone(), so[1].clone()))); + r + }, + ] + }; + + // symmetrized half of the relaxed density, and the D_r-carrying AO/MO objects + let d_r_s: [Tsr; 2] = [ + ((&rdm1_resp[0] + &rdm1_resp[0].t()).mapv(|x| 0.5 * x)).into_contig(ColMajor), + ((&rdm1_resp[1] + &rdm1_resp[1].t()).mapv(|x| 0.5 * x)).into_contig(ColMajor), + ]; + // the exchange (K) part of the SCF functional linearized upon the relaxed density carries + // the SCF hybrid coefficient (forge `C1^s = cx C D_r^{S,s}`; unity for the MP2/HF reference) + let cx_scf = mol_obj.xc_data.dfa_hybrid_scf; + let c1_dr: [Tsr; 2] = [ + (mo_coeff[0].view() % d_r_s[0].view()).mapv(|x| cx_scf * x).into_contig(ColMajor), + (mo_coeff[1].view() % d_r_s[1].view()).mapv(|x| cx_scf * x).into_contig(ColMajor), + ]; + let d_r_ao: [Tsr; 2] = [ + (mo_coeff[0].view() % d_r_s[0].view() % mo_coeff[0].view().t()).into_contig(ColMajor), + (mo_coeff[1].view() % d_r_s[1].view() % mo_coeff[1].view().t()).into_contig(ColMajor), + ]; + // per-spin total relaxed densities (SCF + relaxed correlation) and their fchk sections + // (total and spin) for the density dump of force jobs; only materialized when the fchk + // output is requested (`SCF::dh_rdm1_resp`) + if mol_obj.ctrl.outputs.iter().any(|output| output.eq("fchk")) { + let dm_alpha = &d_hf[0] + &d_r_ao[0]; + let dm_beta = &d_hf[1] + &d_r_ao[1]; + self.rdm1_dump = Some(vec![ + to_matrix_full(&dm_alpha + &dm_beta), + to_matrix_full(&dm_alpha - &dm_beta), + ]); + } + let axd_dr = { + let resp_ao = resp_objs.get_response_rdm(&[d_r_ao[0].view(), d_r_ao[1].view()], true); + [ + (mo_coeff[0].view().t() % resp_ao[0].view() % mo_coeff[0].view()).into_contig(ColMajor), + (mo_coeff[1].view().t() % resp_ao[1].view() % mo_coeff[1].view()).into_contig(ColMajor), + ] + }; + + // ── W = W_I + W_II + W_III (MO), per spin, from the `gfock_part` blocks ── // + let mut w_ao: [Tsr; 2] = [rt::zeros_like(&d_hf[0]), rt::zeros_like(&d_hf[1])]; + for s in 0..2 { + let eps2: Tsr = 2.0 * mo_energy[s].view(); + let mut w_mo: Tsr = rt::zeros(([nmo, nmo], &device)); + { + let oo = (&gfock_part[s].i((so[s].clone(), so[s].clone())) + - &eps2.i((None, so[s].clone())) * rdm1[s].i((so[s].clone(), so[s].clone()))) + .mapv(|x: f64| -0.5 * x); + let vv = (&gfock_part[s].i((sv[s].clone(), sv[s].clone())) + - &eps2.i((None, sv[s].clone())) * rdm1[s].i((sv[s].clone(), sv[s].clone()))) + .mapv(|x: f64| -0.5 * x); + w_mo.i_mut((so[s].clone(), so[s].clone())).assign(&oo); + w_mo.i_mut((sv[s].clone(), sv[s].clone())).assign(&vv); + w_mo.i_mut((sv[s].clone(), so[s].clone())).assign( + &gfock_part[s].i((so[s].clone(), sv[s].clone())).t().mapv(|x: f64| -x), + ); + } + // forge W_II = -D_r (.) eps with the UNSYMMETRISED relaxed density and the raw energies + *&mut w_mo -= &rdm1_resp[s] * &mo_energy[s].view().i((None, ..)); + *&mut w_mo.i_mut((so[s].clone(), so[s].clone())) -= 0.5 * axd_dr[s].i((so[s].clone(), so[s].clone())); + w_ao[s] = (mo_coeff[s].view() % w_mo.view() % mo_coeff[s].view().t()).into_contig(ColMajor); + } + + // ── de_h and de_s1 (spin-summed) ── // + let mut de_h: Tsr = rt::zeros(([3, natm], &device)); + let mut de_s1: Tsr = rt::zeros(([3, natm], &device)); + let tsr_ipovlp = { + let (out, shape) = mol.integrate("int1e_ipovlp", "s1", None).into(); + rt::asarray((out, shape, &device)) + }; + { + let mut hcore_gen = generator_deriv_hcore(scf_data); + let w_ao_tot = &w_ao[0] + &w_ao[1]; + let d_r_ao_tot = &d_r_ao[0] + &d_r_ao[1]; + let s1_mu = (&tsr_ipovlp * &w_ao_tot.i((.., .., None))).sum_axes(1); + let s1_nu = (&tsr_ipovlp * &w_ao_tot.view().t().i((.., .., None))).sum_axes(1); + let ao_slice = mol_obj.aoslice_by_atom(); + for atm in 0..natm { + let h1 = hcore_gen(atm); + de_h.i_mut((.., atm)).assign(&(&h1 * &d_r_ao_tot.i((.., .., None))).sum_axes([0, 1])); + let [_, _, p0, p1] = ao_slice[atm]; + let part = (s1_mu.i(p0..p1).sum_axes(0) + s1_nu.i(p0..p1).sum_axes(0)).mapv(|x: f64| -x); + de_s1.i_mut((.., atm)).assign(&part); + } + } + + // ── metric objects, following the rimatr convention (ctrl) ── // + let j2c_decomp = get_j2c_decomp(&aux, &device, mol_obj.ctrl.j2c_decomp); + let s_op: Tsr = match &j2c_decomp { + J2CDecompose::Cd { j2c_l, uplo: Upper, .. } => { + let eye = rt::eye((naux, &device)); + let u_inv = rt::linalg::solve_triangular((j2c_l.view(), eye.view(), Upper)); + u_inv.t().into_contig(ColMajor) + }, + J2CDecompose::Eig { j2c_l_inv, .. } => j2c_l_inv.clone(), + _ => panic!("UDHGradient does not support lower-triangular Cholesky factors (developer option)."), + }; + let tsr_int2c2e_ip1 = { + let (out, shape) = aux.integrate("int2c2e_ip1", "s1", None).into(); + rt::asarray((out, shape, &device)) + }; + // strictly-lower mask with 1/2 diagonal of forge's `generator_L_1`; only the Cholesky + // convention uses it per atom (`Eig` folds the weight matrices once instead, below) + let l1_cd = match &j2c_decomp { + J2CDecompose::Cd { j2c_l, .. } => { + let l_lower = j2c_l.t().into_contig(ColMajor); + let mut l_mask: Tsr = rt::zeros(([naux, naux], &device)); + for i in 0..naux { + for j in 0..i { + l_mask[[i, j]] = 1.0; + } + l_mask[[i, i]] = 0.5; + } + Some((l_lower, l_mask)) + }, + J2CDecompose::Eig { .. } => None, + }; + + // ── persistent MO intermediates; density dots of the rimatr by two GEMVs per spin ── // + // true peak of the partner intermediates in construction order per spin: each source + // coexists with its solve transient, the transient with the permuted `[i, Q, x]` layout + // replacing it, and all of them with the partners already built + let mut resident = 0.0f64; + let mut partner_peak = 0.0f64; + for s in 0..2 { + if nocc[s] == 0 { + continue; + } + let (y_s, g_s) = ((nocc[s] * nmo * naux) as f64, (nvir[s] * nocc[s] * naux) as f64); + partner_peak = partner_peak.max(resident + 2.0 * y_s); + partner_peak = partner_peak.max(resident + y_s + 2.0 * g_s); + resident += y_s + g_s; + } + handle_memory_exceed(partner_peak * 8.0 / 1048576.0, mem_avail, abort); + let d_hf_tp: [Tsr; 2] = + [pack_triu_tilde(d_hf[0].view()), pack_triu_tilde(d_hf[1].view())]; + let d_r_ao_tp: [Tsr; 2] = + [pack_triu_tilde(d_r_ao[0].view()), pack_triu_tilde(d_r_ao[1].view())]; + let y_dot_d: [Tsr; 2] = [ + (ederi_utp.view().t() % d_hf_tp[0].view().reshape([nao_tp, 1])).into_shape([naux]), + (ederi_utp.view().t() % d_hf_tp[1].view().reshape([nao_tp, 1])).into_shape([naux]), + ]; + let y_dot_dr: [Tsr; 2] = [ + (ederi_utp.view().t() % d_r_ao_tp[0].view().reshape([nao_tp, 1])).into_shape([naux]), + (ederi_utp.view().t() % d_r_ao_tp[1].view().reshape([nao_tp, 1])).into_shape([naux]), + ]; + let mut y_ip: [Tsr; 2] = [ + rt::zeros(([nocc[0], nmo, naux], &device)), + rt::zeros(([nocc[1], nmo, naux], &device)), + ]; + let mut r_jk: Tsr = rt::zeros(([naux, naux], &device)); + let mut r_ri: Tsr = rt::zeros(([naux, naux], &device)); + let g2d = [g_vix[0].view().into_shape([nvir[0] * nocc[0], naux]), g_vix[1].view().into_shape([nvir[1] * nocc[1], naux])]; + for p in 0..naux { + let col = ederi_utp.i((.., p)).unpack_tri(Upper, FlagSymm::Sy); + for s in 0..2 { + if nocc[s] == 0 { + continue; + } + y_ip[s].i_mut((.., .., p)).assign(&(c0[s].view().t() % &col % mo_coeff[s].view())); + // forge RI-term metric coupling: R_ri[P, p] = sum_{i a} G[a, i, P] (i a|p), + // spin-summed over the two channels + let m_q = (c0[s].view().t() % &col % cv[s].view()).into_contig(ColMajor); + let mt = m_q.view().t().into_contig(ColMajor); + let part = g2d[s].t() % mt.view().reshape(nvir[s] * nocc[s]); + if s == 0 { + r_ri.i_mut((.., p)).assign(&part); + } else { + *&mut r_ri.i_mut((.., p)) += ∂ + } + } + } + for s in 0..2 { + for i in 0..nocc[s] { + let y2t = y_ip[s].i((i, .., ..)).into_contig(ColMajor); + let a_i = y2t.view().t() % d_r_s[s].view(); + *&mut r_jk += &a_i.view() % y2t.view(); + } + } + *&mut r_jk *= cx_scf; + // `Eig` metric-derivative weights: the y-dot cross terms and the exchange weight form + // `M_jk`, the RI weight `M_ri`; each is folded once by the Fréchet derivative of + // `J^-1/2`, so the atom loop contracts the raw metric derivatives directly + let l1_eig = match &j2c_decomp { + J2CDecompose::Eig { j2c_e: Some(j2c_e), j2c_v: Some(j2c_v), .. } => { + let thresh = mol_obj.ctrl.j2c_decomp.threshold.unwrap_or(J2C_THRESH); + let y_dot_d_tot: Tsr = (&y_dot_d[0] + &y_dot_d[1]).into_owned(); + let y_dot_dr_tot: Tsr = (&y_dot_dr[0] + &y_dot_dr[1]).into_owned(); + let outer: Tsr = y_dot_dr_tot.view().i((.., None)) % y_dot_d_tot.view().i((None, ..)); + let m_jk: Tsr = (&outer + &outer.t()).mapv(|x: f64| -x) + 2.0 * r_jk.view(); + let m_ri: Tsr = -r_ri.view(); + Some(( + l1_fold_eig(m_jk.view(), j2c_e.view(), j2c_v.view(), thresh), + l1_fold_eig(m_ri.view(), j2c_e.view(), j2c_v.view(), thresh), + )) + }, + J2CDecompose::Eig { .. } => panic!( + "The eigen-decomposed 2c-2e ERI carries no stored eigenpairs; cannot build the DH metric-derivative kernels." + ), + J2CDecompose::Cd { .. } => None, + }; + + // ── pre-contractions `Ȳ = partner % S` in the single kept layout `[i, Q, x]`, per spin ── + // one stored layout per partner: the solve result `[(i m), Q]` / `[(a i), Q]` is permuted + // once into `[i, Q, m]` / `[i, Q, a]`, and each spin source is dropped before its permute + let mut ybar_jk: [Option; 2] = [None, None]; + let mut ybar_ri: [Option; 2] = [None, None]; + let mut y_src = y_ip.map(Some); + let mut g_src = g_vix.map(Some); + for s in 0..2 { + if nocc[s] == 0 { + continue; + } + let src = y_src[s].take().unwrap(); + let solved: Tsr = src.view().into_shape([nocc[s] * nmo, naux]) % s_op.view(); // [(i m), Q] + drop(src); + ybar_jk[s] = Some(solved.view().into_shape([nocc[s], nmo, naux]).transpose([0, 2, 1]).into_contig(ColMajor)); // [i, Q, m] + let src = g_src[s].take().unwrap(); + let solved: Tsr = src.view().into_shape([nvir[s] * nocc[s], naux]) % s_op.view(); // [(a i), Q] + drop(src); + ybar_ri[s] = Some(solved.view().into_shape([nvir[s], nocc[s], naux]).transpose([1, 2, 0]).into_contig(ColMajor)); // [i, Q, a] + } + + // ── derivative-integral passes: shell batches over all atoms, atom scatter in-batch ── // + let mol_loc = mol.ao_loc(); + let aux_loc = aux.ao_loc(); + let mol_ao_slice = mol_obj.aoslice_by_atom(); + let aux_slice = mol_obj.make_auxmol_fake().aoslice_by_atom(); + let nbas = mol.nbas(); + let nauxbas = aux.nbas(); + let ncomp = 3usize; + let nvir_max = nvir[0].max(nvir[1]); + let nocc_max = nocc[0].max(nocc[1]); + let ip1_unit = 6 * nao * naux * ncomp + 2 * naux * nmo; + let shell_batch = calc_batch_size::(ip1_unit, mem_avail, None, None).max(1); + let ip2_unit = 3 * nao_tp + 3 * nao * nao + 3 * nocc_max * (nao + nmo + nvir_max); + let aux_batch = calc_batch_size::(ip2_unit, mem_avail, None, None).max(1); + handle_memory_exceed(ip1_unit as f64 * 8.0 / 1048576.0, mem_avail, abort); + handle_memory_exceed(ip2_unit as f64 * 8.0 / 1048576.0, mem_avail, abort); + + let mut de_jk: Tsr = rt::zeros(([3, natm], &device)); + let mut de_rint: Tsr = rt::zeros(([3, natm], &device)); + let mut de_jk_acc = vec![0.0f64; 3 * natm]; + let mut de_rint_acc = vec![0.0f64; 3 * natm]; + let mut ydd_sol_atoms: [Vec; 2] = [Vec::with_capacity(natm), Vec::with_capacity(natm)]; + let mut ydr_sol_atoms: [Vec; 2] = [Vec::with_capacity(natm), Vec::with_capacity(natm)]; + + // ip1 side: u-slot / v-slot of the derivative integrals, per atom and per spin + for atm in 0..natm { + let [shl_a0, shl_a1, _, _] = mol_ao_slice[atm]; + let mut ydd_raw: [Tsr; 2] = + [rt::zeros(([naux, ncomp], &device)), rt::zeros(([naux, ncomp], &device))]; + let mut ydr_raw: [Tsr; 2] = + [rt::zeros(([naux, ncomp], &device)), rt::zeros(([naux, ncomp], &device))]; + let mut t4_3c = [0.0f64; 3]; + let mut ri_3c = [0.0f64; 3]; + let shl_batches = blocksize_partition(&mol_loc[shl_a0..=shl_a1], shell_batch) + .into_iter() + .map(|[i0, i1]| [shl_a0 + i0, shl_a0 + i1]); + for [s0, s1] in shl_batches { + let ip1 = { + let shl_slices = [[s0, s1], [0, nbas], [0, nauxbas]]; + let (out, shape) = + CInt::integrate_cross("int3c2e_ip1", [&mol, &mol, &aux], "s1", shl_slices).into(); + rt::asarray((out, shape, &device)) + }; // [nub, naov, naux, ncomp] + let nub = ip1.shape()[0]; + let naov = ip1.shape()[1]; + let ip1_mat = ip1.view().into_shape([nub * naov, naux * ncomp]); + let r = mol_loc[s0]..mol_loc[s1]; + for s in 0..2 { + if nocc[s] == 0 { + continue; + } + let d_rows = d_hf[s].i((r.clone(), ..)).into_contig(ColMajor); + let ydd_seg = ip1_mat.view().t() % d_rows.view().reshape([nub * naov, 1]); + *&mut ydd_raw[s] += &ydd_seg.into_shape([naux, ncomp]); + let d_rows = d_r_ao[s].i((r.clone(), ..)).into_contig(ColMajor); + let ydr_seg = ip1_mat.view().t() % d_rows.view().reshape([nub * naov, 1]); + *&mut ydr_raw[s] += &ydr_seg.into_shape([naux, ncomp]); + } + // transposed copy `[naov, (nub naux ncomp)]`: the full-AO axis becomes a GEMM index + let ip1_t = ip1.view().transpose([1, 0, 2, 3]).into_contig(ColMajor); + let ip1_t2 = ip1_t.view().into_shape([naov, nub * naux * ncomp]); + for s in 0..2 { + if nocc[s] == 0 { + continue; + } + // u-slot / v-slot of t4 (exchange upon D_r) and of the RI term, per spin + t4_3c = add_fold_ip1_batched( + ip1_t2.view(), + nub, + c0[s].i((r.clone(), ..)), + c1_dr[s].view(), + ybar_jk[s].as_ref().unwrap().view(), + 2.0, + t4_3c, + ); + t4_3c = add_fold_ip1_batched_v( + ip1_t2.view(), + nub, + c1_dr[s].i((r.clone(), ..)), + c0[s].view(), + ybar_jk[s].as_ref().unwrap().view(), + 2.0, + t4_3c, + ); + ri_3c = add_fold_ip1_batched( + ip1_t2.view(), + nub, + c0[s].i((r.clone(), ..)), + cv[s].view(), + ybar_ri[s].as_ref().unwrap().view(), + -1.0, + ri_3c, + ); + ri_3c = add_fold_ip1_batched_v( + ip1_t2.view(), + nub, + cv[s].i((r.clone(), ..)), + c0[s].view(), + ybar_ri[s].as_ref().unwrap().view(), + -1.0, + ri_3c, + ); + } + } + for s in 0..2 { + ydd_sol_atoms[s].push(s_op.view() % ydd_raw[s].view()); + ydr_sol_atoms[s].push(s_op.view() % ydr_raw[s].view()); + } + for t in 0..ncomp { + de_jk_acc[t * natm + atm] += t4_3c[t]; + de_rint_acc[t * natm + atm] += ri_3c[t]; + } + } + + // ip2 side: Q-slot batches spanning all aux shells, scattered by atom afterwards + let mut wdd_all: [Tsr; 2] = + [rt::zeros(([naux, ncomp], &device)), rt::zeros(([naux, ncomp], &device))]; + let mut wdr_all: [Tsr; 2] = + [rt::zeros(([naux, ncomp], &device)), rt::zeros(([naux, ncomp], &device))]; + let q_batches = blocksize_partition(&aux_loc[0..=nauxbas], aux_batch); + for [q0, q1] in q_batches { + let ip2 = { + let shl_slices = [[0, nbas], [0, nbas], [q0, q1]]; + let (out, shape) = + CInt::integrate_cross("int3c2e_ip2", [&mol, &mol, &aux], "s2ij", shl_slices).into(); + rt::asarray((out, shape, &device)) + }; // [nao_tp, nb_q, ncomp] + let nb_q = ip2.shape()[1]; + let ip2_2d = ip2.view().into_shape([nao_tp, nb_q * ncomp]); + let wdd_batch: Vec = (0..2) + .map(|s| { + (ip2_2d.view().t() % d_hf_tp[s].view().reshape([nao_tp, 1])).into_shape([nb_q, ncomp]) + }) + .collect(); + let wdr_batch: Vec = (0..2) + .map(|s| { + (ip2_2d.view().t() % d_r_ao_tp[s].view().reshape([nao_tp, 1])).into_shape([nb_q, ncomp]) + }) + .collect(); + let (qa, qb) = (aux_loc[q0], aux_loc[q1]); + for s in 0..2 { + wdd_all[s].i_mut((qa..qb, ..)).assign(&wdd_batch[s]); + wdr_all[s].i_mut((qa..qb, ..)).assign(&wdr_batch[s]); + } + let ip2_unp = ip2_2d.unpack_tri(Upper, FlagSymm::Sy); // [mu, nu, (q t)] + for s in 0..2 { + if nocc[s] == 0 { + continue; + } + let w1 = c0[s].view().t() % &ip2_unp; // [i, nu, (q t)] + let f_jk = (w1.view() % c1_dr[s].view()).into_contig(ColMajor); + let f_ri = (w1.view() % cv[s].view()).into_contig(ColMajor); + for t in 0..ncomp { + // the partner is stored once, so the two contraction axes are paired + // individually (`(i m)` is not merged) in both halves of the reduction + let r_jk_q = rt::vecdot( + &ybar_jk[s].as_ref().unwrap().view().i((.., qa..qb, ..)), + &f_jk.i((.., .., t * nb_q..(t + 1) * nb_q)).transpose([0, 2, 1]), + ([0, 2], [0, 2]), + ); + let r_ri_q = rt::vecdot( + &ybar_ri[s].as_ref().unwrap().view().i((.., qa..qb, ..)), + &f_ri.i((.., .., t * nb_q..(t + 1) * nb_q)).transpose([0, 2, 1]), + ([0, 2], [0, 2]), + ); + for atm in 0..natm { + let [_, _, auq0, auq1] = aux_slice[atm]; + let (lo, hi) = (qa.max(auq0), qb.min(auq1)); + if lo >= hi { + continue; + } + let f0 = lo - qa; + let f1 = hi - qa; + de_jk_acc[t * natm + atm] += 2.0 * r_jk_q.i(f0..f1).sum_all(); + de_rint_acc[t * natm + atm] -= r_ri_q.i(f0..f1).sum_all(); + } + } + } + } + + // ── per-atom metric couplings, L_1 terms and assembly ── // + for atm in 0..natm { + let [_, _, auq0, auq1] = aux_slice[atm]; + let s_op_cols = s_op.i((.., auq0..auq1)); + let ydd2: [Tsr; 2] = [ + &s_op_cols % wdd_all[0].i((auq0..auq1, ..)), + &s_op_cols % wdd_all[1].i((auq0..auq1, ..)), + ]; + let ydr2: [Tsr; 2] = [ + &s_op_cols % wdr_all[0].i((auq0..auq1, ..)), + &s_op_cols % wdr_all[1].i((auq0..auq1, ..)), + ]; + let mut ydd_l1: [Tsr; 2] = + [rt::zeros(([naux, ncomp], &device)), rt::zeros(([naux, ncomp], &device))]; + let mut ydr_l1: [Tsr; 2] = + [rt::zeros(([naux, ncomp], &device)), rt::zeros(([naux, ncomp], &device))]; + let mut t4_l1 = [0.0f64; 3]; + let mut ri_l1 = [0.0f64; 3]; + // `Cd`: forge's L_1 generator of this atom, contracted with the solve operator; + // `Eig`: the raw metric derivative contracts against the folded weights directly + for t in 0..ncomp { + match &l1_cd { + Some((l_lower, l_mask)) => { + let m0 = &s_op_cols % tsr_int2c2e_ip1.i((auq0..auq1, .., t)); + let mut m: Tsr = &m0 % s_op.view().t(); + let m_t = m.view().t().into_owned(); + *&mut m += &m_t; + let l1_t = (l_lower.view() % &(l_mask * &m)).mapv(|x: f64| -x); + let l1di = s_op.view() % l1_t.view(); + for s in 0..2 { + ydd_l1[s].i_mut((.., t)).assign(&(l1di.view() % y_dot_d[s].view())); + ydr_l1[s].i_mut((.., t)).assign(&(l1di.view() % y_dot_dr[s].view())); + } + t4_l1[t] = 2.0 * rt::vecdot(&l1di, r_jk.view(), -1).sum_all(); + ri_l1[t] = -rt::vecdot(&l1di, r_ri.view(), -1).sum_all(); + }, + None => { + let (k_jk, k_ri) = + l1_eig.as_ref().expect("the metric-derivative kernels must exist for a non-Cholesky policy."); + let dj = tsr_int2c2e_ip1.i((auq0..auq1, .., t)); + t4_l1[t] = rt::vecdot(&dj, &k_jk.i((auq0..auq1, ..)), -1).sum_all(); + ri_l1[t] = rt::vecdot(&dj, &k_ri.i((auq0..auq1, ..)), -1).sum_all(); + }, + } + } + + // forge t1 + t2 (all four spin pairs) + t4 (per spin), and the RI term + let mut grad = [0.0f64; 3]; + let mut grad_ri = [0.0f64; 3]; + for t in 0..ncomp { + let c_ydd = (&ydd2[0].i((.., t)) + &ydd_l1[0].i((.., t)) + + &(&ydd2[1].i((.., t)) + &ydd_l1[1].i((.., t)))) + .into_owned(); + let c_ydr = (&ydr2[0].i((.., t)) + &ydr_l1[0].i((.., t)) + + &(&ydr2[1].i((.., t)) + &ydr_l1[1].i((.., t)))) + .into_owned(); + let v_ydd = &ydd_sol_atoms[0][atm].i((.., t)) + &ydd_sol_atoms[1][atm].i((.., t)); + let v_ydr = &ydr_sol_atoms[0][atm].i((.., t)) + &ydr_sol_atoms[1][atm].i((.., t)); + let y_dot_d_tot = &y_dot_d[0] + &y_dot_d[1]; + let y_dot_dr_tot = &y_dot_dr[0] + &y_dot_dr[1]; + let s1: f64 = rt::vecdot(&y_dot_d_tot, &(&c_ydr + &(2.0 * &v_ydr)).view(), -1).sum_all(); + let s2: f64 = rt::vecdot(&y_dot_dr_tot, &(&c_ydd + &(2.0 * &v_ydd)).view(), -1).sum_all(); + grad[t] = -s1 - s2 + t4_l1[t] + de_jk_acc[t * natm + atm]; + grad_ri[t] = ri_l1[t] + de_rint_acc[t * natm + atm]; + } + de_jk.i_mut((.., atm)).assign(&rt::asarray((grad.as_slice(), [3], &device))); + de_rint.i_mut((.., atm)).assign(&rt::asarray((grad_ri.as_slice(), [3], &device))); + } + + // ── doubly-hybrid (final-functional) terms, all absent for the pure-MP2 family ── // + // `de_k_dh`: exchange-derivative difference upon the SCF density (`factor_k = + // dfa_hybrid_pos - dfa_hybrid_scf`; the J part is unchanged and covered by the "SCF" + // entry), i.e. the D_scf part of forge `get_gradient_jk`'s `0.5 cx_n C D_mo` piece; + // `de_xc_n`: XC skeleton difference `xc_n - xc` upon D_scf, contracted per spin (forge + // `_get_gradient_gga` explicit part; the SCF-functional one is the "SCF" entry's + // `de_xc`); + // `de_resp`: `` per spin, the explicit AO-derivative of the SCF + // functional's XC potential (forge `_get_gradient_gga` response part; its fxc-on-D_r + // half enters the CP-SCF kernel of the Z-vector instead); + // `de_ovlp_dh`: occupied-Pulay difference `Tr(S1, W_n^s - dme0^s)` with + // `W_n^s = C_o^s F_n^s(D_scf)_oo C_o^sT` (no occupancy factor; forge + // `_get_gradient_enfunc`), against the "SCF" entry's `dme0` term. + let mut de_k_dh: Tsr = rt::zeros(([3, natm], &device)); + let mut de_xc_n: Tsr = rt::zeros(([3, natm], &device)); + let mut de_resp: Tsr = rt::zeros(([3, natm], &device)); + let mut de_ovlp_dh: Tsr = rt::zeros(([3, natm], &device)); + if has_final_functional || add_resp { + if add_delta_k { + let mut grad_helper = uhf_grad_helper(scf_data, self.mpi_operator); + grad_helper.flags.factor_k = Some(delta_hyb); + grad_helper.calc_de_jk(); + let de_k = &grad_helper.result["de_k"]; + de_k_dh.assign(&rt::asarray((&de_k.data, de_k.size, &device))); + // the auxiliary-basis-response part of the exchange gradient + if let Some(de_kaux) = grad_helper.result.get("de_kaux") { + *&mut de_k_dh += &rt::asarray((&de_kaux.data, de_kaux.size, &device)); + } + } + if add_xc_n { + // the DH total energy carries the *final* functional's XC skeleton explicitly + // (forge `grad_gga`), while the "SCF" entry contributes the SCF functional's one; + // add their difference, contracted with the per-spin SCF density + let grids = scf_data.grids.as_ref().unwrap(); + let dao_xc_n = xc_skeleton_dao(scf_data, self.mpi_operator, false, grids); + let dao_xc_scf = xc_skeleton_dao(scf_data, self.mpi_operator, true, grids); + let mut dao_xc: Tsr = rt::zeros(([nao, 3], &device)); + for s in 0..2 { + *&mut dao_xc += + &get_grad_dao_ovlp((&dao_xc_n[s] - &dao_xc_scf[s]).view(), d_hf[s].view()); + } + for atm in 0..natm { + let [_, _, p0, p1] = mol_ao_slice[atm]; + *&mut de_xc_n.i_mut((.., atm)) += &dao_xc.i(p0..p1).sum_axes(0); + } + } + if add_resp { + let (coords, weights, atm_idx_grids, quadrature_weights) = grids_regrouped.as_ref().unwrap(); + let dm0 = [ + get_dm0_restricted(mo_coeff[0].view(), mo_occ[0].view()), + get_dm0_restricted(mo_coeff[1].view(), mo_occ[1].view()), + ]; + let mut ni = NIMatmul::new(&mol, coords, weights, atm_idx_grids, quadrature_weights); + // the explicit AO-derivative of the SCF functional's XC potential (forge + // `hessian.uks._get_vxc_deriv1`), contracted with the relaxed density on the fly; + // the grid-shift terms are absent, following the forge reference + de_resp = de_xc_response_term(&mol, &scf_xc_func_list_uks(scf_data), &mut ni, &dm0, &d_r_ao); + } + if has_final_functional { + let fock_n = fock_n.as_ref().unwrap(); + let mut w_n: Tsr = rt::zeros_like(&d_hf[0]); + for s in 0..2 { + let f_n_oo = (c0[s].view().t() % fock_n[s].view() % c0[s].view()).into_contig(ColMajor); + *&mut w_n += &(c0[s].view() % &f_n_oo % c0[s].view().t()); + } + // the "SCF" entry's `de_ovlp` carries the per-spin `dme0 = C (n eps) C^T` term, + // while the final-functional Pulay reads `-Tr(S1_oo^s, F_n^s_oo)` per spin; add + // the difference only (`2` = the two-half contraction pattern) + let dme0 = get_dme0(mo_coeff[0].view(), mo_occ[0].view(), mo_energy[0].view()) + + get_dme0(mo_coeff[1].view(), mo_occ[1].view(), mo_energy[1].view()); + let dao_ovlp: Tsr = get_grad_dao_ovlp(tsr_ipovlp.view(), (&w_n - &dme0).view()); + for atm in 0..natm { + let [_, _, p0, p1] = mol_ao_slice[atm]; + *&mut de_ovlp_dh.i_mut((.., atm)) += &dao_ovlp.i(p0..p1).sum_axes(0); + } + } + } + + // ── assemble the parts ── // + let de = de_h.view() + de_s1.view() + de_jk.view() + de_rint.view() + de_k_dh.view() + + de_xc_n.view() + de_resp.view() + de_ovlp_dh.view(); + for (key, val) in [ + ("de_h", de_h), + ("de_s1", de_s1), + ("de_jk", de_jk), + ("de_rint", de_rint), + ("de_k_dh", de_k_dh), + ("de_xc_n", de_xc_n), + ("de_resp", de_resp), + ("de_ovlp_dh", de_ovlp_dh), + ("de", de.into_owned()), + ] { + self.result.insert(key.to_string(), to_matrix_full(val)); + } + self + } +} + +/// Device tensor (e.g. a `[3, natm]` gradient part or a `[nao, nao]` density) into a +/// `MatrixFull` of the same shape. +fn to_matrix_full(tsr: Tsr) -> MatrixFull { + let shape: [usize; 2] = tsr.shape().to_vec().try_into().unwrap(); + MatrixFull::from_vec(shape, tsr.into_shape(-1).into_raw()).unwrap() +} + +/// Fold a metric-derivative weight matrix for the eigen convention: the `L_1` term +/// `sum_weights . raw d(J^-1/2)` equals `` with `F` the (self-adjoint) Fréchet +/// derivative of `J^-1/2`, i.e. `V (D ∘ (V^T M^T V)) V^T` with +/// `D_ij = -1 / (sqrt(e_i) (sqrt(e_i) + sqrt(e_j)))`; the returned `G + G^T` pairs with one +/// atom block of the raw `int2c2e_ip1` (its transposed half included). Eigenpairs below +/// `threshold` are discarded, exactly as in the decomposition itself. Unrestricted twin of the +/// helper in [`crate::grad::rdh`]. +fn l1_fold_eig(weight: TsrView, j2c_e: TsrView, j2c_v: TsrView, threshold: f64) -> Tsr { + let n = j2c_e.view().less(threshold).sum(); + let sr = j2c_e.i(n..).pow(0.5); + let v = j2c_v.i((.., n..)).into_contig(ColMajor); + let ssum = &sr.i((.., None)) + &sr.i((None, ..)); + let d: Tsr = (&sr.i((.., None)) * &ssum).mapv(|x: f64| -1.0 / x); + let p = v.view().t() % (weight.t() % v.view()); // V^T M^T V + let g: Tsr = v.view() % &(&d * &p) % v.view().t(); + (&g + &g.t()).into_owned() +} + +/// Whether the SCF-iteration functional carries a non-HF XC part (the gate of the +/// response-density XC term; the kernels handle LDA/GGA/MGGA alike, unlike forge's GGA-only +/// gate). Unrestricted twin of the helper in [`crate::grad::rdh`], duplicated to keep the +/// restricted module untouched. +fn has_scf_xc(scf_data: &SCF) -> bool { + !scf_data.mol.xc_data.dfa_compnt_scf.is_empty() +} + +/// A bare unrestricted gradient helper carrying the SCF control flags, for reuse of the UKS +/// gradient kernels (`calc_de_jk` with custom J/K factors, `get_vxc_rayon_new`); the unrestricted +/// twin of the helper in [`crate::grad::rdh`]. +fn uhf_grad_helper<'a>(scf_data: &'a SCF, mpi_operator: &'a Option) -> RIUHFGradient<'a> { + let mol_obj = &scf_data.mol; + let flags = RIHFGradientFlagsBuilder::default() + .print_level(mol_obj.ctrl.print_level) + .max_memory(mol_obj.ctrl.max_memory) + .auxbasis_response(mol_obj.ctrl.auxbasis_response) + .factor_j(None) + .factor_k(None) + .build() + .unwrap(); + RIUHFGradient { scf_data, flags, mpi_operator, result: HashMap::new() } +} + +/// Per-spin grid contribution of the XC skeleton derivative, `vmat_s[nao, nao, 3]` (the +/// row-scatter, contracted with the per-spin SCF density and the two-half factor, is the +/// skeleton-gradient contribution); `scf` selects the SCF-iteration functional, else the final +/// (`pos`) one. Unrestricted twin of the helper in [`crate::grad::rdh`]. +fn xc_skeleton_dao( + scf_data: &SCF, + mpi_operator: &Option, + scf: bool, + grids: &crate::dft::Grids, +) -> Vec { + let mol_obj = &scf_data.mol; + let xc_data = &mol_obj.xc_data; + let (code, param) = if scf { + (&xc_data.dfa_compnt_scf, &xc_data.dfa_paramr_scf) + } else { + ( + xc_data.dfa_compnt_pos.as_ref().expect("the final functional has no XC component list."), + xc_data.dfa_paramr_pos.as_ref().expect("the final functional has no XC parameter list."), + ) + }; + let xc_type = if code.iter().any(|&c| xc_data.init_libxc(&c).needs_tau()) { + XCType::MGGA + } else if code + .iter() + .any(|&c| !matches!(xc_data.init_libxc(&c).family(), libxc::enums::LibXCFamily::LDA | libxc::enums::LibXCFamily::HybLDA)) + { + XCType::GGA + } else { + XCType::LDA + }; + let xc_data_sel = XCData { + xc_type, + xc_code: code, + xc_params: param, + device: DeviceOpenBLAS::default(), + dm: &scf_data.density_matrix, + mo_coeffs: None, + occ: None, + }; + let mut grids_clone = grids.clone(); + crate::grad::uks::get_vxc_rayon_new(&uhf_grad_helper(scf_data, mpi_operator), &xc_data_sel, &mut grids_clone, mol_obj, 16) +} + +/// Response-density XC term `` of the DH gradient (forge +/// `_get_gradient_gga` response part, coefficient 1 per spin), with `F_1[A, t, mu, nu]` the +/// explicit AO-derivative of the SCF functional's XC potential (the gradient-level analogue of +/// forge `hessian.uks._get_vxc_deriv1`, including its fxc-upon-skeleton-density part; the +/// grid-shift terms are absent, following the forge reference). +/// +/// Chunk-parallel over the (atom-grouped) grid, with the relaxed density contracted immediately, +/// so only the Hessian-level kernels of [`crate::dft::numint_matmul::hess_uks`] are reused and no +/// full-grid `[nao, nao, 3, natm]` accumulator is held. Unrestricted twin of the helper in +/// [`crate::grad::rdh`]. +fn de_xc_response_term( + mol: &CInt, + xc_func_list: &[(f64, libxc::functional::LibXCFunctional)], + ni: &mut NIMatmul, + dm0: &[Tsr; 2], + d_r_ao: &[Tsr; 2], +) -> Tsr { + use crate::dft::numint_matmul::hess_rks::{get_drho, get_vmat_ip, get_vmat_vxc}; + use crate::dft::numint_matmul::hess_uks::{get_rho_exc_vxc_fxc_uks, get_vmat_fxc_uks}; + use crate::dft::xceff::prelude::XCDenType::*; + + let natm = mol.natm(); + let nao = mol.nao(); + let ngrids = ni.weights.len(); + let device = dm0[0].device().clone(); + let aoslices = mol.aoslice_by_atom(); + let func_refs: Vec<&libxc::functional::LibXCFunctional> = xc_func_list.iter().map(|(_, f)| f).collect(); + let xc_type = determine_den_type_from_list(&func_refs); + // gradient-level AO derivative order (pyscf `grad.uks` convention; the Hessian needs one more) + let ao_deriv = match xc_type { + RHO => 1, + _ => 2, + }; + let ncomp_ao_dm0 = xc_type.num_ao_comp(); + + let de_resp: Tsr = rt::zeros(([3, natm].f(), &device)); + let nchunk_target = (rayon::current_num_threads() * 8).max(1); + let nsplit = ((ngrids + nchunk_target - 1) / nchunk_target).max(1); + let chunks = (0..ngrids).step_by(nsplit).map(|s| (s, (s + nsplit).min(ngrids))).collect::>(); + let guard = std::sync::Mutex::new(()); + use rayon::prelude::*; + let weights_full = ni.weights.clone(); + let d_r_ao_ref = d_r_ao; + let dm0_owned = [dm0[0].clone(), dm0[1].clone()]; + let ni_ref = &*ni; + let device_ref = &device; + let aoslices_ref = &aoslices; + chunks.into_par_iter().for_each(|(start, end)| { + let mut ni_chunk = ni_ref.split_batch(start, end); + let weights = rt::asarray((&weights_full[start..end], device_ref)); + let ao = ni_chunk.get_cached_ao(ao_deriv); + let ao_dm0_a = ao.i((Ellipsis, ..ncomp_ao_dm0)) % &dm0_owned[0]; + let ao_dm0_b = ao.i((Ellipsis, ..ncomp_ao_dm0)) % &dm0_owned[1]; + let (_rho, _exc, vxc, fxc) = + get_rho_exc_vxc_fxc_uks(xc_func_list, ao.view(), ao_dm0_a.view(), ao_dm0_b.view()); + let wv_a = &weights * &vxc.i((.., .., 0)); + let wv_b = &weights * &vxc.i((.., .., 1)); + let wf = &weights * &fxc; + let drho_a = get_drho(xc_type, ao.view(), ao_dm0_a.view(), aoslices_ref); + let drho_b = get_drho(xc_type, ao.view(), ao_dm0_b.view(), aoslices_ref); + let vmat_ip_a = get_vmat_ip(xc_type, ao.view(), wv_a.view()); + let vmat_ip_b = get_vmat_ip(xc_type, ao.view(), wv_b.view()); + let (vmat_fxc_a, vmat_fxc_b) = + get_vmat_fxc_uks(xc_type, ao.view(), drho_a.view(), drho_b.view(), wf.view()); + let vmat_a = &vmat_fxc_a + &get_vmat_vxc(vmat_ip_a.view(), aoslices_ref); + let vmat_b = &vmat_fxc_b + &get_vmat_vxc(vmat_ip_b.view(), aoslices_ref); + let mut de_chunk: Tsr = rt::zeros(([3, natm].f(), &device)); + for atm in 0..natm { + for t in 0..3 { + let v0: f64 = rt::vecdot( + &d_r_ao_ref[0].view().into_shape([nao * nao]), + &vmat_a.i((.., .., t, atm)).into_shape([nao * nao]), + -1, + ) + .sum_all(); + let v1: f64 = rt::vecdot( + &d_r_ao_ref[1].view().into_shape([nao * nao]), + &vmat_b.i((.., .., t, atm)).into_shape([nao * nao]), + -1, + ) + .sum_all(); + de_chunk[[t, atm]] = v0 + v1; + } + } + let _lock = guard.lock().unwrap(); + unsafe { *&mut de_resp.force_mut() += &de_chunk }; + }); + de_resp +} + +/// One ip1 derivative-integral pass on an AO-shell batch, reduced against the pre-contracted +/// partner `Ȳ` (the metric solve is hoisted out of the atom loop); shared kernel with +/// [`crate::grad::rdh`], duplicated here to keep the restricted module untouched. +fn add_fold_ip1_batched( + ip1_t2: TsrView, // [naov, nub naux ncomp] + nub: usize, + c_row: TsrView, // [nub, x], the atom-side MO weight + c_col: TsrView, // [naov, y], the full-AO-side MO weight + ybar: TsrView, // [x, naux, y], pre-contracted partner + factor: f64, + mut acc: [f64; 3], +) -> [f64; 3] { + let device = ip1_t2.device().clone(); + let naov = ip1_t2.shape()[0]; + let naux = ip1_t2.shape()[1] / (nub * 3); + let x = c_row.shape()[1]; + let y = c_col.shape()[1]; + let mut k_buf: Tsr = rt::zeros(([nub, naux * y], &device)); + k_buf.view_mut().matmul_from(&c_row, &ybar.reshape([x, naux * y]), 1.0, 0.0); + let k_buf_view = k_buf.view(); + let k2 = k_buf_view.reshape([nub * naux, y]); + let mut g_buf: Tsr = rt::zeros(([nub * naux, y], &device)); + for t in 0..3 { + let blk = ip1_t2.i((.., t * nub * naux..(t + 1) * nub * naux)); + g_buf.view_mut().matmul_from(&blk.t(), &c_col, 1.0, 0.0); + let s: f64 = rt::vecdot(&g_buf, &k2, -1).sum_all(); + acc[t] += factor * s; + } + acc +} + +/// [`add_fold_ip1_batched`] counterpart for the `v` slot, where the partner `Ȳ[i, Q, m]` is +/// contracted over its slow `m` axis instead: the same two GEMMs are kept, but the partner +/// enters on the left of a single `[nocc naux, nmo] % [nmo, nub]` fold while the full-AO-side +/// fold is assembled from one `[y, naov] % [naov, naux]` GEMM per batch row `u` (each landing +/// contiguously in the paired buffer). Both stay within the `u` slot's GEMM structure and the +/// reduction remains a `rt::vecdot` over matching layouts, so no second partner layout is +/// stored; unrestricted twin of the kernel in [`crate::grad::rdh`]. +fn add_fold_ip1_batched_v( + ip1_t2: TsrView, // [naov, nub naux ncomp] + nub: usize, + c_row: TsrView, // [nub, x], the atom-side MO weight (contracted partner axis) + c_col: TsrView, // [naov, y], the full-AO-side MO weight (partner's leading axis) + ybar: TsrView, // [y, naux, x], pre-contracted partner in the single kept layout + factor: f64, + mut acc: [f64; 3], +) -> [f64; 3] { + let device = ip1_t2.device().clone(); + let naov = ip1_t2.shape()[0]; + let naux = ip1_t2.shape()[1] / (nub * 3); + let x = c_row.shape()[1]; + let y = c_col.shape()[1]; + let mut k_buf: Tsr = rt::zeros(([y * naux, nub], &device)); + k_buf.view_mut().matmul_from(&ybar.reshape([y * naux, x]), &c_row.t(), 1.0, 0.0); + let mut g_buf: Tsr = rt::zeros(([y, naux, nub], &device)); + let c_col_t = c_col.t(); + let ip1_t3 = ip1_t2.reshape([naov, nub, naux * 3]); // [naov, u, (Q t)] + for t in 0..3 { + for u in 0..nub { + let blk = ip1_t3.i((.., u, t * naux..(t + 1) * naux)); + g_buf.i_mut((.., .., u)).matmul_from(&c_col_t, &blk, 1.0, 0.0); + } + let s: f64 = rt::vecdot(&k_buf, &g_buf.view().reshape([y * naux, nub]), 0).sum_all(); + acc[t] += factor * s; + } + acc +} + +impl GradAPI for UDHGradient<'_> { + fn get_gradient(&self) -> MatrixFull { + self.result["de"].clone() + } + fn get_energy(&self) -> f64 { + self.energy + } +} diff --git a/src/grad/uks.rs b/src/grad/uks.rs index 9eb1cc9fd92a80389207c8a17c8f14371c826a04..d7beb36fcc5429292cd101c79442bc6eced99d95 100644 --- a/src/grad/uks.rs +++ b/src/grad/uks.rs @@ -433,7 +433,14 @@ fn get_vxc_rayon(gradient_method: &RIUHFGradient, xc_data: &XCData, grids: &mut } -fn get_vxc_rayon_new(gradient_method: &RIUHFGradient, xc_data: &XCData, grids: &mut Grids, mol: &Molecule, max_memory: usize) -> Vec> { +/// Grid contribution of the UKS XC skeleton derivative, `vmat_s[nao, nao, 3]` per spin (the +/// row-scatter, times `2`, contracted with the per-spin density, is the XC skeleton gradient), for +/// the XC functional carried by `xc_data`. +/// +/// Public so that post-SCF (doubly-hybrid) gradients can evaluate the skeleton derivative of +/// the *final* functional, whose component list is not the SCF one (`xc_data` is built by the +/// caller, e.g. from `dfa_compnt_pos` / `dfa_paramr_pos`). +pub fn get_vxc_rayon_new(gradient_method: &RIUHFGradient, xc_data: &XCData, grids: &mut Grids, mol: &Molecule, max_memory: usize) -> Vec> { let default_omp_num_threads = omp_get_num_threads_wrapper(); let num_grids = grids.weights.len(); diff --git a/src/main_driver.rs b/src/main_driver.rs index bfd8a1a9258daea4900f13c27646a41b821ac6ae..63dc36e565cf4e2900c98eb585c8801702842646 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -855,11 +855,14 @@ fn eval_force(scf_data: &mut SCF, time_mark: &mut utilities::TimeRecords, mpi_op // current available analytical gradients methods: // 1) analytical RHF, UHF force // 2) dftd force - // - // disallow post-scf calculations for force - if scf_data.mol.xc_data.is_fifth_dfa() { + // 3) doubly-hybrid (PT2-family) correlation gradient, as a separate entry + let is_dh = matches!(scf_data.mol.xc_data.dfa_family_pos, Some(crate::dft::DFAFamily::PT2)); + if scf_data.mol.xc_data.is_fifth_dfa() && !is_dh { panic!("Analytic Gradient calculation is currently not available for post-SCF methods."); } + if is_dh && mpi_operator.is_some() { + panic!("Analytic Gradient calculation of post-SCF methods is not MPI-parallelized yet."); + } if scf_data.mol.ctrl.print_level > 1 { println!("Gradient evaluation using Analytical differentiation"); @@ -899,6 +902,25 @@ fn eval_force(scf_data: &mut SCF, time_mark: &mut utilities::TimeRecords, mpi_op // list of (gradient name, gradient data) let mut grad_data_list: Vec<(String, Box)> = vec![]; + // the DH (PT2-family) entry hands the total relaxed density back for the fchk density + // dump (`SCF::dh_rdm1_resp`, written out by the `outputs` pass) + let mut dh_rdm1_resp: Option>> = None; + + // post-SCF (doubly-hybrid) gradients need the DFT grids again whenever the DH entry runs + // on a grid (final-functional XC part / hybrid difference, or the SCF functional's own XC + // response); `xdh_calculations` frees them, so regenerate (pure MP2 needs none) + let is_dh = matches!(scf_data.mol.xc_data.dfa_family_pos, Some(crate::dft::DFAFamily::PT2)); + let dh_needs_grids = { + let xc_data = &scf_data.mol.xc_data; + let final_xc = xc_data.dfa_compnt_pos.as_ref().map_or(false, |v| !v.is_empty()); + let delta_hyb = + xc_data.dfa_hybrid_pos.unwrap_or(xc_data.dfa_hybrid_scf) - xc_data.dfa_hybrid_scf; + final_xc || delta_hyb.abs() > 1.0e-10 || !xc_data.dfa_compnt_scf.is_empty() + }; + if is_dh && dh_needs_grids && scf_data.grids.is_none() { + scf_data.grids = Some(crate::dft::Grids::build(&mut scf_data.mol)); + } + // 1. self-consistent gradient data let grad_data_scf: Box = { if !scf_data.mol.ctrl.spin_polarization { @@ -925,6 +947,21 @@ fn eval_force(scf_data: &mut SCF, time_mark: &mut utilities::TimeRecords, mpi_op }; grad_data_list.push(("SCF".into(), grad_data_scf)); + // 1.5 doubly-hybrid (RI-PT2) correlation gradient; requires the decomposed rimatr + if is_dh { + if !scf_data.mol.ctrl.spin_polarization { + let mut grad_data_dh = crate::grad::rdh::RDHGradient::new(&scf_data, mpi_operator); + grad_data_dh.calc(); + dh_rdm1_resp = grad_data_dh.rdm1_dump.take(); + grad_data_list.push(("DH".into(), Box::new(grad_data_dh))); + } else { + let mut grad_data_dh = crate::grad::udh::UDHGradient::new(&scf_data, mpi_operator); + grad_data_dh.calc(); + dh_rdm1_resp = grad_data_dh.rdm1_dump.take(); + grad_data_list.push(("DH".into(), Box::new(grad_data_dh))); + } + } + // 2. dftd gradient data // we will force to evaluate dftd gradient, since dftd3 is not bottleneck for small to medium molecules use crate::dftd::grad::DFTDGrad; @@ -947,6 +984,13 @@ fn eval_force(scf_data: &mut SCF, time_mark: &mut utilities::TimeRecords, mpi_op } } + // the gradient objects borrow `scf_data`; release them before handing the DH relaxed + // density to the fchk output pass — force jobs only + drop(grad_data_list); + if is_dh && matches!(scf_data.mol.ctrl.job_type, JobType::Force) { + scf_data.dh_rdm1_resp = dh_rdm1_resp; + } + // TDDFT analytic-gradient response for the requested excited state. if let Some(tddft_ctrl) = scf_data.mol.ctrl.tddft.clone() { if tddft_ctrl.tddft_grad_state > 0 { @@ -995,7 +1039,8 @@ fn eval_force(scf_data: &mut SCF, time_mark: &mut utilities::TimeRecords, mpi_op println!("{}", formated_force(&gradient, &scf_data.mol.geom.elem)); println!("------------------------------------"); - (scf_data.scf_energy, gradient) + // the total energy (fifth-DFA/rpa/ai correction included), consistent with the gradient + (collect_total_energy(scf_data), gradient) }; time_mark.count("force"); diff --git a/src/ri_pt2/mod.rs b/src/ri_pt2/mod.rs index d823fde42b9c877cffd20b04ca51d99b454e15c2..34f1996802a659265530b3f4bce3c9548a7975cd 100644 --- a/src/ri_pt2/mod.rs +++ b/src/ri_pt2/mod.rs @@ -45,8 +45,37 @@ pub mod pt2_pair_eng; pub mod torch_pt2_pair_eng; pub mod pure_pt2_r_elecderiv; +pub mod pure_pt2_u_elecderiv; pub mod rgfock_pt2; +/// Batch size (number of outer occupied indices) of the amplitude-streaming electronic +/// derivatives ([`pure_pt2_r_elecderiv`] / [`pure_pt2_u_elecderiv`]). +/// +/// Each streamed batch costs `elems_per_step` elements per outer occupied index and must fit both +/// the available memory (`mem_avail_mb`, when given) and `headroom` times the working class +/// `class_elems` (the `nvir * nocc * naux` transform/`G` objects); the result is capped at `nocc` +/// and floored at 1. Amplitude storage is therefore bounded by the working class (times +/// `headroom`), with a single window -- the full `nocc^2 nvir^2` tensor -- only when that tensor +/// is itself within the bound (small systems). +pub fn occ_batch_step(nocc: usize, elems_per_step: usize, class_elems: usize, mem_avail_mb: Option, headroom: f64) -> usize { + let elems_per_step = elems_per_step.max(1); + let by_class = (headroom * class_elems as f64 / elems_per_step as f64).floor() as usize; + let by_mem = match mem_avail_mb { + Some(mb) => (mb.max(0.0) * 1048576.0 / (elems_per_step as f64 * 8.0)).floor() as usize, + None => usize::MAX, + }; + by_class.min(by_mem).clamp(1, nocc.max(1)) +} + +/// Window boundaries `[0, step, 2*step, ..., nocc]` of the amplitude-streaming occupied batches +/// (see [`occ_batch_step`]). +pub fn occ_batch_index(nocc: usize, step: usize) -> Vec { + let step = step.max(1); + let mut index: Vec = (0..nocc).step_by(step).collect(); + index.push(nocc); + index +} + #[derive(Clone)] pub struct PT2 { pub pt2_type: usize, diff --git a/src/ri_pt2/pure_pt2_r_elecderiv.rs b/src/ri_pt2/pure_pt2_r_elecderiv.rs index df800a5039efde79d6e1307b100a8776e8ec7391..5510ca6cb03aef854e85110c28dce5d8e1848872 100644 --- a/src/ri_pt2/pure_pt2_r_elecderiv.rs +++ b/src/ri_pt2/pure_pt2_r_elecderiv.rs @@ -28,19 +28,31 @@ where pub struct RPT2ElecDerivIncoreArg { pub c_os: f64, pub c_ss: f64, + /// Evaluate the full four-block generalized-Fock contribution (OO+OV+VO+VV, including the + /// orbital-energy terms) instead of the W3/W4 off-diagonal blocks only. + pub full_gfock: bool, + /// Additionally dump the amplitude intermediate $G_{ai}^\mathtt{P}$ (contracted with the + /// cderi over the virtual axis) as an output; costs an extra `nvir * nocc * naux` tensor. + pub dump_g_vix: bool, } pub struct RPT2ElecDerivIncoreOut { /// MP2 correlation energy. This is side product from property evaluation. pub e_corr: f64, - /// **Partial** generalized Fock matrix contribution. Shape `(nmo, nmo)`. + /// Generalized Fock matrix contribution of the PT2 part, shape `(nmo, nmo)`. /// - /// Note this lacks SCF response upon rdm1_corr contribution (the - /// $A_{ai, pq} D_{pq}^{\mathrm{RDM}}$ term of the Lagrangian, which requires a response - /// object of the underlying SCF). + /// By default only the OV and VO blocks are filled (the W3/W4 terms); with `full_gfock` + /// all four blocks are filled, including the $2 \varepsilon \cdot \mathrm{rdm1}^{\mathrm{corr}}$ + /// orbital-energy terms of the diagonal blocks. In both modes this lacks the SCF response + /// upon rdm1_corr contribution (the $A_{pq, \lambda} D_{\lambda}^{\mathrm{RDM}}$ term, + /// which requires a response object of the underlying SCF). pub gfock_part: Tsr, /// 1-RDM in MO basis (correlation contribution, unrelaxed). Shape `(nmo, nmo)`. pub rdm1_corr: Tsr, + /// Amplitude intermediate $G_{ai}^\mathtt{P} = T_{ij,ab} (i a | \mathtt{P})$ (summed over the + /// occupied index), only when requested by [`RPT2ElecDerivIncoreArg::dump_g_vix`]. + /// Shape `(nvir, nocc, naux)`. + pub g_vix: Option>, } pub fn get_rpt2_elec_deriv_incore( @@ -59,9 +71,12 @@ where let RPT2ElecDerivIncoreInp { cderi, cderi_vox, occ_coeff, vir_coeff, occ_energy, vir_energy, index_occ_outer_vec } = input; - let RPT2ElecDerivIncoreArg { c_os, c_ss } = arg; + let RPT2ElecDerivIncoreArg { c_os, c_ss, full_gfock, dump_g_vix } = arg; + let full_gfock = *full_gfock; + let dump_g_vix = *dump_g_vix; - let bi1_scale = O::from_f64(2.0 * c_os).unwrap(); + // biorthogonal amplitude `T = (c_os + c_ss) t - c_ss t^T` (forge `restricted_biorthogonalize`) + let bi1_scale = O::from_f64(c_os + c_ss).unwrap(); let bi2_scale = O::from_f64(*c_ss).unwrap(); let device = cderi.device().clone(); @@ -89,8 +104,30 @@ where let eng_corr_double: Arc> = Arc::new(Mutex::new(0.0)); let mut gfock_part = rt::zeros(([nmo, nmo].f(), &device)); let mut rdm1_corr = rt::zeros(([nmo, nmo].f(), &device)); + let mut g_vix_dump: Option> = if dump_g_vix { Some(rt::zeros(([nvir, nocc, naux].f(), &device))) } else { None }; let w3: Arc>> = Arc::new(Mutex::new(rt::zeros(([nocc, nvir].f(), &device)))); - let w4: Arc>> = Arc::new(Mutex::new(rt::zeros(([nvir, nocc].f(), &device)))); + let w4: Option>>> = if full_gfock { + None + } else { + Some(Arc::new(Mutex::new(rt::zeros(([nvir, nocc].f(), &device))))) + }; + // full-mode gfock block accumulators (the 4-contribution parts; orbital-energy terms and + // the response term are added outside the auxiliary loop) + let g_oo: Option>>> = if full_gfock { + Some(Arc::new(Mutex::new(rt::zeros(([nocc, nocc].f(), &device))))) + } else { + None + }; + let g_vv: Option>>> = if full_gfock { + Some(Arc::new(Mutex::new(rt::zeros(([nvir, nvir].f(), &device))))) + } else { + None + }; + let g_vo: Option>>> = if full_gfock { + Some(Arc::new(Mutex::new(rt::zeros(([nvir, nocc].f(), &device))))) + } else { + None + }; // --- buffer allocation --- // @@ -184,40 +221,84 @@ where VAL_1, VAL_0, ); + if let Some(g_vix_dump) = g_vix_dump.as_mut() { + g_vix_dump + .i_mut((.., io_slice[0]..io_slice[1], ..)) + .assign(&G_vix.mapv(|x| x.to_f64().unwrap())); + } // --- block-4 --- // let mocc_batch = occ_coeff.i((.., io_slice[0]..io_slice[1])); - let w4_bi: Arc>> = Arc::new(Mutex::new(rt::zeros(([nao, nstep_occ].f(), &device)))); + let [io_lo, io_hi] = [io_slice[0], io_slice[1]]; + let w4_bi: Option>>> = if !full_gfock { + Some(Arc::new(Mutex::new(rt::zeros(([nao, nstep_occ].f(), &device))))) + } else { + None + }; (0..naux).into_par_iter().for_each(|p| { // cderi[oix]: ao2mo let cderi_sy = cderi.i((.., p)).unpack_tri(Upper, FlagSymm::Sy).mapv(&cast); let cderi_bi = &cderi_sy % &mocc_batch; let cderi_oi = occ_coeff.t() % &cderi_bi; - // w3[ov] = cderi[oi, x] % G[vi, x]' + // w3[ov] = cderi[oi, x] % G[vi, x]'; OV block contribution in both modes let w3_part = cderi_oi % G_vix.i((.., .., p)).t(); { let mut w3 = w3.lock().unwrap(); *w3 += &w3_part; } - // w4[bi] = cderi[bb, x] % G[bi, x] - let G_bi = &vir_coeff % G_vix.i((.., .., p)); - let w4_bi_part = &cderi_sy % &G_bi; - { - let mut w4_bi = w4_bi.lock().unwrap(); - *w4_bi += &w4_bi_part; + if full_gfock { + // remaining blocks from the same per-aux column: VO from Y_vv[ab, x] built on + // the fly, OO/VV from the G[vi, x] / cderi_vox[vi, x] pair + let y_vv = vir_coeff.t() % (&cderi_sy % &vir_coeff); + let vo_part = &y_vv % G_vix.i((.., .., p)); + let oo_part = G_vix.i((.., .., p)).t() % cderi_vox.i((.., .., p)); + let vv_part = G_vix.i((.., .., p)) % cderi_vox.i((.., io_lo..io_hi, p)).t(); + { + let mut g_vo = g_vo.as_ref().unwrap().lock().unwrap(); + *&mut g_vo.i_mut((.., io_lo..io_hi)) += &vo_part; + let mut g_oo = g_oo.as_ref().unwrap().lock().unwrap(); + *&mut g_oo.i_mut((io_lo..io_hi, ..)) += &oo_part; + let mut g_vv = g_vv.as_ref().unwrap().lock().unwrap(); + *g_vv += &vv_part; + } + } else { + // w4[bi] = cderi[bb, x] % G[bi, x] + let G_bi = &vir_coeff % G_vix.i((.., .., p)); + let w4_bi_part = &cderi_sy % &G_bi; + { + let mut w4_bi = w4_bi.as_ref().unwrap().lock().unwrap(); + *w4_bi += &w4_bi_part; + } } }); - let w4_bi = w4_bi.lock().unwrap(); - w4.lock().unwrap().i_mut((.., io_slice[0]..io_slice[1])).assign(vir_coeff.t() % w4_bi.view()); + if !full_gfock { + let w4_bi = w4_bi.as_ref().unwrap().lock().unwrap(); + w4.as_ref().unwrap().lock().unwrap().i_mut((.., io_slice[0]..io_slice[1])).assign(vir_coeff.t() % w4_bi.view()); + } } let e_corr = *eng_corr_double.lock().unwrap(); let w3 = w3.lock().unwrap(); - let w4 = w4.lock().unwrap(); gfock_part.i_mut((so, sv)).assign(w3.mapv(|x| 4.0 * x.to_f64().unwrap())); - gfock_part.i_mut((sv, so)).assign(w4.mapv(|x| 4.0 * x.to_f64().unwrap())); + if full_gfock { + let g_oo = g_oo.as_ref().unwrap().lock().unwrap(); + let g_vv = g_vv.as_ref().unwrap().lock().unwrap(); + let g_vo = g_vo.as_ref().unwrap().lock().unwrap(); + *&mut gfock_part.i_mut((so, so)) += g_oo.mapv(|x| 4.0 * x.to_f64().unwrap()); + gfock_part.i_mut((sv, so)).assign(g_vo.mapv(|x| 4.0 * x.to_f64().unwrap())); + *&mut gfock_part.i_mut((sv, sv)) += g_vv.mapv(|x| 4.0 * x.to_f64().unwrap()); + // orbital-energy terms of the diagonal blocks: F[i, j] += 2 * eps_j * rdm1_corr[i, j]; + // energies as (1, n), broadcast down the (fast) row axis in col-major + let scale_occ = occ_energy.mapv(|x| 2.0 * x.to_f64().unwrap()); + let scale_vir = vir_energy.mapv(|x| 2.0 * x.to_f64().unwrap()); + *&mut gfock_part.i_mut((so, so)) += scale_occ.i((None, ..)) * rdm1_corr.i((so, so)); + *&mut gfock_part.i_mut((sv, sv)) += scale_vir.i((None, ..)) * rdm1_corr.i((sv, sv)); + } else { + let w4 = w4.as_ref().unwrap().lock().unwrap(); + gfock_part.i_mut((sv, so)).assign(w4.mapv(|x| 4.0 * x.to_f64().unwrap())); + } - RPT2ElecDerivIncoreOut { e_corr, gfock_part, rdm1_corr } + RPT2ElecDerivIncoreOut { e_corr, gfock_part, rdm1_corr, g_vix: g_vix_dump } } diff --git a/src/ri_pt2/pure_pt2_u_elecderiv.rs b/src/ri_pt2/pure_pt2_u_elecderiv.rs new file mode 100644 index 0000000000000000000000000000000000000000..1ef44b77f925ad24e209dbe2e4f73d8dafdfea9b --- /dev/null +++ b/src/ri_pt2/pure_pt2_u_elecderiv.rs @@ -0,0 +1,338 @@ +//! Pure functions for the unrestricted RI-PT2 electronic derivative (generalized Fock and +//! related quantities), the unrestricted twin of +//! [`crate::ri_pt2::pure_pt2_r_elecderiv`]. +//! +//! Spin blocks follow pyscf-forge `dh/resp.py::_prepare_pt2_u` (per spin $\sigma$, with the +//! opposite spin labelled $\varsigma$): +//! +//! - amplitudes $t^{\sigma\varsigma}_{ijab} = (i a_\sigma | j b_\varsigma) / \Delta$, with the +//! biorthogonal combinations $T^{\sigma\sigma} = \frac{1}{2} c_{ss} (t - t^{\top})$ (the +//! transpose over the virtual pair only; the compound $(i,a) \leftrightarrow (j,b)$ +//! antisymmetry then holds through the amplitude pair symmetry, cf. forge +//! `hermi_sum_last2dim`) and $T^{\alpha\beta} = c_{os} t$ (single-counted); +//! - correlation rdm1 per spin, $D^{\sigma}_{oo} = -2 \Sigma T^{\sigma\sigma} t^{\sigma\sigma} +//! - \Sigma T^{\alpha\beta} t^{\alpha\beta}$ (and $+ / \varsigma$-transposed contracts for the +//! virtual block); +//! - the amplitude intermediate $G^\sigma_{ia}[\mathtt{P}] = 4 \Sigma T^{\sigma\sigma} +//! (j b_\sigma | \mathtt{P}) + 2 \Sigma T^{\alpha\beta} (j b_\varsigma | \mathtt{P})$; +//! - the (full four-block) generalized Fock contribution with coefficient $1$ on the +//! $G$-contractions (the restricted kernel scales the same blocks by $4$; with the +//! $G$ normalization above the two conventions give identical $W$ blocks), plus the +//! $2 \varepsilon \cdot \mathrm{rdm1}$ orbital-energy terms of the diagonal blocks; +//! - the correlation energy $\Sigma T g$ over all three spin blocks, i.e. +//! $c_{os} E_{os} + c_{ss} E_{ss}$ already in the coefficient-weighted form. +//! +//! Every block streams its amplitudes over windows ([`UPT2ElecDerivIncoreInp::index_occ_outer_vec`] +//! / [`UPT2ElecDerivIncoreInp::index_vir_outer_vec`], cf. +//! [`occ_batch_index`](crate::ri_pt2::occ_batch_index)): only window-sized buffers exist +//! transiently and are discarded after their rdm1/$G$ accumulation, so amplitude storage is +//! bounded by the `nvir * nocc * naux` class (times the caller's headroom); a single window -- +//! the full `nocc^2 nvir^2` tensor -- only occurs when that tensor is within the bound. The +//! same-spin blocks use the pair-symmetry structure of +//! the restricted kernel (triangular pair loop with the transpose fill) and window their outer +//! occupied index. The $\alpha\beta$ block stores only the plain amplitudes ($T = c_{os} t$ +//! enters every contraction as a factor); each of its four rdm1 Grams must keep its own index +//! pair complete in the working set, and the four pair up as α-side / β-side, so the amplitudes +//! are streamed in two passes (α- then β-virtual windowed) — each a per-window partial-trace +//! GEMM pair, with no `[.]^2`-sized intermediate. + +use crate::ri_jk::pure_ao2mo::get_ao2mo_s2ij_to_s1_notrans; +use crate::utilities::rstsr_util::{Tsr, TsrView}; +use rayon::prelude::*; +use rstsr::prelude::*; +use std::sync::{Arc, Mutex}; + +pub struct UPT2ElecDerivIncoreInp<'a> { + /// Cholesky-decomposed 3c2e ERI, shape `[nao_tp, naux]`. + pub cderi: TsrView<'a>, + /// Optionally pre-transformed (vir, occ, aux) three-center integrals per spin. + pub cderi_vox: [Option; 2], + /// Occupied MO coefficients, shape `[nao, nocc_s]` per spin. + pub occ_coeff: [TsrView<'a>; 2], + /// Virtual MO coefficients, shape `[nao, nvir_s]` per spin. + pub vir_coeff: [TsrView<'a>; 2], + /// Occupied orbital energies, shape `[nocc_s]` per spin. + pub occ_energy: [TsrView<'a>; 2], + /// Virtual orbital energies, shape `[nvir_s]` per spin. + pub vir_energy: [TsrView<'a>; 2], + /// Occupancy-window boundaries of the streamed same-spin blocks, per spin (cf. + /// [`occ_batch_index`](crate::ri_pt2::occ_batch_index)): the σσ block of spin `s` streams its + /// outer occupied index through `index_occ_outer_vec[s]`. Each vector starts at 0 and ends at + /// the spin's `nocc`. + pub index_occ_outer_vec: [Vec; 2], + /// Virtual-window boundaries of the streamed opposite-spin block, per spin: entry `0` windows + /// the α virtual index of the first (energy/$G$/β-side rdm1) pass, entry `1` the β virtual + /// index of the second (α-side rdm1) pass. Each vector starts at 0 and ends at the spin's + /// `nvir`. + pub index_vir_outer_vec: [Vec; 2], +} + +pub struct UPT2ElecDerivIncoreArg { + /// Opposite-spin correlation factor $c_{os}$. + pub c_os: f64, + /// Same-spin correlation factor $c_{ss}$. + pub c_ss: f64, +} + +pub struct UPT2ElecDerivIncoreOut { + /// PT2 correlation energy $c_{os} E_{os} + c_{ss} E_{ss}$. + pub e_corr: f64, + /// Unrelaxed correlation rdm1, shape `[nmo_s, nmo_s]` per spin. + pub rdm1_corr: [Tsr; 2], + /// Full four-block generalized Fock contribution, shape `[nmo_s, nmo_s]` per spin, with the + /// orbital-energy terms of the diagonal blocks included. + pub gfock_part: [Tsr; 2], + /// Amplitude intermediate $G^\sigma$, shape `[nvir_s, nocc_s, naux]` per spin. + pub g_vix: [Tsr; 2], +} + +pub fn get_rupt2_elec_deriv_incore( + input: &UPT2ElecDerivIncoreInp, + arg: &UPT2ElecDerivIncoreArg, +) -> UPT2ElecDerivIncoreOut { + let UPT2ElecDerivIncoreInp { cderi, cderi_vox, occ_coeff, vir_coeff, occ_energy, vir_energy, index_occ_outer_vec, index_vir_outer_vec } = input; + let UPT2ElecDerivIncoreArg { c_os, c_ss } = arg; + let (c_os, c_ss) = (*c_os, *c_ss); + let device = cderi.device().clone(); + let naux = cderi.shape()[1]; + let nocc = [occ_coeff[0].shape()[1], occ_coeff[1].shape()[1]]; + let nvir = [vir_coeff[0].shape()[1], vir_coeff[1].shape()[1]]; + let nmo = [nocc[0] + nvir[0], nocc[1] + nvir[1]]; + for s in 0..2 { + assert_eq!(index_occ_outer_vec[s].first(), Some(&0), "index_occ_outer_vec[{s}] must start with 0"); + assert_eq!(index_occ_outer_vec[s].last(), Some(&nocc[s]), "index_occ_outer_vec[{s}] must end with nocc"); + assert!(index_occ_outer_vec[s].is_sorted(), "index_occ_outer_vec[{s}] must be sorted"); + } + for s in 0..2 { + assert_eq!(index_vir_outer_vec[s].first(), Some(&0), "index_vir_outer_vec[{s}] must start with 0"); + assert_eq!(index_vir_outer_vec[s].last(), Some(&nvir[s]), "index_vir_outer_vec[{s}] must end with nvir"); + assert!(index_vir_outer_vec[s].is_sorted(), "index_vir_outer_vec[{s}] must be sorted"); + } + + let cderi_vox: [Tsr; 2] = + [0, 1].map(|s| match &cderi_vox[s] { + Some(x) => x.to_owned(), + None => get_ao2mo_s2ij_to_s1_notrans( + cderi.view(), + Upper, + &[vir_coeff[s].view()], + &[occ_coeff[s].view()], + |x| x, + ) + .remove(0), + }); + + let e_corr = Arc::new(Mutex::new(0.0)); + let mut rdm1 = [ + rt::zeros(([nmo[0], nmo[0]].f(), &device)), + rt::zeros(([nmo[1], nmo[1]].f(), &device)), + ]; + let mut g_vix = [ + rt::zeros(([nvir[0], nocc[0], naux].f(), &device)), + rt::zeros(([nvir[1], nocc[1], naux].f(), &device)), + ]; + let d_vv_outer = [ + (-vir_energy[0].i((.., None)) - vir_energy[0].i((None, ..))).into_dim::(), + (-vir_energy[1].i((.., None)) - vir_energy[1].i((None, ..))).into_dim::(), + ]; + + // ── same-spin blocks: pair-symmetrized amplitudes streamed over occupied windows ── // + for s in 0..2 { + // same-spin blocks contribute nothing when the same-spin factor vanishes (xygj-os etc.) + if nocc[s] == 0 || nvir[s] == 0 || c_ss.abs() < 1.0e-12 { + continue; + } + let (nv, no) = (nvir[s], nocc[s]); + let so = 0..no; + let sv = no..nmo[s]; + // biorthogonal coefficients: T = bi1 t - bi2 t^T with T = 1/2 c_ss (t - t^T) + let (bi1, bi2) = (0.5 * c_ss, 0.5 * c_ss); + for io_slice in index_occ_outer_vec[s].windows(2) { + let nstep = io_slice[1] - io_slice[0]; + let mut t_vivo: Tsr = rt::zeros(([nv, nstep, nv, no].f(), &device)); + let mut T_vivo: Tsr = rt::zeros(([nv, nstep, nv, no].f(), &device)); + let mut g_vix_b: Tsr = rt::zeros(([nv, nstep, naux].f(), &device)); + // pairs (i, j): j <= i inside the window, all j outside it (each pair once) + let mut pair_ij: Vec<[usize; 2]> = Vec::new(); + for i in io_slice[0]..io_slice[1] { + let range_1 = io_slice[0]..=i; + let range_2 = 0..io_slice[0]; + let range_3 = io_slice[1]..no; + pair_ij.extend(range_1.chain(range_2).chain(range_3).map(|j| [i, j])); + } + let e_corr_s = e_corr.clone(); + pair_ij.into_par_iter().for_each(|[i, j]| { + let d_ij = occ_energy[s][[i]] + occ_energy[s][[j]]; + let g_vv = cderi_vox[s].i((.., i, ..)) % cderi_vox[s].i((.., j, ..)).t(); + let d_vv = &d_vv_outer[s] + d_ij; + let t_vv = &g_vv / &d_vv; + let t_vv_ = T_vv_full(&t_vv, bi1, bi2); + if j <= i { + let diag_scale = if i == j { 1.0 } else { 2.0 }; + *e_corr_s.lock().unwrap() += diag_scale * (&t_vv_ * &g_vv).sum(); + } + let mut t_vivo = unsafe { t_vivo.force_mut() }; + let mut t_vivo_T = unsafe { T_vivo.force_mut() }; + t_vivo.i_mut((.., i - io_slice[0], .., j)).assign(&t_vv); + t_vivo_T.i_mut((.., i - io_slice[0], .., j)).assign(&t_vv_); + if (io_slice[0] <= j) && (j < i) { + t_vivo.i_mut((.., j - io_slice[0], .., i)).assign(t_vv.t()); + t_vivo_T.i_mut((.., j - io_slice[0], .., i)).assign(t_vv_.t()); + } + }); + // D[oo] -= 2 sum T t; D[vv] += 2 sum T t (per-window partial sums over the i rows) + let scr = t_vivo.view().reshape((-1, no)).t() % T_vivo.view().reshape((-1, no)); + *&mut rdm1[s].i_mut((so.clone(), so.clone())) -= &(scr * 2.0); + let scr = t_vivo.view().reshape((nv, -1)) % T_vivo.view().reshape((nv, -1)).t(); + *&mut rdm1[s].i_mut((sv.clone(), sv.clone())) += &(scr * 2.0); + // G^s[same-spin part] = 4 T[viv, o] % cderi_vox[vo, x] + let mut g_vix_b_2d = rt::asarray((g_vix_b.raw_mut(), [nv * nstep, naux].f(), &device)); + g_vix_b_2d.matmul_from( + &T_vivo.view().into_shape([nv * nstep, nv * no]), + &cderi_vox[s].view().into_shape([nv * no, naux]), + 4.0, + 0.0, + ); + *&mut g_vix[s].i_mut((.., io_slice[0]..io_slice[1], ..)) += &g_vix_b; + } + } + + // ── opposite-spin block, two streaming passes ── // + // Each rdm1 piece is a Gram (outer product summed over the remaining three indices) whose + // own index pair must be complete in the working set; the four Grams pair up as α-side + // (α-occ / α-virt pairs) and β-side (β-occ / β-virt pairs), so one pass per side — each + // windowing its own α respectively β virtual index — covers all four without a `nocc^2 + // nvir^2` tensor. The energy and both G intermediates ride the first pass. + if nocc[0] > 0 && nocc[1] > 0 && nvir[0] > 0 && nvir[1] > 0 { + let (nva, nvb, noa, nob) = (nvir[0], nvir[1], nocc[0], nocc[1]); + let d_vv_ab = + (-vir_energy[0].i((.., None)) - vir_energy[1].i((None, ..))).into_dim::(); + + // --- pass 1: α-virtual-windowed, β transparent; energy, G^α/G^β, β-side Grams --- + // permuted copy [i, a, P] keeps the `(i a)` row order mergeable for the `G^β` contraction + let cderi_vox_0_iao = + cderi_vox[0].view().into_shape([nva, noa, naux]).transpose([1, 0, 2]).into_contig(ColMajor); + for io_slice in index_vir_outer_vec[0].windows(2) { + let [a0, a1] = [io_slice[0], io_slice[1]]; + let nva_b = a1 - a0; + let mut t_x: Tsr = rt::zeros(([nva_b, noa, nvb, nob].f(), &device)); + let pair_ik: Vec<[usize; 2]> = + (0..noa).flat_map(|i| (0..nob).map(move |k| [i, k])).collect(); + let e_corr_x = e_corr.clone(); + pair_ik.into_par_iter().for_each(|[i, k]| { + let d_ik = occ_energy[0][[i]] + occ_energy[1][[k]]; + let g_ab = cderi_vox[0].i((a0..a1, i, ..)) % cderi_vox[1].i((.., k, ..)).t(); + let d_ab = &d_vv_ab.i((a0..a1, ..)) + d_ik; + let t_ab = &g_ab / &d_ab; + *e_corr_x.lock().unwrap() += c_os * (&t_ab * &g_ab).sum(); + let mut t_x = unsafe { t_x.force_mut() }; + t_x.i_mut((.., i, .., k)).assign(&t_ab); + }); + // G^α[(a i), P] += 2 c_os t[(a i), (b k)] (b k|P) + let t2_a = t_x.view().into_shape([nva_b * noa, nvb * nob]); + let g_a = (t2_a.view() % cderi_vox[1].view().into_shape([nvb * nob, naux])) * (2.0 * c_os); + *&mut g_vix[0].i_mut((a0..a1, .., ..)) += &g_a.into_shape([nva_b, noa, naux]); + // G^β[(b k), P] += 2 c_os Σ_{a, i} t[(i a), (b k)] (i a|P) over the window's α + let t_ia = t_x.view().transpose([1, 0, 2, 3]).into_contig(ColMajor); + let t2_ia = t_ia.view().into_shape([noa * nva_b, nvb * nob]); + let c_ia_2d = cderi_vox_0_iao.i((.., a0..a1, ..)).view().into_shape([noa * nva_b, naux]); + let g_b = (t2_ia.view().t() % c_ia_2d.view()) * (2.0 * c_os); + *&mut g_vix[1] += &g_b.into_shape([nvb, nob, naux]); + // β-side Grams: complete β-occupied / β-virtual pairs, partial sums over the α window + let t3 = t_x.view().into_shape([nva_b * noa * nvb, nob]); // [(a i b), k] + let m_b_oo = t3.view().t() % t3.view(); + *&mut rdm1[1].i_mut((0..nob, 0..nob)) -= &(m_b_oo * c_os); + let w = t_x.view().transpose([2, 0, 1, 3]).into_contig(ColMajor); + let w = w.view().into_shape([nvb, nva_b * noa * nob]); // [b, (a i k)] + let m_b_vv = w.view() % w.view().t(); + *&mut rdm1[1].i_mut((nob..nmo[1], nob..nmo[1])) += &(m_b_vv * c_os); + } + + // --- pass 2: β-virtual-windowed, α transparent; α-side Grams (no energy/G) --- + let cderi_vox_0_2d = cderi_vox[0].view().into_shape([nva * noa, naux]); // [(a i), P] + let d_alpha_full = (-vir_energy[0].i((.., None)) + occ_energy[0].i((None, ..))).into_dim::(); + for jo_slice in index_vir_outer_vec[1].windows(2) { + let [b0, b1] = [jo_slice[0], jo_slice[1]]; + let nvb_b = b1 - b0; + let d_beta_b = (-vir_energy[1].i((b0..b1, None)) + occ_energy[1].i((None, ..))).into_dim::(); + let yb = cderi_vox[1] + .i((b0..b1, .., ..)) + .into_contig(ColMajor) + .view() + .into_shape([nvb_b * nob, naux]); + let g_all = cderi_vox_0_2d.view() % yb.view().t(); // [(a i), (b j)] + let d_all = &d_alpha_full.view().into_shape([nva * noa, 1]) + &d_beta_b.view().into_shape([1, nvb_b * nob]); + let t_x = (g_all / d_all).into_contig(ColMajor); + let t_x = t_x.view().into_shape([nva, noa, nvb_b, nob]); + // α-occupied Gram: [(a b j), i] partial over the β window + let a_oo = t_x.view().transpose([0, 2, 3, 1]).into_contig(ColMajor); + let a_oo = a_oo.view().into_shape([nva * nvb_b * nob, noa]); + let m_a_oo = a_oo.view().t() % a_oo.view(); + *&mut rdm1[0].i_mut((0..noa, 0..noa)) -= &(m_a_oo * c_os); + // α-virtual Gram: [a, (i b j)] partial over the β window + let v = t_x.view().into_shape([nva, noa * nvb_b * nob]); + let m_a_vv = v.view() % v.view().t(); + *&mut rdm1[0].i_mut((noa..nmo[0], noa..nmo[0])) += &(m_a_vv * c_os); + } + } + + // ── generalized Fock blocks: per-aux-column contractions upon the complete G^s ── // + let mut gfock = [ + rt::zeros(([nmo[0], nmo[0]].f(), &device)), + rt::zeros(([nmo[1], nmo[1]].f(), &device)), + ]; + for s in 0..2 { + if nocc[s] == 0 || nvir[s] == 0 { + continue; + } + let (nv, no) = (nvir[s], nocc[s]); + let (mocc, mvir) = (&occ_coeff[s], &vir_coeff[s]); + let g_oo: Arc> = Arc::new(Mutex::new(rt::zeros(([no, no].f(), &device)))); + let g_ov: Arc> = Arc::new(Mutex::new(rt::zeros(([no, nv].f(), &device)))); + let g_vo: Arc> = Arc::new(Mutex::new(rt::zeros(([nv, no].f(), &device)))); + let g_vv: Arc> = Arc::new(Mutex::new(rt::zeros(([nv, nv].f(), &device)))); + (0..naux).into_par_iter().for_each(|p| { + let cderi_sy = cderi.i((.., p)).unpack_tri(Upper, FlagSymm::Sy); + let cderi_oi = mocc.view().t() % (&cderi_sy % mocc.view()); + let g_p = g_vix[s].i((.., .., p)); + let vox_p = cderi_vox[s].i((.., .., p)); + let ov_part = &cderi_oi % g_p.t(); + let y_vv = mvir.view().t() % (&cderi_sy % mvir.view()); + let vo_part = &y_vv % &g_p; + let oo_part = g_p.t() % &vox_p; + let vv_part = &g_p % vox_p.t(); + *g_ov.lock().unwrap() += &ov_part; + *g_vo.lock().unwrap() += &vo_part; + *g_oo.lock().unwrap() += &oo_part; + *g_vv.lock().unwrap() += &vv_part; + }); + let g_oo = g_oo.lock().unwrap(); + let g_ov = g_ov.lock().unwrap(); + let g_vo = g_vo.lock().unwrap(); + let g_vv = g_vv.lock().unwrap(); + // orbital-energy terms of the diagonal blocks: F[i, j] += 2 eps_j rdm1[i, j] + let scale_occ = occ_energy[s].mapv(|x| 2.0 * x); + let scale_vir = vir_energy[s].mapv(|x| 2.0 * x); + let so = 0..no; + let sv = no..nmo[s]; + *&mut gfock[s].i_mut((so.clone(), so.clone())) += + &(&*g_oo + &scale_occ.i((None, ..)) * rdm1[s].i((so.clone(), so.clone()))); + gfock[s].i_mut((so.clone(), sv.clone())).assign(&*g_ov); + gfock[s].i_mut((sv.clone(), so.clone())).assign(&*g_vo); + *&mut gfock[s].i_mut((sv.clone(), sv.clone())) += + &(&*g_vv + &scale_vir.i((None, ..)) * rdm1[s].i((sv.clone(), sv.clone()))); + } + + let e_corr = *e_corr.lock().unwrap(); + UPT2ElecDerivIncoreOut { + e_corr, + rdm1_corr: rdm1, + gfock_part: gfock, + g_vix, + } +} + +/// Biorthogonal same-spin amplitude `T = bi1 t - bi2 t^T`. +fn T_vv_full(t_vv: &Tsr, bi1: f64, bi2: f64) -> Tsr { + t_vv.mapv(|x| bi1 * x) - t_vv.t().mapv(|x| bi2 * x) +} diff --git a/src/ri_pt2/rgfock_pt2.rs b/src/ri_pt2/rgfock_pt2.rs index 1d7656d8f27fc26f005f0aab4a77da2000a22bac..cd9993713e735c849c1fecc03a895a7e571df625 100644 --- a/src/ri_pt2/rgfock_pt2.rs +++ b/src/ri_pt2/rgfock_pt2.rs @@ -3,11 +3,15 @@ //! This module implements [`RGFockAPI`] for the RI-PT2 (RI-MP2-type) correlation contribution, //! along with the Z-vector (relaxed density) machinery: //! -//! - The generalized Fock contribution of PT2 has only the off-diagonal blocks filled: -//! $\mathscr{F}_{ia} = - (W^\texttt{3}_{ai})^\top$ (OV) and $\mathscr{F}_{ai} = -//! W^\texttt{4}_{ai}$ (VO), given by [`crate::ri_pt2::pure_pt2_r_elecderiv`]; the SCF -//! response upon the correlation rdm1 ($A_{ai, pq} D_{pq}^{\mathrm{RDM}}$ term of the -//! Lagrangian) is added to the VO block. +//! - The generalized Fock contribution of PT2 comes in two modes. The default fills only the +//! off-diagonal blocks: $\mathscr{F}_{ia} = - (W^\texttt{3}_{ai})^\top$ (OV) and +//! $\mathscr{F}_{ai} = W^\texttt{4}_{ai}$ (VO), given by +//! [`crate::ri_pt2::pure_pt2_r_elecderiv`]. Requests containing the OO or VV parts switch to +//! the full four-block evaluation (batched over the auxiliary index), which carries the +//! $2 \varepsilon \cdot \mathrm{rdm1}^{\mathrm{corr}}$ orbital-energy terms in the diagonal +//! blocks. In both modes the SCF response upon the correlation rdm1 +//! ($A_{ai, pq} D_{pq}^{\mathrm{RDM}}$ term of the Lagrangian) is added to the VO block (and, +//! in the full mode, to the OO block as well). //! - The Lagrangian $L_{ai}$ is the antisymmetrized generalized Fock, which gives //! $L_{ai} = W^\texttt{3}_{ai} + W^\texttt{4}_{ai} + A_{ai, pq} D_{pq}^{\mathrm{RDM}}$. //! - The Z-vector is obtained by solving the Z-vector (CP-SCF-type) equation @@ -70,18 +74,24 @@ where /// allocation of a cast copy. pub cderi_vox: Option>, /// Batching (outer slice) indices of occupied orbitals; must start with `0` and end with - /// `nocc`. Degenerate occupied orbitals must stay within one batch. + /// `nocc`. The kernels are window-agnostic: any monotonic partition works, and degenerate + /// occupied orbitals may be split across windows (each unordered pair is visited once and the + /// per-window partial sums accumulate). pub index_occ_outer_vec: Vec, /// Opposite-spin correlation factor $c_\mathrm{OS}$. pub c_os: f64, /// Same-spin correlation factor $c_\mathrm{SS}$. pub c_ss: f64, - /// Generalized Fock matrix in MO basis, accumulating the already-evaluated parts (only OV - /// and VO are implemented); a zero matrix `[nmo, nmo]` at creation. - pub gfock: Tsr, - /// The parts of `gfock` that have been evaluated (only OV and VO can ever be set). - pub gfock_flags: BitFlags, - /// Cached results, keyed by tensor name. + /// Whether to additionally cache the amplitude intermediate $G_{ai}^\mathtt{P}$ under the + /// result key `"g_vix"` (used by the analytic-gradient driver); costs an extra + /// `nvir * nocc * naux` tensor. Must be set before the first evaluation (a cached kernel is + /// not re-run; enabling it later panics rather than omitting `"g_vix"`). + pub dump_g_vix: bool, + /// Cached results, keyed by tensor name. Keys used by this struct: `"rdm1"`, + /// `"gfock_part"`/`"gfock_part_full"` (electronic-derivative gfock contribution in the two + /// modes of [`Self::make_elec_deriv`]), `"gfock_oo"`/`"gfock_ov"`/`"gfock_vo"`/`"gfock_vv"` + /// (blocks of the generalized Fock), `"axd"`, `"axd_vo"`, `"lagrangian"`, `"z_vector"`, + /// `"rdm1_resp"`. pub result: HashMap, /// Correlation energy of the most recent electronic-derivative evaluation, if any. pub e_corr: Option, @@ -104,8 +114,6 @@ where c_os: f64, c_ss: f64, ) -> Self { - let nmo = mo_occ.shape()[0]; - let device = mo_coeff.device().clone(); Self { mo_coeff, mo_occ, @@ -115,8 +123,7 @@ where index_occ_outer_vec, c_os, c_ss, - gfock: rt::zeros(([nmo, nmo].f(), &device)), - gfock_flags: BitFlags::empty(), + dump_g_vix: false, result: HashMap::new(), e_corr: None, timing: Vec::new(), @@ -127,7 +134,22 @@ where /// partial generalized Fock with its OV/VO blocks filled by the energy-weighted terms), and /// cache the results. Repeated calls directly return the cached results. pub fn make_elec_deriv(&mut self) { - if self.result.contains_key("rdm1") { + self.make_elec_deriv_with(false); + } + + /// [`Self::make_elec_deriv`] with the gfock mode selected: `full_gfock` evaluates the full + /// four-block contribution (OO+OV+VO+VV with the orbital-energy terms) in place of the + /// W3/W4 off-diagonal blocks. The two modes cache their `gfock_part` under separate keys + /// (`"gfock_part"`/`"gfock_part_full"`); a cached legacy evaluation is rerun on a full + /// request (only `gfock_part` differs; e_corr and rdm1 are mode-independent), and a cached + /// full evaluation serves any request. + fn make_elec_deriv_with(&mut self, full_gfock: bool) { + let key = if full_gfock { "gfock_part_full" } else { "gfock_part" }; + if self.result.contains_key(key) || (!full_gfock && self.result.contains_key("gfock_part_full")) { + assert!( + !self.dump_g_vix || self.result.contains_key("g_vix"), + "dump_g_vix was enabled after the cached electronic-derivative evaluation; set it before the first make_elec_deriv call." + ); return; } let t0 = std::time::Instant::now(); @@ -148,23 +170,33 @@ where vir_energy: evir.view(), index_occ_outer_vec: &self.index_occ_outer_vec, }; - let arg = RPT2ElecDerivIncoreArg { c_os: self.c_os, c_ss: self.c_ss }; + let arg = RPT2ElecDerivIncoreArg { c_os: self.c_os, c_ss: self.c_ss, full_gfock, dump_g_vix: self.dump_g_vix }; let out = get_rpt2_elec_deriv_incore(&input, &arg, |x| O::from_f64(x).unwrap()); self.e_corr = Some(out.e_corr); - self.result.insert("gfock_part".to_string(), out.gfock_part); + self.result.insert(key.to_string(), out.gfock_part); self.result.insert("rdm1".to_string(), out.rdm1_corr); + if let Some(g_vix) = out.g_vix { + self.result.insert("g_vix".to_string(), g_vix); + } self.timing.push(("in RGFockPT2, make_elec_deriv".to_string(), t0.elapsed().as_secs_f64())); } - /// SCF response upon the (full MO-space) correlation rdm1, contracted to the vir-occupied - /// block: $A_{ai, pq} D_{pq}^{\mathrm{RDM}}$. Shape `[nvir, nocc]`. + /// The generalized-Fock contribution cached by [`Self::make_elec_deriv`]: the full + /// four-block result if that mode ran (a cached full evaluation serves any request), + /// otherwise the legacy W3/W4 off-diagonal result. + fn gfock_part_cached(&self) -> &Tsr { + self.result.get("gfock_part_full").or_else(|| self.result.get("gfock_part")).unwrap() + } + + /// SCF response upon the (full MO-space) correlation rdm1, in MO basis: + /// $A_{pq, \lambda} D_{\lambda}^{\mathrm{RDM}}$. Shape `[nmo, nmo]`. /// /// # Parameters /// /// - `resp` : response objects of the SCF-iteration functional. - pub fn make_axd_vo<'r>(&mut self, resp: &mut (dyn RRespAPI + 'r)) -> Tsr { - if !self.result.contains_key("axd_vo") { + pub fn make_axd<'r>(&mut self, resp: &mut (dyn RRespAPI + 'r)) -> Tsr { + if !self.result.contains_key("axd") { let t0 = std::time::Instant::now(); self.make_elec_deriv(); let rdm1 = self.result["rdm1"].view(); @@ -173,13 +205,29 @@ where // (energy-derivative) side and evaluates on the SCF-grade resource; a future input // keyword may open the low-precision (`prec = false`) window here as well let resp_ao = resp.get_response_rdm(dm_ao.view(), true); - let axd_mo = self.mo_coeff.view().t() % resp_ao % self.mo_coeff.view(); + let axd_mo = self.mo_coeff.t() % resp_ao % self.mo_coeff.view(); + self.result.insert("axd".to_string(), axd_mo); + self.timing.push(("in RGFockPT2, make_axd".to_string(), t0.elapsed().as_secs_f64())); + } + self.result["axd"].to_owned() + } + + /// Vir-occupied block of [`Self::make_axd`]: $A_{ai, pq} D_{pq}^{\mathrm{RDM}}$. Shape + /// `[nvir, nocc]`. + /// + /// # Parameters + /// + /// - `resp` : response objects of the SCF-iteration functional. + pub fn make_axd_vo<'r>(&mut self, resp: &mut (dyn RRespAPI + 'r)) -> Tsr { + if !self.result.contains_key("axd_vo") { + // `make_axd` performs the contraction on the SCF-grade resource (`prec = true`), as + // this generalized-Fock side term is energy-derivative, not CP-SCF + let axd = self.make_axd(resp); let nocc = self.nocc(); let nmo = self.nmo(); let so = rt::slice!(0, nocc); let sv = rt::slice!(nocc, nmo); - self.result.insert("axd_vo".to_string(), axd_mo.i((sv, so)).into_contig(ColMajor)); - self.timing.push(("in RGFockPT2, make_axd_vo".to_string(), t0.elapsed().as_secs_f64())); + self.result.insert("axd_vo".to_string(), axd.i((sv, so)).into_contig(ColMajor)); } self.result["axd_vo"].to_owned() } @@ -201,7 +249,7 @@ where let nmo = self.nmo(); let so = rt::slice!(0, nocc); let sv = rt::slice!(nocc, nmo); - let gfock_part = self.result["gfock_part"].view(); + let gfock_part = self.gfock_part_cached().view(); let mut lag = gfock_part.i((sv, so)).to_owned() - gfock_part.i((so, sv)).t(); lag += self.make_axd_vo(resp).view(); self.result.insert("lagrangian".to_string(), lag); @@ -259,53 +307,74 @@ impl RGFockAPI for RGFockPT2<'_, O> where O: BlasFloat + ToPrimitive + FromPrimitive + NumAssignOps + 'static, { - /// Generalized Fock of the PT2 contribution. Only the OV and VO parts are implemented - /// ($\mathscr{F}_{ia} = - (W^\texttt{3}_{ai})^\top$ and $\mathscr{F}_{ai} = - /// W^\texttt{4}_{ai}$, plus the SCF response upon the correlation rdm1 - /// $A_{ai, pq} D_{pq}^{\mathrm{RDM}}$ in the VO block); requesting the OO or VV parts is - /// rejected. Parts that are not requested are left zero. Evaluated parts are accumulated in - /// `self.gfock` (tracked by `self.gfock_flags`) and directly reused on later calls. The - /// response object is required only when the VO part is actually evaluated. + /// Generalized Fock of the PT2 contribution. Requests within OV/VO are served by the + /// W3/W4 off-diagonal blocks; any request containing the OO or VV parts switches to the + /// full four-block evaluation (batched over the auxiliary index), whose diagonal blocks + /// carry the $2 \varepsilon \cdot \mathrm{rdm1}^{\mathrm{corr}}$ orbital-energy terms. + /// The SCF response upon the correlation rdm1 ($A_{ai, pq} D_{pq}^{\mathrm{RDM}}$) enters + /// the OO and VO blocks (the density does not respond to virtual-coefficient changes); + /// OV and VV carry no response term. Parts that are not requested are left zero. + /// + /// Each block is evaluated once and cached in `self.result` (keys `"gfock_oo"`/ + /// `"gfock_ov"`/`"gfock_vo"`/`"gfock_vv"`); later calls, with whatever `flags`, directly + /// reassemble the result from the cached blocks. The `resp` object is not part of the + /// cache: it is consulted only while a block is actually being evaluated, and ignored + /// (not compared, not re-contracted) on calls served from the cache. It is required only + /// when the OO or VO part is evaluated and the SCF response term is not yet cached. fn make_gfock<'r>(&mut self, resp: Option<&mut (dyn RRespAPI + 'r)>, flags: BitFlags) -> Tsr { - assert!( - !flags.contains(GFockFlags::OO) && !flags.contains(GFockFlags::VV), - "RI-PT2 generalized Fock has only the OV and VO parts implemented; OO and VV are rejected" - ); - let t0 = std::time::Instant::now(); + let full = flags.contains(GFockFlags::OO) || flags.contains(GFockFlags::VV); + self.make_elec_deriv_with(full); + let nocc = self.nocc(); let nmo = self.nmo(); let so = rt::slice!(0, nocc); let sv = rt::slice!(nocc, nmo); - // evaluate the requested parts that are not yet in `self.gfock` - if flags.contains(GFockFlags::OV) && !self.gfock_flags.contains(GFockFlags::OV) { - self.make_elec_deriv(); - *&mut self.gfock.i_mut((so, sv)) += &self.result["gfock_part"].view().i((so, sv)); - self.gfock_flags.insert(GFockFlags::OV); + // the SCF response term, required by the OO and VO blocks being evaluated + let need_axd = (flags.contains(GFockFlags::OO) && !self.result.contains_key("gfock_oo")) + || (flags.contains(GFockFlags::VO) && !self.result.contains_key("gfock_vo")); + if need_axd && !self.result.contains_key("axd") { + self.make_axd(resp.expect( + "RI-PT2 generalized Fock OO/VO parts require the response object (for the SCF response upon rdm1_corr)", + )); } - if flags.contains(GFockFlags::VO) && !self.gfock_flags.contains(GFockFlags::VO) { - self.make_elec_deriv(); - let resp = resp.expect( - "RI-PT2 generalized Fock VO part requires the response object (for the SCF response upon rdm1_corr)", - ); - // W^4 term, plus the SCF response upon rdm1_corr, in the vir-occupied block - *&mut self.gfock.i_mut((sv, so)) += &self.result["gfock_part"].view().i((sv, so)); - let axd_vo = self.make_axd_vo(resp); - *&mut self.gfock.i_mut((sv, so)) += &axd_vo; - self.gfock_flags.insert(GFockFlags::VO); + + // evaluate the requested blocks that are not yet cached + let mut evaluated = false; + for (flag, key, rows, cols, with_axd) in [ + (GFockFlags::OO, "gfock_oo", so, so, true), + (GFockFlags::OV, "gfock_ov", so, sv, false), + (GFockFlags::VO, "gfock_vo", sv, so, true), + (GFockFlags::VV, "gfock_vv", sv, sv, false), + ] { + if !flags.contains(flag) || self.result.contains_key(key) { + continue; + } + let mut block = self.gfock_part_cached().view().i((rows, cols)).into_contig(ColMajor); + if with_axd { + block += self.result["axd"].view().i((rows, cols)); + } + self.result.insert(key.to_string(), block); + evaluated = true; } - // re-construct the matrix of the requested parts, not the accumulated `self.gfock` + // re-construct the matrix of the requested parts from the cached blocks let device = self.mo_coeff.device().clone(); let mut gfock: Tsr = rt::zeros(([nmo, nmo].f(), &device)); - if flags.contains(GFockFlags::OV) { - gfock.i_mut((so, sv)).assign(&self.gfock.i((so, sv))); + for (flag, key, rows, cols, _) in [ + (GFockFlags::OO, "gfock_oo", so, so, true), + (GFockFlags::OV, "gfock_ov", so, sv, false), + (GFockFlags::VO, "gfock_vo", sv, so, true), + (GFockFlags::VV, "gfock_vv", sv, sv, false), + ] { + if flags.contains(flag) { + gfock.i_mut((rows, cols)).assign(&self.result[key].view()); + } } - if flags.contains(GFockFlags::VO) { - gfock.i_mut((sv, so)).assign(&self.gfock.i((sv, so))); + if evaluated { + self.timing.push(("in RGFockPT2, make_gfock".to_string(), t0.elapsed().as_secs_f64())); } - self.timing.push(("in RGFockPT2, make_gfock".to_string(), t0.elapsed().as_secs_f64())); gfock } diff --git a/src/scf_io/fchk.rs b/src/scf_io/fchk.rs index 4c5f46c7fa56e49202c6618d189f3652dd565f55..82c89092747d8d03123da31c6682391576dccc09 100644 --- a/src/scf_io/fchk.rs +++ b/src/scf_io/fchk.rs @@ -8,6 +8,7 @@ use crate::external_libs::py2fch; use crate::scf_io::SCF; use crate::scf_io::SCFType; use crate::x2c::RelativisticMethod; +use tensors::MatrixFull; macro_rules! dump_real_r2f { @@ -47,6 +48,13 @@ impl SCF { println!("\nSave fch file"); self.create_fchk_head(); self.fchk_write_mo(); + // the DH density dump (see `SCF::dh_rdm1_resp`) + if let Some(sections) = self.dh_rdm1_resp.as_ref() { + let labels = ["Total MP2 Density", "Spin MP2 Density"]; + for (label, dm) in labels.iter().zip(sections.iter()) { + self.fchk_append_density_ao(label, dm); + } + } } pub fn create_fchk_head(&self) { @@ -357,29 +365,29 @@ impl SCF { } // leave MOs blank, it will be written by librest2fch let n_dm = nbf*(nbf+1)/2; - // Total SCF density: real data for restricted closed-shell calculations (the - // lower-triangular, column-packed AO density, permuted to Gaussian's AO order so it - // agrees with the MO coefficients written by librest2fch); a zero placeholder - // elsewhere (librest2fch only requires the section header to exist). librest2fch - // copies the section content verbatim when regenerating the MO coefficients. - // - // NOTE: the unrestricted (`spin_channel == 2`) case still writes the zero placeholder; - // dumping real alpha/beta ("Total SCF Density" plus spin-polarized sections in - // Gaussian's convention) is left as future work. - let mut packed_dm = vec![0.0f64; n_dm]; - if self.mol.spin_channel == 1 { - let dm = &self.density_matrix[0]; - let perm = self.gaussian_ao_permutation(); - let mut i_index = 0; - for j in 0..nbf { - for i in 0..=j { - packed_dm[i_index] = dm.data[perm[i] + perm[j] * nbf]; - i_index += 1; - } - } - } + // Total SCF density: the lower-triangular, column-packed AO density, permuted to + // Gaussian's AO order so it agrees with the MO coefficients written by librest2fch. For + // open-shell UHF runs the alpha/beta split is written in Gaussian's convention as + // `Total SCF Density` (alpha + beta) plus `Spin SCF Density` (alpha - beta); a zero + // placeholder is kept for ROHF (its fractional occupations have no alpha/beta split). + // librest2fch only requires the section header to exist and copies the section content + // verbatim when regenerating the MO coefficients. + let packed_dm = match self.scftype { + SCFType::UHF => { + let dm_total = self.density_matrix[0].clone() + self.density_matrix[1].clone(); + self.fchk_pack_density_ao(&dm_total) + }, + _ if self.mol.spin_channel == 1 => self.fchk_pack_density_ao(&self.density_matrix[0]), + _ => vec![0.0f64; n_dm], + }; write!(input, "Total SCF Density R N={:12}\n", n_dm); dump_real_r2f!(input, packed_dm); + if let SCFType::UHF = self.scftype { + let dm_spin = self.density_matrix[0].clone() - self.density_matrix[1].clone(); + let packed_spin = self.fchk_pack_density_ao(&dm_spin); + write!(input, "Spin SCF Density R N={:12}\n", n_dm); + dump_real_r2f!(input, packed_spin); + } input.sync_all().unwrap(); } @@ -457,4 +465,31 @@ impl SCF { input.sync_all().unwrap(); } + /// Permute an AO-basis density matrix into Gaussian's AO order (consistent with the MO + /// coefficients of the same fchk file) and pack its lower triangle by column — the layout + /// of the fchk density sections (same as `Total SCF Density`). + fn fchk_pack_density_ao(&self, dm_ao: &MatrixFull) -> Vec { + let nbf = dm_ao.size[0]; + assert_eq!(dm_ao.size[1], nbf, "The packed AO density must be square."); + let perm = self.gaussian_ao_permutation(); + assert_eq!(perm.len(), nbf, "AO permutation size mismatch."); + let mut packed = Vec::with_capacity(nbf * (nbf + 1) / 2); + for pj in 0..nbf { + let j = perm[pj]; + for pi in 0..=pj { + packed.push(dm_ao.data[perm[pi] + j * nbf]); + } + } + packed + } + + /// Append a symmetric AO-basis density as the Gaussian fchk density section `label` + /// (e.g. `Total MP2 Density`): the density is permuted into Gaussian's AO order, packed + /// lower-triangular by column and appended through [`SCF::fchk_append_density_section`] + /// (which documents the file/label requirements). + pub fn fchk_append_density_ao(&self, label: &str, dm_ao: &MatrixFull) { + let packed = self.fchk_pack_density_ao(dm_ao); + self.fchk_append_density_section(label, &packed); + } + } \ No newline at end of file diff --git a/src/scf_io/mod.rs b/src/scf_io/mod.rs index 4b14b811296e888a7d7378f6d2692c7063fdec57..4567480e25173854b61aa475a4e289c0e85b73eb 100644 --- a/src/scf_io/mod.rs +++ b/src/scf_io/mod.rs @@ -138,6 +138,12 @@ pub struct SCF { /// `(excitation energy, eigenvector)` in the solver's ordering. Used by /// the TDDFT analytic-gradient driver. pub tddft_excitations: Option)>>, + /// Total relaxed (response) density of the most recent analytic double-hybrid (PT2-family) + /// force evaluation, prepared for the fchk density dump: the section matrices in + /// `[Total MP2 Density]` (restricted) / `[Total, Spin MP2 Density]` (unrestricted) order. + /// Written out by the `outputs` pass when it writes the fchk (`post_scf_output`); `None` + /// otherwise. + pub dh_rdm1_resp: Option>>, /// State restored from a GW/evGW checkpoint when /// `resume_from_checkpoint = true` (see `crate::fileop::gw_checkpoint`). /// @@ -215,6 +221,7 @@ impl SCF { solvent_scf: None, scf_converged: false, tddft_excitations: None, + dh_rdm1_resp: None, gw_checkpoint_state: None, }; diff --git a/tests/analdrv/response/pt2_rgfock.rs b/tests/analdrv/response/pt2_rgfock.rs index a16ae566d85f98008ba8935b4acaa6891a28a5e4..f40bffbfe4f285503c3679b310b6f2dfff69b4ee 100644 --- a/tests/analdrv/response/pt2_rgfock.rs +++ b/tests/analdrv/response/pt2_rgfock.rs @@ -3,11 +3,14 @@ use pyrest::analdrv::response::rgfock_interface::solve_z_vector; use pyrest::analdrv::response::rresp_interface::rscf_resp_interface; use pyrest::analdrv::response::trait_rgfock::{GFockFlags, RGFockAPI}; use pyrest::molecule_io::Molecule; +use pyrest::ri_jk::decompose::generate_rimatr_bare_on_cint; use pyrest::ri_jk::get_ao2mo_s2ij_to_s1_notrans; -use pyrest::ri_jk::util::get_cint_mol; +use pyrest::ri_jk::util::{get_cint_aux_on_path, get_cint_mol}; use pyrest::ri_pt2::rgfock_pt2::RGFockPT2; use pyrest::scf_io::{self, scf_without_build}; -use pyrest::utilities::rstsr_util::{RestTensorToRstsrTsrAPI, RestTensorToRstsrViewAPI}; +use pyrest::utilities::rstsr_util::{ + fingerprint_f64, RestTensorToRstsrTsrAPI, RestTensorToRstsrViewAPI, +}; use pyrest::{ctrl_io, ri_pt2}; use rstsr::prelude::*; @@ -102,7 +105,7 @@ fn test_nh3() { vir_energy: vir_energy.view(), index_occ_outer_vec: &[0, 2, 5], }; - let arg = RPT2ElecDerivIncoreArg { c_os: 1.0, c_ss: 1.0 }; + let arg = RPT2ElecDerivIncoreArg { c_os: 1.0, c_ss: 1.0, full_gfock: false, dump_g_vix: false }; let output = get_rpt2_elec_deriv_incore(&input, &arg, |x| x); println!("MP2 correlation energy: {}", output.e_corr); let e_corr_ref = -0.245426806393; @@ -226,7 +229,7 @@ fn test_nh3_fp32() { vir_energy: vir_energy.view(), index_occ_outer_vec: &[0, 2, 5], }; - let arg = RPT2ElecDerivIncoreArg { c_os: 1.0, c_ss: 1.0 }; + let arg = RPT2ElecDerivIncoreArg { c_os: 1.0, c_ss: 1.0, full_gfock: false, dump_g_vix: false }; let output = get_rpt2_elec_deriv_incore(&input, &arg, |x| x as f32); let e_corr_ref = -0.245426806393; println!("MP2 correlation energy (f32): {}", output.e_corr); @@ -293,4 +296,103 @@ fn test_nh3_fp32() { assert!(rt::allclose(&dip_resp, &dip_resp_ref, (1e-4, 1e-6))); } +static INPUT_NH3_SVP: &str = r##" +[ctrl] + print_level = 2 + num_threads = 4 + xc = "mp2" + basis_path = "def2-svp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 1.0 + spin_polarization = false + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 1.0 0.0 0.0 + H -0.34202014 0.0 0.93969262 + H -0.34202014 0.81379768 -0.46984631 + """ +"##; + +/// Full four-block (OO+OV+VO+VV) generalized Fock of RI-PT2 against the pyscf reference +/// pipeline `MP2-gfock-final.py` (NH3/def2-SVP): the PT2 part runs on the mp2fit auxiliary +/// basis `def2-svp-rifit` (generated here on separate CInt integrators), while the SCF and +/// the response A-terms use the SCF's `def2-universal-jkfit` (element-wise identical to +/// pyscf's `def2-svp-jkfit` for N/H). The fingerprint is evaluated on the transposed view: +/// the matrix is non-symmetric, and `fingerprint_f64` ravels column-major while pyscf's +/// `lib.fingerprint` ravels row-major. +#[test] +fn test_nh3_gfock_full() { + let keys = toml::from_str::(&INPUT_NH3_SVP[..]).unwrap(); + let (ctrl, geom) = ctrl_io::parse_ctl_from_json(&keys).unwrap(); + let mol = Molecule::build_native(ctrl, geom, None).unwrap(); + let mut scf_data = scf_io::SCF::build(mol, &None); + scf_without_build(&mut scf_data, &None); + + let device = rt::DeviceBLAS::default(); + let mo_energy = (&scf_data.eigenvalues[0]).to_rstsr(&device); + let mut mo_coeff = (&scf_data.eigenvectors[0]).to_rstsr(&device); + + // canonicalize MO signs (largest-|element| of each column positive): eigenvector signs are + // solver-arbitrary, and the pyscf reference orbitals of `MP2-gfock-final.py` follow this + // convention, which makes the elementwise MO-basis fingerprint comparison well-defined + for j in 0..mo_coeff.shape()[1] { + let col = mo_coeff.i((.., j)); + let (imax, _) = col.iter().enumerate().max_by(|a, b| a.1.abs().partial_cmp(&b.1.abs()).unwrap()).unwrap(); + let scale = if col[[imax]] < 0.0 { -1.0 } else { 1.0 }; + let scaled = col.to_owned() * scale; + mo_coeff.i_mut((.., j)).assign(&scaled); + } + + // decomposed ERI of the mp2fit auxiliary basis for the PT2 part + let cderi_mp2_raw = generate_rimatr_bare_on_cint( + &get_cint_mol(&scf_data.mol), + &get_cint_aux_on_path(&scf_data.mol, "def2-svp-rifit"), + scf_data.mol.ctrl.j2c_decomp, + None, + ); + let cderi_mp2 = cderi_mp2_raw.to_rstsr_view(&device); + + let idx_core = scf_data.mol.start_mo; + let idx_lumo = scf_data.lumo[0]; + let num_mo = mo_energy.size(); + let mut mo_occ = rt::zeros(([num_mo].f(), &device)); + mo_occ.i_mut(idx_core..idx_lumo).fill(2.0); + + let config = AnalDrvConfig::default(); + let mut resp_objs = rscf_resp_interface(&scf_data, &config); + let mut rgfock = RGFockPT2::::new( + mo_coeff.to_owned(), + mo_occ, + mo_energy.to_owned(), + cderi_mp2.into_cow(), + None, + vec![0, 2, 5], + 1.0, + 1.0, + ); + resp_objs.make_cpscf_preparation(mo_coeff.view(), rgfock.mo_occ.view(), mo_energy.view()); + + let gfock = rgfock.make_gfock( + Some(&mut resp_objs), + GFockFlags::OO | GFockFlags::OV | GFockFlags::VO | GFockFlags::VV, + ); + println!("PT2 correlation energy: {}", rgfock.e_corr.unwrap()); + let gfock_t = gfock.view().t().into_contig(ColMajor); + let fingerprint = fingerprint_f64(gfock_t.view()); + println!("full gfock fingerprint: {fingerprint:.17} (ref 0.24156063812028558)"); + // tolerance floor: the reference carries the pyscf default-SCF convergence noise of its own + // recorded run (a tight-convergence rerun of the same script differs by ~1.6e-8, and the + // present result agrees with that rerun to ~5e-9) + assert!( + (fingerprint - 0.24156063812028558).abs() < 1e-7, + "full gfock fingerprint mismatch" + ); +} + // --- following is utilities for developing dipole evaluation --- // diff --git a/tests/grad/mod.rs b/tests/grad/mod.rs index 2fd363ba7c8e8219f42a08294c5ab0c6f6342f5b..de5fd44531286c527b3f04a6cab0ea600c464c2e 100644 --- a/tests/grad/mod.rs +++ b/tests/grad/mod.rs @@ -1,3 +1,5 @@ +mod rdh; +mod udh; mod ri_jk_rhf; mod ri_jk_rks; mod ri_jk_uhf; diff --git a/tests/grad/rdh.rs b/tests/grad/rdh.rs new file mode 100644 index 0000000000000000000000000000000000000000..512e2f3b9fa95323b536b995d31bf2f83b3a8cbc --- /dev/null +++ b/tests/grad/rdh.rs @@ -0,0 +1,190 @@ +//! Doubly-hybrid (`xc = "mp2"`) analytic gradient of distorted NH3 (restricted RI-MP2). +//! +//! Reference: pyscf-forge `DFDH(xc="MP2")` analytic gradient (pyscf 2.14.0 + pyscf-forge), +//! def2-SVP, def2-universal-jkfit for both the SCF JK fitting and the RI (3-center) part, +//! CPKS tolerance 1e-8. The JSON reference was dumped per atom in Hartree/bohr (row-major +//! `[natom, 3]`); REST stores the gradient as `[3, natm]`. +//! The correlation part below is the `"DH"` entry of `eval_force` (the `"SCF"` entry covers +//! the HF part). Both `[ctrl.j2c_decomp]` policies (`cd`, `eig`) are exercised; the eig +//! gradients agree with the cd ones to roundoff (~1e-13; the auxiliary basis is +//! well-conditioned), so the same references are asserted with the same tolerances. + +use pyrest::grad::rdh::RDHGradient; +use pyrest::grad::traits::GradAPI; +use pyrest::molecule_io::Molecule; +use pyrest::ri_pt2; +use pyrest::scf_io::{self, scf_without_build}; +use pyrest::ctrl_io; +use rstsr::prelude::*; + +static INPUT_NH3: &str = r##" +[ctrl] + print_level = 0 + num_threads = 4 + xc = "mp2" + basis_path = "def2-SVP" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 1.0 + spin_polarization = false + +[ctrl.j2c_decomp] +policy = "POLICY" + +[ctrl.ri_pt2] +new_driver = true + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0000000000 0.0000000000 0.0000000000 + H 0.0000000000 0.0000000000 1.1500000000 + H 0.9674007198 0.0000000000 -0.2902341250 + H -0.4917018117 0.8516525201 -0.3062961203 + """ +"##; + +static INPUT_XYG3: &str = r##" +[ctrl] + print_level = 0 + num_threads = 4 + xc = "xyg3" + basis_path = "def2-SVP" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 1.0 + spin_polarization = false + +[ctrl.j2c_decomp] +policy = "POLICY" + +[ctrl.ri_pt2] +new_driver = true + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0000000000 0.0000000000 0.0000000000 + H 0.0000000000 0.0000000000 1.1500000000 + H 0.9674007198 0.0000000000 -0.2902341250 + H -0.4917018117 0.8516525201 -0.3062961203 + """ +"##; + +fn test_nh3_rmp2_grad_with(policy: &str, tol: f64) { + let keys = toml::from_str::(&INPUT_NH3.replace("POLICY", policy)).unwrap(); + let (ctrl, geom) = ctrl_io::parse_ctl_from_json(&keys).unwrap(); + let mol = Molecule::build_native(ctrl, geom, None).unwrap(); + let mut scf_data = scf_io::SCF::build(mol, &None); + scf_without_build(&mut scf_data, &None); + + // total energy (RI-MP2); also puts the SCF into the post-SCF state the gradient expects + let eng = ri_pt2::xdh_calculations(&mut scf_data, &None).unwrap(); + println!("RI-MP2 total energy ({policy}): {eng}"); + assert!((eng - (-56.32508171804932)).abs() < 2e-6, "RI-MP2 total energy mismatch"); + + let mut grad = RDHGradient::new(&scf_data, &None); + grad.calc(); + let de = grad.result.get("de").unwrap().clone(); + let de = rt::asarray((de.data, de.size)); + + // pyscf-forge `gradient_corr_only_hartree_per_bohr` (natom, 3) -> REST [3, natm] + #[rustfmt::skip] + let de_ref = vec![ + 0.0065780504, 0.0111234158, 0.0057757713, + -0.0007435066, -0.0011691955, -0.0095653436, + -0.0100640547, -0.0008857724, 0.0018050477, + 0.0042295109, -0.0090684480, 0.0019845245, + ]; + let de_ref = rt::asarray((de_ref, [3, 4])); + println!("DH correlation gradient ({policy}) (REST): {:16.9}", de.t()); + println!("DH correlation gradient (ref) : {:16.9}", de_ref.t()); + println!("Maximum Error {:?}", (&de_ref - &de).abs().max_all()); + + // parts are also cached for inspection + for key in ["de_h", "de_s1", "de_jk", "de_rint"] { + let part = grad.result.get(key).unwrap().clone(); + println!("part {key}: {:16.9}", rt::asarray((part.data, part.size)).t()); + } + assert!((&de_ref - &de).abs().max_all() < tol); + + // PT2 correlation energy is carried along as the entry energy + assert!((grad.get_energy() - (-0.19179525779814563)).abs() < 2e-6, "PT2 correlation energy mismatch"); +} + +#[test] +fn test_nh3_rmp2_grad() { + test_nh3_rmp2_grad_with("cd", 1.0e-5); +} + +#[test] +fn test_nh3_rmp2_grad_eig() { + test_nh3_rmp2_grad_with("eig", 1.0e-5); +} + +/// Doubly-hybrid (`xc = "xyg3"`) analytic gradient of distorted NH3 (restricted). +/// +/// Reference: pyscf-forge `DFDH(xc="XYG3")` analytic gradient (`ref_analytic_rxyg3.json` of the +/// working notes), def2-SVP, def2-universal-jkfit for both the SCF JK fitting and the RI part, +/// grid level 3 (REST's default `grid_gen_level`, same point set as the pyscf reference), CPKS +/// tolerance 1e-8. The asserted total is the `"SCF"` (DF-B3LYP) entry plus the `"DH"` entry; +/// the residual against pyscf is at the same quadrature/RI noise level as pyscf's own +/// finite-difference check (6.8e-7). +fn test_nh3_rxyg3_grad_with(policy: &str, tol: f64) { + let keys = toml::from_str::(&INPUT_XYG3.replace("POLICY", policy)).unwrap(); + let (ctrl, geom) = ctrl_io::parse_ctl_from_json(&keys).unwrap(); + let mol = Molecule::build_native(ctrl, geom, None).unwrap(); + let mut scf_data = scf_io::SCF::build(mol, &None); + scf_without_build(&mut scf_data, &None); + + let eng = ri_pt2::xdh_calculations(&mut scf_data, &None).unwrap(); + println!("XYG3 total energy ({policy}): {eng}"); + assert!((eng - (-56.445621659465374)).abs() < 2e-6, "XYG3 total energy mismatch"); + + // regenerate the common SCF grids (xdh_calculations frees them for memory) + scf_data.grids = Some(pyrest::dft::Grids::build(&mut scf_data.mol)); + + use pyrest::grad::rhf::RIRHFGradient; + let mut grad_scf = RIRHFGradient::new(&scf_data, &None); + grad_scf.calc_rks(); + + let mut grad = RDHGradient::new(&scf_data, &None); + grad.calc(); + + let de = rt::asarray((&grad.result["de"].data, grad.result["de"].size)) + + rt::asarray((&grad_scf.result["de"].data, grad_scf.result["de"].size)); + + // pyscf-forge XYG3 total gradient (natom, 3) -> REST [3, natm] + #[rustfmt::skip] + let de_ref = vec![ + 0.0169586567, 0.0023709868, -0.0760256914, + -0.0014295009, -0.0040509894, 0.0822119912, + -0.0013968304, -0.0076827463, 0.0000330828, + -0.0141324624, 0.0093628946, -0.0062191910, + ]; + let de_ref = rt::asarray((de_ref, [3, 4])); + println!("XYG3 total gradient ({policy}) (REST): {:16.9}", de.t()); + println!("XYG3 total gradient (ref) : {:16.9}", de_ref.t()); + println!("Maximum Error {:?}", (&de_ref - &de).abs().max_all()); + + // parts are also cached for inspection + for key in ["de_h", "de_s1", "de_jk", "de_rint", "de_k_dh", "de_xc_n", "de_resp", "de_ovlp_dh"] { + let part = grad.result.get(key).unwrap().clone(); + println!("part {key}: {:16.9}", rt::asarray((part.data, part.size)).t()); + } + assert!((&de_ref - &de).abs().max_all() < tol, "XYG3 total gradient mismatch"); +} + +#[test] +fn test_nh3_rxyg3_grad() { + test_nh3_rxyg3_grad_with("cd", 1.0e-6); +} + +#[test] +fn test_nh3_rxyg3_grad_eig() { + test_nh3_rxyg3_grad_with("eig", 1.0e-6); +} diff --git a/tests/grad/udh.rs b/tests/grad/udh.rs new file mode 100644 index 0000000000000000000000000000000000000000..2aebf8535c118c4edfd1e8146692b934667f0d1c --- /dev/null +++ b/tests/grad/udh.rs @@ -0,0 +1,215 @@ +//! Doubly-hybrid analytic gradient of distorted NH3+ (unrestricted): the pure-MP2 family +//! (`xc = "mp2"`) and the XYG3 final-functional terms. +//! +//! Reference: pyscf-forge `UDFDH` analytic gradients (pyscf 2.14.0 + pyscf-forge), +//! def2-SVP, def2-universal-jkfit for both the SCF JK fitting and the RI (3-center) part, +//! CPKS tolerance 1e-8. The JSON reference was dumped per atom in Hartree/bohr (row-major +//! `[natom, 3]`); REST stores the gradient as `[3, natm]`. +//! The correlation part below is the `"DH"` entry of `eval_force` (the `"SCF"` entry covers +//! the UHF part). Both `[ctrl.j2c_decomp]` policies (`cd`, `eig`) are exercised; the eig +//! gradients agree with the cd ones to roundoff (~1e-13; the auxiliary basis is +//! well-conditioned), so the same references are asserted with the same tolerances. + +use pyrest::ctrl_io; +use pyrest::grad::traits::GradAPI; +use pyrest::grad::udh::UDHGradient; +use pyrest::molecule_io::Molecule; +use pyrest::ri_pt2; +use pyrest::scf_io::{self, scf_without_build}; +use rstsr::prelude::*; + +static INPUT_NH3: &str = r##" +[ctrl] + print_level = 0 + num_threads = 4 + xc = "mp2" + basis_path = "def2-SVP" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 1.0 + spin = 2.0 + spin_polarization = true + +[ctrl.j2c_decomp] +policy = "POLICY" + +[ctrl.ri_pt2] +new_driver = true + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0000000000 0.0000000000 0.0000000000 + H 0.0000000000 0.0000000000 1.1500000000 + H 0.9674007198 0.0000000000 -0.2902341250 + H -0.4917018117 0.8516525201 -0.3062961203 + """ +"##; + +static INPUT_XYG3_U: &str = r##" +[ctrl] + print_level = 0 + num_threads = 4 + xc = "xyg3" + basis_path = "def2-SVP" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 1.0 + spin = 2.0 + spin_polarization = true + +[ctrl.j2c_decomp] +policy = "POLICY" + +[ctrl.ri_pt2] +new_driver = true + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0000000000 0.0000000000 0.0000000000 + H 0.0000000000 0.0000000000 1.1500000000 + H 0.9674007198 0.0000000000 -0.2902341250 + H -0.4917018117 0.8516525201 -0.3062961203 + """ +"##; + +fn test_nh3_ump2_grad_with(policy: &str, tol: f64) { + let keys = toml::from_str::(&INPUT_NH3.replace("POLICY", policy)).unwrap(); + let (ctrl, geom) = ctrl_io::parse_ctl_from_json(&keys).unwrap(); + let mol = Molecule::build_native(ctrl, geom, None).unwrap(); + let mut scf_data = scf_io::SCF::build(mol, &None); + scf_without_build(&mut scf_data, &None); + + // total energy (RI-UMP2); also puts the SCF into the post-SCF state the gradient expects + let eng = ri_pt2::xdh_calculations(&mut scf_data, &None).unwrap(); + println!("RI-UMP2 total energy ({policy}): {eng}"); + assert!((eng - (-55.95341829313824)).abs() < 2e-6, "RI-UMP2 total energy mismatch"); + + let mut grad = UDHGradient::new(&scf_data, &None); + grad.calc(); + let de = grad.result.get("de").unwrap().clone(); + let de = rt::asarray((de.data, de.size)); + + // pyscf-forge `gradient_corr_only_hartree_per_bohr` (natom, 3) -> REST [3, natm] + #[rustfmt::skip] + let de_ref = vec![ + 0.0069061512, 0.0118056890, 0.0063273326, + -0.0009358131, -0.0015731210, -0.0092567167, + -0.0095325681, -0.0013691064, 0.0013982791, + 0.0035622300, -0.0088634616, 0.0015311050, + ]; + let de_ref = rt::asarray((de_ref, [3, 4])); + println!("DH correlation gradient ({policy}) (REST): {:16.9}", de.t()); + println!("DH correlation gradient (ref) : {:16.9}", de_ref.t()); + println!("Maximum Error {:?}", (&de_ref - &de).abs().max_all()); + + // parts are also cached for inspection + for key in ["de_h", "de_s1", "de_jk", "de_rint"] { + let part = grad.result.get(key).unwrap().clone(); + println!("part {key}: {:16.9}", rt::asarray((part.data, part.size)).t()); + } + assert!((&de_ref - &de).abs().max_all() < tol); + + // PT2 correlation energy is carried along as the entry energy + assert!((grad.get_energy() - (-0.1488233019937465)).abs() < 2e-6, "PT2 correlation energy mismatch"); + + // total gradient: the "SCF" (UHF) entry plus the "DH" entry, against pyscf-forge's total + use pyrest::grad::uhf::RIUHFGradient; + let mut grad_scf = RIUHFGradient::new(&scf_data, &None); + grad_scf.calc(); + let de_tot = rt::asarray((grad.result["de"].data.clone(), grad.result["de"].size)) + + rt::asarray((grad_scf.result["de"].data.clone(), grad_scf.result["de"].size)); + + // pyscf-forge `gradient_hartree_per_bohr` (natom, 3) -> REST [3, natm] + #[rustfmt::skip] + let de_tot_ref = vec![ + -0.0109435544, -0.0477856558, -0.1027269995, + 0.0114885212, 0.0186658009, 0.0664058305, + -0.0155263679, 0.0185574053, 0.0212383810, + 0.0149814011, 0.0105624497, 0.0150827880, + ]; + let de_tot_ref = rt::asarray((de_tot_ref, [3, 4])); + println!("UMP2 total gradient ({policy}) (REST): {:16.9}", de_tot.t()); + println!("UMP2 total gradient (ref) : {:16.9}", de_tot_ref.t()); + println!("Maximum Error {:?}", (&de_tot_ref - &de_tot).abs().max_all()); + assert!((&de_tot_ref - &de_tot).abs().max_all() < tol, "UMP2 total gradient mismatch"); +} + +#[test] +fn test_nh3_ump2_grad() { + test_nh3_ump2_grad_with("cd", 1.0e-5); +} + +#[test] +fn test_nh3_ump2_grad_eig() { + test_nh3_ump2_grad_with("eig", 1.0e-5); +} + +/// Doubly-hybrid (`xc = "xyg3"`) analytic gradient of distorted NH3+ (unrestricted). +/// +/// Reference: pyscf-forge `UDFDH(xc="XYG3")` analytic gradient (`ref_analytic_uxyg3.json` of the +/// working notes), def2-SVP, def2-universal-jkfit for both the SCF JK fitting and the RI part, +/// grid level 3 (REST's default `grid_gen_level`, same point set as the pyscf reference), CPKS +/// tolerance 1e-8; pyscf's own finite-difference spot check on this system is 5.1e-7. The +/// asserted total is the `"SCF"` (DF-B3LYP) entry plus the `"DH"` entry; the residual against +/// pyscf sits at the quadrature-noise floor of the unrestricted grid. +fn test_nh3_uxyg3_grad_with(policy: &str, tol: f64) { + let keys = toml::from_str::(&INPUT_XYG3_U.replace("POLICY", policy)).unwrap(); + let (ctrl, geom) = ctrl_io::parse_ctl_from_json(&keys).unwrap(); + let mol = Molecule::build_native(ctrl, geom, None).unwrap(); + let mut scf_data = scf_io::SCF::build(mol, &None); + scf_without_build(&mut scf_data, &None); + + let eng = ri_pt2::xdh_calculations(&mut scf_data, &None).unwrap(); + println!("UXYG3 total energy ({policy}): {eng}"); + assert!((eng - (-56.06868136593759)).abs() < 2e-6, "UXYG3 total energy mismatch"); + + // regenerate the common SCF grids (xdh_calculations frees them for memory) + scf_data.grids = Some(pyrest::dft::Grids::build(&mut scf_data.mol)); + + use pyrest::grad::uhf::RIUHFGradient; + let mut grad_scf = RIUHFGradient::new(&scf_data, &None); + grad_scf.calc_uks(); + + let mut grad = UDHGradient::new(&scf_data, &None); + grad.calc(); + + let de = rt::asarray((&grad.result["de"].data, grad.result["de"].size)) + + rt::asarray((&grad_scf.result["de"].data, grad_scf.result["de"].size)); + + // pyscf-forge UXYG3 total gradient (natom, 3) -> REST [3, natm] + #[rustfmt::skip] + let de_ref = vec![ + -0.0130662926, -0.0509920020, -0.1027765100, + 0.0117354582, 0.0191023970, 0.0673133968, + -0.0128975312, 0.0190512659, 0.0207786570, + 0.0142281688, 0.0128385585, 0.0146845393, + ]; + let de_ref = rt::asarray((de_ref, [3, 4])); + println!("UXYG3 total gradient ({policy}) (REST): {:16.9}", de.t()); + println!("UXYG3 total gradient (ref) : {:16.9}", de_ref.t()); + println!("Maximum Error {:?}", (&de_ref - &de).abs().max_all()); + + // parts are also cached for inspection + for key in ["de_h", "de_s1", "de_jk", "de_rint", "de_k_dh", "de_xc_n", "de_resp", "de_ovlp_dh"] { + let part = grad.result.get(key).unwrap().clone(); + println!("part {key}: {:16.9}", rt::asarray((part.data, part.size)).t()); + } + assert!((&de_ref - &de).abs().max_all() < tol, "UXYG3 total gradient mismatch"); + + // PT2 correlation energy is carried along as the entry energy + assert!((grad.get_energy() - (-0.06697732359480862)).abs() < 2e-6, "PT2 correlation energy mismatch"); +} + +#[test] +fn test_nh3_uxyg3_grad() { + test_nh3_uxyg3_grad_with("cd", 1.0e-6); +} + +#[test] +fn test_nh3_uxyg3_grad_eig() { + test_nh3_uxyg3_grad_with("eig", 1.0e-6); +}