From b48d9f533fc6ee630e17e79749d8a2df16b6379d Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 10:06:00 +0800 Subject: [PATCH 01/39] port: hessian_backup, numint_matmul, xceff --- Cargo.toml | 3 +- src/dft/mod.rs | 2 + .../docs/get_rho_from_dm_with_output.md | 48 + ...rho_from_homogeneous_braket_with_output.md | 38 + ...t_rho_from_one_bra_mult_ket_with_output.md | 44 + src/dft/numint_matmul/docs/mod.md | 64 + src/dft/numint_matmul/hess_rks.rs | 873 +++++++++++++ src/dft/numint_matmul/hess_uks.rs | 744 +++++++++++ src/dft/numint_matmul/mod.rs | 107 ++ src/dft/numint_matmul/nimatmul.rs | 503 ++++++++ src/dft/numint_matmul/pure_eval_rho.rs | 465 +++++++ src/dft/numint_matmul/pure_xcpot.rs | 1028 +++++++++++++++ src/dft/xceff/flags.rs | 107 ++ src/dft/xceff/libxc-xcfun-trans.md | 131 ++ src/dft/xceff/libxc_wrap.rs | 247 ++++ src/dft/xceff/mod.rs | 24 + src/dft/xceff/xc_deriv.rs | 427 +++++++ src/hessian_backup/cint_handling.rs | 115 ++ src/hessian_backup/hcore.rs | 231 ++++ src/hessian_backup/krylov_block.rs | 258 ++++ src/hessian_backup/mod.rs | 45 + src/hessian_backup/nuc_repl.rs | 80 ++ src/hessian_backup/ovlp.rs | 141 ++ src/hessian_backup/rscf.rs | 506 ++++++++ src/hessian_backup/trait_rhess.rs | 190 +++ src/hessian_backup/trait_uhess.rs | 172 +++ src/hessian_backup/trait_util.rs | 18 + src/hessian_backup/uscf.rs | 540 ++++++++ src/hessian_backup/vib.rs | 1131 +++++++++++++++++ src/lib.rs | 3 + src/ri_jk/prelude_dev.rs | 4 +- src/ri_jk/util.rs | 44 + src/utilities/rstsr_util.rs | 9 + 33 files changed, 8339 insertions(+), 3 deletions(-) create mode 100644 src/dft/numint_matmul/docs/get_rho_from_dm_with_output.md create mode 100644 src/dft/numint_matmul/docs/get_rho_from_homogeneous_braket_with_output.md create mode 100644 src/dft/numint_matmul/docs/get_rho_from_one_bra_mult_ket_with_output.md create mode 100644 src/dft/numint_matmul/docs/mod.md create mode 100644 src/dft/numint_matmul/hess_rks.rs create mode 100644 src/dft/numint_matmul/hess_uks.rs create mode 100644 src/dft/numint_matmul/mod.rs create mode 100644 src/dft/numint_matmul/nimatmul.rs create mode 100644 src/dft/numint_matmul/pure_eval_rho.rs create mode 100644 src/dft/numint_matmul/pure_xcpot.rs create mode 100644 src/dft/xceff/flags.rs create mode 100644 src/dft/xceff/libxc-xcfun-trans.md create mode 100644 src/dft/xceff/libxc_wrap.rs create mode 100644 src/dft/xceff/mod.rs create mode 100644 src/dft/xceff/xc_deriv.rs create mode 100644 src/hessian_backup/cint_handling.rs create mode 100644 src/hessian_backup/hcore.rs create mode 100644 src/hessian_backup/krylov_block.rs create mode 100644 src/hessian_backup/mod.rs create mode 100644 src/hessian_backup/nuc_repl.rs create mode 100644 src/hessian_backup/ovlp.rs create mode 100644 src/hessian_backup/rscf.rs create mode 100644 src/hessian_backup/trait_rhess.rs create mode 100644 src/hessian_backup/trait_uhess.rs create mode 100644 src/hessian_backup/trait_util.rs create mode 100644 src/hessian_backup/uscf.rs create mode 100644 src/hessian_backup/vib.rs diff --git a/Cargo.toml b/Cargo.toml index 76e3ce9469..35a0421d51 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ anyhow = "1" time = "0.1.12" clap = "3.0.0-beta.2" regex = "1.*" -itertools = "0.8" +itertools = "0.15" serde_json = "1.0" serde = {version="1.0",features=["derive"]} serde-inline-default = "1.0" @@ -57,6 +57,7 @@ rstsr-core = { version = "0.7", default-features = false } rstsr-openblas = { version = "0.7", features = ["openmp"] } tblis = { version = "0.2", features = ["dynamic_loading"] } derive_builder = { version = "0.20" } +indexmap = { version = "2.14" } dftd3 = { version = "0.2", optional = true } dftd4 = { version = "0.2", optional = true } diff --git a/src/dft/mod.rs b/src/dft/mod.rs index 31d3e5e79a..9a798acbd7 100644 --- a/src/dft/mod.rs +++ b/src/dft/mod.rs @@ -7,6 +7,8 @@ pub mod xc_deriv; pub mod num_int; pub mod parse_xc; pub mod response; +pub mod numint_matmul; +pub mod xceff; #[cfg(feature = "mpi")] use mpi::collective::SystemOperation; diff --git a/src/dft/numint_matmul/docs/get_rho_from_dm_with_output.md b/src/dft/numint_matmul/docs/get_rho_from_dm_with_output.md new file mode 100644 index 0000000000..ad3026f5c4 --- /dev/null +++ b/src/dft/numint_matmul/docs/get_rho_from_dm_with_output.md @@ -0,0 +1,48 @@ +## Equations and Concepts + +- `RHO`: + + $$ + \rho_g = \sum_{\mu \nu} \phi_{g \mu} D_{\mu \nu} \phi_{g \nu} + $$ + +- `SIGMA`: + + $$ + \begin{aligned} + \sigma_{g t} &= \sum_{\mu \nu} \left( \phi_{g \mu, t} D_{\mu \nu} \phi_{g \nu} + \phi_{g \mu} D_{\mu \nu} \phi_{g \nu, t} \right) + \\\\ + &= 2 \sum_{\mu \nu} \phi_{g \mu} D_{\mu \nu} \phi_{g \nu, t} \quad \text{(symm applied)} + \end{aligned} + $$ + + Note we have already assumed the symmetry of the density matrix, so the two terms in the summation are equal. In the actual code, we only compute one of them and multiply by 2. + +- `TAU`: + + $$ + \tau_g = \frac{1}{2} \sum_{\mu \nu} \sum_{t} \phi_{g \mu, t} D_{\mu \nu} \phi_{g \nu, t} + $$ + +- `LAPL`: + + $$ + \begin{aligned} + \nabla^2 \rho_g + &= \sum_{\mu \nu} \left( \sum_{t} \phi_{g \mu, tt} D_{\mu \nu} \phi_{g \nu} + 2 \sum_{t} \phi_{g \mu, t} D_{\mu \nu} \phi_{g \nu, t} + \sum_{t} \phi_{g \mu} D_{\mu \nu} \phi_{g \nu, tt} \right) + \\\\ + &= 4 \tau_g + 2 \sum_{\mu \nu} \sum_{t} \phi_{g \mu} D_{\mu \nu} \phi_{g \nu, tt} + \quad \text{(symm applied)} + \end{aligned} + $$ + + Similar to `SIGMA`, we have already assumed the symmetry of the density matrix, so the first and third terms in the summation are equal. In the actual code, we only compute one of them and multiply by 2. + +## Usage tips + +For MatMul driver, we are not going to exploit the sparsity of atomic grids or density matrix. It is better to use `braket` versions to exploit the low-rank nature of the density matrix. + +Usually, in only the following three cases, use this function: +- density comes from post-SCF reduced density matrix, which is usually non-zero-definite. +- prototype validation. +- basis set is small, number of occupation is large. diff --git a/src/dft/numint_matmul/docs/get_rho_from_homogeneous_braket_with_output.md b/src/dft/numint_matmul/docs/get_rho_from_homogeneous_braket_with_output.md new file mode 100644 index 0000000000..1c8360b94a --- /dev/null +++ b/src/dft/numint_matmul/docs/get_rho_from_homogeneous_braket_with_output.md @@ -0,0 +1,38 @@ +# Equations and Concepts + +Please note that since we assumed the homogeneous braket form, the density matrix is symmetric, and some factors of two is applicable to `SIGMA` and `LAPL`. + +- Orbital grids are defined as + + $$ + \begin{aligned} + \phi_{g i} &= \sum_{\mu} \phi_{g \mu} C_{\mu i} + \\\\ + \phi_{g i, t} &= \sum_{\mu} \phi_{g \mu, t} C_{\mu i} + \end{aligned} + $$ + +- `RHO`: + + $$ + \rho_g = \sum_{i} \phi_{g i}^2 + $$ + +- `SIGMA`: + + $$ + \sigma_{g t} = 2 \sum_{i} \phi_{g i} \phi_{g i, t} + $$ + +- `TAU`: + + $$ + \tau_g = \frac{1}{2} \sum_{i} \sum_{t} \phi_{g i, t}^2 + $$ + +- `LAPL`: + + $$ + \nabla^2 \rho_g + = 4 \tau_g + 2 \sum_{i} \sum_{t} \phi_{g i} \phi_{g i, tt} + $$ diff --git a/src/dft/numint_matmul/docs/get_rho_from_one_bra_mult_ket_with_output.md b/src/dft/numint_matmul/docs/get_rho_from_one_bra_mult_ket_with_output.md new file mode 100644 index 0000000000..f8bbd3f5f4 --- /dev/null +++ b/src/dft/numint_matmul/docs/get_rho_from_one_bra_mult_ket_with_output.md @@ -0,0 +1,44 @@ +# Equations and Concepts + +Since bra and ket may differ, the density matrix $D_{\mu\nu} = \sum_i \mathrm{bra}_{\mu i} \mathrm{ket}_{\nu i}$ is not necessarily symmetric, so both cross terms must be computed explicitly. + +- Orbital grids are defined as + + $$ + \begin{aligned} + \phi_{g i}^{\mathrm{bra}} &= \sum_{\mu} \phi_{g \mu} \mathrm{bra}_{\mu i} + \\\\ + \phi_{g i}^{\mathrm{ket}} &= \sum_{\mu} \phi_{g \mu} \mathrm{ket}_{\mu i} + \\\\ + \phi_{g i, t}^{\mathrm{bra}} &= \sum_{\mu} \phi_{g \mu, t} \mathrm{bra}_{\mu i} + \\\\ + \phi_{g i, t}^{\mathrm{ket}} &= \sum_{\mu} \phi_{g \mu, t} \mathrm{ket}_{\mu i} + \end{aligned} + $$ + +- `RHO`: + + $$ + \rho_g = \sum_{i} \phi_{g i}^{\mathrm{bra}} \phi_{g i}^{\mathrm{ket}} + $$ + +- `SIGMA`: + + $$ + \sigma_{g t} = \sum_{i} \left( \phi_{g i}^{\mathrm{bra}} \phi_{g i, t}^{\mathrm{ket}} + \phi_{g i, t}^{\mathrm{bra}} \phi_{g i}^{\mathrm{ket}} \right) + $$ + + Unlike the symmetric case, the two terms are not equal in general so both must be computed. + +- `TAU`: + + $$ + \tau_g = \frac{1}{2} \sum_{i} \sum_{t} \phi_{g i, t}^{\mathrm{bra}} \phi_{g i, t}^{\mathrm{ket}} + $$ + +- `LAPL`: + + $$ + \nabla^2 \rho_g + = 4 \tau_g + \sum_{i} \sum_{t} \left( \phi_{g i}^{\mathrm{bra}} \phi_{g i, tt}^{\mathrm{ket}} + \phi_{g i, tt}^{\mathrm{bra}} \phi_{g i}^{\mathrm{ket}} \right) + $$ diff --git a/src/dft/numint_matmul/docs/mod.md b/src/dft/numint_matmul/docs/mod.md new file mode 100644 index 0000000000..dfd76a9a0c --- /dev/null +++ b/src/dft/numint_matmul/docs/mod.md @@ -0,0 +1,64 @@ +# Introduction of `MatMul` DFT grid driver + +This driver is the naive DFT driver. The code effort is minimized, yet still be efficient to small systems. + +As the name suggests, we fully utilize BLAS3 GEMM to perform the DFT grid operations. + +## `ao`: atomic orbital values + +The AO tensor `ao` has shape `[ngrids, nao, ncomp]` where the component dimension is ordered as: + +| Index | Deriv | Component | +|--|--|--| +| 0 | 0 | $\phi_{g \mu}$ | +| 1–3 | 1 | $\phi_{g \mu, x}, \phi_{g \mu, y}, \phi_{g \mu, z}$ | +| 4–9 | 2 | $\phi_{g \mu,xx}, \phi_{g \mu,xy}, \phi_{g \mu,xz}, \phi_{g \mu,yy}, \phi_{g \mu,yz}, \phi_{g \mu,zz}$ | + +Additional to the above table, for notation simplicity, we denote + +- orbital notation + + $$ + \phi_{g \mu} = \phi_{\mu} (\bm{r}_g) + $$ + + Where subscript $g$ denotes grid, $\bm{r}_g$ denotes the coordinate of grid point $g$, and $\mu$ denotes the AO index. + +- orbital derivative notation + + $$ + \phi_{g \mu, x} = \partial_x \phi_{\mu} (\bm{r}_g) + $$ + + The partial derivative is taken with respect to the electron coordinate. + +- We usually use $t, s, r \in \{ x, y, z \}$ to denote the subscript of the coordinate. + +## `rho`: density values + +We will introduce 4 types of density values: + +| property | [`RHO`] | [`SIGMA`] | [`TAU`] | [`LAPL`] | +|--|--|--|--|--| +| notation | $\rho$ | $\sigma$ | $\tau$ | $\nabla^2 \rho$ | +| [`num_rho_comp`](NIDenType::num_rho_comp) | 1 | 4 | 5 | 6 | +| [`num_ao_deriv`](NIDenType::num_ao_deriv) | 0 | 1 | 1 | 2 | +| [`num_ao_comp`](NIDenType::num_ao_comp) | 1 | 4 | 4 | 10 | +| usual XC type | LDA | GGA | mGGA | mGGA | + +The density to be input in this driver is `[ngirds, num_rho_comp]` for spin unpolarized case, and `[ngrids, num_rho_comp, 2]` for spin polarized case. The density is computed by contracting the AO tensor with the density matrix, which is the same as the usual DFT grid driver. + +The density components are ordered as: + +$$ +\rho, \rho_x, \rho_y, \rho_z, \tau, \nabla^2 \rho +$$ + +For laplacian density functional, whatever the xc functional uses tau, currently we enforce the evaluation of $\tau$ to be the 4th component of the density. + +## Important pure functions + +**Density Grid Evaluation** + +- [`get_rho_from_dm_with_output`] +- [`get_rho_from_homogeneous_braket_with_output] diff --git a/src/dft/numint_matmul/hess_rks.rs b/src/dft/numint_matmul/hess_rks.rs new file mode 100644 index 0000000000..556ddd9dfa --- /dev/null +++ b/src/dft/numint_matmul/hess_rks.rs @@ -0,0 +1,873 @@ +// see also pyhessref/nimatmul/rks.py + +use super::prelude::*; +use crate::hessian_backup::prelude::*; + +use XCDenType::*; + +/* #region const dimensions/indices definition */ + +const O: usize = 0; +const X: usize = 1; +const Y: usize = 2; +const Z: usize = 3; +const XX: usize = 4; +const XY: usize = 5; +const XZ: usize = 6; +const YX: usize = 5; +const YY: usize = 7; +const YZ: usize = 8; +const ZX: usize = 6; +const ZY: usize = 8; +const ZZ: usize = 9; +const XXX: usize = 10; +const XXY: usize = 11; +const XXZ: usize = 12; +const XYY: usize = 13; +const XYZ: usize = 14; +const XZZ: usize = 15; +const YYY: usize = 16; +const YYZ: usize = 17; +const YZZ: usize = 18; +const ZZZ: usize = 19; + +const IDX_AO_DERIV2: [[usize; 3]; 3] = [[XX, XY, XZ], [XY, YY, YZ], [XZ, YZ, ZZ]]; + +pub const fn get_hess_ao_deriv(xc_type: XCDenType) -> usize { + match xc_type { + RHO => 2, + SIGMA => 3, + TAU => 3, + LAPL => unimplemented!(), + } +} + +pub const fn get_hess_ncomp_ao_dm0(xc_type: XCDenType) -> usize { + match xc_type { + RHO => 1, + SIGMA => 4, + TAU => 4, + LAPL => unimplemented!(), + } +} + +/* #endregion */ + +/* #region macro for indexing last dimension */ + +macro_rules! index { + ($tsr: ident, $($idx:expr),*) => { + $tsr.i((Ellipsis, $($idx),*)) + }; +} + +macro_rules! index_mut { + ($tsr: ident, $($idx:expr),*) => { + (*&mut $tsr.i_mut((Ellipsis, $($idx),*))) + }; +} + +/* #endregion */ + +/* #region basic pure functions of skeleton hessian evaluation */ + +pub fn make_hessian_setup( + mol: &CInt, + xc_func_list: &[(f64, LibXCFunctional)], + ni: &mut NIMatmul, + dm0: TsrView, + atm_list: Option<&[usize]>, +) -> (HashMap<&'static str, Tsr>, IndexMap<&'static str, f64>) { + assert!(!xc_func_list.is_empty(), "xc_func_list must not be empty"); + let atm_list = atm_list.map_or_else(|| (0..mol.natm()).collect_vec(), |lst| lst.to_vec()); + // ao slices indexed by `atm_list` + let aoslices_full = mol.aoslice_by_atom(); + let aoslices = atm_list.iter().map(|&iatm| aoslices_full[iatm]).collect_vec(); + // overall xc_type is the strictest (max nvar) one across the functionals; + let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); + + let device = dm0.device().clone(); + let weights = rt::asarray((ni.weights.clone(), &device)); + + 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); + }; + + // --- ao, rho, vxc, fxc --- // + + // ao [ngrids, nao, ncomp] + // ao_dm0 [ngrids, nao, ncomp_ao_dm0] + // rho [ngrids, nvar] + // vxc [ngrids, nvar] + // fxc [ngrids, nvar, nvar] + + let t0 = std::time::Instant::now(); + let ao = ni.get_cached_ao(get_hess_ao_deriv(xc_type)); + tic("ao", t0); + + let t0 = std::time::Instant::now(); + let ncomp_ao_dm0 = get_hess_ncomp_ao_dm0(xc_type); + let ao_dm0 = index!(ao, ..ncomp_ao_dm0) % &dm0; + tic("ao_dm0", t0); + + let t0 = std::time::Instant::now(); + let (rho, vxc, fxc) = get_rho_vxc_fxc(xc_func_list, ao.view(), ao_dm0.view()); + let wv = &weights * &vxc; + let wf = &weights * &fxc; + tic("rho, vxc, fxc", t0); + + // --- drho --- // + + // drho [ngrids, nvar, 3, natm] + let t0 = std::time::Instant::now(); + let drho = get_drho(xc_type, ao.view(), ao_dm0.view(), &aoslices); + tic("drho", t0); + + // --- de_fxc --- // + + // de_fxc [3, 3, natm, natm] + let t0 = std::time::Instant::now(); + let de_fxc = get_de_fxc(wf.view(), drho.view()); + tic("de_fxc", t0); + + // --- de_vxc_diag --- // + + // de_vxc_diag [3, 3, natm, natm] + let t0 = std::time::Instant::now(); + let de_vxc_diag = get_de_vxc_diag(xc_type, ao.view(), ao_dm0.view(), wv.view(), &aoslices); + tic("de_vxc_diag", t0); + + // --- de_vxc_off --- // + + // de_vxc_off [3, 3, natm, natm] + let t0 = std::time::Instant::now(); + let de_vxc_off = get_de_vxc_off(xc_type, ao.view(), dm0.view(), wv.view(), &aoslices); + tic("de_vxc_off", t0); + + // --- vmat_ip --- // + + // vmat_ip [nao, nao, 3] + let t0 = std::time::Instant::now(); + let vmat_ip = get_vmat_ip(xc_type, ao.view(), wv.view()); + tic("vmat_ip", t0); + + // --- vmat_deriv1 --- // + + // vmat_deriv1 [nao, nao, 3, natm] + let t0 = std::time::Instant::now(); + let vmat_deriv1 = get_vmat_deriv1(xc_type, ao.view(), drho.view(), wf.view(), vmat_ip.view(), &aoslices); + tic("vmat_deriv1", t0); + + let result = HashMap::from([ + ("rho", rho), + ("vxc", vxc), + ("fxc", fxc), + ("de_fxc", de_fxc), + ("de_vxc_diag", de_vxc_diag), + ("de_vxc_off", de_vxc_off), + ("vmat_ip", vmat_ip), + ("vmat_deriv1", vmat_deriv1), + ]); + (result, timing) +} + +pub fn get_rho_vxc_fxc(xc_func_list: &[(f64, LibXCFunctional)], ao: TsrView, ao_dm0: TsrView) -> (Tsr, Tsr, Tsr) { + // see also pyhessref/nimatmul/rks.py + + assert!(!xc_func_list.is_empty(), "xc_func_list must not be empty"); + // overall xc_type is the strictest one across the functionals; partial + // contributions from looser families are added into the leading slice. + 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 = ao.shape()[0]; + let device = ao.device().clone(); + + let mut rho = rt::zeros(([ngrids, nvar], &device)); + index_mut!(rho, 0) += rt::vecdot(index!(ao, 0), index!(ao_dm0, O), 1); + if matches!(xc_type, SIGMA | TAU) { + index_mut!(rho, X) += 2 * rt::vecdot(index!(ao, X), index!(ao_dm0, O), 1); + index_mut!(rho, Y) += 2 * rt::vecdot(index!(ao, Y), index!(ao_dm0, O), 1); + index_mut!(rho, Z) += 2 * rt::vecdot(index!(ao, Z), index!(ao_dm0, O), 1); + } + if matches!(xc_type, TAU) { + index_mut!(rho, 4) += 0.5 + * (rt::vecdot(index!(ao, X), index!(ao_dm0, X), 1) + + rt::vecdot(index!(ao, Y), index!(ao_dm0, Y), 1) + + rt::vecdot(index!(ao, Z), index!(ao_dm0, Z), 1)) + } + + let mut vxc = rt::zeros(([ngrids, nvar], &device)); + let mut fxc = rt::zeros(([ngrids, nvar, nvar], &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(); + // each sub-functional consumes only the leading `nvar_i` rho components. + 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(); + // accumulate into the leading slice of the (possibly larger) global tensors. + *&mut vxc.i_mut((.., ..nvar_i)) += *scale * vxc_i; + *&mut fxc.i_mut((.., ..nvar_i, ..nvar_i)) += *scale * fxc_i; + } + + (rho, vxc, fxc) +} + +pub fn get_drho(xc_type: XCDenType, ao: TsrView, ao_dm0: TsrView, aoslices: &[[usize; 4]]) -> Tsr { + // direct transformation of `pyhessref/nimatmul/rks.py`, function `_make_drho` + + let ngrids = ao.shape()[0]; + let nvar = xc_type.num_nvar(); + let natm = aoslices.len(); + let device = ao.device().clone(); + + let mut drho = rt::zeros(([ngrids, nvar, 3, natm], &device)); + + // components: [rho_var, t_direction, cbra, cket] + // tuple that contribute to each rho component + // for symmetric components, result is multiplied by 2 at the end + + // RHO part + let mut components = vec![(0, 0, X, O), (0, 1, Y, O), (0, 2, Z, O)]; + // SIGMA part + if matches!(xc_type, SIGMA | TAU) { + // bra deriv2, ket 0 + let sigma_bra2_ket0 = [ + [(1, 0, XX, O), (2, 0, XY, O), (3, 0, XZ, O)], + [(1, 1, YX, O), (2, 1, YY, O), (3, 1, YZ, O)], + [(1, 2, ZX, O), (2, 2, ZY, O), (3, 2, ZZ, O)], + ]; + components.extend(sigma_bra2_ket0.concat()); + // bra deriv1, ket deriv1 + let sigma_bra1_ket1 = [ + [(1, 0, X, X), (2, 0, X, Y), (3, 0, X, Z)], + [(1, 1, Y, X), (2, 1, Y, Y), (3, 1, Y, Z)], + [(1, 2, Z, X), (2, 2, Z, Y), (3, 2, Z, Z)], + ]; + components.extend(sigma_bra1_ket1.concat()); + } + // TAU part + if matches!(xc_type, TAU) { + // bra deriv2, ket deriv1 + let tau_bra2_ket1 = [ + [(4, 0, XX, X), (4, 0, XY, Y), (4, 0, XZ, Z)], + [(4, 1, YX, X), (4, 1, YY, Y), (4, 1, YZ, Z)], + [(4, 2, ZX, X), (4, 2, ZY, Y), (4, 2, ZZ, Z)], + ]; + components.extend(tau_bra2_ket1.concat()); + } + + for (A, &[_, _, p0, p1]) in aoslices.iter().enumerate() { + let slc = rt::slice!(p0, p1); + for &(v, t, cbra, cket) in &components { + *&mut drho.i_mut((.., v, t, A)) -= rt::vecdot(ao.i((.., slc, cbra)), ao_dm0.i((.., slc, cket)), 1); + } + } + + // RHO and SIGMA part multiply factor 2 + // TAU part does not multiply factor 2 because of the 0.5 factor in rho_tau + match xc_type { + RHO => *&mut drho.i_mut((.., 0..1)) *= 2.0, + SIGMA | TAU => *&mut drho.i_mut((.., 0..4)) *= 2.0, + LAPL => unimplemented!(), + } + drho +} + +pub fn get_de_fxc(wf: TsrView, drho: TsrView) -> Tsr { + // wf , drho, drho -> de_fxc + // gxy, gxtA, gysB -> tsAB + + let [ngrids, nvar, _, natm] = drho.shape().iter().cloned().collect_array().unwrap(); + + // wf * drho -> tmp + // gxy.. * gx.tA -> gytA + let tmp1 = rt::vecdot(wf.i((.., .., .., None, None)), drho.i((.., .., None, .., ..)), 1); + + // tmp1 * drho -> tmp2 + // gytA * gysB -> tAsB + let tmp1 = tmp1.reshape([ngrids * nvar, natm * 3]); + let drho = drho.reshape([ngrids * nvar, natm * 3]); + let tmp2 = tmp1.t() % drho; + + // transpose tmp2 to get de_fxc + // tAsB -> tsAB + tmp2.reshape([3, natm, 3, natm]).transpose([0, 2, 1, 3]).into_contig(ColMajor) +} + +pub fn get_de_vxc_diag(xc_type: XCDenType, ao: TsrView, ao_dm0: TsrView, wv: TsrView, aoslices: &[[usize; 4]]) -> Tsr { + const TRIPLE_SIGMA_DIAG: [[usize; 3]; 6] = + [[XXX, XXY, XXZ], [XXY, XYY, XYZ], [XXZ, XYZ, XZZ], [XYY, YYY, YYZ], [XYZ, YYZ, YZZ], [XZZ, YZZ, ZZZ]]; + const TRIPLE_TAU_DIAG: [[usize; 6]; 3] = + [[XXX, XXY, XXZ, XYY, XYZ, XZZ], [XXY, XYY, XYZ, YYY, YYZ, YZZ], [XXZ, XYZ, XZZ, YYZ, YZZ, ZZZ]]; + + let natm = aoslices.len(); + let nao = ao.shape()[1]; + let device = ao.device().clone(); + + let mut dao_vxc_diag: Tsr = rt::zeros(([nao, 6], &device)); + + // contribution 1: lda/gga ao deriv 2 + let mut aow = index!(ao_dm0, O) * index!(wv, 0); + if matches!(xc_type, SIGMA | TAU) { + aow += index!(ao_dm0, X) * index!(wv, X); + aow += index!(ao_dm0, Y) * index!(wv, Y); + aow += index!(ao_dm0, Z) * index!(wv, Z); + } + for (idx_ts, its) in [XX, XY, XZ, YY, YZ, ZZ].into_iter().enumerate() { + index_mut!(dao_vxc_diag, idx_ts) += 2 * rt::vecdot(index!(ao, its), &aow, 0); + } + + // contribution 2: gga ao deriv 3 + if matches!(xc_type, SIGMA | TAU) { + for (idx_ts, &[i3x, i3y, i3z]) in TRIPLE_SIGMA_DIAG.iter().enumerate() { + let aow = + index!(ao, i3x) * index!(wv, X) + index!(ao, i3y) * index!(wv, Y) + index!(ao, i3z) * index!(wv, Z); + index_mut!(dao_vxc_diag, idx_ts) += 2 * rt::vecdot(&aow, index!(ao_dm0, O), 0); + } + } + + // contribution 3: tau ao deriv 3 + if matches!(xc_type, TAU) { + for (r, &idx_tri) in TRIPLE_TAU_DIAG.iter().enumerate() { + let aow = index!(ao_dm0, r + 1) * index!(wv, 4); + for (idx_ts, &i3) in idx_tri.iter().enumerate() { + index_mut!(dao_vxc_diag, idx_ts) += rt::vecdot(index!(ao, i3), &aow, 0); + } + } + } + + // reduce ao-wise contributions to atom-wise contributions + let mut de_vxc_diag = rt::zeros(([6, natm, natm], &device)); + for (A, &[_, _, p0, p1]) in aoslices.iter().enumerate() { + let slc = rt::slice!(p0, p1); + de_vxc_diag.i_mut((.., A, A)).assign(dao_vxc_diag.i(slc).sum_axes(0)); + } + // symmetrize and expand de_vxc_diag from [6, natm, natm] to [3, 3, natm, natm] + de_vxc_diag.index_select(0, [0, 1, 2, 1, 3, 4, 2, 4, 5]).into_shape([3, 3, natm, natm]) +} + +pub fn get_de_vxc_off(xc_type: XCDenType, ao: TsrView, dm0: TsrView, wv: TsrView, aoslices: &[[usize; 4]]) -> Tsr { + let natm = aoslices.len(); + let nao = ao.shape()[1]; + let device = ao.device().clone(); + + let mut dao_vxc_off: Tsr = rt::zeros(([nao, nao, 3, 3], &device)); + + if matches!(xc_type, RHO) { + for t in 0..3 { + let aowv = index!(wv, 0) * index!(ao, t + 1); + for s in 0..3 { + index_mut!(dao_vxc_off, t, s).matmul_from(aowv.t(), index!(ao, s + 1), 1.0, 1.0); + } + } + } + + if matches!(xc_type, SIGMA | TAU) { + for t in 0..3 { + let mut aowv: Tsr = 0.5 * index!(wv, 0) * index!(ao, t + 1); + for r in 0..3 { + aowv += index!(wv, r + 1) * index!(ao, IDX_AO_DERIV2[t][r]); + } + for s in 0..3 { + index_mut!(dao_vxc_off, t, s).matmul_from(aowv.t(), index!(ao, s + 1), 2.0, 1.0); + } + } + } + + if matches!(xc_type, TAU) { + let mut dao_vxc_tau: Tsr = rt::zeros(([nao, nao, 3, 3], &device)); + for r in 0..3 { + for t in 0..3 { + let aowv: Tsr = 0.5 * index!(wv, 4) * index!(ao, IDX_AO_DERIV2[t][r]); + for s in 0..t + 1 { + index_mut!(dao_vxc_tau, t, s).matmul_from(aowv.t(), index!(ao, IDX_AO_DERIV2[s][r]), 1.0, 1.0); + } + } + } + + for t in 0..3 { + for s in 0..t + 1 { + index_mut!(dao_vxc_off, t, s) += &index!(dao_vxc_tau, t, s); + } + for s in 0..t { + index_mut!(dao_vxc_off, s, t) += &index!(dao_vxc_tau, t, s).t(); + } + } + } + + // add transposition + let dao_vxc_off = &dao_vxc_off + dao_vxc_off.transpose([1, 0, 3, 2]); + + // reduce ao-wise contributions to atom-wise contributions + let mut de_vxc_off = rt::zeros(([3, 3, natm, natm], &device)); + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + for (B, &[_, _, p0B, p1B]) in aoslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + let contrib = rt::vecdot(dao_vxc_off.i((slcA, slcB)), dm0.i((slcA, slcB)), ([0, 1], [0, 1])); + de_vxc_off.i_mut((.., .., A, B)).assign(&contrib); + de_vxc_off.i_mut((.., .., B, A)).assign(contrib.t()); + } + } + + de_vxc_off +} + +pub fn get_vmat_ip(xc_type: XCDenType, ao: TsrView, wv: TsrView) -> Tsr { + // direct transformation of `pyhessref/nimatmul/rks.py`, function `_vmat_ip` + + let nao = ao.shape()[1]; + let device = ao.device().clone(); + + let mut vmat_ip = rt::zeros(([nao, nao, 3], &device)); + + if matches!(xc_type, RHO) { + // bra-on-A and ket-on-A halves are identical for LDA + // (both equal 0.5 * wv[0] * ao[t+1]^T @ ao[O]); folded into a single contraction. + let aow: Tsr = index!(wv, 0) * index!(ao, O); + for t in 0..3 { + index_mut!(vmat_ip, t).matmul_from(&index!(ao, t + 1).t(), &aow, 1.0, 1.0); + } + return vmat_ip; + } + + // GGA + MGGA share the same SIGMA structure + if matches!(xc_type, SIGMA | TAU) { + let mut aow: Tsr = 0.5 * index!(wv, 0) * index!(ao, O); + for r in 0..3 { + aow += index!(wv, r + 1) * index!(ao, r + 1); + } + for t in 0..3 { + index_mut!(vmat_ip, t).matmul_from(&index!(ao, t + 1).t(), &aow, 1.0, 1.0); + } + + for t in 0..3 { + let mut aow_d: Tsr = 0.5 * index!(wv, 0) * index!(ao, t + 1); + for r in 0..3 { + aow_d += index!(wv, r + 1) * index!(ao, IDX_AO_DERIV2[t][r]); + } + index_mut!(vmat_ip, t).matmul_from(&aow_d.t(), &index!(ao, O), 1.0, 1.0); + } + } + + // MGGA tau channel + if matches!(xc_type, TAU) { + for r in 0..3 { + let aow: Tsr = 0.5 * index!(wv, 4) * index!(ao, r + 1); + for t in 0..3 { + index_mut!(vmat_ip, t).matmul_from(&index!(ao, IDX_AO_DERIV2[t][r]).t(), &aow, 1.0, 1.0); + } + } + } + + vmat_ip +} + +pub fn get_vmat_deriv1( + xc_type: XCDenType, + ao: TsrView, + drho: TsrView, + wf: TsrView, + vmat_ip: TsrView, + aoslices: &[[usize; 4]], +) -> Tsr { + let natm = aoslices.len(); + let nao = ao.shape()[1]; + + let mut vmat_deriv1: Tsr = rt::zeros(([nao, nao, 3, natm], ao.device())); + + for (A, &[_, _, p0, p1]) in aoslices.iter().enumerate() { + let slc = rt::slice!(p0, p1); + + if matches!(xc_type, RHO) { + // wf_rho: [ngrids, 3] + let wf_rho: Tsr = 0.5 * index!(wf, O, O) * drho.i((.., O, .., A)); + for t in 0..3 { + let aow = index!(wf_rho, t) * index!(ao, O); + index_mut!(vmat_deriv1, t, A).matmul_from(aow.t(), index!(ao, O), 1.0, 1.0); + } + } + + if matches!(xc_type, SIGMA | TAU) { + // wf_rho = np.einsum("gxy, gxt -> gyt", wf, drho[A]) + // wf_rho[:, 0] *= 0.5, wf_rho[:, 4] *= 0.25 + let mut wf_rho = rt::vecdot(&wf, drho.i((.., .., None, .., A)), 1); + *&mut wf_rho.i_mut((.., 0)) *= 0.5; + if matches!(xc_type, TAU) { + *&mut wf_rho.i_mut((.., 4)) *= 0.25; + } + for t in 0..3 { + let aow = rt::vecdot(wf_rho.i((.., None, ..4, t)), ao.i((.., .., ..4)), 2); + index_mut!(vmat_deriv1, t, A).matmul_from(aow.t(), index!(ao, O), 1.0, 1.0); + } + + if matches!(xc_type, TAU) { + for r in [X, Y, Z] { + for t in 0..3 { + let aow = wf_rho.i((.., 4, t)) * index!(ao, r); + index_mut!(vmat_deriv1, t, A).matmul_from(aow.t(), index!(ao, r), 1.0, 1.0); + } + } + } + } + + *&mut vmat_deriv1.i_mut((slc, .., .., A)) -= vmat_ip.i((slc, .., ..)); + } + + &vmat_deriv1 + vmat_deriv1.swapaxes(0, 1) +} + +/* #endregion */ + +/* #region response */ + +pub fn get_rks_response_bra( + ni: &mut NIMatmul, + den_type: XCDenType, + fxc_eff: TsrView, + mo1_bra: TsrView, + mocc: TsrView, +) -> (Tsr, IndexMap<&'static str, f64>) { + let nao = mo1_bra.shape()[0]; + let nocc = mo1_bra.shape()[1]; + let mo1_bra_shape = mo1_bra.shape().to_vec(); + let mo1_bra = mo1_bra.reshape((nao, nocc, -1)); + let mo1_bra_list = mo1_bra.axes_iter(-1).collect_vec(); + + 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); + + let t0 = std::time::Instant::now(); + let rho1 = ni.make_rho_from_one_bra_mult_ket(mocc.view(), &mo1_bra_list, den_type); + tic("rho1", t0); + + let t0 = std::time::Instant::now(); + let resp = ni.make_rks_fxc_pot_with_eff_bra_trans(fxc_eff, rho1.view(), mocc.view(), den_type); + tic("resp", t0); + + // The 4.0 times is a trick of closed-shell coefficient + let resp = 4.0 * resp.into_shape(mo1_bra_shape); + (resp, timing) +} + +/* #endregion */ + +/* #region parallel/batch wrapper */ + +pub fn make_hessian_setup_batched( + mol: &CInt, + xc_func_list: &[(f64, LibXCFunctional)], + ni: &mut NIMatmul, + dm0: TsrView, + atm_list: Option<&[usize]>, + verbose: bool, +) -> (HashMap<&'static str, Tsr>, IndexMap<&'static str, f64>) { + // batch for grids + // - except for rho, vxc, fxc; other tensors can be added (reduced) + // - rho, vxc, fxc requires concation + + // outer iter: batch by nbatch (limit memory usage) + // inner iter: batch by nchunk (for parallel) + + let ngrids = ni.weights.len(); + let nbatch = ni.nbatch; + let nchunk = ni.nchunk; + let device = dm0.device().clone(); + let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); + let nvar = xc_type.num_nvar(); + let deriv_level = get_hess_ao_deriv(xc_type); + let natm = atm_list.map_or_else(|| mol.natm(), |lst| lst.len()); + let nao = mol.nao(); + + let rho: Tsr = rt::zeros(([ngrids, nvar], &device)); + let vxc: Tsr = rt::zeros(([ngrids, nvar], &device)); + let fxc: 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)); + let vmat_ip: Tsr = rt::zeros(([nao, nao, 3], &device)); + let vmat_deriv1: Tsr = rt::zeros(([nao, nao, 3, natm], &device)); + + let timing = Arc::new(Mutex::new(IndexMap::from([ + ("ao", 0.0), + ("ao_dm0", 0.0), + ("rho, vxc, fxc", 0.0), + ("drho", 0.0), + ("de_fxc", 0.0), + ("de_vxc_diag", 0.0), + ("de_vxc_off", 0.0), + ("vmat_ip", 0.0), + ("vmat_deriv1", 0.0), + ("total", 0.0), + ]))); + let time_total = std::time::Instant::now(); + + // atomic guard to avoid racing write + let guard = Mutex::new(()); + + for start_batch in (0..ngrids).step_by(nbatch) { + let end_batch = (start_batch + nbatch).min(ngrids); + + // handle AO integral at batch level + let t0 = std::time::Instant::now(); + let mut ni_batch = ni.split_batch(start_batch, end_batch); + ni_batch.get_cached_ao(deriv_level); + { + let mut timing = timing.lock().unwrap(); + timing["ao"] += t0.elapsed().as_secs_f64(); + } + + // other parts can be parallelized by chunk, with atomic guard for reduction + (start_batch..end_batch).into_par_iter().step_by(nchunk).for_each(|start| { + let end = (start + nchunk).min(end_batch); + let mut ni_chunk = ni_batch.split_batch(start - start_batch, end - start_batch); + let (result_chunk, timing_chunk) = + make_hessian_setup(mol, xc_func_list, &mut ni_chunk, dm0.view(), atm_list); + // fill rho, vxc, fxc + // this is assumed to be not racing, so no guard at here + unsafe { + let rho_slc = rho.i(start..end); + let vxc_slc = vxc.i(start..end); + let fxc_slc = fxc.i(start..end); + let mut rho_slc = rho_slc.force_mut(); + let mut vxc_slc = vxc_slc.force_mut(); + let mut fxc_slc = fxc_slc.force_mut(); + rho_slc.assign(&result_chunk["rho"]); + vxc_slc.assign(&result_chunk["vxc"]); + fxc_slc.assign(&result_chunk["fxc"]); + } + // add up other tensors + unsafe { + let lock = guard.lock().unwrap(); + *&mut de_fxc.force_mut() += &result_chunk["de_fxc"]; + *&mut de_vxc_diag.force_mut() += &result_chunk["de_vxc_diag"]; + *&mut de_vxc_off.force_mut() += &result_chunk["de_vxc_off"]; + *&mut vmat_ip.force_mut() += &result_chunk["vmat_ip"]; + *&mut vmat_deriv1.force_mut() += &result_chunk["vmat_deriv1"]; + drop(lock); + } + // add up timing + { + let mut timing = timing.lock().unwrap(); + for (key, value) in timing_chunk { + *timing.get_mut(key).unwrap() += value; + } + } + }); + + // fill total time outside parallel loop + { + let mut timing = timing.lock().unwrap(); + timing.insert("total", time_total.elapsed().as_secs_f64()); + } + + // verbose print of timing and batch info + if verbose { + let timing = timing.lock().unwrap(); + println!("In make_hessian_setup_batched, Batch {start_batch}..{end_batch}"); + println!(" Elapsed time from start (Wall time): {:.4} sec", timing["total"]); + } + } + + let timing = timing.lock().unwrap(); + if verbose { + println!("Finished make_hessian_setup_batched"); + println!(" Total elapsed time (Wall time): {:.4} sec", timing["total"]); + println!(" Timing breakdown (CPU time for others, Wall time for `ao`):"); + for (key, value) in timing.iter() { + if *key != "total" { + println!(" {key:>20}: {value:.4} sec"); + } + } + } + + let result = HashMap::from([ + ("rho", rho), + ("vxc", vxc), + ("fxc", fxc), + ("de_fxc", de_fxc), + ("de_vxc_diag", de_vxc_diag), + ("de_vxc_off", de_vxc_off), + ("vmat_ip", vmat_ip), + ("vmat_deriv1", vmat_deriv1), + ]); + + (result, timing.clone()) +} + +pub fn get_rks_response_bra_batched( + ni: &mut NIMatmul, + den_type: XCDenType, + fxc_eff: TsrView, + mo1_bra: TsrView, + mocc: TsrView, + verbose: bool, +) -> (Tsr, IndexMap<&'static str, f64>) { + let ngrids = ni.weights.len(); + let nbatch = ni.nbatch; + let mo1_bra_shape = mo1_bra.shape().to_vec(); + let device = mo1_bra.device().clone(); + let mut resp = rt::zeros((mo1_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, timing_batch) = + get_rks_response_bra(&mut ni_batch, den_type, fxc_eff.i(start..end), mo1_bra.view(), mocc.view()); + 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_rks_response_bra_batched, Batch {start}..{end}"); + println!(" Elapsed time from start (Wall time): {:.4} sec", duration); + } + } + + if verbose { + println!("Finished get_rks_response_bra_batched"); + println!(" Total elapsed time (Wall time): {:.4} sec", timing["total"]); + println!(" Timing breakdown (Wall time):"); + for (key, value) in timing.iter() { + if *key != "total" { + println!(" {key:>20}: {value:.4} sec"); + } + } + } + + (resp, timing) +} + +/* #endregion */ + +/* #region final implementation of RKS Hessian */ + +pub struct RHessKSNIMatmul<'a> { + pub mol: CInt, + pub xc_func_list: &'a [(f64, LibXCFunctional)], + pub ni: NIMatmul<'a>, + pub ni_cpks: Option>, + pub verbose: bool, + pub intmd: HashMap, + pub result: HashMap, +} + +impl<'a> RHessKSNIMatmul<'a> { + pub fn new(mol: &CInt, xc_func_list: &'a [(f64, LibXCFunctional)], ni: NIMatmul<'a>, verbose: bool) -> Self { + Self { + mol: mol.clone(), + xc_func_list, + ni, + ni_cpks: None, + verbose, + intmd: HashMap::new(), + result: HashMap::new(), + } + } + + /// Perform the Hessian setup for RKS calculations. + /// + /// Some intermediates: + /// - `vxc`, `fxc`: These will be renamed to `cpks_vxc` and `cpks_fxc` if CP-KS specific grid + /// (`ni_cpks`) is not specified by user. If CP-KS specific grid is specified, then `cpks_vxc` + /// and `cpks_fxc` will be computed after in trait implementation + /// [`RHessElecInteractAPI::make_response_preparation`]. + /// - `rho`: discarded. + /// - `de_fxc`, `de_vxc_diag`, `de_vxc_off`, `vmat_ip`, `vmat_deriv1`: These are the main + /// results of the Hessian setup and will be stored in `self.intmd` for later use in the + /// Hessian calculation. + pub fn make_hessian_setup(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) { + // run RKS hessian setup + let dm0 = get_dm0_restricted(mo_coeff, mo_occ); + let (result, _timing) = + make_hessian_setup_batched(&self.mol, self.xc_func_list, &mut self.ni, dm0.view(), atm_list, self.verbose); + + // handling intermediates and results + for (key, val) in result.into_iter() { + if key == "vxc" || key == "fxc" { + // vxc, fxc storation is actually for cp-ks. + // If `ni_cpks` is not specified, then we can use the vxc/fxc from the hessian setup for + // cp-ks as well. + if self.ni_cpks.is_none() { + let key_to_store = format!("cpks_{key}"); + self.intmd.insert(key_to_store, val); + } + } else if ["rho"].contains(&key) { + // some keys to be discarded + continue; + } else { + self.intmd.insert(key.to_string(), val); + } + } + } + + /// Check if the Hessian setup is done by verifying the presence of the "de_fxc" key in the + /// intermediate results. + pub fn is_hessian_setup_done(&self) -> bool { + self.intmd.contains_key("de_fxc") + } +} + +impl<'a> HessUtilAPI for RHessKSNIMatmul<'a> {} + +impl<'a> RHessElecInteractAPI for RHessKSNIMatmul<'a> { + fn make_skeleton_hess(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr { + if !self.is_hessian_setup_done() { + self.make_hessian_setup(mo_coeff, mo_occ, atm_list); + } + &self.intmd["de_fxc"] + &self.intmd["de_vxc_diag"] + &self.intmd["de_vxc_off"] + } + + fn get_deriv1_ao(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr { + if !self.is_hessian_setup_done() { + self.make_hessian_setup(mo_coeff, mo_occ, atm_list); + } + self.intmd["vmat_deriv1"].to_owned() + } + + fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView) { + self.intmd.insert("mo_coeff".to_string(), mo_coeff.into_contig(ColMajor)); + self.intmd.insert("mo_occ".to_string(), mo_occ.into_contig(ColMajor)); + } + + fn get_response_bra(&mut self, bra: TsrView) -> Tsr { + let ni_cpks = self.ni_cpks.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 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_cpks, + determine_den_type_from_list(&self.xc_func_list.iter().map(|(_, f)| f).collect_vec()), + fxc_eff.view(), + bra, + mocc.view(), + self.verbose, + ); + resp + } +} + +/* #endregion */ diff --git a/src/dft/numint_matmul/hess_uks.rs b/src/dft/numint_matmul/hess_uks.rs new file mode 100644 index 0000000000..612e8fe7e7 --- /dev/null +++ b/src/dft/numint_matmul/hess_uks.rs @@ -0,0 +1,744 @@ +// see also pyhessref/nimatmul/uks.py + +use super::hess_rks::{get_de_vxc_diag, get_de_vxc_off, get_drho, get_vmat_ip}; +use super::prelude::*; +use crate::hessian_backup::prelude::*; + +use XCDenType::*; + +/* #region const dimensions/indices definition */ + +const O: usize = 0; +const X: usize = 1; +const Y: usize = 2; +const Z: usize = 3; + +#[allow(non_upper_case_globals)] +const α: usize = 0; +#[allow(non_upper_case_globals)] +const β: usize = 1; + +/* #endregion */ + +/* #region macro for indexing last dimension */ + +macro_rules! index { + ($tsr: ident, $($idx:expr),*) => { + $tsr.i((Ellipsis, $($idx),*)) + }; +} + +macro_rules! index_mut { + ($tsr: ident, $($idx:expr),*) => { + (*&mut $tsr.i_mut((Ellipsis, $($idx),*))) + }; +} + +/* #endregion */ + +/* #region basic pure functions of skeleton hessian evaluation */ + +const fn get_hess_ao_deriv(xc_type: XCDenType) -> usize { + match xc_type { + RHO => 2, + SIGMA => 3, + TAU => 3, + LAPL => unimplemented!(), + } +} + +const fn get_hess_ncomp_ao_dm0(xc_type: XCDenType) -> usize { + match xc_type { + RHO => 1, + SIGMA => 4, + TAU => 4, + LAPL => unimplemented!(), + } +} + +pub fn get_rho_vxc_fxc_uks( + xc_func_list: &[(f64, LibXCFunctional)], + ao: TsrView, + ao_dm0α: TsrView, + ao_dm0β: TsrView, +) -> (Tsr, Tsr, Tsr) { + // Returns: (rhoα, rhoβ, vxc, fxc) + // rhoα, rhoβ: [ngrids, nvar] + // vxc: [ngrids, nvar, 2] + // fxc: [ngrids, nvar, 2, nvar, 2] + + 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 = ao.shape()[0]; + let device = ao.device().clone(); + + // Compute rho + let mut rho = rt::zeros(([ngrids, nvar, 2], &device)); + for (σ, ao_dm0σ) in [(α, &ao_dm0α), (β, &ao_dm0β)] { + index_mut!(rho, 0, σ) += rt::vecdot(index!(ao, 0), index!(ao_dm0σ, O), 1); + if matches!(xc_type, SIGMA | TAU) { + index_mut!(rho, X, σ) += 2 * rt::vecdot(index!(ao, X), index!(ao_dm0σ, O), 1); + index_mut!(rho, Y, σ) += 2 * rt::vecdot(index!(ao, Y), index!(ao_dm0σ, O), 1); + index_mut!(rho, Z, σ) += 2 * rt::vecdot(index!(ao, Z), index!(ao_dm0σ, O), 1); + } + if matches!(xc_type, TAU) { + index_mut!(rho, 4, σ) += 0.5 + * (rt::vecdot(index!(ao, X), index!(ao_dm0σ, X), 1) + + rt::vecdot(index!(ao, Y), index!(ao_dm0σ, Y), 1) + + rt::vecdot(index!(ao, Z), index!(ao_dm0σ, Z), 1)) + } + } + + // Evaluate vxc and fxc with spin-polarized libxc + 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; + } + + (rho, vxc, fxc) +} + +pub fn get_drho_uks( + xc_type: XCDenType, + ao: TsrView, + ao_dm0α: TsrView, + ao_dm0β: TsrView, + aoslices: &[[usize; 4]], +) -> (Tsr, Tsr) { + let drhoα = get_drho(xc_type, ao.view(), ao_dm0α.view(), aoslices); + let drhoβ = get_drho(xc_type, ao.view(), ao_dm0β.view(), aoslices); + (drhoα, drhoβ) +} + +fn get_de_fxc_uks_inner(wf_block: TsrView, drho1: TsrView, drho2: TsrView) -> Tsr { + // wf_block: [ngrids, nvar, nvar] (a single spin block of wf) + // drho1, drho2: [ngrids, nvar, 3, natm] + // result: [3, 3, natm, natm] + + let [ngrids, nvar, _, natm] = drho1.shape().iter().cloned().collect_array().unwrap(); + + let tmp1 = rt::vecdot(wf_block.i((.., .., .., None, None)), drho1.i((.., .., None, .., ..)), 1); + let tmp1 = tmp1.reshape([ngrids * nvar, natm * 3]); + let drho2 = drho2.reshape([ngrids * nvar, natm * 3]); + let tmp2 = tmp1.t() % drho2; + + tmp2.reshape([3, natm, 3, natm]).transpose([0, 2, 1, 3]).into_contig(ColMajor) +} + +pub fn get_de_fxc_uks(wf: TsrView, drhoα: TsrView, drhoβ: TsrView) -> Tsr { + // wf: [ngrids, nvar, 2, nvar, 2] + // drhoα, drhoβ: [ngrids, nvar, 3, natm] + // result: [3, 3, natm, natm] + + let de_αα = get_de_fxc_uks_inner(wf.i((.., .., α, .., α)), drhoα.view(), drhoα.view()); + let de_αβ = get_de_fxc_uks_inner(wf.i((.., .., α, .., β)), drhoα.view(), drhoβ.view()); + let de_βα = get_de_fxc_uks_inner(wf.i((.., .., β, .., α)), drhoβ.view(), drhoα.view()); + let de_ββ = get_de_fxc_uks_inner(wf.i((.., .., β, .., β)), drhoβ.view(), drhoβ.view()); + + &de_αα + &de_αβ + &de_βα + &de_ββ +} + +#[allow(clippy::too_many_arguments)] +pub fn get_vmat_deriv1_uks( + xc_type: XCDenType, + ao: TsrView, + drhoα: TsrView, + drhoβ: TsrView, + wf: TsrView, + vmat_ip_α: TsrView, + vmat_ip_β: TsrView, + aoslices: &[[usize; 4]], +) -> (Tsr, Tsr) { + let natm = aoslices.len(); + let nao = ao.shape()[1]; + let device = ao.device(); + + let mut vmatα_deriv1: Tsr = rt::zeros(([nao, nao, 3, natm], device)); + let mut vmatβ_deriv1: Tsr = rt::zeros(([nao, nao, 3, natm], device)); + + for (A, &[_, _, p0, p1]) in aoslices.iter().enumerate() { + let slc = rt::slice!(p0, p1); + + if matches!(xc_type, RHO) { + // LDA: fxc is [G, 1, 2, 1, 2], extract scalar spin blocks + // Alpha output (s2=α): fxc_αα @ drho_α + fxc_βα @ drho_β + let wf_αα_00: Tsr = 0.5 * wf.i((.., O, α, O, α)); // [G], s1=α, s2=α + let wf_βα_00: Tsr = 0.5 * wf.i((.., O, β, O, α)); // [G], s1=β, s2=α + + let wvα_f: Tsr = + wf_αα_00.i((.., None)) * drhoα.i((.., O, .., A)) + wf_βα_00.i((.., None)) * drhoβ.i((.., O, .., A)); + + // Beta output (s2=β): fxc_αβ @ drho_α + fxc_ββ @ drho_β + let wf_αβ_00: Tsr = 0.5 * wf.i((.., O, α, O, β)); // [G], s1=α, s2=β + let wf_ββ_00: Tsr = 0.5 * wf.i((.., O, β, O, β)); // [G], s1=β, s2=β + + let wvβ_f: Tsr = + wf_αβ_00.i((.., None)) * drhoα.i((.., O, .., A)) + wf_ββ_00.i((.., None)) * drhoβ.i((.., O, .., A)); + + for t in 0..3 { + let aowα = wvα_f.i((.., t)) * index!(ao, O); + index_mut!(vmatα_deriv1, t, A).matmul_from(aowα.t(), index!(ao, O), 1.0, 1.0); + let aowβ = wvβ_f.i((.., t)) * index!(ao, O); + index_mut!(vmatβ_deriv1, t, A).matmul_from(aowβ.t(), index!(ao, O), 1.0, 1.0); + } + } + + if matches!(xc_type, SIGMA | TAU) { + let wf_αα = wf.i((.., .., α, .., α)); // [G, x, y] + let wf_αβ = wf.i((.., .., α, .., β)); // [G, x, y] + let wf_βα = wf.i((.., .., β, .., α)); // [G, x, y] + let wf_ββ = wf.i((.., .., β, .., β)); // [G, x, y] + + let drhoα_A = drhoα.i((.., .., .., A)); // [G, x, 3] + let drhoβ_A = drhoβ.i((.., .., .., A)); // [G, x, 3] + + // Python: wvα_f[y,t,g] = sum_x wf_αα[x,y,g]*drhoα[A,t,x,g] + wf_βα[x,y,g]*drhoβ[A,t,x,g] + // For each direction t, compute per-grid contraction: + // wvα_f_t[g, y] = wf_αα[g, :, y]^T @ drhoα_A[g, :, t] + wf_βα[g, :, y]^T @ drhoβ_A[g, :, t] + + for t in 0..3 { + // drhoα_A[:,:,t] shape: [G, x], reshape → [G, x, 1] + let drhoα_t = drhoα_A.i((.., .., t)); + let drhoβ_t = drhoβ_A.i((.., .., t)); + + // vecdot on axis 1: [G, x, y] @ [G, x, 1] → contract x + // Remaining: [G, y] and [G, 1] → col-major broadcast → [G, y] + // Alpha output (s2=α): fxc_αα @ drho_α + fxc_βα @ drho_β + // Beta output (s2=β): fxc_αβ @ drho_α + fxc_ββ @ drho_β + let wf_rho_αα = rt::vecdot(wf_αα.view(), drhoα_t.view(), 1); + let wf_rho_βα = rt::vecdot(wf_βα.view(), drhoβ_t.view(), 1); + let wf_rho_αβ = rt::vecdot(wf_αβ.view(), drhoα_t.view(), 1); + let wf_rho_ββ = rt::vecdot(wf_ββ.view(), drhoβ_t.view(), 1); + let mut wf_rho_α = &wf_rho_αα + &wf_rho_βα; // [G, y] + let mut wf_rho_β = &wf_rho_αβ + &wf_rho_ββ; // [G, y] + + *&mut wf_rho_α.i_mut((.., 0)) *= 0.5; + *&mut wf_rho_β.i_mut((.., 0)) *= 0.5; + if matches!(xc_type, TAU) { + *&mut wf_rho_α.i_mut((.., 4)) *= 0.25; + *&mut wf_rho_β.i_mut((.., 4)) *= 0.25; + } + + // Contract with ao: aow = sum_c wvα_f_t[:,c] * ao[c] + for c in 0..4 { + let aowα = wf_rho_α.i((.., c)) * index!(ao, c); // [G, nao] + index_mut!(vmatα_deriv1, t, A).matmul_from(aowα.t(), index!(ao, O), 1.0, 1.0); + let aowβ = wf_rho_β.i((.., c)) * index!(ao, c); // [G, nao] + index_mut!(vmatβ_deriv1, t, A).matmul_from(aowβ.t(), index!(ao, O), 1.0, 1.0); + } + + if matches!(xc_type, TAU) { + for r in [X, Y, Z] { + let aowα = wf_rho_α.i((.., 4)) * index!(ao, r); // [G, nao] + index_mut!(vmatα_deriv1, t, A).matmul_from(aowα.t(), index!(ao, r), 1.0, 1.0); + let aowβ = wf_rho_β.i((.., 4)) * index!(ao, r); // [G, nao] + index_mut!(vmatβ_deriv1, t, A).matmul_from(aowβ.t(), index!(ao, r), 1.0, 1.0); + } + } + } + } + + *&mut vmatα_deriv1.i_mut((slc, .., .., A)) -= vmat_ip_α.i((slc, .., ..)); + *&mut vmatβ_deriv1.i_mut((slc, .., .., A)) -= vmat_ip_β.i((slc, .., ..)); + } + + let vmatα_deriv1 = &vmatα_deriv1 + vmatα_deriv1.swapaxes(0, 1); + let vmatβ_deriv1 = &vmatβ_deriv1 + vmatβ_deriv1.swapaxes(0, 1); + + (vmatα_deriv1, vmatβ_deriv1) +} + +pub fn make_hessian_setup_uks( + mol: &CInt, + xc_func_list: &[(f64, LibXCFunctional)], + ni: &mut NIMatmul, + dm0α: TsrView, + dm0β: TsrView, + atm_list: Option<&[usize]>, +) -> (HashMap<&'static str, Tsr>, IndexMap<&'static str, f64>) { + assert!(!xc_func_list.is_empty(), "xc_func_list must not be empty"); + let atm_list = atm_list.map_or_else(|| (0..mol.natm()).collect_vec(), |lst| lst.to_vec()); + let aoslices_full = mol.aoslice_by_atom(); + let aoslices = atm_list.iter().map(|&iatm| aoslices_full[iatm]).collect_vec(); + let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); + + let device = dm0α.device().clone(); + let weights = rt::asarray((ni.weights.clone(), &device)); + + 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); + }; + + // --- ao, rho, vxc, fxc --- // + let t0 = std::time::Instant::now(); + let ao = ni.get_cached_ao(get_hess_ao_deriv(xc_type)); + tic("ao", t0); + + let t0 = std::time::Instant::now(); + let ncomp_ao_dm0 = get_hess_ncomp_ao_dm0(xc_type); + let ao_dm0α = index!(ao, ..ncomp_ao_dm0) % &dm0α; + let ao_dm0β = index!(ao, ..ncomp_ao_dm0) % &dm0β; + tic("ao_dm0", t0); + + let t0 = std::time::Instant::now(); + let (rho, vxc, fxc) = get_rho_vxc_fxc_uks(xc_func_list, ao.view(), ao_dm0α.view(), ao_dm0β.view()); + let wvα = &weights * vxc.i((.., .., α)); // [ngrids, nvar] + let wvβ = &weights * vxc.i((.., .., β)); // [ngrids, nvar] + let wf = &weights * &fxc; // [ngrids, nvar, 2, nvar, 2] + tic("rho, vxc, fxc", t0); + + // --- drho --- // + let t0 = std::time::Instant::now(); + let (drhoα, drhoβ) = get_drho_uks(xc_type, ao.view(), ao_dm0α.view(), ao_dm0β.view(), &aoslices); + tic("drho", t0); + + // --- de_fxc --- // + let t0 = std::time::Instant::now(); + let de_fxc = get_de_fxc_uks(wf.view(), drhoα.view(), drhoβ.view()); + tic("de_fxc", t0); + + // --- de_vxc_diag (per spin) --- // + let t0 = std::time::Instant::now(); + let de_vxc_diag_α = get_de_vxc_diag(xc_type, ao.view(), ao_dm0α.view(), wvα.view(), &aoslices); + let de_vxc_diag_β = get_de_vxc_diag(xc_type, ao.view(), ao_dm0β.view(), wvβ.view(), &aoslices); + tic("de_vxc_diag", t0); + + // --- de_vxc_off (per spin) --- // + let t0 = std::time::Instant::now(); + let de_vxc_off_α = get_de_vxc_off(xc_type, ao.view(), dm0α.view(), wvα.view(), &aoslices); + let de_vxc_off_β = get_de_vxc_off(xc_type, ao.view(), dm0β.view(), wvβ.view(), &aoslices); + tic("de_vxc_off", t0); + + // --- vmat_ip (per spin) --- // + let t0 = std::time::Instant::now(); + let vmat_ip_α = get_vmat_ip(xc_type, ao.view(), wvα.view()); + let vmat_ip_β = get_vmat_ip(xc_type, ao.view(), wvβ.view()); + tic("vmat_ip", t0); + + // --- vmat_deriv1 (UKS spin-coupled) --- // + let t0 = std::time::Instant::now(); + let (vmat_deriv1_α, vmat_deriv1_β) = get_vmat_deriv1_uks( + xc_type, + ao.view(), + drhoα.view(), + drhoβ.view(), + wf.view(), + vmat_ip_α.view(), + vmat_ip_β.view(), + &aoslices, + ); + tic("vmat_deriv1", t0); + + let result = HashMap::from([ + ("rho", rho), + ("vxc", vxc), + ("fxc", fxc), + ("de_fxc", de_fxc), + ("de_vxc_diag_a", de_vxc_diag_α), + ("de_vxc_diag_b", de_vxc_diag_β), + ("de_vxc_off_a", de_vxc_off_α), + ("de_vxc_off_b", de_vxc_off_β), + ("vmat_ip_a", vmat_ip_α), + ("vmat_ip_b", vmat_ip_β), + ("vmat_deriv1_a", vmat_deriv1_α), + ("vmat_deriv1_b", vmat_deriv1_β), + ]); + (result, timing) +} + +/* #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) +} + +/* #endregion */ + +/* #region parallel/batch wrapper */ + +pub fn make_hessian_setup_batched_uks( + mol: &CInt, + xc_func_list: &[(f64, LibXCFunctional)], + ni: &mut NIMatmul, + dm0α: TsrView, + dm0β: TsrView, + atm_list: Option<&[usize]>, + verbose: bool, +) -> (HashMap<&'static str, Tsr>, IndexMap<&'static str, f64>) { + let ngrids = ni.weights.len(); + let nbatch = ni.nbatch; + let nchunk = ni.nchunk; + let device = dm0α.device().clone(); + let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); + let nvar = xc_type.num_nvar(); + let deriv_level = get_hess_ao_deriv(xc_type); + let natm = atm_list.map_or_else(|| mol.natm(), |lst| lst.len()); + let nao = mol.nao(); + + let rhoα: Tsr = rt::zeros(([ngrids, nvar], &device)); + let rhoβ: Tsr = rt::zeros(([ngrids, nvar], &device)); + let vxc: Tsr = rt::zeros(([ngrids, nvar, 2], &device)); + let fxc: 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)); + let de_vxc_off_α: Tsr = rt::zeros(([3, 3, natm, natm], &device)); + let de_vxc_off_β: Tsr = rt::zeros(([3, 3, natm, natm], &device)); + let vmat_ip_α: Tsr = rt::zeros(([nao, nao, 3], &device)); + let vmat_ip_β: Tsr = rt::zeros(([nao, nao, 3], &device)); + let vmat_deriv1_α: Tsr = rt::zeros(([nao, nao, 3, natm], &device)); + let vmat_deriv1_β: Tsr = rt::zeros(([nao, nao, 3, natm], &device)); + + let timing = Arc::new(Mutex::new(IndexMap::from([ + ("ao", 0.0), + ("ao_dm0", 0.0), + ("rho, vxc, fxc", 0.0), + ("drho", 0.0), + ("de_fxc", 0.0), + ("de_vxc_diag", 0.0), + ("de_vxc_off", 0.0), + ("vmat_ip", 0.0), + ("vmat_deriv1", 0.0), + ("total", 0.0), + ]))); + let time_total = std::time::Instant::now(); + let guard = Mutex::new(()); + + for start_batch in (0..ngrids).step_by(nbatch) { + let end_batch = (start_batch + nbatch).min(ngrids); + + let t0 = std::time::Instant::now(); + let mut ni_batch = ni.split_batch(start_batch, end_batch); + ni_batch.get_cached_ao(deriv_level); + { + let mut timing = timing.lock().unwrap(); + timing["ao"] += t0.elapsed().as_secs_f64(); + } + + (start_batch..end_batch).into_par_iter().step_by(nchunk).for_each(|start| { + let end = (start + nchunk).min(end_batch); + let mut ni_chunk = ni_batch.split_batch(start - start_batch, end - start_batch); + let (result_chunk, timing_chunk) = + make_hessian_setup_uks(mol, xc_func_list, &mut ni_chunk, dm0α.view(), dm0β.view(), atm_list); + + unsafe { + let rhoα_slc = rhoα.i(start..end); + let rhoβ_slc = rhoβ.i(start..end); + let vxc_slc = vxc.i(start..end); + let fxc_slc = fxc.i(start..end); + let mut rhoα_slc = rhoα_slc.force_mut(); + let mut rhoβ_slc = rhoβ_slc.force_mut(); + let mut vxc_slc = vxc_slc.force_mut(); + let mut fxc_slc = fxc_slc.force_mut(); + rhoα_slc.assign(&result_chunk["rho"].i((.., .., α))); + rhoβ_slc.assign(&result_chunk["rho"].i((.., .., β))); + vxc_slc.assign(&result_chunk["vxc"]); + fxc_slc.assign(&result_chunk["fxc"]); + } + unsafe { + let lock = guard.lock().unwrap(); + *&mut de_fxc.force_mut() += &result_chunk["de_fxc"]; + *&mut de_vxc_diag_α.force_mut() += &result_chunk["de_vxc_diag_a"]; + *&mut de_vxc_diag_β.force_mut() += &result_chunk["de_vxc_diag_b"]; + *&mut de_vxc_off_α.force_mut() += &result_chunk["de_vxc_off_a"]; + *&mut de_vxc_off_β.force_mut() += &result_chunk["de_vxc_off_b"]; + *&mut vmat_ip_α.force_mut() += &result_chunk["vmat_ip_a"]; + *&mut vmat_ip_β.force_mut() += &result_chunk["vmat_ip_b"]; + *&mut vmat_deriv1_α.force_mut() += &result_chunk["vmat_deriv1_a"]; + *&mut vmat_deriv1_β.force_mut() += &result_chunk["vmat_deriv1_b"]; + drop(lock); + } + { + let mut timing = timing.lock().unwrap(); + for (key, value) in timing_chunk { + *timing.get_mut(key).unwrap() += value; + } + } + }); + + { + let mut timing = timing.lock().unwrap(); + timing.insert("total", time_total.elapsed().as_secs_f64()); + } + if verbose { + let timing = timing.lock().unwrap(); + println!("In make_hessian_setup_batched_uks, Batch {start_batch}..{end_batch}"); + println!(" Elapsed time from start (Wall time): {:.4} sec", timing["total"]); + } + } + + let timing = timing.lock().unwrap(); + if verbose { + println!("Finished make_hessian_setup_batched_uks"); + 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"); + } + } + } + + let result = HashMap::from([ + ("rhoa", rhoα), + ("rhob", rhoβ), + ("vxc", vxc), + ("fxc", fxc), + ("de_fxc", de_fxc), + ("de_vxc_diag_a", de_vxc_diag_α), + ("de_vxc_diag_b", de_vxc_diag_β), + ("de_vxc_off_a", de_vxc_off_α), + ("de_vxc_off_b", de_vxc_off_β), + ("vmat_ip_a", vmat_ip_α), + ("vmat_ip_b", vmat_ip_β), + ("vmat_deriv1_a", vmat_deriv1_α), + ("vmat_deriv1_b", vmat_deriv1_β), + ]); + + (result, timing.clone()) +} + +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 */ + +pub struct UHessKSNIMatmul<'a> { + pub mol: CInt, + pub xc_func_list: &'a [(f64, LibXCFunctional)], + pub ni: NIMatmul<'a>, + pub ni_cpks: Option>, + pub verbose: bool, + pub intmd: HashMap, + pub result: HashMap, +} + +impl<'a> UHessKSNIMatmul<'a> { + pub fn new(mol: &CInt, xc_func_list: &'a [(f64, LibXCFunctional)], ni: NIMatmul<'a>, verbose: bool) -> Self { + Self { + mol: mol.clone(), + xc_func_list, + ni, + ni_cpks: None, + verbose, + intmd: HashMap::new(), + result: HashMap::new(), + } + } + + 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(), mo_occ[β].view().greater(0).into_vec()]; + let mocc_α = mo_coeff[α].bool_select(-1, &occidx[α]); + let mocc_β = mo_coeff[β].bool_select(-1, &occidx[β]); + let dm0α = &mocc_α % mocc_α.t(); + let dm0β = &mocc_β % mocc_β.t(); + + let (result, _timing) = make_hessian_setup_batched_uks( + &self.mol, + self.xc_func_list, + &mut self.ni, + dm0α.view(), + dm0β.view(), + atm_list, + self.verbose, + ); + + for (key, val) in result.into_iter() { + if key == "vxc" || key == "fxc" { + if self.ni_cpks.is_none() { + let key_to_store = format!("cpks_{key}"); + self.intmd.insert(key_to_store, val); + } + } else if ["rhoa", "rhob"].contains(&key) { + continue; + } else { + self.intmd.insert(key.to_string(), val); + } + } + } + + pub fn is_hessian_setup_done(&self) -> bool { + self.intmd.contains_key("de_fxc") + } +} + +impl<'a> HessUtilAPI for UHessKSNIMatmul<'a> {} + +impl<'a> UHessElecInteractAPI for UHessKSNIMatmul<'a> { + fn make_skeleton_hess( + &mut self, + mo_coeff: &[TsrView; 2], + mo_occ: &[TsrView; 2], + atm_list: Option<&[usize]>, + ) -> Tsr { + if !self.is_hessian_setup_done() { + self.make_hessian_setup(mo_coeff, mo_occ, atm_list); + } + &self.intmd["de_fxc"] + + &self.intmd["de_vxc_diag_a"] + + &self.intmd["de_vxc_off_a"] + + &self.intmd["de_vxc_diag_b"] + + &self.intmd["de_vxc_off_b"] + } + + fn get_deriv1_ao( + &mut self, + mo_coeff: &[TsrView; 2], + mo_occ: &[TsrView; 2], + atm_list: Option<&[usize]>, + ) -> [Tsr; 2] { + if !self.is_hessian_setup_done() { + self.make_hessian_setup(mo_coeff, mo_occ, atm_list); + } + [self.intmd["vmat_deriv1_a"].to_owned(), self.intmd["vmat_deriv1_b"].to_owned()] + } + + fn make_response_preparation(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) { + 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)); + } + + 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 new file mode 100644 index 0000000000..78b6ae1618 --- /dev/null +++ b/src/dft/numint_matmul/mod.rs @@ -0,0 +1,107 @@ +//! Naive matrix multiplication driver for DFT numerical integration. +//! +//! Though saying "naive", it should be sufficiently good for dense GTO grids - basis pairs (small systems). +//! For large molecules, we do not exploit sparsity here for code simplicity. +#![doc = include_str!("docs/mod.md")] + +pub mod hess_rks; +pub mod hess_uks; +pub mod nimatmul; +pub mod pure_eval_rho; +pub mod pure_xcpot; + +#[allow(unused)] +pub mod prelude { + use super::*; + + pub(super) use indexmap::IndexMap; + pub(super) use libxc::prelude::*; + pub(super) use rest_libcint::prelude::*; + pub(super) use std::collections::HashMap; + pub(super) use std::sync::{Arc, Mutex}; + pub(super) use itertools::Itertools; + pub(super) use rayon::prelude::*; + + pub(super) use super::nimatmul::*; + pub(super) use super::pure_eval_rho::*; + pub(super) use super::pure_xcpot::*; + pub(super) use crate::dft::xceff::prelude::*; + pub(super) use crate::ni_check_shape; + pub(super) use crate::ri_jk::util::get_dm0_restricted; + pub(super) use crate::utilities::rstsr_util::prelude::*; + pub(super) use crate::utilities::buffer_pool::BufferPool; + + pub(super) type TsrView<'a, T = f64> = TensorView<'a, T, DeviceBLAS>; + pub(super) type Tsr = Tensor; +} + +/* #region utility ni_check_shape */ + +pub trait NIIntoUsizeVec { + fn into_usize_vec(self) -> Vec; +} + +impl NIIntoUsizeVec for i32 { + fn into_usize_vec(self) -> Vec { + vec![self as usize] + } +} + +impl NIIntoUsizeVec for usize { + fn into_usize_vec(self) -> Vec { + vec![self] + } +} + +impl NIIntoUsizeVec for &[usize] { + fn into_usize_vec(self) -> Vec { + self.to_vec() + } +} + +impl NIIntoUsizeVec for [usize; N] { + fn into_usize_vec(self) -> Vec { + self.to_vec() + } +} + +impl NIIntoUsizeVec for Vec { + fn into_usize_vec(self) -> Vec { + self + } +} + +impl NIIntoUsizeVec for &Vec { + fn into_usize_vec(self) -> Vec { + self.clone() + } +} + + +#[macro_export] +macro_rules! ni_check_shape { + ($actual:expr, $expected:expr, $msg:expr) => {{ + use $crate::dft::numint_matmul::NIIntoUsizeVec; + if $actual.into_usize_vec() != $expected.into_usize_vec() { + let str_actual = stringify!($actual); + let str_expected = stringify!($expected); + panic!( + "Shape mismatch: expected {} = {:?}, but got {} = {:?}; message: {}", + str_expected, + $expected.into_usize_vec(), + str_actual, + $actual.into_usize_vec(), + $msg + ); + } + }}; + + ($cond:expr, $msg:expr) => {{ + if !$cond { + let str_cond = stringify!($cond); + panic!("Condition failed: {}; message: {}", str_cond, $msg); + } + }}; +} + +/* #endregion */ diff --git a/src/dft/numint_matmul/nimatmul.rs b/src/dft/numint_matmul/nimatmul.rs new file mode 100644 index 0000000000..f7f6d2a733 --- /dev/null +++ b/src/dft/numint_matmul/nimatmul.rs @@ -0,0 +1,503 @@ +use super::prelude::*; + +/// Numerical integration driver using matrix-multiplication. +/// +/// Holds molecular coordinates, grid weights, and caches AO integral evaluations +/// to avoid recomputation across multiple density/XCPot evaluations. +#[derive(Clone, Debug)] +pub struct NIMatmul<'a> { + pub cint: CInt, + pub coords: Vec<[f64; 3]>, + pub weights: Vec, + + /// Cache for computed AO values, keyed by derivative order (e.g., "deriv0", "deriv1", etc.). + /// + /// This cache is not designed to be modified by API caller in usual cases. + pub cache_tensor: HashMap>, + + /// Number of grid points to process in one chunk. + /// + /// Relations of size: full-grid > batch > chunk > per-grid = 1. + /// + /// This value is better set to KC of micro-kernel (256-512 for usual x86 server). + /// Default to be 384. + pub nchunk: usize, + + /// Number of grid points to process in one batch. + /// + /// Relations of size: full-grid > batch > chunk > per-grid = 1. + /// + /// Full grid requires `[ngrids, nao, ncomp]` AO tensor, which can be too large to fit in memory + /// for big systems. That's why we need to split the grid into batches. + /// + /// This value is better set to a proper size, not exceeding available memory, and be multiple + /// of `nchunk` for better performance. + /// Default to be 384 * 4 * nthreads. nthreads is determined at runtime by rayon. + pub nbatch: usize, +} + +impl<'a> NIMatmul<'a> { + /// Creates a new instance with the given integral engine, grid coordinates, and weights. + pub fn new(cint: &CInt, coords: &[[f64; 3]], weights: &[f64]) -> Self { + assert!(coords.len() == weights.len(), "Number of coordinates must match number of weights"); + let nchunk = 384; + let nbatch = nchunk * 4 * rayon::current_num_threads(); + Self { + cint: cint.clone(), + coords: coords.to_vec(), + weights: weights.to_vec(), + cache_tensor: HashMap::new(), + nchunk, + nbatch, + } + } + + /// Clone everything, except the cached tensors. + pub fn duplicate(&self) -> Self { + Self { + cint: self.cint.clone(), + coords: self.coords.clone(), + weights: self.weights.clone(), + cache_tensor: HashMap::new(), + nchunk: self.nchunk, + nbatch: self.nbatch, + } + } + + /// Creates a new instance for a subset of the grid points, slicing from `start` to `end` + /// (exclusive). + /// + /// Cached AO tensors for the full grid are sliced accordingly and cached for the batch. + /// + /// This is useful for processing the grid in batches to reduce memory usage. + pub fn split_batch(&self, start: usize, end: usize) -> NIMatmul<'_> { + assert!(start < end && end <= self.coords.len(), "Invalid batch range"); + let mut new = self.duplicate(); + new.coords = self.coords[start..end].to_vec(); + new.weights = self.weights[start..end].to_vec(); + // if AO cached for the full grid exists, slice and cache the AO for the batch + let mut cached_tensors = HashMap::new(); + for keys in self.cache_tensor.keys() { + if keys.strip_prefix("ao_deriv").and_then(|s| s.parse::().ok()).is_some() { + let ao_full = self.cache_tensor.get(keys).unwrap(); + let ao_batch = ao_full.i(start..end); + cached_tensors.insert(keys.clone(), ao_batch.into_cow()); + } + } + new.cache_tensor = cached_tensors; + new + } + + /// Evaluates AO integrals for the given derivative order and returns as a tensor. + /// + /// The returned tensor has shape `[ngrids, nao, ncomp]` where `ncomp = AO_DERIV_DIM[deriv]`. + pub fn prepare_ao(&self, deriv: usize) -> Tsr { + let eval_name = format!("deriv{}", deriv); + let (out, shape) = self.cint.eval_gto(&eval_name, &self.coords).into(); + let device = DeviceBLAS::default(); + rt::asarray((out, shape, &device)) + } + + /// Returns cached AO values for the given derivative order, computing and caching if needed. + /// + /// When a higher derivative order has already been cached, the required subset is returned + /// directly without recomputation. + pub fn get_cached_ao(&mut self, deriv: usize) -> TsrView<'_> { + assert!( + deriv < AO_DERIV_DIM.len(), + "Derivative order {deriv} is too high, max supported is {}", + AO_DERIV_DIM.len() - 1 + ); + // determine the maximum ao deriv that have already been computed and cached + let filter_closure = |k: &String| k.strip_prefix("ao_deriv").and_then(|s| s.parse::().ok()); + let max_cached_deriv = self.cache_tensor.keys().filter_map(filter_closure).max(); + // if the requested deriv is already cached, return it + if let Some(max_deriv) = max_cached_deriv { + if max_deriv >= deriv { + let cache_key = format!("ao_deriv{}", max_deriv); + return self.cache_tensor.get(&cache_key).unwrap().i((.., .., ..AO_DERIV_DIM[deriv])); + } + } + + // otherwise, compute and cache all missing ao deriv up to the requested one + let key = format!("ao_deriv{}", deriv); + self.cache_tensor.insert(key.clone(), self.prepare_ao(deriv).into_cow()); + self.cache_tensor.get(&key).unwrap().view() + } + + /// Evaluates density from density matrices. + /// + /// # Parameters + /// + /// - `dm_list` : density matrices, each of shape `[nao, nao]`; one per set + /// - `den_type` : which density components to compute + /// + /// # Returns + /// + /// Density tensor of shape `[ngrids, nvar, nset]`. + pub fn make_rho_from_dm(&mut self, dm_list: &[TsrView], den_type: XCDenType) -> Tsr { + let nchunk = self.nchunk; + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + + let ngrids = ao.shape()[0]; + let nao = ao.shape()[1]; + for dm in dm_list { + ni_check_shape!(dm.ndim(), 2, "Each density matrix must be 2-dim"); + ni_check_shape!(dm.shape()[0..2], [nao, nao], "Density matrix must match AO dimension"); + } + let nset = dm_list.len(); + + let out_shape = [ngrids, den_type.num_nvar(), nset]; + let device = ao.device().clone(); + let mut out = rt::zeros((out_shape.f(), &device)); + get_rho_from_dm_with_output(ao, dm_list, den_type, out.view_mut(), nchunk); + out + } + + /// Evaluates density from orbital coefficients where bra and ket are the same. + /// + /// # Parameters + /// + /// - `bra_list` : orbital coefficient matrices, each of shape `[nao, nocc_i]` + /// - `den_type` : which density components to compute + /// + /// # Returns + /// + /// Density tensor of shape `[ngrids, nvar, nset]`. + pub fn make_rho_from_homogeneous_braket(&mut self, bra_list: &[TsrView], den_type: XCDenType) -> Tsr { + let nchunk = self.nchunk; + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + + let ngrids = ao.shape()[0]; + let nao = ao.shape()[1]; + for bra in bra_list { + ni_check_shape!(bra.ndim(), 2, "Each braket must be 2-dim"); + ni_check_shape!(bra.shape()[0], nao, "Bra's first dimension must match AO dimension"); + } + let nset = bra_list.len(); + let out_shape = [ngrids, den_type.num_nvar(), nset]; + let device = ao.device().clone(); + let mut out = rt::zeros((out_shape.f(), &device)); + get_rho_from_homogeneous_braket_with_output(ao, bra_list, den_type, out.view_mut(), nchunk); + out + } + + /// Evaluates density from one shared bra and multiple kets. + /// + /// # Parameters + /// + /// - `bra` : shared orbital coefficient matrix, shape `[nao, nocc]` + /// - `ket_list` : orbital coefficient matrices, each of shape `[nao, nocc]` + /// - `den_type` : which density components to compute + /// + /// # Returns + /// + /// Density tensor of shape `[ngrids, nvar, nset]`. + pub fn make_rho_from_one_bra_mult_ket(&mut self, bra: TsrView, ket_list: &[TsrView], den_type: XCDenType) -> Tsr { + let nchunk = self.nchunk; + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + + let ngrids = ao.shape()[0]; + let nao = ao.shape()[1]; + ni_check_shape!(bra.ndim(), 2, "Bra must be 2-dim"); + ni_check_shape!(bra.shape()[0], nao, "Bra first dimension must match AO dimension"); + let nocc = bra.shape()[1]; + for ket in ket_list { + ni_check_shape!(ket.ndim(), 2, "Each ket must be 2-dim"); + ni_check_shape!(ket.shape()[0], nao, "Ket first dimension must match AO dimension"); + ni_check_shape!(ket.shape()[1], nocc, "Ket second dimension must match bra"); + } + let nset = ket_list.len(); + + let out_shape = [ngrids, den_type.num_nvar(), nset]; + let device = ao.device().clone(); + let mut out = rt::zeros((out_shape.f(), &device)); + get_rho_from_one_bra_mult_ket_with_output(ao, bra, ket_list, den_type, out.view_mut(), nchunk); + out + } + + /// Evaluates density from multiple bra-ket pairs. + /// + /// # Parameters + /// + /// - `bra_list` : orbital coefficient matrices for bra, each of shape `[nao, nocc_i]` + /// - `ket_list` : orbital coefficient matrices for ket, each of shape `[nao, nocc_i]` + /// - `den_type` : which density components to compute + /// + /// # Returns + /// + /// Density tensor of shape `[ngrids, nvar, nset]`. + pub fn make_rho_from_mult_bra_mult_ket( + &mut self, + bra_list: &[TsrView], + ket_list: &[TsrView], + den_type: XCDenType, + ) -> Tsr { + let nchunk = self.nchunk; + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + + let ngrids = ao.shape()[0]; + let nao = ao.shape()[1]; + for (bra, ket) in bra_list.iter().zip(ket_list.iter()) { + ni_check_shape!(bra.ndim(), 2, "Each bra must be 2-dim"); + ni_check_shape!(ket.ndim(), 2, "Each ket must be 2-dim"); + ni_check_shape!(nao, bra.shape()[0], "Bra first dimension must match AO dimension"); + ni_check_shape!(nao, ket.shape()[0], "Ket first dimension must match AO dimension"); + ni_check_shape!(bra.shape()[1], ket.shape()[1], "Bra and ket occupation must match"); + } + let nset = bra_list.len(); + + let out_shape = [ngrids, den_type.num_nvar(), nset]; + let device = ao.device().clone(); + let mut out = rt::zeros((out_shape.f(), &device)); + get_rho_from_mult_bra_mult_ket_with_output(ao, bra_list, ket_list, den_type, out.view_mut(), nchunk); + out + } + + /// Evaluates XC potential (1st order) with vxc_eff. + /// + /// # Parameters + /// + /// - `vxc_eff` : effective XC potential, shape `[ngrids, nvar]` for RKS, `[ngrids, nvar, 2]` + /// for UKS + /// - `den_type` : which density components to compute + /// - `spin` : spin polarization or not + /// + /// # Returns + /// + /// XC potential, shape `[nao, nao]` for RKS, `[nao, nao, 2]` for UKS. + pub fn make_vxc_pot_with_eff(&mut self, vxc_eff: TsrView, den_type: XCDenType, spin: XCSpin) -> Tsr { + let nchunk = self.nchunk; + let weights_data = self.weights.clone(); + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + let nao = ao.shape()[1]; + let device = ao.device().clone(); + let weights_tsr = rt::asarray((weights_data.clone(), [weights_data.len()], &device)); + + match spin { + XCSpin::Unpolarized => { + let mut out = rt::zeros(([nao, nao], &device)); + rks_vxc_pot_with_eff_with_output(den_type, vxc_eff, ao, weights_tsr.view(), out.view_mut(), nchunk); + out + }, + XCSpin::Polarized => { + let mut out = rt::zeros(([nao, nao, 2], &device)); + uks_vxc_pot_with_eff_with_output(den_type, vxc_eff, ao, weights_tsr.view(), out.view_mut(), nchunk); + out + }, + } + } + + /// Evaluates XC potential (2nd order) with fxc_eff. + /// + /// # Parameters + /// + /// - `fxc_eff` : effective XC kernel, shape `[ngrids, nvar, nvar]` for RKS, `[ngrids, nvar, 2, + /// nvar, 2]` for UKS + /// - `rho1` : first-order density response, shape `[ngrids, nvar, nset]` for RKS, `[ngrids, + /// nvar, 2, nset]` for UKS + /// - `den_type` : which density components to compute + /// - `spin` : spin polarization or not + /// + /// # Returns + /// + /// Second-order XC potential, shape `[nao, nao, nset]` for RKS, `[nao, nao, 2, nset]` for UKS. + pub fn make_fxc_pot_with_eff(&mut self, fxc_eff: TsrView, rho1: TsrView, den_type: XCDenType, spin: XCSpin) -> Tsr { + let nchunk = self.nchunk; + let weights_data = self.weights.clone(); + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + let nao = ao.shape()[1]; + let device = ao.device().clone(); + let weights_tsr = rt::asarray((weights_data.clone(), [weights_data.len()], &device)); + + match spin { + XCSpin::Unpolarized => { + let nset = rho1.shape()[2]; + let mut out = rt::zeros(([nao, nao, nset], &device)); + rks_fxc_pot_with_eff_with_output( + den_type, + fxc_eff, + rho1, + ao, + weights_tsr.view(), + out.view_mut(), + nchunk, + ); + out + }, + XCSpin::Polarized => { + let nset = rho1.shape()[3]; + let mut out = rt::zeros(([nao, nao, 2, nset], &device)); + uks_fxc_pot_with_eff_with_output( + den_type, + fxc_eff, + rho1, + ao, + weights_tsr.view(), + out.view_mut(), + nchunk, + ); + out + }, + } + } + + /// Evaluates XC potential (2nd order) with fxc_eff, applying bra transformation. + /// + /// This function only works for RKS (spin-unpolarized case). + /// + /// Bra is usually the occupied orbital coefficient, which can lower the computational cost. + /// + /// # Parameters + /// + /// - `fxc_eff` : effective XC kernel, shape `[ngrids, nvar, nvar]` + /// - `rho1` : first-order density response, shape `[ngrids, nvar, nset]` + /// - `bra` : bra orbital coefficients, shape `[nao, nocc]` + /// - `den_type` : which density components to compute + /// + /// # Returns + /// + /// Bra-transformed XC potential, shape `[nao, nocc, nset]`. + pub fn make_rks_fxc_pot_with_eff_bra_trans( + &mut self, + fxc_eff: TsrView, + rho1: TsrView, + bra: TsrView, + den_type: XCDenType, + ) -> Tsr { + let nchunk = self.nchunk; + let weights_data = self.weights.clone(); + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + let nao = ao.shape()[1]; + let device = ao.device().clone(); + let weights_tsr = rt::asarray((weights_data.clone(), [weights_data.len()], &device)); + + let nset = rho1.shape()[2]; + let nocc = bra.shape()[1]; + let mut out = rt::zeros(([nao, nocc, nset], &device)); + rks_fxc_pot_with_eff_bra_trans_with_output( + den_type, + fxc_eff, + rho1, + ao, + weights_tsr.view(), + bra, + out.view_mut(), + nchunk, + ); + out + } + + /// Evaluates XC potential (2nd order) with fxc_eff, applying bra transformation. + /// + /// This function only works for UKS (spin-polarized case). + /// + /// Bra is usually the occupied orbital coefficient, which can lower the computational cost. + /// + /// # Parameters + /// + /// - `fxc_eff` : effective XC kernel, shape `[ngrids, nvar, 2, nvar, 2]` + /// - `rho1` : first-order density response, shape `[ngrids, nvar, 2, nset]` + /// - `bra` : bra orbital coefficients, first shape `[nao, nocc_alpha]`, second shape `[nao, + /// nocc_beta]` + /// - `den_type` : which density components to compute + pub fn make_uks_fxc_pot_with_eff_bra_trans( + &mut self, + fxc_eff: TsrView, + rho1: TsrView, + bra: &[TsrView; 2], + den_type: XCDenType, + ) -> [Tsr; 2] { + let nchunk = self.nchunk; + let weights_data = self.weights.clone(); + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + let nao = ao.shape()[1]; + let device = ao.device().clone(); + let weights_tsr = rt::asarray((weights_data.clone(), [weights_data.len()], &device)); + + let nset = rho1.shape()[3]; + let nocc_alpha = bra[0].shape()[1]; + let nocc_beta = bra[1].shape()[1]; + let mut out_alpha = rt::zeros(([nao, nocc_alpha, nset], &device)); + let mut out_beta = rt::zeros(([nao, nocc_beta, nset], &device)); + uks_fxc_pot_with_eff_bra_trans_with_output( + den_type, + fxc_eff, + rho1, + ao, + weights_tsr.view(), + bra, + &mut [out_alpha.view_mut(), out_beta.view_mut()], + nchunk, + ); + [out_alpha, out_beta] + } + + /// Evaluates XC potential (3rd order) with kxc_eff. + /// + /// # Parameters + /// + /// - `kxc_eff` : effective XC kernel, shape `[ngrids, nvar, nvar, nvar]` for RKS, `[ngrids, + /// nvar, 2, nvar, 2, nvar, 2]` for UKS + /// - `rho1` : first-order density response, shape `[ngrids, nvar, nset1]` for RKS, `[ngrids, + /// nvar, 2, nset1]` for UKS + /// - `rho2` : second-order density response, shape `[ngrids, nvar, nset2]` for RKS, `[ngrids, + /// nvar, 2, nset2]` for UKS + /// - `den_type` : which density components to compute + /// - `spin` : spin polarization or not + /// + /// # Returns + /// + /// Third-order XC potential, shape `[nao, nao, nset1, nset2]` for RKS, + /// `[nao, nao, 2, nset1, nset2]` for UKS. + pub fn make_kxc_pot_with_eff( + &mut self, + kxc_eff: TsrView, + rho1: TsrView, + rho2: TsrView, + den_type: XCDenType, + spin: XCSpin, + ) -> Tsr { + let nchunk = self.nchunk; + let weights_data = self.weights.clone(); + let ao = self.get_cached_ao(den_type.num_ao_deriv()); + let nao = ao.shape()[1]; + let device = ao.device().clone(); + let weights_tsr = rt::asarray((weights_data.clone(), [weights_data.len()], &device)); + + match spin { + XCSpin::Unpolarized => { + let nset1 = rho1.shape()[2]; + let nset2 = rho2.shape()[2]; + let mut out = rt::zeros(([nao, nao, nset1, nset2], &device)); + rks_kxc_pot_with_eff_with_output( + den_type, + kxc_eff, + rho1, + rho2, + ao, + weights_tsr.view(), + out.view_mut(), + nchunk, + ); + out + }, + XCSpin::Polarized => { + let nset1 = rho1.shape()[3]; + let nset2 = rho2.shape()[3]; + let mut out = rt::zeros(([nao, nao, 2, nset1, nset2], &device)); + uks_kxc_pot_with_eff_with_output( + den_type, + kxc_eff, + rho1, + rho2, + ao, + weights_tsr.view(), + out.view_mut(), + nchunk, + ); + out + }, + } + } +} diff --git a/src/dft/numint_matmul/pure_eval_rho.rs b/src/dft/numint_matmul/pure_eval_rho.rs new file mode 100644 index 0000000000..7440ae35bf --- /dev/null +++ b/src/dft/numint_matmul/pure_eval_rho.rs @@ -0,0 +1,465 @@ +//! Density evaluation (parallel enhanced) + +use super::prelude::*; +use XCDenType::*; + +/// Evaluate density from density matrices (parallel enhanced). +/// +/// # Parameters +/// +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `dm_list` : density matrices, each of shape `[nao, nao]`; one per set +/// - `den_type` : which density components to compute +/// - `out` : output buffer, shape `[ngrids, num_rho_comp, nset]` +/// - `ngrids_chunk` : number of grid points to process in one chunk +pub fn get_rho_from_dm_with_output(ao: TsrView, dm_list: &[TsrView], den_type: XCDenType, out: TsrMut, nchunk: usize) { + ni_check_shape!(ao.ndim(), 3, "AO values must be 3-dim"); + let nao = ao.shape()[1]; + + for dm in dm_list { + ni_check_shape!(dm.ndim(), 2, "Each density matrix must be 2-dim"); + ni_check_shape!(dm.shape()[0..2], [nao, nao], "Density matrix must match AO dimension"); + } + let nset = dm_list.len(); + let ngrids = ao.shape()[0]; + let nvar = den_type.num_nvar(); + let device = ao.device().clone(); + + ni_check_shape!(out.shape().clone(), [ngrids, nvar, nset], "Output shape mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * nchunk * (nao + + // nvar) f64 + let scr_pool = BufferPool::new(|| vec![0.0; nchunk * nao]); + let out_pool = BufferPool::new(|| vec![0.0; nchunk * nvar]); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = nset; + let ntask = ntask_grid * ntask_i; + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let iset = itask % ntask_i; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + let chunk_size = end - start; + + let dm = &dm_list[iset]; + let ao_chunk = ao.i(start..end); + + // get buffers from pool + let mut scr_buf = scr_pool.get(); + let mut out_buf = out_pool.get(); + out_buf.fill(0.0); + + let mut scr = rt::asarray((&mut scr_buf, [chunk_size, nao].f(), &device)); + let mut out_local = rt::asarray((&mut out_buf, [chunk_size, nvar].f(), &device)); + + // rho part + scr.matmul_from(ao_chunk.i((.., .., 0)), dm, 1.0, 0.0); + out_local.i_mut((.., 0)).vecdot_from(&scr, ao_chunk.i((.., .., 0)), 1); + // sigma part + if matches!(den_type, SIGMA | TAU | LAPL) { + out_local.i_mut((.., 1..4)).vecdot_from(&scr.i((.., .., None)), &ao_chunk.i((.., .., 1..4)), 1); + *&mut out_local.i_mut((.., 1..4)) *= 2.0; + } + // lapl part (second derivative of AO) + if matches!(den_type, LAPL) { + for t in [4, 7, 9] { + *&mut out_local.i_mut((.., 5)) += 2.0 * rt::vecdot(&scr, ao_chunk.i((.., .., t)), 1); + } + } + // tau part + if matches!(den_type, TAU | LAPL) { + for t in 1..4 { + scr.matmul_from(ao_chunk.i((.., .., t)), dm, 0.5, 0.0); + *&mut out_local.i_mut((.., 4)) += rt::vecdot(&scr, ao_chunk.i((.., .., t)), 1); + } + } + // lapl part (tau contribution) + if matches!(den_type, LAPL) { + let tau_contrib = 4.0 * out_local.i((.., 4)).to_owned(); + *&mut out_local.i_mut((.., 5)) += tau_contrib; + } + + // write back (should not race by design) + let mut out = unsafe { out.force_mut() }; + out.i_mut((start..end, .., iset)).assign(&out_local); + + // return buffers to pool + scr_pool.put(scr_buf); + out_pool.put(out_buf); + }); +} + +/// Evaluate density from homogeneous bra-ket (parallel enhanced). +/// +/// # Parameters +/// +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `bra_list` : orbital coefficient matrices, each of shape `[nao, nocc_i]` +/// - `den_type` : which density components to compute +/// - `out` : output buffer, shape `[ngrids, num_rho_comp, nset]` +/// - `nchunk` : number of grid points to process in one chunk +pub fn get_rho_from_homogeneous_braket_with_output( + ao: TsrView, + bra_list: &[TsrView], + den_type: XCDenType, + out: TsrMut, + nchunk: usize, +) { + ni_check_shape!(ao.ndim(), 3, "AO values must be 3-dim"); + let nao = ao.shape()[1]; + + for bra in bra_list { + ni_check_shape!(bra.ndim(), 2, "Each bra must be 2-dim"); + ni_check_shape!(nao, bra.shape()[0], "AO dimension must match braket dimension"); + } + let nocc_max = bra_list.iter().map(|bra| bra.shape()[1]).max().unwrap_or(0); + + let nset = bra_list.len(); + let ngrids = ao.shape()[0]; + let nvar = den_type.num_nvar(); + let device = ao.device().clone(); + + ni_check_shape!(out.shape().clone(), [ngrids, nvar, nset], "Output shape mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * nchunk * (2 * + // nocc_max + nvar) f64 + let scr1_pool = BufferPool::new(|| vec![0.0; nchunk * nocc_max]); + let scr2_pool = BufferPool::new(|| vec![0.0; nchunk * nocc_max]); + let out_pool = BufferPool::new(|| vec![0.0; nchunk * nvar]); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = nset; + let ntask = ntask_grid * ntask_i; + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let iset = itask % ntask_i; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + let chunk_size = end - start; + + let bra = &bra_list[iset]; + let nocc = bra.shape()[1]; + let ao_chunk = ao.i(start..end); + + // get buffers from pool + let mut scr1_buf = scr1_pool.get(); + let mut scr2_buf = scr2_pool.get(); + let mut out_buf = out_pool.get(); + out_buf.fill(0.0); + + let mut scr1 = rt::asarray((&mut scr1_buf, [chunk_size, nocc].f(), &device)); + let mut scr2 = rt::asarray((&mut scr2_buf, [chunk_size, nocc].f(), &device)); + let mut out_local = rt::asarray((&mut out_buf, [chunk_size, nvar].f(), &device)); + + // rho part + scr1.matmul_from(ao_chunk.i((.., .., 0)), bra, 1.0, 0.0); + out_local.i_mut((.., 0)).vecdot_from(&scr1, &scr1, 1); + if matches!(den_type, SIGMA | TAU | LAPL) { + for t in 1..4 { + scr2.matmul_from(ao_chunk.i((.., .., t)), bra, 1.0, 0.0); + // sigma part + out_local.i_mut((.., t)).vecdot_from(&scr1, &scr2, 1); + *&mut out_local.i_mut((.., t)) *= 2.; + // tau part + if matches!(den_type, TAU | LAPL) { + *&mut out_local.i_mut((.., 4)) += 0.5 * rt::vecdot(&scr2, &scr2, 1); + } + } + } + if matches!(den_type, LAPL) { + // lapl part (second derivative of AO) + for t in [4, 7, 9] { + scr2.matmul_from(ao_chunk.i((.., .., t)), bra, 1.0, 0.0); + *&mut out_local.i_mut((.., 5)) += 2.0 * rt::vecdot(&scr1, &scr2, 1); + } + // lapl part (tau contribution) + let tau_contrib = 4.0 * out_local.i((.., 4)).to_owned(); + *&mut out_local.i_mut((.., 5)) += tau_contrib; + } + + // write back (should not race by design) + let mut out = unsafe { out.force_mut() }; + out.i_mut((start..end, .., iset)).assign(&out_local); + + // return buffers to pool + scr1_pool.put(scr1_buf); + scr2_pool.put(scr2_buf); + out_pool.put(out_buf); + }); +} + +/// Evaluate density from one bra with multiple kets (parallel enhanced). +/// +/// # Parameters +/// +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `bra` : shared orbital coefficient matrix, shape `[nao, nocc]` +/// - `ket_list` : orbital coefficient matrices, each of shape `[nao, nocc]` +/// - `den_type` : which density components to compute +/// - `out` : output buffer, shape `[ngrids, num_rho_comp, nset]` +/// - `nchunk` : number of grid points to process in one chunk +pub fn get_rho_from_one_bra_mult_ket_with_output( + ao: TsrView, + bra: TsrView, + ket_list: &[TsrView], + den_type: XCDenType, + out: TsrMut, + nchunk: usize, +) { + ni_check_shape!(ao.ndim(), 3, "AO values must be 3-dim"); + let nao = ao.shape()[1]; + + ni_check_shape!(bra.ndim(), 2, "Bra must be 2-dim"); + ni_check_shape!(nao, bra.shape()[0], "Bra first dimension must match AO dimension"); + let nocc = bra.shape()[1]; + + for ket in ket_list { + ni_check_shape!(ket.ndim(), 2, "Each ket must be 2-dim"); + ni_check_shape!(nao, ket.shape()[0], "Ket first dimension must match AO dimension"); + ni_check_shape!(ket.shape()[1], nocc, "Ket second dimension must match bra"); + } + let nset = ket_list.len(); + let ngrids = ao.shape()[0]; + let nvar = den_type.num_nvar(); + let device = ao.device().clone(); + + ni_check_shape!(out.shape().clone(), [ngrids, nvar, nset], "Output shape mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * nchunk * (3 * + // nocc + nvar) f64 + let scr1_pool = BufferPool::new(|| vec![0.0; nchunk * nocc]); + let scr2_pool = BufferPool::new(|| vec![0.0; nchunk * nocc]); + let scr3_pool = BufferPool::new(|| vec![0.0; nchunk * nocc]); + let out_pool = BufferPool::new(|| vec![0.0; nchunk * nvar]); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = nset; + let ntask = ntask_grid * ntask_i; + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let iset = itask % ntask_i; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + let chunk_size = end - start; + + let ket = &ket_list[iset]; + let ao_chunk = ao.i(start..end); + + // get buffers from pool + let mut scr1_buf = scr1_pool.get(); + let mut scr2_buf = scr2_pool.get(); + let mut scr3_buf = scr3_pool.get(); + let mut out_buf = out_pool.get(); + out_buf.fill(0.0); + + let mut scr1 = rt::asarray((&mut scr1_buf, [chunk_size, nocc].f(), &device)); + let mut scr2 = rt::asarray((&mut scr2_buf, [chunk_size, nocc].f(), &device)); + let mut scr3 = rt::asarray((&mut scr3_buf, [chunk_size, nocc].f(), &device)); + let mut out_local = rt::asarray((&mut out_buf, [chunk_size, nvar].f(), &device)); + + // Pre-compute scr1 = ao_0_chunk @ bra + scr1.matmul_from(ao_chunk.i((.., .., 0)), &bra, 1.0, 0.0); + + // rho part + scr2.matmul_from(ao_chunk.i((.., .., 0)), ket, 1.0, 0.0); + out_local.i_mut((.., 0)).vecdot_from(&scr1, &scr2, 1); + + // sigma part + if matches!(den_type, SIGMA | TAU | LAPL) { + for t in 1..4 { + scr3.matmul_from(ao_chunk.i((.., .., t)), ket, 1.0, 0.0); + out_local.i_mut((.., t)).vecdot_from(&scr1, &scr3, 1); + scr3.matmul_from(ao_chunk.i((.., .., t)), &bra, 1.0, 0.0); + *&mut out_local.i_mut((.., t)) += rt::vecdot(&scr3, &scr2, 1); + } + } + + // lapl part (second derivative of AO), must come before tau which overwrites scr2 + if matches!(den_type, LAPL) { + for t in [4, 7, 9] { + scr3.matmul_from(ao_chunk.i((.., .., t)), ket, 1.0, 0.0); + *&mut out_local.i_mut((.., 5)) += rt::vecdot(&scr1, &scr3, 1); + scr3.matmul_from(ao_chunk.i((.., .., t)), &bra, 1.0, 0.0); + *&mut out_local.i_mut((.., 5)) += rt::vecdot(&scr3, &scr2, 1); + } + } + + // tau part (overwrites scr2, which is no longer needed for sigma/lapl) + if matches!(den_type, TAU | LAPL) { + for t in 1..4 { + scr2.matmul_from(ao_chunk.i((.., .., t)), ket, 1.0, 0.0); + scr3.matmul_from(ao_chunk.i((.., .., t)), &bra, 1.0, 0.0); + *&mut out_local.i_mut((.., 4)) += 0.5 * rt::vecdot(&scr3, &scr2, 1); + } + } + + // lapl part (tau contribution) + if matches!(den_type, LAPL) { + let tau_contrib = 4.0 * out_local.i((.., 4)).to_owned(); + *&mut out_local.i_mut((.., 5)) += tau_contrib; + } + + // write back (should not race by design) + let mut out = unsafe { out.force_mut() }; + out.i_mut((start..end, .., iset)).assign(&out_local); + + // return buffers to pool + scr1_pool.put(scr1_buf); + scr2_pool.put(scr2_buf); + scr3_pool.put(scr3_buf); + out_pool.put(out_buf); + }); +} + +/// Evaluate density from multiple bra-ket pairs (parallel enhanced). +/// +/// # Parameters +/// +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `bra_list` : orbital coefficient matrices for bra +/// - `ket_list` : orbital coefficient matrices for ket +/// - `den_type` : which density components to compute +/// - `out` : output buffer, shape `[ngrids, num_rho_comp, nset]` +/// - `nchunk` : number of grid points to process in one chunk +pub fn get_rho_from_mult_bra_mult_ket_with_output( + ao: TsrView, + bra_list: &[TsrView], + ket_list: &[TsrView], + den_type: XCDenType, + out: TsrMut, + nchunk: usize, +) { + ni_check_shape!(ao.ndim(), 3, "AO values must be 3-dim"); + let nao = ao.shape()[1]; + + ni_check_shape!(bra_list.len(), ket_list.len(), "bra_list and ket_list must have same length"); + let nocc_max = bra_list.iter().map(|bra| bra.shape()[1]).max().unwrap_or(0); + + for (bra, ket) in bra_list.iter().zip(ket_list.iter()) { + ni_check_shape!(bra.ndim(), 2, "Each bra must be 2-dim"); + ni_check_shape!(ket.ndim(), 2, "Each ket must be 2-dim"); + ni_check_shape!(nao, bra.shape()[0], "Bra first dimension must match AO dimension"); + ni_check_shape!(nao, ket.shape()[0], "Ket first dimension must match AO dimension"); + ni_check_shape!(bra.shape()[1], ket.shape()[1], "Bra and ket occupation must match"); + } + let nset = bra_list.len(); + let ngrids = ao.shape()[0]; + let nvar = den_type.num_nvar(); + let device = ao.device().clone(); + + ni_check_shape!(out.shape().clone(), [ngrids, nvar, nset], "Output shape mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * nchunk * (3 * + // nocc_max + nvar) f64 + let scr1_pool = BufferPool::new(|| vec![0.0; nchunk * nocc_max]); + let scr2_pool = BufferPool::new(|| vec![0.0; nchunk * nocc_max]); + let scr3_pool = BufferPool::new(|| vec![0.0; nchunk * nocc_max]); + let out_pool = BufferPool::new(|| vec![0.0; nchunk * nvar]); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = nset; + let ntask = ntask_grid * ntask_i; + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let iset = itask % ntask_i; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + let chunk_size = end - start; + + let bra = &bra_list[iset]; + let ket = &ket_list[iset]; + let nocc = bra.shape()[1]; + let ao_chunk = ao.i(start..end); + + // get buffers from pool + let mut scr1_buf = scr1_pool.get(); + let mut scr2_buf = scr2_pool.get(); + let mut scr3_buf = scr3_pool.get(); + let mut out_buf = out_pool.get(); + out_buf.fill(0.0); + + let mut scr1 = rt::asarray((&mut scr1_buf, [chunk_size, nocc].f(), &device)); + let mut scr2 = rt::asarray((&mut scr2_buf, [chunk_size, nocc].f(), &device)); + let mut scr3 = rt::asarray((&mut scr3_buf, [chunk_size, nocc].f(), &device)); + let mut out_local = rt::asarray((&mut out_buf, [chunk_size, nvar].f(), &device)); + + // rho part + scr1.matmul_from(ao_chunk.i((.., .., 0)), bra, 1.0, 0.0); + scr2.matmul_from(ao_chunk.i((.., .., 0)), ket, 1.0, 0.0); + out_local.i_mut((.., 0)).vecdot_from(&scr1, &scr2, 1); + + // sigma part + if matches!(den_type, SIGMA | TAU | LAPL) { + for t in 1..4 { + scr3.matmul_from(ao_chunk.i((.., .., t)), ket, 1.0, 0.0); + out_local.i_mut((.., t)).vecdot_from(&scr1, &scr3, 1); + scr3.matmul_from(ao_chunk.i((.., .., t)), bra, 1.0, 0.0); + *&mut out_local.i_mut((.., t)) += rt::vecdot(&scr3, &scr2, 1); + } + } + + // lapl part (second derivative of AO), must come before tau which overwrites scr1/scr2 + if matches!(den_type, LAPL) { + for t in [4, 7, 9] { + scr3.matmul_from(ao_chunk.i((.., .., t)), ket, 1.0, 0.0); + *&mut out_local.i_mut((.., 5)) += rt::vecdot(&scr1, &scr3, 1); + scr3.matmul_from(ao_chunk.i((.., .., t)), bra, 1.0, 0.0); + *&mut out_local.i_mut((.., 5)) += rt::vecdot(&scr3, &scr2, 1); + } + } + + // tau part (overwrites scr1/scr2, which are no longer needed for sigma/lapl) + if matches!(den_type, TAU | LAPL) { + for t in 1..4 { + scr1.matmul_from(ao_chunk.i((.., .., t)), bra, 1.0, 0.0); + scr2.matmul_from(ao_chunk.i((.., .., t)), ket, 1.0, 0.0); + *&mut out_local.i_mut((.., 4)) += 0.5 * rt::vecdot(&scr1, &scr2, 1); + } + } + + // lapl part (tau contribution) + if matches!(den_type, LAPL) { + let tau_contrib = 4.0 * out_local.i((.., 4)).to_owned(); + *&mut out_local.i_mut((.., 5)) += tau_contrib; + } + + // write back (should not race by design) + let mut out = unsafe { out.force_mut() }; + out.i_mut((start..end, .., iset)).assign(&out_local); + + // return buffers to pool + scr1_pool.put(scr1_buf); + scr2_pool.put(scr2_buf); + scr3_pool.put(scr3_buf); + out_pool.put(out_buf); + }); +} diff --git a/src/dft/numint_matmul/pure_xcpot.rs b/src/dft/numint_matmul/pure_xcpot.rs new file mode 100644 index 0000000000..bf74574e19 --- /dev/null +++ b/src/dft/numint_matmul/pure_xcpot.rs @@ -0,0 +1,1028 @@ +//! XC Potential generation (parallel enhanced) + +use super::prelude::*; +use XCDenType::*; + +/// Contract AO values with a weight vector to produce a symmetric matrix. +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `wv` : weight vector, shape `[ngrids, nvar]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `out` : output buffer, shape `[nao, nao]` +/// - `buf` : scratch buffer of length at least `ngrids * nao` +fn contract_ao_wv_without_symmetrize(den_type: XCDenType, wv: TsrView, ao: TsrView, mut out: TsrMut, buf: &mut [f64]) { + ni_check_shape!(wv.ndim(), 2, "Weight vector must be 2-dim"); + let nvar = wv.shape()[1]; + let ngrids = wv.shape()[0]; + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + let nao = ao.shape()[1]; + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + + ni_check_shape!(out.shape(), [nao, nao], "Output shape mismatch"); + ni_check_shape!(buf.len() >= ngrids * nao, "Buffer length insufficient"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // clean notation of slice, just for readability + + /// Macro for slicing AO tensor with the last index. + /// Returns [ngrids, nao] view for the specified component index. + macro_rules! ao_ { + [$idx:expr] => { + ao.i((.., .., $idx)) + }; + } + /// Macro for slicing weight vector with the last index. + /// Returns [ngrids] view for the specified component index. + macro_rules! wv_ { + [$idx:expr] => { + wv.i((.., $idx)) + }; + } + + let device = out.device().clone(); + let mut scr = rt::asarray((buf, [ngrids, nao], &device)); + + // RHO contribution + // out += 0.5 * ao_![0].t() % (wv_![0] * ao_![0]); + rt::mul_with_output(ao_![0], wv_![0], scr.view_mut()); + out.matmul_from(ao_![0].t(), scr.view(), 0.5, 0.0); + // SIGMA contribution + if matches!(den_type, SIGMA | TAU) { + // out += ao_![1].t() % (wv_![1] * ao_![0]); + rt::mul_with_output(ao_![0], wv_![1], scr.view_mut()); + out.matmul_from(ao_![1].t(), scr.view(), 1.0, 1.0); + // out += ao_![2].t() % (wv_![2] * ao_![0]); + rt::mul_with_output(ao_![0], wv_![2], scr.view_mut()); + out.matmul_from(ao_![2].t(), scr.view(), 1.0, 1.0); + // out += ao_![3].t() % (wv_![3] * ao_![0]); + rt::mul_with_output(ao_![0], wv_![3], scr.view_mut()); + out.matmul_from(ao_![3].t(), scr.view(), 1.0, 1.0); + } + // TAU contribution + if matches!(den_type, TAU) { + // out += 0.25 * ao_![1].t() % (wv_![4] * ao_![1]); + rt::mul_with_output(ao_![1], wv_![4], scr.view_mut()); + out.matmul_from(ao_![1].t(), scr.view(), 0.25, 1.0); + // out += 0.25 * ao_![2].t() % (wv_![4] * ao_![2]); + rt::mul_with_output(ao_![2], wv_![4], scr.view_mut()); + out.matmul_from(ao_![2].t(), scr.view(), 0.25, 1.0); + // out += 0.25 * ao_![3].t() % (wv_![4] * ao_![3]); + rt::mul_with_output(ao_![3], wv_![4], scr.view_mut()); + out.matmul_from(ao_![3].t(), scr.view(), 0.25, 1.0); + } +} + +/// Evaluate XC potential (1st order) with vxc_eff. +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `vxc_eff` : effective potential for XC, shape `[ngrids, nvar]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `weights` : grid weights, shape `[ngrids]` +/// - `vxc` : output vxc, shape `[nao, nao]` +/// - `nchunk` : number of grid points to process in one chunk +pub fn rks_vxc_pot_with_eff_with_output( + den_type: XCDenType, + vxc_eff: TsrView, + ao: TsrView, + weights: TsrView, + mut vxc: TsrMut, + nchunk: usize, +) { + ni_check_shape!(vxc_eff.ndim(), 2, "Effective potential must be 2-dim"); + let nvar = vxc_eff.shape()[1]; + let ngrids = vxc_eff.shape()[0]; + ni_check_shape!(weights.shape(), [ngrids], "Weights shape mismatch"); + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + let nao = ao.shape()[1]; + ni_check_shape!(vxc.shape(), [nao, nao], "Output shape mismatch"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // vxc_eff contraction + let vxc_contracted = weights * vxc_eff; + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * sum of init sizes f64 + let buffer_init = || vec![0.0; nchunk * nao]; + let buffer_pool = BufferPool::new(buffer_init); + let vxc_init = || vec![0.0; nao * nao]; + let vxc_pool = BufferPool::new(vxc_init); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask = ntask_grid; + + // atomic guard to avoid racing write + let guard = Mutex::new(()); + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let igrid = itask; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + + // get buffer from pool + let mut buf = buffer_pool.get(); + let mut vxc_buf = vxc_pool.get(); + let mut vxc_local = rt::asarray((&mut vxc_buf, [nao, nao], ao.device())); + + // perform actual evaulation + let vxc_contracted_chunk = vxc_contracted.i(start..end); + let ao_chunk = ao.i(start..end); + contract_ao_wv_without_symmetrize( + den_type, + vxc_contracted_chunk.view(), + ao_chunk.view(), + vxc_local.view_mut(), + &mut buf, + ); + + // write back with lock + let lock = guard.lock().unwrap(); + let mut vxc = unsafe { vxc.force_mut() }; + *&mut vxc += &vxc_local; + drop(lock); + + // return buffer to pool + buffer_pool.put(buf); + vxc_pool.put(vxc_buf); + }); + + // finally symmetrize the output + let vxc_buf = vxc.swapaxes(0, 1).to_owned(); + *&mut vxc += vxc_buf; +} + +/// Evaluate XC potential (2nd order, RKS) with fxc_eff (parallel enhanced). +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `fxc_eff` : effective XC kernel, shape `[ngrids, nvar, nvar]` +/// - `rho1` : first-order density response, shape `[ngrids, nvar, nset]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `weights` : grid weights, shape `[ngrids]` +/// - `fxc` : output fxc, shape `[nao, nao, nset]` +/// - `nchunk` : number of grid points to process in one chunk +pub fn rks_fxc_pot_with_eff_with_output( + den_type: XCDenType, + fxc_eff: TsrView, + rho1: TsrView, + ao: TsrView, + weights: TsrView, + mut fxc: TsrMut, + nchunk: usize, +) { + ni_check_shape!(rho1.ndim(), 3, "rho1 tensor must be 3-dim"); + let nset = rho1.shape()[2]; + let nvar = rho1.shape()[1]; + let ngrids = rho1.shape()[0]; + ni_check_shape!(fxc_eff.shape(), [ngrids, nvar, nvar], "fxc_eff shape mismatch"); + ni_check_shape!(weights.shape(), [ngrids], "Weights shape mismatch"); + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + let nao = ao.shape()[1]; + ni_check_shape!(fxc.shape(), [nao, nao, nset], "Output shape mismatch"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // fxc_eff contraction + let fxc_eff_weighted = &weights * &fxc_eff; + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * sum of init sizes f64 + let buffer_init = || vec![0.0; nchunk * nao]; + let buffer_pool = BufferPool::new(buffer_init); + let fxc_init = || vec![0.0; nao * nao]; + let fxc_pool = BufferPool::new(fxc_init); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = nset; + let ntask = ntask_grid * ntask_i; + + // atomic guard to avoid racing write + let guard = (0..ntask_i).map(|_| Mutex::new(())).collect_vec(); + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let i = itask % ntask_i; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + + // get buffer from pool + let mut buf = buffer_pool.get(); + let mut fxc_buf = fxc_pool.get(); + let mut fxc_local = rt::asarray((&mut fxc_buf, [nao, nao], ao.device())); + + // perform actual evaluation + let rho1_chunk = rho1.i((start..end, .., None, i)); + let fxc_eff_weighted_chunk = fxc_eff_weighted.i(start..end); + let fxc_contracted_chunk = (&fxc_eff_weighted_chunk * rho1_chunk).sum_axes(1); + let ao_chunk = ao.i(start..end); + contract_ao_wv_without_symmetrize( + den_type, + fxc_contracted_chunk.view(), + ao_chunk.view(), + fxc_local.view_mut(), + &mut buf, + ); + + // write back with lock + let lock = guard[i].lock().unwrap(); + let mut fxc = unsafe { fxc.force_mut() }; + *&mut fxc.i_mut((.., .., i)) += &fxc_local; + drop(lock); + + // return buffer to pool + buffer_pool.put(buf); + fxc_pool.put(fxc_buf); + }); + + // finally symmetrize the output + let mut fxc_buf: Tsr = rt::zeros(([nao, nao], fxc.device())); + for i in 0..nset { + fxc_buf.assign(&fxc.i((.., .., i)).t()); + *&mut fxc.i_mut((.., .., i)) += &fxc_buf; + } +} + +/// Contract AO with wv for RHO/SIGMA/TAU, bra-transformed variant (parallel enhanced). +/// +/// This produces an asymmetric `[nao, nocc]` output (no symmetrization needed). +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `wv` : weight vector, shape `[ngrids, nvar]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `ao_bra` : bra-transformed AO values, shape `[ngrids, nocc, ncomp]` +/// - `out` : output buffer, shape `[nao, nocc]` +/// - `buf` : scratch buffer of length at least `ngrids * nocc` +fn contract_ao_wv_bra( + den_type: XCDenType, + wv: TsrView, + ao: TsrView, + ao_bra: TsrView, + mut out: TsrMut, + buf: &mut [f64], +) { + ni_check_shape!(wv.ndim(), 2, "Weight vector must be 2-dim"); + let nvar = wv.shape()[1]; + let ngrids = wv.shape()[0]; + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + let nao = ao.shape()[1]; + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + ni_check_shape!(ao_bra.ndim(), 3, "ao_bra tensor must be 3-dim"); + ni_check_shape!(ao_bra.shape()[0], ngrids, "ao_bra grids dimension mismatch"); + let nocc = ao_bra.shape()[1]; + ni_check_shape!(ao_bra.shape()[2] >= den_type.num_ao_comp(), "ao_bra component dimension insufficient"); + ni_check_shape!(out.shape(), [nao, nocc], "Output shape mismatch"); + ni_check_shape!(buf.len() >= ngrids * nocc, "Buffer length insufficient"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + macro_rules! ao_ { + [$idx:expr] => { + ao.i((.., .., $idx)) + }; + } + macro_rules! ao_bra_ { + [$idx:expr] => { + ao_bra.i((.., .., $idx)) + }; + } + macro_rules! wv_ { + [$idx:expr] => { + wv.i((.., $idx)) + }; + } + + let device = out.device().clone(); + let mut scr = rt::asarray((buf, [ngrids, nocc], &device)); + + // RHO contribution (coefficient 1.0, not 0.5 — no symmetrization) + rt::mul_with_output(ao_bra_![0], wv_![0], scr.view_mut()); + out.matmul_from(ao_![0].t(), scr.view(), 1.0, 0.0); + + // SIGMA contribution (6 terms: ao_bra[t]*wv[t]@ao[0].T + ao_bra[0]*wv[t]@ao[t].T) + if matches!(den_type, SIGMA | TAU) { + for t in 1..4 { + rt::mul_with_output(ao_bra_![t], wv_![t], scr.view_mut()); + out.matmul_from(ao_![0].t(), scr.view(), 1.0, 1.0); + rt::mul_with_output(ao_bra_![0], wv_![t], scr.view_mut()); + out.matmul_from(ao_![t].t(), scr.view(), 1.0, 1.0); + } + } + + // TAU contribution (coefficient 0.5, not 0.25 — no symmetrization) + if matches!(den_type, TAU) { + for t in 1..4 { + rt::mul_with_output(ao_bra_![t], wv_![4], scr.view_mut()); + out.matmul_from(ao_![t].t(), scr.view(), 0.5, 1.0); + } + } +} + +/// Evaluate XC potential (2nd order, RKS) with fxc_eff, bra transformed (parallel enhanced). +/// +/// Bra is usually the occupied orbital coefficient (row-major applied to $\mu$, col-major applied +/// to $\nu$), which can lower the computational cost. +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `fxc_eff` : effective XC kernel, shape `[ngrids, nvar, nvar]` +/// - `rho1` : first-order density response, shape `[ngrids, nvar, nset]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `weights` : grid weights, shape `[ngrids]` +/// - `bra` : bra orbital coefficients, shape `[nao, nocc]` +/// - `fxc` : output fxc (bra transformed), shape `[nao, nocc, nset]` +/// - `nchunk` : number of grid points to process in one chunk +#[allow(clippy::too_many_arguments)] +pub fn rks_fxc_pot_with_eff_bra_trans_with_output( + den_type: XCDenType, + fxc_eff: TsrView, + rho1: TsrView, + ao: TsrView, + weights: TsrView, + bra: TsrView, + fxc: TsrMut, + nchunk: usize, +) { + ni_check_shape!(rho1.ndim(), 3, "rho1 tensor must be 3-dim"); + let nset = rho1.shape()[2]; + let nvar = rho1.shape()[1]; + let ngrids = rho1.shape()[0]; + ni_check_shape!(fxc_eff.shape(), [ngrids, nvar, nvar], "fxc_eff shape mismatch"); + ni_check_shape!(weights.shape(), [ngrids], "Weights shape mismatch"); + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + let nao = ao.shape()[1]; + ni_check_shape!(bra.ndim(), 2, "bra must be 2-dim"); + ni_check_shape!(bra.shape()[0], nao, "bra first dimension must match nao"); + let nocc = bra.shape()[1]; + ni_check_shape!(fxc.shape(), [nao, nocc, nset], "Output shape mismatch"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // Pre-compute ao_bra: [ngrids, nocc, ncomp] + let device = ao.device().clone(); + let ncomp = den_type.num_ao_comp(); + let mut ao_bra = rt::zeros(([ngrids, nocc, ncomp], &device)); + for c in 0..ncomp { + ao_bra.i_mut((.., .., c)).matmul_from(ao.i((.., .., c)), &bra, 1.0, 0.0); + } + + // fxc_eff contraction + let fxc_eff_weighted = &weights * &fxc_eff; + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * sum of init sizes f64 + let buffer_init = || vec![0.0; nchunk * nocc]; + let buffer_pool = BufferPool::new(buffer_init); + let fxc_init = || vec![0.0; nao * nocc]; + let fxc_pool = BufferPool::new(fxc_init); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = nset; + let ntask = ntask_grid * ntask_i; + + // atomic guard to avoid racing write + let guard = (0..ntask_i).map(|_| Mutex::new(())).collect_vec(); + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let i = itask % ntask_i; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + + // get buffer from pool + let mut buf = buffer_pool.get(); + let mut fxc_buf = fxc_pool.get(); + let mut fxc_local = rt::asarray((&mut fxc_buf, [nao, nocc], ao.device())); + + // perform actual evaluation + let rho1_chunk = rho1.i((start..end, .., None, i)); + let fxc_eff_weighted_chunk = fxc_eff_weighted.i(start..end); + let fxc_contracted_chunk = (&fxc_eff_weighted_chunk * rho1_chunk).sum_axes(1); + let ao_chunk = ao.i(start..end); + let ao_bra_chunk = ao_bra.i(start..end); + contract_ao_wv_bra( + den_type, + fxc_contracted_chunk.view(), + ao_chunk.view(), + ao_bra_chunk.view(), + fxc_local.view_mut(), + &mut buf, + ); + + // write back with lock + let lock = guard[i].lock().unwrap(); + let mut fxc = unsafe { fxc.force_mut() }; + *&mut fxc.i_mut((.., .., i)) += &fxc_local; + drop(lock); + + // return buffer to pool + buffer_pool.put(buf); + fxc_pool.put(fxc_buf); + }); +} + +/// Evaluate XC potential (3rd order, RKS) with kxc_eff (parallel enhanced). +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `kxc_eff` : effective XC kernel, shape `[ngrids, nvar, nvar, nvar]` +/// - `rho1` : first-order density response, shape `[ngrids, nvar, nset1]` +/// - `rho2` : second-order density response, shape `[ngrids, nvar, nset2]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `weights` : grid weights, shape `[ngrids]` +/// - `kxc` : output kxc, shape `[nao, nao, nset1, nset2]` +/// - `nchunk` : number of grid points to process in one chunk +#[allow(clippy::too_many_arguments)] +pub fn rks_kxc_pot_with_eff_with_output( + den_type: XCDenType, + kxc_eff: TsrView, + rho1: TsrView, + rho2: TsrView, + ao: TsrView, + weights: TsrView, + mut kxc: TsrMut, + nchunk: usize, +) { + ni_check_shape!(rho1.ndim(), 3, "rho1 tensor must be 3-dim"); + ni_check_shape!(rho2.ndim(), 3, "rho2 tensor must be 3-dim"); + let nset1 = rho1.shape()[2]; + let nset2 = rho2.shape()[2]; + let nvar = rho1.shape()[1]; + let ngrids = rho1.shape()[0]; + ni_check_shape!(kxc_eff.shape(), [ngrids, nvar, nvar, nvar], "kxc_eff shape mismatch"); + ni_check_shape!(rho2.shape()[0..2], [ngrids, nvar], "rho2 shape mismatch"); + ni_check_shape!(weights.shape(), [ngrids], "Weights shape mismatch"); + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + let nao = ao.shape()[1]; + ni_check_shape!(kxc.shape(), [nao, nao, nset1, nset2], "Output shape mismatch"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // kxc_eff contraction + let kxc_eff_weighted = &weights * &kxc_eff; + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * sum of init sizes f64 + let buffer_init = || vec![0.0; nchunk * nao]; + let buffer_pool = BufferPool::new(buffer_init); + let kxc_init = || vec![0.0; nao * nao]; + let kxc_pool = BufferPool::new(kxc_init); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = nset1 * nset2; + let ntask = ntask_grid * ntask_i; + + // atomic guard to avoid racing write + let guard = (0..ntask_i).map(|_| Mutex::new(())).collect_vec(); + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let j = itask % ntask_i; + let i1 = j % nset1; + let i2 = j / nset1; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + + // get buffer from pool + let mut buf = buffer_pool.get(); + let mut kxc_buf = kxc_pool.get(); + let mut kxc_local = rt::asarray((&mut kxc_buf, [nao, nao], ao.device())); + + // perform actual evaluation + let rho1_chunk = rho1.i((start..end, .., None, i1)); + let rho2_chunk = rho2.i((start..end, .., None, i2)); + let kxc_eff_weighted_chunk = kxc_eff_weighted.i(start..end); + // Two-step contraction: first with rho1, then with rho2 + let temp = (&kxc_eff_weighted_chunk * rho1_chunk).sum_axes(1); + let kxc_contracted_chunk = (&temp * rho2_chunk).sum_axes(1); + let ao_chunk = ao.i(start..end); + contract_ao_wv_without_symmetrize( + den_type, + kxc_contracted_chunk.view(), + ao_chunk.view(), + kxc_local.view_mut(), + &mut buf, + ); + + // write back with lock + let lock = guard[j].lock().unwrap(); + let mut kxc = unsafe { kxc.force_mut() }; + *&mut kxc.i_mut((.., .., i1, i2)) += &kxc_local; + drop(lock); + + // return buffer to pool + buffer_pool.put(buf); + kxc_pool.put(kxc_buf); + }); + + // finally symmetrize the output + let mut kxc_buf: Tsr = rt::zeros(([nao, nao], kxc.device())); + for i2 in 0..nset2 { + for i1 in 0..nset1 { + kxc_buf.assign(&kxc.i((.., .., i1, i2)).t()); + *&mut kxc.i_mut((.., .., i1, i2)) += &kxc_buf; + } + } +} + +/// Evaluate XC potential (1st order, UKS) with vxc_eff (parallel enhanced). +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `vxc_eff` : effective XC potential, shape `[ngrids, nvar, 2]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `weights` : grid weights, shape `[ngrids]` +/// - `vxc` : output vxc, shape `[nao, nao, 2]` +/// - `nchunk` : number of grid points to process in one chunk +pub fn uks_vxc_pot_with_eff_with_output( + den_type: XCDenType, + vxc_eff: TsrView, + ao: TsrView, + weights: TsrView, + mut vxc: TsrMut, + nchunk: usize, +) { + ni_check_shape!(vxc_eff.ndim(), 3, "Effective potential must be 3-dim"); + let nvar = vxc_eff.shape()[1]; + let ngrids = vxc_eff.shape()[0]; + ni_check_shape!(vxc_eff.shape()[2], 2, "vxc_eff must have 2 spin channels"); + ni_check_shape!(weights.shape(), [ngrids], "Weights shape mismatch"); + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + let nao = ao.shape()[1]; + ni_check_shape!(vxc.shape(), [nao, nao, 2], "Output shape mismatch"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // vxc_eff contraction + let vxc_eff_weighted = &weights * &vxc_eff; + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * sum of init sizes f64 + let buffer_init = || vec![0.0; nchunk * nao]; + let buffer_pool = BufferPool::new(buffer_init); + let vxc_init = || vec![0.0; nao * nao]; + let vxc_pool = BufferPool::new(vxc_init); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = 2; + let ntask = ntask_grid * ntask_i; + + // atomic guard to avoid racing write + let guard = (0..ntask_i).map(|_| Mutex::new(())).collect_vec(); + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let s = itask % ntask_i; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + + // get buffer from pool + let mut buf = buffer_pool.get(); + let mut vxc_buf = vxc_pool.get(); + let mut vxc_local = rt::asarray((&mut vxc_buf, [nao, nao], ao.device())); + + // perform actual evaluation + let vxc_contracted_chunk = vxc_eff_weighted.i((start..end, .., s)); + let ao_chunk = ao.i(start..end); + contract_ao_wv_without_symmetrize( + den_type, + vxc_contracted_chunk.view(), + ao_chunk.view(), + vxc_local.view_mut(), + &mut buf, + ); + + // write back with lock + let lock = guard[s].lock().unwrap(); + let mut vxc = unsafe { vxc.force_mut() }; + *&mut vxc.i_mut((.., .., s)) += &vxc_local; + drop(lock); + + // return buffer to pool + buffer_pool.put(buf); + vxc_pool.put(vxc_buf); + }); + + // finally symmetrize the output + let mut vxc_buf: Tsr = rt::zeros(([nao, nao], vxc.device())); + for s in 0..2 { + vxc_buf.assign(&vxc.i((.., .., s)).t()); + *&mut vxc.i_mut((.., .., s)) += &vxc_buf; + } +} + +/// Evaluate XC potential (2nd order, UKS) with fxc_eff (parallel enhanced). +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `fxc_eff` : effective XC kernel, shape `[ngrids, nvar, 2, nvar, 2]` +/// - `rho1` : first-order density response, shape `[ngrids, nvar, 2, nset]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `weights` : grid weights, shape `[ngrids]` +/// - `fxc` : output fxc, shape `[nao, nao, 2, nset]` +/// - `nchunk` : number of grid points to process in one chunk +pub fn uks_fxc_pot_with_eff_with_output( + den_type: XCDenType, + fxc_eff: TsrView, + rho1: TsrView, + ao: TsrView, + weights: TsrView, + mut fxc: TsrMut, + nchunk: usize, +) { + ni_check_shape!(rho1.ndim(), 4, "rho1 tensor must be 4-dim"); + let nset = rho1.shape()[3]; + let nvar = rho1.shape()[1]; + let ngrids = rho1.shape()[0]; + ni_check_shape!(rho1.shape()[2], 2, "rho1 must have 2 spin channels"); + ni_check_shape!(fxc_eff.shape(), [ngrids, nvar, 2, nvar, 2], "fxc_eff shape mismatch"); + ni_check_shape!(weights.shape(), [ngrids], "Weights shape mismatch"); + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + let nao = ao.shape()[1]; + ni_check_shape!(fxc.shape(), [nao, nao, 2, nset], "Output shape mismatch"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // fxc_eff contraction + let fxc_eff_weighted = &weights * &fxc_eff; + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * sum of init sizes f64 + let buffer_init = || vec![0.0; nchunk * nao]; + let buffer_pool = BufferPool::new(buffer_init); + let fxc_init = || vec![0.0; nao * nao]; + let fxc_pool = BufferPool::new(fxc_init); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = 2 * nset; + let ntask = ntask_grid * ntask_i; + + // atomic guard to avoid racing write + let guard = (0..ntask_i).map(|_| Mutex::new(())).collect_vec(); + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let j = itask % ntask_i; + let s = j % 2; + let i = j / 2; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + + // get buffer from pool + let mut buf = buffer_pool.get(); + let mut fxc_buf = fxc_pool.get(); + let mut fxc_local = rt::asarray((&mut fxc_buf, [nao, nao], ao.device())); + + // perform actual evaluation + let rho1_chunk = rho1.i((start..end, .., .., None, i)); + let fxc_eff_weighted_chunk = fxc_eff_weighted.i(start..end); + // Contract over the inner spin+var pair (axes 1 and 2) + let fxc_contracted_chunk = (&fxc_eff_weighted_chunk.i((.., .., .., .., s)) * rho1_chunk).sum_axes([1, 2]); + let ao_chunk = ao.i(start..end); + contract_ao_wv_without_symmetrize( + den_type, + fxc_contracted_chunk.view(), + ao_chunk.view(), + fxc_local.view_mut(), + &mut buf, + ); + + // write back with lock + let lock = guard[j].lock().unwrap(); + let mut fxc = unsafe { fxc.force_mut() }; + *&mut fxc.i_mut((.., .., s, i)) += &fxc_local; + drop(lock); + + // return buffer to pool + buffer_pool.put(buf); + fxc_pool.put(fxc_buf); + }); + + // finally symmetrize the output + let mut fxc_buf: Tsr = rt::zeros(([nao, nao], fxc.device())); + for i in 0..nset { + for s in 0..2 { + fxc_buf.assign(&fxc.i((.., .., s, i)).t()); + *&mut fxc.i_mut((.., .., s, i)) += &fxc_buf; + } + } +} + +/// Evaluate XC potential (3rd order, UKS) with kxc_eff (parallel enhanced). +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `kxc_eff` : effective XC kernel, shape `[ngrids, nvar, 2, nvar, 2, nvar, 2]` +/// - `rho1` : first-order density response, shape `[ngrids, nvar, 2, nset1]` +/// - `rho2` : second-order density response, shape `[ngrids, nvar, 2, nset2]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `weights` : grid weights, shape `[ngrids]` +/// - `kxc` : output kxc, shape `[nao, nao, 2, nset1, nset2]` +/// - `nchunk` : number of grid points to process in one chunk +#[allow(clippy::too_many_arguments)] +pub fn uks_kxc_pot_with_eff_with_output( + den_type: XCDenType, + kxc_eff: TsrView, + rho1: TsrView, + rho2: TsrView, + ao: TsrView, + weights: TsrView, + mut kxc: TsrMut, + nchunk: usize, +) { + ni_check_shape!(rho1.ndim(), 4, "rho1 tensor must be 4-dim"); + ni_check_shape!(rho2.ndim(), 4, "rho2 tensor must be 4-dim"); + let nset1 = rho1.shape()[3]; + let nset2 = rho2.shape()[3]; + let nvar = rho1.shape()[1]; + let ngrids = rho1.shape()[0]; + ni_check_shape!(rho1.shape()[2], 2, "rho1 must have 2 spin channels"); + ni_check_shape!(rho2.shape()[0..3], [ngrids, nvar, 2], "rho2 shape mismatch"); + ni_check_shape!(kxc_eff.shape(), [ngrids, nvar, 2, nvar, 2, nvar, 2], "kxc_eff shape mismatch"); + ni_check_shape!(weights.shape(), [ngrids], "Weights shape mismatch"); + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + let nao = ao.shape()[1]; + ni_check_shape!(kxc.shape(), [nao, nao, 2, nset1, nset2], "Output shape mismatch"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // kxc_eff contraction + let kxc_eff_weighted = &weights * &kxc_eff; + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * sum of init sizes f64 + let buffer_init = || vec![0.0; nchunk * nao]; + let buffer_pool = BufferPool::new(buffer_init); + let kxc_init = || vec![0.0; nao * nao]; + let kxc_pool = BufferPool::new(kxc_init); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = 2 * nset1 * nset2; + let ntask = ntask_grid * ntask_i; + + // atomic guard to avoid racing write + let guard = (0..ntask_i).map(|_| Mutex::new(())).collect_vec(); + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let j = itask % ntask_i; + let s = j % 2; + let i1 = (j / 2) % nset1; + let i2 = (j / 2) / nset1; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + + // get buffer from pool + let mut buf = buffer_pool.get(); + let mut kxc_buf = kxc_pool.get(); + let mut kxc_local = rt::asarray((&mut kxc_buf, [nao, nao], ao.device())); + + // perform actual evaluation + let rho1_chunk = rho1.i((start..end, .., .., None, None, None, i1)); + let rho2_chunk = rho2.i((start..end, .., .., None, i2)); + let kxc_eff_weighted_chunk = kxc_eff_weighted.i(start..end); + // Two-step contraction for UKS kxc + let kxc_slice = kxc_eff_weighted_chunk.i((.., .., .., .., .., .., s)); + let temp = (&kxc_slice * rho1_chunk).sum_axes([1, 2]); + let kxc_contracted_chunk = (&temp * rho2_chunk).sum_axes([1, 2]); + let ao_chunk = ao.i(start..end); + contract_ao_wv_without_symmetrize( + den_type, + kxc_contracted_chunk.view(), + ao_chunk.view(), + kxc_local.view_mut(), + &mut buf, + ); + + // write back with lock + let lock = guard[j].lock().unwrap(); + let mut kxc = unsafe { kxc.force_mut() }; + *&mut kxc.i_mut((.., .., s, i1, i2)) += &kxc_local; + drop(lock); + + // return buffer to pool + buffer_pool.put(buf); + kxc_pool.put(kxc_buf); + }); + + // finally symmetrize the output + let mut kxc_buf: Tsr = rt::zeros(([nao, nao], kxc.device())); + for i2 in 0..nset2 { + for i1 in 0..nset1 { + for s in 0..2 { + kxc_buf.assign(&kxc.i((.., .., s, i1, i2)).t()); + *&mut kxc.i_mut((.., .., s, i1, i2)) += &kxc_buf; + } + } + } +} + +/// Evaluate XC potential (2nd order, UKS) with fxc_eff, bra transformed (parallel enhanced). +/// +/// Bra is usually the occupied orbital coefficient (row-major applied to $\mu$, col-major applied +/// to $\nu$), which can lower the computational cost. +/// +/// # Parameters +/// +/// - `den_type`: the type of density to compute. Can be `RHO`, `SIGMA`, `TAU`. +/// - `fxc_eff` : effective XC kernel, shape `[ngrids, nvar, 2, nvar, 2]` +/// - `rho1` : first-order density response, shape `[ngrids, nvar, 2, nset]` +/// - `ao` : AO values and derivatives, shape `[ngrids, nao, ncomp]` +/// - `weights` : grid weights, shape `[ngrids]` +/// - `bra` : bra orbital coefficients, first shape `[nao, nocc_alpha]`, second shape `[nao, +/// nocc_beta]` +/// - `fxc` : output fxc (bra transformed), first shape `[nao, nocc_alpha, nset]`, second shape +/// `[nao, nocc_beta, nset]` +/// - `nchunk` : number of grid points to process in one chunk +#[allow(clippy::too_many_arguments)] +pub fn uks_fxc_pot_with_eff_bra_trans_with_output( + den_type: XCDenType, + fxc_eff: TsrView, + rho1: TsrView, + ao: TsrView, + weights: TsrView, + bra: &[TsrView; 2], + fxc: &mut [TsrMut; 2], + nchunk: usize, +) { + ni_check_shape!(rho1.ndim(), 4, "rho1 tensor must be 4-dim"); + let nset = rho1.shape()[3]; + let nvar = rho1.shape()[1]; + let ngrids = rho1.shape()[0]; + ni_check_shape!(rho1.shape()[2], 2, "rho1 must have 2 spin channels"); + ni_check_shape!(fxc_eff.shape(), [ngrids, nvar, 2, nvar, 2], "fxc_eff shape mismatch"); + ni_check_shape!(weights.shape(), [ngrids], "Weights shape mismatch"); + ni_check_shape!(den_type.num_nvar(), nvar, "Dimension mismatch for input density type"); + ni_check_shape!(ao.ndim(), 3, "AO tensor must be 3-dim"); + ni_check_shape!(ao.shape()[0], ngrids, "AO grids dimension mismatch"); + ni_check_shape!(ao.shape()[2] >= den_type.num_ao_comp(), "AO component dimension insufficient"); + let nao = ao.shape()[1]; + for bra_spin in bra.iter() { + ni_check_shape!(bra_spin.ndim(), 2, "bra tensors must be 2-dim"); + ni_check_shape!(bra_spin.shape()[0], nao, "bra first dimension must match nao"); + } + let nocc_alpha = bra[0].shape()[1]; + let nocc_beta = bra[1].shape()[1]; + let nocc_max = nocc_alpha.max(nocc_beta); + ni_check_shape!(fxc[0].shape(), [nao, nocc_alpha, nset], "Output shape mismatch"); + ni_check_shape!(fxc[1].shape(), [nao, nocc_beta, nset], "Output shape mismatch"); + + if den_type == LAPL { + panic!("Contracting AO with LAPL density type is not supported"); + } + + // Pre-compute ao_bra: [ngrids, nocc, 2, ncomp] + let device = ao.device().clone(); + let ncomp = den_type.num_ao_comp(); + let mut ao_bra_alpha = rt::zeros(([ngrids, nocc_alpha, ncomp], &device)); + let mut ao_bra_beta = rt::zeros(([ngrids, nocc_beta, ncomp], &device)); + for c in 0..ncomp { + ao_bra_alpha.i_mut((.., .., c)).matmul_from(ao.i((.., .., c)), &bra[0], 1.0, 0.0); + ao_bra_beta.i_mut((.., .., c)).matmul_from(ao.i((.., .., c)), &bra[1], 1.0, 0.0); + } + + // fxc_eff contraction + let fxc_eff_weighted = &weights * &fxc_eff; + + // buffer pool initialization + // Each BufferPool lazily creates per-thread buffers; peak usage = nthreads * sum of init sizes f64 + let buffer_init = || vec![0.0; nchunk * nao]; + let buffer_pool = BufferPool::new(buffer_init); + let fxc_init = || vec![0.0; nao * nocc_max]; + let fxc_pool = BufferPool::new(fxc_init); + + // task numbers + let ntask_grid = ngrids.div_ceil(nchunk); + let ntask_i = 2 * nset; + let ntask = ntask_grid * ntask_i; + + // atomic guard to avoid racing write + let guard = (0..ntask_i).map(|_| Mutex::new(())).collect_vec(); + + (0..ntask).into_par_iter().for_each(|itask| { + // determine task configuration + let j = itask % ntask_i; + let s = j % 2; + let i = j / 2; + let igrid = itask / ntask_i; + + // determine the grid chunk for this task + let start = igrid * nchunk; + let end = ((igrid + 1) * nchunk).min(ngrids); + + // get buffer from pool + let nocc = if s == 0 { nocc_alpha } else { nocc_beta }; + let mut buf = buffer_pool.get(); + let mut fxc_buf = fxc_pool.get(); + let mut fxc_local = rt::asarray((&mut fxc_buf, [nao, nocc], ao.device())); + + // perform actual evaluation + let rho1_chunk = rho1.i((start..end, .., .., None, i)); + let fxc_eff_weighted_chunk = fxc_eff_weighted.i(start..end); + // Contract over the inner spin+var pair (axes 1 and 2) + let fxc_contracted_chunk = (&fxc_eff_weighted_chunk.i((.., .., .., .., s)) * rho1_chunk).sum_axes([1, 2]); + let ao_chunk = ao.i(start..end); + let ao_bra_chunk = if s == 0 { ao_bra_alpha.i(start..end) } else { ao_bra_beta.i(start..end) }; + contract_ao_wv_bra( + den_type, + fxc_contracted_chunk.view(), + ao_chunk.view(), + ao_bra_chunk.view(), + fxc_local.view_mut(), + &mut buf, + ); + + // write back with lock + let lock = guard[j].lock().unwrap(); + let mut fxc_spin = unsafe { fxc[s].force_mut() }; + *&mut fxc_spin.i_mut((.., .., i)) += &fxc_local; + drop(lock); + + // return buffer to pool + buffer_pool.put(buf); + fxc_pool.put(fxc_buf); + }); +} diff --git a/src/dft/xceff/flags.rs b/src/dft/xceff/flags.rs new file mode 100644 index 0000000000..6ab361e6bb --- /dev/null +++ b/src/dft/xceff/flags.rs @@ -0,0 +1,107 @@ +//! Flags for DFT evaluation. +//! +//! This file is currently in xceff module, which only LibXC implementation uses. But it should be +//! independent to some specific XC evaluation implementation, and can be used by other +//! implementations as well. So this file is free to be moved to a more general place in the future. + +pub const AO_DERIV_DIM: [usize; 5] = [1, 4, 10, 20, 35]; + +/// Density type for XC functionals. +/// +/// - RHO: only density +/// - SIGMA: density + gradient +/// - TAU: density + gradient + kinetic energy density +/// - LAPL: density + gradient + kinetic energy density + laplacian +/// +/// Note for this enum, each higher-level density type also contains all components of the +/// lower-level types. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum XCDenType { + RHO, + SIGMA, + TAU, + LAPL, +} + +impl XCDenType { + /// Returns the number of components in the output density for this XC type. + /// + /// - RHO: 1 component (density) + /// - SIGMA: 4 components (density + 3 gradient components) + /// - TAU: 5 components (density + 3 gradient components + kinetic energy density) + /// - LAPL: 6 components (density + 3 gradient components + kinetic energy density + laplacian) + pub fn num_nvar(&self) -> usize { + match self { + XCDenType::RHO => 1, + XCDenType::SIGMA => 4, + XCDenType::TAU => 5, + XCDenType::LAPL => 6, + } + } + + /// Returns the required AO derivative level for this XC type. + /// + /// - RHO: 0th order + /// - SIGMA: 1st order (gradient) + /// - TAU: 1st order (gradient) + /// - LAPL: 2nd order (Laplacian) + pub fn num_ao_deriv(&self) -> usize { + match self { + XCDenType::RHO => 0, + XCDenType::SIGMA => 1, + XCDenType::TAU => 1, + XCDenType::LAPL => 2, + } + } + + /// Returns the number of AO components needed for this XC type + /// + /// - RHO: 1 component (AO value) + /// - SIGMA: 4 components (AO value + 3 gradient components) + /// - TAU: 4 components (AO value + 3 gradient components) [ + /// - LAPL: 10 components (AO value + 3 gradient components + 6 second derivative components) + pub fn num_ao_comp(&self) -> usize { + AO_DERIV_DIM[self.num_ao_deriv()] + } +} + +/// Parallelization strategy for XC evaluation. +/// +/// This enum allows three kinds of parallelization strategies by `From` trait implementations: +/// +/// - usize number : parallel with given chunk size; +/// - None : Use default chunk size determined by the implementation function; +/// - bool : parallel with auto-chunking if true, or serial if false. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum XCPar { + Par { chunk_size: Option }, + Serial, +} + +impl From for XCPar { + fn from(chunk_size: usize) -> Self { + XCPar::Par { chunk_size: Some(chunk_size) } + } +} + +impl From> for XCPar { + fn from(chunk_size: Option) -> Self { + XCPar::Par { chunk_size } + } +} + +impl From for XCPar { + fn from(parallel: bool) -> Self { + if parallel { + XCPar::Par { chunk_size: None } + } else { + XCPar::Serial + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum XCSpin { + Unpolarized, + Polarized, +} diff --git a/src/dft/xceff/libxc-xcfun-trans.md b/src/dft/xceff/libxc-xcfun-trans.md new file mode 100644 index 0000000000..2279115c3c --- /dev/null +++ b/src/dft/xceff/libxc-xcfun-trans.md @@ -0,0 +1,131 @@ +# Notes on LibXC/XCFun convention transformation + +## Explanation by specific example + +LibXC and XCFun have different conventions on how to handle the spin-polarized cases. + +We will use GGA spin-polarized second-order derivative (fxc) as example to explain that. + +As the first step, we state that for the energy part (zk) and first-order derivative part (vxc), LibXC and XCFun have the same conventions. The order is + +| Index | Notation | +|--|--| +| 0 | zk | +| 1 | r_u | +| 2 | r_d | +| 3 | s_uu | +| 4 | s_ud | +| 5 | s_dd | + +For the notation in above table, +- Rho is spin-polarization two-component, `r_u` means $\rho^\uparrow$, `r_d` means $\rho^\downarrow$. +- Sigma is spin-polarization three-component, `s_uu` means $\sigma^{\uparrow \uparrow}$, `s_ud` means $\sigma^{\uparrow \downarrow}$, `s_dd` means $\sigma^{\downarrow \downarrow}$. + +For the fxc part, there will involve two densities as variables. The LibXC and XCFun have fundamental difference at + +- LibXC respect which type of density first (sorted by rho, sigma), then its spin; +- XCFun respect the spin-component first (sorted by r_u, r_d, s_uu, s_ud, s_dd). + +It is more favorable to use XCFun-style for future DFT evaluation. However, LibXC supports more functionals, with better API design and popularity. So an index transform mapping is required. For this specific task, + +| Index | Notation
LibXC | Notation
XCFun | Map | +|-:|--|--|-:| +| 6 | r_u / r_u | r_u / r_u | 6 | +| 7 | r_u / r_d | r_u / r_d | 7 | +| 8 | r_d / r_d | r_u / s_uu | 9 | +| 9 | r_u / s_uu | r_u / s_ud | 10 | +| 10 | r_u / s_ud | r_u / s_dd | 11 | +| 11 | r_u / s_dd | r_d / r_d | 8 | +| 12 | r_d / s_uu | r_d / s_uu | 12 | +| 13 | r_d / s_ud | r_d / s_ud | 13 | +| 14 | r_d / s_dd | r_d / s_dd | 14 | +| 15 | s_uu / s_uu | s_uu / s_uu | 15 | +| 16 | s_uu / s_ud | s_uu / s_ud | 16 | +| 17 | s_uu / s_dd | s_uu / s_dd | 17 | +| 18 | s_ud / s_ud | s_ud / s_ud | 18 | +| 19 | s_ud / s_dd | s_ud / s_dd | 19 | +| 20 | s_dd / s_dd | s_dd / s_dd | 20 | + +We can see that + +- LibXC will first category `r/r` (6--8), then `r/s` (9--14), then `s/s` (16--20). For each category, sort by spin. +- XCFun will first category `r_u` (6--10), then `r_d` (11--14), then `s_uu` (15--17), then `s_ud` (18--19), then `s_dd` (20). For each category, sort by the same way in first category. + +## Extension to the example + +- **Higher derivative**: We may encounter more higher derivatives (usually up to 4th derivative, but can be more). +- **More kinds of density**: We may use more density types. The priority is RHO > SIGMA > TAU > LAPL. + - TAU: `t_u`, `t_d` + - LAPL: `l_u`, `l_d` +- We assume the RHO (LDA) only inputs RHO; SIGMA (GGA) inputs both RHO and SIGMA; TAU (some meta-GGA) inputs RHO SIGMA TAU, and LAPL (some meta-GGA) inputs all RHO SIGMA TAU LAPL (though some LAPL meta-GGAs does not actually input tau, but we still require the TAU to be available for simplicity). + +## Code for Automatic Index Map Generation + +```python +from itertools import combinations_with_replacement +from math import comb + + +def libxc_to_xcfun_indices_map(den_type: str, deriv: int) -> list[int]: + """Spin-Polarized Indices Map from LibXC to XCFun. + + Parameters + ---------- + den_type : str + Density Type. Supports `rho`, `sigma`, `tau`, `lapl`. + deriv : int + Derivative level. + + Example + ------- + >>> libxc_to_xcfun_indices_map("sigma", 2) + [6, 7, 9, 10, 11, 8, 12, 13, 14, 15, 16, 17, 18, 19, 20] + """ + # Each variable: (type_priority, spin_index) + # RHO has 2 spin components, SIGMA has 3, TAU has 2, LAPL has 2 + group_specs = [ + ("rho", 0, 2), + ("sigma", 1, 3), + ("tau", 2, 2), + ("lapl", 3, 2), + ] + + type_map = { + "rho": ["rho"], + "sigma": ["rho", "sigma"], + "tau": ["rho", "sigma", "tau"], + "lapl": ["rho", "sigma", "tau", "lapl"], + } + + if den_type not in type_map: + raise ValueError(f"Unknown den_type: {den_type}") + + active_groups = set(type_map[den_type]) + + # Build variable list: each variable is (type_priority, spin_index) + variables = [] + for group_name, priority, n_spin in group_specs: + if group_name in active_groups: + for spin in range(n_spin): + variables.append((priority, spin)) + + d = len(variables) + + # Generate all non-decreasing multi-indices of length deriv + # combinations_with_replacement yields them in lexicographic order = XCFun order + xcfun_order = list(combinations_with_replacement(range(d), deriv)) + + # LibXC order: sort by density type signature first, then by variable indices + def libxc_key(mi): + return tuple(variables[i][0] for i in mi) + mi + + libxc_order = sorted(xcfun_order, key=libxc_key) + + # Build reverse lookup: multi-index -> LibXC position + libxc_pos = {mi: pos for pos, mi in enumerate(libxc_order)} + + # Base offset = sum of outputs for all previous derivative levels + base_offset = sum(comb(d + i - 1, i) for i in range(deriv)) + + return [base_offset + libxc_pos[mi] for mi in xcfun_order] +``` diff --git a/src/dft/xceff/libxc_wrap.rs b/src/dft/xceff/libxc_wrap.rs new file mode 100644 index 0000000000..e45abf3204 --- /dev/null +++ b/src/dft/xceff/libxc_wrap.rs @@ -0,0 +1,247 @@ +use super::prelude::*; +use super::xc_deriv::{libxc_transform_xcfun_indices, transform_xc_inner}; +use libxc::compute_cpu::LibXCCpuInput; +use libxc::prelude::*; + +use LibXCSpin::*; +use XCDenType::*; + +/// Determine the density type required by the given XC functional. +pub fn determine_den_type(xc_func: &LibXCFunctional) -> XCDenType { + match xc_func.family() { + LibXCFamily::LDA | LibXCFamily::HybLDA => RHO, + LibXCFamily::GGA | LibXCFamily::HybGGA => SIGMA, + LibXCFamily::MGGA | LibXCFamily::HybMGGA => { + if xc_func.needs_laplacian() { + LAPL + } else { + TAU + } + }, + _ => panic!("Unsupported functional family: {:?}", xc_func.family()), + } +} + +pub fn determine_den_type_from_list(xc_func_list: &[&LibXCFunctional]) -> XCDenType { + xc_func_list.iter().map(|f| determine_den_type(f)).max_by_key(|&den_type| den_type.num_nvar()).unwrap() +} + +/// Evaluate the XC energy/potential, to LibXC raw output. +pub fn libxc_eval_inner(xc_func: &LibXCFunctional, rho: TsrView, deriv: usize) -> (Vec, LibXCOutputLayout) { + // sanity check + // rho must be either [ngrids, 1/4/5/6] or [ngrids, 1/4/5/6, 2] + match xc_func.spin() { + Unpolarized => ni_check_shape!(rho.ndim(), 2, "rho for unpolarized functionals must be a 2-dim tensor"), + Polarized => { + ni_check_shape!(rho.ndim(), 3, "rho for polarized functionals must be a 3-dim tensor"); + ni_check_shape!(rho.shape()[2], 2, "rho for polarized functionals must have last dimension of size 2"); + }, + } + // we do not support laplacian currently + if xc_func.needs_laplacian() { + panic!("Laplacian-dependent functionals are not supported yet"); + } + let den_type = determine_den_type(xc_func); + ni_check_shape!(rho.shape()[1] >= den_type.num_nvar(), "Input density does not have enough components"); + let do_rho = matches!(den_type, RHO | SIGMA | TAU | LAPL); + let do_sigma = matches!(den_type, SIGMA | TAU | LAPL); + let do_tau = matches!(den_type, TAU | LAPL); + + if xc_func.spin() == Unpolarized { + // build up owned components + let xc_rho = do_rho.then(|| rho.i((.., 0)).to_vec()); + let xc_sigma = do_sigma.then(|| rt::vecdot(rho.i((.., 1..4)), rho.i((.., 1..4)), 1).into_vec()); + let xc_tau = do_tau.then(|| rho.i((.., 4)).to_vec()); + // construct xc_input + let mut xc_input = LibXCCpuInput::new(); + for (key, value) in [("rho", xc_rho.as_ref()), ("sigma", xc_sigma.as_ref()), ("tau", xc_tau.as_ref())] { + value.map(|v| xc_input.insert(key.to_string(), v.as_slice())); + } + xc_func.compute_xc(&xc_input, deriv).unwrap_or_else(|e| panic!("LibXC compute_xc failed: {e}")) + } else { + // build up owned components + // note libxc's convention is [2, ngrids] for rho, not [ngrids, 2]. + // we need to transpose the input rho before feeding into libxc. + let xc_rho = do_rho.then(|| rho.i((.., 0, ..)).t().into_shape(-1).to_vec()); + let xc_tau = do_tau.then(|| rho.i((.., 4, ..)).t().into_shape(-1).to_vec()); + // sigma is more complicated, as it requires (uu ud dd) components. + let xc_sigma = do_sigma.then(|| { + let ngrids = rho.shape()[0]; + let mut sigma = rt::zeros(([ngrids, 3], &rho.device().clone())); + sigma.i_mut((.., 0)).vecdot_from(rho.i((.., 1..4, 0)), rho.i((.., 1..4, 0)), 1); + sigma.i_mut((.., 1)).vecdot_from(rho.i((.., 1..4, 0)), rho.i((.., 1..4, 1)), 1); + sigma.i_mut((.., 2)).vecdot_from(rho.i((.., 1..4, 1)), rho.i((.., 1..4, 1)), 1); + sigma.t().into_shape(-1).to_vec() + }); + let mut xc_input = LibXCCpuInput::new(); + for (key, value) in [("rho", xc_rho.as_ref()), ("sigma", xc_sigma.as_ref()), ("tau", xc_tau.as_ref())] { + value.map(|v| xc_input.insert(key.to_string(), v.as_slice())); + } + xc_func.compute_xc(&xc_input, deriv).unwrap_or_else(|e| panic!("LibXC compute_xc failed: {e}")) + } +} + +/// Transpose matrix with buffer. +/// +/// This will perform inplace, but algorithm is naive. Should be able to optimize but that's too +/// hard for me. +fn transpose_with_buffer(slc: &mut [f64], m: usize, n: usize, buf: &mut [f64]) { + // currently it is a naive implementation + assert!(slc.len() >= n * m); + assert!(buf.len() >= n * m); + for j in 0..m { + for i in 0..n { + buf[j * n + i] = slc[i * m + j]; + } + } + slc[..n * m].copy_from_slice(&buf[..n * m]); +} + +/// Evaluate effective XC potential from LibXC functional and density, in serial. +pub fn libxc_eval_eff_serial(xc_func: &LibXCFunctional, rho: TsrView, deriv: usize) -> Vec { + let den_type = determine_den_type(xc_func); + let device = rho.device().clone(); + let (mut xc_val, xc_layout) = libxc_eval_inner(xc_func, rho.view(), deriv); + // transpose the spin-related components + // first find the largest intermediate size + let ngrids = rho.shape()[0]; + let buf_size = xc_layout.iter_to_range().map(|(_, r)| r.end - r.start).max().unwrap_or(ngrids); + let mut buf = vec![0.0; buf_size]; + if xc_func.spin() == Polarized { + for (_, r) in xc_layout.iter_to_range() { + if r.end - r.start != ngrids { + let ncomp = (r.end - r.start) / ngrids; + transpose_with_buffer(&mut xc_val[r], ncomp, ngrids, &mut buf); + } + } + } + let ngrids = rho.shape()[0]; + let xlen = xc_val.len() / ngrids; + let xc_val = rt::asarray((xc_val, [ngrids, xlen].f(), &device)); + let xc_val = libxc_transform_xcfun_indices(xc_val.view(), den_type, xc_func.spin(), deriv); + (0..=deriv).map(|order| transform_xc_inner(rho.view(), xc_val.view(), den_type, xc_func.spin(), order)).collect() +} + +/// Evaluate effective XC potential from LibXC functional and density, in parallel. +pub fn libxc_eval_eff_parallel( + xc_func: &LibXCFunctional, + rho: TsrView, + deriv: usize, + par_chunk_size: Option, +) -> Vec { + // if in threadpool (thread-index is Some), we use 1 thread to avoid nested parallelism. + let nthreads = rayon::current_thread_index().map_or(rayon::current_num_threads(), |_| 1); + if nthreads == 1 { + return libxc_eval_eff_serial(xc_func, rho, deriv); + } + + // determine chunk size + // this setting is probably good? anyway, user can set by argument. + let den_type = determine_den_type(xc_func); + let spin = xc_func.spin(); + let par_chunk_size = par_chunk_size.unwrap_or(match (den_type, spin) { + (RHO, Unpolarized) => 16384, + (RHO, Polarized) => 6144, + (SIGMA, _) => 384, + (TAU | LAPL, _) => 256, + }); + let ngrids = rho.shape()[0]; + let par_chunk_size = ngrids.div_ceil(nthreads).max(par_chunk_size); + if par_chunk_size >= ngrids { + // if chunk size is larger than total grids, just do serial computation. + return libxc_eval_eff_serial(xc_func, rho, deriv); + } + + // determine output shape at this stage + let nvar = den_type.num_nvar(); + // generate shapes + let out_shapes = (0..=deriv) + .map(|order| { + // unpolarized: [ngrids, [nvar] * deriv] + // polarized: [ngrids, [nvar, 2] * deriv] + let mut shape = vec![ngrids]; + for _ in 0..order { + match spin { + Unpolarized => shape.push(nvar), + Polarized => shape.extend_from_slice(&[nvar, 2]), + } + } + shape + }) + .collect_vec(); + // generate tensors + let xc_eff = out_shapes.iter().map(|shape| rt::zeros((shape.to_vec(), &rho.device().clone()))).collect_vec(); + + // parallel computation + (0..ngrids).into_par_iter().step_by(par_chunk_size).for_each(|start| { + let stop = (start + par_chunk_size).min(ngrids); + let rho_chunk = rho.i(start..stop); + let xc_eff_chunk = libxc_eval_eff_serial(xc_func, rho_chunk, deriv); + for (order, xc_eff_order) in xc_eff_chunk.into_iter().enumerate() { + let xc_eff_orig = xc_eff[order].i(start..stop); + let mut xc_eff_orig = unsafe { xc_eff_orig.force_mut() }; + xc_eff_orig.assign(&xc_eff_order); + } + }); + + xc_eff +} + +/// Evaluate effective XC potential from LibXC functional and density, with parallel option. +/// +/// # Parameter +/// +/// - `xc_func`: The LibXC functional to evaluate. +/// +/// The user should make sure to initialize the LibXC functional with **proper +/// spin-polarization**, and possibly **proper omega/parameter settings**. +/// +/// - `rho`: The density tensor. The shape must be +/// +/// - Spin Unpolarized: `[ngrids, nvar]` +/// - Spin Polarized: `[ngrids, nvar, 2]` +/// +/// Where `nvar` is the number of variables: +/// +/// - RHO: 1 ($\rho$) +/// - SIGMA: 4 ($\rho$, $\rho_x$, $\rho_y$, $\rho_z$) +/// - TAU: 5 ($\rho$, $\rho_x$, $\rho_y$, $\rho_z$, $\tau$) +/// +/// Note we slightly differ to the PySCF's convention. RHO's dimension `nvar=1` cannot be +/// squeezed out. This is to make the code more consistent and easier to implement. +/// +/// - `deriv`: The maximum derivative order to evaluate. Note the smaller derivates will also be +/// evaluated and returned. +/// +/// - `par`: How to parallelize the evaluation. +/// +/// - true/false: Whether to parallelize or not. If true, it will use the default parallelization +/// strategy. +/// - usize: The chunk size to be parallelized. If None, it will use the default chunk size. +/// +/// The default chunksize depends on density type and spin: +/// +/// - RHO, Unpolarized: 16384 +/// - RHO, Polarized: 6144 +/// - SIGMA: 384 +/// - TAU: 256 +/// +/// # Input/Output Shapes +/// +/// Please note the output will contain all derivatives up to `deriv` user specified. For example, +/// if `deriv = 1` for GGA (SIGMA) and unpolarized, the output will contain two tensors, first of +/// shape `[ngrids]`, second of shape `[ngrids, 3]`. +/// +/// | deriv | Output `xc_eff`
Unpolarized | Output `xc_eff`
Polarized | +/// |-------|--------------------------------|---------------------------------------| +/// | 0 | `[ngrids]` | `[ngrids]` | +/// | 1 | `[ngrids, nvar]` | `[ngrids, nvar, 2]` | +/// | 2 | `[ngrids, nvar, nvar]` | `[ngrids, nvar, 2, nvar, 2]` | +/// | 3 | `[ngrids, nvar, nvar, nvar]` | `[ngrids, nvar, 2, nvar, 2, nvar, 2]` | +pub fn libxc_eval_eff(xc_func: &LibXCFunctional, rho: TsrView, deriv: usize, par: impl Into) -> Vec { + let par = par.into(); + match par { + XCPar::Par { chunk_size } => libxc_eval_eff_parallel(xc_func, rho, deriv, chunk_size), + XCPar::Serial => libxc_eval_eff_serial(xc_func, rho, deriv), + } +} diff --git a/src/dft/xceff/mod.rs b/src/dft/xceff/mod.rs new file mode 100644 index 0000000000..99aeed52ce --- /dev/null +++ b/src/dft/xceff/mod.rs @@ -0,0 +1,24 @@ +//! Generate effective XC potentials by LibXC driver. +//! +//! Functionality of this module is similar to `libxc_itrf.rs`, function `eval_xc_eff`. +//! We may handle refactor and merge in the future. + +pub mod flags; +pub mod libxc_wrap; +pub mod xc_deriv; + +pub mod prelude { + use super::*; + + pub use flags::{XCDenType, XCPar, XCSpin, AO_DERIV_DIM}; + pub use libxc_wrap::{determine_den_type, determine_den_type_from_list, libxc_eval_eff}; + + pub(super) use crate::ni_check_shape; + pub(super) use crate::utilities::rstsr_util::prelude::*; + pub(super) use itertools::Itertools; + pub(super) use libxc::prelude::*; + pub(super) use rayon::prelude::*; + + pub(super) type TsrView<'a, T = f64> = TensorView<'a, T, DeviceBLAS>; + pub(super) type Tsr = Tensor; +} diff --git a/src/dft/xceff/xc_deriv.rs b/src/dft/xceff/xc_deriv.rs new file mode 100644 index 0000000000..6abfad4181 --- /dev/null +++ b/src/dft/xceff/xc_deriv.rs @@ -0,0 +1,427 @@ +use super::prelude::*; + +use LibXCSpin::*; +use XCDenType::*; + +// https://stackoverflow.com/a/65563202/7740992 +pub fn count_combinations(n: usize, r: usize) -> usize { + if r > n { + 0 + } else { + (1..=r).fold(1, |acc, val| acc * (n - val + 1) / val) + } +} + +/* #region libxc-to-xcfun index convention change */ + +pub const fn libxc_to_xcfun_mapping_parts(den_type: XCDenType, deriv: usize) -> Option<&'static [usize]> { + match (den_type, deriv) { + (RHO, _) => None, + (SIGMA, 0) => Some(&[0]), + (SIGMA, 1) => Some(&[1, 2, 3, 4, 5]), + (SIGMA, 2) => Some(&[6, 7, 9, 10, 11, 8, 12, 13, 14, 15, 16, 17, 18, 19, 20]), + (SIGMA, 3) => Some(&[ + 21, 22, 25, 26, 27, 23, 28, 29, 30, 34, 35, 36, 37, 38, 39, 24, 31, 32, 33, 40, 41, 42, 43, 44, 45, 46, 47, + 48, 49, 50, 51, 52, 53, 54, 55, + ]), + (SIGMA, 4) => Some(&[ + 56, 57, 61, 62, 63, 58, 64, 65, 66, 73, 74, 75, 76, 77, 78, 59, 67, 68, 69, 79, 80, 81, 82, 83, 84, 91, 92, + 93, 94, 95, 96, 97, 98, 99, 100, 60, 70, 71, 72, 85, 86, 87, 88, 89, 90, 101, 102, 103, 104, 105, 106, 107, + 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, + ]), + (TAU, 0) => Some(&[0]), + (TAU, 1) => Some(&[1, 2, 3, 4, 5, 6, 7]), + (TAU, 2) => Some(&[ + 8, 9, 11, 12, 13, 17, 18, 10, 14, 15, 16, 19, 20, 21, 22, 23, 27, 28, 24, 25, 29, 30, 26, 31, 32, 33, 34, + 35, + ]), + (TAU, 3) => Some(&[ + 36, 37, 40, 41, 42, 49, 50, 38, 43, 44, 45, 51, 52, 55, 56, 57, 67, 68, 58, 59, 69, 70, 60, 71, 72, 79, 80, + 81, 39, 46, 47, 48, 53, 54, 61, 62, 63, 73, 74, 64, 65, 75, 76, 66, 77, 78, 82, 83, 84, 85, 86, 87, 95, 96, + 88, 89, 97, 98, 90, 99, 100, 107, 108, 109, 91, 92, 101, 102, 93, 103, 104, 110, 111, 112, 94, 105, 106, + 113, 114, 115, 116, 117, 118, 119, + ]), + (TAU, 4) => Some(&[ + 120, 121, 125, 126, 127, 137, 138, 122, 128, 129, 130, 139, 140, 145, 146, 147, 163, 164, 148, 149, 165, + 166, 150, 167, 168, 181, 182, 183, 123, 131, 132, 133, 141, 142, 151, 152, 153, 169, 170, 154, 155, 171, + 172, 156, 173, 174, 184, 185, 186, 190, 191, 192, 210, 211, 193, 194, 212, 213, 195, 214, 215, 234, 235, + 236, 196, 197, 216, 217, 198, 218, 219, 237, 238, 239, 199, 220, 221, 240, 241, 242, 252, 253, 254, 255, + 124, 134, 135, 136, 143, 144, 157, 158, 159, 175, 176, 160, 161, 177, 178, 162, 179, 180, 187, 188, 189, + 200, 201, 202, 222, 223, 203, 204, 224, 225, 205, 226, 227, 243, 244, 245, 206, 207, 228, 229, 208, 230, + 231, 246, 247, 248, 209, 232, 233, 249, 250, 251, 256, 257, 258, 259, 260, 261, 262, 275, 276, 263, 264, + 277, 278, 265, 279, 280, 295, 296, 297, 266, 267, 281, 282, 268, 283, 284, 298, 299, 300, 269, 285, 286, + 301, 302, 303, 313, 314, 315, 316, 270, 271, 287, 288, 272, 289, 290, 304, 305, 306, 273, 291, 292, 307, + 308, 309, 317, 318, 319, 320, 274, 293, 294, 310, 311, 312, 321, 322, 323, 324, 325, 326, 327, 328, 329, + ]), + // Laplacian-dependent functionals are not supported yet + (LAPL, _) => unimplemented!(), + // We currently only support up to 4th order derivatives (exc, vxc, kxc, fxc, lxc) + (_, 5..) => unimplemented!(), + } +} + +#[doc = include_str!("libxc-xcfun-trans.md")] +pub fn libxc_to_xcfun_mapping(den_type: XCDenType, spin: LibXCSpin, deriv: usize) -> Option> { + // some cases that do not require reordering + if deriv <= 1 || spin == Unpolarized || den_type == RHO { + return None; + } + // now we assume spin polarized and order >= 2 + Some((0..=deriv).flat_map(|x| libxc_to_xcfun_mapping_parts(den_type, x).unwrap().to_vec()).collect()) +} + +pub fn libxc_transform_xcfun_indices( + xc0: TsrView<'_>, + den_type: XCDenType, + spin: LibXCSpin, + deriv: usize, +) -> TsrCow<'_> { + // sanity check + assert!(xc0.ndim() == 2, "xc0 must be a 2-dim tensor"); + let indices = libxc_to_xcfun_mapping(den_type, spin, deriv); + if let Some(indices) = indices { + xc0.index_select(-1, &indices).into_cow() + } else { + xc0.into_cow() + } +} + +/* #endregion libxc-to-xcfun index convention change */ + +/* #region xlen utility */ + +/// Get the number of components (xlen) for a given density type and spin polarization. +/// +/// This value should match the first derivative length. +pub const fn get_xc_xlen(den_type: XCDenType, spin: LibXCSpin) -> usize { + match (den_type, spin) { + (RHO, Unpolarized) => 1, + (RHO, Polarized) => 2, + (SIGMA, Unpolarized) => 2, + (SIGMA, Polarized) => 5, + (TAU, Unpolarized) => 3, + (TAU, Polarized) => 7, + (LAPL, _) => unimplemented!(), + } +} + +/// Generates raveled unique indices for the Cartesian product of a given number +/// of variables and order. +pub fn product_uniq_indices(xlen: usize, order: usize) -> Vec { + // Generate all unique combinations with replacement + let uniq_idx: Vec> = + (0..xlen).combinations_with_replacement(order).map(|v| v.into_iter().collect()).collect(); + + // Create a mapping from sorted indices to their position in uniq_idx + let mut index_map = std::collections::HashMap::new(); + for (pos, indices) in uniq_idx.iter().enumerate() { + index_map.insert(indices.clone(), pos); + } + + // Generate all possible Cartesian product indices + let cartesian_product = (0..order).map(|_| 0..xlen).multi_cartesian_product(); + + // For each index in the Cartesian product, find its sorted version and lookup + // the unique position + cartesian_product + .map(|indices| { + let mut sorted = indices.clone(); + sorted.sort(); + *index_map.get(&sorted).unwrap() + }) + .collect() +} + +/* #endregion xlen utility */ + +/* #region unfold sigma */ + +pub fn vxc_unfold_sigma_spin0( + frho: &mut [f64], + fsigma: &[f64], + rho: &[f64], + ncounts: usize, + nvar: usize, + ngrids: usize, +) { + let ncg = ncounts * ngrids; + let nvg = nvar * ngrids; + + // Define accessor macros matching the C version's pattern + macro_rules! fr_at { + // f_rho index + ($g:expr, $x:expr, $n:expr) => { + frho[$g + $x * ngrids + $n * nvg] + }; + } + macro_rules! fs_at { + // f_rho index + ($g:expr, $n:expr, $x:expr) => { + fsigma[$g + $n * ngrids + $x * ncg] + }; + } + macro_rules! rho_at { + ($g:expr, $x:expr) => { + rho[$g + $x * ngrids] + }; + } + + for n in 0..ncounts { + for g in 0..ngrids { + // Main computation block + fr_at!(g, 0, n) = fs_at!(g, n, 0); + fr_at!(g, 1, n) = fs_at!(g, n, 1) * rho_at!(g, 1) * 2.0; + fr_at!(g, 2, n) = fs_at!(g, n, 1) * rho_at!(g, 2) * 2.0; + fr_at!(g, 3, n) = fs_at!(g, n, 1) * rho_at!(g, 3) * 2.0; + } + } + + if nvar > 4 { + assert_eq!(nvar, 5, "MGGA case requires exactly 5 variables"); + for n in 0..ncounts { + for g in 0..ngrids { + fr_at!(g, 4, n) = fs_at!(g, n, 2); + } + } + } +} + +pub fn vxc_unfold_sigma_spin1( + frho: &mut [f64], + fsigma: &[f64], + rho: &[f64], + ncounts: usize, + nvar: usize, + ngrids: usize, +) { + let ncg = ncounts * ngrids; + let nvg = nvar * ngrids; + + // Helper macros to access the arrays by indices + macro_rules! fr_at { + // f_rho index + ($g:expr, $x:expr, $a:expr, $n:expr) => { + frho[$g + $x * ngrids + ($a + $n * 2) * nvg] + }; + } + macro_rules! fs_at { + // f_sigma index + ($g:expr, $n:expr, $x:expr) => { + fsigma[$g + $n * ngrids + $x * ncg] + }; + } + macro_rules! rho_at { + ($g:expr, $x:expr, $a:expr) => { + rho[$g + $x * ngrids + $a * nvg] + }; + } + + for n in 0..ncounts { + for g in 0..ngrids { + // Main computation block + fr_at!(g, 0, 0, n) = fs_at!(g, n, 0); + fr_at!(g, 0, 1, n) = fs_at!(g, n, 1); + fr_at!(g, 1, 0, n) = fs_at!(g, n, 2) * rho_at!(g, 1, 0) * 2.0 + fs_at!(g, n, 3) * rho_at!(g, 1, 1); + fr_at!(g, 1, 1, n) = fs_at!(g, n, 3) * rho_at!(g, 1, 0) + 2.0 * fs_at!(g, n, 4) * rho_at!(g, 1, 1); + fr_at!(g, 2, 0, n) = fs_at!(g, n, 2) * rho_at!(g, 2, 0) * 2.0 + fs_at!(g, n, 3) * rho_at!(g, 2, 1); + fr_at!(g, 2, 1, n) = fs_at!(g, n, 3) * rho_at!(g, 2, 0) + 2.0 * fs_at!(g, n, 4) * rho_at!(g, 2, 1); + fr_at!(g, 3, 0, n) = fs_at!(g, n, 2) * rho_at!(g, 3, 0) * 2.0 + fs_at!(g, n, 3) * rho_at!(g, 3, 1); + fr_at!(g, 3, 1, n) = fs_at!(g, n, 3) * rho_at!(g, 3, 0) + 2.0 * fs_at!(g, n, 4) * rho_at!(g, 3, 1); + } + } + + if nvar > 4 { + assert_eq!(nvar, 5, "MGGA case requires exactly 5 variables"); + for n in 0..ncounts { + for g in 0..ngrids { + fr_at!(g, 4, 0, n) = fs_at!(g, n, 5); + fr_at!(g, 4, 1, n) = fs_at!(g, n, 6); + } + } + } +} + +pub fn unfold_sigma( + rho: TsrView, + xc_val: TsrView, + spin: LibXCSpin, + order: usize, + nvar: usize, + xlen: usize, + reserve: usize, +) -> Tsr { + assert!(nvar >= 4); + let nvar_spin = if spin == Unpolarized { nvar } else { 2 * nvar }; + let ngrids = rho.shape()[0]; + // check dimensions + assert!(xc_val.shape()[0] == ngrids, "xc_val length mismatch"); + assert!(xc_val.ndim() == 2, "xc_val must be a 2-dim tensor"); + match spin { + Unpolarized => assert!(rho.shape() == &[ngrids, nvar]), + Polarized => assert!(rho.shape() == &[ngrids, nvar, 2]), + }; + + let n_transform = order - reserve; + let mut xc_tensor_shape = vec![ngrids]; + match spin { + Unpolarized => xc_tensor_shape.extend(vec![nvar; n_transform]), + Polarized => xc_tensor_shape.extend(vec![[nvar, 2]; n_transform].iter().flatten()), + } + xc_tensor_shape.extend(vec![xlen; reserve]); + let mut xc_tensor: Tsr = unsafe { rt::empty((xc_tensor_shape, xc_val.device())) }; + + let idx = product_uniq_indices(xlen, order); + for (it, &io) in idx.iter().enumerate() { + // please note that currently, RSTSR's `tensor.raw()` returns the pointer + // (slice) of original data, instead of offsetted pointer points to the first + // element of tensor. + // So we need additionally define an offsetted slice. + let xc_val_offsetted = &xc_val.raw()[xc_val.offset()..]; + xc_tensor.raw_mut()[it * ngrids..(it + 1) * ngrids] + .copy_from_slice(&xc_val_offsetted[io * ngrids..(io + 1) * ngrids]); + } + + // also note the raw usage, rho is not assured to be offset-zero. + let rho_raw = &rho.raw()[rho.offset()..]; + + let mut buf = unsafe { xc_tensor.empty_like() }; + for i in 0..n_transform { + std::mem::swap(&mut xc_tensor, &mut buf); + let ncounts = xlen.pow((order - 1 - i) as u32) * nvar_spin.pow(i as u32); + + match spin { + Unpolarized => vxc_unfold_sigma_spin0(xc_tensor.raw_mut(), buf.raw(), rho_raw, ncounts, nvar, ngrids), + Polarized => vxc_unfold_sigma_spin1(xc_tensor.raw_mut(), buf.raw(), rho_raw, ncounts, nvar, ngrids), + } + } + + xc_tensor +} + +/* #endregion unfold sigma */ + +#[allow(clippy::deref_addrof)] +pub fn transform_xc_inner(rho: TsrView, xc_val: TsrView, den_type: XCDenType, spin: LibXCSpin, order: usize) -> Tsr { + if order >= 4 { + panic!("currently only support order < 4 (exc, vxc, kxc, fxc). You specified order {order}"); + } + + // sanity check for dimensions + let ngrids = rho.shape()[0]; + let nvar = den_type.num_nvar(); + let xlen = get_xc_xlen(den_type, spin); + // check dimensions + ni_check_shape!(xc_val.shape()[0], ngrids, "xc_val length (grids) mismatch"); + ni_check_shape!(xc_val.ndim(), 2, "xc_val must be a 2-dim tensor"); + // check shape [ngrids, nvar, nspin if exist], otherwise panic + match spin { + Unpolarized => ni_check_shape!(rho.ndim(), 2, "rho must be a 2-dim tensor"), + Polarized => { + ni_check_shape!(rho.ndim(), 3, "rho must be a 3-dim tensor"); + ni_check_shape!(rho.shape()[2], 2, "rho last dimension should be 2 for polarized case"); + }, + }; + ni_check_shape!(rho.shape()[0], ngrids, "rho first dimension must be grids"); + ni_check_shape!(rho.shape()[1] >= nvar, "rho second dimension (variables) should be larger than {nvar}"); + let rho = rho.change_contig(ColMajor); + let xc_val = xc_val.to_contig(ColMajor); + // double check input tensor + // since we are using some raw functionality to get the raw slice, the offset must be zero. + + // offsets of xc_val + let mut offsets = vec![0]; + offsets.extend((0..=order).map(|o| count_combinations(xlen + o, o))); + let offset_max = offsets.last().unwrap(); + ni_check_shape!(xc_val.shape()[1] >= *offset_max, "xc_val length should be larger than {offset_max}"); + + // offsets match current order + let (p0, p1) = (offsets[order], offsets[order + 1]); + + // quick return for LDA + if den_type == RHO { + let xc_out = xc_val.i((.., p0..p1)); + if spin == Unpolarized { + // shape: [ngrids, 1, 1, ..., 1] + // | [1]*order | + let mut shape = vec![ngrids]; + shape.extend(vec![1; order]); + return xc_out.into_shape(shape); + } else { + let indices = product_uniq_indices(xlen, order); + let xc_out = xc_out.index_select(-1, &indices); + // shape: [ngrids, 1, 2, 1, 2, ..., 1, 2] + // | [1, 2] * order | + let mut shape = vec![ngrids]; + shape.extend(vec![[1, 2]; order].into_iter().flatten()); + return xc_out.into_shape(shape); + } + } + + let mut xc_tensor = unfold_sigma(rho.view(), xc_val.i((.., p0..p1)), spin, order, nvar, xlen, 0); + + if order <= 1 { + // quick return for 0/1-order derivatives, which does not involve pair derivatives of sigma + return xc_tensor; + } + + if spin == Unpolarized { + // currently we can only handle order = 2, 3 cases + // for order > 3, following code is not correct + let n_pairs = 1; // only correct for order = 2, 3 + let (p0, p1) = (offsets[order - n_pairs], offsets[order - n_pairs + 1]); + let xc_sub = unfold_sigma(rho.view(), xc_val.i((.., p0..p1)), spin, order - n_pairs, nvar, xlen, n_pairs); + let xc_sub: Tsr = 2.0 * xc_sub.i((Ellipsis, 1)); + match order { + 2 => *&mut xc_tensor.i_mut((.., 1..4, 1..4)).diagonal_mut((0, -1, -2)) += xc_sub, + 3 => { + let permute_order_list = [[0, 1, 2, 3], [0, 2, 3, 1], [0, 3, 1, 2]]; + for permute_order in permute_order_list { + let mut xc_tensor_perm = xc_tensor.view_mut().into_transpose(&permute_order); + *&mut xc_tensor_perm.i_mut((Ellipsis, 1..4, 1..4)).diagonal_mut((0, -1, -2)) += &xc_sub; + } + }, + _ => unreachable!(), + } + } else { + // currently we can only handle order = 2, 3 cases + // for order > 3, following code is not correct + let n_pairs = 1; // only correct for order = 2, 3 + let (p0, p1) = (offsets[order - n_pairs], offsets[order - n_pairs + 1]); + let xc_sub = unfold_sigma(rho.view(), xc_val.i((.., p0..p1)), spin, order - n_pairs, nvar, xlen, n_pairs); + // just the sigma components, spin expanded + let xc_sub = xc_sub.i((Ellipsis, 2..5)); + let xc_sub = xc_sub.index_select(-1, &[0, 1, 1, 2]); + let xc_sub_shape = { + let mut xc_sub_shape = xc_sub.shape().clone(); + xc_sub_shape.pop(); + xc_sub_shape.extend(vec![2, 2]); + xc_sub_shape + }; + let mut xc_sub = xc_sub.into_shape(xc_sub_shape); + *&mut xc_sub.i_mut((Ellipsis, 0, 0)) *= 2.0; + *&mut xc_sub.i_mut((Ellipsis, 1, 1)) *= 2.0; + match order { + 2 => { + let permute_spin = [0, 2, 4, 1, 3]; + let mut xc_tensor_spin = xc_tensor.view_mut().into_transpose(&permute_spin); + // the case of order=2 does not require xc_sub to permute by spin indices + *&mut xc_tensor_spin.i_mut((Ellipsis, 1..4, 1..4)).diagonal_mut((0, -1, -2)) += &xc_sub; + }, + 3 => { + let xc_tensor_permute_spin = [0, 2, 4, 6, 1, 3, 5]; + let mut xc_tensor_spin = xc_tensor.view_mut().into_transpose(&xc_tensor_permute_spin); + let xc_sub_permute_spin = [0, 2, 3, 4, 1]; + let xc_sub_spin = xc_sub.transpose(&xc_sub_permute_spin); + + let permute_order_list = [[0, 1, 2, 3, 4, 5, 6], [0, 2, 3, 1, 5, 6, 4], [0, 3, 1, 2, 6, 4, 5]]; + for permute_order in permute_order_list { + let mut xc_tensor_perm = xc_tensor_spin.view_mut().into_transpose(&permute_order); + *&mut xc_tensor_perm.i_mut((Ellipsis, 1..4, 1..4)).diagonal_mut((0, -1, -2)) += &xc_sub_spin; + } + }, + _ => unreachable!(), + } + } + + xc_tensor +} diff --git a/src/hessian_backup/cint_handling.rs b/src/hessian_backup/cint_handling.rs new file mode 100644 index 0000000000..a5a3605d86 --- /dev/null +++ b/src/hessian_backup/cint_handling.rs @@ -0,0 +1,115 @@ +use super::prelude::*; +use rest_libcint::util::ShlsSlice; + +/// A wrapper around [`CInt::integrate`] that directly transform the output and shape to tensor. +/// +/// This only handles single molecule integrals. For cross integrals, see [`hess_intor_cross`]. +/// +/// # Notes +/// +/// The following notes also apply to [`hess_intor_cross`]. +/// +/// Note that this wrapper **only works in hessian-related tasks**. We will transform integrators +/// like `int2c2e_ipip1` original shape `[naux, naux, 9]` to [naux, naux, 3, 3]` to make it more +/// intuitive to use. For other integrators, the shape will be unchanged. +/// +/// Also note that, for second order derivative, for example of `int3c2e_ip1ip2` $(\partial_t \mu +/// \nu | \partial_s P)$, the returned shape is `[nao, nao, naux, 3, 3]`, denoting the indices of +/// $(\mu, \nu, P, s, t)$. Please be very careful about the last two dimensions, which are of +/// indices `[s, t]` for column major. +pub fn hess_intor( + mol: &CInt, + intor_name: &str, + symm: &str, + shls_slice: impl Into, + device: &DeviceBLAS, +) -> Tsr { + let (out, shape) = mol.integrate(intor_name, symm, shls_slice.into()).into(); + let ip_matches = intor_name.matches("ip").count(); + let shape = match ip_matches { + 0 | 1 => shape, + 2 => { + // check last dimension is 9 + assert_eq!(shape.last(), Some(&9), "For integrator with 2 'ip' in name, the last dimension should be 9."); + // transform last dimension from 9 to (3, 3) + let mut new_shape = shape.clone(); + new_shape.pop(); + new_shape.push(3); + new_shape.push(3); + new_shape + }, + _ => panic!("Unsupported integrator with more than 2 'ip' in name."), + }; + rt::asarray((out, shape, device)) +} + +/// A wrapper around [`CInt::integrate_cross`] that directly transform the output and shape to +/// tensor. +/// +/// Notes see also [`hess_intor`]. +pub fn hess_intor_cross( + mol_list: &[&CInt], + intor_name: &str, + symm: &str, + shls_slice: impl Into, + device: &DeviceBLAS, +) -> Tsr { + let (out, shape) = CInt::integrate_cross(intor_name, mol_list, symm, shls_slice.into()).into(); + let ip_matches = intor_name.matches("ip").count(); + let shape = match ip_matches { + 0 | 1 => shape, + 2 => { + // check last dimension is 9 + assert_eq!(shape.last(), Some(&9), "For integrator with 2 'ip' in name, the last dimension should be 9."); + // transform last dimension from 9 to (3, 3) + let mut new_shape = shape.clone(); + new_shape.pop(); + new_shape.push(3); + new_shape.push(3); + new_shape + }, + _ => panic!("Unsupported integrator with more than 2 'ip' in name."), + }; + rt::asarray((out, shape, device)) +} + +/// A wrapper that generates 3c-2e ERIs. +/// +/// This returns a closure that takes `shls_aux` as input. So, the batch is always on the auxiliary +/// shell dimension. Also note input is shell index, not AO index. +pub fn generator_hess_intor_j3c_by_aux<'a>( + mol: &'a CInt, + aux: &'a CInt, + intor_name: &'a str, + symm: &'a str, + device: &DeviceBLAS, +) -> impl Fn([usize; 2]) -> Tsr + 'a { + let shls_mol = [0, mol.nbas()]; + let device = device.clone(); + move |shls_aux: [usize; 2]| { + hess_intor_cross(&[mol, mol, aux], intor_name, symm, [shls_mol, shls_mol, shls_aux], &device) + } +} + +pub fn get_ecp_atoms(mol: &CInt) -> Vec { + const ATOM_OF: usize = rest_libcint::ffi::cint_ffi::ATOM_OF as usize; + // remove duplicates and sort + mol.ecpbas.iter().map(|&ecpbas| ecpbas[ATOM_OF] as usize).sorted().dedup().collect_vec() +} + +/// Filter the atom-slice array according to an optional list of atom indices. +/// +/// When `atm_list` is `None`, returns the full slices (length `mol.natm()`) and all indices. +/// When `atm_list` is `Some(&[i, j, ...])`, returns slices for only those atoms (length +/// `list.len()`) and the same list as the index mapping (local → global). +pub fn filter_aoslices(mol: &CInt, atm_list: Option<&[usize]>) -> (Vec<[usize; 4]>, Vec) { + let full_slices = mol.aoslice_by_atom(); + match atm_list { + None => (full_slices, (0..mol.natm()).collect()), + Some(list) => { + let slices = list.iter().map(|&i| full_slices[i]).collect(); + let indices = list.to_vec(); + (slices, indices) + }, + } +} diff --git a/src/hessian_backup/hcore.rs b/src/hessian_backup/hcore.rs new file mode 100644 index 0000000000..a60b1842c2 --- /dev/null +++ b/src/hessian_backup/hcore.rs @@ -0,0 +1,231 @@ +use super::prelude::*; + +/// Generator for second derivatives of the core Hamiltonian (skeleton derivative). +/// +/// # Parameters +/// +/// - `mol` : [`CInt`]. The molecule object. +/// - `device` : [`DeviceBLAS`]. The device on which the returned tensor is allocated. +/// +/// # Returns +/// +/// - `FnMut(A: usize, B: usize) -> Tsr`. A function that computes the second derivative of the core +/// Hamiltonian with respect to the nuclear coordinates. The returned array has shape [nao, nao, +/// 3, 3]. +pub fn generator_hcore_deriv2(mol: &CInt, device: &DeviceBLAS) -> impl FnMut(usize, usize) -> Tsr { + // preparation + let device = device.clone(); + let mut mol = mol.clone(); + let nao = mol.nao(); + let nbas = mol.nbas(); + let ecp_atoms = get_ecp_atoms(&mol); + let aoslices = mol.aoslice_by_atom(); + + // we need to prepare some integrals, to somehow avoid redundant calculations in the loop + // - aa: Hamiltonian derivative to only the first basis + // - ab: Hamiltonian derivative to the first and second basis + // all integrals are of shape [nao, nao, 3, 3] + let mut h2_aa = hess_intor(&mol, "int1e_ipipkin", "s1", None, &device); + let mut h2_ab = hess_intor(&mol, "int1e_ipkinip", "s1", None, &device); + h2_aa += hess_intor(&mol, "int1e_ipipnuc", "s1", None, &device); + h2_ab += hess_intor(&mol, "int1e_ipnucip", "s1", None, &device); + if mol.has_ecp() { + h2_aa += hess_intor(&mol, "ECPscalar_ipipnuc", "s1", None, &device); + h2_ab += hess_intor(&mol, "ECPscalar_ipnucip", "s1", None, &device); + } + + move |A: usize, B: usize| { + let [sh0A, sh1A, p0A, p1A] = aoslices[A]; + let [sh0B, sh1B, p0B, p1B] = aoslices[B]; + let slcA = rt::slice!(p0A, p1A); + let slcB = rt::slice!(p0B, p1B); + let zA = mol.atom_charge(A); + let zB = mol.atom_charge(B); + + let mut hcore_deriv: Tsr = rt::zeros(([nao, nao, 3, 3], &device)); + + if A == B { + mol.with_rinv_at_nucleus(A, |mol| { + let mut rinv_aa = -zA * hess_intor(mol, "int1e_ipiprinv", "s1", None, &device); + let mut rinv_ab = -zA * hess_intor(mol, "int1e_iprinvip", "s1", None, &device); + if ecp_atoms.contains(&A) { + rinv_aa += hess_intor(mol, "ECPscalar_ipiprinv", "s1", None, &device); + rinv_ab += hess_intor(mol, "ECPscalar_iprinvip", "s1", None, &device); + } + hcore_deriv += &rinv_aa; + hcore_deriv += &rinv_ab; + *&mut hcore_deriv.i_mut((slcA, ..)) -= rinv_aa.i((slcA, ..)); + *&mut hcore_deriv.i_mut((slcA, ..)) -= rinv_ab.i((slcA, ..)); + *&mut hcore_deriv.i_mut((.., slcA)) -= rinv_aa.i((slcA, ..)).swapaxes(0, 1); + *&mut hcore_deriv.i_mut((.., slcA)) -= rinv_ab.i((.., slcA)); + }); + *&mut hcore_deriv.i_mut((slcA, ..)) += h2_aa.i((slcA, ..)); + *&mut hcore_deriv.i_mut((slcA, slcA)) += h2_ab.i((slcA, slcA)); + } else { + *&mut hcore_deriv.i_mut((slcA, slcB)) += h2_ab.i((slcA, slcB)); + mol.with_rinv_at_nucleus(A, |mol| { + let shls_slice = [[sh0B, sh1B], [0, nbas]]; + let mut rinv_atom_aa = -zA * hess_intor(mol, "int1e_ipiprinv", "s1", shls_slice, &device); + let mut rinv_atom_ab = -zA * hess_intor(mol, "int1e_iprinvip", "s1", shls_slice, &device); + if ecp_atoms.contains(&A) { + rinv_atom_aa += hess_intor(mol, "ECPscalar_ipiprinv", "s1", shls_slice, &device); + rinv_atom_ab += hess_intor(mol, "ECPscalar_iprinvip", "s1", shls_slice, &device); + } + *&mut hcore_deriv.i_mut((slcB, ..)) -= &rinv_atom_aa; + *&mut hcore_deriv.i_mut((slcB, ..)) -= &rinv_atom_ab.swapaxes(-1, -2); + }); + mol.with_rinv_at_nucleus(B, |mol| { + let shls_slice = [[sh0A, sh1A], [0, nbas]]; + let mut rinv_atom_aa = -zB * hess_intor(mol, "int1e_ipiprinv", "s1", shls_slice, &device); + let mut rinv_atom_ab = -zB * hess_intor(mol, "int1e_iprinvip", "s1", shls_slice, &device); + if ecp_atoms.contains(&B) { + rinv_atom_aa += hess_intor(mol, "ECPscalar_ipiprinv", "s1", shls_slice, &device); + rinv_atom_ab += hess_intor(mol, "ECPscalar_iprinvip", "s1", shls_slice, &device); + } + *&mut hcore_deriv.i_mut((slcA, ..)) -= &rinv_atom_aa; + *&mut hcore_deriv.i_mut((slcA, ..)) -= &rinv_atom_ab; + }); + } + + &hcore_deriv + hcore_deriv.swapaxes(0, 1) + } +} + +/// Hessian contribution from the core Hamiltonian. +/// +/// # Parameters +/// +/// - `mol` : [`CInt`]. The molecule object. +/// - `dm0` : shape [nao, nao]. The SCF density matrix. +/// - `atm_list` : optional list of atom indices to compute the Hessian for. If `None`, all atoms +/// are computed. +/// +/// # Returns +/// +/// - `de_hcore` : shape `[3, 3, natm, natm]`. The Hessian contribution from the core Hamiltonian. +pub fn get_hess_hcore(mol: &CInt, dm0: TsrView, atm_list: Option<&[usize]>) -> Tsr { + let device = dm0.device(); + let nao = mol.nao(); + let (_aoslices, atm_indices) = filter_aoslices(mol, atm_list); + let natm = atm_indices.len(); + + assert_eq!(dm0.shape(), &[nao, nao], "density matrix shape not correct."); + + let mut de_hcore = rt::zeros(([3, 3, natm, natm], device)); + let mut gen_hcore_deriv2 = generator_hcore_deriv2(mol, device); + + for A in 0..natm { + let A_glob = atm_indices[A]; + for B in 0..=A { + let B_glob = atm_indices[B]; + let hcore_deriv2 = gen_hcore_deriv2(A_glob, B_glob); + *&mut de_hcore.i_mut((.., .., B, A)) += (hcore_deriv2 * &dm0).sum_axes((0, 1)); + } + for B in 0..A { + let de_to_copy = de_hcore.i((.., .., B, A)).t().to_owned(); + *&mut de_hcore.i_mut((.., .., A, B)) += de_to_copy; + } + } + de_hcore +} + +/// Generator for first derivatives of the core Hamiltonian (skeleton derivative). +/// +/// # Parameters +/// +/// - `mol` : [`CInt`]. The molecule object. +/// - `device` : [`DeviceBLAS`]. The device object. +/// +/// # Returns +/// +/// - `get_hcore_deriv_at_atoms` : `FnMut(A: usize) -> Tsr`. A function that computes the first +/// derivative of the core Hamiltonian with respect to the nuclear coordinates. Input is the atom +/// index A. The returned array has shape `[nao, nao, 3]`. +pub fn generator_hcore_deriv1(mol: &CInt, device: &DeviceBLAS) -> impl FnMut(usize) -> Tsr { + // preparation + let device = device.clone(); + let mut mol = mol.clone(); + let nao = mol.nao(); + let ecp_atoms = get_ecp_atoms(&mol); + let aoslices = mol.aoslice_by_atom(); + + let mut h1_a = -hess_intor(&mol, "int1e_ipkin", "s1", None, &device); + h1_a -= hess_intor(&mol, "int1e_ipnuc", "s1", None, &device); + if mol.has_ecp() { + h1_a -= hess_intor(&mol, "ECPscalar_ipnuc", "s1", None, &device); + } + + move |A: usize| { + let [_, _, p0A, p1A] = aoslices[A]; + let slcA = rt::slice!(p0A, p1A); + let zA = mol.atom_charge(A); + + let mut hcore_deriv: Tsr = rt::zeros(([nao, nao, 3], &device)); + *&mut hcore_deriv.i_mut(slcA) += h1_a.i(slcA); + mol.with_rinv_at_nucleus(A, |mol| { + let mut rinv_a = -zA * hess_intor(mol, "int1e_iprinv", "s1", None, &device); + if ecp_atoms.contains(&A) { + rinv_a += hess_intor(mol, "ECPscalar_iprinv", "s1", None, &device); + } + hcore_deriv += &rinv_a; + }); + &hcore_deriv + hcore_deriv.swapaxes(0, 1) + } +} + +/// Hessian contribution from the core Hamiltonian. +pub struct RHessHcore { + pub mol: CInt, + pub device: DeviceBLAS, +} + +impl RHessHcore { + pub fn new(mol: &CInt, device: &DeviceBLAS) -> Self { + Self { mol: mol.clone(), device: device.clone() } + } +} + +impl HessUtilAPI for RHessHcore {} + +impl RHessCoreAPI for RHessHcore { + fn make_skeleton_hess(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr { + let dm0 = get_dm0_restricted(mo_coeff, mo_occ); + get_hess_hcore(&self.mol, dm0.view(), atm_list) + } + + fn generator_deriv1(&self) -> Box Tsr> { + Box::new(generator_hcore_deriv1(&self.mol, &self.device)) + } +} + +/// Hessian contribution from the core Hamiltonian (unrestricted version). +pub struct UHessHcore { + pub mol: CInt, + pub device: DeviceBLAS, +} + +impl UHessHcore { + pub fn new(mol: &CInt, device: &DeviceBLAS) -> Self { + Self { mol: mol.clone(), device: device.clone() } + } +} + +impl HessUtilAPI for UHessHcore {} + +impl UHessCoreAPI for UHessHcore { + fn make_skeleton_hess( + &mut self, + mo_coeff: &[TsrView; 2], + mo_occ: &[TsrView; 2], + atm_list: Option<&[usize]>, + ) -> Tsr { + let [α, β] = [0, 1]; + let dm0 = get_dm0_restricted(mo_coeff[α].view(), mo_occ[α].view()) + + get_dm0_restricted(mo_coeff[β].view(), mo_occ[β].view()); + get_hess_hcore(&self.mol, dm0.view(), atm_list) + } + + fn generator_deriv1(&self) -> Box Tsr> { + Box::new(generator_hcore_deriv1(&self.mol, &self.device)) + } +} diff --git a/src/hessian_backup/krylov_block.rs b/src/hessian_backup/krylov_block.rs new file mode 100644 index 0000000000..22b24bf848 --- /dev/null +++ b/src/hessian_backup/krylov_block.rs @@ -0,0 +1,258 @@ +//! Block Krylov subspace solver for `(1 + A) x = b`. +//! +//! The convergence behavior should approximately matches PySCF's `lib.krylov` block algorithm: each +//! cycle adds the surviving trial directions to the subspace, and the basis is kept non-normalized +//! so that the squared norms of the new trial vectors act as the convergence signal without +//! requiring an explicit residual evaluation. +//! +//! To bound memory the subspace is capped at `max_space` cycles. When that cap is hit without +//! convergence the solver performs a **hard restart** (GMRES(m)-style): the projected system is +//! solved to obtain the current best approximation, that approximation is folded into a running +//! `x_accum`, the subspace is reset, and the residual `b - (I+A) x_accum` becomes the new RHS. +//! +//! Layout convention is col-major: each right-hand side / basis vector is a **column** of the +//! corresponding matrix. So `b` is shaped `[n, nset]`, the operator maps `[n, nblock] -> [n, +//! nblock]`, and the basis matrices `xs`, `ax` are `[n, nd]` with new vectors appended along axis +//! 1. +//! +//! # Note +//! +//! This file is generated by AI. Though this function has been tested, I have not carefully +//! verified all details. + +use super::prelude::*; + +/// Solve `(I + aop) x = b` by a block Krylov subspace method with hard restarts. +/// +/// # Parameters +/// +/// - `aop` : Linear operator. Given a `[n, nblock]` input it must return a `[n, nblock]` output +/// (the action of `A` applied column-wise). +/// - `b` : Right-hand sides, shape `[n, nset]`. Each column is one RHS. +/// - `x0` : Optional initial guess, shape `[n, nset]`. Zero initial guess is used if not provided. +/// - `tol` : Convergence tolerance on `max(||new_trial_vec_i||)`. +/// - `max_cycle` : Maximum **total** number of inner cycles, summed across restarts. Recommended +/// value is 54, and is better to be a multiple (or much larger) than `max_space`. +/// - `max_space` : Maximum subspace size in cycles before a hard restart is triggered. Typical +/// values are 6..=20, recommended 14 for CP-HF problems. With `max_space >= max_cycle` no restart +/// ever happens (matches the pre-restart behavior). Storage is `O(n * nset * (max_space + 1))`. +/// - `lindep` : Vectors with `||v||^2 < lindep` are dropped from the subspace. +/// +/// # Returns +/// +/// `x` of shape `[n, nset]`, an approximate solution of `(I + aop) x = b`. +pub fn krylov_block( + mut aop: impl FnMut(TsrView) -> Tsr, + b: TsrView, + x0: Option, + tol: f64, + max_cycle: usize, + max_space: usize, + lindep: f64, +) -> Tsr { + let device = b.device().clone(); + let n = b.shape()[0]; + let nset = b.shape()[1]; + + let b_orig = b.to_owned(); + + // x_accum plays the role of a running initial guess that is refined on each + // hard restart. After every restart we re-form the residual b - (I+A) x_accum + // and rebuild the Krylov subspace from scratch. + let mut x_accum: Tsr = match x0.as_ref() { + Some(x0v) => x0v.to_owned(), + None => rt::zeros(([n, nset], &device)), + }; + + // Pre-allocate basis storage at the bounded restart size. The slabs are + // overwritten in place on each restart by resetting `nd = 0`. + let max_basis = nset * (max_space + 1); + let mut xs: Tsr = rt::zeros(([n, max_basis], &device)); + let mut ax: Tsr = rt::zeros(([n, max_basis], &device)); + let mut all_innerprod: Vec = Vec::with_capacity(max_basis); + + let conv_thresh = lindep.max(tol * tol); + + let mut total_cycles: usize = 0; + let mut restart_idx: usize = 0; + + // We must remember the last completed inner loop's subspace and the residual + // that produced it, so that the final projected solve uses a consistent pair. + let mut last_nd: usize = 0; + let mut last_b: Tsr = b_orig.clone(); + // Track the most recent convergence signal so we can detect non-convergence + // at the end. Initialised to 0.0 so the "empty residual" early-exit at the + // top of the outer loop is treated as already-converged. + let mut last_max_innerprod: f64 = 0.0; + + while total_cycles < max_cycle { + // Form this restart's RHS: b - (I+A) x_accum. On the first pass with no + // initial guess this is just b (skip the extra aop evaluation). + let b_residual: Tsr = + if restart_idx == 0 && x0.is_none() { b_orig.clone() } else { &b_orig - (&x_accum + aop(x_accum.view())) }; + + // Orthogonalize the columns of the residual. + let (mut x1, mut innerprod) = orth_block(b_residual.view(), lindep); + + if x1.shape()[1] == 0 { + // Residual is already (numerically) zero; x_accum is the answer. + last_nd = 0; + last_b = b_residual; + break; + } + + // Reset the subspace for this restart. + all_innerprod.clear(); + let mut nd: usize = 0; + let mut inner_converged = false; + + for inner in 0..max_space { + if total_cycles >= max_cycle { + break; + } + total_cycles += 1; + let nblock = x1.shape()[1]; + + let axt = aop(x1.view()); + + xs.i_mut((.., nd..nd + nblock)).assign(&x1); + ax.i_mut((.., nd..nd + nblock)).assign(&axt); + all_innerprod.extend_from_slice(&innerprod); + nd += nblock; + + // Orthogonalize axt against the full subspace; same algebra as before. + let xs_slc = xs.i((.., ..nd)); + let ip_vec = rt::asarray((&all_innerprod, &device)); + let coeffs = (xs_slc.t() % &axt) / ip_vec.i((.., None)); + let x1_new = axt - &xs_slc % &coeffs; + + // MGS new directions: keep at the numerical-zero floor; defer + // the user's lindep filter so borderline vectors still feed + // the projected solve via xs / ax. + let (mut next_x1, mut next_ip) = orth_block(x1_new.view(), lindep); + + // Convergence test uses the unfiltered max innerprod. + let max_innerprod = next_ip.iter().copied().fold(0.0_f64, f64::max); + let r = max_innerprod.sqrt(); + + // Filter the next-iteration directions to those above conv_thresh. + let keep_mask: Vec = next_ip.iter().map(|&ip| ip > conv_thresh).collect(); + if keep_mask.iter().any(|&b| !b) { + next_x1 = next_x1.bool_select(1, &keep_mask); + next_ip = keep_mask + .iter() + .zip(next_ip.iter()) + .filter_map(|(&m, &ip)| if m { Some(ip) } else { None }) + .collect(); + } + + // Per-element diagnostics on the new trial block. + let l2_per_elem = (max_innerprod / (n as f64)).sqrt(); + let max_abs: f64 = x1_new.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs())); + + println!( + "restart {} inner {} (total cycle {}): max(||v||^2) = {:.3e}, max(||v||) = {:.3e}, per-elem L2 = {:.3e}, max-abs = {:.3e}", + restart_idx, inner + 1, total_cycles, max_innerprod, r, l2_per_elem, max_abs, + ); + + x1 = next_x1; + innerprod = next_ip; + last_max_innerprod = max_innerprod; + + if max_innerprod < conv_thresh { + inner_converged = true; + break; + } + } + + last_nd = nd; + last_b = b_residual; + + if inner_converged || total_cycles >= max_cycle { + break; + } + + // Hard restart: solve the projected system, fold x_partial into x_accum, + // discard the subspace, and continue the outer loop. + let x_partial = projected_solve(xs.i((.., ..nd)), ax.i((.., ..nd)), &all_innerprod, last_b.view()); + x_accum += &x_partial; + restart_idx += 1; + println!("---- restart {restart_idx}: x_accum refined, subspace reset ----"); + } + + // Final projected solve on the last subspace, using the matching residual. + let x_out: Tsr = if last_nd == 0 { + x_accum + } else { + let x_final = projected_solve(xs.i((.., ..last_nd)), ax.i((.., ..last_nd)), &all_innerprod, last_b.view()); + x_accum + x_final + }; + + // Hard fail on non-convergence: a silently unconverged x propagates as a + // subtly wrong Hessian downstream, which is worse than crashing here. + if last_max_innerprod >= conv_thresh { + panic!( + "krylov_block failed to converge: max(||v||^2) = {:.3e} >= tol^2 = {:.3e} after {} cycles \ + ({} restarts, max_cycle = {}, max_space = {}). Increase max_cycle or check the operator.", + last_max_innerprod, conv_thresh, total_cycles, restart_idx, max_cycle, max_space, + ); + } + + x_out +} + +/// Solve the projected `(I + A_proj) c = g` system and reconstruct `Xs c`. +/// +/// `xs_slc` and `ax_slc` are the subspace and its image under A, both `[n, nd]`. `inner` is the +/// per-column squared norm of `xs_slc` (length `nd`). `b_proj_src` is the RHS that produced this +/// subspace; the returned tensor has shape `[n, nset]`. +fn projected_solve(xs_slc: TsrView, ax_slc: TsrView, inner: &[f64], b_proj_src: TsrView) -> Tsr { + let nd = xs_slc.shape()[1]; + // h[i, j] = (xs.T @ ax)[i, j] + delta_ij * ||xs[:, i]||^2 + let mut h: Tsr = xs_slc.t() % &ax_slc; + for i in 0..nd { + h[[i, i]] += inner[i]; + } + // g[i, k] = (xs.T @ b)[i, k] + let g: Tsr = xs_slc.t() % &b_proj_src; + let c = rt::linalg::solve_general((h, g)); + &xs_slc % &c +} + +/// Modified Gram-Schmidt over the **columns** of `vec`, keeping non-normalized +/// orthogonal vectors and tracking their squared norms. Mirrors `_orth_block` +/// in the Python prototype (which operated on rows; here we operate on columns +/// because the rest of the solver is col-major). +/// +/// Columns whose remaining squared norm falls below `lindep` are dropped. +/// +/// Returns `(out, norms_sq)` where `out` has shape `[n, m]` with `m <= nblock`. +fn orth_block(vec: TsrView, lindep: f64) -> (Tsr, Vec) { + let device = vec.device().clone(); + let n = vec.shape()[0]; + let nblock = vec.shape()[1]; + + let mut result: Vec = Vec::with_capacity(nblock); + let mut norms_sq: Vec = Vec::with_capacity(nblock); + + for i in 0..nblock { + let mut vi: Tsr = vec.i((.., i)).to_owned(); + for j in 0..result.len() { + let coeff = (&vi % &result[j]).to_scalar() / norms_sq[j]; + vi -= coeff * &result[j]; + } + let nsq = (&vi % &vi).to_scalar(); + if nsq > lindep { + result.push(vi); + norms_sq.push(nsq); + } + } + + // stack vectors to output matrix if any survived + if result.is_empty() { + // rt::stack does not allow zero-length, where numpy also disallowed. + (rt::zeros(([n, 0], &device)), norms_sq) + } else { + (rt::stack((result, -1)), norms_sq) + } +} diff --git a/src/hessian_backup/mod.rs b/src/hessian_backup/mod.rs new file mode 100644 index 0000000000..7e8ec57818 --- /dev/null +++ b/src/hessian_backup/mod.rs @@ -0,0 +1,45 @@ +// trait definitions +pub mod trait_rhess; +pub mod trait_uhess; +pub mod trait_util; + +// core hess implementations +pub mod hcore; +pub mod nuc_repl; + +// overlap hess implementations +pub mod ovlp; + +// total hess implementations +pub mod rscf; +pub mod uscf; + +// vibrational analysis +pub mod vib; + +// utilities +pub mod cint_handling; +pub mod krylov_block; + +#[allow(unused_imports)] +pub mod prelude { + use super::*; + + pub use hcore::{RHessHcore, UHessHcore}; + pub use krylov_block::krylov_block; + pub use nuc_repl::HessNucRepl; + pub use ovlp::{RHessOvlp, UHessOvlp}; + pub use rscf::{HessSCFConfig, RHessSCF}; + pub use trait_rhess::{HessNucAPI, RHessCoreAPI, RHessElecInteractAPI}; + pub use trait_uhess::{UHessCoreAPI, UHessElecInteractAPI}; + pub use trait_util::HessUtilAPI; + pub use uscf::UHessSCF; + + pub(super) use crate::ri_jk::util::{get_dm0_restricted, get_dme0_restricted}; + pub(super) use crate::utilities::rstsr_util::prelude::*; + pub(super) use cint_handling::*; + pub(super) use itertools::Itertools; + pub(super) use rayon::prelude::*; + pub(super) use rest_libcint::prelude::*; + pub(super) use std::collections::HashMap; +} diff --git a/src/hessian_backup/nuc_repl.rs b/src/hessian_backup/nuc_repl.rs new file mode 100644 index 0000000000..08c2dfca33 --- /dev/null +++ b/src/hessian_backup/nuc_repl.rs @@ -0,0 +1,80 @@ +use super::prelude::*; + +/// Hessian contribution from nuclear repulsion. +/// +/// # Parameters +/// +/// - `mol` : [`CInt`]. The molecule object. +/// - `device` : [`DeviceBLAS`]. The device on which the returned tensor is allocated. +/// - `atm_list` : optional list of atom indices for the Hessian. If `None`, all atoms are computed. +/// +/// # Returns +/// +/// - `de_nuc` : shape `[3, 3, natm, natm]`. The nuclear repulsion Hessian. +pub fn get_nuc_repl_hess(mol: &CInt, device: &DeviceBLAS, atm_list: Option<&[usize]>) -> Tsr { + // Note this is the number of atoms in the original molecule, not the selected atoms. We will select + // the sub-block at the end. + let natm = mol.natm(); + + // de_nuc: shape [3, 3, natm, natm]. The nuclear repulsion Hessian. + let mut de_nuc: Tsr = rt::zeros(([3, 3, natm, natm], device)); + // qs: shape [natm]. The atomic charges. + let qs = rt::asarray((mol.atom_charges(), device)); + // rs: shape [3, natm]. The atomic coordinates. + let rs = rt::asarray((mol.atom_coords(), device)).into_unpack_array(0); + + for A in 0..natm { + // r12: shape [3, natm]. The vector from atom A to other atoms. + let r12 = rs.i((.., A)) - &rs; + // s12: shape [natm]. The distance from atom A to other atoms. + // this value will be divided, so we set zero distance to inf to avoid NaN. + let mut s12 = r12.l2_norm_axes(0); + s12[[A]] = f64::INFINITY; + + // tmp1: shape [natm] + let tmp1 = qs[[A]] * &qs / s12.pow(3); + // prefactor: shape [natm] + let prefactor = -3.0 * qs[[A]] * &qs / s12.pow(5); + // tmp2: shape [3, 3, natm] + let tmp2 = prefactor.i((None, None, ..)) * r12.i((.., None, ..)) * r12.i((None, .., ..)); + + // diagonal block + let tmp1_sum = tmp1.sum(); // scalar + let tmp2_sum = tmp2.sum_axes(-1); // shape [3, 3] + de_nuc[[0, 0, A, A]] -= tmp1_sum; + de_nuc[[1, 1, A, A]] -= tmp1_sum; + de_nuc[[2, 2, A, A]] -= tmp1_sum; + *&mut de_nuc.i_mut((.., .., A, A)) -= tmp2_sum; + + // off-diagonal blocks + *&mut de_nuc.i_mut((0, 0, A, ..)) += &tmp1; + *&mut de_nuc.i_mut((1, 1, A, ..)) += &tmp1; + *&mut de_nuc.i_mut((2, 2, A, ..)) += &tmp1; + *&mut de_nuc.i_mut((.., .., A, ..)) += &tmp2; + } + + match atm_list { + None => de_nuc, + Some(list) => de_nuc.index_select(-1, list).index_select(-2, list), + } +} + +/// Hessian contribution from nuclear repulsion. +pub struct HessNucRepl { + pub mol: CInt, + pub device: DeviceBLAS, +} + +impl HessNucRepl { + pub fn new(mol: &CInt, device: &DeviceBLAS) -> Self { + Self { mol: mol.clone(), device: device.clone() } + } +} + +impl HessUtilAPI for HessNucRepl {} + +impl HessNucAPI for HessNucRepl { + fn make_skeleton_hess(&mut self, atm_list: Option<&[usize]>) -> Tsr { + get_nuc_repl_hess(&self.mol, &self.device, atm_list) + } +} diff --git a/src/hessian_backup/ovlp.rs b/src/hessian_backup/ovlp.rs new file mode 100644 index 0000000000..fa272270dc --- /dev/null +++ b/src/hessian_backup/ovlp.rs @@ -0,0 +1,141 @@ +use super::prelude::*; + +/// Hessian contribution from overlap matrix derivative. +/// +/// # Notes +/// +/// Please be aware that the overlap matrix derivative is **NOT skeleton derivative**. +/// +/// It's true origin is the application of Hellmann-Feynman theorem, that converts part of the +/// response of density matrix to the response of basis functions. +/// +/// # Parameters +/// +/// - `mol` : [`CInt`]. The molecule object. +/// - `dme0` : shape `[nao, nao]`. The energy-weighted density matrix for current SCF component. +/// - `atm_list` : optional list of atom indices to compute the Hessian for. If `None`, all atoms +/// are computed. +/// +/// Returns +/// ------- +/// - `de_ovlp` : shape `[3, 3, natm, natm]`. The Hessian contribution from the overlap matrix +/// derivative. +pub fn get_hess_ovlp(mol: &CInt, dme0: TsrView, atm_list: Option<&[usize]>) -> Tsr { + let device = dme0.device(); + let nao = mol.nao(); + let (aoslices, _atm_indices) = filter_aoslices(mol, atm_list); + let natm = aoslices.len(); + + assert_eq!(dme0.shape(), &[nao, nao], "density matrix shape not correct."); + + let s2_aa = hess_intor(mol, "int1e_ipipovlp", "s1", None, device); + let s2_ab = hess_intor(mol, "int1e_ipovlpip", "s1", None, device); + + let mut de_ovlp = rt::zeros(([3, 3, natm, natm], device)); + for A in 0..natm { + let [_, _, p0A, p1A] = aoslices[A]; + let slcA = rt::slice!(p0A, p1A); + let scr = -2 * (s2_aa.i(slcA) * dme0.i(slcA)).sum_axes([0, 1]); + *&mut de_ovlp.i_mut((.., .., A, A)) += scr; + + for B in 0..=A { + let [_, _, p0B, p1B] = aoslices[B]; + let slcB = rt::slice!(p0B, p1B); + let scr = -2 * (s2_ab.i((slcA, slcB)) * dme0.i((slcA, slcB))).sum_axes([0, 1]); + *&mut de_ovlp.i_mut((.., .., B, A)) += scr; + } + for B in 0..A { + let de_to_copy = de_ovlp.i((.., .., B, A)).t().to_owned(); + *&mut de_ovlp.i_mut((.., .., A, B)) += de_to_copy; + } + } + de_ovlp +} + +/// Generator for the first derivative of overlap matrix. +/// +/// # Parameters +/// +/// - `mol` : [`CInt`]. The molecule object. +/// - `device` : [`DeviceBLAS`]. The device on which the returned tensor is allocated. +/// +/// # Returns +/// +/// - `FnMut(A: usize) -> Tsr`. A function that computes the first derivative of the overlap matrix +/// with respect to the nuclear coordinates. Input is the global atom index A. The returned array +/// has shape `[nao, nao, 3]`. +pub fn generator_ovlp_deriv1(mol: &CInt, device: &DeviceBLAS) -> impl FnMut(usize) -> Tsr { + // preparation + let device = device.clone(); + let nao = mol.nao(); + let aoslices = mol.aoslice_by_atom(); + + let int1e_ipovlp = hess_intor(mol, "int1e_ipovlp", "s1", None, &device); + + move |A: usize| { + let [_, _, p0, p1] = aoslices[A]; + let slc = rt::slice!(p0, p1); + let mut s1ao = rt::zeros(([nao, nao, 3], &device)); + *&mut s1ao.i_mut((slc, ..)) -= int1e_ipovlp.i(slc); + *&mut s1ao.i_mut((.., slc)) -= int1e_ipovlp.i(slc).swapaxes(0, 1); + s1ao + } +} + +/// Hessian contribution from overlap matrix derivative. +/// +/// Note that overlap is special to the SCF part, in that +/// - The contribution of hessian from overlap is not skeleton, so we do not derive this class from +/// [`RHessCoreAPI`]. +/// - The CP-HF requires both first order derivative of hcore and ovlp, but their roles are +/// different. +/// +/// Due to these reasons, although it has the similar interface to [`RHessCoreAPI`], +/// [`RHessOvlp`] is designed as a standalone class, without inheriting from any abstract class. +pub struct RHessOvlp { + pub mol: CInt, + pub device: DeviceBLAS, +} + +impl RHessOvlp { + pub fn new(mol: &CInt, device: &DeviceBLAS) -> Self { + Self { mol: mol.clone(), device: device.clone() } + } + + pub fn make_hess(&self, dme0: TsrView, atm_list: Option<&[usize]>) -> Tsr { + get_hess_ovlp(&self.mol, dme0, atm_list) + } + + pub fn generator_deriv1(&self) -> impl FnMut(usize) -> Tsr { + generator_ovlp_deriv1(&self.mol, &self.device) + } + + pub fn natm(&self) -> usize { + self.mol.natm() + } +} + +/// Hessian contribution from overlap matrix derivative for unrestricted SCF. +pub struct UHessOvlp { + pub mol: CInt, + pub device: DeviceBLAS, +} + +impl UHessOvlp { + pub fn new(mol: &CInt, device: &DeviceBLAS) -> Self { + Self { mol: mol.clone(), device: device.clone() } + } + + pub fn make_hess(&self, dme0: [TsrView; 2], atm_list: Option<&[usize]>) -> Tsr { + let [α, β] = [0, 1]; + get_hess_ovlp(&self.mol, (&dme0[α] + &dme0[β]).view(), atm_list) + } + + pub fn generator_deriv1(&self) -> impl FnMut(usize) -> Tsr { + generator_ovlp_deriv1(&self.mol, &self.device) + } + + pub fn natm(&self) -> usize { + self.mol.natm() + } +} diff --git a/src/hessian_backup/rscf.rs b/src/hessian_backup/rscf.rs new file mode 100644 index 0000000000..456b178348 --- /dev/null +++ b/src/hessian_backup/rscf.rs @@ -0,0 +1,506 @@ +//! Hessian implementations for restricted SCF. + +use super::prelude::*; + +#[derive(Debug, Clone, PartialEq)] +pub struct HessSCFConfig { + pub level_shift: f64, + pub cphf_tol: f64, + pub cphf_max_cycle: usize, + pub cphf_max_space: usize, + pub cphf_lindep: f64, +} + +impl Default for HessSCFConfig { + fn default() -> Self { + Self { level_shift: 0.0, cphf_tol: 1e-8, cphf_max_cycle: 42, cphf_max_space: 14, cphf_lindep: 1e-14 } + } +} + +/// Working solver and maintainer of all hessian components for restricted SCF method. +pub struct RHessSCF<'a> { + pub mo_coeff: Tsr, + pub mo_occ: Tsr, + pub mo_energy: Tsr, + pub ovlp_obj: RHessOvlp, + pub nuc_list: Vec<&'a mut dyn HessNucAPI>, + pub core_list: Vec<&'a mut dyn RHessCoreAPI>, + pub el_list: Vec<&'a mut dyn RHessElecInteractAPI>, + pub config: HessSCFConfig, + pub atm_list: Option>, + pub result: HashMap, + /// Timing information. Represented by wall time in second. + pub timing: Vec<(String, f64)>, +} + +impl<'a> RHessSCF<'a> { + #[allow(clippy::too_many_arguments)] + pub fn new( + mo_coeff: Tsr, + mo_occ: Tsr, + mo_energy: Tsr, + ovlp_obj: RHessOvlp, + nuc_list: Vec<&'a mut dyn HessNucAPI>, + core_list: Vec<&'a mut dyn RHessCoreAPI>, + el_list: Vec<&'a mut dyn RHessElecInteractAPI>, + config: HessSCFConfig, + atm_list: Option<&[usize]>, + ) -> Self { + Self { + mo_coeff, + mo_occ, + mo_energy, + ovlp_obj, + nuc_list, + core_list, + el_list, + config, + atm_list: atm_list.map(|x| x.to_vec()), + result: HashMap::new(), + timing: Vec::new(), + } + } + + /// Number of atoms over which the Hessian is computed. This is `atm_list.len()` if + /// `atm_list` is `Some`, otherwise the total number of atoms in the molecule. + pub fn natm(&self) -> usize { + match &self.atm_list { + Some(list) => list.len(), + None => self.ovlp_obj.natm(), + } + } + + /// Return the list of (global) atom indices the Hessian is computed for, ordered the same + /// way as the local indexing used in the returned Hessian. + pub fn atm_indices(&self) -> Vec { + match &self.atm_list { + Some(list) => list.clone(), + None => (0..self.ovlp_obj.natm()).collect(), + } + } + + /// Compute the dimensionless CPHF right-hand side, along with necessary intermediates for later + /// steps. + /// + /// Note there are some differences compared to usual CP-HF: + /// - Usual CP-HF is `(ea - ei) U - AU = B`, where now we handle something like `U + (A / (ea - + /// ei)) U = - B / (ea - ei)` + /// - We now handle the U in all-occ block, instead of standard vir-occ block; this will omit + /// the response evaluation during rhs (B), making the rhs evaluation cheap, but we also need + /// to carefully handle the CP-HF equation. this behavior should be similar to PySCF's + /// `solve_withs1`. + /// + /// Note **dimless** here means the CP-HF equation is of no quantity dimension (量纲), but not + /// the tensor structure is dimensionless. + /// + /// # Returns + /// + /// A dictionary containing: + /// - `rhs : shape `[nmo, nocc, 3, natm]`. The dimensionless CPHF right-hand side. + /// - `f1mo` : shape `[nmo, nocc, 3, natm]`. The first-order derivative of the Fock matrix in MO + /// basis. + /// - `s1mo` : shape `[nmo, nocc, 3, natm]`. The first-order derivative of the overlap matrix in + /// MO basis. + pub fn compute_dimless_cphf_rhs(&mut self) -> HashMap<&'static str, Tsr> { + // setups + let t0 = std::time::Instant::now(); + let mo_coeff = &self.mo_coeff; + let mo_occ = &self.mo_occ; + let mo_energy = &self.mo_energy; + let level_shift = self.config.level_shift; + let device = mo_coeff.device().clone(); + + let [nao, nmo] = mo_coeff.shape().to_vec().try_into().unwrap(); + let occidx = mo_occ.view().greater(0).into_vec(); + let viridx = occidx.iter().map(|&x| !x).collect_vec(); + let mocc = mo_coeff.bool_select(-1, &occidx); + let eocc = mo_energy.bool_select(-1, &occidx); + let evir = mo_energy.bool_select(-1, &viridx); + let nocc = occidx.iter().filter(|&&x| x).count(); + let natm = self.natm(); + let atm_indices = self.atm_indices(); + let atm_list = self.atm_list.as_deref(); + + let e_ai = evir.i((.., None)) - eocc.i((None, ..)); + let e_ai_shift = &e_ai + level_shift; + + // --- f1mo --- // + + // fock skeleton derivative (core contribution) + let mut f1ao_core: Tsr = rt::zeros(([nao, nao, 3, natm], &device)); + for core_obj in self.core_list.iter() { + let t1 = std::time::Instant::now(); + let mut gen_core_deriv1 = core_obj.generator_deriv1(); + for (A_loc, &A_glob) in atm_indices.iter().enumerate() { + *&mut f1ao_core.i_mut((Ellipsis, A_loc)) += gen_core_deriv1(A_glob); + } + self.timing.push(( + format!("in compute_dimless_cphf_rhs, f1ao_core_{}", core_obj.get_type_name()), + t1.elapsed().as_secs_f64(), + )); + } + + // fock skeleton derivative (electron interaction contribution, half-transformed to bra) + let mut f1bra_el: Tsr = rt::zeros(([nao, nocc, 3, natm], &device)); + for el_obj in self.el_list.iter_mut() { + let t1 = std::time::Instant::now(); + f1bra_el += el_obj.get_deriv1_bra(mo_coeff.view(), mo_occ.view(), atm_list); + self.timing.push(( + format!("in compute_dimless_cphf_rhs, f1bra_el_{}", el_obj.get_type_name()), + t1.elapsed().as_secs_f64(), + )); + } + + // construct whole f1mo + let f1bra = f1bra_el + f1ao_core % &mocc; + let f1mo = mo_coeff.t() % f1bra; + + // --- s1mo --- // + + let t1 = std::time::Instant::now(); + + let mut gen_ovlp_deriv1 = self.ovlp_obj.generator_deriv1(); + let mut s1ao: Tsr = rt::zeros(([nao, nao, 3, natm], &device)); + for (A_loc, &A_glob) in atm_indices.iter().enumerate() { + *&mut s1ao.i_mut((Ellipsis, A_loc)) += gen_ovlp_deriv1(A_glob); + } + let s1mo = mo_coeff.t() % (&s1ao % &mocc); + + self.timing.push(("in compute_dimless_cphf_rhs, s1mo".to_string(), t1.elapsed().as_secs_f64())); + + // --- dimensionless cphf rhs --- // + + let so = rt::slice!(0, nocc); + let sv = rt::slice!(nocc, nmo); + let b1mo = &f1mo - &s1mo * eocc.i((None, ..)); + let mut rhs = rt::zeros(([nmo, nocc, 3, natm], &device)); + *&mut rhs.i_mut(sv) += -b1mo.i(sv) / &e_ai_shift; + *&mut rhs.i_mut(so) += -0.5 * s1mo.i(so); + + self.timing.push(("compute_dimless_cphf_rhs".to_string(), t0.elapsed().as_secs_f64())); + HashMap::from([("f1mo", f1mo), ("s1mo", s1mo), ("rhs", rhs)]) + } + + /// Prepare the response for CPHF calculation. + /// + /// This involves all electron-interaction objects. + pub fn make_response_preparation(&mut self) { + let t0 = std::time::Instant::now(); + for el_obj in self.el_list.iter_mut() { + let t1 = std::time::Instant::now(); + el_obj.make_response_preparation(self.mo_coeff.view(), self.mo_occ.view()); + self.timing.push(( + format!("in make_response_preparation, {}", el_obj.get_type_name()), + t1.elapsed().as_secs_f64(), + )); + } + 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 CPHF. + /// + /// # Parameters + /// + /// - `mo1` : shape `[nmo, nocc, ...]`. The perturbation in MO space. + /// + /// # Returns + /// + /// - `resp` : shape `[nmo, nocc, ...]`. The response in MO space. + pub fn response_mo(&mut self, mo1: TsrView) -> Tsr { + let mo_coeff = self.mo_coeff.view(); + let ubra = &mo_coeff % &mo1; + let mut resp = rt::zeros_like(&mo1); + for el_obj in self.el_list.iter_mut() { + let t1 = std::time::Instant::now(); + resp += mo_coeff.t() % el_obj.get_response_bra(ubra.view()); + self.timing.push((format!("in response_mo, {}", el_obj.get_type_name()), t1.elapsed().as_secs_f64())); + } + resp + } + + /// Compute the dimensionless response for CP-HF calculation. + /// + /// Compared to usual CP-HF response, this additionally handles + /// - 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). + /// + /// # Parameters + /// + /// - `mo1` : shape `[nmo, nocc, ...]`. The perturbation in MO space. + /// + /// # Returns + /// + /// - `resp` : shape `[nmo, nocc, ...]`. The dimensionless response in MO space. + pub fn response_dimless_cphf(&mut self, mo1: TsrView) -> Tsr { + let t0 = std::time::Instant::now(); + let mo_occ = self.mo_occ.view(); + let mo_energy = self.mo_energy.view(); + let level_shift = self.config.level_shift; + let occidx = mo_occ.view().greater(0).into_vec(); + let viridx = occidx.iter().map(|&x| !x).collect_vec(); + let nocc = occidx.iter().filter(|&&x| x).count(); + let nmo = mo_occ.shape()[0]; + let eocc = mo_energy.bool_select(-1, &occidx); + let evir = mo_energy.bool_select(-1, &viridx); + let so = rt::slice!(0, nocc); + let sv = rt::slice!(nocc, nmo); + let e_ai = evir.i((.., None)) - eocc.i((None, ..)); + let e_ai_shift = &e_ai + level_shift; + + let mut resp = self.response_mo(mo1.view()); + + // handle dimensionless denominator and force handle virtual-part only + if level_shift != 0.0 { + resp -= level_shift * &mo1; + } + *&mut resp.i_mut(sv) /= &e_ai_shift; + resp.i_mut(so).fill(0.0); + self.timing.push(("response_dimless_cphf".to_string(), t0.elapsed().as_secs_f64())); + resp + } + + /// Solve the dimensionless CP-HF equation using a Krylov solver. + /// + /// This should solves `U + resp(U) = rhs`. Note difference of standard CP-HF equation as + /// mentioned in functions above. + /// + /// # Parameters + /// + /// - `rhs` : shape `[nmo, nocc, ...]`. Dimensionless right-hand side. + /// + /// # Returns + /// + /// - `mo1` : shape `[nmo, nocc, ...]`. Perturbation in MO space that solves the dimensionless + /// CP-HF equation. + pub fn solve_dimless_cphf(&mut self, rhs: TsrView) -> Tsr { + let t0 = std::time::Instant::now(); + let rhs_shape = rhs.shape().to_vec(); + let nmo = rhs.shape()[0]; + let nocc = rhs.shape()[1]; + let rhs = rhs.reshape((nmo * nocc, -1)); + + let tol = self.config.cphf_tol; + let max_cycle = self.config.cphf_max_cycle; + let max_space = self.config.cphf_max_space; + let lindep = self.config.cphf_lindep; + + let response_cphf_flattened = |x: TsrView| -> Tsr { + let x = x.reshape((nmo, nocc, -1)); + let y = self.response_dimless_cphf(x.view()); + y.into_shape((nmo * nocc, -1)) + }; + let mo1 = krylov_block(response_cphf_flattened, rhs.view(), None, tol, max_cycle, max_space, lindep); + let mo1 = mo1.into_shape(rhs_shape); + + self.timing.push(("solve_dimless_cphf".to_string(), t0.elapsed().as_secs_f64())); + mo1 + } + + /// Finalize the CP-HF calculation by computing necessary intermediates for Hessian assembly. + /// + /// This includes: + /// - Re-computing the mo1 (as post-iteration computation), as well as removing the level shift. + /// - Computing the derivative of occupied orbital energy with respect to perturbation (mo_e1). + /// Note occupied orbital energy (shape [nocc]) is diagonal of Fock, and Fock matrix is + /// diagonal. However, with the definition that `U[occ, occ] = -0.5 S1[occ, occ]`, the + /// off-diagonal part of derivative of Fock in occupied-occupied block is not zero. That's why + /// this term is actually matrix. + /// + /// # Parameters + /// + /// - `f1mo` : shape `[nmo, nocc, 3, natm]`. The first-order derivative of the Fock matrix in MO + /// basis, obtained from [`Self::compute_dimless_cphf_rhs`]. + /// - `s1mo` : shape `[nmo, nocc, 3, natm]`. The first-order derivative of the overlap matrix in + /// MO basis, obtained from [`Self::compute_dimless_cphf_rhs`]. + /// - `mo1` : shape `[nmo, nocc, 3, natm]`. The perturbation in MO space obtained from Krylov + /// solver. + /// + /// # Returns + /// + /// `HashMap<&str, Tsr>` + /// + /// - `mo1` : shape `[nmo, nocc, 3, natm]`. The finalized perturbation in MO space. + /// - `mo_e1` : shape `[nocc, nocc, 3, natm]`. The derivative of occupied orbital energies (Fock + /// matrix) with respect to perturbation. + pub fn finalize_cphf(&mut self, f1mo: TsrView, s1mo: TsrView, mo1: TsrView) -> HashMap<&'static str, Tsr> { + let t0 = std::time::Instant::now(); + let mo_occ = self.mo_occ.view(); + let mo_energy = self.mo_energy.view(); + let occidx = mo_occ.view().greater(0).into_vec(); + let viridx = occidx.iter().map(|&x| !x).collect_vec(); + let nocc = occidx.iter().filter(|&&x| x).count(); + let nmo = mo_occ.shape()[0]; + let eocc = mo_energy.bool_select(-1, &occidx); + let evir = mo_energy.bool_select(-1, &viridx); + let so = rt::slice!(0, nocc); + let sv = rt::slice!(nocc, nmo); + let e_ai = evir.i((.., None)) - eocc.i((None, ..)); + let e_ij = eocc.i((.., None)) - eocc.i((None, ..)); + + // last-iter the cp-hf equation, and remove the level-shift + let b1mo = f1mo - s1mo * eocc.i((None, ..)) + self.response_mo(mo1.view()); + let mut mo1 = mo1.to_owned(); + mo1.i_mut(sv).assign(-b1mo.i(sv) / e_ai); + + // get the derivative of fock matrix in occ-occ block (derivative of orbital energy with rotation) + let mo_e1 = b1mo.i(so) + mo1.i(so) * e_ij; + + self.timing.push(("finalize_cphf".to_string(), t0.elapsed().as_secs_f64())); + HashMap::from([("mo1", mo1), ("mo_e1", mo_e1)]) + } + + /// Compute the CP-HF contribution to the Hessian using the finalized CP-HF results. + /// + /// # Parameters + /// + /// - `f1mo` : shape `[nmo, nocc, 3, natm]`. The first-order derivative of the Fock matrix in MO + /// basis, obtained from [`Self::compute_dimless_cphf_rhs`]. + /// - `s1mo` : shape `[nmo, nocc, natm, 3]`. The first-order skeleton derivative of the overlap + /// matrix in MO basis, obtained from [`Self::compute_dimless_cphf_rhs`]. + /// - `mo1` : shape `[nmo, nocc, natm, 3]`. The finalized perturbation in MO space obtained from + /// [`Self::finalize_cphf`]. + /// - `mo_e1` : shape `[nocc, nocc, 3, natm]`. The derivative of occupied orbital energies (Fock + /// matrix) with respect to perturbation, obtained from [`Self::finalize_cphf`]. + /// + /// # Returns + /// + /// - `de_cphf` : shape `[3, 3, natm, natm]`. The CP-HF contribution to the Hessian. + pub fn get_cphf_hess(&self, f1mo: TsrView, s1mo: TsrView, mo1: TsrView, mo_e1: TsrView) -> Tsr { + let natm = self.natm(); + let mo_occ = self.mo_occ.view(); + let mo_energy = self.mo_energy.view(); + let occidx = mo_occ.view().greater(0).into_vec(); + let nocc = occidx.iter().filter(|&&x| x).count(); + let eocc = mo_energy.bool_select(-1, &occidx); + let so = rt::slice!(0, nocc); + let device = mo1.device().clone(); + + let s1oo = s1mo.i(so); + let mut de_cphf = rt::zeros(([3, 3, natm, natm], &device)); + // well, code style is ruined by rustfmt ... + for A in 0..natm { + for B in 0..=A { + let mut de_BA = de_cphf.i_mut((.., .., B, A)); + de_BA += 4 * (f1mo.i((.., .., None, .., A)) * mo1.i((.., .., .., None, B))).sum_axes([0, 1]); + de_BA -= 4 + * (s1mo.i((.., .., None, .., A)) * mo1.i((.., .., .., None, B)) * eocc.i((None, ..))) + .sum_axes([0, 1]); + de_BA -= 2 * (s1oo.i((.., .., None, .., A)) * mo_e1.i((.., .., .., None, B))).sum_axes([0, 1]); + } + for B in 0..A { + let de_to_copy = de_cphf.i((.., .., B, A)).t().to_owned(); + *&mut de_cphf.i_mut((.., .., A, B)) += de_to_copy; + } + } + de_cphf + } + + /// Compute the CP-HF contribution to the Hessian by running through the entire CP-HF workflow. + /// + /// - Compute the dimensionless CPHF right-hand side and necessary intermediates. + /// - Prepare the response for CPHF calculation. + /// - Solve the dimensionless CP-HF equation using a Krylov solver. + /// - Finalize the CP-HF results by computing necessary intermediates for Hessian assembly. + /// - Compute the CP-HF contribution to the Hessian using the finalized CP-HF results. + /// + /// # Returns + /// + /// - `de_cphf` : shape `[3, 3, natm, natm]`. The CP-HF contribution to the Hessian. + pub fn make_cphf_hess(&mut self) -> Tsr { + let pre_cphf_dict = self.compute_dimless_cphf_rhs(); + let f1mo = pre_cphf_dict["f1mo"].view(); + let s1mo = pre_cphf_dict["s1mo"].view(); + let rhs = pre_cphf_dict["rhs"].view(); + + self.make_response_preparation(); + let mo1 = self.solve_dimless_cphf(rhs.view()); + let finalize_dict = self.finalize_cphf(f1mo.view(), s1mo.view(), mo1.view()); + let mo1 = finalize_dict["mo1"].view(); + let mo_e1 = finalize_dict["mo_e1"].view(); + + self.get_cphf_hess(f1mo.view(), s1mo.view(), mo1.view(), mo_e1.view()) + } + + /// Compute the total skeleton contribution to the Hessian. + /// + /// **Total** means that we sum over all skeleton contributions from both core and + /// electron-interaction objects. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo]`. The molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo]`. The orbital occupations. + /// + /// # Returns + /// + /// - `de_skeleton` : shape `[3, 3, natm, natm]`. The total skeleton contribution to the + /// Hessian. + pub fn make_skeleton_hess(&mut self) -> Tsr { + let natm = self.natm(); + let mo_coeff = self.mo_coeff.view(); + let mo_occ = self.mo_occ.view(); + let atm_list = self.atm_list.as_deref(); + + let device = self.mo_coeff.device().clone(); + let mut de_skeleton = rt::zeros(([3, 3, natm, natm], &device)); + for nuc_obj in self.nuc_list.iter_mut() { + let t0 = std::time::Instant::now(); + let de_nuc = nuc_obj.make_skeleton_hess(atm_list); + let nuc_obj_name = nuc_obj.get_type_name(); + self.result.insert(format!("de_skeleton_{}", nuc_obj_name), de_nuc.to_owned()); + self.timing.push((format!("de_skeleton_{}", nuc_obj_name,), t0.elapsed().as_secs_f64())); + de_skeleton += de_nuc; + } + for core_obj in self.core_list.iter_mut() { + let t0 = std::time::Instant::now(); + let de_core = core_obj.make_skeleton_hess(mo_coeff.view(), mo_occ.view(), atm_list); + let core_obj_name = core_obj.get_type_name(); + self.result.insert(format!("de_skeleton_{}", core_obj_name), de_core.to_owned()); + self.timing.push((format!("de_skeleton_{}", core_obj_name,), t0.elapsed().as_secs_f64())); + de_skeleton += de_core; + } + for el_obj in self.el_list.iter_mut() { + let t0 = std::time::Instant::now(); + let de_el = el_obj.make_skeleton_hess(mo_coeff.view(), mo_occ.view(), atm_list); + let el_obj_name = el_obj.get_type_name(); + self.result.insert(format!("de_skeleton_{}", el_obj_name), de_el.to_owned()); + self.timing.push((format!("de_skeleton_{}", el_obj_name,), t0.elapsed().as_secs_f64())); + de_skeleton += de_el; + } + de_skeleton + } + + /// Compute the total Hessian by summing over skeleton, overlap, and CP-HF contributions. + /// + /// # Returns + /// + /// - `de_hess` : shape `[3, 3, natm, natm]`. The total Hessian. + pub fn make_hess(&mut self) -> Tsr { + let t0 = std::time::Instant::now(); + let mo_coeff = self.mo_coeff.view(); + let mo_occ = self.mo_occ.view(); + let mo_energy = self.mo_energy.view(); + let dme0 = get_dme0_restricted(mo_coeff, mo_occ, mo_energy); + let atm_list = self.atm_list.clone(); + + let de_skeleton = self.make_skeleton_hess(); + + let t1 = std::time::Instant::now(); + let de_ovlp = self.ovlp_obj.make_hess(dme0.view(), atm_list.as_deref()); + self.result.insert("de_ovlp".to_string(), de_ovlp.to_owned()); + self.timing.push(("de_ovlp".to_string(), t1.elapsed().as_secs_f64())); + + let t1 = std::time::Instant::now(); + let de_cphf = self.make_cphf_hess(); + self.result.insert("de_cphf".to_string(), de_cphf.to_owned()); + self.timing.push(("de_cphf".to_string(), t1.elapsed().as_secs_f64())); + + let de_tot = de_skeleton + de_ovlp + de_cphf; + self.result.insert("de_tot".to_string(), de_tot.to_owned()); + self.timing.push(("de_tot".to_string(), t0.elapsed().as_secs_f64())); + de_tot + } +} diff --git a/src/hessian_backup/trait_rhess.rs b/src/hessian_backup/trait_rhess.rs new file mode 100644 index 0000000000..7326d4da77 --- /dev/null +++ b/src/hessian_backup/trait_rhess.rs @@ -0,0 +1,190 @@ +use super::prelude::*; + +/// Abstract class for Hessian-related API for nuclear repulsion contribution. +/// +/// # Term Explanation +/// +/// **Nuc component** here actually means the term is of zero-order with right of (electron) density +/// matrix. Usually this is only related to nuclear repulsion, but should denote all components that +/// do not depend on density matrix (like DFT-D3). +/// +/// This hessian contribution is not related to density, so no RHF/UHF/GHF distinguishment required. +pub trait HessNucAPI: HessUtilAPI { + fn make_skeleton_hess(&mut self, atm_list: Option<&[usize]>) -> Tsr; +} + +/// Abstract class for Hessian-related API for restricted SCF core components. +/// +/// # Term Explanation +/// +/// **Core component** here actually means the term is of first-order with right of (electron) +/// density matrix. +/// +/// - Core Hamiltonian is first-order (linear to density matrix). +/// - External field may have nuclear and electronic contributions. For dipole field, as an example, +/// the electronic contribution is of first-order, and can be counted in core-hamiltonian in some +/// frameworks. +/// +/// We have function `make_skeleton_hess` here to count the **skeleton** contribution of the +/// Hessian. We do not handle derivative of density matrix here, which is the responsibility of CPHF +/// solver. +pub trait RHessCoreAPI: HessUtilAPI { + /// Generate the **skeleton** contribution of Hessian for current SCF component. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. In usual cases, the + /// occupied orbitals should have occupation 2, and virtual orbitals should have occupation 0. + /// - `atm_list` : optional list of atom indices to compute the Hessian for. If `None`, all + /// atoms are computed. + /// + /// # Returns + /// + /// - `hess` : shape `[3, 3, natm, natm]`. The Hessian matrix for current SCF component, where + /// `natm = atm_list.len()` if `atm_list` is `Some`, else `mol.natm()`. + /// + /// Note the hessian should be of indices `[s, t, B, A]` for column major. + fn make_skeleton_hess(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr; + + /// Generate the function to compute the first-order derivative of core component. + /// + /// This function only works for first-order density matrix contribution (like hcore). If this + /// component does not contribute (like nuclear repulsion), return None. + /// + /// # Parameters (in closure) + /// + /// - `A` : usize. The atom index (global, in original molecule) for which the derivative is + /// taken. + /// + /// # Returns (in closure) + /// + /// - `deriv1` : shape `[nao, nao, 3]`. The first-order derivative of core component with + /// respect to the position of atom `A`. + fn generator_deriv1(&self) -> Box Tsr>; +} + +/// Abstract class for Hessian-related API for restricted SCF electronic interaction components. +/// +/// # Term Explanation +/// +/// **Electronic interaction** here actually means the term is of two-order (or higher-order) with +/// right of (electron) density matrix. +/// +/// - J/K contribution from Hartree-Fock is exactly two-order. +/// - DFT contribution is non-linear to density matrix, and should be counted as infinity-order. +/// - Implicit-solvent/VV10 is probably categorized here. +/// +/// In SCF iteration, introducing two-order (or higher-order) contribution requires the program to +/// make some modification to Fock matrix construction. This kind of terms is substentially +/// different from zero/one-order core components, and should be handled separately. +pub trait RHessElecInteractAPI: HessUtilAPI { + /// Generate the **skeleton** contribution of Hessian for current SCF component. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. In usual cases, the + /// occupied orbitals should have occupation 2, and virtual orbitals should have occupation 0. + /// - `atm_list` : optional list of atom indices to compute the Hessian for. If `None`, all + /// atoms are computed. + /// + /// # Returns + /// + /// - `hess` : shape `[3, 3, natm, natm]`. The Hessian matrix for current SCF component. + /// + /// Note the hessian should be of indices `[s, t, B, A]` for column major. + fn make_skeleton_hess(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr; + + /// First order skeleton derivative in AO basis. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. + /// - `atm_list` : optional list of atom indices over which derivatives are computed. + /// + /// # Returns + /// + /// - `deriv_ao` : shape `[nao, nao, 3, natm]`. The first-order skeleton derivative in AO basis. + fn get_deriv1_ao(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr; + + /// First order skeleton derivative in half-transformed MO basis. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. + /// - `atm_list` : optional list of atom indices over which derivatives are computed. + /// + /// # Returns + /// + /// - `deriv_bra` : shape `[nao, nocc, 3, natm]`. The first-order skeleton derivative in + /// half-transformed MO basis. Note that this function will handle the order of occupied + /// orbitals. If occupation number is not sorted contiguously, you may be extra cautious to + /// this function. + /// + /// # Notes + /// + /// If [`get_deriv1_ao`] implemented, this function should behave like `deriv_bra = deriv_ao @ + /// mocc`, where `mocc` is the occupied molecular coefficients (as ket). + /// + /// However, in some cases, it is probably better to skip the usage of [`get_deriv1_ao`] and + /// directly use this function. By ket half-transformation, some RI-JK or DFT methods will + /// benefit from boost by using low-rank occupied orbitals, instead of using full AO basis. + /// + /// # See also + /// + /// [`get_deriv1_ao`] + /// + /// [`get_deriv1_ao`]: Self::get_deriv1_ao + fn get_deriv1_bra(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr { + let occidx = mo_occ.view().greater(0).into_vec(); + let mocc = mo_coeff.bool_select(-1, &occidx); + self.get_deriv1_ao(mo_coeff, mo_occ, atm_list) % mocc + } + + /// Prepare the data for response calculation. + /// + /// Response (related to second order of density matrix derivative to energy) will be called + /// multiple-times in CP-HF 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]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. + fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView); + + /// Get the response contribution for current SCF component. + /// + /// This function will be called multiple-times in CP-HF solver and other places. + /// Call [`make_response_preparation`] before this function to make sure the data is ready. + /// + /// Also, this function will not pass in the MO coefficients and occupation numbers. + /// If you need them, you should store them in the object by function + /// `make_response_preparation`. + /// + /// # Parameters + /// + /// - `bra` : shape `[nao, nocc, ...]`. The bra part for response calculation. This is usually + /// the derivative of MO coefficients (like $U_{\mu i}^\mathbb{A}$ given by CP-HF). + /// + /// # Returns + /// + /// - `resp_bra` : shape `[nao, nocc, ...]`. The response potential (related to second order of + /// density matrix derivative to energy). + /// + /// # 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) -> Tsr; +} diff --git a/src/hessian_backup/trait_uhess.rs b/src/hessian_backup/trait_uhess.rs new file mode 100644 index 0000000000..5fa31cd9e2 --- /dev/null +++ b/src/hessian_backup/trait_uhess.rs @@ -0,0 +1,172 @@ +use super::prelude::*; + +/// Abstract class for Hessian-related API for unrestricted SCF core components. +/// +/// Difference to [`RHessCoreAPI`] is that we may need different signature. Basic ideas are exactly +/// the same. +pub trait UHessCoreAPI: HessUtilAPI { + /// Generate the **skeleton** contribution of Hessian for current SCF component. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo_α]` and `[nao, nmo_β]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo_α]` and `[nmo_β]`. Molecular orbital occupation numbers. In usual + /// cases, the occupied orbitals should have occupation 1 (unrestricted), and virtual orbitals + /// should have occupation 0. + /// - `atm_list` : optional list of atom indices to compute the Hessian for. If `None`, all + /// atoms are computed. + /// + /// # Returns + /// + /// - `hess` : shape `[3, 3, natm, natm]`. The Hessian matrix for current SCF component, where + /// `natm = atm_list.len()` if `atm_list` is `Some`, else `mol.natm()`. + /// + /// Note the hessian should be of indices `[s, t, B, A]` for column major. + /// + /// # See also + /// + /// [`RHessCoreAPI::make_skeleton_hess`]. Signature difference: `mo_coeff` and `mo_occ` type + /// different. + fn make_skeleton_hess(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], atm_list: Option<&[usize]>) + -> Tsr; + + /// Generate the function to compute the first-order derivative of core component. + /// + /// This function works in atomic basis, so should have same implementation to restricted case. + /// + /// # Parameters (in closure) + /// + /// - `A` : usize. The atom index (global, in original molecule) for which the derivative is + /// taken. + /// + /// # Returns (in closure) + /// + /// - `deriv1` : shape `[nao, nao, 3]`. The first-order derivative of core component with + /// respect to the position of atom `A`. + /// + /// # See also + /// + /// [`RHessCoreAPI::generator_deriv1`]. Signature difference: no difference. + fn generator_deriv1(&self) -> Box Tsr>; +} + +/// Abstract class for Hessian-related API for restricted SCF electronic interaction components. +/// +/// Difference to [`RHessElecInteractAPI`] is that we may need different signature. Basic ideas are +/// exactly the same. +pub trait UHessElecInteractAPI: HessUtilAPI { + /// Generate the **skeleton** contribution of Hessian for current SCF component. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo_α]` and `[nao, nmo_β]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo_α]` and `[nmo_β]`. Molecular orbital occupation numbers. In usual + /// cases, the occupied orbitals should have occupation 1, and virtual orbitals should have + /// occupation 0. + /// - `atm_list` : optional list of atom indices to compute the Hessian for. If `None`, all + /// atoms are computed. + /// + /// # Returns + /// + /// - `hess` : shape `[3, 3, natm, natm]`. The Hessian matrix for current SCF component. + /// + /// Note the hessian should be of indices `[s, t, B, A]` for column major. + /// + /// # See also + /// + /// [`RHessElecInteractAPI::make_skeleton_hess`]. Signature difference: `mo_coeff` and `mo_occ` + /// type different. + fn make_skeleton_hess(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], atm_list: Option<&[usize]>) + -> Tsr; + + /// First order skeleton derivative in AO basis. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo_α]` and `[nao, nmo_β]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo_α]` and `[nmo_β]`. Molecular orbital occupation numbers. + /// - `atm_list` : optional list of atom indices over which derivatives are computed. + /// + /// # Returns + /// + /// - `deriv_ao` : shape `[nao, nao, 3, natm]` and `[nao, nao, 3, natm]`. The first-order + /// skeleton derivative in AO basis. + /// + /// # See also + /// + /// [`RHessElecInteractAPI::get_deriv1_ao`]. Signature difference: `mo_coeff` and `mo_occ` type + /// different, output shape different. + fn get_deriv1_ao(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], atm_list: Option<&[usize]>) + -> [Tsr; 2]; + + /// First order skeleton derivative in half-transformed MO basis. + /// + /// # Parameters + /// + /// - `mo_coeff` : shape `[nao, nmo_α]` and `[nao, nmo_β]`. Molecular orbital coefficients. + /// - `mo_occ` : shape `[nmo_α]` and `[nmo_β]`. Molecular orbital occupation numbers. + /// - `atm_list` : optional list of atom indices over which derivatives are computed. + /// + /// # Returns + /// + /// - `deriv_bra` : shape `[nao, nocc_α, 3, natm]` and `[nao, nocc_β, 3, natm]`. The first-order + /// skeleton derivative in half-transformed MO basis. Note that this function will handle the + /// order of occupied orbitals. If occupation number is not sorted contiguously, you may be + /// extra cautious to this function. + /// + /// # See also + /// + /// [`RHessElecInteractAPI::get_deriv1_bra`]. Signature difference: `mo_coeff` and `mo_occ` type + /// different, output type different. + fn get_deriv1_bra( + &mut self, + mo_coeff: &[TsrView; 2], + mo_occ: &[TsrView; 2], + atm_list: Option<&[usize]>, + ) -> [Tsr; 2] { + // Using `a` and `b` may be confusing for spin description. Using non-ASCII may be better. + // Decision: + // - In code or comments, we may use Greek `α`, `β` (not Cyrillic). This violates UTS #39 and + // RFCS-2457 (also related to rust lint `mixed_script_confusables`). But I think it's prettier for + // reading, and more clear then 0/1 in meaning. + // - In dictionary keys, we use `0` and `1` to avoid any potential issues. + let [α, β] = [0, 1]; + + let occidx = [mo_occ[α].view().greater(0).into_vec(), mo_occ[β].view().greater(0).into_vec()]; + let mocc = [mo_coeff[α].bool_select(-1, &occidx[α]), mo_coeff[β].bool_select(-1, &occidx[β])]; + 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/hessian_backup/trait_util.rs b/src/hessian_backup/trait_util.rs new file mode 100644 index 0000000000..2948f7ece5 --- /dev/null +++ b/src/hessian_backup/trait_util.rs @@ -0,0 +1,18 @@ +//! Utility trait for Hessian computation. + +pub trait HessUtilAPI { + /// Get the type name of the implementor. + /// + /// The full type name is too verbose. We will use the short type name for display usage. + fn get_type_name(&self) -> String { + // the full type name is probably something like + // crate::mod::struct + // we will only preserve the `struct` part. + std::any::type_name::().split("::").last().unwrap().split('<').next().unwrap().to_string() + } + + /// Get the full type name of the implementor. For debugging usage. + fn get_full_type_name(&self) -> String { + std::any::type_name::().to_string() + } +} diff --git a/src/hessian_backup/uscf.rs b/src/hessian_backup/uscf.rs new file mode 100644 index 0000000000..13357e08e2 --- /dev/null +++ b/src/hessian_backup/uscf.rs @@ -0,0 +1,540 @@ +//! Hessian implementations for unrestricted SCF. + +use super::prelude::*; +/// Working solver and maintainer of all hessian components for unrestricted SCF method. +pub struct UHessSCF<'a> { + pub mo_coeff: [Tsr; 2], + pub mo_occ: [Tsr; 2], + pub mo_energy: [Tsr; 2], + pub ovlp_obj: 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>, + pub config: HessSCFConfig, + pub atm_list: Option>, + pub result: HashMap, + /// Timing information. Represented by wall time in second. + pub timing: Vec<(String, f64)>, +} + +impl<'a> UHessSCF<'a> { + #[allow(clippy::too_many_arguments)] + pub fn new( + mo_coeff: [Tsr; 2], + mo_occ: [Tsr; 2], + mo_energy: [Tsr; 2], + ovlp_obj: UHessOvlp, + nuc_list: Vec<&'a mut dyn HessNucAPI>, + core_list: Vec<&'a mut dyn UHessCoreAPI>, + el_list: Vec<&'a mut dyn UHessElecInteractAPI>, + config: HessSCFConfig, + atm_list: Option>, + ) -> Self { + Self { + mo_coeff, + mo_occ, + mo_energy, + ovlp_obj, + nuc_list, + core_list, + el_list, + config, + atm_list, + result: HashMap::new(), + timing: Vec::new(), + } + } + + /// Number of atoms over which the Hessian is computed. This is `atm_list.len()` if + /// `atm_list` is `Some`, otherwise the total number of atoms in the molecule. + pub fn natm(&self) -> usize { + match &self.atm_list { + Some(list) => list.len(), + None => self.ovlp_obj.natm(), + } + } + + /// Return the list of (global) atom indices the Hessian is computed for, ordered the same + /// way as the local indexing used in the returned Hessian. + pub fn atm_indices(&self) -> Vec { + match &self.atm_list { + Some(list) => list.clone(), + None => (0..self.ovlp_obj.natm()).collect(), + } + } + + /// Compute the dimensionless CPHF right-hand side, along with necessary intermediates for later + /// steps. + /// + /// # Returns + /// + /// A dictionary containing: + /// - `rhs : shape `[nmo, nocc_α, 3, natm]` and `[nmo, nocc_β, 3, natm]`. The dimensionless CPHF + /// right-hand side. + /// - `f1mo` : shape `[nmo, nocc_α, 3, natm]` and `[nmo, nocc_β, 3, natm]`. The first-order + /// derivative of the Fock matrix in MO basis. + /// - `s1mo` : shape `[nmo, nocc_α, 3, natm]` and `[nmo, nocc_β, 3, natm]`. The first-order + /// derivative of the overlap matrix in MO basis. + pub fn compute_dimless_cphf_rhs(&mut self) -> HashMap<&'static str, Tsr> { + // setups + let t0 = std::time::Instant::now(); + let [α, β] = [0, 1]; + 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.config.level_shift; + let device = mo_coeff[α].device().clone(); + + let nao = mo_coeff[α].shape()[0]; + let nmo = [mo_coeff[α].shape()[1], mo_coeff[β].shape()[1]]; + 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 mocc = [mo_coeff[α].bool_select(-1, &occidx[α]), mo_coeff[β].bool_select(-1, &occidx[β])]; + 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 nocc = [mocc[α].shape()[1], mocc[β].shape()[1]]; + let natm = self.natm(); + let atm_indices = self.atm_indices(); + let atm_list = self.atm_list.as_deref(); + + let e_ai = [evir[α].i((.., None)) - eocc[α].i((None, ..)), evir[β].i((.., None)) - eocc[β].i((None, ..))]; + let e_ai_shift = [&e_ai[α] + level_shift, &e_ai[β] + level_shift]; + + // --- f1mo --- // + + // fock skeleton derivative (core contribution) + let mut f1ao_core: Tsr = rt::zeros(([nao, nao, 3, natm], &device)); + for core_obj in self.core_list.iter() { + let t1 = std::time::Instant::now(); + let mut gen_core_deriv1 = core_obj.generator_deriv1(); + for (A_loc, &A_glob) in atm_indices.iter().enumerate() { + *&mut f1ao_core.i_mut((Ellipsis, A_loc)) += gen_core_deriv1(A_glob); + } + self.timing.push(( + format!("in compute_dimless_cphf_rhs, f1ao_core_{}", core_obj.get_type_name()), + t1.elapsed().as_secs_f64(), + )); + } + + // fock skeleton derivative (electron interaction contribution, half-transformed to bra) + let mut f1bra_el: [Tsr; 2] = + [rt::zeros(([nao, nocc[α], 3, natm], &device)), rt::zeros(([nao, nocc[β], 3, natm], &device))]; + for el_obj in self.el_list.iter_mut() { + let t1 = std::time::Instant::now(); + let bra = el_obj.get_deriv1_bra(&mo_coeff, &mo_occ, atm_list); + f1bra_el[α] += &bra[α]; + f1bra_el[β] += &bra[β]; + self.timing.push(( + format!("in compute_dimless_cphf_rhs, f1bra_el_{}", el_obj.get_type_name()), + t1.elapsed().as_secs_f64(), + )); + } + + // construct whole f1mo + let f1mo_α = mo_coeff[α].t() % (&f1ao_core % &mocc[α] + &f1bra_el[α]); + let f1mo_β = mo_coeff[β].t() % (&f1ao_core % &mocc[β] + &f1bra_el[β]); + + // --- s1mo --- // + + let t1 = std::time::Instant::now(); + + let mut gen_ovlp_deriv1 = self.ovlp_obj.generator_deriv1(); + let mut s1ao: Tsr = rt::zeros(([nao, nao, 3, natm], &device)); + for (A_loc, &A_glob) in atm_indices.iter().enumerate() { + *&mut s1ao.i_mut((Ellipsis, A_loc)) += gen_ovlp_deriv1(A_glob); + } + let s1mo_α = mo_coeff[α].t() % (&s1ao % &mocc[α]); + let s1mo_β = mo_coeff[β].t() % (&s1ao % &mocc[β]); + + self.timing.push(("in compute_dimless_cphf_rhs, s1mo".to_string(), t1.elapsed().as_secs_f64())); + + // --- dimensionless rhs --- // + + let so = [rt::slice!(0, nocc[α]), rt::slice!(0, nocc[β])]; + let sv = [rt::slice!(nocc[α], nmo[α]), rt::slice!(nocc[β], nmo[β])]; + let b1mo_α = &f1mo_α - &s1mo_α * eocc[α].i((None, ..)); + let b1mo_β = &f1mo_β - &s1mo_β * eocc[β].i((None, ..)); + let mut rhs_α = rt::zeros(([nmo[α], nocc[α], 3, natm], &device)); + let mut rhs_β = rt::zeros(([nmo[β], nocc[β], 3, natm], &device)); + *&mut rhs_α.i_mut(sv[α]) += -b1mo_α.i(sv[α]) / &e_ai_shift[α]; + *&mut rhs_β.i_mut(sv[β]) += -b1mo_β.i(sv[β]) / &e_ai_shift[β]; + *&mut rhs_α.i_mut(so[α]) += -0.5 * s1mo_α.i(so[α]); + *&mut rhs_β.i_mut(so[β]) += -0.5 * s1mo_β.i(so[β]); + + self.timing.push(("compute_dimless_cphf_rhs".to_string(), t0.elapsed().as_secs_f64())); + HashMap::from([ + ("f1mo_0", f1mo_α), + ("f1mo_1", f1mo_β), + ("s1mo_0", s1mo_α), + ("s1mo_1", s1mo_β), + ("rhs_0", rhs_α), + ("rhs_1", rhs_β), + ]) + } + + /// Prepare the response for CPHF calculation. + /// + /// This involves all electron-interaction objects. + 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(), + )); + } + 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 CPHF. + /// + /// # 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-HF 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_cphf(&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.config.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_cphf".to_string(), t0.elapsed().as_secs_f64())); + resp + } + + /// Solve the dimensionless CP-HF 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-HF equation. + pub fn solve_dimless_cphf(&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.config.cphf_tol; + let max_cycle = self.config.cphf_max_cycle; + let max_space = self.config.cphf_max_space; + let lindep = self.config.cphf_lindep; + + 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_cphf_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_cphf(&[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_cphf_flattened, rhs_packed.view(), None, tol, max_cycle, max_space, lindep); + 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_cphf".to_string(), t0.elapsed().as_secs_f64())); + [mo1_α, mo1_β] + } + + /// Finalize the CP-HF calculation by computing necessary intermediates for Hessian assembly. + /// + /// + /// # Parameters + /// + /// - `f1mo` : shape `[nmo_α, nocc_α, 3, natm]` and `[nmo_β, nocc_β, 3, natm]`. The first-order + /// derivative of the Fock matrix in MO basis, obtained from + /// [`Self::compute_dimless_cphf_rhs`]. + /// - `s1mo` : shape `[nmo_α, nocc_α, 3, natm]` and `[nmo_β, nocc_β, 3, natm]`. The first-order + /// derivative of the overlap matrix in MO basis, obtained from + /// [`Self::compute_dimless_cphf_rhs`]. + /// - `mo1` : shape `[nmo_α, nocc_α, 3, natm]` and `[nmo_β, nocc_β, 3, natm]`. The perturbation + /// in MO space obtained from Krylov solver. + /// + /// # Returns + /// + /// `HashMap<&str, Tsr>` + /// + /// - `mo1_0`, `mo1_1` : shape `[nmo_α, nocc_α, 3, natm]` and `[nmo_β, nocc_β, 3, natm]`. The + /// finalized perturbation in MO space. + /// - `mo_e1_0`, `mo_e1_1` : shape `[nocc_α, nocc_α, 3, natm]` and `[nocc_β, nocc_β, 3, natm]`. + /// The derivative of occupied orbital energies (Fock matrix) with respect to perturbation. + pub fn finalize_cphf( + &mut self, + f1mo: &[TsrView; 2], + s1mo: &[TsrView; 2], + mo1: &[TsrView; 2], + ) -> HashMap<&'static str, Tsr> { + 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 so = [rt::slice!(0, nocc[α]), rt::slice!(0, nocc[β])]; + let sv = [rt::slice!(nocc[α], nmo[α]), rt::slice!(nocc[β], nmo[β])]; + let e_ai = [evir[α].i((.., None)) - eocc[α].i((None, ..)), evir[β].i((.., None)) - eocc[β].i((None, ..))]; + 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 b1mo_α = &f1mo[α] - &s1mo[α] * eocc[α].i((None, ..)) + &last_resp[α]; + let b1mo_β = &f1mo[β] - &s1mo[β] * eocc[β].i((None, ..)) + &last_resp[β]; + let mut mo1_α = mo1[α].to_owned(); + let mut mo1_β = mo1[β].to_owned(); + mo1_α.i_mut(sv[α]).assign(-b1mo_α.i(sv[α]) / &e_ai[α]); + mo1_β.i_mut(sv[β]).assign(-b1mo_β.i(sv[β]) / &e_ai[β]); + + // get the derivative of fock matrix in occ-occ block (derivative of orbital energy with rotation) + let mo_e1_α = b1mo_α.i(so[α]) + &mo1_α.i(so[α]) * &e_ij[α]; + let mo_e1_β = b1mo_β.i(so[β]) + &mo1_β.i(so[β]) * &e_ij[β]; + + self.timing.push(("finalize_cphf".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_β)]) + } + + pub fn get_cphf_hess( + &self, + f1mo: &[TsrView; 2], + s1mo: &[TsrView; 2], + mo1: &[TsrView; 2], + mo_e1: &[TsrView; 2], + ) -> Tsr { + let [α, β] = [0, 1]; + let natm = self.natm(); + 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 nocc = [occidx[α].iter().filter(|&&x| x).count(), occidx[β].iter().filter(|&&x| x).count()]; + let eocc = [ + self.mo_energy[α].view().bool_select(-1, &occidx[α]), + self.mo_energy[β].view().bool_select(-1, &occidx[β]), + ]; + let so = [rt::slice!(0, nocc[α]), rt::slice!(0, nocc[β])]; + let s1oo = [s1mo[α].i(so[α]), s1mo[β].i(so[β])]; + let device = f1mo[α].device().clone(); + + let mut de_cphf = rt::zeros(([3, 3, natm, natm], &device)); + for A in 0..natm { + for B in 0..=A { + let mut de_BA = de_cphf.i_mut((.., .., B, A)); + for σ in [α, β] { + de_BA += 2 * (f1mo[σ].i((.., .., None, .., A)) * mo1[σ].i((.., .., .., None, B))).sum_axes([0, 1]); + de_BA -= 2 + * (s1mo[σ].i((.., .., None, .., A)) * mo1[σ].i((.., .., .., None, B)) * eocc[σ].i((None, ..))) + .sum_axes([0, 1]); + de_BA -= (s1oo[σ].i((.., .., None, .., A)) * mo_e1[σ].i((.., .., .., None, B))).sum_axes([0, 1]); + } + } + for B in 0..A { + let de_to_copy = de_cphf.i((.., .., B, A)).t().to_owned(); + *&mut de_cphf.i_mut((.., .., A, B)) += de_to_copy; + } + } + de_cphf + } + + pub fn make_cphf_hess(&mut self) -> Tsr { + let pre_cphf_dict = self.compute_dimless_cphf_rhs(); + let f1mo = [pre_cphf_dict.get("f1mo_0").unwrap().view(), pre_cphf_dict.get("f1mo_1").unwrap().view()]; + let s1mo = [pre_cphf_dict.get("s1mo_0").unwrap().view(), pre_cphf_dict.get("s1mo_1").unwrap().view()]; + let rhs = [pre_cphf_dict.get("rhs_0").unwrap().view(), pre_cphf_dict.get("rhs_1").unwrap().view()]; + + self.make_response_preparation(); + let mo1 = self.solve_dimless_cphf(&rhs); + let mo1_view = [mo1[0].view(), mo1[1].view()]; + let finalize_dict = self.finalize_cphf(&f1mo, &s1mo, &mo1_view); + let mo1 = [finalize_dict.get("mo1_0").unwrap().view(), finalize_dict.get("mo1_1").unwrap().view()]; + let mo_e1 = [finalize_dict.get("mo_e1_0").unwrap().view(), finalize_dict.get("mo_e1_1").unwrap().view()]; + + self.get_cphf_hess(&f1mo, &s1mo, &mo1, &mo_e1) + } + + /// Compute the total skeleton contribution to the Hessian. + /// + /// **Total** means that we sum over all skeleton contributions from both core and + /// electron-interaction objects. + /// + /// # Returns + /// + /// - `de_skeleton` : shape `[3, 3, natm, natm]`. The total skeleton contribution to the + /// Hessian. + pub fn make_skeleton_hess(&mut self) -> Tsr { + let [α, β] = [0, 1]; + let natm = self.natm(); + let mo_coeff = [self.mo_coeff[α].view(), self.mo_coeff[β].view()]; + let mo_occ = [self.mo_occ[α].view(), self.mo_occ[β].view()]; + let atm_list = self.atm_list.as_deref(); + + let device = self.mo_coeff[α].device().clone(); + let mut de_skeleton = rt::zeros(([3, 3, natm, natm], &device)); + for nuc_obj in self.nuc_list.iter_mut() { + let t0 = std::time::Instant::now(); + let de_nuc = nuc_obj.make_skeleton_hess(atm_list); + let nuc_obj_name = nuc_obj.get_type_name(); + self.result.insert(format!("de_skeleton_{}", nuc_obj_name), de_nuc.to_owned()); + self.timing.push((format!("de_skeleton_{}", nuc_obj_name,), t0.elapsed().as_secs_f64())); + de_skeleton += de_nuc; + } + for core_obj in self.core_list.iter_mut() { + let t0 = std::time::Instant::now(); + let de_core = core_obj.make_skeleton_hess(&mo_coeff, &mo_occ, atm_list); + let core_obj_name = core_obj.get_type_name(); + self.result.insert(format!("de_skeleton_{}", core_obj_name), de_core.to_owned()); + self.timing.push((format!("de_skeleton_{}", core_obj_name,), t0.elapsed().as_secs_f64())); + de_skeleton += de_core; + } + for el_obj in self.el_list.iter_mut() { + let t0 = std::time::Instant::now(); + let de_el = el_obj.make_skeleton_hess(&mo_coeff, &mo_occ, atm_list); + let el_obj_name = el_obj.get_type_name(); + self.result.insert(format!("de_skeleton_{}", el_obj_name), de_el.to_owned()); + self.timing.push((format!("de_skeleton_{}", el_obj_name,), t0.elapsed().as_secs_f64())); + de_skeleton += de_el; + } + de_skeleton + } + + /// Compute the total Hessian by summing over skeleton, overlap, and CP-HF contributions. + /// + /// # Returns + /// + /// - `de_hess` : shape `[3, 3, natm, natm]`. The total Hessian. + pub fn make_hess(&mut self) -> Tsr { + let t0 = std::time::Instant::now(); + let [α, β] = [0, 1]; + let dme0 = [ + get_dme0_restricted(self.mo_coeff[α].view(), self.mo_occ[α].view(), self.mo_energy[α].view()), + get_dme0_restricted(self.mo_coeff[β].view(), self.mo_occ[β].view(), self.mo_energy[β].view()), + ]; + let atm_list = self.atm_list.clone(); + + let de_skeleton = self.make_skeleton_hess(); + + let t1 = std::time::Instant::now(); + let de_ovlp = self.ovlp_obj.make_hess([dme0[α].view(), dme0[β].view()], atm_list.as_deref()); + self.result.insert("de_ovlp".to_string(), de_ovlp.to_owned()); + self.timing.push(("de_ovlp".to_string(), t1.elapsed().as_secs_f64())); + + let t1 = std::time::Instant::now(); + let de_cphf = self.make_cphf_hess(); + self.result.insert("de_cphf".to_string(), de_cphf.to_owned()); + self.timing.push(("de_cphf".to_string(), t1.elapsed().as_secs_f64())); + + let de_tot = de_skeleton + de_ovlp + de_cphf; + self.result.insert("de_tot".to_string(), de_tot.to_owned()); + self.timing.push(("de_tot".to_string(), t0.elapsed().as_secs_f64())); + de_tot + } +} diff --git a/src/hessian_backup/vib.rs b/src/hessian_backup/vib.rs new file mode 100644 index 0000000000..29c019c1fb --- /dev/null +++ b/src/hessian_backup/vib.rs @@ -0,0 +1,1131 @@ +//! Vibrational (harmonic) analysis module. +//! +//! All tensors are column-major. Geometry is stored as `[3, natm]`, masses as +//! `[natm]`, Hessian as `[3*natm, 3*natm]`. +//! +//! No IR intensity / dipole-derivative terms at this stage. +//! +//! # Note +//! +//! This module is direct transformation from psi4 (psi4/psi4/driver/qcdb/vib.py). +//! This file contains AI assisted code, and not fully reviewed by human. + +use super::prelude::*; + +// --------------------------------------------------------------------------- +// Physical constants (CODATA2014, matching the Python implementation) +// --------------------------------------------------------------------------- +const NA: f64 = 6.022140857e23; +const HARTREE2J: f64 = 4.35974465e-18; +const C: f64 = 299792458.0; +const BOHR2ANG: f64 = 0.52917721067; +const H: f64 = 6.62607004e-34; +const KB: f64 = 1.38064852e-23; +const R_GAS: f64 = 8.3144598; +const HARTREE2KJMOL: f64 = 2625.4996382852164; +const HARTREE2KCALMOL: f64 = 627.5094737775374; +const HARTREE2WAVENUMBERS: f64 = 219474.6313702; +const AMU2KG: f64 = 1.66053904e-27; + +/// Tolerance for detecting nearly-linear geometries in `get_tr_space`. +pub const LINEAR_A_TOL: f64 = 1.0e-2; + +/// cm⁻¹ conversion factor from force-constant eigenvalues (atomic units). +fn uconv_cm1() -> f64 { + (NA * HARTREE2J * 1.0e19).sqrt() / (2.0 * std::f64::consts::PI * C * BOHR2ANG) +} + +// --------------------------------------------------------------------------- +// Translation / rotation space +// --------------------------------------------------------------------------- + +/// Idealized translation + rotation basis vectors. +/// +/// # Parameters +/// - `mass` : `[natm]` atomic masses [u]. +/// - `geom` : `[3, natm]` Cartesian geometry [a₀] (column-major). +/// - `space` : `"T"`, `"R"`, or `"TR"`. +/// +/// # Returns +/// `tr` : `[3*natm, nrt]` orthonormal basis (each column is a TR vector). +pub fn get_tr_space(mass: TsrView, geom: TsrView, space: &str) -> Tsr { + _get_tr_space(mass, geom, space, None) +} + +/// Internal helper with custom SVD tolerance. When `tol` is `Some(t)`, that +/// absolute value is used as the singular-value cutoff; when `None`, the default +/// machine-epsilon-based tolerance (`ndof × max(s) × ε`) is used. +fn _get_tr_space(mass: TsrView, geom: TsrView, space: &str, tol_user: Option) -> Tsr { + let device = geom.device().clone(); + let natm = geom.shape()[1]; + let ndof = 3 * natm; + + let mass_vec = mass.reshape(-1).to_vec(); + + // sqrtmmm: [3*natm], interleaved per-atom (m_a repeated 3×): [m0,m0,m0, m1,m1,m1, ...] + let mut sqrtmmm = Vec::with_capacity(ndof); + let mut xxx = Vec::with_capacity(ndof); + let mut yyy = Vec::with_capacity(ndof); + let mut zzz = Vec::with_capacity(ndof); + for a in 0..natm { + let sm = mass_vec[a].sqrt(); + let gx = geom[[0, a]]; + let gy = geom[[1, a]]; + let gz = geom[[2, a]]; + for _ in 0..3 { + sqrtmmm.push(sm); + xxx.push(gx); + yyy.push(gy); + zzz.push(gz); + } + } + let sqrtmmm = rt::asarray((sqrtmmm, &device)); + let xxx = rt::asarray((xxx, &device)); + let yyy = rt::asarray((yyy, &device)); + let zzz = rt::asarray((zzz, &device)); + + // unit vectors ux/uy/uz each [3*natm]: per-atom [1,0,0]/[0,1,0]/[0,0,1] + let mut ux = Vec::with_capacity(ndof); + let mut uy = Vec::with_capacity(ndof); + let mut uz = Vec::with_capacity(ndof); + for _ in 0..natm { + ux.extend_from_slice(&[1.0, 0.0, 0.0]); + uy.extend_from_slice(&[0.0, 1.0, 0.0]); + uz.extend_from_slice(&[0.0, 0.0, 1.0]); + } + let ux = rt::asarray((ux, &device)); + let uy = rt::asarray((uy, &device)); + let uz = rt::asarray((uz, &device)); + + // translation [3*natm] + let t1 = &sqrtmmm * &ux; + let t2 = &sqrtmmm * &uy; + let t3 = &sqrtmmm * &uz; + // rotation [3*natm] + let r4 = &sqrtmmm * (&yyy * &uz - &zzz * &uy); + let r5 = &sqrtmmm * (&zzz * &ux - &xxx * &uz); + let r6 = &sqrtmmm * (&xxx * &uy - &yyy * &ux); + + let mut cols: Vec = Vec::new(); + if space.contains('T') { + cols.push(t1); + cols.push(t2); + cols.push(t3); + } + if space.contains('R') { + cols.push(r4); + cols.push(r5); + cols.push(r6); + } + if cols.is_empty() { + cols.push(rt::asarray((vec![0.0_f64; ndof], &device))); + } + + let tr_raw: Tsr = rt::stack((cols, -1)); // [3*natm, n_raw] + + // orthonormal basis for the column space via SVD: tr_raw = U S Vh, Q = U[:, :num] + let (u, s, _vh): (Tsr, Tsr, Tsr) = rt::linalg::svd(tr_raw.view()).into(); + let svec = s.reshape(-1).to_vec(); + let tol = match tol_user { + Some(t) => t, + None => { + let smax = svec.iter().copied().fold(0.0_f64, f64::max); + (ndof as f64) * smax * f64::EPSILON + }, + }; + let num = svec.iter().filter(|&&x| x > tol).count(); + u.i((.., ..num)).to_owned() +} + +// --------------------------------------------------------------------------- +// Rotation constants and rotor type +// --------------------------------------------------------------------------- + +/// Rotational constants. +/// +/// # Parameters +/// - `mass` : `[natm]` atomic masses [u]. +/// - `atom_coords` : `[3, natm]` **mass-centred** geometry [a₀]. +/// - `unit` : `"GHz"` or `"wavenumber"`. +/// +/// # Returns +/// `e` : `[3]` rotational constants (sorted ascending). +pub fn rotation_const(mass: TsrView, atom_coords: TsrView, unit: &str) -> Tsr { + let device = mass.device().clone(); + // im = Σ_a m_a r_{a,r} r_{a,s} -> [3, 3]; atom_coords is [3, natm] + let weighted = &atom_coords * mass.i((None, ..)); // [3, natm] + let im_rr = &weighted % atom_coords.t(); // [3, 3] + let trace = im_rr.diagonal(None).sum(); + // I = trace*I - Σ m r r^T + let eye: Tsr = rt::eye((3, &device)); + let im = &eye * trace - &im_rr; + + let mut e = rt::linalg::eigvalsh(im.view()); + let mut evec = e.reshape(-1).to_vec(); + for v in evec.iter_mut() { + if v.abs() < 1e-9 { + *v = 0.0; + } + } + e = rt::asarray((evec, &device)); + + let unit_im = AMU2KG * 5.2917721092e-11_f64.powi(2); + let unit_hz = 1.0545718001391127e-34 / (4.0 * std::f64::consts::PI * unit_im); + + let unit_lower = unit.to_lowercase(); + let conv = if unit_lower == "ghz" { + 1e-9 + } else if unit_lower == "wavenumber" { + 1.0 / C * 1e-2 + } else { + panic!("Unsupported unit {}", unit); + }; + let mut out = Vec::with_capacity(3); + for &x in e.reshape(-1).to_vec().iter() { + if x.abs() < 1e-30 { + out.push(f64::INFINITY); + } else { + out.push(unit_hz / x * conv); + } + } + rt::asarray((out, &device)) +} + +/// Classify rotor type from rotational constants [GHz]. +/// +/// Returns `"ATOM"`, `"LINEAR"`, or `"REGULAR"`. +pub fn get_rotor_type(rot_const_ghz: TsrView) -> &'static str { + let v = rot_const_ghz.reshape(-1).to_vec(); + if v.iter().all(|&x| x > 1e8) { + "ATOM" + } else if v[0] > 1e8 && (v[1] - v[2]).abs() < 1e-3 { + "LINEAR" + } else { + "REGULAR" + } +} + +// --------------------------------------------------------------------------- +// VibInfo result container +// --------------------------------------------------------------------------- + +/// Output of [`harmonic_analysis`]. +/// +/// All tensors are column-major. Modes are stored as **columns** of `q`/`w`/`x` +/// (each `[ndof, ndof]`). Per-mode quantities (`omega`, `mu`, `k`, ...) are +/// `[ndof]`. +#[derive(Clone)] +pub struct VibInfo { + /// Number of degrees of freedom (3 × natm). + pub ndof: usize, + /// Frequency `[ndof]`, stored as `(real, imag)` pairs (imag > 0 ⇒ imaginary mode). + pub omega: Vec, + /// Imaginary flag per mode (true if `imag > real`). + pub imag: Vec, + /// Mass-weighted normal modes, flattened `[ndof × ndof]` in col-major order + /// (columns = modes): element `(row, col)` = `q[row + col * ndof]`. + pub q: Vec, + /// Un-mass-weighted normal modes, same flat layout `[ndof × ndof]`. + pub w: Vec, + /// Normalized un-mass-weighted normal modes, same flat layout `[ndof × ndof]`. + pub x: Vec, + /// Degeneracy count per mode `[ndof]` (as `i64`). + pub degeneracy: Vec, + /// TR/V classification per mode: `"TR"`, `"V"`, or `"-"`. + pub trv: Vec<&'static str>, + /// Reduced mass `[ndof]` [u]. + pub mu: Vec, + /// Force constant `[ndof]` [mDyne/Å]. + pub k: Vec, + /// RMS deviation v=0 `[ndof]` [a₀·u^½]. + pub dq0: Vec, + /// Turning point v=0 (mass-weighted) `[ndof]` [a₀·u^½]. + pub qtp0: Vec, + /// Turning point v=0 (Cartesian) `[ndof]` [a₀]. + pub xtp0: Vec, + /// Characteristic vibrational temperature `[ndof]` [K]. + pub theta_vib: Vec, +} + +impl VibInfo { + /// Frequency for mode `i` as a signed real number: positive for real modes, + /// negative for imaginary modes (the "imaginary freq as negative" convention). + /// For a real mode `omega[i]` is the frequency; for an imaginary mode + /// `omega[i]` holds the magnitude and this returns its negation. + pub fn freq_signed(&self, i: usize) -> f64 { + if self.imag[i] { + -self.omega[i] + } else { + self.omega[i] + } + } + + /// Indices of vibrational modes (`TRV == "V"`). + pub fn vib_indices(&self) -> Vec { + self.trv.iter().enumerate().filter(|(_, &t)| t == "V").map(|(i, _)| i).collect() + } + + /// Number of atoms. + pub fn nat(&self) -> usize { + self.ndof / 3 + } + + /// Returns element `(row, col)` of a flat col-major `[ndof × ndof]` normal- + /// coordinate array (`q`, `w`, or `x`). + #[inline] + pub fn normco_index(data: &[f64], ndof: usize, row: usize, col: usize) -> f64 { + data[row + col * ndof] + } + + /// Reconstruct a col-major `[ndof, ndof]` tensor from a flat normal-mode + /// slice (e.g. `&self.q`, `&self.w`, `&self.x`) — for testing / advanced use. + pub fn normco_matrix(&self, data: &[f64], device: &DeviceBLAS) -> Tsr { + rt::asarray((data.to_vec(), [self.ndof, self.ndof].f(), device)) + } +} + +// --------------------------------------------------------------------------- +// Helper: standardize column phases so extreme element is positive +// --------------------------------------------------------------------------- + +/// Return a copy of `q` (column-major `[n, m]`) where each column is scaled so +/// that its element of maximum absolute value is positive (tol 1e-2). +fn phase_cols_to_max_element(q: TsrView, tol: f64) -> Tsr { + let (n, m) = (q.shape()[0], q.shape()[1]); + let mut out = q.to_owned(); + for v in 0..m { + let mut vextreme = 0.0_f64; + for r in 0..n { + vextreme = vextreme.max(out[[r, v]].abs()); + } + // first index whose fabs equals vextreme within tol + let mut iextreme = 0; + for r in 0..n { + if (vextreme - out[[r, v]].abs()) < tol { + iextreme = r; + break; + } + } + if out[[iextreme, v]] < 0.0 { + for r in 0..n { + out[[r, v]] = -out[[r, v]]; + } + } + } + out +} + +/// Check whether vector `vec` (length `n`) lies in the subspace spanned by the +/// columns of `space` (`[n, nrt]`), via SVD: `vec` is in the space iff stacking +/// it as an extra column does not increase the rank. +fn vec_in_space(vec: &[f64], space: TsrView, tol: f64) -> bool { + let device = space.device().clone(); + let nrt = space.shape()[1]; + let vec_t = rt::asarray((vec.to_vec(), &device)); // [n] + let mut cols: Vec = Vec::with_capacity(nrt + 1); + for c in 0..nrt { + cols.push(space.i((.., c)).to_owned()); // [n] + } + cols.push(vec_t); + let merged: Tsr = rt::stack((cols, -1)); // [n, nrt+1] + let (_u, s, _vh): (Tsr, Tsr, Tsr) = rt::linalg::svd(merged.view()).into(); + let svec = s.reshape(-1).to_vec(); + svec.last().copied().unwrap_or(0.0) < tol +} + +// --------------------------------------------------------------------------- +// harmonic_analysis +// --------------------------------------------------------------------------- + +/// Extract frequencies, normal modes, and other properties from an electronic +/// Hessian. Rust port of `pyhessref.vib.harmonic_analysis`. +/// +/// # Parameters +/// - `hess` : `[3*natm, 3*natm]` non-mass-weighted Cartesian Hessian [E_h/a₀²]. +/// - `geom` : `[3, natm]` Cartesian geometry [a₀] (column-major). +/// - `mass` : `[natm]` atomic masses [u]. +/// - `project_trans` : project out idealized translations. +/// - `project_rot` : project out idealized rotations. +/// +/// # Returns +/// [`VibInfo`] with all 3×natm modes (TR + V). Geometry must be in **bohr**. +pub fn harmonic_analysis( + hess: TsrView, + geom: TsrView, + mass: TsrView, + project_trans: bool, + project_rot: bool, +) -> VibInfo { + let device = hess.device().clone(); + let natm = mass.shape()[0]; + let ndof = 3 * natm; + assert_eq!(geom.shape().as_slice(), &[3, natm], "geom must be [3, natm]"); + assert_eq!(hess.shape().as_slice(), &[ndof, ndof], "hess must be [3*natm, 3*natm]"); + + let nmwhess = hess.into_contig(ColMajor); + + // --------------- translation / rotation projector --------------- + let space = format!("{}{}", if project_trans { "T" } else { "" }, if project_rot { "R" } else { "" }); + // use LINEAR_A_TOL so nearly-linear geometries correctly drop to 5 TR dof + let tr_space = _get_tr_space(mass.view(), geom.view(), &space, Some(LINEAR_A_TOL)); // [ndof, nrt] + + // projector P = I - Σ |tr⟩⟨tr| + let nrt = tr_space.shape()[1]; + let mut p: Tsr = rt::zeros(([ndof, ndof], &device)); + for i in 0..ndof { + p[[i, i]] = 1.0; + } + for c in 0..nrt { + // outer(tr_col, tr_col): [ndof, ndof] + let trc = tr_space.i((.., c)).to_owned(); // [ndof] + let outer = &trc.i((.., None)) * &trc.i((None, ..)); // [ndof, ndof] + *&mut p -= &outer; + } + + // --------------- mass-weight & solve --------------- + // sqrtmmm / sqrtmmminv : [ndof], interleaved per-atom + let mass_vec = mass.reshape(-1).to_vec(); + let mut sqrtmmm = Vec::with_capacity(ndof); + let mut sqrtmmminv = Vec::with_capacity(ndof); + for a in 0..natm { + let sm = mass_vec[a].sqrt(); + for _ in 0..3 { + sqrtmmm.push(sm); + sqrtmmminv.push(1.0 / sm); + } + } + let sqrtmmminv_t = rt::asarray((&sqrtmmminv, &device)); + + // mwhess[i,j] = hess[i,j] / sqrt(m_i * m_j) + // numpy: (sqrtmmminv[:,None] * nmwhess) * sqrtmmminv[None,:] + // col-major: nmwhess * sqrtmmminv[:, None] broadcasts axis-0 ; then * sqrtmmminv[None,:] axis-1 + let mwhess = (&nmwhess * sqrtmmminv_t.i((.., None))) * sqrtmmminv_t.i((None, ..)); + + // project & diagonalise: mwhess_proj = P^T mwhess P (P symmetric so P^T=P) + let mwhess_proj = p.t() % (&mwhess % &p); + + let (fc_au, qL_raw): (Tsr, Tsr) = rt::linalg::eigh(mwhess_proj.view()).into(); + // eigh returns ascending eigenvalues already; eigenvectors are columns. + + // sort ascending (eigh already ascending, but be explicit/safe) + let fc_vec = fc_au.reshape(-1).to_vec(); + let mut order: Vec = (0..ndof).collect(); + order.sort_by(|&a, &b| fc_vec[a].partial_cmp(&fc_vec[b]).unwrap()); + + // reorder eigenvalues and eigenvector columns + let fc_sorted: Vec = order.iter().map(|&i| fc_vec[i]).collect(); + let mut qL: Tsr = rt::zeros(([ndof, ndof], &device)); + for (new, &old) in order.iter().enumerate() { + for r in 0..ndof { + qL[[r, new]] = qL_raw[[r, old]]; + } + } + // phase convention + let qL = phase_cols_to_max_element(qL.view(), 1.0e-2); + + // --------------- frequencies (complex sqrt) --------------- + let uconv_cm = uconv_cm1(); + // omega = sqrt(fc) * uconv ; imaginary if fc < 0 + let mut omega = Vec::with_capacity(ndof); + let mut imag = Vec::with_capacity(ndof); + for &fc in fc_sorted.iter() { + if fc < 0.0 { + let mag = (-fc).sqrt() * uconv_cm; + omega.push(mag); // store magnitude; imag flag set + imag.push(true); + } else { + omega.push(fc.sqrt() * uconv_cm); + imag.push(false); + } + } + let omega_real: Vec = (0..ndof).map(|i| if imag[i] { 0.0 } else { omega[i] }).collect(); + + // --------------- degeneracies (group by round(omega_real, 1)) --------------- + // Note: Python uses round(frequency_cm_1, 1) on the complex array (rounds real part). + let mut degeneracy = vec![0i64; ndof]; + { + // group indices by rounded real frequency + let mut keys: Vec<(i64, usize)> = (0..ndof).map(|i| ((omega_real[i] * 10.0).round() as i64, i)).collect(); + keys.sort_by_key(|&(k, _)| k); + let mut start = 0; + while start < keys.len() { + let k = keys[start].0; + let mut end = start + 1; + while end < keys.len() && keys[end].0 == k { + end += 1; + } + let count = (end - start) as i64; + for j in start..end { + degeneracy[keys[j].1] = count; + } + start = end; + } + } + + // --------------- TR / V classification --------------- + // vec_in_space(qL[:, i], tr_space rows) + let mut trv: Vec<&'static str> = Vec::with_capacity(ndof); + for i in 0..ndof { + let qcol: Vec = (0..ndof).map(|r| qL[[r, i]]).collect(); + if vec_in_space(&qcol, tr_space.view(), 1.0e-4) { + trv.push("TR"); + } else if omega_real[i].abs() < 1.0e-3 { + trv.push("-"); + } else { + trv.push("V"); + } + } + + // --------------- conversion factors --------------- + let uconv_mdyne_a = 0.1 * (2.0 * std::f64::consts::PI * C).powi(2) / NA; + let uconv_S = ((C * (2.0 * std::f64::consts::PI * BOHR2ANG).powi(2)) / (H * NA * 1.0e21)).sqrt(); + + // --------------- normal modes & reduced mass --------------- + // w = m^{-1/2} q ; w[a,i] = q[a,i] / sqrt(m_a) + // numpy: sqrtmmminv[:,None] * q (broadcast rows) + // col-major: q * sqrtmmminv[:, None] + let wL = &qL * sqrtmmminv_t.i((.., None)); + // column-wise L2 norm: l2_norm_axes(0) sums over axis 0 (rows) per column + let w_norm = wL.l2_norm_axes(0); // [ndof] + let w_norm_vec = w_norm.reshape(-1).to_vec(); + let mu: Vec = w_norm_vec.iter().map(|&n| 1.0 / (n * n)).collect(); + + // x = sqrt(mu) * w ; numpy: sqrt(mu) * w (mu is [ndof], broadcasts over columns/axis-1) + // col-major: w * sqrt(mu)[None, :] → broadcasts over axis 1 + let sqrt_mu: Vec = mu.iter().map(|&m| m.sqrt()).collect(); + let sqrt_mu_t = rt::asarray((&sqrt_mu, &device)); + let xL = &wL * sqrt_mu_t.i((None, ..)); + + // --------------- force constants --------------- + // k = mu * omega^2 * uconv_mdyne_a (uses real part of omega) + let k: Vec = (0..ndof).map(|i| mu[i] * omega_real[i] * omega_real[i] * uconv_mdyne_a).collect(); + + // --------------- turning points (v=0) --------------- + let tp_rnc = (2.0 * 0.0 + 1.0_f64).sqrt(); // = 1 + let mut qtp0 = vec![0.0_f64; ndof]; + let mut xtp0 = vec![0.0_f64; ndof]; + for i in 0..ndof { + let denom_q = omega_real[i].sqrt() * uconv_S; + qtp0[i] = if denom_q == 0.0 || !denom_q.is_finite() { 0.0 } else { tp_rnc / denom_q }; + let denom_x = (omega_real[i] * mu[i]).sqrt() * uconv_S; + xtp0[i] = if denom_x == 0.0 || !denom_x.is_finite() { 0.0 } else { tp_rnc / denom_x }; + } + let dq0: Vec = qtp0.iter().map(|&q| q / 2.0_f64.sqrt()).collect(); + + // --------------- characteristic vibrational temperature --------------- + let uconv_K = 100.0 * H * C / KB; + let theta_vib: Vec = omega_real.iter().map(|&w| w * uconv_K).collect(); + + let q_vec = qL.reshape(-1).to_vec(); + let w_vec = wL.reshape(-1).to_vec(); + let x_vec = xL.reshape(-1).to_vec(); + + VibInfo { ndof, omega, imag, q: q_vec, w: w_vec, x: x_vec, degeneracy, trv, mu, k, dq0, qtp0, xtp0, theta_vib } +} + +// --------------------------------------------------------------------------- +// Rotor type enum for thermo +// --------------------------------------------------------------------------- + +/// Rotor classification for thermochemistry. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RotorType { + Atom, + Linear, + Regular, +} + +impl RotorType { + /// Parse from a string (`"RT_ATOM"` / `"RT_LINEAR"` / `"ATOM"` / `"LINEAR"` / + /// anything else ⇒ `Regular`). + pub fn parse(s: &str) -> Self { + match s { + "RT_ATOM" | "ATOM" => RotorType::Atom, + "RT_LINEAR" | "LINEAR" => RotorType::Linear, + _ => RotorType::Regular, + } + } + + /// Infer from rotational constants [GHz] (uses [`get_rotor_type`]). + pub fn from_rot_const_ghz(rot_const_ghz: TsrView) -> Self { + match get_rotor_type(rot_const_ghz) { + "ATOM" => RotorType::Atom, + "LINEAR" => RotorType::Linear, + _ => RotorType::Regular, + } + } +} + +// --------------------------------------------------------------------------- +// ThermoInfo result container +// --------------------------------------------------------------------------- + +/// Output of [`thermo`]. All energy/heat-capacity values in atomic units +/// (S/Cv/Cp in [mEh/K], ZPE/E/H/G in [Eh]). +#[derive(Clone, Debug)] +pub struct ThermoInfo { + // conditions + pub e0: f64, + pub b: [f64; 3], // rotational constants [cm⁻¹] + pub b_ghz: [f64; 3], // rotational constants [GHz] + pub rotor_type: String, // "ATOM" / "LINEAR" / "REGULAR" + pub sigma: i64, + pub t: f64, + pub p: f64, + // component contributions (each 4: elec, trans, rot, vib) + pub s: [f64; 4], // [mEh/K] + pub cv: [f64; 4], + pub cp: [f64; 4], + pub zpe: [f64; 4], // [Eh]; elec/trans/rot entries are 0 except vib + pub e: [f64; 4], + pub h: [f64; 4], + pub g: [f64; 4], + // totals + pub s_tot: f64, + pub cv_tot: f64, + pub cp_tot: f64, + pub zpe_corr: f64, + pub e_corr: f64, + pub h_corr: f64, + pub g_corr: f64, + pub zpe_tot: f64, + pub e_tot: f64, + pub h_tot: f64, + pub g_tot: f64, +} + +/// Indices for the four thermo components. +pub const ELEC: usize = 0; +pub const TRANS: usize = 1; +pub const ROT: usize = 2; +pub const VIB: usize = 3; + +/// Thermochemical analysis from harmonic vibrational output. +/// +/// # Parameters +/// - `vib` : [`VibInfo`] from [`harmonic_analysis`]. +/// - `t` : temperature [K]. +/// - `p` : pressure [Pa]. +/// - `multiplicity` : spin multiplicity. +/// - `molecular_mass` : total molecular mass [u]. +/// - `e0` : electronic energy at well bottom [Eh]. +/// - `sigma` : rotational (external) symmetry number. +/// - `rot_const` : `[3]` rotational constants [cm⁻¹]. +/// - `rotor_type` : [`RotorType`]; use `None`-equivalent by passing the result of +/// [`RotorType::from_rot_const_ghz`]. +#[allow(clippy::too_many_arguments)] +pub fn thermo( + vib: &VibInfo, + t: f64, + p: f64, + multiplicity: i64, + molecular_mass: f64, + e0: f64, + sigma: i64, + rot_const: &[f64], + rotor_type: RotorType, +) -> ThermoInfo { + // sm[(quantity, term)] before unit conversion: S/Cv/Cp unitless, ZPE/E/H/G in [K] + let mut s = [0.0_f64; 4]; + let mut cv = [0.0_f64; 4]; + let mut cp = [0.0_f64; 4]; + let mut zpe = [0.0_f64; 4]; + let mut e = [0.0_f64; 4]; + let mut h = [0.0_f64; 4]; + let mut g = [0.0_f64; 4]; + + // ---------- electronic ---------- + s[ELEC] = (multiplicity as f64).ln(); + + // ---------- translational ---------- + let beta = 1.0 / (KB * t); + let q_trans = (2.0 * std::f64::consts::PI * molecular_mass * AMU2KG / (beta * H * H)).powf(1.5) * NA / (beta * p); + s[TRANS] = 2.5 + (q_trans / NA).ln(); + cv[TRANS] = 1.5; + cp[TRANS] = 2.5; + e[TRANS] = 1.5 * t; + h[TRANS] = 2.5 * t; + + // ---------- rotational ---------- + match rotor_type { + RotorType::Atom => {}, + RotorType::Linear => { + let q_rot = 1.0 / (beta * (sigma as f64) * 100.0 * C * H * rot_const[1]); + s[ROT] = 1.0 + q_rot.ln(); + cv[ROT] = 1.0; + cp[ROT] = 1.0; + e[ROT] = t; + }, + RotorType::Regular => { + let phi = [ + rot_const[0] * 100.0 * C * H / KB, + rot_const[1] * 100.0 * C * H / KB, + rot_const[2] * 100.0 * C * H / KB, + ]; + let q_rot = + std::f64::consts::PI.sqrt() * t.powf(1.5) / ((sigma as f64) * (phi[0] * phi[1] * phi[2]).sqrt()); + s[ROT] = 1.5 + q_rot.ln(); + cv[ROT] = 1.5; + cp[ROT] = 1.5; + e[ROT] = 1.5 * t; + }, + } + h[ROT] = e[ROT]; + + // ---------- vibrational ---------- + // vib-only modes, exclude imaginary + let vib_idx = vib.vib_indices(); + let mut filtered_theta: Vec = Vec::new(); + for &i in &vib_idx { + if !vib.imag[i] { + filtered_theta.push(vib.theta_vib[i]); + } + } + let t_safe = t.max(1e-14); + let rT: Vec = filtered_theta.iter().map(|&th| th / t_safe).collect(); + + // S_vib = Σ [ rT/(e^rT - 1) - ln(1 - e^-rT) ] + let s_vib: f64 = rT.iter().map(|&r| r / r.exp_m1() - (1.0 - (-r).exp()).ln()).sum(); + // Cv_vib = Σ [ e^rT * (rT/(e^rT - 1))^2 ] + let cv_vib: f64 = rT + .iter() + .map(|&r| { + let denom = r.exp_m1(); + r.exp() * (r / denom).powi(2) + }) + .sum(); + let zpe_vib = rT.iter().sum::() * t / 2.0; + let e_vib = zpe_vib + rT.iter().map(|&r| r * t / r.exp_m1()).sum::(); + + s[VIB] = s_vib; + cv[VIB] = cv_vib; + cp[VIB] = cv_vib; + zpe[VIB] = zpe_vib; + e[VIB] = e_vib; + h[VIB] = e_vib; + + // ---------- Gibbs: G = H - T*S ---------- + for i in 0..4 { + g[i] = h[i] - t * s[i]; + } + + // ---------- convert to atomic units ---------- + let uconv_r_ehk = R_GAS / HARTREE2KJMOL; // R [Eh/K] (×1000 → mEh/K) + for i in 0..4 { + s[i] *= uconv_r_ehk; // [mEh/K] + cv[i] *= uconv_r_ehk; + cp[i] *= uconv_r_ehk; + zpe[i] *= uconv_r_ehk * 0.001; // [Eh] + e[i] *= uconv_r_ehk * 0.001; + h[i] *= uconv_r_ehk * 0.001; + g[i] *= uconv_r_ehk * 0.001; + } + + // ---------- totals ---------- + let s_tot: f64 = s.iter().sum(); + let cv_tot: f64 = cv.iter().sum(); + let cp_tot: f64 = cp.iter().sum(); + let zpe_corr: f64 = zpe.iter().sum(); + let e_corr: f64 = e.iter().sum(); + let h_corr: f64 = h.iter().sum(); + let g_corr: f64 = g.iter().sum(); + let zpe_tot = e0 + zpe_corr; + let e_tot = e0 + e_corr; + let h_tot = e0 + h_corr; + let g_tot = e0 + g_corr; + + let b_ghz = [rot_const[0] * 29.9792458, rot_const[1] * 29.9792458, rot_const[2] * 29.9792458]; + let rotor_str = match rotor_type { + RotorType::Atom => "ATOM", + RotorType::Linear => "LINEAR", + RotorType::Regular => "REGULAR", + }; + + ThermoInfo { + e0, + b: [rot_const[0], rot_const[1], rot_const[2]], + b_ghz, + rotor_type: rotor_str.to_string(), + sigma, + t, + p, + s, + cv, + cp, + zpe, + e, + h, + g, + s_tot, + cv_tot, + cp_tot, + zpe_corr, + e_corr, + h_corr, + g_corr, + zpe_tot, + e_tot, + h_tot, + g_tot, + } +} + +// --------------------------------------------------------------------------- +// Pretty-printer (port of pyhessref.vib.print_vibs) +// --------------------------------------------------------------------------- + +/// Which normal-coordinate definition to print. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum NormCo { + /// mass-weighted (`q`) + Q, + /// un-mass-weighted (`w`) + W, + /// normalized un-mass-weighted (`x`, default) + X, +} + +impl NormCo { + /// Returns a col-major `[ndof, ndof]` tensor view of the chosen normal + /// coordinate, reconstructed from the flat `Vec` storage. + fn select<'a>(&self, vib: &'a VibInfo) -> TsrView<'a> { + let data: &[f64] = match self { + NormCo::Q => &vib.q, + NormCo::W => &vib.w, + NormCo::X => &vib.x, + }; + rt::asarray((data, [vib.ndof, vib.ndof], &DeviceBLAS::default())) + } +} + +/// Right-justify `s` to width `w`, padding left with spaces. +fn rjust(s: &str, w: usize) -> String { + let len = s.chars().count(); + if len >= w { + s.to_string() + } else { + format!("{}{}", " ".repeat(w - len), s) + } +} + +/// Left-justify `s` to width `w`, padding right with spaces. +fn ljust(s: &str, w: usize) -> String { + let len = s.chars().count(); + if len >= w { + s.to_string() + } else { + format!("{}{}", s, " ".repeat(w - len)) + } +} + +/// Center `s` to width `w`. +fn center(s: &str, w: usize) -> String { + let len = s.chars().count(); + if len >= w { + s.to_string() + } else { + let total = w - len; + let left = total / 2; + let right = total - left; + format!("{}{}{}", " ".repeat(left), s, " ".repeat(right)) + } +} + +/// Pretty-print vibrational analysis results. Rust port of +/// `pyhessref.vib.print_vibs`. +/// +/// Only vibrational modes (`TRV == "V"`) are printed. +/// +/// # Parameters +/// - `vib` : [`VibInfo`] from [`harmonic_analysis`]. +/// - `atom_lbl` : atomic symbols; if empty, integers are used. +/// - `normco` : which normal coordinate to print ([`NormCo`]). +/// - `shortlong` : `true` for `(nat, 3)` layout, `false` for `(3*nat, 1)`. +/// - `groupby` : modes per row (`Some(n)`); `None` ⇒ 3 (short) / 6 (long); `Some(usize::MAX)` is +/// treated as "all" (clamped to active count). +/// - `prec` : decimal places for scalar properties. +/// - `ncprec` : decimal places for normal coordinates (`None` ⇒ 2 short / 4 long). +pub fn print_vibs( + vib: &VibInfo, + atom_lbl: &[&str], + normco: NormCo, + shortlong: bool, + groupby: Option, + prec: usize, + ncprec: Option, +) -> String { + let nat = vib.ndof / 3; + let active: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] == "V").collect(); + + let presp = 2; + let prewidth = 24usize; + let colsp = 2usize; + let (groupby, ncprec, width) = if shortlong { + let g = groupby.unwrap_or(3); + let n = ncprec.unwrap_or(2); + (g, n, (n + 4) * 3) + } else { + let g = groupby.unwrap_or(6); + let n = ncprec.unwrap_or(4); + (g, n, n + 8) + }; + let groupby = if groupby == usize::MAX { active.len().max(1) } else { groupby }; + + let normco_t = normco.select(vib); + + let omega_str: Vec = (0..vib.ndof) + .map(|i| if vib.imag[i] { format!("{:.*}i", prec, vib.omega[i]) } else { format!("{:.*}", prec, vib.omega[i]) }) + .collect(); + + let mut lines: Vec = Vec::new(); + + // running 1-based vibrational mode number (renumbered starting from 1) + let mut vib_num = 1usize; + let mut iter = active.iter(); + loop { + let mut row: Vec = Vec::new(); + for _ in 0..groupby { + match iter.next() { + Some(&v) => row.push(v), + None => break, + } + } + if row.is_empty() { + break; + } + + // Vibration header — centered integer, colsp trailing + let mut line = format!("{}{}", " ".repeat(presp), ljust("Vibration", prewidth)); + for &_ in &row { + line.push_str(&format!("{}{}", center(&vib_num.to_string(), width), " ".repeat(colsp))); + vib_num += 1; + } + lines.push(line); + + // Freq — centered, 2 trailing spaces (= colsp) + let mut line = format!("{}{}", " ".repeat(presp), ljust("Freq [cm^-1]", prewidth)); + for &vib in &row { + line.push_str(&format!("{} ", center(&omega_str[vib], width))); + } + lines.push(line); + + // Irrep — skip (no irrep/symmetry in our implementation) + // scalar property rows — centered, colsp trailing + let labels: [&str; 5] = + ["Reduced mass [u]", "Force const [mDyne/A]", "Turning point v=0 [a0]", "RMS dev v=0 [a0 u^1/2]", "Char temp [K]"]; + let val_slices: [&[f64]; 5] = [&vib.mu, &vib.k, &vib.xtp0, &vib.dq0, &vib.theta_vib]; + for (label, vals) in labels.iter().zip(val_slices.iter()) { + let mut l = format!("{}{}", " ".repeat(presp), ljust(label, prewidth)); + for &vib in &row { + l.push_str(&format!( + "{}{}", + center(&format!("{:.*}", prec, vals[vib]), width), + " ".repeat(colsp) + )); + } + lines.push(l); + } + + // separator + let sep_len = prewidth + groupby * (width + colsp) - colsp; + lines.push(format!("{}{}", " ".repeat(presp), "-".repeat(sep_len))); + + // normal coordinate values + if shortlong { + let cell = width / 3; + for at in 0..nat { + let lbl = if at < atom_lbl.len() { atom_lbl[at] } else { "" }; + let mut l = format!("{}{:5} {}", " ".repeat(presp), at + 1, ljust(lbl, prewidth - 8)); + for &vib in &row { + let vx = normco_t[[3 * at, vib]]; + let vy = normco_t[[3 * at + 1, vib]]; + let vz = normco_t[[3 * at + 2, vib]]; + for v in [vx, vy, vz] { + l.push_str(¢er(&format!("{:.*}", ncprec, v), cell)); + } + l.push_str(&" ".repeat(colsp)); + } + lines.push(l); + } + } else { + for at in 0..nat { + let lbl = if at < atom_lbl.len() { atom_lbl[at] } else { "" }; + for xyz in 0..3 { + let axis = ['X', 'Y', 'Z'][xyz]; + let mut l = format!( + "{}{:5} {} {}", + " ".repeat(presp), + at + 1, + axis, + ljust(lbl, prewidth - 14) + ); + for &vib in &row { + let v = normco_t[[3 * at + xyz, vib]]; + l.push_str(&format!( + "{}{}", + center(&format!("{:.*}", ncprec, v), width), + " ".repeat(colsp) + )); + } + lines.push(l); + } + } + } + } + + lines.join("\n") +} + +// --------------------------------------------------------------------------- +// Thermochemistry pretty-printer (port of Psi4 thermo display) +// --------------------------------------------------------------------------- + +/// Pretty-print thermochemistry results. Rust port of Psi4's `thermo` display +/// section. Uses the same three-column layout (cal/(mol K) / J/(mol K) / mEh/K +/// for S/Cv/Cp; kcal/mol / kJ/mol / Eh for E/H/G/ZPE). +/// +/// Not required to be byte-identical to Psi4, but sufficiently similar. +pub fn print_thermo(th: &ThermoInfo, multiplicity: i64, molecular_mass: f64) -> String { + let t = th.t; + let p = th.p; + let sigma = th.sigma; + let e0 = th.e0; + + // helper: format S/Cv/Cp row + let fmt_scv = |label: &str, val_mEhK: f64| -> String { + format!( + " {:<36}{:11.3} [cal/(mol K)] {:11.3} [J/(mol K)] {:15.8} [mEh/K]", + label, + val_mEhK * HARTREE2KCALMOL, + val_mEhK * HARTREE2KJMOL, + val_mEhK, + ) + }; + // helper: format E/H/G/ZPE row + let fmt_ehg = |label: &str, val_Eh: f64| -> String { + format!( + " {:<36}{:11.3} [kcal/mol] {:11.3} [kJ/mol] {:15.8} [Eh]", + label, + val_Eh * HARTREE2KCALMOL, + val_Eh * HARTREE2KJMOL, + val_Eh, + ) + }; + + let component_names = ["Electronic", "Translational", "Rotational", "Vibrational"]; + + let mut lines: Vec = Vec::new(); + + // ---- header ---- + lines.push("\n ==> Thermochemistry Components <==".to_string()); + + // ---- rotational constants ---- + lines.push(format!("\n Rotational constants {:11.6} {:11.6} {:11.6} [cm^-1]", th.b[0], th.b[1], th.b[2],)); + lines.push(format!( + " {:11.4} {:11.4} {:11.4} [GHz]", + th.b_ghz[0], th.b_ghz[1], th.b_ghz[2], + )); + lines.push(format!(" Rotor type {}", th.rotor_type)); + + // ---- Entropy S ---- + lines.push("\n\n Entropy, S".to_string()); + for i in 0..4 { + let mut l = fmt_scv(&format!(" {} S", component_names[i]), th.s[i]); + match i { + ELEC => l.push_str(&format!(" (multiplicity = {})", multiplicity)), + TRANS => l.push_str(&format!(" (mol. weight = {:.4} [u], P = {:.2} [Pa])", molecular_mass, p)), + ROT => l.push_str(&format!(" (symmetry no. = {})", sigma)), + _ => {}, + } + lines.push(l); + } + lines.push(fmt_scv("Total S", th.s_tot)); + lines.push(fmt_scv("Correction S", th.s_tot)); + + // ---- Cv ---- + lines.push("\n\n Constant volume heat capacity, Cv".to_string()); + for i in 0..4 { + lines.push(fmt_scv(&format!(" {} Cv", component_names[i]), th.cv[i])); + } + lines.push(fmt_scv("Total Cv", th.cv_tot)); + lines.push(fmt_scv("Correction Cv", th.cv_tot)); + + // ---- Cp ---- + lines.push("\n\n Constant pressure heat capacity, Cp".to_string()); + for i in 0..4 { + lines.push(fmt_scv(&format!(" {} Cp", component_names[i]), th.cp[i])); + } + lines.push(fmt_scv("Total Cp", th.cp_tot)); + lines.push(fmt_scv("Correction Cp", th.cp_tot)); + + // ---- Energy Analysis ---- + lines.push("\n\n ==> Thermochemistry Energy Analysis <==".to_string()); + + // raw E_e + let raw_e_e_label = " Total E_e, Electronic energy at well bottom"; + lines.push("\n\n Raw electronic energy, E_e".to_string()); + lines.push(format!("{} {:>15.8} [Eh]", rjust(raw_e_e_label, 85), e0)); + + // ---- ZPVE ---- + lines.push("\n\n Zero-point vibrational energy, ZPVE = Sum_i omega_i / 2, E_0 = E_e + ZPVE".to_string()); + { + let mut l = fmt_ehg(" Vibrational ZPVE", th.zpe[VIB]); + l.push_str(&format!(" {:15.3} [cm^-1]", th.zpe[VIB] * HARTREE2WAVENUMBERS)); + lines.push(l); + } + { + let mut l = fmt_ehg(" Correction ZPVE to E_e", th.zpe_corr); + l.push_str(&format!(" {:15.3} [cm^-1]", th.zpe_corr * HARTREE2WAVENUMBERS)); + lines.push(l); + } + let zpe_tot_label = " Total E_0, Enthalpy at 0 [K]"; + lines.push(format!("{} {:>15.8} [Eh]", rjust(zpe_tot_label, 85), th.zpe_tot)); + lines.push(" *** Absolute enthalpy, not an enthalpy of formation ***".to_string()); + + // ---- Thermal energy E ---- + lines.push("\n\n Thermal (internal) energy, E (includes ZPVE and finite-temperature corrections)".to_string()); + for i in 0..4 { + let label = if i == ELEC { + format!(" {} contrib to E beyond E_e", component_names[i]) + } else { + format!(" {} contrib to E", component_names[i]) + }; + lines.push(fmt_ehg(&label, th.e[i])); + } + lines.push(fmt_ehg(" Correction E", th.e_corr)); + lines.push(format!( + " Total E, Thermal (internal) energy at {:7.2} [K]{} {:>15.8} [Eh]", + t, + " ".repeat(47), + th.e_tot + )); + + // ---- Enthalpy H ---- + lines.push("\n\n Enthalpy, H_trans = E_trans + k_B * T = E_trans + P * V".to_string()); + for i in 0..4 { + let label = if i == ELEC { + format!(" {} contrib to H beyond E_e", component_names[i]) + } else { + format!(" {} contrib to H", component_names[i]) + }; + lines.push(fmt_ehg(&label, th.h[i])); + } + lines.push(fmt_ehg(" Correction H", th.h_corr)); + lines.push(format!(" Total H, Enthalpy at {:7.2} [K]{} {:>15.8} [Eh]", t, " ".repeat(53), th.h_tot)); + lines.push(" *** Absolute enthalpy, not an enthalpy of formation ***".to_string()); + + // ---- Gibbs G ---- + lines.push("\n\n Gibbs free energy, G = H - T * S".to_string()); + for i in 0..4 { + let label = if i == ELEC { + format!(" {} contrib to G beyond E_e", component_names[i]) + } else { + format!(" {} contrib to G", component_names[i]) + }; + lines.push(fmt_ehg(&label, th.g[i])); + } + lines.push(fmt_ehg(" Correction G", th.g_corr)); + lines.push(format!(" Total G, Gibbs energy at {:7.2} [K]{} {:>15.8} [Eh]\n", t, " ".repeat(53), th.g_tot)); + + lines.join("\n") +} diff --git a/src/lib.rs b/src/lib.rs index 8d789d7587..3990383670 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -48,6 +48,8 @@ //! #![allow(unused)] #![allow(non_snake_case)] +#![allow(mixed_script_confusables)] +#![allow(confusable_idents)] extern crate rest_tensors as tensors; extern crate hdf5_metno as hdf5; extern crate chrono as time; @@ -81,6 +83,7 @@ pub mod ri_tddft; pub mod fileop; pub mod ri_cphf; pub mod lib_rint; +pub mod hessian_backup; //extern crate rest; diff --git a/src/ri_jk/prelude_dev.rs b/src/ri_jk/prelude_dev.rs index 00e14037d0..3820069bec 100644 --- a/src/ri_jk/prelude_dev.rs +++ b/src/ri_jk/prelude_dev.rs @@ -11,8 +11,8 @@ pub use rstsr_core::prelude_dev::uninitialized_vec; pub use rt::blas::{BlasFloat, LapackDriverAPI}; -pub type Tsr = Tensor; -pub type TsrView<'a, T> = TensorView<'a, T, DeviceBLAS, IxD>; +pub type Tsr = Tensor; +pub type TsrView<'a, T = f64> = TensorView<'a, T, DeviceBLAS, IxD>; // utilities and logic-related imports pub(super) use super::util; diff --git a/src/ri_jk/util.rs b/src/ri_jk/util.rs index 377467d5c6..4bcb893ac1 100644 --- a/src/ri_jk/util.rs +++ b/src/ri_jk/util.rs @@ -28,3 +28,47 @@ pub fn retrive_tp_dim(nao_tp: usize) -> Result { } Ok(nao) } + +/// Generate the density matrix for current SCF component. +/// +/// # Parameters +/// +/// - `mo_coeff` : shape `[nao, nmo]`. Molecular orbital coefficients. +/// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. +/// +/// # Returns +/// +/// - `dm0` : shape `[nao, nao]`. The density matrix for current SCF component. +pub fn get_dm0_restricted(mo_coeff: TsrView, mo_occ: TsrView) -> Tsr { + let [_nao, nmo] = mo_coeff.shape().to_vec().try_into().unwrap(); + assert_eq!(mo_occ.shape(), &[nmo], "mo_occ shape not correct."); + + 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); + &mocc * occ.i((None, ..)) % &mocc.t() +} + +/// Generate the orbital-energy weighted density matrix for current SCF component. +/// +/// # Parameters +/// +/// - `mo_coeff` : shape `[nao, nmo]`. Molecular orbital coefficients. +/// - `mo_occ` : shape `[nmo]`. Molecular orbital occupation numbers. +/// - `mo_energy` : shape `[nmo]`. Molecular orbital energies. +/// +/// # Returns +/// +/// - `dme0` : shape `[nao, nao]`. The orbital-energy weighted density matrix for current SCF +/// component. +pub fn get_dme0_restricted(mo_coeff: TsrView, mo_occ: TsrView, mo_energy: TsrView) -> Tsr { + let [_nao, nmo] = mo_coeff.shape().to_vec().try_into().unwrap(); + assert_eq!(mo_occ.shape(), &[nmo], "mo_occ shape not correct."); + assert_eq!(mo_energy.shape(), &[nmo], "mo_energy shape not correct."); + + 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 eocc = mo_energy.bool_select(-1, &occidx); + &mocc * (occ * eocc).i((None, ..)) % &mocc.t() +} diff --git a/src/utilities/rstsr_util.rs b/src/utilities/rstsr_util.rs index 87754194d0..46b41b8f8b 100644 --- a/src/utilities/rstsr_util.rs +++ b/src/utilities/rstsr_util.rs @@ -5,6 +5,15 @@ use rstsr::prelude::*; use rstsr_core::{prelude_dev::OpAssignAPI, storage::creation::DeviceCreationAnyAPI}; use tensors::{BasicMatrix, MatrixFull}; +pub mod prelude { + pub use rstsr::prelude::*; + + pub type Tsr = Tensor; + pub type TsrView<'a, T = f64> = TensorView<'a, T, DeviceBLAS, IxD>; + pub type TsrMut<'a, T = f64> = TensorMut<'a, T, DeviceBLAS, IxD>; + pub type TsrCow<'a, T = f64> = TensorCow<'a, T, DeviceBLAS, IxD>; +} + /* #region interchange between rstsr and rest_tensor */ // In REST, we always use DeviceBLAS as backend in most cases. -- Gitee From 5ab8586d53327fde312f29ef9577427a86895f1f Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 10:24:04 +0800 Subject: [PATCH 02/39] port: ri_jk (hessian part and related utilities) --- src/dft/numint_matmul/mod.rs | 8 +- src/dft/xceff/mod.rs | 3 +- src/hessian_backup/mod.rs | 3 +- src/ri_jk/hess_r.rs | 1956 ++++++++++++++++++++++++++++++++++ src/ri_jk/hess_u.rs | 314 ++++++ src/ri_jk/mod.rs | 4 + src/ri_jk/prelude_dev.rs | 4 + src/ri_jk/pure_decompose.rs | 245 +++++ src/utilities/rstsr_util.rs | 17 +- 9 files changed, 2536 insertions(+), 18 deletions(-) create mode 100644 src/ri_jk/hess_r.rs create mode 100644 src/ri_jk/hess_u.rs diff --git a/src/dft/numint_matmul/mod.rs b/src/dft/numint_matmul/mod.rs index 78b6ae1618..1359e557bc 100644 --- a/src/dft/numint_matmul/mod.rs +++ b/src/dft/numint_matmul/mod.rs @@ -15,12 +15,13 @@ pub mod prelude { use super::*; pub(super) use indexmap::IndexMap; + pub(super) use itertools::Itertools; pub(super) use libxc::prelude::*; + pub(super) use rayon::prelude::*; pub(super) use rest_libcint::prelude::*; + pub(super) use rstsr::prelude::*; pub(super) use std::collections::HashMap; pub(super) use std::sync::{Arc, Mutex}; - pub(super) use itertools::Itertools; - pub(super) use rayon::prelude::*; pub(super) use super::nimatmul::*; pub(super) use super::pure_eval_rho::*; @@ -28,8 +29,8 @@ pub mod prelude { pub(super) use crate::dft::xceff::prelude::*; pub(super) use crate::ni_check_shape; pub(super) use crate::ri_jk::util::get_dm0_restricted; - pub(super) use crate::utilities::rstsr_util::prelude::*; pub(super) use crate::utilities::buffer_pool::BufferPool; + pub(super) use crate::utilities::rstsr_util::*; pub(super) type TsrView<'a, T = f64> = TensorView<'a, T, DeviceBLAS>; pub(super) type Tsr = Tensor; @@ -77,7 +78,6 @@ impl NIIntoUsizeVec for &Vec { } } - #[macro_export] macro_rules! ni_check_shape { ($actual:expr, $expected:expr, $msg:expr) => {{ diff --git a/src/dft/xceff/mod.rs b/src/dft/xceff/mod.rs index 99aeed52ce..7385cfbf38 100644 --- a/src/dft/xceff/mod.rs +++ b/src/dft/xceff/mod.rs @@ -14,10 +14,11 @@ pub mod prelude { pub use libxc_wrap::{determine_den_type, determine_den_type_from_list, libxc_eval_eff}; pub(super) use crate::ni_check_shape; - pub(super) use crate::utilities::rstsr_util::prelude::*; + pub(super) use crate::utilities::rstsr_util::*; pub(super) use itertools::Itertools; pub(super) use libxc::prelude::*; pub(super) use rayon::prelude::*; + pub(super) use rstsr::prelude::*; pub(super) type TsrView<'a, T = f64> = TensorView<'a, T, DeviceBLAS>; pub(super) type Tsr = Tensor; diff --git a/src/hessian_backup/mod.rs b/src/hessian_backup/mod.rs index 7e8ec57818..5cff5b9653 100644 --- a/src/hessian_backup/mod.rs +++ b/src/hessian_backup/mod.rs @@ -36,10 +36,11 @@ pub mod prelude { pub use uscf::UHessSCF; pub(super) use crate::ri_jk::util::{get_dm0_restricted, get_dme0_restricted}; - pub(super) use crate::utilities::rstsr_util::prelude::*; + pub(super) use crate::utilities::rstsr_util::*; pub(super) use cint_handling::*; pub(super) use itertools::Itertools; pub(super) use rayon::prelude::*; pub(super) use rest_libcint::prelude::*; + pub(super) use rstsr::prelude::*; pub(super) use std::collections::HashMap; } diff --git a/src/ri_jk/hess_r.rs b/src/ri_jk/hess_r.rs new file mode 100644 index 0000000000..1e2f333b72 --- /dev/null +++ b/src/ri_jk/hess_r.rs @@ -0,0 +1,1956 @@ +//! Optimized RI-JK Hessian computation. +//! +//! Algorithm is somehow optimized, and of no reference in current state. +//! Author of first implementation: Andrew J. Zhu +//! +//! The optimization route is not following any article, or codebase that I (ajz34) know, and is not +//! following AI assistance. If there is coincidence, it is accidental. +//! +//! The correct value is compared to PySCF (written by Qiming Sun). Reference article: +//! > Alchemy: A Quantum Chemistry Dataset for Benchmarking AI Models +//! > Chen, et al. arXiv:1906.09427 +//! +//! PySCF code referenced the ORCA's implementation: +//! > Efficient implementation of the analytic second derivatives of Hartree-Fock and hybrid DFT +//! > energies: a detailed analysis of different approximations +//! > Bykov, et al. Mol Phys. 113, 1961 (2015). DOI: 10.1080/00268976.2015.1025114 + +use super::prelude_dev::*; +use crate::grad::rhf::pack_triu_tilde; +use crate::hessian_backup::cint_handling::*; +use crate::hessian_backup::prelude::*; +use crate::ri_jk::util::*; + +use crate::ri_jk::decompose::*; +use crate::ri_jk::pure_decompose::{get_j2c_decomp, solve_by_j2c, solve_by_j2c_mut}; + +/* #region skeleton derivative keys */ + +pub const KEYS_J20: [&str; 3] = ["de_J20_1", "de_J20_2", "de_J20_3"]; +pub const KEYS_K20: [&str; 4] = ["de_K20_1a", "de_K20_1b", "de_K20_2", "de_K20_3"]; +pub const KEYS_J11: [&str; 4] = ["de_J11_1", "de_J11_2", "de_J11_3", "de_J11_4"]; +pub const KEYS_K11: [&str; 4] = ["de_K11_1", "de_K11_2", "de_K11_3", "de_K11_4"]; +pub const KEYS_J02: [&str; 9] = + ["de_J02_1", "de_J02_2", "de_J02_3a", "de_J02_3b", "de_J02_4", "de_J02_5", "de_J02_6", "de_J02_7", "de_J02_8"]; +pub const KEYS_K02: [&str; 9] = + ["de_K02_1", "de_K02_2", "de_K02_3a", "de_K02_3b", "de_K02_4", "de_K02_5", "de_K02_6", "de_K02_7", "de_K02_8"]; + +pub const KEYS_J1AO: [&str; 5] = ["j1ao_aux0", "j1ao_aux1_1", "j1ao_aux1_2", "j1ao_aux1_3", "j1ao_aux1_4"]; +pub const KEYS_K1BRA: [&str; 8] = [ + "k1bra_aux0_1", + "k1bra_aux0_2", + "k1bra_aux0_3", + "k1bra_aux0_4", + "k1bra_aux1_1", + "k1bra_aux1_2", + "k1bra_aux1_3", + "k1bra_aux1_4", +]; + +/* #endregion */ + +/* #region response */ + +/// Separated J/K response-bra core, shared by RHF and UHF. +/// +/// # Shapes +/// +/// - `cderi`: `[nao_tp, naux]` +/// - `mo_coeff[s]`: `[nao, nmo_s]`, `mo_occ[s]`: `[nmo_s]`, `bra[s]`: `[nao, nocc_s, ...]` (the +/// trailing dimensions, collectively `nprop`, must agree across spins) +/// +/// # Returns +/// +/// A tuple `(j_ao, k_bras)`: +/// - `j_ao`: `Option` of shape `[nao, nao, nprop]` — the **spin-independent** Coulomb response +/// operator in AO basis, built from the total density response `sum_s bra_s @ mocc_s.T` (already +/// carrying the internal factor `2.0` from the symmetric cderi contraction; the consumer applies +/// `scale_j` and the per-spin right half-transform `... @ mocc_s`). `None` if `do_j` is false. +/// - `k_bras`: `Vec` (one entry per spin) of shape `[nao, nocc_s, nprop]` — the same-spin +/// exchange response in bra form (already carrying its internal sign/scale; the consumer applies +/// `scale_k`). Empty if `do_k` is false. +/// +/// # Convention notes +/// +/// - J sees the **total** density response, so a single AO operator is produced and shared across +/// spins; this is why UHF can reuse the RHF J path verbatim. +/// - K is strictly same-spin; each spin's bra form is produced independently. +/// - The internal factors (`2.0` on J, the two-term symmetrized sum on K) match the existing RHF +/// optimized response; the per-method `scale_j` / `scale_k` and the RHF `0.5` vs UHF `1.0` +/// exchange prefactor are applied by the consumer, not here. +#[allow(clippy::too_many_arguments)] +pub fn get_rijk_response_bra_separated( + cderi: TsrView, + mo_coeff: &[TsrView], + mo_occ: &[TsrView], + bra: &[TsrView], + do_j: bool, + do_k: bool, + nbatch_aux: usize, +) -> (Option, Vec) { + // notes on shape + // - cderi: [nao_tp, naux] + // - mo_coeff[s]: [nao, nmo_s] + // - mo_occ[s]: [nmo_s] + // - bra[s]: [nao, nocc_s, ...] (trailing dims collectively `nprop`, same across spins) + + let nset = mo_coeff.len(); + assert_eq!(mo_occ.len(), nset); + assert_eq!(bra.len(), nset); + assert!(nset >= 1); + + let nao = mo_coeff[0].shape()[0]; + let naux = cderi.shape()[1]; + let nao_tp = nao * (nao + 1) / 2; + assert_eq!(cderi.shape()[0], nao_tp); + let device = cderi.device().clone(); + + // per-spin occupied coefficients and reshaped bras + let mocc: Vec = (0..nset) + .map(|s| { + let occidx = mo_occ[s].view().greater(0).into_vec(); + mo_coeff[s].view().bool_select(-1, &occidx) + }) + .collect(); + let nocc: Vec = mocc.iter().map(|m| m.shape()[1]).collect(); + let bra_shape_orig: Vec> = bra.iter().map(|b| b.shape().to_vec()).collect(); + let bra: Vec = (0..nset).map(|s| bra[s].view().reshape((nao, nocc[s], -1)).into_contig(ColMajor)).collect(); + let nprop = bra[0].shape()[2]; + for s in 0..nset { + assert_eq!(bra[s].shape()[2], nprop, "bra trailing dim (nprop) must agree across spins"); + } + + let mut j_ao: Option = None; + let mut k_bras: Vec = Vec::new(); + + // --- J contribution (spin-independent, AO form, from total density response) --- // + + if do_j { + // dm1_total = sum_s (bra_s @ mocc_s.T), then symmetrize; pack with tilde; the symmetric + // cderi contraction carries the internal factor 2.0 (matches the RHF optimized response). + let mut dm1: Tsr = rt::zeros(([nao, nao, nprop], &device)); + for s in 0..nset { + dm1 += &bra[s] % &mocc[s].t(); + } + let dm1 = &dm1 + &dm1.swapaxes(0, 1); + let dm1_tp = pack_triu_tilde(dm1.view()); + let itm_j_aux = cderi.t() % &dm1_tp; + let resp_tp_j: Tsr = 2.0 * &cderi % itm_j_aux; + j_ao = Some(resp_tp_j.unpack_tri(Upper, FlagSymm::Sy)); + } + + // --- K contribution (same-spin, bra form, two symmetrized terms) --- // + + if do_k { + for s in 0..nset { + let mocc_s = &mocc[s]; + let bra_s = &bra[s]; + let mut resp_bra_k: Tsr = rt::zeros_like(bra_s); + for iaux_start in (0..naux).step_by(nbatch_aux) { + let iaux_end = (iaux_start + nbatch_aux).min(naux); + let slc = rt::slice!(iaux_start, iaux_end); + // note: the following `naux` is the batch size, shadowing the outer one for brevity + let naux = iaux_end - iaux_start; + + // - cderi: [nao, nao, naux] + // - cderi_bxo: [nao, naux, nocc] + // - cderi_oxo: [nocc, naux, nocc] + // - cderi_box: [nao, nocc, naux] + let cderi = cderi.i((.., slc)).unpack_tri(Upper, FlagSymm::Sy); + let cderi_bxo = (cderi.reshape([nao, nao * naux]).t() % mocc_s).into_shape([nao, naux, nocc[s]]); + let cderi_oxo = + (mocc_s.t() % cderi_bxo.reshape([nao, naux * nocc[s]])).into_shape([nocc[s], naux, nocc[s]]); + + for a in 0..nprop { + let bra_sa = bra_s.i((.., .., a)); + let mut respka = resp_bra_k.i_mut((.., .., a)); + // k contribution part 0: uPj, iPj -> ui + let cderi_bxo_1 = (cderi.reshape([nao, nao * naux]).t() % &bra_sa).into_shape([nao, naux, nocc[s]]); + respka -= + cderi_bxo_1.reshape([nao, naux * nocc[s]]) % cderi_oxo.reshape([nocc[s], naux * nocc[s]]).t(); + // k contribution part 1: uPj, iPj -> ui (i from mocc, j from bra) + let cderi_oxo_1 = + (mocc_s.t() % cderi_bxo_1.reshape([nao, naux * nocc[s]])).into_shape([nocc[s], naux, nocc[s]]); + respka -= + cderi_bxo.reshape([nao, naux * nocc[s]]) % cderi_oxo_1.reshape([nocc[s], naux * nocc[s]]).t(); + } + } + // restore original trailing shape for this spin's bra + let mut shape = bra_shape_orig[s].clone(); + shape[0] = nao; + k_bras.push(resp_bra_k.into_shape(shape)); + } + } + + (j_ao, k_bras) +} + +/* #endregion */ + +/* #region skeleton */ + +macro_rules! tic { + ($timing:expr, $t0:expr, $msg:expr) => { + let t1 = std::time::Instant::now(); + let dt = t1.duration_since($t0).as_secs_f64(); + $timing.push(($msg.to_string(), dt)); + }; +} + +#[allow(clippy::too_many_arguments)] +#[allow(clippy::type_complexity)] +pub fn get_rijk_skeleton_decomposed_separated( + mol: &CInt, + aux: &CInt, + mo_coeff: &[TsrView], + mo_occ: &[TsrView], + cderi: TsrView, + j2c_decomp: &J2CDecompose, + do_j: bool, + do_k: bool, + nbatch_aux: usize, + atm_list: Option<&[usize]>, + dm0: Option, +) -> (Option>, Vec>, Vec<(String, f64)>) { + let device = cderi.device().clone(); + let mut timing = vec![]; + let time_full = std::time::Instant::now(); + + // --- basic checks --- // + if mo_coeff.is_empty() || mo_occ.is_empty() { + panic!("mo_coeff and mo_occ must be non-empty"); + } + if mo_coeff.len() != mo_occ.len() { + panic!("mo_coeff and mo_occ must have the same length"); + } + let nset = mo_coeff.len(); + + // --- prepare shared --- // + let t0 = std::time::Instant::now(); + let (dims, aoslices, auxslices, aux_ranges, shared, solve_aux) = + prepare_shared(mol, aux, j2c_decomp, nbatch_aux, atm_list, &device); + tic!(timing, t0, "prepare_shared"); + + // --- prepare j --- // + let j_in = do_j.then(|| { + let t0 = std::time::Instant::now(); + // - dm0: [nao, nao]; note this density is total density instead of spin-separated density + let dm0 = dm0.map_or_else( + || { + let nao = dims["nao"]; + let mut dm0 = rt::zeros(([nao, nao], &device)); + for iset in 0..nset { + dm0 += get_dm0_restricted(mo_coeff[iset].view(), mo_occ[iset].view()) + } + dm0 + }, + |dm0| dm0.to_owned(), + ); + let j_in = prepare_j(&solve_aux, &dims, dm0.view(), cderi.view()); + tic!(timing, t0, "prepare_j"); + j_in + }); + let mut j_out = do_j.then(HashMap::new); + let mut j_intmd = do_j.then(HashMap::new); + + // --- prepare k --- // + + let mut k_ins = vec![]; + if do_k { + for iset in 0..nset { + let t0 = std::time::Instant::now(); + let k_in = prepare_k(&solve_aux, &dims, mo_coeff[iset].view(), mo_occ[iset].view(), cderi.view()); + k_ins.push(k_in); + tic!(timing, t0, &format!("prepare_k {iset}")); + } + } + let mut k_outs = (0..k_ins.len()).map(|_| HashMap::new()).collect_vec(); + let mut k_intmds = (0..k_ins.len()).map(|_| HashMap::new()).collect_vec(); + + // --- evaluate oneshot --- // + + let t0 = std::time::Instant::now(); + let timing_oneshot = evaluate_oneshot( + &dims, + mol, + aux, + &aoslices, + &auxslices, + &aux_ranges, + &device, + j_in.as_ref(), + &k_ins, + j_out.as_mut(), + &mut k_outs, + ); + timing.extend(timing_oneshot); + tic!(timing, t0, "evaluate_oneshot"); + + // --- evaluate j2c-derivatives-only terms --- // + + let t0 = std::time::Instant::now(); + let timing_j2c_deriv_only = evaluate_j2c_deriv_only( + &dims, + &shared, + aux, + &auxslices, + &device, + j_in.as_ref(), + &k_ins, + j_out.as_mut(), + &mut k_outs, + ); + timing.extend(timing_j2c_deriv_only); + tic!(timing, t0, "evaluate_j2c_deriv_only"); + + // --- evaluate jk1 j2c-skeleton terms --- // + + let t0 = std::time::Instant::now(); + let timing_jk1_j2c_deriv = evaluate_jk1_j2c_deriv( + &dims, + &shared, + cderi.view(), + &auxslices, + &device, + &solve_aux, + j_in.as_ref(), + &k_ins, + j_out.as_mut(), + &mut k_outs, + ); + timing.extend(timing_jk1_j2c_deriv); + tic!(timing, t0, "evaluate_jk1_j2c_deriv"); + + // --- evaluate j3c-ip2 related terms --- // + + let t0 = std::time::Instant::now(); + let timing_j3c_ip2 = evaluate_j3c_ip2( + &dims, + &shared, + cderi.view(), + mol, + aux, + &auxslices, + &aux_ranges, + &device, + &solve_aux, + j_in.as_ref(), + &mut k_ins, + j_out.as_mut(), + &mut k_outs, + j_intmd.as_mut(), + &mut k_intmds, + ); + timing.extend(timing_j3c_ip2); + tic!(timing, t0, "evaluate_j3c_ip2"); + + // --- evaluate j3c-ip1 related terms --- // + + let t0 = std::time::Instant::now(); + let timing_j3c_ip1 = evaluate_j3c_ip1( + &dims, + &shared, + cderi.view(), + mol, + aux, + &aoslices, + &auxslices, + &aux_ranges, + &device, + &solve_aux, + j_in.as_ref(), + &mut k_ins, + j_out.as_mut(), + &mut k_outs, + j_intmd.as_mut(), + &mut k_intmds, + ); + timing.extend(timing_j3c_ip1); + tic!(timing, t0, "evaluate_j3c_ip1"); + + tic!(timing, time_full, "get_rijk_skeleton_decomposed_separated"); + (j_out, k_outs, timing) +} + +pub type FnSolveAux<'a> = Box; +type PrepareSharedOutput<'a> = ( + HashMap<&'static str, usize>, // dims + Vec<[usize; 4]>, // aoslices + Vec<[usize; 4]>, // auxslices + Vec<[usize; 4]>, // aux_ranges + HashMap<&'static str, Tsr>, // shared (intermediates) + FnSolveAux<'a>, // solve_aux(tsr_mut, left/right, do_flip) +); + +pub fn prepare_shared<'a>( + mol: &CInt, + aux: &CInt, + j2c_decomp: &'a J2CDecompose, + nbatch_aux: usize, + atm_list: Option<&[usize]>, + device: &DeviceBLAS, +) -> PrepareSharedOutput<'a> { + // aoslices, auxslices + let atm_list = atm_list.map_or_else(|| (0..mol.natm()).collect_vec(), |list| list.to_vec()); + let natm = atm_list.len(); + let aoslices = mol.aoslice_by_atom(); + let auxslices = aux.aoslice_by_atom(); + let aoslices = atm_list.iter().map(|&i| aoslices[i]).collect_vec(); + let auxslices = atm_list.iter().map(|&i| auxslices[i]).collect_vec(); + + // aux_ranges + let aux_balance = aux.balance_partition(nbatch_aux); + let mut p0 = 0; + let aux_ranges = aux_balance + .into_iter() + .map(|[sh0, sh1, size]| { + let range = [sh0, sh1, p0, p0 + size]; + p0 += size; + range + }) + .collect_vec(); + + let solve_aux = + |tsr_mut: TsrMut, side: FlagSide, do_flip: bool| solve_by_j2c_mut(tsr_mut, j2c_decomp, side, do_flip); + + let j2c_ip1 = hess_intor(aux, "int2c2e_ip1", "s1", None, device); + let mut rcd_j2c_ip1 = j2c_ip1.to_owned(); + rcd_j2c_ip1.axes_iter_mut(-1).for_each(|m| solve_aux(m, Right, false)); + let mut rrcd_j2c_ip1 = rcd_j2c_ip1.to_owned(); + rrcd_j2c_ip1.axes_iter_mut(-1).for_each(|m| solve_aux(m, Right, true)); + + let naux = aux.nao(); + let j2c_inv = { + let mut eye = rt::eye((naux, device)); + solve_aux(eye.view_mut(), Right, true); + let mut out = eye.t().into_contig(ColMajor); + solve_aux(out.view_mut(), Right, true); + out + }; + + let dims = HashMap::from([("natm", natm), ("nao", mol.nao()), ("naux", aux.nao())]); + let shared = HashMap::from([ + ("j2c_ip1", j2c_ip1), + ("rcd_j2c_ip1", rcd_j2c_ip1), + ("rrcd_j2c_ip1", rrcd_j2c_ip1), + ("j2c_inv", j2c_inv), + ]); + + (dims, aoslices, auxslices, aux_ranges, shared, Box::new(solve_aux)) +} + +pub fn prepare_j( + solve_aux: &FnSolveAux, + dims: &HashMap<&'static str, usize>, + dm0: TsrView, + cderi: TsrView, +) -> HashMap<&'static str, Tsr> { + let nao = dims["nao"]; + assert_eq!(dm0.shape(), &[nao, nao], "dm0 in prepare_j must be a single square matrix of shape [nao, nao]"); + + // pack density matrix with scaled (off-diag, diag) + let dm0_tp = pack_triu_tilde(dm0.view()); + let rrcd_eri_aux = { + let mut out = dm0_tp.view() % cderi; + solve_aux(out.view_mut(), Right, true); + out + }; + HashMap::from([("dm0", dm0.to_owned()), ("dm0_tp", dm0_tp), ("rrcd_eri_aux", rrcd_eri_aux)]) +} + +pub fn prepare_k( + solve_aux: &FnSolveAux, + dims: &HashMap<&'static str, usize>, + mo_coeff: TsrView, + mo_occ: TsrView, + cderi: TsrView, +) -> HashMap<&'static str, Tsr> { + let nao = dims["nao"]; + let naux = dims["naux"]; + let nmo = mo_coeff.shape()[1]; + let device = cderi.device().clone(); + assert_eq!(mo_coeff.shape(), &[nao, nmo], "mo_coeff in prepare_k must be of shape [nao, nmo]"); + assert_eq!(mo_occ.shape(), &[nmo], "mo_occ in prepare_k must be of shape [nmo]"); + + let occidx = mo_occ.view().greater(0).into_vec(); + let nocc = occidx.iter().filter(|&&x| x).count(); + let mocc = mo_coeff.bool_select(-1, &occidx); + let occ = mo_occ.bool_select(-1, &occidx); + let mocc_2 = &mocc * occ.view().sqrt().i((None, ..)); + let occ_invsqrt = occ.view().pow(-0.5); + + // orbital transformation + let rcd_eri_bra: Tsr = rt::zeros(([nao, nocc, naux], &device)); + let rcd_eri_occ: Tsr = rt::zeros(([nocc, nocc, naux], &device)); + (0..naux).into_par_iter().for_each(|p| { + let rcd_eri_bra_p = rcd_eri_bra.i((.., .., p)); + let rcd_eri_occ_p = rcd_eri_occ.i((.., .., p)); + let mut rcd_eri_bra_p = unsafe { rcd_eri_bra_p.force_mut() }; + let mut rcd_eri_occ_p = unsafe { rcd_eri_occ_p.force_mut() }; + let cderi_p = cderi.i((.., p)).unpack_tri(Upper, FlagSymm::Sy); + rcd_eri_bra_p.matmul_from(&cderi_p, &mocc_2, 1.0, 0.0); + rcd_eri_occ_p.matmul_from(&mocc_2.t(), &rcd_eri_bra_p, 1.0, 0.0); + }); + + // j2c solve (consumes previous results in-place) + let mut rrcd_eri_bra = rcd_eri_bra; + solve_aux(rrcd_eri_bra.view_mut(), Right, true); + let mut rrcd_eri_occ = rcd_eri_occ; + solve_aux(rrcd_eri_occ.view_mut(), Right, true); + + let fold_eri_bra = &mocc_2 % &rrcd_eri_occ; + + HashMap::from([ + ("mocc", mocc), + ("mocc_2", mocc_2), + ("occ_invsqrt", occ_invsqrt), + ("rrcd_eri_bra", rrcd_eri_bra), + ("rrcd_eri_occ", rrcd_eri_occ), + ("fold_eri_bra", fold_eri_bra), + ]) +} + +#[allow(clippy::too_many_arguments)] +pub fn evaluate_oneshot( + dims: &HashMap<&'static str, usize>, + mol: &CInt, + aux: &CInt, + aoslices: &[[usize; 4]], + auxslices: &[[usize; 4]], + aux_ranges: &[[usize; 4]], + device: &DeviceBLAS, + j_in: Option<&HashMap<&'static str, Tsr>>, + k_ins: &[HashMap<&'static str, Tsr>], + j_out: Option<&mut HashMap<&'static str, Tsr>>, + k_outs: &mut [HashMap<&'static str, Tsr>], +) -> Vec<(String, f64)> { + let mut timing = vec![]; + + let nao = dims["nao"]; + let naux = dims["naux"]; + let natm = dims["natm"]; + let do_j = j_in.is_some(); + let nset_k = k_outs.len(); + + // --- integral generators --- // + + let gen_j3c_ipvip1 = generator_hess_intor_j3c_by_aux(mol, aux, "int3c2e_ipvip1", "s1", device); + let gen_j3c_ipip1 = generator_hess_intor_j3c_by_aux(mol, aux, "int3c2e_ipip1", "s1", device); + let gen_j3c_ip1ip2 = generator_hess_intor_j3c_by_aux(mol, aux, "int3c2e_ip1ip2", "s1", device); + let gen_j3c_ipip2 = generator_hess_intor_j3c_by_aux(mol, aux, "int3c2e_ipip2", "s1", device); + + // --- dbas allocations --- // + + let mut dbas_j: HashMap<&str, Tsr> = HashMap::new(); + let mut dbas_ks: Vec> = (0..nset_k).map(|_| HashMap::new()).collect_vec(); + if do_j { + dbas_j.insert("J20_2", rt::zeros(([nao, nao, 3, 3], device))); + dbas_j.insert("J20_3", rt::zeros(([nao, nao, 3, 3], device))); + dbas_j.insert("J11_1", rt::zeros(([nao, naux, 3, 3], device))); + dbas_j.insert("J02_1", rt::zeros(([naux, 3, 3], device))); + } + for i in 0..nset_k { + let dbas_k = &mut dbas_ks[i]; + dbas_k.insert("K20_2", rt::zeros(([nao, nao, 3, 3], device))); + dbas_k.insert("K20_3", rt::zeros(([nao, nao, 3, 3], device))); + dbas_k.insert("K11_1", rt::zeros(([nao, naux, 3, 3], device))); + dbas_k.insert("K02_1", rt::zeros(([naux, 3, 3], device))); + } + + // --- dbas evaluation --- // + + for &[sh0, sh1, p0, p1] in aux_ranges { + // --- common tensors --- // + + let t0 = std::time::Instant::now(); + + // j-part + let rrcd_eri_aux = j_in.map(|j_in| j_in["rrcd_eri_aux"].i(p0..p1)); + + // k-part + let mut tmps_k_ao = vec![]; + for iset in 0..nset_k { + let k_in = &k_ins[iset]; + let mocc_2 = &k_in["mocc_2"]; + let rrcd_eri_occ = k_in["rrcd_eri_occ"].i((.., .., p0..p1)); + tmps_k_ao.push(mocc_2 % rrcd_eri_occ % mocc_2.t()); + } + + tic!(timing, t0, &format!("evaluate_oneshot, common aux({p0}:{p1})")); + + // --- 20-2 (ipvip1) --- // + + let t0 = std::time::Instant::now(); + let j3c_ipvip1 = gen_j3c_ipvip1([sh0, sh1]); + tic!(timing, t0, &format!("evaluate_oneshot, gen_j3c_ipvip1 aux({p0}:{p1})")); + + let t0 = std::time::Instant::now(); + if do_j { + let rrcd_eri_aux = rrcd_eri_aux.as_ref().unwrap().view(); + let dm0 = j_in.as_ref().unwrap()["dm0"].view(); + let tmp1 = rt::vecdot(&j3c_ipvip1, rrcd_eri_aux.i((None, None, ..)), 2); + *dbas_j.get_mut("J20_2").unwrap() += tmp1 * dm0; + } + for iset in 0..nset_k { + let tmp_k_ao = &tmps_k_ao[iset]; + let tmp1 = rt::vecdot(&j3c_ipvip1, tmp_k_ao, 2); + *dbas_ks[iset].get_mut("K20_2").unwrap() += tmp1; + } + tic!(timing, t0, &format!("evaluate_oneshot, dbas 20-2 aux({p0}:{p1})")); + drop(j3c_ipvip1); + + // --- 20-3 (ipip1) --- // + + let t0 = std::time::Instant::now(); + let j3c_ipip1 = gen_j3c_ipip1([sh0, sh1]); + tic!(timing, t0, &format!("evaluate_oneshot, gen_j3c_ipip1 aux({p0}:{p1})")); + + let t0 = std::time::Instant::now(); + if do_j { + let rrcd_eri_aux = rrcd_eri_aux.as_ref().unwrap().view(); + let dm0 = j_in.as_ref().unwrap()["dm0"].view(); + let tmp1 = rt::vecdot(&j3c_ipip1, rrcd_eri_aux.i((None, None, ..)), 2); + *dbas_j.get_mut("J20_3").unwrap() += tmp1 * dm0; + } + for iset in 0..nset_k { + let tmp_k_ao = &tmps_k_ao[iset]; + let tmp1 = rt::vecdot(&j3c_ipip1, tmp_k_ao, 2); + *dbas_ks[iset].get_mut("K20_3").unwrap() += tmp1; + } + tic!(timing, t0, &format!("evaluate_oneshot, dbas 20-3 aux({p0}:{p1})")); + drop(j3c_ipip1); + + // --- 11-1 (ip1ip2) --- // + + let t0 = std::time::Instant::now(); + let j3c_ip1ip2 = gen_j3c_ip1ip2([sh0, sh1]); + tic!(timing, t0, &format!("evaluate_oneshot, gen_j3c_ip1ip2 aux({p0}:{p1})")); + + let t0 = std::time::Instant::now(); + if do_j { + let rrcd_eri_aux = rrcd_eri_aux.as_ref().unwrap().view(); + let dm0 = j_in.as_ref().unwrap()["dm0"].view(); + let tmp1 = rt::vecdot(&j3c_ip1ip2, &dm0, 1) * rrcd_eri_aux.i((None, ..)); + dbas_j.get_mut("J11_1").unwrap().i_mut((.., p0..p1)).assign(tmp1); + } + for iset in 0..nset_k { + let tmp_k_ao = &tmps_k_ao[iset]; + let tmp1 = rt::vecdot(&j3c_ip1ip2, tmp_k_ao, 1); + dbas_ks[iset].get_mut("K11_1").unwrap().i_mut((.., p0..p1)).assign(tmp1); + } + tic!(timing, t0, &format!("evaluate_oneshot, dbas 11-1 aux({p0}:{p1})")); + drop(j3c_ip1ip2); + + // --- 02-1 (ipip2) --- // + + let t0 = std::time::Instant::now(); + let j3c_ipip2 = gen_j3c_ipip2([sh0, sh1]); + tic!(timing, t0, &format!("evaluate_oneshot, gen_j3c_ipip2 aux({p0}:{p1})")); + + let t0 = std::time::Instant::now(); + if do_j { + let rrcd_eri_aux = rrcd_eri_aux.as_ref().unwrap().view(); + let dm0 = j_in.as_ref().unwrap()["dm0"].view(); + let tmp1 = rt::vecdot(&j3c_ipip2, &dm0, ([0, 1], [0, 1])) * rrcd_eri_aux; + dbas_j.get_mut("J02_1").unwrap().i_mut(p0..p1).assign(tmp1); + } + for iset in 0..nset_k { + let tmp_k_ao = &tmps_k_ao[iset]; + let tmp1 = rt::vecdot(&j3c_ipip2, tmp_k_ao, ([0, 1], [0, 1])); + dbas_ks[iset].get_mut("K02_1").unwrap().i_mut(p0..p1).assign(tmp1); + } + tic!(timing, t0, &format!("evaluate_oneshot, dbas 02-1 aux({p0}:{p1})")); + drop(j3c_ipip2); + } + + // --- reduce to hessian contribution --- // + + let t0 = std::time::Instant::now(); + + if do_j { + let j_out = j_out.unwrap(); + + let mut de_J20_2: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J20_3: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J11_1: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J02_1: Tsr = rt::zeros(([3, 3, natm, natm], device)); + + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + + let tmp = dbas_j["J20_3"].i(slcA).sum_axes([0, 1]); + de_J20_3.i_mut((.., .., A, A)).assign(tmp); + + for (B, &[_, _, p0B, p1B]) in aoslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + let tmp = dbas_j["J20_2"].i((slcA, slcB)).sum_axes([0, 1]); + de_J20_2.i_mut((.., .., B, A)).assign(tmp); + } + + for (B, &[_, _, p0B, p1B]) in auxslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + let tmp = dbas_j["J11_1"].i((slcA, slcB)).sum_axes([0, 1]); + de_J11_1.i_mut((.., .., B, A)).assign(tmp); + } + } + + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + let tmp = dbas_j["J02_1"].i(slcA).sum_axes(0); + de_J02_1.i_mut((.., .., A, A)).assign(tmp); + } + + let scale_J20_2 = 1.0; + let scale_J20_3 = 1.0; + let scale_J11_1 = 2.0; + let scale_J02_1 = 0.5; + j_out.insert("de_J20_2", scale_J20_2 * (&de_J20_2 + &de_J20_2.transpose([1, 0, 3, 2]))); + j_out.insert("de_J20_3", scale_J20_3 * (&de_J20_3 + &de_J20_3.transpose([1, 0, 3, 2]))); + j_out.insert("de_J11_1", scale_J11_1 * (&de_J11_1 + &de_J11_1.transpose([1, 0, 3, 2]))); + j_out.insert("de_J02_1", scale_J02_1 * (&de_J02_1 + &de_J02_1.transpose([1, 0, 3, 2]))); + } + + for iset in 0..nset_k { + let k_out = &mut k_outs[iset]; + + let mut de_K20_2: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K20_3: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K11_1: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K02_1: Tsr = rt::zeros(([3, 3, natm, natm], device)); + + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + + let tmp = dbas_ks[iset]["K20_3"].i(slcA).sum_axes([0, 1]); + de_K20_3.i_mut((.., .., A, A)).assign(tmp); + + for (B, &[_, _, p0B, p1B]) in aoslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + let tmp = dbas_ks[iset]["K20_2"].i((slcA, slcB)).sum_axes([0, 1]); + de_K20_2.i_mut((.., .., B, A)).assign(tmp); + } + + for (B, &[_, _, p0B, p1B]) in auxslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + let tmp = dbas_ks[iset]["K11_1"].i((slcA, slcB)).sum_axes([0, 1]); + de_K11_1.i_mut((.., .., B, A)).assign(tmp); + } + } + + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + let tmp = dbas_ks[iset]["K02_1"].i(slcA).sum_axes(0); + de_K02_1.i_mut((.., .., A, A)).assign(tmp); + } + + let scale_K20_2 = 1.0; + let scale_K20_3 = 1.0; + let scale_K11_1 = 2.0; + let scale_K02_1 = 0.5; + k_out.insert("de_K20_2", scale_K20_2 * (&de_K20_2 + &de_K20_2.transpose([1, 0, 3, 2]))); + k_out.insert("de_K20_3", scale_K20_3 * (&de_K20_3 + &de_K20_3.transpose([1, 0, 3, 2]))); + k_out.insert("de_K11_1", scale_K11_1 * (&de_K11_1 + &de_K11_1.transpose([1, 0, 3, 2]))); + k_out.insert("de_K02_1", scale_K02_1 * (&de_K02_1 + &de_K02_1.transpose([1, 0, 3, 2]))); + } + + tic!(timing, t0, "evaluate_oneshot, reduce to hessian"); + + timing +} + +#[allow(clippy::too_many_arguments)] +pub fn evaluate_j2c_deriv_only( + dims: &HashMap<&'static str, usize>, + shared: &HashMap<&'static str, Tsr>, + aux: &CInt, + auxslices: &[[usize; 4]], + device: &DeviceBLAS, + j_in: Option<&HashMap<&'static str, Tsr>>, + k_ins: &[HashMap<&'static str, Tsr>], + j_out: Option<&mut HashMap<&'static str, Tsr>>, + k_outs: &mut [HashMap<&'static str, Tsr>], +) -> Vec<(String, f64)> { + let mut timing = vec![]; + + let naux = dims["naux"]; + let natm = dims["natm"]; + let do_j = j_in.is_some(); + let nset_k = k_outs.len(); + + let j2c_ip1 = shared["j2c_ip1"].view(); + let rcd_j2c_ip1 = shared["rcd_j2c_ip1"].view(); + let rrcd_j2c_ip1 = shared["rrcd_j2c_ip1"].view(); + let j2c_inv = shared["j2c_inv"].view(); + + // --- integral evaluation --- // + + let t0 = std::time::Instant::now(); + let j2c_ipip1 = hess_intor(aux, "int2c2e_ipip1", "s1", None, device); + let j2c_ip1ip2 = hess_intor(aux, "int2c2e_ip1ip2", "s1", None, device); + tic!(timing, t0, "evaluate_j2c_deriv_only, integration"); + + // --- evaluation j-part --- // + + if do_j { + let t0 = std::time::Instant::now(); + + let j_out = j_out.unwrap(); + let rrcd_eri_aux = j_in.as_ref().unwrap()["rrcd_eri_aux"].view(); + + // --- dbas evaluation --- // + + let dbas_J02_2 = rt::vecdot(&j2c_ipip1, &rrcd_eri_aux, 0) * &rrcd_eri_aux; + let dbas_J02_3a = &j2c_ip1ip2 * &rrcd_eri_aux * rrcd_eri_aux.i((None, ..)); + let tmp1 = &rcd_j2c_ip1 * &rrcd_eri_aux; + let dbas_J02_3b = tmp1.i((.., .., None, ..)) % tmp1.i((.., .., .., None)).swapaxes(0, 1); + let tmp1 = rt::vecdot(&j2c_ip1, &rrcd_eri_aux, 0); + let dbas_J02_6 = tmp1.i((.., None, None, ..)) * &j2c_inv * tmp1.i((None, .., .., None)); + let tmp2 = &rrcd_j2c_ip1 * &rrcd_eri_aux; + let dbas_J02_8 = tmp1.i((None, .., .., None)) * tmp2.i((.., .., None, ..)); + + // --- reduce to hessian contribution --- // + + let mut de_J02_2: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J02_3a: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J02_3b: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J02_6: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J02_8: Tsr = rt::zeros(([3, 3, natm, natm], device)); + + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + + let tmp = dbas_J02_2.i(slcA).sum_axes(0); + de_J02_2.i_mut((.., .., A, A)).assign(tmp); + + for (B, &[_, _, p0B, p1B]) in auxslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + + let tmp = dbas_J02_3a.i((slcA, slcB)).sum_axes([0, 1]); + de_J02_3a.i_mut((.., .., B, A)).assign(tmp); + let tmp = dbas_J02_3b.i((slcA, slcB)).sum_axes([0, 1]); + de_J02_3b.i_mut((.., .., B, A)).assign(tmp); + let tmp = dbas_J02_6.i((slcA, slcB)).sum_axes([0, 1]); + de_J02_6.i_mut((.., .., B, A)).assign(tmp); + let tmp = dbas_J02_8.i((slcA, slcB)).sum_axes([0, 1]); + de_J02_8.i_mut((.., .., B, A)).assign(tmp); + } + } + + let scale_J02_2 = -0.5; + let scale_J02_3a = -0.5; + let scale_J02_3b = 0.5; + let scale_J02_6 = 0.5; + let scale_J02_8 = -1; + j_out.insert("de_J02_2", scale_J02_2 * (&de_J02_2 + &de_J02_2.transpose([1, 0, 3, 2]))); + j_out.insert("de_J02_3a", scale_J02_3a * (&de_J02_3a + &de_J02_3a.transpose([1, 0, 3, 2]))); + j_out.insert("de_J02_3b", scale_J02_3b * (&de_J02_3b + &de_J02_3b.transpose([1, 0, 3, 2]))); + j_out.insert("de_J02_6", scale_J02_6 * (&de_J02_6 + &de_J02_6.transpose([1, 0, 3, 2]))); + j_out.insert("de_J02_8", scale_J02_8 * (&de_J02_8 + &de_J02_8.transpose([1, 0, 3, 2]))); + + tic!(timing, t0, "evaluate_j2c_deriv_only, evaluate j-part"); + } + + // --- evaluation (K) --- // + + for iset in 0..nset_k { + let t0 = std::time::Instant::now(); + + let k_in = &k_ins[iset]; + let k_out = &mut k_outs[iset]; + let rrcd_eri_occ = k_in["rrcd_eri_occ"].view(); + let fold_eri_aux = rrcd_eri_occ.reshape((-1, naux)).t() % rrcd_eri_occ.reshape((-1, naux)); + + // --- dbas evaluation --- // + + let dbas_K02_2 = rt::vecdot(&j2c_ipip1, &fold_eri_aux, 0); + let dbas_K02_3a = &j2c_ip1ip2 * &fold_eri_aux; + let dbas_K02_3b = + rcd_j2c_ip1.i((.., .., None, ..)) % rcd_j2c_ip1.i((.., .., .., None)).swapaxes(0, 1) * &fold_eri_aux; + let dbas_K02_6 = j2c_ip1.i((.., .., None, ..)) % &fold_eri_aux % j2c_ip1.i((.., .., .., None)) * &j2c_inv; + let dbas_K02_8 = fold_eri_aux % j2c_ip1.i((.., .., .., None)) * rrcd_j2c_ip1.i((.., .., None, ..)); + + // --- reduce to hessian contribution --- // + + let mut de_K02_2: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K02_3a: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K02_3b: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K02_6: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K02_8: Tsr = rt::zeros(([3, 3, natm, natm], device)); + + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + + let tmp = dbas_K02_2.i(slcA).sum_axes(0); + de_K02_2.i_mut((.., .., A, A)).assign(tmp); + + for (B, &[_, _, p0B, p1B]) in auxslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + + let tmp = dbas_K02_3a.i((slcA, slcB)).sum_axes([0, 1]); + de_K02_3a.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_K02_3b.i((slcA, slcB)).sum_axes([0, 1]); + de_K02_3b.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_K02_6.i((slcA, slcB)).sum_axes([0, 1]); + de_K02_6.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_K02_8.i((slcA, slcB)).sum_axes([0, 1]); + de_K02_8.i_mut((.., .., B, A)).assign(tmp); + } + } + + let scale_K02_2 = -0.5; + let scale_K02_3a = -0.5; + let scale_K02_3b = 0.5; + let scale_K02_6 = -0.5; + let scale_K02_8 = -1.0; + k_out.insert("de_K02_2", scale_K02_2 * (&de_K02_2 + &de_K02_2.transpose([1, 0, 3, 2]))); + k_out.insert("de_K02_3a", scale_K02_3a * (&de_K02_3a + &de_K02_3a.transpose([1, 0, 3, 2]))); + k_out.insert("de_K02_3b", scale_K02_3b * (&de_K02_3b + &de_K02_3b.transpose([1, 0, 3, 2]))); + k_out.insert("de_K02_6", scale_K02_6 * (&de_K02_6 + &de_K02_6.transpose([1, 0, 3, 2]))); + k_out.insert("de_K02_8", scale_K02_8 * (&de_K02_8 + &de_K02_8.transpose([1, 0, 3, 2]))); + + tic!(timing, t0, &format!("evaluate_j2c_deriv_only, evaluate k-part {iset}")); + } + + timing +} + +#[allow(clippy::too_many_arguments)] +pub fn evaluate_jk1_j2c_deriv( + dims: &HashMap<&'static str, usize>, + shared: &HashMap<&'static str, Tsr>, + cderi: TsrView, + auxslices: &[[usize; 4]], + device: &DeviceBLAS, + solve_aux: &FnSolveAux, + j_in: Option<&HashMap<&'static str, Tsr>>, + k_ins: &[HashMap<&'static str, Tsr>], + j_out: Option<&mut HashMap<&'static str, Tsr>>, + k_outs: &mut [HashMap<&'static str, Tsr>], +) -> Vec<(String, f64)> { + let mut timing = vec![]; + + let nao = dims["nao"]; + let naux = dims["naux"]; + let natm = dims["natm"]; + let do_j = j_in.is_some(); + let nset_k = k_outs.len(); + let j2c_ip1 = shared["j2c_ip1"].view(); + let rcd_j2c_ip1 = shared["rcd_j2c_ip1"].view(); + + // --- evaluation j-part --- // + + if do_j { + let t0 = std::time::Instant::now(); + + let j_out = j_out.unwrap(); + let rrcd_eri_aux = j_in.as_ref().unwrap()["rrcd_eri_aux"].view(); + + // --- j1ao_aux1_3 --- // + + let tmp1 = -rt::vecdot(&j2c_ip1, &rrcd_eri_aux, 0); + let mut tmp2: Tsr = rt::zeros(([naux, 3, natm], device)); + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + tmp2.i_mut((slcA, .., A)).assign(tmp1.i(slcA)); + } + solve_aux(tmp2.view_mut(), Left, true); + let j1ao_aux1_3_tp = &cderi % tmp2.reshape((naux, -1)); + let j1ao_aux1_3 = j1ao_aux1_3_tp.reshape((-1, 3, natm)).unpack_tri(Upper, FlagSymm::Sy); + j_out.insert("j1ao_aux1_3", j1ao_aux1_3); + + // --- j1ao_aux1_4 --- // + + let tmp1 = &rcd_j2c_ip1 * &rrcd_eri_aux; + let mut tmp2: Tsr = rt::zeros(([naux, 3, natm], device)); + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + tmp2.i_mut((.., .., A)).assign(tmp1.i(slcA).sum_axes(0)); + } + let j1ao_aux1_4_tp = &cderi % tmp2.reshape((naux, -1)); + let j1ao_aux1_4 = j1ao_aux1_4_tp.reshape((-1, 3, natm)).unpack_tri(Upper, FlagSymm::Sy); + j_out.insert("j1ao_aux1_4", j1ao_aux1_4); + + tic!(timing, t0, "evaluate_jk1_j2c_deriv, j-part"); + } + + // --- evaluation k-part --- // + + for iset in 0..nset_k { + let t0 = std::time::Instant::now(); + + let k_in = &k_ins[iset]; + let k_out = &mut k_outs[iset]; + let rrcd_eri_occ = k_in["rrcd_eri_occ"].view(); + let rrcd_eri_bra = k_in["rrcd_eri_bra"].view(); + let occ_invsqrt = k_in["occ_invsqrt"].view(); + let nocc = occ_invsqrt.shape()[0]; + + // --- k1bra_aux1_3 --- // + + let mut k1bra_aux1_3 = rt::zeros(([nao, nocc, 3, natm], device)); + for t in 0..3 { + let tmp1 = rrcd_eri_bra.reshape([nao * nocc, naux]) % j2c_ip1.i((.., .., t)).t(); + let tmp1 = tmp1.into_shape([nao, nocc, naux]); + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + let tmp = + tmp1.i((.., .., slcA)).reshape((nao, -1)) % rrcd_eri_occ.i((.., .., slcA)).reshape((nocc, -1)).t(); + k1bra_aux1_3.i_mut((.., .., t, A)).assign(tmp); + } + } + k1bra_aux1_3 *= occ_invsqrt.i((None, ..)); + k_out.insert("k1bra_aux1_3", k1bra_aux1_3); + + // --- k1bra_aux1_4 --- // + + let mut k1bra_aux1_4 = rt::zeros(([nao, nocc, 3, natm], device)); + for t in 0..3 { + let tmp1 = rrcd_eri_occ.reshape([nocc * nocc, naux]) % j2c_ip1.i((.., .., t)).t(); + let tmp1 = tmp1.into_shape([nocc, nocc, naux]); + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + let tmp = + rrcd_eri_bra.i((.., .., slcA)).reshape((nao, -1)) % tmp1.i((.., .., slcA)).reshape((nocc, -1)).t(); + k1bra_aux1_4.i_mut((.., .., t, A)).assign(tmp); + } + } + k1bra_aux1_4 *= occ_invsqrt.i((None, ..)); + k_out.insert("k1bra_aux1_4", k1bra_aux1_4); + + tic!(timing, t0, &format!("evaluate_jk1_j2c_deriv, k-part {iset}")); + } + + timing +} + +#[allow(clippy::too_many_arguments)] +pub fn evaluate_j3c_ip2( + dims: &HashMap<&'static str, usize>, + shared: &HashMap<&'static str, Tsr>, + cderi: TsrView, + mol: &CInt, + aux: &CInt, + auxslices: &[[usize; 4]], + aux_ranges: &[[usize; 4]], + device: &DeviceBLAS, + solve_aux: &FnSolveAux, + j_in: Option<&HashMap<&'static str, Tsr>>, + k_ins: &mut [HashMap<&'static str, Tsr>], + mut j_out: Option<&mut HashMap<&'static str, Tsr>>, + k_outs: &mut [HashMap<&'static str, Tsr>], + mut j_intmd: Option<&mut HashMap<&'static str, Tsr>>, + k_intmds: &mut [HashMap<&'static str, Tsr>], +) -> Vec<(String, f64)> { + let mut timing = vec![]; + + let nao = dims["nao"]; + let naux = dims["naux"]; + let natm = dims["natm"]; + let j2c_ip1 = shared["j2c_ip1"].view(); + let rrcd_j2c_ip1 = shared["rrcd_j2c_ip1"].view(); + let j2c_inv = shared["j2c_inv"].view(); + let do_j = j_in.is_some(); + let nset_k = k_outs.len(); + + let gen_j3c_ip2 = generator_hess_intor_j3c_by_aux(mol, aux, "int3c2e_ip2", "s1", device); + + // --- integral contract --- // + + // for this part, we will evaluate integrals by batch and generate the intermediates; so that in + // future we will not evaluate these integrals again. + + if let Some(j_intmd) = j_intmd.as_mut() { + j_intmd.insert("j3c_ip2_aux", rt::zeros(([naux, 3], device))); + j_intmd.insert("j1ao_aux1_1", rt::zeros(([nao, nao, 3, natm], device))); + } + for (k_in, k_intmd) in k_ins.iter_mut().zip(k_intmds.iter_mut()) { + let nocc = k_in["occ_invsqrt"].shape()[0]; + k_intmd.insert("j3c_ip2_occ", rt::zeros(([nocc, nocc, naux, 3], device))); + k_intmd.insert("k1bra_aux1_2", rt::zeros(([nao, nocc, 3, natm], device))); + } + + for &[sh0, sh1, p0, p1] in aux_ranges.iter() { + let t0 = std::time::Instant::now(); + let j3c_ip2_batch = gen_j3c_ip2([sh0, sh1]); + tic!(timing, t0, &format!("evaluate_j3c_ip2, j3c_ip2 aux({p0}:{p1})")); + + let t0 = std::time::Instant::now(); + if do_j { + let dm0 = j_in.as_ref().unwrap()["dm0"].view(); + let rrcd_eri_aux = j_in.as_ref().unwrap()["rrcd_eri_aux"].view(); + + // --- j3c_ip2_aux --- // + + let mut j3c_ip2_aux = j_intmd.as_mut().unwrap().get_mut("j3c_ip2_aux").unwrap().view_mut(); + let tmp = rt::vecdot(&j3c_ip2_batch, &dm0, ([0, 1], [0, 1])); + j3c_ip2_aux.i_mut(p0..p1).assign(tmp); + + // --- j1ao_aux1_1 --- // + + let mut j1ao_aux1_1 = j_intmd.as_mut().unwrap().get_mut("j1ao_aux1_1").unwrap().view_mut(); + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let start = p0.max(p0A); + let end = p1.min(p1A); + if start >= end { + continue; + } + let slc_batch = rt::slice!(start - p0, end - p0); + let slc_full = rt::slice!(start, end); + + let tmp = -rt::vecdot(j3c_ip2_batch.i((.., .., slc_batch)), rrcd_eri_aux.i((None, None, slc_full)), 2); + *&mut j1ao_aux1_1.i_mut((.., .., .., A)) += tmp; + } + } + tic!(timing, t0, &format!("evaluate_j3c_ip2, intor contract j-part aux({p0}:{p1})")); + + for iset in 0..nset_k { + let t0 = std::time::Instant::now(); + + let k_in = &k_ins[iset]; + let fold_eri_bra = k_in["fold_eri_bra"].view(); + let mocc_2 = k_in["mocc_2"].view(); + let occ_invsqrt = k_in["occ_invsqrt"].view(); + let nocc = occ_invsqrt.shape()[0]; + + // --- j3c_ip2_occ --- // + + let mut j3c_ip2_occ = k_intmds[iset].get_mut("j3c_ip2_occ").unwrap().view_mut(); + let tmp = mocc_2.t() % &j3c_ip2_batch % &mocc_2; + j3c_ip2_occ.i_mut((.., .., p0..p1)).assign(tmp); + + // --- k1bra_aux1_2 --- // + + let mut k1bra_aux1_2 = k_intmds[iset].get_mut("k1bra_aux1_2").unwrap().view_mut(); + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let start = p0.max(p0A); + let end = p1.min(p1A); + if start >= end { + continue; + } + let slc_batch = rt::slice!(start - p0, end - p0); + let slc_full = rt::slice!(start, end); + + // NOTE: the following reshape copys the data + let m1 = fold_eri_bra.i((.., .., slc_full)).swapaxes(0, 1).into_shape((nocc, -1)); + for t in 0..3 { + let m2 = j3c_ip2_batch.i((.., .., slc_batch, t)).change_shape((nao, -1)); + *&mut k1bra_aux1_2.i_mut((.., .., t, A)) -= &m2 % m1.t(); + } + } + + tic!(timing, t0, &format!("evaluate_j3c_ip2, intor contract k-part {iset} aux({p0}:{p1})")); + } + } + + // --- evaluation j-part --- // + + if do_j { + let t0 = std::time::Instant::now(); + + let j_out = j_out.as_mut().unwrap(); + let rrcd_eri_aux = j_in.as_ref().unwrap()["rrcd_eri_aux"].view(); + let j3c_ip2_aux = j_intmd.as_ref().unwrap()["j3c_ip2_aux"].view(); + + // --- dbas evaluation --- // + + let tmp1 = rt::vecdot(&j2c_ip1, &rrcd_eri_aux, 0); + let dbas_J02_4 = j3c_ip2_aux.i((.., None, None, ..)) * tmp1.i((None, .., .., None)) * &j2c_inv; + let dbas_J02_5 = j3c_ip2_aux.i((.., None, None, ..)) * j3c_ip2_aux.i((None, .., .., None)) * &j2c_inv; + let tmp1 = &rrcd_j2c_ip1 * &rrcd_eri_aux; + let dbas_J02_7 = j3c_ip2_aux.i((.., None, None, ..)) * tmp1.swapaxes(0, 1); + + // --- reduce to hessian contribution --- // + + let mut de_J02_4: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J02_5: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J02_7: Tsr = rt::zeros(([3, 3, natm, natm], device)); + + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + for (B, &[_, _, p0B, p1B]) in auxslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + + let tmp = dbas_J02_4.i((slcA, slcB)).sum_axes([0, 1]); + de_J02_4.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_J02_5.i((slcA, slcB)).sum_axes([0, 1]); + de_J02_5.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_J02_7.i((slcA, slcB)).sum_axes([0, 1]); + de_J02_7.i_mut((.., .., B, A)).assign(tmp); + } + } + + let scale_J02_4 = 1.0; + let scale_J02_5 = 0.5; + let scale_J02_7 = -1.0; + j_out.insert("de_J02_4", scale_J02_4 * (&de_J02_4 + &de_J02_4.transpose([1, 0, 3, 2]))); + j_out.insert("de_J02_5", scale_J02_5 * (&de_J02_5 + &de_J02_5.transpose([1, 0, 3, 2]))); + j_out.insert("de_J02_7", scale_J02_7 * (&de_J02_7 + &de_J02_7.transpose([1, 0, 3, 2]))); + + // --- j1ao aux1_2 --- // + + let mut tmp1: Tsr = rt::zeros(([naux, 3, natm], device)); + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + tmp1.i_mut((slcA, .., A)).assign(j3c_ip2_aux.i(slcA)); + } + solve_aux(tmp1.view_mut(), Left, true); + let j1ao_aux1_2_tp = &cderi % tmp1.reshape((naux, natm * 3)); + let j1ao_aux1_2 = -j1ao_aux1_2_tp.reshape((-1, 3, natm)).unpack_tri(Upper, FlagSymm::Sy); + j_out.insert("j1ao_aux1_2", j1ao_aux1_2); + + tic!(timing, t0, "evaluate_j3c_ip2, evaluation j-part"); + } + + // --- evaluation k-part --- // + + for iset in 0..nset_k { + let t0 = std::time::Instant::now(); + + let k_in = &k_ins[iset]; + let k_out = &mut k_outs[iset]; + let rrcd_eri_occ = k_in["rrcd_eri_occ"].view(); + let rrcd_eri_bra = k_in["rrcd_eri_bra"].view(); + let occ_invsqrt = k_in["occ_invsqrt"].view(); + let j3c_ip2_occ = k_intmds[iset]["j3c_ip2_occ"].view(); + let nocc = occ_invsqrt.shape()[0]; + + // --- dbas evaluation --- // + + let tmp1 = rrcd_eri_occ.reshape([nocc * nocc, naux]) % &j2c_ip1; + let tmp1 = tmp1.into_shape([nocc, nocc, naux, 3]); + let j3c_ip2_occ_2d = j3c_ip2_occ.reshape((nocc * nocc, naux, 3)); + let mut dbas_K02_4: Tsr = rt::zeros(([naux, naux, 3, 3], device)); + for t in 0..3 { + let tmp1_t = tmp1.i((.., .., .., t)); + let tmp1_t = tmp1_t.reshape((nocc * nocc, naux)); + let tmp = tmp1_t.t() % &j3c_ip2_occ_2d * &j2c_inv; + dbas_K02_4.i_mut((.., .., .., t)).assign(tmp); + } + + let mut dbas_K02_5: Tsr = rt::zeros(([naux, naux, 3, 3], device)); + for t in 0..3 { + let tmp1_t = j3c_ip2_occ.i((.., .., .., t)); + let tmp1_t = tmp1_t.reshape((nocc * nocc, naux)); + let tmp = tmp1_t.t() % &j3c_ip2_occ_2d * &j2c_inv; + dbas_K02_5.i_mut((.., .., .., t)).assign(tmp); + } + + let rrcd_eri_occ_2d = rrcd_eri_occ.reshape((nocc * nocc, naux)); + let tmp1 = rrcd_eri_occ_2d.t() % j3c_ip2_occ_2d; + let dbas_K02_7 = tmp1.i((.., .., .., None)) * rrcd_j2c_ip1.i((.., .., None, ..)); + + // --- reduce to hessian contribution --- // + + let mut de_K02_4: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K02_5: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K02_7: Tsr = rt::zeros(([3, 3, natm, natm], device)); + + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + for (B, &[_, _, p0B, p1B]) in auxslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + + let tmp = dbas_K02_4.i((slcA, slcB)).sum_axes([0, 1]); + de_K02_4.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_K02_5.i((slcA, slcB)).sum_axes([0, 1]); + de_K02_5.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_K02_7.i((slcA, slcB)).sum_axes([0, 1]); + de_K02_7.i_mut((.., .., B, A)).assign(tmp); + } + } + + let scale_K02_4 = 1.0; + let scale_K02_5 = 0.5; + let scale_K02_7 = -1.0; + k_out.insert("de_K02_4", scale_K02_4 * (&de_K02_4 + &de_K02_4.transpose([1, 0, 3, 2]))); + k_out.insert("de_K02_5", scale_K02_5 * (&de_K02_5 + &de_K02_5.transpose([1, 0, 3, 2]))); + k_out.insert("de_K02_7", scale_K02_7 * (&de_K02_7 + &de_K02_7.transpose([1, 0, 3, 2]))); + + // --- k1bra aux1_1 --- // + + let mut k1bra_aux1_1 = rt::zeros(([nao, nocc, 3, natm], device)); + for (A, &[_, _, p0A, p1A]) in auxslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + for t in 0..3 { + let m1 = rrcd_eri_bra.i((.., .., slcA)).change_shape((nao, -1)); + let m2 = j3c_ip2_occ.i((.., .., slcA, t)).change_shape((nocc, -1)); + k1bra_aux1_1.i_mut((.., .., t, A)).assign(-(&m1 % m2.t())); + } + } + k1bra_aux1_1 *= occ_invsqrt.i((None, ..)); + k_out.insert("k1bra_aux1_1", k1bra_aux1_1); + + tic!(timing, t0, &format!("evaluate_j3c_ip2, evaluation k-part {iset}")); + } + + // --- move some intermediates to output --- // + + if do_j { + let j_out = j_out.as_mut().unwrap(); + let j1ao_aux1_1 = j_intmd.unwrap().remove("j1ao_aux1_1").unwrap(); + j_out.insert("j1ao_aux1_1", j1ao_aux1_1); + } + for (iset, k_intmd) in k_intmds.iter_mut().enumerate() { + let k_out = &mut k_outs[iset]; + let occ_invsqrt = k_ins[iset]["occ_invsqrt"].view(); + let mut k1bra_aux1_2 = k_intmd.remove("k1bra_aux1_2").unwrap(); + *&mut k1bra_aux1_2 *= occ_invsqrt.i((None, ..)); + k_out.insert("k1bra_aux1_2", k1bra_aux1_2); + } + + timing +} + +#[allow(clippy::too_many_arguments)] +pub fn evaluate_j3c_ip1( + dims: &HashMap<&'static str, usize>, + shared: &HashMap<&'static str, Tsr>, + cderi: TsrView, + mol: &CInt, + aux: &CInt, + aoslices: &[[usize; 4]], + auxslices: &[[usize; 4]], + aux_ranges: &[[usize; 4]], + device: &DeviceBLAS, + solve_aux: &FnSolveAux, + j_in: Option<&HashMap<&'static str, Tsr>>, + k_ins: &mut [HashMap<&'static str, Tsr>], + mut j_out: Option<&mut HashMap<&'static str, Tsr>>, + k_outs: &mut [HashMap<&'static str, Tsr>], + mut j_intmd: Option<&mut HashMap<&'static str, Tsr>>, + k_intmds: &mut [HashMap<&'static str, Tsr>], +) -> Vec<(String, f64)> { + let mut timing = vec![]; + + let nao = dims["nao"]; + let naux = dims["naux"]; + let natm = dims["natm"]; + let do_j = j_in.is_some(); + let nset_k = k_outs.len(); + let j2c_ip1 = shared["j2c_ip1"].view(); + + // --- integral contract --- // + + // this part will also handle intermediates, as before in evaluate_j3c_ip2. + + if let Some(j_intmd) = j_intmd.as_mut() { + j_intmd.insert("j3c_ip1_aux", rt::zeros(([nao, naux, 3], device))); + j_intmd.insert("j3c_ip1_j1ao_tmp", rt::zeros(([nao, nao, 3], device))); + } + for (k_in, k_intmd) in k_ins.iter_mut().zip(k_intmds.iter_mut()) { + let nocc = k_in["occ_invsqrt"].shape()[0]; + k_intmd.insert("j3c_ip1_bra", rt::zeros(([nao, nocc, naux, 3], device))); + k_intmd.insert("j3c_ip1_k1ao_tmp", rt::zeros(([nao, nao, 3], device))); + k_intmd.insert("k1bra_aux0_4", rt::zeros(([nao, nocc, 3, natm], device))); + } + + for &[sh0, sh1, p0, p1] in aux_ranges.iter() { + let t0 = std::time::Instant::now(); + let j3c_ip1_batch = generator_hess_intor_j3c_by_aux(mol, aux, "int3c2e_ip1", "s1", device)([sh0, sh1]); + tic!(timing, t0, &format!("evaluate_j3c_ip1, j3c_ip1 aux({p0}:{p1})")); + + let t0 = std::time::Instant::now(); + if do_j { + let dm0 = j_in.as_ref().unwrap()["dm0"].view(); + let rrcd_eri_aux = j_in.as_ref().unwrap()["rrcd_eri_aux"].i(p0..p1); + + // --- j3c_ip1_aux --- // + + let mut j3c_ip1_aux = j_intmd.as_mut().unwrap().get_mut("j3c_ip1_aux").unwrap().view_mut(); + let tmp = rt::vecdot(&j3c_ip1_batch, &dm0, 1); + j3c_ip1_aux.i_mut((.., p0..p1)).assign(tmp); + + // --- j3c_ip1_j1ao_tmp --- // + + let mut j3c_ip1_j1ao_tmp = j_intmd.as_mut().unwrap().get_mut("j3c_ip1_j1ao_tmp").unwrap().view_mut(); + j3c_ip1_j1ao_tmp += rt::vecdot(&j3c_ip1_batch, rrcd_eri_aux.i((None, None, ..)), 2); + } + tic!(timing, t0, &format!("evaluate_j3c_ip1, intor contract j-part aux({p0}:{p1})")); + + for iset in 0..nset_k { + let t0 = std::time::Instant::now(); + + let k_in = &k_ins[iset]; + let k_intmd = &mut k_intmds[iset]; + let rrcd_eri_bra = k_in["rrcd_eri_bra"].view(); + let fold_eri_bra = k_in["fold_eri_bra"].view(); + let mocc_2 = k_in["mocc_2"].view(); + let nocc = k_in["occ_invsqrt"].shape()[0]; + + // some tensors are consequently used, need get_disjoint_mut to avoid borrow checker error + let [j3c_ip1_bra, j3c_ip1_k1ao_tmp] = k_intmd.get_disjoint_mut(["j3c_ip1_bra", "j3c_ip1_k1ao_tmp"]); + let j3c_ip1_bra = j3c_ip1_bra.unwrap(); + let j3c_ip1_k1ao_tmp = j3c_ip1_k1ao_tmp.unwrap(); + + // --- j3c_ip1_bra --- // + + let tmp = &j3c_ip1_batch % &mocc_2; + j3c_ip1_bra.i_mut((.., .., p0..p1)).assign(tmp); + + // --- j3c_ip1_k1ao_tmp --- // + + let m1 = j3c_ip1_bra.i((.., .., p0..p1, ..)).change_shape((nao, -1, 3)); + let m2 = rrcd_eri_bra.i((.., .., p0..p1)).change_shape((nao, -1)); + *j3c_ip1_k1ao_tmp += m1 % m2.t(); + + // --- k1bra_aux0_4 --- // + + let mut k1bra_aux0_4 = k_intmd.get_mut("k1bra_aux0_4").unwrap().view_mut(); + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + // NOTE: reshape with data copy + let tmp1 = fold_eri_bra.i((slcA, .., p0..p1)).swapaxes(0, 1).into_shape((nocc, -1)); + for t in 0..3 { + // NOTE: reshape with data copy + let tmp2 = j3c_ip1_batch.i((slcA, .., .., t)).into_swapaxes(0, 1).into_shape((nao, -1)); + *&mut k1bra_aux0_4.i_mut((.., .., t, A)) -= &tmp2 % tmp1.t(); + } + } + + tic!(timing, t0, &format!("evaluate_j3c_ip1, intor contract k-part {iset} aux({p0}:{p1})")); + } + } + + // --- evaluation j-part --- // + + if do_j { + let t0 = std::time::Instant::now(); + + let j_out = j_out.as_mut().unwrap(); + let j3c_ip1_j1ao_tmp = j_intmd.as_mut().unwrap().remove("j3c_ip1_j1ao_tmp").unwrap(); + let j3c_ip1_aux = j_intmd.as_mut().unwrap().remove("j3c_ip1_aux").unwrap(); + let j3c_ip2_aux = j_intmd.as_mut().unwrap().remove("j3c_ip2_aux").unwrap(); + let rrcd_eri_aux = j_in.as_ref().unwrap()["rrcd_eri_aux"].view(); + + let mut rcd_j3c_ip1_aux = j3c_ip1_aux; + for t in 0..3 { + solve_aux(rcd_j3c_ip1_aux.i_mut((.., .., t)), Right, false); + } + + // --- j1ao aux0 --- // + + let mut j1ao_aux0 = rt::zeros(([nao, nao, 3, natm], device)); + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + *&mut j1ao_aux0.i_mut((.., slcA, .., A)) -= j3c_ip1_j1ao_tmp.i(slcA).swapaxes(0, 1); + *&mut j1ao_aux0.i_mut((slcA, .., .., A)) -= j3c_ip1_j1ao_tmp.i(slcA); + let tmp1 = rcd_j3c_ip1_aux.i(slcA).sum_axes(0); + let tmp2 = &cderi % &tmp1; + *&mut j1ao_aux0.i_mut((Ellipsis, A)) -= 2 * tmp2.unpack_tri(Upper, FlagSymm::Sy); + } + + j_out.insert("j1ao_aux0", j1ao_aux0); + + // --- dbas evaluation --- // + + let mut dbas_J20_1: Tsr = rt::zeros(([nao, nao, 3, 3], device)); + for t in 0..3 { + for s in 0..3 { + let tmp = rcd_j3c_ip1_aux.i((.., .., t)) % rcd_j3c_ip1_aux.i((.., .., s)).t(); + dbas_J20_1.i_mut((.., .., s, t)).assign(tmp); + } + } + + let mut rrcd_j3c_ip1_aux = rcd_j3c_ip1_aux; + for t in 0..3 { + solve_aux(rrcd_j3c_ip1_aux.i_mut((.., .., t)), Right, true); + } + + let mut dbas_J11_2: Tsr = rt::zeros(([nao, naux, 3, 3], device)); + for t in 0..3 { + for s in 0..3 { + let tmp = rrcd_j3c_ip1_aux.i((.., .., t)) % j2c_ip1.i((.., .., s)).t() * rrcd_eri_aux.i((None, ..)); + dbas_J11_2.i_mut((.., .., s, t)).assign(tmp); + } + } + + let tmp1 = rt::vecdot(&j2c_ip1, &rrcd_eri_aux, 0); + let dbas_J11_3: Tsr = rrcd_j3c_ip1_aux.i((.., .., None, ..)) * tmp1.i((None, .., .., None)); + let dbas_J11_4: Tsr = rrcd_j3c_ip1_aux.i((.., .., None, ..)) * j3c_ip2_aux.i((None, .., .., None)); + + // --- reduce to hessian contribution --- // + + let mut de_J20_1: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J11_2: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J11_3: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_J11_4: Tsr = rt::zeros(([3, 3, natm, natm], device)); + + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + + for (B, &[_, _, p0B, p1B]) in aoslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + + let tmp = dbas_J20_1.i((slcA, slcB)).sum_axes([0, 1]); + de_J20_1.i_mut((.., .., B, A)).assign(tmp); + } + + for (B, &[_, _, p0B, p1B]) in auxslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + + let tmp = dbas_J11_2.i((slcA, slcB)).sum_axes([0, 1]); + de_J11_2.i_mut((.., .., B, A)).assign(tmp); + let tmp = dbas_J11_3.i((slcA, slcB)).sum_axes([0, 1]); + de_J11_3.i_mut((.., .., B, A)).assign(tmp); + let tmp = dbas_J11_4.i((slcA, slcB)).sum_axes([0, 1]); + de_J11_4.i_mut((.., .., B, A)).assign(tmp); + } + } + + let scale_J20_1 = 2.0; + let scale_J11_2 = -2.0; + let scale_J11_3 = 2.0; + let scale_J11_4 = 2.0; + j_out.insert("de_J20_1", scale_J20_1 * (&de_J20_1 + &de_J20_1.transpose([1, 0, 3, 2]))); + j_out.insert("de_J11_2", scale_J11_2 * (&de_J11_2 + &de_J11_2.transpose([1, 0, 3, 2]))); + j_out.insert("de_J11_3", scale_J11_3 * (&de_J11_3 + &de_J11_3.transpose([1, 0, 3, 2]))); + j_out.insert("de_J11_4", scale_J11_4 * (&de_J11_4 + &de_J11_4.transpose([1, 0, 3, 2]))); + + tic!(timing, t0, "evaluate_j3c_ip1, evaluation j-part"); + } + + // --- evaluation k-part --- // + + for iset in 0..nset_k { + let t0 = std::time::Instant::now(); + + let k_in = &k_ins[iset]; + let k_out = &mut k_outs[iset]; + let rrcd_eri_bra = k_in["rrcd_eri_bra"].view(); + let occ_invsqrt = k_in["occ_invsqrt"].view(); + let fold_eri_bra = k_in["fold_eri_bra"].view(); + let mocc = k_in["mocc"].view(); + let mocc_2 = k_in["mocc_2"].view(); + let nocc = occ_invsqrt.shape()[0]; + + let j3c_ip2_occ = k_intmds[iset].remove("j3c_ip2_occ").unwrap(); + let j3c_ip1_bra = k_intmds[iset].remove("j3c_ip1_bra").unwrap(); + let j3c_ip1_k1ao_tmp = k_intmds[iset].remove("j3c_ip1_k1ao_tmp").unwrap(); + + // --- k1bra --- // + + let t1 = std::time::Instant::now(); + + let mut k1ao_aux_1: Tsr = rt::zeros(([nao, nao, 3, natm], device)); + let mut k1ao_aux_2: Tsr = rt::zeros(([nao, nao, 3, natm], device)); + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + *&mut k1ao_aux_1.i_mut((.., slcA, .., A)) -= j3c_ip1_k1ao_tmp.i(slcA).swapaxes(0, 1); + *&mut k1ao_aux_2.i_mut((slcA, .., .., A)) -= j3c_ip1_k1ao_tmp.i(slcA); + } + let k1bra_aux0_1 = k1ao_aux_1 % &mocc; + let k1bra_aux0_2 = k1ao_aux_2 % &mocc; + k_out.insert("k1bra_aux0_1", k1bra_aux0_1); + k_out.insert("k1bra_aux0_2", k1bra_aux0_2); + + let mut k1bra_aux0_3: Tsr = rt::zeros(([nao, nocc, 3, natm], device)); + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + for t in 0..3 { + let tmp1 = j3c_ip1_bra.i((slcA, .., .., t)).swapaxes(0, 1) % mocc.i(slcA); + let tmp2 = rrcd_eri_bra.reshape((nao, -1)) % tmp1.reshape((nocc, -1)).t(); + k1bra_aux0_3.i_mut((.., .., t, A)).assign(-tmp2); + } + } + k_out.insert("k1bra_aux0_3", k1bra_aux0_3); + + tic!(timing, t1, &format!("evaluate_j3c_ip1, k1bra-1/2/3 {iset}")); + + // --- dbas evaluation rcd_j3c_ip1_bra --- // + + let mut rcd_j3c_ip1_bra = j3c_ip1_bra; + for t in 0..3 { + solve_aux(rcd_j3c_ip1_bra.i_mut((.., .., .., t)), Right, false); + } + + let dm = &mocc_2 % mocc_2.t(); // spin density, not total density + + let t1 = std::time::Instant::now(); + let mut dbas_K20_1a: Tsr = rt::zeros(([nao, nao, 3, 3], device)); + let tmp1 = rcd_j3c_ip1_bra.reshape((nao, nocc * naux, 3)); + for t in 0..3 { + for s in 0..=t { + let tmp2 = tmp1.i((.., .., t)) % tmp1.i((.., .., s)).t() * &dm; + dbas_K20_1a.i_mut((.., .., s, t)).assign(&tmp2); + // apply symmetric trick + if t != s { + dbas_K20_1a.i_mut((.., .., t, s)).assign(tmp2.t()); + } + } + } + tic!(timing, t1, &format!("evaluate_j3c_ip1, dbas_K20_1a {iset}")); + + // de_K20_1b is special. This term is better to be pre-contracted to hessian. + let t1 = std::time::Instant::now(); + let mut de_K20_1b: Tsr = rt::zeros(([3, 3, natm, natm], device)); + for &[_, _, p0, p1] in aux_ranges { + let slcP = rt::slice!(p0, p1); + let nbatch = p1 - p0; + let mut fold_occ: Tsr = rt::zeros(([nocc, nocc, nbatch, 3, natm], device)); + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + let tmp1 = rcd_j3c_ip1_bra.i((slcA, .., slcP, ..)); + for t in 0..3 { + let tmp = mocc_2.i(slcA).t() % tmp1.i((.., .., .., t)); + fold_occ.i_mut((.., .., .., t, A)).assign(tmp); + } + } + // non-trivial transpose + let fold_occ_swap = fold_occ.swapaxes(0, 1).into_contig(ColMajor); + // handle shape conversions and hessian output + let fold_occ = fold_occ.reshape([nocc * nocc * nbatch, 3 * natm]); + let fold_occ_swap = fold_occ_swap.reshape([nocc * nocc * nbatch, 3 * natm]); + let de_swap = fold_occ.t() % fold_occ_swap; // [s, B, t, A] + let de_increment = de_swap.reshape([3, natm, 3, natm]).into_swapaxes(1, 2); // [s, t, B, A] + de_K20_1b += de_increment; + } + let scale_K20_1b = 1.0; + k_out.insert("de_K20_1b", scale_K20_1b * (&de_K20_1b + &de_K20_1b.transpose([1, 0, 3, 2]))); + tic!(timing, t1, &format!("evaluate_j3c_ip1, de_K20_1b {iset}")); + + // --- dbas evaluation rrcd_j3c_ip1_bra --- // + + let mut rrcd_j3c_ip1_bra = rcd_j3c_ip1_bra; + for t in 0..3 { + solve_aux(rrcd_j3c_ip1_bra.i_mut((.., .., .., t)), Right, true); + } + + let t1 = std::time::Instant::now(); + let mut dbas_K11_2: Tsr = rt::zeros(([nao, naux, 3, 3], device)); + for t in 0..3 { + for s in 0..3 { + let tmp1 = rrcd_j3c_ip1_bra.i((.., .., .., t)).change_shape([nao * nocc, naux]); + let tmp2 = (tmp1 % j2c_ip1.i((.., .., s))).into_shape([nao, nocc, naux]); + let tmp3 = rt::vecdot(&fold_eri_bra, tmp2, 1); + dbas_K11_2.i_mut((.., .., s, t)).assign(&tmp3); + } + } + tic!(timing, t1, &format!("evaluate_j3c_ip1, dbas_K11_2 {iset}")); + + let t1 = std::time::Instant::now(); + let mut dbas_K11_3: Tsr = rt::zeros(([nao, naux, 3, 3], device)); + for s in 0..3 { + let tmp1 = fold_eri_bra.reshape([nao * nocc, naux]) % j2c_ip1.i((.., .., s)); + let tmp1 = tmp1.into_shape([nao, nocc, naux]); + for t in 0..3 { + let tmp = rt::vecdot(rrcd_j3c_ip1_bra.i((.., .., .., t)), &tmp1, 1); + dbas_K11_3.i_mut((.., .., s, t)).assign(tmp); + } + } + tic!(timing, t1, &format!("evaluate_j3c_ip1, dbas_K11_3 {iset}")); + + let t1 = std::time::Instant::now(); + let mut dbas_K11_4: Tsr = rt::zeros(([nao, naux, 3, 3], device)); + for s in 0..3 { + let tmp1 = &mocc_2 % j3c_ip2_occ.i((.., .., .., s)); + for t in 0..3 { + let tmp = rt::vecdot(rrcd_j3c_ip1_bra.i((.., .., .., t)), &tmp1, 1); + dbas_K11_4.i_mut((.., .., s, t)).assign(tmp); + } + } + tic!(timing, t1, &format!("evaluate_j3c_ip1, dbas_K11_4 {iset}")); + + // --- reduce to hessian contribution --- // + + let mut de_K20_1a: Tsr = rt::zeros(([3, 3, natm, natm], device)); + // let mut de_K20_1b: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K11_2: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K11_3: Tsr = rt::zeros(([3, 3, natm, natm], device)); + let mut de_K11_4: Tsr = rt::zeros(([3, 3, natm, natm], device)); + + for (A, &[_, _, p0A, p1A]) in aoslices.iter().enumerate() { + let slcA = rt::slice!(p0A, p1A); + + for (B, &[_, _, p0B, p1B]) in aoslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + + let tmp = dbas_K20_1a.i((slcA, slcB)).sum_axes([0, 1]); + de_K20_1a.i_mut((.., .., B, A)).assign(tmp); + + // let tmp = dbas_K20_1b.i((slcA, slcB)).sum_axes([0, 1]); + // de_K20_1b.i_mut((.., .., B, A)).assign(tmp); + } + + for (B, &[_, _, p0B, p1B]) in auxslices.iter().enumerate() { + let slcB = rt::slice!(p0B, p1B); + + let tmp = dbas_K11_2.i((slcA, slcB)).sum_axes([0, 1]); + de_K11_2.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_K11_3.i((slcA, slcB)).sum_axes([0, 1]); + de_K11_3.i_mut((.., .., B, A)).assign(tmp); + + let tmp = dbas_K11_4.i((slcA, slcB)).sum_axes([0, 1]); + de_K11_4.i_mut((.., .., B, A)).assign(tmp); + } + } + + let scale_K20_1a = 1.0; + let scale_K11_2 = 2.0; + let scale_K11_3 = 2.0; + let scale_K11_4 = 2.0; + k_out.insert("de_K20_1a", scale_K20_1a * (&de_K20_1a + &de_K20_1a.transpose([1, 0, 3, 2]))); + k_out.insert("de_K11_2", scale_K11_2 * (&de_K11_2 + &de_K11_2.transpose([1, 0, 3, 2]))); + k_out.insert("de_K11_3", scale_K11_3 * (&de_K11_3 + &de_K11_3.transpose([1, 0, 3, 2]))); + k_out.insert("de_K11_4", scale_K11_4 * (&de_K11_4 + &de_K11_4.transpose([1, 0, 3, 2]))); + + tic!(timing, t0, &format!("evaluate_j3c_ip1, evaluation k-part {iset}")); + } + + // --- move some intermediates to output --- // + + for (iset, k_intmd) in k_intmds.iter_mut().enumerate() { + let k_out = &mut k_outs[iset]; + let occ_invsqrt = k_ins[iset]["occ_invsqrt"].view(); + let mut k1bra_aux0_4 = k_intmd.remove("k1bra_aux0_4").unwrap(); + *&mut k1bra_aux0_4 *= occ_invsqrt.i((None, ..)); + k_out.insert("k1bra_aux0_4", k1bra_aux0_4); + } + + timing +} + +/* #endregion */ + +/* #region impl */ + +/// Generate cderi and decomposition. +pub fn generate_cderi_with_decomp( + mol: &CInt, + aux: &CInt, + j2c_decomp_option: J2CDecompOption, + device: &DeviceBLAS, +) -> (Tsr, J2CDecompose) { + let j3c = hess_intor_cross(&[mol, mol, aux], "int3c2e", "s2ij", None, device); + let j2c_decomp = get_j2c_decomp(aux, device, j2c_decomp_option); + let cderi = solve_by_j2c(j3c, &j2c_decomp, Right, false); + (cderi, j2c_decomp) +} + +pub struct RHessRIJK<'a> { + pub mol: CInt, + pub aux: CInt, + pub scale_j: f64, + pub scale_k: f64, + pub cderi: TsrCow<'a>, + pub j2c_decomp: J2CDecompose, + pub intmd: HashMap, // intermediates + pub result: HashMap<&'static str, Tsr>, + pub timing: Vec<(String, f64)>, + pub is_skeleton_ready: bool, +} + +impl<'a> RHessRIJK<'a> { + pub fn new_without_cderi(mol: &CInt, aux: &CInt, scale_j: f64, scale_k: f64) -> Self { + let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(1e-14), uplo: Upper }; + // note: the following two options are also valid + // let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(1e-14), + // uplo: Lower }; + // let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Eig, threshold: Some(1e-14), + // uplo: Upper }; + let device = DeviceBLAS::default(); + let (cderi, j2c_decomp) = generate_cderi_with_decomp(mol, aux, j2c_decomp_option, &device); + Self { + mol: mol.clone(), + aux: aux.clone(), + scale_j, + scale_k, + cderi: cderi.into_cow(), + j2c_decomp, + intmd: HashMap::new(), + result: HashMap::new(), + timing: Vec::new(), + is_skeleton_ready: false, + } + } + + pub fn new_with_cderi( + mol: &CInt, + aux: &CInt, + scale_j: f64, + scale_k: f64, + cderi: TsrCow<'a>, + j2c_decomp: J2CDecompose, + ) -> Self { + Self { + mol: mol.clone(), + aux: aux.clone(), + scale_j, + scale_k, + cderi: cderi.into_cow(), + j2c_decomp, + intmd: HashMap::new(), + result: HashMap::new(), + timing: Vec::new(), + is_skeleton_ready: false, + } + } + + pub fn ensure_skeleton(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) { + if self.is_skeleton_ready { + return; + } + let (j_out, k_outs, timing) = get_rijk_skeleton_decomposed_separated( + &self.mol, + &self.aux, + &[mo_coeff], + &[mo_occ], + self.cderi.view(), + &self.j2c_decomp, + self.scale_j != 0.0, + self.scale_k != 0.0, + 72, // TODO: batch size `72` should be tunable by max-memory. + atm_list, + None, + ); + self.timing.extend(timing); + + if let Some(j_out) = j_out { + for (key, value) in j_out.into_iter() { + self.intmd.insert(key.to_string(), value); + } + }; + + for (iset, k_out) in k_outs.into_iter().enumerate() { + // note the keys can clash for output of k. + // for storage of intermediates, we append `` to the key name. + for (key, value) in k_out.into_iter() { + self.intmd.insert(format!("{key}"), value); + } + } + + self.is_skeleton_ready = true; + } +} + +impl<'a> HessUtilAPI for RHessRIJK<'a> {} + +impl<'a> RHessElecInteractAPI for RHessRIJK<'a> { + fn make_skeleton_hess(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr { + self.ensure_skeleton(mo_coeff, mo_occ, atm_list); + let intmd = &self.intmd; + + let device = self.cderi.device(); + let natm = atm_list.map_or_else(|| self.mol.natm(), |list| list.len()); + let hess_init = || -> Tsr { rt::zeros(([3, 3, natm, natm], device)) }; + + let mut de = hess_init(); + if self.scale_j != 0.0 { + let de_J20 = KEYS_J20.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); + let de_J11 = KEYS_J11.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); + let de_J02 = KEYS_J02.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); + let de_J = &de_J20 + &de_J11 + &de_J02; + de += self.scale_j * &de_J; + self.result.insert("de_J20", de_J20); + self.result.insert("de_J11", de_J11); + self.result.insert("de_J02", de_J02); + self.result.insert("de_J", de_J); + } + if self.scale_k != 0.0 { + // rhf only have one spin + let de_K20 = + KEYS_K20.iter().map(|&key| &intmd[&format!("{key}")]).fold(hess_init(), |acc, x| acc + x); + let de_K11 = + KEYS_K11.iter().map(|&key| &intmd[&format!("{key}")]).fold(hess_init(), |acc, x| acc + x); + let de_K02 = + KEYS_K02.iter().map(|&key| &intmd[&format!("{key}")]).fold(hess_init(), |acc, x| acc + x); + let de_K = &de_K20 + &de_K11 + &de_K02; + de -= 0.5 * self.scale_k * &de_K; + self.result.insert("de_K20", de_K20); + self.result.insert("de_K11", de_K11); + self.result.insert("de_K02", de_K02); + self.result.insert("de_K", de_K); + } + self.result.insert("de_skeleton", de.clone()); + de + } + + fn get_deriv1_ao(&mut self, _mo_coeff: TsrView, _mo_occ: TsrView, _atm_list: Option<&[usize]>) -> Tsr { + unimplemented!("This function is not implemented for optimized RI-JK hessian. Use `get_deriv1_bra` instead.") + } + + fn get_deriv1_bra(&mut self, mo_coeff: TsrView, mo_occ: TsrView, atm_list: Option<&[usize]>) -> Tsr { + self.ensure_skeleton(mo_coeff.view(), mo_occ.view(), atm_list); + let intmd = &self.intmd; + + let device = self.cderi.device(); + let natm = atm_list.map_or_else(|| self.mol.natm(), |list| list.len()); + let nao = mo_coeff.shape()[0]; + let occidx = mo_occ.greater(0.0).into_vec(); + let nocc = occidx.iter().filter(|&&x| x).count(); + let mocc = mo_coeff.bool_select(-1, &occidx).into_contig(ColMajor); + + let deriv1_ao_init = || -> Tsr { rt::zeros(([nao, nao, 3, natm], device)) }; + let deriv1_bra_init = || -> Tsr { rt::zeros(([nao, nocc, 3, natm], device)) }; + + let mut deriv1_bra = deriv1_bra_init(); + if self.scale_j != 0.0 { + let j1ao = KEYS_J1AO.iter().map(|&key| &intmd[key]).fold(deriv1_ao_init(), |acc, x| acc + x); + deriv1_bra += self.scale_j * (&j1ao % &mocc); + self.result.insert("j1ao", j1ao); + } + if self.scale_k != 0.0 { + let k1bra = KEYS_K1BRA + .iter() + .map(|&key| &intmd[&format!("{key}")]) + .fold(deriv1_bra_init(), |acc, x| acc + x); + deriv1_bra -= 0.5 * self.scale_k * &k1bra; + self.result.insert("k1bra", k1bra); + } + self.result.insert("deriv1_bra", deriv1_bra.clone()); + deriv1_bra + } + + fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView) { + 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 { + let mo_coeff = self.intmd["mo_coeff"].view(); + let mo_occ = self.intmd["mo_occ"].view(); + let cderi = self.cderi.view(); + + // RHF (single spin) assembly of the separated J/K response core. + // - J (AO form, from total density) contracted with `mocc` and scaled by `scale_j`. + // - K (same-spin bra form) scaled by `scale_k`; the core already bakes in the exchange sign. + // - RHF exchange prefactor (occ = 2) is folded into `scale_k`, matching the naive convention. + let shape_bra = bra.shape().to_vec(); + let nao = mo_coeff.shape()[0]; + let device = mo_coeff.device(); + let occidx = mo_occ.view().greater(0).into_vec(); + let mocc = mo_coeff.bool_select(-1, &occidx); + let nocc = mocc.shape()[1]; + let nprop: usize = shape_bra[2..].iter().product(); + + // TODO: batch size `72` should be tunable by max-memory. + let (j_ao, k_bras) = get_rijk_response_bra_separated( + cderi, + &[mo_coeff.view()], + &[mo_occ.view()], + &[bra.view()], + self.scale_j != 0.0, + self.scale_k != 0.0, + 72, + ); + + let mut resp: Tsr = rt::zeros(([nao, nocc, nprop], device)); + if let Some(resp_ao_j) = j_ao { + resp += self.scale_j * (resp_ao_j % &mocc); + } + if let Some(k_bra) = k_bras.first() { + // K bra is returned in the original trailing shape; flatten trailing dims to (nao, nocc, nprop). + resp += self.scale_k * k_bra.view().reshape((nao, nocc, nprop)); + } + resp.into_shape(shape_bra) + } +} + +/* #endregion */ diff --git a/src/ri_jk/hess_u.rs b/src/ri_jk/hess_u.rs new file mode 100644 index 0000000000..49ec9293f5 --- /dev/null +++ b/src/ri_jk/hess_u.rs @@ -0,0 +1,314 @@ +//! Optimized RI-JK Hessian computation for unrestricted SCF (UHF). +//! +//! This module reuses the J/K-separated skeleton driver +//! [`crate::ri_jk::hess_r::get_rijk_skeleton_decomposed_separated`] (which natively handles an +//! arbitrary number of spin sets) to produce J (from the total density) and both K spins +//! (K^alpha, K^beta) in a **single** pass that shares every 3c-integral batch. +//! +//! - J (Coulomb) is spin-independent: built from the total density ``D^alpha + D^beta``. +//! - K (exchange) is ``K^alpha + K^beta``; each spin channel is built from its own occupied +//! orbitals (UHF occ = 1, so ``mocc_2 = mocc``). +//! +//! Scaling difference to the restricted ([`crate::ri_jk::hess_r::RHessRIJK`]) counterpart: +//! - Skeleton: ``scale_j * de_J - scale_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`]). +//! - First derivative (bra form): ``scale_j * (j1ao @ mocc_s) - scale_k * k1bra_s`` per spin (no +//! ``0.5`` factor), again matching the naive UHF convention. +//! +//! The response (`get_response_bra`) reuses the separated J/K response core +//! [`crate::ri_jk::hess_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). + +use super::prelude_dev::*; +use crate::hessian_backup::prelude::*; +use crate::ri_jk::util::*; + +use crate::ri_jk::decompose::*; +use crate::ri_jk::hess_r::{ + generate_cderi_with_decomp, get_rijk_response_bra_separated, get_rijk_skeleton_decomposed_separated, KEYS_J02, + KEYS_J11, KEYS_J1AO, KEYS_J20, KEYS_K02, KEYS_K11, KEYS_K1BRA, KEYS_K20, +}; + +/* #region impl */ + +pub struct UHessRIJK<'a> { + pub mol: CInt, + pub aux: CInt, + pub scale_j: f64, + pub scale_k: f64, + pub cderi: TsrCow<'a>, + pub j2c_decomp: J2CDecompose, + pub intmd: HashMap, // intermediates + pub result: HashMap<&'static str, Tsr>, + pub timing: Vec<(String, f64)>, + pub is_skeleton_ready: bool, +} + +impl<'a> UHessRIJK<'a> { + pub fn new_without_cderi(mol: &CInt, aux: &CInt, scale_j: f64, scale_k: f64) -> Self { + let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(1e-14), uplo: Upper }; + let device = DeviceBLAS::default(); + let (cderi, j2c_decomp) = generate_cderi_with_decomp(mol, aux, j2c_decomp_option, &device); + Self { + mol: mol.clone(), + aux: aux.clone(), + scale_j, + scale_k, + cderi: cderi.into_cow(), + j2c_decomp, + intmd: HashMap::new(), + result: HashMap::new(), + timing: Vec::new(), + is_skeleton_ready: false, + } + } + + pub fn new_with_cderi( + mol: &CInt, + aux: &CInt, + scale_j: f64, + scale_k: f64, + cderi: TsrCow<'a>, + j2c_decomp: J2CDecompose, + ) -> Self { + Self { + mol: mol.clone(), + aux: aux.clone(), + scale_j, + scale_k, + cderi: cderi.into_cow(), + j2c_decomp, + intmd: HashMap::new(), + result: HashMap::new(), + timing: Vec::new(), + is_skeleton_ready: false, + } + } + + /// Build the total density ``D^alpha + D^beta`` from the per-spin mo_coeff / mo_occ. + /// + /// This is only used for an explicit sanity check; the skeleton driver builds the same total + /// density internally when `dm0` is passed as `None`. + #[allow(dead_code)] + fn get_total_dm0(&self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2]) -> Tsr { + let nao = mo_coeff[0].shape()[0]; + let device = self.cderi.device(); + let mut dm0 = rt::zeros(([nao, nao], device)); + for s in 0..2 { + dm0 += get_dm0_restricted(mo_coeff[s].view(), mo_occ[s].view()); + } + dm0 + } + + pub fn ensure_skeleton(&mut self, mo_coeff: &[TsrView; 2], mo_occ: &[TsrView; 2], atm_list: Option<&[usize]>) { + if self.is_skeleton_ready { + return; + } + // nset = 2 (alpha, beta); the driver builds the total density for J internally and produces + // one K output per spin set, sharing every 3c-integral batch. + let mo_coeff_slice: &[TsrView] = &[mo_coeff[0].view(), mo_coeff[1].view()]; + let mo_occ_slice: &[TsrView] = &[mo_occ[0].view(), mo_occ[1].view()]; + let (j_out, k_outs, timing) = get_rijk_skeleton_decomposed_separated( + &self.mol, + &self.aux, + mo_coeff_slice, + mo_occ_slice, + self.cderi.view(), + &self.j2c_decomp, + self.scale_j != 0.0, + self.scale_k != 0.0, + 72, // TODO: batch size `72` should be tunable by max-memory. + atm_list, + None, + ); + self.timing.extend(timing); + + if let Some(j_out) = j_out { + for (key, value) in j_out.into_iter() { + self.intmd.insert(key.to_string(), value); + } + }; + + for (iset, k_out) in k_outs.into_iter().enumerate() { + // note the keys can clash for output of k across spin sets. + // for storage of intermediates, we append `` to the key name. + for (key, value) in k_out.into_iter() { + self.intmd.insert(format!("{key}"), value); + } + } + + self.is_skeleton_ready = true; + } +} + +impl<'a> HessUtilAPI for UHessRIJK<'a> {} + +impl<'a> UHessElecInteractAPI for UHessRIJK<'a> { + fn make_skeleton_hess( + &mut self, + mo_coeff: &[TsrView; 2], + mo_occ: &[TsrView; 2], + atm_list: Option<&[usize]>, + ) -> Tsr { + self.ensure_skeleton(mo_coeff, mo_occ, atm_list); + let intmd = &self.intmd; + + let device = self.cderi.device(); + let natm = atm_list.map_or_else(|| self.mol.natm(), |list| list.len()); + let hess_init = || -> Tsr { rt::zeros(([3, 3, natm, natm], device)) }; + + // helper: sum a set of K keys over both spin channels + let sum_k_keys = |keys: &[&'static str]| -> Tsr { + let mut acc = hess_init(); + for s in 0..2 { + for &key in keys { + acc += &intmd[&format!("{key}")]; + } + } + acc + }; + + let mut de = hess_init(); + if self.scale_j != 0.0 { + let de_J20 = KEYS_J20.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); + let de_J11 = KEYS_J11.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); + let de_J02 = KEYS_J02.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); + let de_J = &de_J20 + &de_J11 + &de_J02; + de += self.scale_j * &de_J; + self.result.insert("de_J20", de_J20); + self.result.insert("de_J11", de_J11); + self.result.insert("de_J02", de_J02); + self.result.insert("de_J", de_J); + } + if self.scale_k != 0.0 { + let de_K20 = sum_k_keys(&KEYS_K20); + let de_K11 = sum_k_keys(&KEYS_K11); + let de_K02 = sum_k_keys(&KEYS_K02); + let de_K = &de_K20 + &de_K11 + &de_K02; + // UHF: K coefficient is -1 (not -0.5 as in RHF) because de_K already includes the spin sum. + de -= self.scale_k * &de_K; + self.result.insert("de_K20", de_K20); + self.result.insert("de_K11", de_K11); + self.result.insert("de_K02", de_K02); + self.result.insert("de_K", de_K); + } + self.result.insert("de_skeleton", de.clone()); + de + } + + fn get_deriv1_ao( + &mut self, + _mo_coeff: &[TsrView; 2], + _mo_occ: &[TsrView; 2], + _atm_list: Option<&[usize]>, + ) -> [Tsr; 2] { + unimplemented!("This function is not implemented for optimized RI-JK hessian. Use `get_deriv1_bra` instead.") + } + + fn get_deriv1_bra( + &mut self, + mo_coeff: &[TsrView; 2], + mo_occ: &[TsrView; 2], + atm_list: Option<&[usize]>, + ) -> [Tsr; 2] { + self.ensure_skeleton(mo_coeff, mo_occ, atm_list); + let intmd = &self.intmd; + + let device = self.cderi.device(); + let natm = atm_list.map_or_else(|| self.mol.natm(), |list| list.len()); + 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 nocc = [occidx[0].iter().filter(|&&x| x).count(), occidx[1].iter().filter(|&&x| x).count()]; + 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 deriv1_ao_init = || -> Tsr { rt::zeros(([nao, nao, 3, natm], device)) }; + let deriv1_bra_init = |s: usize| -> Tsr { rt::zeros(([nao, nocc[s], 3, natm], device)) }; + + let mut deriv1_bra = [deriv1_bra_init(0), deriv1_bra_init(1)]; + + // J is spin-independent (held in AO form, shared across spins); right half-transform per spin. + if self.scale_j != 0.0 { + let j1ao = KEYS_J1AO.iter().map(|&key| &intmd[key]).fold(deriv1_ao_init(), |acc, x| acc + x); + for s in 0..2 { + deriv1_bra[s] += self.scale_j * (&j1ao % &mocc[s]); + } + self.result.insert("j1ao", j1ao); + } + // K is spin-resolved; k1bra^s is stored as the right half-transform ``k1ao^s @ mocc_s`` + // (shape ``[nao, nocc_s, 3, natm]``), matching the restricted optimized convention. + // Note: per-spin k1bra have different `nocc_s` and cannot be summed across spins. + if self.scale_k != 0.0 { + for s in 0..2 { + let ks = KEYS_K1BRA + .iter() + .map(|&key| &intmd[&format!("{key}")]) + .fold(deriv1_bra_init(s), |acc, x| acc + x); + // UHF: no 0.5 factor (occ = 1), unlike RHF. + deriv1_bra[s] -= self.scale_k * &ks; + self.result.insert(if s == 0 { "k1bra_0" } else { "k1bra_1" }, ks); + } + } + self.result.insert("deriv1_bra_0", deriv1_bra[0].clone()); + 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.scale_j != 0.0, + self.scale_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.scale_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.scale_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 7f267c8dd2..c8b28a9fa9 100644 --- a/src/ri_jk/mod.rs +++ b/src/ri_jk/mod.rs @@ -29,6 +29,10 @@ pub mod decompose; pub mod direct; pub mod incore; +// hessian implementations +pub mod hess_r; +pub mod hess_u; + // exports pub use ao2mo::*; pub use decompose::*; diff --git a/src/ri_jk/prelude_dev.rs b/src/ri_jk/prelude_dev.rs index 3820069bec..44cde08bed 100644 --- a/src/ri_jk/prelude_dev.rs +++ b/src/ri_jk/prelude_dev.rs @@ -8,6 +8,7 @@ pub use rayon::prelude::*; pub use rest_libcint::prelude::*; pub use rstsr::prelude::*; pub use rstsr_core::prelude_dev::uninitialized_vec; +pub use std::collections::HashMap; pub use rt::blas::{BlasFloat, LapackDriverAPI}; @@ -20,3 +21,6 @@ pub use crate::molecule_io::Molecule; pub use crate::utilities::memory_batch::*; pub use core::fmt::Write; pub use rest_tensors::{MatrixFull, MatrixUpper}; + +pub use FlagSide::L as Left; +pub use FlagSide::R as Right; diff --git a/src/ri_jk/pure_decompose.rs b/src/ri_jk/pure_decompose.rs index 220e760ba8..d82f042212 100644 --- a/src/ri_jk/pure_decompose.rs +++ b/src/ri_jk/pure_decompose.rs @@ -1,6 +1,9 @@ use super::decompose::*; use super::prelude_dev::*; +use FlagSide::L as Left; +use FlagSide::R as Right; + /// Decompose 2c-2e ERI matrix using eigen decomposition. /// /// Eigenvalues smaller than the threshold will be discarded, and the corresponding eigenvectors @@ -76,6 +79,12 @@ pub fn get_j2c_decomp(mol: &CInt, device: &DeviceBLAS, j2c_decomp_option: J2CDec /// - In some cases where gradient response evaluation is involved, we may need to solve /// the `(J^-1/2)^T * cderi`, which requires `flip_uplo = true` (the `cderi` is already solved). /// This option should not affect eigen decomposition since it's already symmetric. +/// +/// # TODO +/// +/// This function may probably be deprecated, and is better to use [`solve_by_j2c`] and +/// [`solve_by_j2c_mut`] instead. [`solve_by_j2c`] introduces a `side` parameter, which can +/// handle both left and right solve. This function only implements the right solve. pub fn get_solved_j3c(j3c: Tsr, j2c_decomp: &J2CDecompose, flip_uplo: bool) -> Tsr where T: BlasFloat + FromPrimitive + 'static, @@ -142,6 +151,242 @@ where } } +/// Transform 3c-2e ERI (j3c), use solve/inv-matmul to decomposed 3c-2e ERI (cderi). +/// +/// The function name was previously `get_solved_j3c`. However, this function actually works as +/// transformation to auxiliary basis, independent to what physical nature (3c-ERI or 2c-ERI). So +/// finally the name was changed to `solve_by_j2c` to reflect the actual mathematical operation. +/// +/// - `j3c`: The 3c-2e ERI, of shape (..., naux) in column major order if `side = Right`, or (naux, +/// ...) if `side = Left`. +/// - The remaining dimensions should be contiguous if memory and efficiency is of concern. +/// - `j2c_decomp`: The decomposed 2c-2e ERI, either from Cholesky or eigen decomposition. +/// - `side`: The side of the matrix to solve for. +/// - `flip_uplo`: Whether to flip the uplo in computation. Only affects Cholesky decomposition. +/// - Usual j3c solve `J^-1/2 * j3c` should be `flip_uplo = false`. +/// - In some cases where gradient response evaluation is involved, we may need to solve the +/// `(J^-1/2)^T * cderi`, which requires `flip_uplo = true` (the `cderi` is already solved). +/// This option should not affect eigen decomposition since it's already symmetric. +/// +/// # Note on column-major +/// +/// This function does not behave similarly in row-major program (in python with PySCF), compared to +/// the column-major program (in Rust with RSTSR). This is because we handle 3c-2e ERI differently. +/// +/// Following discussion assumes that auxiliary basis is the most dis-contiguous dimension. +/// +/// - Note non-flip behavior is defined by: 3c-2e ERI -> cholesky decomposed ERI. This progress is +/// considered as not flipped. And a further step, cderi -> j2c-inversed ERI, is considered as +/// flipped. +/// - Col-major (REST) should use `(uv|P)` in most cases, so side=Right and uplo=Upper is default. +/// - Row-major (PySCF) should use `(P|uv)` in most cases, so side=Left and uplo=Lower is default. +pub fn solve_by_j2c(mut j3c: Tsr, j2c_decomp: &J2CDecompose, side: FlagSide, flip_uplo: bool) -> Tsr +where + T: BlasFloat + FromPrimitive + 'static, + DeviceBLAS: LapackDriverAPI, +{ + // if fortran-contiguous, we can directly perform every operation in-place + if j3c.f_contig() { + let j3c_mut = j3c.view_mut(); + solve_by_j2c_mut(j3c_mut, j2c_decomp, side, flip_uplo); + return j3c; + } else { + eprintln!("Input j3c is not column-major (Fortran-contiguous). It may cost more memory and time due to explicit transposition.") + } + let naux = match j2c_decomp { + J2CDecompose::Cd { j2c_l, .. } => j2c_l.shape()[0], + J2CDecompose::Eig { j2c_l_inv, .. } => j2c_l_inv.shape()[0], + }; + let j3c_shape = j3c.shape().clone(); + match side { + Left => assert_eq!( + *j3c_shape.first().unwrap(), + naux, + "First dimension of j3c should match the shape of j2c_l (both to be naux)." + ), + Right => assert_eq!( + *j3c_shape.last().unwrap(), + naux, + "Last dimension of j3c should match the shape of j2c_l (both to be naux)." + ), + }; + + match j2c_decomp { + J2CDecompose::Cd { j2c_l, uplo, .. } => { + // cast type anyway, this is not bottleneck + let j2c_l = j2c_l.mapv(|x| T::from_f64(x).unwrap()); + // get j3c shape, and reshape to 2d for triangular solve; + // note we assume j3c to be something similar to (x, x, naux). + let j3c_2d = match side { + Right => j3c.into_shape((-1, naux)), + Left => j3c.into_shape((naux, -1)), + }; + let mut j3c_2d = match (side, uplo, flip_uplo) { + (Right, Upper, false) => rt::linalg::solve_triangular((j2c_l.t(), j3c_2d.into_reverse_axes(), Lower)), + (Right, Lower, false) => rt::linalg::solve_triangular((j2c_l, j3c_2d.into_reverse_axes(), Lower)), + (Right, Upper, true) => rt::linalg::solve_triangular((j2c_l, j3c_2d.into_reverse_axes(), Upper)), + (Right, Lower, true) => rt::linalg::solve_triangular((j2c_l.t(), j3c_2d.into_reverse_axes(), Upper)), + (Left, Upper, false) => rt::linalg::solve_triangular((j2c_l, j3c_2d, Upper)), + (Left, Lower, false) => rt::linalg::solve_triangular((j2c_l.t(), j3c_2d, Upper)), + (Left, Upper, true) => rt::linalg::solve_triangular((j2c_l.t(), j3c_2d, Lower)), + (Left, Lower, true) => rt::linalg::solve_triangular((j2c_l, j3c_2d, Lower)), + }; + if side == Right { + j3c_2d = j3c_2d.into_reverse_axes(); + } + j3c_2d.into_shape(j3c_shape) // reverse back and reshape back + }, + J2CDecompose::Eig { j2c_l_inv, .. } => { + // we need to perform inplace matmul at this case + // however, inplace matmul is not integrated at BLAS level, we need to batch it manually + + // batch size is currently fixed, not related to available menory at this time: + // > max(1/25 remaining size, naux) + // - we assume the API caller leaves at least 4% of memory for storing j3c; + // - we assume size requirement of j2c copy is acceptable. + + // cast type anyway, this is not bottleneck + let j2c_l_inv = j2c_l_inv.mapv(|x| T::from_f64(x).unwrap()); + let device = j3c.device().clone(); + match side { + Right => { + let mut j3c_2d = j3c.into_shape((-1, naux)); // not transposed + let n = j3c_2d.shape()[0]; + // determine batch size + let nbatch = ((n as f64 * 0.04).ceil() as usize).max(naux); + let mut scratch_vec: Vec = vec![T::zero(); nbatch * naux]; + // perform batched inplace-matmul + for start in (0..n).step_by(nbatch) { + let end = (start + nbatch).min(n); + let mut j3c_batch = j3c_2d.i_mut((start..end, ..)); + let mut scratch = rt::asarray((&mut scratch_vec, [end - start, naux].f(), &device)); + scratch.matmul_from(j3c_batch.view(), j2c_l_inv.view(), T::one(), T::zero()); + j3c_batch.assign(&scratch); + } + j3c_2d.into_shape(j3c_shape) // reshape back + }, + Left => { + let mut j3c_2d = j3c.into_shape((naux, -1)); // not transposed + let n = j3c_2d.shape()[1]; + // determine batch size + let nbatch = ((n as f64 * 0.04).ceil() as usize).max(naux); + let mut scratch_vec: Vec = vec![T::zero(); nbatch * naux]; + // perform batched inplace-matmul + for start in (0..n).step_by(nbatch) { + let end = (start + nbatch).min(n); + let mut j3c_batch = j3c_2d.i_mut((.., start..end)); + let mut scratch = rt::asarray((&mut scratch_vec, [naux, end - start].f(), &device)); + scratch.matmul_from(j2c_l_inv.view(), j3c_batch.view(), T::one(), T::zero()); + j3c_batch.assign(&scratch); + } + j3c_2d.into_shape(j3c_shape) // reshape back + }, + } + }, + } +} + +pub fn solve_by_j2c_mut(mut j3c: TsrMut, j2c_decomp: &J2CDecompose, side: FlagSide, flip_uplo: bool) +where + T: BlasFloat + FromPrimitive + 'static, + DeviceBLAS: LapackDriverAPI, +{ + if !j3c.f_contig() { + panic!("Input j3c must be column-major (Fortran-contiguous) for in-place solve_by_j2c_mut.") + } + let device = j3c.device().clone(); + let j3c_shape = j3c.shape().clone(); + let naux = match j2c_decomp { + J2CDecompose::Cd { j2c_l, .. } => j2c_l.shape()[0], + J2CDecompose::Eig { j2c_l_inv, .. } => j2c_l_inv.shape()[0], + }; + match side { + Left => assert_eq!( + *j3c_shape.first().unwrap(), + naux, + "First dimension of j3c should match the shape of j2c_l (both to be naux)." + ), + Right => assert_eq!( + *j3c_shape.last().unwrap(), + naux, + "Last dimension of j3c should match the shape of j2c_l (both to be naux)." + ), + }; + + match j2c_decomp { + J2CDecompose::Cd { j2c_l, uplo, .. } => { + // cast type anyway, this is not bottleneck + let j2c_l = j2c_l.mapv(|x| T::from_f64(x).unwrap()); + // get j3c shape, and reshape to 2d for triangular solve; + // note we assume j3c to be something similar to (x, x, naux). + let n = j3c.size() / naux; + let j3c_offset = j3c.offset(); + let j3c_raw = &mut j3c.raw_mut()[j3c_offset..]; + let j3c_2d = match side { + Left => rt::asarray((j3c_raw, [naux, n].f(), &device)), + Right => rt::asarray((j3c_raw, [n, naux].f(), &device)), + }; + match (side, uplo, flip_uplo) { + (Right, Upper, false) => rt::linalg::solve_triangular((j2c_l.t(), j3c_2d.into_reverse_axes(), Lower)), + (Right, Lower, false) => rt::linalg::solve_triangular((j2c_l, j3c_2d.into_reverse_axes(), Lower)), + (Right, Upper, true) => rt::linalg::solve_triangular((j2c_l, j3c_2d.into_reverse_axes(), Upper)), + (Right, Lower, true) => rt::linalg::solve_triangular((j2c_l.t(), j3c_2d.into_reverse_axes(), Upper)), + (Left, Upper, false) => rt::linalg::solve_triangular((j2c_l, j3c_2d, Upper)), + (Left, Lower, false) => rt::linalg::solve_triangular((j2c_l.t(), j3c_2d, Upper)), + (Left, Upper, true) => rt::linalg::solve_triangular((j2c_l.t(), j3c_2d, Lower)), + (Left, Lower, true) => rt::linalg::solve_triangular((j2c_l, j3c_2d, Lower)), + }; + }, + J2CDecompose::Eig { j2c_l_inv, .. } => { + // we need to perform inplace matmul at this case + // however, inplace matmul is not integrated at BLAS level, we need to batch it manually + + // batch size is currently fixed, not related to available menory at this time: + // > max(1/25 remaining size, naux) + // - we assume the API caller leaves at least 4% of memory for storing j3c; + // - we assume size requirement of j2c copy is acceptable. + + // cast type anyway, this is not bottleneck + let j2c_l_inv = j2c_l_inv.mapv(|x| T::from_f64(x).unwrap()); + // get j3c shape, and reshape to 2d for matmul; + // note we assume j3c to be something similar to (x, x, naux). + let n = j3c.size() / naux; + let j3c_offset = j3c.offset(); + let j3c_raw = &mut j3c.raw_mut()[j3c_offset..]; + match side { + Right => { + let mut j3c_2d = rt::asarray((j3c_raw, [n, naux].f(), &device)); + // determine batch size + let nbatch = ((n as f64 * 0.04).ceil() as usize).max(naux); + let mut scratch_vec: Vec = vec![T::zero(); nbatch * naux]; + // perform batched inplace-matmul + for start in (0..n).step_by(nbatch) { + let end = (start + nbatch).min(n); + let mut j3c_batch = j3c_2d.i_mut((start..end, ..)); + let mut scratch = rt::asarray((&mut scratch_vec, [end - start, naux].f(), &device)); + scratch.matmul_from(j3c_batch.view(), j2c_l_inv.view(), T::one(), T::zero()); + j3c_batch.assign(&scratch); + } + }, + Left => { + let mut j3c_2d = rt::asarray((j3c_raw, [naux, n].f(), &device)); + // determine batch size + let nbatch = ((n as f64 * 0.04).ceil() as usize).max(naux); + let mut scratch_vec: Vec = vec![T::zero(); nbatch * naux]; + // perform batched inplace-matmul + for start in (0..n).step_by(nbatch) { + let end = (start + nbatch).min(n); + let mut j3c_batch = j3c_2d.i_mut((.., start..end)); + let mut scratch = rt::asarray((&mut scratch_vec, [naux, end - start].f(), &device)); + scratch.matmul_from(j2c_l_inv.view(), j3c_batch.view(), T::one(), T::zero()); + j3c_batch.assign(&scratch); + } + }, + } + }, + } +} + #[test] fn test_check_recoverable() { // this will just test if we can recover the original matrix by the decomposed factors. diff --git a/src/utilities/rstsr_util.rs b/src/utilities/rstsr_util.rs index 46b41b8f8b..414282f9cb 100644 --- a/src/utilities/rstsr_util.rs +++ b/src/utilities/rstsr_util.rs @@ -1,25 +1,18 @@ //! rest_tensor to RSTSR interchange. -use rayon::prelude::*; use rstsr::prelude::*; +use rayon::prelude::*; use rstsr_core::{prelude_dev::OpAssignAPI, storage::creation::DeviceCreationAnyAPI}; use tensors::{BasicMatrix, MatrixFull}; -pub mod prelude { - pub use rstsr::prelude::*; - - pub type Tsr = Tensor; - pub type TsrView<'a, T = f64> = TensorView<'a, T, DeviceBLAS, IxD>; - pub type TsrMut<'a, T = f64> = TensorMut<'a, T, DeviceBLAS, IxD>; - pub type TsrCow<'a, T = f64> = TensorCow<'a, T, DeviceBLAS, IxD>; -} +pub type Tsr = Tensor; +pub type TsrView<'a, T = f64> = TensorView<'a, T, DeviceBLAS, IxD>; +pub type TsrMut<'a, T = f64> = TensorMut<'a, T, DeviceBLAS, IxD>; +pub type TsrCow<'a, T = f64> = TensorCow<'a, T, DeviceBLAS, IxD>; /* #region interchange between rstsr and rest_tensor */ // In REST, we always use DeviceBLAS as backend in most cases. -pub type Tsr = Tensor; -pub type TsrView<'a, T> = TensorView<'a, T, DeviceBLAS, IxD>; -pub type TsrMut<'a, T> = TensorMut<'a, T, DeviceBLAS, IxD>; pub trait RestTensorToRstsrTsrAPI { fn to_rstsr(&self, device: &DeviceBLAS) -> Tsr; -- Gitee From afdfc7a921b401089b461b126144e5ec1be82466 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 12:28:27 +0800 Subject: [PATCH 03/39] interface: rhf/rks completed --- src/dft/numint_matmul/hess_rks.rs | 6 +- src/grad/rhf.rs | 27 ++++++-- src/hessian_backup/config.rs | 35 +++++++++++ src/hessian_backup/mod.rs | 23 ++++++- src/hessian_backup/rscf.rs | 36 +++-------- src/hessian_backup/rscf_interface.rs | 93 ++++++++++++++++++++++++++++ src/hessian_backup/uscf.rs | 13 ++-- src/ri_jk/hess_r.rs | 68 ++++++++++---------- src/ri_jk/hess_u.rs | 50 +++++++-------- tests/hessian_backup/mod.rs | 2 + tests/hessian_backup/rhf.rs | 47 ++++++++++++++ tests/hessian_backup/rks_b3lyp.rs | 47 ++++++++++++++ tests/test_hessian_backup.rs | 1 + 13 files changed, 344 insertions(+), 104 deletions(-) create mode 100644 src/hessian_backup/config.rs create mode 100644 src/hessian_backup/rscf_interface.rs create mode 100644 tests/hessian_backup/mod.rs create mode 100644 tests/hessian_backup/rhf.rs create mode 100644 tests/hessian_backup/rks_b3lyp.rs create mode 100644 tests/test_hessian_backup.rs diff --git a/src/dft/numint_matmul/hess_rks.rs b/src/dft/numint_matmul/hess_rks.rs index 556ddd9dfa..836894b1c4 100644 --- a/src/dft/numint_matmul/hess_rks.rs +++ b/src/dft/numint_matmul/hess_rks.rs @@ -764,7 +764,7 @@ pub fn get_rks_response_bra_batched( pub struct RHessKSNIMatmul<'a> { pub mol: CInt, - pub xc_func_list: &'a [(f64, LibXCFunctional)], + pub xc_func_list: Vec<(f64, LibXCFunctional)>, pub ni: NIMatmul<'a>, pub ni_cpks: Option>, pub verbose: bool, @@ -773,7 +773,7 @@ pub struct RHessKSNIMatmul<'a> { } impl<'a> RHessKSNIMatmul<'a> { - pub fn new(mol: &CInt, xc_func_list: &'a [(f64, LibXCFunctional)], ni: NIMatmul<'a>, verbose: bool) -> Self { + pub fn new(mol: &CInt, xc_func_list: Vec<(f64, LibXCFunctional)>, ni: NIMatmul<'a>, verbose: bool) -> Self { Self { mol: mol.clone(), xc_func_list, @@ -800,7 +800,7 @@ impl<'a> RHessKSNIMatmul<'a> { // run RKS hessian setup let dm0 = get_dm0_restricted(mo_coeff, mo_occ); let (result, _timing) = - make_hessian_setup_batched(&self.mol, self.xc_func_list, &mut self.ni, dm0.view(), atm_list, self.verbose); + make_hessian_setup_batched(&self.mol, &self.xc_func_list, &mut self.ni, dm0.view(), atm_list, self.verbose); // handling intermediates and results for (key, val) in result.into_iter() { diff --git a/src/grad/rhf.rs b/src/grad/rhf.rs index 3f4d060d2f..c73a52cc4a 100644 --- a/src/grad/rhf.rs +++ b/src/grad/rhf.rs @@ -820,9 +820,8 @@ pub fn calc_de_nuc(mol: &Molecule) -> MatrixFull { return de_nuc; } -pub fn pack_triu_tilde(dm: TsrView) -> Tsr { - // Pack the lower triangular part of a matrix into a 1D array - // and non-diagonal values are multiplied by 2. +/// Pack the lower triangular part of a matrix into a 1D array, multiply non-diagonal values by 2. +pub fn pack_triu_tilde_2d(dm: TsrView) -> Tsr { assert_eq!(dm.ndim(), 2); assert_eq!(dm.shape()[0], dm.shape()[1]); let nao = dm.shape()[0]; @@ -830,7 +829,27 @@ pub fn pack_triu_tilde(dm: TsrView) -> Tsr { for i in 0..nao { dm_triu[[(i + 2) * (i + 1) / 2 - 1]] *= 0.5; } - return dm_triu; + dm_triu +} + +/// Pack the lower triangular part of a multi-dimensional array into a smaller-one dimension array, +/// multiply non-diagonal values by 2. +pub fn pack_triu_tilde(dm: TsrView) -> Tsr { + if dm.ndim() == 2 { + return pack_triu_tilde_2d(dm); + } + assert!(dm.ndim() > 2); + assert_eq!(dm.shape()[0], dm.shape()[1]); + let shape_remaining = &dm.shape()[2..]; + let nao = dm.shape()[0]; + let nao_tp = nao * (nao + 1) / 2; + let dm = dm.reshape((nao, nao, -1)); + let mut out = rt::zeros(([nao_tp, dm.shape()[2]], dm.device())); + for (i, dm_i) in dm.axes_iter(-1).enumerate() { + out.i_mut((.., i)).assign(&pack_triu_tilde_2d(dm_i)); + } + let shape_recap = [nao_tp].iter().chain(shape_remaining.iter()).copied().collect::>(); + out.into_shape(shape_recap) } pub fn get_dme0(mo_coeff: TsrView, mo_occ: TsrView, mo_energy: TsrView) -> Tsr { diff --git a/src/hessian_backup/config.rs b/src/hessian_backup/config.rs new file mode 100644 index 0000000000..d942c261cc --- /dev/null +++ b/src/hessian_backup/config.rs @@ -0,0 +1,35 @@ +use serde::{Deserialize, Serialize}; +use serde_inline_default::serde_inline_default; + +#[serde_inline_default] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct HessSCFConfig { + #[serde_inline_default(0.0)] + pub level_shift: f64, + #[serde_inline_default(1e-8)] + pub cphf_tol: f64, + #[serde_inline_default(42)] + pub cphf_max_cycle: usize, + #[serde_inline_default(14)] + pub cphf_max_space: usize, + #[serde_inline_default(1e-14)] + pub cphf_lindep: f64, + #[serde_inline_default(None)] + pub verbose: Option, + #[serde_inline_default(None)] + pub atm_list: Option>, +} + +impl Default for HessSCFConfig { + fn default() -> Self { + Self { + level_shift: 0.0, + cphf_tol: 1e-8, + cphf_max_cycle: 42, + cphf_max_space: 14, + cphf_lindep: 1e-14, + verbose: None, + atm_list: None, + } + } +} diff --git a/src/hessian_backup/mod.rs b/src/hessian_backup/mod.rs index 5cff5b9653..d5650ae5ef 100644 --- a/src/hessian_backup/mod.rs +++ b/src/hessian_backup/mod.rs @@ -1,3 +1,17 @@ +//! Hessian module for REST (currently as backup implementation at PR!101). +//! +//! This module should work in most cases, but is not completely finished and tested. +//! +//! This module does not contain extensive difficult implementation. We defined traits, +//! some common implementations (hcore, nuc, ovlp), total hessian, 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 handle interface to REST before next pull request, in this module. +//! +//! This module currently does not handle post-SCF derivatives. + // trait definitions pub mod trait_rhess; pub mod trait_uhess; @@ -14,22 +28,26 @@ pub mod ovlp; pub mod rscf; pub mod uscf; +// total hess interface to REST +pub mod rscf_interface; + // vibrational analysis pub mod vib; // utilities pub mod cint_handling; +pub mod config; pub mod krylov_block; #[allow(unused_imports)] pub mod prelude { use super::*; + pub use config::HessSCFConfig; pub use hcore::{RHessHcore, UHessHcore}; - pub use krylov_block::krylov_block; pub use nuc_repl::HessNucRepl; pub use ovlp::{RHessOvlp, UHessOvlp}; - pub use rscf::{HessSCFConfig, RHessSCF}; + pub use rscf::RHessSCF; pub use trait_rhess::{HessNucAPI, RHessCoreAPI, RHessElecInteractAPI}; pub use trait_uhess::{UHessCoreAPI, UHessElecInteractAPI}; pub use trait_util::HessUtilAPI; @@ -39,6 +57,7 @@ pub mod prelude { pub(super) use crate::utilities::rstsr_util::*; pub(super) use cint_handling::*; pub(super) use itertools::Itertools; + pub(super) use krylov_block::krylov_block; pub(super) use rayon::prelude::*; pub(super) use rest_libcint::prelude::*; pub(super) use rstsr::prelude::*; diff --git a/src/hessian_backup/rscf.rs b/src/hessian_backup/rscf.rs index 456b178348..a06daf2d9d 100644 --- a/src/hessian_backup/rscf.rs +++ b/src/hessian_backup/rscf.rs @@ -2,32 +2,16 @@ use super::prelude::*; -#[derive(Debug, Clone, PartialEq)] -pub struct HessSCFConfig { - pub level_shift: f64, - pub cphf_tol: f64, - pub cphf_max_cycle: usize, - pub cphf_max_space: usize, - pub cphf_lindep: f64, -} - -impl Default for HessSCFConfig { - fn default() -> Self { - Self { level_shift: 0.0, cphf_tol: 1e-8, cphf_max_cycle: 42, cphf_max_space: 14, cphf_lindep: 1e-14 } - } -} - /// Working solver and maintainer of all hessian components for restricted SCF method. pub struct RHessSCF<'a> { pub mo_coeff: Tsr, pub mo_occ: Tsr, pub mo_energy: Tsr, - pub ovlp_obj: RHessOvlp, + pub ovlp_obj: &'a mut RHessOvlp, pub nuc_list: Vec<&'a mut dyn HessNucAPI>, pub core_list: Vec<&'a mut dyn RHessCoreAPI>, pub el_list: Vec<&'a mut dyn RHessElecInteractAPI>, pub config: HessSCFConfig, - pub atm_list: Option>, pub result: HashMap, /// Timing information. Represented by wall time in second. pub timing: Vec<(String, f64)>, @@ -39,12 +23,11 @@ impl<'a> RHessSCF<'a> { mo_coeff: Tsr, mo_occ: Tsr, mo_energy: Tsr, - ovlp_obj: RHessOvlp, + ovlp_obj: &'a mut RHessOvlp, nuc_list: Vec<&'a mut dyn HessNucAPI>, core_list: Vec<&'a mut dyn RHessCoreAPI>, el_list: Vec<&'a mut dyn RHessElecInteractAPI>, - config: HessSCFConfig, - atm_list: Option<&[usize]>, + config: &HessSCFConfig, ) -> Self { Self { mo_coeff, @@ -54,8 +37,7 @@ impl<'a> RHessSCF<'a> { nuc_list, core_list, el_list, - config, - atm_list: atm_list.map(|x| x.to_vec()), + config: config.clone(), result: HashMap::new(), timing: Vec::new(), } @@ -64,7 +46,7 @@ impl<'a> RHessSCF<'a> { /// Number of atoms over which the Hessian is computed. This is `atm_list.len()` if /// `atm_list` is `Some`, otherwise the total number of atoms in the molecule. pub fn natm(&self) -> usize { - match &self.atm_list { + match &self.config.atm_list { Some(list) => list.len(), None => self.ovlp_obj.natm(), } @@ -73,7 +55,7 @@ impl<'a> RHessSCF<'a> { /// Return the list of (global) atom indices the Hessian is computed for, ordered the same /// way as the local indexing used in the returned Hessian. pub fn atm_indices(&self) -> Vec { - match &self.atm_list { + match &self.config.atm_list { Some(list) => list.clone(), None => (0..self.ovlp_obj.natm()).collect(), } @@ -119,7 +101,7 @@ impl<'a> RHessSCF<'a> { let nocc = occidx.iter().filter(|&&x| x).count(); let natm = self.natm(); let atm_indices = self.atm_indices(); - let atm_list = self.atm_list.as_deref(); + let atm_list = self.config.atm_list.as_deref(); let e_ai = evir.i((.., None)) - eocc.i((None, ..)); let e_ai_shift = &e_ai + level_shift; @@ -442,7 +424,7 @@ impl<'a> RHessSCF<'a> { let natm = self.natm(); let mo_coeff = self.mo_coeff.view(); let mo_occ = self.mo_occ.view(); - let atm_list = self.atm_list.as_deref(); + let atm_list = self.config.atm_list.as_deref(); let device = self.mo_coeff.device().clone(); let mut de_skeleton = rt::zeros(([3, 3, natm, natm], &device)); @@ -484,7 +466,7 @@ impl<'a> RHessSCF<'a> { let mo_occ = self.mo_occ.view(); let mo_energy = self.mo_energy.view(); let dme0 = get_dme0_restricted(mo_coeff, mo_occ, mo_energy); - let atm_list = self.atm_list.clone(); + let atm_list = self.config.atm_list.clone(); let de_skeleton = self.make_skeleton_hess(); diff --git a/src/hessian_backup/rscf_interface.rs b/src/hessian_backup/rscf_interface.rs new file mode 100644 index 0000000000..37eb321429 --- /dev/null +++ b/src/hessian_backup/rscf_interface.rs @@ -0,0 +1,93 @@ +use super::prelude::*; +use crate::dft::numint_matmul::hess_rks::RHessKSNIMatmul; +use crate::dft::numint_matmul::nimatmul::NIMatmul; +use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; +use crate::SCF; + +pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) { + let device = DeviceBLAS::default(); + + // --- basic preparation --- // + let mo_coeff = { + let mo_coeff = &scf_data.eigenvectors[0]; + rt::asarray((&mo_coeff.data, mo_coeff.size, &device)).into_contig(ColMajor) + }; + let mo_occ = { + let mo_occ = &scf_data.occupation[0]; + rt::asarray((mo_occ, [mo_occ.len()], &device)).into_contig(ColMajor) + }; + let mo_energy = { + let mo_energy = &scf_data.eigenvalues[0]; + rt::asarray((mo_energy, [mo_energy.len()], &device)).into_contig(ColMajor) + }; + + let mol_obj = &scf_data.mol; + let mol = get_cint_mol(mol_obj); + let aux = get_cint_aux(mol_obj); + + let mut hess_ovlp_obj = RHessOvlp::new(&mol, &device); + let mut hess_nuc_repl_obj = HessNucRepl::new(&mol, &device); + let mut hess_hcore_obj = RHessHcore::new(&mol, &device); + + let hess_nuc_list: Vec<&mut dyn HessNucAPI> = vec![&mut hess_nuc_repl_obj]; + let hess_hcore_list: Vec<&mut dyn RHessCoreAPI> = vec![&mut hess_hcore_obj]; + let mut hess_el_list: Vec<&mut dyn RHessElecInteractAPI> = Vec::new(); + + // --- RI-JK --- // + + use crate::ri_jk::hess_r::RHessRIJK; + + let is_hf = scf_data.mol.xc_data.dfa_compnt_scf.is_empty(); + let scale_j = 1.0; + let scale_k = match is_hf { + true => 1.0, + false => scf_data.mol.xc_data.dfa_hybrid_scf, + }; + let j2c_decomp_option = &scf_data.mol.ctrl.j2c_decomp; + let j2c_decomp = crate::ri_jk::get_j2c_decomp(&aux, &device, *j2c_decomp_option); + + let mut hess_rijk_obj = if let Some((rimatr, _, _)) = &scf_data.rimatr { + let cderi = rimatr.to_rstsr_view(&device).into_cow(); + RHessRIJK::new_with_cderi(&mol, &aux, scale_j, scale_k, cderi, j2c_decomp) + } else { + panic!( + "This implementation requires cholesky decomposed ERI (or rimatr) to be available and stored in memory." + ); + }; + + hess_el_list.push(&mut hess_rijk_obj); + + // --- DFT --- // + + let mut hess_nimatmul_obj = (!is_hf).then(|| { + use crate::dft::numint_matmul::prelude::*; + use libxc::prelude::*; + + let grid_coords = &scf_data.grids.as_ref().unwrap().coordinates; + let grid_weights = &scf_data.grids.as_ref().unwrap().weights; + let ni = NIMatmul::new(&mol, grid_coords, grid_weights); + + 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::Unpolarized))) + .collect_vec() + }; + let verbose = scf_data.mol.ctrl.print_level >= 2; + RHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose) + }); + if let Some(ref mut hess_nimatmul_obj) =hess_nimatmul_obj { + hess_el_list.push(hess_nimatmul_obj); + } + + // --- run hessian --- // + + let mut hess_scf = + RHessSCF::new(mo_coeff, mo_occ, mo_energy, &mut hess_ovlp_obj, hess_nuc_list, hess_hcore_list, hess_el_list, config); + + let de_hess = hess_scf.make_hess(); + println!("=== HESSIAN ===\n{:12.6}", de_hess.t()) +} diff --git a/src/hessian_backup/uscf.rs b/src/hessian_backup/uscf.rs index 13357e08e2..066e5ca311 100644 --- a/src/hessian_backup/uscf.rs +++ b/src/hessian_backup/uscf.rs @@ -11,7 +11,6 @@ pub struct UHessSCF<'a> { pub core_list: Vec<&'a mut dyn UHessCoreAPI>, pub el_list: Vec<&'a mut dyn UHessElecInteractAPI>, pub config: HessSCFConfig, - pub atm_list: Option>, pub result: HashMap, /// Timing information. Represented by wall time in second. pub timing: Vec<(String, f64)>, @@ -28,7 +27,6 @@ impl<'a> UHessSCF<'a> { core_list: Vec<&'a mut dyn UHessCoreAPI>, el_list: Vec<&'a mut dyn UHessElecInteractAPI>, config: HessSCFConfig, - atm_list: Option>, ) -> Self { Self { mo_coeff, @@ -39,7 +37,6 @@ impl<'a> UHessSCF<'a> { core_list, el_list, config, - atm_list, result: HashMap::new(), timing: Vec::new(), } @@ -48,7 +45,7 @@ impl<'a> UHessSCF<'a> { /// Number of atoms over which the Hessian is computed. This is `atm_list.len()` if /// `atm_list` is `Some`, otherwise the total number of atoms in the molecule. pub fn natm(&self) -> usize { - match &self.atm_list { + match &self.config.atm_list { Some(list) => list.len(), None => self.ovlp_obj.natm(), } @@ -57,7 +54,7 @@ impl<'a> UHessSCF<'a> { /// Return the list of (global) atom indices the Hessian is computed for, ordered the same /// way as the local indexing used in the returned Hessian. pub fn atm_indices(&self) -> Vec { - match &self.atm_list { + match &self.config.atm_list { Some(list) => list.clone(), None => (0..self.ovlp_obj.natm()).collect(), } @@ -95,7 +92,7 @@ impl<'a> UHessSCF<'a> { let nocc = [mocc[α].shape()[1], mocc[β].shape()[1]]; let natm = self.natm(); let atm_indices = self.atm_indices(); - let atm_list = self.atm_list.as_deref(); + let atm_list = self.config.atm_list.as_deref(); let e_ai = [evir[α].i((.., None)) - eocc[α].i((None, ..)), evir[β].i((.., None)) - eocc[β].i((None, ..))]; let e_ai_shift = [&e_ai[α] + level_shift, &e_ai[β] + level_shift]; @@ -475,7 +472,7 @@ impl<'a> UHessSCF<'a> { let natm = self.natm(); let mo_coeff = [self.mo_coeff[α].view(), self.mo_coeff[β].view()]; let mo_occ = [self.mo_occ[α].view(), self.mo_occ[β].view()]; - let atm_list = self.atm_list.as_deref(); + let atm_list = self.config.atm_list.as_deref(); let device = self.mo_coeff[α].device().clone(); let mut de_skeleton = rt::zeros(([3, 3, natm, natm], &device)); @@ -518,7 +515,7 @@ impl<'a> UHessSCF<'a> { get_dme0_restricted(self.mo_coeff[α].view(), self.mo_occ[α].view(), self.mo_energy[α].view()), get_dme0_restricted(self.mo_coeff[β].view(), self.mo_occ[β].view(), self.mo_energy[β].view()), ]; - let atm_list = self.atm_list.clone(); + let atm_list = self.config.atm_list.clone(); let de_skeleton = self.make_skeleton_hess(); diff --git a/src/ri_jk/hess_r.rs b/src/ri_jk/hess_r.rs index 1e2f333b72..e0fd92be6f 100644 --- a/src/ri_jk/hess_r.rs +++ b/src/ri_jk/hess_r.rs @@ -65,10 +65,10 @@ pub const KEYS_K1BRA: [&str; 8] = [ /// - `j_ao`: `Option` of shape `[nao, nao, nprop]` — the **spin-independent** Coulomb response /// operator in AO basis, built from the total density response `sum_s bra_s @ mocc_s.T` (already /// carrying the internal factor `2.0` from the symmetric cderi contraction; the consumer applies -/// `scale_j` and the per-spin right half-transform `... @ mocc_s`). `None` if `do_j` is false. +/// `factor_j` and the per-spin right half-transform `... @ mocc_s`). `None` if `do_j` is false. /// - `k_bras`: `Vec` (one entry per spin) of shape `[nao, nocc_s, nprop]` — the same-spin /// exchange response in bra form (already carrying its internal sign/scale; the consumer applies -/// `scale_k`). Empty if `do_k` is false. +/// `factor_k`). Empty if `do_k` is false. /// /// # Convention notes /// @@ -76,7 +76,7 @@ pub const KEYS_K1BRA: [&str; 8] = [ /// spins; this is why UHF can reuse the RHF J path verbatim. /// - K is strictly same-spin; each spin's bra form is produced independently. /// - The internal factors (`2.0` on J, the two-term symmetrized sum on K) match the existing RHF -/// optimized response; the per-method `scale_j` / `scale_k` and the RHF `0.5` vs UHF `1.0` +/// optimized response; the per-method `factor_j` / `factor_k` and the RHF `0.5` vs UHF `1.0` /// exchange prefactor are applied by the consumer, not here. #[allow(clippy::too_many_arguments)] pub fn get_rijk_response_bra_separated( @@ -1737,8 +1737,8 @@ pub fn generate_cderi_with_decomp( pub struct RHessRIJK<'a> { pub mol: CInt, pub aux: CInt, - pub scale_j: f64, - pub scale_k: f64, + pub factor_j: f64, + pub factor_k: f64, pub cderi: TsrCow<'a>, pub j2c_decomp: J2CDecompose, pub intmd: HashMap, // intermediates @@ -1748,20 +1748,18 @@ pub struct RHessRIJK<'a> { } impl<'a> RHessRIJK<'a> { - pub fn new_without_cderi(mol: &CInt, aux: &CInt, scale_j: f64, scale_k: f64) -> Self { - let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(1e-14), uplo: Upper }; - // note: the following two options are also valid - // let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(1e-14), - // uplo: Lower }; - // let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Eig, threshold: Some(1e-14), - // uplo: Upper }; + pub fn new_without_cderi(mol: &CInt, aux: &CInt, factor_j: f64, factor_k: f64, j2c_decomp_option: J2CDecompOption) -> Self { + // note: following options should be valid + // let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(1e-14), uplo: Upper }; + // let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(1e-14), uplo: Lower }; + // let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Eig, threshold: Some(1e-14), uplo: Upper }; let device = DeviceBLAS::default(); let (cderi, j2c_decomp) = generate_cderi_with_decomp(mol, aux, j2c_decomp_option, &device); Self { mol: mol.clone(), aux: aux.clone(), - scale_j, - scale_k, + factor_j, + factor_k, cderi: cderi.into_cow(), j2c_decomp, intmd: HashMap::new(), @@ -1774,16 +1772,16 @@ impl<'a> RHessRIJK<'a> { pub fn new_with_cderi( mol: &CInt, aux: &CInt, - scale_j: f64, - scale_k: f64, + factor_j: f64, + factor_k: f64, cderi: TsrCow<'a>, j2c_decomp: J2CDecompose, ) -> Self { Self { mol: mol.clone(), aux: aux.clone(), - scale_j, - scale_k, + factor_j, + factor_k, cderi: cderi.into_cow(), j2c_decomp, intmd: HashMap::new(), @@ -1804,8 +1802,8 @@ impl<'a> RHessRIJK<'a> { &[mo_occ], self.cderi.view(), &self.j2c_decomp, - self.scale_j != 0.0, - self.scale_k != 0.0, + self.factor_j != 0.0, + self.factor_k != 0.0, 72, // TODO: batch size `72` should be tunable by max-memory. atm_list, None, @@ -1842,18 +1840,18 @@ impl<'a> RHessElecInteractAPI for RHessRIJK<'a> { let hess_init = || -> Tsr { rt::zeros(([3, 3, natm, natm], device)) }; let mut de = hess_init(); - if self.scale_j != 0.0 { + if self.factor_j != 0.0 { let de_J20 = KEYS_J20.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); let de_J11 = KEYS_J11.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); let de_J02 = KEYS_J02.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); let de_J = &de_J20 + &de_J11 + &de_J02; - de += self.scale_j * &de_J; + de += self.factor_j * &de_J; self.result.insert("de_J20", de_J20); self.result.insert("de_J11", de_J11); self.result.insert("de_J02", de_J02); self.result.insert("de_J", de_J); } - if self.scale_k != 0.0 { + if self.factor_k != 0.0 { // rhf only have one spin let de_K20 = KEYS_K20.iter().map(|&key| &intmd[&format!("{key}")]).fold(hess_init(), |acc, x| acc + x); @@ -1862,7 +1860,7 @@ impl<'a> RHessElecInteractAPI for RHessRIJK<'a> { let de_K02 = KEYS_K02.iter().map(|&key| &intmd[&format!("{key}")]).fold(hess_init(), |acc, x| acc + x); let de_K = &de_K20 + &de_K11 + &de_K02; - de -= 0.5 * self.scale_k * &de_K; + de -= 0.5 * self.factor_k * &de_K; self.result.insert("de_K20", de_K20); self.result.insert("de_K11", de_K11); self.result.insert("de_K02", de_K02); @@ -1891,17 +1889,17 @@ impl<'a> RHessElecInteractAPI for RHessRIJK<'a> { let deriv1_bra_init = || -> Tsr { rt::zeros(([nao, nocc, 3, natm], device)) }; let mut deriv1_bra = deriv1_bra_init(); - if self.scale_j != 0.0 { + if self.factor_j != 0.0 { let j1ao = KEYS_J1AO.iter().map(|&key| &intmd[key]).fold(deriv1_ao_init(), |acc, x| acc + x); - deriv1_bra += self.scale_j * (&j1ao % &mocc); + deriv1_bra += self.factor_j * (&j1ao % &mocc); self.result.insert("j1ao", j1ao); } - if self.scale_k != 0.0 { + if self.factor_k != 0.0 { let k1bra = KEYS_K1BRA .iter() .map(|&key| &intmd[&format!("{key}")]) .fold(deriv1_bra_init(), |acc, x| acc + x); - deriv1_bra -= 0.5 * self.scale_k * &k1bra; + deriv1_bra -= 0.5 * self.factor_k * &k1bra; self.result.insert("k1bra", k1bra); } self.result.insert("deriv1_bra", deriv1_bra.clone()); @@ -1919,9 +1917,9 @@ impl<'a> RHessElecInteractAPI for RHessRIJK<'a> { let cderi = self.cderi.view(); // RHF (single spin) assembly of the separated J/K response core. - // - J (AO form, from total density) contracted with `mocc` and scaled by `scale_j`. - // - K (same-spin bra form) scaled by `scale_k`; the core already bakes in the exchange sign. - // - RHF exchange prefactor (occ = 2) is folded into `scale_k`, matching the naive convention. + // - J (AO form, from total density) contracted with `mocc` and scaled by `factor_j`. + // - K (same-spin bra form) scaled by `factor_k`; the core already bakes in the exchange sign. + // - RHF exchange prefactor (occ = 2) is folded into `factor_k`, matching the naive convention. let shape_bra = bra.shape().to_vec(); let nao = mo_coeff.shape()[0]; let device = mo_coeff.device(); @@ -1936,18 +1934,18 @@ impl<'a> RHessElecInteractAPI for RHessRIJK<'a> { &[mo_coeff.view()], &[mo_occ.view()], &[bra.view()], - self.scale_j != 0.0, - self.scale_k != 0.0, + self.factor_j != 0.0, + self.factor_k != 0.0, 72, ); let mut resp: Tsr = rt::zeros(([nao, nocc, nprop], device)); if let Some(resp_ao_j) = j_ao { - resp += self.scale_j * (resp_ao_j % &mocc); + resp += self.factor_j * (resp_ao_j % &mocc); } if let Some(k_bra) = k_bras.first() { // K bra is returned in the original trailing shape; flatten trailing dims to (nao, nocc, nprop). - resp += self.scale_k * k_bra.view().reshape((nao, nocc, nprop)); + resp += self.factor_k * k_bra.view().reshape((nao, nocc, nprop)); } resp.into_shape(shape_bra) } diff --git a/src/ri_jk/hess_u.rs b/src/ri_jk/hess_u.rs index 49ec9293f5..13fa7fa2e1 100644 --- a/src/ri_jk/hess_u.rs +++ b/src/ri_jk/hess_u.rs @@ -10,10 +10,10 @@ //! orbitals (UHF occ = 1, so ``mocc_2 = mocc``). //! //! Scaling difference to the restricted ([`crate::ri_jk::hess_r::RHessRIJK`]) counterpart: -//! - Skeleton: ``scale_j * de_J - scale_k * de_K`` (K coefficient ``-1``, not ``-0.5`` as in RHF), +//! - 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`]). -//! - First derivative (bra form): ``scale_j * (j1ao @ mocc_s) - scale_k * k1bra_s`` per spin (no +//! - 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. //! //! The response (`get_response_bra`) reuses the separated J/K response core @@ -36,8 +36,8 @@ use crate::ri_jk::hess_r::{ pub struct UHessRIJK<'a> { pub mol: CInt, pub aux: CInt, - pub scale_j: f64, - pub scale_k: f64, + pub factor_j: f64, + pub factor_k: f64, pub cderi: TsrCow<'a>, pub j2c_decomp: J2CDecompose, pub intmd: HashMap, // intermediates @@ -47,15 +47,15 @@ pub struct UHessRIJK<'a> { } impl<'a> UHessRIJK<'a> { - pub fn new_without_cderi(mol: &CInt, aux: &CInt, scale_j: f64, scale_k: f64) -> Self { + pub fn new_without_cderi(mol: &CInt, aux: &CInt, factor_j: f64, factor_k: f64) -> Self { let j2c_decomp_option = J2CDecompOption { policy: J2CDecompPolicy::Cd, threshold: Some(1e-14), uplo: Upper }; let device = DeviceBLAS::default(); let (cderi, j2c_decomp) = generate_cderi_with_decomp(mol, aux, j2c_decomp_option, &device); Self { mol: mol.clone(), aux: aux.clone(), - scale_j, - scale_k, + factor_j, + factor_k, cderi: cderi.into_cow(), j2c_decomp, intmd: HashMap::new(), @@ -68,16 +68,16 @@ impl<'a> UHessRIJK<'a> { pub fn new_with_cderi( mol: &CInt, aux: &CInt, - scale_j: f64, - scale_k: f64, + factor_j: f64, + factor_k: f64, cderi: TsrCow<'a>, j2c_decomp: J2CDecompose, ) -> Self { Self { mol: mol.clone(), aux: aux.clone(), - scale_j, - scale_k, + factor_j, + factor_k, cderi: cderi.into_cow(), j2c_decomp, intmd: HashMap::new(), @@ -117,8 +117,8 @@ impl<'a> UHessRIJK<'a> { mo_occ_slice, self.cderi.view(), &self.j2c_decomp, - self.scale_j != 0.0, - self.scale_k != 0.0, + self.factor_j != 0.0, + self.factor_k != 0.0, 72, // TODO: batch size `72` should be tunable by max-memory. atm_list, None, @@ -171,24 +171,24 @@ impl<'a> UHessElecInteractAPI for UHessRIJK<'a> { }; let mut de = hess_init(); - if self.scale_j != 0.0 { + if self.factor_j != 0.0 { let de_J20 = KEYS_J20.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); let de_J11 = KEYS_J11.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); let de_J02 = KEYS_J02.iter().map(|&key| &intmd[key]).fold(hess_init(), |acc, x| acc + x); let de_J = &de_J20 + &de_J11 + &de_J02; - de += self.scale_j * &de_J; + de += self.factor_j * &de_J; self.result.insert("de_J20", de_J20); self.result.insert("de_J11", de_J11); self.result.insert("de_J02", de_J02); self.result.insert("de_J", de_J); } - if self.scale_k != 0.0 { + if self.factor_k != 0.0 { let de_K20 = sum_k_keys(&KEYS_K20); let de_K11 = sum_k_keys(&KEYS_K11); let de_K02 = sum_k_keys(&KEYS_K02); let de_K = &de_K20 + &de_K11 + &de_K02; // UHF: K coefficient is -1 (not -0.5 as in RHF) because de_K already includes the spin sum. - de -= self.scale_k * &de_K; + de -= self.factor_k * &de_K; self.result.insert("de_K20", de_K20); self.result.insert("de_K11", de_K11); self.result.insert("de_K02", de_K02); @@ -232,24 +232,24 @@ impl<'a> UHessElecInteractAPI for UHessRIJK<'a> { let mut deriv1_bra = [deriv1_bra_init(0), deriv1_bra_init(1)]; // J is spin-independent (held in AO form, shared across spins); right half-transform per spin. - if self.scale_j != 0.0 { + if self.factor_j != 0.0 { let j1ao = KEYS_J1AO.iter().map(|&key| &intmd[key]).fold(deriv1_ao_init(), |acc, x| acc + x); for s in 0..2 { - deriv1_bra[s] += self.scale_j * (&j1ao % &mocc[s]); + deriv1_bra[s] += self.factor_j * (&j1ao % &mocc[s]); } self.result.insert("j1ao", j1ao); } // K is spin-resolved; k1bra^s is stored as the right half-transform ``k1ao^s @ mocc_s`` // (shape ``[nao, nocc_s, 3, natm]``), matching the restricted optimized convention. // Note: per-spin k1bra have different `nocc_s` and cannot be summed across spins. - if self.scale_k != 0.0 { + if self.factor_k != 0.0 { for s in 0..2 { let ks = KEYS_K1BRA .iter() .map(|&key| &intmd[&format!("{key}")]) .fold(deriv1_bra_init(s), |acc, x| acc + x); // UHF: no 0.5 factor (occ = 1), unlike RHF. - deriv1_bra[s] -= self.scale_k * &ks; + deriv1_bra[s] -= self.factor_k * &ks; self.result.insert(if s == 0 { "k1bra_0" } else { "k1bra_1" }, ks); } } @@ -276,8 +276,8 @@ impl<'a> UHessElecInteractAPI for UHessRIJK<'a> { &mo_coeff, &mo_occ, bra, - self.scale_j != 0.0, - self.scale_k != 0.0, + self.factor_j != 0.0, + self.factor_k != 0.0, 72, // TODO: batch size `72` should be tunable by max-memory. ); @@ -298,12 +298,12 @@ impl<'a> UHessElecInteractAPI for UHessRIJK<'a> { // 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.scale_j * (j_ao.view() % &mocc[s]); + 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.scale_k * k_bra.view().reshape((nao, nocc[s], nprop)); + r += self.factor_k * k_bra.view().reshape((nao, nocc[s], nprop)); } resp[s] = Some(r.into_shape(shape)); } diff --git a/tests/hessian_backup/mod.rs b/tests/hessian_backup/mod.rs new file mode 100644 index 0000000000..b8cce74c25 --- /dev/null +++ b/tests/hessian_backup/mod.rs @@ -0,0 +1,2 @@ +pub mod rhf; +pub mod rks_b3lyp; \ No newline at end of file diff --git a/tests/hessian_backup/rhf.rs b/tests/hessian_backup/rhf.rs new file mode 100644 index 0000000000..15f156033a --- /dev/null +++ b/tests/hessian_backup/rhf.rs @@ -0,0 +1,47 @@ +use pyrest::hessian_backup::rscf_interface::rscf_hess_interface; +use pyrest::hessian_backup::config::HessSCFConfig; + +use pyrest::ctrl_io; +use pyrest::molecule_io::Molecule; +use pyrest::scf_io::{self, scf_without_build}; + +static INPUT_NH3: &str = r##" +[ctrl] + print_level = 2 + num_threads = 16 + xc = "hf" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 1.0 + spin_polarization = false + auxbasis_response = true + mixer = "diis" + num_max_diis = 8 + start_diis_cycle = 3 + mix_param = 0.8 + max_scf_cycle = 100 + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 1.0 0.1 0.2 + H 0.3 1.1 0.2 + H 0.1 0.1 1.2 + """ +"##; + +#[test] +fn test_nh3() { + 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); + + let config = HessSCFConfig::default(); + rscf_hess_interface(&mut scf_data, &config); +} \ No newline at end of file diff --git a/tests/hessian_backup/rks_b3lyp.rs b/tests/hessian_backup/rks_b3lyp.rs new file mode 100644 index 0000000000..50f8c8d7d3 --- /dev/null +++ b/tests/hessian_backup/rks_b3lyp.rs @@ -0,0 +1,47 @@ +use pyrest::hessian_backup::rscf_interface::rscf_hess_interface; +use pyrest::hessian_backup::config::HessSCFConfig; + +use pyrest::ctrl_io; +use pyrest::molecule_io::Molecule; +use pyrest::scf_io::{self, scf_without_build}; + +static INPUT_NH3: &str = r##" +[ctrl] + print_level = 2 + num_threads = 16 + xc = "b3lyp" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 1.0 + spin_polarization = false + auxbasis_response = true + mixer = "diis" + num_max_diis = 8 + start_diis_cycle = 3 + mix_param = 0.8 + max_scf_cycle = 100 + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 1.0 0.1 0.2 + H 0.3 1.1 0.2 + H 0.1 0.1 1.2 + """ +"##; + +#[test] +fn test_nh3() { + 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); + + let config = HessSCFConfig::default(); + rscf_hess_interface(&mut scf_data, &config); +} \ No newline at end of file diff --git a/tests/test_hessian_backup.rs b/tests/test_hessian_backup.rs new file mode 100644 index 0000000000..53cf145a45 --- /dev/null +++ b/tests/test_hessian_backup.rs @@ -0,0 +1 @@ +pub mod hessian_backup; \ No newline at end of file -- Gitee From c144c50bf1ea4a265613ba98f00205082db71a39 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 13:21:58 +0800 Subject: [PATCH 04/39] nimatmul: change to a more conservative chunk scheme, add note on avoiding mkl in nimatmul. --- src/dft/numint_matmul/nimatmul.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/dft/numint_matmul/nimatmul.rs b/src/dft/numint_matmul/nimatmul.rs index f7f6d2a733..5b49346c89 100644 --- a/src/dft/numint_matmul/nimatmul.rs +++ b/src/dft/numint_matmul/nimatmul.rs @@ -1,3 +1,9 @@ +// Note on linkage: +// +// This algorithm implementation strictly requires linkage of openblas. +// If mkl is also linked and precedence is given to mkl, the algorithm will be very slow due to +// threading conflict. Use `patchelf` to remove mkl linkage from the binary if necessary. + use super::prelude::*; /// Numerical integration driver using matrix-multiplication. @@ -19,8 +25,8 @@ pub struct NIMatmul<'a> { /// /// Relations of size: full-grid > batch > chunk > per-grid = 1. /// - /// This value is better set to KC of micro-kernel (256-512 for usual x86 server). - /// Default to be 384. + /// This value is better set around KC of micro-kernel (256-512 for usual x86 server). + /// Default to be 1536 (3-6 KCs). pub nchunk: usize, /// Number of grid points to process in one batch. @@ -32,7 +38,7 @@ pub struct NIMatmul<'a> { /// /// This value is better set to a proper size, not exceeding available memory, and be multiple /// of `nchunk` for better performance. - /// Default to be 384 * 4 * nthreads. nthreads is determined at runtime by rayon. + /// Default to be 1536 * 1 * nthreads. nthreads is determined at runtime by rayon. pub nbatch: usize, } @@ -40,8 +46,8 @@ impl<'a> NIMatmul<'a> { /// Creates a new instance with the given integral engine, grid coordinates, and weights. pub fn new(cint: &CInt, coords: &[[f64; 3]], weights: &[f64]) -> Self { assert!(coords.len() == weights.len(), "Number of coordinates must match number of weights"); - let nchunk = 384; - let nbatch = nchunk * 4 * rayon::current_num_threads(); + let nchunk = 1536; + let nbatch = nchunk * 1 * rayon::current_num_threads(); Self { cint: cint.clone(), coords: coords.to_vec(), -- Gitee From 366c70229d3c770fe483b1c14d050ea0fbe6f20d Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 14:06:08 +0800 Subject: [PATCH 05/39] interface: add vibrational analysis --- src/hessian_backup/rscf_interface.rs | 65 +++++++++++++++++++++++++--- src/hessian_backup/vib.rs | 8 ++++ tests/hessian_backup/rhf.rs | 15 ++++++- 3 files changed, 82 insertions(+), 6 deletions(-) diff --git a/src/hessian_backup/rscf_interface.rs b/src/hessian_backup/rscf_interface.rs index 37eb321429..78e2234767 100644 --- a/src/hessian_backup/rscf_interface.rs +++ b/src/hessian_backup/rscf_interface.rs @@ -1,10 +1,12 @@ use super::prelude::*; use crate::dft::numint_matmul::hess_rks::RHessKSNIMatmul; use crate::dft::numint_matmul::nimatmul::NIMatmul; +use crate::geom_io::get_mass_charge; +use crate::hessian_backup::vib::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; -pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) { +pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, VibInfo, ThermoInfo) { let device = DeviceBLAS::default(); // --- basic preparation --- // @@ -79,15 +81,68 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) { let verbose = scf_data.mol.ctrl.print_level >= 2; RHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose) }); - if let Some(ref mut hess_nimatmul_obj) =hess_nimatmul_obj { + if let Some(ref mut hess_nimatmul_obj) = hess_nimatmul_obj { hess_el_list.push(hess_nimatmul_obj); } // --- run hessian --- // - let mut hess_scf = - RHessSCF::new(mo_coeff, mo_occ, mo_energy, &mut hess_ovlp_obj, hess_nuc_list, hess_hcore_list, hess_el_list, config); + let mut hess_scf = RHessSCF::new( + mo_coeff, + mo_occ, + mo_energy, + &mut hess_ovlp_obj, + hess_nuc_list, + hess_hcore_list, + hess_el_list, + config, + ); let de_hess = hess_scf.make_hess(); - println!("=== HESSIAN ===\n{:12.6}", de_hess.t()) + println!("=== HESSIAN ===\n{:12.6}", de_hess.t()); + + // 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); + } + } + + // --- perform vibrational analysis --- // + + // first transpose hessian to [3*natm, 3*natm] + let natm = de_hess.shape()[3]; + let hess = de_hess.transpose((0, 2, 1, 3)).into_shape((3 * natm, 3 * natm)); + let atm_list = config.atm_list.clone().unwrap_or_else(|| (0..mol.natm()).collect_vec()); + let elems = atm_list.iter().map(|&i| scf_data.mol.geom.elem[i].clone()).collect_vec(); + + let mass_charge = get_mass_charge(&elems); + let mass = mass_charge.iter().map(|(m, _)| *m).collect_vec(); + let mass_rt = rt::asarray((&mass, &device)); + + let geom = atm_list.iter().map(|&i| mol.atom_coord(i)).collect_vec(); + let geom_rt = rt::asarray((geom, &device)).into_unpack_array(0); + let vib = harmonic_analysis(hess.view(), geom_rt.view(), mass_rt.view(), true, true); + + println!("=============== Vibrational Analysis ==============="); + let elems_ref = elems.iter().map(|e| e.as_str()).collect_vec(); + let msg_vib = print_vibs(&vib, &elems_ref, NormCo::X, true, Some(3), 4, None); + println!("{}", msg_vib); + + let geom_c = mass_centred_geom(geom_rt.view(), mass_rt.view()); + let rc_cm = rotation_const(mass_rt.view(), geom_c.view(), "wavenumber").to_vec(); + let rc_ghz = rotation_const(mass_rt.view(), geom_c.view(), "GHz"); + let rotor = RotorType::from_rot_const_ghz(rc_ghz.view()); + let mass_sum = mass.iter().sum::(); + let e0 = scf_data.scf_energy; + let multiplicity = scf_data.mol.ctrl.spin; + // TODO: sigma (rotational symmetry number) is not yet determined. For now, we set it to 1. + let th = thermo(&vib, 298.15, 101325.0, multiplicity as _, mass_sum, e0, 1, &rc_cm, rotor); + + println!("=============== Thermo Analysis ==============="); + let msg_th = print_thermo(&th, multiplicity as _, mass_sum); + println!("{}", msg_th); + + (de_hess.into_shape(-1).into_vec(), vib, th) } diff --git a/src/hessian_backup/vib.rs b/src/hessian_backup/vib.rs index 29c019c1fb..2c7c0a8835 100644 --- a/src/hessian_backup/vib.rs +++ b/src/hessian_backup/vib.rs @@ -35,6 +35,14 @@ fn uconv_cm1() -> f64 { (NA * HARTREE2J * 1.0e19).sqrt() / (2.0 * std::f64::consts::PI * C * BOHR2ANG) } +/// Mass-centred geometry `[3, natm]`. +pub fn mass_centred_geom(geom: TsrView, mass: TsrView) -> Tsr { + let mass_sum = mass.to_vec().iter().sum::(); + let mc = (&geom * mass.i((None, ..))).sum_axes(1) / mass_sum; // [3] + let mc_vec = mc.to_vec(); + &geom - rt::asarray((mc_vec, [3, 1].c(), geom.device())) +} + // --------------------------------------------------------------------------- // Translation / rotation space // --------------------------------------------------------------------------- diff --git a/tests/hessian_backup/rhf.rs b/tests/hessian_backup/rhf.rs index 15f156033a..752337c82c 100644 --- a/tests/hessian_backup/rhf.rs +++ b/tests/hessian_backup/rhf.rs @@ -43,5 +43,18 @@ fn test_nh3() { scf_without_build(&mut scf_data, &None); let config = HessSCFConfig::default(); - rscf_hess_interface(&mut scf_data, &config); + let (_, vib, th) = rscf_hess_interface(&mut scf_data, &config); + + let ref_freqs = [1263.343780, 1367.102321, 1424.072405, 2132.997526, 2443.140863, 3517.051480]; + for (k, &i) in vib.vib_indices().iter().enumerate() { + assert!((vib.omega[i] - ref_freqs[k]).abs() < 1e-2, "freq {}: {} != {}", i, vib.omega[i], ref_freqs[k]); + } + + pub const TRANS: usize = 1; + pub const VIB: usize = 3; + + assert!((th.zpe[VIB] - 0.027674515751).abs() < 1e-6); + assert!((th.s[TRANS] - 0.054886124656).abs() < 1e-6); + assert!((th.cv_tot - 0.010125123641).abs() < 1e-6); + assert!((th.cp_tot - 0.013291934132).abs() < 1e-6); } \ No newline at end of file -- Gitee From aa41df5cda7bdce977b90d8ba85e63211626a7d2 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 14:38:31 +0800 Subject: [PATCH 06/39] interface: uhf/uks completed --- src/dft/numint_matmul/hess_uks.rs | 6 +- src/hessian_backup/mod.rs | 1 + src/hessian_backup/rscf_interface.rs | 4 +- src/hessian_backup/uscf.rs | 8 +- src/hessian_backup/uscf_interface.rs | 157 +++++++++++++++++++++++++++ tests/hessian_backup/mod.rs | 4 +- tests/hessian_backup/rhf.rs | 8 +- tests/hessian_backup/rks_b3lyp.rs | 8 +- tests/hessian_backup/uhf.rs | 53 +++++++++ tests/hessian_backup/uks_tpss0.rs | 54 +++++++++ 10 files changed, 291 insertions(+), 12 deletions(-) create mode 100644 src/hessian_backup/uscf_interface.rs create mode 100644 tests/hessian_backup/uhf.rs create mode 100644 tests/hessian_backup/uks_tpss0.rs diff --git a/src/dft/numint_matmul/hess_uks.rs b/src/dft/numint_matmul/hess_uks.rs index 612e8fe7e7..1467130f9b 100644 --- a/src/dft/numint_matmul/hess_uks.rs +++ b/src/dft/numint_matmul/hess_uks.rs @@ -619,7 +619,7 @@ pub fn get_uks_response_bra_batched( pub struct UHessKSNIMatmul<'a> { pub mol: CInt, - pub xc_func_list: &'a [(f64, LibXCFunctional)], + pub xc_func_list: Vec<(f64, LibXCFunctional)>, pub ni: NIMatmul<'a>, pub ni_cpks: Option>, pub verbose: bool, @@ -628,7 +628,7 @@ pub struct UHessKSNIMatmul<'a> { } impl<'a> UHessKSNIMatmul<'a> { - pub fn new(mol: &CInt, xc_func_list: &'a [(f64, LibXCFunctional)], ni: NIMatmul<'a>, verbose: bool) -> Self { + pub fn new(mol: &CInt, xc_func_list: Vec<(f64, LibXCFunctional)>, ni: NIMatmul<'a>, verbose: bool) -> Self { Self { mol: mol.clone(), xc_func_list, @@ -649,7 +649,7 @@ impl<'a> UHessKSNIMatmul<'a> { let (result, _timing) = make_hessian_setup_batched_uks( &self.mol, - self.xc_func_list, + &self.xc_func_list, &mut self.ni, dm0α.view(), dm0β.view(), diff --git a/src/hessian_backup/mod.rs b/src/hessian_backup/mod.rs index d5650ae5ef..3234e73335 100644 --- a/src/hessian_backup/mod.rs +++ b/src/hessian_backup/mod.rs @@ -30,6 +30,7 @@ pub mod uscf; // total hess interface to REST pub mod rscf_interface; +pub mod uscf_interface; // vibrational analysis pub mod vib; diff --git a/src/hessian_backup/rscf_interface.rs b/src/hessian_backup/rscf_interface.rs index 78e2234767..4226d38338 100644 --- a/src/hessian_backup/rscf_interface.rs +++ b/src/hessian_backup/rscf_interface.rs @@ -1,6 +1,4 @@ use super::prelude::*; -use crate::dft::numint_matmul::hess_rks::RHessKSNIMatmul; -use crate::dft::numint_matmul::nimatmul::NIMatmul; use crate::geom_io::get_mass_charge; use crate::hessian_backup::vib::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; @@ -62,6 +60,8 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, // --- DFT --- // let mut hess_nimatmul_obj = (!is_hf).then(|| { + use crate::dft::numint_matmul::hess_rks::RHessKSNIMatmul; + use crate::dft::numint_matmul::nimatmul::NIMatmul; use crate::dft::numint_matmul::prelude::*; use libxc::prelude::*; diff --git a/src/hessian_backup/uscf.rs b/src/hessian_backup/uscf.rs index 066e5ca311..202003cce4 100644 --- a/src/hessian_backup/uscf.rs +++ b/src/hessian_backup/uscf.rs @@ -6,7 +6,7 @@ pub struct UHessSCF<'a> { pub mo_coeff: [Tsr; 2], pub mo_occ: [Tsr; 2], pub mo_energy: [Tsr; 2], - pub ovlp_obj: UHessOvlp, + 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>, @@ -22,11 +22,11 @@ impl<'a> UHessSCF<'a> { mo_coeff: [Tsr; 2], mo_occ: [Tsr; 2], mo_energy: [Tsr; 2], - ovlp_obj: UHessOvlp, + 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>, - config: HessSCFConfig, + config: &HessSCFConfig, ) -> Self { Self { mo_coeff, @@ -36,7 +36,7 @@ impl<'a> UHessSCF<'a> { nuc_list, core_list, el_list, - config, + config: config.clone(), result: HashMap::new(), timing: Vec::new(), } diff --git a/src/hessian_backup/uscf_interface.rs b/src/hessian_backup/uscf_interface.rs new file mode 100644 index 0000000000..2cd9f9048e --- /dev/null +++ b/src/hessian_backup/uscf_interface.rs @@ -0,0 +1,157 @@ +use super::prelude::*; +use crate::geom_io::get_mass_charge; +use crate::hessian_backup::vib::*; +use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; +use crate::SCF; + +pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, VibInfo, ThermoInfo) { + let device = DeviceBLAS::default(); + + // --- basic preparation --- // + let mo_coeff = { + let mo_coeff_0 = &scf_data.eigenvectors[0]; + let mo_coeff_1 = &scf_data.eigenvectors[1]; + let mo_coeff_0 = rt::asarray((&mo_coeff_0.data, mo_coeff_0.size, &device)).into_contig(ColMajor); + let mo_coeff_1 = rt::asarray((&mo_coeff_1.data, mo_coeff_1.size, &device)).into_contig(ColMajor); + [mo_coeff_0, mo_coeff_1] + }; + let mo_occ = { + let mo_occ_0 = &scf_data.occupation[0]; + let mo_occ_1 = &scf_data.occupation[1]; + let mo_occ_0 = rt::asarray((mo_occ_0, [mo_occ_0.len()], &device)).into_contig(ColMajor); + let mo_occ_1 = rt::asarray((mo_occ_1, [mo_occ_1.len()], &device)).into_contig(ColMajor); + [mo_occ_0, mo_occ_1] + }; + let mo_energy = { + let mo_energy_0 = &scf_data.eigenvalues[0]; + let mo_energy_1 = &scf_data.eigenvalues[1]; + let mo_energy_0 = rt::asarray((mo_energy_0, [mo_energy_0.len()], &device)).into_contig(ColMajor); + let mo_energy_1 = rt::asarray((mo_energy_1, [mo_energy_1.len()], &device)).into_contig(ColMajor); + [mo_energy_0, mo_energy_1] + }; + + let mol_obj = &scf_data.mol; + let mol = get_cint_mol(mol_obj); + let aux = get_cint_aux(mol_obj); + + let mut hess_ovlp_obj = UHessOvlp::new(&mol, &device); + let mut hess_nuc_repl_obj = HessNucRepl::new(&mol, &device); + let mut hess_hcore_obj = UHessHcore::new(&mol, &device); + + let hess_nuc_list: Vec<&mut dyn HessNucAPI> = vec![&mut hess_nuc_repl_obj]; + let hess_hcore_list: Vec<&mut dyn UHessCoreAPI> = vec![&mut hess_hcore_obj]; + let mut hess_el_list: Vec<&mut dyn UHessElecInteractAPI> = Vec::new(); + + // --- RI-JK --- // + + 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; + let scale_k = match is_hf { + true => 1.0, + false => scf_data.mol.xc_data.dfa_hybrid_scf, + }; + let j2c_decomp_option = &scf_data.mol.ctrl.j2c_decomp; + let j2c_decomp = crate::ri_jk::get_j2c_decomp(&aux, &device, *j2c_decomp_option); + + let mut hess_rijk_obj = if let Some((rimatr, _, _)) = &scf_data.rimatr { + let cderi = rimatr.to_rstsr_view(&device).into_cow(); + UHessRIJK::new_with_cderi(&mol, &aux, scale_j, scale_k, cderi, j2c_decomp) + } else { + panic!( + "This implementation requires cholesky decomposed ERI (or rimatr) to be available and stored in memory." + ); + }; + + hess_el_list.push(&mut hess_rijk_obj); + + // --- DFT --- // + + let mut hess_nimatmul_obj = (!is_hf).then(|| { + use crate::dft::numint_matmul::hess_uks::UHessKSNIMatmul; + use crate::dft::numint_matmul::nimatmul::NIMatmul; + use crate::dft::numint_matmul::prelude::*; + use libxc::prelude::*; + + let grid_coords = &scf_data.grids.as_ref().unwrap().coordinates; + let grid_weights = &scf_data.grids.as_ref().unwrap().weights; + let ni = NIMatmul::new(&mol, grid_coords, grid_weights); + + 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; + UHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose) + }); + if let Some(ref mut hess_nimatmul_obj) = hess_nimatmul_obj { + hess_el_list.push(hess_nimatmul_obj); + } + + // --- 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, + config, + ); + + let de_hess = hess_scf.make_hess(); + println!("=== HESSIAN ===\n{:12.6}", de_hess.t()); + + // 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); + } + } + + // --- perform vibrational analysis --- // + + // first transpose hessian to [3*natm, 3*natm] + let natm = de_hess.shape()[3]; + let hess = de_hess.transpose((0, 2, 1, 3)).into_shape((3 * natm, 3 * natm)); + let atm_list = config.atm_list.clone().unwrap_or_else(|| (0..mol.natm()).collect_vec()); + let elems = atm_list.iter().map(|&i| scf_data.mol.geom.elem[i].clone()).collect_vec(); + + let mass_charge = get_mass_charge(&elems); + let mass = mass_charge.iter().map(|(m, _)| *m).collect_vec(); + let mass_rt = rt::asarray((&mass, &device)); + + let geom = atm_list.iter().map(|&i| mol.atom_coord(i)).collect_vec(); + let geom_rt = rt::asarray((geom, &device)).into_unpack_array(0); + let vib = harmonic_analysis(hess.view(), geom_rt.view(), mass_rt.view(), true, true); + + println!("=============== Vibrational Analysis ==============="); + let elems_ref = elems.iter().map(|e| e.as_str()).collect_vec(); + let msg_vib = print_vibs(&vib, &elems_ref, NormCo::X, true, Some(3), 4, None); + println!("{}", msg_vib); + + let geom_c = mass_centred_geom(geom_rt.view(), mass_rt.view()); + let rc_cm = rotation_const(mass_rt.view(), geom_c.view(), "wavenumber").to_vec(); + let rc_ghz = rotation_const(mass_rt.view(), geom_c.view(), "GHz"); + let rotor = RotorType::from_rot_const_ghz(rc_ghz.view()); + let mass_sum = mass.iter().sum::(); + let e0 = scf_data.scf_energy; + let multiplicity = scf_data.mol.ctrl.spin; + // TODO: sigma (rotational symmetry number) is not yet determined. For now, we set it to 1. + let th = thermo(&vib, 298.15, 101325.0, multiplicity as _, mass_sum, e0, 1, &rc_cm, rotor); + + println!("=============== Thermo Analysis ==============="); + let msg_th = print_thermo(&th, multiplicity as _, mass_sum); + println!("{}", msg_th); + + (de_hess.into_shape(-1).into_vec(), vib, th) +} diff --git a/tests/hessian_backup/mod.rs b/tests/hessian_backup/mod.rs index b8cce74c25..97a6a0e275 100644 --- a/tests/hessian_backup/mod.rs +++ b/tests/hessian_backup/mod.rs @@ -1,2 +1,4 @@ pub mod rhf; -pub mod rks_b3lyp; \ No newline at end of file +pub mod rks_b3lyp; +pub mod uhf; +pub mod uks_tpss0; diff --git a/tests/hessian_backup/rhf.rs b/tests/hessian_backup/rhf.rs index 752337c82c..7f8e62a7a4 100644 --- a/tests/hessian_backup/rhf.rs +++ b/tests/hessian_backup/rhf.rs @@ -5,6 +5,8 @@ use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; use pyrest::scf_io::{self, scf_without_build}; +use rstsr::prelude::*; + static INPUT_NH3: &str = r##" [ctrl] print_level = 2 @@ -43,7 +45,7 @@ fn test_nh3() { scf_without_build(&mut scf_data, &None); let config = HessSCFConfig::default(); - let (_, vib, th) = rscf_hess_interface(&mut scf_data, &config); + let (de, vib, th) = rscf_hess_interface(&mut scf_data, &config); let ref_freqs = [1263.343780, 1367.102321, 1424.072405, 2132.997526, 2443.140863, 3517.051480]; for (k, &i) in vib.vib_indices().iter().enumerate() { @@ -57,4 +59,8 @@ fn test_nh3() { assert!((th.s[TRANS] - 0.054886124656).abs() < 1e-6); assert!((th.cv_tot - 0.010125123641).abs() < 1e-6); assert!((th.cp_tot - 0.013291934132).abs() < 1e-6); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian:\n{:12.6}", de.t()); } \ No newline at end of file diff --git a/tests/hessian_backup/rks_b3lyp.rs b/tests/hessian_backup/rks_b3lyp.rs index 50f8c8d7d3..529e153c03 100644 --- a/tests/hessian_backup/rks_b3lyp.rs +++ b/tests/hessian_backup/rks_b3lyp.rs @@ -5,6 +5,8 @@ use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; use pyrest::scf_io::{self, scf_without_build}; +use rstsr::prelude::*; + static INPUT_NH3: &str = r##" [ctrl] print_level = 2 @@ -43,5 +45,9 @@ fn test_nh3() { scf_without_build(&mut scf_data, &None); let config = HessSCFConfig::default(); - rscf_hess_interface(&mut scf_data, &config); + let (de, _, _) = rscf_hess_interface(&mut scf_data, &config); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian:\n{:12.6}", de.t()); } \ No newline at end of file diff --git a/tests/hessian_backup/uhf.rs b/tests/hessian_backup/uhf.rs new file mode 100644 index 0000000000..5efd91735a --- /dev/null +++ b/tests/hessian_backup/uhf.rs @@ -0,0 +1,53 @@ +use pyrest::hessian_backup::uscf_interface::uscf_hess_interface; +use pyrest::hessian_backup::config::HessSCFConfig; + +use pyrest::ctrl_io; +use pyrest::molecule_io::Molecule; +use pyrest::scf_io::{self, scf_without_build}; + +use rstsr::prelude::*; + +static INPUT_NH3: &str = r##" +[ctrl] + print_level = 2 + num_threads = 16 + xc = "hf" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 2.0 + spin = 3.0 + spin_polarization = true + auxbasis_response = true + mixer = "diis" + num_max_diis = 8 + start_diis_cycle = 3 + mix_param = 0.8 + max_scf_cycle = 100 + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 1.0 0.1 0.2 + H 0.3 1.1 0.2 + H 0.1 0.1 1.2 + """ +"##; + +#[test] +fn test_nh3() { + 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); + + let config = HessSCFConfig::default(); + let (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian:\n{:12.6}", de.t()); +} \ No newline at end of file diff --git a/tests/hessian_backup/uks_tpss0.rs b/tests/hessian_backup/uks_tpss0.rs new file mode 100644 index 0000000000..0a00aeb17c --- /dev/null +++ b/tests/hessian_backup/uks_tpss0.rs @@ -0,0 +1,54 @@ +use pyrest::hessian_backup::uscf_interface::uscf_hess_interface; +use pyrest::hessian_backup::config::HessSCFConfig; + +use pyrest::ctrl_io; +use pyrest::molecule_io::Molecule; +use pyrest::scf_io::{self, scf_without_build}; + +use rstsr::prelude::*; + +static INPUT_NH3: &str = r##" +[ctrl] + print_level = 2 + num_threads = 16 + xc = "TPSS0" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 2.0 + spin = 3.0 + spin_polarization = true + auxbasis_response = true + mixer = "diis" + num_max_diis = 8 + start_diis_cycle = 3 + mix_param = 0.8 + max_scf_cycle = 100 + xc_parser = "parse_xc" + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 1.0 0.1 0.2 + H 0.3 1.1 0.2 + H 0.1 0.1 1.2 + """ +"##; + +#[test] +fn test_nh3() { + 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); + + let config = HessSCFConfig::default(); + let (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian:\n{:12.6}", de.t()); +} \ No newline at end of file -- Gitee From 521363d5b2e1406c971b6a776ceb8f4c06ddd320 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 14:58:38 +0800 Subject: [PATCH 07/39] vib: updated some formatting --- src/hessian_backup/vib.rs | 45 +++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/src/hessian_backup/vib.rs b/src/hessian_backup/vib.rs index 2c7c0a8835..c9e3a8f6d9 100644 --- a/src/hessian_backup/vib.rs +++ b/src/hessian_backup/vib.rs @@ -863,8 +863,18 @@ pub fn print_vibs( ncprec: Option, ) -> String { let nat = vib.ndof / 3; - let active: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] == "V").collect(); - + let active: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] != "TR").collect(); + + // print inactive modes (TR) first, if any + println!("Inactive modes:"); + let inactive: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] == "TR").collect(); + for &i in &inactive { + let trv = vib.trv[i]; + let freq = if vib.imag[i] { format!("{:.4}i", vib.omega[i]) } else { format!("{:.4}", vib.omega[i]) }; + println!("{:<2} Freq [cm^-1]: {}", trv, freq); + } + println!(""); + println!("Active modes:"); let presp = 2; let prewidth = 24usize; let colsp = 2usize; @@ -919,17 +929,18 @@ pub fn print_vibs( // Irrep — skip (no irrep/symmetry in our implementation) // scalar property rows — centered, colsp trailing - let labels: [&str; 5] = - ["Reduced mass [u]", "Force const [mDyne/A]", "Turning point v=0 [a0]", "RMS dev v=0 [a0 u^1/2]", "Char temp [K]"]; + let labels: [&str; 5] = [ + "Reduced mass [u]", + "Force const [mDyne/A]", + "Turning point v=0 [a0]", + "RMS dev v=0 [a0 u^1/2]", + "Char temp [K]", + ]; let val_slices: [&[f64]; 5] = [&vib.mu, &vib.k, &vib.xtp0, &vib.dq0, &vib.theta_vib]; for (label, vals) in labels.iter().zip(val_slices.iter()) { let mut l = format!("{}{}", " ".repeat(presp), ljust(label, prewidth)); for &vib in &row { - l.push_str(&format!( - "{}{}", - center(&format!("{:.*}", prec, vals[vib]), width), - " ".repeat(colsp) - )); + l.push_str(&format!("{}{}", center(&format!("{:.*}", prec, vals[vib]), width), " ".repeat(colsp))); } lines.push(l); } @@ -960,25 +971,17 @@ pub fn print_vibs( let lbl = if at < atom_lbl.len() { atom_lbl[at] } else { "" }; for xyz in 0..3 { let axis = ['X', 'Y', 'Z'][xyz]; - let mut l = format!( - "{}{:5} {} {}", - " ".repeat(presp), - at + 1, - axis, - ljust(lbl, prewidth - 14) - ); + let mut l = + format!("{}{:5} {} {}", " ".repeat(presp), at + 1, axis, ljust(lbl, prewidth - 14)); for &vib in &row { let v = normco_t[[3 * at + xyz, vib]]; - l.push_str(&format!( - "{}{}", - center(&format!("{:.*}", ncprec, v), width), - " ".repeat(colsp) - )); + l.push_str(&format!("{}{}", center(&format!("{:.*}", ncprec, v), width), " ".repeat(colsp))); } lines.push(l); } } } + lines.push("".to_string()); // blank line after each group } lines.join("\n") -- Gitee From d31b26d6c050869538cdcf5ae903d17525ace992 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 15:10:16 +0800 Subject: [PATCH 08/39] remove some markdown documents --- .../docs/get_rho_from_dm_with_output.md | 48 ------------- ...rho_from_homogeneous_braket_with_output.md | 38 ---------- ...t_rho_from_one_bra_mult_ket_with_output.md | 44 ------------ src/dft/numint_matmul/docs/mod.md | 64 ----------------- src/dft/numint_matmul/mod.rs | 1 - src/dft/xceff/libxc-xcfun-trans.md | 71 ------------------- 6 files changed, 266 deletions(-) delete mode 100644 src/dft/numint_matmul/docs/get_rho_from_dm_with_output.md delete mode 100644 src/dft/numint_matmul/docs/get_rho_from_homogeneous_braket_with_output.md delete mode 100644 src/dft/numint_matmul/docs/get_rho_from_one_bra_mult_ket_with_output.md delete mode 100644 src/dft/numint_matmul/docs/mod.md diff --git a/src/dft/numint_matmul/docs/get_rho_from_dm_with_output.md b/src/dft/numint_matmul/docs/get_rho_from_dm_with_output.md deleted file mode 100644 index ad3026f5c4..0000000000 --- a/src/dft/numint_matmul/docs/get_rho_from_dm_with_output.md +++ /dev/null @@ -1,48 +0,0 @@ -## Equations and Concepts - -- `RHO`: - - $$ - \rho_g = \sum_{\mu \nu} \phi_{g \mu} D_{\mu \nu} \phi_{g \nu} - $$ - -- `SIGMA`: - - $$ - \begin{aligned} - \sigma_{g t} &= \sum_{\mu \nu} \left( \phi_{g \mu, t} D_{\mu \nu} \phi_{g \nu} + \phi_{g \mu} D_{\mu \nu} \phi_{g \nu, t} \right) - \\\\ - &= 2 \sum_{\mu \nu} \phi_{g \mu} D_{\mu \nu} \phi_{g \nu, t} \quad \text{(symm applied)} - \end{aligned} - $$ - - Note we have already assumed the symmetry of the density matrix, so the two terms in the summation are equal. In the actual code, we only compute one of them and multiply by 2. - -- `TAU`: - - $$ - \tau_g = \frac{1}{2} \sum_{\mu \nu} \sum_{t} \phi_{g \mu, t} D_{\mu \nu} \phi_{g \nu, t} - $$ - -- `LAPL`: - - $$ - \begin{aligned} - \nabla^2 \rho_g - &= \sum_{\mu \nu} \left( \sum_{t} \phi_{g \mu, tt} D_{\mu \nu} \phi_{g \nu} + 2 \sum_{t} \phi_{g \mu, t} D_{\mu \nu} \phi_{g \nu, t} + \sum_{t} \phi_{g \mu} D_{\mu \nu} \phi_{g \nu, tt} \right) - \\\\ - &= 4 \tau_g + 2 \sum_{\mu \nu} \sum_{t} \phi_{g \mu} D_{\mu \nu} \phi_{g \nu, tt} - \quad \text{(symm applied)} - \end{aligned} - $$ - - Similar to `SIGMA`, we have already assumed the symmetry of the density matrix, so the first and third terms in the summation are equal. In the actual code, we only compute one of them and multiply by 2. - -## Usage tips - -For MatMul driver, we are not going to exploit the sparsity of atomic grids or density matrix. It is better to use `braket` versions to exploit the low-rank nature of the density matrix. - -Usually, in only the following three cases, use this function: -- density comes from post-SCF reduced density matrix, which is usually non-zero-definite. -- prototype validation. -- basis set is small, number of occupation is large. diff --git a/src/dft/numint_matmul/docs/get_rho_from_homogeneous_braket_with_output.md b/src/dft/numint_matmul/docs/get_rho_from_homogeneous_braket_with_output.md deleted file mode 100644 index 1c8360b94a..0000000000 --- a/src/dft/numint_matmul/docs/get_rho_from_homogeneous_braket_with_output.md +++ /dev/null @@ -1,38 +0,0 @@ -# Equations and Concepts - -Please note that since we assumed the homogeneous braket form, the density matrix is symmetric, and some factors of two is applicable to `SIGMA` and `LAPL`. - -- Orbital grids are defined as - - $$ - \begin{aligned} - \phi_{g i} &= \sum_{\mu} \phi_{g \mu} C_{\mu i} - \\\\ - \phi_{g i, t} &= \sum_{\mu} \phi_{g \mu, t} C_{\mu i} - \end{aligned} - $$ - -- `RHO`: - - $$ - \rho_g = \sum_{i} \phi_{g i}^2 - $$ - -- `SIGMA`: - - $$ - \sigma_{g t} = 2 \sum_{i} \phi_{g i} \phi_{g i, t} - $$ - -- `TAU`: - - $$ - \tau_g = \frac{1}{2} \sum_{i} \sum_{t} \phi_{g i, t}^2 - $$ - -- `LAPL`: - - $$ - \nabla^2 \rho_g - = 4 \tau_g + 2 \sum_{i} \sum_{t} \phi_{g i} \phi_{g i, tt} - $$ diff --git a/src/dft/numint_matmul/docs/get_rho_from_one_bra_mult_ket_with_output.md b/src/dft/numint_matmul/docs/get_rho_from_one_bra_mult_ket_with_output.md deleted file mode 100644 index f8bbd3f5f4..0000000000 --- a/src/dft/numint_matmul/docs/get_rho_from_one_bra_mult_ket_with_output.md +++ /dev/null @@ -1,44 +0,0 @@ -# Equations and Concepts - -Since bra and ket may differ, the density matrix $D_{\mu\nu} = \sum_i \mathrm{bra}_{\mu i} \mathrm{ket}_{\nu i}$ is not necessarily symmetric, so both cross terms must be computed explicitly. - -- Orbital grids are defined as - - $$ - \begin{aligned} - \phi_{g i}^{\mathrm{bra}} &= \sum_{\mu} \phi_{g \mu} \mathrm{bra}_{\mu i} - \\\\ - \phi_{g i}^{\mathrm{ket}} &= \sum_{\mu} \phi_{g \mu} \mathrm{ket}_{\mu i} - \\\\ - \phi_{g i, t}^{\mathrm{bra}} &= \sum_{\mu} \phi_{g \mu, t} \mathrm{bra}_{\mu i} - \\\\ - \phi_{g i, t}^{\mathrm{ket}} &= \sum_{\mu} \phi_{g \mu, t} \mathrm{ket}_{\mu i} - \end{aligned} - $$ - -- `RHO`: - - $$ - \rho_g = \sum_{i} \phi_{g i}^{\mathrm{bra}} \phi_{g i}^{\mathrm{ket}} - $$ - -- `SIGMA`: - - $$ - \sigma_{g t} = \sum_{i} \left( \phi_{g i}^{\mathrm{bra}} \phi_{g i, t}^{\mathrm{ket}} + \phi_{g i, t}^{\mathrm{bra}} \phi_{g i}^{\mathrm{ket}} \right) - $$ - - Unlike the symmetric case, the two terms are not equal in general so both must be computed. - -- `TAU`: - - $$ - \tau_g = \frac{1}{2} \sum_{i} \sum_{t} \phi_{g i, t}^{\mathrm{bra}} \phi_{g i, t}^{\mathrm{ket}} - $$ - -- `LAPL`: - - $$ - \nabla^2 \rho_g - = 4 \tau_g + \sum_{i} \sum_{t} \left( \phi_{g i}^{\mathrm{bra}} \phi_{g i, tt}^{\mathrm{ket}} + \phi_{g i, tt}^{\mathrm{bra}} \phi_{g i}^{\mathrm{ket}} \right) - $$ diff --git a/src/dft/numint_matmul/docs/mod.md b/src/dft/numint_matmul/docs/mod.md deleted file mode 100644 index dfd76a9a0c..0000000000 --- a/src/dft/numint_matmul/docs/mod.md +++ /dev/null @@ -1,64 +0,0 @@ -# Introduction of `MatMul` DFT grid driver - -This driver is the naive DFT driver. The code effort is minimized, yet still be efficient to small systems. - -As the name suggests, we fully utilize BLAS3 GEMM to perform the DFT grid operations. - -## `ao`: atomic orbital values - -The AO tensor `ao` has shape `[ngrids, nao, ncomp]` where the component dimension is ordered as: - -| Index | Deriv | Component | -|--|--|--| -| 0 | 0 | $\phi_{g \mu}$ | -| 1–3 | 1 | $\phi_{g \mu, x}, \phi_{g \mu, y}, \phi_{g \mu, z}$ | -| 4–9 | 2 | $\phi_{g \mu,xx}, \phi_{g \mu,xy}, \phi_{g \mu,xz}, \phi_{g \mu,yy}, \phi_{g \mu,yz}, \phi_{g \mu,zz}$ | - -Additional to the above table, for notation simplicity, we denote - -- orbital notation - - $$ - \phi_{g \mu} = \phi_{\mu} (\bm{r}_g) - $$ - - Where subscript $g$ denotes grid, $\bm{r}_g$ denotes the coordinate of grid point $g$, and $\mu$ denotes the AO index. - -- orbital derivative notation - - $$ - \phi_{g \mu, x} = \partial_x \phi_{\mu} (\bm{r}_g) - $$ - - The partial derivative is taken with respect to the electron coordinate. - -- We usually use $t, s, r \in \{ x, y, z \}$ to denote the subscript of the coordinate. - -## `rho`: density values - -We will introduce 4 types of density values: - -| property | [`RHO`] | [`SIGMA`] | [`TAU`] | [`LAPL`] | -|--|--|--|--|--| -| notation | $\rho$ | $\sigma$ | $\tau$ | $\nabla^2 \rho$ | -| [`num_rho_comp`](NIDenType::num_rho_comp) | 1 | 4 | 5 | 6 | -| [`num_ao_deriv`](NIDenType::num_ao_deriv) | 0 | 1 | 1 | 2 | -| [`num_ao_comp`](NIDenType::num_ao_comp) | 1 | 4 | 4 | 10 | -| usual XC type | LDA | GGA | mGGA | mGGA | - -The density to be input in this driver is `[ngirds, num_rho_comp]` for spin unpolarized case, and `[ngrids, num_rho_comp, 2]` for spin polarized case. The density is computed by contracting the AO tensor with the density matrix, which is the same as the usual DFT grid driver. - -The density components are ordered as: - -$$ -\rho, \rho_x, \rho_y, \rho_z, \tau, \nabla^2 \rho -$$ - -For laplacian density functional, whatever the xc functional uses tau, currently we enforce the evaluation of $\tau$ to be the 4th component of the density. - -## Important pure functions - -**Density Grid Evaluation** - -- [`get_rho_from_dm_with_output`] -- [`get_rho_from_homogeneous_braket_with_output] diff --git a/src/dft/numint_matmul/mod.rs b/src/dft/numint_matmul/mod.rs index 1359e557bc..6b03701adf 100644 --- a/src/dft/numint_matmul/mod.rs +++ b/src/dft/numint_matmul/mod.rs @@ -2,7 +2,6 @@ //! //! Though saying "naive", it should be sufficiently good for dense GTO grids - basis pairs (small systems). //! For large molecules, we do not exploit sparsity here for code simplicity. -#![doc = include_str!("docs/mod.md")] pub mod hess_rks; pub mod hess_uks; diff --git a/src/dft/xceff/libxc-xcfun-trans.md b/src/dft/xceff/libxc-xcfun-trans.md index 2279115c3c..60af99a4a9 100644 --- a/src/dft/xceff/libxc-xcfun-trans.md +++ b/src/dft/xceff/libxc-xcfun-trans.md @@ -58,74 +58,3 @@ We can see that - TAU: `t_u`, `t_d` - LAPL: `l_u`, `l_d` - We assume the RHO (LDA) only inputs RHO; SIGMA (GGA) inputs both RHO and SIGMA; TAU (some meta-GGA) inputs RHO SIGMA TAU, and LAPL (some meta-GGA) inputs all RHO SIGMA TAU LAPL (though some LAPL meta-GGAs does not actually input tau, but we still require the TAU to be available for simplicity). - -## Code for Automatic Index Map Generation - -```python -from itertools import combinations_with_replacement -from math import comb - - -def libxc_to_xcfun_indices_map(den_type: str, deriv: int) -> list[int]: - """Spin-Polarized Indices Map from LibXC to XCFun. - - Parameters - ---------- - den_type : str - Density Type. Supports `rho`, `sigma`, `tau`, `lapl`. - deriv : int - Derivative level. - - Example - ------- - >>> libxc_to_xcfun_indices_map("sigma", 2) - [6, 7, 9, 10, 11, 8, 12, 13, 14, 15, 16, 17, 18, 19, 20] - """ - # Each variable: (type_priority, spin_index) - # RHO has 2 spin components, SIGMA has 3, TAU has 2, LAPL has 2 - group_specs = [ - ("rho", 0, 2), - ("sigma", 1, 3), - ("tau", 2, 2), - ("lapl", 3, 2), - ] - - type_map = { - "rho": ["rho"], - "sigma": ["rho", "sigma"], - "tau": ["rho", "sigma", "tau"], - "lapl": ["rho", "sigma", "tau", "lapl"], - } - - if den_type not in type_map: - raise ValueError(f"Unknown den_type: {den_type}") - - active_groups = set(type_map[den_type]) - - # Build variable list: each variable is (type_priority, spin_index) - variables = [] - for group_name, priority, n_spin in group_specs: - if group_name in active_groups: - for spin in range(n_spin): - variables.append((priority, spin)) - - d = len(variables) - - # Generate all non-decreasing multi-indices of length deriv - # combinations_with_replacement yields them in lexicographic order = XCFun order - xcfun_order = list(combinations_with_replacement(range(d), deriv)) - - # LibXC order: sort by density type signature first, then by variable indices - def libxc_key(mi): - return tuple(variables[i][0] for i in mi) + mi - - libxc_order = sorted(xcfun_order, key=libxc_key) - - # Build reverse lookup: multi-index -> LibXC position - libxc_pos = {mi: pos for pos, mi in enumerate(libxc_order)} - - # Base offset = sum of outputs for all previous derivative levels - base_offset = sum(comb(d + i - 1, i) for i in range(deriv)) - - return [base_offset + libxc_pos[mi] for mi in xcfun_order] -``` -- Gitee From 217f56ed6c7a2d146848896bdd9e05960cc19a30 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 24 Jun 2026 15:19:54 +0800 Subject: [PATCH 09/39] add keyword and main_driver --- src/ctrl_io/mod.rs | 8 ++++++++ src/main_driver.rs | 26 +++++++++++++++++++++++++- 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index 3cb6ddfd13..12535815bf 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -361,6 +361,8 @@ pub struct InputKeywords { /// Use the optimized fxc_matvec_opt (rayon + pre-allocated workspace). #[pyo3(get, set)] pub use_fxc_opt: bool, + /// Use module `hessian_backup` instead of `hessian` for Hessian calculation. + pub use_hessian_backup: bool, } impl Default for InputKeywords { @@ -520,6 +522,7 @@ impl InputKeywords { tddft: None, cphf: None, use_fxc_opt: false, + use_hessian_backup: false, } } @@ -1813,6 +1816,11 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result { tmp_str.to_lowercase() }, other => String::from("legacy"), }; + tmp_input.use_hessian_backup = match tmp_ctrl.get("use_hessian_backup").unwrap_or(&serde_json::Value::Null) { + serde_json::Value::Bool(tmp_bool) => *tmp_bool, + serde_json::Value::String(tmp_str) => tmp_str.to_lowercase().parse().unwrap_or(false), + _ => false, + }; //=========================================================== // Global check of ctrl keywords and futher modification diff --git a/src/main_driver.rs b/src/main_driver.rs index 2dd1a5b6eb..4da3c4a978 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -243,7 +243,31 @@ pub fn main_driver() -> anyhow::Result<()> { }, // UNVERIFIED NORMAL MODES CALCULATION JobType::NormalModes => { - eval_normal_modes(&mut scf_data, &mut time_mark, &mpi_operator); + if !scf_data.mol.ctrl.use_hessian_backup { + eval_normal_modes(&mut scf_data, &mut time_mark, &mpi_operator); + } else { + use crate::hessian_backup::rscf_interface::rscf_hess_interface; + use crate::hessian_backup::uscf_interface::uscf_hess_interface; + use crate::hessian_backup::config::HessSCFConfig; + + eprintln!("[WARN] You are using the hessian_backup module, which is still under development."); + + // simple guard, but currently many methods (solvent, range-separate, dftd are not supported) + if scf_data.mol.xc_data.is_fifth_dfa() { + panic!("Normal modes calculation is currently not available for post-SCF methods."); + } + + let config = HessSCFConfig::default(); + match scf_data.scftype { + SCFType::RHF => { + rscf_hess_interface(&scf_data, &config); + } + SCFType::UHF => { + uscf_hess_interface(&scf_data, &config); + }, + _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types.") + } + } }, // ------------ _ => {} -- Gitee From af0ea1d34473a2b2e929dfff9c2928f2235c16ad Mon Sep 17 00:00:00 2001 From: ajz34 Date: Fri, 26 Jun 2026 15:56:16 +0800 Subject: [PATCH 10/39] refactor(vib): drive TR/V classification by rotor type (PySCF way) Replace Psi4's numerical TR-mask approach (SVD rank cutoff on the TR basis + per-mode _vec_in_space membership test with an arbitrary 1e-4 tolerance) with PySCF's physical, rotor-type-driven approach. - nrt (number of TR dof) is now derived from rotor type via _get_rotor_type(rotation_const(mass, geom_cm, 'GHz')): ATOM->0, LINEAR->5, REGULAR->6, reduced for project_trans/project_rot. - _get_TR_space gains an nrt param: when given, select the top-nrt singular vectors so the projection rank is self-consistent with the TR count, bypassing the LINEAR_A_TOL SVD cutoff. - TRV classification is count-based on the spectral gap: after projection P = I - QQ^T, TR directions sit in null(P) with machine-zero force constants (~1e-13), separated from any real vibration (~2e-9 for 10 cm^-1) by ~8 orders. The nrt modes with smallest |force_constant| are tagged 'TR'; the rest 'V' (with '-' kept as a safety net). - geom passed to the projector is not mass-centred (only the local geom_cm for the rotation_const call), preserving the exact REF_OMEGA reference values. - Rust: drop the now-private vec_in_space; mirror the nrt/count logic. Applied to both pyhessref/vib.py and src/hessian/vib.rs. Tests: 60 py + 21 rs pass unchanged (NH3, acetaldehyde both REGULAR). Co-authored-by: Claude Code Co-authored-by: glm-5.2 --- src/hessian_backup/vib.rs | 149 +++++++++++++++++++++++--------------- 1 file changed, 92 insertions(+), 57 deletions(-) diff --git a/src/hessian_backup/vib.rs b/src/hessian_backup/vib.rs index c9e3a8f6d9..3a0515be56 100644 --- a/src/hessian_backup/vib.rs +++ b/src/hessian_backup/vib.rs @@ -8,6 +8,7 @@ //! # Note //! //! This module is direct transformation from psi4 (psi4/psi4/driver/qcdb/vib.py). +//! Some modifications are from PySCF (pyscf/hessian/thermo.py). //! This file contains AI assisted code, and not fully reviewed by human. use super::prelude::*; @@ -57,13 +58,20 @@ pub fn mass_centred_geom(geom: TsrView, mass: TsrView) -> Tsr { /// # Returns /// `tr` : `[3*natm, nrt]` orthonormal basis (each column is a TR vector). pub fn get_tr_space(mass: TsrView, geom: TsrView, space: &str) -> Tsr { - _get_tr_space(mass, geom, space, None) + _get_tr_space(mass, geom, space, None, None) } -/// Internal helper with custom SVD tolerance. When `tol` is `Some(t)`, that -/// absolute value is used as the singular-value cutoff; when `None`, the default -/// machine-epsilon-based tolerance (`ndof × max(s) × ε`) is used. -fn _get_tr_space(mass: TsrView, geom: TsrView, space: &str, tol_user: Option) -> Tsr { +/// Internal helper. When `nrt_user` is `Some(k)`, the top-`k` singular +/// vectors are selected (rank set by the caller from rotor type). Otherwise +/// `tol_user` (`Some(t)` absolute cutoff, or `None` for the default +/// `ndof × max(s) × ε` tolerance) determines the rank. +fn _get_tr_space( + mass: TsrView, + geom: TsrView, + space: &str, + tol_user: Option, + nrt_user: Option, +) -> Tsr { let device = geom.device().clone(); let natm = geom.shape()[1]; let ndof = 3 * natm; @@ -134,14 +142,19 @@ fn _get_tr_space(mass: TsrView, geom: TsrView, space: &str, tol_user: Option t, + let num = match nrt_user { + Some(k) => k, None => { - let smax = svec.iter().copied().fold(0.0_f64, f64::max); - (ndof as f64) * smax * f64::EPSILON + let tol = match tol_user { + Some(t) => t, + None => { + let smax = svec.iter().copied().fold(0.0_f64, f64::max); + (ndof as f64) * smax * f64::EPSILON + }, + }; + svec.iter().filter(|&&x| x > tol).count() }, }; - let num = svec.iter().filter(|&&x| x > tol).count(); u.i((.., ..num)).to_owned() } @@ -323,24 +336,6 @@ fn phase_cols_to_max_element(q: TsrView, tol: f64) -> Tsr { out } -/// Check whether vector `vec` (length `n`) lies in the subspace spanned by the -/// columns of `space` (`[n, nrt]`), via SVD: `vec` is in the space iff stacking -/// it as an extra column does not increase the rank. -fn vec_in_space(vec: &[f64], space: TsrView, tol: f64) -> bool { - let device = space.device().clone(); - let nrt = space.shape()[1]; - let vec_t = rt::asarray((vec.to_vec(), &device)); // [n] - let mut cols: Vec = Vec::with_capacity(nrt + 1); - for c in 0..nrt { - cols.push(space.i((.., c)).to_owned()); // [n] - } - cols.push(vec_t); - let merged: Tsr = rt::stack((cols, -1)); // [n, nrt+1] - let (_u, s, _vh): (Tsr, Tsr, Tsr) = rt::linalg::svd(merged.view()).into(); - let svec = s.reshape(-1).to_vec(); - svec.last().copied().unwrap_or(0.0) < tol -} - // --------------------------------------------------------------------------- // harmonic_analysis // --------------------------------------------------------------------------- @@ -372,10 +367,41 @@ pub fn harmonic_analysis( let nmwhess = hess.into_contig(ColMajor); + // --------------- rotor type → number of TR dof --------------- + // The TR count is determined physically from the rotor type (3 for an + // atom, 5 for a linear molecule, 6 otherwise), rather than by an SVD + // rank cutoff on the (numerically truncated) TR basis. This mirrors + // PySCF's thermo.harmonic_analysis and makes the number of projected-out + // directions self-consistent with the TR/V classification below. + let geom_cm = { + let mass_sum: f64 = mass.reshape(-1).to_vec().iter().sum(); + // mc = Σ m_a r_a / Σ m_a -> [3]; reshape to [3, 1] to broadcast over natm + let mc = (geom.view() * mass.i((None, ..))).sum_axes(1) / mass_sum; // [3] + geom.view() - mc.i((.., None)) // [3, natm] + }; + let rc_ghz = rotation_const(mass.view(), geom_cm.view(), "GHz"); + let rotor_type = get_rotor_type(rc_ghz.view()); + let mut nrt = match rotor_type { + "ATOM" => 0, + "LINEAR" => 5, + _ => 6, + }; + // translation is always 3 dof; rotations are 0/2/3 for ATOM/LINEAR/REGULAR. + if !project_trans { + nrt = (nrt - 3).max(0); + } + if !project_rot { + let nrot = match rotor_type { + "ATOM" => 0, + "LINEAR" => 2, + _ => 3, + }; + nrt = (nrt - nrot).max(0); + } + // --------------- translation / rotation projector --------------- let space = format!("{}{}", if project_trans { "T" } else { "" }, if project_rot { "R" } else { "" }); - // use LINEAR_A_TOL so nearly-linear geometries correctly drop to 5 TR dof - let tr_space = _get_tr_space(mass.view(), geom.view(), &space, Some(LINEAR_A_TOL)); // [ndof, nrt] + let tr_space = _get_tr_space(mass.view(), geom.view(), &space, None, Some(nrt)); // [ndof, nrt] // projector P = I - Σ |tr⟩⟨tr| let nrt = tr_space.shape()[1]; @@ -471,13 +497,24 @@ pub fn harmonic_analysis( } // --------------- TR / V classification --------------- - // vec_in_space(qL[:, i], tr_space rows) + // After projection P = I - Σ|tr⟩⟨tr|, the TR directions lie in null(P) + // and have machine-zero force constants (~1e-13), separated from every + // genuine vibrational eigenvalue (a 10 cm⁻¹ mode is already ~2e-9) by + // many orders of magnitude. So the nrt modes with smallest + // |force_constant| are the TR modes; the rest are vibrations. This is + // more robust than the per-mode SVD membership test (vec_in_space) and + // its arbitrary 1e-4 tolerance. + let mut order_by_fc: Vec = (0..ndof).collect(); + order_by_fc.sort_by(|&a, &b| fc_sorted[a].abs().partial_cmp(&fc_sorted[b].abs()).unwrap()); + let mut is_tr = vec![false; ndof]; + for &i in order_by_fc.iter().take(nrt) { + is_tr[i] = true; + } let mut trv: Vec<&'static str> = Vec::with_capacity(ndof); for i in 0..ndof { - let qcol: Vec = (0..ndof).map(|r| qL[[r, i]]).collect(); - if vec_in_space(&qcol, tr_space.view(), 1.0e-4) { + if is_tr[i] { trv.push("TR"); - } else if omega_real[i].abs() < 1.0e-3 { + } else if fc_sorted[i].abs() < 1.0e-3 { trv.push("-"); } else { trv.push("V"); @@ -863,18 +900,8 @@ pub fn print_vibs( ncprec: Option, ) -> String { let nat = vib.ndof / 3; - let active: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] != "TR").collect(); - - // print inactive modes (TR) first, if any - println!("Inactive modes:"); - let inactive: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] == "TR").collect(); - for &i in &inactive { - let trv = vib.trv[i]; - let freq = if vib.imag[i] { format!("{:.4}i", vib.omega[i]) } else { format!("{:.4}", vib.omega[i]) }; - println!("{:<2} Freq [cm^-1]: {}", trv, freq); - } - println!(""); - println!("Active modes:"); + let active: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] == "V").collect(); + let presp = 2; let prewidth = 24usize; let colsp = 2usize; @@ -929,18 +956,17 @@ pub fn print_vibs( // Irrep — skip (no irrep/symmetry in our implementation) // scalar property rows — centered, colsp trailing - let labels: [&str; 5] = [ - "Reduced mass [u]", - "Force const [mDyne/A]", - "Turning point v=0 [a0]", - "RMS dev v=0 [a0 u^1/2]", - "Char temp [K]", - ]; + let labels: [&str; 5] = + ["Reduced mass [u]", "Force const [mDyne/A]", "Turning point v=0 [a0]", "RMS dev v=0 [a0 u^1/2]", "Char temp [K]"]; let val_slices: [&[f64]; 5] = [&vib.mu, &vib.k, &vib.xtp0, &vib.dq0, &vib.theta_vib]; for (label, vals) in labels.iter().zip(val_slices.iter()) { let mut l = format!("{}{}", " ".repeat(presp), ljust(label, prewidth)); for &vib in &row { - l.push_str(&format!("{}{}", center(&format!("{:.*}", prec, vals[vib]), width), " ".repeat(colsp))); + l.push_str(&format!( + "{}{}", + center(&format!("{:.*}", prec, vals[vib]), width), + " ".repeat(colsp) + )); } lines.push(l); } @@ -971,11 +997,20 @@ pub fn print_vibs( let lbl = if at < atom_lbl.len() { atom_lbl[at] } else { "" }; for xyz in 0..3 { let axis = ['X', 'Y', 'Z'][xyz]; - let mut l = - format!("{}{:5} {} {}", " ".repeat(presp), at + 1, axis, ljust(lbl, prewidth - 14)); + let mut l = format!( + "{}{:5} {} {}", + " ".repeat(presp), + at + 1, + axis, + ljust(lbl, prewidth - 14) + ); for &vib in &row { let v = normco_t[[3 * at + xyz, vib]]; - l.push_str(&format!("{}{}", center(&format!("{:.*}", ncprec, v), width), " ".repeat(colsp))); + l.push_str(&format!( + "{}{}", + center(&format!("{:.*}", ncprec, v), width), + " ".repeat(colsp) + )); } lines.push(l); } -- Gitee From 6f7c72a2b4e1b07185b5ae75850c33a14f37d99e Mon Sep 17 00:00:00 2001 From: ajz34 Date: Fri, 26 Jun 2026 18:04:43 +0800 Subject: [PATCH 11/39] vib: fix non-active mode classification, print non-active frequencies --- src/hessian_backup/vib.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/hessian_backup/vib.rs b/src/hessian_backup/vib.rs index 3a0515be56..7f899eb447 100644 --- a/src/hessian_backup/vib.rs +++ b/src/hessian_backup/vib.rs @@ -514,7 +514,7 @@ pub fn harmonic_analysis( for i in 0..ndof { if is_tr[i] { trv.push("TR"); - } else if fc_sorted[i].abs() < 1.0e-3 { + } else if fc_sorted[i].abs().sqrt() < 1.0e-3 { trv.push("-"); } else { trv.push("V"); @@ -902,6 +902,17 @@ pub fn print_vibs( let nat = vib.ndof / 3; let active: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] == "V").collect(); + // print inactive modes (TR) first, if any + println!("Inactive modes:"); + let inactive: Vec = (0..vib.ndof).filter(|&i| vib.trv[i] != "V").collect(); + for &i in &inactive { + let trv = vib.trv[i]; + let freq = if vib.imag[i] { format!("{:.4}i", vib.omega[i]) } else { format!("{:.4}", vib.omega[i]) }; + println!("{:<2} Freq [cm^-1]: {}", trv, freq); + } + println!(""); + println!("Active modes:"); + let presp = 2; let prewidth = 24usize; let colsp = 2usize; -- Gitee From 312f525856250606526845a58e07a973f00d8466 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Fri, 26 Jun 2026 20:50:35 +0800 Subject: [PATCH 12/39] hessian_backup: add ctrl-tunable config and cphf/skeleton grid levels Make HessSCFConfig control-input tunable via a `hessian_backup` table in ctrl.in (serde path, mirroring j2c_decomp), and wire the two new grid-level options: - grid_level_cphf: activates a dedicated (coarser) CP-KS grid for the iterative response; default = grid_gen_level.max(3) - 2. vxc/fxc are recomputed on this grid from the occupied MOs via a bra-ket contraction (lean make_cpks_vxc_fxc[_uks], no full dm0, no skeleton intermediates). - grid_level_skeleton: regenerates the skeleton grid; LDA/GGA reuse the SCF DFT grid, MGGA (TAU) adds 2 levels by default. Add Grids::build_with_level and tests covering the ser/de parse, the grid-level paths (RKS B3LYP, UKS TPSS0/TPSSh), and a braket-vs-dm0 equivalence check for the lean cpks vxc/fxc. Co-authored-by: Claude Code Co-authored-by: glm-5 --- src/ctrl_io/mod.rs | 4 + src/dft/mod.rs | 59 +++++++ src/dft/numint_matmul/hess_rks.rs | 81 +++++++++- src/dft/numint_matmul/hess_uks.rs | 83 +++++++++- src/hessian_backup/config.rs | 14 ++ src/hessian_backup/rscf_interface.rs | 37 ++++- src/hessian_backup/uscf_interface.rs | 37 ++++- src/main_driver.rs | 7 +- tests/hessian_backup/ctrl_parse.rs | 55 +++++++ tests/hessian_backup/mod.rs | 5 +- tests/hessian_backup/rks_b3lyp_grid_levels.rs | 144 ++++++++++++++++++ tests/hessian_backup/uks_tpss0_grid_levels.rs | 80 ++++++++++ .../{uks_tpss0.rs => uks_tpssh.rs} | 18 ++- 13 files changed, 603 insertions(+), 21 deletions(-) create mode 100644 tests/hessian_backup/ctrl_parse.rs create mode 100644 tests/hessian_backup/rks_b3lyp_grid_levels.rs create mode 100644 tests/hessian_backup/uks_tpss0_grid_levels.rs rename tests/hessian_backup/{uks_tpss0.rs => uks_tpssh.rs} (71%) diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index 12535815bf..701033d47e 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -8,6 +8,7 @@ use std::{fs, sync::Arc}; 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::hessian_backup::config::HessSCFConfig; use crate::ctrl_io::tddft_parameters::parse_tddft_keywords; use crate::ctrl_io::cphf_parameters::parse_cphf_keywords; use crate::{check_norm::force_state_occupation::ForceStateOccupation}; @@ -363,6 +364,7 @@ pub struct InputKeywords { pub use_fxc_opt: bool, /// Use module `hessian_backup` instead of `hessian` for Hessian calculation. pub use_hessian_backup: bool, + pub hessian_backup: HessSCFConfig, } impl Default for InputKeywords { @@ -523,6 +525,7 @@ impl InputKeywords { cphf: None, use_fxc_opt: false, use_hessian_backup: false, + hessian_backup: HessSCFConfig::default(), } } @@ -1821,6 +1824,7 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result tmp_str.to_lowercase().parse().unwrap_or(false), _ => false, }; + tmp_input.hessian_backup = tmp_ctrl.get("hessian_backup").map(serde_from_value).unwrap_or_default(); //=========================================================== // Global check of ctrl keywords and futher modification diff --git a/src/dft/mod.rs b/src/dft/mod.rs index 9a798acbd7..baf6723734 100644 --- a/src/dft/mod.rs +++ b/src/dft/mod.rs @@ -3476,6 +3476,65 @@ impl Grids { } } + /// Build a grid using the same atom-grid machinery as [`Grids::build`], but with an explicit + /// `level` override (instead of `mol.ctrl.grid_gen_level`). + /// + /// This is used by the `hessian_backup` module to generate skeleton / cphf grids at levels + /// that may differ from the SCF DFT grid. Unlike [`Grids::build`], this skips the + /// `external_grids` path and MPI distribution, and does not apply the round-robin permutation. + pub fn build_with_level(mol: &Molecule, level: usize) -> Grids { + let radial_precision = mol.ctrl.radial_precision; + let min_num_angular_points: usize = mol.ctrl.min_num_angular_points; + let max_num_angular_points: usize = mol.ctrl.max_num_angular_points; + let hardness: usize = mol.ctrl.hardness; + let pruning: String = mol.ctrl.pruning.clone(); + let rad_grid_method: String = mol.ctrl.rad_grid_method.clone(); + + let mass_charge = get_mass_charge(&mol.geom.rg_elem); + let proton_charges: Vec = mass_charge.iter().map(|value| value.1 as i32).collect(); + let center_coordinates_bohr = mol.geom.to_numgrid_io(); + let mut alpha_max: Vec = vec![]; + let mut alpha_min: Vec> = vec![]; + mol.basis4elem.iter().for_each(|value| { + let (tmp_alpha_min, tmp_alpha_max) = value.to_numgrid_io(); + alpha_max.push(tmp_alpha_max); + alpha_min.push(tmp_alpha_min); + }); + + let mut coordinates: Vec<[f64; 3]> = vec![]; + let mut weights: Vec = vec![]; + alpha_min.iter().zip(alpha_max.iter()).enumerate().for_each(|(center_index, value)| { + let (rs_atom, ws_atom) = gen_grids::atom_grid( + value.0.clone(), + value.1.clone(), + radial_precision, + min_num_angular_points, + max_num_angular_points, + proton_charges.clone(), + center_index, + center_coordinates_bohr.clone(), + hardness, + pruning.clone(), + rad_grid_method.clone(), + level, + ); + coordinates.extend(rs_atom.iter().map(|value| [value.0, value.1, value.2])); + weights.extend(ws_atom); + }); + + Grids { + weights, + coordinates, + ao: None, + aop: None, + parallel_balancing: Vec::new(), + non0tab: None, + ao_cutoff: 0.0, + ao_compressed: None, + aop_compressed: None, + } + } + pub fn formated_output(&self) { self.coordinates.iter().zip(self.weights.iter()).for_each(|value| { println!("r: ({:6.3},{:6.3},{:6.3}), w: {:16.8}",value.0[0],value.0[1],value.0[2],value.1); diff --git a/src/dft/numint_matmul/hess_rks.rs b/src/dft/numint_matmul/hess_rks.rs index 836894b1c4..22499bd8e9 100644 --- a/src/dft/numint_matmul/hess_rks.rs +++ b/src/dft/numint_matmul/hess_rks.rs @@ -203,6 +203,28 @@ pub fn get_rho_vxc_fxc(xc_func_list: &[(f64, LibXCFunctional)], ao: TsrView, ao_ + rt::vecdot(index!(ao, Z), index!(ao_dm0, Z), 1)) } + let (vxc, fxc) = eval_vxc_fxc_from_rho(xc_func_list, rho.view()); + (rho, vxc, fxc) +} + +/// Evaluate the (spin-unpolarized) XC potential `vxc` and kernel `fxc` from a given ground-state +/// `rho`, by summing each sub-functional's `libxc_eval_eff` (deriv = 2) contribution. +/// +/// `rho` : shape `[ngrids, nvar]`, where `nvar` is the strictest density-variable count across the +/// functional list. Extracted from [`get_rho_vxc_fxc`] so that callers which already have `rho` +/// (e.g. obtained directly from occupied orbitals via a bra-ket contraction) can skip the +/// `ao_dm0`-based density formation. +pub fn eval_vxc_fxc_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], &device)); let mut fxc = rt::zeros(([ngrids, nvar, nvar], &device)); for (scale, xc_func) in xc_func_list { @@ -216,8 +238,42 @@ pub fn get_rho_vxc_fxc(xc_func_list: &[(f64, LibXCFunctional)], ao: TsrView, ao_ *&mut vxc.i_mut((.., ..nvar_i)) += *scale * vxc_i; *&mut fxc.i_mut((.., ..nvar_i, ..nvar_i)) += *scale * fxc_i; } + (vxc, fxc) +} - (rho, 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. +/// +/// Compared to [`make_hessian_setup_batched`], this: +/// - skips all skeleton-Hessian intermediates (`de_fxc`, `de_vxc_diag`, `de_vxc_off`, `vmat_ip`, +/// `vmat_deriv1`); +/// - evaluates AO at the minimum derivative order needed to form the density +/// (`xc_type.num_ao_deriv()`, i.e. 0 for LDA / 1 for GGA, MGGA) instead of the Hessian derivative +/// order (`get_hess_ao_deriv`, 2 / 3); +/// - forms the ground-state density `rho` directly from the occupied MO coefficients via a bra-ket +/// contraction ([`NIMatmul::make_rho_from_homogeneous_braket`]) instead of building the full +/// `[nao, nao]` `dm0` matrix, exploiting the low rank of the occupied space (consistent with the +/// CP-KS response path in [`get_rks_response_bra`]). +/// +/// The CP-KS grid is assumed coarse enough that the full-grid AO tensor fits in memory; no batching +/// is performed. +pub fn make_cpks_vxc_fxc( + xc_func_list: &[(f64, LibXCFunctional)], + ni: &mut NIMatmul, + mo_coeff: TsrView, + mo_occ: TsrView, +) -> (Tsr, Tsr) { + let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); + // bake the occupation into the occupied coefficients (sqrt so the bra-ket square reproduces the + // occ-weighted density for every density component: rho, sigma, tau are all bilinear in phi). + 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, 1] (single set) -> [ngrids, nvar] + let rho = ni.make_rho_from_homogeneous_braket(&[mocc_2.view()], xc_type); + eval_vxc_fxc_from_rho(xc_func_list, rho.i((.., .., 0))) } pub fn get_drho(xc_type: XCDenType, ao: TsrView, ao_dm0: TsrView, aoslices: &[[usize; 4]]) -> Tsr { @@ -785,6 +841,16 @@ impl<'a> RHessKSNIMatmul<'a> { } } + /// 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 + } + /// Perform the Hessian setup for RKS calculations. /// /// Some intermediates: @@ -848,6 +914,19 @@ impl<'a> RHessElecInteractAPI for RHessKSNIMatmul<'a> { fn make_response_preparation(&mut self, mo_coeff: TsrView, mo_occ: TsrView) { self.intmd.insert("mo_coeff".to_string(), mo_coeff.into_contig(ColMajor)); self.intmd.insert("mo_occ".to_string(), mo_occ.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 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). + if let Some(ni_cpks) = self.ni_cpks.as_mut() { + 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_cpks, mo_coeff, 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) -> Tsr { diff --git a/src/dft/numint_matmul/hess_uks.rs b/src/dft/numint_matmul/hess_uks.rs index 1467130f9b..1ef5435ba2 100644 --- a/src/dft/numint_matmul/hess_uks.rs +++ b/src/dft/numint_matmul/hess_uks.rs @@ -95,6 +95,27 @@ pub fn get_rho_vxc_fxc_uks( } // Evaluate vxc and fxc with spin-polarized libxc + let (vxc, fxc) = eval_vxc_fxc_uks_from_rho(xc_func_list, rho.view()); + (rho, 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 [`get_rho_vxc_fxc_uks`] 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 { @@ -106,8 +127,46 @@ pub fn get_rho_vxc_fxc_uks( *&mut vxc.i_mut((.., ..nvar_i, ..)) += *scale * vxc_i; *&mut fxc.i_mut((.., ..nvar_i, .., ..nvar_i, ..)) += *scale * fxc_i; } + (vxc, fxc) +} - (rho, 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::hess_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()); + let nvar = xc_type.num_nvar(); + + // 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()) } pub fn get_drho_uks( @@ -640,6 +699,16 @@ impl<'a> UHessKSNIMatmul<'a> { } } + /// 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 + } + 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(), mo_occ[β].view().greater(0).into_vec()]; let mocc_α = mo_coeff[α].bool_select(-1, &occidx[α]); @@ -712,6 +781,18 @@ impl<'a> UHessElecInteractAPI for UHessKSNIMatmul<'a> { 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() { + let (vxc, fxc) = make_cpks_vxc_fxc_uks(&self.xc_func_list, ni_cpks, mo_coeff, 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] { diff --git a/src/hessian_backup/config.rs b/src/hessian_backup/config.rs index d942c261cc..47341d228b 100644 --- a/src/hessian_backup/config.rs +++ b/src/hessian_backup/config.rs @@ -18,6 +18,18 @@ pub struct HessSCFConfig { pub verbose: Option, #[serde_inline_default(None)] pub atm_list: Option>, + /// Grid level for the CP-KS response calculation. + /// + /// By default, the CP-KS grid level is set to `grid_gen_level.max(3) - 2` (much coarser than the SCF grid). + #[serde_inline_default(None)] + pub grid_level_cphf: Option, + /// Grid level for the skeleton grid used to evaluate the XC potential and kernel. + /// + /// By default, the skeleton grid level is set to + /// - `grid_gen_level` for LDA/GGA functionals + /// - `grid_gen_level + 2` for MGGA (TAU) functionals. + #[serde_inline_default(None)] + pub grid_level_skeleton: Option, } impl Default for HessSCFConfig { @@ -30,6 +42,8 @@ impl Default for HessSCFConfig { cphf_lindep: 1e-14, verbose: None, atm_list: None, + grid_level_cphf: None, + grid_level_skeleton: None, } } } diff --git a/src/hessian_backup/rscf_interface.rs b/src/hessian_backup/rscf_interface.rs index 4226d38338..4f6efc220f 100644 --- a/src/hessian_backup/rscf_interface.rs +++ b/src/hessian_backup/rscf_interface.rs @@ -63,12 +63,10 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, use crate::dft::numint_matmul::hess_rks::RHessKSNIMatmul; use crate::dft::numint_matmul::nimatmul::NIMatmul; use crate::dft::numint_matmul::prelude::*; + use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; + use crate::dft::Grids; use libxc::prelude::*; - let grid_coords = &scf_data.grids.as_ref().unwrap().coordinates; - let grid_weights = &scf_data.grids.as_ref().unwrap().weights; - let ni = NIMatmul::new(&mol, grid_coords, grid_weights); - 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; @@ -79,7 +77,36 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, .collect_vec() }; let verbose = scf_data.mol.ctrl.print_level >= 2; - RHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose) + + // Determine skeleton / cphf grid levels. + // - skeleton: LDA/GGA use the SCF DFT grid; MGGA (TAU) adds 2 levels. + // - cphf: grid_gen_level.max(3) - 2 (coarser, for the iterative CP-KS 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 sk_level = config.grid_level_skeleton.unwrap_or(if is_mgga { grid_gen_level + 2 } else { grid_gen_level }); + let cphf_level = config.grid_level_cphf.unwrap_or(grid_gen_level.max(3) - 2); + + // skeleton grid: reuse the SCF grid when the level matches, else regenerate. + let ni = if sk_level == grid_gen_level { + let grid_coords = &scf_data.grids.as_ref().unwrap().coordinates; + let grid_weights = &scf_data.grids.as_ref().unwrap().weights; + NIMatmul::new(&mol, grid_coords, grid_weights) + } else { + let sk_grid = Grids::build_with_level(mol_obj, sk_level); + NIMatmul::new(&mol, &sk_grid.coordinates, &sk_grid.weights) + }; + + // cphf 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 cphf_level == sk_level { + RHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose) + } else { + let cphf_grid = Grids::build_with_level(mol_obj, cphf_level); + let ni_cpks = NIMatmul::new(&mol, &cphf_grid.coordinates, &cphf_grid.weights); + RHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose).set_ni_cpks(ni_cpks) + }; + hess_nimatmul_obj }); if let Some(ref mut hess_nimatmul_obj) = hess_nimatmul_obj { hess_el_list.push(hess_nimatmul_obj); diff --git a/src/hessian_backup/uscf_interface.rs b/src/hessian_backup/uscf_interface.rs index 2cd9f9048e..f2466889a8 100644 --- a/src/hessian_backup/uscf_interface.rs +++ b/src/hessian_backup/uscf_interface.rs @@ -72,12 +72,10 @@ pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, use crate::dft::numint_matmul::hess_uks::UHessKSNIMatmul; use crate::dft::numint_matmul::nimatmul::NIMatmul; use crate::dft::numint_matmul::prelude::*; + use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; + use crate::dft::Grids; use libxc::prelude::*; - let grid_coords = &scf_data.grids.as_ref().unwrap().coordinates; - let grid_weights = &scf_data.grids.as_ref().unwrap().weights; - let ni = NIMatmul::new(&mol, grid_coords, grid_weights); - 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; @@ -88,7 +86,36 @@ pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, .collect_vec() }; let verbose = scf_data.mol.ctrl.print_level >= 2; - UHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose) + + // Determine skeleton / cphf grid levels. + // - skeleton: LDA/GGA use the SCF DFT grid; MGGA (TAU) adds 2 levels. + // - cphf: grid_gen_level.max(3) - 2 (coarser, for the iterative CP-KS 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 sk_level = config.grid_level_skeleton.unwrap_or(if is_mgga { grid_gen_level + 2 } else { grid_gen_level }); + let cphf_level = config.grid_level_cphf.unwrap_or(grid_gen_level.max(3) - 2); + + // skeleton grid: reuse the SCF grid when the level matches, else regenerate. + let ni = if sk_level == grid_gen_level { + let grid_coords = &scf_data.grids.as_ref().unwrap().coordinates; + let grid_weights = &scf_data.grids.as_ref().unwrap().weights; + NIMatmul::new(&mol, grid_coords, grid_weights) + } else { + let sk_grid = Grids::build_with_level(mol_obj, sk_level); + NIMatmul::new(&mol, &sk_grid.coordinates, &sk_grid.weights) + }; + + // cphf 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 cphf_level == sk_level { + UHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose) + } else { + let cphf_grid = Grids::build_with_level(mol_obj, cphf_level); + let ni_cpks = NIMatmul::new(&mol, &cphf_grid.coordinates, &cphf_grid.weights); + UHessKSNIMatmul::new(&mol, xc_func_list, ni, verbose).set_ni_cpks(ni_cpks) + }; + hess_nimatmul_obj }); if let Some(ref mut hess_nimatmul_obj) = hess_nimatmul_obj { hess_el_list.push(hess_nimatmul_obj); diff --git a/src/main_driver.rs b/src/main_driver.rs index 4da3c4a978..f5bf07bf8f 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -248,7 +248,6 @@ pub fn main_driver() -> anyhow::Result<()> { } else { use crate::hessian_backup::rscf_interface::rscf_hess_interface; use crate::hessian_backup::uscf_interface::uscf_hess_interface; - use crate::hessian_backup::config::HessSCFConfig; eprintln!("[WARN] You are using the hessian_backup module, which is still under development."); @@ -257,13 +256,13 @@ pub fn main_driver() -> anyhow::Result<()> { panic!("Normal modes calculation is currently not available for post-SCF methods."); } - let config = HessSCFConfig::default(); + let config = &scf_data.mol.ctrl.hessian_backup; match scf_data.scftype { SCFType::RHF => { - rscf_hess_interface(&scf_data, &config); + rscf_hess_interface(&scf_data, config); } SCFType::UHF => { - uscf_hess_interface(&scf_data, &config); + uscf_hess_interface(&scf_data, config); }, _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types.") } diff --git a/tests/hessian_backup/ctrl_parse.rs b/tests/hessian_backup/ctrl_parse.rs new file mode 100644 index 0000000000..24c617f2ec --- /dev/null +++ b/tests/hessian_backup/ctrl_parse.rs @@ -0,0 +1,55 @@ +//! Test that the `hessian_backup` keyword in `ctrl.in` is parsed into +//! `InputKeywords.hessian_backup` via the serde path (mirroring `j2c_decomp`). +//! +//! Only exercises the ctrl-block parser (`parse_ctrl_keywords`); deliberately does *not* build a +//! full `Molecule`, to avoid global thread-pool / basis-pool side effects that can perturb other +//! tests sharing the test binary. + +use pyrest::ctrl_io; + +#[test] +fn test_parse_hessian_backup_table() { + let input = r##" +[ctrl] + xc = "b3lyp" + basis_path = "def2-svp" + num_threads = 16 + use_hessian_backup = true + +[ctrl.hessian_backup] + cphf_tol = 1e-7 + cphf_max_cycle = 7 + grid_level_cphf = 2 + grid_level_skeleton = 4 + atm_list = [0, 1] +"##; + + let keys = toml::from_str::(input).unwrap(); + let ctrl = ctrl_io::parse_ctrl_keywords(&keys).unwrap(); + + assert!(ctrl.use_hessian_backup, "use_hessian_backup should be parsed"); + let cfg = &ctrl.hessian_backup; + assert!((cfg.cphf_tol - 1e-7).abs() < 1e-12); + assert_eq!(cfg.cphf_max_cycle, 7); + assert_eq!(cfg.grid_level_cphf, Some(2)); + assert_eq!(cfg.grid_level_skeleton, Some(4)); + assert_eq!(cfg.atm_list.as_deref(), Some(&[0usize, 1][..])); +} + +#[test] +fn test_hessian_backup_default_when_absent() { + let input = r##" +[ctrl] + xc = "b3lyp" + basis_path = "def2-svp" + num_threads = 16 +"##; + let keys = toml::from_str::(input).unwrap(); + let ctrl = ctrl_io::parse_ctrl_keywords(&keys).unwrap(); + + assert!(!ctrl.use_hessian_backup); + // defaults from HessSCFConfig::default() + assert_eq!(ctrl.hessian_backup.grid_level_cphf, None); + assert_eq!(ctrl.hessian_backup.grid_level_skeleton, None); + assert_eq!(ctrl.hessian_backup.cphf_max_cycle, 42); +} diff --git a/tests/hessian_backup/mod.rs b/tests/hessian_backup/mod.rs index 97a6a0e275..faf7e0ea35 100644 --- a/tests/hessian_backup/mod.rs +++ b/tests/hessian_backup/mod.rs @@ -1,4 +1,7 @@ +pub mod ctrl_parse; pub mod rhf; pub mod rks_b3lyp; +pub mod rks_b3lyp_grid_levels; pub mod uhf; -pub mod uks_tpss0; +pub mod uks_tpss0_grid_levels; +pub mod uks_tpssh; diff --git a/tests/hessian_backup/rks_b3lyp_grid_levels.rs b/tests/hessian_backup/rks_b3lyp_grid_levels.rs new file mode 100644 index 0000000000..59bfdc9e4f --- /dev/null +++ b/tests/hessian_backup/rks_b3lyp_grid_levels.rs @@ -0,0 +1,144 @@ +//! RKS (B3LYP, GGA) Hessian with explicit `grid_level_cphf` / `grid_level_skeleton`. +//! +//! Exercises the dedicated CP-KS grid path (`ni_cpks = Some`) and the skeleton-grid +//! regeneration path for a GGA functional. + +use pyrest::hessian_backup::config::HessSCFConfig; +use pyrest::hessian_backup::rscf_interface::rscf_hess_interface; + +use pyrest::ctrl_io; +use pyrest::dft::numint_matmul::hess_rks::{get_hess_ncomp_ao_dm0, get_rho_vxc_fxc, make_cpks_vxc_fxc}; +use pyrest::dft::numint_matmul::nimatmul::NIMatmul; +use pyrest::dft::xceff::prelude::determine_den_type_from_list; +use pyrest::molecule_io::Molecule; +use pyrest::scf_io::{self, scf_without_build}; + +use rstsr::prelude::*; + +static INPUT_NH3: &str = r##" +[ctrl] + print_level = 2 + num_threads = 16 + xc = "b3lyp" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 1.0 + spin_polarization = false + auxbasis_response = true + mixer = "diis" + num_max_diis = 8 + start_diis_cycle = 3 + mix_param = 0.8 + max_scf_cycle = 100 + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 1.0 0.1 0.2 + H 0.3 1.1 0.2 + H 0.1 0.1 1.2 + """ +"##; + +fn run_with_config(config: HessSCFConfig) -> Vec { + 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); + let (de, _, _) = rscf_hess_interface(&mut scf_data, &config); + de +} + +#[test] +fn test_nh3_explicit_grid_levels() { + // For GGA the default skeleton level equals grid_gen_level (3); force both grids to differ + // from each other and from the SCF grid so the regeneration + ni_cpks paths are exercised. + let config = HessSCFConfig { grid_level_skeleton: Some(4), grid_level_cphf: Some(2), ..Default::default() }; + let de = run_with_config(config); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian (explicit grid levels):\n{:12.6}", de.t()); + assert_eq!(de.shape().to_vec(), vec![3, 3, natm, natm]); +} + +#[test] +fn test_nh3_default_grid_levels() { + // Default config: skeleton = grid_gen_level (3), cphf = grid_gen_level.max(3) - 2 = 1. + // This activates the ni_cpks path with a coarse cphf grid. + let config = HessSCFConfig::default(); + let de = run_with_config(config); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian (default grid levels):\n{:12.6}", de.t()); + assert_eq!(de.shape().to_vec(), vec![3, 3, natm, natm]); +} + +#[test] +fn test_nh3_cphf_equals_skeleton() { + // cphf level == skeleton level (both 3) -> ni_cpks = None fast path (reuse skeleton vxc/fxc). + let config = HessSCFConfig { grid_level_skeleton: Some(3), grid_level_cphf: Some(3), ..Default::default() }; + let de = run_with_config(config); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian (cphf == skeleton):\n{:12.6}", de.t()); + assert_eq!(de.shape().to_vec(), vec![3, 3, natm, natm]); +} + +/// Numerical-equivalence check: the lean braket-based [`make_cpks_vxc_fxc`] (which forms rho from +/// occupied MOs via `make_rho_from_homogeneous_braket`) must produce the same `vxc` / `fxc` as the +/// dm0-based [`get_rho_vxc_fxc`] on the same grid. +#[test] +fn test_cpks_vxc_fxc_matches_dm0_path() { + use itertools::Itertools; + use libxc::prelude::*; + use pyrest::ri_jk::util::get_dm0_restricted; + + 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); + + let device = DeviceBLAS::default(); + let mo_coeff = + rt::asarray((&scf_data.eigenvectors[0].data, scf_data.eigenvectors[0].size, &device)).into_contig(ColMajor); + let mo_occ = rt::asarray((&scf_data.occupation[0], [scf_data.occupation[0].len()], &device)).into_contig(ColMajor); + + let mol_obj = &scf_data.mol; + let cint_mol = pyrest::ri_jk::util::get_cint_mol(mol_obj); + let grid_coords = &scf_data.grids.as_ref().unwrap().coordinates; + let grid_weights = &scf_data.grids.as_ref().unwrap().weights; + let mut ni = NIMatmul::new(&cint_mol, grid_coords, grid_weights); + + let xc_func_list: Vec<(f64, LibXCFunctional)> = scf_data + .mol + .xc_data + .dfa_compnt_scf + .iter() + .zip(scf_data.mol.xc_data.dfa_paramr_scf.iter()) + .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, LibXCSpin::Unpolarized))) + .collect(); + + // reference: dm0 -> ao_dm0 -> get_rho_vxc_fxc + let xc_type = determine_den_type_from_list(&xc_func_list.iter().map(|(_, f)| f).collect_vec()); + let ncomp_ao_dm0 = get_hess_ncomp_ao_dm0(xc_type); + let dm0 = get_dm0_restricted(mo_coeff.view(), mo_occ.view()); + let ao = ni.get_cached_ao(xc_type.num_ao_deriv()); + let ao_dm0 = ao.i((Ellipsis, ..ncomp_ao_dm0)) % &dm0; + let (_rho_ref, vxc_ref, fxc_ref) = get_rho_vxc_fxc(&xc_func_list, ao.view(), ao_dm0.view()); + + // lean: braket from mo_coeff/mo_occ (clear AO cache so the braket path re-evaluates freshly) + ni.cache_tensor.clear(); + let (vxc, fxc) = make_cpks_vxc_fxc(&xc_func_list, &mut ni, mo_coeff.view(), mo_occ.view()); + + assert!(rt::allclose(&vxc, &vxc_ref, None)); + assert!(rt::allclose(&fxc, &fxc_ref, None)); +} diff --git a/tests/hessian_backup/uks_tpss0_grid_levels.rs b/tests/hessian_backup/uks_tpss0_grid_levels.rs new file mode 100644 index 0000000000..492e2c549a --- /dev/null +++ b/tests/hessian_backup/uks_tpss0_grid_levels.rs @@ -0,0 +1,80 @@ +//! UKS (TPSS0, MGGA) Hessian with explicit `grid_level_cphf` / `grid_level_skeleton`. +//! +//! Exercises the MGGA skeleton-grid regeneration (`grid_gen_level + 2` default) and the +//! dedicated CP-KS grid path on the UKS side. + +use pyrest::hessian_backup::config::HessSCFConfig; +use pyrest::hessian_backup::uscf_interface::uscf_hess_interface; + +use pyrest::ctrl_io; +use pyrest::molecule_io::Molecule; +use pyrest::scf_io::{self, scf_without_build}; + +use rstsr::prelude::*; + +static INPUT_NH3: &str = r##" +[ctrl] + print_level = 2 + num_threads = 16 + xc = "TPSS0" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 2.0 + spin = 3.0 + spin_polarization = true + auxbasis_response = true + mixer = "diis" + num_max_diis = 8 + start_diis_cycle = 3 + mix_param = 0.8 + max_scf_cycle = 100 + xc_parser = "parse_xc" + +[geom] + name = "NH3" + unit = "Angstrom" + position = """ + N 0.0 0.0 0.0 + H 1.0 0.1 0.2 + H 0.3 1.1 0.2 + H 0.1 0.1 1.2 + """ +"##; + +fn run_with_config(config: HessSCFConfig) -> Vec { + 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); + let (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config); + de +} + +#[test] +fn test_nh3_mgga_default_skeleton() { + // Default skeleton level for MGGA = grid_gen_level + 2 = 5; cphf = 1 -> ni_cpks path. + let config = HessSCFConfig::default(); + let de = run_with_config(config); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian (MGGA default grid levels):\n{:12.6}", de.t()); + assert_eq!(de.shape().to_vec(), vec![3, 3, natm, natm]); +} + +#[test] +fn test_nh3_mgga_explicit_grid_levels() { + let config = HessSCFConfig { + grid_level_skeleton: Some(5), + grid_level_cphf: Some(2), + ..Default::default() + }; + let de = run_with_config(config); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian (MGGA explicit grid levels):\n{:12.6}", de.t()); + assert_eq!(de.shape().to_vec(), vec![3, 3, natm, natm]); +} diff --git a/tests/hessian_backup/uks_tpss0.rs b/tests/hessian_backup/uks_tpssh.rs similarity index 71% rename from tests/hessian_backup/uks_tpss0.rs rename to tests/hessian_backup/uks_tpssh.rs index 0a00aeb17c..ba7715e608 100644 --- a/tests/hessian_backup/uks_tpss0.rs +++ b/tests/hessian_backup/uks_tpssh.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_backup::uscf_interface::uscf_hess_interface; use pyrest::hessian_backup::config::HessSCFConfig; +use pyrest::hessian_backup::uscf_interface::uscf_hess_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; @@ -11,7 +11,7 @@ static INPUT_NH3: &str = r##" [ctrl] print_level = 2 num_threads = 16 - xc = "TPSS0" + xc = "TPSSh" basis_path = "def2-tzvp" auxbas_path = "def2-universal-jkfit" eri_type = "ri-v" @@ -39,6 +39,7 @@ static INPUT_NH3: &str = r##" #[test] fn test_nh3() { + // Note this test requires 10+ seconds to run. 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(); @@ -46,9 +47,18 @@ fn test_nh3() { scf_without_build(&mut scf_data, &None); let config = HessSCFConfig::default(); - let (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config); + let (de, vib, _th) = uscf_hess_interface(&mut scf_data, &config); let natm = 4; let de = rt::asarray((&de, [3, 3, natm, natm])); println!("Hessian:\n{:12.6}", de.t()); -} \ No newline at end of file + + // reference freqs from gaussian + // we allow positive frequencies to be in 2cm^-1 error for this case + let ref_freqs = [-5082.3895, -742.4607, 1195.1428, 2035.0250, 2250.3643, 3429.0955]; + for (k, &i) in vib.vib_indices().iter().enumerate() { + if !vib.imag[i] { + assert!((vib.omega[i] - ref_freqs[k]).abs() < 2.0, "freq {}: {} != {}", i, vib.omega[i], ref_freqs[k]); + } + } +} -- Gitee From 62395627a2bda32e3fe468b600955c2d964e893d Mon Sep 17 00:00:00 2001 From: ajz34 Date: Sun, 28 Jun 2026 15:22:59 +0800 Subject: [PATCH 13/39] hessian backup: ecp case tested --- tests/hessian_backup/rhf.rs | 59 +++++++++++++++++++++++++++++++++++-- 1 file changed, 57 insertions(+), 2 deletions(-) diff --git a/tests/hessian_backup/rhf.rs b/tests/hessian_backup/rhf.rs index 7f8e62a7a4..94824130a3 100644 --- a/tests/hessian_backup/rhf.rs +++ b/tests/hessian_backup/rhf.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_backup::rscf_interface::rscf_hess_interface; use pyrest::hessian_backup::config::HessSCFConfig; +use pyrest::hessian_backup::rscf_interface::rscf_hess_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; @@ -63,4 +63,59 @@ fn test_nh3() { let natm = 4; let de = rt::asarray((&de, [3, 3, natm, natm])); println!("Hessian:\n{:12.6}", de.t()); -} \ No newline at end of file +} + +static INPUT_SBH3_HBR: &str = r##" +[ctrl] + print_level = 2 + num_threads = 16 + xc = "hf" + basis_path = "def2-tzvp" + auxbas_path = "def2-universal-jkfit" + eri_type = "ri-v" + charge = 0.0 + spin = 1.0 + spin_polarization = false + auxbasis_response = true + mixer = "diis" + num_max_diis = 8 + start_diis_cycle = 3 + mix_param = 0.8 + max_scf_cycle = 100 + +[geom] + name = "BIH3" + unit = "Angstrom" + position = """ +Sb 0.00000000 0.00000000 0.71474217 +H 0.00000000 0.00000000 -2.10597083 +Br 0.00000000 0.00000000 -3.52882583 +H 0.69982699 1.21213591 1.64001815 +H 0.69982699 -1.21213591 1.64001815 +H -1.39965399 0.00000000 1.64001815 + """ +"##; + +#[test] +fn test_sbh3_hbr() { + let keys = toml::from_str::(&INPUT_SBH3_HBR[..]).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 config = HessSCFConfig::default(); + let (_, vib, _) = rscf_hess_interface(&mut scf_data, &config); + + // reference value from gaussian 16 + // we allow 1 cm^-1 difference due to RI-JK/conventional difference + let ref_freqs = [ + -169.4952, -169.4951, 65.6069, 333.1334, 333.1334, 871.2882, 922.0154, 922.0154, 2152.3976, 2152.3977, + 2166.9802, 2645.4606, + ]; + for (k, &i) in vib.vib_indices().iter().enumerate() { + // imag freq multiplies by negative + let f = if vib.imag[i] { -vib.omega[i] } else { vib.omega[i] }; + assert!((f - ref_freqs[k]).abs() < 1.0, "freq {}: {} != {}", i, f, ref_freqs[k]); + } +} -- Gitee From d34451b0ce9d4afbc367952a21118bdc90da0e6d Mon Sep 17 00:00:00 2001 From: ajz34 Date: Sun, 28 Jun 2026 18:31:41 +0800 Subject: [PATCH 14/39] hessian-backup: add symmetry detection module Co-authored-by: Claude Code Co-authored-by: glm-5.2 --- src/hessian_backup/mod.rs | 1 + src/hessian_backup/point_group_detect/bits.rs | 118 +++ .../point_group_detect/detect.rs | 34 + .../point_group_detect/elements.rs | 90 ++ src/hessian_backup/point_group_detect/geom.rs | 138 +++ .../point_group_detect/interface_to_rest.rs | 12 + .../point_group_detect/linalg.rs | 126 +++ .../point_group_detect/matrix.rs | 158 +++ src/hessian_backup/point_group_detect/mod.rs | 43 + .../point_group_detect/molecule.rs | 945 ++++++++++++++++++ src/hessian_backup/point_group_detect/vec3.rs | 158 +++ 11 files changed, 1823 insertions(+) create mode 100644 src/hessian_backup/point_group_detect/bits.rs create mode 100644 src/hessian_backup/point_group_detect/detect.rs create mode 100644 src/hessian_backup/point_group_detect/elements.rs create mode 100644 src/hessian_backup/point_group_detect/geom.rs create mode 100644 src/hessian_backup/point_group_detect/interface_to_rest.rs create mode 100644 src/hessian_backup/point_group_detect/linalg.rs create mode 100644 src/hessian_backup/point_group_detect/matrix.rs create mode 100644 src/hessian_backup/point_group_detect/mod.rs create mode 100644 src/hessian_backup/point_group_detect/molecule.rs create mode 100644 src/hessian_backup/point_group_detect/vec3.rs diff --git a/src/hessian_backup/mod.rs b/src/hessian_backup/mod.rs index 3234e73335..264a7fad19 100644 --- a/src/hessian_backup/mod.rs +++ b/src/hessian_backup/mod.rs @@ -39,6 +39,7 @@ pub mod vib; pub mod cint_handling; pub mod config; pub mod krylov_block; +pub mod point_group_detect; #[allow(unused_imports)] pub mod prelude { diff --git a/src/hessian_backup/point_group_detect/bits.rs b/src/hessian_backup/point_group_detect/bits.rs new file mode 100644 index 0000000000..1fda4dd0cc --- /dev/null +++ b/src/hessian_backup/point_group_detect/bits.rs @@ -0,0 +1,118 @@ +//! Bitwise encoding of the D2h subgroups and symmetry operations. Verbatim +//! port of `psi4/driver/qcdb/libmintspointgrp.py` (`SymmOps`, `PointGroups`, +//! `similar`, `bits_to_basic_name`). +//! +//! Each of the 8 D2h operations is a bit: +//! `E=0, C2_z=1, C2_y=2, C2_x=4, i=8, σ_xy=16, σ_xz=32, σ_yz=64` (ID=128 sentinel). +//! A point group is the OR of its members' bits. + +#![allow(dead_code)] + +pub mod symm_ops { + pub const E: u8 = 0; + pub const C2_Z: u8 = 1; + pub const C2_Y: u8 = 2; + pub const C2_X: u8 = 4; + pub const I: u8 = 8; + pub const SIGMA_XY: u8 = 16; + pub const SIGMA_XZ: u8 = 32; + pub const SIGMA_YZ: u8 = 64; + pub const ID: u8 = 128; +} + +pub mod point_groups { + use super::symm_ops::*; + pub const C1: u8 = E; + pub const CI: u8 = E | I; + pub const C2X: u8 = E | C2_X; + pub const C2Y: u8 = E | C2_Y; + pub const C2Z: u8 = E | C2_Z; + pub const CSZ: u8 = E | SIGMA_XY; + pub const CSY: u8 = E | SIGMA_XZ; + pub const CSX: u8 = E | SIGMA_YZ; + pub const D2: u8 = E | C2_X | C2_Y | C2_Z; + pub const C2VX: u8 = E | C2_X | SIGMA_XY | SIGMA_XZ; + pub const C2VY: u8 = E | C2_Y | SIGMA_XY | SIGMA_YZ; + pub const C2VZ: u8 = E | C2_Z | SIGMA_XZ | SIGMA_YZ; + pub const C2HX: u8 = E | C2_X | SIGMA_YZ | I; + pub const C2HY: u8 = E | C2_Y | SIGMA_XZ | I; + pub const C2HZ: u8 = E | C2_Z | SIGMA_XY | I; + pub const D2H: u8 = E | C2_X | C2_Y | C2_Z | I | SIGMA_XY | SIGMA_XZ | SIGMA_YZ; +} + +/// The 7 non-identity operations tested by `find_highest_point_group`, paired +/// with the diagonal of their 3×3 matrix `[d00, d11, d22]` (used for the +/// element-wise `naivemult` application). Order matches the reference. +pub fn tested_ops() -> [(u8, [f64; 3]); 7] { + use symm_ops::*; + [ + (C2_Z, [-1.0, -1.0, 1.0]), + (C2_Y, [-1.0, 1.0, -1.0]), + (C2_X, [1.0, -1.0, -1.0]), + (I, [-1.0, -1.0, -1.0]), + (SIGMA_XY, [1.0, 1.0, -1.0]), + (SIGMA_XZ, [1.0, -1.0, 1.0]), + (SIGMA_YZ, [-1.0, 1.0, 1.0]), + ] +} + +/// Simple (non-directional) Schoenflies symbol from bits, lowercase. +/// Reference: `bits_to_basic_name`. Returns one of +/// `c1, ci, c2, cs, d2, c2v, c2h, d2h`. +pub fn bits_to_basic_name(bits: u8) -> &'static str { + use point_groups::*; + match bits { + C1 => "c1", + CI => "ci", + C2X | C2Y | C2Z => "c2", + CSX | CSY | CSZ => "cs", + D2 => "d2", + C2VX | C2VY | C2VZ => "c2v", + C2HX | C2HY | C2HZ => "c2h", + D2H => "d2h", + _ => "c1", // reference raises; we degrade to c1 for unexpected combos + } +} + +/// All directionally-equivalent bit patterns for a given directional subgroup. +/// Reference: `similar`. Used only when a user overrides symmetry (not on the +/// bare detection path), but ported for completeness. +pub fn similar(bits: u8) -> Vec { + use point_groups::*; + let cs = [CSX, CSY, CSZ]; + let c2v = [C2VZ, C2VY, C2VX]; + let c2h = [C2HZ, C2HY, C2HX]; + let c2 = [C2Z, C2Y, C2X]; + if cs.contains(&bits) { + cs.to_vec() + } else if c2v.contains(&bits) { + c2v.to_vec() + } else if c2h.contains(&bits) { + c2h.to_vec() + } else if c2.contains(&bits) { + c2.to_vec() + } else { + vec![bits] // d2h, d2, ci, c1 are self-only + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn basic_names() { + use point_groups::*; + assert_eq!(bits_to_basic_name(C1), "c1"); + assert_eq!(bits_to_basic_name(C2VZ), "c2v"); + assert_eq!(bits_to_basic_name(D2H), "d2h"); + assert_eq!(bits_to_basic_name(CSX), "cs"); + } + + #[test] + fn similar_cs() { + use point_groups::*; + let s = similar(CSX); + assert!(s.contains(&CSX) && s.contains(&CSY) && s.contains(&CSZ)); + } +} diff --git a/src/hessian_backup/point_group_detect/detect.rs b/src/hessian_backup/point_group_detect/detect.rs new file mode 100644 index 0000000000..f9f1047ec9 --- /dev/null +++ b/src/hessian_backup/point_group_detect/detect.rs @@ -0,0 +1,34 @@ +//! Public detection API. + +use super::molecule::{SymmMolecule, Tmpl}; + +/// A detected point group. +#[derive(Debug, Clone)] +pub struct PointGroup { + /// Schoenflies symbol, e.g. `"C2v"`, `"D3d"`, `"Td"`, `"D_inf_h"`, `"S4"`. + pub full_name: String, + /// Rotational symmetry number σ (Sn yields n/2, so this may be fractional). + pub sigma: f64, + pub(crate) template: Tmpl, + pub(crate) n: u32, +} + +impl SymmMolecule { + /// Detect the point group and rotational symmetry number. + pub fn detect(&self) -> PointGroup { + let (template, n) = self.detect_inner(); + PointGroup { + full_name: template.full_name(n), + sigma: template.sigma(n), + template, + n, + } + } +} + +/// Convenience: detect the point-group symbol for a geometry (Bohr) given +/// element symbols. Uses default masses and tolerance. +pub fn detect_point_group(symbols: &[&str], geom: &[[f64; 3]]) -> String { + let m = SymmMolecule::new(symbols.iter().map(|s| s.to_string()).collect(), geom.to_vec()); + m.detect().full_name +} diff --git a/src/hessian_backup/point_group_detect/elements.rs b/src/hessian_backup/point_group_detect/elements.rs new file mode 100644 index 0000000000..ecea99e7f7 --- /dev/null +++ b/src/hessian_backup/point_group_detect/elements.rs @@ -0,0 +1,90 @@ +//! Element symbol ↔ atomic number, and most-abundant-isotope masses. +//! Masses are from qcelemental's NIST 2011 table (same source Psi4 uses). +//! +//! NOTE: per-atom mass overrides supplied via [`super::Molecule::with_masses`] +//! always take precedence; this table is only a fallback when no mass is +//! supplied. The symmetry test fixtures supply explicit masses. + +/// (symbol, Z, most-abundant-isotope mass) +const TABLE: &[(&str, u8, f64)] = &[ + ("H", 1, 1.00782503223), + ("He", 2, 4.00260325413), + ("Li", 3, 7.0160034366), + ("Be", 4, 9.012183065), + ("B", 5, 11.00930536), + ("C", 6, 12.0), + ("N", 7, 14.00307400443), + ("O", 8, 15.99491461957), + ("F", 9, 18.99840316273), + ("Ne", 10, 19.9924401762), + ("Na", 11, 22.989769282), + ("Mg", 12, 23.985041697), + ("Al", 13, 26.98153853), + ("Si", 14, 27.97692653465), + ("P", 15, 30.97376199842), + ("S", 16, 31.9720711744), + ("Cl", 17, 34.968852682), + ("Ar", 18, 39.9623831237), + ("K", 19, 38.9637064864), + ("Ca", 20, 39.962590863), + ("Sc", 21, 44.95590828), + ("Ti", 22, 47.94794198), + ("V", 23, 50.94395704), + ("Cr", 24, 51.94050623), + ("Mn", 25, 54.93804391), + ("Fe", 26, 55.93493633), + ("Co", 27, 58.93319429), + ("Ni", 28, 57.93534241), + ("Cu", 29, 62.92959772), + ("Zn", 30, 63.92914201), + ("Ga", 31, 68.9255735), + ("Ge", 32, 73.921177761), + ("As", 33, 74.92159457), + ("Se", 34, 79.9165218), + ("Br", 35, 78.9183376), + ("Kr", 36, 83.9114977282), + ("I", 53, 126.9044719), +]; + +/// Normalize a user-supplied symbol to the table form (first letter upper, +/// rest lower), matching Psi4's element lookup. "X" → dummy (Z=0). +fn normalize(s: &str) -> String { + let mut chars = s.chars(); + match chars.next() { + Some(first) => first.to_ascii_uppercase().to_string() + &chars.as_str().to_ascii_lowercase(), + None => String::new(), + } +} + +/// Atomic number for an element symbol, or `None` if unknown. "X" → 0 (dummy). +pub fn symbol_to_z(s: &str) -> Option { + let n = normalize(s); + if n == "X" { + return Some(0); + } + TABLE.iter().find(|(sym, _, _)| *sym == n).map(|(_, z, _)| *z) +} + +/// Default most-abundant-isotope mass for atomic number `z`, or `0.0` if +/// unknown (dummy `z==0` → 0.0). +pub fn z_to_mass(z: u8) -> f64 { + TABLE.iter().find(|(_, zz, _)| *zz == z).map(|(_, _, m)| *m).unwrap_or(0.0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lookups() { + assert_eq!(symbol_to_z("C"), Some(6)); + assert_eq!(symbol_to_z("c"), Some(6)); + assert_eq!(symbol_to_z("CL"), Some(17)); + assert_eq!(symbol_to_z("br"), Some(35)); + assert_eq!(symbol_to_z("X"), Some(0)); + assert_eq!(symbol_to_z("Uuo"), None); + assert!((z_to_mass(1) - 1.00782503223).abs() < 1e-12); + assert!((z_to_mass(8) - 15.99491461957).abs() < 1e-12); + assert_eq!(z_to_mass(0), 0.0); + } +} diff --git a/src/hessian_backup/point_group_detect/geom.rs b/src/hessian_backup/point_group_detect/geom.rs new file mode 100644 index 0000000000..8a2883c4e5 --- /dev/null +++ b/src/hessian_backup/point_group_detect/geom.rs @@ -0,0 +1,138 @@ +//! Geometry-level helpers operating on `N×3` point sets. Port of the free +//! functions in `psi4/driver/qcdb/libmintsmolecule.py`: +//! `matrix_3d_rotation`, `matrix_3d_rotation_Cn`, `equal_but_for_row_order`, +//! `atom_present_in_geom`, and the `atom_at_position` lookup. + +use super::matrix::{householder, points_apply, rodrigues}; + +/// Rotate a set of row-vectors `mat` about `axis` by `phi` radians. If `sn` is +/// true, additionally reflect through the plane perpendicular to `axis` +/// (improper rotation). Reference: `matrix_3d_rotation` (libmintsmolecule.py:3167). +pub fn matrix_3d_rotation(mat: &[[f64; 3]], axis: &[f64; 3], phi: f64, sn: bool) -> Vec<[f64; 3]> { + let r = rodrigues(axis, phi); + let mut rotated = points_apply(mat, &r); + if sn { + let h = householder(axis); + rotated = points_apply(&rotated, &h); + } + rotated +} + +/// Find the highest `n` such that a `Cn` (proper rotation if `reflect=false`, +/// improper `Sn` if `reflect=true`) about `axis` maps `coord` onto itself. +/// `max_cn_to_check` of 0 means "probe up to `coord.len()`". Reference: +/// `matrix_3d_rotation_Cn` (libmintsmolecule.py:3148). +pub fn matrix_3d_rotation_cn( + coord: &[[f64; 3]], + axis: &[f64; 3], + reflect: bool, + tol: f64, + max_cn_to_check: usize, +) -> usize { + let max_possible = if max_cn_to_check == 0 { coord.len() } else { max_cn_to_check }; + let mut cn = 1; // C1 always present + let two_pi = 2.0 * std::f64::consts::PI; + for n in 2..=max_possible { + let rotated = matrix_3d_rotation(coord, axis, two_pi / n as f64, reflect); + if equal_but_for_row_order(coord, &rotated, tol) { + cn = n; + } + } + cn +} + +/// Set equality of two `N×3` geometries ignoring row order: every row of `mat` +/// must match *some* row of `rhs` element-wise within `tol`. Reference: +/// `equal_but_for_row_order` (libmintsmolecule.py:3228). +pub fn equal_but_for_row_order(mat: &[[f64; 3]], rhs: &[[f64; 3]], tol: f64) -> bool { + 'outer: for m in 0..mat.len() { + for m_rhs in 0..rhs.len() { + let mut matched = true; + for n in 0..mat[m].len() { + if (mat[m][n] - rhs[m_rhs][n]).abs() > tol { + matched = false; + break; + } + } + if matched { + continue 'outer; // found a matching row for row m + } + } + return false; // no matching row for row m + } + true +} + +/// Is there an atom (row) in `geom` within `tol` of point `b` (Euclidean)? +/// Reference: `atom_present_in_geom` (libmintsmolecule.py:3136). +pub fn atom_present_in_geom(geom: &[[f64; 3]], b: &[f64; 3], tol: f64) -> bool { + for a in geom { + let d0 = a[0] - b[0]; + let d1 = a[1] - b[1]; + let d2 = a[2] - b[2]; + if (d0 * d0 + d1 * d1 + d2 * d2).sqrt() < tol { + return true; + } + } + false +} + +/// Index of the nearest atom to `b`, if within `tol`. Reference: +/// `Molecule.atom_at_position` (libmintsmolecule.py:1153) — nearest-neighbor +/// with squared-distance cutoff `tol*tol`. Returns `None` if none within tol. +pub fn atom_at_position(geom: &[[f64; 3]], b: &[f64; 3], tol: f64) -> Option { + if geom.is_empty() { + return None; + } + let mut best = 0usize; + let mut best_d2 = f64::INFINITY; + for (i, a) in geom.iter().enumerate() { + let d0 = a[0] - b[0]; + let d1 = a[1] - b[1]; + let d2 = a[2] - b[2]; + let d2 = d0 * d0 + d1 * d1 + d2 * d2; + if d2 < best_d2 { + best_d2 = d2; + best = i; + } + } + if best_d2 < tol * tol { + Some(best) + } else { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn row_order_equality() { + let a = [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]; + let b = [[2.0, 2.0, 2.0], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]; + assert!(equal_but_for_row_order(&a, &b, 1e-9)); + let c = [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [9.0, 9.0, 9.0]]; + assert!(!equal_but_for_row_order(&a, &c, 1e-9)); + } + + #[test] + fn cn_about_z_square() { + // four points on the unit circle in xy: C4 about z + let pts = [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, -1.0, 0.0], + ]; + let cn = matrix_3d_rotation_cn(&pts, &[0.0, 0.0, 1.0], false, 1e-9, 0); + assert_eq!(cn, 4); + } + + #[test] + fn atom_at_position_nearest() { + let g = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]; + assert_eq!(atom_at_position(&g, &[1.05, 0.0, 0.0], 0.1), Some(1)); + assert_eq!(atom_at_position(&g, &[5.0, 5.0, 5.0], 0.1), None); + } +} diff --git a/src/hessian_backup/point_group_detect/interface_to_rest.rs b/src/hessian_backup/point_group_detect/interface_to_rest.rs new file mode 100644 index 0000000000..1140ae8ded --- /dev/null +++ b/src/hessian_backup/point_group_detect/interface_to_rest.rs @@ -0,0 +1,12 @@ +use super::SymmMolecule; + +pub fn get_full_point_group_for_vib( + symbols: &[String], + masses: &[f64], + coords: &[[f64; 3]], + tol: f64, +) -> (String, f64) { + let mol = SymmMolecule::new(symbols.to_vec(), coords.to_vec()).with_masses(masses.to_vec()).with_tol(tol); + let pg = mol.detect(); + (pg.full_name, pg.sigma) +} diff --git a/src/hessian_backup/point_group_detect/linalg.rs b/src/hessian_backup/point_group_detect/linalg.rs new file mode 100644 index 0000000000..f34ee4e514 --- /dev/null +++ b/src/hessian_backup/point_group_detect/linalg.rs @@ -0,0 +1,126 @@ +//! 3×3 real-symmetric eigensolver. Verbatim port of Psi4's +//! `vecutil.diagonalize3x3symmat` (cyclic Jacobi, 50-iteration cap). +//! Pure std. Used to diagonalize the inertia tensor for principal axes. + +use super::matrix::Matrix3; + +/// Diagonalize a real symmetric 3×3 matrix. Returns `(eigenvalues, Q)` where +/// the **columns** of `Q` are the eigenvectors: eigenvector `i` is +/// `[Q[0][i], Q[1][i], Q[2][i]]`. (Matches the reference: `Q[r][p]`.) +pub fn diagonalize3x3symmat(m: &Matrix3) -> ([f64; 3], Matrix3) { + let mut a = *m; // working copy (only upper triangle is used) + let mut eig = super::matrix::identity(); + let mut w = [a[0][0], a[1][1], a[2][2]]; + + // (reference computes SQR(tr(A)) as `sd` but never uses it; omitted.) + + for n_iter in 0..50 { + // test for convergence: sum of |off-diagonal| + let mut so = 0.0; + for p in 0..3 { + for q in (p + 1)..3 { + so += a[p][q].abs(); + } + } + if so == 0.0 { + return (w, eig); + } + + let thresh = if n_iter < 4 { 0.2 * so / 9.0 } else { 0.0 }; + + for p in 0..3 { + for q in (p + 1)..3 { + let g = 100.0 * a[p][q].abs(); + if n_iter > 4 + && (w[p].abs() + g == w[p].abs()) + && (w[q].abs() + g == w[q].abs()) + { + a[p][q] = 0.0; + } else if a[p][q].abs() > thresh { + let h = w[q] - w[p]; + let t; + if h.abs() + g == h.abs() { + t = a[p][q] / h; + } else { + let theta = 0.5 * h / a[p][q]; + if theta < 0.0 { + t = -1.0 / ((1.0 + theta * theta).sqrt() - theta); + } else { + t = 1.0 / ((1.0 + theta * theta).sqrt() + theta); + } + } + let c = 1.0 / (1.0 + t * t).sqrt(); + let s = t * c; + let z = t * a[p][q]; + + a[p][q] = 0.0; + w[p] -= z; + w[q] += z; + + for r in 0..p { + let tr = a[r][p]; + a[r][p] = c * tr - s * a[r][q]; + a[r][q] = s * tr + c * a[r][q]; + } + for r in (p + 1)..q { + let tr = a[p][r]; + a[p][r] = c * tr - s * a[r][q]; + a[r][q] = s * tr + c * a[r][q]; + } + for r in (q + 1)..3 { + let tr = a[p][r]; + a[p][r] = c * tr - s * a[q][r]; + a[q][r] = s * tr + c * a[q][r]; + } + for r in 0..3 { + let tr = eig[r][p]; + eig[r][p] = c * tr - s * eig[r][q]; + eig[r][q] = s * tr + c * eig[r][q]; + } + } + } + } + } + + // not converged within 50 iterations; return best-effort (reference returns None) + (w, eig) +} + +#[cfg(test)] +mod tests { + use super::*; + use super::super::matrix::matmul; + + #[test] + fn diagonal_diag_matrix() { + let m = [[2.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, -3.0]]; + let (w, q) = diagonalize3x3symmat(&m); + // eigenvalues (unordered): 5, 2, -3 + let mut ws = w; + ws.sort_by(|a, b| a.partial_cmp(b).unwrap()); + assert!((ws[0] - (-3.0)).abs() < 1e-12); + assert!((ws[1] - 2.0).abs() < 1e-12); + assert!((ws[2] - 5.0).abs() < 1e-12); + // Q is orthogonal + let qt = super::super::matrix::transpose(&q); + let prod = matmul(&q, &qt); + for i in 0..3 { + for j in 0..3 { + let target = if i == j { 1.0 } else { 0.0 }; + assert!((prod[i][j] - target).abs() < 1e-12, "Q not orthogonal"); + } + } + } + + #[test] + fn diagonal_known_symmetric() { + // [[2,1,0],[1,2,0],[0,0,3]] -> eigenvalues 1, 3, 3 + let m = [[2.0, 1.0, 0.0], [1.0, 2.0, 0.0], [0.0, 0.0, 3.0]]; + let (w, _q) = diagonalize3x3symmat(&m); + let mut ws = w; + ws.sort_by(|a, b| a.partial_cmp(b).unwrap()); + assert!((ws[0] - 1.0).abs() < 1e-10); + assert!((ws[1] - 3.0).abs() < 1e-10); + assert!((ws[2] - 3.0).abs() < 1e-10); + } +} diff --git a/src/hessian_backup/point_group_detect/matrix.rs b/src/hessian_backup/point_group_detect/matrix.rs new file mode 100644 index 0000000000..0d019ca562 --- /dev/null +++ b/src/hessian_backup/point_group_detect/matrix.rs @@ -0,0 +1,158 @@ +//! 3×3 matrix helpers. Port of the matrix utilities used by Psi4's libmints +//! (`vecutil.py`: `zero`, `identity`, `mult`, `transpose`, `matadd`) plus the +//! Rodrigues rotation matrix and Householder reflection matrix used by +//! `matrix_3d_rotation`. Pure std. + +pub type Matrix3 = [[f64; 3]; 3]; + +#[inline] +pub fn zero() -> Matrix3 { + [[0.0; 3]; 3] +} + +#[inline] +pub fn identity() -> Matrix3 { + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] +} + +#[inline] +pub fn transpose(m: &Matrix3) -> Matrix3 { + [ + [m[0][0], m[1][0], m[2][0]], + [m[0][1], m[1][1], m[2][1]], + [m[0][2], m[1][2], m[2][2]], + ] +} + +/// 3×3 · 3×3 matmul. Reference: `vecutil.mult`. +pub fn matmul(a: &Matrix3, b: &Matrix3) -> Matrix3 { + let mut r = zero(); + for i in 0..3 { + for j in 0..3 { + let mut s = 0.0; + for k in 0..3 { + s += a[i][k] * b[k][j]; + } + r[i][j] = s; + } + } + r +} + +/// N×3 (rows are points) · 3×3 -> N×3. Reference: `vecutil.mult` with the +/// `TypeError` vector fallback not needed since we always pass 3 columns. +pub fn points_matmul(points: &[[f64; 3]], m: &Matrix3) -> Vec<[f64; 3]> { + let mut r = vec![[0.0; 3]; points.len()]; + for i in 0..points.len() { + for j in 0..3 { + let mut s = 0.0; + for k in 0..3 { + s += points[i][k] * m[k][j]; + } + r[i][j] = s; + } + } + r +} + +/// N×3 (rows) · 3×3ᵀ -> N×3, i.e. each row `r` becomes `m · r` (treating rows +/// as column vectors transformed by `m`). This is `coord · Rᵀ` in the reference +/// (`matrix_3d_rotation`), where `R` is the row-vector rotation matrix. +pub fn points_apply(points: &[[f64; 3]], m: &Matrix3) -> Vec<[f64; 3]> { + let mut r = vec![[0.0; 3]; points.len()]; + for i in 0..points.len() { + for j in 0..3 { + let mut s = 0.0; + for k in 0..3 { + s += m[j][k] * points[i][k]; + } + r[i][j] = s; + } + } + r +} + +/// Determinant of a 3×3 matrix. Reference: `vecutil.determinant`. +pub fn determinant(m: &Matrix3) -> f64 { + m[0][0] * m[1][1] * m[2][2] - m[0][2] * m[1][1] * m[2][0] + + m[0][1] * m[1][2] * m[2][0] + - m[0][1] * m[1][0] * m[2][2] + + m[0][2] * m[1][0] * m[2][1] + - m[0][0] * m[1][2] * m[2][1] +} + +/// Rodrigues rotation matrix for rotation by `phi` about `axis` (need not be +/// normalized). Returns the row-vector form `R` from `matrix_3d_rotation` +/// (`libmintsmolecule.py:3167`); the caller applies it as `coord · Rᵀ` via +/// [`points_apply`]. +/// +/// **Careful:** the off-diagonal signs follow the reference's asymmetric +/// pattern, not the symmetric textbook form. +pub fn rodrigues(axis: &[f64; 3], phi: f64) -> Matrix3 { + let w = super::vec3::normalize(axis); + let (wx, wy, wz) = (w[0], w[1], w[2]); + let cp = 1.0 - phi.cos(); + let s = phi.sin(); + let c = phi.cos(); + [ + [wx * wx * cp + c, wx * wy * cp - s * wz, wx * wz * cp + s * wy], + [wy * wx * cp + s * wz, wy * wy * cp + c, wy * wz * cp - s * wx], + [wz * wx * cp - s * wy, wz * wy * cp + s * wx, wz * wz * cp + c], + ] +} + +/// Householder reflection matrix `H = I - 2 w wᵀ` for reflection through the +/// plane perpendicular to `axis` (normalized internally). Used for the improper +/// (Sn) branch of `matrix_3d_rotation`. Reference: `libmintsmolecule.py:3201`. +pub fn householder(axis: &[f64; 3]) -> Matrix3 { + let w = super::vec3::normalize(axis); + let (wx, wy, wz) = (w[0], w[1], w[2]); + let mut h = identity(); + h[0][0] -= 2.0 * wx * wx; + h[1][1] -= 2.0 * wy * wy; + h[2][2] -= 2.0 * wz * wz; + h[1][0] += 2.0 * wx * wy; + h[2][0] += 2.0 * wx * wz; + h[2][1] += 2.0 * wy * wz; + h[0][1] += 2.0 * wx * wy; + h[0][2] += 2.0 * wx * wz; + h[1][2] += 2.0 * wy * wz; + h +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn matmul_identity() { + let m = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]; + let id = identity(); + let p = matmul(&m, &id); + for i in 0..3 { + for j in 0..3 { + assert!((p[i][j] - m[i][j]).abs() < 1e-15); + } + } + } + + #[test] + fn rodrigues_about_z_90() { + // rotating the x-axis (as a row) about z by 90° should give +y + let r = rodrigues(&[0.0, 0.0, 1.0], std::f64::consts::FRAC_PI_2); + let out = points_apply(&[[1.0, 0.0, 0.0]], &r); + assert!((out[0][0]).abs() < 1e-15); + assert!((out[0][1] - 1.0).abs() < 1e-15); + assert!((out[0][2]).abs() < 1e-15); + } + + #[test] + fn householder_z_flips_xy() { + // reflection through plane perp to z (the xy-plane) flips z, keeps x,y + let h = householder(&[0.0, 0.0, 1.0]); + let out = points_apply(&[[1.0, 2.0, 3.0]], &h); + assert!((out[0][0] - 1.0).abs() < 1e-15); + assert!((out[0][1] - 2.0).abs() < 1e-15); + assert!((out[0][2] + 3.0).abs() < 1e-15); + } +} diff --git a/src/hessian_backup/point_group_detect/mod.rs b/src/hessian_backup/point_group_detect/mod.rs new file mode 100644 index 0000000000..c2c671660a --- /dev/null +++ b/src/hessian_backup/point_group_detect/mod.rs @@ -0,0 +1,43 @@ +//! Molecular point-group detection — a pure-std port of Psi4's libmints +//! algorithm (`psi4/driver/qcdb/libmintsmolecule.py`). +//! +//! # Note on AI usage +//! +//! This module is purely written by AI (GLM-5.2 with Claude Code). Human review is virtually nonexistent. +//! +//! I (ajz34) notices some redundant implementations (for examples elements). +//! These files can be refactored in the future, or kept as independent implementations for clarity. +//! +//! As long as functionality retains, I'm not aware if anyone will modify this module freely. +//! +//! Unittests are not included in this module. For more information, see +//! . +//! +//! # Example +//! ``` +//! use point_group_detection::{SymmMolecule, detect_point_group}; +//! +//! // H2O (Bohr) -> C2v +//! let mol = SymmMolecule::new( +//! vec!["O".into(), "H".into(), "H".into()], +//! vec![[0.0, 0.0, 0.12], [0.0, 1.54, -0.97], [0.0, -1.54, -0.97]], +//! ); +//! let pg = mol.detect(); +//! assert_eq!(pg.full_name, "C2v"); +//! ``` +//! +//! See `NOTES.md` and `psi4_ref/` for the reference algorithm and test suite. + +mod bits; +mod detect; +mod elements; +mod geom; +mod linalg; +mod matrix; +mod molecule; +mod vec3; + +mod interface_to_rest; + +pub use detect::{detect_point_group, PointGroup}; +pub use molecule::{SymmMolecule, Tmpl}; diff --git a/src/hessian_backup/point_group_detect/molecule.rs b/src/hessian_backup/point_group_detect/molecule.rs new file mode 100644 index 0000000000..5d80dbc0c1 --- /dev/null +++ b/src/hessian_backup/point_group_detect/molecule.rs @@ -0,0 +1,945 @@ +//! The `Molecule` detection core. Port of `psi4/driver/qcdb/libmintsmolecule.py` +//! (`center_of_mass`, `inertia_tensor`/`molecule.py`, `rotor_type`, +//! `is_linear_planar`, `has_inversion`, `is_plane`, `is_axis`, `like_world_axis`, +//! `symmetry_frame`, `find_highest_point_group`, `set_full_point_group`) and +//! `molecule.py::rotational_symmetry_number`. + +use super::bits; +use super::elements; +use super::geom::{atom_at_position, atom_present_in_geom, matrix_3d_rotation, matrix_3d_rotation_cn}; +use super::linalg::diagonalize3x3symmat; +use super::matrix::{points_matmul, Matrix3}; +use super::vec3::{self, Vec3, ZERO}; + +pub const DEFAULT_SYM_TOL: f64 = 1.0e-8; +pub const FULL_PG_TOL: f64 = 1.0e-8; +pub const NOISY_ZERO: f64 = 1.0e-8; + +/// Full point-group template (the "n" is substituted later). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Tmpl { + Atom, + DinfH, + CinfV, + Td, + Oh, + Ih, + C1, + Cs, + Ci, + Cn, + Cnv, + Cnh, + Dn, + Dnd, + Dnh, + Sn, +} + +impl Tmpl { + /// Schoenflies symbol with `n` substituted. Reference: `get_full_point_group`. + pub fn full_name(self, n: u32) -> String { + use Tmpl::*; + match self { + Atom => "ATOM".into(), + DinfH => "D_inf_h".into(), + CinfV => "C_inf_v".into(), + Td => "Td".into(), + Oh => "Oh".into(), + Ih => "Ih".into(), + C1 => "C1".into(), + Cs => "Cs".into(), + Ci => "Ci".into(), + Cn => format!("C{}", n), + Cnv => format!("C{}v", n), + Cnh => format!("C{}h", n), + Dn => format!("D{}", n), + Dnd => format!("D{}d", n), + Dnh => format!("D{}h", n), + Sn => format!("S{}", n), + } + } + + /// Rotational symmetry number σ. Reference: `rotational_symmetry_number`. + pub fn sigma(self, n: u32) -> f64 { + use Tmpl::*; + match self { + Atom | C1 | Ci | Cs | CinfV => 1.0, + DinfH => 2.0, + Td => 12.0, + Oh => 24.0, + Ih => 60.0, + Cn | Cnv | Cnh => n as f64, + Dn | Dnd | Dnh => 2.0 * n as f64, + Sn => n as f64 / 2.0, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Rotor { + Atom, + Linear, + Spherical, + Prolate, + Oblate, + Asymmetric, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum AxisLike { + X, + Y, + Z, +} + +/// A molecule for point-group detection. Geometry is in **Bohr**. +pub struct SymmMolecule { + pub symbols: Vec, + z: Vec, + pub geom: Vec<[f64; 3]>, + masses: Vec, + pub tol: f64, +} + +impl SymmMolecule { + pub fn new(symbols: Vec, geom: Vec<[f64; 3]>) -> Self { + assert_eq!(symbols.len(), geom.len(), "symbols and geom length mismatch"); + let z: Vec = symbols.iter().map(|s| elements::symbol_to_z(s).unwrap_or(0)).collect(); + let masses = z.iter().map(|&z| elements::z_to_mass(z)).collect(); + Self { symbols, z, geom, masses, tol: DEFAULT_SYM_TOL } + } + + /// Override per-atom masses (e.g. isotopic substitution to break symmetry). + /// Length must equal the atom count. + pub fn with_masses(mut self, masses: Vec) -> Self { + assert_eq!(masses.len(), self.z.len(), "masses length mismatch"); + self.masses = masses; + self + } + + pub fn with_tol(mut self, tol: f64) -> Self { + self.tol = tol; + self + } + + fn natom(&self) -> usize { + self.z.len() + } + + /// Reference `is_equivalent_to`: same Z and same mass (exact). Ghost status + /// is irrelevant here (no ghosts in the test set). + fn equiv(&self, i: usize, j: usize) -> bool { + self.z[i] == self.z[j] && self.masses[i] == self.masses[j] + } + + /// Mass-weighted center of mass. Reference: `center_of_mass`. + fn center_of_mass(&self, g: &[[f64; 3]]) -> Vec3 { + let mut ret = [0.0; 3]; + let mut total = 0.0; + for i in 0..g.len() { + let m = self.masses[i]; + ret = vec3::add(&ret, &vec3::scale(&g[i], m)); + total += m; + } + vec3::scale(&ret, 1.0 / total) + } + + /// Mass-weighted inertia tensor. Reference: `inertia_tensor_partial`. + fn inertia_tensor(&self, g: &[[f64; 3]]) -> Matrix3 { + let mut t = [[0.0; 3]; 3]; + for i in 0..g.len() { + let (x, y, z) = (g[i][0], g[i][1], g[i][2]); + let m = self.masses[i]; + t[0][0] += m * (y * y + z * z); + t[1][1] += m * (x * x + z * z); + t[2][2] += m * (x * x + y * y); + t[0][1] -= m * x * y; + t[0][2] -= m * x * z; + t[1][2] -= m * y * z; + } + t[1][0] = t[0][1]; + t[2][0] = t[0][2]; + t[2][1] = t[1][2]; + for r in 0..3 { + for c in 0..3 { + if t[r][c].abs() < ZERO { + t[r][c] = 0.0; + } + } + } + t + } + + /// Reference: installed `rotor_type` (1.10.2). + fn rotor_type(&self, g: &[[f64; 3]]) -> Rotor { + let (mut evals, _evecs) = diagonalize3x3symmat(&self.inertia_tensor(g)); + evals.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let rot_const: [f64; 3] = [ + if evals[0] > 1.0e-6 { 1.0 / evals[0] } else { 0.0 }, + if evals[1] > 1.0e-6 { 1.0 / evals[1] } else { 0.0 }, + if evals[2] > 1.0e-6 { 1.0 / evals[2] } else { 0.0 }, + ]; + + let mut degen = 0; + for i in 0..2 { + for j in (i + 1)..3 { + if degen >= 2 { + continue; + } + let rabs = (rot_const[i] - rot_const[j]).abs(); + let tmp = rot_const[i].max(rot_const[j]); + let rel = if rabs > ZERO { rabs / tmp } else { 0.0 }; + if rel < self.tol { + degen += 1; + } + } + } + + if self.natom() == 1 { + Rotor::Atom + } else if rot_const[0] == 0.0 { + Rotor::Linear + } else if degen == 2 { + Rotor::Spherical + } else if degen == 1 { + if (rot_const[1] - rot_const[2]) < 1.0e-6 { + Rotor::Prolate + } else if (rot_const[0] - rot_const[1]) < 1.0e-6 { + Rotor::Oblate + } else { + Rotor::Asymmetric + } + } else { + Rotor::Asymmetric + } + } + + /// Reference: `is_linear_planar`. + fn is_linear_planar(&self, g: &[[f64; 3]], tol: f64) -> (bool, bool) { + if self.natom() < 3 { + return (true, true); + } + let a = g[0]; + let b = g[1]; + let ba = vec3::normalize(&vec3::sub(&b, &a)); + let mut ca: Vec3 = [0.0; 3]; + let mut min_ba_dot_ca = 1.0; + for i in 2..self.natom() { + let tmp = vec3::normalize(&vec3::sub(&g[i], &a)); + if vec3::dot(&ba, &tmp).abs() < min_ba_dot_ca { + ca = tmp; + min_ba_dot_ca = vec3::dot(&ba, &tmp).abs(); + } + } + if min_ba_dot_ca >= 1.0 - tol { + return (true, true); + } + let linear = false; + if self.natom() < 4 { + return (linear, true); + } + let baxca = vec3::normalize(&vec3::cross(&ba, &ca)); + for i in 2..self.natom() { + let tmp = vec3::sub(&g[i], &a); + if vec3::dot(&tmp, &baxca).abs() > tol { + return (linear, false); + } + } + (linear, true) + } + + /// Reference: `has_inversion`. + fn has_inversion(&self, g: &[[f64; 3]], origin: &Vec3, tol: f64) -> bool { + for at in 0..self.natom() { + let inverted = vec3::sub( + &vec3::scale(origin, 2.0), + &g[at], + ); + match atom_at_position(g, &inverted, tol) { + Some(idx) if self.equiv(idx, at) => {} + _ => return false, + } + } + true + } + + /// Reference: `is_plane` (reflection through plane with normal `uperp` at `origin`). + fn is_plane(&self, g: &[[f64; 3]], origin: &Vec3, uperp: &Vec3, tol: f64) -> bool { + for i in 0..self.natom() { + let a = vec3::sub(&g[i], origin); + let apar = vec3::scale(uperp, vec3::dot(uperp, &a)); + let aperp = vec3::sub(&a, &apar); + let reflected = vec3::add(&vec3::sub(&aperp, &apar), origin); + match atom_at_position(g, &reflected, tol) { + Some(idx) if self.equiv(idx, i) => {} + _ => return false, + } + } + true + } + + /// Reference: `is_axis`. + fn is_axis(&self, g: &[[f64; 3]], origin: &Vec3, axis: &Vec3, order: usize, tol: f64) -> bool { + let two_pi = 2.0 * std::f64::consts::PI; + for i in 0..self.natom() { + let a = vec3::sub(&g[i], origin); + for j in 1..order { + let r = vec3::rotate(&a, j as f64 * two_pi / order as f64, axis); + let r = vec3::add(&r, origin); + match atom_at_position(g, &r, tol) { + Some(idx) if self.equiv(idx, i) => {} + _ => return false, + } + } + } + true + } + + /// Is the vertical plane containing the z-axis at azimuth `theta` a mirror + /// plane? Reflection across it maps `(x,y,z) -> (x cos2θ + y sin2θ, + /// x sin2θ - y cos2θ, z)`. Every atom must map to an equivalent atom. + /// (Replaces the reference's single-pivot reflect-x test, which depended on + /// a favorable orientation — see `has_vertical_mirror`.) + fn is_vertical_plane(&self, g: &[[f64; 3]], theta: f64, tol: f64) -> bool { + let c2 = (2.0 * theta).cos(); + let s2 = (2.0 * theta).sin(); + for i in 0..self.natom() { + let (x, y, z) = (g[i][0], g[i][1], g[i][2]); + let refl = [x * c2 + y * s2, x * s2 - y * c2, z]; + match atom_at_position(g, &refl, tol) { + Some(j) if self.equiv(j, i) => {} + _ => return false, + } + } + true + } + + /// Robust σv detection: scan candidate vertical mirror planes and return + /// true if any is a mirror. Fixes the orientation dependence of the + /// reference's pivot approach (an `acos` sign ambiguity caused skew-rotated + /// C3v molecules to lose their σv and report C3). + /// + /// A vertical plane at azimuth θ reflects an atom's xy-angle α to 2θ−α, so + /// the only θ values that can be mirrors are: θ = αᵢ (an atom lying on the + /// plane) and θ = (αᵢ+αⱼ)/2 (a mirror pair bisected by the plane), for + /// equivalent atoms at equal distance from the z-axis. Planes are periodic + /// mod π, so candidates are normalized to [0, π). + fn has_vertical_mirror(&self, g: &[[f64; 3]], tol: f64) -> bool { + let pi = std::f64::consts::PI; + // off-axis atoms: (azimuth, radius², index) + let off: Vec<(f64, f64, usize)> = (0..self.natom()) + .filter_map(|i| { + let r2 = g[i][0] * g[i][0] + g[i][1] * g[i][1]; + if r2 > tol * tol { + Some((g[i][1].atan2(g[i][0]), r2, i)) + } else { + None + } + }) + .collect(); + if off.is_empty() { + return false; // no off-axis atoms: no vertical mirror (linear case handled elsewhere) + } + + let mut cands: Vec = Vec::new(); + for a in 0..off.len() { + let (ang_a, r2_a, ia) = off[a]; + cands.push(ang_a); // on-plane candidate (atom lies on its own σv) + for b in (a + 1)..off.len() { + let (ang_b, r2_b, ib) = off[b]; + if !self.equiv(ia, ib) { + continue; + } + // mirror images are equidistant from the z-axis + if (r2_a - r2_b).abs() > 1.0e-6 { + continue; + } + cands.push(0.5 * (ang_a + ang_b)); // bisector candidate + } + } + + // dedupe mod π + let mut uniq: Vec = Vec::new(); + for c in cands { + let c = c.rem_euclid(pi); + let mut dup = false; + for u in &uniq { + let d = (c - u).abs().min((c - u + pi).abs()).min((c - u - pi).abs()); + if d < 1.0e-7 { + dup = true; + break; + } + } + if !dup { + uniq.push(c); + } + } + + for theta in uniq { + if self.is_vertical_plane(g, theta, tol) { + return true; + } + } + false + } + + /// Robust rotation-axis detection: does the molecule have a `order`-fold + /// rotation axis (in any direction) through `origin`? Scans candidate axes + /// derived from the atom positions: each atom direction `normalize(Aᵢ)` + /// (axes through opposite atoms, e.g. C4 of an octahedron or C5 of an + /// icosahedron) and each equivalent-pair bisector `normalize(Aᵢ+Aⱼ)` (axes + /// through edge midpoints). Used to distinguish Oh (has C4) from Ih (no C4) + /// independent of orientation — the reference tested S4 about the z-axis + /// only, which depended on `symmetry_frame` aligning a C4 axis to z. + fn has_rotation_axis(&self, g: &[[f64; 3]], origin: &Vec3, order: usize, tol: f64) -> bool { + let positions: Vec = (0..self.natom()).map(|i| vec3::sub(&g[i], origin)).collect(); + // candidate axes: atom directions + for i in 0..self.natom() { + let a = positions[i]; + let na = vec3::dot(&a, &a); + if na < tol * tol { + continue; + } + let axis = vec3::normalize(&a); + if self.is_axis(g, origin, &axis, order, tol) { + return true; + } + } + // candidate axes: equivalent-pair bisectors (axes through edge midpoints) + for i in 0..self.natom() { + let a = positions[i]; + let adota = vec3::dot(&a, &a); + for j in 0..i { + if !self.equiv(i, j) { + continue; + } + let b = positions[j]; + if (adota - vec3::dot(&b, &b)).abs() > 1.0e-6 { + continue; + } + let axis = vec3::add(&a, &b); + if vec3::norm(&axis) < 1.0e-12 { + continue; + } + let axis = vec3::normalize(&axis); + if self.is_axis(g, origin, &axis, order, tol) { + return true; + } + } + } + false + } + + /// Reference: `like_world_axis`. + fn like_world_axis(axis: &Vec3) -> (AxisLike, Vec3) { + let worldx = [1.0, 0.0, 0.0]; + let worldy = [0.0, 1.0, 0.0]; + let worldz = [0.0, 0.0, 1.0]; + let xlike = vec3::dot(axis, &worldx).abs(); + let ylike = vec3::dot(axis, &worldy).abs(); + let zlike = vec3::dot(axis, &worldz).abs(); + if (xlike - ylike) > 1.0e-12 && (xlike - zlike) > 1.0e-12 { + let a = if vec3::dot(axis, &worldx) < 0.0 { vec3::scale(axis, -1.0) } else { *axis }; + (AxisLike::X, a) + } else if (ylike - zlike) > 1.0e-12 { + let a = if vec3::dot(axis, &worldy) < 0.0 { vec3::scale(axis, -1.0) } else { *axis }; + (AxisLike::Y, a) + } else { + let a = if vec3::dot(axis, &worldz) < 0.0 { vec3::scale(axis, -1.0) } else { *axis }; + (AxisLike::Z, a) + } + } + + /// Tier A: highest D2h subgroup as a bit field. Reference: + /// `find_highest_point_group`. + fn find_highest_point_group(&self, g: &[[f64; 3]], tol: f64) -> u8 { + let mut pg_bits = 0u8; + for (op_bit, diag) in bits::tested_ops() { + let mut found = true; + for at in 0..self.natom() { + let pos = vec3::naivemult(&g[at], &diag); + match atom_at_position(g, &pos, tol) { + Some(idx) if self.equiv(idx, at) => {} + _ => { + found = false; + break; + } + } + } + if found { + pg_bits |= op_bit; + } + } + pg_bits + } + + /// Reference: `symmetry_frame`. Returns the 3×3 frame (columns = xaxis, + /// yaxis, zaxis) that reorients the molecule so its symmetry axes align + /// with the world axes. + fn symmetry_frame(&self, g: &[[f64; 3]], tol: f64) -> Matrix3 { + let com = self.center_of_mass(g); + let shifted: Vec<[f64; 3]> = g.iter().map(|p| vec3::sub(p, &com)).collect(); + let worldx = [1.0, 0.0, 0.0]; + let worldy = [0.0, 1.0, 0.0]; + let worldz = [0.0, 0.0, 1.0]; + + let mut sigma = [0.0; 3]; + let mut sigmav = [0.0; 3]; + let mut c2axis = [0.0; 3]; + let mut c2axisperp = [0.0; 3]; + + let (linear, planar) = self.is_linear_planar(g, tol); + let have_inversion = self.has_inversion(g, &com, tol); + + // --- check for C2 axis --- + let mut have_c2axis = false; + if self.natom() < 2 { + have_c2axis = true; + c2axis = [0.0, 0.0, 1.0]; + } else if linear { + have_c2axis = true; + c2axis = vec3::normalize(&vec3::sub(&g[1], &g[0])); + } else if planar && have_inversion { + let ba = vec3::normalize(&vec3::sub(&g[1], &g[0])); + for i in 2..self.natom() { + let ca = vec3::normalize(&vec3::sub(&g[i], &g[0])); + let baxca = vec3::cross(&ba, &ca); + if vec3::norm(&baxca) > tol { + have_c2axis = true; + c2axis = vec3::normalize(&baxca); + break; + } + } + } else { + 'outer: for i in 0..self.natom() { + let a = shifted[i]; + let adota = vec3::dot(&a, &a); + for j in 0..=i { + if !self.equiv(i, j) { + continue; + } + let b = shifted[j]; + if (adota - vec3::dot(&b, &b)).abs() > tol { + continue; + } + let axis = vec3::add(&a, &b); + if vec3::norm(&axis) < tol { + continue; + } + let axis = vec3::normalize(&axis); + if self.is_axis(g, &com, &axis, 2, tol) { + have_c2axis = true; + c2axis = axis; + break 'outer; + } + } + } + } + + let mut c2like = AxisLike::Z; + if have_c2axis { + let (like, axis) = Self::like_world_axis(&c2axis); + c2like = like; + c2axis = axis; + } + + // --- check for C2 axis perpendicular to first --- + let mut have_c2axisperp = false; + if have_c2axis { + if self.natom() < 2 { + have_c2axisperp = true; + c2axisperp = [1.0, 0.0, 0.0]; + } else if linear { + if have_inversion { + have_c2axisperp = true; + c2axisperp = vec3::perp_unit(&c2axis, &[0.0, 0.0, 1.0]); + } + } else { + 'outer: for i in 0..self.natom() { + let a = vec3::sub(&g[i], &com); + let adota = vec3::dot(&a, &a); + for j in 0..i { + if !self.equiv(i, j) { + continue; + } + let b = vec3::sub(&g[j], &com); + if (adota - vec3::dot(&b, &b)).abs() > tol { + continue; + } + let axis = vec3::add(&a, &b); + if vec3::norm(&axis) < tol { + continue; + } + let axis = vec3::normalize(&axis); + if vec3::dot(&axis, &c2axis).abs() > tol { + continue; + } + if self.is_axis(g, &com, &axis, 2, tol) { + have_c2axisperp = true; + c2axisperp = axis; + break 'outer; + } + } + } + } + } + + if have_c2axisperp { + let (mut c2perplike, snapped) = Self::like_world_axis(&c2axisperp); + c2axisperp = snapped; + // try to make c2axis the z axis + if c2perplike == AxisLike::Z { + std::mem::swap(&mut c2axisperp, &mut c2axis); + c2perplike = c2like; + c2like = AxisLike::Z; + } + if c2like != AxisLike::Z { + c2axis = if c2like == AxisLike::X { + vec3::cross(&c2axis, &c2axisperp) + } else { + vec3::cross(&c2axisperp, &c2axis) + }; + let (like, axis) = Self::like_world_axis(&c2axis); + c2like = like; + c2axis = axis; + } + if c2perplike == AxisLike::Y { + c2axisperp = vec3::cross(&c2axisperp, &c2axis); + let (_like, axis) = Self::like_world_axis(&c2axisperp); + c2axisperp = axis; + } + } + + // --- check for vertical plane --- + let mut have_sigmav = false; + if have_c2axis { + if self.natom() < 2 { + have_sigmav = true; + sigmav = c2axisperp; + } else if linear { + have_sigmav = true; + if have_c2axisperp { + sigmav = c2axisperp; + } else { + sigmav = vec3::perp_unit(&c2axis, &[0.0, 0.0, 1.0]); + } + } else { + 'outer: for i in 0..self.natom() { + let a = vec3::sub(&g[i], &com); + let adota = vec3::dot(&a, &a); + for j in 0..=i { + if !self.equiv(i, j) { + continue; + } + let b = vec3::sub(&g[j], &com); + if (adota - vec3::dot(&b, &b)).abs() > tol { + continue; + } + let inplane = vec3::add(&b, &a); + let norm_inplane = vec3::norm(&inplane); + if norm_inplane < tol { + continue; + } + let inplane = vec3::scale(&inplane, 1.0 / norm_inplane); + let perp = vec3::cross(&c2axis, &inplane); + let norm_perp = vec3::norm(&perp); + if norm_perp < tol { + continue; + } + let perp = vec3::scale(&perp, 1.0 / norm_perp); + if self.is_plane(g, &com, &perp, tol) { + have_sigmav = true; + sigmav = perp; + break 'outer; + } + } + } + } + } + + if have_sigmav { + let (sigmavlike, axis) = Self::like_world_axis(&sigmav); + sigmav = axis; + if c2like == AxisLike::Z && sigmavlike == AxisLike::Y { + sigmav = vec3::cross(&sigmav, &c2axis); + } else if c2like == AxisLike::Y && sigmavlike == AxisLike::Z { + sigmav = vec3::cross(&c2axis, &sigmav); + } + } + + // --- check for any sigma plane (only if no inversion and no c2 axis) --- + let mut have_sigma = false; + if !have_inversion && !have_c2axis { + if planar { + let ba = vec3::normalize(&vec3::sub(&g[1], &g[0])); + for i in 2..self.natom() { + let ca = vec3::normalize(&vec3::sub(&g[i], &g[0])); + let baxca = vec3::cross(&ba, &ca); + if vec3::norm(&baxca) > tol { + have_sigma = true; + sigma = vec3::normalize(&baxca); + break; + } + } + } else { + 'outer: for i in 0..self.natom() { + let a = vec3::sub(&g[i], &com); + let adota = vec3::dot(&a, &a); + for j in 0..i { + if !self.equiv(i, j) { + continue; + } + let b = vec3::sub(&g[j], &com); + let bdotb = vec3::dot(&b, &b); + if (adota - bdotb).abs() > tol { + continue; + } + let perp = vec3::sub(&b, &a); + let norm_perp = vec3::norm(&perp); + if norm_perp < tol { + continue; + } + let perp = vec3::scale(&perp, 1.0 / norm_perp); + if self.is_plane(g, &com, &perp, tol) { + have_sigma = true; + sigma = perp; + break 'outer; + } + } + } + } + } + + if have_sigma { + let xlikeness = vec3::dot(&sigma, &worldx).abs(); + let ylikeness = vec3::dot(&sigma, &worldy).abs(); + let zlikeness = vec3::dot(&sigma, &worldz).abs(); + if xlikeness > ylikeness && xlikeness > zlikeness { + if vec3::dot(&sigma, &worldx) < 0.0 { + sigma = vec3::scale(&sigma, -1.0); + } + } else if ylikeness > zlikeness { + if vec3::dot(&sigma, &worldy) < 0.0 { + sigma = vec3::scale(&sigma, -1.0); + } + } else if vec3::dot(&sigma, &worldz) < 0.0 { + sigma = vec3::scale(&sigma, -1.0); + } + } + + // --- assemble the three axes --- + let mut xaxis = worldx; + let mut zaxis = worldz; + if have_c2axis { + zaxis = c2axis; + if have_sigmav { + xaxis = sigmav; + } else if have_c2axisperp { + xaxis = c2axisperp; + } else { + xaxis = vec3::perp_unit(&zaxis, &zaxis); + } + } else if have_sigma { + zaxis = sigma; + xaxis = vec3::perp_unit(&zaxis, &zaxis); + } + + for k in 0..3 { + if zaxis[k].abs() < NOISY_ZERO { + zaxis[k] = 0.0; + } + if xaxis[k].abs() < NOISY_ZERO { + xaxis[k] = 0.0; + } + } + let yaxis = vec3::scale(&vec3::cross(&xaxis, &zaxis), -1.0); + + let mut frame = super::matrix::zero(); + for i in 0..3 { + frame[i][0] = xaxis[i]; + frame[i][1] = yaxis[i]; + frame[i][2] = zaxis[i]; + } + frame + } + + /// Tier B: determine the full point group template and `n`. Reference: + /// `set_full_point_group`. `g` must already be COM-recentered and + /// symmetry-frame rotated (the state of `self.geometry()` after + /// `update_geometry`). + fn set_full_point_group(&self, g: &[[f64; 3]], tol: f64) -> (Tmpl, u32) { + // local recentered copy (g is already COM-recentered; com ~ 0) + let com = self.center_of_mass(g); + let mut geom: Vec<[f64; 3]> = g.iter().map(|p| vec3::sub(p, &com)).collect(); + + let rotor = self.rotor_type(g); + let pg_bits = self.find_highest_point_group(g, tol); + let d2h_subgroup = bits::bits_to_basic_name(pg_bits); + let op_i = self.has_inversion(g, &[0.0, 0.0, 0.0], tol); + + let z_axis = [0.0, 0.0, 1.0]; + + match rotor { + Rotor::Atom => (Tmpl::Atom, 0), + Rotor::Linear => { + if op_i { + (Tmpl::DinfH, 0) + } else { + (Tmpl::CinfV, 0) + } + } + Rotor::Spherical => { + if !op_i { + (Tmpl::Td, 3) + } else { + // Oh has a C4 axis (anywhere); Ih does not (its highest is + // C5). Test for a C4 axis in any direction rather than S4 + // about z only — the latter depended on `symmetry_frame` + // aligning a C4 axis to z and mis-reported Oh as Ih for + // generically-oriented molecules. + let origin = [0.0, 0.0, 0.0]; + if self.has_rotation_axis(g, &origin, 4, tol) { + (Tmpl::Oh, 4) + } else { + (Tmpl::Ih, 5) + } + } + } + Rotor::Asymmetric => { + let (tmpl, n) = match d2h_subgroup { + "c1" => (Tmpl::C1, 1u32), + "ci" => (Tmpl::Ci, 1), + "c2" => (Tmpl::Cn, 2), + "cs" => (Tmpl::Cs, 1), + "d2" => (Tmpl::Dn, 2), + "c2v" => (Tmpl::Cnv, 2), + "c2h" => (Tmpl::Cnh, 2), + "d2h" => (Tmpl::Dnh, 2), + _ => (Tmpl::C1, 1), + }; + (tmpl, n) + } + Rotor::Prolate | Rotor::Oblate => { + // symmetric top + let (evals, evecs) = diagonalize3x3symmat(&self.inertia_tensor(g)); + // zip eigenvalues with eigenvectors (columns of evecs -> rows), sort ascending + let mut ev_list: Vec<(f64, Vec3)> = (0..3) + .map(|i| (evals[i], [evecs[0][i], evecs[1][i], evecs[2][i]])) + .collect(); + ev_list.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + let i_evals: [f64; 3] = [ev_list[0].0, ev_list[1].0, ev_list[2].0]; + let i_evecs: [Vec3; 3] = [ev_list[0].1, ev_list[1].1, ev_list[2].1]; + + let mut unique_axis = 1usize; + if (i_evals[0] - i_evals[1]).abs() < tol { + unique_axis = 2; + } else if (i_evals[1] - i_evals[2]).abs() < tol { + unique_axis = 0; + } + let old_axis = i_evecs[unique_axis]; + + let ddot = vec3::dot(&z_axis, &old_axis); + let phi = if (ddot - 1.0).abs() < 1.0e-10 { + 0.0 + } else if (ddot + 1.0).abs() < 1.0e-10 { + std::f64::consts::PI + } else { + ddot.acos() + }; + + if phi.abs() > 1.0e-14 { + let rot_axis = vec3::cross(&z_axis, &old_axis); + geom = matrix_3d_rotation(&geom, &rot_axis, phi, false); + } + + let cn_z = matrix_3d_rotation_cn(&geom, &z_axis, false, tol, 0); + let sn_z = matrix_3d_rotation_cn(&geom, &z_axis, true, tol, 0); + + // sigma_h (xy plane): reflect z + let mut op_sigma_h = true; + for at in 0..self.natom() { + if geom[at][2].abs() < tol { + continue; + } + let test_atom = [geom[at][0], geom[at][1], -geom[at][2]]; + if !atom_present_in_geom(&geom, &test_atom, tol) { + op_sigma_h = false; + break; + } + } + + // sigma_v: robust scan of candidate vertical mirror planes. + // (Reference pivoted one atom into the yz plane and tested + // reflect-x; that depended on the pivot's xy-angle sign and + // failed for generically-oriented Cnv molecules such as NH3.) + let op_sigma_v = self.has_vertical_mirror(&geom, tol); + + // perpendicular C2's (pair by Z only, per reference). This scan + // is invariant under rotations about z, so it needs no pivot. + let mut is_d = false; + 'pair: for i in 0..self.natom() { + let a = geom[i]; + let adota = vec3::dot(&a, &a); + for j in 0..i { + if self.z[i] != self.z[j] { + continue; + } + let b = geom[j]; + if (adota - vec3::dot(&b, &b)).abs() > 1.0e-6 { + continue; + } + let axis = vec3::add(&a, &b); + if vec3::norm(&axis) < 1.0e-12 { + continue; + } + let axis = vec3::normalize(&axis); + if vec3::dot(&axis, &z_axis).abs() > 1.0e-6 { + continue; + } + if matrix_3d_rotation_cn(&geom, &axis, false, tol, 2) == 2 { + is_d = true; + break 'pair; + } + } + } + + let cn = cn_z as u32; + let sn = sn_z as u32; + if sn == 2 * cn && !is_d { + return (Tmpl::Sn, sn); + } + if is_d { + if op_sigma_h && op_sigma_v { + return (Tmpl::Dnh, cn); + } else if sn == 2 * cn { + return (Tmpl::Dnd, cn); + } else { + return (Tmpl::Dn, cn); + } + } else { + if op_sigma_h && sn == cn { + return (Tmpl::Cnh, cn); + } else if op_sigma_v { + return (Tmpl::Cnv, cn); + } else { + return (Tmpl::Cn, cn); + } + } + } + } + } + + /// Run the full detection pipeline. Returns `(template, n)`. + pub(crate) fn detect_inner(&self) -> (Tmpl, u32) { + let tol = self.tol; + // 1. COM recenter + let com = self.center_of_mass(&self.geom); + let mut g: Vec<[f64; 3]> = self.geom.iter().map(|p| vec3::sub(p, &com)).collect(); + // 2. symmetry_frame reorientation + let frame = self.symmetry_frame(&g, tol); + g = points_matmul(&g, &frame); + // 3. Tier A (inside set_full_point_group) + Tier B + self.set_full_point_group(&g, tol) + } +} diff --git a/src/hessian_backup/point_group_detect/vec3.rs b/src/hessian_backup/point_group_detect/vec3.rs new file mode 100644 index 0000000000..144234bb1b --- /dev/null +++ b/src/hessian_backup/point_group_detect/vec3.rs @@ -0,0 +1,158 @@ +//! 3D vector helpers. Direct port of `psi4/driver/qcdb/vecutil.py` vector +//! routines. Vectors are `[f64; 3]` (length-3 throughout, matching the +//! reference). Pure std. + +pub type Vec3 = [f64; 3]; + +pub const ZERO: f64 = 1.0e-14; + +#[inline] +pub fn norm(v: &Vec3) -> f64 { + (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt() +} + +#[inline] +pub fn add(v: &Vec3, u: &Vec3) -> Vec3 { + [v[0] + u[0], v[1] + u[1], v[2] + u[2]] +} + +#[inline] +pub fn sub(v: &Vec3, u: &Vec3) -> Vec3 { + [v[0] - u[0], v[1] - u[1], v[2] - u[2]] +} + +#[inline] +pub fn dot(v: &Vec3, u: &Vec3) -> f64 { + v[0] * u[0] + v[1] * u[1] + v[2] * u[2] +} + +#[inline] +pub fn scale(v: &Vec3, d: f64) -> Vec3 { + [v[0] * d, v[1] * d, v[2] * d] +} + +/// Element-wise (component-wise) product. Reference: `naivemult`. +#[inline] +pub fn naivemult(v: &Vec3, u: &Vec3) -> Vec3 { + [v[0] * u[0], v[1] * u[1], v[2] * u[2]] +} + +#[inline] +pub fn normalize(v: &Vec3) -> Vec3 { + let m = norm(v); + [v[0] / m, v[1] / m, v[2] / m] +} + +#[inline] +pub fn distance(v: &Vec3, u: &Vec3) -> f64 { + let d0 = v[0] - u[0]; + let d1 = v[1] - u[1]; + let d2 = v[2] - u[2]; + (d0 * d0 + d1 * d1 + d2 * d2).sqrt() +} + +#[inline] +pub fn cross(v: &Vec3, u: &Vec3) -> Vec3 { + [ + v[1] * u[2] - v[2] * u[1], + v[2] * u[0] - v[0] * u[2], + v[0] * u[1] - v[1] * u[0], + ] +} + +/// Unit vector perpendicular to length-3 vectors `u` and `v`. +/// Reference: `vecutil.py:perp_unit`. Handles the degenerate cross-product +/// case by choosing a vector perpendicular to the larger of `u`, `v` in the +/// plane of their two largest components. +pub fn perp_unit(u: &Vec3, v: &Vec3) -> Vec3 { + // try cross product + let result = cross(u, v); + let rdotr = dot(&result, &result); + + if rdotr < 1.0e-16 { + // cross product too small to normalize: pick the larger of u, v + let (d, dotprodd) = if dot(u, u) < dot(v, v) { + (*v, dot(v, v)) + } else { + (*u, dot(u, u)) + }; + + if dotprodd < 1.0e-16 { + // both tiny -> arbitrary + return [1.0, 0.0, 0.0]; + } + + // choose a vector perpendicular to d, in the plane of d's two largest + // components (90° rotation within that plane). + let absd = [d[0].abs(), d[1].abs(), d[2].abs()]; + let (axis0, axis1) = if (absd[1] - absd[0]) > 1.0e-12 { + if (absd[2] - absd[0]) > 1.0e-12 { + (1, 2) + } else { + (1, 0) + } + } else if (absd[2] - absd[1]) > 1.0e-12 { + (0, 2) + } else { + (0, 1) + }; + let mut r = [0.0, 0.0, 0.0]; + r[axis0] = d[axis1]; + r[axis1] = -d[axis0]; + normalize(&r) + } else { + scale(&result, 1.0 / rdotr.sqrt()) + } +} + +/// Rotate vector `v` about `axis` by `theta` radians. +/// Reference: `vecutil.py:rotate`. Used by `is_axis`; not on the main +/// `set_full_point_group` path (which uses the matrix form in `geom.rs`). +pub fn rotate(v: &Vec3, theta: f64, axis: &Vec3) -> Vec3 { + // (reference computes `unitaxis = normalize(axis)` but never uses it; omitted) + // parallel component along axis + let parallel = scale(axis, dot(v, axis) / dot(axis, axis)); + let perpendicular = sub(v, ¶llel); + // third orthonormal axis + let mut third = perp_unit(axis, &perpendicular); + third = scale(&third, norm(&perpendicular)); + + let mut result = add(¶llel, &add(&scale(&perpendicular, theta.cos()), &scale(&third, theta.sin()))); + for item in result.iter_mut() { + if item.abs() < ZERO { + *item = 0.0; + } + } + result +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn vec_basics() { + assert!((norm(&[3.0, 4.0, 0.0]) - 5.0).abs() < 1e-15); + assert_eq!(dot(&[1.0, 2.0, 3.0], &[4.0, -5.0, 6.0]), 4.0 - 10.0 + 18.0); + assert_eq!(cross(&[1.0, 0.0, 0.0], &[0.0, 1.0, 0.0]), [0.0, 0.0, 1.0]); + let n = normalize(&[0.0, 0.0, 5.0]); + assert!((n[2] - 1.0).abs() < 1e-15); + } + + #[test] + fn perp_unit_basic() { + let p = perp_unit(&[0.0, 0.0, 1.0], &[1.0, 0.0, 0.0]); + // perpendicular to both -> ±y + assert!((dot(&p, &[0.0, 0.0, 1.0]).abs()) < 1e-15); + assert!((dot(&p, &[1.0, 0.0, 0.0]).abs()) < 1e-15); + assert!((norm(&p) - 1.0).abs() < 1e-15); + } + + #[test] + fn rotate_about_z() { + let r = rotate(&[1.0, 0.0, 0.0], std::f64::consts::FRAC_PI_2, &[0.0, 0.0, 1.0]); + assert!((r[0] - 0.0).abs() < 1e-15); + assert!((r[1] - 1.0).abs() < 1e-15); + assert!((r[2] - 0.0).abs() < 1e-15); + } +} -- Gitee From 00d5a2e371567c1f7d02fb28fbf59084cf17326d Mon Sep 17 00:00:00 2001 From: ajz34 Date: Sun, 28 Jun 2026 18:22:55 +0800 Subject: [PATCH 15/39] hessian-backup: interface symmetry detection to hessian thermo analysis --- src/hessian_backup/config.rs | 10 ++++++++-- src/hessian_backup/point_group_detect/mod.rs | 18 +++++++++--------- src/hessian_backup/rscf.rs | 4 ++-- src/hessian_backup/rscf_interface.rs | 12 +++++++++--- src/hessian_backup/uscf.rs | 4 ++-- src/hessian_backup/uscf_interface.rs | 12 +++++++++--- 6 files changed, 39 insertions(+), 21 deletions(-) diff --git a/src/hessian_backup/config.rs b/src/hessian_backup/config.rs index 47341d228b..32a0a1e8a3 100644 --- a/src/hessian_backup/config.rs +++ b/src/hessian_backup/config.rs @@ -5,7 +5,7 @@ use serde_inline_default::serde_inline_default; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct HessSCFConfig { #[serde_inline_default(0.0)] - pub level_shift: f64, + pub cphf_level_shift: f64, #[serde_inline_default(1e-8)] pub cphf_tol: f64, #[serde_inline_default(42)] @@ -30,12 +30,17 @@ pub struct HessSCFConfig { /// - `grid_gen_level + 2` for MGGA (TAU) functionals. #[serde_inline_default(None)] pub grid_level_skeleton: Option, + /// Tolerance for point group detection in vibrational analysis. Default to 1e-5 Bohr. + /// + /// Note that this tolerance will be divided by sqrt(1 + natm). + #[serde_inline_default(1.0e-5)] + pub tol_point_group: f64, } impl Default for HessSCFConfig { fn default() -> Self { Self { - level_shift: 0.0, + cphf_level_shift: 0.0, cphf_tol: 1e-8, cphf_max_cycle: 42, cphf_max_space: 14, @@ -44,6 +49,7 @@ impl Default for HessSCFConfig { atm_list: None, grid_level_cphf: None, grid_level_skeleton: None, + tol_point_group: 1.0e-5, } } } diff --git a/src/hessian_backup/point_group_detect/mod.rs b/src/hessian_backup/point_group_detect/mod.rs index c2c671660a..a0e3fb7142 100644 --- a/src/hessian_backup/point_group_detect/mod.rs +++ b/src/hessian_backup/point_group_detect/mod.rs @@ -28,16 +28,16 @@ //! //! See `NOTES.md` and `psi4_ref/` for the reference algorithm and test suite. -mod bits; -mod detect; -mod elements; -mod geom; -mod linalg; -mod matrix; -mod molecule; -mod vec3; +pub mod bits; +pub mod detect; +pub mod elements; +pub mod geom; +pub mod linalg; +pub mod matrix; +pub mod molecule; +pub mod vec3; -mod interface_to_rest; +pub mod interface_to_rest; pub use detect::{detect_point_group, PointGroup}; pub use molecule::{SymmMolecule, Tmpl}; diff --git a/src/hessian_backup/rscf.rs b/src/hessian_backup/rscf.rs index a06daf2d9d..398d54b334 100644 --- a/src/hessian_backup/rscf.rs +++ b/src/hessian_backup/rscf.rs @@ -89,7 +89,7 @@ impl<'a> RHessSCF<'a> { let mo_coeff = &self.mo_coeff; let mo_occ = &self.mo_occ; let mo_energy = &self.mo_energy; - let level_shift = self.config.level_shift; + let level_shift = self.config.cphf_level_shift; let device = mo_coeff.device().clone(); let [nao, nmo] = mo_coeff.shape().to_vec().try_into().unwrap(); @@ -220,7 +220,7 @@ impl<'a> RHessSCF<'a> { let t0 = std::time::Instant::now(); let mo_occ = self.mo_occ.view(); let mo_energy = self.mo_energy.view(); - let level_shift = self.config.level_shift; + let level_shift = self.config.cphf_level_shift; let occidx = mo_occ.view().greater(0).into_vec(); let viridx = occidx.iter().map(|&x| !x).collect_vec(); let nocc = occidx.iter().filter(|&&x| x).count(); diff --git a/src/hessian_backup/rscf_interface.rs b/src/hessian_backup/rscf_interface.rs index 4f6efc220f..5e02964286 100644 --- a/src/hessian_backup/rscf_interface.rs +++ b/src/hessian_backup/rscf_interface.rs @@ -149,7 +149,7 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, let mass_rt = rt::asarray((&mass, &device)); let geom = atm_list.iter().map(|&i| mol.atom_coord(i)).collect_vec(); - let geom_rt = rt::asarray((geom, &device)).into_unpack_array(0); + let geom_rt = rt::asarray((&geom, &device)).into_unpack_array(0); let vib = harmonic_analysis(hess.view(), geom_rt.view(), mass_rt.view(), true, true); println!("=============== Vibrational Analysis ==============="); @@ -164,10 +164,16 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, let mass_sum = mass.iter().sum::(); let e0 = scf_data.scf_energy; let multiplicity = scf_data.mol.ctrl.spin; - // TODO: sigma (rotational symmetry number) is not yet determined. For now, we set it to 1. - let th = thermo(&vib, 298.15, 101325.0, multiplicity as _, mass_sum, e0, 1, &rc_cm, rotor); + + use super::point_group_detect::interface_to_rest::get_full_point_group_for_vib; + let tol_pg = config.tol_point_group / (1.0 + natm as f64).sqrt(); + let (pg_name, pg_sigma) = get_full_point_group_for_vib(&elems, &mass, &geom, tol_pg); + + let th = thermo(&vib, 298.15, 101325.0, multiplicity as _, mass_sum, e0, pg_sigma as _, &rc_cm, rotor); println!("=============== Thermo Analysis ==============="); + println!(""); + println!("Point group: {}, sigma (rotation symmetry number): {}", pg_name, pg_sigma); let msg_th = print_thermo(&th, multiplicity as _, mass_sum); println!("{}", msg_th); diff --git a/src/hessian_backup/uscf.rs b/src/hessian_backup/uscf.rs index 202003cce4..27e87e2843 100644 --- a/src/hessian_backup/uscf.rs +++ b/src/hessian_backup/uscf.rs @@ -79,7 +79,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.config.level_shift; + let level_shift = self.config.cphf_level_shift; let device = mo_coeff[α].device().clone(); let nao = mo_coeff[α].shape()[0]; @@ -241,7 +241,7 @@ impl<'a> UHessSCF<'a> { 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.config.level_shift; + let level_shift = self.config.cphf_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[β])]; diff --git a/src/hessian_backup/uscf_interface.rs b/src/hessian_backup/uscf_interface.rs index f2466889a8..ffe46b19d7 100644 --- a/src/hessian_backup/uscf_interface.rs +++ b/src/hessian_backup/uscf_interface.rs @@ -158,7 +158,7 @@ pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, let mass_rt = rt::asarray((&mass, &device)); let geom = atm_list.iter().map(|&i| mol.atom_coord(i)).collect_vec(); - let geom_rt = rt::asarray((geom, &device)).into_unpack_array(0); + let geom_rt = rt::asarray((&geom, &device)).into_unpack_array(0); let vib = harmonic_analysis(hess.view(), geom_rt.view(), mass_rt.view(), true, true); println!("=============== Vibrational Analysis ==============="); @@ -173,10 +173,16 @@ pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, let mass_sum = mass.iter().sum::(); let e0 = scf_data.scf_energy; let multiplicity = scf_data.mol.ctrl.spin; - // TODO: sigma (rotational symmetry number) is not yet determined. For now, we set it to 1. - let th = thermo(&vib, 298.15, 101325.0, multiplicity as _, mass_sum, e0, 1, &rc_cm, rotor); + + use super::point_group_detect::interface_to_rest::get_full_point_group_for_vib; + let tol_pg = config.tol_point_group / (1.0 + natm as f64).sqrt(); + let (pg_name, pg_sigma) = get_full_point_group_for_vib(&elems, &mass, &geom, tol_pg); + + let th = thermo(&vib, 298.15, 101325.0, multiplicity as _, mass_sum, e0, pg_sigma as _, &rc_cm, rotor); println!("=============== Thermo Analysis ==============="); + println!(""); + println!("Point group: {}, sigma (rotation symmetry number): {}", pg_name, pg_sigma); let msg_th = print_thermo(&th, multiplicity as _, mass_sum); println!("{}", msg_th); -- Gitee From 6dd3ac6d2a682418a34a0d6882e6f3ec6afe683f Mon Sep 17 00:00:00 2001 From: ajz34 Date: Sun, 28 Jun 2026 19:36:55 +0800 Subject: [PATCH 16/39] hessian-backup: somehow better printing --- src/hessian_backup/rscf_interface.rs | 22 +++++++++++++++++++++- src/hessian_backup/uscf_interface.rs | 22 +++++++++++++++++++++- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/src/hessian_backup/rscf_interface.rs b/src/hessian_backup/rscf_interface.rs index 5e02964286..5a2bfc8254 100644 --- a/src/hessian_backup/rscf_interface.rs +++ b/src/hessian_backup/rscf_interface.rs @@ -126,7 +126,27 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, ); let de_hess = hess_scf.make_hess(); - println!("=== HESSIAN ===\n{:12.6}", de_hess.t()); + + 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 { diff --git a/src/hessian_backup/uscf_interface.rs b/src/hessian_backup/uscf_interface.rs index ffe46b19d7..5d5a950862 100644 --- a/src/hessian_backup/uscf_interface.rs +++ b/src/hessian_backup/uscf_interface.rs @@ -135,7 +135,27 @@ pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, ); let de_hess = hess_scf.make_hess(); - println!("=== HESSIAN ===\n{:12.6}", de_hess.t()); + + 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 { -- Gitee From fd75c754d7b32f844ebc23fadc4504db9d2da7b2 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Sun, 28 Jun 2026 20:02:43 +0800 Subject: [PATCH 17/39] hessian-backup: somehow better printing --- tests/hessian_backup/uhf.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/hessian_backup/uhf.rs b/tests/hessian_backup/uhf.rs index 5efd91735a..7fa0450aa8 100644 --- a/tests/hessian_backup/uhf.rs +++ b/tests/hessian_backup/uhf.rs @@ -34,6 +34,9 @@ static INPUT_NH3: &str = r##" H 0.3 1.1 0.2 H 0.1 0.1 1.2 """ + +[hessian_backup] +atm_list = [0, 1, 3] "##; #[test] @@ -47,7 +50,7 @@ fn test_nh3() { let config = HessSCFConfig::default(); let (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config); - let natm = 4; + let natm = 3; let de = rt::asarray((&de, [3, 3, natm, natm])); println!("Hessian:\n{:12.6}", de.t()); } \ No newline at end of file -- Gitee From 6e912d5b6b72d669aecceb51ad3ec6385a628a5f Mon Sep 17 00:00:00 2001 From: ajz34 Date: Sun, 28 Jun 2026 20:15:53 +0800 Subject: [PATCH 18/39] vib: use `crate::constants` instead of defining physical constants --- src/constants/mod.rs | 17 +++++++- src/hessian_backup/vib.rs | 81 ++++++++++++++++----------------------- 2 files changed, 47 insertions(+), 51 deletions(-) diff --git a/src/constants/mod.rs b/src/constants/mod.rs index 7b8aceaf3e..ead236c200 100644 --- a/src/constants/mod.rs +++ b/src/constants/mod.rs @@ -316,7 +316,7 @@ pub const KR_SHELL: [f64; 18] = // 1s 2s 2p 3s 3p 4s 3d 4p 5s 4d 5p 4f 5d 6s 6p pub const NELE_IN_SHELLS: [f64; 15] = [2.0, 2.0, 6.0, 2.0, 6.0, 2.0, 10.0, 6.0, 2.0, 10.0, 6.0, 14.0, 10.0, 2.0, 6.0]; -// +// =========== physical constants =================================== pub const LIGHT_SPEED: f64 = 137.03599967994; // http://physics.nist.gov/cgi-bin/cuu/Value?alph // BOHR = .529 177 210 92(17) e-10m // http://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0 pub const BOHR: f64 = 0.52917721092; // Angstroms @@ -328,8 +328,21 @@ pub const AVOGADRO: f64 = 6.022140857e23; // https://physics.nist.gov/cgi- pub const PLANCK: f64 = 6.626070040e-34; // J*s http://physics.nist.gov/cgi-bin/cuu/Value?h pub const E_CHARGE: f64 = 1.6021766208e-19; pub const DEBYE:f64 = 3.335641e-30; // C*m = 1e-18/LIGHT_SPEED_SI https://cccbdb.nist.gov/debye.asp -pub const AU2DEBYE:f64 = E_CHARGE * BOHR*1e-10 / DEBYE; // 2.541746 +// CODATA 2022 +// https://docs.scipy.org/doc/scipy/reference/constants.html +pub const SPEED_OF_LIGHT: f64 = 299792458.0; // m s^-1 +pub const BOLTZMANN: f64 = 1.380649e-23; // J K^-1 +pub const R_GAS: f64 = 8.31446261815324; // J mol^-1 K^-1 + +// =========== unit conversion =================================== +pub const BOHR2ANG: f64 = BOHR; +pub const AU2DEBYE:f64 = E_CHARGE * BOHR*1e-10 / DEBYE; // 2.541746 +pub const HARTREE2J: f64 = 4.35974465e-18; +pub const HARTREE2KJMOL: f64 = 2625.4996382852164; +pub const HARTREE2KCALMOL: f64 = 627.5094737775374; +pub const HARTREE2WAVENUMBERS: f64 = 219474.6313702; +pub const AMU2KG: f64 = 1.66053904e-27; pub const MPI_CHUNK:usize = 134217728; // around 1 GB diff --git a/src/hessian_backup/vib.rs b/src/hessian_backup/vib.rs index 7f899eb447..3be62337b9 100644 --- a/src/hessian_backup/vib.rs +++ b/src/hessian_backup/vib.rs @@ -16,24 +16,18 @@ use super::prelude::*; // --------------------------------------------------------------------------- // Physical constants (CODATA2014, matching the Python implementation) // --------------------------------------------------------------------------- -const NA: f64 = 6.022140857e23; -const HARTREE2J: f64 = 4.35974465e-18; -const C: f64 = 299792458.0; -const BOHR2ANG: f64 = 0.52917721067; -const H: f64 = 6.62607004e-34; -const KB: f64 = 1.38064852e-23; -const R_GAS: f64 = 8.3144598; -const HARTREE2KJMOL: f64 = 2625.4996382852164; -const HARTREE2KCALMOL: f64 = 627.5094737775374; -const HARTREE2WAVENUMBERS: f64 = 219474.6313702; -const AMU2KG: f64 = 1.66053904e-27; + +use crate::constants::{ + AMU2KG, AVOGADRO, BOHR2ANG, BOLTZMANN, HARTREE2J, HARTREE2KCALMOL, HARTREE2KJMOL, HARTREE2WAVENUMBERS, PLANCK, + R_GAS, SPEED_OF_LIGHT, +}; /// Tolerance for detecting nearly-linear geometries in `get_tr_space`. pub const LINEAR_A_TOL: f64 = 1.0e-2; /// cm⁻¹ conversion factor from force-constant eigenvalues (atomic units). fn uconv_cm1() -> f64 { - (NA * HARTREE2J * 1.0e19).sqrt() / (2.0 * std::f64::consts::PI * C * BOHR2ANG) + (AVOGADRO * HARTREE2J * 1.0e19).sqrt() / (2.0 * std::f64::consts::PI * SPEED_OF_LIGHT * BOHR2ANG) } /// Mass-centred geometry `[3, natm]`. @@ -65,13 +59,7 @@ pub fn get_tr_space(mass: TsrView, geom: TsrView, space: &str) -> Tsr { /// vectors are selected (rank set by the caller from rotor type). Otherwise /// `tol_user` (`Some(t)` absolute cutoff, or `None` for the default /// `ndof × max(s) × ε` tolerance) determines the rank. -fn _get_tr_space( - mass: TsrView, - geom: TsrView, - space: &str, - tol_user: Option, - nrt_user: Option, -) -> Tsr { +fn _get_tr_space(mass: TsrView, geom: TsrView, space: &str, tol_user: Option, nrt_user: Option) -> Tsr { let device = geom.device().clone(); let natm = geom.shape()[1]; let ndof = 3 * natm; @@ -197,7 +185,7 @@ pub fn rotation_const(mass: TsrView, atom_coords: TsrView, unit: &str) -> Tsr { let conv = if unit_lower == "ghz" { 1e-9 } else if unit_lower == "wavenumber" { - 1.0 / C * 1e-2 + 1.0 / SPEED_OF_LIGHT * 1e-2 } else { panic!("Unsupported unit {}", unit); }; @@ -522,8 +510,9 @@ pub fn harmonic_analysis( } // --------------- conversion factors --------------- - let uconv_mdyne_a = 0.1 * (2.0 * std::f64::consts::PI * C).powi(2) / NA; - let uconv_S = ((C * (2.0 * std::f64::consts::PI * BOHR2ANG).powi(2)) / (H * NA * 1.0e21)).sqrt(); + let uconv_mdyne_a = 0.1 * (2.0 * std::f64::consts::PI * SPEED_OF_LIGHT).powi(2) / AVOGADRO; + let uconv_S = + ((SPEED_OF_LIGHT * (2.0 * std::f64::consts::PI * BOHR2ANG).powi(2)) / (PLANCK * AVOGADRO * 1.0e21)).sqrt(); // --------------- normal modes & reduced mass --------------- // w = m^{-1/2} q ; w[a,i] = q[a,i] / sqrt(m_a) @@ -558,7 +547,7 @@ pub fn harmonic_analysis( let dq0: Vec = qtp0.iter().map(|&q| q / 2.0_f64.sqrt()).collect(); // --------------- characteristic vibrational temperature --------------- - let uconv_K = 100.0 * H * C / KB; + let uconv_K = 100.0 * PLANCK * SPEED_OF_LIGHT / BOLTZMANN; let theta_vib: Vec = omega_real.iter().map(|&w| w * uconv_K).collect(); let q_vec = qL.reshape(-1).to_vec(); @@ -683,9 +672,11 @@ pub fn thermo( s[ELEC] = (multiplicity as f64).ln(); // ---------- translational ---------- - let beta = 1.0 / (KB * t); - let q_trans = (2.0 * std::f64::consts::PI * molecular_mass * AMU2KG / (beta * H * H)).powf(1.5) * NA / (beta * p); - s[TRANS] = 2.5 + (q_trans / NA).ln(); + let beta = 1.0 / (BOLTZMANN * t); + let q_trans = (2.0 * std::f64::consts::PI * molecular_mass * AMU2KG / (beta * PLANCK * PLANCK)).powf(1.5) + * AVOGADRO + / (beta * p); + s[TRANS] = 2.5 + (q_trans / AVOGADRO).ln(); cv[TRANS] = 1.5; cp[TRANS] = 2.5; e[TRANS] = 1.5 * t; @@ -695,7 +686,7 @@ pub fn thermo( match rotor_type { RotorType::Atom => {}, RotorType::Linear => { - let q_rot = 1.0 / (beta * (sigma as f64) * 100.0 * C * H * rot_const[1]); + let q_rot = 1.0 / (beta * (sigma as f64) * 100.0 * SPEED_OF_LIGHT * PLANCK * rot_const[1]); s[ROT] = 1.0 + q_rot.ln(); cv[ROT] = 1.0; cp[ROT] = 1.0; @@ -703,9 +694,9 @@ pub fn thermo( }, RotorType::Regular => { let phi = [ - rot_const[0] * 100.0 * C * H / KB, - rot_const[1] * 100.0 * C * H / KB, - rot_const[2] * 100.0 * C * H / KB, + rot_const[0] * 100.0 * SPEED_OF_LIGHT * PLANCK / BOLTZMANN, + rot_const[1] * 100.0 * SPEED_OF_LIGHT * PLANCK / BOLTZMANN, + rot_const[2] * 100.0 * SPEED_OF_LIGHT * PLANCK / BOLTZMANN, ]; let q_rot = std::f64::consts::PI.sqrt() * t.powf(1.5) / ((sigma as f64) * (phi[0] * phi[1] * phi[2]).sqrt()); @@ -967,17 +958,18 @@ pub fn print_vibs( // Irrep — skip (no irrep/symmetry in our implementation) // scalar property rows — centered, colsp trailing - let labels: [&str; 5] = - ["Reduced mass [u]", "Force const [mDyne/A]", "Turning point v=0 [a0]", "RMS dev v=0 [a0 u^1/2]", "Char temp [K]"]; + let labels: [&str; 5] = [ + "Reduced mass [u]", + "Force const [mDyne/A]", + "Turning point v=0 [a0]", + "RMS dev v=0 [a0 u^1/2]", + "Char temp [K]", + ]; let val_slices: [&[f64]; 5] = [&vib.mu, &vib.k, &vib.xtp0, &vib.dq0, &vib.theta_vib]; for (label, vals) in labels.iter().zip(val_slices.iter()) { let mut l = format!("{}{}", " ".repeat(presp), ljust(label, prewidth)); for &vib in &row { - l.push_str(&format!( - "{}{}", - center(&format!("{:.*}", prec, vals[vib]), width), - " ".repeat(colsp) - )); + l.push_str(&format!("{}{}", center(&format!("{:.*}", prec, vals[vib]), width), " ".repeat(colsp))); } lines.push(l); } @@ -1008,20 +1000,11 @@ pub fn print_vibs( let lbl = if at < atom_lbl.len() { atom_lbl[at] } else { "" }; for xyz in 0..3 { let axis = ['X', 'Y', 'Z'][xyz]; - let mut l = format!( - "{}{:5} {} {}", - " ".repeat(presp), - at + 1, - axis, - ljust(lbl, prewidth - 14) - ); + let mut l = + format!("{}{:5} {} {}", " ".repeat(presp), at + 1, axis, ljust(lbl, prewidth - 14)); for &vib in &row { let v = normco_t[[3 * at + xyz, vib]]; - l.push_str(&format!( - "{}{}", - center(&format!("{:.*}", ncprec, v), width), - " ".repeat(colsp) - )); + l.push_str(&format!("{}{}", center(&format!("{:.*}", ncprec, v), width), " ".repeat(colsp))); } lines.push(l); } -- Gitee From d9ae86bf75ff2092e42e7df2c4ef117fa1597fdc Mon Sep 17 00:00:00 2001 From: ajz34 Date: Mon, 29 Jun 2026 10:59:12 +0800 Subject: [PATCH 19/39] vib: fix saturated minus of usize value --- src/hessian_backup/vib.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hessian_backup/vib.rs b/src/hessian_backup/vib.rs index 3be62337b9..688c20c001 100644 --- a/src/hessian_backup/vib.rs +++ b/src/hessian_backup/vib.rs @@ -376,7 +376,7 @@ pub fn harmonic_analysis( }; // translation is always 3 dof; rotations are 0/2/3 for ATOM/LINEAR/REGULAR. if !project_trans { - nrt = (nrt - 3).max(0); + nrt = nrt.max(3) - 3; } if !project_rot { let nrot = match rotor_type { @@ -384,7 +384,7 @@ pub fn harmonic_analysis( "LINEAR" => 2, _ => 3, }; - nrt = (nrt - nrot).max(0); + nrt = nrt.max(nrot) - nrot; } // --------------- translation / rotation projector --------------- -- Gitee From 7b732038abd9107181ffbcec14b66dadae368898 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Mon, 29 Jun 2026 13:36:17 +0800 Subject: [PATCH 20/39] point_group_detection: removed functions not used in actual pg detection --- src/hessian_backup/point_group_detect/bits.rs | 29 ------------------- .../point_group_detect/detect.rs | 13 +-------- .../point_group_detect/matrix.rs | 9 ------ src/hessian_backup/point_group_detect/mod.rs | 8 ++--- src/hessian_backup/point_group_detect/vec3.rs | 8 ----- 5 files changed, 5 insertions(+), 62 deletions(-) diff --git a/src/hessian_backup/point_group_detect/bits.rs b/src/hessian_backup/point_group_detect/bits.rs index 1fda4dd0cc..e36c8dca4b 100644 --- a/src/hessian_backup/point_group_detect/bits.rs +++ b/src/hessian_backup/point_group_detect/bits.rs @@ -74,28 +74,6 @@ pub fn bits_to_basic_name(bits: u8) -> &'static str { } } -/// All directionally-equivalent bit patterns for a given directional subgroup. -/// Reference: `similar`. Used only when a user overrides symmetry (not on the -/// bare detection path), but ported for completeness. -pub fn similar(bits: u8) -> Vec { - use point_groups::*; - let cs = [CSX, CSY, CSZ]; - let c2v = [C2VZ, C2VY, C2VX]; - let c2h = [C2HZ, C2HY, C2HX]; - let c2 = [C2Z, C2Y, C2X]; - if cs.contains(&bits) { - cs.to_vec() - } else if c2v.contains(&bits) { - c2v.to_vec() - } else if c2h.contains(&bits) { - c2h.to_vec() - } else if c2.contains(&bits) { - c2.to_vec() - } else { - vec![bits] // d2h, d2, ci, c1 are self-only - } -} - #[cfg(test)] mod tests { use super::*; @@ -108,11 +86,4 @@ mod tests { assert_eq!(bits_to_basic_name(D2H), "d2h"); assert_eq!(bits_to_basic_name(CSX), "cs"); } - - #[test] - fn similar_cs() { - use point_groups::*; - let s = similar(CSX); - assert!(s.contains(&CSX) && s.contains(&CSY) && s.contains(&CSZ)); - } } diff --git a/src/hessian_backup/point_group_detect/detect.rs b/src/hessian_backup/point_group_detect/detect.rs index f9f1047ec9..b010ff486f 100644 --- a/src/hessian_backup/point_group_detect/detect.rs +++ b/src/hessian_backup/point_group_detect/detect.rs @@ -1,6 +1,6 @@ //! Public detection API. -use super::molecule::{SymmMolecule, Tmpl}; +use super::molecule::SymmMolecule; /// A detected point group. #[derive(Debug, Clone)] @@ -9,8 +9,6 @@ pub struct PointGroup { pub full_name: String, /// Rotational symmetry number σ (Sn yields n/2, so this may be fractional). pub sigma: f64, - pub(crate) template: Tmpl, - pub(crate) n: u32, } impl SymmMolecule { @@ -20,15 +18,6 @@ impl SymmMolecule { PointGroup { full_name: template.full_name(n), sigma: template.sigma(n), - template, - n, } } } - -/// Convenience: detect the point-group symbol for a geometry (Bohr) given -/// element symbols. Uses default masses and tolerance. -pub fn detect_point_group(symbols: &[&str], geom: &[[f64; 3]]) -> String { - let m = SymmMolecule::new(symbols.iter().map(|s| s.to_string()).collect(), geom.to_vec()); - m.detect().full_name -} diff --git a/src/hessian_backup/point_group_detect/matrix.rs b/src/hessian_backup/point_group_detect/matrix.rs index 0d019ca562..e67f732923 100644 --- a/src/hessian_backup/point_group_detect/matrix.rs +++ b/src/hessian_backup/point_group_detect/matrix.rs @@ -72,15 +72,6 @@ pub fn points_apply(points: &[[f64; 3]], m: &Matrix3) -> Vec<[f64; 3]> { r } -/// Determinant of a 3×3 matrix. Reference: `vecutil.determinant`. -pub fn determinant(m: &Matrix3) -> f64 { - m[0][0] * m[1][1] * m[2][2] - m[0][2] * m[1][1] * m[2][0] - + m[0][1] * m[1][2] * m[2][0] - - m[0][1] * m[1][0] * m[2][2] - + m[0][2] * m[1][0] * m[2][1] - - m[0][0] * m[1][2] * m[2][1] -} - /// Rodrigues rotation matrix for rotation by `phi` about `axis` (need not be /// normalized). Returns the row-vector form `R` from `matrix_3d_rotation` /// (`libmintsmolecule.py:3167`); the caller applies it as `coord · Rᵀ` via diff --git a/src/hessian_backup/point_group_detect/mod.rs b/src/hessian_backup/point_group_detect/mod.rs index a0e3fb7142..dac1d38811 100644 --- a/src/hessian_backup/point_group_detect/mod.rs +++ b/src/hessian_backup/point_group_detect/mod.rs @@ -14,8 +14,8 @@ //! . //! //! # Example -//! ``` -//! use point_group_detection::{SymmMolecule, detect_point_group}; +//! ```ignore +//! use rest::hessian_backup::point_group_detect::SymmMolecule; //! //! // H2O (Bohr) -> C2v //! let mol = SymmMolecule::new( @@ -39,5 +39,5 @@ pub mod vec3; pub mod interface_to_rest; -pub use detect::{detect_point_group, PointGroup}; -pub use molecule::{SymmMolecule, Tmpl}; +pub use detect::PointGroup; +pub use molecule::SymmMolecule; diff --git a/src/hessian_backup/point_group_detect/vec3.rs b/src/hessian_backup/point_group_detect/vec3.rs index 144234bb1b..d1dbc6fd31 100644 --- a/src/hessian_backup/point_group_detect/vec3.rs +++ b/src/hessian_backup/point_group_detect/vec3.rs @@ -43,14 +43,6 @@ pub fn normalize(v: &Vec3) -> Vec3 { [v[0] / m, v[1] / m, v[2] / m] } -#[inline] -pub fn distance(v: &Vec3, u: &Vec3) -> f64 { - let d0 = v[0] - u[0]; - let d1 = v[1] - u[1]; - let d2 = v[2] - u[2]; - (d0 * d0 + d1 * d1 + d2 * d2).sqrt() -} - #[inline] pub fn cross(v: &Vec3, u: &Vec3) -> Vec3 { [ -- Gitee From d2ffa619ebbc535e40e8129e872882d5e789e678 Mon Sep 17 00:00:00 2001 From: Shirong_Wang Date: Wed, 8 Jul 2026 19:31:11 +0800 Subject: [PATCH 21/39] unify bohr --- src/constants/element.rs | 380 +++++++++++++++++++++++++++++ src/constants/mod.rs | 396 +------------------------------ src/dft/gen_grids/atom.rs | 2 +- src/dft/gen_grids/atom_old.rs | 2 +- src/dft/gen_grids/parameters.rs | 3 - src/dft/gen_grids/prune.rs | 5 +- src/dft/gen_grids/radial.rs | 3 +- src/fileop/chkfile.rs | 2 +- src/geom_io/mod.rs | 24 +- src/grad/mod.rs | 4 +- src/lib_rint/mod.rs | 2 +- src/main_driver.rs | 8 +- src/post_scf_analysis/mod.rs | 4 +- src/post_scf_analysis/rrs_pbc.rs | 6 +- src/ri_bse/dipoles.rs | 2 +- src/ri_bse/response.rs | 2 +- src/solvent/smd_cds.rs | 4 +- 17 files changed, 425 insertions(+), 424 deletions(-) create mode 100644 src/constants/element.rs diff --git a/src/constants/element.rs b/src/constants/element.rs new file mode 100644 index 0000000000..40185e0283 --- /dev/null +++ b/src/constants/element.rs @@ -0,0 +1,380 @@ +use lazy_static::lazy_static; + +pub const SPECIES_NAME: [&str; 118] = ["H", "He", + "Li","Be","B", "C","N", "O", "F", "Ne", + "Na","Mg","Al","Si","P", "S", "Cl","Ar", + "K", "Ca","Sc","Ti","V", "Cr","Mn","Fe","Co","Ni","Cu","Zn","Ga","Ge","As","Se","Br","Kr", + "Rb","Sr","Y", "Zr","Nb","Mo","Tc","Ru","Rh","Pd","Ag","Cd","In","Sn","Sb","Te","I", "Xe", + "Cs","Ba","La","Ce","Pr","Nd","Pm","Sm","Eu","Gd","Tb","Dy","Ho","Er","Tm","Yb","Lu","Hf","Ta","W", "Re","Os","Ir","Pt","Au","Hg","Tl","Pb","Bi","Po","At","Rn", + "Fr","Ra","Ac","Th","Pa","U", "Np","Pu","Am","Cm","Bk","Cf","Es","Fm","Md","No","Lr","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Cn","Nh","Fl","Mc","Lv","Ts","Og" + ]; + +// IUPAC 2021 standard atomic weights (Prohaska et al., Pure Appl. Chem., 2022) +// Table 1 column 7: abridged to 5 significant figures; interval elements use conventional single values +// 14 interval elements: H, Li, B, C, N, O, Mg, Si, S, Cl, Ar, Br, Tl, Pb +// Tc, Pm, Po-At, Rn-Ac, Np and beyond: most stable isotope mass +// AME2020 / NUBASE2020 via IUPAC 2021 Table 2 + +pub const MASS_CHARGE: [(f64,f64);118] = [ + (1.0080,1.0), (4.0026,2.0), + + (6.94,3.0), (9.0122,4.0), + (10.81,5.0), (12.011,6.0), (14.007,7.0), (15.999,8.0), + (18.998,9.0), (20.180,10.0), + + (22.990,11.0), (24.305,12.0), + (26.982,13.0), (28.085,14.0), (30.974,15.0), (32.06,16.0), + (35.45,17.0), (39.95,18.0), + + (39.098,19.0), (40.078,20.0), + (44.956,21.0), (47.867,22.0), (50.942,23.0), (51.996,24.0), + (54.938,25.0), (55.845,26.0), (58.933,27.0), (58.693,28.0), + (63.546,29.0), (65.38,30.0), + + (69.723,31.0), (72.630,32.0), + (74.922,33.0), (78.971,34.0), (79.904,35.0), (83.798,36.0), + + (85.468,37.0), (87.62,38.0), (88.906,39.0), (91.224,40.0), + (92.906,41.0), (95.95,42.0), (97.90721,43.0), (101.07,44.0), + (102.91,45.0), (106.42,46.0), (107.87,47.0), (112.41,48.0), + + (114.82,49.0), (118.71,50.0), (121.76,51.0), (127.60,52.0), + (126.90,53.0), (131.29,54.0), + + (132.91,55.0), (137.33,56.0), + + (138.91,57.0), (140.12,58.0), (140.91,59.0), (144.24,60.0), + (144.91276,61.0), (150.36,62.0), (151.96,63.0), (157.25,64.0), + (158.93,65.0), (162.50,66.0), (164.93,67.0), (167.26,68.0), + (168.93,69.0), (173.05,70.0), + + (174.97,71.0), (178.49,72.0), + (180.95,73.0), (183.84,74.0), (186.21,75.0), (190.23,76.0), + (192.22,77.0), (195.08,78.0), (196.97,79.0), (200.59,80.0), + + (204.38,81.0), (207.2,82.0), (208.98,83.0), (208.98243,84.0), + (209.98715,85.0), (222.01758,86.0), + + (223.01973,87.0), (226.02541,88.0), + (227.02775,89.0), (232.04,90.0), (231.04,91.0), (238.03,92.0), + + (237.04817,93.0), (244.06420,94.0), (243.06138,95.0), (247.07035,96.0), + (247.07031,97.0), (251.07959,98.0), (252.08298,99.0), (257.09511,100.0), + (258.09843,101.0), (259.10100,102.0), + + (262.10962,103.0), (267.12179,104.0), + (268.12567,105.0), (271.13378,106.0), (270.13337,107.0), (269.13365,108.0), + (278.15649,109.0), (281.16455,110.0), (282.16934,111.0), (285.17723,112.0), + + (286.18246,113.0), (289.19052,114.0), (289.19397,115.0), (293.20458,116.0), + (294.21084,117.0), (294.21398,118.0), +]; + +pub const ELEM1ST: [&str;2] = ["H", "He"]; +pub const ELEM2ND: [&str;8] = ["Li","Be","B", "C", "N", "O", "F", "Ne"]; +pub const ELEM3RD: [&str;8] = ["Na","Mg","Al","Si","P", "S", "Cl","Ar"]; +pub const ELEM4TH: [&str;18] = ["K", "Ca","Sc","Ti","V", "Cr","Mn","Fe","Co","Ni","Cu","Zn","Ga","Ge","As","Se","Br","Kr"]; +pub const ELEM5TH: [&str;18] = ["Rb","Sr","Y", "Zr","Nb","Mo","Tc","Ru","Rh","Pd","Ag","Cd","In","Sn","Sb","Te","I", "Xe"]; +pub const ELEM6TH: [&str;32] = ["Cs","Ba","La","Ce","Pr","Nd","Pm","Sm","Eu","Gd","Tb","Dy","Ho","Er","Tm","Yb","Lu","Hf","Ta","W", "Re","Os","Ir","Pt","Au","Hg","Tl","Pb","Bi","Po","At","Rn"]; +pub const ELEM7TH: [&str;32] = ["Fr","Ra","Ac","Th","Pa","U", "Np","Pu","Am","Cm","Bk","Cf","Es","Fm","Md","No","Lr","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Cn","Nh","Fl","Mc","Lv","Ts","Og"]; + +pub const ELEMTMS: [&str;40] = [ + "Sc","Ti","V", "Cr","Mn","Fe","Co","Ni","Cu","Zn", + "Y", "Zr","Nb","Mo","Tc","Ru","Rh","Pd","Ag","Cd", + "La","Hf","Ta","W", "Re","Os","Ir","Pt","Au","Hg", + "Ac","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Cn" +]; + + +lazy_static!{ + pub static ref SPECIES_INFO: HashMap<&'static str, &'static (f64,f64)> = { + let mut m = HashMap::new(); + SPECIES_NAME.iter().zip(MASS_CHARGE.iter()).for_each(|(name,info)| { + m.insert(*name,info); + }); + m + }; +} + +pub const ATOM_CONFIGURATION: [[usize; 4]; 119] = [ + [ 0, 0, 0, 0], // 0 GHOST + [ 1, 0, 0, 0], // 1 H + [ 2, 0, 0, 0], // 2 He + [ 3, 0, 0, 0], // 3 Li + [ 4, 0, 0, 0], // 4 Be + [ 4, 1, 0, 0], // 5 B + [ 4, 2, 0, 0], // 6 C + [ 4, 3, 0, 0], // 7 N + [ 4, 4, 0, 0], // 8 O + [ 4, 5, 0, 0], // 9 F + [ 4, 6, 0, 0], // 10 Ne + [ 5, 6, 0, 0], // 11 Na + [ 6, 6, 0, 0], // 12 Mg + [ 6, 7, 0, 0], // 13 Al + [ 6, 8, 0, 0], // 14 Si + [ 6, 9, 0, 0], // 15 P + [ 6,10, 0, 0], // 16 S + [ 6,11, 0, 0], // 17 Cl + [ 6,12, 0, 0], // 18 Ar + [ 7,12, 0, 0], // 19 K + [ 8,12, 0, 0], // 20 Ca + [ 8,12, 1, 0], // 21 Sc + [ 8,12, 2, 0], // 22 Ti + [ 8,12, 3, 0], // 23 V + [ 7,12, 5, 0], // 24 Cr + [ 8,12, 5, 0], // 25 Mn + [ 8,12, 6, 0], // 26 Fe + [ 8,12, 7, 0], // 27 Co + [ 8,12, 8, 0], // 28 Ni + [ 7,12,10, 0], // 29 Cu + [ 8,12,10, 0], // 30 Zn + [ 8,13,10, 0], // 31 Ga + [ 8,14,10, 0], // 32 Ge + [ 8,15,10, 0], // 33 As + [ 8,16,10, 0], // 34 Se + [ 8,17,10, 0], // 35 Br + [ 8,18,10, 0], // 36 Kr + [ 9,18,10, 0], // 37 Rb + [10,18,10, 0], // 38 Sr + [10,18,11, 0], // 39 Y + [10,18,12, 0], // 40 Zr + [ 9,18,14, 0], // 41 Nb + [ 9,18,15, 0], // 42 Mo + [10,18,15, 0], // 43 Tc + [ 9,18,17, 0], // 44 Ru + [ 9,18,18, 0], // 45 Rh + [ 8,18,20, 0], // 46 Pd + [ 9,18,20, 0], // 47 Ag + [10,18,20, 0], // 48 Cd + [10,19,20, 0], // 49 In + [10,20,20, 0], // 50 Sn + [10,21,20, 0], // 51 Sb + [10,22,20, 0], // 52 Te + [10,23,20, 0], // 53 I + [10,24,20, 0], // 54 Xe + [11,24,20, 0], // 55 Cs + [12,24,20, 0], // 56 Ba + [12,24,21, 0], // 57 La + [12,24,21, 1], // 58 Ce + [12,24,20, 3], // 59 Pr + [12,24,20, 4], // 60 Nd + [12,24,20, 5], // 61 Pm + [12,24,20, 6], // 62 Sm + [12,24,20, 7], // 63 Eu + [12,24,21, 7], // 64 Gd + [12,24,21, 8], // 65 Tb + [12,24,20,10], // 66 Dy + [12,24,20,11], // 67 Ho + [12,24,20,12], // 68 Er + [12,24,20,13], // 69 Tm + [12,24,20,14], // 70 Yb + [12,24,21,14], // 71 Lu + [12,24,22,14], // 72 Hf + [12,24,23,14], // 73 Ta + [12,24,24,14], // 74 W + [12,24,25,14], // 75 Re + [12,24,26,14], // 76 Os + [12,24,27,14], // 77 Ir + [11,24,29,14], // 78 Pt + [11,24,30,14], // 79 Au + [12,24,30,14], // 80 Hg + [12,25,30,14], // 81 Tl + [12,26,30,14], // 82 Pb + [12,27,30,14], // 83 Bi + [12,28,30,14], // 84 Po + [12,29,30,14], // 85 At + [12,30,30,14], // 86 Rn + [13,30,30,14], // 87 Fr + [14,30,30,14], // 88 Ra + [14,30,31,14], // 89 Ac + [14,30,32,14], // 90 Th + [14,30,31,16], // 91 Pa + [14,30,31,17], // 92 U + [14,30,31,18], // 93 Np + [14,30,30,20], // 94 Pu + [14,30,30,21], // 95 Am + [14,30,31,21], // 96 Cm + [14,30,31,22], // 97 Bk + [14,30,30,24], // 98 Cf + [14,30,30,25], // 99 Es + [14,30,30,26], //100 Fm + [14,30,30,27], //101 Md + [14,30,30,28], //102 No + [14,30,31,28], //103 Lr + [14,30,32,28], //104 Rf + [14,30,33,28], //105 Db + [14,30,34,28], //106 Sg + [14,30,35,28], //107 Bh + [14,30,36,28], //108 Hs + [14,30,37,28], //109 Mt + [14,30,38,28], //110 Ds + [14,30,39,28], //111 Rg + [14,30,40,28], //112 Cn + [14,31,40,28], //113 Nh + [14,32,40,28], //114 Fl + [14,33,40,28], //115 Mc + [14,34,40,28], //116 Lv + [14,35,40,28], //117 Ts + [14,36,40,28], //118 Og +]; + +// SAD and ECP configuration +pub const S_SHELL: [f64; 1] = [2.0]; +pub const P_SHELL: [f64; 3] = [2.0, 2.0, 2.0]; +pub const D_SHELL: [f64; 5] = [2.0, 2.0, 2.0, 2.0, 2.0]; +pub const F_SHELL: [f64; 7] = [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]; +pub const XE_SHELL: [f64; 27] = +// 1s 2s 2p 2p 2p 3s 3p 3p 3p 4s 3d 3d 3d 3d 3d 4p 4p 4p 5s 4d 4d 4d 4d 4d 5p 5p 5p + [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]; +pub const KR_SHELL: [f64; 18] = +// 1s 2s 2p 2p 2p 3s 3p 3p 3p 4s 3d 3d 3d 3d 3d 4p 4p 4p + [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]; + +// 1s 2s 2p 3s 3p 4s 3d 4p 5s 4d 5p 4f 5d 6s 6p +pub const NELE_IN_SHELLS: [f64; 15] = [2.0, 2.0, 6.0, 2.0, 6.0, 2.0, 10.0, 6.0, 2.0, 10.0, 6.0, 14.0, 10.0, 2.0, 6.0]; + +lazy_static!{ + pub static ref ATOMIC_RADII: HashMap<&'static str, &'static f64> = { + /// 获取元素周期表中所有元素的原子半径（单位：Å） + /// 数据来源：Clementi-Raimondi半径、共价半径和范德华半径的综合 + let mut radii = HashMap::new(); + + // 第1周期 + radii.insert( "H", &0.53); // 氢 + radii.insert("He", &0.31); // 氦 + + // 第2周期 + radii.insert("Li", &1.67); // 锂 + radii.insert("Be", &1.12); // 铍 + radii.insert( "B", &0.87); // 硼 + radii.insert( "C", &0.67); // 碳 + radii.insert( "N", &0.56); // 氮 + radii.insert( "O", &0.48); // 氧 + radii.insert( "F", &0.42); // 氟 + radii.insert("Ne", &0.38); // 氖 + + // 第3周期 + radii.insert("Na", &1.90); // 钠 + radii.insert("Mg", &1.45); // 镁 + radii.insert("Al", &1.18); // 铝 + radii.insert("Si", &1.11); // 硅 + radii.insert( "P", &0.98); // 磷 + radii.insert( "S", &0.88); // 硫 + radii.insert("Cl", &0.79); // 氯 + radii.insert("Ar", &0.71); // 氩 + + // 第4周期 + radii.insert( "K", &2.43); // 钾 + radii.insert("Ca", &1.94); // 钙 + radii.insert("Sc", &1.84); // 钪 + radii.insert("Ti", &1.76); // 钛 + radii.insert( "V", &1.71); // 钒 + radii.insert("Cr", &1.66); // 铬 + radii.insert("Mn", &1.61); // 锰 + radii.insert("Fe", &1.56); // 铁 + radii.insert("Co", &1.52); // 钴 + radii.insert("Ni", &1.49); // 镍 + radii.insert("Cu", &1.45); // 铜 + radii.insert("Zn", &1.42); // 锌 + radii.insert("Ga", &1.36); // 镓 + radii.insert("Ge", &1.25); // 锗 + radii.insert("As", &1.14); // 砷 + radii.insert("Se", &1.03); // 硒 + radii.insert("Br", &0.94); // 溴 + radii.insert("Kr", &0.88); // 氪 + + // 第5周期 + radii.insert("Rb", &2.65); // 铷 + radii.insert("Sr", &2.19); // 锶 + radii.insert( "Y", &2.12); // 钇 + radii.insert("Zr", &2.06); // 锆 + radii.insert("Nb", &1.98); // 铌 + radii.insert("Mo", &1.90); // 钼 + radii.insert("Tc", &1.83); // 锝 + radii.insert("Ru", &1.78); // 钌 + radii.insert("Rh", &1.73); // 铑 + radii.insert("Pd", &1.69); // 钯 + radii.insert("Ag", &1.65); // 银 + radii.insert("Cd", &1.61); // 镉 + radii.insert("In", &1.56); // 铟 + radii.insert("Sn", &1.45); // 锡 + radii.insert("Sb", &1.33); // 锑 + radii.insert("Te", &1.23); // 碲 + radii.insert( "I", &1.15); // 碘 + radii.insert("Xe", &1.08); // 氙 + + // 第6周期 + radii.insert("Cs", &2.98); // 铯 + radii.insert("Ba", &2.53); // 钡 + radii.insert("La", &2.50); // 镧 + radii.insert("Ce", &2.48); // 铈 + radii.insert("Pr", &2.47); // 镨 + radii.insert("Nd", &2.45); // 钕 + radii.insert("Pm", &2.43); // 钷 + radii.insert("Sm", &2.42); // 钐 + radii.insert("Eu", &2.40); // 铕 + radii.insert("Gd", &2.38); // 钆 + radii.insert("Tb", &2.37); // 铽 + radii.insert("Dy", &2.35); // 镝 + radii.insert("Ho", &2.33); // 钬 + radii.insert("Er", &2.32); // 铒 + radii.insert("Tm", &2.30); // 铥 + radii.insert("Yb", &2.28); // 镱 + radii.insert("Lu", &2.27); // 镥 + radii.insert("Hf", &2.25); // 铪 + radii.insert("Ta", &2.20); // 钽 + radii.insert( "W", &2.10); // 钨 + radii.insert("Re", &2.05); // 铼 + radii.insert("Os", &2.00); // 锇 + radii.insert("Ir", &1.97); // 铱 + radii.insert("Pt", &1.92); // 铂 + radii.insert("Au", &1.87); // 金 + radii.insert("Hg", &1.75); // 汞 + radii.insert("Tl", &1.70); // 铊 + radii.insert("Pb", &1.54); // 铅 + radii.insert("Bi", &1.43); // 铋 + radii.insert("Po", &1.35); // 钋 + radii.insert("At", &1.27); // 砹 + radii.insert("Rn", &1.20); // 氡 + + // 第7周期 + radii.insert("Fr", &3.00); // 钫 + radii.insert("Ra", &2.70); // 镭 + radii.insert("Ac", &2.60); // 锕 + radii.insert("Th", &2.50); // 钍 + radii.insert("Pa", &2.40); // 镤 + radii.insert( "U", &2.30); // 铀 + radii.insert("Np", &2.30); // 镎 + radii.insert("Pu", &2.30); // 钚 + radii.insert("Am", &2.30); // 镅 + radii.insert("Cm", &2.30); // 锔 + radii.insert("Bk", &2.30); // 锫 + radii.insert("Cf", &2.30); // 锎 + radii.insert("Es", &2.30); // 锿 + radii.insert("Fm", &2.30); // 镄 + radii.insert("Md", &2.30); // 钔 + radii.insert("No", &2.30); // 锘 + radii.insert("Lr", &2.30); // 铹 + radii.insert("Rf", &2.30); // 卢瑟福 + radii.insert("Db", &2.30); // 𨧀 + radii.insert("Sg", &2.30); // 𨭎 + radii.insert("Bh", &2.30); // 𨨏 + radii.insert("Hs", &2.30); // 𨭆 + radii.insert("Mt", &2.30); // 䥑 + radii.insert("Ds", &2.30); // 鐽 + radii.insert("Rg", &2.30); // 錀 + radii.insert("Cn", &2.30); // 鎶 + radii.insert("Nh", &2.30); // 鉨 + radii.insert("Fl", &2.30); // 鈇 + radii.insert("Mc", &2.30); // 鏌 + radii.insert("Lv", &2.30); // 鉝 + radii.insert("Ts", &2.30); // 鿬 + radii.insert("Og", &2.30); // 鿫 + + // 添加一些常见的同位素和特殊表示 + radii.insert("D", &0.53); // 氘 + radii.insert("T", &0.53); // 氚 + + radii + }; +} \ No newline at end of file diff --git a/src/constants/mod.rs b/src/constants/mod.rs index d92d992271..b5d61f973a 100644 --- a/src/constants/mod.rs +++ b/src/constants/mod.rs @@ -10,119 +10,15 @@ use lazy_static::lazy_static; pub use crate::constants::vsap::*; pub use crate::constants::c2s::*; pub use crate::constants::cartesian_gto::*; +pub use crate::constants::element::*; -//struct Matrix3x3 { -// size -// -//} - -pub const SPECIES_NAME: [&str; 118] = ["H", "He", - "Li","Be","B", "C","N", "O", "F", "Ne", - "Na","Mg","Al","Si","P", "S", "Cl","Ar", - "K", "Ca","Sc","Ti","V", "Cr","Mn","Fe","Co","Ni","Cu","Zn","Ga","Ge","As","Se","Br","Kr", - "Rb","Sr","Y", "Zr","Nb","Mo","Tc","Ru","Rh","Pd","Ag","Cd","In","Sn","Sb","Te","I", "Xe", - "Cs","Ba","La","Ce","Pr","Nd","Pm","Sm","Eu","Gd","Tb","Dy","Ho","Er","Tm","Yb","Lu","Hf","Ta","W", "Re","Os","Ir","Pt","Au","Hg","Tl","Pb","Bi","Po","At","Rn", - "Fr","Ra","Ac","Th","Pa","U", "Np","Pu","Am","Cm","Bk","Cf","Es","Fm","Md","No","Lr","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Cn","Nh","Fl","Mc","Lv","Ts","Og" - ]; - -// IUPAC 2021 standard atomic weights (Prohaska et al., Pure Appl. Chem., 2022) -// Table 1 column 7: abridged to 5 significant figures; interval elements use conventional single values -// 14 interval elements: H, Li, B, C, N, O, Mg, Si, S, Cl, Ar, Br, Tl, Pb -// Tc, Pm, Po-At, Rn-Ac, Np and beyond: most stable isotope mass -// AME2020 / NUBASE2020 via IUPAC 2021 Table 2 - -pub const MASS_CHARGE: [(f64,f64);118] = [ - (1.0080,1.0), (4.0026,2.0), - - (6.94,3.0), (9.0122,4.0), - (10.81,5.0), (12.011,6.0), (14.007,7.0), (15.999,8.0), - (18.998,9.0), (20.180,10.0), - - (22.990,11.0), (24.305,12.0), - (26.982,13.0), (28.085,14.0), (30.974,15.0), (32.06,16.0), - (35.45,17.0), (39.95,18.0), - - (39.098,19.0), (40.078,20.0), - (44.956,21.0), (47.867,22.0), (50.942,23.0), (51.996,24.0), - (54.938,25.0), (55.845,26.0), (58.933,27.0), (58.693,28.0), - (63.546,29.0), (65.38,30.0), - - (69.723,31.0), (72.630,32.0), - (74.922,33.0), (78.971,34.0), (79.904,35.0), (83.798,36.0), - - (85.468,37.0), (87.62,38.0), (88.906,39.0), (91.224,40.0), - (92.906,41.0), (95.95,42.0), (97.90721,43.0), (101.07,44.0), - (102.91,45.0), (106.42,46.0), (107.87,47.0), (112.41,48.0), - - (114.82,49.0), (118.71,50.0), (121.76,51.0), (127.60,52.0), - (126.90,53.0), (131.29,54.0), - - (132.91,55.0), (137.33,56.0), - (138.91,57.0), (140.12,58.0), (140.91,59.0), (144.24,60.0), - (144.91276,61.0), (150.36,62.0), (151.96,63.0), (157.25,64.0), - (158.93,65.0), (162.50,66.0), (164.93,67.0), (167.26,68.0), - (168.93,69.0), (173.05,70.0), - - (174.97,71.0), (178.49,72.0), - (180.95,73.0), (183.84,74.0), (186.21,75.0), (190.23,76.0), - (192.22,77.0), (195.08,78.0), (196.97,79.0), (200.59,80.0), - - (204.38,81.0), (207.2,82.0), (208.98,83.0), (208.98243,84.0), - (209.98715,85.0), (222.01758,86.0), - - (223.01973,87.0), (226.02541,88.0), - (227.02775,89.0), (232.04,90.0), (231.04,91.0), (238.03,92.0), - - (237.04817,93.0), (244.06420,94.0), (243.06138,95.0), (247.07035,96.0), - (247.07031,97.0), (251.07959,98.0), (252.08298,99.0), (257.09511,100.0), - (258.09843,101.0), (259.10100,102.0), - - (262.10962,103.0), (267.12179,104.0), - (268.12567,105.0), (271.13378,106.0), (270.13337,107.0), (269.13365,108.0), - (278.15649,109.0), (281.16455,110.0), (282.16934,111.0), (285.17723,112.0), - - (286.18246,113.0), (289.19052,114.0), (289.19397,115.0), (293.20458,116.0), - (294.21084,117.0), (294.21398,118.0), -]; - -pub const ELEM1ST: [&str;2] = ["H", "He"]; -pub const ELEM2ND: [&str;8] = ["Li","Be","B", "C", "N", "O", "F", "Ne"]; -pub const ELEM3RD: [&str;8] = ["Na","Mg","Al","Si","P", "S", "Cl","Ar"]; -pub const ELEM4TH: [&str;18] = ["K", "Ca","Sc","Ti","V", "Cr","Mn","Fe","Co","Ni","Cu","Zn","Ga","Ge","As","Se","Br","Kr"]; -pub const ELEM5TH: [&str;18] = ["Rb","Sr","Y", "Zr","Nb","Mo","Tc","Ru","Rh","Pd","Ag","Cd","In","Sn","Sb","Te","I", "Xe"]; -pub const ELEM6TH: [&str;32] = ["Cs","Ba","La","Ce","Pr","Nd","Pm","Sm","Eu","Gd","Tb","Dy","Ho","Er","Tm","Yb","Lu","Hf","Ta","W", "Re","Os","Ir","Pt","Au","Hg","Tl","Pb","Bi","Po","At","Rn"]; -pub const ELEM7TH: [&str;32] = ["Fr","Ra","Ac","Th","Pa","U", "Np","Pu","Am","Cm","Bk","Cf","Es","Fm","Md","No","Lr","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Cn","Nh","Fl","Mc","Lv","Ts","Og"]; - -pub const ELEMTMS: [&str;40] = [ - "Sc","Ti","V", "Cr","Mn","Fe","Co","Ni","Cu","Zn", - "Y", "Zr","Nb","Mo","Tc","Ru","Rh","Pd","Ag","Cd", - "La","Hf","Ta","W", "Re","Os","Ir","Pt","Au","Hg", - "Ac","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Cn" -]; - - -lazy_static!{ - pub static ref SPECIES_INFO: HashMap<&'static str, &'static (f64,f64)> = { - let mut m = HashMap::new(); - SPECIES_NAME.iter().zip(MASS_CHARGE.iter()).for_each(|(name,info)| { - m.insert(*name,info); - }); - m - }; -} // use for the inverse (sqrt inverse) of the auxiliary coulomb matrix pub const AUXBAS_THRESHOLD: f64 = 1.0e-10; pub const INVERSE_THRESHOLD: f64 = 1.0e-10; pub const SQRT_THRESHOLD: f64 = 1.0e-10; -pub const CM: f64 = 8065.541; -pub const ANG:f64 = 0.5291772083; -pub const EV: f64 = 27.2113845; -pub const FQ: f64 = 1822.888; -pub const E: f64 = std::f64::consts::E; -pub const PI: f64 = std::f64::consts::PI; pub const E5: f64 = 1.0e5; pub const E6: f64 = 1.0e6; @@ -136,127 +32,7 @@ pub struct DMatrix { indicing:[usize;2], data: [f64; L] } -pub const ATOM_CONFIGURATION: [[usize; 4]; 119] = [ - [ 0, 0, 0, 0], // 0 GHOST - [ 1, 0, 0, 0], // 1 H - [ 2, 0, 0, 0], // 2 He - [ 3, 0, 0, 0], // 3 Li - [ 4, 0, 0, 0], // 4 Be - [ 4, 1, 0, 0], // 5 B - [ 4, 2, 0, 0], // 6 C - [ 4, 3, 0, 0], // 7 N - [ 4, 4, 0, 0], // 8 O - [ 4, 5, 0, 0], // 9 F - [ 4, 6, 0, 0], // 10 Ne - [ 5, 6, 0, 0], // 11 Na - [ 6, 6, 0, 0], // 12 Mg - [ 6, 7, 0, 0], // 13 Al - [ 6, 8, 0, 0], // 14 Si - [ 6, 9, 0, 0], // 15 P - [ 6,10, 0, 0], // 16 S - [ 6,11, 0, 0], // 17 Cl - [ 6,12, 0, 0], // 18 Ar - [ 7,12, 0, 0], // 19 K - [ 8,12, 0, 0], // 20 Ca - [ 8,12, 1, 0], // 21 Sc - [ 8,12, 2, 0], // 22 Ti - [ 8,12, 3, 0], // 23 V - [ 7,12, 5, 0], // 24 Cr - [ 8,12, 5, 0], // 25 Mn - [ 8,12, 6, 0], // 26 Fe - [ 8,12, 7, 0], // 27 Co - [ 8,12, 8, 0], // 28 Ni - [ 7,12,10, 0], // 29 Cu - [ 8,12,10, 0], // 30 Zn - [ 8,13,10, 0], // 31 Ga - [ 8,14,10, 0], // 32 Ge - [ 8,15,10, 0], // 33 As - [ 8,16,10, 0], // 34 Se - [ 8,17,10, 0], // 35 Br - [ 8,18,10, 0], // 36 Kr - [ 9,18,10, 0], // 37 Rb - [10,18,10, 0], // 38 Sr - [10,18,11, 0], // 39 Y - [10,18,12, 0], // 40 Zr - [ 9,18,14, 0], // 41 Nb - [ 9,18,15, 0], // 42 Mo - [10,18,15, 0], // 43 Tc - [ 9,18,17, 0], // 44 Ru - [ 9,18,18, 0], // 45 Rh - [ 8,18,20, 0], // 46 Pd - [ 9,18,20, 0], // 47 Ag - [10,18,20, 0], // 48 Cd - [10,19,20, 0], // 49 In - [10,20,20, 0], // 50 Sn - [10,21,20, 0], // 51 Sb - [10,22,20, 0], // 52 Te - [10,23,20, 0], // 53 I - [10,24,20, 0], // 54 Xe - [11,24,20, 0], // 55 Cs - [12,24,20, 0], // 56 Ba - [12,24,21, 0], // 57 La - [12,24,21, 1], // 58 Ce - [12,24,20, 3], // 59 Pr - [12,24,20, 4], // 60 Nd - [12,24,20, 5], // 61 Pm - [12,24,20, 6], // 62 Sm - [12,24,20, 7], // 63 Eu - [12,24,21, 7], // 64 Gd - [12,24,21, 8], // 65 Tb - [12,24,20,10], // 66 Dy - [12,24,20,11], // 67 Ho - [12,24,20,12], // 68 Er - [12,24,20,13], // 69 Tm - [12,24,20,14], // 70 Yb - [12,24,21,14], // 71 Lu - [12,24,22,14], // 72 Hf - [12,24,23,14], // 73 Ta - [12,24,24,14], // 74 W - [12,24,25,14], // 75 Re - [12,24,26,14], // 76 Os - [12,24,27,14], // 77 Ir - [11,24,29,14], // 78 Pt - [11,24,30,14], // 79 Au - [12,24,30,14], // 80 Hg - [12,25,30,14], // 81 Tl - [12,26,30,14], // 82 Pb - [12,27,30,14], // 83 Bi - [12,28,30,14], // 84 Po - [12,29,30,14], // 85 At - [12,30,30,14], // 86 Rn - [13,30,30,14], // 87 Fr - [14,30,30,14], // 88 Ra - [14,30,31,14], // 89 Ac - [14,30,32,14], // 90 Th - [14,30,31,16], // 91 Pa - [14,30,31,17], // 92 U - [14,30,31,18], // 93 Np - [14,30,30,20], // 94 Pu - [14,30,30,21], // 95 Am - [14,30,31,21], // 96 Cm - [14,30,31,22], // 97 Bk - [14,30,30,24], // 98 Cf - [14,30,30,25], // 99 Es - [14,30,30,26], //100 Fm - [14,30,30,27], //101 Md - [14,30,30,28], //102 No - [14,30,31,28], //103 Lr - [14,30,32,28], //104 Rf - [14,30,33,28], //105 Db - [14,30,34,28], //106 Sg - [14,30,35,28], //107 Bh - [14,30,36,28], //108 Hs - [14,30,37,28], //109 Mt - [14,30,38,28], //110 Ds - [14,30,39,28], //111 Rg - [14,30,40,28], //112 Cn - [14,31,40,28], //113 Nh - [14,32,40,28], //114 Fl - [14,33,40,28], //115 Mc - [14,34,40,28], //116 Lv - [14,35,40,28], //117 Ts - [14,36,40,28], //118 Og -]; + // =========== libcint =================================== // for the bas index - libcint @@ -300,25 +76,16 @@ pub const NUC_FRAC_CHARGE: i32 = 3; pub const ENV_PRT_START: usize = 20; - -// SAD and ECP configuration -pub const S_SHELL: [f64; 1] = [2.0]; -pub const P_SHELL: [f64; 3] = [2.0, 2.0, 2.0]; -pub const D_SHELL: [f64; 5] = [2.0, 2.0, 2.0, 2.0, 2.0]; -pub const F_SHELL: [f64; 7] = [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]; -pub const XE_SHELL: [f64; 27] = -// 1s 2s 2p 2p 2p 3s 3p 3p 3p 4s 3d 3d 3d 3d 3d 4p 4p 4p 5s 4d 4d 4d 4d 4d 5p 5p 5p - [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]; -pub const KR_SHELL: [f64; 18] = -// 1s 2s 2p 2p 2p 3s 3p 3p 3p 4s 3d 3d 3d 3d 3d 4p 4p 4p - [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]; - -// 1s 2s 2p 3s 3p 4s 3d 4p 5s 4d 5p 4f 5d 6s 6p -pub const NELE_IN_SHELLS: [f64; 15] = [2.0, 2.0, 6.0, 2.0, 6.0, 2.0, 10.0, 6.0, 2.0, 10.0, 6.0, 14.0, 10.0, 2.0, 6.0]; - +// math, physics +pub const CM: f64 = 8065.541; +pub const EV: f64 = 27.2113845; +pub const FQ: f64 = 1822.888; +pub const E: f64 = std::f64::consts::E; +pub const PI: f64 = std::f64::consts::PI; // pub const LIGHT_SPEED: f64 = 137.03599967994; // http://physics.nist.gov/cgi-bin/cuu/Value?alph // BOHR = .529 177 210 92(17) e-10m // http://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0 +// source: CODATA 2010 (https://physics.nist.gov/cuu/Constants/ArchiveASCII/allascii_2010.txt) pub const BOHR: f64 = 0.52917721092; // Angstroms pub const BOHR_SI: f64 = BOHR * 1e-10; @@ -336,148 +103,3 @@ pub const AU2DEBYE:f64 = E_CHARGE * BOHR*1e-10 / DEBYE; // 2.541746 pub const MPI_CHUNK:usize = 134217728; // around 1 GB -lazy_static!{ - pub static ref ATOMIC_RADII: HashMap<&'static str, &'static f64> = { - /// 获取元素周期表中所有元素的原子半径（单位：Å） - /// 数据来源：Clementi-Raimondi半径、共价半径和范德华半径的综合 - let mut radii = HashMap::new(); - - // 第1周期 - radii.insert( "H", &0.53); // 氢 - radii.insert("He", &0.31); // 氦 - - // 第2周期 - radii.insert("Li", &1.67); // 锂 - radii.insert("Be", &1.12); // 铍 - radii.insert( "B", &0.87); // 硼 - radii.insert( "C", &0.67); // 碳 - radii.insert( "N", &0.56); // 氮 - radii.insert( "O", &0.48); // 氧 - radii.insert( "F", &0.42); // 氟 - radii.insert("Ne", &0.38); // 氖 - - // 第3周期 - radii.insert("Na", &1.90); // 钠 - radii.insert("Mg", &1.45); // 镁 - radii.insert("Al", &1.18); // 铝 - radii.insert("Si", &1.11); // 硅 - radii.insert( "P", &0.98); // 磷 - radii.insert( "S", &0.88); // 硫 - radii.insert("Cl", &0.79); // 氯 - radii.insert("Ar", &0.71); // 氩 - - // 第4周期 - radii.insert( "K", &2.43); // 钾 - radii.insert("Ca", &1.94); // 钙 - radii.insert("Sc", &1.84); // 钪 - radii.insert("Ti", &1.76); // 钛 - radii.insert( "V", &1.71); // 钒 - radii.insert("Cr", &1.66); // 铬 - radii.insert("Mn", &1.61); // 锰 - radii.insert("Fe", &1.56); // 铁 - radii.insert("Co", &1.52); // 钴 - radii.insert("Ni", &1.49); // 镍 - radii.insert("Cu", &1.45); // 铜 - radii.insert("Zn", &1.42); // 锌 - radii.insert("Ga", &1.36); // 镓 - radii.insert("Ge", &1.25); // 锗 - radii.insert("As", &1.14); // 砷 - radii.insert("Se", &1.03); // 硒 - radii.insert("Br", &0.94); // 溴 - radii.insert("Kr", &0.88); // 氪 - - // 第5周期 - radii.insert("Rb", &2.65); // 铷 - radii.insert("Sr", &2.19); // 锶 - radii.insert( "Y", &2.12); // 钇 - radii.insert("Zr", &2.06); // 锆 - radii.insert("Nb", &1.98); // 铌 - radii.insert("Mo", &1.90); // 钼 - radii.insert("Tc", &1.83); // 锝 - radii.insert("Ru", &1.78); // 钌 - radii.insert("Rh", &1.73); // 铑 - radii.insert("Pd", &1.69); // 钯 - radii.insert("Ag", &1.65); // 银 - radii.insert("Cd", &1.61); // 镉 - radii.insert("In", &1.56); // 铟 - radii.insert("Sn", &1.45); // 锡 - radii.insert("Sb", &1.33); // 锑 - radii.insert("Te", &1.23); // 碲 - radii.insert( "I", &1.15); // 碘 - radii.insert("Xe", &1.08); // 氙 - - // 第6周期 - radii.insert("Cs", &2.98); // 铯 - radii.insert("Ba", &2.53); // 钡 - radii.insert("La", &2.50); // 镧 - radii.insert("Ce", &2.48); // 铈 - radii.insert("Pr", &2.47); // 镨 - radii.insert("Nd", &2.45); // 钕 - radii.insert("Pm", &2.43); // 钷 - radii.insert("Sm", &2.42); // 钐 - radii.insert("Eu", &2.40); // 铕 - radii.insert("Gd", &2.38); // 钆 - radii.insert("Tb", &2.37); // 铽 - radii.insert("Dy", &2.35); // 镝 - radii.insert("Ho", &2.33); // 钬 - radii.insert("Er", &2.32); // 铒 - radii.insert("Tm", &2.30); // 铥 - radii.insert("Yb", &2.28); // 镱 - radii.insert("Lu", &2.27); // 镥 - radii.insert("Hf", &2.25); // 铪 - radii.insert("Ta", &2.20); // 钽 - radii.insert( "W", &2.10); // 钨 - radii.insert("Re", &2.05); // 铼 - radii.insert("Os", &2.00); // 锇 - radii.insert("Ir", &1.97); // 铱 - radii.insert("Pt", &1.92); // 铂 - radii.insert("Au", &1.87); // 金 - radii.insert("Hg", &1.75); // 汞 - radii.insert("Tl", &1.70); // 铊 - radii.insert("Pb", &1.54); // 铅 - radii.insert("Bi", &1.43); // 铋 - radii.insert("Po", &1.35); // 钋 - radii.insert("At", &1.27); // 砹 - radii.insert("Rn", &1.20); // 氡 - - // 第7周期 - radii.insert("Fr", &3.00); // 钫 - radii.insert("Ra", &2.70); // 镭 - radii.insert("Ac", &2.60); // 锕 - radii.insert("Th", &2.50); // 钍 - radii.insert("Pa", &2.40); // 镤 - radii.insert( "U", &2.30); // 铀 - radii.insert("Np", &2.30); // 镎 - radii.insert("Pu", &2.30); // 钚 - radii.insert("Am", &2.30); // 镅 - radii.insert("Cm", &2.30); // 锔 - radii.insert("Bk", &2.30); // 锫 - radii.insert("Cf", &2.30); // 锎 - radii.insert("Es", &2.30); // 锿 - radii.insert("Fm", &2.30); // 镄 - radii.insert("Md", &2.30); // 钔 - radii.insert("No", &2.30); // 锘 - radii.insert("Lr", &2.30); // 铹 - radii.insert("Rf", &2.30); // 卢瑟福 - radii.insert("Db", &2.30); // 𨧀 - radii.insert("Sg", &2.30); // 𨭎 - radii.insert("Bh", &2.30); // 𨨏 - radii.insert("Hs", &2.30); // 𨭆 - radii.insert("Mt", &2.30); // 䥑 - radii.insert("Ds", &2.30); // 鐽 - radii.insert("Rg", &2.30); // 錀 - radii.insert("Cn", &2.30); // 鎶 - radii.insert("Nh", &2.30); // 鉨 - radii.insert("Fl", &2.30); // 鈇 - radii.insert("Mc", &2.30); // 鏌 - radii.insert("Lv", &2.30); // 鉝 - radii.insert("Ts", &2.30); // 鿬 - radii.insert("Og", &2.30); // 鿫 - - // 添加一些常见的同位素和特殊表示 - radii.insert("D", &0.53); // 氘 - radii.insert("T", &0.53); // 氚 - - radii - }; -} \ No newline at end of file diff --git a/src/dft/gen_grids/atom.rs b/src/dft/gen_grids/atom.rs index eb7b74435d..97570c81f9 100644 --- a/src/dft/gen_grids/atom.rs +++ b/src/dft/gen_grids/atom.rs @@ -121,7 +121,7 @@ pub fn atom_grid( //println!("radial num = {}", default_radial_num(proton_charges[center_index] as usize)); // factors match DIRAC code //println!("rs = {:?}, w = {:?}", rs, weights_radial); - let rb = bragg::get_bragg_angstrom(proton_charges[center_index]) / (5.0 * 0.529177249); + let rb = bragg::get_bragg_angstrom(proton_charges[center_index]) / (5.0 * crate::constants::BOHR); let mut coordinates = Vec::new(); let mut weights = Vec::new(); diff --git a/src/dft/gen_grids/atom_old.rs b/src/dft/gen_grids/atom_old.rs index 87a6ef95aa..24b68f05ec 100644 --- a/src/dft/gen_grids/atom_old.rs +++ b/src/dft/gen_grids/atom_old.rs @@ -91,7 +91,7 @@ pub fn atom_grid( //println!("radial num = {}", default_radial_num(proton_charges[center_index] as usize)); // factors match DIRAC code //println!("rs = {:?}, w = {:?}", rs, weights_radial); - let rb = bragg::get_bragg_angstrom(proton_charges[center_index]) / (5.0 * 0.529177249); + let rb = bragg::get_bragg_angstrom(proton_charges[center_index]) / (5.0 * crate::constants::BOHR); let mut coordinates = Vec::new(); let mut weights = Vec::new(); diff --git a/src/dft/gen_grids/parameters.rs b/src/dft/gen_grids/parameters.rs index 265a7dfdc9..897af95671 100644 --- a/src/dft/gen_grids/parameters.rs +++ b/src/dft/gen_grids/parameters.rs @@ -7,9 +7,6 @@ pub static SG1RADII: [f64; 19] = [0.0, 1.0000, 0.5882, //H, He 3.0769, 2.0513, 1.5385, 1.2308, 1.0256, 0.8791, 0.7692, 0.6838, //2nd Period 4.0909, 3.1579, 2.5714, 2.1687, 1.8750, 1.6514, 1.4754, 1.3333]; //3rd Period -/// Bohr radius. -pub const BOHR: f64 = 0.52917721092; - /// Bragg radii for elements. pub static BRAGG0: [f64; 131] = [0.0, // Ghost atom 0.35, 1.40, // 1s diff --git a/src/dft/gen_grids/prune.rs b/src/dft/gen_grids/prune.rs index 89004b4358..34baac0ba1 100644 --- a/src/dft/gen_grids/prune.rs +++ b/src/dft/gen_grids/prune.rs @@ -7,9 +7,10 @@ use num_traits::{ToPrimitive}; use tensors::MatrixFull; -//use super::parameters::{SG1RADII, BOHR, BRAGG0, LEBEDEV_NGRID}; +//use super::parameters::{SG1RADII, BRAGG0, LEBEDEV_NGRID}; use crate::{dft::Grids, utilities::balancing}; -use super::{parameters::{SG1RADII, BOHR, BRAGG0, LEBEDEV_NGRID}, atom::default_angular_num}; +use crate::constants::BOHR; +use super::{parameters::{SG1RADII, BRAGG0, LEBEDEV_NGRID}, atom::default_angular_num}; /// Standard Grid 1 according to _P. M. W. Gill, B. G. Johnson, J. A. Pople. Chemical Physics Letters 209, 506-512 (1993)_.
/// Reference can be found [here](https://doi.org/10.1016/0009-2614(93)80125-9). diff --git a/src/dft/gen_grids/radial.rs b/src/dft/gen_grids/radial.rs index 4a6e57b7bb..25f1205cd9 100644 --- a/src/dft/gen_grids/radial.rs +++ b/src/dft/gen_grids/radial.rs @@ -28,7 +28,8 @@ use std::f64::consts::PI; use super::bragg; use super::bse; use super::parameters; -use super::parameters::{BOHR, BRAGG0}; +use super::parameters::BRAGG0; +use crate::constants::BOHR; use statrs::function::gamma; /// Krack-Koster radial grid according to _M. Krack, A. M. Köster. The Journal of Chemical Physics 108, 3226-3234 (1998)_, eqs. 9-13.
diff --git a/src/fileop/chkfile.rs b/src/fileop/chkfile.rs index aea7d1929d..77d1c2c40e 100644 --- a/src/fileop/chkfile.rs +++ b/src/fileop/chkfile.rs @@ -183,7 +183,7 @@ pub fn save_overlap(scf_data: &SCF) { } pub fn save_geometry(scf_data: &SCF) { - let ang = crate::constants::ANG; + let ang = crate::constants::BOHR; let chkfile= &scf_data.mol.ctrl.chkfile; let path = Path::new(chkfile); let file = if path.exists() { diff --git a/src/geom_io/mod.rs b/src/geom_io/mod.rs index 7143f60fa3..b2591fff6f 100644 --- a/src/geom_io/mod.rs +++ b/src/geom_io/mod.rs @@ -14,7 +14,7 @@ use serde_json::Value; //use tensors::Tensors; use crate::basis_io::Basis4Elem; -use crate::constants::{ANG, ATOMIC_RADII, MASS_CHARGE, SPECIES_INFO, SPECIES_NAME}; +use crate::constants::{BOHR, ATOMIC_RADII, MASS_CHARGE, SPECIES_INFO, SPECIES_NAME}; use crate::external_field::ExtField; mod pyrest_geom_io; @@ -343,7 +343,7 @@ impl GeomCell { } pub fn geom_update(&mut self, new_position:&[f64], unit: GeomUnit) { let factor = match unit { - GeomUnit::Angstrom => ANG, + GeomUnit::Angstrom => BOHR, GeomUnit::Bohr => 1.0, }; if self.position.data.len() != new_position.len() { @@ -423,7 +423,7 @@ impl GeomCell { let mut tmp_pos_tensor = MatrixFull::from_vec(tmp_size, tmp_pos).unwrap(); if let GeomUnit::Angstrom = unit { // To store the geometry position in "Bohr" according to the convention of quantum chemistry. - tmp_pos_tensor.self_multiple(ANG.powf(-1.0)); + tmp_pos_tensor.self_multiple(BOHR.powf(-1.0)); }; Ok((tmp_ele, tmp_fix, tmp_pos_tensor, tmp_nfree)) } @@ -457,7 +457,7 @@ impl GeomCell { let mut tmp_lat = unsafe{MatrixFull::from_vec_unchecked([3,3],tmp_vec)}; if let GeomUnit::Angstrom = unit { // To store the lattice vector in "Bohr" according to the convention of quantum chemistry. - tmp_lat.self_multiple(ANG.powf(-1.0)); + tmp_lat.self_multiple(BOHR.powf(-1.0)); }; Ok(tmp_lat) //if frac_bool { @@ -534,7 +534,7 @@ impl GeomCell { let mut tmp_pos_tensor = MatrixFull::from_vec(tmp_size, tmp_pos).unwrap(); if let GeomUnit::Angstrom = unit { // To store the geometry position in "Bohr" according to the convention of quantum chemistry. - tmp_pos_tensor.self_multiple(ANG.powf(-1.0)); + tmp_pos_tensor.self_multiple(BOHR.powf(-1.0)); }; Ok((tmp_ele, tmp_fix, tmp_pos_tensor, tmp_nfree)) } @@ -606,7 +606,7 @@ impl GeomCell { let mut tmp_pos_tensor = MatrixFull::from_vec(tmp_size, tmp_pos).unwrap(); if let GeomUnit::Angstrom = unit { // To store the geometry position in "Bohr" according to the convention of quantum chemistry. - tmp_pos_tensor.self_multiple(ANG.powf(-1.0)); + tmp_pos_tensor.self_multiple(BOHR.powf(-1.0)); }; Ok((tmp_ele, tmp_fix, tmp_pos_tensor, tmp_nfree)) @@ -700,7 +700,7 @@ impl GeomCell { let mut tmp_pos_tensor = MatrixFull::from_vec(tmp_size, tmp_bs_pos).unwrap(); if let GeomUnit::Angstrom = unit { // To store the geometry position in "Bohr" according to the convention of quantum chemistry. - tmp_pos_tensor.self_multiple(ANG.powf(-1.0)); + tmp_pos_tensor.self_multiple(BOHR.powf(-1.0)); }; Some((tmp_bs_ele, tmp_pos_tensor)) }; @@ -712,7 +712,7 @@ impl GeomCell { let mut tmp_pos_tensor = MatrixFull::from_vec(tmp_size, tmp_pc_pos).unwrap(); if let GeomUnit::Angstrom = unit { // To store the geometry position in "Bohr" according to the convention of quantum chemistry. - tmp_pos_tensor.self_multiple(ANG.powf(-1.0)); + tmp_pos_tensor.self_multiple(BOHR.powf(-1.0)); }; Some((tmp_pc_chg, tmp_pos_tensor)) }; @@ -724,7 +724,7 @@ impl GeomCell { let mut tmp_pos_tensor = MatrixFull::from_vec(tmp_size, tmp_ep_pos).unwrap(); if let GeomUnit::Angstrom = unit { // To store the geometry position in "Bohr" according to the convention of quantum chemistry. - tmp_pos_tensor.self_multiple(ANG.powf(-1.0)); + tmp_pos_tensor.self_multiple(BOHR.powf(-1.0)); }; Some((tmp_ep_pth, tmp_pos_tensor)) }; @@ -738,7 +738,7 @@ impl GeomCell { } pub fn to_xyz(&self, filename: String) { - let ang = crate::constants::ANG; + let ang = crate::constants::BOHR; let mut input = fs::File::create(&filename).unwrap(); write!(input, "{}\n\n", self.elem.len()); self.position.iter_columns_full().zip(self.elem.iter()).for_each(|(pos, elem)| { @@ -747,7 +747,7 @@ impl GeomCell { } pub fn formated_geometry(&self) -> String { - let ang = crate::constants::ANG; + let ang = crate::constants::BOHR; let mut input = String::new(); //write!(input, "{}\n\n", self.elem.len()); self.position.iter_columns_full().zip(self.elem.iter()).for_each(|(pos, elem)| { @@ -808,7 +808,7 @@ impl GeomCell { let rj = self.position.iter_column(j); let mut dd = ri.zip(rj) .fold(0.0,|acc,(ri,rj)| acc + (ri-rj).powf(2.0)).sqrt(); - dd *= crate::constants::ANG.powf(-1.0); + dd *= crate::constants::BOHR.powf(-1.0); dd } diff --git a/src/grad/mod.rs b/src/grad/mod.rs index d341dc17b4..787c25ffd7 100644 --- a/src/grad/mod.rs +++ b/src/grad/mod.rs @@ -8,7 +8,7 @@ pub mod traits; use std::io::{self, Write}; use crate::main_driver::{collect_total_energy, performance_essential_calculations}; -use crate::{constants::{ANG, EV}, scf_io::{initialize_scf, scf_without_build, SCF}, utilities}; +use crate::{constants::{BOHR, EV}, scf_io::{initialize_scf, scf_without_build, SCF}, utilities}; use crate::mpi_io::{MPIData, MPIOperator}; use tensors::MatrixFull; @@ -132,7 +132,7 @@ pub fn formated_force(force: &MatrixFull, elem: &Vec) -> String { pub fn formated_force_ev(force: &MatrixFull, elem: &Vec) -> String { let mut output = String::new(); force.iter_columns_full().zip(elem.iter()).for_each(|(force, elem)| { - output = format!("{}{:3}{:16.8}{:16.8}{:16.8}\n", output, elem, force[0]*EV/ANG,force[1]*EV/ANG,force[2]*EV/ANG); + output = format!("{}{:3}{:16.8}{:16.8}{:16.8}\n", output, elem, force[0]*EV/BOHR,force[1]*EV/BOHR,force[2]*EV/BOHR); }); output diff --git a/src/lib_rint/mod.rs b/src/lib_rint/mod.rs index a3394cc8f8..79358bcc61 100644 --- a/src/lib_rint/mod.rs +++ b/src/lib_rint/mod.rs @@ -11728,7 +11728,7 @@ position = [ scf_without_build(&mut scf_data, &None); let cov_a = fragment_dipole_fluctuation_tensor(&scf_data, &[0]); let cov_b = fragment_dipole_fluctuation_tensor(&scf_data, &[1]); - let distance_bohr = distance_ang / crate::constants::ANG; + let distance_bohr = distance_ang / crate::constants::BOHR; let x = fragment_connected_dipole_x(cov_a, cov_b, distance_bohr); println!( "occ-closure dipole X_disp(He2): R={distance_ang:.3} Ang ({distance_bohr:.6} bohr) X={x:.16e} logR={:.8} log|X|={:.8}", diff --git a/src/main_driver.rs b/src/main_driver.rs index 3f4b4edb29..8904a1de46 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -13,7 +13,7 @@ use num_traits::Pow; use pyo3::prelude::*; //use autocxx::prelude::*; use crate::ctrl_io::JobType; -use crate::constants::{ANG, AU2DEBYE}; +use crate::constants::{BOHR, AU2DEBYE}; use crate::scf_io::{scf_without_build, SCFType, SCF}; use tensors::{MathMatrix, MatrixFull}; use tensors::matrix_blas_lapack::_dsyevd; @@ -180,7 +180,7 @@ pub fn main_driver() -> anyhow::Result<()> { if scf_data.mol.ctrl.print_level>0 { println!("Geometry optimization invoked"); } - let displace = scf_data.mol.ctrl.nforce_displacement/ANG; + let displace = scf_data.mol.ctrl.nforce_displacement/BOHR; let mut position = scf_data.mol.geom.position.iter().map(|x| *x).collect::>(); lbfgs().minimize( @@ -637,7 +637,7 @@ fn eval_force(scf_data: &mut SCF, time_mark: &mut utilities::TimeRecords, mpi_op if scf_data.mol.ctrl.print_level > 1 { println!("Gradient evaluation using numerical differentiation"); } - let displace = scf_data.mol.ctrl.nforce_displacement / ANG; + let displace = scf_data.mol.ctrl.nforce_displacement / BOHR; let (energy, nforce) = numerical_force(&scf_data, displace, &mpi_operator); println!("------ Output gradient [a.u.] ------"); println!("{}", formated_force(&nforce, &scf_data.mol.geom.elem)); @@ -755,7 +755,7 @@ fn eval_normal_modes( let num_atoms = scf_data.mol.geom.nfree; let dim = num_atoms * 3; let displace_ang = scf_data.mol.ctrl.nhessian_displacement; - let displace = displace_ang / ANG; // convert Angstrom to Bohr + let displace = displace_ang / BOHR; // convert Angstrom to Bohr if scf_data.mol.ctrl.print_level > 0 { println!(""); diff --git a/src/post_scf_analysis/mod.rs b/src/post_scf_analysis/mod.rs index 360109fbbe..7f9753a08a 100644 --- a/src/post_scf_analysis/mod.rs +++ b/src/post_scf_analysis/mod.rs @@ -10,7 +10,7 @@ pub mod spin_correction; use rest_libcint::prelude::rest_libcint_wrapper::int1e_r; use tensors::{MathMatrix, RIFull}; -use crate::constants::{ANG, AU2DEBYE}; +use crate::constants::{BOHR, AU2DEBYE}; use crate::grad::{formated_force, formated_force_ev, numerical_force}; use crate::mpi_io::MPIOperator; use crate::ri_pt2::sbge2::{close_shell_sbge2_rayon, open_shell_sbge2_rayon}; @@ -130,7 +130,7 @@ pub fn post_scf_output(scf_data: &SCF, mpi_operator: &Option) { } } else if output_type.eq("num_force") { let displace = match scf_data.mol.geom.unit { - crate::geom_io::GeomUnit::Angstrom => scf_data.mol.ctrl.nforce_displacement/ANG, + crate::geom_io::GeomUnit::Angstrom => scf_data.mol.ctrl.nforce_displacement/BOHR, crate::geom_io::GeomUnit::Bohr => scf_data.mol.ctrl.nforce_displacement, }; let (energy, num_force) = numerical_force(scf_data, displace, mpi_operator); diff --git a/src/post_scf_analysis/rrs_pbc.rs b/src/post_scf_analysis/rrs_pbc.rs index 83c8655038..d4152fe647 100644 --- a/src/post_scf_analysis/rrs_pbc.rs +++ b/src/post_scf_analysis/rrs_pbc.rs @@ -7,7 +7,7 @@ use std::f64::consts::PI; use rayon::prelude::*; use approx::abs_diff_eq; use std::collections::HashMap; -use crate::constants::ANG; +use crate::constants::BOHR; use crate::geom_io::{GeomUnit,get_charge}; use crate::molecule_io::Molecule; @@ -126,7 +126,7 @@ pub fn rrs_pbc_match(mol: &Molecule) -> (HashMap<(i32,i32,i32),Vec>, Vec< let vec_rrs_pbc_vec = rrs_pbc_vec.data; let mut tot_pbc_vec = vec![vec![0.0;3];3]; //[3,3] let factor = match unit { - GeomUnit::Angstrom => ANG, + GeomUnit::Angstrom => BOHR, GeomUnit::Bohr => 1.0, }; for i in 0..rrs_pbc_dim { @@ -194,7 +194,7 @@ pub fn rrs_pbc_new(scf: &SCF,index_map: HashMap<(i32,i32,i32),Vec>) -> (V let rrs_pbc_dim = scf.mol.geom.pbc_dim; let unit = scf.mol.geom.unit.clone(); let factor = match unit { - GeomUnit::Angstrom => ANG, + GeomUnit::Angstrom => BOHR, GeomUnit::Bohr => 1.0, }; let mut vec_rrs_pbc_vec = rrs_pbc_vec.data.clone(); diff --git a/src/ri_bse/dipoles.rs b/src/ri_bse/dipoles.rs index 408eadbfb9..f18e513121 100644 --- a/src/ri_bse/dipoles.rs +++ b/src/ri_bse/dipoles.rs @@ -3,7 +3,7 @@ use crate::mpi_io::MPIOperator; use std::path::Path; use rest_libcint::prelude::rest_libcint_wrapper::int1e_r; use tensors::{MathMatrix, MatrixFull, RIFull,MatrixFullSlice}; -use crate::constants::{ANG, AU2DEBYE, SPECIES_INFO}; +use crate::constants::{AU2DEBYE, SPECIES_INFO}; use itertools::Itertools; use crate::ri_gw::get_occupation_parameters; use rest_tensors::matrix::matrix_blas_lapack::{_dgemv}; diff --git a/src/ri_bse/response.rs b/src/ri_bse/response.rs index 758b150743..7b7afbb327 100644 --- a/src/ri_bse/response.rs +++ b/src/ri_bse/response.rs @@ -3,7 +3,7 @@ use crate::mpi_io::MPIOperator; use std::path::Path; use rest_libcint::rest_libcint_wrapper::int1e_r; use tensors::{MathMatrix, MatrixFull, RIFull,MatrixFullSlice}; -use crate::constants::{ANG, AU2DEBYE, SPECIES_INFO}; +use crate::constants::{AU2DEBYE, SPECIES_INFO}; use itertools::Itertools; use crate::ri_gw::get_occupation_parameters; use std::f64::consts::SQRT_2; diff --git a/src/solvent/smd_cds.rs b/src/solvent/smd_cds.rs index c2d60443ce..881356d552 100644 --- a/src/solvent/smd_cds.rs +++ b/src/solvent/smd_cds.rs @@ -62,7 +62,7 @@ //! //! ## Unit Conversions //! -//! - `TO_ANGS = 0.529177` Bohr → Å +//! - `TO_ANGS = crate::constants::BOHR` Bohr → Å //! - `TO_KCAL = 627.509` Hartree → kcal/mol //! - Internal computation in Å and kcal/mol; public API input/output in Bohr and Hartree. //! @@ -1880,7 +1880,7 @@ fn cds_eg( coords: &[[f64; 3]], atomic_numbers: &[usize], sigma: &[f64; 151], hsigma: &[f64; 151], rad: &[f64], ) -> (f64, f64, Vec<[f64; 3]>) { - const TO_ANGS: f64 = 0.52917724924; // Bohr → Å + const TO_ANGS: f64 = crate::constants::BOHR; // Bohr → Å const TO_KCAL: f64 = 627.509451; // Hartree → kcal/mol // ---- Build rlio (pairwise distances, Å) and urlio (unit vectors) ---- -- Gitee From ca4c9c323353c9012f2c8862c8e9ebd25b7ed36b Mon Sep 17 00:00:00 2001 From: Shirong_Wang Date: Wed, 8 Jul 2026 20:08:17 +0800 Subject: [PATCH 22/39] unify all phys cons --- src/constants/element.rs | 1 + src/constants/mod.rs | 39 +++++++++++++++++--------------- src/hessian/rhf.rs | 2 +- src/post_scf_analysis/rrs_pbc.rs | 6 ++--- src/ri_bse/dynamicbse.rs | 4 ++-- src/ri_bse/nonlinbse.rs | 6 ++--- src/ri_gw/mod.rs | 5 ++-- src/ri_tddft/response.rs | 4 ++-- src/solvent/smd_cds.rs | 8 +++---- src/x2c/mod.rs | 2 +- 10 files changed, 40 insertions(+), 37 deletions(-) diff --git a/src/constants/element.rs b/src/constants/element.rs index 40185e0283..227874225d 100644 --- a/src/constants/element.rs +++ b/src/constants/element.rs @@ -1,4 +1,5 @@ use lazy_static::lazy_static; +use std::collections::HashMap; pub const SPECIES_NAME: [&str; 118] = ["H", "He", "Li","Be","B", "C","N", "O", "F", "Ne", diff --git a/src/constants/mod.rs b/src/constants/mod.rs index b5d61f973a..73b6a57653 100644 --- a/src/constants/mod.rs +++ b/src/constants/mod.rs @@ -2,10 +2,11 @@ pub mod vsap; pub mod c2s; pub mod solvent; mod cartesian_gto; +pub mod element; -use std::collections::HashMap; +// use std::collections::HashMap; -use lazy_static::lazy_static; +// use lazy_static::lazy_static; pub use crate::constants::vsap::*; pub use crate::constants::c2s::*; @@ -77,28 +78,30 @@ pub const NUC_FRAC_CHARGE: i32 = 3; pub const ENV_PRT_START: usize = 20; // math, physics -pub const CM: f64 = 8065.541; -pub const EV: f64 = 27.2113845; -pub const FQ: f64 = 1822.888; -pub const E: f64 = std::f64::consts::E; -pub const PI: f64 = std::f64::consts::PI; +// NOTE: these constants come from several different CODATA releases (see per-line sources). +pub const EV: f64 = 27.2113845; // Hartree energy in eV, CODATA 2002 +pub const HARTREE2KCAL: f64 = 627.509451; // Hartree -> kcal/mol; CODATA 2002 (Hartree energy in J / (kcal * Avogadro)) +pub const HARTREE2WAVENUMBER: f64 = 219474.63; // Hartree -> cm^-1 (hartree-inverse meter relationship / 100); CODATA (~2.194746313705e7 m^-1) +pub const FQ: f64 = 1822.888; // u/m_e ratio (= 1 / electron mass in u); year-insensitive (stable across CODATA releases) +pub const E: f64 = std::f64::consts::E; // math constant (std); year-insensitive +pub const PI: f64 = std::f64::consts::PI; // math constant (std); year-insensitive // -pub const LIGHT_SPEED: f64 = 137.03599967994; // http://physics.nist.gov/cgi-bin/cuu/Value?alph +pub const LIGHT_SPEED: f64 = 137.03599967994; // inverse fine-structure constant, CODATA 2006; http://physics.nist.gov/cgi-bin/cuu/Value?alph // BOHR = .529 177 210 92(17) e-10m // http://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0 // source: CODATA 2010 (https://physics.nist.gov/cuu/Constants/ArchiveASCII/allascii_2010.txt) pub const BOHR: f64 = 0.52917721092; // Angstroms pub const BOHR_SI: f64 = BOHR * 1e-10; -pub const G_ELECTRON: f64 = 2.00231930436182; // http://physics.nist.gov/cgi-bin/cuu/Value?gem -pub const E_MASS: f64 = 9.10938356e-31; // kg https://physics.nist.gov/cgi-bin/cuu/Value?me -pub const AVOGADRO: f64 = 6.022140857e23; // https://physics.nist.gov/cgi-bin/cuu/Value?na -pub const PLANCK: f64 = 6.626070040e-34; // J*s http://physics.nist.gov/cgi-bin/cuu/Value?h -pub const BOLTZMANN: f64 = 1.380649e-23; // J/K https://physics.nist.gov/cgi-bin/cuu/Value?k -pub const CLIGHT_CMS: f64 = 2.99792458e10; // speed of light, cm/s -pub const R_GAS: f64 = BOLTZMANN * AVOGADRO; // J/(mol*K) ideal gas constant -pub const E_CHARGE: f64 = 1.6021766208e-19; -pub const DEBYE:f64 = 3.335641e-30; // C*m = 1e-18/LIGHT_SPEED_SI https://cccbdb.nist.gov/debye.asp -pub const AU2DEBYE:f64 = E_CHARGE * BOHR*1e-10 / DEBYE; // 2.541746 +pub const G_ELECTRON: f64 = 2.00231930436182; // CODATA 2014; http://physics.nist.gov/cgi-bin/cuu/Value?gem +pub const E_MASS: f64 = 9.10938356e-31; // kg, CODATA 2014; https://physics.nist.gov/cgi-bin/cuu/Value?me +pub const AVOGADRO: f64 = 6.022140857e23; // CODATA 2014; https://physics.nist.gov/cgi-bin/cuu/Value?na +pub const PLANCK: f64 = 6.626070040e-34; // J*s, CODATA 2014; http://physics.nist.gov/cgi-bin/cuu/Value?h +pub const BOLTZMANN: f64 = 1.380649e-23; // J/K, CODATA 2018 (exact, SI 2019); https://physics.nist.gov/cgi-bin/cuu/Value?k +pub const CLIGHT_CMS: f64 = 2.99792458e10; // speed of light, cm/s; exact by SI definition (year-insensitive) +pub const R_GAS: f64 = BOLTZMANN * AVOGADRO; // J/(mol*K) ideal gas constant; derived +pub const E_CHARGE: f64 = 1.6021766208e-19; // C, CODATA 2014 +pub const DEBYE:f64 = 3.335641e-30; // C*m = 1e-18/LIGHT_SPEED_SI; defined via c, year-insensitive; https://cccbdb.nist.gov/debye.asp +pub const AU2DEBYE:f64 = E_CHARGE * BOHR*1e-10 / DEBYE; // 2.541746; derived pub const MPI_CHUNK:usize = 134217728; // around 1 GB diff --git a/src/hessian/rhf.rs b/src/hessian/rhf.rs index b2cb7f9832..a04d23cee1 100644 --- a/src/hessian/rhf.rs +++ b/src/hessian/rhf.rs @@ -4247,7 +4247,7 @@ pub fn compute_frequencies_from_hessian( // Convert eigenvalues to frequencies in cm⁻¹ // ω² = λ (in a.u.), ω = √λ (a.u.), ν = ω/(2π) (a.u.) // 1 Hartree = 219474.63 cm⁻¹ - let conv = 219474.63 / 1822.8885f64.sqrt(); // = 5140.49 + let conv = crate::constants::HARTREE2WAVENUMBER / 1822.8885f64.sqrt(); // = 5140.49 let mut freqs = vec![0.0; n3]; for i in 0..n3 { let lambda = eigvals[i]; diff --git a/src/post_scf_analysis/rrs_pbc.rs b/src/post_scf_analysis/rrs_pbc.rs index d4152fe647..1b79252076 100644 --- a/src/post_scf_analysis/rrs_pbc.rs +++ b/src/post_scf_analysis/rrs_pbc.rs @@ -7,7 +7,7 @@ use std::f64::consts::PI; use rayon::prelude::*; use approx::abs_diff_eq; use std::collections::HashMap; -use crate::constants::BOHR; +use crate::constants::{BOHR, EV}; use crate::geom_io::{GeomUnit,get_charge}; use crate::molecule_io::Molecule; @@ -52,12 +52,12 @@ pub fn rrs_pbc_output(scf_data: &SCF, unit_cell_elem: Vec, index_map: Ha } let mut homo_lumo_gap = vec![0.0;tot_k_points]; homo_lumo_gap.iter_mut().zip(pbc_result.iter()).for_each(|(h,p)| { - *h = (p[(cell_charge/2.0) as usize] - p[(cell_charge/2.0) as usize - 1])*27.2113863; + *h = (p[(cell_charge/2.0) as usize] - p[(cell_charge/2.0) as usize - 1])*EV; }); let min_hlg = homo_lumo_gap.iter().filter(|&&x| !x.is_nan()).min_by(|a, b| a.partial_cmp(b).unwrap()); println!("HOCO-LUCO gap: {:?} eV",min_hlg.unwrap()); let elec = scf_data.mol.num_elec[0]; - let hlg = (scf_data.eigenvalues[0][(elec/2.0) as usize] - scf_data.eigenvalues[0][(elec/2.0) as usize - 1])*27.2113863; + let hlg = (scf_data.eigenvalues[0][(elec/2.0) as usize] - scf_data.eigenvalues[0][(elec/2.0) as usize - 1])*EV; println!("HOMO-LUMO gap: {:?} eV",hlg); } diff --git a/src/ri_bse/dynamicbse.rs b/src/ri_bse/dynamicbse.rs index 6b32cc4173..b50d377397 100644 --- a/src/ri_bse/dynamicbse.rs +++ b/src/ri_bse/dynamicbse.rs @@ -833,7 +833,7 @@ pub fn dynamic_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { println!(" # Excitation energy (eV) ‖T(λ)x‖"); println!(" ─── ───────────────────── ──────────"); for k in 0..result.n_found { - let lam_ev = result.eigenvalues[k] * 27.2114; + let lam_ev = result.eigenvalues[k] * crate::constants::EV; println!(" {:>3} {:>12.6} eV {:>9.2e}", k, lam_ev, result.residuals[k]); } @@ -842,7 +842,7 @@ pub fn dynamic_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { for k in 0..result.n_found { let xv: Vec = (0..n).map(|i| result.eigenvectors[[i, k]]).collect(); println!("\n Excitation #{}: λ = {:.6} Ha = {:.6} eV", - k, result.eigenvalues[k], result.eigenvalues[k] * 27.2114); + k, result.eigenvalues[k], result.eigenvalues[k] * crate::constants::EV); super::leading_components(&xv, occ_size, vir_size); } } diff --git a/src/ri_bse/nonlinbse.rs b/src/ri_bse/nonlinbse.rs index 89af13a0ab..f8d4f740c7 100644 --- a/src/ri_bse/nonlinbse.rs +++ b/src/ri_bse/nonlinbse.rs @@ -1218,7 +1218,7 @@ pub fn nlfeast_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { println!(" # Excitation energy (eV) ‖T(λ)x‖"); println!(" ─── ───────────────────── ──────────"); for k in 0..result.n_found { - let lam_ev = result.eigenvalues[k] * 27.2114; + let lam_ev = result.eigenvalues[k] * crate::constants::EV; println!(" {:>3} {:>12.6} eV {:>9.2e}", k, lam_ev, result.residuals[k]); } @@ -1227,7 +1227,7 @@ pub fn nlfeast_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { for k in 0..result.n_found { let xv: Vec = (0..n).map(|i| result.eigenvectors[[i, k]]).collect(); println!("\n Excitation #{}: λ = {:.6} Ha = {:.6} eV", - k, result.eigenvalues[k], result.eigenvalues[k] * 27.2114); + k, result.eigenvalues[k], result.eigenvalues[k] * crate::constants::EV); super::leading_components(&xv, occ_size, vir_size); } } @@ -1659,7 +1659,7 @@ pub fn nlfeast_dynamical_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { println!(" # Excitation energy (eV) ‖T(λ)x‖"); println!(" ─── ───────────────────── ──────────"); for k in 0..result.n_found { - let lam_ev = result.eigenvalues[k] * 27.2114; + let lam_ev = result.eigenvalues[k] * crate::constants::EV; println!(" {:>3} {:>12.6} eV {:>9.2e}", k, lam_ev, result.residuals[k]); } } diff --git a/src/ri_gw/mod.rs b/src/ri_gw/mod.rs index f5e010ed8f..6f8f3dcb83 100644 --- a/src/ri_gw/mod.rs +++ b/src/ri_gw/mod.rs @@ -141,7 +141,6 @@ pub fn gw_calculations(scf_data:&mut SCF,num_freq:usize,vxc_nn:&Vec,cancel_ let v_matrix=v_matrix(&scf_data,&ri_mat); let (start_mo,num_state,occ_size,vir_size,homo,lumo)=get_occupation_parameters(scf_data,'Y'); let (start_mo,num_state_cutoff,occ_size,vir_size_cutoff,homo,lumo)=get_occupation_parameters(scf_data,'N'); - let hatree=27.2113863; let spin_channel=scf_data.mol.ctrl.spin_channel; let mut quasiparticle_energies_g:Vec=Vec::new(); let mut quasiparticle_energies_w:Vec=Vec::new(); @@ -1483,7 +1482,7 @@ pub fn generate_real_axis_vchiv( let grid_label = if grid_type == 1 { "quadratic (power-law)" } else { "linear" }; println!("Low-rank contour: Maximum real frequency needed for v*chi*v = {:.6} Ha = {:.6} eV", - de_max, de_max * 27.2113863); + de_max, de_max * crate::constants::EV); println!("Low-rank contour: Computing sqrt(v)*chi*sqrt(v) at {} real-axis grid points ({})", nomega_chi_real, grid_label); if print_level > 2 { @@ -1508,7 +1507,7 @@ pub fn generate_real_axis_vchiv( }; if print_level > 2 { println!("[DEBUG] Real-axis freq {} / {} : omega = {:.10} Ha = {:.6} eV (t={:.6})", - i_omega + 1, nomega_chi_real, omega_real, omega_real * 27.2113863, + i_omega + 1, nomega_chi_real, omega_real, omega_real * crate::constants::EV, (i_omega as f64) / ((nomega_chi_real - 1) as f64)); } let lr = low_rank_vchi_vsqrt( diff --git a/src/ri_tddft/response.rs b/src/ri_tddft/response.rs index 01ab240484..c14e9cbb9a 100644 --- a/src/ri_tddft/response.rs +++ b/src/ri_tddft/response.rs @@ -222,7 +222,7 @@ fn print_response_tddft_results( let pol_im: f64 = density_imag.iter().zip(mu_z_vec.iter()).map(|(d, m)| d * m).sum(); println!(" --- Response TDDFT Results ---"); - println!(" Frequency ω = {:.8} Ha ({:.4} eV)", omega, omega * 27.2114); + println!(" Frequency ω = {:.8} Ha ({:.4} eV)", omega, omega * crate::constants::EV); println!(" Lifetime γ = {:.8} Ha", gamma); println!(" Re[α_zz(ω)] = {:.12e} a.u.", pol_re); println!(" Im[α_zz(ω)] = {:.12e} a.u.", pol_im); @@ -685,7 +685,7 @@ pub fn response_tddft(scf: &mut SCF) -> Result<(), String> { println!(" Method: Response (frequency-domain)"); println!(" Spin: {}", if xlet == 'S' { "Singlet" } else if xlet == 'T' { "Triplet" } else { "Generic" }); println!(" occ_size={}, vir_size={}, dim={}", occ_size, vir_size, dim); - println!(" ω = {:.8} Ha ({:.4} eV)", omega, omega * 27.2114); + println!(" ω = {:.8} Ha ({:.4} eV)", omega, omega * crate::constants::EV); println!(" γ = {:.8} Ha", gamma); // Prepare fxc data diff --git a/src/solvent/smd_cds.rs b/src/solvent/smd_cds.rs index 881356d552..923584b5a2 100644 --- a/src/solvent/smd_cds.rs +++ b/src/solvent/smd_cds.rs @@ -63,7 +63,7 @@ //! ## Unit Conversions //! //! - `TO_ANGS = crate::constants::BOHR` Bohr → Å -//! - `TO_KCAL = 627.509` Hartree → kcal/mol +//! - `TO_KCAL = crate::constants::HARTREE2KCAL` Hartree → kcal/mol //! - Internal computation in Å and kcal/mol; public API input/output in Bohr and Hartree. //! //! ## References @@ -1881,7 +1881,7 @@ fn cds_eg( sigma: &[f64; 151], hsigma: &[f64; 151], rad: &[f64], ) -> (f64, f64, Vec<[f64; 3]>) { const TO_ANGS: f64 = crate::constants::BOHR; // Bohr → Å - const TO_KCAL: f64 = 627.509451; // Hartree → kcal/mol + const TO_KCAL: f64 = crate::constants::HARTREE2KCAL; // Hartree → kcal/mol // ---- Build rlio (pairwise distances, Å) and urlio (unit vectors) ---- let ncot = nat * (nat + 1) / 2; @@ -2135,12 +2135,12 @@ pub fn compute_cds( cds_print_scalar("gcds_kcal (CDS energy, kcal/mol)", gcds_kcal); cds_print_scalar("tarea (total SASA, A^2)", tarea); cds_print_grad("dcds (CDS gradient, Hartree/Bohr)", &dcds); - const TO_KCAL: f64 = 627.509451; + const TO_KCAL: f64 = crate::constants::HARTREE2KCAL; cds_print_scalar("gcds_hartree (CDS energy, Hartree)", gcds_kcal / TO_KCAL); } // kcal/mol → Hartree (gradient already in Hartree/Bohr from cds_eg) - const TO_KCAL: f64 = 627.509451; + const TO_KCAL: f64 = crate::constants::HARTREE2KCAL; let gcds = gcds_kcal / TO_KCAL; (gcds, tarea, dcds) diff --git a/src/x2c/mod.rs b/src/x2c/mod.rs index 7f54f3ddb1..95e07c408f 100644 --- a/src/x2c/mod.rs +++ b/src/x2c/mod.rs @@ -168,7 +168,7 @@ impl Molecule { pub fn generate_sfx2c_hamiltonian(&self) -> MatrixUpper { - let light_speed: f64 = 137.03599967994; + let light_speed: f64 = crate::constants::LIGHT_SPEED; let (xmol, contr_coeff) = self.decontract_basis(); let n_contracted = self.num_basis; -- Gitee From d5dcfcab27e5651be3257b0e735227c76ea4a2f3 Mon Sep 17 00:00:00 2001 From: Shirong_Wang Date: Wed, 8 Jul 2026 20:46:34 +0800 Subject: [PATCH 23/39] proper import --- src/dft/gen_grids/atom.rs | 3 ++- src/fileop/chkfile.rs | 3 ++- src/geom_io/mod.rs | 6 +++--- src/hessian/rhf.rs | 3 ++- src/lib_rint/mod.rs | 4 ++-- src/main_driver.rs | 1 - src/ri_bse/dynamicbse.rs | 5 +++-- src/ri_bse/nonlinbse.rs | 7 ++++--- src/ri_gw/mod.rs | 6 +++--- src/ri_tddft/response.rs | 5 +++-- src/solvent/smd_cds.rs | 13 +++++++------ src/x2c/mod.rs | 4 ++-- 12 files changed, 33 insertions(+), 27 deletions(-) diff --git a/src/dft/gen_grids/atom.rs b/src/dft/gen_grids/atom.rs index 97570c81f9..ac294f9f04 100644 --- a/src/dft/gen_grids/atom.rs +++ b/src/dft/gen_grids/atom.rs @@ -11,6 +11,7 @@ use super::lebedev; use super::prune::*; use super::radial; use super::parameters::LEBEDEV_NGRID; +use crate::constants::BOHR; /// Generate atomic grids based on given arguments.
/// Download basis set information from www.basissetexchange.org to determine alpha_max and alpha_min. @@ -121,7 +122,7 @@ pub fn atom_grid( //println!("radial num = {}", default_radial_num(proton_charges[center_index] as usize)); // factors match DIRAC code //println!("rs = {:?}, w = {:?}", rs, weights_radial); - let rb = bragg::get_bragg_angstrom(proton_charges[center_index]) / (5.0 * crate::constants::BOHR); + let rb = bragg::get_bragg_angstrom(proton_charges[center_index]) / (5.0 * BOHR); let mut coordinates = Vec::new(); let mut weights = Vec::new(); diff --git a/src/fileop/chkfile.rs b/src/fileop/chkfile.rs index 77d1c2c40e..d06f274d8b 100644 --- a/src/fileop/chkfile.rs +++ b/src/fileop/chkfile.rs @@ -7,6 +7,7 @@ use crate::scf_io::{SCF, SCFType}; use crate::geom_io::get_mass_charge; use crate::basis_io::Basis4Elem; use rest_libcint::CintType; +use crate::constants::BOHR; pub fn write_scf_attribute(group: &hdf5::Group, dataset_name: &str, value: &[T]) where @@ -183,7 +184,7 @@ pub fn save_overlap(scf_data: &SCF) { } pub fn save_geometry(scf_data: &SCF) { - let ang = crate::constants::BOHR; + let ang = BOHR; let chkfile= &scf_data.mol.ctrl.chkfile; let path = Path::new(chkfile); let file = if path.exists() { diff --git a/src/geom_io/mod.rs b/src/geom_io/mod.rs index b2591fff6f..6c280d5da4 100644 --- a/src/geom_io/mod.rs +++ b/src/geom_io/mod.rs @@ -738,7 +738,7 @@ impl GeomCell { } pub fn to_xyz(&self, filename: String) { - let ang = crate::constants::BOHR; + let ang = BOHR; let mut input = fs::File::create(&filename).unwrap(); write!(input, "{}\n\n", self.elem.len()); self.position.iter_columns_full().zip(self.elem.iter()).for_each(|(pos, elem)| { @@ -747,7 +747,7 @@ impl GeomCell { } pub fn formated_geometry(&self) -> String { - let ang = crate::constants::BOHR; + let ang = BOHR; let mut input = String::new(); //write!(input, "{}\n\n", self.elem.len()); self.position.iter_columns_full().zip(self.elem.iter()).for_each(|(pos, elem)| { @@ -808,7 +808,7 @@ impl GeomCell { let rj = self.position.iter_column(j); let mut dd = ri.zip(rj) .fold(0.0,|acc,(ri,rj)| acc + (ri-rj).powf(2.0)).sqrt(); - dd *= crate::constants::BOHR.powf(-1.0); + dd *= BOHR.powf(-1.0); dd } diff --git a/src/hessian/rhf.rs b/src/hessian/rhf.rs index a04d23cee1..c10d3af042 100644 --- a/src/hessian/rhf.rs +++ b/src/hessian/rhf.rs @@ -10,6 +10,7 @@ use tensors::matrix_blas_lapack::_power_rayon_for_symmetric_matrix; use tensors::MatrixFull; use crate::utilities::rstsr_util::*; +use crate::constants::HARTREE2WAVENUMBER; /// Routing for each optimizable G-term in calc_ej_ek(). #[derive(Clone, Copy, PartialEq, Debug)] @@ -4247,7 +4248,7 @@ pub fn compute_frequencies_from_hessian( // Convert eigenvalues to frequencies in cm⁻¹ // ω² = λ (in a.u.), ω = √λ (a.u.), ν = ω/(2π) (a.u.) // 1 Hartree = 219474.63 cm⁻¹ - let conv = crate::constants::HARTREE2WAVENUMBER / 1822.8885f64.sqrt(); // = 5140.49 + let conv = HARTREE2WAVENUMBER / 1822.8885f64.sqrt(); // = 5140.49 let mut freqs = vec![0.0; n3]; for i in 0..n3 { let lambda = eigvals[i]; diff --git a/src/lib_rint/mod.rs b/src/lib_rint/mod.rs index 79358bcc61..2ea4c7f41d 100644 --- a/src/lib_rint/mod.rs +++ b/src/lib_rint/mod.rs @@ -1,6 +1,6 @@ use crate::basis_io::etb::{etb_gen_for_atom_list, get_etb_elem}; use crate::basis_io::Basis4Elem; -use crate::constants::{AUXBAS_THRESHOLD, SQRT_THRESHOLD}; +use crate::constants::{AUXBAS_THRESHOLD, BOHR, SQRT_THRESHOLD}; use crate::geom_io::GeomCell; use crate::molecule_io::Molecule; use crate::scf_io::{ @@ -11728,7 +11728,7 @@ position = [ scf_without_build(&mut scf_data, &None); let cov_a = fragment_dipole_fluctuation_tensor(&scf_data, &[0]); let cov_b = fragment_dipole_fluctuation_tensor(&scf_data, &[1]); - let distance_bohr = distance_ang / crate::constants::BOHR; + let distance_bohr = distance_ang / BOHR; let x = fragment_connected_dipole_x(cov_a, cov_b, distance_bohr); println!( "occ-closure dipole X_disp(He2): R={distance_ang:.3} Ang ({distance_bohr:.6} bohr) X={x:.16e} logR={:.8} log|X|={:.8}", diff --git a/src/main_driver.rs b/src/main_driver.rs index 8904a1de46..b64a868a48 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -954,7 +954,6 @@ mod geometric_pyo3_impl { pyo3::prepare_freethreaded_python(); let elem = scf_data.mol.geom.elem.iter().map(|x| x.as_str()).collect::>(); - const BOHR: f64 = crate::constants::BOHR; let xyz = scf_data.mol.geom.position.data.iter().map(|x| x * BOHR).collect::>(); //let xyz = scf_data.mol.geom.position.iter().map(|x| *x).collect::>(); let xyzs = vec![xyz]; diff --git a/src/ri_bse/dynamicbse.rs b/src/ri_bse/dynamicbse.rs index b50d377397..4ee720608f 100644 --- a/src/ri_bse/dynamicbse.rs +++ b/src/ri_bse/dynamicbse.rs @@ -17,6 +17,7 @@ use num::Complex; use rand::Rng; use std::time::Instant; use tensors::{MathMatrix, MatrixFull}; +use crate::constants::EV; use rest_tensors::matrix::matrix_blas_lapack::{ _dgemm_full, _dgemm_scaled, _dgeev, }; @@ -833,7 +834,7 @@ pub fn dynamic_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { println!(" # Excitation energy (eV) ‖T(λ)x‖"); println!(" ─── ───────────────────── ──────────"); for k in 0..result.n_found { - let lam_ev = result.eigenvalues[k] * crate::constants::EV; + let lam_ev = result.eigenvalues[k] * EV; println!(" {:>3} {:>12.6} eV {:>9.2e}", k, lam_ev, result.residuals[k]); } @@ -842,7 +843,7 @@ pub fn dynamic_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { for k in 0..result.n_found { let xv: Vec = (0..n).map(|i| result.eigenvectors[[i, k]]).collect(); println!("\n Excitation #{}: λ = {:.6} Ha = {:.6} eV", - k, result.eigenvalues[k], result.eigenvalues[k] * crate::constants::EV); + k, result.eigenvalues[k], result.eigenvalues[k] * EV); super::leading_components(&xv, occ_size, vir_size); } } diff --git a/src/ri_bse/nonlinbse.rs b/src/ri_bse/nonlinbse.rs index f8d4f740c7..e512223492 100644 --- a/src/ri_bse/nonlinbse.rs +++ b/src/ri_bse/nonlinbse.rs @@ -23,6 +23,7 @@ use rand::Rng; use rayon::prelude::*; use std::time::Instant; use tensors::{MathMatrix, MatrixFull}; +use crate::constants::EV; use rest_tensors::matrix::matrix_blas_lapack::{ _dgemm_full, _dgemm_scaled, _dgeev, }; @@ -1218,7 +1219,7 @@ pub fn nlfeast_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { println!(" # Excitation energy (eV) ‖T(λ)x‖"); println!(" ─── ───────────────────── ──────────"); for k in 0..result.n_found { - let lam_ev = result.eigenvalues[k] * crate::constants::EV; + let lam_ev = result.eigenvalues[k] * EV; println!(" {:>3} {:>12.6} eV {:>9.2e}", k, lam_ev, result.residuals[k]); } @@ -1227,7 +1228,7 @@ pub fn nlfeast_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { for k in 0..result.n_found { let xv: Vec = (0..n).map(|i| result.eigenvectors[[i, k]]).collect(); println!("\n Excitation #{}: λ = {:.6} Ha = {:.6} eV", - k, result.eigenvalues[k], result.eigenvalues[k] * crate::constants::EV); + k, result.eigenvalues[k], result.eigenvalues[k] * EV); super::leading_components(&xv, occ_size, vir_size); } } @@ -1659,7 +1660,7 @@ pub fn nlfeast_dynamical_bse_main(scf_data: &SCF, qp_ctrl: &QuasiParticle) { println!(" # Excitation energy (eV) ‖T(λ)x‖"); println!(" ─── ───────────────────── ──────────"); for k in 0..result.n_found { - let lam_ev = result.eigenvalues[k] * crate::constants::EV; + let lam_ev = result.eigenvalues[k] * EV; println!(" {:>3} {:>12.6} eV {:>9.2e}", k, lam_ev, result.residuals[k]); } } diff --git a/src/ri_gw/mod.rs b/src/ri_gw/mod.rs index 6f8f3dcb83..1c1b947423 100644 --- a/src/ri_gw/mod.rs +++ b/src/ri_gw/mod.rs @@ -1,5 +1,5 @@ //use std::simd::num; -use crate::constants::PI; +use crate::constants::{EV, PI}; use itertools::Itertools; use std::ops::Range; use crate::utilities; @@ -1482,7 +1482,7 @@ pub fn generate_real_axis_vchiv( let grid_label = if grid_type == 1 { "quadratic (power-law)" } else { "linear" }; println!("Low-rank contour: Maximum real frequency needed for v*chi*v = {:.6} Ha = {:.6} eV", - de_max, de_max * crate::constants::EV); + de_max, de_max * EV); println!("Low-rank contour: Computing sqrt(v)*chi*sqrt(v) at {} real-axis grid points ({})", nomega_chi_real, grid_label); if print_level > 2 { @@ -1507,7 +1507,7 @@ pub fn generate_real_axis_vchiv( }; if print_level > 2 { println!("[DEBUG] Real-axis freq {} / {} : omega = {:.10} Ha = {:.6} eV (t={:.6})", - i_omega + 1, nomega_chi_real, omega_real, omega_real * crate::constants::EV, + i_omega + 1, nomega_chi_real, omega_real, omega_real * EV, (i_omega as f64) / ((nomega_chi_real - 1) as f64)); } let lr = low_rank_vchi_vsqrt( diff --git a/src/ri_tddft/response.rs b/src/ri_tddft/response.rs index c14e9cbb9a..c573e40422 100644 --- a/src/ri_tddft/response.rs +++ b/src/ri_tddft/response.rs @@ -20,6 +20,7 @@ use std::fs::OpenOptions; use std::io::Write; use rest_tensors::MatrixFull; +use crate::constants::EV; use rest_tensors::matrix::matrix_blas_lapack::{_dsolve, _dgemm_full}; use crate::ri_bse::response::{ @@ -222,7 +223,7 @@ fn print_response_tddft_results( let pol_im: f64 = density_imag.iter().zip(mu_z_vec.iter()).map(|(d, m)| d * m).sum(); println!(" --- Response TDDFT Results ---"); - println!(" Frequency ω = {:.8} Ha ({:.4} eV)", omega, omega * crate::constants::EV); + println!(" Frequency ω = {:.8} Ha ({:.4} eV)", omega, omega * EV); println!(" Lifetime γ = {:.8} Ha", gamma); println!(" Re[α_zz(ω)] = {:.12e} a.u.", pol_re); println!(" Im[α_zz(ω)] = {:.12e} a.u.", pol_im); @@ -685,7 +686,7 @@ pub fn response_tddft(scf: &mut SCF) -> Result<(), String> { println!(" Method: Response (frequency-domain)"); println!(" Spin: {}", if xlet == 'S' { "Singlet" } else if xlet == 'T' { "Triplet" } else { "Generic" }); println!(" occ_size={}, vir_size={}, dim={}", occ_size, vir_size, dim); - println!(" ω = {:.8} Ha ({:.4} eV)", omega, omega * crate::constants::EV); + println!(" ω = {:.8} Ha ({:.4} eV)", omega, omega * EV); println!(" γ = {:.8} Ha", gamma); // Prepare fxc data diff --git a/src/solvent/smd_cds.rs b/src/solvent/smd_cds.rs index 923584b5a2..aa54da59cd 100644 --- a/src/solvent/smd_cds.rs +++ b/src/solvent/smd_cds.rs @@ -62,8 +62,8 @@ //! //! ## Unit Conversions //! -//! - `TO_ANGS = crate::constants::BOHR` Bohr → Å -//! - `TO_KCAL = crate::constants::HARTREE2KCAL` Hartree → kcal/mol +//! - `TO_ANGS = BOHR` Bohr → Å +//! - `TO_KCAL = HARTREE2KCAL` Hartree → kcal/mol //! - Internal computation in Å and kcal/mol; public API input/output in Bohr and Hartree. //! //! ## References @@ -74,6 +74,7 @@ use std::f64::consts::PI; use tensors::MatrixFull; +use crate::constants::{BOHR, HARTREE2KCAL}; // ============================================================================ // Debug printing helper (controlled by env var REST_CDS_DEBUG=1) @@ -1880,8 +1881,8 @@ fn cds_eg( coords: &[[f64; 3]], atomic_numbers: &[usize], sigma: &[f64; 151], hsigma: &[f64; 151], rad: &[f64], ) -> (f64, f64, Vec<[f64; 3]>) { - const TO_ANGS: f64 = crate::constants::BOHR; // Bohr → Å - const TO_KCAL: f64 = crate::constants::HARTREE2KCAL; // Hartree → kcal/mol + const TO_ANGS: f64 = BOHR; // Bohr → Å + const TO_KCAL: f64 = HARTREE2KCAL; // Hartree → kcal/mol // ---- Build rlio (pairwise distances, Å) and urlio (unit vectors) ---- let ncot = nat * (nat + 1) / 2; @@ -2135,12 +2136,12 @@ pub fn compute_cds( cds_print_scalar("gcds_kcal (CDS energy, kcal/mol)", gcds_kcal); cds_print_scalar("tarea (total SASA, A^2)", tarea); cds_print_grad("dcds (CDS gradient, Hartree/Bohr)", &dcds); - const TO_KCAL: f64 = crate::constants::HARTREE2KCAL; + const TO_KCAL: f64 = HARTREE2KCAL; cds_print_scalar("gcds_hartree (CDS energy, Hartree)", gcds_kcal / TO_KCAL); } // kcal/mol → Hartree (gradient already in Hartree/Bohr from cds_eg) - const TO_KCAL: f64 = crate::constants::HARTREE2KCAL; + const TO_KCAL: f64 = HARTREE2KCAL; let gcds = gcds_kcal / TO_KCAL; (gcds, tarea, dcds) diff --git a/src/x2c/mod.rs b/src/x2c/mod.rs index 95e07c408f..7cd75f5713 100644 --- a/src/x2c/mod.rs +++ b/src/x2c/mod.rs @@ -6,7 +6,7 @@ use rest_libcint::prelude::*; use rest_libcint::CINTR2CDATA; use rest_libcint_wrapper::*; use crate::basis_io::BasInfo; -use crate::constants::ENV_PRT_START; +use crate::constants::{ENV_PRT_START, LIGHT_SPEED}; use crate::molecule_io::{Molecule, get_basis_name}; use serde::{Deserialize, Serialize}; @@ -168,7 +168,7 @@ impl Molecule { pub fn generate_sfx2c_hamiltonian(&self) -> MatrixUpper { - let light_speed: f64 = crate::constants::LIGHT_SPEED; + let light_speed: f64 = LIGHT_SPEED; let (xmol, contr_coeff) = self.decontract_basis(); let n_contracted = self.num_basis; -- Gitee From 887d1187fe2054ebe67c5a567f2363019150d074 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 8 Jul 2026 20:57:16 +0800 Subject: [PATCH 24/39] fix !169 possible mis-delete trait --- src/scf_io/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/scf_io/mod.rs b/src/scf_io/mod.rs index 5b97d7f41f..12d6531290 100644 --- a/src/scf_io/mod.rs +++ b/src/scf_io/mod.rs @@ -25,7 +25,7 @@ use mpi::collective::SystemOperation; use pyo3::{pyclass}; use tensors::matrix_blas_lapack::{_dgemm, _dgemm_full, _dgemv, _dspgvx, _dsymm, _dsyrk, _hamiltonian_fast_solver, _power_rayon_for_symmetric_matrix, _dsyevd}; use tensors::{map_upper_to_full, BasicMatrix, ERIFold4, MathMatrix, MatrixFull, MatrixFullSlice, MatrixUpper, MatrixUpperSlice, RIFull, TensorSliceMut}; -use tensors::{TensorOpt,TensorSlice}; +use tensors::{TensorOpt, TensorSlice, BasicMatUp}; use itertools::{Itertools}; use rayon::prelude::*; use std::collections::HashMap; -- Gitee From a86c1d4a33821acf05fa191d61addce8a2501d9b Mon Sep 17 00:00:00 2001 From: Shirong_Wang Date: Wed, 8 Jul 2026 20:58:43 +0800 Subject: [PATCH 25/39] update FQ --- src/constants/mod.rs | 2 +- src/hessian/rhf.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/constants/mod.rs b/src/constants/mod.rs index 73b6a57653..08be236d13 100644 --- a/src/constants/mod.rs +++ b/src/constants/mod.rs @@ -82,7 +82,7 @@ pub const ENV_PRT_START: usize = 20; pub const EV: f64 = 27.2113845; // Hartree energy in eV, CODATA 2002 pub const HARTREE2KCAL: f64 = 627.509451; // Hartree -> kcal/mol; CODATA 2002 (Hartree energy in J / (kcal * Avogadro)) pub const HARTREE2WAVENUMBER: f64 = 219474.63; // Hartree -> cm^-1 (hartree-inverse meter relationship / 100); CODATA (~2.194746313705e7 m^-1) -pub const FQ: f64 = 1822.888; // u/m_e ratio (= 1 / electron mass in u); year-insensitive (stable across CODATA releases) +pub const FQ: f64 = 1822.8884861920776; // u/m_e ratio (= 1 / electron mass in u 5.48579909070e-4), CODATA 2014 pub const E: f64 = std::f64::consts::E; // math constant (std); year-insensitive pub const PI: f64 = std::f64::consts::PI; // math constant (std); year-insensitive // diff --git a/src/hessian/rhf.rs b/src/hessian/rhf.rs index c10d3af042..965b369dc6 100644 --- a/src/hessian/rhf.rs +++ b/src/hessian/rhf.rs @@ -10,7 +10,7 @@ use tensors::matrix_blas_lapack::_power_rayon_for_symmetric_matrix; use tensors::MatrixFull; use crate::utilities::rstsr_util::*; -use crate::constants::HARTREE2WAVENUMBER; +use crate::constants::{FQ, HARTREE2WAVENUMBER}; /// Routing for each optimizable G-term in calc_ej_ek(). #[derive(Clone, Copy, PartialEq, Debug)] @@ -4248,7 +4248,7 @@ pub fn compute_frequencies_from_hessian( // Convert eigenvalues to frequencies in cm⁻¹ // ω² = λ (in a.u.), ω = √λ (a.u.), ν = ω/(2π) (a.u.) // 1 Hartree = 219474.63 cm⁻¹ - let conv = HARTREE2WAVENUMBER / 1822.8885f64.sqrt(); // = 5140.49 + let conv = HARTREE2WAVENUMBER / FQ.sqrt(); // = 5140.49 let mut freqs = vec![0.0; n3]; for i in 0..n3 { let lambda = eigvals[i]; -- Gitee From 80917e0c63561faf9311a108857c61365b4c9e67 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 8 Jul 2026 20:58:43 +0800 Subject: [PATCH 26/39] fix constant duplicate --- src/constants/mod.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/constants/mod.rs b/src/constants/mod.rs index fa23feb0a1..e3e9997ed8 100644 --- a/src/constants/mod.rs +++ b/src/constants/mod.rs @@ -335,8 +335,6 @@ pub const DEBYE:f64 = 3.335641e-30; // C*m = 1e-18/LIGHT_SPEED_SI htt // CODATA 2022 // https://docs.scipy.org/doc/scipy/reference/constants.html pub const SPEED_OF_LIGHT: f64 = 299792458.0; // m s^-1 -pub const BOLTZMANN: f64 = 1.380649e-23; // J K^-1 -pub const R_GAS: f64 = 8.31446261815324; // J mol^-1 K^-1 // =========== unit conversion =================================== pub const BOHR2ANG: f64 = BOHR; -- Gitee From b496210782f20ae0bc5b11fe021fb34fdaa4fe69 Mon Sep 17 00:00:00 2001 From: Shirong_Wang Date: Wed, 8 Jul 2026 21:12:03 +0800 Subject: [PATCH 27/39] fix test --- src/external_field/num_dipole.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/external_field/num_dipole.rs b/src/external_field/num_dipole.rs index 244f4a83e3..9ebb6e7698 100644 --- a/src/external_field/num_dipole.rs +++ b/src/external_field/num_dipole.rs @@ -71,8 +71,8 @@ mod debug { print_level = 2 job_type = "numerical dipole" xc = "hf" - basis_path = "basis-set-pool/def2-TZVP" - auxbas_path = "basis-set-pool/def2-SVP-JKFIT" + basis_path = "def2-TZVP" + auxbas_path = "def2-universal-JKFIT" guessfile = "none" charge = 0.0 spin = 1.0 @@ -102,8 +102,8 @@ mod debug { [ctrl] print_level = 2 xc = "hf" - basis_path = "basis-set-pool/def2-TZVP" - auxbas_path = "basis-set-pool/def2-SV(P)-JKFIT" + basis_path = "def2-TZVP" + auxbas_path = "def2-SV(P)-JKFIT" guessfile = "none" charge = 0.0 spin = 1.0 -- Gitee From b07f9a9de3d8e2ac112a9dfcbc22e0d066ff1ae7 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 8 Jul 2026 21:48:00 +0800 Subject: [PATCH 28/39] rename hessian_backup to hessian_opbase --- src/ctrl_io/mod.rs | 16 ++++++------- src/dft/mod.rs | 2 +- src/dft/numint_matmul/hess_rks.rs | 2 +- src/dft/numint_matmul/hess_uks.rs | 2 +- .../cint_handling.rs | 0 .../config.rs | 0 .../hcore.rs | 0 .../krylov_block.rs | 0 src/{hessian_backup => hessian_opbase}/mod.rs | 24 +++++++++++++++---- .../nuc_repl.rs | 0 .../ovlp.rs | 0 .../point_group_detect/bits.rs | 0 .../point_group_detect/detect.rs | 0 .../point_group_detect/elements.rs | 0 .../point_group_detect/geom.rs | 0 .../point_group_detect/interface_to_rest.rs | 0 .../point_group_detect/linalg.rs | 0 .../point_group_detect/matrix.rs | 0 .../point_group_detect/mod.rs | 2 +- .../point_group_detect/molecule.rs | 0 .../point_group_detect/vec3.rs | 0 .../rscf.rs | 0 .../rscf_interface.rs | 2 +- .../trait_rhess.rs | 0 .../trait_uhess.rs | 0 .../trait_util.rs | 0 .../uscf.rs | 0 .../uscf_interface.rs | 2 +- src/{hessian_backup => hessian_opbase}/vib.rs | 0 src/lib.rs | 2 +- src/main_driver.rs | 10 ++++---- src/ri_jk/hess_r.rs | 4 ++-- src/ri_jk/hess_u.rs | 2 +- .../ctrl_parse.rs | 24 +++++++++---------- .../{hessian_backup => hessian_opbase}/mod.rs | 0 .../{hessian_backup => hessian_opbase}/rhf.rs | 4 ++-- .../rks_b3lyp.rs | 4 ++-- .../rks_b3lyp_grid_levels.rs | 4 ++-- .../{hessian_backup => hessian_opbase}/uhf.rs | 6 ++--- .../uks_tpss0_grid_levels.rs | 4 ++-- .../uks_tpssh.rs | 4 ++-- tests/test_hessian_backup.rs | 1 - tests/test_hessian_opbase.rs | 1 + 43 files changed, 69 insertions(+), 53 deletions(-) rename src/{hessian_backup => hessian_opbase}/cint_handling.rs (100%) rename src/{hessian_backup => hessian_opbase}/config.rs (100%) rename src/{hessian_backup => hessian_opbase}/hcore.rs (100%) rename src/{hessian_backup => hessian_opbase}/krylov_block.rs (100%) rename src/{hessian_backup => hessian_opbase}/mod.rs (52%) rename src/{hessian_backup => hessian_opbase}/nuc_repl.rs (100%) rename src/{hessian_backup => hessian_opbase}/ovlp.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/bits.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/detect.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/elements.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/geom.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/interface_to_rest.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/linalg.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/matrix.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/mod.rs (95%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/molecule.rs (100%) rename src/{hessian_backup => hessian_opbase}/point_group_detect/vec3.rs (100%) rename src/{hessian_backup => hessian_opbase}/rscf.rs (100%) rename src/{hessian_backup => hessian_opbase}/rscf_interface.rs (99%) rename src/{hessian_backup => hessian_opbase}/trait_rhess.rs (100%) rename src/{hessian_backup => hessian_opbase}/trait_uhess.rs (100%) rename src/{hessian_backup => hessian_opbase}/trait_util.rs (100%) rename src/{hessian_backup => hessian_opbase}/uscf.rs (100%) rename src/{hessian_backup => hessian_opbase}/uscf_interface.rs (99%) rename src/{hessian_backup => hessian_opbase}/vib.rs (100%) rename tests/{hessian_backup => hessian_opbase}/ctrl_parse.rs (67%) rename tests/{hessian_backup => hessian_opbase}/mod.rs (100%) rename tests/{hessian_backup => hessian_opbase}/rhf.rs (97%) rename tests/{hessian_backup => hessian_opbase}/rks_b3lyp.rs (92%) rename tests/{hessian_backup => hessian_opbase}/rks_b3lyp_grid_levels.rs (98%) rename tests/{hessian_backup => hessian_opbase}/uhf.rs (91%) rename tests/{hessian_backup => hessian_opbase}/uks_tpss0_grid_levels.rs (95%) rename tests/{hessian_backup => hessian_opbase}/uks_tpssh.rs (94%) delete mode 100644 tests/test_hessian_backup.rs create mode 100644 tests/test_hessian_opbase.rs diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index e6e9beb6f7..4089d0db9c 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -8,7 +8,7 @@ use std::{fs, sync::Arc}; 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::hessian_backup::config::HessSCFConfig; +use crate::hessian_opbase::config::HessSCFConfig; 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; @@ -386,9 +386,9 @@ pub struct InputKeywords { #[pyo3(get, set)] pub use_fxc_opt: bool, pub rel: RelativisticMethod, - /// Use module `hessian_backup` instead of `hessian` for Hessian calculation. - pub use_hessian_backup: bool, - pub hessian_backup: HessSCFConfig, + /// Use module `hessian_opbase` instead of `hessian` for Hessian calculation. + pub use_hessian_opbase: bool, + pub hessian_opbase: HessSCFConfig, } impl Default for InputKeywords { @@ -554,8 +554,8 @@ impl InputKeywords { thermo: None, use_fxc_opt: false, rel: RelativisticMethod::None, - use_hessian_backup: false, - hessian_backup: HessSCFConfig::default(), + use_hessian_opbase: false, + hessian_opbase: HessSCFConfig::default(), } } @@ -1943,12 +1943,12 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result { tmp_str.to_lowercase() }, other => String::from("legacy"), }; - tmp_input.use_hessian_backup = match tmp_ctrl.get("use_hessian_backup").unwrap_or(&serde_json::Value::Null) { + tmp_input.use_hessian_opbase = match tmp_ctrl.get("use_hessian_opbase").unwrap_or(&serde_json::Value::Null) { serde_json::Value::Bool(tmp_bool) => *tmp_bool, serde_json::Value::String(tmp_str) => tmp_str.to_lowercase().parse().unwrap_or(false), _ => false, }; - tmp_input.hessian_backup = tmp_ctrl.get("hessian_backup").map(serde_from_value).unwrap_or_default(); + tmp_input.hessian_opbase = tmp_ctrl.get("hessian_opbase").map(serde_from_value).unwrap_or_default(); //=========================================================== // Global check of ctrl keywords and futher modification diff --git a/src/dft/mod.rs b/src/dft/mod.rs index ea96b37577..ee79785c34 100644 --- a/src/dft/mod.rs +++ b/src/dft/mod.rs @@ -3510,7 +3510,7 @@ impl Grids { /// Build a grid using the same atom-grid machinery as [`Grids::build`], but with an explicit /// `level` override (instead of `mol.ctrl.grid_gen_level`). /// - /// This is used by the `hessian_backup` module to generate skeleton / cphf grids at levels + /// This is used by the `hessian_opbase` module to generate skeleton / cphf grids at levels /// that may differ from the SCF DFT grid. Unlike [`Grids::build`], this skips the /// `external_grids` path and MPI distribution, and does not apply the round-robin permutation. pub fn build_with_level(mol: &Molecule, level: usize) -> Grids { diff --git a/src/dft/numint_matmul/hess_rks.rs b/src/dft/numint_matmul/hess_rks.rs index 22499bd8e9..9187e0bad2 100644 --- a/src/dft/numint_matmul/hess_rks.rs +++ b/src/dft/numint_matmul/hess_rks.rs @@ -1,7 +1,7 @@ // see also pyhessref/nimatmul/rks.py use super::prelude::*; -use crate::hessian_backup::prelude::*; +use crate::hessian_opbase::prelude::*; use XCDenType::*; diff --git a/src/dft/numint_matmul/hess_uks.rs b/src/dft/numint_matmul/hess_uks.rs index 1ef5435ba2..353bab4f26 100644 --- a/src/dft/numint_matmul/hess_uks.rs +++ b/src/dft/numint_matmul/hess_uks.rs @@ -2,7 +2,7 @@ use super::hess_rks::{get_de_vxc_diag, get_de_vxc_off, get_drho, get_vmat_ip}; use super::prelude::*; -use crate::hessian_backup::prelude::*; +use crate::hessian_opbase::prelude::*; use XCDenType::*; diff --git a/src/hessian_backup/cint_handling.rs b/src/hessian_opbase/cint_handling.rs similarity index 100% rename from src/hessian_backup/cint_handling.rs rename to src/hessian_opbase/cint_handling.rs diff --git a/src/hessian_backup/config.rs b/src/hessian_opbase/config.rs similarity index 100% rename from src/hessian_backup/config.rs rename to src/hessian_opbase/config.rs diff --git a/src/hessian_backup/hcore.rs b/src/hessian_opbase/hcore.rs similarity index 100% rename from src/hessian_backup/hcore.rs rename to src/hessian_opbase/hcore.rs diff --git a/src/hessian_backup/krylov_block.rs b/src/hessian_opbase/krylov_block.rs similarity index 100% rename from src/hessian_backup/krylov_block.rs rename to src/hessian_opbase/krylov_block.rs diff --git a/src/hessian_backup/mod.rs b/src/hessian_opbase/mod.rs similarity index 52% rename from src/hessian_backup/mod.rs rename to src/hessian_opbase/mod.rs index 264a7fad19..47f921153e 100644 --- a/src/hessian_backup/mod.rs +++ b/src/hessian_opbase/mod.rs @@ -1,16 +1,32 @@ -//! Hessian module for REST (currently as backup implementation at PR!101). +//! Analytical hessian module for REST (Operator-Based analytical hessian architecture). //! -//! This module should work in most cases, but is not completely finished and tested. +//! The name `opbase` is short for "operator-based". This concept should echo ORCA's hessian implementation, +//! though this is recently found by me (ajz34) from ORCA output. We can split hessian contributions into +//! distinct operators (hcore, nuc, J/K, DFT numint, etc), and use trait to define the interface for each operator. +//! +//! 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. //! -//! This module does not contain extensive difficult implementation. We defined traits, +//! This module should work in most cases, but still requires further testing and efficiency update. +//! +//! This module does not contain extensive detailed implementation. We defined traits, //! some common implementations (hcore, nuc, ovlp), total hessian, 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 handle interface to REST before next pull request, in this module. +//! We will also handle interface to REST. //! //! This module currently does not handle post-SCF derivatives. +//! +//! Some important utilities comes from other programs, and we acknowledge them here. +//! - `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-HF) from PySCF, translated with help by AI, reviewed +//! by extensive testing. // trait definitions pub mod trait_rhess; diff --git a/src/hessian_backup/nuc_repl.rs b/src/hessian_opbase/nuc_repl.rs similarity index 100% rename from src/hessian_backup/nuc_repl.rs rename to src/hessian_opbase/nuc_repl.rs diff --git a/src/hessian_backup/ovlp.rs b/src/hessian_opbase/ovlp.rs similarity index 100% rename from src/hessian_backup/ovlp.rs rename to src/hessian_opbase/ovlp.rs diff --git a/src/hessian_backup/point_group_detect/bits.rs b/src/hessian_opbase/point_group_detect/bits.rs similarity index 100% rename from src/hessian_backup/point_group_detect/bits.rs rename to src/hessian_opbase/point_group_detect/bits.rs diff --git a/src/hessian_backup/point_group_detect/detect.rs b/src/hessian_opbase/point_group_detect/detect.rs similarity index 100% rename from src/hessian_backup/point_group_detect/detect.rs rename to src/hessian_opbase/point_group_detect/detect.rs diff --git a/src/hessian_backup/point_group_detect/elements.rs b/src/hessian_opbase/point_group_detect/elements.rs similarity index 100% rename from src/hessian_backup/point_group_detect/elements.rs rename to src/hessian_opbase/point_group_detect/elements.rs diff --git a/src/hessian_backup/point_group_detect/geom.rs b/src/hessian_opbase/point_group_detect/geom.rs similarity index 100% rename from src/hessian_backup/point_group_detect/geom.rs rename to src/hessian_opbase/point_group_detect/geom.rs diff --git a/src/hessian_backup/point_group_detect/interface_to_rest.rs b/src/hessian_opbase/point_group_detect/interface_to_rest.rs similarity index 100% rename from src/hessian_backup/point_group_detect/interface_to_rest.rs rename to src/hessian_opbase/point_group_detect/interface_to_rest.rs diff --git a/src/hessian_backup/point_group_detect/linalg.rs b/src/hessian_opbase/point_group_detect/linalg.rs similarity index 100% rename from src/hessian_backup/point_group_detect/linalg.rs rename to src/hessian_opbase/point_group_detect/linalg.rs diff --git a/src/hessian_backup/point_group_detect/matrix.rs b/src/hessian_opbase/point_group_detect/matrix.rs similarity index 100% rename from src/hessian_backup/point_group_detect/matrix.rs rename to src/hessian_opbase/point_group_detect/matrix.rs diff --git a/src/hessian_backup/point_group_detect/mod.rs b/src/hessian_opbase/point_group_detect/mod.rs similarity index 95% rename from src/hessian_backup/point_group_detect/mod.rs rename to src/hessian_opbase/point_group_detect/mod.rs index dac1d38811..79e9bcc76e 100644 --- a/src/hessian_backup/point_group_detect/mod.rs +++ b/src/hessian_opbase/point_group_detect/mod.rs @@ -15,7 +15,7 @@ //! //! # Example //! ```ignore -//! use rest::hessian_backup::point_group_detect::SymmMolecule; +//! use rest::hessian_opbase::point_group_detect::SymmMolecule; //! //! // H2O (Bohr) -> C2v //! let mol = SymmMolecule::new( diff --git a/src/hessian_backup/point_group_detect/molecule.rs b/src/hessian_opbase/point_group_detect/molecule.rs similarity index 100% rename from src/hessian_backup/point_group_detect/molecule.rs rename to src/hessian_opbase/point_group_detect/molecule.rs diff --git a/src/hessian_backup/point_group_detect/vec3.rs b/src/hessian_opbase/point_group_detect/vec3.rs similarity index 100% rename from src/hessian_backup/point_group_detect/vec3.rs rename to src/hessian_opbase/point_group_detect/vec3.rs diff --git a/src/hessian_backup/rscf.rs b/src/hessian_opbase/rscf.rs similarity index 100% rename from src/hessian_backup/rscf.rs rename to src/hessian_opbase/rscf.rs diff --git a/src/hessian_backup/rscf_interface.rs b/src/hessian_opbase/rscf_interface.rs similarity index 99% rename from src/hessian_backup/rscf_interface.rs rename to src/hessian_opbase/rscf_interface.rs index 5a2bfc8254..8d185a365d 100644 --- a/src/hessian_backup/rscf_interface.rs +++ b/src/hessian_opbase/rscf_interface.rs @@ -1,6 +1,6 @@ use super::prelude::*; use crate::geom_io::get_mass_charge; -use crate::hessian_backup::vib::*; +use crate::hessian_opbase::vib::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; diff --git a/src/hessian_backup/trait_rhess.rs b/src/hessian_opbase/trait_rhess.rs similarity index 100% rename from src/hessian_backup/trait_rhess.rs rename to src/hessian_opbase/trait_rhess.rs diff --git a/src/hessian_backup/trait_uhess.rs b/src/hessian_opbase/trait_uhess.rs similarity index 100% rename from src/hessian_backup/trait_uhess.rs rename to src/hessian_opbase/trait_uhess.rs diff --git a/src/hessian_backup/trait_util.rs b/src/hessian_opbase/trait_util.rs similarity index 100% rename from src/hessian_backup/trait_util.rs rename to src/hessian_opbase/trait_util.rs diff --git a/src/hessian_backup/uscf.rs b/src/hessian_opbase/uscf.rs similarity index 100% rename from src/hessian_backup/uscf.rs rename to src/hessian_opbase/uscf.rs diff --git a/src/hessian_backup/uscf_interface.rs b/src/hessian_opbase/uscf_interface.rs similarity index 99% rename from src/hessian_backup/uscf_interface.rs rename to src/hessian_opbase/uscf_interface.rs index 5d5a950862..08ce607ed5 100644 --- a/src/hessian_backup/uscf_interface.rs +++ b/src/hessian_opbase/uscf_interface.rs @@ -1,6 +1,6 @@ use super::prelude::*; use crate::geom_io::get_mass_charge; -use crate::hessian_backup::vib::*; +use crate::hessian_opbase::vib::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; diff --git a/src/hessian_backup/vib.rs b/src/hessian_opbase/vib.rs similarity index 100% rename from src/hessian_backup/vib.rs rename to src/hessian_opbase/vib.rs diff --git a/src/lib.rs b/src/lib.rs index 819903674e..64ded5c0fd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -86,7 +86,7 @@ pub mod fileop; pub mod ri_cphf; pub mod lib_rint; pub mod x2c; -pub mod hessian_backup; +pub mod hessian_opbase; //extern crate rest; diff --git a/src/main_driver.rs b/src/main_driver.rs index e307b6712f..28cec314c2 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -243,20 +243,20 @@ pub fn main_driver() -> anyhow::Result<()> { }, // UNVERIFIED NORMAL MODES CALCULATION JobType::NormalModes => { - if !scf_data.mol.ctrl.use_hessian_backup { + if !scf_data.mol.ctrl.use_hessian_opbase { eval_normal_modes(&mut scf_data, &mut time_mark, &mpi_operator); } else { - use crate::hessian_backup::rscf_interface::rscf_hess_interface; - use crate::hessian_backup::uscf_interface::uscf_hess_interface; + use crate::hessian_opbase::rscf_interface::rscf_hess_interface; + use crate::hessian_opbase::uscf_interface::uscf_hess_interface; - eprintln!("[WARN] You are using the hessian_backup module, which is still under development."); + eprintln!("[WARN] You are using the hessian_opbase module, which is still under development."); // simple guard, but currently many methods (solvent, range-separate, dftd are not supported) if scf_data.mol.xc_data.is_fifth_dfa() { panic!("Normal modes calculation is currently not available for post-SCF methods."); } - let config = &scf_data.mol.ctrl.hessian_backup; + let config = &scf_data.mol.ctrl.hessian_opbase; match scf_data.scftype { SCFType::RHF => { rscf_hess_interface(&scf_data, config); diff --git a/src/ri_jk/hess_r.rs b/src/ri_jk/hess_r.rs index e0fd92be6f..a76257f9f7 100644 --- a/src/ri_jk/hess_r.rs +++ b/src/ri_jk/hess_r.rs @@ -17,8 +17,8 @@ use super::prelude_dev::*; use crate::grad::rhf::pack_triu_tilde; -use crate::hessian_backup::cint_handling::*; -use crate::hessian_backup::prelude::*; +use crate::hessian_opbase::cint_handling::*; +use crate::hessian_opbase::prelude::*; use crate::ri_jk::util::*; use crate::ri_jk::decompose::*; diff --git a/src/ri_jk/hess_u.rs b/src/ri_jk/hess_u.rs index 13fa7fa2e1..baa0d4ac1b 100644 --- a/src/ri_jk/hess_u.rs +++ b/src/ri_jk/hess_u.rs @@ -22,7 +22,7 @@ //! per spin in bra form (same-spin only). use super::prelude_dev::*; -use crate::hessian_backup::prelude::*; +use crate::hessian_opbase::prelude::*; use crate::ri_jk::util::*; use crate::ri_jk::decompose::*; diff --git a/tests/hessian_backup/ctrl_parse.rs b/tests/hessian_opbase/ctrl_parse.rs similarity index 67% rename from tests/hessian_backup/ctrl_parse.rs rename to tests/hessian_opbase/ctrl_parse.rs index 24c617f2ec..5af82cb449 100644 --- a/tests/hessian_backup/ctrl_parse.rs +++ b/tests/hessian_opbase/ctrl_parse.rs @@ -1,5 +1,5 @@ -//! Test that the `hessian_backup` keyword in `ctrl.in` is parsed into -//! `InputKeywords.hessian_backup` via the serde path (mirroring `j2c_decomp`). +//! Test that the `hessian_opbase` keyword in `ctrl.in` is parsed into +//! `InputKeywords.hessian_opbase` via the serde path (mirroring `j2c_decomp`). //! //! Only exercises the ctrl-block parser (`parse_ctrl_keywords`); deliberately does *not* build a //! full `Molecule`, to avoid global thread-pool / basis-pool side effects that can perturb other @@ -8,15 +8,15 @@ use pyrest::ctrl_io; #[test] -fn test_parse_hessian_backup_table() { +fn test_parse_hessian_opbase_table() { let input = r##" [ctrl] xc = "b3lyp" basis_path = "def2-svp" num_threads = 16 - use_hessian_backup = true + use_hessian_opbase = true -[ctrl.hessian_backup] +[ctrl.hessian_opbase] cphf_tol = 1e-7 cphf_max_cycle = 7 grid_level_cphf = 2 @@ -27,8 +27,8 @@ fn test_parse_hessian_backup_table() { let keys = toml::from_str::(input).unwrap(); let ctrl = ctrl_io::parse_ctrl_keywords(&keys).unwrap(); - assert!(ctrl.use_hessian_backup, "use_hessian_backup should be parsed"); - let cfg = &ctrl.hessian_backup; + assert!(ctrl.use_hessian_opbase, "use_hessian_opbase should be parsed"); + let cfg = &ctrl.hessian_opbase; assert!((cfg.cphf_tol - 1e-7).abs() < 1e-12); assert_eq!(cfg.cphf_max_cycle, 7); assert_eq!(cfg.grid_level_cphf, Some(2)); @@ -37,7 +37,7 @@ fn test_parse_hessian_backup_table() { } #[test] -fn test_hessian_backup_default_when_absent() { +fn test_hessian_opbase_default_when_absent() { let input = r##" [ctrl] xc = "b3lyp" @@ -47,9 +47,9 @@ fn test_hessian_backup_default_when_absent() { let keys = toml::from_str::(input).unwrap(); let ctrl = ctrl_io::parse_ctrl_keywords(&keys).unwrap(); - assert!(!ctrl.use_hessian_backup); + assert!(!ctrl.use_hessian_opbase); // defaults from HessSCFConfig::default() - assert_eq!(ctrl.hessian_backup.grid_level_cphf, None); - assert_eq!(ctrl.hessian_backup.grid_level_skeleton, None); - assert_eq!(ctrl.hessian_backup.cphf_max_cycle, 42); + assert_eq!(ctrl.hessian_opbase.grid_level_cphf, None); + assert_eq!(ctrl.hessian_opbase.grid_level_skeleton, None); + assert_eq!(ctrl.hessian_opbase.cphf_max_cycle, 42); } diff --git a/tests/hessian_backup/mod.rs b/tests/hessian_opbase/mod.rs similarity index 100% rename from tests/hessian_backup/mod.rs rename to tests/hessian_opbase/mod.rs diff --git a/tests/hessian_backup/rhf.rs b/tests/hessian_opbase/rhf.rs similarity index 97% rename from tests/hessian_backup/rhf.rs rename to tests/hessian_opbase/rhf.rs index 94824130a3..ae401c6778 100644 --- a/tests/hessian_backup/rhf.rs +++ b/tests/hessian_opbase/rhf.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_backup::config::HessSCFConfig; -use pyrest::hessian_backup::rscf_interface::rscf_hess_interface; +use pyrest::hessian_opbase::config::HessSCFConfig; +use pyrest::hessian_opbase::rscf_interface::rscf_hess_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; diff --git a/tests/hessian_backup/rks_b3lyp.rs b/tests/hessian_opbase/rks_b3lyp.rs similarity index 92% rename from tests/hessian_backup/rks_b3lyp.rs rename to tests/hessian_opbase/rks_b3lyp.rs index 529e153c03..612223f473 100644 --- a/tests/hessian_backup/rks_b3lyp.rs +++ b/tests/hessian_opbase/rks_b3lyp.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_backup::rscf_interface::rscf_hess_interface; -use pyrest::hessian_backup::config::HessSCFConfig; +use pyrest::hessian_opbase::rscf_interface::rscf_hess_interface; +use pyrest::hessian_opbase::config::HessSCFConfig; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; diff --git a/tests/hessian_backup/rks_b3lyp_grid_levels.rs b/tests/hessian_opbase/rks_b3lyp_grid_levels.rs similarity index 98% rename from tests/hessian_backup/rks_b3lyp_grid_levels.rs rename to tests/hessian_opbase/rks_b3lyp_grid_levels.rs index 59bfdc9e4f..2557bfaf66 100644 --- a/tests/hessian_backup/rks_b3lyp_grid_levels.rs +++ b/tests/hessian_opbase/rks_b3lyp_grid_levels.rs @@ -3,8 +3,8 @@ //! Exercises the dedicated CP-KS grid path (`ni_cpks = Some`) and the skeleton-grid //! regeneration path for a GGA functional. -use pyrest::hessian_backup::config::HessSCFConfig; -use pyrest::hessian_backup::rscf_interface::rscf_hess_interface; +use pyrest::hessian_opbase::config::HessSCFConfig; +use pyrest::hessian_opbase::rscf_interface::rscf_hess_interface; use pyrest::ctrl_io; use pyrest::dft::numint_matmul::hess_rks::{get_hess_ncomp_ao_dm0, get_rho_vxc_fxc, make_cpks_vxc_fxc}; diff --git a/tests/hessian_backup/uhf.rs b/tests/hessian_opbase/uhf.rs similarity index 91% rename from tests/hessian_backup/uhf.rs rename to tests/hessian_opbase/uhf.rs index 7fa0450aa8..f4a03710c8 100644 --- a/tests/hessian_backup/uhf.rs +++ b/tests/hessian_opbase/uhf.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_backup::uscf_interface::uscf_hess_interface; -use pyrest::hessian_backup::config::HessSCFConfig; +use pyrest::hessian_opbase::uscf_interface::uscf_hess_interface; +use pyrest::hessian_opbase::config::HessSCFConfig; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; @@ -35,7 +35,7 @@ static INPUT_NH3: &str = r##" H 0.1 0.1 1.2 """ -[hessian_backup] +[hessian_opbase] atm_list = [0, 1, 3] "##; diff --git a/tests/hessian_backup/uks_tpss0_grid_levels.rs b/tests/hessian_opbase/uks_tpss0_grid_levels.rs similarity index 95% rename from tests/hessian_backup/uks_tpss0_grid_levels.rs rename to tests/hessian_opbase/uks_tpss0_grid_levels.rs index 492e2c549a..66a2935828 100644 --- a/tests/hessian_backup/uks_tpss0_grid_levels.rs +++ b/tests/hessian_opbase/uks_tpss0_grid_levels.rs @@ -3,8 +3,8 @@ //! Exercises the MGGA skeleton-grid regeneration (`grid_gen_level + 2` default) and the //! dedicated CP-KS grid path on the UKS side. -use pyrest::hessian_backup::config::HessSCFConfig; -use pyrest::hessian_backup::uscf_interface::uscf_hess_interface; +use pyrest::hessian_opbase::config::HessSCFConfig; +use pyrest::hessian_opbase::uscf_interface::uscf_hess_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; diff --git a/tests/hessian_backup/uks_tpssh.rs b/tests/hessian_opbase/uks_tpssh.rs similarity index 94% rename from tests/hessian_backup/uks_tpssh.rs rename to tests/hessian_opbase/uks_tpssh.rs index ba7715e608..ec95d391aa 100644 --- a/tests/hessian_backup/uks_tpssh.rs +++ b/tests/hessian_opbase/uks_tpssh.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_backup::config::HessSCFConfig; -use pyrest::hessian_backup::uscf_interface::uscf_hess_interface; +use pyrest::hessian_opbase::config::HessSCFConfig; +use pyrest::hessian_opbase::uscf_interface::uscf_hess_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; diff --git a/tests/test_hessian_backup.rs b/tests/test_hessian_backup.rs deleted file mode 100644 index 53cf145a45..0000000000 --- a/tests/test_hessian_backup.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod hessian_backup; \ No newline at end of file diff --git a/tests/test_hessian_opbase.rs b/tests/test_hessian_opbase.rs new file mode 100644 index 0000000000..1bd61cdb5f --- /dev/null +++ b/tests/test_hessian_opbase.rs @@ -0,0 +1 @@ +pub mod hessian_opbase; \ No newline at end of file -- Gitee From e7beb3f6f369ce7b76ff0952372153fa492f0d3f Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 8 Jul 2026 23:10:14 +0800 Subject: [PATCH 29/39] hessian_opbase: update some printing utilities --- src/hessian_opbase/krylov_block.rs | 2 +- src/hessian_opbase/mod.rs | 3 + src/hessian_opbase/rscf_interface.rs | 44 +----------- src/hessian_opbase/uscf_interface.rs | 44 +----------- src/hessian_opbase/vib.rs | 3 - src/hessian_opbase/vib_interface.rs | 101 +++++++++++++++++++++++++++ 6 files changed, 109 insertions(+), 88 deletions(-) create mode 100644 src/hessian_opbase/vib_interface.rs diff --git a/src/hessian_opbase/krylov_block.rs b/src/hessian_opbase/krylov_block.rs index 22b24bf848..104eebd184 100644 --- a/src/hessian_opbase/krylov_block.rs +++ b/src/hessian_opbase/krylov_block.rs @@ -151,7 +151,7 @@ pub fn krylov_block( let max_abs: f64 = x1_new.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs())); println!( - "restart {} inner {} (total cycle {}): max(||v||^2) = {:.3e}, max(||v||) = {:.3e}, per-elem L2 = {:.3e}, max-abs = {:.3e}", + "krylov restart {} inner {} (total cycle {}): max(||v||^2) = {:.3e}, max(||v||) = {:.3e}, per-elem L2 = {:.3e}, max-abs = {:.3e}", restart_idx, inner + 1, total_cycles, max_innerprod, r, l2_per_elem, max_abs, ); diff --git a/src/hessian_opbase/mod.rs b/src/hessian_opbase/mod.rs index 47f921153e..cdd9ba53ba 100644 --- a/src/hessian_opbase/mod.rs +++ b/src/hessian_opbase/mod.rs @@ -28,6 +28,8 @@ //! - `krylov_block.rs`: Krylov solver (used in CP-HF) from PySCF, translated with help by AI, reviewed //! by extensive testing. +#![warn(unused)] + // trait definitions pub mod trait_rhess; pub mod trait_uhess; @@ -50,6 +52,7 @@ pub mod uscf_interface; // vibrational analysis pub mod vib; +pub mod vib_interface; // utilities pub mod cint_handling; diff --git a/src/hessian_opbase/rscf_interface.rs b/src/hessian_opbase/rscf_interface.rs index 8d185a365d..859b1fe162 100644 --- a/src/hessian_opbase/rscf_interface.rs +++ b/src/hessian_opbase/rscf_interface.rs @@ -1,6 +1,6 @@ use super::prelude::*; -use crate::geom_io::get_mass_charge; use crate::hessian_opbase::vib::*; +use crate::hessian_opbase::vib_interface::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; @@ -62,7 +62,6 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, let mut hess_nimatmul_obj = (!is_hf).then(|| { use crate::dft::numint_matmul::hess_rks::RHessKSNIMatmul; use crate::dft::numint_matmul::nimatmul::NIMatmul; - use crate::dft::numint_matmul::prelude::*; use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; use crate::dft::Grids; use libxc::prelude::*; @@ -158,44 +157,5 @@ pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, // --- perform vibrational analysis --- // - // first transpose hessian to [3*natm, 3*natm] - let natm = de_hess.shape()[3]; - let hess = de_hess.transpose((0, 2, 1, 3)).into_shape((3 * natm, 3 * natm)); - let atm_list = config.atm_list.clone().unwrap_or_else(|| (0..mol.natm()).collect_vec()); - let elems = atm_list.iter().map(|&i| scf_data.mol.geom.elem[i].clone()).collect_vec(); - - let mass_charge = get_mass_charge(&elems); - let mass = mass_charge.iter().map(|(m, _)| *m).collect_vec(); - let mass_rt = rt::asarray((&mass, &device)); - - let geom = atm_list.iter().map(|&i| mol.atom_coord(i)).collect_vec(); - let geom_rt = rt::asarray((&geom, &device)).into_unpack_array(0); - let vib = harmonic_analysis(hess.view(), geom_rt.view(), mass_rt.view(), true, true); - - println!("=============== Vibrational Analysis ==============="); - let elems_ref = elems.iter().map(|e| e.as_str()).collect_vec(); - let msg_vib = print_vibs(&vib, &elems_ref, NormCo::X, true, Some(3), 4, None); - println!("{}", msg_vib); - - let geom_c = mass_centred_geom(geom_rt.view(), mass_rt.view()); - let rc_cm = rotation_const(mass_rt.view(), geom_c.view(), "wavenumber").to_vec(); - let rc_ghz = rotation_const(mass_rt.view(), geom_c.view(), "GHz"); - let rotor = RotorType::from_rot_const_ghz(rc_ghz.view()); - let mass_sum = mass.iter().sum::(); - let e0 = scf_data.scf_energy; - let multiplicity = scf_data.mol.ctrl.spin; - - use super::point_group_detect::interface_to_rest::get_full_point_group_for_vib; - let tol_pg = config.tol_point_group / (1.0 + natm as f64).sqrt(); - let (pg_name, pg_sigma) = get_full_point_group_for_vib(&elems, &mass, &geom, tol_pg); - - let th = thermo(&vib, 298.15, 101325.0, multiplicity as _, mass_sum, e0, pg_sigma as _, &rc_cm, rotor); - - println!("=============== Thermo Analysis ==============="); - println!(""); - println!("Point group: {}, sigma (rotation symmetry number): {}", pg_name, pg_sigma); - let msg_th = print_thermo(&th, multiplicity as _, mass_sum); - println!("{}", msg_th); - - (de_hess.into_shape(-1).into_vec(), vib, th) + vibration_analysis_interface(scf_data, config, de_hess.view()) } diff --git a/src/hessian_opbase/uscf_interface.rs b/src/hessian_opbase/uscf_interface.rs index 08ce607ed5..4c220ab4ae 100644 --- a/src/hessian_opbase/uscf_interface.rs +++ b/src/hessian_opbase/uscf_interface.rs @@ -1,6 +1,6 @@ use super::prelude::*; -use crate::geom_io::get_mass_charge; use crate::hessian_opbase::vib::*; +use crate::hessian_opbase::vib_interface::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; @@ -71,7 +71,6 @@ pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, let mut hess_nimatmul_obj = (!is_hf).then(|| { use crate::dft::numint_matmul::hess_uks::UHessKSNIMatmul; use crate::dft::numint_matmul::nimatmul::NIMatmul; - use crate::dft::numint_matmul::prelude::*; use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; use crate::dft::Grids; use libxc::prelude::*; @@ -167,44 +166,5 @@ pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, // --- perform vibrational analysis --- // - // first transpose hessian to [3*natm, 3*natm] - let natm = de_hess.shape()[3]; - let hess = de_hess.transpose((0, 2, 1, 3)).into_shape((3 * natm, 3 * natm)); - let atm_list = config.atm_list.clone().unwrap_or_else(|| (0..mol.natm()).collect_vec()); - let elems = atm_list.iter().map(|&i| scf_data.mol.geom.elem[i].clone()).collect_vec(); - - let mass_charge = get_mass_charge(&elems); - let mass = mass_charge.iter().map(|(m, _)| *m).collect_vec(); - let mass_rt = rt::asarray((&mass, &device)); - - let geom = atm_list.iter().map(|&i| mol.atom_coord(i)).collect_vec(); - let geom_rt = rt::asarray((&geom, &device)).into_unpack_array(0); - let vib = harmonic_analysis(hess.view(), geom_rt.view(), mass_rt.view(), true, true); - - println!("=============== Vibrational Analysis ==============="); - let elems_ref = elems.iter().map(|e| e.as_str()).collect_vec(); - let msg_vib = print_vibs(&vib, &elems_ref, NormCo::X, true, Some(3), 4, None); - println!("{}", msg_vib); - - let geom_c = mass_centred_geom(geom_rt.view(), mass_rt.view()); - let rc_cm = rotation_const(mass_rt.view(), geom_c.view(), "wavenumber").to_vec(); - let rc_ghz = rotation_const(mass_rt.view(), geom_c.view(), "GHz"); - let rotor = RotorType::from_rot_const_ghz(rc_ghz.view()); - let mass_sum = mass.iter().sum::(); - let e0 = scf_data.scf_energy; - let multiplicity = scf_data.mol.ctrl.spin; - - use super::point_group_detect::interface_to_rest::get_full_point_group_for_vib; - let tol_pg = config.tol_point_group / (1.0 + natm as f64).sqrt(); - let (pg_name, pg_sigma) = get_full_point_group_for_vib(&elems, &mass, &geom, tol_pg); - - let th = thermo(&vib, 298.15, 101325.0, multiplicity as _, mass_sum, e0, pg_sigma as _, &rc_cm, rotor); - - println!("=============== Thermo Analysis ==============="); - println!(""); - println!("Point group: {}, sigma (rotation symmetry number): {}", pg_name, pg_sigma); - let msg_th = print_thermo(&th, multiplicity as _, mass_sum); - println!("{}", msg_th); - - (de_hess.into_shape(-1).into_vec(), vib, th) + vibration_analysis_interface(scf_data, config, de_hess.view()) } diff --git a/src/hessian_opbase/vib.rs b/src/hessian_opbase/vib.rs index 688c20c001..75fc6a5200 100644 --- a/src/hessian_opbase/vib.rs +++ b/src/hessian_opbase/vib.rs @@ -22,9 +22,6 @@ use crate::constants::{ R_GAS, SPEED_OF_LIGHT, }; -/// Tolerance for detecting nearly-linear geometries in `get_tr_space`. -pub const LINEAR_A_TOL: f64 = 1.0e-2; - /// cm⁻¹ conversion factor from force-constant eigenvalues (atomic units). fn uconv_cm1() -> f64 { (AVOGADRO * HARTREE2J * 1.0e19).sqrt() / (2.0 * std::f64::consts::PI * SPEED_OF_LIGHT * BOHR2ANG) diff --git a/src/hessian_opbase/vib_interface.rs b/src/hessian_opbase/vib_interface.rs new file mode 100644 index 0000000000..0bf33ddcb8 --- /dev/null +++ b/src/hessian_opbase/vib_interface.rs @@ -0,0 +1,101 @@ +use super::prelude::*; +use crate::geom_io::get_mass_charge; +use crate::hessian_opbase::vib::*; +use crate::ri_jk::util::get_cint_mol; +use crate::SCF; + +/// Vibrational analysis interface for REST. +/// +/// `de_hess` should be of shape `[3, 3, natm, natm]`, and is the hessian matrix in atomic units. +pub fn vibration_analysis_interface( + scf_data: &SCF, + config: &HessSCFConfig, + de_hess: TsrView, +) -> (Vec, VibInfo, ThermoInfo) { + let device = de_hess.device().clone(); + let mol_obj = &scf_data.mol; + let mol = get_cint_mol(mol_obj); + + println!("=============== Vibrational Analysis (in hessian_opbase) ==============="); + + // first transpose hessian to [3*natm, 3*natm] + let natm = de_hess.shape()[3]; + let hess = de_hess.transpose((0, 2, 1, 3)).into_shape((3 * natm, 3 * natm)); + let atm_list = config.atm_list.clone().unwrap_or_else(|| (0..mol.natm()).collect_vec()); + let elems = atm_list.iter().map(|&i| scf_data.mol.geom.elem[i].clone()).collect_vec(); + + let mass_charge = get_mass_charge(&elems); + let mass = mass_charge.iter().map(|(m, _)| *m).collect_vec(); + let mass_rt = rt::asarray((&mass, &device)); + + let geom = atm_list.iter().map(|&i| mol.atom_coord(i)).collect_vec(); + let geom_rt = rt::asarray((&geom, &device)).into_unpack_array(0); + let vib = harmonic_analysis(hess.view(), geom_rt.view(), mass_rt.view(), true, true); + println!(""); + let elems_ref = elems.iter().map(|e| e.as_str()).collect_vec(); + let msg_vib = print_vibs(&vib, &elems_ref, NormCo::X, true, Some(3), 4, None); + println!("{}", msg_vib); + println!(""); + + println!("=============== End of Vibrational Analysis (in hessian_opbase) ==============="); + + println!(""); + + println!("=============== Thermo Analysis (Usual Style in hessian_opbase) ==============="); + println!(""); + println!("Note: This is usual thermo analysis. To activate Shermo-style thermo analysis, please not use hessian_opbase module at this time."); + + let geom_c = mass_centred_geom(geom_rt.view(), mass_rt.view()); + let rc_cm = rotation_const(mass_rt.view(), geom_c.view(), "wavenumber").to_vec(); + let rc_ghz = rotation_const(mass_rt.view(), geom_c.view(), "GHz"); + let rotor = RotorType::from_rot_const_ghz(rc_ghz.view()); + let mass_sum = mass.iter().sum::(); + let e0 = scf_data.scf_energy; + let multiplicity = scf_data.mol.ctrl.spin; + + use super::point_group_detect::interface_to_rest::get_full_point_group_for_vib; + let tol_pg = config.tol_point_group / (1.0 + natm as f64).sqrt(); + let (pg_name, pg_sigma) = get_full_point_group_for_vib(&elems, &mass, &geom, tol_pg); + + let thermo_ctrl = scf_data.mol.ctrl.thermo.as_ref(); + let temperature = thermo_ctrl.map(|t| t.temperature).unwrap_or(298.15); + let pressure = thermo_ctrl.map(|t| t.pressure * 101325.0).unwrap_or(101325.0); + // Electronic energy can be zero, which means we use SCF energy as electronic energy. + let mut electronic_energy = thermo_ctrl.map(|t| t.electronic_energy).unwrap_or(e0); + if electronic_energy == 0.0 { + electronic_energy = e0; + } + // Symmetric number will only set to user-defined value, if it is not 1. + // This is due to current implementation of thermo keyword default will be 1. Value 1 does not mean user specified 1. + let mut symmetry_number = thermo_ctrl.map(|t| t.symmetry_number as i64).unwrap_or(0); + if symmetry_number == 1 || symmetry_number == 0 { + symmetry_number = pg_sigma as i64; + } + if symmetry_number != pg_sigma as i64 { + println!( + "Warning: symmetry_number in thermo section ({}) is different from detected point group sigma ({}).", + symmetry_number, pg_sigma + ); + println!(" We will use the user-specified symmetry_number in thermo section for thermo analysis."); + } + + let th = thermo( + &vib, + temperature, + pressure, + multiplicity as _, + mass_sum, + electronic_energy, + symmetry_number, + &rc_cm, + rotor, + ); + println!("Point group: {}, sigma (rotation symmetry number): {}", pg_name, pg_sigma); + let msg_th = print_thermo(&th, multiplicity as _, mass_sum); + println!("{}", msg_th); + println!(""); + + println!("=============== End of Thermo Analysis (Usual Style in hessian_opbase) ==============="); + + (de_hess.into_shape(-1).into_vec(), vib, th) +} -- Gitee From 91d4ce08eed0fa33f5f3194565687d5f09912590 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Wed, 8 Jul 2026 23:39:17 +0800 Subject: [PATCH 30/39] hessian_opbase: update main_driver and ctrl logic --- src/ctrl_io/mod.rs | 16 +++------ src/main_driver.rs | 50 ++++++++++++++------------- tests/hessian_opbase/ctrl_parse.rs | 55 ------------------------------ tests/hessian_opbase/mod.rs | 1 - 4 files changed, 31 insertions(+), 91 deletions(-) delete mode 100644 tests/hessian_opbase/ctrl_parse.rs diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index 4089d0db9c..78ee660c89 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -86,6 +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.hessian_opbase = tmp_keys.get("hessian_opbase").map(serde_from_value); Ok((tmp_input,tmp_geomcell)) } @@ -386,9 +387,8 @@ pub struct InputKeywords { #[pyo3(get, set)] pub use_fxc_opt: bool, pub rel: RelativisticMethod, - /// Use module `hessian_opbase` instead of `hessian` for Hessian calculation. - pub use_hessian_opbase: bool, - pub hessian_opbase: HessSCFConfig, + /// + pub hessian_opbase: Option, } impl Default for InputKeywords { @@ -554,8 +554,7 @@ impl InputKeywords { thermo: None, use_fxc_opt: false, rel: RelativisticMethod::None, - use_hessian_opbase: false, - hessian_opbase: HessSCFConfig::default(), + hessian_opbase: None, } } @@ -1943,12 +1942,7 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result { tmp_str.to_lowercase() }, other => String::from("legacy"), }; - tmp_input.use_hessian_opbase = match tmp_ctrl.get("use_hessian_opbase").unwrap_or(&serde_json::Value::Null) { - serde_json::Value::Bool(tmp_bool) => *tmp_bool, - serde_json::Value::String(tmp_str) => tmp_str.to_lowercase().parse().unwrap_or(false), - _ => false, - }; - tmp_input.hessian_opbase = tmp_ctrl.get("hessian_opbase").map(serde_from_value).unwrap_or_default(); + println!("[DEBUG]: {:?}", tmp_ctrl.get("hessian_opbase")); //=========================================================== // Global check of ctrl keywords and futher modification diff --git a/src/main_driver.rs b/src/main_driver.rs index 28cec314c2..85dc681720 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -243,30 +243,7 @@ pub fn main_driver() -> anyhow::Result<()> { }, // UNVERIFIED NORMAL MODES CALCULATION JobType::NormalModes => { - if !scf_data.mol.ctrl.use_hessian_opbase { - eval_normal_modes(&mut scf_data, &mut time_mark, &mpi_operator); - } else { - use crate::hessian_opbase::rscf_interface::rscf_hess_interface; - use crate::hessian_opbase::uscf_interface::uscf_hess_interface; - - eprintln!("[WARN] You are using the hessian_opbase module, which is still under development."); - - // simple guard, but currently many methods (solvent, range-separate, dftd are not supported) - if scf_data.mol.xc_data.is_fifth_dfa() { - panic!("Normal modes calculation is currently not available for post-SCF methods."); - } - - let config = &scf_data.mol.ctrl.hessian_opbase; - match scf_data.scftype { - SCFType::RHF => { - rscf_hess_interface(&scf_data, config); - } - SCFType::UHF => { - uscf_hess_interface(&scf_data, config); - }, - _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types.") - } - } + eval_normal_modes(&mut scf_data, &mut time_mark, &mpi_operator); }, // ------------ _ => {} @@ -407,6 +384,31 @@ pub fn main_driver() -> anyhow::Result<()> { if let Some(ref hess_ctrl) = scf_data.mol.ctrl.hessian { crate::hessian::rhf_hessian_main(&scf_data, hess_ctrl, &mut time_mark); } + + if let Some(ref hessian_opbase_ctrl) = scf_data.mol.ctrl.hessian_opbase { + time_mark.new_item("HessianOpBase", "analytical Hessian (hessian_opbase module)"); + time_mark.count_start("HessianOpBase"); + use crate::hessian_opbase::rscf_interface::rscf_hess_interface; + use crate::hessian_opbase::uscf_interface::uscf_hess_interface; + + eprintln!("[WARN] You are using the hessian_opbase module, which is still under development."); + + // simple guard, but currently many methods (solvent, range-separate, dftd are not supported) + if scf_data.mol.xc_data.is_fifth_dfa() { + panic!("Normal modes calculation is currently not available for post-SCF methods."); + } + + match scf_data.scftype { + SCFType::RHF => { + rscf_hess_interface(&scf_data, hessian_opbase_ctrl); + } + SCFType::UHF => { + uscf_hess_interface(&scf_data, hessian_opbase_ctrl); + }, + _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types.") + } + time_mark.count("HessianOpBase"); + } time_mark.count("Overall"); diff --git a/tests/hessian_opbase/ctrl_parse.rs b/tests/hessian_opbase/ctrl_parse.rs deleted file mode 100644 index 5af82cb449..0000000000 --- a/tests/hessian_opbase/ctrl_parse.rs +++ /dev/null @@ -1,55 +0,0 @@ -//! Test that the `hessian_opbase` keyword in `ctrl.in` is parsed into -//! `InputKeywords.hessian_opbase` via the serde path (mirroring `j2c_decomp`). -//! -//! Only exercises the ctrl-block parser (`parse_ctrl_keywords`); deliberately does *not* build a -//! full `Molecule`, to avoid global thread-pool / basis-pool side effects that can perturb other -//! tests sharing the test binary. - -use pyrest::ctrl_io; - -#[test] -fn test_parse_hessian_opbase_table() { - let input = r##" -[ctrl] - xc = "b3lyp" - basis_path = "def2-svp" - num_threads = 16 - use_hessian_opbase = true - -[ctrl.hessian_opbase] - cphf_tol = 1e-7 - cphf_max_cycle = 7 - grid_level_cphf = 2 - grid_level_skeleton = 4 - atm_list = [0, 1] -"##; - - let keys = toml::from_str::(input).unwrap(); - let ctrl = ctrl_io::parse_ctrl_keywords(&keys).unwrap(); - - assert!(ctrl.use_hessian_opbase, "use_hessian_opbase should be parsed"); - let cfg = &ctrl.hessian_opbase; - assert!((cfg.cphf_tol - 1e-7).abs() < 1e-12); - assert_eq!(cfg.cphf_max_cycle, 7); - assert_eq!(cfg.grid_level_cphf, Some(2)); - assert_eq!(cfg.grid_level_skeleton, Some(4)); - assert_eq!(cfg.atm_list.as_deref(), Some(&[0usize, 1][..])); -} - -#[test] -fn test_hessian_opbase_default_when_absent() { - let input = r##" -[ctrl] - xc = "b3lyp" - basis_path = "def2-svp" - num_threads = 16 -"##; - let keys = toml::from_str::(input).unwrap(); - let ctrl = ctrl_io::parse_ctrl_keywords(&keys).unwrap(); - - assert!(!ctrl.use_hessian_opbase); - // defaults from HessSCFConfig::default() - assert_eq!(ctrl.hessian_opbase.grid_level_cphf, None); - assert_eq!(ctrl.hessian_opbase.grid_level_skeleton, None); - assert_eq!(ctrl.hessian_opbase.cphf_max_cycle, 42); -} diff --git a/tests/hessian_opbase/mod.rs b/tests/hessian_opbase/mod.rs index faf7e0ea35..9927b8629c 100644 --- a/tests/hessian_opbase/mod.rs +++ b/tests/hessian_opbase/mod.rs @@ -1,4 +1,3 @@ -pub mod ctrl_parse; pub mod rhf; pub mod rks_b3lyp; pub mod rks_b3lyp_grid_levels; -- Gitee From 9123bbc2962c6ba4e8f3f190605528be97ffec3a Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 11:06:06 +0800 Subject: [PATCH 31/39] rename hessian_opbase to analdrv --- .../cint_handling.rs | 0 src/{hessian_opbase => analdrv}/config.rs | 0 src/{hessian_opbase => analdrv}/hcore.rs | 0 .../krylov_block.rs | 0 src/{hessian_opbase => analdrv}/mod.rs | 13 ++++++----- src/{hessian_opbase => analdrv}/nuc_repl.rs | 0 src/{hessian_opbase => analdrv}/ovlp.rs | 0 .../point_group_detect/bits.rs | 0 .../point_group_detect/detect.rs | 0 .../point_group_detect/elements.rs | 0 .../point_group_detect/geom.rs | 0 .../point_group_detect/interface_to_rest.rs | 0 .../point_group_detect/linalg.rs | 0 .../point_group_detect/matrix.rs | 0 .../point_group_detect/mod.rs | 2 +- .../point_group_detect/molecule.rs | 0 .../point_group_detect/vec3.rs | 0 src/{hessian_opbase => analdrv}/rscf.rs | 0 .../rscf_interface.rs | 4 ++-- .../trait_rhess.rs | 0 .../trait_uhess.rs | 0 src/{hessian_opbase => analdrv}/trait_util.rs | 0 src/{hessian_opbase => analdrv}/uscf.rs | 0 .../uscf_interface.rs | 4 ++-- src/{hessian_opbase => analdrv}/vib.rs | 0 .../vib_interface.rs | 12 +++++----- src/ctrl_io/mod.rs | 10 ++++----- src/dft/mod.rs | 2 +- src/dft/numint_matmul/hess_rks.rs | 2 +- src/dft/numint_matmul/hess_uks.rs | 2 +- src/lib.rs | 2 +- src/main_driver.rs | 22 +++++++++++-------- src/ri_jk/hess_r.rs | 4 ++-- src/ri_jk/hess_u.rs | 2 +- tests/{hessian_opbase => analdrv}/mod.rs | 0 tests/{hessian_opbase => analdrv}/rhf.rs | 4 ++-- .../{hessian_opbase => analdrv}/rks_b3lyp.rs | 4 ++-- .../rks_b3lyp_grid_levels.rs | 4 ++-- tests/{hessian_opbase => analdrv}/uhf.rs | 6 ++--- .../uks_tpss0_grid_levels.rs | 4 ++-- .../{hessian_opbase => analdrv}/uks_tpssh.rs | 4 ++-- tests/test_hessian_opbase.rs | 2 +- 42 files changed, 58 insertions(+), 51 deletions(-) rename src/{hessian_opbase => analdrv}/cint_handling.rs (100%) rename src/{hessian_opbase => analdrv}/config.rs (100%) rename src/{hessian_opbase => analdrv}/hcore.rs (100%) rename src/{hessian_opbase => analdrv}/krylov_block.rs (100%) rename src/{hessian_opbase => analdrv}/mod.rs (84%) rename src/{hessian_opbase => analdrv}/nuc_repl.rs (100%) rename src/{hessian_opbase => analdrv}/ovlp.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/bits.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/detect.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/elements.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/geom.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/interface_to_rest.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/linalg.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/matrix.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/mod.rs (95%) rename src/{hessian_opbase => analdrv}/point_group_detect/molecule.rs (100%) rename src/{hessian_opbase => analdrv}/point_group_detect/vec3.rs (100%) rename src/{hessian_opbase => analdrv}/rscf.rs (100%) rename src/{hessian_opbase => analdrv}/rscf_interface.rs (98%) rename src/{hessian_opbase => analdrv}/trait_rhess.rs (100%) rename src/{hessian_opbase => analdrv}/trait_uhess.rs (100%) rename src/{hessian_opbase => analdrv}/trait_util.rs (100%) rename src/{hessian_opbase => analdrv}/uscf.rs (100%) rename src/{hessian_opbase => analdrv}/uscf_interface.rs (98%) rename src/{hessian_opbase => analdrv}/vib.rs (100%) rename src/{hessian_opbase => analdrv}/vib_interface.rs (90%) rename tests/{hessian_opbase => analdrv}/mod.rs (100%) rename tests/{hessian_opbase => analdrv}/rhf.rs (97%) rename tests/{hessian_opbase => analdrv}/rks_b3lyp.rs (92%) rename tests/{hessian_opbase => analdrv}/rks_b3lyp_grid_levels.rs (98%) rename tests/{hessian_opbase => analdrv}/uhf.rs (91%) rename tests/{hessian_opbase => analdrv}/uks_tpss0_grid_levels.rs (95%) rename tests/{hessian_opbase => analdrv}/uks_tpssh.rs (94%) diff --git a/src/hessian_opbase/cint_handling.rs b/src/analdrv/cint_handling.rs similarity index 100% rename from src/hessian_opbase/cint_handling.rs rename to src/analdrv/cint_handling.rs diff --git a/src/hessian_opbase/config.rs b/src/analdrv/config.rs similarity index 100% rename from src/hessian_opbase/config.rs rename to src/analdrv/config.rs diff --git a/src/hessian_opbase/hcore.rs b/src/analdrv/hcore.rs similarity index 100% rename from src/hessian_opbase/hcore.rs rename to src/analdrv/hcore.rs diff --git a/src/hessian_opbase/krylov_block.rs b/src/analdrv/krylov_block.rs similarity index 100% rename from src/hessian_opbase/krylov_block.rs rename to src/analdrv/krylov_block.rs diff --git a/src/hessian_opbase/mod.rs b/src/analdrv/mod.rs similarity index 84% rename from src/hessian_opbase/mod.rs rename to src/analdrv/mod.rs index cdd9ba53ba..24436ca3b7 100644 --- a/src/hessian_opbase/mod.rs +++ b/src/analdrv/mod.rs @@ -1,11 +1,14 @@ -//! Analytical hessian module for REST (Operator-Based analytical hessian architecture). -//! -//! The name `opbase` is short for "operator-based". This concept should echo ORCA's hessian implementation, -//! though this is recently found by me (ajz34) from ORCA output. We can split hessian contributions into -//! distinct operators (hcore, nuc, J/K, DFT numint, etc), and use trait to define the interface for each operator. +//! 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 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. +//! +//! Currently only hessian property is implemented. This module may be full refactored in future to +//! support other types of derivatives and properties. //! //! This module should work in most cases, but still requires further testing and efficiency update. //! diff --git a/src/hessian_opbase/nuc_repl.rs b/src/analdrv/nuc_repl.rs similarity index 100% rename from src/hessian_opbase/nuc_repl.rs rename to src/analdrv/nuc_repl.rs diff --git a/src/hessian_opbase/ovlp.rs b/src/analdrv/ovlp.rs similarity index 100% rename from src/hessian_opbase/ovlp.rs rename to src/analdrv/ovlp.rs diff --git a/src/hessian_opbase/point_group_detect/bits.rs b/src/analdrv/point_group_detect/bits.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/bits.rs rename to src/analdrv/point_group_detect/bits.rs diff --git a/src/hessian_opbase/point_group_detect/detect.rs b/src/analdrv/point_group_detect/detect.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/detect.rs rename to src/analdrv/point_group_detect/detect.rs diff --git a/src/hessian_opbase/point_group_detect/elements.rs b/src/analdrv/point_group_detect/elements.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/elements.rs rename to src/analdrv/point_group_detect/elements.rs diff --git a/src/hessian_opbase/point_group_detect/geom.rs b/src/analdrv/point_group_detect/geom.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/geom.rs rename to src/analdrv/point_group_detect/geom.rs diff --git a/src/hessian_opbase/point_group_detect/interface_to_rest.rs b/src/analdrv/point_group_detect/interface_to_rest.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/interface_to_rest.rs rename to src/analdrv/point_group_detect/interface_to_rest.rs diff --git a/src/hessian_opbase/point_group_detect/linalg.rs b/src/analdrv/point_group_detect/linalg.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/linalg.rs rename to src/analdrv/point_group_detect/linalg.rs diff --git a/src/hessian_opbase/point_group_detect/matrix.rs b/src/analdrv/point_group_detect/matrix.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/matrix.rs rename to src/analdrv/point_group_detect/matrix.rs diff --git a/src/hessian_opbase/point_group_detect/mod.rs b/src/analdrv/point_group_detect/mod.rs similarity index 95% rename from src/hessian_opbase/point_group_detect/mod.rs rename to src/analdrv/point_group_detect/mod.rs index 79e9bcc76e..8f4e9f8452 100644 --- a/src/hessian_opbase/point_group_detect/mod.rs +++ b/src/analdrv/point_group_detect/mod.rs @@ -15,7 +15,7 @@ //! //! # Example //! ```ignore -//! use rest::hessian_opbase::point_group_detect::SymmMolecule; +//! use rest::analdrv::point_group_detect::SymmMolecule; //! //! // H2O (Bohr) -> C2v //! let mol = SymmMolecule::new( diff --git a/src/hessian_opbase/point_group_detect/molecule.rs b/src/analdrv/point_group_detect/molecule.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/molecule.rs rename to src/analdrv/point_group_detect/molecule.rs diff --git a/src/hessian_opbase/point_group_detect/vec3.rs b/src/analdrv/point_group_detect/vec3.rs similarity index 100% rename from src/hessian_opbase/point_group_detect/vec3.rs rename to src/analdrv/point_group_detect/vec3.rs diff --git a/src/hessian_opbase/rscf.rs b/src/analdrv/rscf.rs similarity index 100% rename from src/hessian_opbase/rscf.rs rename to src/analdrv/rscf.rs diff --git a/src/hessian_opbase/rscf_interface.rs b/src/analdrv/rscf_interface.rs similarity index 98% rename from src/hessian_opbase/rscf_interface.rs rename to src/analdrv/rscf_interface.rs index 859b1fe162..80c0a17b64 100644 --- a/src/hessian_opbase/rscf_interface.rs +++ b/src/analdrv/rscf_interface.rs @@ -1,6 +1,6 @@ use super::prelude::*; -use crate::hessian_opbase::vib::*; -use crate::hessian_opbase::vib_interface::*; +use crate::analdrv::vib::*; +use crate::analdrv::vib_interface::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; diff --git a/src/hessian_opbase/trait_rhess.rs b/src/analdrv/trait_rhess.rs similarity index 100% rename from src/hessian_opbase/trait_rhess.rs rename to src/analdrv/trait_rhess.rs diff --git a/src/hessian_opbase/trait_uhess.rs b/src/analdrv/trait_uhess.rs similarity index 100% rename from src/hessian_opbase/trait_uhess.rs rename to src/analdrv/trait_uhess.rs diff --git a/src/hessian_opbase/trait_util.rs b/src/analdrv/trait_util.rs similarity index 100% rename from src/hessian_opbase/trait_util.rs rename to src/analdrv/trait_util.rs diff --git a/src/hessian_opbase/uscf.rs b/src/analdrv/uscf.rs similarity index 100% rename from src/hessian_opbase/uscf.rs rename to src/analdrv/uscf.rs diff --git a/src/hessian_opbase/uscf_interface.rs b/src/analdrv/uscf_interface.rs similarity index 98% rename from src/hessian_opbase/uscf_interface.rs rename to src/analdrv/uscf_interface.rs index 4c220ab4ae..9c866767cf 100644 --- a/src/hessian_opbase/uscf_interface.rs +++ b/src/analdrv/uscf_interface.rs @@ -1,6 +1,6 @@ use super::prelude::*; -use crate::hessian_opbase::vib::*; -use crate::hessian_opbase::vib_interface::*; +use crate::analdrv::vib::*; +use crate::analdrv::vib_interface::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; diff --git a/src/hessian_opbase/vib.rs b/src/analdrv/vib.rs similarity index 100% rename from src/hessian_opbase/vib.rs rename to src/analdrv/vib.rs diff --git a/src/hessian_opbase/vib_interface.rs b/src/analdrv/vib_interface.rs similarity index 90% rename from src/hessian_opbase/vib_interface.rs rename to src/analdrv/vib_interface.rs index 0bf33ddcb8..9649b837a7 100644 --- a/src/hessian_opbase/vib_interface.rs +++ b/src/analdrv/vib_interface.rs @@ -1,6 +1,6 @@ use super::prelude::*; use crate::geom_io::get_mass_charge; -use crate::hessian_opbase::vib::*; +use crate::analdrv::vib::*; use crate::ri_jk::util::get_cint_mol; use crate::SCF; @@ -16,7 +16,7 @@ pub fn vibration_analysis_interface( let mol_obj = &scf_data.mol; let mol = get_cint_mol(mol_obj); - println!("=============== Vibrational Analysis (in hessian_opbase) ==============="); + println!("=============== Vibrational Analysis (in analdrv) ==============="); // first transpose hessian to [3*natm, 3*natm] let natm = de_hess.shape()[3]; @@ -37,13 +37,13 @@ pub fn vibration_analysis_interface( println!("{}", msg_vib); println!(""); - println!("=============== End of Vibrational Analysis (in hessian_opbase) ==============="); + println!("=============== End of Vibrational Analysis (in analdrv) ==============="); println!(""); - println!("=============== Thermo Analysis (Usual Style in hessian_opbase) ==============="); + println!("=============== Thermo Analysis (Usual Style in analdrv) ==============="); println!(""); - println!("Note: This is usual thermo analysis. To activate Shermo-style thermo analysis, please not use hessian_opbase module at this time."); + println!("Note: This is usual thermo analysis. To activate Shermo-style thermo analysis, please not use analdrv module at this time."); let geom_c = mass_centred_geom(geom_rt.view(), mass_rt.view()); let rc_cm = rotation_const(mass_rt.view(), geom_c.view(), "wavenumber").to_vec(); @@ -95,7 +95,7 @@ pub fn vibration_analysis_interface( println!("{}", msg_th); println!(""); - println!("=============== End of Thermo Analysis (Usual Style in hessian_opbase) ==============="); + println!("=============== End of Thermo Analysis (Usual Style in analdrv) ==============="); (de_hess.into_shape(-1).into_vec(), vib, th) } diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index 78ee660c89..1d929c8971 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -8,7 +8,7 @@ use std::{fs, sync::Arc}; 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::hessian_opbase::config::HessSCFConfig; +use crate::analdrv::config::HessSCFConfig; 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; @@ -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.hessian_opbase = tmp_keys.get("hessian_opbase").map(serde_from_value); + tmp_input.analdrv = tmp_keys.get("analdrv").map(serde_from_value); Ok((tmp_input,tmp_geomcell)) } @@ -388,7 +388,7 @@ pub struct InputKeywords { pub use_fxc_opt: bool, pub rel: RelativisticMethod, /// - pub hessian_opbase: Option, + pub analdrv: Option, } impl Default for InputKeywords { @@ -554,7 +554,7 @@ impl InputKeywords { thermo: None, use_fxc_opt: false, rel: RelativisticMethod::None, - hessian_opbase: None, + analdrv: None, } } @@ -1942,7 +1942,7 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result { tmp_str.to_lowercase() }, other => String::from("legacy"), }; - println!("[DEBUG]: {:?}", tmp_ctrl.get("hessian_opbase")); + println!("[DEBUG]: {:?}", tmp_ctrl.get("analdrv")); //=========================================================== // Global check of ctrl keywords and futher modification diff --git a/src/dft/mod.rs b/src/dft/mod.rs index ee79785c34..3b344b6b95 100644 --- a/src/dft/mod.rs +++ b/src/dft/mod.rs @@ -3510,7 +3510,7 @@ impl Grids { /// Build a grid using the same atom-grid machinery as [`Grids::build`], but with an explicit /// `level` override (instead of `mol.ctrl.grid_gen_level`). /// - /// This is used by the `hessian_opbase` module to generate skeleton / cphf grids at levels + /// This is used by the `analdrv` module to generate skeleton / cphf grids at levels /// that may differ from the SCF DFT grid. Unlike [`Grids::build`], this skips the /// `external_grids` path and MPI distribution, and does not apply the round-robin permutation. pub fn build_with_level(mol: &Molecule, level: usize) -> Grids { diff --git a/src/dft/numint_matmul/hess_rks.rs b/src/dft/numint_matmul/hess_rks.rs index 9187e0bad2..937593659d 100644 --- a/src/dft/numint_matmul/hess_rks.rs +++ b/src/dft/numint_matmul/hess_rks.rs @@ -1,7 +1,7 @@ // see also pyhessref/nimatmul/rks.py use super::prelude::*; -use crate::hessian_opbase::prelude::*; +use crate::analdrv::prelude::*; use XCDenType::*; diff --git a/src/dft/numint_matmul/hess_uks.rs b/src/dft/numint_matmul/hess_uks.rs index 353bab4f26..ed82425582 100644 --- a/src/dft/numint_matmul/hess_uks.rs +++ b/src/dft/numint_matmul/hess_uks.rs @@ -2,7 +2,7 @@ use super::hess_rks::{get_de_vxc_diag, get_de_vxc_off, get_drho, get_vmat_ip}; use super::prelude::*; -use crate::hessian_opbase::prelude::*; +use crate::analdrv::prelude::*; use XCDenType::*; diff --git a/src/lib.rs b/src/lib.rs index 64ded5c0fd..7c57f09dd0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -86,7 +86,7 @@ pub mod fileop; pub mod ri_cphf; pub mod lib_rint; pub mod x2c; -pub mod hessian_opbase; +pub mod analdrv; //extern crate rest; diff --git a/src/main_driver.rs b/src/main_driver.rs index 85dc681720..1c9715c739 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -385,13 +385,17 @@ pub fn main_driver() -> anyhow::Result<()> { crate::hessian::rhf_hessian_main(&scf_data, hess_ctrl, &mut time_mark); } - if let Some(ref hessian_opbase_ctrl) = scf_data.mol.ctrl.hessian_opbase { - time_mark.new_item("HessianOpBase", "analytical Hessian (hessian_opbase module)"); - time_mark.count_start("HessianOpBase"); - use crate::hessian_opbase::rscf_interface::rscf_hess_interface; - use crate::hessian_opbase::uscf_interface::uscf_hess_interface; + //=================================== + // Analytical derivative module calculations + //=================================== + if let Some(ref analdrv_ctrl) = scf_data.mol.ctrl.analdrv { + time_mark.new_item("AnalDrv", "analytical derivative module"); + time_mark.count_start("AnalDrv"); + use crate::analdrv::rscf_interface::rscf_hess_interface; + use crate::analdrv::uscf_interface::uscf_hess_interface; - eprintln!("[WARN] You are using the hessian_opbase module, which is still under development."); + eprintln!("[WARN] You are using analdrv module, which is still under development."); + eprintln!("[WARN] Keywords of analdrv will be updated in future versions."); // simple guard, but currently many methods (solvent, range-separate, dftd are not supported) if scf_data.mol.xc_data.is_fifth_dfa() { @@ -400,14 +404,14 @@ pub fn main_driver() -> anyhow::Result<()> { match scf_data.scftype { SCFType::RHF => { - rscf_hess_interface(&scf_data, hessian_opbase_ctrl); + rscf_hess_interface(&scf_data, analdrv_ctrl); } SCFType::UHF => { - uscf_hess_interface(&scf_data, hessian_opbase_ctrl); + uscf_hess_interface(&scf_data, analdrv_ctrl); }, _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types.") } - time_mark.count("HessianOpBase"); + time_mark.count("AnalDrv"); } time_mark.count("Overall"); diff --git a/src/ri_jk/hess_r.rs b/src/ri_jk/hess_r.rs index a76257f9f7..e107191ad0 100644 --- a/src/ri_jk/hess_r.rs +++ b/src/ri_jk/hess_r.rs @@ -17,8 +17,8 @@ use super::prelude_dev::*; use crate::grad::rhf::pack_triu_tilde; -use crate::hessian_opbase::cint_handling::*; -use crate::hessian_opbase::prelude::*; +use crate::analdrv::cint_handling::*; +use crate::analdrv::prelude::*; use crate::ri_jk::util::*; use crate::ri_jk::decompose::*; diff --git a/src/ri_jk/hess_u.rs b/src/ri_jk/hess_u.rs index baa0d4ac1b..7c7ed61ef7 100644 --- a/src/ri_jk/hess_u.rs +++ b/src/ri_jk/hess_u.rs @@ -22,7 +22,7 @@ //! per spin in bra form (same-spin only). use super::prelude_dev::*; -use crate::hessian_opbase::prelude::*; +use crate::analdrv::prelude::*; use crate::ri_jk::util::*; use crate::ri_jk::decompose::*; diff --git a/tests/hessian_opbase/mod.rs b/tests/analdrv/mod.rs similarity index 100% rename from tests/hessian_opbase/mod.rs rename to tests/analdrv/mod.rs diff --git a/tests/hessian_opbase/rhf.rs b/tests/analdrv/rhf.rs similarity index 97% rename from tests/hessian_opbase/rhf.rs rename to tests/analdrv/rhf.rs index ae401c6778..9d9cd76fe0 100644 --- a/tests/hessian_opbase/rhf.rs +++ b/tests/analdrv/rhf.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_opbase::config::HessSCFConfig; -use pyrest::hessian_opbase::rscf_interface::rscf_hess_interface; +use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::rscf_interface::rscf_hess_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; diff --git a/tests/hessian_opbase/rks_b3lyp.rs b/tests/analdrv/rks_b3lyp.rs similarity index 92% rename from tests/hessian_opbase/rks_b3lyp.rs rename to tests/analdrv/rks_b3lyp.rs index 612223f473..611d207075 100644 --- a/tests/hessian_opbase/rks_b3lyp.rs +++ b/tests/analdrv/rks_b3lyp.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_opbase::rscf_interface::rscf_hess_interface; -use pyrest::hessian_opbase::config::HessSCFConfig; +use pyrest::analdrv::rscf_interface::rscf_hess_interface; +use pyrest::analdrv::config::HessSCFConfig; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; diff --git a/tests/hessian_opbase/rks_b3lyp_grid_levels.rs b/tests/analdrv/rks_b3lyp_grid_levels.rs similarity index 98% rename from tests/hessian_opbase/rks_b3lyp_grid_levels.rs rename to tests/analdrv/rks_b3lyp_grid_levels.rs index 2557bfaf66..430f877b0e 100644 --- a/tests/hessian_opbase/rks_b3lyp_grid_levels.rs +++ b/tests/analdrv/rks_b3lyp_grid_levels.rs @@ -3,8 +3,8 @@ //! Exercises the dedicated CP-KS grid path (`ni_cpks = Some`) and the skeleton-grid //! regeneration path for a GGA functional. -use pyrest::hessian_opbase::config::HessSCFConfig; -use pyrest::hessian_opbase::rscf_interface::rscf_hess_interface; +use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::rscf_interface::rscf_hess_interface; use pyrest::ctrl_io; use pyrest::dft::numint_matmul::hess_rks::{get_hess_ncomp_ao_dm0, get_rho_vxc_fxc, make_cpks_vxc_fxc}; diff --git a/tests/hessian_opbase/uhf.rs b/tests/analdrv/uhf.rs similarity index 91% rename from tests/hessian_opbase/uhf.rs rename to tests/analdrv/uhf.rs index f4a03710c8..6f1908f3ca 100644 --- a/tests/hessian_opbase/uhf.rs +++ b/tests/analdrv/uhf.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_opbase::uscf_interface::uscf_hess_interface; -use pyrest::hessian_opbase::config::HessSCFConfig; +use pyrest::analdrv::uscf_interface::uscf_hess_interface; +use pyrest::analdrv::config::HessSCFConfig; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; @@ -35,7 +35,7 @@ static INPUT_NH3: &str = r##" H 0.1 0.1 1.2 """ -[hessian_opbase] +[analdrv] atm_list = [0, 1, 3] "##; diff --git a/tests/hessian_opbase/uks_tpss0_grid_levels.rs b/tests/analdrv/uks_tpss0_grid_levels.rs similarity index 95% rename from tests/hessian_opbase/uks_tpss0_grid_levels.rs rename to tests/analdrv/uks_tpss0_grid_levels.rs index 66a2935828..1e7d3f8561 100644 --- a/tests/hessian_opbase/uks_tpss0_grid_levels.rs +++ b/tests/analdrv/uks_tpss0_grid_levels.rs @@ -3,8 +3,8 @@ //! Exercises the MGGA skeleton-grid regeneration (`grid_gen_level + 2` default) and the //! dedicated CP-KS grid path on the UKS side. -use pyrest::hessian_opbase::config::HessSCFConfig; -use pyrest::hessian_opbase::uscf_interface::uscf_hess_interface; +use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::uscf_interface::uscf_hess_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; diff --git a/tests/hessian_opbase/uks_tpssh.rs b/tests/analdrv/uks_tpssh.rs similarity index 94% rename from tests/hessian_opbase/uks_tpssh.rs rename to tests/analdrv/uks_tpssh.rs index ec95d391aa..2c83a8c632 100644 --- a/tests/hessian_opbase/uks_tpssh.rs +++ b/tests/analdrv/uks_tpssh.rs @@ -1,5 +1,5 @@ -use pyrest::hessian_opbase::config::HessSCFConfig; -use pyrest::hessian_opbase::uscf_interface::uscf_hess_interface; +use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::uscf_interface::uscf_hess_interface; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; diff --git a/tests/test_hessian_opbase.rs b/tests/test_hessian_opbase.rs index 1bd61cdb5f..ab0f9246bf 100644 --- a/tests/test_hessian_opbase.rs +++ b/tests/test_hessian_opbase.rs @@ -1 +1 @@ -pub mod hessian_opbase; \ No newline at end of file +pub mod analdrv; \ No newline at end of file -- Gitee From 85a215348571c2883f15ebbe50aab2b2f2c88eea Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 11:40:08 +0800 Subject: [PATCH 32/39] analdrv: stage changes to main driver, changed configuration struct name --- src/analdrv/config.rs | 21 ++++++++++++++++--- src/analdrv/interface.rs | 28 ++++++++++++++++++++++++++ src/analdrv/mod.rs | 9 +++++---- src/analdrv/rscf.rs | 4 ++-- src/analdrv/rscf_interface.rs | 2 +- src/analdrv/uscf.rs | 4 ++-- src/analdrv/uscf_interface.rs | 2 +- src/analdrv/vib_interface.rs | 2 +- src/ctrl_io/geometric_pyo3_io.rs | 7 ++++++- src/ctrl_io/mod.rs | 4 ++-- src/main_driver.rs | 23 +++------------------ tests/analdrv/rhf.rs | 6 +++--- tests/analdrv/rks_b3lyp.rs | 4 ++-- tests/analdrv/rks_b3lyp_grid_levels.rs | 10 ++++----- tests/analdrv/uhf.rs | 4 ++-- tests/analdrv/uks_tpss0_grid_levels.rs | 8 ++++---- tests/analdrv/uks_tpssh.rs | 4 ++-- 17 files changed, 87 insertions(+), 55 deletions(-) create mode 100644 src/analdrv/interface.rs diff --git a/src/analdrv/config.rs b/src/analdrv/config.rs index 32a0a1e8a3..f7a44b7839 100644 --- a/src/analdrv/config.rs +++ b/src/analdrv/config.rs @@ -1,9 +1,24 @@ use serde::{Deserialize, Serialize}; use serde_inline_default::serde_inline_default; +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub enum AnalDrvTask { + /// Analytical Hessian matrix. + #[serde( + alias = "hessian", + alias = "hess", + alias = "frequency", + alias = "freq", + alias = "vibration", + alias = "vib", + alias = "thermo" + )] + Hessian, +} + #[serde_inline_default] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct HessSCFConfig { +pub struct AnalDrvConfig { #[serde_inline_default(0.0)] pub cphf_level_shift: f64, #[serde_inline_default(1e-8)] @@ -31,13 +46,13 @@ pub struct HessSCFConfig { #[serde_inline_default(None)] pub grid_level_skeleton: Option, /// Tolerance for point group detection in vibrational analysis. Default to 1e-5 Bohr. - /// + /// /// Note that this tolerance will be divided by sqrt(1 + natm). #[serde_inline_default(1.0e-5)] pub tol_point_group: f64, } -impl Default for HessSCFConfig { +impl Default for AnalDrvConfig { fn default() -> Self { Self { cphf_level_shift: 0.0, diff --git a/src/analdrv/interface.rs b/src/analdrv/interface.rs new file mode 100644 index 0000000000..85503315bf --- /dev/null +++ b/src/analdrv/interface.rs @@ -0,0 +1,28 @@ +//! Interface to REST other programs. + +use crate::analdrv::config::AnalDrvConfig; +use crate::analdrv::vib::{ThermoInfo, VibInfo}; +use crate::scf_io::{SCFType, SCF}; + +pub fn hess_interface(scf_data: &SCF, analdrv_ctrl: &AnalDrvConfig) -> (Vec, VibInfo, ThermoInfo) { + use crate::analdrv::rscf_interface::rscf_hess_interface; + use crate::analdrv::uscf_interface::uscf_hess_interface; + + eprintln!("[WARN] You are using analdrv module, which is still under development."); + eprintln!("[WARN] Keywords of analdrv will be updated in future versions."); + + // simple guard, but currently many methods (solvent, range-separate, dftd are not supported) + if scf_data.mol.xc_data.is_fifth_dfa() { + panic!("Normal modes calculation is currently not available for post-SCF methods."); + } + + match scf_data.scftype { + SCFType::RHF => { + rscf_hess_interface(&scf_data, analdrv_ctrl) + }, + SCFType::UHF => { + uscf_hess_interface(&scf_data, analdrv_ctrl) + }, + _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types."), + } +} diff --git a/src/analdrv/mod.rs b/src/analdrv/mod.rs index 24436ca3b7..27b52ecaa1 100644 --- a/src/analdrv/mod.rs +++ b/src/analdrv/mod.rs @@ -1,12 +1,12 @@ //! 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 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. -//! +//! //! Currently only hessian property is implemented. This module may be full refactored in future to //! support other types of derivatives and properties. //! @@ -50,6 +50,7 @@ pub mod rscf; pub mod uscf; // total hess interface to REST +pub mod interface; pub mod rscf_interface; pub mod uscf_interface; @@ -67,7 +68,7 @@ pub mod point_group_detect; pub mod prelude { use super::*; - pub use config::HessSCFConfig; + pub use config::AnalDrvConfig; pub use hcore::{RHessHcore, UHessHcore}; pub use nuc_repl::HessNucRepl; pub use ovlp::{RHessOvlp, UHessOvlp}; diff --git a/src/analdrv/rscf.rs b/src/analdrv/rscf.rs index 398d54b334..80b8d7ede7 100644 --- a/src/analdrv/rscf.rs +++ b/src/analdrv/rscf.rs @@ -11,7 +11,7 @@ pub struct RHessSCF<'a> { pub nuc_list: Vec<&'a mut dyn HessNucAPI>, pub core_list: Vec<&'a mut dyn RHessCoreAPI>, pub el_list: Vec<&'a mut dyn RHessElecInteractAPI>, - pub config: HessSCFConfig, + pub config: AnalDrvConfig, pub result: HashMap, /// Timing information. Represented by wall time in second. pub timing: Vec<(String, f64)>, @@ -27,7 +27,7 @@ impl<'a> RHessSCF<'a> { nuc_list: Vec<&'a mut dyn HessNucAPI>, core_list: Vec<&'a mut dyn RHessCoreAPI>, el_list: Vec<&'a mut dyn RHessElecInteractAPI>, - config: &HessSCFConfig, + config: &AnalDrvConfig, ) -> Self { Self { mo_coeff, diff --git a/src/analdrv/rscf_interface.rs b/src/analdrv/rscf_interface.rs index 80c0a17b64..75421c682e 100644 --- a/src/analdrv/rscf_interface.rs +++ b/src/analdrv/rscf_interface.rs @@ -4,7 +4,7 @@ use crate::analdrv::vib_interface::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; -pub fn rscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, VibInfo, ThermoInfo) { +pub fn rscf_hess_interface(scf_data: &SCF, config: &AnalDrvConfig) -> (Vec, VibInfo, ThermoInfo) { let device = DeviceBLAS::default(); // --- basic preparation --- // diff --git a/src/analdrv/uscf.rs b/src/analdrv/uscf.rs index 27e87e2843..77ccb98629 100644 --- a/src/analdrv/uscf.rs +++ b/src/analdrv/uscf.rs @@ -10,7 +10,7 @@ pub struct UHessSCF<'a> { pub nuc_list: Vec<&'a mut dyn HessNucAPI>, pub core_list: Vec<&'a mut dyn UHessCoreAPI>, pub el_list: Vec<&'a mut dyn UHessElecInteractAPI>, - pub config: HessSCFConfig, + pub config: AnalDrvConfig, pub result: HashMap, /// Timing information. Represented by wall time in second. pub timing: Vec<(String, f64)>, @@ -26,7 +26,7 @@ impl<'a> UHessSCF<'a> { nuc_list: Vec<&'a mut dyn HessNucAPI>, core_list: Vec<&'a mut dyn UHessCoreAPI>, el_list: Vec<&'a mut dyn UHessElecInteractAPI>, - config: &HessSCFConfig, + config: &AnalDrvConfig, ) -> Self { Self { mo_coeff, diff --git a/src/analdrv/uscf_interface.rs b/src/analdrv/uscf_interface.rs index 9c866767cf..539ce4c908 100644 --- a/src/analdrv/uscf_interface.rs +++ b/src/analdrv/uscf_interface.rs @@ -4,7 +4,7 @@ use crate::analdrv::vib_interface::*; use crate::ri_jk::util::{get_cint_aux, get_cint_mol}; use crate::SCF; -pub fn uscf_hess_interface(scf_data: &SCF, config: &HessSCFConfig) -> (Vec, VibInfo, ThermoInfo) { +pub fn uscf_hess_interface(scf_data: &SCF, config: &AnalDrvConfig) -> (Vec, VibInfo, ThermoInfo) { let device = DeviceBLAS::default(); // --- basic preparation --- // diff --git a/src/analdrv/vib_interface.rs b/src/analdrv/vib_interface.rs index 9649b837a7..279090ddd3 100644 --- a/src/analdrv/vib_interface.rs +++ b/src/analdrv/vib_interface.rs @@ -9,7 +9,7 @@ use crate::SCF; /// `de_hess` should be of shape `[3, 3, natm, natm]`, and is the hessian matrix in atomic units. pub fn vibration_analysis_interface( scf_data: &SCF, - config: &HessSCFConfig, + config: &AnalDrvConfig, de_hess: TsrView, ) -> (Vec, VibInfo, ThermoInfo) { let device = de_hess.device().clone(); diff --git a/src/ctrl_io/geometric_pyo3_io.rs b/src/ctrl_io/geometric_pyo3_io.rs index 0d24e90d43..eaa12de6d4 100644 --- a/src/ctrl_io/geometric_pyo3_io.rs +++ b/src/ctrl_io/geometric_pyo3_io.rs @@ -40,7 +40,7 @@ pub struct GeomeTRIC { pub verbose: i32, // 0: concise and default, 1: normal, 2: verbose, 3: low-level function output // Hessian computation pub analytic_hessian: bool, // If true, use REST analytical Hessian instead of geomeTRIC numerical finite-difference - + pub use_analdrv: bool, // Use analytical derivative module (analdrv) for Hessian computation. Default false. } impl Default for GeomeTRIC { @@ -69,6 +69,7 @@ impl Default for GeomeTRIC { prefix: "GeomeTRIC".to_string(), verbose: 0, analytic_hessian: false, + use_analdrv: false, } } } @@ -215,6 +216,10 @@ pub fn parse_geometric_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result< serde_json::Value::Bool(b) => *b, other => false, }; + geometric.use_analdrv = match tmp_ctrl.get("use_analdrv").unwrap_or(&serde_json::Value::Null) { + serde_json::Value::Bool(b) => *b, + other => false, + }; return Ok(Some(geometric)); }, other => { diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index 1d929c8971..5e9696a473 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -8,7 +8,7 @@ use std::{fs, sync::Arc}; 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::HessSCFConfig; +use crate::analdrv::config::AnalDrvConfig; 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; @@ -388,7 +388,7 @@ pub struct InputKeywords { pub use_fxc_opt: bool, pub rel: RelativisticMethod, /// - pub analdrv: Option, + pub analdrv: Option, } impl Default for InputKeywords { diff --git a/src/main_driver.rs b/src/main_driver.rs index 1c9715c739..7d4a6c51af 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -391,26 +391,8 @@ pub fn main_driver() -> anyhow::Result<()> { if let Some(ref analdrv_ctrl) = scf_data.mol.ctrl.analdrv { time_mark.new_item("AnalDrv", "analytical derivative module"); time_mark.count_start("AnalDrv"); - use crate::analdrv::rscf_interface::rscf_hess_interface; - use crate::analdrv::uscf_interface::uscf_hess_interface; - - eprintln!("[WARN] You are using analdrv module, which is still under development."); - eprintln!("[WARN] Keywords of analdrv will be updated in future versions."); - - // simple guard, but currently many methods (solvent, range-separate, dftd are not supported) - if scf_data.mol.xc_data.is_fifth_dfa() { - panic!("Normal modes calculation is currently not available for post-SCF methods."); - } - - match scf_data.scftype { - SCFType::RHF => { - rscf_hess_interface(&scf_data, analdrv_ctrl); - } - SCFType::UHF => { - uscf_hess_interface(&scf_data, analdrv_ctrl); - }, - _ => unimplemented!("Normal modes calculation is only implemented for RHF and UHF SCF types.") - } + use crate::analdrv::interface::hess_interface; + hess_interface(&scf_data, analdrv_ctrl); time_mark.count("AnalDrv"); } @@ -1013,6 +995,7 @@ mod geometric_pyo3_impl { if pl > 0 { println!(" Computing analytical Hessian for geomeTRIC initial guess..."); } + let use_analdrv = scf_data.mol.ctrl.geometric_pyo3.as_ref().map_or(false, |g| g.use_analdrv); let hess_total = crate::hessian::compute_hessian(&*scf_data) .expect("Analytical Hessian computation failed for geometry optimization"); let natm = scf_data.mol.geom.nfree; diff --git a/tests/analdrv/rhf.rs b/tests/analdrv/rhf.rs index 9d9cd76fe0..9531139086 100644 --- a/tests/analdrv/rhf.rs +++ b/tests/analdrv/rhf.rs @@ -1,4 +1,4 @@ -use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::config::AnalDrvConfig; use pyrest::analdrv::rscf_interface::rscf_hess_interface; use pyrest::ctrl_io; @@ -44,7 +44,7 @@ fn test_nh3() { let mut scf_data = scf_io::SCF::build(mol, &None); scf_without_build(&mut scf_data, &None); - let config = HessSCFConfig::default(); + let config = AnalDrvConfig::default(); let (de, vib, th) = rscf_hess_interface(&mut scf_data, &config); let ref_freqs = [1263.343780, 1367.102321, 1424.072405, 2132.997526, 2443.140863, 3517.051480]; @@ -104,7 +104,7 @@ fn test_sbh3_hbr() { let mut scf_data = scf_io::SCF::build(mol, &None); scf_without_build(&mut scf_data, &None); - let config = HessSCFConfig::default(); + let config = AnalDrvConfig::default(); let (_, vib, _) = rscf_hess_interface(&mut scf_data, &config); // reference value from gaussian 16 diff --git a/tests/analdrv/rks_b3lyp.rs b/tests/analdrv/rks_b3lyp.rs index 611d207075..d8c09a1434 100644 --- a/tests/analdrv/rks_b3lyp.rs +++ b/tests/analdrv/rks_b3lyp.rs @@ -1,5 +1,5 @@ use pyrest::analdrv::rscf_interface::rscf_hess_interface; -use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::config::AnalDrvConfig; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; @@ -44,7 +44,7 @@ fn test_nh3() { let mut scf_data = scf_io::SCF::build(mol, &None); scf_without_build(&mut scf_data, &None); - let config = HessSCFConfig::default(); + let config = AnalDrvConfig::default(); let (de, _, _) = rscf_hess_interface(&mut scf_data, &config); let natm = 4; diff --git a/tests/analdrv/rks_b3lyp_grid_levels.rs b/tests/analdrv/rks_b3lyp_grid_levels.rs index 430f877b0e..cb181924c6 100644 --- a/tests/analdrv/rks_b3lyp_grid_levels.rs +++ b/tests/analdrv/rks_b3lyp_grid_levels.rs @@ -3,7 +3,7 @@ //! Exercises the dedicated CP-KS grid path (`ni_cpks = Some`) and the skeleton-grid //! regeneration path for a GGA functional. -use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::config::AnalDrvConfig; use pyrest::analdrv::rscf_interface::rscf_hess_interface; use pyrest::ctrl_io; @@ -44,7 +44,7 @@ static INPUT_NH3: &str = r##" """ "##; -fn run_with_config(config: HessSCFConfig) -> Vec { +fn run_with_config(config: AnalDrvConfig) -> Vec { 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(); @@ -58,7 +58,7 @@ fn run_with_config(config: HessSCFConfig) -> Vec { fn test_nh3_explicit_grid_levels() { // For GGA the default skeleton level equals grid_gen_level (3); force both grids to differ // from each other and from the SCF grid so the regeneration + ni_cpks paths are exercised. - let config = HessSCFConfig { grid_level_skeleton: Some(4), grid_level_cphf: Some(2), ..Default::default() }; + let config = AnalDrvConfig { grid_level_skeleton: Some(4), grid_level_cphf: Some(2), ..Default::default() }; let de = run_with_config(config); let natm = 4; @@ -71,7 +71,7 @@ fn test_nh3_explicit_grid_levels() { fn test_nh3_default_grid_levels() { // Default config: skeleton = grid_gen_level (3), cphf = grid_gen_level.max(3) - 2 = 1. // This activates the ni_cpks path with a coarse cphf grid. - let config = HessSCFConfig::default(); + let config = AnalDrvConfig::default(); let de = run_with_config(config); let natm = 4; @@ -83,7 +83,7 @@ fn test_nh3_default_grid_levels() { #[test] fn test_nh3_cphf_equals_skeleton() { // cphf level == skeleton level (both 3) -> ni_cpks = None fast path (reuse skeleton vxc/fxc). - let config = HessSCFConfig { grid_level_skeleton: Some(3), grid_level_cphf: Some(3), ..Default::default() }; + let config = AnalDrvConfig { grid_level_skeleton: Some(3), grid_level_cphf: Some(3), ..Default::default() }; let de = run_with_config(config); let natm = 4; diff --git a/tests/analdrv/uhf.rs b/tests/analdrv/uhf.rs index 6f1908f3ca..b4ad6eb3ad 100644 --- a/tests/analdrv/uhf.rs +++ b/tests/analdrv/uhf.rs @@ -1,5 +1,5 @@ use pyrest::analdrv::uscf_interface::uscf_hess_interface; -use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::config::AnalDrvConfig; use pyrest::ctrl_io; use pyrest::molecule_io::Molecule; @@ -47,7 +47,7 @@ fn test_nh3() { let mut scf_data = scf_io::SCF::build(mol, &None); scf_without_build(&mut scf_data, &None); - let config = HessSCFConfig::default(); + let config = AnalDrvConfig::default(); let (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config); let natm = 3; diff --git a/tests/analdrv/uks_tpss0_grid_levels.rs b/tests/analdrv/uks_tpss0_grid_levels.rs index 1e7d3f8561..6468bac0c5 100644 --- a/tests/analdrv/uks_tpss0_grid_levels.rs +++ b/tests/analdrv/uks_tpss0_grid_levels.rs @@ -3,7 +3,7 @@ //! Exercises the MGGA skeleton-grid regeneration (`grid_gen_level + 2` default) and the //! dedicated CP-KS grid path on the UKS side. -use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::config::AnalDrvConfig; use pyrest::analdrv::uscf_interface::uscf_hess_interface; use pyrest::ctrl_io; @@ -42,7 +42,7 @@ static INPUT_NH3: &str = r##" """ "##; -fn run_with_config(config: HessSCFConfig) -> Vec { +fn run_with_config(config: AnalDrvConfig) -> Vec { 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(); @@ -55,7 +55,7 @@ fn run_with_config(config: HessSCFConfig) -> Vec { #[test] fn test_nh3_mgga_default_skeleton() { // Default skeleton level for MGGA = grid_gen_level + 2 = 5; cphf = 1 -> ni_cpks path. - let config = HessSCFConfig::default(); + let config = AnalDrvConfig::default(); let de = run_with_config(config); let natm = 4; @@ -66,7 +66,7 @@ fn test_nh3_mgga_default_skeleton() { #[test] fn test_nh3_mgga_explicit_grid_levels() { - let config = HessSCFConfig { + let config = AnalDrvConfig { grid_level_skeleton: Some(5), grid_level_cphf: Some(2), ..Default::default() diff --git a/tests/analdrv/uks_tpssh.rs b/tests/analdrv/uks_tpssh.rs index 2c83a8c632..79a9cc3245 100644 --- a/tests/analdrv/uks_tpssh.rs +++ b/tests/analdrv/uks_tpssh.rs @@ -1,4 +1,4 @@ -use pyrest::analdrv::config::HessSCFConfig; +use pyrest::analdrv::config::AnalDrvConfig; use pyrest::analdrv::uscf_interface::uscf_hess_interface; use pyrest::ctrl_io; @@ -46,7 +46,7 @@ fn test_nh3() { let mut scf_data = scf_io::SCF::build(mol, &None); scf_without_build(&mut scf_data, &None); - let config = HessSCFConfig::default(); + let config = AnalDrvConfig::default(); let (de, vib, _th) = uscf_hess_interface(&mut scf_data, &config); let natm = 4; -- Gitee From f5f92595917aa0d1b709dd2751524437b90d6fac Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 12:00:31 +0800 Subject: [PATCH 33/39] analdrv: add keyword `analdrv_tasks` for usual usage of property evaluation --- src/analdrv/interface.rs | 10 +++++++++- src/ctrl_io/mod.rs | 14 +++++++++++--- src/main_driver.rs | 8 +++++--- 3 files changed, 25 insertions(+), 7 deletions(-) diff --git a/src/analdrv/interface.rs b/src/analdrv/interface.rs index 85503315bf..39ae00a09d 100644 --- a/src/analdrv/interface.rs +++ b/src/analdrv/interface.rs @@ -1,9 +1,17 @@ //! Interface to REST other programs. -use crate::analdrv::config::AnalDrvConfig; +use crate::analdrv::config::{AnalDrvConfig, AnalDrvTask}; use crate::analdrv::vib::{ThermoInfo, VibInfo}; use crate::scf_io::{SCFType, SCF}; +pub fn analdrv_interface(scf_data: &SCF, tasks: &[AnalDrvTask], config: &AnalDrvConfig) { + for task in tasks { + match task { + AnalDrvTask::Hessian => hess_interface(scf_data, config), + }; + } +} + pub fn hess_interface(scf_data: &SCF, analdrv_ctrl: &AnalDrvConfig) -> (Vec, VibInfo, ThermoInfo) { use crate::analdrv::rscf_interface::rscf_hess_interface; use crate::analdrv::uscf_interface::uscf_hess_interface; diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index 5e9696a473..b2a4dc5ff2 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -8,7 +8,7 @@ use std::{fs, sync::Arc}; 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; +use crate::analdrv::config::{AnalDrvConfig, AnalDrvTask}; 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; @@ -387,8 +387,10 @@ pub struct InputKeywords { #[pyo3(get, set)] pub use_fxc_opt: bool, pub rel: RelativisticMethod, - /// + /// Analytical derivative driver configuration. pub analdrv: Option, + /// Analytical derivative tasks to perform. + pub analdrv_tasks: Vec, } impl Default for InputKeywords { @@ -555,6 +557,7 @@ impl InputKeywords { use_fxc_opt: false, rel: RelativisticMethod::None, analdrv: None, + analdrv_tasks: Vec::new(), } } @@ -1942,7 +1945,12 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result { tmp_str.to_lowercase() }, other => String::from("legacy"), }; - println!("[DEBUG]: {:?}", tmp_ctrl.get("analdrv")); + + tmp_input.analdrv_tasks = match tmp_ctrl.get("analdrv_tasks").unwrap_or(&serde_json::Value::Null) { + serde_json::Value::String(tmp_str) => vec![serde_from_value(&tmp_str.to_string().into())], + serde_json::Value::Array(tmp_arr) => tmp_arr.iter().map(|x| serde_from_value(x)).collect(), + _ => panic!("analdrv_tasks must be a string or an array of strings"), + }; //=========================================================== // Global check of ctrl keywords and futher modification diff --git a/src/main_driver.rs b/src/main_driver.rs index 7d4a6c51af..4649403461 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -388,11 +388,13 @@ pub fn main_driver() -> anyhow::Result<()> { //=================================== // Analytical derivative module calculations //=================================== - if let Some(ref analdrv_ctrl) = scf_data.mol.ctrl.analdrv { + if !scf_data.mol.ctrl.analdrv_tasks.is_empty() { time_mark.new_item("AnalDrv", "analytical derivative module"); time_mark.count_start("AnalDrv"); - use crate::analdrv::interface::hess_interface; - hess_interface(&scf_data, analdrv_ctrl); + use crate::analdrv::interface::analdrv_interface; + let tasks = &scf_data.mol.ctrl.analdrv_tasks; + let config = scf_data.mol.ctrl.analdrv.clone().unwrap_or_default(); + analdrv_interface(&scf_data, tasks, &config); time_mark.count("AnalDrv"); } -- Gitee From 7a8d64bc4a5f17cadce47d97147ab1c849c69491 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 13:14:09 +0800 Subject: [PATCH 34/39] analdrv: merge existing geometric process --- src/ctrl_io/mod.rs | 1 + src/main_driver.rs | 29 ++++++++++++++++++++--------- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/src/ctrl_io/mod.rs b/src/ctrl_io/mod.rs index b2a4dc5ff2..ade1fbe282 100644 --- a/src/ctrl_io/mod.rs +++ b/src/ctrl_io/mod.rs @@ -1947,6 +1947,7 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result vec![], serde_json::Value::String(tmp_str) => vec![serde_from_value(&tmp_str.to_string().into())], serde_json::Value::Array(tmp_arr) => tmp_arr.iter().map(|x| serde_from_value(x)).collect(), _ => panic!("analdrv_tasks must be a string or an array of strings"), diff --git a/src/main_driver.rs b/src/main_driver.rs index 4649403461..4b282e0fa8 100644 --- a/src/main_driver.rs +++ b/src/main_driver.rs @@ -983,23 +983,34 @@ mod geometric_pyo3_impl { //"#; //let mut params = tomlstr2py(optimizer_params)?; - let mut params = if let Some(params) = &scf_data.mol.ctrl.geometric_pyo3 { - let params = toml2py(¶ms.to_toml())?; - params - } else { - panic!("For geometric_pyo3, you must specify the parameters in the control file.") - }; + let rust_params = scf_data.mol.ctrl.geometric_pyo3.as_ref().expect("For geometric_pyo3, you must specify the parameters in the control file."); + let params = toml2py(&rust_params.to_toml())?; // ── Compute analytical Hessian before optimization (if requested) ── let hessian_analytic_path: Option = - if scf_data.mol.ctrl.geometric_pyo3.as_ref().map_or(false, |g| g.analytic_hessian) { + if rust_params.analytic_hessian { let pl = scf_data.mol.ctrl.print_level; if pl > 0 { println!(" Computing analytical Hessian for geomeTRIC initial guess..."); } let use_analdrv = scf_data.mol.ctrl.geometric_pyo3.as_ref().map_or(false, |g| g.use_analdrv); - let hess_total = crate::hessian::compute_hessian(&*scf_data) - .expect("Analytical Hessian computation failed for geometry optimization"); + let hess_total = if !use_analdrv { + crate::hessian::compute_hessian(&*scf_data) + .expect("Analytical Hessian computation failed for geometry optimization") + } else { + use crate::analdrv::interface::hess_interface; + use rstsr::prelude::*; + + let config = scf_data.mol.ctrl.analdrv.clone().unwrap_or_default(); + let (hess_raw, _, _) = hess_interface(&scf_data, &config); + let natm = (hess_raw.len() / 9).isqrt(); + assert!(natm * natm * 9 == hess_raw.len(), "Hessian raw data length does not match expected size for {} atoms", natm); + let device = DeviceBLAS::default(); + let hess_rt = rt::asarray((&hess_raw, [3, 3, natm, natm].f(), &device)); + let hess_rt = hess_rt.transpose([0, 2, 1, 3]).into_shape((3 * natm, 3 * natm)); + let hess_vec = hess_rt.into_shape(-1).into_vec(); + MatrixFull::from_vec([3 * natm, 3 * natm], hess_vec).unwrap() + }; let natm = scf_data.mol.geom.nfree; let n3 = natm * 3; let mut out = String::new(); -- Gitee From acc41a9f375d9b55336412e90a888823db6f5fd5 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 15:06:58 +0800 Subject: [PATCH 35/39] analdrv: main readme update --- README.md | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/README.md b/README.md index 7cb1d5795b..874c2c8ff7 100644 --- a/README.md +++ b/README.md @@ -673,6 +673,59 @@ spin = 1 hessian = { solver = "krylov", frequencies = true, verbose = 2 } ``` +## 解析梯度性质模块 `analdrv` 计算相关设置 + +解析梯度模块 `analdrv` 模块是实验性质模块。目前实现了 Hessian (原子核坐标二阶梯度) 功能。 +它实现了不同于 `hessian` 模块的解析 Hessian 计算。目前该模块的 Hessian 功能支持 RHF/RKS/UHF/UKS 方法。对于 DFT，支持 LDA/GGA/mGGA 以及其对应的杂化泛函。该模块的程序有性能优化，与目前顶级的量化程序 (ORCA 等) 有相当或更好的性能。 + +### 设置待计算性质的任务 + +目前仅支持 Hessian 计算。需要在 `[ctrl]` 区块中设置 `analdrv_tasks` 关键词以启动对应性质的计算。 +```toml +[ctrl] +analdrv_tasks = "freq" +``` + +### 解析梯度模块 `analdrv` 模块选项 + +在设置任务后，用户可以在 `[analdrv]` 区块中设置对应的计算选项。该区块的关键词包括： +- `cphf_level_shift`：CPHF 求解时对 $\varepsilon_i - \varepsilon_a$ 的求解偏移。默认为 0，单位 Hartree。 +- `cphf_tol`：CPHF 中的 Krylov 求解阈值。默认 1e-8，无量纲。实际求解阈值也受制于 `cphf_lindep`，且 `cphf_lindep` 经常是更宽松的阈值。 +- `cphf_max_cycle`：CPHF 最大迭代步数。默认为 42 步。CPHF 与 SCF 不同，一般 6-10 步能收敛。这里的最大步数一般不需要设得很大。 +- `cphf_max_space`：CPHF 中 Krylov 空间的数量。默认为 14。该数值不宜设太小，因为超过该数值时，Krylov 求解器会代入最后一次迭代重新作为初猜，重置求解过程。但该数值设太大会对内存产生压力。 +- `cphf_lindep`：CPHF 中一些数值过程的数值精度阈值。默认 1e-14，无量纲。 +- `verbose`：打印强度。默认为 None，使用输入卡 `[ctrl]` 区块的 verbose。 +- `atm_list`：选择一部分原子进行 Hessian 计算。默认为 None，即所有原子参与 Hessian 计算。 +- `grid_level_cphf`：CPHF 的 DFT 格点级别。仅影响 numint_matmul 后端实现。默认为 None，是 `[ctrl]` 中 grid_generation_level 关键词设定值减 2 (SCF 默认格点级别是 3，对应 Hessian 的级别是 1)；最低级别是 1。 +- `grid_level_skeleton`：Skeleton 导数 (包括 2 阶 Hessian 贡献、1 阶 Fock 贡献) 的 DFT 格点级别。仅影响 numint_matmul 后端实现。默认为 None： + - LDA/GGA 使用与 SCF 同样的格点； + - mGGA 将比 `grid_generation_level` 增加 2 级别。 +- `tol_point_group`：振动分析中的点群对称性判断阈值 (用于计算转动对称性，对熵矫正有贡献)。默认 1e-5，单位 Bohr / sqrt(atom)。 + +作为例子，运行 Hessian 计算、增大 CP-HF Krylov 求解器空间到 20、强制 CP-HF 中 DFT 格点积分级别为 2，所需要引入的、相比于能量计算的额外设置如下： +```toml +[ctrl] +analdrv_tasks = "freq" + +[analdrv] +cphf_max_space = 20 +grid_level_cphf = 2 +``` + +### 与 `[thermo]` 模块的交互 + +若使用 `analdrv` 模块进行 Hessian 计算，会自动运行频率与热矫正分析。 +该模块目前维护自己的频率分析与热矫正计算逻辑，独立于 `thermo` 模块 (参考程序 Psi4)。但我们接受输入卡中 `[thermo]` 区块的关键词 `temperature`, `pressure`, `symmetry_number` 与 `electronic_energy` 关键词，这些关键词会在热矫正计算中使用。 + +### 与 `[geometric_pyo3]` 模块的交互 + +使用下述样例，以 `analdrv` 模块 (替换 `hessian` 模块) 进行 geometric-pyo3 所需要的 Hessian 计算： +```toml +[geometric_pyo3] +analytic_hessian = true +use_analdrv = true +``` + # Detailed description of [geom] block in the control file - `name`：取值String类型。分子体系的名称 - `unit`：取值String类型。坐标单位。目前支持：angstrom和bohr @@ -902,6 +955,7 @@ REST 提供两条独立的频率/热化学计算路径，请勿混淆： - "stop"：不做构型优化，只计算初始结构的Hessian矩阵 - "each"：计算构型优化中每一步的Hessian矩阵 - `analytic_hessian`：取值 bool。若设为 `true`，则在调用 geomeTRIC 优化前先用 REST 的解析 Hessian 模块计算结果并注入 geomeTRIC，**完全避免 geomeTRIC 的数值有限差分 Hessian 计算**。解析 Hessian 含 CP-HF 轨道弛豫贡献，精度远优于有限差分，且后续 BFGS 更新不受影响。缺省值：`false`。建议在过渡态搜索（`transition = true`）或 IRC 追踪（`irc = true`）中启用。 +- `use_analdrv`：取值 bool。若设为 `true`，使用 `analdrv` 模块进行 Hessian 计算；若设为 `false`，使用 `hessian` 模块进行 Hessian 计算。缺省值：`false`。 - `frequency`：取值bool，当得到Hessian矩阵后，是否开展频率计算和热化学分析。缺省值：true - `thermo`：取值[f64;2]，提供热力学分析的状态：[温度 (K),压强 (bar)]。缺省值：[300.0, 1.0] - `reset`：取值bool。当近似 Hessian 的特征值低于 `epsilon` 阈值时，是否将其重置回 guess Hessian。对于稳态优化，缺省值为 true。若体系梯度含噪声、BFGS 更新每步失败（出现 "Eigenvalues below ... returning guess"），可设为 false 保留 Hessian 并加对角 shift 继续优化。 -- Gitee From 384a706d88d550b0791eb134a3028a40c4cbc7c4 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 15:22:11 +0800 Subject: [PATCH 36/39] fix trait import in scf_io --- src/scf_io/mod.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/scf_io/mod.rs b/src/scf_io/mod.rs index 5b97d7f41f..faeeeba761 100644 --- a/src/scf_io/mod.rs +++ b/src/scf_io/mod.rs @@ -47,6 +47,9 @@ use self::util::norm; use smear::apply_smearing; use smear::annealed_sigma; +#[allow(unused_imports)] +use tensors::BasicMatUp; + #[pyclass] #[derive(Clone)] pub struct SCF { -- Gitee From de2e4f7fd9d0511318c38e9d4b4abc540fa4f89e Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 15:30:17 +0800 Subject: [PATCH 37/39] Cargo.toml: change mpi version requirement from "0.8.0" to "0.8" (allow auto patch update). --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index df4fc6f31e..66db77692d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,7 +44,7 @@ fuzzy-matcher = "0.3.7" #rust_libecpint = {path="../rust_libecpint"} # seems to pull in libssl and libcrypto which is rejected by manylinux reqwest = { version = "0.10", default-features = false, features = ["blocking", "json", "rustls-tls"] } -mpi = {version = "0.8.0", features = ["user-operations", "derive"], optional = true} +mpi = {version = "0.8", features = ["user-operations", "derive"], optional = true} # for cuda #cudarc = "0.10.0" rest_tensors = {path="../rest_tensors"} -- Gitee From 51db390cc9d655efe713bb968d6508c7a2d2fa92 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 15:38:27 +0800 Subject: [PATCH 38/39] fix import clash from PR!177 --- src/scf_io/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/scf_io/mod.rs b/src/scf_io/mod.rs index d3b0f10a60..e8ae8e6251 100644 --- a/src/scf_io/mod.rs +++ b/src/scf_io/mod.rs @@ -25,7 +25,7 @@ use mpi::collective::SystemOperation; use pyo3::{pyclass}; use tensors::matrix_blas_lapack::{_dgemm, _dgemm_full, _dgemv, _dspgvx, _dsymm, _dsyrk, _hamiltonian_fast_solver, _power_rayon_for_symmetric_matrix, _dsyevd}; use tensors::{map_upper_to_full, BasicMatrix, ERIFold4, MathMatrix, MatrixFull, MatrixFullSlice, MatrixUpper, MatrixUpperSlice, RIFull, TensorSliceMut}; -use tensors::{TensorOpt, TensorSlice, BasicMatUp}; +use tensors::{TensorOpt, TensorSlice}; use itertools::{Itertools}; use rayon::prelude::*; use std::collections::HashMap; -- Gitee From a610cd58214fb3c22621fab99adf87aea6735948 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Thu, 9 Jul 2026 16:05:52 +0800 Subject: [PATCH 39/39] restructured some constant definitions --- src/constants/mod.rs | 34 ++++++++++++++++------------------ 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/src/constants/mod.rs b/src/constants/mod.rs index baf17359b6..739383ee96 100644 --- a/src/constants/mod.rs +++ b/src/constants/mod.rs @@ -80,11 +80,9 @@ pub const ENV_PRT_START: usize = 20; // math, physics // NOTE: these constants come from several different CODATA releases (see per-line sources). pub const EV: f64 = 27.2113845; // Hartree energy in eV, CODATA 2002 -pub const HARTREE2KCAL: f64 = 627.509451; // Hartree -> kcal/mol; CODATA 2002 (Hartree energy in J / (kcal * Avogadro)) -pub const HARTREE2WAVENUMBER: f64 = 219474.63; // Hartree -> cm^-1 (hartree-inverse meter relationship / 100); CODATA (~2.194746313705e7 m^-1) pub const FQ: f64 = 1822.8884861920776; // u/m_e ratio (= 1 / electron mass in u 5.48579909070e-4), CODATA 2014 -pub const E: f64 = std::f64::consts::E; // math constant (std); year-insensitive -pub const PI: f64 = std::f64::consts::PI; // math constant (std); year-insensitive +pub const E: f64 = std::f64::consts::E; // math constant (std); year-insensitive +pub const PI: f64 = std::f64::consts::PI; // math constant (std); year-insensitive // pub const LIGHT_SPEED: f64 = 137.03599967994; // inverse fine-structure constant, CODATA 2006; http://physics.nist.gov/cgi-bin/cuu/Value?alph // BOHR = .529 177 210 92(17) e-10m // http://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0 @@ -92,28 +90,28 @@ pub const LIGHT_SPEED: f64 = 137.03599967994; // inverse fine-structure consta pub const BOHR: f64 = 0.52917721092; // Angstroms pub const BOHR_SI: f64 = BOHR * 1e-10; -pub const G_ELECTRON: f64 = 2.00231930436182; // CODATA 2014; http://physics.nist.gov/cgi-bin/cuu/Value?gem +pub const G_ELECTRON: f64 = 2.00231930436182; // CODATA 2014; http://physics.nist.gov/cgi-bin/cuu/Value?gem pub const E_MASS: f64 = 9.10938356e-31; // kg, CODATA 2014; https://physics.nist.gov/cgi-bin/cuu/Value?me pub const AVOGADRO: f64 = 6.022140857e23; // CODATA 2014; https://physics.nist.gov/cgi-bin/cuu/Value?na pub const PLANCK: f64 = 6.626070040e-34; // J*s, CODATA 2014; http://physics.nist.gov/cgi-bin/cuu/Value?h pub const BOLTZMANN: f64 = 1.380649e-23; // J/K, CODATA 2018 (exact, SI 2019); https://physics.nist.gov/cgi-bin/cuu/Value?k pub const CLIGHT_CMS: f64 = 2.99792458e10; // speed of light, cm/s; exact by SI definition (year-insensitive) pub const R_GAS: f64 = BOLTZMANN * AVOGADRO; // J/(mol*K) ideal gas constant; derived -pub const E_CHARGE: f64 = 1.6021766208e-19; // C, CODATA 2014 -pub const DEBYE:f64 = 3.335641e-30; // C*m = 1e-18/LIGHT_SPEED_SI; defined via c, year-insensitive; https://cccbdb.nist.gov/debye.asp -pub const AU2DEBYE:f64 = E_CHARGE * BOHR*1e-10 / DEBYE; // 2.541746; derived - -// CODATA 2022 -// https://docs.scipy.org/doc/scipy/reference/constants.html -pub const SPEED_OF_LIGHT: f64 = 299792458.0; // m s^-1 +pub const E_CHARGE: f64 = 1.6021766208e-19; // C, CODATA 2014 +pub const DEBYE: f64 = 3.335641e-30; // C*m = 1e-18/LIGHT_SPEED_SI; defined via c, year-insensitive; https://cccbdb.nist.gov/debye.asp +pub const SPEED_OF_LIGHT: f64 = 299792458.0; // CODATA 2022, m s^-1, https://physics.nist.gov/cgi-bin/cuu/Value?c // =========== unit conversion =================================== -pub const BOHR2ANG: f64 = BOHR; -pub const HARTREE2J: f64 = 4.35974465e-18; -pub const HARTREE2KJMOL: f64 = 2625.4996382852164; -pub const HARTREE2KCALMOL: f64 = 627.5094737775374; -pub const HARTREE2WAVENUMBERS: f64 = 219474.6313702; -pub const AMU2KG: f64 = 1.66053904e-27; +pub const CALORIE2J: f64 = 4.184; // calorie -> joule, https://en.wikipedia.org/wiki/Calorie +pub const HARTREE2WAVENUMBER: f64 = 219474.63; // hartree -> cm^-1 (hartree-inverse meter relationship / 100); CODATA (~2.194746313705e7 m^-1) +pub const BOHR2ANG: f64 = BOHR; // bohr -> Angstrom, by definition +pub const AU2DEBYE:f64 = E_CHARGE * BOHR*1e-10 / DEBYE; // AU -> Debye, derived, 2.541746 +pub const HARTREE2J: f64 = 4.359744722206e-18; // hartree -> joule, CODATA 2022, https://physics.nist.gov/cgi-bin/cuu/Value?Ahr, atomic unit of energy +pub const HARTREE2KJMOL: f64 = HARTREE2J * AVOGADRO / 1000.0; // hartree -> kJ/mol, derived, 2625.499681768687 +pub const HARTREE2KCALMOL: f64 = HARTREE2KJMOL / CALORIE2J; // hartree -> kcal/mol, derived, 627.5094841703362 +pub const HARTREE2KCAL: f64 = HARTREE2KCALMOL; // hartree -> kcal/mol, derived, 627.5094841703362 +pub const HARTREE2WAVENUMBERS: f64 = 219474.63136314; // hartree -> cm^-1, CODATA 2022, https://physics.nist.gov/cgi-bin/cuu/Value?hrminv, hartree-inverse meter relationship +pub const AMU2KG: f64 = 1.66053906892e-27; // amu -> kg, CODATA 2022, https://physics.nist.gov/cgi-bin/cuu/Value?ukg, unified atomic mass unit pub const MPI_CHUNK:usize = 134217728; // around 1 GB -- Gitee