From eeef6e186e8bfba629776e9009301c066f8f9436 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Tue, 22 Sep 2026 12:02:15 +0800 Subject: [PATCH 01/21] Refactor UHF response into URespAPI trait with shared RespSCF carrier Add the unrestricted response layer, mirroring the existing restricted one (RRespAPI/RRespSCF): - trait_uresp.rs: URespAPI trait with per-spin ([TsrView; 2] / [Tsr; 2], index 0 = alpha, 1 = beta) signatures; get_fock_rdm and get_response_rdm are reserved for the future ugfock machinery. - uresp_interface.rs: URespSCF interface constructing the unrestricted response object from REST SCF data. - response/mod.rs: RespSCF carrier enum (R/U), built once per SCF type in analdrv_interface (and the geometric-pyo3 hessian hook) and passed by &mut to the hessian and multipole drivers; previously the UHF hessian built its own response machinery internally from resp config. - Move the U-side CP-SCF response machinery out of the hessian drivers: hessian/uscf.rs and uscf_interface.rs slim down; the numint UKS part splits into numint_matmul/resp_uks.rs and the RI-JK part into ri_jk/resp_u.rs. - Update the UHF/UKS hessian tests to the new interface. No behavior change intended; analdrv test gate passes (28 passed). Co-authored-by: Claude Code Co-authored-by: glm-5.3-flash --- src/analdrv/hessian/mod.rs | 39 +- src/analdrv/hessian/rscf.rs | 19 +- src/analdrv/hessian/trait_uhess.rs | 37 +- src/analdrv/hessian/uscf.rs | 221 ++-------- src/analdrv/hessian/uscf_interface.rs | 236 +++++----- src/analdrv/interface.rs | 44 +- src/analdrv/mod.rs | 27 +- src/analdrv/multipole/interface.rs | 22 +- src/analdrv/response/mod.rs | 30 +- src/analdrv/response/trait_uresp.rs | 139 ++++++ src/analdrv/response/uresp_interface.rs | 413 ++++++++++++++++++ src/dft/numint_matmul/hess_uks.rs | 293 +------------ src/dft/numint_matmul/mod.rs | 1 + src/dft/numint_matmul/resp_uks.rs | 324 ++++++++++++++ src/main_driver.rs | 13 +- src/ri_jk/hess_u.rs | 63 +-- src/ri_jk/mod.rs | 1 + src/ri_jk/resp_u.rs | 176 ++++++++ tests/analdrv/hessian/uhf.rs | 4 +- tests/analdrv/hessian/uks_grid_shift.rs | 4 +- .../analdrv/hessian/uks_tpss0_grid_levels.rs | 6 +- tests/analdrv/hessian/uks_tpssh.rs | 4 +- 22 files changed, 1355 insertions(+), 761 deletions(-) create mode 100644 src/analdrv/response/trait_uresp.rs create mode 100644 src/analdrv/response/uresp_interface.rs create mode 100644 src/dft/numint_matmul/resp_uks.rs create mode 100644 src/ri_jk/resp_u.rs diff --git a/src/analdrv/hessian/mod.rs b/src/analdrv/hessian/mod.rs index 275ec7a00..e34b1416b 100644 --- a/src/analdrv/hessian/mod.rs +++ b/src/analdrv/hessian/mod.rs @@ -4,12 +4,12 @@ //! - trait definitions for hessian components ([`trait_rhess`], [`trait_uhess`], [`trait_util`]); //! - component hessian implementations ([`hcore`], [`nuc_repl`], [`ovlp`], and integral handling //! [`cint_handling`]); -//! - total hessian drivers for restricted and unrestricted SCF ([`rscf`], [`uscf`]), which also -//! contain the CP-SCF response machinery; +//! - total hessian drivers for restricted and unrestricted SCF ([`rscf`], [`uscf`]); //! - total hessian interface to REST ([`rscf_interface`], [`uscf_interface`]). //! -//! Note that the CP-SCF parts (the response methods in [`rscf`] and [`uscf`]) are currently -//! embedded in the total hessian drivers; they may be decoupled into a separate module in future. +//! The CP-SCF response machinery lives in the separate +//! [`response`](crate::analdrv::response) module; the hessian drivers consume its response +//! objects through the `resp` field. // trait definitions pub mod trait_rhess; @@ -34,7 +34,7 @@ pub mod rscf_interface; pub mod uscf_interface; use crate::analdrv::config::AnalDrvConfig; -use crate::analdrv::response::rresp_interface::RRespSCF; +use crate::analdrv::response::RespSCF; use crate::scf_io::{SCFType, SCF}; use crate::thermo::ThermoResult; @@ -57,7 +57,7 @@ pub struct HessOutput { pub fn hess_interface<'a>( scf_data: &'a SCF, config: &AnalDrvConfig, - resp_objs: Option<&mut RRespSCF<'a>>, + resp_objs: Option<&mut RespSCF<'a>>, ) -> HessOutput { use crate::analdrv::hessian::rscf_interface::rscf_hess_interface; use crate::analdrv::hessian::uscf_interface::uscf_hess_interface; @@ -70,16 +70,25 @@ pub fn hess_interface<'a>( panic!("Normal modes calculation is currently not available for post-SCF methods."); } - // The RHF hessian consumes the shared response object built by the caller; the UHF hessian - // builds and owns its own (U-side) response machinery internally, so it takes the response - // settings explicitly. + // The hessian consumes the shared response object built by the caller (the `RespSCF` + // carrier, holding the restricted/unrestricted variant per the SCF type). let (hessian, vib, _) = match scf_data.scftype { - SCFType::RHF => rscf_hess_interface( - scf_data, - &config.nucgrad, - resp_objs.expect("internal error: the RHF hessian requires the shared response object"), - ), - SCFType::UHF => uscf_hess_interface(scf_data, &config.nucgrad, &config.resp), + SCFType::RHF => { + let RespSCF::R(resp) = + resp_objs.expect("internal error: the RHF hessian requires the shared response object") + else { + panic!("internal error: the RHF hessian requires the restricted (R) response object variant"); + }; + rscf_hess_interface(scf_data, &config.nucgrad, resp) + }, + SCFType::UHF => { + let RespSCF::U(resp) = + resp_objs.expect("internal error: the UHF hessian requires the shared response object") + else { + panic!("internal error: the UHF hessian requires the unrestricted (U) response object variant"); + }; + uscf_hess_interface(scf_data, &config.nucgrad, resp) + }, _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types."), }; diff --git a/src/analdrv/hessian/rscf.rs b/src/analdrv/hessian/rscf.rs index a97bd7581..b83544610 100644 --- a/src/analdrv/hessian/rscf.rs +++ b/src/analdrv/hessian/rscf.rs @@ -5,14 +5,17 @@ use crate::analdrv::prelude::*; /// Working solver and maintainer of all hessian components for restricted SCF method. /// /// The lifetime parameters decouple the response object's borrow from its data, and the response -/// object's data region from the hessian component objects: `'a` is the region of the SCF data -/// the response object borrows, `'b` the (typically shorter) borrow of `resp` itself, and `'c` -/// the region of the hessian component objects (`ovlp_obj` and the lists). Bundling `'a` with -/// `'b` (`&'a mut RRespSCF<'a>`) is rejected by dropck: the borrow would have to end before the -/// referent's destructor, while that destructor pins the pointee's lifetime until the referent's -/// drop. Keeping `'a` separate from `'c` allows the response object to be shared from outside -/// (built once by the caller, e.g. the analdrv task loop) while the component objects remain -/// function-local. +/// object's data region from the hessian component objects: +/// +/// - `'a`: the region of the SCF data the response object borrows; +/// - `'b`: the (typically shorter) borrow of `resp` itself; +/// - `'c`: the region of the hessian component objects (`ovlp_obj` and the lists). +/// +/// Bundling `'a` with `'b` (`&'a mut RRespSCF<'a>`) is rejected by dropck: the borrow would have +/// to end before the referent's destructor, while that destructor pins the pointee's lifetime +/// until the referent's drop. Keeping `'a` separate from `'c` allows the response object to be +/// shared from outside (built once by the caller, e.g. the analdrv task loop) while the component +/// objects remain function-local. pub struct RHessSCF<'a, 'b, 'c> { pub mo_coeff: Tsr, pub mo_occ: Tsr, diff --git a/src/analdrv/hessian/trait_uhess.rs b/src/analdrv/hessian/trait_uhess.rs index b1a3fadaa..3be45f296 100644 --- a/src/analdrv/hessian/trait_uhess.rs +++ b/src/analdrv/hessian/trait_uhess.rs @@ -54,6 +54,11 @@ pub trait UHessCoreAPI: AnalDrvBaseAPI { /// /// Difference to [`RHessElecInteractAPI`] is that we may need different signature. Basic ideas are /// exactly the same. +/// +/// This trait contains only the hessian-specific skeleton contractions. Response-related +/// functionalities (fock generation, response preparation and contraction) live in the separate +/// [`URespAPI`](crate::analdrv::response::trait_uresp::URespAPI) and its own response objects; +/// drivers that need both maintain separate lists of the two kinds of objects. pub trait UHessElecInteractAPI: AnalDrvBaseAPI { /// Generate the **skeleton** contribution of Hessian for current SCF component. /// @@ -137,36 +142,4 @@ pub trait UHessElecInteractAPI: AnalDrvBaseAPI { let deriv1_ao = self.get_deriv1_ao(mo_coeff, mo_occ, atm_list); [&deriv1_ao[α] % &mocc[α], &deriv1_ao[β] % &mocc[β]] } - - /// Prepare the data for response calculation. - /// - /// # Parameters - /// - /// - `mo_coeff` : shape `[nao, nmo_α]` and `[nao, nmo_β]`. Molecular orbital coefficients. - /// - `mo_occ` : shape `[nmo_α]` and `[nmo_β]`. Molecular orbital occupation numbers. - /// - /// # See also - /// - /// [`RHessElecInteractAPI::make_response_preparation`]. Signature difference: `mo_coeff` and - /// `mo_occ` type different. - fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]); - - /// Get the response contribution for current SCF component. - /// - /// - /// # Parameters - /// - /// - `bra` : shape `[nao, nocc_α, ...]` and `[nao, nocc_β, ...]`. The bra part for response - /// calculation. - /// - /// # Returns - /// - /// - `resp_bra` : shape `[nao, nocc_α, ...]` and `[nao, nocc_β, ...]`. The response potential - /// (related to second order of density matrix derivative to energy). - /// - /// # See also - /// - /// [`RHessElecInteractAPI::get_response_bra`]. Signature difference: input and output type - /// different. - fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2]; } diff --git a/src/analdrv/hessian/uscf.rs b/src/analdrv/hessian/uscf.rs index 291b25269..26844234a 100644 --- a/src/analdrv/hessian/uscf.rs +++ b/src/analdrv/hessian/uscf.rs @@ -1,38 +1,45 @@ //! Hessian implementations for unrestricted SCF. use crate::analdrv::prelude::*; + /// Working solver and maintainer of all hessian components for unrestricted SCF method. -pub struct UHessSCF<'a> { +/// +/// The lifetime parameters mean the same as in +/// [`RHessSCF`](crate::analdrv::hessian::rscf::RHessSCF) (documented there): they decouple the +/// response object's borrow from its data, and the response object's data region from the hessian +/// component objects. +pub struct UHessSCF<'a, 'b, 'c> { pub mo_coeff: [Tsr; 2], pub mo_occ: [Tsr; 2], pub mo_energy: [Tsr; 2], - pub ovlp_obj: &'a mut UHessOvlp, - pub nuc_list: Vec<&'a mut dyn HessNucAPI>, - pub core_list: Vec<&'a mut dyn UHessCoreAPI>, - pub el_list: Vec<&'a mut dyn UHessElecInteractAPI>, - /// The `AnalDrvNucgradCfg` sub-config (`atm_list` etc.). + pub ovlp_obj: &'c mut UHessOvlp, + pub nuc_list: Vec<&'c mut dyn HessNucAPI>, + pub core_list: Vec<&'c mut dyn UHessCoreAPI>, + pub el_list: Vec<&'c mut dyn UHessElecInteractAPI>, + /// Response (fock/response) object for the electron-interaction contributions (the composite + /// `URespSCF` of the SCF); the CP-SCF machinery (`make_response_preparation`/`response_mo`) + /// calls it, separate from the skeleton-hessian objects in `el_list`. + pub resp: &'b mut URespSCF<'a>, + /// The `AnalDrvNucgradCfg` sub-config (`atm_list` etc.); the CP-SCF solver settings + /// (`level_shift` etc.) come from the captured `resp_cfg` on `self.resp`. pub cfg: AnalDrvNucgradCfg, - /// CP-SCF solver and response-grid settings; the U-side response machinery is built and - /// owned internally, so these are passed in explicitly instead of captured on a response - /// object (unlike the restricted side, where `RRespSCF.resp_cfg` is the source). - pub resp_cfg: AnalDrvRespCfg, pub result: HashMap, /// Timing information. Represented by wall time in second. pub timing: Vec<(String, f64)>, } -impl<'a> UHessSCF<'a> { +impl<'a, 'b, 'c> UHessSCF<'a, 'b, 'c> { #[allow(clippy::too_many_arguments)] pub fn new( mo_coeff: [Tsr; 2], mo_occ: [Tsr; 2], mo_energy: [Tsr; 2], - ovlp_obj: &'a mut UHessOvlp, - nuc_list: Vec<&'a mut dyn HessNucAPI>, - core_list: Vec<&'a mut dyn UHessCoreAPI>, - el_list: Vec<&'a mut dyn UHessElecInteractAPI>, + ovlp_obj: &'c mut UHessOvlp, + nuc_list: Vec<&'c mut dyn HessNucAPI>, + core_list: Vec<&'c mut dyn UHessCoreAPI>, + el_list: Vec<&'c mut dyn UHessElecInteractAPI>, + resp: &'b mut URespSCF<'a>, cfg: &AnalDrvNucgradCfg, - resp_cfg: &AnalDrvRespCfg, ) -> Self { Self { mo_coeff, @@ -42,8 +49,8 @@ impl<'a> UHessSCF<'a> { nuc_list, core_list, el_list, + resp, cfg: cfg.clone(), - resp_cfg: resp_cfg.clone(), result: HashMap::new(), timing: Vec::new(), } @@ -71,7 +78,7 @@ impl<'a> UHessSCF<'a> { /// steps. /// /// # Returns - /// + /// /// A dictionary containing: /// - `rhs : shape `[nmo, nocc_α, 3, natm]` and `[nmo, nocc_β, 3, natm]`. The dimensionless CP-SCF /// right-hand side. @@ -86,7 +93,7 @@ impl<'a> UHessSCF<'a> { let mo_coeff = [self.mo_coeff[α].view(), self.mo_coeff[β].view()]; let mo_occ = [self.mo_occ[α].view(), self.mo_occ[β].view()]; let mo_energy = [self.mo_energy[α].view(), self.mo_energy[β].view()]; - let level_shift = self.resp_cfg.level_shift; + let level_shift = self.resp.resp_cfg.level_shift; let device = mo_coeff[α].device().clone(); let nao = mo_coeff[α].shape()[0]; @@ -178,181 +185,17 @@ impl<'a> UHessSCF<'a> { /// Prepare the response for CP-SCF calculation. /// - /// This involves all electron-interaction objects. + /// This involves all response objects, and stores the orbital state in the response object + /// for its inherent CP-SCF machinery. pub fn make_response_preparation(&mut self) { let t0 = std::time::Instant::now(); let mo_coeff = [self.mo_coeff[0].view(), self.mo_coeff[1].view()]; let mo_occ = [self.mo_occ[0].view(), self.mo_occ[1].view()]; - for el_obj in self.el_list.iter_mut() { - let t1 = std::time::Instant::now(); - el_obj.make_response_preparation(&mo_coeff, &mo_occ); - self.timing.push(( - format!("in make_response_preparation, {}", el_obj.get_type_name()), - t1.elapsed().as_secs_f64(), - )); - } + let mo_energy = [self.mo_energy[0].view(), self.mo_energy[1].view()]; + self.resp.make_cpscf_preparation(&mo_coeff, &mo_occ, &mo_energy); self.timing.push(("make_response_preparation".to_string(), t0.elapsed().as_secs_f64())); } - /// Compute the response of the system to a given perturbation in MO space (mo1), which is - /// needed for CP-SCF. - /// - /// # Parameters - /// - /// - `mo1` : shape `[nmo, nocc_α, ...]` and `[nmo, nocc_β, ...]`. The perturbation in MO space. - /// - /// # Returns - /// - /// - `resp` : shape `[nmo, nocc_α, ...]` and `[nmo, nocc_β, ...]`. The response in MO space. - pub fn response_mo(&mut self, mo1: &[TsrView; 2]) -> [Tsr; 2] { - let [α, β] = [0, 1]; - let ubra_α = &self.mo_coeff[α] % &mo1[α]; - let ubra_β = &self.mo_coeff[β] % &mo1[β]; - let mut resp_α = rt::zeros_like(&ubra_α); - let mut resp_β = rt::zeros_like(&ubra_β); - - for el_obj in self.el_list.iter_mut() { - let t1 = std::time::Instant::now(); - let el_resp = el_obj.get_response_bra(&[ubra_α.view(), ubra_β.view()]); - resp_α += self.mo_coeff[α].t() % &el_resp[α]; - resp_β += self.mo_coeff[β].t() % &el_resp[β]; - self.timing.push((format!("in response_mo, {}", el_obj.get_type_name()), t1.elapsed().as_secs_f64())); - } - [resp_α, resp_β] - } - - /// Compute the dimensionless response for CP-SCF calculation. - /// - /// # Parameters - /// - /// - `mo1` : shape `[nmo, nocc_α, ...]` and `[nmo, nocc_β, ...]`. The perturbation in MO space. - /// - /// # Returns - /// - /// - `resp` : shape `[nmo, nocc_α, ...]` and `[nmo, nocc_β, ...]`. The dimensionless response - /// in MO space. - pub fn response_dimless_cpscf(&mut self, mo1: &[TsrView; 2]) -> [Tsr; 2] { - let t0 = std::time::Instant::now(); - let [α, β] = [0, 1]; - let mo_occ = [self.mo_occ[α].view(), self.mo_occ[β].view()]; - let occidx = [mo_occ[α].view().greater(0).into_vec(), mo_occ[β].view().greater(0).into_vec()]; - let viridx = [occidx[α].iter().map(|&x| !x).collect_vec(), occidx[β].iter().map(|&x| !x).collect_vec()]; - let nocc = [occidx[α].iter().filter(|&&x| x).count(), occidx[β].iter().filter(|&&x| x).count()]; - let nmo = [mo_occ[α].shape()[0], mo_occ[β].shape()[0]]; - let eocc = [ - self.mo_energy[α].view().bool_select(-1, &occidx[α]), - self.mo_energy[β].view().bool_select(-1, &occidx[β]), - ]; - let evir = [ - self.mo_energy[α].view().bool_select(-1, &viridx[α]), - self.mo_energy[β].view().bool_select(-1, &viridx[β]), - ]; - let e_ai = [evir[α].i((.., None)) - eocc[α].i((None, ..)), evir[β].i((.., None)) - eocc[β].i((None, ..))]; - let level_shift = self.resp_cfg.level_shift; - let e_ai_shift = [&e_ai[0] + level_shift, &e_ai[1] + level_shift]; - let so = [rt::slice!(0, nocc[α]), rt::slice!(0, nocc[β])]; - let sv = [rt::slice!(nocc[α], nmo[α]), rt::slice!(nocc[β], nmo[β])]; - - let mut resp = self.response_mo(mo1); - - // handle dimension less denominator and occupied response part - if level_shift != 0.0 { - *&mut resp[α] -= level_shift * &mo1[α]; - *&mut resp[β] -= level_shift * &mo1[β]; - } - *&mut resp[α].i_mut(sv[α]) /= &e_ai_shift[α]; - *&mut resp[β].i_mut(sv[β]) /= &e_ai_shift[β]; - resp[α].i_mut(so[α]).fill(0.0); - resp[β].i_mut(so[β]).fill(0.0); - self.timing.push(("response_dimless_cpscf".to_string(), t0.elapsed().as_secs_f64())); - resp - } - - /// Solve the dimensionless CP-SCF equation using a Krylov solver. - /// - /// # Parameters - /// - /// - `rhs` : shape `[nmo, nocc_α, ...]` and `[nmo, nocc_β, ...]`. Dimensionless right-hand - /// side. - /// - /// # Returns - /// - /// - `mo1` : shape `[nmo, nocc_α, ...]` and `[nmo, nocc_β, ...]`. Perturbation in MO space that - /// solves the dimensionless CP-SCF equation. - pub fn solve_dimless_cpscf(&mut self, rhs: &[TsrView; 2]) -> [Tsr; 2] { - let t0 = std::time::Instant::now(); - let [α, β] = [0, 1]; - let rhs_shape = [rhs[α].shape().to_vec(), rhs[β].shape().to_vec()]; - let nmo = [rhs[α].shape()[0], rhs[β].shape()[0]]; - let nocc = [rhs[α].shape()[1], rhs[β].shape()[1]]; - let rhs = [rhs[α].reshape((nmo[α], nocc[α], -1)), rhs[β].reshape((nmo[β], nocc[β], -1))]; - let device = rhs[α].device().clone(); - - let tol = self.resp_cfg.tol; - let max_cycle = self.resp_cfg.max_cycle; - let max_space = self.resp_cfg.max_space; - let lindep = self.resp_cfg.lindep; - let tol_inflation = self.resp_cfg.tol_inflation; - - let pack_flattened = |x: &[TsrView; 2]| -> Tsr { - // original: [nmo_α, nocc_α, nprop] and [nmo_β, nocc_β, nprop] - // target: [nmo_α * nocc_α + nmo_β * nocc_β, nprop] - assert_eq!(x[α].ndim(), 3, "Expected x[α] to have shape [nmo_α, nocc_α, nprop]"); - assert_eq!(x[β].ndim(), 3, "Expected x[β] to have shape [nmo_β, nocc_β, nprop]"); - let nprop = x[α].shape()[2]; - let mut x_flattened = rt::zeros(([nmo[α] * nocc[α] + nmo[β] * nocc[β], nprop], &device)); - for A in 0..nprop { - x_flattened.i_mut((..nmo[α] * nocc[α], A)).assign(x[α].i((.., .., A)).reshape(-1)); - x_flattened.i_mut((nmo[α] * nocc[α].., A)).assign(x[β].i((.., .., A)).reshape(-1)); - } - x_flattened - }; - - let unpack_flattened = |x: TsrView| -> [Tsr; 2] { - // original: [nmo_α * nocc_α + nmo_β * nocc_β, nprop] - // target: [nmo_α, nocc_α, nprop] and [nmo_β, nocc_β, nprop] - assert_eq!(x.ndim(), 2, "Expected x to have shape [nmo_α * nocc_α + nmo_β * nocc_β, nprop]"); - let nprop = x.shape()[1]; - let idx_split = nmo[α] * nocc[α]; - let mut x_α = rt::zeros(([nmo[α], nocc[α], nprop], &device)); - let mut x_β = rt::zeros(([nmo[β], nocc[β], nprop], &device)); - for A in 0..nprop { - x_α.i_mut((.., .., A)).assign(x.i((..idx_split, A)).reshape((nmo[α], nocc[α]))); - x_β.i_mut((.., .., A)).assign(x.i((idx_split.., A)).reshape((nmo[β], nocc[β]))); - } - [x_α, x_β] - }; - - let response_cpscf_flattened = |x: TsrView| -> Tsr { - // split x by spin and reshape to original shape - let [x_α, x_β] = unpack_flattened(x); - // compute response by usual means - let resp = self.response_dimless_cpscf(&[x_α.view(), x_β.view()]); - // flatten resp to shape (nmo*nocc, nprop) - let resp_view = resp.iter().map(|r| r.view()).collect_array().unwrap(); - pack_flattened(&resp_view) - }; - - let rhs_view = rhs.iter().map(|r| r.view()).collect_array().unwrap(); - let rhs_packed = pack_flattened(&rhs_view); - let mo1_flattened = krylov_block( - response_cpscf_flattened, - rhs_packed.view(), - None, - tol, - max_cycle, - max_space, - lindep, - tol_inflation, - ); - let [mo1_α, mo1_β] = unpack_flattened(mo1_flattened.view()); - let mo1_α = mo1_α.into_shape(rhs_shape[α].to_vec()); - let mo1_β = mo1_β.into_shape(rhs_shape[β].to_vec()); - - self.timing.push(("solve_dimless_cpscf".to_string(), t0.elapsed().as_secs_f64())); - [mo1_α, mo1_β] - } - /// Finalize the CP-SCF calculation by computing necessary intermediates for Hessian assembly. /// /// @@ -402,7 +245,7 @@ impl<'a> UHessSCF<'a> { let e_ij = [eocc[α].i((.., None)) - eocc[α].i((None, ..)), eocc[β].i((.., None)) - eocc[β].i((None, ..))]; // last-iter the cp-hf equation, and remove the level-shift - let last_resp = self.response_mo(mo1); + let last_resp = self.resp.response_mo(mo1); let b1mo_α = &f1mo[α] - &s1mo[α] * eocc[α].i((None, ..)) + &last_resp[α]; let b1mo_β = &f1mo[β] - &s1mo[β] * eocc[β].i((None, ..)) + &last_resp[β]; let mut mo1_α = mo1[α].to_owned(); @@ -469,7 +312,7 @@ impl<'a> UHessSCF<'a> { let rhs = [pre_cpscf_dict.get("rhs_0").unwrap().view(), pre_cpscf_dict.get("rhs_1").unwrap().view()]; self.make_response_preparation(); - let mo1 = self.solve_dimless_cpscf(&rhs); + let mo1 = self.resp.solve_dimless_cpscf(&rhs); let mo1_view = [mo1[0].view(), mo1[1].view()]; let finalize_dict = self.finalize_cpscf(&f1mo, &s1mo, &mo1_view); let mo1 = [finalize_dict.get("mo1_0").unwrap().view(), finalize_dict.get("mo1_1").unwrap().view()]; diff --git a/src/analdrv/hessian/uscf_interface.rs b/src/analdrv/hessian/uscf_interface.rs index 0a8167b75..eff0e3482 100644 --- a/src/analdrv/hessian/uscf_interface.rs +++ b/src/analdrv/hessian/uscf_interface.rs @@ -1,12 +1,29 @@ use crate::analdrv::prelude::*; +use crate::analdrv::response::rresp_interface::scf_jk_factors; +use crate::analdrv::response::uresp_interface::{scf_xc_func_list_uks, URespSCF}; use crate::analdrv::vibration::vib::*; use crate::analdrv::vibration::vib_interface::*; use crate::dft::gen_grids::RadiiAdjust; +use crate::dft::numint_matmul::nimatmul::{regroup_grids_by_atom, NIMatmul}; +use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; +use crate::dft::Grids; use crate::dftd::hess::HessDFTD; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; -pub fn uscf_hess_interface(scf_data: &SCF, cfg: &AnalDrvNucgradCfg, resp_cfg: &AnalDrvRespCfg) -> (Vec, VibInfo, Option) { +/// Unrestricted SCF hessian: assembles the driver (overlap/nuclear/core/electronic parts plus +/// the response object), evaluates the hessian and performs the vibrational analysis. +/// +/// `resp_obj` is the SCF response object (composite [`URespSCF`], RI-JK plus XC parts) that +/// provides the CP-SCF solve for the relaxed density. It is not built here: initialize it +/// beforehand with [`crate::analdrv::response::uresp_interface::uscf_resp_interface`] on the +/// same `scf_data`. The driver mutates it (response preparation and cached intermediates), and +/// a shared object can be reused by further tasks afterwards. +pub fn uscf_hess_interface<'a>( + scf_data: &'a SCF, + cfg: &AnalDrvNucgradCfg, + resp_obj: &mut URespSCF<'a>, +) -> (Vec, VibInfo, Option) { let device = DeviceBLAS::default(); // --- basic preparation --- // @@ -58,16 +75,7 @@ pub fn uscf_hess_interface(scf_data: &SCF, cfg: &AnalDrvNucgradCfg, resp_cfg: &A use crate::ri_jk::hess_u::UHessRIJK; - let is_hf = scf_data.mol.xc_data.dfa_compnt_scf.is_empty(); - let scale_j = 1.0; - // For range-separated hybrids the exchange is decomposed as (same as the SCF Fock assembly) - // K = alpha * K_full - (alpha - hyb) * K_erfc, - // so the full-range object carries the long-range coefficient `alpha`, and the short-range - // correction is a separate object evaluated below. - let scale_k = match is_hf { - true => 1.0, - false => scf_data.mol.xc_data.rsh_alpha().unwrap_or(scf_data.mol.xc_data.dfa_hybrid_scf), - }; + let (scale_j, scale_k, rsh) = scf_jk_factors(scf_data); let j2c_decomp_option = &scf_data.mol.ctrl.j2c_decomp; let j2c_decomp = crate::ri_jk::get_j2c_decomp(&aux, &device, *j2c_decomp_option); @@ -89,9 +97,7 @@ pub fn uscf_hess_interface(scf_data: &SCF, cfg: &AnalDrvNucgradCfg, resp_cfg: &A // carries mol/aux clones holding the (negative) omega, so all its derivative integrals // are short-range (erfc kernel), and the cderi is the incore `rimatr_sr` built in the SCF. // There is no short-range Coulomb contribution (factor_j = 0). - let mut hess_rijk_sr_obj = scf_data.mol.xc_data.omega().map(|omega| { - let alpha = scf_data.mol.xc_data.rsh_alpha().unwrap(); - let hyb = scf_data.mol.xc_data.dfa_hybrid_scf; + let mut hess_rijk_sr_obj = rsh.map(|(omega, factor_k_sr)| { let (mut mol_sr, mut aux_sr) = (mol.clone(), aux.clone()); // negative omega in libcint's convention evaluates the short-range erfc kernel mol_sr.set_omega(-omega); @@ -99,7 +105,7 @@ pub fn uscf_hess_interface(scf_data: &SCF, cfg: &AnalDrvNucgradCfg, resp_cfg: &A if let Some((rimatr_sr, _, _)) = &scf_data.rimatr_sr { let cderi_sr = rimatr_sr.to_rstsr_view(&device).into_cow(); let j2c_decomp_sr = crate::ri_jk::get_j2c_decomp(&aux_sr, &device, *j2c_decomp_option); - UHessRIJK::new_with_cderi(&mol_sr, &aux_sr, 0.0, hyb - alpha, cderi_sr, j2c_decomp_sr) + UHessRIJK::new_with_cderi(&mol_sr, &aux_sr, 0.0, factor_k_sr, cderi_sr, j2c_decomp_sr) } else { panic!( "The range-separated Hessian requires the short-range ERI (rimatr_sr) to be built and stored in memory." @@ -112,81 +118,15 @@ pub fn uscf_hess_interface(scf_data: &SCF, cfg: &AnalDrvNucgradCfg, resp_cfg: &A // --- DFT --- // - let mut hess_nimatmul_obj = (!is_hf).then(|| { + // The skeleton-level grid of the XC hessian object is built by `scf_skeleton_nimatmul_uks` + // (SCF grid reused or regenerated per the skeleton level policy); no grid data is shared + // with the response objects. + let mut hess_nimatmul_obj = scf_skeleton_nimatmul_uks(scf_data, cfg).map(|ni| { use crate::dft::numint_matmul::hess_uks::UHessKSNIMatmul; - use crate::dft::numint_matmul::nimatmul::NIMatmul; - use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; - use crate::dft::Grids; - use libxc::prelude::*; - - let xc_func_list = { - let xc_code = &scf_data.mol.xc_data.dfa_compnt_scf; - let xc_params = &scf_data.mol.xc_data.dfa_paramr_scf; - xc_code - .iter() - .zip(xc_params.iter()) - .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, LibXCSpin::Polarized))) - .collect_vec() - }; - let verbose = scf_data.mol.ctrl.print_level > 2; - // Determine skeleton / cpscf grid levels. - // - skeleton: the SCF DFT grid; only MGGA (TAU) without the grid-shift adds 2 levels - // (the grid-shift terms restore the grid-related accuracy the finer grid compensated). - // - cpscf: grid_gen_level.max(3) - 2 (coarser, for the iterative CP-SCF response). - let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); - let is_mgga = matches!(xc_type, XCDenType::TAU); - let grid_gen_level = scf_data.mol.ctrl.grid_gen_level; + let verbose = scf_data.mol.ctrl.print_level > 2; let grid_shift = cfg.grid_shift_deriv; - let sk_level = cfg.grid_level_skeleton.unwrap_or(if is_mgga && !grid_shift { - grid_gen_level + 2 - } else { - grid_gen_level - }); - let cpscf_level = resp_cfg.grid_level.unwrap_or(grid_gen_level.max(3) - 2); - - // skeleton grid: reuse the SCF grid when the level matches, else regenerate. Either - // way, regroup to atom-grouped order (non-decreasing atm_idx): the SCF grid is - // round-robin permuted for load balancing, while the Becke grid-shift attribution - // requires the ByAtom grouping. The regrouping only permutes, never changes values. - let ni = { - use crate::dft::numint_matmul::nimatmul::regroup_grids_by_atom; - - let (coordinates, weights, atm_idx, quadrature_weights) = if sk_level == grid_gen_level { - let grids = scf_data.grids.as_ref().unwrap(); - ( - grids.coordinates.clone(), - grids.weights.clone(), - grids.atm_idx.clone(), - grids.quadrature_weights.clone(), - ) - } else { - let sk_grid = Grids::build_with_level(mol_obj, sk_level); - (sk_grid.coordinates, sk_grid.weights, sk_grid.atm_idx, sk_grid.quadrature_weights) - }; - let (coordinates, weights, atm_idx, quadrature_weights) = - regroup_grids_by_atom(coordinates, weights, atm_idx, quadrature_weights, mol.natm()); - NIMatmul::new(&mol, &coordinates, &weights, &atm_idx, &quadrature_weights) - .with_radii_adjust(RadiiAdjust::from_str(&mol_obj.ctrl.radii_adjust)) - }; - - // cpscf grid: when it coincides with the skeleton grid, leave `ni_cpks = None` so the - // skeleton's vxc/fxc are reused; otherwise build a dedicated (coarser) grid. - let hess_nimatmul_obj = if cpscf_level == sk_level { - UHessKSNIMatmul::new(&mol, xc_func_list, ni, grid_shift, verbose) - } else { - let cpscf_grid = Grids::build_with_level(mol_obj, cpscf_level); - let ni_cpks = NIMatmul::new( - &mol, - &cpscf_grid.coordinates, - &cpscf_grid.weights, - &cpscf_grid.atm_idx, - &cpscf_grid.quadrature_weights, - ) - .with_radii_adjust(RadiiAdjust::from_str(&mol_obj.ctrl.radii_adjust)); - UHessKSNIMatmul::new(&mol, xc_func_list, ni, grid_shift, verbose).set_ni_cpks(ni_cpks) - }; - hess_nimatmul_obj + UHessKSNIMatmul::new(&mol, scf_xc_func_list_uks(scf_data), ni, grid_shift, verbose) }); if let Some(ref mut hess_nimatmul_obj) = hess_nimatmul_obj { hess_el_list.push(hess_nimatmul_obj); @@ -194,50 +134,96 @@ pub fn uscf_hess_interface(scf_data: &SCF, cfg: &AnalDrvNucgradCfg, resp_cfg: &A // --- run hessian --- // - let mut hess_scf = UHessSCF::new( - mo_coeff, - mo_occ, - mo_energy, - &mut hess_ovlp_obj, - hess_nuc_list, - hess_hcore_list, - hess_el_list, - cfg, - resp_cfg, - ); - - let de_hess = hess_scf.make_hess(); - - if scf_data.mol.ctrl.print_level >= 2 { - println!("=== HESSIAN ==="); - println!("Print hessian in [tA, sB] format (component xyz first, atom then)"); - println!(""); - // print hessian matrix [t, s, A, B] -> [tA, sB] - let natm = de_hess.shape()[3]; - let hess_mat = de_hess.transpose([0, 2, 1, 3]).into_shape((3 * natm, 3 * natm)); - // print 6 columns at a time, with index header - for j in (0..3 * natm).step_by(6) { - let j_end = (j + 6).min(3 * natm); - let col_header = " ".repeat(6) + &(j..j_end).map(|i| format!("{:>12}", i)).collect::(); - println!("{}", col_header); - for i in 0..3 * natm { - let row_str = - format!("{i:>4} ") + &(j..j_end).map(|j| format!("{:12.6}", hess_mat[[i, j]])).collect::(); - println!("{}", row_str); - } + // The driver and its borrowed lists live in an inner scope, so that the borrows of the + // hessian/response objects provably end before the objects themselves are dropped. + let de_hess = { + let mut hess_scf = UHessSCF::new( + mo_coeff, + mo_occ, + mo_energy, + &mut hess_ovlp_obj, + hess_nuc_list, + hess_hcore_list, + hess_el_list, + resp_obj, + cfg, + ); + + let de_hess = hess_scf.make_hess(); + + if scf_data.mol.ctrl.print_level >= 2 { + println!("=== HESSIAN ==="); + println!("Print hessian in [tA, sB] format (component xyz first, atom then)"); println!(""); + // print hessian matrix [t, s, A, B] -> [tA, sB] + let natm = de_hess.shape()[3]; + let hess_mat = de_hess.transpose([0, 2, 1, 3]).into_shape((3 * natm, 3 * natm)); + // print 6 columns at a time, with index header + for j in (0..3 * natm).step_by(6) { + let j_end = (j + 6).min(3 * natm); + let col_header = " ".repeat(6) + &(j..j_end).map(|i| format!("{:>12}", i)).collect::(); + println!("{}", col_header); + for i in 0..3 * natm { + let row_str = format!("{i:>4} ") + + &(j..j_end).map(|j| format!("{:12.6}", hess_mat[[i, j]])).collect::(); + println!("{}", row_str); + } + println!(""); + } } - } - // print timing information - if scf_data.mol.ctrl.print_level >= 2 { - println!("Timing in Hessian calculation:"); - for (key, value) in hess_scf.timing.iter() { - println!(" {:60}: {:10.6} seconds", key, value); + // print timing information + if scf_data.mol.ctrl.print_level >= 2 { + println!("Timing in Hessian calculation:"); + for (key, value) in hess_scf.timing.iter() { + println!(" {:60}: {:10.6} seconds", key, value); + } } - } + + de_hess + }; // --- perform vibrational analysis --- // vibration_analysis_interface(scf_data, cfg, de_hess.view()) } + +/// The skeleton-level `NIMatmul` for the DFT XC hessian contribution, `None` for pure HF +/// (unrestricted sibling of `scf_skeleton_nimatmul` in +/// [`rscf_interface`](crate::analdrv::hessian::rscf_interface)). +/// +/// The skeleton grid level is the SCF grid level, raised by 2 for MGGA (TAU) functionals when +/// the grid-shift derivative terms are off (the grid-shift terms restore the grid-related +/// accuracy the finer grid compensated). The grid reuses the SCF grid when the level matches, +/// else is regenerated; either way it is regrouped to atom-grouped order (non-decreasing +/// `atm_idx`): the SCF grid is round-robin permuted for load balancing, while the Becke +/// grid-shift attribution requires the ByAtom grouping. The regrouping only permutes, never +/// changes values. +fn scf_skeleton_nimatmul_uks<'a>(scf_data: &'a SCF, cfg: &AnalDrvNucgradCfg) -> Option> { + if scf_data.mol.xc_data.dfa_compnt_scf.is_empty() { + return None; + } + let mol = get_cint_mol(&scf_data.mol); + + let xc_func_list = scf_xc_func_list_uks(scf_data); + let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); + let is_mgga = matches!(xc_type, XCDenType::TAU); + let grid_gen_level = scf_data.mol.ctrl.grid_gen_level; + let grid_shift = cfg.grid_shift_deriv; + let sk_level = + cfg.grid_level_skeleton.unwrap_or(if is_mgga && !grid_shift { grid_gen_level + 2 } else { grid_gen_level }); + + let (coordinates, weights, atm_idx, quadrature_weights) = if sk_level == grid_gen_level { + let grids = scf_data.grids.as_ref().unwrap(); + (grids.coordinates.clone(), grids.weights.clone(), grids.atm_idx.clone(), grids.quadrature_weights.clone()) + } else { + let sk_grid = Grids::build_with_level(&scf_data.mol, sk_level); + (sk_grid.coordinates, sk_grid.weights, sk_grid.atm_idx, sk_grid.quadrature_weights) + }; + let (coordinates, weights, atm_idx, quadrature_weights) = + regroup_grids_by_atom(coordinates, weights, atm_idx, quadrature_weights, mol.natm()); + // the Becke grid-shift terms of the hessian rebuild the partition from the atomic radii, so + // the driver must know the radii adjustment scheme the grid was generated with + let radii_adjust = RadiiAdjust::from_str(&scf_data.mol.ctrl.radii_adjust); + Some(NIMatmul::new(&mol, &coordinates, &weights, &atm_idx, &quadrature_weights).with_radii_adjust(radii_adjust)) +} diff --git a/src/analdrv/interface.rs b/src/analdrv/interface.rs index 5e24cdd8a..2cd8900a8 100644 --- a/src/analdrv/interface.rs +++ b/src/analdrv/interface.rs @@ -6,6 +6,8 @@ use crate::analdrv::multipole::interface::{multipole_interface, MultipoleOutput} use crate::analdrv::response::rgfock_interface; use crate::analdrv::response::rresp_interface; use crate::analdrv::response::rresp_interface::rscf_resp_interface; +use crate::analdrv::response::uresp_interface::uscf_resp_interface; +use crate::analdrv::response::RespSCF; use crate::scf_io::{SCFType, SCF}; use crate::thermo::ThermoResult; @@ -54,20 +56,19 @@ pub struct AnaldrvOutput { /// /// # Shared objects /// -/// The RHF-level response object ([`crate::analdrv::response::rresp_interface::RRespSCF`]) is -/// built once and shared (sequentially, by `&mut`) between the task arms: the RHF hessian always -/// needs it, and the multipole task needs it for the relaxed DH increments. The UHF hessian -/// builds and owns its own (U-side) response machinery internally. +/// The SCF response object (the [`RespSCF`] carrier — [`RRespSCF`] for RHF, +/// [`URespSCF`] for UHF, per the SCF type) is built once and shared (sequentially, by `&mut`) +/// between the task arms: the hessian always needs it, and the multipole task needs it for the +/// relaxed DH increments. /// /// The multipole task on post-SCF (fifth-DFA) methods needs the DFT grids (for both the unrelaxed /// and the relaxed DH increments), which `xdh_calculations` frees after the energy evaluation; /// they are regenerated here if absent. This is grid generation only: the AO tabulation is /// prepared by the respective consumers, not by the grid build. -pub fn analdrv_interface( - scf_data: &mut SCF, - tasks: &[AnalDrvTask], - config: &AnalDrvConfig, -) -> AnaldrvOutput { +/// +/// [`RRespSCF`]: crate::analdrv::response::rresp_interface::RRespSCF +/// [`URespSCF`]: crate::analdrv::response::uresp_interface::URespSCF +pub fn analdrv_interface(scf_data: &mut SCF, tasks: &[AnalDrvTask], config: &AnalDrvConfig) -> AnaldrvOutput { let mut output = AnaldrvOutput::default(); // --- preparation (the only phase taking scf_data mutably) --- // @@ -77,8 +78,7 @@ pub fn analdrv_interface( let is_fifth = scf_data.mol.xc_data.is_fifth_dfa(); // the relaxed (Z-vector) DH multipole increments are the only multipole requirement on the // response object; the unrelaxed increments still need the DFT grids (regenerated below) - let multipole_relaxed_dh = - has_multipole && is_fifth && config.multipole.rdm1_relax == MultipoleRdm1Relax::Relaxed; + let multipole_relaxed_dh = has_multipole && is_fifth && config.multipole.rdm1_relax == MultipoleRdm1Relax::Relaxed; // the DFT grids are needed only when a numint (DFT) consumer exists: the final-energy // functional's XC part, or, in the relaxed mode, the SCF functional's XC response. A pure // MP2-family method (HF functional, no final-functional DFT part) consumes no grid at all. @@ -92,8 +92,12 @@ pub fn analdrv_interface( } let need_resp = has_hessian || multipole_relaxed_dh; - let mut resp_objs = if need_resp && matches!(scf_data.scftype, SCFType::RHF) { - Some(rscf_resp_interface(scf_data, config)) + let mut resp_obj = if need_resp { + match scf_data.scftype { + SCFType::RHF => Some(RespSCF::R(rscf_resp_interface(scf_data, config))), + SCFType::UHF => Some(RespSCF::U(uscf_resp_interface(scf_data, config))), + _ => None, + } } else { None }; @@ -103,13 +107,13 @@ pub fn analdrv_interface( for task in tasks { match task { AnalDrvTask::Hessian => { - let hess = hess_interface(scf_data, config, resp_objs.as_mut()); + let hess = hess_interface(scf_data, config, resp_obj.as_mut()); output.frequencies_cm = Some(hess.frequencies_cm); output.modes_trv = Some(hess.modes_trv); output.thermo = hess.thermo; }, AnalDrvTask::Multipole => { - let resp = if multipole_relaxed_dh { resp_objs.as_mut() } else { None }; + let resp = if multipole_relaxed_dh { resp_obj.as_mut() } else { None }; output.multipole = Some(multipole_interface(scf_data, config, resp)); }, } @@ -126,15 +130,9 @@ pub fn analdrv_interface( /// under the legacy top-level `"thermo"` key when present. pub fn analdrv_json_interface(output: &AnaldrvOutput) -> HashMap { let mut json = HashMap::new(); - json.insert( - "analdrv".to_string(), - serde_json::to_value(output).unwrap_or(serde_json::Value::Null), - ); + json.insert("analdrv".to_string(), serde_json::to_value(output).unwrap_or(serde_json::Value::Null)); if let Some(thermo) = &output.thermo { - json.insert( - "thermo".to_string(), - serde_json::to_value(thermo).unwrap_or(serde_json::Value::Null), - ); + json.insert("thermo".to_string(), serde_json::to_value(thermo).unwrap_or(serde_json::Value::Null)); } json } diff --git a/src/analdrv/mod.rs b/src/analdrv/mod.rs index a698deb4f..f13dc6556 100644 --- a/src/analdrv/mod.rs +++ b/src/analdrv/mod.rs @@ -1,8 +1,9 @@ //! Analytical derivative module for REST. //! -//! This should serve as semi-independent module, handling property computations that requires (only) -//! derivatives. The energy components are considered to be linearly added, so that components serve -//! as independent operators (term of ORCA); trait design of this module will fully support this idea. +//! This should serve as semi-independent module, handling property computations that requires +//! (only) derivatives. The energy components are considered to be linearly added, so that +//! components serve as independent operators (term of ORCA); trait design of this module will fully +//! support this idea. //! //! This is a more modular and flexible design, and should be easier to maintain and extend. //! The API design document is not written at this time, but will be available in the future. @@ -19,21 +20,23 @@ //! through the DH density increments. //! //! This module does not contain extensive detailed implementation. Please refer to the [`hessian`] -//! submodule for hessian traits, component implementations, total hessian drivers (including CP-SCF), -//! and important utilities. +//! submodule for hessian traits, component implementations, total hessian drivers (including +//! CP-SCF), and important utilities. //! - For optimized RI-JK implementation, please refer to [`crate::ri_jk`] module. //! - For DFT matmul implementation, please refer to [`crate::dft::numint_matmul`] module. //! //! We will also handle interface to REST. //! //! Some important utilities comes from other programs, and we acknowledge them here. -//! - `vibration/vib.rs`: Vibration analysis from Psi4, partially translated by AI, not fully reviewed by human. -//! - TR/V (translation-rotation and vibration classification) is different to Psi4. We will use rotor-type -//! to determine number of degrees of freedom (TR mode). -//! - `point_group_detect`: Point group detection from Psi4, translated by AI, not reviewed by human but have been tested. +//! - `vibration/vib.rs`: Vibration analysis from Psi4, partially translated by AI, not fully +//! reviewed by human. +//! - TR/V (translation-rotation and vibration classification) is different to Psi4. We will use +//! rotor-type to determine number of degrees of freedom (TR mode). +//! - `point_group_detect`: Point group detection from Psi4, translated by AI, not reviewed by human +//! but have been tested. //! - Note some point group detection is minorly different (such as C3v). -//! - `krylov_block.rs`: Krylov solver (used in CP-SCF) from PySCF, translated with help by AI, reviewed -//! by extensive testing. +//! - `krylov_block.rs`: Krylov solver (used in CP-SCF) from PySCF, translated with help by AI, +//! reviewed by extensive testing. #![warn(unused)] @@ -77,6 +80,8 @@ pub mod prelude { pub use multipole::trait_multipole::MultipoleNucAPI; pub use response::rresp_interface::RRespSCF; pub use response::trait_rresp::RRespAPI; + pub use response::trait_uresp::URespAPI; + pub use response::uresp_interface::URespSCF; pub use trait_util::AnalDrvBaseAPI; pub(super) use crate::ri_jk::util::{get_dm0_restricted, get_dme0_restricted}; diff --git a/src/analdrv/multipole/interface.rs b/src/analdrv/multipole/interface.rs index e882b9681..9a2977d18 100644 --- a/src/analdrv/multipole/interface.rs +++ b/src/analdrv/multipole/interface.rs @@ -3,7 +3,7 @@ use crate::analdrv::config::{AnalDrvConfig, MultipoleRdm1Relax}; use crate::analdrv::multipole::rmultipole::RMultipoleDH; use crate::analdrv::response::rgfock_interface::rgfock_dh_interface; -use crate::analdrv::response::rresp_interface::RRespSCF; +use crate::analdrv::response::RespSCF; use crate::ri_jk::util::get_cint_mol; use crate::scf_io::{SCFType, SCF}; use crate::utilities::rstsr_util::RestTensorToRstsrTsrAPI; @@ -64,12 +64,14 @@ pub struct MultipoleOutput { /// (Z-vector) increment as well; whether `resp_objs` is passed is decided by the caller from /// `multipole_rdm1_relax`. /// -/// `resp_objs` is only used for the relaxed DH increments, and requires the DFT grids to be -/// present; the caller ([`crate::analdrv::interface::analdrv_interface`]) regenerates them. +/// `resp_objs` is the [`RespSCF`] carrier (the multipole task itself is restricted-only for now; +/// the unrestricted variant is rejected here until the U arm is wired). It is only used for the +/// relaxed DH increments, and requires the DFT grids to be present; the caller +/// ([`crate::analdrv::interface::analdrv_interface`]) regenerates them. pub fn multipole_interface<'a>( scf_data: &'a SCF, config: &AnalDrvConfig, - resp_objs: Option<&mut RRespSCF<'a>>, + resp_objs: Option<&mut RespSCF<'a>>, ) -> MultipoleOutput { match scf_data.scftype { SCFType::RHF => {}, @@ -107,7 +109,17 @@ pub fn multipole_interface<'a>( None }; - let mut rmultipole = RMultipoleDH::new(&mol_cint, mo_coeff, mo_occ, origin, rgfock.as_mut(), resp_objs); + // The relaxed (Z-vector) increments consume the restricted response object; the U variant of + // the carrier is rejected here — the unrestricted multipole task is not wired yet (the task + // is also restricted-only at the SCF-type check above). + let resp = resp_objs.map(|resp| match resp { + RespSCF::R(r) => r, + RespSCF::U(_) => panic!( + "Multipole evaluation is currently only implemented for restricted (RHF/RKS) calculations; an unrestricted response object was given." + ), + }); + + let mut rmultipole = RMultipoleDH::new(&mol_cint, mo_coeff, mo_occ, origin, rgfock.as_mut(), resp); // evaluate the requested orders (each order caches itself; the DH density increments are // shared between the orders through the gfock result cache) diff --git a/src/analdrv/response/mod.rs b/src/analdrv/response/mod.rs index a1f62ec5b..35c007fbc 100644 --- a/src/analdrv/response/mod.rs +++ b/src/analdrv/response/mod.rs @@ -1,14 +1,36 @@ //! Energy response to orbitals and density matrix. -// fock and response (restricted) +// fock and response pub mod trait_rresp; -// generalized fock (restricted) -pub mod trait_rgfock; +pub mod trait_uresp; -// interface to REST SCF data (restricted) +// interface to REST SCF data (fock and response) pub mod rresp_interface; +pub mod uresp_interface; + +// generalized fock (restricted) +pub mod trait_rgfock; // generalized fock interface and driver for double-hybrid type methods (restricted) pub mod rgfock_hcore; pub mod rgfock_interface; +/// Carrier of the SCF response object of either spin treatment. +/// +/// Built once by the task loop ([`crate::analdrv::interface::analdrv_interface`], per the SCF +/// type) and passed by `&mut` to the property drivers (hessian, multipole), so that the response +/// preparation and cached intermediates are shared across tasks. +/// +/// This is a plain carrier with no façade methods: the restricted and unrestricted response +/// types ([`RRespSCF`]/[`URespSCF`]) have different signatures (spin arrays) and no unified driver +/// exists, so every consumer matches on the variant it supports. +/// +/// [`RRespSCF`]: rresp_interface::RRespSCF +/// [`URespSCF`]: uresp_interface::URespSCF +pub enum RespSCF<'a> { + /// Restricted response object ([`rscf_resp_interface`](rresp_interface::rscf_resp_interface)). + R(rresp_interface::RRespSCF<'a>), + /// Unrestricted response object + /// ([`uscf_resp_interface`](uresp_interface::uscf_resp_interface)). + U(uresp_interface::URespSCF<'a>), +} diff --git a/src/analdrv/response/trait_uresp.rs b/src/analdrv/response/trait_uresp.rs new file mode 100644 index 000000000..95f895d71 --- /dev/null +++ b/src/analdrv/response/trait_uresp.rs @@ -0,0 +1,139 @@ +//! Trait for Fock and response matrix (unrestricted). + +use crate::analdrv::prelude::*; + +/// Abstract class for Fock and response matrix (unrestricted). +/// +/// Unrestricted sibling of [`RRespAPI`](crate::analdrv::response::trait_rresp::RRespAPI): every +/// orbital-space quantity exists per spin, as `[TsrView; 2]` (input) / `[Tsr; 2]` (output) arrays +/// with index 0 = α and index 1 = β. +/// +/// # Term Explanation +/// +/// - **Fock**: the first-order derivative to energy wrt density matrix. +/// - **Resp** (response): the second-order derivative to energy wrt density matrix, contracted by +/// input rdm/bra. +/// - **rdm1**: reduced one-particle density matrix +/// - **bra**: Bra-ket (half side), usually refers to occupied molecular coefficients (as input) or +/// contracted fock/response matrix that is half-transformed by occupied molecular coefficients +/// (as output). +/// +/// # Unrestricted conventions +/// +/// - The Coulomb response sees the **total** (α+β) density: a single AO operator is produced from +/// `sum_s bra_s @ mocc_s.T` internally; the consumer applies the per-spin right +/// half-transformation and a `0.5` prefactor against the restricted convention (which carries +/// the occupation-2 factor). Exchange is strictly same-spin. +/// - The DFT XC kernel `fxc` carries its spin blocks `[nvar, 2, nvar, 2]`; the response carries the +/// UHF factor `2.0` (hermitian symmetry only, no spin degeneracy) against the restricted `4.0`. +/// - The hessian prefactors assembled from the CP-SCF solution are `2*` per spin against the +/// restricted `4*` (see the `get_cpscf_hess` of the hessian driver). +pub trait URespAPI: AnalDrvBaseAPI { + /// Generate Fock matrix by density matrix (rdm, reduced one-particle density matrix). + /// + /// This is the canonical way to generate Fock matrix, however, in many cases it is not the most + /// efficient way. Refer to [`URespAPI::get_fock_coeff`] for better way to generate Fock matrix. + /// + /// # Parameters + /// + /// - `rdm` : shape `[nao, nao]` per spin. Reduced one-particle (spin) density matrix. + /// + /// # Returns + /// + /// - `fock` : shape `[nao, nao]` per spin. Fock matrix; the α/β blocks share the Coulomb + /// potential of the total density and differ in the same-spin exchange. + /// + /// # Reserved for ugfock + /// + /// No unrestricted driver consumes this method yet; it is part of the trait surface reserved + /// for the future unrestricted generalized-Fock machinery, mirroring how the restricted + /// [`RRespAPI::get_fock_rdm`] serves [`RGFockAPI`](crate::analdrv::response::trait_rgfock::RGFockAPI). + fn get_fock_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2]; + + /// Generate Fock matrix from molecular coefficients and occupation numbers. + /// + /// Override this function to leverage the algorithmic advantage by occupation number over + /// molecular orbital number. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo_s]` per spin. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo_s]` per spin. Molecular orbital occupation numbers. + /// + /// Note `nmo_s` can be set to `nocc_s` if only occupied orbitals are considered. This can save + /// memory and speed up the calculation. + /// + /// # Returns + /// + /// - `fock` : shape `[nao, nao]` per spin. Fock matrix (the operator in AO basis, not + /// contracted by `mo_coeff`). + fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) -> [Tsr; 2] { + let rdm = [get_dm0_restricted(mo_coeff[0].view(), mo_occ[0].view()), get_dm0_restricted(mo_coeff[1].view(), mo_occ[1].view())]; + self.get_fock_rdm(&[rdm[0].view(), rdm[1].view()]) + } + + /// Prepare the data for response calculation. + /// + /// Response (related to second order of density matrix derivative to energy) will be called + /// multiple-times in CP-SCF solver and other places. + /// + /// Some methods (especially DFT) may be helpful to prepare some data for response calculation, + /// and store them in the object. + /// + /// For Hartree-Fock methods, they usually also need to store the `mo_coeff` and `mo_occ`, so to + /// make sure [`get_response_bra`](Self::get_response_bra) can be called with only bra as input. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo_s]` per spin. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo_s]` per spin. Molecular orbital occupation numbers. + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]); + + /// Generate response matrix. + /// + /// Call [`make_response_preparation`] before this function to make sure the data is ready. + /// + /// # Parameters + /// + /// - `rdm` : shape `[nao, nao, ...]` per spin. Reduced one-particle (spin) density matrix list. + /// + /// # Returns + /// + /// - `resp` : shape `[nao, nao, ...]` per spin. Response matrix. + /// + /// # Reserved for ugfock + /// + /// No unrestricted implementation provides this rdm-form entry yet; the method is part of the + /// trait surface reserved for the future unrestricted generalized-Fock machinery, mirroring + /// [`RRespAPI::get_response_rdm`] (which stayed unimplemented until the RI-PT2 generalized + /// Fock needed it). + /// + /// [`make_response_preparation`]: Self::make_response_preparation + fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2]; + + /// Generate response matrix in half-transformed MO basis. + /// + /// Override this function to leverage the algorithmic advantage by occupation number over + /// molecular orbital number. + /// Call [`make_response_preparation`] before this function to make sure the data is ready. + /// + /// # Parameters + /// + /// - `bra` : shape `[nao, nocc_s, ...]` per spin. Bra-ket (half side), usually refers to + /// occupied molecular coefficients (as input) or contracted fock/response matrix that is + /// half-transformed by occupied molecular coefficients (as output). This is usually the + /// derivative of MO coefficients (like $U_{\mu i}^\mathbb{A}$ given by CP-SCF). + /// + /// # Returns + /// + /// - `resp_bra` : shape `[nao, nocc_s, ...]` per spin. Response matrix in half-transformed MO + /// basis, carrying the unrestricted prefactor structure documented on the trait. + /// + /// # Notes + /// + /// This function may not work for fractional occupation. + /// We have not prepared to propose a good API for fractional occupation. + /// + /// [`make_response_preparation`]: Self::make_response_preparation + fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2]; +} diff --git a/src/analdrv/response/uresp_interface.rs b/src/analdrv/response/uresp_interface.rs new file mode 100644 index 000000000..de2ffa670 --- /dev/null +++ b/src/analdrv/response/uresp_interface.rs @@ -0,0 +1,413 @@ +//! Response (fock/response) objects for unrestricted SCF, built from SCF data. +//! +//! This is the unrestricted counterpart of [`rscf_resp_interface`](crate::analdrv::response:: +//! rresp_interface::rscf_resp_interface): it assembles the `URespAPI` objects for all +//! electron-interaction contributions of a converged unrestricted SCF. The hessian interface calls +//! it and feeds the returned objects to `UHessSCF`; a future standalone unrestricted +//! response/property driver can call it directly, without any hessian machinery. +//! +//! The DFT XC response object evaluates the fock path directly on the SCF grid (in its native +//! round-robin atom-interleaved order — the fock path is a plain quadrature sum, so no +//! atom-regrouping is applied), and the response path on a dedicated (usually coarser) cpscf +//! grid. The hessian-side skeleton grid policy (including the MGGA level bump) lives in the +//! hessian interface; no grid identity or grid data is shared between the hessian and response +//! subsystems. + +use crate::analdrv::config::AnalDrvRespCfg; +use crate::analdrv::prelude::*; +use crate::analdrv::response::rresp_interface::scf_jk_factors; +use crate::dft::numint_matmul::nimatmul::NIMatmul; +use crate::dft::numint_matmul::resp_uks::URespKSNIMatmul; +use crate::dft::Grids; +use crate::ri_jk::resp_u::URespRIJK; +use crate::ri_jk::util::get_cint_mol; +use crate::SCF; + +use libxc::prelude::*; + +/// List of `(scale, functional)` pairs of the SCF XC functional (spin-polarized). +/// +/// Spin-polarized sibling of [`scf_xc_func_list`](crate::analdrv::response::rresp_interface:: +/// scf_xc_func_list): the functionals are created with [`LibXCSpin::Polarized`], as required by +/// the UKS rho/vxc/fxc spin-block layout. +pub fn scf_xc_func_list_uks(scf_data: &SCF) -> Vec<(f64, LibXCFunctional)> { + let xc_code = &scf_data.mol.xc_data.dfa_compnt_scf; + let xc_params = &scf_data.mol.xc_data.dfa_paramr_scf; + 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. +/// +/// The type is itself a composite `URespAPI`: every trait method fans out to the elements, with +/// the matrix-valued outputs (`get_fock_rdm`/`get_fock_coeff`/`get_response_rdm`/ +/// `get_response_bra`) summed over the contributions per spin. Drivers therefore hold `&mut +/// URespSCF` as their single response object. +/// +/// On top of the trait, the type carries the CP-SCF machinery shared by the drivers +/// ([`Self::response_mo`]/[`Self::response_dimless_cpscf`]/[`Self::solve_dimless_cpscf`]): the +/// solver settings are captured at build time, and the orbital state is stored by +/// [`Self::make_cpscf_preparation`]. The CP-SCF solve is genuinely **coupled** across spins: the +/// α and β unknowns are flattened into one block-Krylov space (`nmo_α·nocc_α + nmo_β·nocc_β`), +/// and every response contraction sees both spins in a single call. +/// +/// The list is open-ended: further contributions can be appended to [`Self::resp_list`] without +/// changing this type. Should a driver require the concrete type of an entry, manual downcast +/// through utilities on `AnalDrvBaseAPI` is the escape hatch. +pub struct URespSCF<'a> { + /// Response objects of all electron-interaction contributions: for examples, RI-JK, RSH + /// exchange, DFT XC NIMatmul. + pub resp_list: Vec>, + /// CP-SCF solver settings, captured at build time by [`uscf_resp_interface`]. + pub resp_cfg: AnalDrvRespCfg, + /// Orbital state stored by [`Self::make_cpscf_preparation`]; the CP-SCF machinery panics + /// until it is set. + cpscf_state: Option, +} + +/// Orbital state cached by [`URespSCF::make_cpscf_preparation`] for the inherent CP-SCF +/// machinery, per spin (index 0 = α, index 1 = β). +pub struct UCpscfState { + /// Molecular orbital coefficients, shape `[nao, nmo_s]` per spin. + pub mo_coeff: [Tsr; 2], + /// Level-shifted orbital-energy differences `e_a - e_i + shift`, shape `[nvir_s, nocc_s]` + /// per spin. + pub e_ai_shift: [Tsr; 2], + /// Number of occupied orbitals per spin. + pub nocc: [usize; 2], + /// Total number of molecular orbitals per spin. + pub nmo: [usize; 2], +} + +impl<'a> AnalDrvBaseAPI for URespSCF<'a> {} + +impl<'a> URespAPI for URespSCF<'a> { + fn get_fock_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + let mut fock = [rt::zeros_like(&rdm[0]), rt::zeros_like(&rdm[1])]; + for resp_obj in self.resp_list.iter_mut() { + let fock_obj = resp_obj.get_fock_rdm(&[rdm[0].view(), rdm[1].view()]); + fock[0] += &fock_obj[0]; + fock[1] += &fock_obj[1]; + } + fock + } + + fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) -> [Tsr; 2] { + // delegated per element (not through the default rdm route), so that elements overriding + // this function for efficiency keep their advantage + let mut fock: Option<[Tsr; 2]> = None; + for resp_obj in self.resp_list.iter_mut() { + let fock_obj = resp_obj + .get_fock_coeff(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()]); + fock = Some(match fock { + Some([f0, f1]) => [f0 + &fock_obj[0], f1 + &fock_obj[1]], + None => fock_obj, + }); + } + fock.expect("URespSCF must hold at least one response object") + } + + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { + for resp_obj in self.resp_list.iter_mut() { + resp_obj.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[ + mo_occ[0].view(), + mo_occ[1].view(), + ]); + } + } + + fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + let mut resp = [rt::zeros_like(&rdm[0]), rt::zeros_like(&rdm[1])]; + for resp_obj in self.resp_list.iter_mut() { + let resp_obj_rdm = resp_obj.get_response_rdm(&[rdm[0].view(), rdm[1].view()]); + resp[0] += &resp_obj_rdm[0]; + resp[1] += &resp_obj_rdm[1]; + } + resp + } + + fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2] { + let mut resp = [rt::zeros_like(&bra[0]), rt::zeros_like(&bra[1])]; + for resp_obj in self.resp_list.iter_mut() { + let resp_obj_bra = resp_obj.get_response_bra(&[bra[0].view(), bra[1].view()]); + resp[0] += &resp_obj_bra[0]; + resp[1] += &resp_obj_bra[1]; + } + resp + } +} + +impl<'a> URespSCF<'a> { + /// Store the orbital state and prepare all response objects for the CP-SCF calculation. + /// + /// This wraps the trait-level [`URespAPI::make_response_preparation`] fan-out, and + /// additionally caches the per-spin orbital coefficients and the level-shifted + /// orbital-energy differences, so that [`Self::response_mo`], [`Self::response_dimless_cpscf`] + /// and [`Self::solve_dimless_cpscf`] can be called with the perturbation only. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo_s]` per spin. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo_s]` per spin. Molecular orbital occupation numbers. + /// - `mo_energy` : shape `[nmo_s]` per spin. Molecular orbital energies. + pub fn make_cpscf_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], mo_energy: &[TsrView; 2]) { + let [α, β] = [0, 1]; + for resp_obj in self.resp_list.iter_mut() { + resp_obj.make_response_preparation(&[mo_coeff[α].view(), mo_coeff[β].view()], &[ + mo_occ[α].view(), + mo_occ[β].view(), + ]); + } + + let occidx = [mo_occ[α].view().greater(0).into_vec(), mo_occ[β].view().greater(0).into_vec()]; + let viridx = [occidx[α].iter().map(|&x| !x).collect_vec(), occidx[β].iter().map(|&x| !x).collect_vec()]; + let eocc = [mo_energy[α].bool_select(-1, &occidx[α]), mo_energy[β].bool_select(-1, &occidx[β])]; + let evir = [mo_energy[α].bool_select(-1, &viridx[α]), mo_energy[β].bool_select(-1, &viridx[β])]; + let e_ai = [evir[α].i((.., None)) - eocc[α].i((None, ..)), evir[β].i((.., None)) - eocc[β].i((None, ..))]; + let e_ai_shift = [&e_ai[α] + self.resp_cfg.level_shift, &e_ai[β] + self.resp_cfg.level_shift]; + + self.cpscf_state = Some(UCpscfState { + mo_coeff: [mo_coeff[α].to_owned(), mo_coeff[β].to_owned()], + e_ai_shift, + nocc: [occidx[α].iter().filter(|&&x| x).count(), occidx[β].iter().filter(|&&x| x).count()], + nmo: [mo_occ[α].shape()[0], mo_occ[β].shape()[0]], + }); + } + + /// The CP-SCF orbital state cached by [`Self::make_cpscf_preparation`], for reuse by drivers + /// building on the CP-SCF machinery. Panics before the preparation. + /// + /// Note the inherent CP-SCF methods of this type access the state as a direct field, keeping + /// the borrow disjoint from `resp_list`; external callers use this accessor. + pub fn cpscf_state(&self) -> &UCpscfState { + self.cpscf_state.as_ref().expect("Call `URespSCF::make_cpscf_preparation` before the CP-SCF machinery.") + } + + /// Compute the response in MO space to a perturbation in MO space. + /// + /// Half-transforms the per-spin perturbation to the AO bra, contracts the response of all + /// contributions (each contraction sees both spins at once — the Coulomb and XC kernels + /// couple them), and transforms back to MO space. + /// Call [`Self::make_cpscf_preparation`] before this function to make sure the data is ready. + /// + /// # Parameters + /// + /// - `mo1` : shape `[nmo_s, nocc_s, ...]` per spin. The perturbation in MO space. + /// + /// # Returns + /// + /// - `resp` : shape `[nmo_s, nocc_s, ...]` per spin. The response in MO space. + pub fn response_mo(&mut self, mo1: &[TsrView; 2]) -> [Tsr; 2] { + let [α, β] = [0, 1]; + let state = + self.cpscf_state.as_ref().expect("Call `URespSCF::make_cpscf_preparation` before the CP-SCF machinery."); + let mo_coeff = [state.mo_coeff[α].view(), state.mo_coeff[β].view()]; + let ubra_α = &mo_coeff[α] % &mo1[α]; + let ubra_β = &mo_coeff[β] % &mo1[β]; + let mut resp_α = rt::zeros_like(&ubra_α); + let mut resp_β = rt::zeros_like(&ubra_β); + + // the get_response_bra fan-out is inlined (not the composite trait method), keeping the + // `cpscf_state` borrow disjoint from the `resp_list` iteration + for resp_obj in self.resp_list.iter_mut() { + let el_resp = resp_obj.get_response_bra(&[ubra_α.view(), ubra_β.view()]); + resp_α += mo_coeff[α].t() % &el_resp[α]; + resp_β += mo_coeff[β].t() % &el_resp[β]; + } + [resp_α, resp_β] + } + + /// Compute the response in the dimensionless form used by the CP-SCF solve. + /// + /// Compared to [`Self::response_mo`], this additionally handles per spin + /// - the level shift in denominator + /// - the zeroing of occupied-part response (we use `mo1[occ, occ]` part for evaluating + /// `resp[vir, occ]`, but we actually only want to solve the `mo1[vir, occ]` part and freeze + /// `mo1[occ, occ]` part to always be 0.5 times of ovlp_deriv1). + /// Call [`Self::make_cpscf_preparation`] before this function to make sure the data is ready. + /// + /// # Parameters + /// + /// - `mo1` : shape `[nmo_s, nocc_s, ...]` per spin. The perturbation in MO space. + /// + /// # Returns + /// + /// - `resp` : shape `[nmo_s, nocc_s, ...]` per spin. The dimensionless response in MO space. + pub fn response_dimless_cpscf(&mut self, mo1: &[TsrView; 2]) -> [Tsr; 2] { + let [α, β] = [0, 1]; + let mut resp = self.response_mo(mo1); + + let state = + self.cpscf_state.as_ref().expect("Call `URespSCF::make_cpscf_preparation` before the CP-SCF machinery."); + let level_shift = self.resp_cfg.level_shift; + let so = [rt::slice!(0, state.nocc[α]), rt::slice!(0, state.nocc[β])]; + let sv = [rt::slice!(state.nocc[α], state.nmo[α]), rt::slice!(state.nocc[β], state.nmo[β])]; + + // handle dimensionless denominator and force handle virtual-part only + if level_shift != 0.0 { + resp[α] -= level_shift * &mo1[α]; + resp[β] -= level_shift * &mo1[β]; + } + *&mut resp[α].i_mut(sv[α]) /= &state.e_ai_shift[α]; + *&mut resp[β].i_mut(sv[β]) /= &state.e_ai_shift[β]; + resp[α].i_mut(so[α]).fill(0.0); + resp[β].i_mut(so[β]).fill(0.0); + resp + } + + /// Solve the dimensionless CP-SCF equation using a Krylov solver. + /// + /// This solves `U + resp(U) = rhs` for both spins **coupled**: the α and β unknowns are + /// flattened into one block-Krylov space (`nmo_α·nocc_α + nmo_β·nocc_β` per property), so + /// the spin-coupled response (total-density Coulomb, spin-block XC kernel) is applied + /// exactly. Note difference of standard CP-SCF equation as mentioned in the hessian driver. + /// Call [`Self::make_cpscf_preparation`] before this function to make sure the data is ready. + /// + /// # Parameters + /// + /// - `rhs` : shape `[nmo_s, nocc_s, ...]` per spin. Dimensionless right-hand side. + /// + /// # Returns + /// + /// - `mo1` : shape `[nmo_s, nocc_s, ...]` per spin. Perturbation in MO space that solves the + /// dimensionless CP-SCF equation. + pub fn solve_dimless_cpscf(&mut self, rhs: &[TsrView; 2]) -> [Tsr; 2] { + let [α, β] = [0, 1]; + let rhs_shape = [rhs[α].shape().to_vec(), rhs[β].shape().to_vec()]; + let nmo = [rhs[α].shape()[0], rhs[β].shape()[0]]; + let nocc = [rhs[α].shape()[1], rhs[β].shape()[1]]; + let rhs = [rhs[α].reshape((nmo[α], nocc[α], -1)), rhs[β].reshape((nmo[β], nocc[β], -1))]; + let device = rhs[α].device().clone(); + + let tol = self.resp_cfg.tol; + let max_cycle = self.resp_cfg.max_cycle; + let max_space = self.resp_cfg.max_space; + let lindep = self.resp_cfg.lindep; + let tol_inflation = self.resp_cfg.tol_inflation; + + let pack_flattened = |x: &[TsrView; 2]| -> Tsr { + // original: [nmo_α, nocc_α, nprop] and [nmo_β, nocc_β, nprop] + // target: [nmo_α * nocc_α + nmo_β * nocc_β, nprop] + assert_eq!(x[α].ndim(), 3, "Expected x[α] to have shape [nmo_α, nocc_α, nprop]"); + assert_eq!(x[β].ndim(), 3, "Expected x[β] to have shape [nmo_β, nocc_β, nprop]"); + let nprop = x[α].shape()[2]; + let mut x_flattened = rt::zeros(([nmo[α] * nocc[α] + nmo[β] * nocc[β], nprop], &device)); + for A in 0..nprop { + x_flattened.i_mut((..nmo[α] * nocc[α], A)).assign(x[α].i((.., .., A)).reshape(-1)); + x_flattened.i_mut((nmo[α] * nocc[α].., A)).assign(x[β].i((.., .., A)).reshape(-1)); + } + x_flattened + }; + + let unpack_flattened = |x: TsrView| -> [Tsr; 2] { + // original: [nmo_α * nocc_α + nmo_β * nocc_β, nprop] + // target: [nmo_α, nocc_α, nprop] and [nmo_β, nocc_β, nprop] + assert_eq!(x.ndim(), 2, "Expected x to have shape [nmo_α * nocc_α + nmo_β * nocc_β, nprop]"); + let nprop = x.shape()[1]; + let idx_split = nmo[α] * nocc[α]; + let mut x_α = rt::zeros(([nmo[α], nocc[α], nprop], &device)); + let mut x_β = rt::zeros(([nmo[β], nocc[β], nprop], &device)); + for A in 0..nprop { + x_α.i_mut((.., .., A)).assign(x.i((..idx_split, A)).reshape((nmo[α], nocc[α]))); + x_β.i_mut((.., .., A)).assign(x.i((idx_split.., A)).reshape((nmo[β], nocc[β]))); + } + [x_α, x_β] + }; + + let response_cpscf_flattened = |x: TsrView| -> Tsr { + // split x by spin and reshape to original shape + let [x_α, x_β] = unpack_flattened(x); + // compute response by usual means + let resp = self.response_dimless_cpscf(&[x_α.view(), x_β.view()]); + // flatten resp to shape (nmo*nocc, nprop) + let resp_view = resp.iter().map(|r| r.view()).collect_array().unwrap(); + pack_flattened(&resp_view) + }; + + let rhs_view = rhs.iter().map(|r| r.view()).collect_array().unwrap(); + let rhs_packed = pack_flattened(&rhs_view); + let mo1_flattened = krylov_block( + response_cpscf_flattened, + rhs_packed.view(), + None, + tol, + max_cycle, + max_space, + lindep, + tol_inflation, + ); + let [mo1_α, mo1_β] = unpack_flattened(mo1_flattened.view()); + let mo1_α = mo1_α.into_shape(rhs_shape[α].to_vec()); + let mo1_β = mo1_β.into_shape(rhs_shape[β].to_vec()); + [mo1_α, mo1_β] + } +} + +/// Build the response (fock/response) objects for a converged unrestricted SCF. +pub fn uscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> URespSCF<'a> { + let device = DeviceBLAS::default(); + let mut resp_list: Vec> = Vec::new(); + + // --- RI-JK --- // + + let (factor_j, factor_k, rsh) = scf_jk_factors(scf_data); + + { + let (rimatr, _, _) = scf_data.rimatr.as_ref().expect( + "This implementation requires cholesky decomposed ERI (or rimatr) to be available and stored in memory.", + ); + let cderi = rimatr.to_rstsr_view(&device).into_cow(); + resp_list.push(Box::new(URespRIJK::new_with_cderi(factor_j, factor_k, cderi))); + } + + // The short-range exchange correction (range-separated hybrids) is a separate response + // object reusing the full-range implementation; it evaluates on the short-range `rimatr_sr` + // ERI, with no Coulomb part (factor_j = 0). + if let Some((_omega, factor_k_sr)) = rsh { + let (rimatr_sr, _, _) = scf_data.rimatr_sr.as_ref().expect( + "The range-separated response requires the short-range ERI (rimatr_sr) to be built and stored in memory.", + ); + let cderi_sr = rimatr_sr.to_rstsr_view(&device).into_cow(); + resp_list.push(Box::new(URespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr))); + } + + // --- DFT --- // + + let is_hf = scf_data.mol.xc_data.dfa_compnt_scf.is_empty(); + if !is_hf { + let mol = get_cint_mol(&scf_data.mol); + let xc_func_list = scf_xc_func_list_uks(scf_data); + let verbose = scf_data.mol.ctrl.print_level > 2; + + // common grid (fock path): the SCF grid as is; the fock path is a plain quadrature sum, + // so it does not require the atom-grouped (ByAtom) ordering. + let grids = scf_data.grids.as_ref().unwrap(); + let ni = NIMatmul::new(&mol, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); + + // response grid: when it coincides with the SCF grid, leave `ni_resp = None`; the + // response then evaluates (and caches) on the common grid. Otherwise build a dedicated + // (usually coarser) grid. + let grid_gen_level = scf_data.mol.ctrl.grid_gen_level; + let grid_resp_level = config.resp.grid_level.unwrap_or(grid_gen_level.max(3) - 2); + let resp_obj = if grid_resp_level == grid_gen_level { + URespKSNIMatmul::new(xc_func_list, ni, verbose) + } else { + let cpscf_grid = Grids::build_with_level(&scf_data.mol, grid_resp_level); + let ni_resp = NIMatmul::new( + &mol, + &cpscf_grid.coordinates, + &cpscf_grid.weights, + &cpscf_grid.atm_idx, + &cpscf_grid.quadrature_weights, + ); + URespKSNIMatmul::new(xc_func_list, ni, verbose).set_ni_resp(ni_resp) + }; + resp_list.push(Box::new(resp_obj)); + } + + URespSCF { resp_list, resp_cfg: config.resp.clone(), cpscf_state: None } +} diff --git a/src/dft/numint_matmul/hess_uks.rs b/src/dft/numint_matmul/hess_uks.rs index b47c75e92..e93ec9105 100644 --- a/src/dft/numint_matmul/hess_uks.rs +++ b/src/dft/numint_matmul/hess_uks.rs @@ -135,75 +135,6 @@ pub fn get_rho_exc_vxc_fxc_uks( (rho, exc, vxc, fxc) } -/// Evaluate the spin-polarized XC potential `vxc` and kernel `fxc` from a given ground-state -/// `rho`, by summing each sub-functional's spin-polarized `libxc_eval_eff` (deriv = 2) contribution. -/// -/// `rho` : shape `[ngrids, nvar, 2]`. Extracted from the rho formation so that callers which -/// already have `rho` (e.g. obtained directly from occupied spin-orbitals via a bra-ket -/// contraction) can skip the `ao_dm0`-based density formation. -pub fn eval_vxc_fxc_uks_from_rho(xc_func_list: &[(f64, LibXCFunctional)], rho: TsrView) -> (Tsr, Tsr) { - assert!(!xc_func_list.is_empty(), "xc_func_list must not be empty"); - let xc_type = xc_func_list - .iter() - .map(|(_, f)| determine_den_type(f)) - .max_by_key(|t| t.num_nvar()) - .expect("xc_func_list must not be empty"); - let nvar = xc_type.num_nvar(); - let ngrids = rho.shape()[0]; - let device = rho.device().clone(); - - let mut vxc = rt::zeros(([ngrids, nvar, 2], &device)); - let mut fxc = rt::zeros(([ngrids, nvar, 2, nvar, 2], &device)); - for (scale, xc_func) in xc_func_list { - let xc_type_i = determine_den_type(xc_func); - let nvar_i = xc_type_i.num_nvar(); - let rho_i = rho.i((.., ..nvar_i, ..)); - let xc_eff = libxc_eval_eff(xc_func, rho_i, 2, false); - let [_, vxc_i, fxc_i] = xc_eff.into_iter().collect_array().unwrap(); - *&mut vxc.i_mut((.., ..nvar_i, ..)) += *scale * vxc_i; - *&mut fxc.i_mut((.., ..nvar_i, .., ..nvar_i, ..)) += *scale * fxc_i; - } - (vxc, fxc) -} - -/// Lean evaluation of only `vxc` and `fxc` on a given grid, for use as the CP-KS `cpks_vxc` / -/// `cpks_fxc` when a dedicated (coarser) CP-KS grid is attached (UKS counterpart of -/// [`make_cpks_vxc_fxc`](super::resp_rks::make_cpks_vxc_fxc)). -/// -/// Skips all skeleton-Hessian intermediates, evaluates AO at the minimum derivative order needed -/// to form the spin densities (`xc_type.num_ao_deriv()`), and forms each spin density `rhoσ` -/// directly from the occupied spin-orbital coefficients via a bra-ket contraction -/// ([`NIMatmul::make_rho_from_homogeneous_braket`]) instead of building the full `[nao, nao]` spin -/// density matrices (consistent with the CP-KS response path in [`get_uks_response_bra`]). -pub fn make_cpks_vxc_fxc_uks( - xc_func_list: &[(f64, LibXCFunctional)], - ni: &mut NIMatmul, - mo_coeff: &[TsrView; 2], - mo_occ: &[TsrView; 2], -) -> (Tsr, Tsr) { - let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); - - // For each spin, bake the occupation into the occupied spin coefficients (sqrt so the bra-ket - // square reproduces the occ-weighted spin density for every component) and form rhoσ directly - // via a bra-ket contraction. For UKS the spin occupation is 1, so this is a no-op in the common - // case but stays correct for fractional occupation. - let occidx_α = mo_occ[α].view().greater(0).into_vec(); - let mocc_α = mo_coeff[α].bool_select(-1, &occidx_α); - let occ_α = mo_occ[α].bool_select(-1, &occidx_α); - let occ_sqrt_α = occ_α.mapv(f64::sqrt); - let mocc_2_α = &mocc_α * occ_sqrt_α.i((None, ..)); - - let occidx_β = mo_occ[β].view().greater(0).into_vec(); - let mocc_β = mo_coeff[β].bool_select(-1, &occidx_β); - let occ_β = mo_occ[β].bool_select(-1, &occidx_β); - let occ_sqrt_β = occ_β.mapv(f64::sqrt); - let mocc_2_β = &mocc_β * occ_sqrt_β.i((None, ..)); - - // rho : [ngrids, nvar, 2] - let rho = ni.make_rho_from_homogeneous_braket(&[mocc_2_α.view(), mocc_2_β.view()], xc_type); - eval_vxc_fxc_uks_from_rho(xc_func_list, rho.view()) -} - /// Single spin-pair fxc contraction `einsum("g, Atxg, xyg, Bsyg -> ABts", /// weights, drho1, fxc_block, drho2)`, the inner loop of [`get_de_fxc_uks`]. /// @@ -234,8 +165,7 @@ fn get_de_fxc_uks_inner(wf_block: TsrView, drho1: TsrView, drho2: TsrView) -> Ts /// /// - `wf` : shape `[ngrids, nvar, 2, nvar, 2]` (g, x, s1, y, s2). Grid-weighted spin-polarized fxc /// kernel. -/// - `drhoα`, `drhoβ` : shape `[ngrids, nvar, 3, natm]`, per-spin outputs of -/// [`get_drho`]. +/// - `drhoα`, `drhoβ` : shape `[ngrids, nvar, 3, natm]`, per-spin outputs of [`get_drho`]. /// /// # Returns /// @@ -261,8 +191,7 @@ pub fn get_de_fxc_uks(wf: TsrView, drhoα: TsrView, drhoβ: TsrView) -> Tsr { /// - `xc_type` : density family. /// - `ao` : shape `[ngrids, nao, ncomp]`; reads the leading `num_ao_comp` channels (1 for RHO, 4 /// for SIGMA/TAU), through [`xc_fock_stack`]. -/// - `drhoα`, `drhoβ` : shape `[ngrids, nvar, 3, natm]`, per-spin outputs of -/// [`get_drho`]. +/// - `drhoα`, `drhoβ` : shape `[ngrids, nvar, 3, natm]`, per-spin outputs of [`get_drho`]. /// - `wf` : shape `[ngrids, nvar, 2, nvar, 2]`. Grid-weighted spin-polarized fxc kernel. /// /// # Returns @@ -496,8 +425,7 @@ pub fn get_de_becke_vxc_parts_uks( /// - `w` : shape `[ngrids]`. Grid weights of the chunk. /// - `dw` : shape `[ngrids, 3, natm]` (g, t, A). Grid-first Becke `dw` (see /// [`get_de_becke_atom_2_uks`]). -/// - `vmat_ip_α`, `vmat_ip_β` : shape `[nao, nao, 3]`, per-spin outputs of -/// [`get_vmat_ip`]. +/// - `vmat_ip_α`, `vmat_ip_β` : shape `[nao, nao, 3]`, per-spin outputs of [`get_vmat_ip`]. /// /// # Returns /// @@ -571,8 +499,8 @@ pub fn get_vmat_becke_parts_uks( /// - `ni` : numerical-integration driver restricted to the chunk's grids. /// - `dm0α`, `dm0β` : shape `[nao, nao]`. Per-spin reference density matrices in AO basis. /// - `atm_idx` : atom that generated the chunk's grids. -/// - `tables` : precomputed molecular tables of the Becke partition (shared across chunks); must -/// be `Some` when `grid_shift` is set. +/// - `tables` : precomputed molecular tables of the Becke partition (shared across chunks); must be +/// `Some` when `grid_shift` is set. /// - `hardness` : Becke switch-function hardness. /// - `grid_shift` : evaluate the Becke grid-shift terms. /// @@ -665,7 +593,8 @@ pub fn make_hessian_setup_chunk_becke_uks( if grid_shift { // --- becke partition: dw in full, ddw only through the cddw contraction --- // - // cddw (nset = 1): the only ddw consumer is de_becke_full_2 = ddw . (exc * (rhoa[0] + rhob[0])) + // cddw (nset = 1): the only ddw consumer is de_becke_full_2 = ddw . (exc * (rhoa[0] + + // rhob[0])) let cddw = (&exc * (rho.i((.., 0, α)) + rho.i((.., 0, β)))).into_vec(); let tables = tables.expect("BeckeMolTables must be given when grid_shift is set"); @@ -1049,113 +978,6 @@ pub fn make_hessian_setup_becke_uks( /* #endregion */ -/* #region response */ - -pub fn get_uks_response_bra( - ni: &mut NIMatmul, - den_type: XCDenType, - fxc_eff: TsrView, - bra: &[TsrView; 2], - mocc: &[TsrView; 2], -) -> ([Tsr; 2], IndexMap<&'static str, f64>) { - let nao = bra[α].shape()[0]; - let nocc_α = bra[α].shape()[1]; - let nocc_β = bra[β].shape()[1]; - let bra_α_shape = bra[α].shape().to_vec(); - let bra_β_shape = bra[β].shape().to_vec(); - let bra_α = bra[α].reshape((nao, nocc_α, -1)); - let bra_β = bra[β].reshape((nao, nocc_β, -1)); - let nset = bra_α.shape()[2]; - - let mut timing = IndexMap::new(); - let mut tic = |label: &'static str, t0: std::time::Instant| { - let elapsed = t0.elapsed().as_secs_f64(); - timing.insert(label, elapsed); - }; - - let t0 = std::time::Instant::now(); - ni.get_cached_ao(den_type.num_ao_deriv()); - tic("ao", t0); - - // Compute per-spin rho1 - let t0 = std::time::Instant::now(); - let bra_α_list = bra_α.axes_iter(-1).collect_vec(); - let bra_β_list = bra_β.axes_iter(-1).collect_vec(); - let rho1α = ni.make_rho_from_one_bra_mult_ket(mocc[α].view(), &bra_α_list, den_type); - let rho1β = ni.make_rho_from_one_bra_mult_ket(mocc[β].view(), &bra_β_list, den_type); - // Stack into [ngrids, nvar, 2, nset] - let ngrids = rho1α.shape()[0]; - let nvar = den_type.num_nvar(); - let device = rho1α.device().clone(); - let mut rho1 = rt::zeros(([ngrids, nvar, 2, nset], &device)); - rho1.i_mut((.., .., α, ..)).assign(&rho1α); - rho1.i_mut((.., .., β, ..)).assign(&rho1β); - tic("rho1", t0); - - // Compute UKS fxc bra-trans response - let t0 = std::time::Instant::now(); - let resp = ni.make_uks_fxc_pot_with_eff_bra_trans(fxc_eff, rho1.view(), mocc, den_type); - tic("resp", t0); - - // UKS CPHF factor: 2.0 (hermitian symmetry only, no spin degeneracy) - let [resp_α, resp_β] = resp; - let resp_α = 2.0 * resp_α.into_shape(bra_α_shape); - let resp_β = 2.0 * resp_β.into_shape(bra_β_shape); - ([resp_α, resp_β], timing) -} - -pub fn get_uks_response_bra_batched( - ni: &mut NIMatmul, - den_type: XCDenType, - fxc_eff: TsrView, - bra: &[TsrView; 2], - mocc: &[TsrView; 2], - verbose: bool, -) -> ([Tsr; 2], IndexMap<&'static str, f64>) { - let ngrids = ni.weights.len(); - let nbatch = ni.nbatch; - let bra_α_shape = bra[α].shape().to_vec(); - let bra_β_shape = bra[β].shape().to_vec(); - let device = bra[α].device().clone(); - let mut resp_α = rt::zeros((bra_α_shape, &device)); - let mut resp_β = rt::zeros((bra_β_shape, &device)); - let mut timing = IndexMap::from([("ao", 0.0), ("rho1", 0.0), ("resp", 0.0), ("total", 0.0)]); - - let t0 = std::time::Instant::now(); - for start in (0..ngrids).step_by(nbatch) { - let end = (start + nbatch).min(ngrids); - let mut ni_batch = ni.split_batch(start, end); - let ([resp_batch_α, resp_batch_β], timing_batch) = - get_uks_response_bra(&mut ni_batch, den_type, fxc_eff.i(start..end), bra, mocc); - resp_α += resp_batch_α; - resp_β += resp_batch_β; - for (key, value) in timing_batch { - *timing.get_mut(key).unwrap() += value; - } - let duration = t0.elapsed().as_secs_f64(); - timing.insert("total", duration); - if verbose { - println!("In get_uks_response_bra_batched, Batch {start}..{end}"); - println!(" Elapsed time from start (Wall time): {:.4} sec", duration); - } - } - - if verbose { - println!("Finished get_uks_response_bra_batched"); - println!(" Total elapsed time (Wall time): {:.4} sec", timing["total"]); - println!(" Timing breakdown:"); - for (key, value) in timing.iter() { - if *key != "total" { - println!(" {key:>20}: {value:.4} sec"); - } - } - } - - ([resp_α, resp_β], timing) -} - -/* #endregion */ - /* #region final implementation of UKS Hessian */ /// UKS Hessian XC component with the Becke grid-shift, the unrestricted @@ -1172,19 +994,17 @@ pub struct UHessKSNIMatmul<'a> { pub mol: CInt, /// List of `(scale, functional)` pairs of the XC functional. pub xc_func_list: Vec<(f64, LibXCFunctional)>, - /// Numerical-integration driver over the (atom-grouped) Hessian grid. + /// Numerical-integration driver over the (atom-grouped) Hessian grid, built from the same + /// grid data as the response object [`URespKSNIMatmul`] but an independent instance. + /// + /// [`URespKSNIMatmul`]: super::resp_uks::URespKSNIMatmul pub ni: NIMatmul<'a>, - /// Optional separate grid for the CP-KS response; `None` reuses `ni`. - pub ni_cpks: Option>, /// Evaluate the Becke grid-shift terms; requires full atom attribution of /// the grids. pub grid_shift: bool, /// Print per-chunk progress of the Hessian setup. pub verbose: bool, - /// Intermediates of the Hessian setup: all keys of - /// [`make_hessian_setup_becke_uks`] (with `fxc` renamed to `cpks_fxc` - /// unless a CP-KS-specific grid is given), plus `mo_coeff_0/1 [nao, nmo]`, - /// `mo_occ_0/1 [nmo]` from [`UHessElecInteractAPI::make_response_preparation`]. + /// Skeleton-Hessian intermediates: all keys of [`make_hessian_setup_becke_uks`]. pub intmd: HashMap, } @@ -1205,24 +1025,12 @@ impl<'a> UHessKSNIMatmul<'a> { grid_shift: bool, verbose: bool, ) -> Self { - Self { mol: mol.clone(), xc_func_list, ni, ni_cpks: None, grid_shift, verbose, intmd: HashMap::new() } - } - - /// Attach a dedicated CP-KS numerical-integration grid (`ni_cpks`). - /// - /// When set, the CP-KS response (`get_response_bra`) is evaluated on this grid instead of the - /// skeleton grid, and `cpks_vxc` / `cpks_fxc` are recomputed on it during - /// [`make_response_preparation`](Self::make_response_preparation). - pub fn set_ni_cpks(mut self, ni_cpks: NIMatmul<'a>) -> Self { - self.ni_cpks = Some(ni_cpks); - self + Self { mol: mol.clone(), xc_func_list, ni, grid_shift, verbose, intmd: HashMap::new() } } /// Perform the Hessian setup for UKS calculations. /// - /// `fxc` is stored as `cpks_fxc` (unless a CP-KS-specific grid is given) - /// for the response; `de_xc_skeleton` and `vmat_deriv1_grid_a/b` are the - /// main results. + /// `de_xc_skeleton` and `vmat_deriv1_grid_a/b` are the main results. pub fn make_hessian_setup(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], atm_list: Option<&[usize]>) { let occidx_α = mo_occ[α].view().greater(0).into_vec(); let occidx_β = mo_occ[β].view().greater(0).into_vec(); @@ -1243,13 +1051,7 @@ impl<'a> UHessKSNIMatmul<'a> { ); for (key, val) in result.into_iter() { - if key == "fxc" { - if self.ni_cpks.is_none() { - self.intmd.insert("cpks_fxc".to_string(), val); - } - } else { - self.intmd.insert(key.to_string(), val); - } + self.intmd.insert(key.to_string(), val); } } @@ -1286,71 +1088,6 @@ impl<'a> UHessElecInteractAPI for UHessKSNIMatmul<'a> { } [self.intmd["vmat_deriv1_grid_a"].to_owned(), self.intmd["vmat_deriv1_grid_b"].to_owned()] } - - /// Cached on first call for fixed inputs: the CP-KS `vxc`/`fxc` evaluation is skipped when - /// this object was already prepared with the same orbitals. - fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { - // cache check (before the orbital refresh): whether the expensive CP-KS-grid evaluation - // below was already performed with the same orbitals; `cpks_fxc` on the CP-KS grid is - // only written here (the `make_hessian_setup` one lives on the skeleton grid, and only - // when `ni_cpks` is absent), so with `ni_cpks` set the key coexists with the orbitals - let already_prepared = self.ni_cpks.is_some() - && self.intmd.contains_key("cpks_fxc") - && is_same_tensor(self.intmd["mo_coeff_0"].view(), mo_coeff[α].view()) - && is_same_tensor(self.intmd["mo_coeff_1"].view(), mo_coeff[β].view()) - && is_same_tensor(self.intmd["mo_occ_0"].view(), mo_occ[α].view()) - && is_same_tensor(self.intmd["mo_occ_1"].view(), mo_occ[β].view()); - self.intmd.insert("mo_coeff_0".to_string(), mo_coeff[α].view().into_contig(ColMajor)); - self.intmd.insert("mo_coeff_1".to_string(), mo_coeff[β].view().into_contig(ColMajor)); - self.intmd.insert("mo_occ_0".to_string(), mo_occ[α].view().into_contig(ColMajor)); - self.intmd.insert("mo_occ_1".to_string(), mo_occ[β].view().into_contig(ColMajor)); - - // When a dedicated CP-KS grid is set, `cpks_vxc` / `cpks_fxc` were NOT stored during - // `make_hessian_setup` (the skeleton grid's vxc/fxc live on a different grid and must not - // be reused). Recompute them here on the CP-KS grid from the ground-state spin densities, - // using the lean [`make_cpks_vxc_fxc_uks`] (no skeleton intermediates, minimal AO - // derivative order, spin densities formed from occupied spin-orbitals via a bra-ket - // contraction rather than full spin dm0 matrices). - if let Some(ni_cpks) = self.ni_cpks.as_mut() { - if already_prepared { - return; - } - let mo_coeff_α = self.intmd["mo_coeff_0"].view(); - let mo_coeff_β = self.intmd["mo_coeff_1"].view(); - let mo_occ_α = self.intmd["mo_occ_0"].view(); - let mo_occ_β = self.intmd["mo_occ_1"].view(); - let (vxc, fxc) = - make_cpks_vxc_fxc_uks(&self.xc_func_list, ni_cpks, &[mo_coeff_α, mo_coeff_β], &[mo_occ_α, mo_occ_β]); - self.intmd.insert("cpks_vxc".to_string(), vxc); - self.intmd.insert("cpks_fxc".to_string(), fxc); - } - } - - fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2] { - let ni_cpks = self.ni_cpks.as_mut().unwrap_or(&mut self.ni); - let mo_coeff_α = self.intmd["mo_coeff_0"].view(); - let mo_coeff_β = self.intmd["mo_coeff_1"].view(); - let mo_occ_α = self.intmd["mo_occ_0"].view(); - let mo_occ_β = self.intmd["mo_occ_1"].view(); - let fxc_eff = self.intmd["cpks_fxc"].view(); - - let occidx_α = mo_occ_α.view().greater(0).into_vec(); - let occidx_β = mo_occ_β.view().greater(0).into_vec(); - let mocc_α = mo_coeff_α.bool_select(-1, &occidx_α); - let mocc_β = mo_coeff_β.bool_select(-1, &occidx_β); - - let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); - - let ([resp_α, resp_β], _timing) = get_uks_response_bra_batched( - ni_cpks, - den_type, - fxc_eff.view(), - bra, - &[mocc_α.view(), mocc_β.view()], - self.verbose, - ); - [resp_α, resp_β] - } } /* #endregion */ diff --git a/src/dft/numint_matmul/mod.rs b/src/dft/numint_matmul/mod.rs index a1b5d8119..3b1f36162 100644 --- a/src/dft/numint_matmul/mod.rs +++ b/src/dft/numint_matmul/mod.rs @@ -10,6 +10,7 @@ pub mod nimatmul; pub mod pure_eval_rho; pub mod pure_xcpot; pub mod resp_rks; +pub mod resp_uks; #[allow(unused)] pub mod prelude { diff --git a/src/dft/numint_matmul/resp_uks.rs b/src/dft/numint_matmul/resp_uks.rs new file mode 100644 index 000000000..17afa3d62 --- /dev/null +++ b/src/dft/numint_matmul/resp_uks.rs @@ -0,0 +1,324 @@ +//! Response (Fock and response matrix) implementation for the UKS numint-matmul XC component. +//! +//! The response object [`URespKSNIMatmul`] is a standalone object for the fock/response +//! functionality of the UKS numerical-integration XC contribution, independent of the hessian +//! machinery of `UHessKSNIMatmul` (which owns its own grid, built from the same grid data). +//! +//! The two grids are separated by purpose: the fock path (`get_fock_rdm`/`get_fock_coeff`) always +//! evaluates on the common (large) grid `ni`, while the response path +//! (`make_response_preparation`/`get_response_bra`) uses the small response grid `ni_resp` when +//! one is attached. The returned fock is the DFT numerical-integration XC contribution only, never +//! the full SCF fock. +//! +//! Main-thread only (not `Send`/`Sync`). + +use super::prelude::*; +use crate::analdrv::prelude::*; + +#[allow(non_upper_case_globals)] +const α: usize = 0; +#[allow(non_upper_case_globals)] +const β: usize = 1; + +/// Evaluate the spin-polarized XC potential `vxc` and kernel `fxc` from a given ground-state +/// `rho`, by summing each sub-functional's spin-polarized `libxc_eval_eff` (deriv = 2) contribution. +/// +/// `rho` : shape `[ngrids, nvar, 2]`. Extracted from the rho formation so that callers which +/// already have `rho` (e.g. obtained directly from occupied spin-orbitals via a bra-ket +/// contraction) can skip the `ao_dm0`-based density formation. +pub fn eval_vxc_fxc_uks_from_rho(xc_func_list: &[(f64, LibXCFunctional)], rho: TsrView) -> (Tsr, Tsr) { + assert!(!xc_func_list.is_empty(), "xc_func_list must not be empty"); + let xc_type = xc_func_list + .iter() + .map(|(_, f)| determine_den_type(f)) + .max_by_key(|t| t.num_nvar()) + .expect("xc_func_list must not be empty"); + let nvar = xc_type.num_nvar(); + let ngrids = rho.shape()[0]; + let device = rho.device().clone(); + + let mut vxc = rt::zeros(([ngrids, nvar, 2], &device)); + let mut fxc = rt::zeros(([ngrids, nvar, 2, nvar, 2], &device)); + for (scale, xc_func) in xc_func_list { + let xc_type_i = determine_den_type(xc_func); + let nvar_i = xc_type_i.num_nvar(); + let rho_i = rho.i((.., ..nvar_i, ..)); + let xc_eff = libxc_eval_eff(xc_func, rho_i, 2, false); + let [_, vxc_i, fxc_i] = xc_eff.into_iter().collect_array().unwrap(); + *&mut vxc.i_mut((.., ..nvar_i, ..)) += *scale * vxc_i; + *&mut fxc.i_mut((.., ..nvar_i, .., ..nvar_i, ..)) += *scale * fxc_i; + } + (vxc, fxc) +} + +/// Lean evaluation of only `vxc` and `fxc` on a given grid, for use as the CP-KS `cpks_vxc` / +/// `cpks_fxc` when a dedicated (coarser) CP-KS grid is attached (UKS counterpart of +/// [`make_cpks_vxc_fxc`](super::resp_rks::make_cpks_vxc_fxc)). +/// +/// Skips all skeleton-Hessian intermediates, evaluates AO at the minimum derivative order needed +/// to form the spin densities (`xc_type.num_ao_deriv()`), and forms each spin density `rhoσ` +/// directly from the occupied spin-orbital coefficients via a bra-ket contraction +/// ([`NIMatmul::make_rho_from_homogeneous_braket`]) instead of building the full `[nao, nao]` spin +/// density matrices (consistent with the CP-KS response path in [`get_uks_response_bra`]). +pub fn make_cpks_vxc_fxc_uks( + xc_func_list: &[(f64, LibXCFunctional)], + ni: &mut NIMatmul, + mo_coeff: &[TsrView; 2], + mo_occ: &[TsrView; 2], +) -> (Tsr, Tsr) { + let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); + + // For each spin, bake the occupation into the occupied spin coefficients (sqrt so the bra-ket + // square reproduces the occ-weighted spin density for every component) and form rhoσ directly + // via a bra-ket contraction. For UKS the spin occupation is 1, so this is a no-op in the common + // case but stays correct for fractional occupation. + let occidx_α = mo_occ[α].view().greater(0).into_vec(); + let mocc_α = mo_coeff[α].bool_select(-1, &occidx_α); + let occ_α = mo_occ[α].bool_select(-1, &occidx_α); + let occ_sqrt_α = occ_α.mapv(f64::sqrt); + let mocc_2_α = &mocc_α * occ_sqrt_α.i((None, ..)); + + let occidx_β = mo_occ[β].view().greater(0).into_vec(); + let mocc_β = mo_coeff[β].bool_select(-1, &occidx_β); + let occ_β = mo_occ[β].bool_select(-1, &occidx_β); + let occ_sqrt_β = occ_β.mapv(f64::sqrt); + let mocc_2_β = &mocc_β * occ_sqrt_β.i((None, ..)); + + // rho : [ngrids, nvar, 2] + let rho = ni.make_rho_from_homogeneous_braket(&[mocc_2_α.view(), mocc_2_β.view()], xc_type); + eval_vxc_fxc_uks_from_rho(xc_func_list, rho.view()) +} + +/* #region response */ + +pub fn get_uks_response_bra( + ni: &mut NIMatmul, + den_type: XCDenType, + fxc_eff: TsrView, + bra: &[TsrView; 2], + mocc: &[TsrView; 2], +) -> ([Tsr; 2], IndexMap<&'static str, f64>) { + let nao = bra[α].shape()[0]; + let nocc_α = bra[α].shape()[1]; + let nocc_β = bra[β].shape()[1]; + let bra_α_shape = bra[α].shape().to_vec(); + let bra_β_shape = bra[β].shape().to_vec(); + let bra_α = bra[α].reshape((nao, nocc_α, -1)); + let bra_β = bra[β].reshape((nao, nocc_β, -1)); + let nset = bra_α.shape()[2]; + + let mut timing = IndexMap::new(); + let mut tic = |label: &'static str, t0: std::time::Instant| { + let elapsed = t0.elapsed().as_secs_f64(); + timing.insert(label, elapsed); + }; + + let t0 = std::time::Instant::now(); + ni.get_cached_ao(den_type.num_ao_deriv()); + tic("ao", t0); + + // Compute per-spin rho1 + let t0 = std::time::Instant::now(); + let bra_α_list = bra_α.axes_iter(-1).collect_vec(); + let bra_β_list = bra_β.axes_iter(-1).collect_vec(); + let rho1α = ni.make_rho_from_one_bra_mult_ket(mocc[α].view(), &bra_α_list, den_type); + let rho1β = ni.make_rho_from_one_bra_mult_ket(mocc[β].view(), &bra_β_list, den_type); + // Stack into [ngrids, nvar, 2, nset] + let ngrids = rho1α.shape()[0]; + let nvar = den_type.num_nvar(); + let device = rho1α.device().clone(); + let mut rho1 = rt::zeros(([ngrids, nvar, 2, nset], &device)); + rho1.i_mut((.., .., α, ..)).assign(&rho1α); + rho1.i_mut((.., .., β, ..)).assign(&rho1β); + tic("rho1", t0); + + // Compute UKS fxc bra-trans response + let t0 = std::time::Instant::now(); + let resp = ni.make_uks_fxc_pot_with_eff_bra_trans(fxc_eff, rho1.view(), mocc, den_type); + tic("resp", t0); + + // UKS CPHF factor: 2.0 (hermitian symmetry only, no spin degeneracy) + let [resp_α, resp_β] = resp; + let resp_α = 2.0 * resp_α.into_shape(bra_α_shape); + let resp_β = 2.0 * resp_β.into_shape(bra_β_shape); + ([resp_α, resp_β], timing) +} + +pub fn get_uks_response_bra_batched( + ni: &mut NIMatmul, + den_type: XCDenType, + fxc_eff: TsrView, + bra: &[TsrView; 2], + mocc: &[TsrView; 2], + verbose: bool, +) -> ([Tsr; 2], IndexMap<&'static str, f64>) { + let ngrids = ni.weights.len(); + let nbatch = ni.nbatch; + let bra_α_shape = bra[α].shape().to_vec(); + let bra_β_shape = bra[β].shape().to_vec(); + let device = bra[α].device().clone(); + let mut resp_α = rt::zeros((bra_α_shape, &device)); + let mut resp_β = rt::zeros((bra_β_shape, &device)); + let mut timing = IndexMap::from([("ao", 0.0), ("rho1", 0.0), ("resp", 0.0), ("total", 0.0)]); + + let t0 = std::time::Instant::now(); + for start in (0..ngrids).step_by(nbatch) { + let end = (start + nbatch).min(ngrids); + let mut ni_batch = ni.split_batch(start, end); + let ([resp_batch_α, resp_batch_β], timing_batch) = + get_uks_response_bra(&mut ni_batch, den_type, fxc_eff.i(start..end), bra, mocc); + resp_α += resp_batch_α; + resp_β += resp_batch_β; + for (key, value) in timing_batch { + *timing.get_mut(key).unwrap() += value; + } + let duration = t0.elapsed().as_secs_f64(); + timing.insert("total", duration); + if verbose { + println!("In get_uks_response_bra_batched, Batch {start}..{end}"); + println!(" Elapsed time from start (Wall time): {:.4} sec", duration); + } + } + + if verbose { + println!("Finished get_uks_response_bra_batched"); + println!(" Total elapsed time (Wall time): {:.4} sec", timing["total"]); + println!(" Timing breakdown:"); + for (key, value) in timing.iter() { + if *key != "total" { + println!(" {key:>20}: {value:.4} sec"); + } + } + } + + ([resp_α, resp_β], timing) +} + +/* #endregion */ + +/// Response (fock/response matrix) object for the UKS numint-matmul XC contribution. +pub struct URespKSNIMatmul<'a> { + /// List of `(scale, functional)` pairs of the XC functional (spin-polarized). + pub xc_func_list: Vec<(f64, LibXCFunctional)>, + /// Common (large) numerical-integration grid, used by the fock path. + pub ni: NIMatmul<'a>, + /// Optional separate small grid for the response path; `None` reuses `ni`. + pub ni_resp: Option>, + /// Print progress of the response evaluation. + pub verbose: bool, + /// Response intermediates: `mo_coeff_0/1 [nao, nmo_s]` and `mo_occ_0/1 [nmo_s]` from + /// [`URespAPI::make_response_preparation`], `cpks_vxc [ngrids, nvar, 2]` and `cpks_fxc + /// [ngrids, nvar, 2, nvar, 2]` on the selected response grid. + pub intmd: HashMap, +} + +impl<'a> URespKSNIMatmul<'a> { + /// Create a new UKS response object. + /// + /// # Parameters + /// + /// - `xc_func_list` : list of `(scale, functional)` pairs (spin-polarized functionals). + /// - `ni` : numerical-integration driver over the common (large) grid. + /// - `verbose` : print progress. + pub fn new(xc_func_list: Vec<(f64, LibXCFunctional)>, ni: NIMatmul<'a>, verbose: bool) -> Self { + Self { xc_func_list, ni, ni_resp: None, verbose, intmd: HashMap::new() } + } + + /// Attach a dedicated small numerical-integration grid (`ni_resp`) for the response path. + /// + /// When set, the response (`get_response_bra`) is evaluated on this grid instead of the common + /// (large) grid, and `cpks_vxc` / `cpks_fxc` are computed on it during + /// [`make_response_preparation`](URespAPI::make_response_preparation). + pub fn set_ni_resp(mut self, ni_resp: NIMatmul<'a>) -> Self { + self.ni_resp = Some(ni_resp); + self + } +} + +impl<'a> AnalDrvBaseAPI for URespKSNIMatmul<'a> {} + +impl<'a> URespAPI for URespKSNIMatmul<'a> { + /// Fock (the XC numint contribution only) per spin from spin density matrices, on the common + /// grid `ni`. + fn get_fock_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + for rdm_s in rdm { + assert_eq!(rdm_s.ndim(), 2, "rdm must have 2 dimensions"); + } + let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); + // rho : [ngrids, nvar, 2] (the spin is the set axis) + let rho = self.ni.make_rho_from_dm(&[rdm[α].view(), rdm[β].view()], den_type); + let (vxc, _fxc) = eval_vxc_fxc_uks_from_rho(&self.xc_func_list, rho.view()); + let pot = self.ni.make_vxc_pot_with_eff(vxc.view(), den_type, XCSpin::Polarized); + [pot.i((.., .., α)).to_owned(), pot.i((.., .., β)).to_owned()] + } + + // note: `get_fock_coeff` keeps the trait default (dm route via `get_dm0_restricted` per spin). + + /// rdm-form entry of the unrestricted response kernel; reserved for the future ugfock + /// machinery (no consumer yet), mirroring the restricted side before RI-PT2. + fn get_response_rdm(&mut self, _rdm: &[TsrView; 2]) -> [Tsr; 2] { + unimplemented!("get_response_rdm is not implemented for URespKSNIMatmul; reserved for the future ugfock machinery.") + } + + /// Cached on first call for fixed inputs: the CP-KS `vxc`/`fxc` evaluation is skipped when + /// this object was already prepared with the same orbitals (e.g. repeated calls from + /// multi-order property evaluations directly reuse the stored results). + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { + // cache check (before the orbital refresh): whether the expensive evaluation below was + // already performed with the same orbitals (`mo_coeff`/`mo_occ` are always inserted + // together with `cpks_vxc`/`cpks_fxc`, so the keys coexist) + let already_prepared = self.intmd.contains_key("cpks_fxc") + && is_same_tensor(self.intmd["mo_coeff_0"].view(), mo_coeff[α].view()) + && is_same_tensor(self.intmd["mo_coeff_1"].view(), mo_coeff[β].view()) + && is_same_tensor(self.intmd["mo_occ_0"].view(), mo_occ[α].view()) + && is_same_tensor(self.intmd["mo_occ_1"].view(), mo_occ[β].view()); + self.intmd.insert("mo_coeff_0".to_string(), mo_coeff[α].view().into_contig(ColMajor)); + self.intmd.insert("mo_coeff_1".to_string(), mo_coeff[β].view().into_contig(ColMajor)); + self.intmd.insert("mo_occ_0".to_string(), mo_occ[α].view().into_contig(ColMajor)); + self.intmd.insert("mo_occ_1".to_string(), mo_occ[β].view().into_contig(ColMajor)); + if already_prepared { + return; + } + + // Compute `cpks_vxc` / `cpks_fxc` on the selected response grid (the small `ni_resp` when + // attached, else the common grid) from the ground-state spin densities, using the lean + // [`make_cpks_vxc_fxc_uks`] (no skeleton intermediates, minimal AO derivative order, spin + // densities formed from occupied spin-orbitals via a bra-ket contraction rather than full + // spin dm0 matrices). + let ni_resp = self.ni_resp.as_mut().unwrap_or(&mut self.ni); + let mo_coeff_α = self.intmd["mo_coeff_0"].view(); + let mo_coeff_β = self.intmd["mo_coeff_1"].view(); + let mo_occ_α = self.intmd["mo_occ_0"].view(); + let mo_occ_β = self.intmd["mo_occ_1"].view(); + let (vxc, fxc) = + make_cpks_vxc_fxc_uks(&self.xc_func_list, ni_resp, &[mo_coeff_α, mo_coeff_β], &[mo_occ_α, mo_occ_β]); + self.intmd.insert("cpks_vxc".to_string(), vxc); + self.intmd.insert("cpks_fxc".to_string(), fxc); + } + + fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2] { + let ni_resp = self.ni_resp.as_mut().unwrap_or(&mut self.ni); + let mo_coeff_α = self.intmd["mo_coeff_0"].view(); + let mo_coeff_β = self.intmd["mo_coeff_1"].view(); + let mo_occ_α = self.intmd["mo_occ_0"].view(); + let mo_occ_β = self.intmd["mo_occ_1"].view(); + let fxc_eff = self.intmd["cpks_fxc"].view(); + + let occidx_α = mo_occ_α.view().greater(0).into_vec(); + let occidx_β = mo_occ_β.view().greater(0).into_vec(); + let mocc_α = mo_coeff_α.bool_select(-1, &occidx_α); + let mocc_β = mo_coeff_β.bool_select(-1, &occidx_β); + + let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); + + let ([resp_α, resp_β], _timing) = get_uks_response_bra_batched( + ni_resp, + den_type, + fxc_eff.view(), + bra, + &[mocc_α.view(), mocc_β.view()], + self.verbose, + ); + [resp_α, resp_β] + } +} diff --git a/src/main_driver.rs b/src/main_driver.rs index 1afa7b5cd..5a61d634f 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -1209,15 +1209,16 @@ mod geometric_pyo3_impl { } else { use crate::analdrv::hessian::hess_interface; use crate::analdrv::response::rresp_interface::rscf_resp_interface; + use crate::analdrv::response::uresp_interface::uscf_resp_interface; + use crate::analdrv::response::RespSCF; use rstsr::prelude::*; let config = scf_data.mol.ctrl.analdrv.clone().unwrap_or_default(); - // shared RHF response object for the hessian. - // UHF builds its own internally, and will implement the UHF response interface in the future. - let mut resp_objs = if matches!(scf_data.scftype, crate::scf_io::SCFType::RHF) { - Some(rscf_resp_interface(&scf_data, &config)) - } else { - None + // shared response object for the hessian, per the SCF type + let mut resp_objs = match scf_data.scftype { + crate::scf_io::SCFType::RHF => Some(RespSCF::R(rscf_resp_interface(&scf_data, &config))), + crate::scf_io::SCFType::UHF => Some(RespSCF::U(uscf_resp_interface(&scf_data, &config))), + _ => None, }; let hess_out = hess_interface(&scf_data, &config, resp_objs.as_mut()); let hess_raw = hess_out.hessian; diff --git a/src/ri_jk/hess_u.rs b/src/ri_jk/hess_u.rs index d46142adf..7250e21ef 100644 --- a/src/ri_jk/hess_u.rs +++ b/src/ri_jk/hess_u.rs @@ -11,15 +11,11 @@ //! //! Scaling difference to the restricted ([`crate::ri_jk::hess_r::RHessRIJK`]) counterpart: //! - Skeleton: ``factor_j * de_J - factor_k * de_K`` (K coefficient ``-1``, not ``-0.5`` as in RHF), -//! because UHF ``de_K = K^alpha + K^beta`` already absorbs the spin sum (matches -//! [`crate::ri_jk::hess_u_naive::UHessRIJKNaive`]). +//! because UHF ``de_K = K^alpha + K^beta`` already absorbs the spin sum. //! - First derivative (bra form): ``factor_j * (j1ao @ mocc_s) - factor_k * k1bra_s`` per spin (no -//! ``0.5`` factor), again matching the naive UHF convention. +//! ``0.5`` factor). //! -//! The response (`get_response_bra`) reuses the separated J/K response core -//! [`crate::ri_jk::resp_r::get_rijk_response_bra_separated`] shared with RHF: J is produced once -//! in AO form from the total density response and right half-transformed per spin; K is produced -//! per spin in bra form (same-spin only). +//! The response (`URespRIJK`) lives in the separate [`crate::ri_jk::resp_u`] module. use super::prelude_dev::*; use crate::analdrv::prelude::*; @@ -30,7 +26,6 @@ use crate::ri_jk::hess_r::{ generate_cderi_with_decomp, get_rijk_skeleton_decomposed_separated, KEYS_J02, KEYS_J11, KEYS_J1AO, KEYS_J20, KEYS_K02, KEYS_K11, KEYS_K1BRA, KEYS_K20, }; -use crate::ri_jk::resp_r::get_rijk_response_bra_separated; /* #region impl */ @@ -258,58 +253,6 @@ impl<'a> UHessElecInteractAPI for UHessRIJK<'a> { self.result.insert("deriv1_bra_1", deriv1_bra[1].clone()); deriv1_bra } - - fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { - self.intmd.insert("mo_coeff_0".to_string(), mo_coeff[0].view().into_contig(RowMajor)); - self.intmd.insert("mo_coeff_1".to_string(), mo_coeff[1].view().into_contig(RowMajor)); - self.intmd.insert("mo_occ_0".to_string(), mo_occ[0].to_owned()); - self.intmd.insert("mo_occ_1".to_string(), mo_occ[1].to_owned()); - } - - fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2] { - let mo_coeff = [self.intmd["mo_coeff_0"].view(), self.intmd["mo_coeff_1"].view()]; - let mo_occ = [self.intmd["mo_occ_0"].view(), self.intmd["mo_occ_1"].view()]; - let cderi = self.cderi.view(); - let device = mo_coeff[0].device(); - // Shared separated J/K response core: J (AO form, from total density) + per-spin K (bra form). - let (j_ao, k_bras) = get_rijk_response_bra_separated( - cderi, - &mo_coeff, - &mo_occ, - bra, - self.factor_j != 0.0, - self.factor_k != 0.0, - 72, // TODO: batch size `72` should be tunable by max-memory. - ); - - let nao = mo_coeff[0].shape()[0]; - let occidx = [mo_occ[0].view().greater(0.0).into_vec(), mo_occ[1].view().greater(0.0).into_vec()]; - let mocc = [ - mo_coeff[0].view().bool_select(-1, &occidx[0]).into_contig(ColMajor), - mo_coeff[1].view().bool_select(-1, &occidx[1]).into_contig(ColMajor), - ]; - let nocc = [mocc[0].shape()[1], mocc[1].shape()[1]]; - - let mut resp = [None, None]; - for s in 0..2 { - let shape = bra[s].shape().to_vec(); - let nprop: usize = shape[2..].iter().product(); - let mut r = rt::zeros(([nao, nocc[s], nprop], device)); - // J: spin-independent AO operator, right half-transformed by this spin's mocc. - // The shared `j_ao` carries the RHF symmetrization factor (effective `4 * J1`); UHF - // naive J uses `2 * J1`, so an extra `0.5` prefactor is applied here (occ = 1 vs 2). - if let Some(j_ao) = j_ao.as_ref() { - r += 0.5 * self.factor_j * (j_ao.view() % &mocc[s]); - } - // K: same-spin bra form (UHF occ = 1, so no 0.5 factor — unlike RHF). The core already - // bakes in the exchange sign, so this is an additive contribution. - if let Some(k_bra) = k_bras.get(s) { - r += self.factor_k * k_bra.view().reshape((nao, nocc[s], nprop)); - } - resp[s] = Some(r.into_shape(shape)); - } - [resp[0].take().unwrap(), resp[1].take().unwrap()] - } } /* #endregion */ diff --git a/src/ri_jk/mod.rs b/src/ri_jk/mod.rs index 9d5b18a10..89854babb 100644 --- a/src/ri_jk/mod.rs +++ b/src/ri_jk/mod.rs @@ -37,6 +37,7 @@ pub mod hess_u; // response implementations pub mod resp_r; +pub mod resp_u; // generalized fock implementations pub mod gfock_r; diff --git a/src/ri_jk/resp_u.rs b/src/ri_jk/resp_u.rs new file mode 100644 index 000000000..0fe24b3c5 --- /dev/null +++ b/src/ri_jk/resp_u.rs @@ -0,0 +1,176 @@ +//! Response (Fock and response matrix) implementation for RI-JK, unrestricted. +//! +//! The response object [`URespRIJK`] is a standalone object holding the minimal state for the +//! fock/response functionality of the RI-JK electronic-interaction contribution of an unrestricted +//! SCF; it is fully independent of the skeleton-hessian machinery of `UHessRIJK` (which owns its +//! own copy of the factors and `cderi`). +//! +//! The response-bra core delegates to the shared separated J/K kernel +//! [`get_rijk_response_bra_separated`](super::resp_r::get_rijk_response_bra_separated) (also used +//! by the RHF response object): J is produced once in AO form from the total (α+β) density +//! response and right half-transformed per spin with a `0.5` prefactor; K is produced per spin in +//! bra form (same-spin only). + +use super::prelude_dev::*; +use crate::analdrv::prelude::*; +use crate::ri_jk::pure_incore::{get_vj_ri_incore, get_vk_ri_incore_coeff, get_vk_ri_incore_dm}; +use crate::ri_jk::resp_r::get_rijk_response_bra_separated; +use crate::ri_jk::util::get_dm0_restricted; + +/// Response (fock/response matrix) object for RI-JK, unrestricted. +/// +/// This object deliberately holds only what the response contractions need: the Cholesky +/// decomposed ERI, the J/K factors, and the response intermediates. It does **not** need +/// `mol`/`aux`/`j2c_decomp` — those belong to the skeleton-hessian machinery of the hessian +/// object. +/// +/// Main-thread only (not `Send`/`Sync`). +pub struct URespRIJK<'a> { + /// Coulomb factor, absorbed into the returned fock/response tensors. + pub factor_j: f64, + /// Exchange factor, absorbed into the returned fock/response tensors. + pub factor_k: f64, + /// Cholesky decomposed 3c-2e ERI, shape `[nao_tp, naux]`. + pub cderi: TsrCow<'a>, + /// Response intermediates: `mo_coeff_s [nao, nmo_s]` (RowMajor) and `mo_occ_s [nmo_s]`, stored + /// by [`URespAPI::make_response_preparation`]. + pub intmd: HashMap, +} + +impl<'a> URespRIJK<'a> { + /// Create from a Cholesky-decomposed ERI (`cderi`). + pub fn new_with_cderi(factor_j: f64, factor_k: f64, cderi: TsrCow<'a>) -> Self { + Self { factor_j, factor_k, cderi, intmd: HashMap::new() } + } +} + +impl<'a> AnalDrvBaseAPI for URespRIJK<'a> {} + +impl<'a> URespAPI for URespRIJK<'a> { + /// Fock per spin from spin density matrices: `factor_j * J(rdm_α + rdm_β) - factor_k * + /// K(rdm_s)`. + /// + /// The Coulomb block sees the total density; the exchange block is same-spin. Unlike the + /// restricted [`RRespRIJK`](super::resp_r::RRespRIJK) there is no `0.5` on the exchange: each + /// spin density carries occupation 1. + /// + /// The `rdm` is assumed to be symmetric, as required by the pure in-core functions. + fn get_fock_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + let [α, β] = [0, 1]; + for rdm_s in rdm { + assert_eq!(rdm_s.ndim(), 2, "rdm must have 2 dimensions"); + } + let nao = rdm[α].shape()[0]; + let device = self.cderi.device(); + + let mut fock = [rt::zeros(([nao, nao], device)), rt::zeros(([nao, nao], device))]; + if self.factor_j != 0.0 { + // J: total (α+β) density, spin-independent + let dm_total = &rdm[α] + &rdm[β]; + let dms = dm_total.into_contig(ColMajor).into_shape((nao, nao, 1)); + let vj = get_vj_ri_incore(self.cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); + fock[α] += self.factor_j * &vj; + fock[β] += self.factor_j * &vj; + } + if self.factor_k != 0.0 { + // K: same-spin exchange + for (s, fock_s) in fock.iter_mut().enumerate() { + let dms = rdm[s].view().into_contig(ColMajor).into_shape((nao, nao, 1)); + // TODO: batch size `72` should be tunable by max-memory. + let vk = get_vk_ri_incore_dm(self.cderi.view(), dms.view(), 72).i((.., .., 0)).into_contig(ColMajor); + *fock_s -= self.factor_k * &vk; + } + } + fock + } + + /// Fock per spin from molecular coefficients: J by the dm route (no coeff-native vj exists), K + /// by the coeff-native pure function per spin. + fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) -> [Tsr; 2] { + let [α, β] = [0, 1]; + let nao = mo_coeff[α].shape()[0]; + let device = self.cderi.device(); + + let mut fock = [rt::zeros(([nao, nao], device)), rt::zeros(([nao, nao], device))]; + if self.factor_j != 0.0 { + // J: total (α+β) density, spin-independent + let dm_total = get_dm0_restricted(mo_coeff[α].view(), mo_occ[α].view()) + + get_dm0_restricted(mo_coeff[β].view(), mo_occ[β].view()); + let dms = dm_total.into_contig(ColMajor).into_shape((nao, nao, 1)); + let vj = get_vj_ri_incore(self.cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); + fock[α] += self.factor_j * &vj; + fock[β] += self.factor_j * &vj; + } + if self.factor_k != 0.0 { + // K: same-spin exchange, coeff-native + for (s, fock_s) in fock.iter_mut().enumerate() { + let nmo = mo_coeff[s].shape()[1]; + let coeff = mo_coeff[s].view().into_contig(ColMajor).into_shape((nao, nmo, 1)); + let occ = mo_occ[s].view().into_contig(ColMajor).into_shape((nmo, 1)); + // TODO: batch size `72` should be tunable by max-memory. + let vk = + get_vk_ri_incore_coeff(self.cderi.view(), coeff.view(), occ.view(), 72).i((.., .., 0)).into_contig(ColMajor); + *fock_s -= self.factor_k * &vk; + } + } + fock + } + + fn get_response_rdm(&mut self, _rdm: &[TsrView; 2]) -> [Tsr; 2] { + // rdm-form entry of the unrestricted response kernel; reserved for the future ugfock + // machinery (no consumer yet), mirroring the restricted side before RI-PT2 + unimplemented!("get_response_rdm is not implemented for URespRIJK; reserved for the future ugfock machinery.") + } + + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { + self.intmd.insert("mo_coeff_0".to_string(), mo_coeff[0].view().into_contig(RowMajor)); + self.intmd.insert("mo_coeff_1".to_string(), mo_coeff[1].view().into_contig(RowMajor)); + self.intmd.insert("mo_occ_0".to_string(), mo_occ[0].to_owned()); + self.intmd.insert("mo_occ_1".to_string(), mo_occ[1].to_owned()); + } + + fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2] { + let mo_coeff = [self.intmd["mo_coeff_0"].view(), self.intmd["mo_coeff_1"].view()]; + let mo_occ = [self.intmd["mo_occ_0"].view(), self.intmd["mo_occ_1"].view()]; + let cderi = self.cderi.view(); + let device = mo_coeff[0].device(); + // Shared separated J/K response core: J (AO form, from total density) + per-spin K (bra form). + let (j_ao, k_bras) = get_rijk_response_bra_separated( + cderi, + &mo_coeff, + &mo_occ, + bra, + self.factor_j != 0.0, + self.factor_k != 0.0, + 72, // TODO: batch size `72` should be tunable by max-memory. + ); + + let nao = mo_coeff[0].shape()[0]; + let occidx = [mo_occ[0].view().greater(0.0).into_vec(), mo_occ[1].view().greater(0.0).into_vec()]; + let mocc = [ + mo_coeff[0].view().bool_select(-1, &occidx[0]).into_contig(ColMajor), + mo_coeff[1].view().bool_select(-1, &occidx[1]).into_contig(ColMajor), + ]; + let nocc = [mocc[0].shape()[1], mocc[1].shape()[1]]; + + let mut resp = [None, None]; + for s in 0..2 { + let shape = bra[s].shape().to_vec(); + let nprop: usize = shape[2..].iter().product(); + let mut r = rt::zeros(([nao, nocc[s], nprop], device)); + // J: spin-independent AO operator, right half-transformed by this spin's mocc. + // The shared `j_ao` carries the RHF symmetrization factor (effective `4 * J1`); UHF + // naive J uses `2 * J1`, so an extra `0.5` prefactor is applied here (occ = 1 vs 2). + if let Some(j_ao) = j_ao.as_ref() { + r += 0.5 * self.factor_j * (j_ao.view() % &mocc[s]); + } + // K: same-spin bra form (UHF occ = 1, so no 0.5 factor — unlike RHF). The core already + // bakes in the exchange sign, so this is an additive contribution. + if let Some(k_bra) = k_bras.get(s) { + r += self.factor_k * k_bra.view().reshape((nao, nocc[s], nprop)); + } + resp[s] = Some(r.into_shape(shape)); + } + [resp[0].take().unwrap(), resp[1].take().unwrap()] + } +} diff --git a/tests/analdrv/hessian/uhf.rs b/tests/analdrv/hessian/uhf.rs index e42f5bcc7..6aed87a12 100644 --- a/tests/analdrv/hessian/uhf.rs +++ b/tests/analdrv/hessian/uhf.rs @@ -1,4 +1,5 @@ use pyrest::analdrv::hessian::uscf_interface::uscf_hess_interface; +use pyrest::analdrv::response::uresp_interface::uscf_resp_interface; use pyrest::analdrv::config::AnalDrvConfig; use pyrest::ctrl_io; @@ -42,7 +43,8 @@ fn test_nh3() { scf_without_build(&mut scf_data, &None); let config = AnalDrvConfig::default(); - let (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config.nucgrad, &config.resp); + let mut resp_obj = uscf_resp_interface(&scf_data, &config); + let (de, _vib, _th) = uscf_hess_interface(&scf_data, &config.nucgrad, &mut resp_obj); let natm = 3; let de = rt::asarray((&de, [3, 3, natm, natm])); diff --git a/tests/analdrv/hessian/uks_grid_shift.rs b/tests/analdrv/hessian/uks_grid_shift.rs index 9f665f2c9..ee7064726 100644 --- a/tests/analdrv/hessian/uks_grid_shift.rs +++ b/tests/analdrv/hessian/uks_grid_shift.rs @@ -4,6 +4,7 @@ use pyrest::analdrv::config::{AnalDrvConfig, AnalDrvNucgradCfg}; use pyrest::analdrv::hessian::uscf_interface::uscf_hess_interface; +use pyrest::analdrv::response::uresp_interface::uscf_resp_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; use pyrest::scf_io::{self, scf_without_build}; @@ -69,7 +70,8 @@ fn run_scf(input: &str) -> scf_io::SCF { } fn run_hessian(scf_data: &mut scf_io::SCF, config: &AnalDrvConfig) -> Tensor { - let (de, _, _) = uscf_hess_interface(scf_data, &config.nucgrad, &config.resp); + let mut resp_obj = uscf_resp_interface(scf_data, config); + let (de, _, _) = uscf_hess_interface(scf_data, &config.nucgrad, &mut resp_obj); let natm = 4; rt::asarray((de, [3, 3, natm, natm], &DeviceBLAS::default())) } diff --git a/tests/analdrv/hessian/uks_tpss0_grid_levels.rs b/tests/analdrv/hessian/uks_tpss0_grid_levels.rs index 020675f17..031f0f12d 100644 --- a/tests/analdrv/hessian/uks_tpss0_grid_levels.rs +++ b/tests/analdrv/hessian/uks_tpss0_grid_levels.rs @@ -5,6 +5,7 @@ use pyrest::analdrv::config::{AnalDrvConfig, AnalDrvNucgradCfg, AnalDrvRespCfg}; use pyrest::analdrv::hessian::uscf_interface::uscf_hess_interface; +use pyrest::analdrv::response::uresp_interface::uscf_resp_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; @@ -42,13 +43,14 @@ fn run_with_config(config: AnalDrvConfig) -> Vec { 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 (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config.nucgrad, &config.resp); + let mut resp_obj = uscf_resp_interface(&scf_data, &config); + let (de, _vib, _th) = uscf_hess_interface(&scf_data, &config.nucgrad, &mut resp_obj); de } #[test] fn test_nh3_mgga_default_skeleton() { - // Default skeleton level for MGGA = grid_gen_level + 2 = 5; cpscf = 1 -> ni_cpks path. + // Default skeleton level for MGGA = grid_gen_level + 2 = 5; response grid = 1 -> ni_resp path. let config = AnalDrvConfig::default(); let de = run_with_config(config); diff --git a/tests/analdrv/hessian/uks_tpssh.rs b/tests/analdrv/hessian/uks_tpssh.rs index 958dd5b2e..a34ba588a 100644 --- a/tests/analdrv/hessian/uks_tpssh.rs +++ b/tests/analdrv/hessian/uks_tpssh.rs @@ -1,5 +1,6 @@ use pyrest::analdrv::config::AnalDrvConfig; use pyrest::analdrv::hessian::uscf_interface::uscf_hess_interface; +use pyrest::analdrv::response::uresp_interface::uscf_resp_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; @@ -41,7 +42,8 @@ fn test_nh3() { scf_without_build(&mut scf_data, &None); let config = AnalDrvConfig::default(); - let (de, vib, _th) = uscf_hess_interface(&mut scf_data, &config.nucgrad, &config.resp); + let mut resp_obj = uscf_resp_interface(&scf_data, &config); + let (de, vib, _th) = uscf_hess_interface(&scf_data, &config.nucgrad, &mut resp_obj); let natm = 4; let de = rt::asarray((&de, [3, 3, natm, natm])); -- Gitee From 159ecfe448baf7046ff8583476ff48c502252598 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Tue, 22 Sep 2026 16:22:06 +0800 Subject: [PATCH 02/21] Add unrestricted SCF multipole moments and U response rdm form Electric multipole moments now support UHF/UKS calculations: - New UMultipoleDH (src/analdrv/multipole/umultipole.rs), the unrestricted sibling of RMultipoleDH: nuclear plus SCF-density (alpha+beta) contractions for orders 1-4 with the traceless quadrupole, at SCF level only (unrestricted post-SCF fifth-DFA increments are rejected at the interface). - multipole_interface dispatches on the SCF type: the restricted arm (RHF/RKS plus the PT2-family relaxed/unrelaxed DH increments) and a new unrestricted arm (SCF-level UHF/UKS); ROHF is rejected, and order validation is shared. Shared helpers in rmultipole become pub(crate); UMultipoleDH is exported through the analdrv prelude. - Implement the rdm-form response kernels consumed by the task (and reserved for the future ugfock machinery): URespRIJK::get_response_rdm (2 J(rdm_a + rdm_b) - 2 K(rdm_s) per spin, closed-shell limit matching the restricted 4J - 2K) and URespKSNIMatmul::get_response_rdm (2 fxc per spin on the common grid, cached as fxc_common_grid). - Tests: NH2 radical UHF/UKS multipoles against pyscf, the NH2 interface-level task run (incl. post-SCF rejection), and the unrestricted rdm-form kernels pinned against the bra form and the closed-shell limit. - Drop machine-local reference paths from analdrv test provenance notes, keeping the tool and version information. Co-authored-by: Claude Code Co-authored-by: glm-5.3-flash --- src/analdrv/config.rs | 10 +- src/analdrv/mod.rs | 5 +- src/analdrv/multipole/interface.rs | 159 ++++++++++----- src/analdrv/multipole/mod.rs | 13 +- src/analdrv/multipole/rmultipole.rs | 4 +- src/analdrv/multipole/umultipole.rs | 197 +++++++++++++++++++ src/analdrv/response/trait_uresp.rs | 7 +- src/dft/numint_matmul/resp_uks.rs | 68 ++++++- src/ri_jk/resp_u.rs | 46 ++++- tests/analdrv/multipole/co2_hf.rs | 3 +- tests/analdrv/multipole/h2o_hf.rs | 17 +- tests/analdrv/multipole/h2o_rks.rs | 4 +- tests/analdrv/multipole/hi_rmp2.rs | 3 +- tests/analdrv/multipole/interface_nh2_uhf.rs | 139 +++++++++++++ tests/analdrv/multipole/mod.rs | 3 + tests/analdrv/multipole/nh2_uhf.rs | 121 ++++++++++++ tests/analdrv/multipole/nh2_uks.rs | 79 ++++++++ tests/analdrv/multipole/nh3_xyg3.rs | 9 +- tests/analdrv/response/mod.rs | 1 + tests/analdrv/response/uresp_gfock.rs | 163 +++++++++++++++ tests/analdrv/response/xyg3_rgfock.rs | 3 +- 21 files changed, 962 insertions(+), 92 deletions(-) create mode 100644 src/analdrv/multipole/umultipole.rs create mode 100644 tests/analdrv/multipole/interface_nh2_uhf.rs create mode 100644 tests/analdrv/multipole/nh2_uhf.rs create mode 100644 tests/analdrv/multipole/nh2_uks.rs create mode 100644 tests/analdrv/response/uresp_gfock.rs diff --git a/src/analdrv/config.rs b/src/analdrv/config.rs index aefd1d00a..42533b40b 100644 --- a/src/analdrv/config.rs +++ b/src/analdrv/config.rs @@ -166,7 +166,9 @@ impl Default for AnalDrvNucgradCfg { /// Relaxation treatment of the double-hybrid (DH) density increments for the multipole moments. /// -/// Only meaningful for PT2-family post-SCF (fifth-DFA) methods; silently ignored otherwise. +/// Only meaningful for PT2-family post-SCF (fifth-DFA) restricted methods; silently ignored +/// otherwise (the unrestricted multipole task is SCF-level only — post-SCF unrestricted methods +/// are rejected there). #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum MultipoleRdm1Relax { /// Relaxed density: solve the Z-vector (CP-SCF) and include the response increment. @@ -186,8 +188,10 @@ impl Default for MultipoleRdm1Relax { /// Settings of the electric multipole moment evaluation. /// /// These keywords control what is evaluated in the [`Multipole`](AnalDrvTask::Multipole) task. -/// All moments are evaluated in atomic units; the default origin is the coordinate origin -/// `[0, 0, 0]` (Bohr), the same print convention as Gaussian and pyscf. +/// Supported for RHF/RKS and UHF/UKS calculations, with the PT2-family post-SCF (fifth-DFA) +/// density increments restricted-only. All moments are evaluated in atomic units; the default +/// origin is the coordinate origin `[0, 0, 0]` (Bohr), the same print convention as Gaussian +/// and pyscf. #[serde_inline_default] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AnalDrvMultipoleCfg { diff --git a/src/analdrv/mod.rs b/src/analdrv/mod.rs index f13dc6556..9fa270447 100644 --- a/src/analdrv/mod.rs +++ b/src/analdrv/mod.rs @@ -16,8 +16,8 @@ //! This module should work in most cases, but still requires further testing and efficiency update. //! //! The nuclear-coordinate derivative properties (hessian) are currently SCF-level only; the -//! electric multipole moments additionally support the PT2-family post-SCF (fifth-DFA) methods -//! through the DH density increments. +//! electric multipole moments support restricted and unrestricted SCF, and additionally the +//! PT2-family post-SCF (fifth-DFA) methods (restricted only) through the DH density increments. //! //! This module does not contain extensive detailed implementation. Please refer to the [`hessian`] //! submodule for hessian traits, component implementations, total hessian drivers (including @@ -78,6 +78,7 @@ pub mod prelude { pub use multipole::nuc_charge::MultipoleNucCharge; pub use multipole::rmultipole::RMultipoleDH; pub use multipole::trait_multipole::MultipoleNucAPI; + pub use multipole::umultipole::UMultipoleDH; pub use response::rresp_interface::RRespSCF; pub use response::trait_rresp::RRespAPI; pub use response::trait_uresp::URespAPI; diff --git a/src/analdrv/multipole/interface.rs b/src/analdrv/multipole/interface.rs index 9a2977d18..0ca7540f9 100644 --- a/src/analdrv/multipole/interface.rs +++ b/src/analdrv/multipole/interface.rs @@ -1,7 +1,9 @@ //! Interface of the multipole task to the analdrv driver. -use crate::analdrv::config::{AnalDrvConfig, MultipoleRdm1Relax}; +use crate::analdrv::config::{AnalDrvConfig, AnalDrvMultipoleCfg, MultipoleRdm1Relax}; use crate::analdrv::multipole::rmultipole::RMultipoleDH; +use crate::analdrv::multipole::umultipole::UMultipoleDH; +use crate::analdrv::prelude::Tsr; use crate::analdrv::response::rgfock_interface::rgfock_dh_interface; use crate::analdrv::response::RespSCF; use crate::ri_jk::util::get_cint_mol; @@ -10,6 +12,7 @@ use crate::utilities::rstsr_util::RestTensorToRstsrTsrAPI; use rstsr::prelude::*; use serde::Serialize; +use std::collections::HashMap; /// Moments of one multipole order, split into contributions. /// @@ -55,30 +58,28 @@ pub struct MultipoleOutput { pub hexadecapole: Option, } -/// Run the electric multipole moment task of the analdrv driver. +/// Run the electric multipole moment task of the analdrv driver, dispatching on the SCF type. /// -/// For usual SCF methods (HF/DFT/hybrid-DFT, where the energy functional is the SCF functional), -/// neither the DH composite nor the response object is given to the driver: the moments are the -/// plain nuclear plus SCF-density contractions. For PT2-family post-SCF (fifth-DFA) methods, the -/// unrelaxed correlation rdm1 increment is added, and, when `resp_objs` is passed, the relaxed -/// (Z-vector) increment as well; whether `resp_objs` is passed is decided by the caller from -/// `multipole_rdm1_relax`. -/// -/// `resp_objs` is the [`RespSCF`] carrier (the multipole task itself is restricted-only for now; -/// the unrestricted variant is rejected here until the U arm is wired). It is only used for the -/// relaxed DH increments, and requires the DFT grids to be present; the caller -/// ([`crate::analdrv::interface::analdrv_interface`]) regenerates them. +/// The restricted arm ([`multipole_interface_r`]) covers RHF/RKS and the PT2-family post-SCF +/// (fifth-DFA) methods: for usual SCF methods neither the DH composite nor the response object +/// is given to the driver, while for post-SCF methods the unrelaxed correlation rdm1 increment +/// is added, and, when `resp_objs` is passed, the relaxed (Z-vector) increment as well; whether +/// `resp_objs` is passed is decided by the caller from `multipole_rdm1_relax`. The unrestricted +/// arm ([`multipole_interface_u`]) covers SCF-level UHF/UKS only. ROHF is not supported. pub fn multipole_interface<'a>( scf_data: &'a SCF, config: &AnalDrvConfig, resp_objs: Option<&mut RespSCF<'a>>, ) -> MultipoleOutput { match scf_data.scftype { - SCFType::RHF => {}, - _ => panic!("Multipole evaluation is currently only implemented for restricted (RHF/RKS) calculations."), + SCFType::RHF => multipole_interface_r(scf_data, config, resp_objs), + SCFType::UHF => multipole_interface_u(scf_data, config, resp_objs), + SCFType::ROHF => panic!("Multipole evaluation is not implemented for ROHF calculations."), } +} - let mp_cfg = &config.multipole; +/// Validate `multipole_orders` and return it. +fn checked_orders(mp_cfg: &AnalDrvMultipoleCfg) -> &[usize] { assert!(!mp_cfg.orders.is_empty(), "multipole_orders cannot be empty"); for &order in &mp_cfg.orders { assert!( @@ -86,7 +87,23 @@ pub fn multipole_interface<'a>( "multipole_orders entries must be 1 (dipole) to 4 (hexadecapole), got {order}" ); } + &mp_cfg.orders +} +/// Restricted arm of [`multipole_interface`]: RHF/RKS (SCF level) and PT2-family post-SCF +/// (fifth-DFA) methods. +/// +/// `resp_objs` is the [`RespSCF`] carrier, used only for the relaxed DH increments (the U +/// variant of the carrier is rejected here; the unrestricted task has its own arm and consumes +/// no response object). It requires the DFT grids to be present; the caller +/// ([`crate::analdrv::interface::analdrv_interface`]) regenerates them. +fn multipole_interface_r<'a>( + scf_data: &'a SCF, + config: &AnalDrvConfig, + resp_objs: Option<&mut RespSCF<'a>>, +) -> MultipoleOutput { + let mp_cfg = &config.multipole; + let orders = checked_orders(mp_cfg); let origin = mp_cfg.origin.unwrap_or([0.0; 3]); let device = DeviceBLAS::default(); @@ -110,8 +127,7 @@ pub fn multipole_interface<'a>( }; // The relaxed (Z-vector) increments consume the restricted response object; the U variant of - // the carrier is rejected here — the unrestricted multipole task is not wired yet (the task - // is also restricted-only at the SCF-type check above). + // the carrier cannot appear here (the unrestricted task never requests a response object). let resp = resp_objs.map(|resp| match resp { RespSCF::R(r) => r, RespSCF::U(_) => panic!( @@ -123,22 +139,14 @@ pub fn multipole_interface<'a>( // evaluate the requested orders (each order caches itself; the DH density increments are // shared between the orders through the gfock result cache) - for &order in &mp_cfg.orders { + for &order in orders { match order { - 1 => { - rmultipole.make_dipole(); - }, - 2 => { - rmultipole.make_quadrupole(); - }, - 3 => { - rmultipole.make_octupole(); - }, - 4 => { - rmultipole.make_hexadecapole(); - }, + 1 => rmultipole.make_dipole(), + 2 => rmultipole.make_quadrupole(), + 3 => rmultipole.make_octupole(), + 4 => rmultipole.make_hexadecapole(), _ => unreachable!("multipole order range checked above"), - } + }; } let print_level = scf_data.mol.ctrl.print_level; @@ -187,29 +195,88 @@ pub fn multipole_interface<'a>( scf_data.fchk_append_density_section("Total MP2 Density", &packed); println!( " total density: {}", - if relaxed { - "relaxed (SCF + corr. + Z-vector response)" - } else { - "unrelaxed (SCF + corr.)" - } + if relaxed { "relaxed (SCF + corr. + Z-vector response)" } else { "unrelaxed (SCF + corr.)" } ); } } MultipoleOutput { origin, - dipole: order_parts(&rmultipole, "dip", false), - quadrupole: order_parts(&rmultipole, "quad", true), - octupole: order_parts(&rmultipole, "oct", false), - hexadecapole: order_parts(&rmultipole, "hex", false), + dipole: order_parts(&rmultipole.result, "dip", false), + quadrupole: order_parts(&rmultipole.result, "quad", true), + octupole: order_parts(&rmultipole.result, "oct", false), + hexadecapole: order_parts(&rmultipole.result, "hex", false), + } +} + +/// Unrestricted arm of [`multipole_interface`]: SCF-level UHF/UKS only. +/// +/// The post-SCF (fifth-DFA) density increments are not implemented for unrestricted methods — +/// both the unrelaxed correlation rdm1 and the relaxed (Z-vector) increment would need the +/// unrestricted PT2 generalized-Fock machinery — so fifth-DFA calculations are rejected +/// outright. No response object is consumed: `multipole_rdm1_relax` is silently ignored +/// (SCF-level moments carry no relaxation), and `multipole_rdm1_dump` is silently ignored +/// (post-SCF restricted methods only). +fn multipole_interface_u<'a>( + scf_data: &'a SCF, + config: &AnalDrvConfig, + resp_objs: Option<&mut RespSCF<'a>>, +) -> MultipoleOutput { + if scf_data.mol.xc_data.is_fifth_dfa() { + panic!( + "Multipole evaluation is not implemented for unrestricted post-SCF (double-hybrid, \ + fifth-DFA) methods yet; only SCF-level UHF/UKS is supported." + ); + } + if resp_objs.is_some() { + panic!( + "The unrestricted multipole evaluation does not consume a response object; do not pass `resp_objs` for UHF/UKS calculations." + ); + } + + let mp_cfg = &config.multipole; + let orders = checked_orders(mp_cfg); + let origin = mp_cfg.origin.unwrap_or([0.0; 3]); + + let device = DeviceBLAS::default(); + let mo_coeff = [(&scf_data.eigenvectors[0]).to_rstsr(&device), (&scf_data.eigenvectors[1]).to_rstsr(&device)]; + let mo_occ = [(&scf_data.occupation[0]).to_rstsr(&device), (&scf_data.occupation[1]).to_rstsr(&device)]; + let mol_cint = get_cint_mol(&scf_data.mol); + + let mut umultipole = UMultipoleDH::new(&mol_cint, mo_coeff, mo_occ, origin); + + // evaluate the requested orders (each order caches itself) + for &order in orders { + match order { + 1 => umultipole.make_dipole(), + 2 => umultipole.make_quadrupole(), + 3 => umultipole.make_octupole(), + 4 => umultipole.make_hexadecapole(), + _ => unreachable!("multipole order range checked above"), + }; + } + + let print_level = scf_data.mol.ctrl.print_level; + umultipole.print_multipole(print_level); + if print_level >= 2 { + for (label, t) in &umultipole.timing { + println!("Timing {label}: {t:.3} s"); + } + } + + MultipoleOutput { + origin, + dipole: order_parts(&umultipole.result, "dip", false), + quadrupole: order_parts(&umultipole.result, "quad", true), + octupole: order_parts(&umultipole.result, "oct", false), + hexadecapole: order_parts(&umultipole.result, "hex", false), } } -/// Assemble the per-order output from the driver result map; `None` if the order was not +/// Assemble the per-order output from a driver result map; `None` if the order was not /// evaluated. `traceless` includes the `quad_tot_traceless` entry (quadrupole only). -fn order_parts(driver: &RMultipoleDH, prefix: &str, traceless: bool) -> Option { - let flat = - |suffix: &str| driver.result.get(&format!("{prefix}_{suffix}")).map(|v| v.view().into_shape(-1).into_vec()); +fn order_parts(result: &HashMap, prefix: &str, traceless: bool) -> Option { + let flat = |suffix: &str| result.get(&format!("{prefix}_{suffix}")).map(|v| v.view().into_shape(-1).into_vec()); Some(MultipoleOrderParts { nuc: flat("nuc")?, scf: flat("scf")?, @@ -217,7 +284,7 @@ fn order_parts(driver: &RMultipoleDH, prefix: &str, traceless: bool) -> Option Tsr { } /// Intor name and `result`/`timing` key prefix of a multipole order. -fn multipole_order_spec(order: usize) -> (&'static str, &'static str) { +pub(crate) fn multipole_order_spec(order: usize) -> (&'static str, &'static str) { match order { 1 => ("int1e_r", "dip"), 2 => ("int1e_rr", "quad"), @@ -64,7 +64,7 @@ fn multipole_order_spec(order: usize) -> (&'static str, &'static str) { /// Sorted (non-decreasing) Cartesian multi-indices enumerating the unique components of a /// permutation-symmetric order-`order` moment tensor (6 for order 2, 10 for order 3, 15 for /// order 4). -fn sorted_multi_indices(order: usize) -> Vec> { +pub(crate) fn sorted_multi_indices(order: usize) -> Vec> { let mut out = vec![vec![]]; for _ in 0..order { let mut next = Vec::new(); diff --git a/src/analdrv/multipole/umultipole.rs b/src/analdrv/multipole/umultipole.rs new file mode 100644 index 000000000..00eec820b --- /dev/null +++ b/src/analdrv/multipole/umultipole.rs @@ -0,0 +1,197 @@ +//! Multipole moment driver for unrestricted SCF methods. + +use super::rmultipole::{multipole_intor, multipole_order_spec, quadrupole_to_traceless, sorted_multi_indices}; +use crate::analdrv::prelude::*; + +/// Working solver and maintainer of multipole moment evaluation for unrestricted SCF methods. +/// +/// The unrestricted sibling of [`RMultipoleDH`](super::rmultipole::RMultipoleDH), at the SCF +/// level only: the moments are the nuclear plus SCF-density (α+β) contractions. The post-SCF +/// (fifth-DFA) density increments are not implemented for unrestricted methods and are rejected +/// at the interface, so this driver holds no gfock/response objects; the `DH` of the name only +/// mirrors the restricted driver, where the corresponding fields will slot in when the +/// unrestricted double-hybrid machinery exists. The orders 1-4 (dipole, quadrupole, octupole, +/// hexadecapole) share a single generic evaluation core, and the traceless quadrupole is stored +/// alongside the raw second moments. +/// +/// # Structure +/// +/// - `mol` : the molecule object, for the electronic multipole integrals. +/// - `mo_coeff` : shape `[nao, nmo_s]` per spin; `mo_occ` : shape `[nmo_s]` per spin. Converged +/// orbitals of the SCF-iteration functional (`nmo_α` and `nmo_β` may differ). Orbital +/// energies are not needed for property evaluation, and are not stored. +/// - `origin` : the origin of the multipole evaluation, fixed at construction (default `[0.0; 3]`). +/// - `nuc_obj` : the non-electronic (nuclear charge) contribution object. +/// - `result` : evaluated contributions, keyed `_` with prefixes `dip`/`quad`/ +/// `oct`/`hex` and parts `nuc`/`scf`/`tot` (plus `quad_tot_traceless`). +/// - `timing` : wall-time information, keyed similarly to `result`. +pub struct UMultipoleDH { + pub mol: CInt, + pub mo_coeff: [Tsr; 2], + pub mo_occ: [Tsr; 2], + pub origin: [f64; 3], + pub nuc_obj: MultipoleNucCharge, + pub result: HashMap, + /// Timing information. Represented by wall time in second. + pub timing: Vec<(String, f64)>, +} + +impl UMultipoleDH { + /// Create the multipole driver. + /// + /// # Parameters + /// + /// - `mol` : [`CInt`]. The molecule object. + /// - `mo_coeff` : shape `[nao, nmo_s]` per spin. Converged SCF molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo_s]` per spin. Converged SCF occupation numbers. + /// - `origin` : the origin of the multipole evaluation. + pub fn new(mol: &CInt, mo_coeff: [Tsr; 2], mo_occ: [Tsr; 2], origin: [f64; 3]) -> Self { + let device = mo_coeff[0].device().clone(); + Self { + mol: mol.clone(), + mo_coeff, + mo_occ, + origin, + nuc_obj: MultipoleNucCharge::new(mol, &device), + result: HashMap::new(), + timing: Vec::new(), + } + } + + /// Generic evaluation core of the raw Cartesian moment of the given order (1 = dipole, 2 = + /// quadrupole, 3 = octupole, 4 = hexadecapole): the nuclear and the SCF-density (α+β) + /// contributions, stored in `result`. + /// + /// Cached on first call per order (the origin and the orbitals are fixed at construction); + /// repeated evaluation of the same order directly returns the stored total. + fn make_multipole_core(&mut self, order: usize) -> Tsr { + let [α, β] = [0, 1]; + let (intor, prefix) = multipole_order_spec(order); + if let Some(tot) = self.result.get(&format!("{prefix}_tot")) { + return tot.to_owned(); + } + let device = self.mo_coeff[α].device().clone(); + + // multipole integrals at the fixed origin: shape [nao, nao, 3, ..., 3] (`order` axes of 3) + let t0 = std::time::Instant::now(); + let int = multipole_intor(&mut self.mol, intor, self.origin, &device); + self.timing.push((format!("{prefix}_intor"), t0.elapsed().as_secs_f64())); + + // nuclear contribution + let t0 = std::time::Instant::now(); + let m_nuc = match order { + 1 => self.nuc_obj.make_dipole_nuc(self.origin), + 2 => self.nuc_obj.make_quadrupole_nuc(self.origin), + 3 => self.nuc_obj.make_octupole_nuc(self.origin), + 4 => self.nuc_obj.make_hexadecapole_nuc(self.origin), + _ => unreachable!("order range checked in multipole_order_spec"), + }; + self.timing.push((format!("{prefix}_nuc"), t0.elapsed().as_secs_f64())); + + // SCF density contribution (the total α+β density) + let t0 = std::time::Instant::now(); + let dm0 = get_dm0_restricted(self.mo_coeff[α].view(), self.mo_occ[α].view()) + + get_dm0_restricted(self.mo_coeff[β].view(), self.mo_occ[β].view()); + let m_scf = -(&int * &dm0).sum_axes([0, 1]); + self.timing.push((format!("{prefix}_scf"), t0.elapsed().as_secs_f64())); + + let m_tot = &m_nuc + &m_scf; + self.result.insert(format!("{prefix}_nuc"), m_nuc); + self.result.insert(format!("{prefix}_scf"), m_scf); + + if order == 2 { + let m_tot_traceless = quadrupole_to_traceless(&m_tot); + self.result.insert("quad_tot_traceless".to_string(), m_tot_traceless); + } + self.result.insert(format!("{prefix}_tot"), m_tot.clone()); + m_tot + } + + /// Evaluate the dipole moment: all contributions, stored in `result`. + /// + /// # Returns + /// + /// - `dip_tot` : shape `[3]`. The total dipole moment, the sum of all contributions. + pub fn make_dipole(&mut self) -> Tsr { + self.make_multipole_core(1) + } + + /// Evaluate the raw quadrupole moment: all contributions, stored in `result` (together with + /// the traceless total under `quad_tot_traceless`). + /// + /// # Returns + /// + /// - `quad_tot` : shape `[3, 3]`. The total raw (second-moment) quadrupole moment. + pub fn make_quadrupole(&mut self) -> Tsr { + self.make_multipole_core(2) + } + + /// Evaluate the raw octupole moment: all contributions, stored in `result`. + /// + /// # Returns + /// + /// - `oct_tot` : shape `[3, 3, 3]`. The total raw octupole moment. + pub fn make_octupole(&mut self) -> Tsr { + self.make_multipole_core(3) + } + + /// Evaluate the raw hexadecapole moment: all contributions, stored in `result`. + /// + /// # Returns + /// + /// - `hex_tot` : shape `[3, 3, 3, 3]`. The total raw hexadecapole moment. + pub fn make_hexadecapole(&mut self) -> Tsr { + self.make_multipole_core(4) + } + + /// Print the evaluated multipole moments in the analdrv output style. + /// + /// Only sections whose entries exist in `result` are printed (so a dipole-only evaluation + /// does not print an empty quadrupole section). All quantities are in atomic units. The + /// print is gated by `print_level >= 1` (the timing output stays at `>= 2`). + pub fn print_multipole(&self, print_level: usize) { + if print_level < 1 { + return; + } + println!("=============== Multipole Moments (in analdrv) ==============="); + println!("Origin (a.u.): [{:14.8}, {:14.8}, {:14.8}]", self.origin[0], self.origin[1], self.origin[2]); + + if self.result.contains_key("dip_tot") { + println!("Electric dipole moment (a.u.):"); + for (key, label) in [("dip_nuc", "nuclear charges"), ("dip_scf", "SCF density"), ("dip_tot", "total")] { + if let Some(v) = self.result.get(key) { + println!(" {:<24}: {:16.12}", label, v); + } + } + // total dipole additionally in Debye (1 a.u. = 2.54174623 Debye) + if let Some(v) = self.result.get("dip_tot") { + let au2debye = crate::constants::AU2DEBYE; + println!( + " {:<24}: {:16.12} {:16.12} {:16.12}", + "total, Debye (X Y Z)", + v[[0]] * au2debye, + v[[1]] * au2debye, + v[[2]] * au2debye + ); + } + } + + self.print_moment_section("Electric quadrupole moment, raw second moments (a.u.)", "quad_tot", 2); + self.print_moment_section("Electric quadrupole moment, traceless (a.u.)", "quad_tot_traceless", 2); + self.print_moment_section("Electric octupole moment (a.u.)", "oct_tot", 3); + self.print_moment_section("Electric hexadecapole moment (a.u.)", "hex_tot", 4); + } + + /// Print one moment section from the `result` entry `key`, listing the unique components of + /// the permutation-symmetric order-`order` tensor. Does nothing if the key is absent. + fn print_moment_section(&self, title: &str, key: &str, order: usize) { + let Some(tot) = self.result.get(key) else { return }; + let flat = tot.clone().into_shape(-1); + println!("{title}:"); + for idx in sorted_multi_indices(order) { + let label: String = idx.iter().map(|&t| ['X', 'Y', 'Z'][t]).collect(); + let c = idx.iter().enumerate().fold(0usize, |acc, (k, &t)| acc + t * 3usize.pow((order - 1 - k) as u32)); + println!(" {label}: {:16.12}", flat[[c]]); + } + } +} diff --git a/src/analdrv/response/trait_uresp.rs b/src/analdrv/response/trait_uresp.rs index 95f895d71..93972c48d 100644 --- a/src/analdrv/response/trait_uresp.rs +++ b/src/analdrv/response/trait_uresp.rs @@ -103,10 +103,9 @@ pub trait URespAPI: AnalDrvBaseAPI { /// /// # Reserved for ugfock /// - /// No unrestricted implementation provides this rdm-form entry yet; the method is part of the - /// trait surface reserved for the future unrestricted generalized-Fock machinery, mirroring - /// [`RRespAPI::get_response_rdm`] (which stayed unimplemented until the RI-PT2 generalized - /// Fock needed it). + /// No unrestricted driver consumes this rdm-form entry yet; it is the form required by the + /// future unrestricted generalized-Fock machinery, mirroring how the restricted + /// [`RRespAPI::get_response_rdm`] serves the RI-PT2 generalized Fock. /// /// [`make_response_preparation`]: Self::make_response_preparation fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2]; diff --git a/src/dft/numint_matmul/resp_uks.rs b/src/dft/numint_matmul/resp_uks.rs index 17afa3d62..08e5998df 100644 --- a/src/dft/numint_matmul/resp_uks.rs +++ b/src/dft/numint_matmul/resp_uks.rs @@ -254,10 +254,70 @@ impl<'a> URespAPI for URespKSNIMatmul<'a> { // note: `get_fock_coeff` keeps the trait default (dm route via `get_dm0_restricted` per spin). - /// rdm-form entry of the unrestricted response kernel; reserved for the future ugfock - /// machinery (no consumer yet), mirroring the restricted side before RI-PT2. - fn get_response_rdm(&mut self, _rdm: &[TsrView; 2]) -> [Tsr; 2] { - unimplemented!("get_response_rdm is not implemented for URespKSNIMatmul; reserved for the future ugfock machinery.") + /// Response matrix per spin from (symmetrizable) full spin density matrices: `2 * fxc` + /// response on the common (large) grid `ni` — the UHF factor `2.0` (hermitian symmetry + /// only, no spin degeneracy) against the restricted `4.0` of + /// [`RRespKSNIMatmul`](super::resp_rks::RRespKSNIMatmul). The per-spin `rdm` carries a + /// leading `[nao, nao]` square followed by an arbitrary (possibly empty) set of trailing + /// dimensions, shared by both spins; the output has the same per-spin shape. + /// + /// This is the rdm-form entry of the unrestricted response kernel, the form required by the + /// future unrestricted generalized-Fock Lagrangian term `A_{ai, pq} D_{pq}`; accordingly, + /// the fxc kernel is built from the ground-state spin densities on the common (large) grid + /// `ni` — not on the dedicated response grid — mirroring + /// [`RRespKSNIMatmul::get_response_rdm`](super::resp_rks::RRespKSNIMatmul::get_response_rdm). + /// Call [`make_response_preparation`](URespAPI::make_response_preparation) before this + /// function to make sure the data is ready. + fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + for rdm_s in rdm { + assert!(rdm_s.ndim() >= 2, "rdm must have at least 2 dimensions"); + } + let rdm_shape = [rdm[α].shape().to_vec(), rdm[β].shape().to_vec()]; + let nao = rdm_shape[α][0]; + for s in [α, β] { + assert_eq!(nao, rdm_shape[s][1], "the first two dimensions of rdm must be equal"); + } + assert_eq!(rdm_shape[α], rdm_shape[β], "the rdm shape must agree across spins"); + let nset: usize = rdm_shape[α][2..].iter().product(); + let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); + + if !self.intmd.contains_key("fxc_common_grid") { + // panic with "no entry" if preparation is skipped + let mo_coeff = [self.intmd["mo_coeff_0"].view(), self.intmd["mo_coeff_1"].view()]; + let mo_occ = [self.intmd["mo_occ_0"].view(), self.intmd["mo_occ_1"].view()]; + let (_, fxc) = make_cpks_vxc_fxc_uks(&self.xc_func_list, &mut self.ni, &mo_coeff, &mo_occ); + self.intmd.insert("fxc_common_grid".to_string(), fxc); + } + // NOTE: unlike `cpks_fxc`, the cached `fxc_common_grid` is NOT re-validated against the + // orbitals on later calls. This is safe under the analdrv driver contract (the orbitals + // are fixed for the lifetime of a driver object), but any future caller that + // re-prepares with different orbitals must invalidate this entry as well. + let fxc = self.intmd["fxc_common_grid"].view(); + + // flatten the trailing dimensions into one set dimension; symmetrize each spin's + // leading [nao, nao] block — the fxc kernel only sees the symmetric part anyway + let rdm_sym = [ + ((&rdm[α] + &rdm[α].swapaxes(0, 1)) * 0.5).into_shape((nao, nao, nset)), + ((&rdm[β] + &rdm[β].swapaxes(0, 1)) * 0.5).into_shape((nao, nao, nset)), + ]; + + // rho1 : [ngrids, nvar, 2, nset], assembled set by set from the per-spin dm lists + let ngrids = self.ni.weights.len(); + let nvar = den_type.num_nvar(); + let device = rdm[α].device().clone(); + let mut rho1 = rt::zeros(([ngrids, nvar, 2, nset], &device)); + for k in 0..nset { + let rho1_k = + self.ni.make_rho_from_dm(&[rdm_sym[α].i((.., .., k)), rdm_sym[β].i((.., .., k))], den_type); + rho1.i_mut((.., .., .., k)).assign(&rho1_k); + } + + let pot = self.ni.make_fxc_pot_with_eff(fxc, rho1.view(), den_type, XCSpin::Polarized); + // restore the input's per-spin trailing shape + [ + (pot.i((.., .., α, ..)) * 2.0_f64).into_shape(rdm_shape[α].to_vec()), + (pot.i((.., .., β, ..)) * 2.0_f64).into_shape(rdm_shape[β].to_vec()), + ] } /// Cached on first call for fixed inputs: the CP-KS `vxc`/`fxc` evaluation is skipped when diff --git a/src/ri_jk/resp_u.rs b/src/ri_jk/resp_u.rs index 0fe24b3c5..3a67484f0 100644 --- a/src/ri_jk/resp_u.rs +++ b/src/ri_jk/resp_u.rs @@ -116,10 +116,48 @@ impl<'a> URespAPI for URespRIJK<'a> { fock } - fn get_response_rdm(&mut self, _rdm: &[TsrView; 2]) -> [Tsr; 2] { - // rdm-form entry of the unrestricted response kernel; reserved for the future ugfock - // machinery (no consumer yet), mirroring the restricted side before RI-PT2 - unimplemented!("get_response_rdm is not implemented for URespRIJK; reserved for the future ugfock machinery.") + /// Response matrix per spin from (symmetrizable) spin density matrices: `2 * J(rdm_α + + /// rdm_β) - 2 * K(rdm_s)`. + /// + /// This is the rdm-form entry of the unrestricted response kernel, consistent with the + /// restricted `4 * J - 2 * K` convention of [`RRespRIJK`](super::resp_r::RRespRIJK) and with + /// this object's own bra form (whose `0.5` Coulomb prefactor carries the same convention): + /// per-spin occupation 1 against the restricted 2 halves the Coulomb scale, while the + /// two-term symmetrized exchange kernel keeps the restricted scale. In the closed-shell + /// limit (`rdm_α = rdm_β = X`), each spin's response equals the restricted kernel on `X`. + /// No preparation is required. + fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + for rdm_s in rdm { + assert_eq!(rdm_s.ndim(), 2, "rdm must have 2 dimensions"); + } + let [α, β] = [0, 1]; + let nao = rdm[α].shape()[0]; + let device = self.cderi.device(); + + // symmetrize (the pure in-core functions require symmetric input); the Coulomb response + // sees the total (α+β) density + let rdm_sym = [ + ((&rdm[α] + &rdm[α].t()) * 0.5).into_contig(ColMajor).into_shape((nao, nao, 1)), + ((&rdm[β] + &rdm[β].t()) * 0.5).into_contig(ColMajor).into_shape((nao, nao, 1)), + ]; + let rdm_total = (&rdm_sym[α] + &rdm_sym[β]).into_contig(ColMajor); + + let mut resp = [rt::zeros(([nao, nao], device)), rt::zeros(([nao, nao], device))]; + if self.factor_j != 0.0 { + let vj = get_vj_ri_incore(self.cderi.view(), rdm_total.view()).i((.., .., 0)).into_contig(ColMajor); + // 2.0 against the restricted 4.0: occupation 1 vs 2 on the total-density Coulomb + resp[α] += 2.0 * self.factor_j * &vj; + resp[β] += 2.0 * self.factor_j * &vj; + } + if self.factor_k != 0.0 { + // K: same-spin exchange, same 2.0 scale as restricted (two-term symmetrized kernel) + for (s, resp_s) in resp.iter_mut().enumerate() { + // TODO: batch size `72` should be tunable by max-memory. + let vk = get_vk_ri_incore_dm(self.cderi.view(), rdm_sym[s].view(), 72).i((.., .., 0)).into_contig(ColMajor); + *resp_s -= 2.0 * self.factor_k * &vk; + } + } + resp } fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { diff --git a/tests/analdrv/multipole/co2_hf.rs b/tests/analdrv/multipole/co2_hf.rs index a88f88dc7..07b57477c 100644 --- a/tests/analdrv/multipole/co2_hf.rs +++ b/tests/analdrv/multipole/co2_hf.rs @@ -1,6 +1,5 @@ //! CO2 (HF, RI-JK, def2-TZVP): the neutral linear molecule has vanishing dipole, so its raw -//! quadrupole is origin-invariant. Quadrupole reference from pyscf -//! (`local-runs/260911-multipole/gen_refs.py`, pyscf 2.14.0). +//! quadrupole is origin-invariant. Quadrupole reference from pyscf 2.14.0. use pyrest::analdrv::multipole::rmultipole::RMultipoleDH; use pyrest::ctrl_io; diff --git a/tests/analdrv/multipole/h2o_hf.rs b/tests/analdrv/multipole/h2o_hf.rs index a6048a024..051a19412 100644 --- a/tests/analdrv/multipole/h2o_hf.rs +++ b/tests/analdrv/multipole/h2o_hf.rs @@ -1,7 +1,6 @@ //! SCF-level (HF, RI-JK) multipole moments of an asymmetric H2O (def2-TZVP): dipole and -//! quadrupole verified against pyscf (`local-runs/260911-multipole/gen_refs.py`, pyscf 2.14.0, -//! RI-JK with def2-universal-jkfit), octupole and hexadecapole against Gaussian 16 RevB.01 -//! exact HF (`local-runs/260911-multipole/g16/`; default `#p` output contains the moments +//! quadrupole verified against pyscf (2.14.0, RI-JK with def2-universal-jkfit), octupole and +//! hexadecapole against Gaussian 16 RevB.01 exact HF (default `#p` output contains the moments //! through hexadecapole, so no extra keyword is needed), plus origin-invariance (dipole) and //! origin-shift-formula (quadrupole) checks. @@ -117,12 +116,12 @@ fn test_h2o() { (1e-6, 1e-7) )); - // octupole and hexadecapole: totals referenced to Gaussian 16 RevB.01 exact HF - // (`local-runs/260911-multipole/g16/`, 4-decimal print; note G16's default output contains - // these sections without any extra keyword, and its "Traceless Quadrupole" uses Q - Tr/3 I - // without the 3/2 factor). Tolerances carry the exact-vs-RI-JK density difference (dipole/ - // quadrupole calibration: 8e-5 / 2e-4 a.u.) plus the G16 print rounding. The nuclear parts - // are closed-form and asserted exactly against an independent computation. + // octupole and hexadecapole: totals referenced to Gaussian 16 RevB.01 exact HF (4-decimal + // print; note G16's default output contains these sections without any extra keyword, and + // its "Traceless Quadrupole" uses Q - Tr/3 I without the 3/2 factor). Tolerances carry the + // exact-vs-RI-JK density difference (dipole/quadrupole calibration: 8e-5 / 2e-4 a.u.) plus + // the G16 print rounding. The nuclear parts are closed-form and asserted exactly against an + // independent computation. let oct_tot = rmultipole.make_octupole(); assert!(rt::allclose( &rmultipole.result["oct_nuc"], diff --git a/tests/analdrv/multipole/h2o_rks.rs b/tests/analdrv/multipole/h2o_rks.rs index 1a0d1b00d..88bdc409f 100644 --- a/tests/analdrv/multipole/h2o_rks.rs +++ b/tests/analdrv/multipole/h2o_rks.rs @@ -1,6 +1,6 @@ //! SCF-level (B3LYP, RI-JK) multipole moments of an asymmetric H2O (def2-TZVP), verified against -//! pyscf (`local-runs/260911-multipole/gen_refs.py`, pyscf 2.14.0, RI-JK with -//! def2-universal-jkfit). Tolerances carry the REST/pyscf grid differences (~1e-5). +//! pyscf (2.14.0, RI-JK with def2-universal-jkfit). Tolerances carry the REST/pyscf grid +//! differences (~1e-5). use pyrest::analdrv::multipole::rmultipole::RMultipoleDH; use pyrest::ctrl_io; diff --git a/tests/analdrv/multipole/hi_rmp2.rs b/tests/analdrv/multipole/hi_rmp2.rs index af17450b3..5b80dc69a 100644 --- a/tests/analdrv/multipole/hi_rmp2.rs +++ b/tests/analdrv/multipole/hi_rmp2.rs @@ -8,8 +8,7 @@ //! //! References: Gaussian 16 Legacy RevB.01, `#p MP2(Full)/def2TZVP Density=MP2 NoSymm` (relaxed //! MP2 density; note G16's default MP2 freezes the I 4s/4p subshells, so `Full` is required to -//! match REST's `start_mo = 0`), geometry H(0,0,0), I(0,0,1.609 Angstrom), run recorded at -//! `local-runs/260912-hi-ecp/g16`: +//! match REST's `start_mo = 0`), geometry H(0,0,0), I(0,0,1.609 Angstrom): //! //! - E(SCF) = -297.242009239, E2 = -0.5911611131, EUMP2 = -297.83317035213 //! - dipole Z = -0.5630 D (G16 prints 4 decimals; 1e-5 a.u. rounding) diff --git a/tests/analdrv/multipole/interface_nh2_uhf.rs b/tests/analdrv/multipole/interface_nh2_uhf.rs new file mode 100644 index 000000000..700142501 --- /dev/null +++ b/tests/analdrv/multipole/interface_nh2_uhf.rs @@ -0,0 +1,139 @@ +//! Interface-level (through `analdrv_interface`) multipole test: SCF-level UHF NH2, plus the +//! rejection of the unrestricted post-SCF (fifth-DFA) path. +//! +//! Drives the full unrestricted task path — task dispatch, explicit `multipole_origin`, stdout +//! print, and the results-JSON export. The moment references are identical to +//! `tests/analdrv/multipole/nh2_uhf.rs` (pyscf 2.14.0, RI-JK with def2-universal-jkfit). + +use pyrest::analdrv::config::{AnalDrvConfig, AnalDrvTask}; +use pyrest::analdrv::interface::{analdrv_interface, analdrv_json_interface}; +use pyrest::ctrl_io; +use pyrest::molecule_io::Molecule; +use pyrest::scf_io::{self, scf_without_build}; + +static INPUT_NH2: &str = r##" +[ctrl] + print_level = 2 + num_threads = 4 + xc = "hf" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 2.0 + spin_polarization = true + +[geom] + name = "NH2" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 0.58 0.20 0.85 + H -0.50 -0.25 0.83 + """ +"##; + +static INPUT_NH2_XYG3: &str = r##" +[ctrl] + print_level = 2 + num_threads = 4 + xc = "xyg3" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 2.0 + spin_polarization = true + +[ctrl.ri_pt2] + new_driver = true + +[geom] + name = "NH2" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 0.58 0.20 0.85 + H -0.50 -0.25 0.83 + """ +"##; + +fn allclose(v: &[f64], ref_: &[f64], tol: f64) { + assert_eq!(v.len(), ref_.len()); + for (a, b) in v.iter().zip(ref_.iter()) { + assert!((a - b).abs() < tol, "component mismatch: {a} vs {b}"); + } +} + +fn build_scf(input: &str) -> scf_io::SCF { + let keys = toml::from_str::(input).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); + scf_data +} + +#[test] +fn test_nh2_uhf_interface() { + let mut scf_data = build_scf(INPUT_NH2); + let tasks = vec![AnalDrvTask::Multipole]; + + // --- explicit origin [0, 0, 0]: identical settings to the driver-level nh2_uhf test --- // + let mut config = AnalDrvConfig::default(); + config.multipole.orders = vec![1, 2]; + config.multipole.origin = Some([0.0; 3]); + let output = analdrv_interface(&mut scf_data, &tasks, &config); + let json_extra = analdrv_json_interface(&output); + + let mp = output.multipole.as_ref().expect("multipole task output"); + assert_eq!(mp.origin, [0.0; 3]); + let dip = mp.dipole.as_ref().unwrap(); + // SCF level: no corr/resp increments at all + assert!(dip.corr.is_none()); + assert!(dip.resp.is_none()); + // reference: pyscf RI-JK UHF (def2-tzvp / def2-universal-jkfit) + allclose(&dip.nuc, &[0.151178089965, -0.094486306228, 3.174739889269], 1e-8); + allclose(&dip.tot, &[0.016656856517, -0.038475737329, 0.915089522375], 1e-6); + allclose( + &mp.quadrupole.as_ref().unwrap().tot, + &[ + -5.193509415837, 0.020912254543, 0.050343803281, 0.020912254543, -5.239324563951, + -0.030545457675, 0.050343803281, -0.030545457675, -4.210872999273, + ], + 1e-6, + ); + // traceless quadrupole: 3/2 (Q - Tr(Q)/3 I), XX component from the raw reference above + let quad_traceless = mp.quadrupole.as_ref().unwrap().traceless.as_ref().unwrap(); + allclose(&quad_traceless[..1], &[-0.468410634225], 1e-6); + + // results-JSON layout: no corr/resp keys for SCF-level unrestricted, no hessian keys + let anal = &json_extra["analdrv"]; + assert!(anal.get("frequencies_cm").is_none()); + let dip_json = &anal["multipole"]["dipole"]; + assert!(dip_json["corr"].is_null()); + assert!(dip_json["resp"].is_null()); + allclose( + &dip_json["tot"].as_array().unwrap().iter().map(|v| v.as_f64().unwrap()).collect::>(), + &[0.016656856517, -0.038475737329, 0.915089522375], + 1e-6, + ); + + // --- default origin: the coordinate origin [0, 0, 0] --- // + let mut config = AnalDrvConfig::default(); + config.multipole.orders = vec![1]; + let output = analdrv_interface(&mut scf_data, &tasks, &config); + let mp = output.multipole.as_ref().unwrap(); + assert_eq!(mp.origin, [0.0; 3]); + allclose(&mp.dipole.as_ref().unwrap().tot, &[0.016656856517, -0.038475737329, 0.915089522375], 1e-6); +} + +#[test] +#[should_panic(expected = "Multipole evaluation is not implemented for unrestricted post-SCF")] +fn test_nh2_uhf_fifth_dfa_panics() { + let mut scf_data = build_scf(INPUT_NH2_XYG3); + let tasks = vec![AnalDrvTask::Multipole]; + let mut config = AnalDrvConfig::default(); + config.multipole.orders = vec![1]; + let _ = analdrv_interface(&mut scf_data, &tasks, &config); +} diff --git a/tests/analdrv/multipole/mod.rs b/tests/analdrv/multipole/mod.rs index a3f4a5b7d..07a5c5bbd 100644 --- a/tests/analdrv/multipole/mod.rs +++ b/tests/analdrv/multipole/mod.rs @@ -3,5 +3,8 @@ pub mod h2o_hf; pub mod h2o_rks; pub mod hi_rmp2; pub mod interface_h2o_rks; +pub mod interface_nh2_uhf; pub mod interface_nh3_xyg3; +pub mod nh2_uhf; +pub mod nh2_uks; pub mod nh3_xyg3; diff --git a/tests/analdrv/multipole/nh2_uhf.rs b/tests/analdrv/multipole/nh2_uhf.rs new file mode 100644 index 000000000..71fdc4e80 --- /dev/null +++ b/tests/analdrv/multipole/nh2_uhf.rs @@ -0,0 +1,121 @@ +//! SCF-level (UHF, RI-JK) multipole moments of the NH2 radical (doublet, ²B1, def2-TZVP), +//! verified against pyscf (2.14.0, RI-JK with def2-universal-jkfit), plus origin-invariance +//! (dipole). NH2 is a non-degenerate open shell, unlike OH whose ²Π SOMO makes UHF/UKS +//! properties ill-conditioned. + +use pyrest::analdrv::multipole::umultipole::UMultipoleDH; +use pyrest::ctrl_io; +use pyrest::molecule_io::Molecule; +use pyrest::ri_jk::util::get_cint_mol; +use pyrest::scf_io::{self, scf_without_build}; +use pyrest::utilities::rstsr_util::RestTensorToRstsrTsrAPI; + +use rstsr::prelude::*; + +static INPUT_NH2: &str = r##" +[ctrl] + print_level = 2 + num_threads = 4 + xc = "hf" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 2.0 + spin_polarization = true + +[geom] + name = "NH2" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 0.58 0.20 0.85 + H -0.50 -0.25 0.83 + """ +"##; + +#[test] +fn test_nh2_uhf() { + let keys = toml::from_str::(&INPUT_NH2[..]).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_coeff = [(&scf_data.eigenvectors[0]).to_rstsr(&device), (&scf_data.eigenvectors[1]).to_rstsr(&device)]; + let mo_occ = [(&scf_data.occupation[0]).to_rstsr(&device), (&scf_data.occupation[1]).to_rstsr(&device)]; + let mol_cint = get_cint_mol(&scf_data.mol); + + let mut umultipole = UMultipoleDH::new(&mol_cint, mo_coeff.clone(), mo_occ.clone(), [0.0; 3]); + let dip_tot = umultipole.make_dipole(); + let quad_tot = umultipole.make_quadrupole(); + umultipole.print_multipole(2); + + // reference: pyscf RI-JK UHF (def2-tzvp / def2-universal-jkfit) + assert!(rt::allclose( + &umultipole.result["dip_nuc"], + &rt::asarray((vec![0.151178089965, -0.094486306228, 3.174739889269], &device)), + (1e-8, 1e-9) + )); + assert!(rt::allclose( + &umultipole.result["dip_scf"], + &rt::asarray((vec![-0.134521233449, 0.056010568899, -2.259650366894], &device)), + (1e-6, 1e-7) + )); + assert!(rt::allclose( + &dip_tot, + &rt::asarray((vec![0.016656856517, -0.038475737329, 0.915089522375], &device)), + (1e-6, 1e-7) + )); + + assert!(rt::allclose( + &umultipole.result["quad_nuc"], + &rt::asarray(( + vec![2.094072413886, 0.860626623033, 0.278543056417, 0.860626623033, 0.366034144651, + -0.133914930970, 0.278543056417, -0.133914930970, 5.040200895224], + &device + )) + .into_shape((3, 3)), + (1e-8, 1e-9) + )); + assert!(rt::allclose( + &umultipole.result["quad_scf"], + &rt::asarray(( + vec![-7.287581829724, -0.839714368491, -0.228199253137, -0.839714368491, -5.605358708602, + 0.103369473295, -0.228199253137, 0.103369473295, -9.251073894497], + &device + )) + .into_shape((3, 3)), + (1e-6, 1e-7) + )); + assert!(rt::allclose( + &quad_tot, + &rt::asarray(( + vec![-5.193509415837, 0.020912254543, 0.050343803281, 0.020912254543, -5.239324563951, + -0.030545457675, 0.050343803281, -0.030545457675, -4.210872999273], + &device + )) + .into_shape((3, 3)), + (1e-6, 1e-7) + )); + assert!(rt::allclose( + &umultipole.result["quad_tot_traceless"], + &rt::asarray(( + vec![-0.468410634225, 0.031368381814, 0.075515704921, 0.031368381814, -0.537133356396, + -0.045818186512, 0.075515704921, -0.045818186512, 1.005543990621], + &device + )) + .into_shape((3, 3)), + (1e-6, 1e-7) + )); + + // shifted origin: the dipole of the neutral radical is origin-independent + let mut umultipole2 = UMultipoleDH::new(&mol_cint, mo_coeff, mo_occ, [0.3, -0.2, 0.15]); + let dip2 = umultipole2.make_dipole(); + assert!(rt::allclose(&dip2, &dip_tot, (1e-9, 1e-10))); + + for (label, t) in &umultipole.timing { + println!("Timing {label}: {t:.3} s"); + } +} diff --git a/tests/analdrv/multipole/nh2_uks.rs b/tests/analdrv/multipole/nh2_uks.rs new file mode 100644 index 000000000..f6ca2a220 --- /dev/null +++ b/tests/analdrv/multipole/nh2_uks.rs @@ -0,0 +1,79 @@ +//! SCF-level (UKS B3LYP, RI-JK) multipole moments of the NH2 radical (doublet, def2-TZVP), +//! verified against pyscf (2.14.0, RI-JK with def2-universal-jkfit). Tolerances carry the +//! REST/pyscf grid differences (~2e-6 on the small off-diagonal quadrupole components). + +use pyrest::analdrv::multipole::umultipole::UMultipoleDH; +use pyrest::ctrl_io; +use pyrest::molecule_io::Molecule; +use pyrest::ri_jk::util::get_cint_mol; +use pyrest::scf_io::{self, scf_without_build}; +use pyrest::utilities::rstsr_util::RestTensorToRstsrTsrAPI; + +use rstsr::prelude::*; + +static INPUT_NH2: &str = r##" +[ctrl] + print_level = 2 + num_threads = 4 + xc = "b3lyp" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 2.0 + spin_polarization = true + +[geom] + name = "NH2" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 0.58 0.20 0.85 + H -0.50 -0.25 0.83 + """ +"##; + +#[test] +fn test_nh2_uks() { + let keys = toml::from_str::(&INPUT_NH2[..]).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_coeff = [(&scf_data.eigenvectors[0]).to_rstsr(&device), (&scf_data.eigenvectors[1]).to_rstsr(&device)]; + let mo_occ = [(&scf_data.occupation[0]).to_rstsr(&device), (&scf_data.occupation[1]).to_rstsr(&device)]; + let mol_cint = get_cint_mol(&scf_data.mol); + + let mut umultipole = UMultipoleDH::new(&mol_cint, mo_coeff, mo_occ, [0.0; 3]); + let dip_tot = umultipole.make_dipole(); + let quad_tot = umultipole.make_quadrupole(); + umultipole.print_multipole(2); + + // reference: pyscf RI-JK UKS B3LYP (def2-tzvp / def2-universal-jkfit) + assert!(rt::allclose( + &umultipole.result["dip_scf"], + &rt::asarray((vec![-0.136575312767, 0.057794212957, -2.312851619633], &device)), + (1e-5, 1e-6) + )); + assert!(rt::allclose( + &dip_tot, + &rt::asarray((vec![0.014602777199, -0.036692093271, 0.861888269636], &device)), + (1e-5, 1e-6) + )); + assert!(rt::allclose( + &quad_tot, + &rt::asarray(( + vec![-5.264151352116, 0.026124544967, 0.047625109070, 0.026124544967, -5.320049075464, + -0.028624095927, 0.047625109070, -0.028624095927, -4.355249946458], + &device + )) + .into_shape((3, 3)), + (1e-5, 5e-6) + )); + + for (label, t) in &umultipole.timing { + println!("Timing {label}: {t:.3} s"); + } +} diff --git a/tests/analdrv/multipole/nh3_xyg3.rs b/tests/analdrv/multipole/nh3_xyg3.rs index 2680be9c7..495fdb591 100644 --- a/tests/analdrv/multipole/nh3_xyg3.rs +++ b/tests/analdrv/multipole/nh3_xyg3.rs @@ -1,11 +1,10 @@ //! XYG3 (xDH double hybrid) multipole moments of NH3 (def2-TZVP), driving the `RMultipoleDH` //! evaluator with the `RGFockDH` composite and the `RRespSCF` response object. //! -//! Dipole decomposition reference values are those recorded by -//! `local-runs/260910-xyg3/xyg3_dipole_contribs.py` (pyscf-forge master @ 0566d43, pyscf 2.14.0) -//! and already asserted contribution-by-contribution in `tests/analdrv/response/xyg3_rgfock.rs`; -//! the quadrupole SCF-density reference is from `local-runs/260911-multipole/gen_refs.py` -//! (pyscf 2.14.0, B3LYP RI-JK). The correlation/response quadrupole increments have no external +//! Dipole decomposition reference values are those obtained with pyscf-forge DFDH (master +//! @ 0566d43, pyscf 2.14.0), already asserted contribution-by-contribution in +//! `tests/analdrv/response/xyg3_rgfock.rs`; the quadrupole SCF-density reference is from pyscf +//! 2.14.0 (B3LYP RI-JK). The correlation/response quadrupole increments have no external //! reference yet and are printed only. use pyrest::analdrv::config::AnalDrvConfig; diff --git a/tests/analdrv/response/mod.rs b/tests/analdrv/response/mod.rs index 00b08467c..1c1fff432 100644 --- a/tests/analdrv/response/mod.rs +++ b/tests/analdrv/response/mod.rs @@ -1,2 +1,3 @@ pub mod pt2_rgfock; +pub mod uresp_gfock; pub mod xyg3_rgfock; \ No newline at end of file diff --git a/tests/analdrv/response/uresp_gfock.rs b/tests/analdrv/response/uresp_gfock.rs new file mode 100644 index 000000000..b42cc0e85 --- /dev/null +++ b/tests/analdrv/response/uresp_gfock.rs @@ -0,0 +1,163 @@ +//! Unrestricted response-kernel rdm form (`get_response_rdm`, the ugfock groundwork) for +//! URespRIJK and URespKSNIMatmul, pinned two ways: +//! - consistency with the hessian-tested bra form (`get_response_bra`): the rdm form on the +//! symmetrized bra-generated densities, right half-transformed by `mocc`, equals the bra form; +//! - closed-shell limit: with both spins carrying the restricted test density, each spin's +//! response equals the restricted kernel (`4 J - 2 K` / `4 fxc`). + +use pyrest::analdrv::response::rresp_interface::{scf_jk_factors, scf_xc_func_list}; +use pyrest::analdrv::response::trait_rresp::RRespAPI; +use pyrest::analdrv::response::trait_uresp::URespAPI; +use pyrest::analdrv::response::uresp_interface::scf_xc_func_list_uks; +use pyrest::ctrl_io; +use pyrest::dft::numint_matmul::nimatmul::NIMatmul; +use pyrest::dft::numint_matmul::resp_rks::RRespKSNIMatmul; +use pyrest::dft::numint_matmul::resp_uks::URespKSNIMatmul; +use pyrest::molecule_io::Molecule; +use pyrest::ri_jk::resp_r::RRespRIJK; +use pyrest::ri_jk::resp_u::URespRIJK; +use pyrest::ri_jk::util::get_cint_mol; +use pyrest::scf_io::{self, scf_without_build}; +use pyrest::utilities::rstsr_util::{RestTensorToRstsrTsrAPI, RestTensorToRstsrViewAPI, Tsr}; + +use rstsr::prelude::*; + +static INPUT_OH_UKS: &str = r##" +[ctrl] + print_level = 2 + num_threads = 4 + xc = "b3lyp" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 2.0 + spin_polarization = true + +[geom] + name = "OH" + unit = "Angstrom" + position = """ + O 0.0 0.0 0.0 + H 0.75 0.55 0.10 + """ +"##; + +static INPUT_H2O_RKS: &str = r##" +[ctrl] + print_level = 2 + num_threads = 4 + xc = "b3lyp" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 1.0 + spin_polarization = false + +[geom] + name = "H2O" + unit = "Angstrom" + position = """ + O 0.0 0.0 0.0 + H 0.75 0.55 0.10 + H -0.70 0.60 -0.05 + """ +"##; + +fn build_scf(input: &str) -> scf_io::SCF { + let keys = toml::from_str::(input).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); + scf_data +} + +fn det_matrix(m: usize, n: usize, seed: usize, device: &DeviceBLAS) -> Tsr { + let data: Vec = (0..m * n).map(|k| (((k + seed) * 7 + 3) % 11) as f64 / 11.0 - 0.5).collect(); + rt::asarray((data, vec![m, n], device)) +} + +#[test] +fn test_oh_uks_response_rdm_consistency() { + let scf_data = build_scf(INPUT_OH_UKS); + let device = rt::DeviceBLAS::default(); + let mo_coeff = [(&scf_data.eigenvectors[0]).to_rstsr(&device), (&scf_data.eigenvectors[1]).to_rstsr(&device)]; + let mo_occ = [(&scf_data.occupation[0]).to_rstsr(&device), (&scf_data.occupation[1]).to_rstsr(&device)]; + let nao = mo_coeff[0].shape()[0]; + + // --- URespRIJK: bra form vs rdm form on the bra-generated densities --- // + let (factor_j, factor_k, _) = scf_jk_factors(&scf_data); + let (rimatr, _, _) = scf_data.rimatr.as_ref().unwrap(); + let mut rijk = URespRIJK::new_with_cderi(factor_j, factor_k, rimatr.to_rstsr_view(&device).into_cow()); + rijk.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()]); + + let occidx = [mo_occ[0].view().greater(0).into_vec(), mo_occ[1].view().greater(0).into_vec()]; + let mocc = [mo_coeff[0].bool_select(-1, &occidx[0]), mo_coeff[1].bool_select(-1, &occidx[1])]; + let bra = [det_matrix(nao, mocc[0].shape()[1], 1, &device), det_matrix(nao, mocc[1].shape()[1], 2, &device)]; + + let resp_bra = rijk.get_response_bra(&[bra[0].view(), bra[1].view()]); + let x = [&bra[0] % mocc[0].t(), &bra[1] % mocc[1].t()]; + let x_sym = [(&x[0] + &x[0].t()) * 0.5, (&x[1] + &x[1].t()) * 0.5]; + let resp_rdm = rijk.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()]); + for s in 0..2 { + let expect = resp_rdm[s].view() % mocc[s].view(); + assert!(rt::allclose(&resp_bra[s], &expect, (1e-8, 1e-9)), "RI-JK bra/rdm mismatch, spin {s}"); + } + + // --- URespKSNIMatmul: the same check on the XC kernel --- // + let grids = scf_data.grids.as_ref().unwrap(); + let mol_cint = get_cint_mol(&scf_data.mol); + let ni = NIMatmul::new(&mol_cint, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); + let mut uks = URespKSNIMatmul::new(scf_xc_func_list_uks(&scf_data), ni, false); + uks.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()]); + + let resp_bra = uks.get_response_bra(&[bra[0].view(), bra[1].view()]); + let resp_rdm = uks.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()]); + for s in 0..2 { + let expect = resp_rdm[s].view() % mocc[s].view(); + assert!(rt::allclose(&resp_bra[s], &expect, (1e-7, 1e-8)), "UKS bra/rdm mismatch, spin {s}"); + } +} + +#[test] +fn test_closed_shell_limit() { + let scf_data = build_scf(INPUT_H2O_RKS); + let device = rt::DeviceBLAS::default(); + let mo_coeff = (&scf_data.eigenvectors[0]).to_rstsr(&device); + let mo_occ = (&scf_data.occupation[0]).to_rstsr(&device); + let nao = mo_coeff.shape()[0]; + + // deterministic symmetric test density + let raw = det_matrix(nao, nao, 5, &device); + let x = (&raw + &raw.t()) * 0.5; + + // --- RI-JK: U(X, X) per spin == R(X) --- // + let (factor_j, factor_k, _) = scf_jk_factors(&scf_data); + let (rimatr, _, _) = scf_data.rimatr.as_ref().unwrap(); + let mut rrijk = RRespRIJK::new_with_cderi(factor_j, factor_k, rimatr.to_rstsr_view(&device).into_cow()); + let mut urijk = URespRIJK::new_with_cderi(factor_j, factor_k, rimatr.to_rstsr_view(&device).into_cow()); + let resp_r = rrijk.get_response_rdm(x.view()); + let resp_u = urijk.get_response_rdm(&[x.view(), x.view()]); + for s in 0..2 { + assert!(rt::allclose(&resp_u[s], &resp_r, (1e-8, 1e-9)), "RI-JK closed-shell mismatch, spin {s}"); + } + + // --- KS XC: the same, with the U object prepared at halved occupation (occ 1 per spin) --- // + let grids = scf_data.grids.as_ref().unwrap(); + let mol_cint = get_cint_mol(&scf_data.mol); + let ni_r = NIMatmul::new(&mol_cint, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); + let ni_u = NIMatmul::new(&mol_cint, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); + let mut rks = RRespKSNIMatmul::new(scf_xc_func_list(&scf_data), ni_r, false); + let mut uks = URespKSNIMatmul::new(scf_xc_func_list_uks(&scf_data), ni_u, false); + rks.make_response_preparation(mo_coeff.view(), mo_occ.view()); + let occ_half = &mo_occ * 0.5_f64; + uks.make_response_preparation(&[mo_coeff.view(), mo_coeff.view()], &[occ_half.view(), occ_half.view()]); + + let resp_r = rks.get_response_rdm(x.view()); + let resp_u = uks.get_response_rdm(&[x.view(), x.view()]); + for s in 0..2 { + assert!(rt::allclose(&resp_u[s], &resp_r, (1e-7, 1e-8)), "UKS closed-shell mismatch, spin {s}"); + } +} diff --git a/tests/analdrv/response/xyg3_rgfock.rs b/tests/analdrv/response/xyg3_rgfock.rs index 34c711a2e..0aae6bba9 100644 --- a/tests/analdrv/response/xyg3_rgfock.rs +++ b/tests/analdrv/response/xyg3_rgfock.rs @@ -1,8 +1,7 @@ //! XYG3 (xDH double hybrid) generalized Fock / relaxed dipole of NH3 (def2-TZVP), verified //! contribution-by-contribution against pyscf-forge DFDH. //! -//! Reference values were recorded by -//! `local-runs/260910-xyg3/xyg3_dipole_contribs.py` (pyscf-forge master @ 0566d43, pyscf 2.14.0): +//! Reference values were obtained with pyscf-forge DFDH (master @ 0566d43, pyscf 2.14.0): //! XYG3 on B3LYP orbitals, basis def2-TZVP, auxiliary basis def2-universal-jkfit for both the JK //! fitting and the RI-MP2 (3-center) part. The Lagrangian decomposition follows pyscf-forge's //! `prepare_lagrangian`: `L_pt2 = W3 + W4 + Ax0(D_rdm1)` and -- Gitee From 642bb9367ca37bceb9ec3cfbe95b82f36f8fb0f5 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 23 Sep 2026 12:44:13 +0800 Subject: [PATCH 03/21] Add resp_auxbas_path keyword for response-specific auxiliary basis MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New analdrv resp configuration key `resp_auxbas_path`: when set, the RI-JK response (A-tensor) objects evaluate on a decomposed ERI freshly built on that auxiliary basis instead of borrowing the SCF rimatr/rimatr_sr, while the energy-derivative (B-side) objects keep the SCF auxiliary basis — the CP-SCF solve becomes a mixed-representation approximation, parallel to the coarser response grid of grid_level_cpscf. - analdrv: keyword validation at interface entry (hard error together with even_tempered_basis; warn when no task of the run uses the response object); config round-trip tests extended. - ri_jk/resp_auxbas (new module): per-component cderi definitions (cderi full-range, cderi_sr short-range RSH with -omega in libcint convention); a path equal to the SCF auxbas_path is treated as unset when the SCF tensor is in memory, avoiding an identical rebuild. - ri_jk/util: get_cint_aux_on_path regenerates the auxiliary CInt from a local keyword copy, so the SCF molecule is only read, never modified. - ri_jk/decompose: generate_rimatr_bare_on_cint builds the rimatr on explicitly given CInt integrators instead of the molecule's own bases. - utilities: RestTensorIntoRstsrTsrAPI for MatrixFull (owned conversion). Co-authored-by: Claude Code Co-authored-by: glm-5.3-flash[1m] --- src/analdrv/config.rs | 22 +++- src/analdrv/interface.rs | 6 ++ src/analdrv/response/rresp_interface.rs | 14 +-- src/analdrv/response/uresp_interface.rs | 14 +-- src/ri_jk/decompose.rs | 26 ++++- src/ri_jk/mod.rs | 3 + src/ri_jk/resp_auxbas.rs | 131 ++++++++++++++++++++++++ src/ri_jk/util.rs | 14 ++- src/utilities/rstsr_util.rs | 11 ++ 9 files changed, 217 insertions(+), 24 deletions(-) create mode 100644 src/ri_jk/resp_auxbas.rs diff --git a/src/analdrv/config.rs b/src/analdrv/config.rs index 42533b40b..bab90a44b 100644 --- a/src/analdrv/config.rs +++ b/src/analdrv/config.rs @@ -48,6 +48,8 @@ impl Default for AnalDrvGeneralCfg { /// `grid_level` is not a solver setting but the DFT grid of the CP-SCF response evaluation (the /// response path; the fock path and the DH generalized Fock stay on the SCF grid), so it is named /// in the `grid_level_*` family of [`AnalDrvNucgradCfg::grid_level_skeleton`] with `cpscf` as the +/// scope. Likewise, `auxbas_path` configures the response objects themselves (not only the +/// CP-SCF solve — every consumer of the response object shares it), so it is named in the `resp_` /// scope. /// /// The legacy `cphf_*` prefixed key names are still accepted as aliases. @@ -79,6 +81,19 @@ pub struct AnalDrvRespCfg { #[serde(rename = "grid_level_cpscf", alias = "grid_level_cphf")] #[serde_inline_default(None)] pub grid_level: Option, + /// Auxiliary basis of the RI-JK response objects (the response/A-tensor side), as a + /// basis-set pool name or an element-JSON directory — the same value format as + /// `ctrl.auxbas_path`. + /// + /// By default `None`: the response objects reuse the SCF auxiliary basis, borrowing the SCF + /// `rimatr` (and `rimatr_sr` for range-separated hybrids) with no copy. When set, a + /// decomposed ERI is freshly built on this basis for the response objects only — the + /// energy-derivative (B-side) objects of the hessian/multipole tasks keep the SCF auxiliary + /// basis, so the CP-SCF solve becomes a mixed-representation approximation (parallel to + /// `grid_level_cpscf`). Not supported together with `even_tempered_basis`. + #[serde(rename = "resp_auxbas_path")] + #[serde_inline_default(None)] + pub auxbas_path: Option, } impl Default for AnalDrvRespCfg { @@ -91,6 +106,7 @@ impl Default for AnalDrvRespCfg { lindep: 1e-15, tol_inflation: 1000.0, grid_level: None, + auxbas_path: None, } } } @@ -283,12 +299,14 @@ mod tests { "gau_thermo": true, "atm_list": [0, 1], "verbose": 3, + "resp_auxbas_path": "def2-universal-jkfit", }); let config: AnalDrvConfig = serde_json::from_value(v).unwrap(); assert_eq!(config.resp.tol, 1.0e-8); assert_eq!(config.resp.max_cycle, 99); assert_eq!(config.resp.level_shift, 0.0); assert_eq!(config.resp.grid_level, Some(2)); + assert_eq!(config.resp.auxbas_path, Some("def2-universal-jkfit".to_string())); assert_eq!(config.nucgrad.grid_level_skeleton, Some(4)); assert!(!config.nucgrad.grid_shift_deriv); assert!(config.nucgrad.gau_thermo); @@ -334,16 +352,18 @@ mod tests { let config: AnalDrvConfig = serde_json::from_value(serde_json::json!({})).unwrap(); assert_eq!(config, AnalDrvConfig::default()); - // serialization stays flat, under the cpscf_* / grid_level_* key names + // serialization stays flat, under the cpscf_* / grid_level_* / resp_* key names let mut config = AnalDrvConfig::default(); config.resp.tol = 1.0e-8; config.resp.max_space = 20; config.resp.grid_level = Some(2); + config.resp.auxbas_path = Some("def2-universal-jkfit".to_string()); config.multipole.rdm1_relax = MultipoleRdm1Relax::Unrelaxed; let v = serde_json::to_value(&config).unwrap(); assert_eq!(v["cpscf_tol"], serde_json::json!(1.0e-8)); assert_eq!(v["cpscf_max_space"], serde_json::json!(20)); assert_eq!(v["grid_level_cpscf"], serde_json::json!(2)); + assert_eq!(v["resp_auxbas_path"], serde_json::json!("def2-universal-jkfit")); assert_eq!(v["multipole_rdm1_relax"], serde_json::json!("unrelaxed")); } } diff --git a/src/analdrv/interface.rs b/src/analdrv/interface.rs index 2cd8900a8..0db06fb9c 100644 --- a/src/analdrv/interface.rs +++ b/src/analdrv/interface.rs @@ -3,6 +3,7 @@ use crate::analdrv::config::{AnalDrvConfig, AnalDrvTask, MultipoleRdm1Relax}; use crate::analdrv::hessian::hess_interface; use crate::analdrv::multipole::interface::{multipole_interface, MultipoleOutput}; +use crate::ri_jk::resp_auxbas::validate_resp_auxbas; use crate::analdrv::response::rgfock_interface; use crate::analdrv::response::rresp_interface; use crate::analdrv::response::rresp_interface::rscf_resp_interface; @@ -92,6 +93,11 @@ pub fn analdrv_interface(scf_data: &mut SCF, tasks: &[AnalDrvTask], config: &Ana } let need_resp = has_hessian || multipole_relaxed_dh; + + // response-specific auxiliary basis (`resp_auxbas_path`): validate it against this run's + // tasks here; the response objects pick it up per component when they are built + validate_resp_auxbas(scf_data, config, need_resp); + let mut resp_obj = if need_resp { match scf_data.scftype { SCFType::RHF => Some(RespSCF::R(rscf_resp_interface(scf_data, config))), diff --git a/src/analdrv/response/rresp_interface.rs b/src/analdrv/response/rresp_interface.rs index 7fabfa586..9ee0af66d 100644 --- a/src/analdrv/response/rresp_interface.rs +++ b/src/analdrv/response/rresp_interface.rs @@ -16,6 +16,7 @@ use crate::analdrv::config::AnalDrvRespCfg; use crate::analdrv::prelude::*; +use crate::ri_jk::resp_auxbas; use crate::dft::numint_matmul::nimatmul::NIMatmul; use crate::dft::numint_matmul::resp_rks::RRespKSNIMatmul; use crate::dft::Grids; @@ -300,18 +301,16 @@ impl<'a> RRespSCF<'a> { /// Build the response (fock/response) objects for a converged restricted SCF. pub fn rscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> RRespSCF<'a> { - let device = DeviceBLAS::default(); let mut resp_list: Vec> = Vec::new(); // --- RI-JK --- // let (factor_j, factor_k, rsh) = scf_jk_factors(scf_data); + // decomposed ERI of the RI-JK response object: its own definition — the freshly built one of + // `resp_auxbas_path` when set (owned), otherwise the SCF rimatr (borrowed, zero copy) { - let (rimatr, _, _) = scf_data.rimatr.as_ref().expect( - "This implementation requires cholesky decomposed ERI (or rimatr) to be available and stored in memory.", - ); - let cderi = rimatr.to_rstsr_view(&device).into_cow(); + let cderi = resp_auxbas::cderi(scf_data, config); resp_list.push(Box::new(RRespRIJK::new_with_cderi(factor_j, factor_k, cderi))); } @@ -319,10 +318,7 @@ pub fn rscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> RRe // object reusing the full-range implementation; it evaluates on the short-range `rimatr_sr` // ERI, with no Coulomb part (factor_j = 0). if let Some((_omega, factor_k_sr)) = rsh { - let (rimatr_sr, _, _) = scf_data.rimatr_sr.as_ref().expect( - "The range-separated response requires the short-range ERI (rimatr_sr) to be built and stored in memory.", - ); - let cderi_sr = rimatr_sr.to_rstsr_view(&device).into_cow(); + let cderi_sr = resp_auxbas::cderi_sr(scf_data, config); resp_list.push(Box::new(RRespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr))); } diff --git a/src/analdrv/response/uresp_interface.rs b/src/analdrv/response/uresp_interface.rs index de2ffa670..15e54608c 100644 --- a/src/analdrv/response/uresp_interface.rs +++ b/src/analdrv/response/uresp_interface.rs @@ -16,6 +16,7 @@ use crate::analdrv::config::AnalDrvRespCfg; use crate::analdrv::prelude::*; use crate::analdrv::response::rresp_interface::scf_jk_factors; +use crate::ri_jk::resp_auxbas; use crate::dft::numint_matmul::nimatmul::NIMatmul; use crate::dft::numint_matmul::resp_uks::URespKSNIMatmul; use crate::dft::Grids; @@ -349,18 +350,16 @@ impl<'a> URespSCF<'a> { /// Build the response (fock/response) objects for a converged unrestricted SCF. pub fn uscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> URespSCF<'a> { - let device = DeviceBLAS::default(); let mut resp_list: Vec> = Vec::new(); // --- RI-JK --- // let (factor_j, factor_k, rsh) = scf_jk_factors(scf_data); + // decomposed ERI of the RI-JK response object: its own definition — the freshly built one of + // `resp_auxbas_path` when set (owned), otherwise the SCF rimatr (borrowed, zero copy) { - let (rimatr, _, _) = scf_data.rimatr.as_ref().expect( - "This implementation requires cholesky decomposed ERI (or rimatr) to be available and stored in memory.", - ); - let cderi = rimatr.to_rstsr_view(&device).into_cow(); + let cderi = resp_auxbas::cderi(scf_data, config); resp_list.push(Box::new(URespRIJK::new_with_cderi(factor_j, factor_k, cderi))); } @@ -368,10 +367,7 @@ pub fn uscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> URe // object reusing the full-range implementation; it evaluates on the short-range `rimatr_sr` // ERI, with no Coulomb part (factor_j = 0). if let Some((_omega, factor_k_sr)) = rsh { - let (rimatr_sr, _, _) = scf_data.rimatr_sr.as_ref().expect( - "The range-separated response requires the short-range ERI (rimatr_sr) to be built and stored in memory.", - ); - let cderi_sr = rimatr_sr.to_rstsr_view(&device).into_cow(); + let cderi_sr = resp_auxbas::cderi_sr(scf_data, config); resp_list.push(Box::new(URespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr))); } diff --git a/src/ri_jk/decompose.rs b/src/ri_jk/decompose.rs index 6d293c709..b1cce0f13 100644 --- a/src/ri_jk/decompose.rs +++ b/src/ri_jk/decompose.rs @@ -136,16 +136,34 @@ pub enum J2CDecompose { /// - by Cholesky way, solve j3c to cderi does not require extra memory /// - by eigen way, the matmul will cost at most `2 * naux^2` or 4% of the final cderi matrix. pub fn generate_rimatr_bare(mol_obj: &Molecule, omega: Option) -> MatrixFull { - let mut mol = util::get_cint_mol(mol_obj); - let mut aux = util::get_cint_aux(mol_obj); + generate_rimatr_bare_on_cint( + &util::get_cint_mol(mol_obj), + &util::get_cint_aux(mol_obj), + mol_obj.ctrl.j2c_decomp, + omega, + ) +} + +/// Generate decomposed 3c-2e ERI (cderi/rimatr) on explicitly given `CInt` integrators. +/// +/// The molecular-orbital side (`mol`) and the auxiliary side (`aux`) are supplied directly, +/// together with the j2c decomposition setting `j2c_decomp_option` and the range-separation +/// parameter `omega`, so the rimatr can be evaluated on basis combinations other than the +/// molecule's own (e.g. the response-specific auxiliary basis of `resp_auxbas_path`). +pub fn generate_rimatr_bare_on_cint( + mol: &CInt, + aux: &CInt, + j2c_decomp_option: J2CDecompOption, + omega: Option, +) -> MatrixFull { + let mut mol = mol.clone(); + let mut aux = aux.clone(); if let Some(omega) = omega { mol.set_omega(omega); aux.set_omega(omega); } - let j2c_decomp_option = mol_obj.ctrl.j2c_decomp; - let device = DeviceBLAS::default(); let j2c_decomp = get_j2c_decomp(&aux, &device, j2c_decomp_option); diff --git a/src/ri_jk/mod.rs b/src/ri_jk/mod.rs index 89854babb..4196ac6ed 100644 --- a/src/ri_jk/mod.rs +++ b/src/ri_jk/mod.rs @@ -39,6 +39,9 @@ pub mod hess_u; pub mod resp_r; pub mod resp_u; +// response-specific auxiliary basis (`resp_auxbas_path`) +pub mod resp_auxbas; + // generalized fock implementations pub mod gfock_r; diff --git a/src/ri_jk/resp_auxbas.rs b/src/ri_jk/resp_auxbas.rs new file mode 100644 index 000000000..22e118208 --- /dev/null +++ b/src/ri_jk/resp_auxbas.rs @@ -0,0 +1,131 @@ +//! Response-specific auxiliary basis of the `resp_auxbas_path` keyword. +//! +//! When set, the RI-JK response objects evaluate on decomposed ERIs freshly built on that +//! auxiliary basis instead of borrowing the SCF `rimatr`/`rimatr_sr`. This is a response-side +//! (A-tensor) substitution only: the energy-derivative (B-side) objects of the hessian/multipole +//! tasks keep the SCF auxiliary basis, so the CP-SCF solve becomes a mixed-representation +//! approximation — parallel to the coarser response grid of `grid_level_cpscf`. +//! +//! Each energy contribution component that requires a decomposed ERI has its own definition +//! here: [`cderi`] for the full-range RI-JK object, [`cderi_sr`] for the RSH short-range +//! object. With the keyword set, the auxiliary-side `CInt` is regenerated on the +//! response basis (the SCF molecule is only read, never modified) and substituted into the rimatr +//! builder, yielding an owned tensor; with the keyword unset, the SCF tensors are borrowed as +//! before (zero copy). + +use super::decompose::generate_rimatr_bare_on_cint; +use super::prelude_dev::*; +use super::util::{get_cint_aux_on_path, get_cint_mol}; +use crate::analdrv::config::AnalDrvConfig; +use crate::SCF; + +use log::debug; + +/// Validate `resp_auxbas_path` before any response object is built. +/// +/// This hard-errors on the unsupported `even_tempered_basis` combination (the even-tempered +/// machinery would silently override the response basis), and warns when the keyword is set but no +/// response (CP-SCF) object will be built (`need_resp = false`). The per-component definitions +/// below call it as well, so paths that build the response objects directly (e.g. the geomopt +/// analytic hessian) are covered too. +pub fn validate_resp_auxbas(scf_data: &SCF, config: &AnalDrvConfig, need_resp: bool) { + if config.resp.auxbas_path.is_none() { + return; + } + if scf_data.mol.ctrl.even_tempered_basis { + panic!("`resp_auxbas_path` is not supported together with `even_tempered_basis`."); + } + if !need_resp { + log::warn!( + "`resp_auxbas_path` is set, but no task of this run uses the response (CP-SCF) object; the keyword will be ignored." + ); + } +} + +/// Decomposed ERI of the full-range RI-JK response object, as a [`TsrCow`]. +/// +/// With `resp_auxbas_path` set, the rimatr is freshly built on that basis: the auxiliary-side +/// `CInt` is regenerated from a local copy of the keywords +/// ([`get_cint_aux_on_path`](crate::ri_jk::util::get_cint_aux_on_path)) and, together with the +/// molecule-side one, handed to the CInt-level rimatr builder +/// ([`generate_rimatr_bare_on_cint`](crate::ri_jk::generate_rimatr_bare_on_cint)), so the SCF +/// molecule stays untouched and the tensor is owned. A `resp_auxbas_path` equal to the SCF +/// `auxbas_path` (exact match) is treated as unset when the SCF `rimatr` is available: the +/// borrowed tensor is used instead of rebuilding an identical one. With the keyword unset, the +/// SCF `rimatr` is borrowed (zero copy). +pub fn cderi<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { + let device = DeviceBLAS::default(); + // a response basis equal to the SCF one, with the SCF tensor available, is treated as unset + let resp_auxbas_path = match config.resp.auxbas_path.as_deref() { + Some(path) if scf_data.mol.ctrl.auxbas_path == path && scf_data.rimatr.is_some() => { + debug!("`resp_auxbas_path` equals the SCF auxiliary basis; treated as unset."); + None + }, + path => path, + }; + match resp_auxbas_path { + Some(resp_auxbas_path) => { + validate_resp_auxbas(scf_data, config, true); + debug!("Building the response rimatr on `resp_auxbas_path` = \"{}\".", resp_auxbas_path); + let mol = get_cint_mol(&scf_data.mol); + let aux = get_cint_aux_on_path(&scf_data.mol, resp_auxbas_path); + let cderi = + generate_rimatr_bare_on_cint(&mol, &aux, scf_data.mol.ctrl.j2c_decomp, None).into_rstsr(&device); + debug!( + "RI-JK response object on the `resp_auxbas_path` decomposed ERI (owned, {} auxiliary functions).", + cderi.shape()[1] + ); + cderi.into_cow() + }, + None => { + debug!("RI-JK response object on the SCF rimatr (borrowed, zero copy)."); + let (rimatr, _, _) = scf_data.rimatr.as_ref().expect( + "This implementation requires cholesky decomposed ERI (or rimatr) to be available and stored in memory.", + ); + rimatr.to_rstsr_view(&device).into_cow() + }, + } +} + +/// Decomposed ERI of the RSH short-range RI-JK response object, as a [`TsrCow`]. +/// +/// Same definition as [`cderi`] on the short-range tensor: with `resp_auxbas_path` set, the +/// SR rimatr is freshly built (with `-omega` in libcint's convention, cf. the SCF integral +/// preparation), and a path equal to the SCF `auxbas_path` is treated as unset when the SCF +/// `rimatr_sr` is available; with the keyword unset, the SCF `rimatr_sr` is borrowed (zero +/// copy). +pub fn cderi_sr<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { + let device = DeviceBLAS::default(); + // a response basis equal to the SCF one, with the SCF tensor available, is treated as unset + let resp_auxbas_path = match config.resp.auxbas_path.as_deref() { + Some(path) if scf_data.mol.ctrl.auxbas_path == path && scf_data.rimatr_sr.is_some() => { + debug!("`resp_auxbas_path` equals the SCF auxiliary basis; treated as unset."); + None + }, + path => path, + }; + match resp_auxbas_path { + Some(resp_auxbas_path) => { + validate_resp_auxbas(scf_data, config, true); + // SR RI omega is negative in libcint's convention (cf. the SCF integral preparation) + let omega = scf_data.mol.xc_data.omega().unwrap(); + debug!( + "Building the short-range response rimatr for RSH (omega = {:.4}) on `resp_auxbas_path` = \"{}\".", + omega, resp_auxbas_path + ); + let mol = get_cint_mol(&scf_data.mol); + let aux = get_cint_aux_on_path(&scf_data.mol, resp_auxbas_path); + let cderi_sr = + generate_rimatr_bare_on_cint(&mol, &aux, scf_data.mol.ctrl.j2c_decomp, Some(-omega)).into_rstsr(&device); + debug!("RSH short-range response object on the `resp_auxbas_path` decomposed ERI (owned)."); + cderi_sr.into_cow() + }, + None => { + debug!("RSH short-range response object on the SCF rimatr_sr (borrowed, zero copy)."); + let (rimatr_sr, _, _) = scf_data.rimatr_sr.as_ref().expect( + "The range-separated response requires the short-range ERI (rimatr_sr) to be built and stored in memory.", + ); + rimatr_sr.to_rstsr_view(&device).into_cow() + }, + } +} diff --git a/src/ri_jk/util.rs b/src/ri_jk/util.rs index 4bcb893ac..2b9faeabd 100644 --- a/src/ri_jk/util.rs +++ b/src/ri_jk/util.rs @@ -6,6 +6,18 @@ pub fn get_cint_mol(mol_obj: &Molecule) -> CInt { } pub fn get_cint_aux(mol_obj: &Molecule) -> CInt { + get_cint_aux_on_path(mol_obj, &mol_obj.ctrl.auxbas_path) +} + +/// [`get_cint_aux`] on an auxiliary basis given by `auxbas_path` instead of the molecule's own +/// `ctrl.auxbas_path` (same value format; pool names are resolved by the basis machinery). +/// +/// The molecule is only read: the auxiliary-basis information is regenerated from a local copy of +/// the keywords, and nothing on `mol_obj` is modified. +pub fn get_cint_aux_on_path(mol_obj: &Molecule, auxbas_path: &str) -> CInt { + let mut ctrl = mol_obj.ctrl.clone(); + ctrl.auxbas_path = auxbas_path.to_owned(); + let etb = if mol_obj.ctrl.even_tempered_basis == true { let etb_elem = get_etb_elem(&mol_obj.geom, &mol_obj.ctrl.etb_start_atom_number); let etb_basis = etb_gen_for_atom_list(&mol_obj, &mol_obj.ctrl.etb_beta, &etb_elem); @@ -14,7 +26,7 @@ pub fn get_cint_aux(mol_obj: &Molecule) -> CInt { None }; - let (_, atm, bas, env, _, _, _) = Molecule::collect_auxbas(&mol_obj.ctrl, &mol_obj.geom, etb); + let (_, atm, bas, env, _, _, _) = Molecule::collect_auxbas(&ctrl, &mol_obj.geom, etb); let mut cint = CInt::new(); cint.initial_r2c(&atm, atm.len() as _, &bas, bas.len() as _, &env); cint.set_cint_type(mol_obj.cint_type); diff --git a/src/utilities/rstsr_util.rs b/src/utilities/rstsr_util.rs index b6ddf53c3..d5e8d000b 100644 --- a/src/utilities/rstsr_util.rs +++ b/src/utilities/rstsr_util.rs @@ -78,6 +78,17 @@ where } } +impl RestTensorIntoRstsrTsrAPI for MatrixFull +where + T: Clone, +{ + fn into_rstsr(self, device: &DeviceBLAS) -> Tsr { + let layout = layout_from_rest_tensor_matrix(&self); + let data = self.data; + rt::asarray((data, layout, device)) + } +} + impl RestTensorToRstsrTsrAPI for &[&MatrixFull] where T: Clone, -- Gitee From 5a79f165380d902957d2e151053a8aef7bb748fd Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 23 Sep 2026 16:03:18 +0800 Subject: [PATCH 04/21] Fix code-review findings on analdrv U response and hessian setup Correctness fixes from the pre-PR code review of the U response / resp_auxbas work; numerics of existing tested paths are unchanged. - analdrv interface and the geomopt analytic-hessian path reject the hessian task on post-SCF (fifth-DFA) methods outright before any response/grid preparation, instead of an opaque unwrap panic on the DFT grids freed by xdh_calculations; the unrestricted fifth-DFA multipole rejection is hoisted likewise. - URespSCF regains the CP-SCF timing breakdown lost in the trait extraction (solve_dimless_cpscf, response_dimless_cpscf, per-object in response_mo); the hessian driver splices it into its report. - URespSCF::response_mo accumulates in MO space (zeros_like of mo1), correct when the MO space is trimmed below nao by linear dependence. - The hessian setup no longer materializes the skeleton-grid fxc kernel (R and U): its chunk entry, full-grid assembly, and returned key were all unconsumed since the CP-KS response moved to the resp_* NIMatmul objects. Co-authored-by: Claude Code Co-authored-by: glm-5.3-flash[1m] --- src/analdrv/hessian/uscf.rs | 4 ++++ src/analdrv/interface.rs | 15 +++++++++++++++ src/analdrv/response/uresp_interface.rs | 19 ++++++++++++++++--- src/dft/numint_matmul/hess_rks.rs | 13 +------------ src/dft/numint_matmul/hess_uks.rs | 13 +------------ src/main_driver.rs | 5 +++++ 6 files changed, 42 insertions(+), 27 deletions(-) diff --git a/src/analdrv/hessian/uscf.rs b/src/analdrv/hessian/uscf.rs index 26844234a..f3059c006 100644 --- a/src/analdrv/hessian/uscf.rs +++ b/src/analdrv/hessian/uscf.rs @@ -257,6 +257,8 @@ impl<'a, 'b, 'c> UHessSCF<'a, 'b, 'c> { let mo_e1_α = b1mo_α.i(so[α]) + &mo1_α.i(so[α]) * &e_ij[α]; let mo_e1_β = b1mo_β.i(so[β]) + &mo1_β.i(so[β]) * &e_ij[β]; + // splice the response-object timing of the last-iteration `response_mo` call + self.timing.append(&mut self.resp.timing); self.timing.push(("finalize_cpscf".to_string(), t0.elapsed().as_secs_f64())); HashMap::from([("mo1_0", mo1_α), ("mo1_1", mo1_β), ("mo_e1_0", mo_e1_α), ("mo_e1_1", mo_e1_β)]) } @@ -313,6 +315,8 @@ impl<'a, 'b, 'c> UHessSCF<'a, 'b, 'c> { self.make_response_preparation(); let mo1 = self.resp.solve_dimless_cpscf(&rhs); + // splice the CP-SCF solve timing entries of the response object into the driver report + self.timing.append(&mut self.resp.timing); let mo1_view = [mo1[0].view(), mo1[1].view()]; let finalize_dict = self.finalize_cpscf(&f1mo, &s1mo, &mo1_view); let mo1 = [finalize_dict.get("mo1_0").unwrap().view(), finalize_dict.get("mo1_1").unwrap().view()]; diff --git a/src/analdrv/interface.rs b/src/analdrv/interface.rs index 0db06fb9c..afe6af464 100644 --- a/src/analdrv/interface.rs +++ b/src/analdrv/interface.rs @@ -77,6 +77,21 @@ pub fn analdrv_interface(scf_data: &mut SCF, tasks: &[AnalDrvTask], config: &Ana let has_multipole = tasks.iter().any(|task| matches!(task, AnalDrvTask::Multipole)); let has_hessian = tasks.iter().any(|task| matches!(task, AnalDrvTask::Hessian)); let is_fifth = scf_data.mol.xc_data.is_fifth_dfa(); + // the hessian task on post-SCF methods is rejected outright (the same guard as inside + // `hess_interface`, for direct callers); reject here already, so that no grid/response + // preparation is wasted on a run that cannot succeed + if is_fifth && has_hessian { + panic!("Normal modes calculation is currently not available for post-SCF methods."); + } + // the unrestricted multipole task is rejected unconditionally for post-SCF methods (in + // `multipole_interface_u`); reject here already, so that no grid/response preparation is + // wasted on a run that cannot succeed + if is_fifth && has_multipole && matches!(scf_data.scftype, SCFType::UHF) { + panic!( + "Multipole evaluation is not implemented for unrestricted post-SCF (double-hybrid, \ + fifth-DFA) methods yet; only SCF-level UHF/UKS is supported." + ); + } // the relaxed (Z-vector) DH multipole increments are the only multipole requirement on the // response object; the unrelaxed increments still need the DFT grids (regenerated below) let multipole_relaxed_dh = has_multipole && is_fifth && config.multipole.rdm1_relax == MultipoleRdm1Relax::Relaxed; diff --git a/src/analdrv/response/uresp_interface.rs b/src/analdrv/response/uresp_interface.rs index 15e54608c..ac8e17073 100644 --- a/src/analdrv/response/uresp_interface.rs +++ b/src/analdrv/response/uresp_interface.rs @@ -68,6 +68,13 @@ pub struct URespSCF<'a> { /// Orbital state stored by [`Self::make_cpscf_preparation`]; the CP-SCF machinery panics /// until it is set. cpscf_state: Option, + /// Timing information. Represented by wall time in second. + /// + /// Entries of the inherent CP-SCF machinery ([`Self::response_mo`], + /// [`Self::response_dimless_cpscf`], [`Self::solve_dimless_cpscf`]), pushed in call order. + /// The drivers splice them into their own timing report (e.g. the hessian driver, after the + /// CP-SCF solve and in `finalize_cpscf`). + pub timing: Vec<(String, f64)>, } /// Orbital state cached by [`URespSCF::make_cpscf_preparation`] for the inherent CP-SCF @@ -209,15 +216,17 @@ impl<'a> URespSCF<'a> { let mo_coeff = [state.mo_coeff[α].view(), state.mo_coeff[β].view()]; let ubra_α = &mo_coeff[α] % &mo1[α]; let ubra_β = &mo_coeff[β] % &mo1[β]; - let mut resp_α = rt::zeros_like(&ubra_α); - let mut resp_β = rt::zeros_like(&ubra_β); + let mut resp_α = rt::zeros_like(&mo1[α]); + let mut resp_β = rt::zeros_like(&mo1[β]); // the get_response_bra fan-out is inlined (not the composite trait method), keeping the // `cpscf_state` borrow disjoint from the `resp_list` iteration for resp_obj in self.resp_list.iter_mut() { + let t1 = std::time::Instant::now(); let el_resp = resp_obj.get_response_bra(&[ubra_α.view(), ubra_β.view()]); resp_α += mo_coeff[α].t() % &el_resp[α]; resp_β += mo_coeff[β].t() % &el_resp[β]; + self.timing.push((format!("in response_mo, {}", resp_obj.get_type_name()), t1.elapsed().as_secs_f64())); } [resp_α, resp_β] } @@ -239,6 +248,7 @@ impl<'a> URespSCF<'a> { /// /// - `resp` : shape `[nmo_s, nocc_s, ...]` per spin. The dimensionless response in MO space. pub fn response_dimless_cpscf(&mut self, mo1: &[TsrView; 2]) -> [Tsr; 2] { + let t0 = std::time::Instant::now(); let [α, β] = [0, 1]; let mut resp = self.response_mo(mo1); @@ -257,6 +267,7 @@ impl<'a> URespSCF<'a> { *&mut resp[β].i_mut(sv[β]) /= &state.e_ai_shift[β]; resp[α].i_mut(so[α]).fill(0.0); resp[β].i_mut(so[β]).fill(0.0); + self.timing.push(("response_dimless_cpscf".to_string(), t0.elapsed().as_secs_f64())); resp } @@ -277,6 +288,7 @@ impl<'a> URespSCF<'a> { /// - `mo1` : shape `[nmo_s, nocc_s, ...]` per spin. Perturbation in MO space that solves the /// dimensionless CP-SCF equation. pub fn solve_dimless_cpscf(&mut self, rhs: &[TsrView; 2]) -> [Tsr; 2] { + let t0 = std::time::Instant::now(); let [α, β] = [0, 1]; let rhs_shape = [rhs[α].shape().to_vec(), rhs[β].shape().to_vec()]; let nmo = [rhs[α].shape()[0], rhs[β].shape()[0]]; @@ -344,6 +356,7 @@ impl<'a> URespSCF<'a> { let [mo1_α, mo1_β] = unpack_flattened(mo1_flattened.view()); let mo1_α = mo1_α.into_shape(rhs_shape[α].to_vec()); let mo1_β = mo1_β.into_shape(rhs_shape[β].to_vec()); + self.timing.push(("solve_dimless_cpscf".to_string(), t0.elapsed().as_secs_f64())); [mo1_α, mo1_β] } } @@ -405,5 +418,5 @@ pub fn uscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> URe resp_list.push(Box::new(resp_obj)); } - URespSCF { resp_list, resp_cfg: config.resp.clone(), cpscf_state: None } + URespSCF { resp_list, resp_cfg: config.resp.clone(), cpscf_state: None, timing: Vec::new() } } diff --git a/src/dft/numint_matmul/hess_rks.rs b/src/dft/numint_matmul/hess_rks.rs index 046df166e..670a19b56 100644 --- a/src/dft/numint_matmul/hess_rks.rs +++ b/src/dft/numint_matmul/hess_rks.rs @@ -969,7 +969,7 @@ pub fn get_vmat_becke_parts( /// across chunks by a plain sum; grid-atom keys carry only the chunk atom's /// contribution and are scattered by [`make_hessian_setup_becke`]: /// -/// - Sum: `fxc [ngrids, nvar, nvar]` (disjoint grid ranges); `de_vxc_diag`, `de_vxc_off`, `de_fxc`, +/// - Sum: `de_vxc_diag`, `de_vxc_off`, `de_fxc`, /// `de_becke_full_1/2` `[3, 3, natm, natm]`; `vmat_ip [nao, nao, 3]`; `vmat_fxc`, `vmat_vxc`, /// `vmat_deriv1`, `vmat_becke_dw` `[nao, nao, 3, natm]`. /// - Scatter into column `B = atm_idx` (direction axes interchanged): `de_becke_atom_1/2`, @@ -1115,7 +1115,6 @@ pub fn make_hessian_setup_chunk_becke( } let mut result = HashMap::from([ - ("fxc", fxc), ("de_vxc_diag", de_vxc_diag), ("de_vxc_off", de_vxc_off), ("de_fxc", de_fxc), @@ -1200,8 +1199,6 @@ pub fn make_hessian_setup_becke( let nao = mol.nao(); let ngrids = ni.weights.len(); let nchunk = ni.nchunk; - let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); - let nvar = xc_type.num_nvar(); let device = dm0.device().clone(); // per-atom grid boundaries, deduced from the (atom-grouped) attribution @@ -1230,7 +1227,6 @@ pub fn make_hessian_setup_becke( let chunks = quad_split_by_atom(&atm_quad_split, ngrids, nchunk); let nchunks = chunks.len(); - let fxc_full: Tsr = rt::zeros(([ngrids, nvar, nvar], &device)); let de_fxc: Tsr = rt::zeros(([3, 3, natm, natm], &device)); let de_vxc_diag: Tsr = rt::zeros(([3, 3, natm, natm], &device)); let de_vxc_off: Tsr = rt::zeros(([3, 3, natm, natm], &device)); @@ -1272,12 +1268,6 @@ pub fn make_hessian_setup_becke( hardness, grid_shift, ); - // fxc: disjoint grid ranges - unsafe { - let fxc_slc = fxc_full.i(start..end); - let mut fxc_slc = fxc_slc.force_mut(); - fxc_slc.assign(&result_chunk["fxc"]); - } // sum the full-grid keys, scatter the grid-atom keys into the chunk // atom's column of the last (B) axis unsafe { @@ -1353,7 +1343,6 @@ pub fn make_hessian_setup_becke( }; let mut result = HashMap::from([ - ("fxc", fxc_full), ("de_vxc_diag", de_vxc_diag), ("de_vxc_off", de_vxc_off), ("de_fxc", de_fxc), diff --git a/src/dft/numint_matmul/hess_uks.rs b/src/dft/numint_matmul/hess_uks.rs index e93ec9105..5071c52ea 100644 --- a/src/dft/numint_matmul/hess_uks.rs +++ b/src/dft/numint_matmul/hess_uks.rs @@ -510,7 +510,7 @@ pub fn get_vmat_becke_parts_uks( /// across chunks by a plain sum; grid-atom keys carry only the chunk atom's /// contribution and are scattered by [`make_hessian_setup_becke_uks`]: /// -/// - Sum: `fxc [ngrids, nvar, 2, nvar, 2]` (disjoint grid ranges); `de_fxc`, `de_vxc_diag_a/b`, +/// - Sum: `de_fxc`, `de_vxc_diag_a/b`, /// `de_vxc_off_a/b`, `de_becke_full_1/2` `[3, 3, natm, natm]`; `vmat_ip_a/b [nao, nao, 3]`; /// `vmat_fxc_a/b`, `vmat_vxc_a/b`, `vmat_deriv1_a/b`, `vmat_becke_dw_a/b` `[nao, nao, 3, natm]`. /// - Scatter into column `B = atm_idx` (direction axes interchanged): `de_becke_atom_1/2`, @@ -690,7 +690,6 @@ pub fn make_hessian_setup_chunk_becke_uks( } let mut result = HashMap::from([ - ("fxc", fxc), ("de_fxc", de_fxc), ("de_vxc_diag_a", de_vxc_diag_α), ("de_vxc_diag_b", de_vxc_diag_β), @@ -777,8 +776,6 @@ pub fn make_hessian_setup_becke_uks( let nao = mol.nao(); let ngrids = ni.weights.len(); let nchunk = ni.nchunk; - let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); - let nvar = xc_type.num_nvar(); let device = dm0α.device().clone(); // per-atom grid boundaries, deduced from the (atom-grouped) attribution @@ -805,7 +802,6 @@ pub fn make_hessian_setup_becke_uks( let chunks = quad_split_by_atom(&atm_quad_split, ngrids, nchunk); let nchunks = chunks.len(); - let fxc_full: Tsr = rt::zeros(([ngrids, nvar, 2, nvar, 2], &device)); let de_fxc: Tsr = rt::zeros(([3, 3, natm, natm], &device)); let de_vxc_diag_α: Tsr = rt::zeros(([3, 3, natm, natm], &device)); let de_vxc_diag_β: Tsr = rt::zeros(([3, 3, natm, natm], &device)); @@ -857,12 +853,6 @@ pub fn make_hessian_setup_becke_uks( hardness, grid_shift, ); - // fxc: disjoint grid ranges - unsafe { - let fxc_slc = fxc_full.i(start..end); - let mut fxc_slc = fxc_slc.force_mut(); - fxc_slc.assign(&result_chunk["fxc"]); - } // sum the full-grid keys, scatter the grid-atom keys into the chunk // atom's column of the last (B) axis unsafe { @@ -951,7 +941,6 @@ pub fn make_hessian_setup_becke_uks( }; let mut result = HashMap::from([ - ("fxc", fxc_full), ("de_fxc", de_fxc), ("de_vxc_diag_a", de_vxc_diag_α), ("de_vxc_diag_b", de_vxc_diag_β), diff --git a/src/main_driver.rs b/src/main_driver.rs index 5a61d634f..821e2fd14 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -1214,6 +1214,11 @@ mod geometric_pyo3_impl { use rstsr::prelude::*; let config = scf_data.mol.ctrl.analdrv.clone().unwrap_or_default(); + // post-SCF (fifth-DFA) methods have no analytic hessian; reject before any + // response object is built (the DFT grids it would need are already freed) + if scf_data.mol.xc_data.is_fifth_dfa() { + panic!("Normal modes calculation is currently not available for post-SCF methods."); + } // shared response object for the hessian, per the SCF type let mut resp_objs = match scf_data.scftype { crate::scf_io::SCFType::RHF => Some(RespSCF::R(rscf_resp_interface(&scf_data, &config))), -- Gitee From 551cc7c3a7d2fbd9ab226ec6fc29d6d7ad45efa7 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 23 Sep 2026 16:40:13 +0800 Subject: [PATCH 05/21] Refactor code-review cleanups on analdrv response and hessian Cleanup pass from the pre-PR review of the U response / resp_auxbas work, applied on top of the correctness fixes; numerics unchanged, analdrv gate passes (response 5, multipole 11, hessian 18, 0 failed). - drop the byte-identical unrestricted skeleton-grid hessian setup; the restricted `scf_skeleton_nimatmul` (now pub(crate)) serves both spin treatments (the density-type determination is family-based) - extract the DFT response-grid policy (the `grid_level_cpscf` default and the reuse-vs-dedicated decision) into a shared `resp_grids`, and the SCF XC functional list into a spin-parametrized `scf_xc_func_list_with_spin` - split `validate_resp_auxbas`: the `even_tempered_basis` hard error is enforced at every response build via `assert_resp_auxbas_supported`; the unused-keyword warn stays task-level only - centralize the `RespSCF` variant extraction into `expect_r_mut`/`expect_u_mut` (hessian and multipole interfaces) - share the multipole print (header, dipole block, moment sections) between the restricted and unrestricted drivers - share the hessian matrix/timing print between the R/U interfaces - form the UKS rdm-form response densities in one batched `make_rho_from_dm` call (spin-interleaved list) instead of per-set allocations and copies Co-authored-by: Claude Code Co-authored-by: glm-5.3-flash[1m] --- src/analdrv/hessian/mod.rs | 52 ++++++++--- src/analdrv/hessian/rscf_interface.rs | 36 ++------ src/analdrv/hessian/uscf_interface.rs | 81 ++--------------- src/analdrv/multipole/interface.rs | 7 +- src/analdrv/multipole/rmultipole.rs | 116 +++++++++++++----------- src/analdrv/multipole/umultipole.rs | 52 +---------- src/analdrv/response/mod.rs | 27 +++++- src/analdrv/response/rresp_interface.rs | 73 +++++++++------ src/analdrv/response/uresp_interface.rs | 40 ++------ src/dft/numint_matmul/resp_uks.rs | 15 ++- src/ri_jk/resp_auxbas.rs | 31 +++++-- 11 files changed, 229 insertions(+), 301 deletions(-) diff --git a/src/analdrv/hessian/mod.rs b/src/analdrv/hessian/mod.rs index e34b1416b..b25507a74 100644 --- a/src/analdrv/hessian/mod.rs +++ b/src/analdrv/hessian/mod.rs @@ -37,6 +37,7 @@ use crate::analdrv::config::AnalDrvConfig; use crate::analdrv::response::RespSCF; use crate::scf_io::{SCFType, SCF}; use crate::thermo::ThermoResult; +use crate::utilities::rstsr_util::Tsr; /// Output of the analytical Hessian task, with the derived vibrational/thermochemical analysis. #[derive(Debug, Clone, Default)] @@ -54,6 +55,41 @@ pub struct HessOutput { pub thermo: Option, } +/// Print the hessian matrix and the driver timing report (shared by the restricted and +/// unrestricted interfaces), at `print_level >= 2`. +/// +/// The matrix is printed in `[tA, sB]` format (component xyz first, atom then), 6 columns at a +/// time with an index header; the timing entries follow in call order. +fn print_hessian_and_timing(de_hess: &Tsr, timing: &[(String, f64)], print_level: usize) { + if print_level < 2 { + return; + } + println!("=== HESSIAN ==="); + println!("Print hessian in [tA, sB] format (component xyz first, atom then)"); + println!(""); + // print hessian matrix [t, s, A, B] -> [tA, sB] + let natm = de_hess.shape()[3]; + let hess_mat = de_hess.transpose([0, 2, 1, 3]).into_shape((3 * natm, 3 * natm)); + // print 6 columns at a time, with index header + for j in (0..3 * natm).step_by(6) { + let j_end = (j + 6).min(3 * natm); + let col_header = " ".repeat(6) + &(j..j_end).map(|i| format!("{:>12}", i)).collect::(); + println!("{}", col_header); + for i in 0..3 * natm { + let row_str = format!("{i:>4} ") + + &(j..j_end).map(|j| format!("{:12.6}", hess_mat[[i, j]])).collect::(); + println!("{}", row_str); + } + println!(""); + } + + // print timing information + println!("Timing in Hessian calculation:"); + for (key, value) in timing.iter() { + println!(" {:60}: {:10.6} seconds", key, value); + } +} + pub fn hess_interface<'a>( scf_data: &'a SCF, config: &AnalDrvConfig, @@ -74,20 +110,12 @@ pub fn hess_interface<'a>( // carrier, holding the restricted/unrestricted variant per the SCF type). let (hessian, vib, _) = match scf_data.scftype { SCFType::RHF => { - let RespSCF::R(resp) = - resp_objs.expect("internal error: the RHF hessian requires the shared response object") - else { - panic!("internal error: the RHF hessian requires the restricted (R) response object variant"); - }; - rscf_hess_interface(scf_data, &config.nucgrad, resp) + let ctx = "internal error: the RHF hessian requires the shared response object"; + rscf_hess_interface(scf_data, &config.nucgrad, resp_objs.expect(ctx).expect_r_mut(ctx)) }, SCFType::UHF => { - let RespSCF::U(resp) = - resp_objs.expect("internal error: the UHF hessian requires the shared response object") - else { - panic!("internal error: the UHF hessian requires the unrestricted (U) response object variant"); - }; - uscf_hess_interface(scf_data, &config.nucgrad, resp) + let ctx = "internal error: the UHF hessian requires the shared response object"; + uscf_hess_interface(scf_data, &config.nucgrad, resp_objs.expect(ctx).expect_u_mut(ctx)) }, _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types."), }; diff --git a/src/analdrv/hessian/rscf_interface.rs b/src/analdrv/hessian/rscf_interface.rs index 10328cb25..19a46353c 100644 --- a/src/analdrv/hessian/rscf_interface.rs +++ b/src/analdrv/hessian/rscf_interface.rs @@ -1,4 +1,5 @@ use crate::analdrv::prelude::*; +use super::print_hessian_and_timing; use crate::analdrv::response::rresp_interface::{scf_jk_factors, scf_xc_func_list}; use crate::analdrv::vibration::vib::*; use crate::analdrv::vibration::vib_interface::*; @@ -141,34 +142,7 @@ pub fn rscf_hess_interface<'a>( let de_hess = hess_scf.make_hess(); - if scf_data.mol.ctrl.print_level >= 2 { - println!("=== HESSIAN ==="); - println!("Print hessian in [tA, sB] format (component xyz first, atom then)"); - println!(""); - // print hessian matrix [t, s, A, B] -> [tA, sB] - let natm = de_hess.shape()[3]; - let hess_mat = de_hess.transpose([0, 2, 1, 3]).into_shape((3 * natm, 3 * natm)); - // print 6 columns at a time, with index header - for j in (0..3 * natm).step_by(6) { - let j_end = (j + 6).min(3 * natm); - let col_header = " ".repeat(6) + &(j..j_end).map(|i| format!("{:>12}", i)).collect::(); - println!("{}", col_header); - for i in 0..3 * natm { - let row_str = format!("{i:>4} ") - + &(j..j_end).map(|j| format!("{:12.6}", hess_mat[[i, j]])).collect::(); - println!("{}", row_str); - } - println!(""); - } - } - - // print timing information - if scf_data.mol.ctrl.print_level >= 2 { - println!("Timing in Hessian calculation:"); - for (key, value) in hess_scf.timing.iter() { - println!(" {:60}: {:10.6} seconds", key, value); - } - } + print_hessian_and_timing(&de_hess, &hess_scf.timing, scf_data.mol.ctrl.print_level); de_hess }; @@ -180,6 +154,10 @@ pub fn rscf_hess_interface<'a>( /// The skeleton-level `NIMatmul` for the DFT XC hessian contribution, `None` for pure HF. /// +/// Shared by the unrestricted hessian interface: the grid policy is spin-independent (the +/// density-type determination depends only on the functional families, and the grid data itself +/// carries no spin treatment). +/// /// The skeleton grid level is the SCF grid level, raised by 2 for MGGA (TAU) functionals when /// the grid-shift derivative terms are off (the grid-shift terms restore the grid-related /// accuracy the finer grid compensated). The grid reuses the SCF grid when the level matches, @@ -187,7 +165,7 @@ pub fn rscf_hess_interface<'a>( /// `atm_idx`): the SCF grid is round-robin permuted for load balancing, while the Becke /// grid-shift attribution requires the ByAtom grouping. The regrouping only permutes, never /// changes values. -fn scf_skeleton_nimatmul<'a>(scf_data: &'a SCF, cfg: &AnalDrvNucgradCfg) -> Option> { +pub(crate) fn scf_skeleton_nimatmul<'a>(scf_data: &'a SCF, cfg: &AnalDrvNucgradCfg) -> Option> { if scf_data.mol.xc_data.dfa_compnt_scf.is_empty() { return None; } diff --git a/src/analdrv/hessian/uscf_interface.rs b/src/analdrv/hessian/uscf_interface.rs index eff0e3482..ba9371798 100644 --- a/src/analdrv/hessian/uscf_interface.rs +++ b/src/analdrv/hessian/uscf_interface.rs @@ -1,12 +1,10 @@ use crate::analdrv::prelude::*; +use super::print_hessian_and_timing; +use crate::analdrv::hessian::rscf_interface::scf_skeleton_nimatmul; use crate::analdrv::response::rresp_interface::scf_jk_factors; use crate::analdrv::response::uresp_interface::{scf_xc_func_list_uks, URespSCF}; use crate::analdrv::vibration::vib::*; use crate::analdrv::vibration::vib_interface::*; -use crate::dft::gen_grids::RadiiAdjust; -use crate::dft::numint_matmul::nimatmul::{regroup_grids_by_atom, NIMatmul}; -use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; -use crate::dft::Grids; use crate::dftd::hess::HessDFTD; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; @@ -118,10 +116,10 @@ pub fn uscf_hess_interface<'a>( // --- DFT --- // - // The skeleton-level grid of the XC hessian object is built by `scf_skeleton_nimatmul_uks` - // (SCF grid reused or regenerated per the skeleton level policy); no grid data is shared + // The skeleton-level grid of the XC hessian object is built by `scf_skeleton_nimatmul` (SCF + // grid reused or regenerated per the skeleton level policy); no grid data is shared // with the response objects. - let mut hess_nimatmul_obj = scf_skeleton_nimatmul_uks(scf_data, cfg).map(|ni| { + let mut hess_nimatmul_obj = scf_skeleton_nimatmul(scf_data, cfg).map(|ni| { use crate::dft::numint_matmul::hess_uks::UHessKSNIMatmul; let verbose = scf_data.mol.ctrl.print_level > 2; @@ -151,34 +149,7 @@ pub fn uscf_hess_interface<'a>( let de_hess = hess_scf.make_hess(); - if scf_data.mol.ctrl.print_level >= 2 { - println!("=== HESSIAN ==="); - println!("Print hessian in [tA, sB] format (component xyz first, atom then)"); - println!(""); - // print hessian matrix [t, s, A, B] -> [tA, sB] - let natm = de_hess.shape()[3]; - let hess_mat = de_hess.transpose([0, 2, 1, 3]).into_shape((3 * natm, 3 * natm)); - // print 6 columns at a time, with index header - for j in (0..3 * natm).step_by(6) { - let j_end = (j + 6).min(3 * natm); - let col_header = " ".repeat(6) + &(j..j_end).map(|i| format!("{:>12}", i)).collect::(); - println!("{}", col_header); - for i in 0..3 * natm { - let row_str = format!("{i:>4} ") - + &(j..j_end).map(|j| format!("{:12.6}", hess_mat[[i, j]])).collect::(); - println!("{}", row_str); - } - println!(""); - } - } - - // print timing information - if scf_data.mol.ctrl.print_level >= 2 { - println!("Timing in Hessian calculation:"); - for (key, value) in hess_scf.timing.iter() { - println!(" {:60}: {:10.6} seconds", key, value); - } - } + print_hessian_and_timing(&de_hess, &hess_scf.timing, scf_data.mol.ctrl.print_level); de_hess }; @@ -187,43 +158,3 @@ pub fn uscf_hess_interface<'a>( vibration_analysis_interface(scf_data, cfg, de_hess.view()) } - -/// The skeleton-level `NIMatmul` for the DFT XC hessian contribution, `None` for pure HF -/// (unrestricted sibling of `scf_skeleton_nimatmul` in -/// [`rscf_interface`](crate::analdrv::hessian::rscf_interface)). -/// -/// The skeleton grid level is the SCF grid level, raised by 2 for MGGA (TAU) functionals when -/// the grid-shift derivative terms are off (the grid-shift terms restore the grid-related -/// accuracy the finer grid compensated). The grid reuses the SCF grid when the level matches, -/// else is regenerated; either way it is regrouped to atom-grouped order (non-decreasing -/// `atm_idx`): the SCF grid is round-robin permuted for load balancing, while the Becke -/// grid-shift attribution requires the ByAtom grouping. The regrouping only permutes, never -/// changes values. -fn scf_skeleton_nimatmul_uks<'a>(scf_data: &'a SCF, cfg: &AnalDrvNucgradCfg) -> Option> { - if scf_data.mol.xc_data.dfa_compnt_scf.is_empty() { - return None; - } - let mol = get_cint_mol(&scf_data.mol); - - let xc_func_list = scf_xc_func_list_uks(scf_data); - let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); - let is_mgga = matches!(xc_type, XCDenType::TAU); - let grid_gen_level = scf_data.mol.ctrl.grid_gen_level; - let grid_shift = cfg.grid_shift_deriv; - let sk_level = - cfg.grid_level_skeleton.unwrap_or(if is_mgga && !grid_shift { grid_gen_level + 2 } else { grid_gen_level }); - - let (coordinates, weights, atm_idx, quadrature_weights) = if sk_level == grid_gen_level { - let grids = scf_data.grids.as_ref().unwrap(); - (grids.coordinates.clone(), grids.weights.clone(), grids.atm_idx.clone(), grids.quadrature_weights.clone()) - } else { - let sk_grid = Grids::build_with_level(&scf_data.mol, sk_level); - (sk_grid.coordinates, sk_grid.weights, sk_grid.atm_idx, sk_grid.quadrature_weights) - }; - let (coordinates, weights, atm_idx, quadrature_weights) = - regroup_grids_by_atom(coordinates, weights, atm_idx, quadrature_weights, mol.natm()); - // the Becke grid-shift terms of the hessian rebuild the partition from the atomic radii, so - // the driver must know the radii adjustment scheme the grid was generated with - let radii_adjust = RadiiAdjust::from_str(&scf_data.mol.ctrl.radii_adjust); - Some(NIMatmul::new(&mol, &coordinates, &weights, &atm_idx, &quadrature_weights).with_radii_adjust(radii_adjust)) -} diff --git a/src/analdrv/multipole/interface.rs b/src/analdrv/multipole/interface.rs index 0ca7540f9..a3b80079e 100644 --- a/src/analdrv/multipole/interface.rs +++ b/src/analdrv/multipole/interface.rs @@ -128,11 +128,8 @@ fn multipole_interface_r<'a>( // The relaxed (Z-vector) increments consume the restricted response object; the U variant of // the carrier cannot appear here (the unrestricted task never requests a response object). - let resp = resp_objs.map(|resp| match resp { - RespSCF::R(r) => r, - RespSCF::U(_) => panic!( - "Multipole evaluation is currently only implemented for restricted (RHF/RKS) calculations; an unrestricted response object was given." - ), + let resp = resp_objs.map(|resp| { + resp.expect_r_mut("Multipole evaluation is currently only implemented for restricted (RHF/RKS) calculations") }); let mut rmultipole = RMultipoleDH::new(&mol_cint, mo_coeff, mo_occ, origin, rgfock.as_mut(), resp); diff --git a/src/analdrv/multipole/rmultipole.rs b/src/analdrv/multipole/rmultipole.rs index 3138134fb..cd2584d4c 100644 --- a/src/analdrv/multipole/rmultipole.rs +++ b/src/analdrv/multipole/rmultipole.rs @@ -81,6 +81,66 @@ pub(crate) fn sorted_multi_indices(order: usize) -> Vec> { out } +/// Print the evaluated multipole moments from a driver `result` map, in the analdrv output style +/// (shared by the restricted and unrestricted drivers). +/// +/// Only sections whose entries exist in `result` are printed (so a dipole-only evaluation does +/// not print an empty quadrupole section, and contribution entries a driver does not hold — e.g. +/// the DH correlation/response increments of an SCF-level evaluation — are skipped). All +/// quantities are in atomic units. The print is gated by `print_level >= 1` (the timing output +/// stays at `>= 2`); the output style of this section is expected to be refactored later. +pub(crate) fn print_multipole_result(result: &HashMap, origin: [f64; 3], print_level: usize) { + if print_level < 1 { + return; + } + println!("=============== Multipole Moments (in analdrv) ==============="); + println!("Origin (a.u.): [{:14.8}, {:14.8}, {:14.8}]", origin[0], origin[1], origin[2]); + + if result.contains_key("dip_tot") { + println!("Electric dipole moment (a.u.):"); + for (key, label) in [ + ("dip_nuc", "nuclear charges"), + ("dip_scf", "SCF density"), + ("dip_corr", "unrelaxed corr. density"), + ("dip_resp", "response (Z-vector)"), + ("dip_tot", "total"), + ] { + if let Some(v) = result.get(key) { + println!(" {:<24}: {:16.12}", label, v); + } + } + // total dipole additionally in Debye (1 a.u. = 2.54174623 Debye) + if let Some(v) = result.get("dip_tot") { + let au2debye = crate::constants::AU2DEBYE; + println!( + " {:<24}: {:16.12} {:16.12} {:16.12}", + "total, Debye (X Y Z)", + v[[0]] * au2debye, + v[[1]] * au2debye, + v[[2]] * au2debye + ); + } + } + + print_moment_section(result, "Electric quadrupole moment, raw second moments (a.u.)", "quad_tot", 2); + print_moment_section(result, "Electric quadrupole moment, traceless (a.u.)", "quad_tot_traceless", 2); + print_moment_section(result, "Electric octupole moment (a.u.)", "oct_tot", 3); + print_moment_section(result, "Electric hexadecapole moment (a.u.)", "hex_tot", 4); +} + +/// Print one moment section from the `result` entry `key`, listing the unique components of the +/// permutation-symmetric order-`order` tensor. Does nothing if the key is absent. +pub(crate) fn print_moment_section(result: &HashMap, title: &str, key: &str, order: usize) { + let Some(tot) = result.get(key) else { return }; + let flat = tot.clone().into_shape(-1); + println!("{title}:"); + for idx in sorted_multi_indices(order) { + let label: String = idx.iter().map(|&t| ['X', 'Y', 'Z'][t]).collect(); + let c = idx.iter().enumerate().fold(0usize, |acc, (k, &t)| acc + t * 3usize.pow((order - 1 - k) as u32)); + println!(" {label}: {:16.12}", flat[[c]]); + } +} + /// Working solver and maintainer of multipole moment evaluation for restricted SCF (and /// double-hybrid, DH) methods. /// @@ -309,59 +369,9 @@ impl<'a, 'b> RMultipoleDH<'a, 'b> { /// Print the evaluated multipole moments in the analdrv output style. /// - /// Only sections whose entries exist in `result` are printed (so a dipole-only evaluation - /// does not print an empty quadrupole section). All quantities are in atomic units. The - /// print is gated by `print_level >= 1` (the timing output stays at `>= 2`); the output - /// style of this section is expected to be refactored later. + /// See [`print_multipole_result`](crate::analdrv::multipole::rmultipole::print_multipole_result) + /// for the format. pub fn print_multipole(&self, print_level: usize) { - if print_level < 1 { - return; - } - println!("=============== Multipole Moments (in analdrv) ==============="); - println!("Origin (a.u.): [{:14.8}, {:14.8}, {:14.8}]", self.origin[0], self.origin[1], self.origin[2]); - - if self.result.contains_key("dip_tot") { - println!("Electric dipole moment (a.u.):"); - for (key, label) in [ - ("dip_nuc", "nuclear charges"), - ("dip_scf", "SCF density"), - ("dip_corr", "unrelaxed corr. density"), - ("dip_resp", "response (Z-vector)"), - ("dip_tot", "total"), - ] { - if let Some(v) = self.result.get(key) { - println!(" {:<24}: {:16.12}", label, v); - } - } - // total dipole additionally in Debye (1 a.u. = 2.54174623 Debye) - if let Some(v) = self.result.get("dip_tot") { - let au2debye = crate::constants::AU2DEBYE; - println!( - " {:<24}: {:16.12} {:16.12} {:16.12}", - "total, Debye (X Y Z)", - v[[0]] * au2debye, - v[[1]] * au2debye, - v[[2]] * au2debye - ); - } - } - - self.print_moment_section("Electric quadrupole moment, raw second moments (a.u.)", "quad_tot", 2); - self.print_moment_section("Electric quadrupole moment, traceless (a.u.)", "quad_tot_traceless", 2); - self.print_moment_section("Electric octupole moment (a.u.)", "oct_tot", 3); - self.print_moment_section("Electric hexadecapole moment (a.u.)", "hex_tot", 4); - } - - /// Print one moment section from the `result` entry `key`, listing the unique components of - /// the permutation-symmetric order-`order` tensor. Does nothing if the key is absent. - fn print_moment_section(&self, title: &str, key: &str, order: usize) { - let Some(tot) = self.result.get(key) else { return }; - let flat = tot.clone().into_shape(-1); - println!("{title}:"); - for idx in sorted_multi_indices(order) { - let label: String = idx.iter().map(|&t| ['X', 'Y', 'Z'][t]).collect(); - let c = idx.iter().enumerate().fold(0usize, |acc, (k, &t)| acc + t * 3usize.pow((order - 1 - k) as u32)); - println!(" {label}: {:16.12}", flat[[c]]); - } + print_multipole_result(&self.result, self.origin, print_level); } } diff --git a/src/analdrv/multipole/umultipole.rs b/src/analdrv/multipole/umultipole.rs index 00eec820b..d1e67205a 100644 --- a/src/analdrv/multipole/umultipole.rs +++ b/src/analdrv/multipole/umultipole.rs @@ -1,6 +1,6 @@ //! Multipole moment driver for unrestricted SCF methods. -use super::rmultipole::{multipole_intor, multipole_order_spec, quadrupole_to_traceless, sorted_multi_indices}; +use super::rmultipole::{multipole_intor, multipole_order_spec, print_multipole_result, quadrupole_to_traceless}; use crate::analdrv::prelude::*; /// Working solver and maintainer of multipole moment evaluation for unrestricted SCF methods. @@ -146,52 +146,10 @@ impl UMultipoleDH { /// Print the evaluated multipole moments in the analdrv output style. /// - /// Only sections whose entries exist in `result` are printed (so a dipole-only evaluation - /// does not print an empty quadrupole section). All quantities are in atomic units. The - /// print is gated by `print_level >= 1` (the timing output stays at `>= 2`). + /// See [`print_multipole_result`](super::rmultipole::print_multipole_result) for the format. + /// This driver holds no DH correlation/response increments, so those dipole entries (which + /// the shared print tolerates being absent) never appear. pub fn print_multipole(&self, print_level: usize) { - if print_level < 1 { - return; - } - println!("=============== Multipole Moments (in analdrv) ==============="); - println!("Origin (a.u.): [{:14.8}, {:14.8}, {:14.8}]", self.origin[0], self.origin[1], self.origin[2]); - - if self.result.contains_key("dip_tot") { - println!("Electric dipole moment (a.u.):"); - for (key, label) in [("dip_nuc", "nuclear charges"), ("dip_scf", "SCF density"), ("dip_tot", "total")] { - if let Some(v) = self.result.get(key) { - println!(" {:<24}: {:16.12}", label, v); - } - } - // total dipole additionally in Debye (1 a.u. = 2.54174623 Debye) - if let Some(v) = self.result.get("dip_tot") { - let au2debye = crate::constants::AU2DEBYE; - println!( - " {:<24}: {:16.12} {:16.12} {:16.12}", - "total, Debye (X Y Z)", - v[[0]] * au2debye, - v[[1]] * au2debye, - v[[2]] * au2debye - ); - } - } - - self.print_moment_section("Electric quadrupole moment, raw second moments (a.u.)", "quad_tot", 2); - self.print_moment_section("Electric quadrupole moment, traceless (a.u.)", "quad_tot_traceless", 2); - self.print_moment_section("Electric octupole moment (a.u.)", "oct_tot", 3); - self.print_moment_section("Electric hexadecapole moment (a.u.)", "hex_tot", 4); - } - - /// Print one moment section from the `result` entry `key`, listing the unique components of - /// the permutation-symmetric order-`order` tensor. Does nothing if the key is absent. - fn print_moment_section(&self, title: &str, key: &str, order: usize) { - let Some(tot) = self.result.get(key) else { return }; - let flat = tot.clone().into_shape(-1); - println!("{title}:"); - for idx in sorted_multi_indices(order) { - let label: String = idx.iter().map(|&t| ['X', 'Y', 'Z'][t]).collect(); - let c = idx.iter().enumerate().fold(0usize, |acc, (k, &t)| acc + t * 3usize.pow((order - 1 - k) as u32)); - println!(" {label}: {:16.12}", flat[[c]]); - } + print_multipole_result(&self.result, self.origin, print_level); } } diff --git a/src/analdrv/response/mod.rs b/src/analdrv/response/mod.rs index 35c007fbc..86d442381 100644 --- a/src/analdrv/response/mod.rs +++ b/src/analdrv/response/mod.rs @@ -21,9 +21,10 @@ pub mod rgfock_interface; /// type) and passed by `&mut` to the property drivers (hessian, multipole), so that the response /// preparation and cached intermediates are shared across tasks. /// -/// This is a plain carrier with no façade methods: the restricted and unrestricted response -/// types ([`RRespSCF`]/[`URespSCF`]) have different signatures (spin arrays) and no unified driver -/// exists, so every consumer matches on the variant it supports. +/// The restricted and unrestricted response types ([`RRespSCF`]/[`URespSCF`]) have different +/// signatures (spin arrays) and no unified driver exists, so every consumer extracts the variant +/// it supports through [`Self::expect_r_mut`]/[`Self::expect_u_mut`]; apart from these +/// variant-accessors this is a plain carrier without façade methods. /// /// [`RRespSCF`]: rresp_interface::RRespSCF /// [`URespSCF`]: uresp_interface::URespSCF @@ -34,3 +35,23 @@ pub enum RespSCF<'a> { /// ([`uscf_resp_interface`](uresp_interface::uscf_resp_interface)). U(uresp_interface::URespSCF<'a>), } + +impl<'a> RespSCF<'a> { + /// The restricted response object ([`RespSCF::R`]), for consumers supporting restricted + /// treatments only; panics with `ctx` on the unrestricted variant. + pub fn expect_r_mut(&mut self, ctx: &str) -> &mut rresp_interface::RRespSCF<'a> { + match self { + RespSCF::R(resp) => resp, + RespSCF::U(_) => panic!("{ctx}: expected the restricted (R) response object variant"), + } + } + + /// The unrestricted response object ([`RespSCF::U`]), for consumers supporting unrestricted + /// treatments only; panics with `ctx` on the restricted variant. + pub fn expect_u_mut(&mut self, ctx: &str) -> &mut uresp_interface::URespSCF<'a> { + match self { + RespSCF::U(resp) => resp, + RespSCF::R(_) => panic!("{ctx}: expected the unrestricted (U) response object variant"), + } + } +} diff --git a/src/analdrv/response/rresp_interface.rs b/src/analdrv/response/rresp_interface.rs index 9ee0af66d..669728ec3 100644 --- a/src/analdrv/response/rresp_interface.rs +++ b/src/analdrv/response/rresp_interface.rs @@ -47,17 +47,23 @@ pub fn scf_jk_factors(scf_data: &SCF) -> (f64, f64, Option<(f64, f64)>) { (factor_j, factor_k, rsh) } -/// List of `(scale, functional)` pairs of the SCF XC functional (spin-unpolarized). -pub fn scf_xc_func_list(scf_data: &SCF) -> Vec<(f64, LibXCFunctional)> { +/// List of `(scale, functional)` pairs of the SCF XC functional, in the given libxc spin +/// treatment. +pub fn scf_xc_func_list_with_spin(scf_data: &SCF, spin: LibXCSpin) -> Vec<(f64, LibXCFunctional)> { let xc_code = &scf_data.mol.xc_data.dfa_compnt_scf; let xc_params = &scf_data.mol.xc_data.dfa_paramr_scf; xc_code .iter() .zip(xc_params.iter()) - .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, LibXCSpin::Unpolarized))) + .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, spin))) .collect_vec() } +/// List of `(scale, functional)` pairs of the SCF XC functional (spin-unpolarized). +pub fn scf_xc_func_list(scf_data: &SCF) -> Vec<(f64, LibXCFunctional)> { + scf_xc_func_list_with_spin(scf_data, LibXCSpin::Unpolarized) +} + /// The response (fock/response) objects of all electron-interaction contributions of a restricted /// SCF, as an owned list of `RRespAPI` trait objects. /// @@ -299,6 +305,40 @@ impl<'a> RRespSCF<'a> { } } +/// The DFT response grids: the common (fock-path) grid and the optional dedicated response grid. +/// +/// Shared by the restricted and unrestricted response interfaces (the grid policy is +/// spin-independent). Returns `(ni, ni_resp)`: +/// +/// - the common grid is the SCF grid as is; the fock path is a plain quadrature sum, so it does +/// not require the atom-grouped (ByAtom) ordering; +/// - the response grid is dedicated (built at `config.resp.grid_level`, by default the coarser +/// `grid_gen_level.max(3) - 2`) when that level differs from the SCF grid generation level, +/// and `None` when they coincide (the response then evaluates and caches on the common grid). +pub fn resp_grids<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> (NIMatmul<'a>, Option>) { + let mol = get_cint_mol(&scf_data.mol); + + let grids = scf_data.grids.as_ref().unwrap(); + let ni = NIMatmul::new(&mol, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); + + let grid_gen_level = scf_data.mol.ctrl.grid_gen_level; + let grid_resp_level = config.resp.grid_level.unwrap_or(grid_gen_level.max(3) - 2); + let ni_resp = if grid_resp_level == grid_gen_level { + None + } else { + let cpscf_grid = Grids::build_with_level(&scf_data.mol, grid_resp_level); + let ni_resp = NIMatmul::new( + &mol, + &cpscf_grid.coordinates, + &cpscf_grid.weights, + &cpscf_grid.atm_idx, + &cpscf_grid.quadrature_weights, + ); + Some(ni_resp) + }; + (ni, ni_resp) +} + /// Build the response (fock/response) objects for a converged restricted SCF. pub fn rscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> RRespSCF<'a> { let mut resp_list: Vec> = Vec::new(); @@ -326,32 +366,13 @@ pub fn rscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> RRe let is_hf = scf_data.mol.xc_data.dfa_compnt_scf.is_empty(); if !is_hf { - let mol = get_cint_mol(&scf_data.mol); let xc_func_list = scf_xc_func_list(scf_data); let verbose = scf_data.mol.ctrl.print_level > 2; - // common grid (fock path): the SCF grid as is; the fock path is a plain quadrature sum, - // so it does not require the atom-grouped (ByAtom) ordering. - let grids = scf_data.grids.as_ref().unwrap(); - let ni = NIMatmul::new(&mol, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); - - // response grid: when it coincides with the SCF grid, leave `ni_resp = None`; the - // response then evaluates (and caches) on the common grid. Otherwise build a dedicated - // (usually coarser) grid. - let grid_gen_level = scf_data.mol.ctrl.grid_gen_level; - let grid_resp_level = config.resp.grid_level.unwrap_or(grid_gen_level.max(3) - 2); - let resp_obj = if grid_resp_level == grid_gen_level { - RRespKSNIMatmul::new(xc_func_list, ni, verbose) - } else { - let cpscf_grid = Grids::build_with_level(&scf_data.mol, grid_resp_level); - let ni_resp = NIMatmul::new( - &mol, - &cpscf_grid.coordinates, - &cpscf_grid.weights, - &cpscf_grid.atm_idx, - &cpscf_grid.quadrature_weights, - ); - RRespKSNIMatmul::new(xc_func_list, ni, verbose).set_ni_resp(ni_resp) + let (ni, ni_resp) = resp_grids(scf_data, config); + let resp_obj = match ni_resp { + Some(ni_resp) => RRespKSNIMatmul::new(xc_func_list, ni, verbose).set_ni_resp(ni_resp), + None => RRespKSNIMatmul::new(xc_func_list, ni, verbose), }; resp_list.push(Box::new(resp_obj)); } diff --git a/src/analdrv/response/uresp_interface.rs b/src/analdrv/response/uresp_interface.rs index ac8e17073..c19401e46 100644 --- a/src/analdrv/response/uresp_interface.rs +++ b/src/analdrv/response/uresp_interface.rs @@ -15,13 +15,10 @@ use crate::analdrv::config::AnalDrvRespCfg; use crate::analdrv::prelude::*; -use crate::analdrv::response::rresp_interface::scf_jk_factors; +use crate::analdrv::response::rresp_interface::{resp_grids, scf_jk_factors}; use crate::ri_jk::resp_auxbas; -use crate::dft::numint_matmul::nimatmul::NIMatmul; use crate::dft::numint_matmul::resp_uks::URespKSNIMatmul; -use crate::dft::Grids; use crate::ri_jk::resp_u::URespRIJK; -use crate::ri_jk::util::get_cint_mol; use crate::SCF; use libxc::prelude::*; @@ -32,13 +29,7 @@ use libxc::prelude::*; /// scf_xc_func_list): the functionals are created with [`LibXCSpin::Polarized`], as required by /// the UKS rho/vxc/fxc spin-block layout. pub fn scf_xc_func_list_uks(scf_data: &SCF) -> Vec<(f64, LibXCFunctional)> { - let xc_code = &scf_data.mol.xc_data.dfa_compnt_scf; - let xc_params = &scf_data.mol.xc_data.dfa_paramr_scf; - xc_code - .iter() - .zip(xc_params.iter()) - .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, LibXCSpin::Polarized))) - .collect_vec() + crate::analdrv::response::rresp_interface::scf_xc_func_list_with_spin(scf_data, LibXCSpin::Polarized) } /// The response (fock/response) objects of all electron-interaction contributions of an @@ -388,32 +379,13 @@ pub fn uscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> URe let is_hf = scf_data.mol.xc_data.dfa_compnt_scf.is_empty(); if !is_hf { - let mol = get_cint_mol(&scf_data.mol); let xc_func_list = scf_xc_func_list_uks(scf_data); let verbose = scf_data.mol.ctrl.print_level > 2; - // common grid (fock path): the SCF grid as is; the fock path is a plain quadrature sum, - // so it does not require the atom-grouped (ByAtom) ordering. - let grids = scf_data.grids.as_ref().unwrap(); - let ni = NIMatmul::new(&mol, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); - - // response grid: when it coincides with the SCF grid, leave `ni_resp = None`; the - // response then evaluates (and caches) on the common grid. Otherwise build a dedicated - // (usually coarser) grid. - let grid_gen_level = scf_data.mol.ctrl.grid_gen_level; - let grid_resp_level = config.resp.grid_level.unwrap_or(grid_gen_level.max(3) - 2); - let resp_obj = if grid_resp_level == grid_gen_level { - URespKSNIMatmul::new(xc_func_list, ni, verbose) - } else { - let cpscf_grid = Grids::build_with_level(&scf_data.mol, grid_resp_level); - let ni_resp = NIMatmul::new( - &mol, - &cpscf_grid.coordinates, - &cpscf_grid.weights, - &cpscf_grid.atm_idx, - &cpscf_grid.quadrature_weights, - ); - URespKSNIMatmul::new(xc_func_list, ni, verbose).set_ni_resp(ni_resp) + let (ni, ni_resp) = resp_grids(scf_data, config); + let resp_obj = match ni_resp { + Some(ni_resp) => URespKSNIMatmul::new(xc_func_list, ni, verbose).set_ni_resp(ni_resp), + None => URespKSNIMatmul::new(xc_func_list, ni, verbose), }; resp_list.push(Box::new(resp_obj)); } diff --git a/src/dft/numint_matmul/resp_uks.rs b/src/dft/numint_matmul/resp_uks.rs index 08e5998df..4fe04a33b 100644 --- a/src/dft/numint_matmul/resp_uks.rs +++ b/src/dft/numint_matmul/resp_uks.rs @@ -301,16 +301,15 @@ impl<'a> URespAPI for URespKSNIMatmul<'a> { ((&rdm[β] + &rdm[β].swapaxes(0, 1)) * 0.5).into_shape((nao, nao, nset)), ]; - // rho1 : [ngrids, nvar, 2, nset], assembled set by set from the per-spin dm lists + // rho1 : [ngrids, nvar, 2, nset], formed in one batched `make_rho_from_dm` call over the + // spin-interleaved dm list: the column-major output splits its last (set) axis as + // `k = spin + 2 * set`, so the list must alternate the spins within each set + let dm_list: Vec = (0..nset) + .flat_map(|k| [rdm_sym[α].i((.., .., k)), rdm_sym[β].i((.., .., k))]) + .collect(); let ngrids = self.ni.weights.len(); let nvar = den_type.num_nvar(); - let device = rdm[α].device().clone(); - let mut rho1 = rt::zeros(([ngrids, nvar, 2, nset], &device)); - for k in 0..nset { - let rho1_k = - self.ni.make_rho_from_dm(&[rdm_sym[α].i((.., .., k)), rdm_sym[β].i((.., .., k))], den_type); - rho1.i_mut((.., .., .., k)).assign(&rho1_k); - } + let rho1 = self.ni.make_rho_from_dm(&dm_list, den_type).into_shape((ngrids, nvar, 2, nset)); let pot = self.ni.make_fxc_pot_with_eff(fxc, rho1.view(), den_type, XCSpin::Polarized); // restore the input's per-spin trailing shape diff --git a/src/ri_jk/resp_auxbas.rs b/src/ri_jk/resp_auxbas.rs index 22e118208..d27b44664 100644 --- a/src/ri_jk/resp_auxbas.rs +++ b/src/ri_jk/resp_auxbas.rs @@ -21,20 +21,33 @@ use crate::SCF; use log::debug; -/// Validate `resp_auxbas_path` before any response object is built. +/// Hard-error on the unsupported `even_tempered_basis` combination. /// -/// This hard-errors on the unsupported `even_tempered_basis` combination (the even-tempered -/// machinery would silently override the response basis), and warns when the keyword is set but no -/// response (CP-SCF) object will be built (`need_resp = false`). The per-component definitions -/// below call it as well, so paths that build the response objects directly (e.g. the geomopt -/// analytic hessian) are covered too. -pub fn validate_resp_auxbas(scf_data: &SCF, config: &AnalDrvConfig, need_resp: bool) { +/// The even-tempered machinery would silently override the response basis, so the keyword +/// combination is rejected. Checked once up front by the task-level +/// [`validate_resp_auxbas`], and again at every response-object build +/// ([`cderi`]/[`cderi_sr`]), so paths that build the response objects directly (e.g. the +/// geomopt analytic hessian) are covered too. +pub fn assert_resp_auxbas_supported(scf_data: &SCF, config: &AnalDrvConfig) { if config.resp.auxbas_path.is_none() { return; } if scf_data.mol.ctrl.even_tempered_basis { panic!("`resp_auxbas_path` is not supported together with `even_tempered_basis`."); } +} + +/// Validate `resp_auxbas_path` against this run's tasks, before any response object is built. +/// +/// This is the task-level check: it warns when the keyword is set but no response (CP-SCF) +/// object will be built (`need_resp = false`), and enforces the hard error of +/// [`assert_resp_auxbas_supported`]. The per-component definitions call the latter themselves, +/// but only this function carries the warn (only the caller knows `need_resp`). +pub fn validate_resp_auxbas(scf_data: &SCF, config: &AnalDrvConfig, need_resp: bool) { + if config.resp.auxbas_path.is_none() { + return; + } + assert_resp_auxbas_supported(scf_data, config); if !need_resp { log::warn!( "`resp_auxbas_path` is set, but no task of this run uses the response (CP-SCF) object; the keyword will be ignored." @@ -65,7 +78,7 @@ pub fn cderi<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { }; match resp_auxbas_path { Some(resp_auxbas_path) => { - validate_resp_auxbas(scf_data, config, true); + assert_resp_auxbas_supported(scf_data, config); debug!("Building the response rimatr on `resp_auxbas_path` = \"{}\".", resp_auxbas_path); let mol = get_cint_mol(&scf_data.mol); let aux = get_cint_aux_on_path(&scf_data.mol, resp_auxbas_path); @@ -106,7 +119,7 @@ pub fn cderi_sr<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { }; match resp_auxbas_path { Some(resp_auxbas_path) => { - validate_resp_auxbas(scf_data, config, true); + assert_resp_auxbas_supported(scf_data, config); // SR RI omega is negative in libcint's convention (cf. the SCF integral preparation) let omega = scf_data.mol.xc_data.omega().unwrap(); debug!( -- Gitee From eb384d3408bacfd371bbae877ee47900ab74faec Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 23 Sep 2026 16:43:22 +0800 Subject: [PATCH 06/21] add readme --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 11647904e..8d57e56e6 100644 --- a/README.md +++ b/README.md @@ -930,6 +930,7 @@ analdrv_tasks = ["multipole", "hessian"] - `cpscf_lindep`：CP-SCF 中一些数值过程的数值精度阈值。默认 1e-15，无量纲。 - `cpscf_tol_inflation`：容忍系数。若 Krylov 真残差 `||r|| < factor * tol`，接受该解而不触发 per-root 求解。缺省为1000.0。 - `grid_level_cpscf`：CP-SCF 中响应路径 (响应/A 张量收缩的 DFT 计算；fock 路径仍用 SCF 格点) 的 DFT 格点级别。仅影响 numint_matmul 后端实现。默认为 None，是 `[ctrl]` 中 grid_generation_level 关键词设定值减 2 (SCF 默认格点级别是 3，对应 Hessian 的级别是 1)；显式取值不设下限。 +- `resp_auxbas_path`：响应计算辅助基路径。默认为 None，即使用 `[ctrl]` 中 auxbas_path 设定的辅助基组。若显式指定，则在 CP-SCF/Z-Vector/response 计算中使用对应辅助基。该选项通常用于指定更小的辅助基以加速 CP-SCF 计算，但作为代价会降低 CP-SCF 的精度且增加内存占用。 这些关键词的旧名称 `cphf_*` 前缀 (`cphf_tol`、`cphf_lindep`、`cphf_tol_inflation` 等) 目前仍然作为别名被接受；响应格点则同时接受其旧名称 `grid_level_cphf`。 -- Gitee From f272190bad00a4f25a431e1c63cff717a6f5d9ee Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 23 Sep 2026 22:52:58 +0800 Subject: [PATCH 07/21] Add per-call precision parameter to response objects for resource selection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a trailing `prec: bool` to the computation methods of RRespAPI/URespAPI: `true` selects the SCF-grade resource (the SCF rimatr for RI-JK, the common grid for DFT), `false` the dedicated low-precision response resource when attached (falling back to the high-precision one otherwise). The CP-SCF machinery (`make_cpscf_preparation`, `response_mo`, including the last-iteration W assembly) contracts in low precision; every other caller passes high — in particular the PT2 rdm-form A-contraction, where a future input keyword may open the low-precision window as well. R/URespRIJK hold the SCF-basis `cderi` plus an optional `cderi_resp` (`resp_auxbas_path`, attached through the `set_cderi_resp` builder), selected per call by `cderi_in_use`. The new `resp_auxbas::cderi_pair`/`cderi_pair_sr` return both tensors, so the fock forms keep the SCF auxiliary basis while only the CP-SCF response contractions evaluate on the response basis. R/URespKSNIMatmul select `ni`/`ni_resp` per call by the same flag, replacing the per-method structural split, and cache under precision-keyed `vxc_{high,low}`/`fxc_{high,low}` entries; the fxc entry is shared between the preparation and the rdm-form contraction (revalidated on orbital change). Add a precision-plumbing test to tests/analdrv/response/uresp_gfock.rs (attached-but-duplicate resource agrees across both precisions and with the unattached object); analdrv gate 35/0. Co-authored-by: Claude Code Co-authored-by: glm-5.3 --- src/analdrv/config.rs | 22 ++--- src/analdrv/response/rresp_interface.rs | 65 +++++++++------ src/analdrv/response/trait_rresp.rs | 31 +++++-- src/analdrv/response/trait_uresp.rs | 25 ++++-- src/analdrv/response/uresp_interface.rs | 87 ++++++++++++-------- src/dft/numint_matmul/resp_rks.rs | 104 ++++++++++++----------- src/dft/numint_matmul/resp_uks.rs | 105 +++++++++++++----------- src/ri_jk/resp_auxbas.rs | 94 +++++++++++---------- src/ri_jk/resp_r.rs | 62 ++++++++++---- src/ri_jk/resp_u.rs | 62 ++++++++++---- src/ri_pt2/rgfock_pt2.rs | 5 +- tests/analdrv/response/uresp_gfock.rs | 94 +++++++++++++++++---- 12 files changed, 488 insertions(+), 268 deletions(-) diff --git a/src/analdrv/config.rs b/src/analdrv/config.rs index bab90a44b..6cf57b2d1 100644 --- a/src/analdrv/config.rs +++ b/src/analdrv/config.rs @@ -46,11 +46,11 @@ impl Default for AnalDrvGeneralCfg { /// /// All these keywords are specific to the CP-SCF-type solve, hence the `cpscf_*` key names. /// `grid_level` is not a solver setting but the DFT grid of the CP-SCF response evaluation (the -/// response path; the fock path and the DH generalized Fock stay on the SCF grid), so it is named -/// in the `grid_level_*` family of [`AnalDrvNucgradCfg::grid_level_skeleton`] with `cpscf` as the -/// scope. Likewise, `auxbas_path` configures the response objects themselves (not only the -/// CP-SCF solve — every consumer of the response object shares it), so it is named in the `resp_` -/// scope. +/// low-precision response path; the fock path and the DH generalized Fock stay on the SCF grid), +/// so it is named in the `grid_level_*` family of +/// [`AnalDrvNucgradCfg::grid_level_skeleton`] with `cpscf` as the scope. Likewise, `auxbas_path` +/// configures the response objects themselves (the low-precision resource, not only the CP-SCF +/// solve), so it is named in the `resp_` scope. /// /// The legacy `cphf_*` prefixed key names are still accepted as aliases. #[serde_inline_default] @@ -81,16 +81,18 @@ pub struct AnalDrvRespCfg { #[serde(rename = "grid_level_cpscf", alias = "grid_level_cphf")] #[serde_inline_default(None)] pub grid_level: Option, - /// Auxiliary basis of the RI-JK response objects (the response/A-tensor side), as a + /// Auxiliary basis of the RI-JK response objects (the low-precision response side), as a /// basis-set pool name or an element-JSON directory — the same value format as /// `ctrl.auxbas_path`. /// /// By default `None`: the response objects reuse the SCF auxiliary basis, borrowing the SCF /// `rimatr` (and `rimatr_sr` for range-separated hybrids) with no copy. When set, a - /// decomposed ERI is freshly built on this basis for the response objects only — the - /// energy-derivative (B-side) objects of the hessian/multipole tasks keep the SCF auxiliary - /// basis, so the CP-SCF solve becomes a mixed-representation approximation (parallel to - /// `grid_level_cpscf`). Not supported together with `even_tempered_basis`. + /// decomposed ERI is freshly built on this basis and attached as the low-precision + /// (`prec = false`) resource: it feeds the CP-SCF response contractions only, while the fock + /// forms (`prec = true`) and the energy-derivative (B-side) objects of the hessian/multipole + /// tasks keep the SCF auxiliary basis, so the CP-SCF solve becomes a mixed-representation + /// approximation (parallel to `grid_level_cpscf`). Not supported together with + /// `even_tempered_basis`. #[serde(rename = "resp_auxbas_path")] #[serde_inline_default(None)] pub auxbas_path: Option, diff --git a/src/analdrv/response/rresp_interface.rs b/src/analdrv/response/rresp_interface.rs index 669728ec3..f05218544 100644 --- a/src/analdrv/response/rresp_interface.rs +++ b/src/analdrv/response/rresp_interface.rs @@ -7,11 +7,14 @@ //! future standalone response/property driver can call it directly, without any hessian //! machinery. //! -//! The DFT XC response object evaluates the fock path directly on the SCF grid (in its native -//! round-robin atom-interleaved order — the fock path is a plain quadrature sum, so no -//! atom-regrouping is applied), and the response path on a dedicated (usually coarser) cpscf -//! grid. The hessian-side skeleton grid policy (including the MGGA level bump) lives in the -//! hessian interface; no grid identity or grid data is shared between the hessian and response +//! The precision of the response contractions is selected per call through the `prec` flag of +//! [`RRespAPI`]: the high-precision (`prec = true`) resource is the SCF-grade one — the SCF +//! `rimatr` for RI-JK, and the SCF grid for DFT (in its native round-robin atom-interleaved +//! order, since the fock path is a plain quadrature sum, so no atom-regrouping is applied) — +//! while the low-precision (`prec = false`) resource of the CP-SCF machinery is the dedicated +//! (usually coarser) response grid / the `resp_auxbas_path` basis when attached. The +//! hessian-side skeleton grid policy (including the MGGA level bump) lives in the hessian +//! interface; no grid identity or grid data is shared between the hessian and response //! subsystems. use crate::analdrv::config::AnalDrvRespCfg; @@ -108,20 +111,20 @@ pub struct RCpscfState { impl<'a> AnalDrvBaseAPI for RRespSCF<'a> {} impl<'a> RRespAPI for RRespSCF<'a> { - fn get_fock_rdm(&mut self, rdm: TsrView) -> Tsr { + fn get_fock_rdm(&mut self, rdm: TsrView, prec: bool) -> Tsr { let mut fock = rt::zeros_like(&rdm); for resp_obj in self.resp_list.iter_mut() { - fock += resp_obj.get_fock_rdm(rdm.view()); + fock += resp_obj.get_fock_rdm(rdm.view(), prec); } fock } - fn get_fock_coeff(&mut self, mo_coeff: TsrView, mo_occ: TsrView) -> Tsr { + fn get_fock_coeff(&mut self, mo_coeff: TsrView, mo_occ: TsrView, prec: bool) -> Tsr { // delegated per element (not through the default rdm route), so that elements overriding // this function for efficiency keep their advantage let mut fock = None; for resp_obj in self.resp_list.iter_mut() { - let fock_obj = resp_obj.get_fock_coeff(mo_coeff.view(), mo_occ.view()); + let fock_obj = resp_obj.get_fock_coeff(mo_coeff.view(), mo_occ.view(), prec); fock = Some(match fock { Some(fock) => fock + fock_obj, None => fock_obj, @@ -130,24 +133,24 @@ impl<'a> RRespAPI for RRespSCF<'a> { fock.expect("RRespSCF must hold at least one response object") } - fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView) { + fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView, prec: bool) { for resp_obj in self.resp_list.iter_mut() { - resp_obj.make_response_preparation(mo_coeff.view(), mo_occ.view()); + resp_obj.make_response_preparation(mo_coeff.view(), mo_occ.view(), prec); } } - fn get_response_rdm(&mut self, rdm: TsrView) -> Tsr { + fn get_response_rdm(&mut self, rdm: TsrView, prec: bool) -> Tsr { let mut resp = rt::zeros_like(&rdm); for resp_obj in self.resp_list.iter_mut() { - resp += resp_obj.get_response_rdm(rdm.view()); + resp += resp_obj.get_response_rdm(rdm.view(), prec); } resp } - fn get_response_bra(&mut self, bra: TsrView) -> Tsr { + fn get_response_bra(&mut self, bra: TsrView, prec: bool) -> Tsr { let mut resp = rt::zeros_like(&bra); for resp_obj in self.resp_list.iter_mut() { - resp += resp_obj.get_response_bra(bra.view()); + resp += resp_obj.get_response_bra(bra.view(), prec); } resp } @@ -167,8 +170,10 @@ impl<'a> RRespSCF<'a> { /// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. /// - `mo_energy` : shape `[nmo]`. Molecular orbital energies. pub fn make_cpscf_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView, mo_energy: TsrView) { + // the CP-SCF core runs the response machinery in low precision (`prec = false`): the + // dedicated response grid / `resp_auxbas_path` basis when attached for resp_obj in self.resp_list.iter_mut() { - resp_obj.make_response_preparation(mo_coeff.view(), mo_occ.view()); + resp_obj.make_response_preparation(mo_coeff.view(), mo_occ.view(), false); } let occidx = mo_occ.view().greater(0).into_vec(); @@ -219,7 +224,10 @@ impl<'a> RRespSCF<'a> { let ubra = &mo_coeff % &mo1; let mut resp = rt::zeros_like(&mo1); for resp_obj in self.resp_list.iter_mut() { - resp += mo_coeff.t() % resp_obj.get_response_bra(ubra.view()); + // low precision: the CP-SCF core contracts the response on the dedicated + // low-precision resource when attached (this includes the last-iteration W-operator + // assembly of the hessian driver) + resp += mo_coeff.t() % resp_obj.get_response_bra(ubra.view(), false); } resp } @@ -347,19 +355,30 @@ pub fn rscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> RRe let (factor_j, factor_k, rsh) = scf_jk_factors(scf_data); - // decomposed ERI of the RI-JK response object: its own definition — the freshly built one of - // `resp_auxbas_path` when set (owned), otherwise the SCF rimatr (borrowed, zero copy) + // decomposed ERIs of the RI-JK response object, per precision: the fock path always borrows + // the SCF rimatr (zero copy); the low-precision response path attaches the freshly built one + // of `resp_auxbas_path` when set (owned), and reuses the SCF rimatr otherwise { - let cderi = resp_auxbas::cderi(scf_data, config); - resp_list.push(Box::new(RRespRIJK::new_with_cderi(factor_j, factor_k, cderi))); + let (cderi, cderi_resp) = resp_auxbas::cderi_pair(scf_data, config); + let resp_obj = match cderi_resp { + Some(cderi_resp) => RRespRIJK::new_with_cderi(factor_j, factor_k, cderi).set_cderi_resp(cderi_resp), + None => RRespRIJK::new_with_cderi(factor_j, factor_k, cderi), + }; + resp_list.push(Box::new(resp_obj)); } // The short-range exchange correction (range-separated hybrids) is a separate response // object reusing the full-range implementation; it evaluates on the short-range `rimatr_sr` // ERI, with no Coulomb part (factor_j = 0). if let Some((_omega, factor_k_sr)) = rsh { - let cderi_sr = resp_auxbas::cderi_sr(scf_data, config); - resp_list.push(Box::new(RRespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr))); + let (cderi_sr, cderi_resp_sr) = resp_auxbas::cderi_pair_sr(scf_data, config); + let resp_obj = match cderi_resp_sr { + Some(cderi_resp_sr) => { + RRespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr).set_cderi_resp(cderi_resp_sr) + }, + None => RRespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr), + }; + resp_list.push(Box::new(resp_obj)); } // --- DFT --- // diff --git a/src/analdrv/response/trait_rresp.rs b/src/analdrv/response/trait_rresp.rs index da3245722..2626eb714 100644 --- a/src/analdrv/response/trait_rresp.rs +++ b/src/analdrv/response/trait_rresp.rs @@ -22,11 +22,14 @@ pub trait RRespAPI: AnalDrvBaseAPI { /// # Parameters /// /// - `rdm` : shape `[nao, nao]`. Reduced one-particle density matrix. + /// - `prec` : precision of the underlying resource: `true` for the high-precision (SCF-grade) + /// one, `false` for the low-precision response resource when attached (falling back to the + /// high-precision one otherwise). /// /// # Returns /// /// - `fock` : shape `[nao, nao]`. Fock matrix. - fn get_fock_rdm(&mut self, rdm: TsrView) -> Tsr; + fn get_fock_rdm(&mut self, rdm: TsrView, prec: bool) -> Tsr; /// Generate Fock matrix from molecular coefficients and occupation numbers. /// @@ -37,6 +40,8 @@ pub trait RRespAPI: AnalDrvBaseAPI { /// /// - `mo_coeff` : shape `[nao, nmo]`. Molecular orbital coefficients. /// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. + /// - `prec` : precision of the underlying resource, forwarded to + /// [`get_fock_rdm`](Self::get_fock_rdm). /// /// Note `nmo` can be set to `nocc` if only occupied orbitals are considered. This can save /// memory and speed up the calculation. @@ -45,9 +50,9 @@ pub trait RRespAPI: AnalDrvBaseAPI { /// /// - `fock` : shape `[nao, nao]`. Fock matrix (the operator in AO basis, not contracted by /// `mo_coeff`). - fn get_fock_coeff(&mut self, mo_coeff: TsrView, mo_occ: TsrView) -> Tsr { + fn get_fock_coeff(&mut self, mo_coeff: TsrView, mo_occ: TsrView, prec: bool) -> Tsr { let rdm = get_dm0_restricted(mo_coeff, mo_occ); - self.get_fock_rdm(rdm.view()) + self.get_fock_rdm(rdm.view(), prec) } /// Prepare the data for response calculation. @@ -65,18 +70,23 @@ pub trait RRespAPI: AnalDrvBaseAPI { /// /// - `mo_coeff` : shape `[nao, nmo]`. Molecular orbital coefficients. /// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. - fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView); + /// - `prec` : precision of the resource the preparation builds on; must match the `prec` of + /// the subsequent response contraction calls (e.g. `false` throughout the CP-SCF solve). + fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView, prec: bool); /// Generate response matrix. /// - /// Refer to [`RRespAPI::get_response_bra`] for better way to leverage the algorithmic advantage by - /// occupation number over molecular orbital number (but also notice the output is different in - /// shape and meaning). + /// Refer to [`RRespAPI::get_response_bra`] for better way to leverage the algorithmic advantage + /// by occupation number over molecular orbital number (but also notice the output is + /// different in shape and meaning). /// Call [`make_response_preparation`] before this function to make sure the data is ready. /// /// # Parameters /// /// - `rdm` : shape `[nao, nao, ...]`. Reduced one-particle density matrix list. + /// - `prec` : precision of the underlying resource: `true` for the high-precision (SCF-grade) + /// one, `false` for the low-precision response resource when attached (falling back to the + /// high-precision one otherwise). /// /// Note on SCF contribution term `mo_coeff` and `mo_occ`: Not useful if energy contribution is /// exactly second-order function of density matrix (such as coulomb or exchange), but is @@ -93,7 +103,7 @@ pub trait RRespAPI: AnalDrvBaseAPI { /// - `resp` : shape `[nao, nao, ...]`. Response matrix. /// /// [`make_response_preparation`]: Self::make_response_preparation - fn get_response_rdm(&mut self, rdm: TsrView) -> Tsr; + fn get_response_rdm(&mut self, rdm: TsrView, prec: bool) -> Tsr; /// Generate response matrix in half-transformed MO basis. /// @@ -107,6 +117,9 @@ pub trait RRespAPI: AnalDrvBaseAPI { /// molecular coefficients (as input) or contracted fock/response matrix that is /// half-transformed by occupied molecular coefficients (as output). This is usually the /// derivative of MO coefficients (like $U_{\mu i}^\mathbb{A}$ given by CP-SCF). + /// - `prec` : precision of the underlying resource: `true` for the high-precision (SCF-grade) + /// one, `false` for the low-precision response resource when attached (falling back to the + /// high-precision one otherwise). The CP-SCF machinery contracts this method with `false`. /// /// # Returns /// @@ -118,5 +131,5 @@ pub trait RRespAPI: AnalDrvBaseAPI { /// We have not prepared to propose a good API for fractional occupation. /// /// [`make_response_preparation`]: Self::make_response_preparation - fn get_response_bra(&mut self, bra: TsrView) -> Tsr; + fn get_response_bra(&mut self, bra: TsrView, prec: bool) -> Tsr; } diff --git a/src/analdrv/response/trait_uresp.rs b/src/analdrv/response/trait_uresp.rs index 93972c48d..eaa83a465 100644 --- a/src/analdrv/response/trait_uresp.rs +++ b/src/analdrv/response/trait_uresp.rs @@ -37,6 +37,9 @@ pub trait URespAPI: AnalDrvBaseAPI { /// # Parameters /// /// - `rdm` : shape `[nao, nao]` per spin. Reduced one-particle (spin) density matrix. + /// - `prec` : precision of the underlying resource: `true` for the high-precision (SCF-grade) + /// one, `false` for the low-precision response resource when attached (falling back to the + /// high-precision one otherwise). /// /// # Returns /// @@ -48,7 +51,7 @@ pub trait URespAPI: AnalDrvBaseAPI { /// No unrestricted driver consumes this method yet; it is part of the trait surface reserved /// for the future unrestricted generalized-Fock machinery, mirroring how the restricted /// [`RRespAPI::get_fock_rdm`] serves [`RGFockAPI`](crate::analdrv::response::trait_rgfock::RGFockAPI). - fn get_fock_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2]; + fn get_fock_rdm(&mut self, rdm: &[TsrView; 2], prec: bool) -> [Tsr; 2]; /// Generate Fock matrix from molecular coefficients and occupation numbers. /// @@ -59,6 +62,8 @@ pub trait URespAPI: AnalDrvBaseAPI { /// /// - `mo_coeff` : shape `[nao, nmo_s]` per spin. Molecular orbital coefficients. /// - `mo_occ` : shape `[nmo_s]` per spin. Molecular orbital occupation numbers. + /// - `prec` : precision of the underlying resource, forwarded to + /// [`get_fock_rdm`](Self::get_fock_rdm). /// /// Note `nmo_s` can be set to `nocc_s` if only occupied orbitals are considered. This can save /// memory and speed up the calculation. @@ -67,9 +72,9 @@ pub trait URespAPI: AnalDrvBaseAPI { /// /// - `fock` : shape `[nao, nao]` per spin. Fock matrix (the operator in AO basis, not /// contracted by `mo_coeff`). - fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) -> [Tsr; 2] { + fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], prec: bool) -> [Tsr; 2] { let rdm = [get_dm0_restricted(mo_coeff[0].view(), mo_occ[0].view()), get_dm0_restricted(mo_coeff[1].view(), mo_occ[1].view())]; - self.get_fock_rdm(&[rdm[0].view(), rdm[1].view()]) + self.get_fock_rdm(&[rdm[0].view(), rdm[1].view()], prec) } /// Prepare the data for response calculation. @@ -87,7 +92,9 @@ pub trait URespAPI: AnalDrvBaseAPI { /// /// - `mo_coeff` : shape `[nao, nmo_s]` per spin. Molecular orbital coefficients. /// - `mo_occ` : shape `[nmo_s]` per spin. Molecular orbital occupation numbers. - fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]); + /// - `prec` : precision of the resource the preparation builds on; must match the `prec` of + /// the subsequent response contraction calls (e.g. `false` throughout the CP-SCF solve). + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], prec: bool); /// Generate response matrix. /// @@ -96,6 +103,9 @@ pub trait URespAPI: AnalDrvBaseAPI { /// # Parameters /// /// - `rdm` : shape `[nao, nao, ...]` per spin. Reduced one-particle (spin) density matrix list. + /// - `prec` : precision of the underlying resource: `true` for the high-precision (SCF-grade) + /// one, `false` for the low-precision response resource when attached (falling back to the + /// high-precision one otherwise). /// /// # Returns /// @@ -108,7 +118,7 @@ pub trait URespAPI: AnalDrvBaseAPI { /// [`RRespAPI::get_response_rdm`] serves the RI-PT2 generalized Fock. /// /// [`make_response_preparation`]: Self::make_response_preparation - fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2]; + fn get_response_rdm(&mut self, rdm: &[TsrView; 2], prec: bool) -> [Tsr; 2]; /// Generate response matrix in half-transformed MO basis. /// @@ -122,6 +132,9 @@ pub trait URespAPI: AnalDrvBaseAPI { /// occupied molecular coefficients (as input) or contracted fock/response matrix that is /// half-transformed by occupied molecular coefficients (as output). This is usually the /// derivative of MO coefficients (like $U_{\mu i}^\mathbb{A}$ given by CP-SCF). + /// - `prec` : precision of the underlying resource: `true` for the high-precision (SCF-grade) + /// one, `false` for the low-precision response resource when attached (falling back to the + /// high-precision one otherwise). The CP-SCF machinery contracts this method with `false`. /// /// # Returns /// @@ -134,5 +147,5 @@ pub trait URespAPI: AnalDrvBaseAPI { /// We have not prepared to propose a good API for fractional occupation. /// /// [`make_response_preparation`]: Self::make_response_preparation - fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2]; + fn get_response_bra(&mut self, bra: &[TsrView; 2], prec: bool) -> [Tsr; 2]; } diff --git a/src/analdrv/response/uresp_interface.rs b/src/analdrv/response/uresp_interface.rs index c19401e46..440016fab 100644 --- a/src/analdrv/response/uresp_interface.rs +++ b/src/analdrv/response/uresp_interface.rs @@ -6,12 +6,15 @@ //! it and feeds the returned objects to `UHessSCF`; a future standalone unrestricted //! response/property driver can call it directly, without any hessian machinery. //! -//! The DFT XC response object evaluates the fock path directly on the SCF grid (in its native -//! round-robin atom-interleaved order — the fock path is a plain quadrature sum, so no -//! atom-regrouping is applied), and the response path on a dedicated (usually coarser) cpscf -//! grid. The hessian-side skeleton grid policy (including the MGGA level bump) lives in the -//! hessian interface; no grid identity or grid data is shared between the hessian and response -//! subsystems. +//! The precision of the response contractions is selected per call through the `prec` flag of +//! [`URespAPI`], mirroring the restricted interface: the high-precision (`prec = true`) resource +//! is the SCF-grade one — the SCF `rimatr` for RI-JK, and the SCF grid for DFT (in its native +//! round-robin atom-interleaved order, since the fock path is a plain quadrature sum, so no +//! atom-regrouping is applied) — while the low-precision (`prec = false`) resource of the CP-SCF +//! machinery is the dedicated (usually coarser) response grid / the `resp_auxbas_path` basis +//! when attached. The hessian-side skeleton grid policy (including the MGGA level bump) lives in +//! the hessian interface; no grid identity or grid data is shared between the hessian and +//! response subsystems. use crate::analdrv::config::AnalDrvRespCfg; use crate::analdrv::prelude::*; @@ -85,23 +88,26 @@ pub struct UCpscfState { impl<'a> AnalDrvBaseAPI for URespSCF<'a> {} impl<'a> URespAPI for URespSCF<'a> { - fn get_fock_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + fn get_fock_rdm(&mut self, rdm: &[TsrView; 2], prec: bool) -> [Tsr; 2] { let mut fock = [rt::zeros_like(&rdm[0]), rt::zeros_like(&rdm[1])]; for resp_obj in self.resp_list.iter_mut() { - let fock_obj = resp_obj.get_fock_rdm(&[rdm[0].view(), rdm[1].view()]); + let fock_obj = resp_obj.get_fock_rdm(&[rdm[0].view(), rdm[1].view()], prec); fock[0] += &fock_obj[0]; fock[1] += &fock_obj[1]; } fock } - fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) -> [Tsr; 2] { + fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], prec: bool) -> [Tsr; 2] { // delegated per element (not through the default rdm route), so that elements overriding // this function for efficiency keep their advantage let mut fock: Option<[Tsr; 2]> = None; for resp_obj in self.resp_list.iter_mut() { - let fock_obj = resp_obj - .get_fock_coeff(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()]); + let fock_obj = resp_obj.get_fock_coeff( + &[mo_coeff[0].view(), mo_coeff[1].view()], + &[mo_occ[0].view(), mo_occ[1].view()], + prec, + ); fock = Some(match fock { Some([f0, f1]) => [f0 + &fock_obj[0], f1 + &fock_obj[1]], None => fock_obj, @@ -110,29 +116,30 @@ impl<'a> URespAPI for URespSCF<'a> { fock.expect("URespSCF must hold at least one response object") } - fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], prec: bool) { for resp_obj in self.resp_list.iter_mut() { - resp_obj.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[ - mo_occ[0].view(), - mo_occ[1].view(), - ]); + resp_obj.make_response_preparation( + &[mo_coeff[0].view(), mo_coeff[1].view()], + &[mo_occ[0].view(), mo_occ[1].view()], + prec, + ); } } - fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + fn get_response_rdm(&mut self, rdm: &[TsrView; 2], prec: bool) -> [Tsr; 2] { let mut resp = [rt::zeros_like(&rdm[0]), rt::zeros_like(&rdm[1])]; for resp_obj in self.resp_list.iter_mut() { - let resp_obj_rdm = resp_obj.get_response_rdm(&[rdm[0].view(), rdm[1].view()]); + let resp_obj_rdm = resp_obj.get_response_rdm(&[rdm[0].view(), rdm[1].view()], prec); resp[0] += &resp_obj_rdm[0]; resp[1] += &resp_obj_rdm[1]; } resp } - fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2] { + fn get_response_bra(&mut self, bra: &[TsrView; 2], prec: bool) -> [Tsr; 2] { let mut resp = [rt::zeros_like(&bra[0]), rt::zeros_like(&bra[1])]; for resp_obj in self.resp_list.iter_mut() { - let resp_obj_bra = resp_obj.get_response_bra(&[bra[0].view(), bra[1].view()]); + let resp_obj_bra = resp_obj.get_response_bra(&[bra[0].view(), bra[1].view()], prec); resp[0] += &resp_obj_bra[0]; resp[1] += &resp_obj_bra[1]; } @@ -155,11 +162,14 @@ impl<'a> URespSCF<'a> { /// - `mo_energy` : shape `[nmo_s]` per spin. Molecular orbital energies. pub fn make_cpscf_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], mo_energy: &[TsrView; 2]) { let [α, β] = [0, 1]; + // the CP-SCF core runs the response machinery in low precision (`prec = false`): the + // dedicated response grid / `resp_auxbas_path` basis when attached for resp_obj in self.resp_list.iter_mut() { - resp_obj.make_response_preparation(&[mo_coeff[α].view(), mo_coeff[β].view()], &[ - mo_occ[α].view(), - mo_occ[β].view(), - ]); + resp_obj.make_response_preparation( + &[mo_coeff[α].view(), mo_coeff[β].view()], + &[mo_occ[α].view(), mo_occ[β].view()], + false, + ); } let occidx = [mo_occ[α].view().greater(0).into_vec(), mo_occ[β].view().greater(0).into_vec()]; @@ -211,10 +221,12 @@ impl<'a> URespSCF<'a> { let mut resp_β = rt::zeros_like(&mo1[β]); // the get_response_bra fan-out is inlined (not the composite trait method), keeping the - // `cpscf_state` borrow disjoint from the `resp_list` iteration + // `cpscf_state` borrow disjoint from the `resp_list` iteration; low precision, as the + // CP-SCF core contracts the response on the dedicated low-precision resource when + // attached (this includes the last-iteration W-operator assembly of the hessian driver) for resp_obj in self.resp_list.iter_mut() { let t1 = std::time::Instant::now(); - let el_resp = resp_obj.get_response_bra(&[ubra_α.view(), ubra_β.view()]); + let el_resp = resp_obj.get_response_bra(&[ubra_α.view(), ubra_β.view()], false); resp_α += mo_coeff[α].t() % &el_resp[α]; resp_β += mo_coeff[β].t() % &el_resp[β]; self.timing.push((format!("in response_mo, {}", resp_obj.get_type_name()), t1.elapsed().as_secs_f64())); @@ -360,19 +372,30 @@ pub fn uscf_resp_interface<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> URe let (factor_j, factor_k, rsh) = scf_jk_factors(scf_data); - // decomposed ERI of the RI-JK response object: its own definition — the freshly built one of - // `resp_auxbas_path` when set (owned), otherwise the SCF rimatr (borrowed, zero copy) + // decomposed ERIs of the RI-JK response object, per precision: the fock path always borrows + // the SCF rimatr (zero copy); the low-precision response path attaches the freshly built one + // of `resp_auxbas_path` when set (owned), and reuses the SCF rimatr otherwise { - let cderi = resp_auxbas::cderi(scf_data, config); - resp_list.push(Box::new(URespRIJK::new_with_cderi(factor_j, factor_k, cderi))); + let (cderi, cderi_resp) = resp_auxbas::cderi_pair(scf_data, config); + let resp_obj = match cderi_resp { + Some(cderi_resp) => URespRIJK::new_with_cderi(factor_j, factor_k, cderi).set_cderi_resp(cderi_resp), + None => URespRIJK::new_with_cderi(factor_j, factor_k, cderi), + }; + resp_list.push(Box::new(resp_obj)); } // The short-range exchange correction (range-separated hybrids) is a separate response // object reusing the full-range implementation; it evaluates on the short-range `rimatr_sr` // ERI, with no Coulomb part (factor_j = 0). if let Some((_omega, factor_k_sr)) = rsh { - let cderi_sr = resp_auxbas::cderi_sr(scf_data, config); - resp_list.push(Box::new(URespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr))); + let (cderi_sr, cderi_resp_sr) = resp_auxbas::cderi_pair_sr(scf_data, config); + let resp_obj = match cderi_resp_sr { + Some(cderi_resp_sr) => { + URespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr).set_cderi_resp(cderi_resp_sr) + }, + None => URespRIJK::new_with_cderi(0.0, factor_k_sr, cderi_sr), + }; + resp_list.push(Box::new(resp_obj)); } // --- DFT --- // diff --git a/src/dft/numint_matmul/resp_rks.rs b/src/dft/numint_matmul/resp_rks.rs index bb73b4d9d..9dff15f8d 100644 --- a/src/dft/numint_matmul/resp_rks.rs +++ b/src/dft/numint_matmul/resp_rks.rs @@ -51,8 +51,8 @@ pub fn eval_vxc_fxc_from_rho(xc_func_list: &[(f64, LibXCFunctional)], rho: TsrVi (vxc, fxc) } -/// Lean evaluation of only `vxc` and `fxc` on a given grid, for use as the CP-KS `cpks_vxc` / -/// `cpks_fxc` when a dedicated (coarser) CP-KS grid is attached. +/// Lean evaluation of only `vxc` and `fxc` on a given grid, for use as the response object's +/// cached `vxc_*` / `fxc_*` when a dedicated (coarser) low-precision grid is attached. /// /// Compared to [`make_hessian_setup_becke`](super::hess_rks::make_hessian_setup_becke), this: /// - skips all skeleton-Hessian intermediates (`de_fxc`, `de_vxc_diag`, `de_vxc_off`, `vmat_ip`, @@ -178,16 +178,17 @@ pub fn get_rks_response_bra_batched( pub struct RRespKSNIMatmul<'a> { /// List of `(scale, functional)` pairs of the XC functional. pub xc_func_list: Vec<(f64, LibXCFunctional)>, - /// Common (large) numerical-integration grid, used by the fock path. + /// High-precision (large, SCF/common) numerical-integration grid, selected by `prec = true`. pub ni: NIMatmul<'a>, - /// Optional separate small grid for the response path; `None` reuses `ni`. + /// Optional separate small grid for the low-precision (`prec = false`) response path; `None` + /// reuses `ni`. pub ni_resp: Option>, /// Print progress of the response evaluation. pub verbose: bool, /// Response intermediates: `mo_coeff [nao, nmo]` and `mo_occ [nmo]` from - /// [`RRespAPI::make_response_preparation`], `cpks_vxc [ngrids, nvar]` and - /// `cpks_fxc [ngrids, nvar, nvar]` on the selected response grid, and `fxc_common_grid - /// [ngrids, nvar, nvar]` on the common grid. + /// [`RRespAPI::make_response_preparation`], and `vxc_{high,low} [ngrids, nvar]` / + /// `fxc_{high,low} [ngrids, nvar, nvar]` cached per precision on the selected grid (the fxc + /// entry is shared between the preparation and the rdm-form contraction). pub intmd: HashMap, } @@ -197,16 +198,18 @@ impl<'a> RRespKSNIMatmul<'a> { /// # Parameters /// /// - `xc_func_list` : list of `(scale, functional)` pairs. - /// - `ni` : numerical-integration driver over the common (large) grid. + /// - `ni` : numerical-integration driver over the high-precision (large) grid. /// - `verbose` : print progress. pub fn new(xc_func_list: Vec<(f64, LibXCFunctional)>, ni: NIMatmul<'a>, verbose: bool) -> Self { Self { xc_func_list, ni, ni_resp: None, verbose, intmd: HashMap::new() } } - /// Attach a dedicated small numerical-integration grid (`ni_resp`) for the response path. + /// Attach a dedicated small numerical-integration grid (`ni_resp`) for the low-precision + /// (`prec = false`) response path. /// - /// When set, the response (`get_response_bra`) is evaluated on this grid instead of the common - /// (large) grid, and `cpks_vxc` / `cpks_fxc` are computed on it during + /// When set, the contractions called with `prec = false` (the CP-SCF machinery) are evaluated + /// on this grid instead of the high-precision (large) grid, and `vxc_low` / `fxc_low` are + /// computed on it during /// [`make_response_preparation`](RRespAPI::make_response_preparation). pub fn set_ni_resp(mut self, ni_resp: NIMatmul<'a>) -> Self { self.ni_resp = Some(ni_resp); @@ -217,55 +220,60 @@ impl<'a> RRespKSNIMatmul<'a> { impl<'a> AnalDrvBaseAPI for RRespKSNIMatmul<'a> {} impl<'a> RRespAPI for RRespKSNIMatmul<'a> { - /// Fock (the XC numint contribution only) from density matrix, on the common grid `ni`. - fn get_fock_rdm(&mut self, rdm: TsrView) -> Tsr { + /// Fock (the XC numint contribution only) from density matrix, on the grid selected by `prec` + /// (the high-precision `ni` for `prec = true`). + fn get_fock_rdm(&mut self, rdm: TsrView, prec: bool) -> Tsr { assert_eq!(rdm.ndim(), 2, "rdm must have 2 dimensions"); let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); - let rho = self.ni.make_rho_from_dm(&[rdm], den_type); + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; + let rho = ni.make_rho_from_dm(&[rdm], den_type); let (vxc, _fxc) = eval_vxc_fxc_from_rho(&self.xc_func_list, rho.i((.., .., 0))); - self.ni.make_vxc_pot_with_eff(vxc.view(), den_type, XCSpin::Unpolarized) + ni.make_vxc_pot_with_eff(vxc.view(), den_type, XCSpin::Unpolarized) } // note: `get_fock_coeff` keeps the trait default (dm route via `get_dm0_restricted`). /// Response matrix from (symmetrizable) full density matrices: `4 * fxc-response` on the - /// common grid. The `rdm` carries a leading `[nao, nao]` square followed by an arbitrary - /// (possibly empty) set of trailing dimensions; the output has the same shape. + /// grid selected by `prec`. The `rdm` carries a leading `[nao, nao]` square followed by an + /// arbitrary (possibly empty) set of trailing dimensions; the output has the same shape. /// /// This is the rdm-form entry of the restricted response kernel, consistent with the /// `4 J - 2 K` convention of [`RRespRIJK`](crate::ri_jk::resp_r::RRespRIJK), and matches /// pyscf's orbital-hessian `vind` (ground-state `_gen_rhf_response`) contracted with `4`. /// It is the form required by the generalized-Fock Lagrangian term `A_{ai, pq} D_{pq}` of - /// post-SCF methods; accordingly, the fxc kernel is built from the ground-state density on - /// the common (large) grid `ni` — not on the dedicated response grid — mirroring the - /// Lagrangian evaluation on the SCF grids in pyscf-forge. - fn get_response_rdm(&mut self, rdm: TsrView) -> Tsr { + /// post-SCF methods; accordingly, the production callers contract it with `prec = true`, so + /// the fxc kernel is built from the ground-state density on the high-precision (large) grid + /// `ni` — mirroring the Lagrangian evaluation on the SCF grids in pyscf-forge. + fn get_response_rdm(&mut self, rdm: TsrView, prec: bool) -> Tsr { assert!(rdm.ndim() >= 2, "rdm must have at least 2 dimensions"); let rdm_shape = rdm.shape().to_vec(); let nao = rdm_shape[0]; assert_eq!(nao, rdm_shape[1], "the first two dimensions of rdm must be equal"); let nset: usize = rdm_shape[2..].iter().product(); let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); + let key_fxc = if prec { "fxc_high" } else { "fxc_low" }; - if !self.intmd.contains_key("fxc_common_grid") { + if !self.intmd.contains_key(key_fxc) { // panic with "no entry" if preparation is skipped let mo_coeff = self.intmd["mo_coeff"].view(); let mo_occ = self.intmd["mo_occ"].view(); - let (_, fxc) = make_cpks_vxc_fxc(&self.xc_func_list, &mut self.ni, mo_coeff, mo_occ); - self.intmd.insert("fxc_common_grid".to_string(), fxc); + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; + let (_, fxc) = make_cpks_vxc_fxc(&self.xc_func_list, ni, mo_coeff, mo_occ); + self.intmd.insert(key_fxc.to_string(), fxc); } - // NOTE: unlike `cpks_fxc`, the cached `fxc_common_grid` is NOT re-validated against the - // orbitals on later calls. This is safe under the analdrv driver contract (the orbitals - // are fixed for the lifetime of a driver object), but any future caller that - // re-prepares with different orbitals must invalidate this entry as well. - let fxc = self.intmd["fxc_common_grid"].view(); + // NOTE: the cached `fxc_*` is written here only when missing; + // [`make_response_preparation`](RRespAPI::make_response_preparation) re-validates it + // against the orbitals and refreshes it on orbital change, so the entry cannot go stale + // under the analdrv driver contract. + let fxc = self.intmd[key_fxc].view(); // flatten the trailing dimensions into one set dimension; assume each set of the rdm can // be symmetrized — the fxc kernel only sees the symmetric part anyway let rdm_sym = ((&rdm + &rdm.swapaxes(0, 1)) * 0.5).into_shape((nao, nao, nset)); let dm_views: Vec = (0..nset).map(|s| rdm_sym.i((.., .., s))).collect(); - let rho1 = self.ni.make_rho_from_dm(&dm_views, den_type); - let resp = self.ni.make_fxc_pot_with_eff(fxc, rho1.view(), den_type, XCSpin::Unpolarized); + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; + let rho1 = ni.make_rho_from_dm(&dm_views, den_type); + let resp = ni.make_fxc_pot_with_eff(fxc, rho1.view(), den_type, XCSpin::Unpolarized); // 4.0 times is a trick of closed-shell coefficient; restore the input's trailing shape (resp * 4.0_f64).into_shape(rdm_shape) } @@ -273,11 +281,12 @@ impl<'a> RRespAPI for RRespKSNIMatmul<'a> { /// Cached on first call for fixed inputs: the CP-KS `vxc`/`fxc` evaluation is skipped when /// this object was already prepared with the same orbitals (e.g. repeated calls from /// multi-order property evaluations directly reuse the stored results). - fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView) { + fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView, prec: bool) { + let (key_vxc, key_fxc) = if prec { ("vxc_high", "fxc_high") } else { ("vxc_low", "fxc_low") }; // cache check (before the orbital refresh): whether the expensive evaluation below was - // already performed with the same orbitals (`mo_coeff`/`mo_occ` are always inserted - // together with `cpks_vxc`/`cpks_fxc`, so the keys coexist) - let already_prepared = self.intmd.contains_key("cpks_fxc") + // already performed with the same orbitals (the `fxc_*` entry is also written by the + // rdm-form contraction on the same grid and orbitals, so it is a valid proxy) + let already_prepared = self.intmd.contains_key(key_fxc) && is_same_tensor(self.intmd["mo_coeff"].view(), mo_coeff.view()) && is_same_tensor(self.intmd["mo_occ"].view(), mo_occ.view()); self.intmd.insert("mo_coeff".to_string(), mo_coeff.into_contig(ColMajor)); @@ -286,28 +295,29 @@ impl<'a> RRespAPI for RRespKSNIMatmul<'a> { return; } - // Compute `cpks_vxc` / `cpks_fxc` on the selected response grid (the small `ni_resp` when - // attached, else the common grid) from the ground-state density, using the lean - // [`make_cpks_vxc_fxc`] (no skeleton intermediates, minimal AO derivative order, density - // formed from occupied MOs via a bra-ket contraction rather than a full dm0). - let ni_resp = self.ni_resp.as_mut().unwrap_or(&mut self.ni); + // Compute the cached `vxc_*` / `fxc_*` on the selected grid (the small `ni_resp` of the + // low-precision path when attached, else the high-precision grid) from the ground-state + // density, using the lean [`make_cpks_vxc_fxc`] (no skeleton intermediates, minimal AO + // derivative order, density formed from occupied MOs via a bra-ket contraction rather + // than a full dm0). + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; let mo_coeff = self.intmd["mo_coeff"].view(); let mo_occ = self.intmd["mo_occ"].view(); - let (vxc, fxc) = make_cpks_vxc_fxc(&self.xc_func_list, ni_resp, mo_coeff, mo_occ); - self.intmd.insert("cpks_vxc".to_string(), vxc); - self.intmd.insert("cpks_fxc".to_string(), fxc); + let (vxc, fxc) = make_cpks_vxc_fxc(&self.xc_func_list, ni, mo_coeff, mo_occ); + self.intmd.insert(key_vxc.to_string(), vxc); + self.intmd.insert(key_fxc.to_string(), fxc); } - fn get_response_bra(&mut self, bra: TsrView) -> Tsr { - let ni_resp = self.ni_resp.as_mut().unwrap_or(&mut self.ni); + fn get_response_bra(&mut self, bra: TsrView, prec: bool) -> Tsr { + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; let mo_coeff = self.intmd.get("mo_coeff").unwrap(); let mo_occ = self.intmd.get("mo_occ").unwrap(); - let fxc_eff = self.intmd.get("cpks_fxc").unwrap(); + let fxc_eff = self.intmd.get(if prec { "fxc_high" } else { "fxc_low" }).unwrap(); let occidx = mo_occ.view().greater(0).into_vec(); let mocc = mo_coeff.bool_select(-1, &occidx); let (resp, _timing) = get_rks_response_bra_batched( - ni_resp, + ni, determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()), fxc_eff.view(), bra, diff --git a/src/dft/numint_matmul/resp_uks.rs b/src/dft/numint_matmul/resp_uks.rs index 4fe04a33b..a485faf54 100644 --- a/src/dft/numint_matmul/resp_uks.rs +++ b/src/dft/numint_matmul/resp_uks.rs @@ -51,9 +51,9 @@ pub fn eval_vxc_fxc_uks_from_rho(xc_func_list: &[(f64, LibXCFunctional)], rho: T (vxc, fxc) } -/// Lean evaluation of only `vxc` and `fxc` on a given grid, for use as the CP-KS `cpks_vxc` / -/// `cpks_fxc` when a dedicated (coarser) CP-KS grid is attached (UKS counterpart of -/// [`make_cpks_vxc_fxc`](super::resp_rks::make_cpks_vxc_fxc)). +/// Lean evaluation of only `vxc` and `fxc` on a given grid, for use as the response object's +/// cached `vxc_*` / `fxc_*` when a dedicated (coarser) low-precision grid is attached (UKS +/// counterpart of [`make_cpks_vxc_fxc`](super::resp_rks::make_cpks_vxc_fxc)). /// /// Skips all skeleton-Hessian intermediates, evaluates AO at the minimum derivative order needed /// to form the spin densities (`xc_type.num_ao_deriv()`), and forms each spin density `rhoσ` @@ -200,15 +200,17 @@ pub fn get_uks_response_bra_batched( pub struct URespKSNIMatmul<'a> { /// List of `(scale, functional)` pairs of the XC functional (spin-polarized). pub xc_func_list: Vec<(f64, LibXCFunctional)>, - /// Common (large) numerical-integration grid, used by the fock path. + /// High-precision (large, SCF/common) numerical-integration grid, selected by `prec = true`. pub ni: NIMatmul<'a>, - /// Optional separate small grid for the response path; `None` reuses `ni`. + /// Optional separate small grid for the low-precision (`prec = false`) response path; `None` + /// reuses `ni`. pub ni_resp: Option>, /// Print progress of the response evaluation. pub verbose: bool, /// Response intermediates: `mo_coeff_0/1 [nao, nmo_s]` and `mo_occ_0/1 [nmo_s]` from - /// [`URespAPI::make_response_preparation`], `cpks_vxc [ngrids, nvar, 2]` and `cpks_fxc - /// [ngrids, nvar, 2, nvar, 2]` on the selected response grid. + /// [`URespAPI::make_response_preparation`], and `vxc_{high,low} [ngrids, nvar, 2]` / + /// `fxc_{high,low} [ngrids, nvar, 2, nvar, 2]` cached per precision on the selected grid + /// (the fxc entry is shared between the preparation and the rdm-form contraction). pub intmd: HashMap, } @@ -218,16 +220,18 @@ impl<'a> URespKSNIMatmul<'a> { /// # Parameters /// /// - `xc_func_list` : list of `(scale, functional)` pairs (spin-polarized functionals). - /// - `ni` : numerical-integration driver over the common (large) grid. + /// - `ni` : numerical-integration driver over the high-precision (large) grid. /// - `verbose` : print progress. pub fn new(xc_func_list: Vec<(f64, LibXCFunctional)>, ni: NIMatmul<'a>, verbose: bool) -> Self { Self { xc_func_list, ni, ni_resp: None, verbose, intmd: HashMap::new() } } - /// Attach a dedicated small numerical-integration grid (`ni_resp`) for the response path. + /// Attach a dedicated small numerical-integration grid (`ni_resp`) for the low-precision + /// (`prec = false`) response path. /// - /// When set, the response (`get_response_bra`) is evaluated on this grid instead of the common - /// (large) grid, and `cpks_vxc` / `cpks_fxc` are computed on it during + /// When set, the contractions called with `prec = false` (the CP-SCF machinery) are evaluated + /// on this grid instead of the high-precision (large) grid, and `vxc_low` / `fxc_low` are + /// computed on it during /// [`make_response_preparation`](URespAPI::make_response_preparation). pub fn set_ni_resp(mut self, ni_resp: NIMatmul<'a>) -> Self { self.ni_resp = Some(ni_resp); @@ -238,24 +242,25 @@ impl<'a> URespKSNIMatmul<'a> { impl<'a> AnalDrvBaseAPI for URespKSNIMatmul<'a> {} impl<'a> URespAPI for URespKSNIMatmul<'a> { - /// Fock (the XC numint contribution only) per spin from spin density matrices, on the common - /// grid `ni`. - fn get_fock_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + /// Fock (the XC numint contribution only) per spin from spin density matrices, on the grid + /// selected by `prec` (the high-precision `ni` for `prec = true`). + fn get_fock_rdm(&mut self, rdm: &[TsrView; 2], prec: bool) -> [Tsr; 2] { for rdm_s in rdm { assert_eq!(rdm_s.ndim(), 2, "rdm must have 2 dimensions"); } let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; // rho : [ngrids, nvar, 2] (the spin is the set axis) - let rho = self.ni.make_rho_from_dm(&[rdm[α].view(), rdm[β].view()], den_type); + let rho = ni.make_rho_from_dm(&[rdm[α].view(), rdm[β].view()], den_type); let (vxc, _fxc) = eval_vxc_fxc_uks_from_rho(&self.xc_func_list, rho.view()); - let pot = self.ni.make_vxc_pot_with_eff(vxc.view(), den_type, XCSpin::Polarized); + let pot = ni.make_vxc_pot_with_eff(vxc.view(), den_type, XCSpin::Polarized); [pot.i((.., .., α)).to_owned(), pot.i((.., .., β)).to_owned()] } // note: `get_fock_coeff` keeps the trait default (dm route via `get_dm0_restricted` per spin). /// Response matrix per spin from (symmetrizable) full spin density matrices: `2 * fxc` - /// response on the common (large) grid `ni` — the UHF factor `2.0` (hermitian symmetry + /// response on the grid selected by `prec` — the UHF factor `2.0` (hermitian symmetry /// only, no spin degeneracy) against the restricted `4.0` of /// [`RRespKSNIMatmul`](super::resp_rks::RRespKSNIMatmul). The per-spin `rdm` carries a /// leading `[nao, nao]` square followed by an arbitrary (possibly empty) set of trailing @@ -263,12 +268,12 @@ impl<'a> URespAPI for URespKSNIMatmul<'a> { /// /// This is the rdm-form entry of the unrestricted response kernel, the form required by the /// future unrestricted generalized-Fock Lagrangian term `A_{ai, pq} D_{pq}`; accordingly, - /// the fxc kernel is built from the ground-state spin densities on the common (large) grid - /// `ni` — not on the dedicated response grid — mirroring + /// the production callers contract it with `prec = true`, so the fxc kernel is built from + /// the ground-state spin densities on the high-precision (large) grid `ni`, mirroring /// [`RRespKSNIMatmul::get_response_rdm`](super::resp_rks::RRespKSNIMatmul::get_response_rdm). /// Call [`make_response_preparation`](URespAPI::make_response_preparation) before this /// function to make sure the data is ready. - fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + fn get_response_rdm(&mut self, rdm: &[TsrView; 2], prec: bool) -> [Tsr; 2] { for rdm_s in rdm { assert!(rdm_s.ndim() >= 2, "rdm must have at least 2 dimensions"); } @@ -280,19 +285,21 @@ impl<'a> URespAPI for URespKSNIMatmul<'a> { assert_eq!(rdm_shape[α], rdm_shape[β], "the rdm shape must agree across spins"); let nset: usize = rdm_shape[α][2..].iter().product(); let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); + let key_fxc = if prec { "fxc_high" } else { "fxc_low" }; - if !self.intmd.contains_key("fxc_common_grid") { + if !self.intmd.contains_key(key_fxc) { // panic with "no entry" if preparation is skipped let mo_coeff = [self.intmd["mo_coeff_0"].view(), self.intmd["mo_coeff_1"].view()]; let mo_occ = [self.intmd["mo_occ_0"].view(), self.intmd["mo_occ_1"].view()]; - let (_, fxc) = make_cpks_vxc_fxc_uks(&self.xc_func_list, &mut self.ni, &mo_coeff, &mo_occ); - self.intmd.insert("fxc_common_grid".to_string(), fxc); + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; + let (_, fxc) = make_cpks_vxc_fxc_uks(&self.xc_func_list, ni, &mo_coeff, &mo_occ); + self.intmd.insert(key_fxc.to_string(), fxc); } - // NOTE: unlike `cpks_fxc`, the cached `fxc_common_grid` is NOT re-validated against the - // orbitals on later calls. This is safe under the analdrv driver contract (the orbitals - // are fixed for the lifetime of a driver object), but any future caller that - // re-prepares with different orbitals must invalidate this entry as well. - let fxc = self.intmd["fxc_common_grid"].view(); + // NOTE: the cached `fxc_*` is written here only when missing; + // [`make_response_preparation`](URespAPI::make_response_preparation) re-validates it + // against the orbitals and refreshes it on orbital change, so the entry cannot go stale + // under the analdrv driver contract. + let fxc = self.intmd[key_fxc].view(); // flatten the trailing dimensions into one set dimension; symmetrize each spin's // leading [nao, nao] block — the fxc kernel only sees the symmetric part anyway @@ -307,11 +314,12 @@ impl<'a> URespAPI for URespKSNIMatmul<'a> { let dm_list: Vec = (0..nset) .flat_map(|k| [rdm_sym[α].i((.., .., k)), rdm_sym[β].i((.., .., k))]) .collect(); - let ngrids = self.ni.weights.len(); + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; + let ngrids = ni.weights.len(); let nvar = den_type.num_nvar(); - let rho1 = self.ni.make_rho_from_dm(&dm_list, den_type).into_shape((ngrids, nvar, 2, nset)); + let rho1 = ni.make_rho_from_dm(&dm_list, den_type).into_shape((ngrids, nvar, 2, nset)); - let pot = self.ni.make_fxc_pot_with_eff(fxc, rho1.view(), den_type, XCSpin::Polarized); + let pot = ni.make_fxc_pot_with_eff(fxc, rho1.view(), den_type, XCSpin::Polarized); // restore the input's per-spin trailing shape [ (pot.i((.., .., α, ..)) * 2.0_f64).into_shape(rdm_shape[α].to_vec()), @@ -322,11 +330,12 @@ impl<'a> URespAPI for URespKSNIMatmul<'a> { /// Cached on first call for fixed inputs: the CP-KS `vxc`/`fxc` evaluation is skipped when /// this object was already prepared with the same orbitals (e.g. repeated calls from /// multi-order property evaluations directly reuse the stored results). - fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], prec: bool) { + let (key_vxc, key_fxc) = if prec { ("vxc_high", "fxc_high") } else { ("vxc_low", "fxc_low") }; // cache check (before the orbital refresh): whether the expensive evaluation below was - // already performed with the same orbitals (`mo_coeff`/`mo_occ` are always inserted - // together with `cpks_vxc`/`cpks_fxc`, so the keys coexist) - let already_prepared = self.intmd.contains_key("cpks_fxc") + // already performed with the same orbitals (the `fxc_*` entry is also written by the + // rdm-form contraction on the same grid and orbitals, so it is a valid proxy) + let already_prepared = self.intmd.contains_key(key_fxc) && is_same_tensor(self.intmd["mo_coeff_0"].view(), mo_coeff[α].view()) && is_same_tensor(self.intmd["mo_coeff_1"].view(), mo_coeff[β].view()) && is_same_tensor(self.intmd["mo_occ_0"].view(), mo_occ[α].view()) @@ -339,29 +348,29 @@ impl<'a> URespAPI for URespKSNIMatmul<'a> { return; } - // Compute `cpks_vxc` / `cpks_fxc` on the selected response grid (the small `ni_resp` when - // attached, else the common grid) from the ground-state spin densities, using the lean - // [`make_cpks_vxc_fxc_uks`] (no skeleton intermediates, minimal AO derivative order, spin - // densities formed from occupied spin-orbitals via a bra-ket contraction rather than full - // spin dm0 matrices). - let ni_resp = self.ni_resp.as_mut().unwrap_or(&mut self.ni); + // Compute the cached `vxc_*` / `fxc_*` on the selected grid (the small `ni_resp` of the + // low-precision path when attached, else the high-precision grid) from the ground-state + // spin densities, using the lean [`make_cpks_vxc_fxc_uks`] (no skeleton intermediates, + // minimal AO derivative order, spin densities formed from occupied spin-orbitals via a + // bra-ket contraction rather than full spin dm0 matrices). + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; let mo_coeff_α = self.intmd["mo_coeff_0"].view(); let mo_coeff_β = self.intmd["mo_coeff_1"].view(); let mo_occ_α = self.intmd["mo_occ_0"].view(); let mo_occ_β = self.intmd["mo_occ_1"].view(); let (vxc, fxc) = - make_cpks_vxc_fxc_uks(&self.xc_func_list, ni_resp, &[mo_coeff_α, mo_coeff_β], &[mo_occ_α, mo_occ_β]); - self.intmd.insert("cpks_vxc".to_string(), vxc); - self.intmd.insert("cpks_fxc".to_string(), fxc); + make_cpks_vxc_fxc_uks(&self.xc_func_list, ni, &[mo_coeff_α, mo_coeff_β], &[mo_occ_α, mo_occ_β]); + self.intmd.insert(key_vxc.to_string(), vxc); + self.intmd.insert(key_fxc.to_string(), fxc); } - fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2] { - let ni_resp = self.ni_resp.as_mut().unwrap_or(&mut self.ni); + fn get_response_bra(&mut self, bra: &[TsrView; 2], prec: bool) -> [Tsr; 2] { + let ni = if prec { &mut self.ni } else { self.ni_resp.as_mut().unwrap_or(&mut self.ni) }; let mo_coeff_α = self.intmd["mo_coeff_0"].view(); let mo_coeff_β = self.intmd["mo_coeff_1"].view(); let mo_occ_α = self.intmd["mo_occ_0"].view(); let mo_occ_β = self.intmd["mo_occ_1"].view(); - let fxc_eff = self.intmd["cpks_fxc"].view(); + let fxc_eff = self.intmd[if prec { "fxc_high" } else { "fxc_low" }].view(); let occidx_α = mo_occ_α.view().greater(0).into_vec(); let occidx_β = mo_occ_β.view().greater(0).into_vec(); @@ -371,7 +380,7 @@ impl<'a> URespAPI for URespKSNIMatmul<'a> { let den_type = determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()); let ([resp_α, resp_β], _timing) = get_uks_response_bra_batched( - ni_resp, + ni, den_type, fxc_eff.view(), bra, diff --git a/src/ri_jk/resp_auxbas.rs b/src/ri_jk/resp_auxbas.rs index d27b44664..c193451a4 100644 --- a/src/ri_jk/resp_auxbas.rs +++ b/src/ri_jk/resp_auxbas.rs @@ -1,14 +1,16 @@ //! Response-specific auxiliary basis of the `resp_auxbas_path` keyword. //! -//! When set, the RI-JK response objects evaluate on decomposed ERIs freshly built on that -//! auxiliary basis instead of borrowing the SCF `rimatr`/`rimatr_sr`. This is a response-side -//! (A-tensor) substitution only: the energy-derivative (B-side) objects of the hessian/multipole -//! tasks keep the SCF auxiliary basis, so the CP-SCF solve becomes a mixed-representation -//! approximation — parallel to the coarser response grid of `grid_level_cpscf`. +//! When set, the RI-JK response objects evaluate their low-precision (`prec = false`) +//! contractions — the CP-SCF response machinery — on decomposed ERIs freshly built on that +//! auxiliary basis, while the fock forms (`prec = true`) keep the SCF `rimatr`/`rimatr_sr`. This +//! is a response-side (A-tensor) substitution only: the energy-derivative (B-side) objects of +//! the hessian/multipole tasks keep the SCF auxiliary basis, so the CP-SCF solve becomes a +//! mixed-representation approximation — parallel to the coarser response grid of +//! `grid_level_cpscf`. //! //! Each energy contribution component that requires a decomposed ERI has its own definition -//! here: [`cderi`] for the full-range RI-JK object, [`cderi_sr`] for the RSH short-range -//! object. With the keyword set, the auxiliary-side `CInt` is regenerated on the +//! here: [`cderi_pair`] for the full-range RI-JK object, [`cderi_pair_sr`] for the RSH +//! short-range object. With the keyword set, the auxiliary-side `CInt` is regenerated on the //! response basis (the SCF molecule is only read, never modified) and substituted into the rimatr //! builder, yielding an owned tensor; with the keyword unset, the SCF tensors are borrowed as //! before (zero copy). @@ -26,7 +28,7 @@ use log::debug; /// The even-tempered machinery would silently override the response basis, so the keyword /// combination is rejected. Checked once up front by the task-level /// [`validate_resp_auxbas`], and again at every response-object build -/// ([`cderi`]/[`cderi_sr`]), so paths that build the response objects directly (e.g. the +/// ([`cderi_pair`]/[`cderi_pair_sr`]), so paths that build the response objects directly (e.g. the /// geomopt analytic hessian) are covered too. pub fn assert_resp_auxbas_supported(scf_data: &SCF, config: &AnalDrvConfig) { if config.resp.auxbas_path.is_none() { @@ -55,18 +57,21 @@ pub fn validate_resp_auxbas(scf_data: &SCF, config: &AnalDrvConfig, need_resp: b } } -/// Decomposed ERI of the full-range RI-JK response object, as a [`TsrCow`]. +/// Decomposed ERIs of the full-range RI-JK response object, per precision, as a pair +/// `(cderi, cderi_resp)` of [`TsrCow`]s. /// -/// With `resp_auxbas_path` set, the rimatr is freshly built on that basis: the auxiliary-side -/// `CInt` is regenerated from a local copy of the keywords -/// ([`get_cint_aux_on_path`](crate::ri_jk::util::get_cint_aux_on_path)) and, together with the -/// molecule-side one, handed to the CInt-level rimatr builder +/// The first entry (the high-precision, `prec = true` resource of the fock forms) is always the +/// SCF `rimatr`, borrowed zero-copy. The second entry (the low-precision, `prec = false` +/// resource of the CP-SCF response contractions) is a freshly built tensor only when +/// `resp_auxbas_path` is set: the auxiliary-side `CInt` is regenerated from a local copy of the +/// keywords ([`get_cint_aux_on_path`](crate::ri_jk::util::get_cint_aux_on_path)) and, together +/// with the molecule-side one, handed to the CInt-level rimatr builder /// ([`generate_rimatr_bare_on_cint`](crate::ri_jk::generate_rimatr_bare_on_cint)), so the SCF /// molecule stays untouched and the tensor is owned. A `resp_auxbas_path` equal to the SCF /// `auxbas_path` (exact match) is treated as unset when the SCF `rimatr` is available: the /// borrowed tensor is used instead of rebuilding an identical one. With the keyword unset, the -/// SCF `rimatr` is borrowed (zero copy). -pub fn cderi<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { +/// second entry is `None` (the response path reuses the SCF tensor). +pub fn cderi_pair<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> (TsrCow<'a>, Option>) { let device = DeviceBLAS::default(); // a response basis equal to the SCF one, with the SCF tensor available, is treated as unset let resp_auxbas_path = match config.resp.auxbas_path.as_deref() { @@ -76,7 +81,7 @@ pub fn cderi<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { }, path => path, }; - match resp_auxbas_path { + let cderi_resp = match resp_auxbas_path { Some(resp_auxbas_path) => { assert_resp_auxbas_supported(scf_data, config); debug!("Building the response rimatr on `resp_auxbas_path` = \"{}\".", resp_auxbas_path); @@ -85,29 +90,30 @@ pub fn cderi<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { let cderi = generate_rimatr_bare_on_cint(&mol, &aux, scf_data.mol.ctrl.j2c_decomp, None).into_rstsr(&device); debug!( - "RI-JK response object on the `resp_auxbas_path` decomposed ERI (owned, {} auxiliary functions).", + "RI-JK low-precision response path on the `resp_auxbas_path` decomposed ERI (owned, {} auxiliary functions).", cderi.shape()[1] ); - cderi.into_cow() + Some(cderi.into_cow()) }, - None => { - debug!("RI-JK response object on the SCF rimatr (borrowed, zero copy)."); - let (rimatr, _, _) = scf_data.rimatr.as_ref().expect( - "This implementation requires cholesky decomposed ERI (or rimatr) to be available and stored in memory.", - ); - rimatr.to_rstsr_view(&device).into_cow() - }, - } + None => None, + }; + debug!("RI-JK fock forms of the response object on the SCF rimatr (borrowed, zero copy)."); + let (rimatr, _, _) = scf_data.rimatr.as_ref().expect( + "This implementation requires cholesky decomposed ERI (or rimatr) to be available and stored in memory.", + ); + let cderi = rimatr.to_rstsr_view(&device).into_cow(); + (cderi, cderi_resp) } -/// Decomposed ERI of the RSH short-range RI-JK response object, as a [`TsrCow`]. +/// Decomposed ERIs of the RSH short-range RI-JK response object, per precision, as a pair +/// `(cderi_sr, cderi_resp_sr)` of [`TsrCow`]s. /// -/// Same definition as [`cderi`] on the short-range tensor: with `resp_auxbas_path` set, the -/// SR rimatr is freshly built (with `-omega` in libcint's convention, cf. the SCF integral -/// preparation), and a path equal to the SCF `auxbas_path` is treated as unset when the SCF -/// `rimatr_sr` is available; with the keyword unset, the SCF `rimatr_sr` is borrowed (zero -/// copy). -pub fn cderi_sr<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { +/// Same definition as [`cderi_pair`] on the short-range tensors: with `resp_auxbas_path` set, +/// the SR rimatr is freshly built (with `-omega` in libcint's convention, cf. the SCF integral +/// preparation) for the low-precision response path, and a path equal to the SCF `auxbas_path` +/// is treated as unset when the SCF `rimatr_sr` is available; the fock forms always borrow the +/// SCF `rimatr_sr`, and with the keyword unset the response path reuses it too (zero copy). +pub fn cderi_pair_sr<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> (TsrCow<'a>, Option>) { let device = DeviceBLAS::default(); // a response basis equal to the SCF one, with the SCF tensor available, is treated as unset let resp_auxbas_path = match config.resp.auxbas_path.as_deref() { @@ -117,7 +123,7 @@ pub fn cderi_sr<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { }, path => path, }; - match resp_auxbas_path { + let cderi_resp_sr = match resp_auxbas_path { Some(resp_auxbas_path) => { assert_resp_auxbas_supported(scf_data, config); // SR RI omega is negative in libcint's convention (cf. the SCF integral preparation) @@ -130,15 +136,15 @@ pub fn cderi_sr<'a>(scf_data: &'a SCF, config: &AnalDrvConfig) -> TsrCow<'a> { let aux = get_cint_aux_on_path(&scf_data.mol, resp_auxbas_path); let cderi_sr = generate_rimatr_bare_on_cint(&mol, &aux, scf_data.mol.ctrl.j2c_decomp, Some(-omega)).into_rstsr(&device); - debug!("RSH short-range response object on the `resp_auxbas_path` decomposed ERI (owned)."); - cderi_sr.into_cow() + debug!("RSH short-range low-precision response path on the `resp_auxbas_path` decomposed ERI (owned)."); + Some(cderi_sr.into_cow()) }, - None => { - debug!("RSH short-range response object on the SCF rimatr_sr (borrowed, zero copy)."); - let (rimatr_sr, _, _) = scf_data.rimatr_sr.as_ref().expect( - "The range-separated response requires the short-range ERI (rimatr_sr) to be built and stored in memory.", - ); - rimatr_sr.to_rstsr_view(&device).into_cow() - }, - } + None => None, + }; + debug!("RSH short-range fock forms of the response object on the SCF rimatr_sr (borrowed, zero copy)."); + let (rimatr_sr, _, _) = scf_data.rimatr_sr.as_ref().expect( + "The range-separated response requires the short-range ERI (rimatr_sr) to be built and stored in memory.", + ); + let cderi_sr = rimatr_sr.to_rstsr_view(&device).into_cow(); + (cderi_sr, cderi_resp_sr) } diff --git a/src/ri_jk/resp_r.rs b/src/ri_jk/resp_r.rs index e4973e70a..86e93a526 100644 --- a/src/ri_jk/resp_r.rs +++ b/src/ri_jk/resp_r.rs @@ -170,8 +170,12 @@ pub struct RRespRIJK<'a> { pub factor_j: f64, /// Exchange factor, absorbed into the returned fock/response tensors. pub factor_k: f64, - /// Cholesky decomposed 3c-2e ERI, shape `[nao_tp, naux]`. + /// Cholesky decomposed 3c-2e ERI of the SCF auxiliary basis (the high-precision resource), + /// shape `[nao_tp, naux]`. pub cderi: TsrCow<'a>, + /// Optional decomposed ERI of the response auxiliary basis (`resp_auxbas_path`, the + /// low-precision resource); `None` reuses `cderi`. + pub cderi_resp: Option>, /// Response intermediates: `mo_coeff [nao, nmo]` (RowMajor) and `mo_occ [nmo]`, stored by /// [`RRespAPI::make_response_preparation`]. pub intmd: HashMap, @@ -180,7 +184,27 @@ pub struct RRespRIJK<'a> { impl<'a> RRespRIJK<'a> { /// Create from a Cholesky-decomposed ERI (`cderi`). pub fn new_with_cderi(factor_j: f64, factor_k: f64, cderi: TsrCow<'a>) -> Self { - Self { factor_j, factor_k, cderi, intmd: HashMap::new() } + Self { factor_j, factor_k, cderi, cderi_resp: None, intmd: HashMap::new() } + } + + /// Attach a dedicated decomposed ERI (`cderi_resp`) of the response auxiliary basis for the + /// low-precision (`prec = false`) path. + /// + /// When set, the response contractions called with `prec = false` (the CP-SCF machinery) + /// evaluate on this tensor instead of the SCF-basis `cderi`. + pub fn set_cderi_resp(mut self, cderi_resp: TsrCow<'a>) -> Self { + self.cderi_resp = Some(cderi_resp); + self + } + + /// The decomposed ERI selected by the precision flag: the SCF-basis `cderi` for `prec = true`, + /// the response-basis [`Self::cderi_resp`] for `prec = false` (falling back to `cderi`). + fn cderi_in_use(&self, prec: bool) -> &TsrCow<'a> { + if prec { + &self.cderi + } else { + self.cderi_resp.as_ref().unwrap_or(&self.cderi) + } } } @@ -190,17 +214,18 @@ impl<'a> RRespAPI for RRespRIJK<'a> { /// Fock from density matrix: `factor_j * J(rdm) - 0.5 * factor_k * K(rdm)`. /// /// The `rdm` is assumed to be symmetric, as required by the pure in-core functions. - fn get_fock_rdm(&mut self, rdm: TsrView) -> Tsr { + fn get_fock_rdm(&mut self, rdm: TsrView, prec: bool) -> Tsr { assert_eq!(rdm.ndim(), 2, "rdm must have 2 dimensions"); let nao = rdm.shape()[0]; - let device = self.cderi.device(); + let cderi = self.cderi_in_use(prec); + let device = cderi.device(); // pure functions require (nao, nao, nset) f-contiguous inputs let dms = rdm.into_contig(ColMajor).into_shape((nao, nao, 1)); // TODO: batch size `72` should be tunable by max-memory. - let vj = get_vj_ri_incore(self.cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); - let vk = get_vk_ri_incore_dm(self.cderi.view(), dms.view(), 72).i((.., .., 0)).into_contig(ColMajor); + let vj = get_vj_ri_incore(cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); + let vk = get_vk_ri_incore_dm(cderi.view(), dms.view(), 72).i((.., .., 0)).into_contig(ColMajor); let mut fock = rt::zeros(([nao, nao], device)); if self.factor_j != 0.0 { @@ -214,20 +239,21 @@ impl<'a> RRespAPI for RRespRIJK<'a> { /// Fock from molecular coefficients: J by the dm route (no coeff-native vj exists), K by the /// coeff-native pure function. - fn get_fock_coeff(&mut self, mo_coeff: TsrView, mo_occ: TsrView) -> Tsr { + fn get_fock_coeff(&mut self, mo_coeff: TsrView, mo_occ: TsrView, prec: bool) -> Tsr { let [nao, nmo] = mo_coeff.shape().to_vec().try_into().unwrap(); - let device = self.cderi.device(); + let cderi = self.cderi_in_use(prec); + let device = cderi.device(); let dm = get_dm0_restricted(mo_coeff.view(), mo_occ.view()); let dms = dm.into_contig(ColMajor).into_shape((nao, nao, 1)); - let vj = get_vj_ri_incore(self.cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); + let vj = get_vj_ri_incore(cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); let coeff = mo_coeff.into_contig(ColMajor).into_shape((nao, nmo, 1)); let occ = mo_occ.into_contig(ColMajor).into_shape((nmo, 1)); // TODO: batch size `72` should be tunable by max-memory. let vk = - get_vk_ri_incore_coeff(self.cderi.view(), coeff.view(), occ.view(), 72).i((.., .., 0)).into_contig(ColMajor); + get_vk_ri_incore_coeff(cderi.view(), coeff.view(), occ.view(), 72).i((.., .., 0)).into_contig(ColMajor); let mut fock = rt::zeros(([nao, nao], device)); if self.factor_j != 0.0 { @@ -239,37 +265,39 @@ impl<'a> RRespAPI for RRespRIJK<'a> { fock } - fn get_response_rdm(&mut self, rdm: TsrView) -> Tsr { + fn get_response_rdm(&mut self, rdm: TsrView, prec: bool) -> Tsr { // assume rdm can be symmetrized, 4 * J - 2 * K assert_eq!(rdm.ndim(), 2, "rdm must have 2 dimensions"); let [nao, nao2] = rdm.shape().to_vec().try_into().unwrap(); assert_eq!(nao, nao2, "rdm must be square"); - let device = self.cderi.device(); + let cderi = self.cderi_in_use(prec); + let device = cderi.device(); let rdm_sym = ((&rdm + &rdm.t()) * 0.5).into_contig(ColMajor).into_shape((nao, nao, 1)); let mut resp = rt::zeros(([nao, nao], device)); if self.factor_j != 0.0 { - let vj = get_vj_ri_incore(self.cderi.view(), rdm_sym.view()).i((.., .., 0)).into_contig(ColMajor); + let vj = get_vj_ri_incore(cderi.view(), rdm_sym.view()).i((.., .., 0)).into_contig(ColMajor); resp += 4.0 * self.factor_j * &vj; } // TODO: batch size `72` should be tunable by max-memory. if self.factor_k != 0.0 { - let vk = get_vk_ri_incore_dm(self.cderi.view(), rdm_sym.view(), 72).i((.., .., 0)).into_contig(ColMajor); + let vk = get_vk_ri_incore_dm(cderi.view(), rdm_sym.view(), 72).i((.., .., 0)).into_contig(ColMajor); resp -= 2.0 * self.factor_k * &vk; } resp } - fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView) { + fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView, _prec: bool) { + // the intermediates are precision-independent (the decomposed ERI is selected per call) self.intmd.insert("mo_coeff".to_string(), mo_coeff.into_contig(RowMajor)); self.intmd.insert("mo_occ".to_string(), mo_occ.to_owned()); } - fn get_response_bra(&mut self, bra: TsrView) -> Tsr { + fn get_response_bra(&mut self, bra: TsrView, prec: bool) -> Tsr { let mo_coeff = self.intmd["mo_coeff"].view(); let mo_occ = self.intmd["mo_occ"].view(); - let cderi = self.cderi.view(); + let cderi = self.cderi_in_use(prec).view(); // RHF (single spin) assembly of the separated J/K response core. // - J (AO form, from total density) contracted with `mocc` and scaled by `factor_j`. diff --git a/src/ri_jk/resp_u.rs b/src/ri_jk/resp_u.rs index 3a67484f0..a935c59ad 100644 --- a/src/ri_jk/resp_u.rs +++ b/src/ri_jk/resp_u.rs @@ -30,8 +30,12 @@ pub struct URespRIJK<'a> { pub factor_j: f64, /// Exchange factor, absorbed into the returned fock/response tensors. pub factor_k: f64, - /// Cholesky decomposed 3c-2e ERI, shape `[nao_tp, naux]`. + /// Cholesky decomposed 3c-2e ERI of the SCF auxiliary basis (the high-precision resource), + /// shape `[nao_tp, naux]`. pub cderi: TsrCow<'a>, + /// Optional decomposed ERI of the response auxiliary basis (`resp_auxbas_path`, the + /// low-precision resource); `None` reuses `cderi`. + pub cderi_resp: Option>, /// Response intermediates: `mo_coeff_s [nao, nmo_s]` (RowMajor) and `mo_occ_s [nmo_s]`, stored /// by [`URespAPI::make_response_preparation`]. pub intmd: HashMap, @@ -40,7 +44,27 @@ pub struct URespRIJK<'a> { impl<'a> URespRIJK<'a> { /// Create from a Cholesky-decomposed ERI (`cderi`). pub fn new_with_cderi(factor_j: f64, factor_k: f64, cderi: TsrCow<'a>) -> Self { - Self { factor_j, factor_k, cderi, intmd: HashMap::new() } + Self { factor_j, factor_k, cderi, cderi_resp: None, intmd: HashMap::new() } + } + + /// Attach a dedicated decomposed ERI (`cderi_resp`) of the response auxiliary basis for the + /// low-precision (`prec = false`) path. + /// + /// When set, the response contractions called with `prec = false` (the CP-SCF machinery) + /// evaluate on this tensor instead of the SCF-basis `cderi`. + pub fn set_cderi_resp(mut self, cderi_resp: TsrCow<'a>) -> Self { + self.cderi_resp = Some(cderi_resp); + self + } + + /// The decomposed ERI selected by the precision flag: the SCF-basis `cderi` for `prec = true`, + /// the response-basis [`Self::cderi_resp`] for `prec = false` (falling back to `cderi`). + fn cderi_in_use(&self, prec: bool) -> &TsrCow<'a> { + if prec { + &self.cderi + } else { + self.cderi_resp.as_ref().unwrap_or(&self.cderi) + } } } @@ -55,20 +79,21 @@ impl<'a> URespAPI for URespRIJK<'a> { /// spin density carries occupation 1. /// /// The `rdm` is assumed to be symmetric, as required by the pure in-core functions. - fn get_fock_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + fn get_fock_rdm(&mut self, rdm: &[TsrView; 2], prec: bool) -> [Tsr; 2] { let [α, β] = [0, 1]; for rdm_s in rdm { assert_eq!(rdm_s.ndim(), 2, "rdm must have 2 dimensions"); } let nao = rdm[α].shape()[0]; - let device = self.cderi.device(); + let cderi = self.cderi_in_use(prec); + let device = cderi.device(); let mut fock = [rt::zeros(([nao, nao], device)), rt::zeros(([nao, nao], device))]; if self.factor_j != 0.0 { // J: total (α+β) density, spin-independent let dm_total = &rdm[α] + &rdm[β]; let dms = dm_total.into_contig(ColMajor).into_shape((nao, nao, 1)); - let vj = get_vj_ri_incore(self.cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); + let vj = get_vj_ri_incore(cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); fock[α] += self.factor_j * &vj; fock[β] += self.factor_j * &vj; } @@ -77,7 +102,7 @@ impl<'a> URespAPI for URespRIJK<'a> { for (s, fock_s) in fock.iter_mut().enumerate() { let dms = rdm[s].view().into_contig(ColMajor).into_shape((nao, nao, 1)); // TODO: batch size `72` should be tunable by max-memory. - let vk = get_vk_ri_incore_dm(self.cderi.view(), dms.view(), 72).i((.., .., 0)).into_contig(ColMajor); + let vk = get_vk_ri_incore_dm(cderi.view(), dms.view(), 72).i((.., .., 0)).into_contig(ColMajor); *fock_s -= self.factor_k * &vk; } } @@ -86,10 +111,11 @@ impl<'a> URespAPI for URespRIJK<'a> { /// Fock per spin from molecular coefficients: J by the dm route (no coeff-native vj exists), K /// by the coeff-native pure function per spin. - fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) -> [Tsr; 2] { + fn get_fock_coeff(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], prec: bool) -> [Tsr; 2] { let [α, β] = [0, 1]; let nao = mo_coeff[α].shape()[0]; - let device = self.cderi.device(); + let cderi = self.cderi_in_use(prec); + let device = cderi.device(); let mut fock = [rt::zeros(([nao, nao], device)), rt::zeros(([nao, nao], device))]; if self.factor_j != 0.0 { @@ -97,7 +123,7 @@ impl<'a> URespAPI for URespRIJK<'a> { let dm_total = get_dm0_restricted(mo_coeff[α].view(), mo_occ[α].view()) + get_dm0_restricted(mo_coeff[β].view(), mo_occ[β].view()); let dms = dm_total.into_contig(ColMajor).into_shape((nao, nao, 1)); - let vj = get_vj_ri_incore(self.cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); + let vj = get_vj_ri_incore(cderi.view(), dms.view()).i((.., .., 0)).into_contig(ColMajor); fock[α] += self.factor_j * &vj; fock[β] += self.factor_j * &vj; } @@ -109,7 +135,7 @@ impl<'a> URespAPI for URespRIJK<'a> { let occ = mo_occ[s].view().into_contig(ColMajor).into_shape((nmo, 1)); // TODO: batch size `72` should be tunable by max-memory. let vk = - get_vk_ri_incore_coeff(self.cderi.view(), coeff.view(), occ.view(), 72).i((.., .., 0)).into_contig(ColMajor); + get_vk_ri_incore_coeff(cderi.view(), coeff.view(), occ.view(), 72).i((.., .., 0)).into_contig(ColMajor); *fock_s -= self.factor_k * &vk; } } @@ -126,13 +152,14 @@ impl<'a> URespAPI for URespRIJK<'a> { /// two-term symmetrized exchange kernel keeps the restricted scale. In the closed-shell /// limit (`rdm_α = rdm_β = X`), each spin's response equals the restricted kernel on `X`. /// No preparation is required. - fn get_response_rdm(&mut self, rdm: &[TsrView; 2]) -> [Tsr; 2] { + fn get_response_rdm(&mut self, rdm: &[TsrView; 2], prec: bool) -> [Tsr; 2] { for rdm_s in rdm { assert_eq!(rdm_s.ndim(), 2, "rdm must have 2 dimensions"); } let [α, β] = [0, 1]; let nao = rdm[α].shape()[0]; - let device = self.cderi.device(); + let cderi = self.cderi_in_use(prec); + let device = cderi.device(); // symmetrize (the pure in-core functions require symmetric input); the Coulomb response // sees the total (α+β) density @@ -144,7 +171,7 @@ impl<'a> URespAPI for URespRIJK<'a> { let mut resp = [rt::zeros(([nao, nao], device)), rt::zeros(([nao, nao], device))]; if self.factor_j != 0.0 { - let vj = get_vj_ri_incore(self.cderi.view(), rdm_total.view()).i((.., .., 0)).into_contig(ColMajor); + let vj = get_vj_ri_incore(cderi.view(), rdm_total.view()).i((.., .., 0)).into_contig(ColMajor); // 2.0 against the restricted 4.0: occupation 1 vs 2 on the total-density Coulomb resp[α] += 2.0 * self.factor_j * &vj; resp[β] += 2.0 * self.factor_j * &vj; @@ -153,24 +180,25 @@ impl<'a> URespAPI for URespRIJK<'a> { // K: same-spin exchange, same 2.0 scale as restricted (two-term symmetrized kernel) for (s, resp_s) in resp.iter_mut().enumerate() { // TODO: batch size `72` should be tunable by max-memory. - let vk = get_vk_ri_incore_dm(self.cderi.view(), rdm_sym[s].view(), 72).i((.., .., 0)).into_contig(ColMajor); + let vk = get_vk_ri_incore_dm(cderi.view(), rdm_sym[s].view(), 72).i((.., .., 0)).into_contig(ColMajor); *resp_s -= 2.0 * self.factor_k * &vk; } } resp } - fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], _prec: bool) { + // the intermediates are precision-independent (the decomposed ERI is selected per call) self.intmd.insert("mo_coeff_0".to_string(), mo_coeff[0].view().into_contig(RowMajor)); self.intmd.insert("mo_coeff_1".to_string(), mo_coeff[1].view().into_contig(RowMajor)); self.intmd.insert("mo_occ_0".to_string(), mo_occ[0].to_owned()); self.intmd.insert("mo_occ_1".to_string(), mo_occ[1].to_owned()); } - fn get_response_bra(&mut self, bra: &[TsrView; 2]) -> [Tsr; 2] { + fn get_response_bra(&mut self, bra: &[TsrView; 2], prec: bool) -> [Tsr; 2] { let mo_coeff = [self.intmd["mo_coeff_0"].view(), self.intmd["mo_coeff_1"].view()]; let mo_occ = [self.intmd["mo_occ_0"].view(), self.intmd["mo_occ_1"].view()]; - let cderi = self.cderi.view(); + let cderi = self.cderi_in_use(prec).view(); let device = mo_coeff[0].device(); // Shared separated J/K response core: J (AO form, from total density) + per-spin K (bra form). let (j_ao, k_bras) = get_rijk_response_bra_separated( diff --git a/src/ri_pt2/rgfock_pt2.rs b/src/ri_pt2/rgfock_pt2.rs index 26d209a2d..1d7656d8f 100644 --- a/src/ri_pt2/rgfock_pt2.rs +++ b/src/ri_pt2/rgfock_pt2.rs @@ -169,7 +169,10 @@ where self.make_elec_deriv(); let rdm1 = self.result["rdm1"].view(); let dm_ao = self.mo_coeff.view() % rdm1 % self.mo_coeff.view().t(); - let resp_ao = resp.get_response_rdm(dm_ao.view()); + // high precision: this rdm-form A-contraction belongs to the generalized-Fock + // (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 nocc = self.nocc(); let nmo = self.nmo(); diff --git a/tests/analdrv/response/uresp_gfock.rs b/tests/analdrv/response/uresp_gfock.rs index b42cc0e85..c46318825 100644 --- a/tests/analdrv/response/uresp_gfock.rs +++ b/tests/analdrv/response/uresp_gfock.rs @@ -1,9 +1,12 @@ //! Unrestricted response-kernel rdm form (`get_response_rdm`, the ugfock groundwork) for -//! URespRIJK and URespKSNIMatmul, pinned two ways: +//! URespRIJK and URespKSNIMatmul, pinned three ways: //! - consistency with the hessian-tested bra form (`get_response_bra`): the rdm form on the //! symmetrized bra-generated densities, right half-transformed by `mocc`, equals the bra form; //! - closed-shell limit: with both spins carrying the restricted test density, each spin's -//! response equals the restricted kernel (`4 J - 2 K` / `4 fxc`). +//! response equals the restricted kernel (`4 J - 2 K` / `4 fxc`); +//! - precision plumbing: with the low-precision resource (`cderi_resp` / `ni_resp`) set to a +//! duplicate of the high-precision one, every form agrees between `prec = true` and +//! `prec = false`. use pyrest::analdrv::response::rresp_interface::{scf_jk_factors, scf_xc_func_list}; use pyrest::analdrv::response::trait_rresp::RRespAPI; @@ -91,16 +94,16 @@ fn test_oh_uks_response_rdm_consistency() { let (factor_j, factor_k, _) = scf_jk_factors(&scf_data); let (rimatr, _, _) = scf_data.rimatr.as_ref().unwrap(); let mut rijk = URespRIJK::new_with_cderi(factor_j, factor_k, rimatr.to_rstsr_view(&device).into_cow()); - rijk.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()]); + rijk.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()], true); let occidx = [mo_occ[0].view().greater(0).into_vec(), mo_occ[1].view().greater(0).into_vec()]; let mocc = [mo_coeff[0].bool_select(-1, &occidx[0]), mo_coeff[1].bool_select(-1, &occidx[1])]; let bra = [det_matrix(nao, mocc[0].shape()[1], 1, &device), det_matrix(nao, mocc[1].shape()[1], 2, &device)]; - let resp_bra = rijk.get_response_bra(&[bra[0].view(), bra[1].view()]); + let resp_bra = rijk.get_response_bra(&[bra[0].view(), bra[1].view()], true); let x = [&bra[0] % mocc[0].t(), &bra[1] % mocc[1].t()]; let x_sym = [(&x[0] + &x[0].t()) * 0.5, (&x[1] + &x[1].t()) * 0.5]; - let resp_rdm = rijk.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()]); + let resp_rdm = rijk.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()], true); for s in 0..2 { let expect = resp_rdm[s].view() % mocc[s].view(); assert!(rt::allclose(&resp_bra[s], &expect, (1e-8, 1e-9)), "RI-JK bra/rdm mismatch, spin {s}"); @@ -111,10 +114,10 @@ fn test_oh_uks_response_rdm_consistency() { let mol_cint = get_cint_mol(&scf_data.mol); let ni = NIMatmul::new(&mol_cint, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); let mut uks = URespKSNIMatmul::new(scf_xc_func_list_uks(&scf_data), ni, false); - uks.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()]); + uks.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()], true); - let resp_bra = uks.get_response_bra(&[bra[0].view(), bra[1].view()]); - let resp_rdm = uks.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()]); + let resp_bra = uks.get_response_bra(&[bra[0].view(), bra[1].view()], true); + let resp_rdm = uks.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()], true); for s in 0..2 { let expect = resp_rdm[s].view() % mocc[s].view(); assert!(rt::allclose(&resp_bra[s], &expect, (1e-7, 1e-8)), "UKS bra/rdm mismatch, spin {s}"); @@ -138,8 +141,8 @@ fn test_closed_shell_limit() { let (rimatr, _, _) = scf_data.rimatr.as_ref().unwrap(); let mut rrijk = RRespRIJK::new_with_cderi(factor_j, factor_k, rimatr.to_rstsr_view(&device).into_cow()); let mut urijk = URespRIJK::new_with_cderi(factor_j, factor_k, rimatr.to_rstsr_view(&device).into_cow()); - let resp_r = rrijk.get_response_rdm(x.view()); - let resp_u = urijk.get_response_rdm(&[x.view(), x.view()]); + let resp_r = rrijk.get_response_rdm(x.view(), true); + let resp_u = urijk.get_response_rdm(&[x.view(), x.view()], true); for s in 0..2 { assert!(rt::allclose(&resp_u[s], &resp_r, (1e-8, 1e-9)), "RI-JK closed-shell mismatch, spin {s}"); } @@ -151,13 +154,76 @@ fn test_closed_shell_limit() { let ni_u = NIMatmul::new(&mol_cint, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); let mut rks = RRespKSNIMatmul::new(scf_xc_func_list(&scf_data), ni_r, false); let mut uks = URespKSNIMatmul::new(scf_xc_func_list_uks(&scf_data), ni_u, false); - rks.make_response_preparation(mo_coeff.view(), mo_occ.view()); + rks.make_response_preparation(mo_coeff.view(), mo_occ.view(), true); let occ_half = &mo_occ * 0.5_f64; - uks.make_response_preparation(&[mo_coeff.view(), mo_coeff.view()], &[occ_half.view(), occ_half.view()]); + uks.make_response_preparation(&[mo_coeff.view(), mo_coeff.view()], &[occ_half.view(), occ_half.view()], true); - let resp_r = rks.get_response_rdm(x.view()); - let resp_u = uks.get_response_rdm(&[x.view(), x.view()]); + let resp_r = rks.get_response_rdm(x.view(), true); + let resp_u = uks.get_response_rdm(&[x.view(), x.view()], true); for s in 0..2 { assert!(rt::allclose(&resp_u[s], &resp_r, (1e-7, 1e-8)), "UKS closed-shell mismatch, spin {s}"); } } + +/// Precision plumbing: with the low-precision resource (`cderi_resp` of `URespRIJK`, `ni_resp` +/// of `URespKSNIMatmul`) set to a duplicate of the high-precision one, every form agrees between +/// `prec = true` and `prec = false` — and the attached object agrees with an unattached one. +#[test] +fn test_precision_selection_plumbing() { + let scf_data = build_scf(INPUT_OH_UKS); + let device = rt::DeviceBLAS::default(); + let mo_coeff = [(&scf_data.eigenvectors[0]).to_rstsr(&device), (&scf_data.eigenvectors[1]).to_rstsr(&device)]; + let mo_occ = [(&scf_data.occupation[0]).to_rstsr(&device), (&scf_data.occupation[1]).to_rstsr(&device)]; + let nao = mo_coeff[0].shape()[0]; + + let occidx = [mo_occ[0].view().greater(0).into_vec(), mo_occ[1].view().greater(0).into_vec()]; + let mocc = [mo_coeff[0].bool_select(-1, &occidx[0]), mo_coeff[1].bool_select(-1, &occidx[1])]; + let bra = [det_matrix(nao, mocc[0].shape()[1], 1, &device), det_matrix(nao, mocc[1].shape()[1], 2, &device)]; + let x = [&bra[0] % mocc[0].t(), &bra[1] % mocc[1].t()]; + let x_sym = [(&x[0] + &x[0].t()) * 0.5, (&x[1] + &x[1].t()) * 0.5]; + + // --- URespRIJK: cderi_resp as a second view of the same rimatr --- // + let (factor_j, factor_k, _) = scf_jk_factors(&scf_data); + let (rimatr, _, _) = scf_data.rimatr.as_ref().unwrap(); + let mut rijk = URespRIJK::new_with_cderi(factor_j, factor_k, rimatr.to_rstsr_view(&device).into_cow()) + .set_cderi_resp(rimatr.to_rstsr_view(&device).into_cow()); + let mut rijk_ref = URespRIJK::new_with_cderi(factor_j, factor_k, rimatr.to_rstsr_view(&device).into_cow()); + + rijk.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()], true); + let bra_high = rijk.get_response_bra(&[bra[0].view(), bra[1].view()], true); + let rdm_high = rijk.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()], true); + let fock_high = rijk.get_fock_rdm(&[x_sym[0].view(), x_sym[1].view()], true); + rijk.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()], false); + let bra_low = rijk.get_response_bra(&[bra[0].view(), bra[1].view()], false); + let rdm_low = rijk.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()], false); + let fock_low = rijk.get_fock_rdm(&[x_sym[0].view(), x_sym[1].view()], false); + + rijk_ref.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()], true); + let bra_ref = rijk_ref.get_response_bra(&[bra[0].view(), bra[1].view()], true); + + for s in 0..2 { + assert!(rt::allclose(&bra_low[s], &bra_high[s], (1e-10, 1e-12)), "RI-JK bra precision mismatch, spin {s}"); + assert!(rt::allclose(&rdm_low[s], &rdm_high[s], (1e-10, 1e-12)), "RI-JK rdm precision mismatch, spin {s}"); + assert!(rt::allclose(&fock_low[s], &fock_high[s], (1e-10, 1e-12)), "RI-JK fock precision mismatch, spin {s}"); + assert!(rt::allclose(&bra_high[s], &bra_ref[s], (1e-10, 1e-12)), "RI-JK attached-vs-plain mismatch, spin {s}"); + } + + // --- URespKSNIMatmul: ni_resp as a second driver over the same grid --- // + let grids = scf_data.grids.as_ref().unwrap(); + let mol_cint = get_cint_mol(&scf_data.mol); + let ni = NIMatmul::new(&mol_cint, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); + let ni_resp = NIMatmul::new(&mol_cint, &grids.coordinates, &grids.weights, &grids.atm_idx, &grids.quadrature_weights); + let mut uks = URespKSNIMatmul::new(scf_xc_func_list_uks(&scf_data), ni, false).set_ni_resp(ni_resp); + + uks.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()], true); + let bra_high = uks.get_response_bra(&[bra[0].view(), bra[1].view()], true); + let rdm_high = uks.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()], true); + uks.make_response_preparation(&[mo_coeff[0].view(), mo_coeff[1].view()], &[mo_occ[0].view(), mo_occ[1].view()], false); + let bra_low = uks.get_response_bra(&[bra[0].view(), bra[1].view()], false); + let rdm_low = uks.get_response_rdm(&[x_sym[0].view(), x_sym[1].view()], false); + + for s in 0..2 { + assert!(rt::allclose(&bra_low[s], &bra_high[s], (1e-10, 1e-12)), "UKS bra precision mismatch, spin {s}"); + assert!(rt::allclose(&rdm_low[s], &rdm_high[s], (1e-10, 1e-12)), "UKS rdm precision mismatch, spin {s}"); + } +} -- Gitee From f0fe6433ab034d7ef9b84e461b4f0468e650bf31 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 24 Sep 2026 10:50:31 +0800 Subject: [PATCH 08/21] Resolve resp_auxbas_path through the basis-path machinery resp_auxbas_path was deserialized raw and handed to get_cint_aux_on_path, which only understands literal element-JSON directories, so pool names and aliases silently missed the pre-installed basis-set pool. Wrap the [analdrv] section parsing in parse_analdrv_keywords (the parse_geom_keywords pattern) and resolve resp_auxbas_path there via path_util::get_valid_basis_path, exactly like auxbas_path. As a side effect, the "equal to the SCF auxiliary basis treated as unset" exact match in cderi_pair/cderi_pair_sr now also fires for pool-name inputs, since both sides hold resolved values. Make ctrl_io::path_util public for this, and drop the unused print_level argument of get_rest_basis_dir. Correct the get_cint_aux_on_path contract: the basis machinery does not resolve pool names; the input must be the resolved form ctrl.auxbas_path holds after parsing. Co-authored-by: Claude Code Co-authored-by: glm-5.3 --- src/analdrv/config.rs | 23 ++++++++++++++++++++++- src/ctrl_io/mod.rs | 8 ++++---- src/ctrl_io/path_util.rs | 2 +- src/ri_jk/util.rs | 4 +++- 4 files changed, 30 insertions(+), 7 deletions(-) diff --git a/src/analdrv/config.rs b/src/analdrv/config.rs index 6cf57b2d1..33d4aa8b2 100644 --- a/src/analdrv/config.rs +++ b/src/analdrv/config.rs @@ -1,5 +1,6 @@ use serde::{Deserialize, Serialize}; use serde_inline_default::serde_inline_default; +use crate::ctrl_io::{path_util, serde_from_value}; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub enum AnalDrvTask { @@ -83,7 +84,8 @@ pub struct AnalDrvRespCfg { pub grid_level: Option, /// Auxiliary basis of the RI-JK response objects (the low-precision response side), as a /// basis-set pool name or an element-JSON directory — the same value format as - /// `ctrl.auxbas_path`. + /// `ctrl.auxbas_path`, and resolved by [`parse_analdrv_keywords`] through the same + /// basis-path machinery (so pool names and aliases work, not only absolute directories). /// /// By default `None`: the response objects reuse the SCF auxiliary basis, borrowing the SCF /// `rimatr` (and `rimatr_sr` for range-separated hybrids) with no copy. When set, a @@ -282,6 +284,25 @@ pub struct AnalDrvConfig { /* #endregion AnalDrvConfig */ +/// Parse the `[analdrv]` section of the control input into an [`AnalDrvConfig`]. +/// +/// `None` when the section is absent; a malformed section panics (the `serde_from_value` +/// convention of the other sections). The `resp_auxbas_path` keyword takes the same value format +/// as `auxbas_path` (a basis-set pool name, an alias, or an element-JSON directory), so it is +/// resolved through the same basis-path machinery here: pool names are looked up in the basis-set +/// pool directories, and anything else is passed through for the basis-set-exchange fallback. +pub fn parse_analdrv_keywords(tmp_keys: &serde_json::Value) -> Option { + let mut config: AnalDrvConfig = tmp_keys.get("analdrv").map(serde_from_value)?; + if let Some(resp_auxbas_path) = &config.resp.auxbas_path { + config.resp.auxbas_path = Some(path_util::get_valid_basis_path( + resp_auxbas_path, + &path_util::get_rest_basis_dir(), + "auxiliary basis", + )); + } + Some(config) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index a646f1fc1..0dafe6888 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -7,7 +7,7 @@ use std::{fs}; use crate::ctrl_io::geometric_pyo3_io::parse_geometric_keywords; use crate::ctrl_io::quasiparticle_methods::parse_quasiparticle_keywords; use crate::ri_jk::decompose::J2CDecompOption; -use crate::analdrv::config::{AnalDrvConfig, AnalDrvTask}; +use crate::analdrv::config::{AnalDrvConfig, AnalDrvTask, parse_analdrv_keywords}; use crate::ctrl_io::tddft_parameters::parse_tddft_keywords; use crate::ctrl_io::hessian_parameters::parse_hessian_keywords; use crate::ctrl_io::thermo_parameters::parse_thermo_keywords; @@ -38,7 +38,7 @@ pub mod tddft_parameters; pub mod hessian_parameters; pub mod thermo_parameters; use geometric_pyo3_io::GeomeTRIC; -mod path_util; +pub mod path_util; use quasiparticle_methods::QuasiParticle; use tddft_parameters::TDDFTParameters; use hessian_parameters::HessianParameters; @@ -86,7 +86,7 @@ pub fn parse_ctl_from_json(tmp_keys: &serde_json::Value) -> anyhow::Result<(Inpu if let Some(tmp_thermo) = &mut tmp_thermo { tmp_input.thermo = Some(std::mem::take(tmp_thermo)); } - tmp_input.analdrv = tmp_keys.get("analdrv").map(serde_from_value); + tmp_input.analdrv = parse_analdrv_keywords(tmp_keys); Ok((tmp_input,tmp_geomcell)) } @@ -922,7 +922,7 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result { diff --git a/src/ctrl_io/path_util.rs b/src/ctrl_io/path_util.rs index 29995375c..15d461f7e 100644 --- a/src/ctrl_io/path_util.rs +++ b/src/ctrl_io/path_util.rs @@ -13,7 +13,7 @@ const DEFAULT_BASIS_PATH: &str = "/opt/rest_workspace/rest/basis-set-pool/"; // }; // } -pub fn get_rest_basis_dir(print_level:usize) -> Vec { +pub fn get_rest_basis_dir() -> Vec { let mut rest_basis_dir:Vec = vec![]; if let Some(val) = env::var("REST_BASIS_DIR").ok() { if !val.is_empty() { diff --git a/src/ri_jk/util.rs b/src/ri_jk/util.rs index 2b9faeabd..7f031d6df 100644 --- a/src/ri_jk/util.rs +++ b/src/ri_jk/util.rs @@ -10,7 +10,9 @@ pub fn get_cint_aux(mol_obj: &Molecule) -> CInt { } /// [`get_cint_aux`] on an auxiliary basis given by `auxbas_path` instead of the molecule's own -/// `ctrl.auxbas_path` (same value format; pool names are resolved by the basis machinery). +/// `ctrl.auxbas_path`. The path must already be resolved by the ctrl-input basis-path machinery +/// (the form `ctrl.auxbas_path` itself holds after parsing): a directory of element JSONs, with +/// missing elements left to the basis-set-exchange fallback of the basis machinery. /// /// The molecule is only read: the auxiliary-basis information is regenerated from a local copy of /// the keywords, and nothing on `mol_obj` is modified. -- Gitee From 28fd680313da48fe4564914a653f7284c784683a Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 24 Sep 2026 11:09:37 +0800 Subject: [PATCH 09/21] fix compile error on !214 edit merge --- src/ri_tddft/tddft.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ri_tddft/tddft.rs b/src/ri_tddft/tddft.rs index 16f7b5f6c..acafeee41 100644 --- a/src/ri_tddft/tddft.rs +++ b/src/ri_tddft/tddft.rs @@ -377,7 +377,7 @@ pub fn prepare_ao_data_with_spin(scf: &SCF, tddft_spin: Option<&str>) -> TDDFTDa .zip(xc_data.dfa_paramr_scf.iter()) .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, LibXCSpin::Polarized))) .collect(); - let (_vxc0, fxc_eff_u) = crate::dft::numint_matmul::hess_uks::eval_vxc_fxc_uks_from_rho( + let (_vxc0, fxc_eff_u) = crate::dft::numint_matmul::resp_uks::eval_vxc_fxc_uks_from_rho( &xc_func_list_pol, rho0_u.view()); (rho0_u, fxc_eff_u) } else { -- Gitee From 9a65c71f68b7c1968852ae3d71fbd34f4ca038bf Mon Sep 17 00:00:00 2001 From: ajz34 Date: Tue, 29 Sep 2026 16:15:28 +0800 Subject: [PATCH 10/21] Add full four-block generalized Fock evaluation to RI-PT2 Kernel (get_rpt2_elec_deriv_incore): the new full_gfock flag of RPT2ElecDerivIncoreArg evaluates the full four-block contribution (OO+OV+VO+VV, batched over the auxiliary index, with the 2*eps*rdm1_corr orbital-energy terms in the diagonal blocks) in place of the W3/W4 off-diagonal blocks; the SCF response term stays outside the kernel in both modes. RGFockPT2: make_gfock now follows the make_ caching convention. The accumulated gfock/gfock_flags/gfock_part_full state fields are dropped in favor of result-map caches: the kernel output per mode ("gfock_part"/ "gfock_part_full"; a legacy cache is upgraded on a full request, a full cache serves any request) and the evaluated blocks ("gfock_oo"/"gfock_ov"/ "gfock_vo"/"gfock_vv", response term included for OO/VO). Repeated calls reassemble directly from the cached blocks; the resp object is consulted only while a block is actually evaluated and is never a cache key. Test: element-wise full-gfock fingerprint against the pyscf reference pipeline (NH3/def2-SVP, MO-sign canonicalization). Co-authored-by: Claude Code Co-authored-by: glm-5.3 --- src/ri_pt2/pure_pt2_r_elecderiv.rs | 102 +++++++++++++--- src/ri_pt2/rgfock_pt2.rs | 170 +++++++++++++++++---------- tests/analdrv/response/pt2_rgfock.rs | 110 ++++++++++++++++- 3 files changed, 301 insertions(+), 81 deletions(-) diff --git a/src/ri_pt2/pure_pt2_r_elecderiv.rs b/src/ri_pt2/pure_pt2_r_elecderiv.rs index df800a503..18542cb30 100644 --- a/src/ri_pt2/pure_pt2_r_elecderiv.rs +++ b/src/ri_pt2/pure_pt2_r_elecderiv.rs @@ -28,16 +28,21 @@ 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, } 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, @@ -59,7 +64,8 @@ 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 } = arg; + let full_gfock = *full_gfock; let bi1_scale = O::from_f64(2.0 * c_os).unwrap(); let bi2_scale = O::from_f64(*c_ss).unwrap(); @@ -90,7 +96,28 @@ where let mut gfock_part = rt::zeros(([nmo, nmo].f(), &device)); let mut rdm1_corr = rt::zeros(([nmo, nmo].f(), &device)); 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 --- // @@ -188,36 +215,75 @@ where // --- 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 } } diff --git a/src/ri_pt2/rgfock_pt2.rs b/src/ri_pt2/rgfock_pt2.rs index 1d7656d8f..8c5644232 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 @@ -76,12 +80,11 @@ where 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. + /// 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 +107,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 +116,6 @@ where index_occ_outer_vec, c_os, c_ss, - gfock: rt::zeros(([nmo, nmo].f(), &device)), - gfock_flags: BitFlags::empty(), result: HashMap::new(), e_corr: None, timing: Vec::new(), @@ -127,7 +126,18 @@ 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")) { return; } let t0 = std::time::Instant::now(); @@ -148,23 +158,30 @@ 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 }; 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); 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 +190,27 @@ 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") { + 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 +232,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 +290,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/tests/analdrv/response/pt2_rgfock.rs b/tests/analdrv/response/pt2_rgfock.rs index a16ae566d..d74f9eac5 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 }; 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 }; 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 --- // -- Gitee From 6204707d663f695bc9456e25b9e71476e69e70c6 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 30 Sep 2026 20:12:03 +0800 Subject: [PATCH 11/21] Add restricted RI-MP2 analytic gradient New src/grad/rdh.rs: RDHGradient implements GradAPI, computing the PT2 correlation gradient of doubly-hybrid methods (for xc = "mp2" the SCF HF part stays the existing "SCF" entry of eval_force). - density side reuses RGFockPT2 (lagrangian + axd) and RRespSCF (Z-vector) to obtain the relaxed correlation density and the amplitude intermediate G (exposed via new dump_g_vix option); - derivative-integral side transforms raw rimatr columns in place against a lower-Cholesky int2c2e decomposition per aux-shell batch; batched intermediates stay within 2 * nocc * nao * naux via calc_batch_size and handle_memory_exceed; no additional rimatr-sized tensor is allocated; - eval_force relaxes the fifth-DFA panic for the PT2 family (non-MPI) and appends the ("DH") gradient entry after "SCF"; - requires j2c_decomp policy = "cd" (the eig convention is not supported by this route and panics with a clear message); - tests/grad/rdh.rs: NH3/def2-SVP/def2-universal-jkfit reference test, agrees 1.3e-8 with pyscf-forge, 6.9e-6 with in-crate numerical_force. Co-authored-by: Claude Code Co-authored-by: glm-5.3 --- src/grad/mod.rs | 1 + src/grad/rdh.rs | 461 +++++++++++++++++++++++++++ src/main_driver.rs | 16 +- src/ri_pt2/pure_pt2_r_elecderiv.rs | 18 +- src/ri_pt2/rgfock_pt2.rs | 10 +- tests/analdrv/response/pt2_rgfock.rs | 4 +- tests/grad/mod.rs | 1 + tests/grad/rdh.rs | 87 +++++ 8 files changed, 590 insertions(+), 8 deletions(-) create mode 100644 src/grad/rdh.rs create mode 100644 tests/grad/rdh.rs diff --git a/src/grad/mod.rs b/src/grad/mod.rs index 189dfff12..920c0add6 100644 --- a/src/grad/mod.rs +++ b/src/grad/mod.rs @@ -1,3 +1,4 @@ +pub mod rdh; 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 000000000..abe0fc6ec --- /dev/null +++ b/src/grad/rdh.rs @@ -0,0 +1,461 @@ +//! Restricted doubly-hybrid (`xc = "mp2"`) analytic gradient: the PT2-correlation part. +//! +//! The SCF (HF) part of the gradient is provided by the existing `"SCF"` entry of +//! [`crate::main_driver::eval_force`]; this entry covers the D_r-carrying (correlation) +//! remainder, following the D_r-dependent blocks of 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 = C D_r^S`; the `0.5 C D_mo` exchange part belongs to the `"SCF"` entry); +//! - `de_rint`: `4 ` (forge `_get_gradient_pt2`, third term). +//! +//! The rimatr must be CD-decomposed (`[ctrl.j2c_decomp] policy = "cd"`), matching the +//! pyscf-forge (and pyscf `df`) convention `cderi[P, uv] = (L^-1 raw3c)[P, uv]`; every +//! derivative contraction left-solves against the same lower-Cholesky factor. + +use crate::analdrv::config::AnalDrvConfig; +use crate::analdrv::response::rresp_interface::{rscf_resp_interface, RRespSCF}; +use crate::analdrv::response::trait_rgfock::{GFockFlags, RGFockAPI}; +use crate::analdrv::response::trait_rresp::RRespAPI; +use crate::grad::traits::GradAPI; +use crate::mpi_io::MPIOperator; +use crate::ri_jk::pure_decompose::solve_by_j2c_mut; +use crate::ri_jk::util; +use crate::ri_jk::{get_j2c_decomp, J2CDecompOption, J2CDecompPolicy, J2CDecompose, J2C_THRESH}; +use crate::ri_pt2::rgfock_pt2::RGFockPT2; +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, +} + +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 } + } + + 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."); + if mol_obj.ctrl.j2c_decomp.policy != J2CDecompPolicy::Cd { + panic!( + "Doubly-hybrid analytic gradient requires the CD-decomposed rimatr. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input." + ); + } + 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 = AnalDrvConfig::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."); + let j3c = rimatr.0.to_rstsr_view(&device); + let mut pt2 = RGFockPT2::::new( + mo_coeff.clone(), + mo_occ, + mo_energy.clone(), + j3c.into_cow(), + None, + vec![0, nocc], + 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 _ = pt2.make_rdm1_resp(&mut resp_objs); + let rdm1_resp = pt2.result["rdm1_resp"].to_owned(); + let gfock_part = pt2.result["gfock_part_full"].to_owned(); + let rdm1 = pt2.result["rdm1"].to_owned(); + let g_vix = pt2.result["g_vix"].to_owned(); // [nvir, nocc, naux] + self.energy = pt2.e_corr.unwrap_or(0.0); + + 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); + let c1_dr = (mo_coeff.view() % d_r_s.view()).into_contig(ColMajor); + let d_r_ao = (mo_coeff.view() % d_r_s.view() % mo_coeff.view().t()).into_contig(ColMajor); + + // 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 mut hcore_gen = crate::grad::rhf::generator_deriv_hcore(scf_data); + let tsr_ipovlp = { + let (out, shape) = mol.integrate("int1e_ipovlp", "s1", None).into(); + rt::asarray((out, shape, &device)) + }; + // 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); + } + } + + // ── CD metric objects ── + let j2c_decomp = get_j2c_decomp( + &aux, + &device, + J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(J2C_THRESH), ..Default::default() }, + ); + let j2c_l = match &j2c_decomp { + J2CDecompose::Cd { j2c_l, uplo: Upper, .. } => j2c_l.clone(), + _ => unreachable!("CD policy checked above"), + }; + // U^-1 solves U X = I (j2c = U^T U); forge's L_inv = inv(L_lower) = (U^-1)^T + let l_inv = { + 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) + }; + let l_lower = j2c_l.t().into_contig(ColMajor); + 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 `generator_L_1` + 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; + } + + // ── persistent MO intermediates (Y_ip, metric couplings R/R2); peak 2 * nocc nao naux ── + 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; + handle_memory_exceed((nocc * nmo * naux + nvir * nocc * naux) as f64 * 8.0 / 1048576.0, mem_avail, abort); + let mut y_ip: Tsr = rt::zeros(([nocc, nmo, naux], &device)); + let mut y_dot_d: Tsr = rt::zeros(([naux], &device)); + let mut y_dot_dr: Tsr = rt::zeros(([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())); + let s_d: f64 = (&col * d_hf.view()).sum_all(); + y_dot_d[[p]] = s_d; + let s_dr: f64 = (&col * d_r_ao.view()).sum_all(); + y_dot_dr[[p]] = s_dr; + // 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(); + } + + // ── per-atom RI-derivative terms ── + 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 shell_batch = calc_batch_size::(3 * nao * naux, mem_avail, None, None).max(1); + let aux_batch = calc_batch_size::(3 * nao_tp, mem_avail, None, None).max(1); + // MO-fold batches: each solved fold `[naux, 3, x_b, y]` stays below the batching budget + let xbatch_jk = calc_batch_size::(6 * naux * nocc, mem_avail, None, None).max(1); + let xbatch_ri = calc_batch_size::(6 * naux * nocc, mem_avail, None, None).max(1); + handle_memory_exceed(3.0 * (nao * naux) 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)); + for atm in 0..natm { + let [shl_a0, shl_a1, au0, au1] = mol_ao_slice[atm]; + let [shl_q0, shl_q1, auq0, auq1] = aux_slice[atm]; + let c0_atom = c0.i((au0..au1, ..)).into_contig(ColMajor); + let cv_atom = cv.i((au0..au1, ..)).into_contig(ColMajor); + let c1_atom = c1_dr.i((au0..au1, ..)).into_contig(ColMajor); + let l_inv_cols = l_inv.i((.., auq0..auq1)).into_contig(ColMajor); // [naux, nb_q] + + // forge L_1 generator of this atom, contracted with L_inv + let mut l1di = Vec::with_capacity(3); + for t in 0..3 { + let m0 = l_inv.i((.., auq0..auq1)) % tsr_int2c2e_ip1.i((auq0..auq1, .., t)); + let mut m: Tsr = m0 % l_inv.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); + l1di.push(l_inv.view() % l1_t.view()); + } + + // AO-shell batches of atom A: density vectors and MO folds of int3c2e_ip1 + let mut ydd_raw: Tsr = rt::zeros(([naux, 3], &device)); + let mut ydr_raw: Tsr = rt::zeros(([naux, 3], &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, nao, naux, 3] + let nub = ip1.shape()[0]; + for u in 0..nub { + let ip1_u = ip1.i((u, .., .., ..)); // [nao, naux, 3] + *&mut ydd_raw += &(&ip1_u * d_hf.i((au0 + u, ..)).i((.., None, None))).sum_axes(0); + *&mut ydr_raw += &(&ip1_u * d_r_ao.i((au0 + u, ..)).i((.., None, None))).sum_axes(0); + } + // u-slot (C0 on the atom rows) and v-slot (C1_dr / C0 swapped), t4 then RI term + t4_3c = add_fold_ip1(ip1.view(), c0_atom.view(), c1_dr.view(), xbatch_jk, &j2c_decomp, y_ip.view(), None, 2.0, t4_3c); + t4_3c = add_fold_ip1(ip1.view(), c1_atom.view(), c0.view(), xbatch_jk, &j2c_decomp, y_ip.view(), Some([1, 0]), 2.0, t4_3c); + ri_3c = add_fold_ip1(ip1.view(), c0_atom.view(), cv.view(), xbatch_ri, &j2c_decomp, g_vix.view(), Some([1, 0]), -4.0, ri_3c); + ri_3c = add_fold_ip1(ip1.view(), cv_atom.view(), c0.view(), xbatch_ri, &j2c_decomp, g_vix.view(), None, -4.0, ri_3c); + } + + // aux-shell batches of atom A: density vectors and MO folds of int3c2e_ip2 + let mut wdd: Tsr = rt::zeros(([auq1 - auq0, 3], &device)); + let mut wdr: Tsr = rt::zeros(([auq1 - auq0, 3], &device)); + 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 q_batches = blocksize_partition(&aux_loc[shl_q0..=shl_q1], aux_batch) + .into_iter() + .map(|[i0, i1]| [shl_q0 + i0, shl_q0 + i1]); + 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, 3] + *&mut wdd += &(&d_hf_tp.i((.., None, None)) * ip2.view()).sum_axes(0); + *&mut wdr += &(&d_r_ao_tp.i((.., None, None)) * ip2.view()).sum_axes(0); + let folded_jk = fold_ip2(ip2.view(), c0.view(), c1_dr.view(), 0..nmo); // [3, nb_q, nocc, nmo] + let solved_jk = metric_aux_atom(l_inv_cols.view(), folded_jk.view()); + t4_3c = add_contract(solved_jk.view(), y_ip.view(), None, 2.0, t4_3c); + let folded_ri = fold_ip2(ip2.view(), c0.view(), cv.view(), 0..nvir); // [3, nb_q, nocc, nvir] + let solved_ri = metric_aux_atom(l_inv_cols.view(), folded_ri.view()); + ri_3c = add_contract(solved_ri.view(), g_vix.view(), Some([1, 0]), -4.0, ri_3c); + } + let ydd2: Tsr = l_inv_cols.view() % wdd.view(); // [naux, 3] + let ydr2: Tsr = l_inv_cols.view() % wdr.view(); + let ydd_sol: Tsr = l_inv.view() % ydd_raw.view(); + let ydr_sol: Tsr = l_inv.view() % ydr_raw.view(); + + // forge L_1 contributions to the Y_1 vectors and to t4 / the RI term + let mut ydd_l1: Tsr = rt::zeros(([naux, 3], &device)); + let mut ydr_l1: Tsr = rt::zeros(([naux, 3], &device)); + let mut t4_l1 = [0.0f64; 3]; + let mut ri_l1 = [0.0f64; 3]; + for t in 0..3 { + ydd_l1.i_mut((.., t)).assign(&(l1di[t].view() % y_dot_d.view())); + ydr_l1.i_mut((.., t)).assign(&(l1di[t].view() % y_dot_dr.view())); + t4_l1[t] = 2.0 * (&l1di[t] * r_jk.view()).sum_all() as f64; + ri_l1[t] = -4.0 * (&l1di[t] * r_ri.view()).sum_all() as f64; + } + + // assemble forge t1 + t2 + t4 (D_r-carrying) and the RI term + let mut grad = [0.0f64; 3]; + let mut grad_ri = [0.0f64; 3]; + for t in 0..3 { + let v_ydd = ydd_sol.i((.., t)); + let v_ydr = ydr_sol.i((.., t)); + let c_ydd = ydd2.i((.., t)) + ydd_l1.i((.., t)); + let c_ydr = 2.0 * &v_ydr + &ydr2.i((.., t)) + &ydr_l1.i((.., t)); + let s1: f64 = (&y_dot_d * &c_ydr).sum_all(); + let s2: f64 = (&y_dot_dr * &(2.0 * &v_ydd + &c_ydd)).sum_all(); + grad[t] = -s1 - s2 + t4_3c[t] + t4_l1[t]; + grad_ri[t] = ri_3c[t] + ri_l1[t]; + } + 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))); + } + + // ── assemble the parts ── + let de = de_h.view() + de_s1.view() + de_jk.view() + de_rint.view(); + for (key, val) in [("de_h", de_h), ("de_s1", de_s1), ("de_jk", de_jk), ("de_rint", de_rint), ("de", de.into_owned())] { + self.result.insert(key.to_string(), to_matrix_full(val)); + } + self + } +} + +/// `[3, natm]` device tensor 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() +} + +/// Contract the atom-row axis `u` of a raw `ip1` batch `[nub, nao, naux, 3]` with `w_row` +/// (atom-A rows of an MO matrix, giving index `x`) and the full AO axis with `w_full` (giving +/// index `y`), left-solve against the CD metric, and accumulate +/// `acc[t] += factor `. +/// `swap` transposes the (x, y) axes of the partner when it is stored as `[y, x, naux]`. +#[allow(clippy::too_many_arguments)] +fn add_fold_ip1( + ip1: TsrView, + w_row: TsrView, + w_full: TsrView, + xbatch: usize, + j2c_decomp: &J2CDecompose, + partner: TsrView, + swap: Option<[usize; 2]>, + factor: f64, + mut acc: [f64; 3], +) -> [f64; 3] { + let device = ip1.device().clone(); + let nub = ip1.shape()[0]; + let naov = ip1.shape()[1]; + let naux = ip1.shape()[2]; + let ncomp = ip1.shape()[3]; + let y = w_full.shape()[1]; + let x = w_row.shape()[1]; + let mut stacked: Tsr = rt::zeros(([nub, y, naux * ncomp], &device)); + for u in 0..nub { + let ip1_u = ip1.i((u, .., .., ..)).into_contig(ColMajor); + let ip1_u = ip1_u.reshape([naov, naux * ncomp]); // (v, P t) + stacked.i_mut((u, .., ..)).assign(&(w_full.t() % ip1_u.view())); // [y, naux 3] + } + let xbatch = xbatch.min(x).max(1); + for x0 in (0..x).step_by(xbatch) { + let x1 = (x0 + xbatch).min(x); + let folded = w_row.i((.., x0..x1)).t() % stacked.view().reshape([nub, y * naux * ncomp]); // [x, (y P t)] + // [x, y, naux, 3] -> [naux, 3, x, y], then left-solve by L^-1 against the CD metric + let mut solved = folded.into_shape([x1 - x0, y, naux, ncomp]).transpose([2, 3, 0, 1]).into_contig(ColMajor); + { + let mut blk = rt::asarray((solved.raw_mut(), [naux, ncomp * (x1 - x0) * y].f(), &device)); + solve_by_j2c_mut(blk.view_mut(), j2c_decomp, FlagSide::L, true); + } + acc = add_contract(solved.view(), partner.view(), swap, factor, acc); + } + acc +} + +/// `acc[t] += factor sum_P `. +fn add_contract(folded: TsrView, partner: TsrView, swap: Option<[usize; 2]>, factor: f64, mut acc: [f64; 3]) -> [f64; 3] { + let naux = folded.shape()[0]; + let ncomp = folded.shape()[1]; + let partner = match swap { + None => partner, + Some([a, b]) => partner.transpose([a, b, 2]), + }; + for t in 0..ncomp { + let mut s = 0.0; + for p in 0..naux { + s += (&folded.i((p, t, .., ..)) * partner.i((.., .., p))).sum_all(); + } + acc[t] += factor * s; + } + acc +} + +/// Fold a raw `ip2` batch `[nao_tp, nb_q, 3]` column-by-column with the MO pair transform +/// `w_left^T col w_right[:, yr]`; returns `[3, nb_q, x, y_len]`. +fn fold_ip2(ip2: TsrView, w_left: TsrView, w_right: TsrView, yr: std::ops::Range) -> Tsr { + let nb_q = ip2.shape()[1]; + let ncomp = ip2.shape()[2]; + let x = w_left.shape()[1]; + let device = ip2.device().clone(); + let mut out = rt::zeros(([ncomp, nb_q, x, yr.len()], &device)); + for q in 0..nb_q { + for t in 0..ncomp { + let col = ip2.i((.., q, t)).unpack_tri(Upper, FlagSymm::Sy); + out.i_mut((t, q, .., ..)).assign(&(w_left.t() % &col % w_right.i((.., yr.clone())))); + } + } + out +} + +/// `sum_{Q in atom} L_inv[P, Q] folded[Q, t, x, y]`; `[3, nb_q, x, y] -> [naux, 3, x, y]`. +fn metric_aux_atom(l_inv_cols: TsrView, folded: TsrView) -> Tsr { + let shape = folded.shape(); + let [ncomp, nb_q, x, y] = [shape[0], shape[1], shape[2], shape[3]]; + let naux = l_inv_cols.shape()[0]; + let folded_t = folded.transpose([1, 0, 2, 3]).into_contig(ColMajor); // [nb_q, 3, x, y] + (l_inv_cols % folded_t.view().reshape([nb_q, ncomp * x * y])).into_shape([naux, ncomp, x, y]) +} + +impl GradAPI for RDHGradient<'_> { + fn get_gradient(&self) -> MatrixFull { + self.result["de"].clone() + } + fn get_energy(&self) -> f64 { + self.energy + } +} diff --git a/src/main_driver.rs b/src/main_driver.rs index a78bc57a0..f84490d7e 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -841,11 +841,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"); @@ -911,6 +914,13 @@ 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 CD-decomposed rimatr + if is_dh { + let mut grad_data_dh = crate::grad::rdh::RDHGradient::new(&scf_data, mpi_operator); + grad_data_dh.calc(); + 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; diff --git a/src/ri_pt2/pure_pt2_r_elecderiv.rs b/src/ri_pt2/pure_pt2_r_elecderiv.rs index 18542cb30..80b32bdb9 100644 --- a/src/ri_pt2/pure_pt2_r_elecderiv.rs +++ b/src/ri_pt2/pure_pt2_r_elecderiv.rs @@ -31,6 +31,9 @@ pub struct RPT2ElecDerivIncoreArg { /// 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 { @@ -46,6 +49,10 @@ pub struct RPT2ElecDerivIncoreOut { 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( @@ -64,8 +71,9 @@ 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, full_gfock } = 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(); let bi2_scale = O::from_f64(*c_ss).unwrap(); @@ -95,6 +103,7 @@ 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: Option>>> = if full_gfock { None @@ -211,6 +220,11 @@ 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 --- // @@ -285,5 +299,5 @@ where 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/rgfock_pt2.rs b/src/ri_pt2/rgfock_pt2.rs index 8c5644232..2d2ea7f56 100644 --- a/src/ri_pt2/rgfock_pt2.rs +++ b/src/ri_pt2/rgfock_pt2.rs @@ -80,6 +80,10 @@ where pub c_os: f64, /// Same-spin correlation factor $c_\mathrm{SS}$. pub c_ss: f64, + /// 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 when enabled. + 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"` @@ -116,6 +120,7 @@ where index_occ_outer_vec, c_os, c_ss, + dump_g_vix: false, result: HashMap::new(), e_corr: None, timing: Vec::new(), @@ -158,12 +163,15 @@ 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, full_gfock }; + 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(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())); } diff --git a/tests/analdrv/response/pt2_rgfock.rs b/tests/analdrv/response/pt2_rgfock.rs index d74f9eac5..f40bffbfe 100644 --- a/tests/analdrv/response/pt2_rgfock.rs +++ b/tests/analdrv/response/pt2_rgfock.rs @@ -105,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, full_gfock: false }; + 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; @@ -229,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, full_gfock: false }; + 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); diff --git a/tests/grad/mod.rs b/tests/grad/mod.rs index 2fd363ba7..d1270e70c 100644 --- a/tests/grad/mod.rs +++ b/tests/grad/mod.rs @@ -1,3 +1,4 @@ +mod rdh; 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 000000000..d8663f6c9 --- /dev/null +++ b/tests/grad/rdh.rs @@ -0,0 +1,87 @@ +//! 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). The rimatr must be CD-decomposed (see `src/grad/rdh.rs`). + +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 = "cd" + +[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 + """ +"##; + +#[test] +fn test_nh3_rmp2_grad() { + let keys = toml::from_str::(&INPUT_NH3[..]).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: {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 (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() < 1.0e-5); + + // PT2 correlation energy is carried along as the entry energy + assert!((grad.get_energy() - (-0.19179525779814563)).abs() < 2e-6, "PT2 correlation energy mismatch"); +} -- Gitee From d8676dd9e9984fb650fd6416152c072e8298a163 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 01:57:43 +0800 Subject: [PATCH 12/21] Restructure rdh.rs gradient data flow for efficiency Vitamin C/def2-TZVP force run 110 s -> 24 s (DH part 22 s, vs pyscf-forge 66 s), NH3 reference unchanged at 1.3e-8, results bit-identical. - hoist the once-per-run metric couplings (Y-bar = partner % solve_op) replacing per-(atom,term) dtrsm solves (60x FLOP cut on that stage); - ip2 route: one batched unpack + two batched half-transforms per aux shell batch, no [naux,3,x,y] strided tensors; - all reductions become contiguous last-axis rt::vecdot on layout-matched [(u Q), y] buffers; ~30 GEMMs per batch instead of thousands of small kernels; - j2c solve pieces come from get_j2c_decomp(ctrl) (convention-agnostic); only the L_1 metric-derivative generator still requires policy = "cd", with the panic re-scoped to that term; - persistent layout-specialized caches ~739 MB, peak RSS bounded by max_memory (identical gradient at 8000 MB). Co-authored-by: Claude Code Co-authored-by: glm-5.3 --- src/grad/rdh.rs | 403 ++++++++++++++++++++++++------------------------ 1 file changed, 203 insertions(+), 200 deletions(-) diff --git a/src/grad/rdh.rs b/src/grad/rdh.rs index abe0fc6ec..548e11118 100644 --- a/src/grad/rdh.rs +++ b/src/grad/rdh.rs @@ -12,9 +12,13 @@ //! (`C1 = C D_r^S`; the `0.5 C D_mo` exchange part belongs to the `"SCF"` entry); //! - `de_rint`: `4 ` (forge `_get_gradient_pt2`, third term). //! -//! The rimatr must be CD-decomposed (`[ctrl.j2c_decomp] policy = "cd"`), matching the -//! pyscf-forge (and pyscf `df`) convention `cderi[P, uv] = (L^-1 raw3c)[P, uv]`; every -//! derivative contraction left-solves against the same lower-Cholesky factor. +//! 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 (last-axis `rt::vecdot`) against `Ȳ`, never materializing +//! `[naux, 3, x, y]` intermediates. The metric-derivative (`L_1`) term, however, follows forge's +//! Cholesky-generator identity and therefore still requires the `Cd` convention (same as the +//! rimatr itself), so a matching `policy` is asserted below. use crate::analdrv::config::AnalDrvConfig; use crate::analdrv::response::rresp_interface::{rscf_resp_interface, RRespSCF}; @@ -22,9 +26,8 @@ use crate::analdrv::response::trait_rgfock::{GFockFlags, RGFockAPI}; use crate::analdrv::response::trait_rresp::RRespAPI; use crate::grad::traits::GradAPI; use crate::mpi_io::MPIOperator; -use crate::ri_jk::pure_decompose::solve_by_j2c_mut; use crate::ri_jk::util; -use crate::ri_jk::{get_j2c_decomp, J2CDecompOption, J2CDecompPolicy, J2CDecompose, J2C_THRESH}; +use crate::ri_jk::{get_j2c_decomp, J2CDecompose, J2C_THRESH}; use crate::ri_pt2::rgfock_pt2::RGFockPT2; use crate::scf_io::SCF; use crate::utilities::memory_batch::{blocksize_partition, calc_batch_size, detect_used_memory_mb, handle_memory_exceed}; @@ -56,11 +59,6 @@ impl<'a> RDHGradient<'a> { 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."); - if mol_obj.ctrl.j2c_decomp.policy != J2CDecompPolicy::Cd { - panic!( - "Doubly-hybrid analytic gradient requires the CD-decomposed rimatr. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input." - ); - } 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.", ); @@ -116,6 +114,7 @@ impl<'a> RDHGradient<'a> { let rdm1 = pt2.result["rdm1"].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); @@ -169,53 +168,55 @@ impl<'a> RDHGradient<'a> { } } - // ── CD metric objects ── - let j2c_decomp = get_j2c_decomp( - &aux, - &device, - J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(J2C_THRESH), ..Default::default() }, - ); - let j2c_l = match &j2c_decomp { - J2CDecompose::Cd { j2c_l, uplo: Upper, .. } => j2c_l.clone(), - _ => unreachable!("CD policy checked above"), - }; - // U^-1 solves U X = I (j2c = U^T U); forge's L_inv = inv(L_lower) = (U^-1)^T - let l_inv = { - 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) + // ── 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 l_lower = j2c_l.t().into_contig(ColMajor); 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 `generator_L_1` - 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; - } + // strictly-lower mask with 1/2 diagonal, of forge `generator_L_1` (Cd only) + let (l_lower, l_mask) = 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; + } + (l_lower, l_mask) + }, + _ => panic!( + "The DH analytic gradient's metric-derivative generator currently supports the CD convention only. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input (the derivative-integral solves themselves follow the rimatr convention)." + ), + }; - // ── persistent MO intermediates (Y_ip, metric couplings R/R2); peak 2 * nocc nao naux ── + // ── 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; handle_memory_exceed((nocc * nmo * naux + nvir * nocc * naux) as f64 * 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 y_dot_d: Tsr = rt::zeros(([naux], &device)); - let mut y_dot_dr: Tsr = rt::zeros(([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())); - let s_d: f64 = (&col * d_hf.view()).sum_all(); - y_dot_d[[p]] = s_d; - let s_dr: f64 = (&col * d_r_ao.view()).sum_all(); - y_dot_dr[[p]] = s_dr; // 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 @@ -228,44 +229,48 @@ impl<'a> RDHGradient<'a> { *&mut r_jk += &a_i.view() % y2t.view(); } - // ── per-atom RI-derivative terms ── + // ── pre-contractions `Ȳ = partner % S` (replace all in-loop metric solves) ── + let ybar_y2: Tsr = y_ip.view().into_shape([nocc * nmo, naux]) % s_op.view(); // [(i m), Q] + let ybar_g2i: Tsr = g_vix.view().transpose([1, 0, 2]).into_contig(ColMajor).view().into_shape([nocc * nvir, naux]) % s_op.view(); // [(i a), Q] + // layout copies `[x, Q, y]` for the ip1 shell-batch reductions + let ybar_jk_u = ybar_y2.view().into_shape([nocc, nmo, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [i, Q, m] + let ybar_jk_v = y_ip.view().transpose([1, 0, 2]).into_contig(ColMajor).view().into_shape([nmo * nocc, naux]) % s_op.view(); // [(m i), Q] + let ybar_jk_v = ybar_jk_v.view().into_shape([nmo, nocc, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [m, Q, i] + let ybar_ri_u = ybar_g2i.view().into_shape([nocc, nvir, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [i, Q, a] + let ybar_ri_v = g_vix.view().into_shape([nvir * nocc, naux]) % s_op.view(); // [(a i), Q] + let ybar_ri_v = ybar_ri_v.view().into_shape([nvir, nocc, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [a, Q, i] + drop(y_ip); + drop(g_vix); + + // ── 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 shell_batch = calc_batch_size::(3 * nao * naux, mem_avail, None, None).max(1); - let aux_batch = calc_batch_size::(3 * nao_tp, mem_avail, None, None).max(1); - // MO-fold batches: each solved fold `[naux, 3, x_b, y]` stays below the batching budget - let xbatch_jk = calc_batch_size::(6 * naux * nocc, mem_avail, None, None).max(1); - let xbatch_ri = calc_batch_size::(6 * naux * nocc, mem_avail, None, None).max(1); - handle_memory_exceed(3.0 * (nao * naux) as f64 * 8.0 / 1048576.0, mem_avail, abort); + 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)); - for atm in 0..natm { - let [shl_a0, shl_a1, au0, au1] = mol_ao_slice[atm]; - let [shl_q0, shl_q1, auq0, auq1] = aux_slice[atm]; - let c0_atom = c0.i((au0..au1, ..)).into_contig(ColMajor); - let cv_atom = cv.i((au0..au1, ..)).into_contig(ColMajor); - let c1_atom = c1_dr.i((au0..au1, ..)).into_contig(ColMajor); - let l_inv_cols = l_inv.i((.., auq0..auq1)).into_contig(ColMajor); // [naux, nb_q] - - // forge L_1 generator of this atom, contracted with L_inv - let mut l1di = Vec::with_capacity(3); - for t in 0..3 { - let m0 = l_inv.i((.., auq0..auq1)) % tsr_int2c2e_ip1.i((auq0..auq1, .., t)); - let mut m: Tsr = m0 % l_inv.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); - l1di.push(l_inv.view() % l1_t.view()); - } + 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); - // AO-shell batches of atom A: density vectors and MO folds of int3c2e_ip1 - let mut ydd_raw: Tsr = rt::zeros(([naux, 3], &device)); - let mut ydr_raw: Tsr = rt::zeros(([naux, 3], &device)); + // 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) @@ -276,72 +281,124 @@ impl<'a> RDHGradient<'a> { 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, nao, naux, 3] + }; // [nub, naov, naux, ncomp] let nub = ip1.shape()[0]; - for u in 0..nub { - let ip1_u = ip1.i((u, .., .., ..)); // [nao, naux, 3] - *&mut ydd_raw += &(&ip1_u * d_hf.i((au0 + u, ..)).i((.., None, None))).sum_axes(0); - *&mut ydr_raw += &(&ip1_u * d_r_ao.i((au0 + u, ..)).i((.., None, None))).sum_axes(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(ip1.view(), c0_atom.view(), c1_dr.view(), xbatch_jk, &j2c_decomp, y_ip.view(), None, 2.0, t4_3c); - t4_3c = add_fold_ip1(ip1.view(), c1_atom.view(), c0.view(), xbatch_jk, &j2c_decomp, y_ip.view(), Some([1, 0]), 2.0, t4_3c); - ri_3c = add_fold_ip1(ip1.view(), c0_atom.view(), cv.view(), xbatch_ri, &j2c_decomp, g_vix.view(), Some([1, 0]), -4.0, ri_3c); - ri_3c = add_fold_ip1(ip1.view(), cv_atom.view(), c0.view(), xbatch_ri, &j2c_decomp, g_vix.view(), None, -4.0, ri_3c); + t4_3c = add_fold_ip1_batched(ip1_t2.view(), nub, c0.i((r.clone(), ..)), c1_dr.view(), ybar_jk_u.view(), 2.0, t4_3c); + t4_3c = add_fold_ip1_batched(ip1_t2.view(), nub, c1_dr.i((r.clone(), ..)), c0.view(), ybar_jk_v.view(), 2.0, t4_3c); + ri_3c = add_fold_ip1_batched(ip1_t2.view(), nub, c0.i((r.clone(), ..)), cv.view(), ybar_ri_u.view(), -4.0, ri_3c); + ri_3c = add_fold_ip1_batched(ip1_t2.view(), nub, cv.i((r, ..)), c0.view(), ybar_ri_v.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]; + } + } - // aux-shell batches of atom A: density vectors and MO folds of int3c2e_ip2 - let mut wdd: Tsr = rt::zeros(([auq1 - auq0, 3], &device)); - let mut wdr: Tsr = rt::zeros(([auq1 - auq0, 3], &device)); - 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 q_batches = blocksize_partition(&aux_loc[shl_q0..=shl_q1], aux_batch) - .into_iter() - .map(|[i0, i1]| [shl_q0 + i0, shl_q0 + i1]); - 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, 3] - *&mut wdd += &(&d_hf_tp.i((.., None, None)) * ip2.view()).sum_axes(0); - *&mut wdr += &(&d_r_ao_tp.i((.., None, None)) * ip2.view()).sum_axes(0); - let folded_jk = fold_ip2(ip2.view(), c0.view(), c1_dr.view(), 0..nmo); // [3, nb_q, nocc, nmo] - let solved_jk = metric_aux_atom(l_inv_cols.view(), folded_jk.view()); - t4_3c = add_contract(solved_jk.view(), y_ip.view(), None, 2.0, t4_3c); - let folded_ri = fold_ip2(ip2.view(), c0.view(), cv.view(), 0..nvir); // [3, nb_q, nocc, nvir] - let solved_ri = metric_aux_atom(l_inv_cols.view(), folded_ri.view()); - ri_3c = add_contract(solved_ri.view(), g_vix.view(), Some([1, 0]), -4.0, ri_3c); + // 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 leading-(i m)/(i a) reshape pairs with the `Ȳ` rows + 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)] + let f_jk2 = f_jk.view().into_shape([nocc * nmo, nb_q * ncomp]); + let f_ri2 = f_ri.view().into_shape([nocc * nvir, nb_q * ncomp]); + for t in 0..ncomp { + // `r_jk_q[q] = sum_{i m} F_jk[(i m), (q t)] Ȳ[(i m), q]` — contracted over (i m) + let r_jk_q = rt::vecdot(&f_jk2.i((.., t * nb_q..(t + 1) * nb_q)), ybar_y2.i((.., qa..qb)), 0); + let r_ri_q = rt::vecdot(&f_ri2.i((.., t * nb_q..(t + 1) * nb_q)), ybar_g2i.i((.., qa..qb)), 0); + 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(); + } } - let ydd2: Tsr = l_inv_cols.view() % wdd.view(); // [naux, 3] - let ydr2: Tsr = l_inv_cols.view() % wdr.view(); - let ydd_sol: Tsr = l_inv.view() % ydd_raw.view(); - let ydr_sol: Tsr = l_inv.view() % ydr_raw.view(); + } - // forge L_1 contributions to the Y_1 vectors and to t4 / the RI term - let mut ydd_l1: Tsr = rt::zeros(([naux, 3], &device)); - let mut ydr_l1: Tsr = rt::zeros(([naux, 3], &device)); + // ── 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]; - for t in 0..3 { - ydd_l1.i_mut((.., t)).assign(&(l1di[t].view() % y_dot_d.view())); - ydr_l1.i_mut((.., t)).assign(&(l1di[t].view() % y_dot_dr.view())); - t4_l1[t] = 2.0 * (&l1di[t] * r_jk.view()).sum_all() as f64; - ri_l1[t] = -4.0 * (&l1di[t] * r_ri.view()).sum_all() as f64; + // forge L_1 generator of this atom, contracted with the solve operator + for t in 0..ncomp { + 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(); } // 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..3 { + for t in 0..ncomp { let v_ydd = ydd_sol.i((.., t)); let v_ydr = ydr_sol.i((.., t)); - let c_ydd = ydd2.i((.., t)) + ydd_l1.i((.., t)); - let c_ydr = 2.0 * &v_ydr + &ydr2.i((.., t)) + &ydr_l1.i((.., t)); - let s1: f64 = (&y_dot_d * &c_ydr).sum_all(); - let s2: f64 = (&y_dot_dr * &(2.0 * &v_ydd + &c_ydd)).sum_all(); - grad[t] = -s1 - s2 + t4_3c[t] + t4_l1[t]; - grad_ri[t] = ri_3c[t] + ri_l1[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))); @@ -362,95 +419,41 @@ fn to_matrix_full(tsr: Tsr) -> MatrixFull { MatrixFull::from_vec(shape, tsr.into_shape(-1).into_raw()).unwrap() } -/// Contract the atom-row axis `u` of a raw `ip1` batch `[nub, nao, naux, 3]` with `w_row` -/// (atom-A rows of an MO matrix, giving index `x`) and the full AO axis with `w_full` (giving -/// index `y`), left-solve against the CD metric, and accumulate -/// `acc[t] += factor `. -/// `swap` transposes the (x, y) axes of the partner when it is stored as `[y, x, naux]`. -#[allow(clippy::too_many_arguments)] -fn add_fold_ip1( - ip1: TsrView, - w_row: TsrView, - w_full: TsrView, - xbatch: usize, - j2c_decomp: &J2CDecompose, - partner: TsrView, - swap: Option<[usize; 2]>, +/// 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). +/// +/// 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.device().clone(); - let nub = ip1.shape()[0]; - let naov = ip1.shape()[1]; - let naux = ip1.shape()[2]; - let ncomp = ip1.shape()[3]; - let y = w_full.shape()[1]; - let x = w_row.shape()[1]; - let mut stacked: Tsr = rt::zeros(([nub, y, naux * ncomp], &device)); - for u in 0..nub { - let ip1_u = ip1.i((u, .., .., ..)).into_contig(ColMajor); - let ip1_u = ip1_u.reshape([naov, naux * ncomp]); // (v, P t) - stacked.i_mut((u, .., ..)).assign(&(w_full.t() % ip1_u.view())); // [y, naux 3] - } - let xbatch = xbatch.min(x).max(1); - for x0 in (0..x).step_by(xbatch) { - let x1 = (x0 + xbatch).min(x); - let folded = w_row.i((.., x0..x1)).t() % stacked.view().reshape([nub, y * naux * ncomp]); // [x, (y P t)] - // [x, y, naux, 3] -> [naux, 3, x, y], then left-solve by L^-1 against the CD metric - let mut solved = folded.into_shape([x1 - x0, y, naux, ncomp]).transpose([2, 3, 0, 1]).into_contig(ColMajor); - { - let mut blk = rt::asarray((solved.raw_mut(), [naux, ncomp * (x1 - x0) * y].f(), &device)); - solve_by_j2c_mut(blk.view_mut(), j2c_decomp, FlagSide::L, true); - } - acc = add_contract(solved.view(), partner.view(), swap, factor, acc); - } - acc -} - -/// `acc[t] += factor sum_P `. -fn add_contract(folded: TsrView, partner: TsrView, swap: Option<[usize; 2]>, factor: f64, mut acc: [f64; 3]) -> [f64; 3] { - let naux = folded.shape()[0]; - let ncomp = folded.shape()[1]; - let partner = match swap { - None => partner, - Some([a, b]) => partner.transpose([a, b, 2]), - }; - for t in 0..ncomp { - let mut s = 0.0; - for p in 0..naux { - s += (&folded.i((p, t, .., ..)) * partner.i((.., .., p))).sum_all(); - } + 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 } -/// Fold a raw `ip2` batch `[nao_tp, nb_q, 3]` column-by-column with the MO pair transform -/// `w_left^T col w_right[:, yr]`; returns `[3, nb_q, x, y_len]`. -fn fold_ip2(ip2: TsrView, w_left: TsrView, w_right: TsrView, yr: std::ops::Range) -> Tsr { - let nb_q = ip2.shape()[1]; - let ncomp = ip2.shape()[2]; - let x = w_left.shape()[1]; - let device = ip2.device().clone(); - let mut out = rt::zeros(([ncomp, nb_q, x, yr.len()], &device)); - for q in 0..nb_q { - for t in 0..ncomp { - let col = ip2.i((.., q, t)).unpack_tri(Upper, FlagSymm::Sy); - out.i_mut((t, q, .., ..)).assign(&(w_left.t() % &col % w_right.i((.., yr.clone())))); - } - } - out -} - -/// `sum_{Q in atom} L_inv[P, Q] folded[Q, t, x, y]`; `[3, nb_q, x, y] -> [naux, 3, x, y]`. -fn metric_aux_atom(l_inv_cols: TsrView, folded: TsrView) -> Tsr { - let shape = folded.shape(); - let [ncomp, nb_q, x, y] = [shape[0], shape[1], shape[2], shape[3]]; - let naux = l_inv_cols.shape()[0]; - let folded_t = folded.transpose([1, 0, 2, 3]).into_contig(ColMajor); // [nb_q, 3, x, y] - (l_inv_cols % folded_t.view().reshape([nb_q, ncomp * x * y])).into_shape([naux, ncomp, x, y]) -} - impl GradAPI for RDHGradient<'_> { fn get_gradient(&self) -> MatrixFull { self.result["de"].clone() -- Gitee From 07754f44f45f55ffa7b56e063ddc747ccef61826 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 05:43:37 +0800 Subject: [PATCH 13/21] Add XYG3-family gradient terms to rdh RDHGradient is now the full doubly-hybrid (fifth-DFA) gradient; new terms follow the M1 data flow: - de_k_dh: Delta-exchange on the SCF density with factor_k = dfa_hybrid_pos - dfa_hybrid_scf, including the kaux part; - de_xc_n: XC skeleton difference (final - SCF functional) on the SCF density, on regenerated grids; - de_resp: fxc-response contraction of the relaxed correlation density with the SCF-functional AO-derivative potential (hessian vmat_deriv1 kernels, 4xfxc convention via the response object); - de_ovlp_dh: S1 contraction with the occ-occ block of the final-functional Fock built on the SCF density; - the Z-vector now solves with the DH-combined Lagrangian (L_pt2 + 4 CvT F_n Co), which is required for a correct D_r whenever the final functional differs from the SCF one; - generalize restricted_biorthogonalize to (c_os + c_ss) t - c_ss t^T (was 2 c_os t - c_ss t^T, valid only for c_os = c_ss; bit-identical for MP2/XYG3, fixes unequal-coefficient DHs); - grid regeneration for the DH gradient in eval_force, gated off for the pure-MP2 family; get_vxc_rayon_new made pub for reuse. NH3/def2-SVP: XYG3 3.3e-7 vs pyscf-forge (test asserts 1e-6), XYGJ-OS 4.4e-7, XYG7 2.6e-7; mp2 path unchanged. Vitamin C XYG3 force run 53.2 s (16 threads). Co-authored-by: Claude Code Co-authored-by: glm-5.3 --- src/grad/rdh.rs | 353 +++++++++++++++++++++++++++-- src/grad/rks.rs | 9 +- src/main_driver.rs | 9 + src/ri_pt2/pure_pt2_r_elecderiv.rs | 3 +- tests/grad/rdh.rs | 83 +++++++ 5 files changed, 438 insertions(+), 19 deletions(-) diff --git a/src/grad/rdh.rs b/src/grad/rdh.rs index 548e11118..63910d365 100644 --- a/src/grad/rdh.rs +++ b/src/grad/rdh.rs @@ -1,16 +1,35 @@ -//! Restricted doubly-hybrid (`xc = "mp2"`) analytic gradient: the PT2-correlation part. +//! Restricted doubly-hybrid (XYG3-family and pure-MP2) analytic gradient. //! -//! The SCF (HF) part of the gradient is provided by the existing `"SCF"` entry of +//! 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, following the D_r-dependent blocks of pyscf-forge `dh/grad/dfdh.py`: +//! 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 = C D_r^S`; the `0.5 C D_mo` exchange part belongs to the `"SCF"` entry); -//! - `de_rint`: `4 ` (forge `_get_gradient_pt2`, third term). +//! (`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 @@ -21,11 +40,20 @@ //! rimatr itself), so a matching `policy` is asserted below. use crate::analdrv::config::AnalDrvConfig; -use crate::analdrv::response::rresp_interface::{rscf_resp_interface, RRespSCF}; +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::{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; @@ -92,10 +120,63 @@ impl<'a> RDHGradient<'a> { .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_gga(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 + }; let j3c = rimatr.0.to_rstsr_view(&device); let mut pt2 = RGFockPT2::::new( mo_coeff.clone(), - mo_occ, + mo_occ.clone(), mo_energy.clone(), j3c.into_cow(), None, @@ -108,17 +189,34 @@ impl<'a> RDHGradient<'a> { // 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 _ = pt2.make_rdm1_resp(&mut resp_objs); - let rdm1_resp = pt2.result["rdm1_resp"].to_owned(); - let gfock_part = pt2.result["gfock_part_full"].to_owned(); 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); - let c1_dr = (mo_coeff.view() % d_r_s.view()).into_contig(ColMajor); + // 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); // SCF response upon the *relaxed* correlation density (forge W_III); the cached `axd` @@ -149,12 +247,12 @@ impl<'a> RDHGradient<'a> { // ── 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); - let tsr_ipovlp = { - let (out, shape) = mol.integrate("int1e_ipovlp", "s1", None).into(); - rt::asarray((out, shape, &device)) - }; // 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); @@ -228,6 +326,7 @@ impl<'a> RDHGradient<'a> { 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; // ── pre-contractions `Ȳ = partner % S` (replace all in-loop metric solves) ── let ybar_y2: Tsr = y_ip.view().into_shape([nocc * nmo, naux]) % s_op.view(); // [(i m), Q] @@ -404,15 +503,169 @@ impl<'a> RDHGradient<'a> { 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 = 2.0 * (&(&dao_xc_n - &dao_xc_scf) * &d_hf.i((.., .., None))).sum_axes(1); + 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 = + 2.0 * (&tsr_ipovlp * &(&(&w_n * 2.0) - &dme0).i((.., .., None))).sum_axes(1); + 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(); - for (key, val) in [("de_h", de_h), ("de_s1", de_s1), ("de_jk", de_jk), ("de_rint", de_rint), ("de", de.into_owned())] { + 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; + 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 +} + /// `[3, natm]` device tensor 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(); @@ -462,3 +715,69 @@ impl GradAPI for RDHGradient<'_> { self.energy } } + +/// Whether the SCF-iteration functional carries a GGA-or-higher XC part (the gate of the +/// response-density XC term; LDA-type functionals have no grid-gradient contribution here, +/// following pyscf-forge's `_xc_type(xc) == "GGA"` gate). +fn has_scf_xc_gga(scf_data: &SCF) -> bool { + let xc_data = &scf_data.mol.xc_data; + if xc_data.dfa_compnt_scf.is_empty() { + return false; + } + let func_list = scf_xc_func_list(scf_data); + let refs: Vec<&libxc::functional::LibXCFunctional> = func_list.iter().map(|(_, f)| f).collect(); + let xc_type = determine_den_type_from_list(&refs); + matches!(xc_type, XCDenType::SIGMA | XCDenType::TAU) +} + +/// 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 df5b45a7d..342f8a832 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/main_driver.rs b/src/main_driver.rs index f84490d7e..5c30e79ce 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -888,6 +888,15 @@ 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![]; + // post-SCF (doubly-hybrid) gradients with a final-functional XC part need the DFT grids + // again, but `xdh_calculations` frees them for memory; regenerate before the SCF gradient + // borrows `scf_data` (the pure-MP2 family has no XC part and needs no grid) + let is_dh = matches!(scf_data.mol.xc_data.dfa_family_pos, Some(crate::dft::DFAFamily::PT2)); + let has_dh_xc = scf_data.mol.xc_data.dfa_compnt_pos.as_ref().map_or(false, |v| !v.is_empty()); + if is_dh && has_dh_xc && 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 { diff --git a/src/ri_pt2/pure_pt2_r_elecderiv.rs b/src/ri_pt2/pure_pt2_r_elecderiv.rs index 80b32bdb9..5510ca6cb 100644 --- a/src/ri_pt2/pure_pt2_r_elecderiv.rs +++ b/src/ri_pt2/pure_pt2_r_elecderiv.rs @@ -75,7 +75,8 @@ where 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(); diff --git a/tests/grad/rdh.rs b/tests/grad/rdh.rs index d8663f6c9..c10b0e868 100644 --- a/tests/grad/rdh.rs +++ b/tests/grad/rdh.rs @@ -44,6 +44,35 @@ new_driver = true """ "##; +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 = "cd" + +[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 + """ +"##; + #[test] fn test_nh3_rmp2_grad() { let keys = toml::from_str::(&INPUT_NH3[..]).unwrap(); @@ -85,3 +114,57 @@ fn test_nh3_rmp2_grad() { // PT2 correlation energy is carried along as the entry energy assert!((grad.get_energy() - (-0.19179525779814563)).abs() < 2e-6, "PT2 correlation energy mismatch"); } + +/// 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). +#[test] +fn test_nh3_rxyg3_grad() { + let keys = toml::from_str::(&INPUT_XYG3[..]).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: {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 (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() < 1.0e-6, "XYG3 total gradient mismatch"); +} -- Gitee From 39d2fd92d1aec0710bb15ec10f401a523588d548 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 09:18:23 +0800 Subject: [PATCH 14/21] Add unrestricted RI-MP2 analytic gradient New src/grad/udh.rs: UDHGradient computes the PT2 correlation gradient for UHF references (xc = "mp2"; the eval_force "DH" entry now dispatches on spin_polarization), mirroring the restricted data flow: - per-spin terms de_h/de_s1/de_jk/de_rint over the shared rimatr, closed-shell UHF verified against the restricted path (de_rint to 1.7e-10, G^a + G^b = 4 G_R to 2e-9); - U-side PT2 electric derivative src/ri_pt2/pure_pt2_u_elecderiv.rs: per-spin amplitudes, T^ss = c_ss (t - t^T)/2 on the virtual pair, T^ab = c_os t, forge-normalized G, full four-block gfock with the 2 eps rdm1 terms; - U Z-vector solved through the coupled URespSCF (no alpha-beta exchange kernel); URespAPI get_fock_* stay stubs (F_n = F_HF collapse for MP2). NH3+ doublet/def2-SVP: total energy 1.9e-9, "DH" entry 4.9e-7 vs pyscf-forge (the reference's own FD-vs-analytic floor is 4.96e-7), in-crate FD 1.05e-6. Restricted path untouched (rdh.rs zero diff). Vitamin C UMP2 force run 44.8 s, 1.06x restricted peak RSS. Co-authored-by: Claude Code Co-authored-by: glm-5.3 --- src/grad/mod.rs | 1 + src/grad/udh.rs | 720 +++++++++++++++++++++++++++++ src/main_driver.rs | 12 +- src/ri_pt2/mod.rs | 1 + src/ri_pt2/pure_pt2_u_elecderiv.rs | 275 +++++++++++ tests/grad/mod.rs | 1 + tests/grad/udh.rs | 108 +++++ 7 files changed, 1115 insertions(+), 3 deletions(-) create mode 100644 src/grad/udh.rs create mode 100644 src/ri_pt2/pure_pt2_u_elecderiv.rs create mode 100644 tests/grad/udh.rs diff --git a/src/grad/mod.rs b/src/grad/mod.rs index 920c0add6..eb6a75b77 100644 --- a/src/grad/mod.rs +++ b/src/grad/mod.rs @@ -1,4 +1,5 @@ pub mod rdh; +pub mod udh; pub mod uhf; pub mod rks; pub mod uks; diff --git a/src/grad/udh.rs b/src/grad/udh.rs new file mode 100644 index 000000000..342a377ad --- /dev/null +++ b/src/grad/udh.rs @@ -0,0 +1,720 @@ +//! 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 Lagrangian +//! $L^\sigma = \mathscr{F}^{\sigma}_{vo} - \mathscr{F}^{\sigma}_{ov\top} + A^\sigma(D^-)_{vo}$. +//! The final-functional (XYG3-type) terms of the restricted module are not implemented yet; the +//! gates below panic until the next milestone, keeping the same gating structure. + +use crate::analdrv::config::AnalDrvConfig; +use crate::analdrv::response::trait_uresp::URespAPI; +use crate::analdrv::response::uresp_interface::{uscf_resp_interface, URespSCF}; +use crate::grad::rhf::{generator_deriv_hcore, pack_triu_tilde}; +use crate::grad::traits::GradAPI; +use crate::mpi_io::MPIOperator; +use crate::ri_jk::util; +use crate::ri_jk::{get_j2c_decomp, J2CDecompose}; +use crate::ri_pt2::pure_pt2_u_elecderiv::{ + get_rupt2_elec_deriv_incore, UPT2ElecDerivIncoreArg, UPT2ElecDerivIncoreInp, +}; +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, +} + +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 } + } + + 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."); + 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; not implemented on the + // unrestricted side yet (next milestone), so the pure-MP2 family is the supported scope + 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 has_final_functional = delta_hyb.abs() > 1.0e-10 + || xc_data.dfa_compnt_pos.as_ref().map_or(false, |v| !v.is_empty()); + if has_final_functional { + panic!("The unrestricted doubly-hybrid gradient currently supports the pure-MP2 family only (the final-functional terms are a separate milestone)."); + } + + // ── 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), + ]; + + // ── response objects, PT2 electronic derivative, relaxed correlation density ── // + let config = AnalDrvConfig::default(); + let mut resp_objs: URespSCF = uscf_resp_interface(scf_data, &config); + 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()], + }, + &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 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); + 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), + ]; + 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)) + }; + let (l_lower, l_mask) = 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; + } + (l_lower, l_mask) + }, + _ => panic!( + "The DH analytic gradient's metric-derivative generator currently supports the CD convention only. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input (the derivative-integral solves themselves follow the rimatr convention)." + ), + }; + + // ── persistent MO intermediates; density dots of the rimatr by two GEMVs per spin ── // + 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; + handle_memory_exceed( + ((nocc[0] + nocc[1]) * nmo * naux + (nvir[0] * nocc[0] + nvir[1] * nocc[1]) * naux) as f64 + * 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; + + // ── pre-contractions `Ȳ = partner % S` (replace all in-loop metric solves), per spin ── // + let mut ybar_y2: [Option; 2] = [None, None]; + let mut ybar_g2i: [Option; 2] = [None, None]; + let mut ybar_jk_u: [Option; 2] = [None, None]; + let mut ybar_jk_v: [Option; 2] = [None, None]; + let mut ybar_ri_u: [Option; 2] = [None, None]; + let mut ybar_ri_v: [Option; 2] = [None, None]; + for s in 0..2 { + if nocc[s] == 0 { + continue; + } + let y2 = y_ip[s].view().into_shape([nocc[s] * nmo, naux]) % s_op.view(); + let g2i = g_vix[s] + .view() + .transpose([1, 0, 2]) + .into_contig(ColMajor) + .view() + .into_shape([nocc[s] * nvir[s], naux]) + % s_op.view(); + ybar_y2[s] = Some(y2); + ybar_g2i[s] = Some(g2i); + ybar_jk_u[s] = Some( + ybar_y2[s] + .as_ref() + .unwrap() + .view() + .into_shape([nocc[s], nmo, naux]) + .transpose([0, 2, 1]) + .into_contig(ColMajor), + ); + let yv = y_ip[s] + .view() + .transpose([1, 0, 2]) + .into_contig(ColMajor) + .view() + .into_shape([nmo * nocc[s], naux]) + % s_op.view(); + ybar_jk_v[s] = Some( + yv.view() + .into_shape([nmo, nocc[s], naux]) + .transpose([0, 2, 1]) + .into_contig(ColMajor), + ); + ybar_ri_u[s] = Some( + ybar_g2i[s] + .as_ref() + .unwrap() + .view() + .into_shape([nocc[s], nvir[s], naux]) + .transpose([0, 2, 1]) + .into_contig(ColMajor), + ); + let gv = g_vix[s].view().into_shape([nvir[s] * nocc[s], naux]) % s_op.view(); + ybar_ri_v[s] = Some( + gv.view() + .into_shape([nvir[s], nocc[s], naux]) + .transpose([0, 2, 1]) + .into_contig(ColMajor), + ); + } + drop(y_ip); + drop(g_vix); + + // ── 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_u[s].as_ref().unwrap().view(), + 2.0, + t4_3c, + ); + t4_3c = add_fold_ip1_batched( + ip1_t2.view(), + nub, + c1_dr[s].i((r.clone(), ..)), + c0[s].view(), + ybar_jk_v[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_u[s].as_ref().unwrap().view(), + -1.0, + ri_3c, + ); + ri_3c = add_fold_ip1_batched( + ip1_t2.view(), + nub, + cv[s].i((r.clone(), ..)), + c0[s].view(), + ybar_ri_v[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); + let f_jk2 = f_jk.view().into_shape([nocc[s] * nmo, nb_q * ncomp]); + let f_ri2 = f_ri.view().into_shape([nocc[s] * nvir[s], nb_q * ncomp]); + for t in 0..ncomp { + let r_jk_q = + rt::vecdot(&f_jk2.i((.., t * nb_q..(t + 1) * nb_q)), ybar_y2[s].as_ref().unwrap().view().i((.., qa..qb)), 0); + let r_ri_q = + rt::vecdot(&f_ri2.i((.., t * nb_q..(t + 1) * nb_q)), ybar_g2i[s].as_ref().unwrap().view().i((.., qa..qb)), 0); + 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]; + for t in 0..ncomp { + 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(); + } + + // 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))); + } + + // ── assemble the parts ── // + let de = de_h.view() + de_s1.view() + de_jk.view() + de_rint.view(); + for (key, val) in [ + ("de_h", de_h), + ("de_s1", de_s1), + ("de_jk", de_jk), + ("de_rint", de_rint), + ("de", de.into_owned()), + ] { + self.result.insert(key.to_string(), to_matrix_full(val)); + } + self + } +} + +/// `[3, natm]` device tensor 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() +} + +/// 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 +} + +impl GradAPI for UDHGradient<'_> { + fn get_gradient(&self) -> MatrixFull { + self.result["de"].clone() + } + fn get_energy(&self) -> f64 { + self.energy + } +} diff --git a/src/main_driver.rs b/src/main_driver.rs index 5c30e79ce..cdd171e72 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -925,9 +925,15 @@ fn eval_force(scf_data: &mut SCF, time_mark: &mut utilities::TimeRecords, mpi_op // 1.5 doubly-hybrid (RI-PT2) correlation gradient; requires the CD-decomposed rimatr if is_dh { - let mut grad_data_dh = crate::grad::rdh::RDHGradient::new(&scf_data, mpi_operator); - grad_data_dh.calc(); - grad_data_list.push(("DH".into(), Box::new(grad_data_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(); + 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(); + grad_data_list.push(("DH".into(), Box::new(grad_data_dh))); + } } // 2. dftd gradient data diff --git a/src/ri_pt2/mod.rs b/src/ri_pt2/mod.rs index 6e893a9fc..a5f82584d 100644 --- a/src/ri_pt2/mod.rs +++ b/src/ri_pt2/mod.rs @@ -45,6 +45,7 @@ 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; #[derive(Clone)] 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 000000000..522199902 --- /dev/null +++ b/src/ri_pt2/pure_pt2_u_elecderiv.rs @@ -0,0 +1,275 @@ +//! 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. +//! +//! Same-spin blocks are evaluated with the pair-symmetry structure of the restricted kernel +//! (single occupied batch, triangular pair loop with the transpose fill). The $\alpha\beta$ +//! block stores only the plain amplitudes ($T = c_{os} t$ enters every contraction as a factor) +//! in a single batch; its four rdm1 contractions are per-index-partial-trace GEMM loops over +//! the contracted pair, so no $[\cdot^2]$-sized intermediate is formed. + +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], +} + +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 } = 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]]; + + 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, rdm1 by 4-d reshapes, G by one GEMM ── // + for s in 0..2 { + if nocc[s] == 0 || nvir[s] == 0 { + 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); + let mut t_vivo: Tsr = rt::zeros(([nv, no, nv, no].f(), &device)); + let mut T_vivo: Tsr = rt::zeros(([nv, no, nv, no].f(), &device)); + let pair_ij: Vec<[usize; 2]> = (0..no).flat_map(|i| (0..=i).map(move |j| [i, j])).collect(); + 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, .., j)).assign(&t_vv); + t_vivo_T.i_mut((.., i, .., j)).assign(&t_vv_); + if j < i { + t_vivo.i_mut((.., j, .., i)).assign(t_vv.t()); + t_vivo_T.i_mut((.., j, .., i)).assign(t_vv_.t()); + } + }); + // D[oo] -= 2 sum T t; D[vv] += 2 sum T t (both sides of the pair matrix summed) + 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 g_part = + T_vivo.view().reshape([nv * no, nv * no]) % cderi_vox[s].view().reshape([nv * no, naux]); + g_vix[s] = (g_part * 4.0 + &g_vix[s].view().into_shape([nv * no, naux])) + .into_shape([nv, no, naux]); + } + + // ── opposite-spin block: all (i_a, k_b) pairs; T = c_os t, so only t is stored ── // + 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::(); + let mut t_x: Tsr = rt::zeros(([nva, 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((.., i, ..)) % cderi_vox[1].i((.., k, ..)).t(); + let d_ab = &d_vv_ab + 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); + }); + let t2 = t_x.view().into_shape([nva * noa, nvb * nob]); // [(a i), (b k)] + // G contributions: rows of T2 = c_os t2 are (a i), columns (b k) + let g_a = (t2.view() % cderi_vox[1].view().into_shape([nvb * nob, naux])) * (2.0 * c_os); + *&mut g_vix[0] += &g_a.into_shape([nva, noa, naux]); + let g_b = (t2.view().t() % cderi_vox[0].view().into_shape([nva * noa, naux])) * (2.0 * c_os); + *&mut g_vix[1] += &g_b.into_shape([nvb, nob, naux]); + // correlation rdm1 pieces; each contracts one amplitude index pair completely (the `oo` + // parts) or takes the partial trace over the paired occupied index (the `vv` parts) + let t3 = t_x.view().into_shape([nva * noa * nvb, nob]); // [(a i b), k] + let t4 = t_x.view().into_shape([nva * noa, nvb, nob]); // [(a i), b, k] + let t4b = t_x.view().into_shape([nva, noa, nvb * nob]); // [a, (i), (k b)]-soluble per spin + let m_a_oo = { + let mut m: Tsr = rt::zeros(([noa, noa].f(), &device)); + let mut ta: Tsr = rt::zeros(([noa, nvb * nob].f(), &device)); + for a in 0..nva { + ta.assign(&t_x.i((a, .., .., ..)).into_shape([noa, nvb * nob])); + *&mut m += &(ta.view() % ta.view().t()); + } + m + }; + *&mut rdm1[0].i_mut((0..noa, 0..noa)) -= &(m_a_oo * c_os); + let m_b_oo = t3.view().t() % t3.view(); + *&mut rdm1[1].i_mut((0..nob, 0..nob)) -= &(m_b_oo * c_os); + let m_a_vv = { + let mut m: Tsr = rt::zeros(([nva, nva].f(), &device)); + for i in 0..noa { + let ti = t4b.i((.., i, ..)); + *&mut m += &(ti.view() % ti.view().t()); + } + m + }; + *&mut rdm1[0].i_mut((noa..nmo[0], noa..nmo[0])) += &(m_a_vv * c_os); + let m_b_vv = { + let mut m: Tsr = rt::zeros(([nvb, nvb].f(), &device)); + for k in 0..nob { + let tk = t4.i((.., .., k)); + *&mut m += &(tk.view().t() % tk.view()); + } + m + }; + *&mut rdm1[1].i_mut((nob..nmo[1], nob..nmo[1])) += &(m_b_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/tests/grad/mod.rs b/tests/grad/mod.rs index d1270e70c..de5fd4453 100644 --- a/tests/grad/mod.rs +++ b/tests/grad/mod.rs @@ -1,4 +1,5 @@ mod rdh; +mod udh; mod ri_jk_rhf; mod ri_jk_rks; mod ri_jk_uhf; diff --git a/tests/grad/udh.rs b/tests/grad/udh.rs new file mode 100644 index 000000000..05e18bcd2 --- /dev/null +++ b/tests/grad/udh.rs @@ -0,0 +1,108 @@ +//! Doubly-hybrid (`xc = "mp2"`) analytic gradient of distorted NH3+ (unrestricted RI-MP2). +//! +//! Reference: pyscf-forge `UDFDH(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 UHF part). The rimatr must be CD-decomposed (see `src/grad/udh.rs`). + +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 = "cd" + +[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 + """ +"##; + +#[test] +fn test_nh3_ump2_grad() { + let keys = toml::from_str::(&INPUT_NH3[..]).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: {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 (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() < 1.0e-5); + + // 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 (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() < 1.0e-5, "UMP2 total gradient mismatch"); +} -- Gitee From 070166b44faf653408610b84349ecc4fda40fe05 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 11:53:26 +0800 Subject: [PATCH 15/21] Add unrestricted XYG3-family gradient terms and stream amplitudes UDHGradient becomes the full fifth-DFA unrestricted gradient; new U terms mirror the restricted accounting: de_k_dh (delta hybrid exchange, kaux included), de_xc_n (per-spin skeleton difference), de_resp (UKS AO-derivative potential response of the correlation density), de_ovlp_dh (per-spin S1 contraction with the final-functional Fock), Z-vector Lagrangian including 2 CvT F_n Co per spin. Enforce the amplitude-storage prohibition on both paths: no nocc^2 nvir^2 tensor may exist, transient or persistent. The U sigma- sigma and alpha-beta blocks stream over occupied windows (alpha-beta needs two passes because each rdm1 Gram keeps its own index pair complete); the restricted driver and the DH gfock interface now pass occupancy-batched index vectors to the existing batched kernel. Multi-window correctness is pinned by single vs irregular-window agreement to 5.6e-17 (temporary test, removed). NH3+: UXYG3 4.7e-7 vs pyscf-forge, UXYGJ-OS 5.9e-7, UXYG7 4.4e-7, in-crate FD 6.4e-6. Tests: test_grad 16/16, test_analdrv 36/36, restricted numbers unchanged. Vitamin C bounded-memory runs now complete down to 2000 MB with gradients identical to 6e-9; batching costs +3-5% wall time. Co-authored-by: Claude Code Co-authored-by: deepseek-v4-flash --- src/analdrv/response/rgfock_interface.rs | 20 +- src/analdrv/response/uresp_interface.rs | 22 ++ src/grad/rdh.rs | 15 +- src/grad/udh.rs | 399 ++++++++++++++++++++++- src/grad/uks.rs | 9 +- src/ri_pt2/mod.rs | 27 ++ src/ri_pt2/pure_pt2_u_elecderiv.rs | 246 ++++++++------ tests/grad/udh.rs | 91 +++++- 8 files changed, 713 insertions(+), 116 deletions(-) diff --git a/src/analdrv/response/rgfock_interface.rs b/src/analdrv/response/rgfock_interface.rs index 6144a2930..d650623a4 100644 --- a/src/analdrv/response/rgfock_interface.rs +++ b/src/analdrv/response/rgfock_interface.rs @@ -457,6 +457,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 (never materializing + // `nocc^2 nvir^2` tensors): size the windows so the per-window transients stay within a few + // times the `nvir * nocc * naux` 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 @@ -467,7 +483,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, ))), @@ -479,7 +495,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 440016fab..0c5d07005 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/rdh.rs b/src/grad/rdh.rs index 63910d365..b569c24d8 100644 --- a/src/grad/rdh.rs +++ b/src/grad/rdh.rs @@ -57,6 +57,7 @@ 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}; @@ -173,6 +174,18 @@ impl<'a> RDHGradient<'a> { } else { None }; + // the PT2 kernel streams its amplitudes over occupied windows (never materializing + // `nocc^2 nvir^2` tensors): 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 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(), @@ -180,7 +193,7 @@ impl<'a> RDHGradient<'a> { mo_energy.clone(), j3c.into_cow(), None, - vec![0, nocc], + index_occ_outer, c_os, c_ss, ); diff --git a/src/grad/udh.rs b/src/grad/udh.rs index 342a377ad..f21717416 100644 --- a/src/grad/udh.rs +++ b/src/grad/udh.rs @@ -17,22 +17,43 @@ //! //! 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 Lagrangian -//! $L^\sigma = \mathscr{F}^{\sigma}_{vo} - \mathscr{F}^{\sigma}_{ov\top} + A^\sigma(D^-)_{vo}$. -//! The final-functional (XYG3-type) terms of the restricted module are not implemented yet; the -//! gates below panic until the next milestone, keeping the same gating structure. +//! $\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::config::AnalDrvConfig; +use crate::analdrv::response::rgfock_interface::dh_jk_factors; use crate::analdrv::response::trait_uresp::URespAPI; -use crate::analdrv::response::uresp_interface::{uscf_resp_interface, URespSCF}; -use crate::grad::rhf::{generator_deriv_hcore, pack_triu_tilde}; +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, 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}; 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, @@ -84,15 +105,39 @@ impl<'a> UDHGradient<'a> { .try_into() .expect("dfa_paramr_adv must have exactly two entries."); - // gates of the final-functional (doubly-hybrid) terms; not implemented on the - // unrestricted side yet (next milestone), so the pure-MP2 family is the supported scope + // 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 has_final_functional = delta_hyb.abs() > 1.0e-10 - || xc_data.dfa_compnt_pos.as_ref().map_or(false, |v| !v.is_empty()); - if has_final_functional { - panic!("The unrestricted doubly-hybrid gradient currently supports the pure-MP2 family only (the final-functional terms are a separate milestone)."); - } + 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_gga(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] = @@ -134,9 +179,77 @@ impl<'a> UDHGradient<'a> { (&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 = AnalDrvConfig::default(); let mut resp_objs: URespSCF = uscf_resp_interface(scf_data, &config); + // the in-core electronic-derivative evaluation streams its amplitudes (never + // materializing `nocc^2 nvir^2` tensors): 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 { @@ -146,6 +259,8 @@ impl<'a> UDHGradient<'a> { 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 }, ); @@ -184,10 +299,15 @@ impl<'a> UDHGradient<'a> { rt::zeros(([nmo, nocc[1]].f(), &device)), ]; for s in 0..2 { - let lag_vo = (&gfock_part[s].i((sv[s].clone(), so[s].clone())) + 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()]); @@ -309,8 +429,6 @@ impl<'a> UDHGradient<'a> { }; // ── persistent MO intermediates; density dots of the rimatr by two GEMVs per spin ── // - 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; handle_memory_exceed( ((nocc[0] + nocc[1]) * nmo * naux + (nvir[0] * nocc[0] + nvir[1] * nocc[1]) * naux) as f64 * 8.0 @@ -658,13 +776,96 @@ impl<'a> UDHGradient<'a> { 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 += + &(2.0 * &(&(&dao_xc_n[s] - &dao_xc_scf[s]) * &d_hf[s].i((.., .., None))).sum_axes(1)); + } + 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 = 2.0 * (&tsr_ipovlp * &(&w_n - &dme0).i((.., .., None))).sum_axes(1); + 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(); + 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)); @@ -679,6 +880,170 @@ fn to_matrix_full(tsr: Tsr) -> MatrixFull { MatrixFull::from_vec(shape, tsr.into_shape(-1).into_raw()).unwrap() } +/// Whether the SCF-iteration functional carries a GGA-or-higher XC part (the gate of the +/// response-density XC term; LDA-type functionals have no grid-gradient contribution here, +/// following pyscf-forge's `_xc_type(xc) == "GGA"` gate). Unrestricted twin of the helper in +/// [`crate::grad::rdh`], duplicated to keep the restricted module untouched. +fn has_scf_xc_gga(scf_data: &SCF) -> bool { + let xc_data = &scf_data.mol.xc_data; + if xc_data.dfa_compnt_scf.is_empty() { + return false; + } + let func_list = scf_xc_func_list_uks(scf_data); + let refs: Vec<&libxc::functional::LibXCFunctional> = func_list.iter().map(|(_, f)| f).collect(); + let xc_type = determine_den_type_from_list(&refs); + matches!(xc_type, XCDenType::SIGMA | XCDenType::TAU) +} + +/// 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; + 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. diff --git a/src/grad/uks.rs b/src/grad/uks.rs index 9eb1cc9fd..d7beb36fc 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/ri_pt2/mod.rs b/src/ri_pt2/mod.rs index a5f82584d..5bcdf8ba7 100644 --- a/src/ri_pt2/mod.rs +++ b/src/ri_pt2/mod.rs @@ -48,6 +48,33 @@ 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`]). +/// +/// The amplitude tensors are never stored at their full `nocc^2 nvir^2` size; the outer occupied +/// index is streamed in windows. Each streamed batch costs `elems_per_step` elements per outer +/// 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. +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_u_elecderiv.rs b/src/ri_pt2/pure_pt2_u_elecderiv.rs index 522199902..082d03484 100644 --- a/src/ri_pt2/pure_pt2_u_elecderiv.rs +++ b/src/ri_pt2/pure_pt2_u_elecderiv.rs @@ -22,11 +22,17 @@ //! - 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. //! -//! Same-spin blocks are evaluated with the pair-symmetry structure of the restricted kernel -//! (single occupied batch, triangular pair loop with the transpose fill). The $\alpha\beta$ -//! block stores only the plain amplitudes ($T = c_{os} t$ enters every contraction as a factor) -//! in a single batch; its four rdm1 contractions are per-index-partial-trace GEMM loops over -//! the contracted pair, so no $[\cdot^2]$-sized intermediate is formed. +//! 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 no `nocc^2 nvir^2` +//! amplitude tensor is ever materialized. 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}; @@ -47,6 +53,16 @@ pub struct UPT2ElecDerivIncoreInp<'a> { 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 { @@ -72,7 +88,7 @@ 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 } = input; + 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(); @@ -80,6 +96,16 @@ pub fn get_rupt2_elec_deriv_incore( 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] { @@ -108,7 +134,7 @@ pub fn get_rupt2_elec_deriv_incore( (-vir_energy[1].i((.., None)) - vir_energy[1].i((None, ..))).into_dim::(), ]; - // ── same-spin blocks: pair-symmetrized amplitudes, rdm1 by 4-d reshapes, G by one GEMM ── // + // ── same-spin blocks: pair-symmetrized amplitudes streamed over occupied windows ── // for s in 0..2 { if nocc[s] == 0 || nvir[s] == 0 { continue; @@ -118,99 +144,133 @@ pub fn get_rupt2_elec_deriv_incore( 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); - let mut t_vivo: Tsr = rt::zeros(([nv, no, nv, no].f(), &device)); - let mut T_vivo: Tsr = rt::zeros(([nv, no, nv, no].f(), &device)); - let pair_ij: Vec<[usize; 2]> = (0..no).flat_map(|i| (0..=i).map(move |j| [i, j])).collect(); - 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(); + 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 mut t_vivo = unsafe { t_vivo.force_mut() }; - let mut t_vivo_T = unsafe { T_vivo.force_mut() }; - t_vivo.i_mut((.., i, .., j)).assign(&t_vv); - t_vivo_T.i_mut((.., i, .., j)).assign(&t_vv_); - if j < i { - t_vivo.i_mut((.., j, .., i)).assign(t_vv.t()); - t_vivo_T.i_mut((.., j, .., i)).assign(t_vv_.t()); - } - }); - // D[oo] -= 2 sum T t; D[vv] += 2 sum T t (both sides of the pair matrix summed) - 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 g_part = - T_vivo.view().reshape([nv * no, nv * no]) % cderi_vox[s].view().reshape([nv * no, naux]); - g_vix[s] = (g_part * 4.0 + &g_vix[s].view().into_shape([nv * no, naux])) - .into_shape([nv, no, naux]); + 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: all (i_a, k_b) pairs; T = c_os t, so only t is stored ── // + // ── 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::(); - let mut t_x: Tsr = rt::zeros(([nva, 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((.., i, ..)) % cderi_vox[1].i((.., k, ..)).t(); - let d_ab = &d_vv_ab + 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); - }); - let t2 = t_x.view().into_shape([nva * noa, nvb * nob]); // [(a i), (b k)] - // G contributions: rows of T2 = c_os t2 are (a i), columns (b k) - let g_a = (t2.view() % cderi_vox[1].view().into_shape([nvb * nob, naux])) * (2.0 * c_os); - *&mut g_vix[0] += &g_a.into_shape([nva, noa, naux]); - let g_b = (t2.view().t() % cderi_vox[0].view().into_shape([nva * noa, naux])) * (2.0 * c_os); - *&mut g_vix[1] += &g_b.into_shape([nvb, nob, naux]); - // correlation rdm1 pieces; each contracts one amplitude index pair completely (the `oo` - // parts) or takes the partial trace over the paired occupied index (the `vv` parts) - let t3 = t_x.view().into_shape([nva * noa * nvb, nob]); // [(a i b), k] - let t4 = t_x.view().into_shape([nva * noa, nvb, nob]); // [(a i), b, k] - let t4b = t_x.view().into_shape([nva, noa, nvb * nob]); // [a, (i), (k b)]-soluble per spin - let m_a_oo = { - let mut m: Tsr = rt::zeros(([noa, noa].f(), &device)); - let mut ta: Tsr = rt::zeros(([noa, nvb * nob].f(), &device)); - for a in 0..nva { - ta.assign(&t_x.i((a, .., .., ..)).into_shape([noa, nvb * nob])); - *&mut m += &(ta.view() % ta.view().t()); - } - m - }; - *&mut rdm1[0].i_mut((0..noa, 0..noa)) -= &(m_a_oo * c_os); - let m_b_oo = t3.view().t() % t3.view(); - *&mut rdm1[1].i_mut((0..nob, 0..nob)) -= &(m_b_oo * c_os); - let m_a_vv = { - let mut m: Tsr = rt::zeros(([nva, nva].f(), &device)); - for i in 0..noa { - let ti = t4b.i((.., i, ..)); - *&mut m += &(ti.view() % ti.view().t()); - } - m - }; - *&mut rdm1[0].i_mut((noa..nmo[0], noa..nmo[0])) += &(m_a_vv * c_os); - let m_b_vv = { - let mut m: Tsr = rt::zeros(([nvb, nvb].f(), &device)); - for k in 0..nob { - let tk = t4.i((.., .., k)); - *&mut m += &(tk.view().t() % tk.view()); - } - m - }; - *&mut rdm1[1].i_mut((nob..nmo[1], nob..nmo[1])) += &(m_b_vv * c_os); + + // --- 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 ── // diff --git a/tests/grad/udh.rs b/tests/grad/udh.rs index 05e18bcd2..d7b4ac491 100644 --- a/tests/grad/udh.rs +++ b/tests/grad/udh.rs @@ -1,6 +1,7 @@ -//! Doubly-hybrid (`xc = "mp2"`) analytic gradient of distorted NH3+ (unrestricted RI-MP2). +//! Doubly-hybrid analytic gradient of distorted NH3+ (unrestricted): the pure-MP2 family +//! (`xc = "mp2"`) and the XYG3 final-functional terms. //! -//! Reference: pyscf-forge `UDFDH(xc="MP2")` analytic gradient (pyscf 2.14.0 + pyscf-forge), +//! 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]`. @@ -44,6 +45,35 @@ new_driver = true """ "##; +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 = "cd" + +[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 + """ +"##; + #[test] fn test_nh3_ump2_grad() { let keys = toml::from_str::(&INPUT_NH3[..]).unwrap(); @@ -106,3 +136,60 @@ fn test_nh3_ump2_grad() { println!("Maximum Error {:?}", (&de_tot_ref - &de_tot).abs().max_all()); assert!((&de_tot_ref - &de_tot).abs().max_all() < 1.0e-5, "UMP2 total gradient mismatch"); } + +/// 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. +#[test] +fn test_nh3_uxyg3_grad() { + let keys = toml::from_str::(&INPUT_XYG3_U[..]).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: {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 (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() < 1.0e-6, "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"); +} -- Gitee From ffba2fc4204bc7ff2330c3b6c184e19c53fb0720 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 13:07:55 +0800 Subject: [PATCH 16/21] Add meta-GGA and DFT-D3(BJ) doubly-hybrid definitions TPSS0_DH and its D3(BJ) variant (Goerigk/Grimme), plus XYG3_D3BJ (xc = B3LYP parent parameters, since XYG3 has no D3 set of its own) and the labelled XYG3_TPSS0 validation construct, which is the only entry exercising the final-functional difference terms on a meta-GGA SCF reference. Validated on NH3/def2-SVP: TPSS0_DH gradient vs energy-FD 2.0e-6 (restricted) / 2.3e-6 (unrestricted), XYG3_TPSS0 7.5e-7; the meta-GGA tau paths (skeleton, response A-tensor, hessian vmat kernels) are exercised. DFT-D3(BJ) enters through the embedded DFTD3 entry, since dftd3 parameter lookup has no doubly-hybrid names; dispersion gradient agrees with pyscf-dispersion to print precision. Co-authored-by: Claude Code Co-authored-by: deepseek-v4-flash --- src/dft/parse_xc/family_bdh.json | 8 ++++++++ src/dft/parse_xc/family_xdh.json | 11 +++++++++++ 2 files changed, 19 insertions(+) diff --git a/src/dft/parse_xc/family_bdh.json b/src/dft/parse_xc/family_bdh.json index 3980fbb29..00a68a0c6 100644 --- a/src/dft/parse_xc/family_bdh.json +++ b/src/dft/parse_xc/family_bdh.json @@ -72,6 +72,14 @@ "code": "0.5*HF + 0.5*PBE, 0.875*PBE + 0.125*MP2", "ref": "10.1063/1.3604569" }, + "TPSS0_DH": { + "code": "0.5*HF + 0.5*TPSS, 0.875*TPSS + 0.125*MP2", + "ref": "10.1021/ct4003527" + }, + "TPSS0_DH_D3BJ": { + "code": "0.5*HF + 0.5*TPSS, 0.875*TPSS + 0.125*MP2 + DFTD3(version=bj, xc=TPSS0)", + "ref": ["10.1021/ct4003527", "10.1021/acs.jpclett.6b00780"] + }, "PBE_QIDH": { "code": "0.6933612744*HF + 0.3066387256*PBE, 0.6666666667*PBE + 0.3333333333*MP2", "ref": "10.1063/1.4890314", diff --git a/src/dft/parse_xc/family_xdh.json b/src/dft/parse_xc/family_xdh.json index c0196a651..38280b8f7 100644 --- a/src/dft/parse_xc/family_xdh.json +++ b/src/dft/parse_xc/family_xdh.json @@ -9,6 +9,17 @@ "code": "0.8033*HF - 0.0140*LDA + 0.2107*B88, 0.6789*LYP + 0.3211*MP2", "ref": "10.1073/pnas.0901093106" }, + "XYG3_D3BJ": { + "code_scf": "B3LYPg", + "code": "0.8033*HF - 0.0140*LDA + 0.2107*B88, 0.6789*LYP + 0.3211*MP2 + DFTD3(version=bj, xc=B3LYP)", + "ref": ["10.1073/pnas.0901093106", "10.1021/acs.jpclett.6b00780"] + }, + "XYG3_TPSS0": { + "code_scf": "0.5*HF + 0.5*TPSS, 0.875*TPSS", + "code": "0.8033*HF - 0.0140*LDA + 0.2107*B88, 0.6789*LYP + 0.3211*MP2", + "ref": "", + "notes": "Validation construct, not a literature functional: the XYG3 final functional on top of a TPSS0 (meta-GGA) SCF reference." + }, "XYGJ_OS": { "code_scf": "B3LYPg", "code": "0.7731*HF + 0.2269*LDA, 0.2309*VWN3 + 0.2754*LYP + 0.4364*MP2_OS", -- Gitee From b207f86e5deb24925c4c7c929bf685798e6e9360 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 14:04:44 +0800 Subject: [PATCH 17/21] Fix review findings in the doubly-hybrid gradient Correctness: - eval_force's analytic branch now returns collect_total_energy, so geometry optimizations follow the energy they differentiate (the DH objective was previously the SCF-functional energy); - refuse range-separated doubly hybrids early (RSH gradients are not supported; previously a silently wrong full-range gradient); - refuse QMMM point charges / solvent with the DH gradient (the correlation-density external-operator terms do not exist; previously a silently incomplete force); - J2C policy/uplo asserts moved to the DH entry (fail fast, was after the whole PT2 + Z-vector cost), and the module doc no longer claims an assertion that did not run; - broaden the driver's grid-regeneration predicate to the modules' actual requirement (custom family JSONs of the HF/MP2 shape no longer panic for missing grids); - de_resp gate now includes LDA SCF functionals (kernels support LDA/GGA/MGGA; the MGGA side was FD-validated), comment rewritten. Robustness / performance: - persistent-memory guard accounts for the co-resident Y-bar partners (~4x under-estimate before); - RGFockPT2 g_vix flag/cache interaction made safe (assert + doc); - empty-grid nsplit no longer panics (step_by(0)); - unresolved [analdrv] control block now passed to the DH response objects (cpscf_* / grid_level_cpscf / resp_auxbas_path were ignored); - skip the unrestricted same-spin pass when c_ss = 0 (OS-only DHs, ~2x wasted work). Docs: amplitude storage bound stated in terms of the nvir nocc naux class; index_occ_outer_vec window-agnostic semantics documented. Also reuse grad::rhf::get_grad_dao_ovlp in de_xc_n/de_ovlp_dh. Re-validated: test_grad 16/36, test_analdrv 36/36, rest_regression DH cases 3/3, geomopt objective matches the force energy to 7e-10, bounded VitC unchanged to 7e-12. Co-authored-by: Claude Code Co-authored-by: deepseek-v4-flash --- src/analdrv/response/rgfock_interface.rs | 6 +-- src/grad/rdh.rs | 65 +++++++++++++++--------- src/grad/udh.rs | 62 +++++++++++++--------- src/main_driver.rs | 19 ++++--- src/ri_pt2/mod.rs | 11 ++-- src/ri_pt2/pure_pt2_u_elecderiv.rs | 9 ++-- src/ri_pt2/rgfock_pt2.rs | 11 +++- 7 files changed, 115 insertions(+), 68 deletions(-) diff --git a/src/analdrv/response/rgfock_interface.rs b/src/analdrv/response/rgfock_interface.rs index d650623a4..6f9066cf4 100644 --- a/src/analdrv/response/rgfock_interface.rs +++ b/src/analdrv/response/rgfock_interface.rs @@ -457,9 +457,9 @@ 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 (never materializing - // `nocc^2 nvir^2` tensors): size the windows so the per-window transients stay within a few - // times the `nvir * nocc * naux` class and inside `max_memory` + // 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")); diff --git a/src/grad/rdh.rs b/src/grad/rdh.rs index b569c24d8..4e5d9c899 100644 --- a/src/grad/rdh.rs +++ b/src/grad/rdh.rs @@ -37,9 +37,8 @@ //! shell batch reduces directly (last-axis `rt::vecdot`) against `Ȳ`, never materializing //! `[naux, 3, x, y]` intermediates. The metric-derivative (`L_1`) term, however, follows forge's //! Cholesky-generator identity and therefore still requires the `Cd` convention (same as the -//! rimatr itself), so a matching `policy` is asserted below. +//! rimatr itself), so `calc` checks the control policy up front (`s_op` itself is trusted). -use crate::analdrv::config::AnalDrvConfig; 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}; @@ -49,13 +48,13 @@ 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::{RIRHFGradient, RIHFGradientFlagsBuilder}; +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_jk::{get_j2c_decomp, J2CDecompPolicy, 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; @@ -88,6 +87,24 @@ impl<'a> RDHGradient<'a> { 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." + ); + // the metric-derivative generator below follows forge's Cholesky-generator identity and + // is CD-only; refuse the other conventions before the PT2/Z-vector work (`s_op` trusted) + assert!( + matches!(scf_data.mol.ctrl.j2c_decomp.policy, J2CDecompPolicy::Cd), + "The DH analytic gradient's metric-derivative generator currently supports the CD convention only. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input (the derivative-integral solves themselves follow the rimatr convention)." + ); + assert!( + !matches!(scf_data.mol.ctrl.j2c_decomp.uplo, Lower), + "RDHGradient does not support lower-triangular Cholesky factors (developer option)." + ); 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.", ); @@ -112,7 +129,7 @@ impl<'a> RDHGradient<'a> { let d_hf = (&scf_data.density_matrix[0]).to_rstsr(&device); // ── response objects, relaxed correlation density, full four-block gfock ── - let config = AnalDrvConfig::default(); + 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 @@ -138,7 +155,7 @@ impl<'a> RDHGradient<'a> { 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_gga(scf_data); + 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 { @@ -174,9 +191,9 @@ impl<'a> RDHGradient<'a> { } else { None }; - // the PT2 kernel streams its amplitudes over occupied windows (never materializing - // `nocc^2 nvir^2` tensors): size the windows so the per-window transients stay within a - // few times the `nvir * nocc * naux` class and inside `max_memory` + // 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; @@ -316,7 +333,13 @@ impl<'a> RDHGradient<'a> { // ── 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; - handle_memory_exceed((nocc * nmo * naux + nvir * nocc * naux) as f64 * 8.0 / 1048576.0, mem_avail, abort); + // true peak of the persistent intermediates: `y_ip`/`g_vix` plus the three layout + // copies of their `Ȳ` partners, all co-resident until the two sources are dropped + handle_memory_exceed( + 4.0 * (nocc * nmo * naux + nvir * nocc * naux) as f64 * 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]); @@ -550,7 +573,7 @@ impl<'a> RDHGradient<'a> { // 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 = 2.0 * (&(&dao_xc_n - &dao_xc_scf) * &d_hf.i((.., .., None))).sum_axes(1); + 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); @@ -574,7 +597,7 @@ impl<'a> RDHGradient<'a> { // 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 = - 2.0 * (&tsr_ipovlp * &(&(&w_n * 2.0) - &dme0).i((.., .., None))).sum_axes(1); + 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); @@ -638,7 +661,7 @@ fn de_xc_response_term( 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; + 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::*; @@ -729,18 +752,10 @@ impl GradAPI for RDHGradient<'_> { } } -/// Whether the SCF-iteration functional carries a GGA-or-higher XC part (the gate of the -/// response-density XC term; LDA-type functionals have no grid-gradient contribution here, -/// following pyscf-forge's `_xc_type(xc) == "GGA"` gate). -fn has_scf_xc_gga(scf_data: &SCF) -> bool { - let xc_data = &scf_data.mol.xc_data; - if xc_data.dfa_compnt_scf.is_empty() { - return false; - } - let func_list = scf_xc_func_list(scf_data); - let refs: Vec<&libxc::functional::LibXCFunctional> = func_list.iter().map(|(_, f)| f).collect(); - let xc_type = determine_den_type_from_list(&refs); - matches!(xc_type, XCDenType::SIGMA | XCDenType::TAU) +/// 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 diff --git a/src/grad/udh.rs b/src/grad/udh.rs index f21717416..4508a4ef3 100644 --- a/src/grad/udh.rs +++ b/src/grad/udh.rs @@ -31,7 +31,6 @@ //! $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::config::AnalDrvConfig; use crate::analdrv::response::rgfock_interface::dh_jk_factors; use crate::analdrv::response::trait_uresp::URespAPI; use crate::analdrv::response::uresp_interface::{ @@ -42,14 +41,16 @@ 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, pack_triu_tilde, RIHFGradientFlagsBuilder}; +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}; +use crate::ri_jk::{get_j2c_decomp, J2CDecompPolicy, J2CDecompose}; use crate::ri_pt2::pure_pt2_u_elecderiv::{ get_rupt2_elec_deriv_incore, UPT2ElecDerivIncoreArg, UPT2ElecDerivIncoreInp, }; @@ -88,6 +89,24 @@ impl<'a> UDHGradient<'a> { 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." + ); + // the metric-derivative generator below follows forge's Cholesky-generator identity and + // is CD-only; refuse the other conventions before the PT2/Z-vector work (`s_op` trusted) + assert!( + matches!(scf_data.mol.ctrl.j2c_decomp.policy, J2CDecompPolicy::Cd), + "The DH analytic gradient's metric-derivative generator currently supports the CD convention only. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input (the derivative-integral solves themselves follow the rimatr convention)." + ); + assert!( + !matches!(scf_data.mol.ctrl.j2c_decomp.uplo, Lower), + "UDHGradient does not support lower-triangular Cholesky factors (developer option)." + ); 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.", ); @@ -121,7 +140,7 @@ impl<'a> UDHGradient<'a> { 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_gga(scf_data); + 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 { @@ -202,10 +221,10 @@ impl<'a> UDHGradient<'a> { }; // ── response objects, PT2 electronic derivative, relaxed correlation density ── // - let config = AnalDrvConfig::default(); + 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 (never - // materializing `nocc^2 nvir^2` tensors): the same-spin blocks window their outer + // 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`. @@ -429,8 +448,10 @@ impl<'a> UDHGradient<'a> { }; // ── persistent MO intermediates; density dots of the rimatr by two GEMVs per spin ── // + // true peak: the `y_ip`/`g_vix` partners plus the three layout copies of their `Ȳ` + // partners per spin, all co-resident until the two sources are dropped handle_memory_exceed( - ((nocc[0] + nocc[1]) * nmo * naux + (nvir[0] * nocc[0] + nvir[1] * nocc[1]) * naux) as f64 + 4.0 * ((nocc[0] + nocc[1]) * nmo * naux + (nvir[0] * nocc[0] + nvir[1] * nocc[1]) * naux) as f64 * 8.0 / 1048576.0, mem_avail, @@ -815,7 +836,7 @@ impl<'a> UDHGradient<'a> { let mut dao_xc: Tsr = rt::zeros(([nao, 3], &device)); for s in 0..2 { *&mut dao_xc += - &(2.0 * &(&(&dao_xc_n[s] - &dao_xc_scf[s]) * &d_hf[s].i((.., .., None))).sum_axes(1)); + &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]; @@ -846,7 +867,7 @@ impl<'a> UDHGradient<'a> { // 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 = 2.0 * (&tsr_ipovlp * &(&w_n - &dme0).i((.., .., None))).sum_axes(1); + 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); @@ -880,19 +901,12 @@ fn to_matrix_full(tsr: Tsr) -> MatrixFull { MatrixFull::from_vec(shape, tsr.into_shape(-1).into_raw()).unwrap() } -/// Whether the SCF-iteration functional carries a GGA-or-higher XC part (the gate of the -/// response-density XC term; LDA-type functionals have no grid-gradient contribution here, -/// following pyscf-forge's `_xc_type(xc) == "GGA"` gate). Unrestricted twin of the helper in -/// [`crate::grad::rdh`], duplicated to keep the restricted module untouched. -fn has_scf_xc_gga(scf_data: &SCF) -> bool { - let xc_data = &scf_data.mol.xc_data; - if xc_data.dfa_compnt_scf.is_empty() { - return false; - } - let func_list = scf_xc_func_list_uks(scf_data); - let refs: Vec<&libxc::functional::LibXCFunctional> = func_list.iter().map(|(_, f)| f).collect(); - let xc_type = determine_den_type_from_list(&refs); - matches!(xc_type, XCDenType::SIGMA | XCDenType::TAU) +/// 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 @@ -991,7 +1005,7 @@ fn de_xc_response_term( 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; + 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::*; diff --git a/src/main_driver.rs b/src/main_driver.rs index cdd171e72..b48e4ea96 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -888,12 +888,18 @@ 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![]; - // post-SCF (doubly-hybrid) gradients with a final-functional XC part need the DFT grids - // again, but `xdh_calculations` frees them for memory; regenerate before the SCF gradient - // borrows `scf_data` (the pure-MP2 family has no XC part and needs no grid) + // 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 has_dh_xc = scf_data.mol.xc_data.dfa_compnt_pos.as_ref().map_or(false, |v| !v.is_empty()); - if is_dh && has_dh_xc && scf_data.grids.is_none() { + 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)); } @@ -1006,7 +1012,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 5bcdf8ba7..2bfe289fa 100644 --- a/src/ri_pt2/mod.rs +++ b/src/ri_pt2/mod.rs @@ -51,11 +51,12 @@ 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`]). /// -/// The amplitude tensors are never stored at their full `nocc^2 nvir^2` size; the outer occupied -/// index is streamed in windows. Each streamed batch costs `elems_per_step` elements per outer -/// 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. +/// 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; diff --git a/src/ri_pt2/pure_pt2_u_elecderiv.rs b/src/ri_pt2/pure_pt2_u_elecderiv.rs index 082d03484..1ef44b77f 100644 --- a/src/ri_pt2/pure_pt2_u_elecderiv.rs +++ b/src/ri_pt2/pure_pt2_u_elecderiv.rs @@ -25,8 +25,10 @@ //! 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 no `nocc^2 nvir^2` -//! amplitude tensor is ever materialized. The same-spin blocks use the pair-symmetry structure of +//! 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 @@ -136,7 +138,8 @@ pub fn get_rupt2_elec_deriv_incore( // ── same-spin blocks: pair-symmetrized amplitudes streamed over occupied windows ── // for s in 0..2 { - if nocc[s] == 0 || nvir[s] == 0 { + // 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]); diff --git a/src/ri_pt2/rgfock_pt2.rs b/src/ri_pt2/rgfock_pt2.rs index 2d2ea7f56..ff86eb5eb 100644 --- a/src/ri_pt2/rgfock_pt2.rs +++ b/src/ri_pt2/rgfock_pt2.rs @@ -74,7 +74,9 @@ 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, @@ -82,7 +84,8 @@ where pub c_ss: f64, /// 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 when enabled. + /// `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 @@ -143,6 +146,10 @@ where 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(); -- Gitee From 8ebd79770ed6d916dbee11f13fabc6924197a239 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 15:37:24 +0800 Subject: [PATCH 18/21] Revert "Add meta-GGA and DFT-D3(BJ) doubly-hybrid definitions" The TPSS0_DH, TPSS0_DH_D3BJ, XYG3_D3BJ and XYG3_TPSS0 family entries did not exist before and are beyond this task's scope; the M12 validation that exercised them is kept as an untracked artifact. This reverts commit ffba2fc4. Co-Authored-By: Claude Code Co-Authored-By: deepseek-v4-flash --- src/dft/parse_xc/family_bdh.json | 8 -------- src/dft/parse_xc/family_xdh.json | 11 ----------- 2 files changed, 19 deletions(-) diff --git a/src/dft/parse_xc/family_bdh.json b/src/dft/parse_xc/family_bdh.json index 00a68a0c6..3980fbb29 100644 --- a/src/dft/parse_xc/family_bdh.json +++ b/src/dft/parse_xc/family_bdh.json @@ -72,14 +72,6 @@ "code": "0.5*HF + 0.5*PBE, 0.875*PBE + 0.125*MP2", "ref": "10.1063/1.3604569" }, - "TPSS0_DH": { - "code": "0.5*HF + 0.5*TPSS, 0.875*TPSS + 0.125*MP2", - "ref": "10.1021/ct4003527" - }, - "TPSS0_DH_D3BJ": { - "code": "0.5*HF + 0.5*TPSS, 0.875*TPSS + 0.125*MP2 + DFTD3(version=bj, xc=TPSS0)", - "ref": ["10.1021/ct4003527", "10.1021/acs.jpclett.6b00780"] - }, "PBE_QIDH": { "code": "0.6933612744*HF + 0.3066387256*PBE, 0.6666666667*PBE + 0.3333333333*MP2", "ref": "10.1063/1.4890314", diff --git a/src/dft/parse_xc/family_xdh.json b/src/dft/parse_xc/family_xdh.json index 38280b8f7..c0196a651 100644 --- a/src/dft/parse_xc/family_xdh.json +++ b/src/dft/parse_xc/family_xdh.json @@ -9,17 +9,6 @@ "code": "0.8033*HF - 0.0140*LDA + 0.2107*B88, 0.6789*LYP + 0.3211*MP2", "ref": "10.1073/pnas.0901093106" }, - "XYG3_D3BJ": { - "code_scf": "B3LYPg", - "code": "0.8033*HF - 0.0140*LDA + 0.2107*B88, 0.6789*LYP + 0.3211*MP2 + DFTD3(version=bj, xc=B3LYP)", - "ref": ["10.1073/pnas.0901093106", "10.1021/acs.jpclett.6b00780"] - }, - "XYG3_TPSS0": { - "code_scf": "0.5*HF + 0.5*TPSS, 0.875*TPSS", - "code": "0.8033*HF - 0.0140*LDA + 0.2107*B88, 0.6789*LYP + 0.3211*MP2", - "ref": "", - "notes": "Validation construct, not a literature functional: the XYG3 final functional on top of a TPSS0 (meta-GGA) SCF reference." - }, "XYGJ_OS": { "code_scf": "B3LYPg", "code": "0.7731*HF + 0.2269*LDA, 0.2309*VWN3 + 0.2754*LYP + 0.4364*MP2_OS", -- Gitee From e159575ac7f5611b8de749454a5ea50484ff0a9c Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 16:25:44 +0800 Subject: [PATCH 19/21] =?UTF-8?q?Cut=20the=20DH=20gradient=20=C8=B2=20part?= =?UTF-8?q?ner=20memory=20to=20two=20layouts?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The restricted/unrestricted DH gradient kept three layouts of each of the two Ȳ data sets co-resident (the metric-solve output [(i m), Q], a permutation copy [i, Q, m] for the u slot and a third [m, Q, i] for the v slot) plus the y_ip/g_vix sources, i.e. 5+5 partner-class tensors at the construction peak (the first bindings of the shadowed `ybar_jk_v`/`ybar_ri_v` are dropped only at the end of the block, so the old guard 4*(Y+G) underestimated the true peak by 25 %). Now each partner is stored in a single layout [i, Q, m] resp. [i, Q, a], built as one metric solve plus one permutation, with the source dropped right after the solve. The u slot keeps its existing kernel; the v slot uses a new kernel that consumes the same layout as the left operand of one [nocc naux, nmo] % [nmo, nub] fold and assembles the full-AO-side fold from one [y, naov] % [naov, naux] GEMM per batch row (same FLOPs and GEMM structure per slot, no second layout); the ip2-side reductions pair the two contraction axes individually. The guard now states the exact peak max(2Y, Y+2G) (udh walks the per-spin construction sequence and releases each spin's source with its own solve transient). Vitamin C (nocc 46, nvir 374, naux 1056; Y = nocc*nao*naux = 155.7 MB, G = 138.6 MB): partner peak 1471.3 MB (5Y+5G) -> 432.9 MB (1Y+2G), persistent footprint 1177.2 MB -> 294.3 MB (Y+G, the two tensors themselves); whole-run peak RSS 14.03 -> 12.22 GB (RMP2) and 17.56 -> 15.99 GB (RXYG3). Force phase: RMP2 26.0 -> 25.6 s, RXYG3 57.1 -> 55.1 s, UMP2 52.9 s (stored 53.5 s). Validation: cargo test --release --test test_grad -- --test-threads=1 -> 16 passed (NH3 RMP2/RXYG3/UMP2/UXYG3 vs the pyscf-forge references); VitC RMP2/RXYG3 gradients match the previous binary to 3e-11 / 9e-12, below the same-binary run-to-run scatter (1.5e-10). Memory findings report: local-runs/260930-vibe-dh-grad/scratch-mem/. Co-authored-by: Claude Code Co-authored-by: deepseek-v4-flash --- src/grad/rdh.rs | 118 +++++++++++++++++++++++++---------- src/grad/udh.rs | 163 +++++++++++++++++++++++++----------------------- 2 files changed, 169 insertions(+), 112 deletions(-) diff --git a/src/grad/rdh.rs b/src/grad/rdh.rs index 4e5d9c899..7c626df25 100644 --- a/src/grad/rdh.rs +++ b/src/grad/rdh.rs @@ -34,10 +34,13 @@ //! 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 (last-axis `rt::vecdot`) against `Ȳ`, never materializing -//! `[naux, 3, x, y]` intermediates. The metric-derivative (`L_1`) term, however, follows forge's -//! Cholesky-generator identity and therefore still requires the `Cd` convention (same as the -//! rimatr itself), so `calc` checks the control policy up front (`s_op` itself is trusted). +//! 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, however, follows forge's Cholesky-generator +//! identity and therefore still requires the `Cd` convention (same as the rimatr itself), so +//! `calc` checks the control policy up front (`s_op` itself is trusted). 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}; @@ -333,13 +336,11 @@ impl<'a> RDHGradient<'a> { // ── 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 persistent intermediates: `y_ip`/`g_vix` plus the three layout - // copies of their `Ȳ` partners, all co-resident until the two sources are dropped - handle_memory_exceed( - 4.0 * (nocc * nmo * naux + nvir * nocc * naux) as f64 * 8.0 / 1048576.0, - mem_avail, - abort, - ); + // 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]); @@ -364,18 +365,21 @@ impl<'a> RDHGradient<'a> { } *&mut r_jk *= cx_scf; - // ── pre-contractions `Ȳ = partner % S` (replace all in-loop metric solves) ── - let ybar_y2: Tsr = y_ip.view().into_shape([nocc * nmo, naux]) % s_op.view(); // [(i m), Q] - let ybar_g2i: Tsr = g_vix.view().transpose([1, 0, 2]).into_contig(ColMajor).view().into_shape([nocc * nvir, naux]) % s_op.view(); // [(i a), Q] - // layout copies `[x, Q, y]` for the ip1 shell-batch reductions - let ybar_jk_u = ybar_y2.view().into_shape([nocc, nmo, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [i, Q, m] - let ybar_jk_v = y_ip.view().transpose([1, 0, 2]).into_contig(ColMajor).view().into_shape([nmo * nocc, naux]) % s_op.view(); // [(m i), Q] - let ybar_jk_v = ybar_jk_v.view().into_shape([nmo, nocc, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [m, Q, i] - let ybar_ri_u = ybar_g2i.view().into_shape([nocc, nvir, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [i, Q, a] - let ybar_ri_v = g_vix.view().into_shape([nvir * nocc, naux]) % s_op.view(); // [(a i), Q] - let ybar_ri_v = ybar_ri_v.view().into_shape([nvir, nocc, naux]).transpose([0, 2, 1]).into_contig(ColMajor); // [a, Q, i] - drop(y_ip); - drop(g_vix); + // ── 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(); @@ -436,10 +440,10 @@ impl<'a> RDHGradient<'a> { 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_u.view(), 2.0, t4_3c); - t4_3c = add_fold_ip1_batched(ip1_t2.view(), nub, c1_dr.i((r.clone(), ..)), c0.view(), ybar_jk_v.view(), 2.0, t4_3c); - ri_3c = add_fold_ip1_batched(ip1_t2.view(), nub, c0.i((r.clone(), ..)), cv.view(), ybar_ri_u.view(), -4.0, ri_3c); - ri_3c = add_fold_ip1_batched(ip1_t2.view(), nub, cv.i((r, ..)), c0.view(), ybar_ri_v.view(), -4.0, ri_3c); + 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()); @@ -473,15 +477,22 @@ impl<'a> RDHGradient<'a> { 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 leading-(i m)/(i a) reshape pairs with the `Ȳ` rows + // 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)] - let f_jk2 = f_jk.view().into_shape([nocc * nmo, nb_q * ncomp]); - let f_ri2 = f_ri.view().into_shape([nocc * nvir, nb_q * ncomp]); for t in 0..ncomp { - // `r_jk_q[q] = sum_{i m} F_jk[(i m), (q t)] Ȳ[(i m), q]` — contracted over (i m) - let r_jk_q = rt::vecdot(&f_jk2.i((.., t * nb_q..(t + 1) * nb_q)), ybar_y2.i((.., qa..qb)), 0); - let r_ri_q = rt::vecdot(&f_ri2.i((.., t * nb_q..(t + 1) * nb_q)), ybar_g2i.i((.., qa..qb)), 0); + // `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)); @@ -709,7 +720,9 @@ fn to_matrix_full(tsr: Tsr) -> MatrixFull { } /// 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). +/// 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 @@ -743,6 +756,43 @@ fn add_fold_ip1_batched( 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() diff --git a/src/grad/udh.rs b/src/grad/udh.rs index 4508a4ef3..1e782cf7c 100644 --- a/src/grad/udh.rs +++ b/src/grad/udh.rs @@ -448,15 +448,21 @@ impl<'a> UDHGradient<'a> { }; // ── persistent MO intermediates; density dots of the rimatr by two GEMVs per spin ── // - // true peak: the `y_ip`/`g_vix` partners plus the three layout copies of their `Ȳ` - // partners per spin, all co-resident until the two sources are dropped - handle_memory_exceed( - 4.0 * ((nocc[0] + nocc[1]) * nmo * naux + (nvir[0] * nocc[0] + nvir[1] * nocc[1]) * naux) as f64 - * 8.0 - / 1048576.0, - mem_avail, - abort, - ); + // 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] = @@ -504,68 +510,26 @@ impl<'a> UDHGradient<'a> { } *&mut r_jk *= cx_scf; - // ── pre-contractions `Ȳ = partner % S` (replace all in-loop metric solves), per spin ── // - let mut ybar_y2: [Option; 2] = [None, None]; - let mut ybar_g2i: [Option; 2] = [None, None]; - let mut ybar_jk_u: [Option; 2] = [None, None]; - let mut ybar_jk_v: [Option; 2] = [None, None]; - let mut ybar_ri_u: [Option; 2] = [None, None]; - let mut ybar_ri_v: [Option; 2] = [None, 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 y2 = y_ip[s].view().into_shape([nocc[s] * nmo, naux]) % s_op.view(); - let g2i = g_vix[s] - .view() - .transpose([1, 0, 2]) - .into_contig(ColMajor) - .view() - .into_shape([nocc[s] * nvir[s], naux]) - % s_op.view(); - ybar_y2[s] = Some(y2); - ybar_g2i[s] = Some(g2i); - ybar_jk_u[s] = Some( - ybar_y2[s] - .as_ref() - .unwrap() - .view() - .into_shape([nocc[s], nmo, naux]) - .transpose([0, 2, 1]) - .into_contig(ColMajor), - ); - let yv = y_ip[s] - .view() - .transpose([1, 0, 2]) - .into_contig(ColMajor) - .view() - .into_shape([nmo * nocc[s], naux]) - % s_op.view(); - ybar_jk_v[s] = Some( - yv.view() - .into_shape([nmo, nocc[s], naux]) - .transpose([0, 2, 1]) - .into_contig(ColMajor), - ); - ybar_ri_u[s] = Some( - ybar_g2i[s] - .as_ref() - .unwrap() - .view() - .into_shape([nocc[s], nvir[s], naux]) - .transpose([0, 2, 1]) - .into_contig(ColMajor), - ); - let gv = g_vix[s].view().into_shape([nvir[s] * nocc[s], naux]) % s_op.view(); - ybar_ri_v[s] = Some( - gv.view() - .into_shape([nvir[s], nocc[s], naux]) - .transpose([0, 2, 1]) - .into_contig(ColMajor), - ); + 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] } - drop(y_ip); - drop(g_vix); // ── derivative-integral passes: shell batches over all atoms, atom scatter in-batch ── // let mol_loc = mol.ao_loc(); @@ -638,16 +602,16 @@ impl<'a> UDHGradient<'a> { nub, c0[s].i((r.clone(), ..)), c1_dr[s].view(), - ybar_jk_u[s].as_ref().unwrap().view(), + ybar_jk[s].as_ref().unwrap().view(), 2.0, t4_3c, ); - t4_3c = add_fold_ip1_batched( + t4_3c = add_fold_ip1_batched_v( ip1_t2.view(), nub, c1_dr[s].i((r.clone(), ..)), c0[s].view(), - ybar_jk_v[s].as_ref().unwrap().view(), + ybar_jk[s].as_ref().unwrap().view(), 2.0, t4_3c, ); @@ -656,16 +620,16 @@ impl<'a> UDHGradient<'a> { nub, c0[s].i((r.clone(), ..)), cv[s].view(), - ybar_ri_u[s].as_ref().unwrap().view(), + ybar_ri[s].as_ref().unwrap().view(), -1.0, ri_3c, ); - ri_3c = add_fold_ip1_batched( + ri_3c = add_fold_ip1_batched_v( ip1_t2.view(), nub, cv[s].i((r.clone(), ..)), c0[s].view(), - ybar_ri_v[s].as_ref().unwrap().view(), + ybar_ri[s].as_ref().unwrap().view(), -1.0, ri_3c, ); @@ -719,13 +683,19 @@ impl<'a> UDHGradient<'a> { 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); - let f_jk2 = f_jk.view().into_shape([nocc[s] * nmo, nb_q * ncomp]); - let f_ri2 = f_ri.view().into_shape([nocc[s] * nvir[s], nb_q * ncomp]); for t in 0..ncomp { - let r_jk_q = - rt::vecdot(&f_jk2.i((.., t * nb_q..(t + 1) * nb_q)), ybar_y2[s].as_ref().unwrap().view().i((.., qa..qb)), 0); - let r_ri_q = - rt::vecdot(&f_ri2.i((.., t * nb_q..(t + 1) * nb_q)), ybar_g2i[s].as_ref().unwrap().view().i((.., qa..qb)), 0); + // 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)); @@ -1089,6 +1059,43 @@ fn add_fold_ip1_batched( 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() -- Gitee From 1a853d45be61984cab929fe1820bbf8762c5bbed Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 16:55:42 +0800 Subject: [PATCH 20/21] Support the eigen J2C decomposition in the DH gradient The doubly-hybrid gradient's metric-derivative (L_1) term no longer requires [ctrl.j2c_decomp] policy = "cd": the CD-only assert and the CD-only generator panic are gone. Under Cd the forge Cholesky generator is unchanged; under Eig (the default policy, which used to panic) the same term is evaluated through the self-adjoint Frechet derivative of J^-1/2: the y-dot/exchange and RI weight matrices are folded once (G = V (D' . (V^T M^T V)) V^T, D'_ij = -1/(sqrt(e_i)(sqrt(e_i)+sqrt(e_j)))) and each atom's raw int2c2e_ip1 block contracts against G + G^T directly, so no per-atom naux^3 work is needed. The fold uses the same retained eigenpairs as the decomposition itself. Both files are mirrored; tests/grad gain eig twins of the RMP2/XYG3/UMP2/UXYG3 cases with the same references/tolerances. Validation: cargo test --release --test test_grad (20 passed, 0 failed); eig vs cd: DH 3.2e-14 (RMP2) / 7.2e-15 (UMP2) / 1.9e-14 (RXYG3) / 1.4e-14 (UXYG3), totals to 5.9e-13; eigen-vs-pyscf-forge residuals unchanged (1.3e-8..5e-7); central-difference check of the total energy agrees to ~5e-6 (FD-noise limited, identical for cd); rest_regression DH cases 3/3; Vitamin C force benchmark no regression (RMP2 29.9 s cd / 29.1 s eig vs 29.8-30.1 s stored at HEAD, RXYG3 69.8 / 69.8 vs 69.2 s stored), cd-vs-eig VitC gradients agree to 2.7e-11. Co-authored-by: Claude Code Co-authored-by: deepseek-v4-flash --- src/grad/rdh.rs | 96 +++++++++++++++++++++++++++++++++------------- src/grad/udh.rs | 92 +++++++++++++++++++++++++++++++++----------- src/main_driver.rs | 2 +- tests/grad/rdh.rs | 50 ++++++++++++++++-------- tests/grad/udh.rs | 54 ++++++++++++++++++-------- 5 files changed, 211 insertions(+), 83 deletions(-) diff --git a/src/grad/rdh.rs b/src/grad/rdh.rs index 7c626df25..8defcf746 100644 --- a/src/grad/rdh.rs +++ b/src/grad/rdh.rs @@ -38,9 +38,10 @@ //! `[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, however, follows forge's Cholesky-generator -//! identity and therefore still requires the `Cd` convention (same as the rimatr itself), so -//! `calc` checks the control policy up front (`s_op` itself is trusted). +//! 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}; @@ -57,7 +58,7 @@ 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, J2CDecompPolicy, J2CDecompose, J2C_THRESH}; +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; @@ -98,12 +99,6 @@ impl<'a> RDHGradient<'a> { 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." ); - // the metric-derivative generator below follows forge's Cholesky-generator identity and - // is CD-only; refuse the other conventions before the PT2/Z-vector work (`s_op` trusted) - assert!( - matches!(scf_data.mol.ctrl.j2c_decomp.policy, J2CDecompPolicy::Cd), - "The DH analytic gradient's metric-derivative generator currently supports the CD convention only. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input (the derivative-integral solves themselves follow the rimatr convention)." - ); assert!( !matches!(scf_data.mol.ctrl.j2c_decomp.uplo, Lower), "RDHGradient does not support lower-triangular Cholesky factors (developer option)." @@ -315,8 +310,9 @@ impl<'a> RDHGradient<'a> { let (out, shape) = aux.integrate("int2c2e_ip1", "s1", None).into(); rt::asarray((out, shape, &device)) }; - // strictly-lower mask with 1/2 diagonal, of forge `generator_L_1` (Cd only) - let (l_lower, l_mask) = match &j2c_decomp { + // 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)); @@ -326,11 +322,9 @@ impl<'a> RDHGradient<'a> { } l_mask[[i, i]] = 0.5; } - (l_lower, l_mask) + Some((l_lower, l_mask)) }, - _ => panic!( - "The DH analytic gradient's metric-derivative generator currently supports the CD convention only. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input (the derivative-integral solves themselves follow the rimatr convention)." - ), + J2CDecompose::Eig { .. } => None, }; // ── persistent MO intermediates; density dots of the rimatr by two GEMVs ── @@ -364,6 +358,25 @@ impl<'a> RDHGradient<'a> { *&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 @@ -517,18 +530,30 @@ impl<'a> RDHGradient<'a> { let mut ydr_l1: Tsr = rt::zeros(([naux, ncomp], &device)); let mut t4_l1 = [0.0f64; 3]; let mut ri_l1 = [0.0f64; 3]; - // forge L_1 generator of this atom, contracted with the solve operator + // `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 { - 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(); + 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 @@ -719,6 +744,23 @@ fn to_matrix_full(tsr: Tsr) -> MatrixFull { 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: diff --git a/src/grad/udh.rs b/src/grad/udh.rs index 1e782cf7c..37843f1aa 100644 --- a/src/grad/udh.rs +++ b/src/grad/udh.rs @@ -50,7 +50,7 @@ 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, J2CDecompPolicy, J2CDecompose}; +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, }; @@ -97,12 +97,6 @@ impl<'a> UDHGradient<'a> { 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." ); - // the metric-derivative generator below follows forge's Cholesky-generator identity and - // is CD-only; refuse the other conventions before the PT2/Z-vector work (`s_op` trusted) - assert!( - matches!(scf_data.mol.ctrl.j2c_decomp.policy, J2CDecompPolicy::Cd), - "The DH analytic gradient's metric-derivative generator currently supports the CD convention only. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input (the derivative-integral solves themselves follow the rimatr convention)." - ); assert!( !matches!(scf_data.mol.ctrl.j2c_decomp.uplo, Lower), "UDHGradient does not support lower-triangular Cholesky factors (developer option)." @@ -430,7 +424,9 @@ impl<'a> UDHGradient<'a> { let (out, shape) = aux.integrate("int2c2e_ip1", "s1", None).into(); rt::asarray((out, shape, &device)) }; - let (l_lower, l_mask) = match &j2c_decomp { + // 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)); @@ -440,11 +436,9 @@ impl<'a> UDHGradient<'a> { } l_mask[[i, i]] = 0.5; } - (l_lower, l_mask) + Some((l_lower, l_mask)) }, - _ => panic!( - "The DH analytic gradient's metric-derivative generator currently supports the CD convention only. Please set\n[ctrl.j2c_decomp]\npolicy = \"cd\"\nin the ctrl input (the derivative-integral solves themselves follow the rimatr convention)." - ), + J2CDecompose::Eig { .. } => None, }; // ── persistent MO intermediates; density dots of the rimatr by two GEMVs per spin ── // @@ -509,6 +503,27 @@ impl<'a> UDHGradient<'a> { } } *&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 @@ -729,19 +744,32 @@ impl<'a> UDHGradient<'a> { [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 { - 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())); + 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(); + }, } - 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(); } // forge t1 + t2 (all four spin pairs) + t4 (per spin), and the RI term @@ -871,6 +899,24 @@ fn to_matrix_full(tsr: Tsr) -> MatrixFull { 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 diff --git a/src/main_driver.rs b/src/main_driver.rs index b48e4ea96..16fdebb68 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -929,7 +929,7 @@ 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 CD-decomposed rimatr + // 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); diff --git a/tests/grad/rdh.rs b/tests/grad/rdh.rs index c10b0e868..512e2f3b9 100644 --- a/tests/grad/rdh.rs +++ b/tests/grad/rdh.rs @@ -5,7 +5,9 @@ //! 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). The rimatr must be CD-decomposed (see `src/grad/rdh.rs`). +//! 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; @@ -28,7 +30,7 @@ static INPUT_NH3: &str = r##" spin_polarization = false [ctrl.j2c_decomp] -policy = "cd" +policy = "POLICY" [ctrl.ri_pt2] new_driver = true @@ -57,7 +59,7 @@ static INPUT_XYG3: &str = r##" spin_polarization = false [ctrl.j2c_decomp] -policy = "cd" +policy = "POLICY" [ctrl.ri_pt2] new_driver = true @@ -73,9 +75,8 @@ new_driver = true """ "##; -#[test] -fn test_nh3_rmp2_grad() { - let keys = toml::from_str::(&INPUT_NH3[..]).unwrap(); +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); @@ -83,7 +84,7 @@ fn test_nh3_rmp2_grad() { // 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: {eng}"); + 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); @@ -100,7 +101,7 @@ fn test_nh3_rmp2_grad() { 0.0042295109, -0.0090684480, 0.0019845245, ]; let de_ref = rt::asarray((de_ref, [3, 4])); - println!("DH correlation gradient (REST): {:16.9}", de.t()); + 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()); @@ -109,12 +110,22 @@ fn test_nh3_rmp2_grad() { 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() < 1.0e-5); + 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 @@ -123,16 +134,15 @@ fn test_nh3_rmp2_grad() { /// 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). -#[test] -fn test_nh3_rxyg3_grad() { - let keys = toml::from_str::(&INPUT_XYG3[..]).unwrap(); +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: {eng}"); + 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) @@ -157,7 +167,7 @@ fn test_nh3_rxyg3_grad() { -0.0141324624, 0.0093628946, -0.0062191910, ]; let de_ref = rt::asarray((de_ref, [3, 4])); - println!("XYG3 total gradient (REST): {:16.9}", de.t()); + 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()); @@ -166,5 +176,15 @@ fn test_nh3_rxyg3_grad() { 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() < 1.0e-6, "XYG3 total gradient mismatch"); + 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 index d7b4ac491..2aebf8535 100644 --- a/tests/grad/udh.rs +++ b/tests/grad/udh.rs @@ -6,7 +6,9 @@ //! 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). The rimatr must be CD-decomposed (see `src/grad/udh.rs`). +//! 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; @@ -29,7 +31,7 @@ static INPUT_NH3: &str = r##" spin_polarization = true [ctrl.j2c_decomp] -policy = "cd" +policy = "POLICY" [ctrl.ri_pt2] new_driver = true @@ -58,7 +60,7 @@ static INPUT_XYG3_U: &str = r##" spin_polarization = true [ctrl.j2c_decomp] -policy = "cd" +policy = "POLICY" [ctrl.ri_pt2] new_driver = true @@ -74,9 +76,8 @@ new_driver = true """ "##; -#[test] -fn test_nh3_ump2_grad() { - let keys = toml::from_str::(&INPUT_NH3[..]).unwrap(); +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); @@ -84,7 +85,7 @@ fn test_nh3_ump2_grad() { // 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: {eng}"); + 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); @@ -101,7 +102,7 @@ fn test_nh3_ump2_grad() { 0.0035622300, -0.0088634616, 0.0015311050, ]; let de_ref = rt::asarray((de_ref, [3, 4])); - println!("DH correlation gradient (REST): {:16.9}", de.t()); + 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()); @@ -110,7 +111,7 @@ fn test_nh3_ump2_grad() { 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() < 1.0e-5); + 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"); @@ -131,10 +132,20 @@ fn test_nh3_ump2_grad() { 0.0149814011, 0.0105624497, 0.0150827880, ]; let de_tot_ref = rt::asarray((de_tot_ref, [3, 4])); - println!("UMP2 total gradient (REST): {:16.9}", de_tot.t()); + 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() < 1.0e-5, "UMP2 total gradient mismatch"); + 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). @@ -145,16 +156,15 @@ fn test_nh3_ump2_grad() { /// 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. -#[test] -fn test_nh3_uxyg3_grad() { - let keys = toml::from_str::(&INPUT_XYG3_U[..]).unwrap(); +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: {eng}"); + 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) @@ -179,7 +189,7 @@ fn test_nh3_uxyg3_grad() { 0.0142281688, 0.0128385585, 0.0146845393, ]; let de_ref = rt::asarray((de_ref, [3, 4])); - println!("UXYG3 total gradient (REST): {:16.9}", de.t()); + 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()); @@ -188,8 +198,18 @@ fn test_nh3_uxyg3_grad() { 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() < 1.0e-6, "UXYG3 total gradient mismatch"); + 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); +} -- Gitee From 2089619d7da2e655b7695740651062748808684f Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 1 Oct 2026 21:59:12 +0800 Subject: [PATCH 21/21] Append the DH response density to the fchk of force evaluations An analytic doubly-hybrid (PT2-family) force job with the fchk output requested (`job_type = "force"`, `outputs = ["fchk"]`) now appends the total relaxed (response) density to `{molecule}.fchk`: `Total MP2 Density` for restricted references, plus `Spin MP2 Density` (alpha - beta) for unrestricted ones - the same density `multipole_rdm1_dump` writes (SCF density plus the relaxed correlation density, i.e. with the Z-vector response in the vir-occ block). No new keyword: the trigger is the fchk output of a force job, as for the multipole dump. The unrestricted fchk additionally gains a real `Total SCF Density` (alpha + beta) and a new `Spin SCF Density` section in place of the previous zero placeholder (ROHF keeps the placeholder). The density is materialized by the DH gradient driver (`RDHGradient/UDHGradient.rdm1_dump`, only when the fchk output is requested) and handed by `eval_force` (return type and structure unchanged) to the transient `SCF::dh_rdm1_resp` for force jobs; `SCF::save_fchk_of_gaussian` writes it right after the MO block via `fchk_append_density_ao`. The multipole dump regenerates the fchk through the two write primitives directly instead of the save, so it does not re-append the DH section (combined runs keep exactly one `Total MP2 Density` section). Validation: the fchk sections match `multipole_rdm1_dump` to 4.6e-8 (Krylov route difference), electron/spin counts Tr(D*S) exact (10.0000000 / 9.0000000 + 1.0000000); gradients with and without the dump agree to 3e-15; G16 `Density=Current` references (RMP2/UMP2/B2PLYP all-electron, def2-SVP) agree to 5.8e-5 / 3.2e-5 / 2.9e-5 max (density-fitting level; 4.6e-6 with a def2-qzvpp-rifit aux); cargo test --release --test test_grad (20 passed) and --test test_analdrv (36 passed); rest_regression DH gradient cases 3/3. Co-authored-by: Claude Code Co-authored-by: deepseek-v4-flash --- README.md | 1 + src/analdrv/multipole/interface.rs | 6 ++- src/grad/rdh.rs | 14 +++++- src/grad/udh.rs | 21 +++++++- src/main_driver.rs | 13 +++++ src/scf_io/fchk.rs | 77 ++++++++++++++++++++++-------- src/scf_io/mod.rs | 7 +++ 7 files changed, 112 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index ed09b2b16..bbe05046d 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 a3b80079e..6eca3ec1b 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/grad/rdh.rs b/src/grad/rdh.rs index 555b40331..fdd74a989 100644 --- a/src/grad/rdh.rs +++ b/src/grad/rdh.rs @@ -77,11 +77,15 @@ pub struct RDHGradient<'a> { 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 } + Self { scf_data, mpi_operator, result: HashMap::new(), energy: 0.0, rdm1_dump: None } } pub fn calc(&mut self) -> &mut Self { @@ -249,6 +253,11 @@ impl<'a> RDHGradient<'a> { 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. @@ -741,7 +750,8 @@ fn de_xc_response_term( de_resp } -/// `[3, natm]` device tensor into a `MatrixFull` of the same shape. +/// 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() diff --git a/src/grad/udh.rs b/src/grad/udh.rs index b7295ed99..8b755380a 100644 --- a/src/grad/udh.rs +++ b/src/grad/udh.rs @@ -74,11 +74,16 @@ pub struct UDHGradient<'a> { 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 } + Self { scf_data, mpi_operator, result: HashMap::new(), energy: 0.0, rdm1_dump: None } } pub fn calc(&mut self) -> &mut Self { @@ -357,6 +362,17 @@ impl<'a> UDHGradient<'a> { (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); [ @@ -896,7 +912,8 @@ impl<'a> UDHGradient<'a> { } } -/// `[3, natm]` device tensor into a `MatrixFull` of the same shape. +/// 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() diff --git a/src/main_driver.rs b/src/main_driver.rs index d75b054af..63dc36e56 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -902,6 +902,10 @@ 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) @@ -948,10 +952,12 @@ fn eval_force(scf_data: &mut SCF, time_mark: &mut utilities::TimeRecords, mpi_op 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))); } } @@ -978,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 { diff --git a/src/scf_io/fchk.rs b/src/scf_io/fchk.rs index 4c5f46c7f..82c890927 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 4b14b8112..4567480e2 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, }; -- Gitee