diff --git a/Cargo.toml b/Cargo.toml index df4fc6f31e80f79f0529f925e3d675fb0059902f..e0d3911c0c70b2e6a1bbbe43f96af16f26eb0488 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" @@ -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"} @@ -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/README.md b/README.md index 7cb1d5795bced623d74e6cc2b362f37f9bb8f0fc..874c2c8ff77a97d096c91ad3ee13d67344b88920 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 继续优化。 diff --git a/src/analdrv/cint_handling.rs b/src/analdrv/cint_handling.rs new file mode 100644 index 0000000000000000000000000000000000000000..a5a3605d86de60b7fe133e5816abefd1d8dd3190 --- /dev/null +++ b/src/analdrv/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/analdrv/config.rs b/src/analdrv/config.rs new file mode 100644 index 0000000000000000000000000000000000000000..f7a44b783960109eb797ab1217ca5b42c42d76cf --- /dev/null +++ b/src/analdrv/config.rs @@ -0,0 +1,70 @@ +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 AnalDrvConfig { + #[serde_inline_default(0.0)] + pub cphf_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>, + /// 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, + /// 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 AnalDrvConfig { + fn default() -> Self { + Self { + cphf_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, + grid_level_cphf: None, + grid_level_skeleton: None, + tol_point_group: 1.0e-5, + } + } +} diff --git a/src/analdrv/hcore.rs b/src/analdrv/hcore.rs new file mode 100644 index 0000000000000000000000000000000000000000..a60b1842c2f10b1bc052f33832f04524f8db121f --- /dev/null +++ b/src/analdrv/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/analdrv/interface.rs b/src/analdrv/interface.rs new file mode 100644 index 0000000000000000000000000000000000000000..39ae00a09d2cc8e2896c9135474c5d4eb7086d7d --- /dev/null +++ b/src/analdrv/interface.rs @@ -0,0 +1,36 @@ +//! Interface to REST other programs. + +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; + + 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/krylov_block.rs b/src/analdrv/krylov_block.rs new file mode 100644 index 0000000000000000000000000000000000000000..104eebd18489ce0b2d698ae1fc5682218ff68afd --- /dev/null +++ b/src/analdrv/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!( + "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, + ); + + 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/analdrv/mod.rs b/src/analdrv/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..27b52ecaa1107c094299ad07522a0947b9f9caa9 --- /dev/null +++ b/src/analdrv/mod.rs @@ -0,0 +1,90 @@ +//! 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. +//! +//! 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 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. + +#![warn(unused)] + +// 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; + +// total hess interface to REST +pub mod interface; +pub mod rscf_interface; +pub mod uscf_interface; + +// vibrational analysis +pub mod vib; +pub mod vib_interface; + +// utilities +pub mod cint_handling; +pub mod config; +pub mod krylov_block; +pub mod point_group_detect; + +#[allow(unused_imports)] +pub mod prelude { + use super::*; + + pub use config::AnalDrvConfig; + pub use hcore::{RHessHcore, UHessHcore}; + pub use nuc_repl::HessNucRepl; + pub use ovlp::{RHessOvlp, UHessOvlp}; + pub use rscf::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::*; + 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::*; + pub(super) use std::collections::HashMap; +} diff --git a/src/analdrv/nuc_repl.rs b/src/analdrv/nuc_repl.rs new file mode 100644 index 0000000000000000000000000000000000000000..08c2dfca333f34408546e7bec59e9182e5519cf0 --- /dev/null +++ b/src/analdrv/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/analdrv/ovlp.rs b/src/analdrv/ovlp.rs new file mode 100644 index 0000000000000000000000000000000000000000..fa272270dcb091fd48271225e821529c354360c0 --- /dev/null +++ b/src/analdrv/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/analdrv/point_group_detect/bits.rs b/src/analdrv/point_group_detect/bits.rs new file mode 100644 index 0000000000000000000000000000000000000000..e36c8dca4ba8e87bf62a1ca3b8f13367990b7403 --- /dev/null +++ b/src/analdrv/point_group_detect/bits.rs @@ -0,0 +1,89 @@ +//! 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 + } +} + +#[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"); + } +} diff --git a/src/analdrv/point_group_detect/detect.rs b/src/analdrv/point_group_detect/detect.rs new file mode 100644 index 0000000000000000000000000000000000000000..b010ff486f75797d2bfc18687b8ea532e45209dc --- /dev/null +++ b/src/analdrv/point_group_detect/detect.rs @@ -0,0 +1,23 @@ +//! Public detection API. + +use super::molecule::SymmMolecule; + +/// 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, +} + +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), + } + } +} diff --git a/src/analdrv/point_group_detect/elements.rs b/src/analdrv/point_group_detect/elements.rs new file mode 100644 index 0000000000000000000000000000000000000000..ecea99e7f7a1133edb950d306ff08f6f532899de --- /dev/null +++ b/src/analdrv/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/analdrv/point_group_detect/geom.rs b/src/analdrv/point_group_detect/geom.rs new file mode 100644 index 0000000000000000000000000000000000000000..8a2883c4e57b58b95de7298d010f19b07a24407f --- /dev/null +++ b/src/analdrv/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/analdrv/point_group_detect/interface_to_rest.rs b/src/analdrv/point_group_detect/interface_to_rest.rs new file mode 100644 index 0000000000000000000000000000000000000000..1140ae8ded77946b23cfc2bb64b66dac9116a360 --- /dev/null +++ b/src/analdrv/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/analdrv/point_group_detect/linalg.rs b/src/analdrv/point_group_detect/linalg.rs new file mode 100644 index 0000000000000000000000000000000000000000..f34ee4e51443a1ad21a821bb903a4805a0e83516 --- /dev/null +++ b/src/analdrv/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/analdrv/point_group_detect/matrix.rs b/src/analdrv/point_group_detect/matrix.rs new file mode 100644 index 0000000000000000000000000000000000000000..e67f7329239cd3b80bfb9967fa19acb0bfbfd6b8 --- /dev/null +++ b/src/analdrv/point_group_detect/matrix.rs @@ -0,0 +1,149 @@ +//! 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 +} + +/// 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/analdrv/point_group_detect/mod.rs b/src/analdrv/point_group_detect/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..8f4e9f84526a2dc760917be8062bbf0dab9afcb1 --- /dev/null +++ b/src/analdrv/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 +//! ```ignore +//! use rest::analdrv::point_group_detect::SymmMolecule; +//! +//! // 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. + +pub mod bits; +pub mod detect; +pub mod elements; +pub mod geom; +pub mod linalg; +pub mod matrix; +pub mod molecule; +pub mod vec3; + +pub mod interface_to_rest; + +pub use detect::PointGroup; +pub use molecule::SymmMolecule; diff --git a/src/analdrv/point_group_detect/molecule.rs b/src/analdrv/point_group_detect/molecule.rs new file mode 100644 index 0000000000000000000000000000000000000000..5d80dbc0c154c350190547091db17292ad8aa2dd --- /dev/null +++ b/src/analdrv/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/analdrv/point_group_detect/vec3.rs b/src/analdrv/point_group_detect/vec3.rs new file mode 100644 index 0000000000000000000000000000000000000000..d1dbc6fd31baa105da5a9d41929d41692aee05f5 --- /dev/null +++ b/src/analdrv/point_group_detect/vec3.rs @@ -0,0 +1,150 @@ +//! 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 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); + } +} diff --git a/src/analdrv/rscf.rs b/src/analdrv/rscf.rs new file mode 100644 index 0000000000000000000000000000000000000000..80b8d7ede77405cad1b60abc097060c63c52672e --- /dev/null +++ b/src/analdrv/rscf.rs @@ -0,0 +1,488 @@ +//! Hessian implementations for restricted SCF. + +use super::prelude::*; + +/// 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: &'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: AnalDrvConfig, + 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: &'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: &AnalDrvConfig, + ) -> Self { + Self { + mo_coeff, + mo_occ, + mo_energy, + ovlp_obj, + nuc_list, + core_list, + el_list, + config: config.clone(), + 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.config.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.config.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.cphf_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.config.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.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(); + 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.config.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.config.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/analdrv/rscf_interface.rs b/src/analdrv/rscf_interface.rs new file mode 100644 index 0000000000000000000000000000000000000000..75421c682e142f171179d3f9c055ee1548924748 --- /dev/null +++ b/src/analdrv/rscf_interface.rs @@ -0,0 +1,161 @@ +use super::prelude::*; +use crate::analdrv::vib::*; +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: &AnalDrvConfig) -> (Vec, VibInfo, ThermoInfo) { + 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::hess_rks::RHessKSNIMatmul; + use crate::dft::numint_matmul::nimatmul::NIMatmul; + use crate::dft::xceff::prelude::{determine_den_type_from_list, XCDenType}; + use crate::dft::Grids; + use libxc::prelude::*; + + let xc_func_list = { + let xc_code = &scf_data.mol.xc_data.dfa_compnt_scf; + let xc_params = &scf_data.mol.xc_data.dfa_paramr_scf; + xc_code + .iter() + .zip(xc_params.iter()) + .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, LibXCSpin::Unpolarized))) + .collect_vec() + }; + let verbose = scf_data.mol.ctrl.print_level >= 2; + + // 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); + } + + // --- 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(); + + if scf_data.mol.ctrl.print_level >= 2 { + println!("=== HESSIAN ==="); + println!("Print hessian in [tA, sB] format (component xyz first, atom then)"); + println!(""); + // print hessian matrix [t, s, A, B] -> [tA, sB] + let natm = de_hess.shape()[3]; + let hess_mat = de_hess.transpose([0, 2, 1, 3]).into_shape((3 * natm, 3 * natm)); + // print 6 columns at a time, with index header + for j in (0..3 * natm).step_by(6) { + let j_end = (j + 6).min(3 * natm); + let col_header = " ".repeat(6) + &(j..j_end).map(|i| format!("{:>12}", i)).collect::(); + println!("{}", col_header); + for i in 0..3 * natm { + let row_str = + format!("{i:>4} ") + &(j..j_end).map(|j| format!("{:12.6}", hess_mat[[i, j]])).collect::(); + println!("{}", row_str); + } + println!(""); + } + } + + // print timing information + if scf_data.mol.ctrl.print_level >= 2 { + println!("Timing in Hessian calculation:"); + for (key, value) in hess_scf.timing.iter() { + println!(" {:60}: {:10.6} seconds", key, value); + } + } + + // --- perform vibrational analysis --- // + + vibration_analysis_interface(scf_data, config, de_hess.view()) +} diff --git a/src/analdrv/trait_rhess.rs b/src/analdrv/trait_rhess.rs new file mode 100644 index 0000000000000000000000000000000000000000..7326d4da77514da2838d7b92771411a7e93ac1f0 --- /dev/null +++ b/src/analdrv/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/analdrv/trait_uhess.rs b/src/analdrv/trait_uhess.rs new file mode 100644 index 0000000000000000000000000000000000000000..5fa31cd9e2b3ea726b9aeabb17cbcd1a9ff36752 --- /dev/null +++ b/src/analdrv/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/analdrv/trait_util.rs b/src/analdrv/trait_util.rs new file mode 100644 index 0000000000000000000000000000000000000000..2948f7ece5d844cc0e877a43aeef3145adbeea36 --- /dev/null +++ b/src/analdrv/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/analdrv/uscf.rs b/src/analdrv/uscf.rs new file mode 100644 index 0000000000000000000000000000000000000000..77ccb98629a5a44f16ada03a2c8796c2a0d200a7 --- /dev/null +++ b/src/analdrv/uscf.rs @@ -0,0 +1,537 @@ +//! 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: &'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>, + pub config: AnalDrvConfig, + 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: &'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: &AnalDrvConfig, + ) -> Self { + Self { + mo_coeff, + mo_occ, + mo_energy, + ovlp_obj, + nuc_list, + core_list, + el_list, + config: config.clone(), + 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.config.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.config.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.cphf_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.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]; + + // --- 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.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[β])]; + + 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.config.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.config.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/analdrv/uscf_interface.rs b/src/analdrv/uscf_interface.rs new file mode 100644 index 0000000000000000000000000000000000000000..539ce4c908715744645ddae7d9d24e21bbb250ec --- /dev/null +++ b/src/analdrv/uscf_interface.rs @@ -0,0 +1,170 @@ +use super::prelude::*; +use crate::analdrv::vib::*; +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: &AnalDrvConfig) -> (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::xceff::prelude::{determine_den_type_from_list, XCDenType}; + use crate::dft::Grids; + use libxc::prelude::*; + + let xc_func_list = { + let xc_code = &scf_data.mol.xc_data.dfa_compnt_scf; + let xc_params = &scf_data.mol.xc_data.dfa_paramr_scf; + xc_code + .iter() + .zip(xc_params.iter()) + .map(|(&code, ¶m)| (param, LibXCFunctional::from_number(code as _, LibXCSpin::Polarized))) + .collect_vec() + }; + let verbose = scf_data.mol.ctrl.print_level >= 2; + + // Determine skeleton / 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); + } + + // --- 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(); + + if scf_data.mol.ctrl.print_level >= 2 { + println!("=== HESSIAN ==="); + println!("Print hessian in [tA, sB] format (component xyz first, atom then)"); + println!(""); + // print hessian matrix [t, s, A, B] -> [tA, sB] + let natm = de_hess.shape()[3]; + let hess_mat = de_hess.transpose([0, 2, 1, 3]).into_shape((3 * natm, 3 * natm)); + // print 6 columns at a time, with index header + for j in (0..3 * natm).step_by(6) { + let j_end = (j + 6).min(3 * natm); + let col_header = " ".repeat(6) + &(j..j_end).map(|i| format!("{:>12}", i)).collect::(); + println!("{}", col_header); + for i in 0..3 * natm { + let row_str = + format!("{i:>4} ") + &(j..j_end).map(|j| format!("{:12.6}", hess_mat[[i, j]])).collect::(); + println!("{}", row_str); + } + println!(""); + } + } + + // print timing information + if scf_data.mol.ctrl.print_level >= 2 { + println!("Timing in Hessian calculation:"); + for (key, value) in hess_scf.timing.iter() { + println!(" {:60}: {:10.6} seconds", key, value); + } + } + + // --- perform vibrational analysis --- // + + vibration_analysis_interface(scf_data, config, de_hess.view()) +} diff --git a/src/analdrv/vib.rs b/src/analdrv/vib.rs new file mode 100644 index 0000000000000000000000000000000000000000..75fc6a5200571f416b117b73ab13faf37f60530e --- /dev/null +++ b/src/analdrv/vib.rs @@ -0,0 +1,1168 @@ +//! 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). +//! Some modifications are from PySCF (pyscf/hessian/thermo.py). +//! This file contains AI assisted code, and not fully reviewed by human. + +use super::prelude::*; + +// --------------------------------------------------------------------------- +// Physical constants (CODATA2014, matching the Python implementation) +// --------------------------------------------------------------------------- + +use crate::constants::{ + AMU2KG, AVOGADRO, BOHR2ANG, BOLTZMANN, HARTREE2J, HARTREE2KCALMOL, HARTREE2KJMOL, HARTREE2WAVENUMBERS, PLANCK, + R_GAS, SPEED_OF_LIGHT, +}; + +/// 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) +} + +/// 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 +// --------------------------------------------------------------------------- + +/// 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, None) +} + +/// 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; + + 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 num = match nrt_user { + Some(k) => k, + None => { + 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() + }, + }; + 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 / SPEED_OF_LIGHT * 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 +} + +// --------------------------------------------------------------------------- +// 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); + + // --------------- 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.max(3) - 3; + } + if !project_rot { + let nrot = match rotor_type { + "ATOM" => 0, + "LINEAR" => 2, + _ => 3, + }; + nrt = nrt.max(nrot) - nrot; + } + + // --------------- translation / rotation projector --------------- + let space = format!("{}{}", if project_trans { "T" } else { "" }, if project_rot { "R" } else { "" }); + 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]; + 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 --------------- + // 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 { + if is_tr[i] { + trv.push("TR"); + } else if fc_sorted[i].abs().sqrt() < 1.0e-3 { + trv.push("-"); + } else { + trv.push("V"); + } + } + + // --------------- conversion factors --------------- + 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) + // 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 * 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(); + 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 / (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; + h[TRANS] = 2.5 * t; + + // ---------- rotational ---------- + match rotor_type { + RotorType::Atom => {}, + RotorType::Linear => { + 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; + e[ROT] = t; + }, + RotorType::Regular => { + let phi = [ + 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()); + 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(); + + // 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; + 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.push("".to_string()); // blank line after each group + } + + 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/analdrv/vib_interface.rs b/src/analdrv/vib_interface.rs new file mode 100644 index 0000000000000000000000000000000000000000..279090ddd3e47762bb89250285af1a8174666ae4 --- /dev/null +++ b/src/analdrv/vib_interface.rs @@ -0,0 +1,101 @@ +use super::prelude::*; +use crate::geom_io::get_mass_charge; +use crate::analdrv::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: &AnalDrvConfig, + 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 analdrv) ==============="); + + // 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 analdrv) ==============="); + + println!(""); + + println!("=============== Thermo Analysis (Usual Style in analdrv) ==============="); + println!(""); + 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(); + 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 analdrv) ==============="); + + (de_hess.into_shape(-1).into_vec(), vib, th) +} diff --git a/src/constants/element.rs b/src/constants/element.rs new file mode 100644 index 0000000000000000000000000000000000000000..227874225d58102fb9f96f9d35b03c309fbd29b8 --- /dev/null +++ b/src/constants/element.rs @@ -0,0 +1,381 @@ +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", + "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 d92d9922716884d4b8df79150c4c5b7e4c72d044..739383ee96d172b136f87cae2ff3c866aa1def6c 100644 --- a/src/constants/mod.rs +++ b/src/constants/mod.rs @@ -2,127 +2,24 @@ 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::*; 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 +33,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,184 +77,41 @@ 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 +// 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 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 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 SPEED_OF_LIGHT: f64 = 299792458.0; // CODATA 2022, m s^-1, https://physics.nist.gov/cgi-bin/cuu/Value?c + +// =========== unit conversion =================================== +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 -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/ctrl_io/geometric_pyo3_io.rs b/src/ctrl_io/geometric_pyo3_io.rs index 0d24e90d43a9afc2b45720f62c4b80f6213bcfac..eaa12de6d4111c142d6cfb2daaed9161078b020f 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 d780aa5583329c4832e3e84ff68f02a9d74f0f92..ade1fbe28240f29dbeefabe1a94c6cd1435a3cb9 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::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; @@ -85,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.analdrv = tmp_keys.get("analdrv").map(serde_from_value); Ok((tmp_input,tmp_geomcell)) } @@ -385,6 +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 { @@ -550,6 +556,8 @@ impl InputKeywords { thermo: None, use_fxc_opt: false, rel: RelativisticMethod::None, + analdrv: None, + analdrv_tasks: Vec::new(), } } @@ -1938,6 +1946,13 @@ pub fn parse_ctrl_keywords(tmp_keys: &serde_json::Value) -> anyhow::Result String::from("legacy"), }; + tmp_input.analdrv_tasks = match tmp_ctrl.get("analdrv_tasks").unwrap_or(&serde_json::Value::Null) { + serde_json::Value::Null => 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"), + }; + //=========================================================== // Global check of ctrl keywords and futher modification //============================================================ diff --git a/src/dft/gen_grids/atom.rs b/src/dft/gen_grids/atom.rs index eb7b74435d881825ca4ddc7e323c33e07ab4724e..ac294f9f04395758c6973b1373e42814062ab9c1 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 * 0.529177249); + 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/dft/gen_grids/atom_old.rs b/src/dft/gen_grids/atom_old.rs index 87a6ef95aa3000a338c36dcb35dd134cd0b83d46..24b68f05ecf60fea3aa486f2b8c18bd66a55d9f8 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 265a7dfdc94e980fc3e25fb6a127ff80e5687aa1..897af95671e6e9ee8441f4fcaa41b5a706bce382 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 89004b4358cd37b0cfbd60188196bce0d5b2b48e..34baac0ba1880cb5c0e7caa166b40a3038aa2f7c 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 4a6e57b7bb735d82a27ff2d9347e85a2517c7a32..25f1205cd98131461e3c9212fc2980d71e6899d8 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/dft/mod.rs b/src/dft/mod.rs index 5510a6775478ad1980f57e04f615a7188533de0d..3b344b6b950087634f15a9cc728b8df05c7b753a 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; @@ -3505,6 +3507,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 `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 { + 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 new file mode 100644 index 0000000000000000000000000000000000000000..937593659d41c77ee1109e51941c17714833c379 --- /dev/null +++ b/src/dft/numint_matmul/hess_rks.rs @@ -0,0 +1,952 @@ +// see also pyhessref/nimatmul/rks.py + +use super::prelude::*; +use crate::analdrv::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 (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 { + 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; + } + (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 { + // 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: Vec<(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: Vec<(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(), + } + } + + /// 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: + /// - `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)); + + // 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 { + 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 0000000000000000000000000000000000000000..ed8242558249c1cc4a31608205547863b989f8e5 --- /dev/null +++ b/src/dft/numint_matmul/hess_uks.rs @@ -0,0 +1,825 @@ +// 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::analdrv::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 (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 { + let xc_type_i = determine_den_type(xc_func); + let nvar_i = xc_type_i.num_nvar(); + let rho_i = rho.i((.., ..nvar_i, ..)); + let xc_eff = libxc_eval_eff(xc_func, rho_i, 2, false); + let [_, vxc_i, fxc_i] = xc_eff.into_iter().collect_array().unwrap(); + *&mut vxc.i_mut((.., ..nvar_i, ..)) += *scale * vxc_i; + *&mut fxc.i_mut((.., ..nvar_i, .., ..nvar_i, ..)) += *scale * fxc_i; + } + (vxc, fxc) +} + +/// Lean evaluation of only `vxc` and `fxc` on a given grid, for use as the CP-KS `cpks_vxc` / +/// `cpks_fxc` when a dedicated (coarser) CP-KS grid is attached (UKS counterpart of +/// [`make_cpks_vxc_fxc`](super::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( + 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: Vec<(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: Vec<(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(), + } + } + + /// 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[α]); + 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)); + + // 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] { + 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 0000000000000000000000000000000000000000..6b03701adfb8a0eebdf8b5a1affd07f1ba58d21b --- /dev/null +++ b/src/dft/numint_matmul/mod.rs @@ -0,0 +1,106 @@ +//! 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. + +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 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 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::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; +} + +/* #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 0000000000000000000000000000000000000000..5b49346c893399cb06db3473d5e6a2aceca4d23e --- /dev/null +++ b/src/dft/numint_matmul/nimatmul.rs @@ -0,0 +1,509 @@ +// 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. +/// +/// 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 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. + /// + /// 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 1536 * 1 * 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 = 1536; + let nbatch = nchunk * 1 * 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 0000000000000000000000000000000000000000..7440ae35bf0b43d2bc32ac4a5cf9b0cbffe2d843 --- /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 0000000000000000000000000000000000000000..bf74574e199e5e276ada849ed9dd9f074d2f363f --- /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 0000000000000000000000000000000000000000..6ab361e6bb8f66e202537064e310cb77fbbf3096 --- /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 0000000000000000000000000000000000000000..60af99a4a90258dd935c98c71cda729f8e4e49e4 --- /dev/null +++ b/src/dft/xceff/libxc-xcfun-trans.md @@ -0,0 +1,60 @@ +# 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). diff --git a/src/dft/xceff/libxc_wrap.rs b/src/dft/xceff/libxc_wrap.rs new file mode 100644 index 0000000000000000000000000000000000000000..e45abf320421c802477926aadd9b15d1f7d569a4 --- /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 0000000000000000000000000000000000000000..7385cfbf38dce4769f604e4ff6f89aa8d18633ab --- /dev/null +++ b/src/dft/xceff/mod.rs @@ -0,0 +1,25 @@ +//! 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::*; + 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/dft/xceff/xc_deriv.rs b/src/dft/xceff/xc_deriv.rs new file mode 100644 index 0000000000000000000000000000000000000000..6abfad4181b855eb4e6fe3e750c133fac843feb1 --- /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/external_field/num_dipole.rs b/src/external_field/num_dipole.rs index 244f4a83e3a7ba8e2a2611acc02098abebb936e2..9ebb6e7698d7dbace1d5935e96b56700c95d2ed3 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 diff --git a/src/fileop/chkfile.rs b/src/fileop/chkfile.rs index aea7d1929d07220cc77c4c08a369a0487dd3bd1e..d06f274d8be786459ce22cf0ea26895f5c1fcb99 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::ANG; + 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 7143f60fa337829cfd88b098f21983d82c6fcfee..6c280d5da4787f0dbd9c6066829f44bad81a8eac 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 = 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 = 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 *= BOHR.powf(-1.0); dd } diff --git a/src/grad/mod.rs b/src/grad/mod.rs index d341dc17b462ac9b70eff9fdce8f65f44ae2e18a..787c25ffd731807d2af3c11ee63a191397f2cd1b 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/grad/rhf.rs b/src/grad/rhf.rs index 3f4d060d2f99f85fc59692075fadd3f1bd4d18a1..c73a52cc4a40563f3ddcf626a9b2c2e7efb60169 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/rhf.rs b/src/hessian/rhf.rs index b2cb7f983248807aacc8b951b1f66d9ce93a0d9b..965b369dc6f95356d2042a010fe66f6fcc85f0b8 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::{FQ, 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 = 219474.63 / 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]; diff --git a/src/lib.rs b/src/lib.rs index 7005c9c9b6f00ddd683aee6426f7b8958f70a0fa..7c57f09dd041d58ba81172a81aba0f9d85982b8a 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; @@ -84,6 +86,7 @@ pub mod fileop; pub mod ri_cphf; pub mod lib_rint; pub mod x2c; +pub mod analdrv; //extern crate rest; diff --git a/src/lib_rint/mod.rs b/src/lib_rint/mod.rs index a3394cc8f838b7a42fa7286bff6dbc0dfefd20a5..2ea4c7f41d5cc510485574789caf616c1c4b8993 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::ANG; + 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 3f4b4edb29637a0a5ce034056821139e817b9702..47c293e7084ec4b332a56dac26af014b767b827d 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( @@ -384,6 +384,19 @@ 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); } + + //=================================== + // Analytical derivative module calculations + //=================================== + 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::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"); + } time_mark.count("Overall"); @@ -637,7 +650,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 +768,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!(""); @@ -954,7 +967,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]; @@ -970,22 +982,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 hess_total = crate::hessian::compute_hessian(&*scf_data) - .expect("Analytical Hessian computation failed for geometry optimization"); + let use_analdrv = scf_data.mol.ctrl.geometric_pyo3.as_ref().map_or(false, |g| g.use_analdrv); + 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(); diff --git a/src/post_scf_analysis/mod.rs b/src/post_scf_analysis/mod.rs index 360109fbbe9049a5f7e23fc6eaf177280d21d136..7f9753a08a4b0094346894d6321e7d81db32f344 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 83c865503847dfda9f971101304dabb9ba3a1d17..1b79252076ecd18c218d776e3bffa87ef039173c 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, 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); } @@ -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 408eadbfb926c2072f1ae6f4ac17ec3563c39085..f18e51312183fb01a26f50939dd1d5f473b3de00 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/dynamicbse.rs b/src/ri_bse/dynamicbse.rs index 6b32cc4173463b8a8dcc68d61ee3798b1c60e2af..4ee720608fe9c959c56932028db7a6b8ac3c9b98 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] * 27.2114; + 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] * 27.2114); + 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 89af13a0ab1ef30712f23cd6f81de1128e04b3cb..e51222349204fafcecb6d1b892cdc78def649918 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] * 27.2114; + 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] * 27.2114); + 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] * 27.2114; + 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_bse/response.rs b/src/ri_bse/response.rs index 758b150743170804adf76912b418152948dafd56..7b7afbb3278d46b20b1134054575558f2cf43dad 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/ri_gw/mod.rs b/src/ri_gw/mod.rs index f5e010ed8f254bb1b7a5b66b474560f3efe8ca40..1c1b947423bf90980c5a3ef6a837875b3671e4d7 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; @@ -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 * 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 * EV, (i_omega as f64) / ((nomega_chi_real - 1) as f64)); } let lr = low_rank_vchi_vsqrt( diff --git a/src/ri_jk/hess_r.rs b/src/ri_jk/hess_r.rs new file mode 100644 index 0000000000000000000000000000000000000000..e107191ad00c8891f8d8e73bb0b8f24ce059a147 --- /dev/null +++ b/src/ri_jk/hess_r.rs @@ -0,0 +1,1954 @@ +//! 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::analdrv::cint_handling::*; +use crate::analdrv::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 +/// `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 +/// `factor_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 `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( + 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 factor_j: f64, + pub factor_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, 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(), + factor_j, + factor_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, + factor_j: f64, + factor_k: f64, + cderi: TsrCow<'a>, + j2c_decomp: J2CDecompose, + ) -> Self { + Self { + mol: mol.clone(), + aux: aux.clone(), + factor_j, + factor_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.factor_j != 0.0, + self.factor_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.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.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.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); + 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.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); + 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.factor_j != 0.0 { + let j1ao = KEYS_J1AO.iter().map(|&key| &intmd[key]).fold(deriv1_ao_init(), |acc, x| acc + x); + deriv1_bra += self.factor_j * (&j1ao % &mocc); + self.result.insert("j1ao", j1ao); + } + 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.factor_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 `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(); + 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.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.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.factor_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 0000000000000000000000000000000000000000..7c7ed61ef73a3e1acab122a503ff8d08d8df5214 --- /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: ``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): ``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 +//! [`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::analdrv::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 factor_j: f64, + pub factor_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, 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(), + factor_j, + factor_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, + factor_j: f64, + factor_k: f64, + cderi: TsrCow<'a>, + j2c_decomp: J2CDecompose, + ) -> Self { + Self { + mol: mol.clone(), + aux: aux.clone(), + factor_j, + factor_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.factor_j != 0.0, + self.factor_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.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.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.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.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); + 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.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.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.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.factor_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.factor_j != 0.0, + self.factor_k != 0.0, + 72, // TODO: batch size `72` should be tunable by max-memory. + ); + + let nao = mo_coeff[0].shape()[0]; + let occidx = [mo_occ[0].view().greater(0.0).into_vec(), mo_occ[1].view().greater(0.0).into_vec()]; + let mocc = [ + mo_coeff[0].view().bool_select(-1, &occidx[0]).into_contig(ColMajor), + mo_coeff[1].view().bool_select(-1, &occidx[1]).into_contig(ColMajor), + ]; + let nocc = [mocc[0].shape()[1], mocc[1].shape()[1]]; + + let mut resp = [None, None]; + for s in 0..2 { + let shape = bra[s].shape().to_vec(); + let nprop: usize = shape[2..].iter().product(); + let mut r = rt::zeros(([nao, nocc[s], nprop], device)); + // J: spin-independent AO operator, right half-transformed by this spin's mocc. + // The shared `j_ao` carries the RHF symmetrization factor (effective `4 * J1`); UHF + // naive J uses `2 * J1`, so an extra `0.5` prefactor is applied here (occ = 1 vs 2). + if let Some(j_ao) = j_ao.as_ref() { + r += 0.5 * self.factor_j * (j_ao.view() % &mocc[s]); + } + // K: same-spin bra form (UHF occ = 1, so no 0.5 factor — unlike RHF). The core already + // bakes in the exchange sign, so this is an additive contribution. + if let Some(k_bra) = k_bras.get(s) { + r += self.factor_k * k_bra.view().reshape((nao, nocc[s], nprop)); + } + resp[s] = Some(r.into_shape(shape)); + } + [resp[0].take().unwrap(), resp[1].take().unwrap()] + } +} + +/* #endregion */ diff --git a/src/ri_jk/mod.rs b/src/ri_jk/mod.rs index 7f267c8dd21d751fbc5eab70ee823c01cca92a6d..c8b28a9fa9c827b56240eb24b12de0fd4e46e418 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 00e14037d0fa8de7876ef0929ebe5b6bef0daf3e..44cde08bed8600fee35a4ca8fe929ebb0424d12a 100644 --- a/src/ri_jk/prelude_dev.rs +++ b/src/ri_jk/prelude_dev.rs @@ -8,11 +8,12 @@ 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}; -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; @@ -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 220e760ba8aace30a6dcd1d3ba9dff0110979fd8..d82f0422121417d1bde3c92a0b2a15f79ce5d755 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/ri_jk/util.rs b/src/ri_jk/util.rs index 377467d5c68c19ccda4743d4afac3b04ceed1f67..4bcb893ac1c5266889c8db61a143a47368c2ddaf 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/ri_tddft/response.rs b/src/ri_tddft/response.rs index 01ab240484ab5aa21b222bd9aa3206aa2ec7d092..c573e40422fffa4af23bd95fdb1a37bd7d637df1 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 * 27.2114); + 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 * 27.2114); + println!(" ω = {:.8} Ha ({:.4} eV)", omega, omega * EV); println!(" γ = {:.8} Ha", gamma); // Prepare fxc data diff --git a/src/scf_io/mod.rs b/src/scf_io/mod.rs index 5b97d7f41ff5874572c94ebd44c3b25028646163..e8ae8e6251073bd8bdeb685060500db7992f96ac 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}; use itertools::{Itertools}; use rayon::prelude::*; use std::collections::HashMap; @@ -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 { diff --git a/src/solvent/smd_cds.rs b/src/solvent/smd_cds.rs index c2d60443ce937e8478e1665a8425c72b2d2cc87f..aa54da59cda80c99ebdaaa3ccbe9deb6d8ea1fa7 100644 --- a/src/solvent/smd_cds.rs +++ b/src/solvent/smd_cds.rs @@ -62,8 +62,8 @@ //! //! ## Unit Conversions //! -//! - `TO_ANGS = 0.529177` Bohr → Å -//! - `TO_KCAL = 627.509` 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 = 0.52917724924; // Bohr → Å - const TO_KCAL: f64 = 627.509451; // 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 = 627.509451; + 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 = 627.509451; + const TO_KCAL: f64 = HARTREE2KCAL; let gcds = gcds_kcal / TO_KCAL; (gcds, tarea, dcds) diff --git a/src/utilities/rstsr_util.rs b/src/utilities/rstsr_util.rs index 87754194d02706078084bfa965cecd29af68c82f..414282f9cb9856c7b4744ad9e13542f816576b8e 100644 --- a/src/utilities/rstsr_util.rs +++ b/src/utilities/rstsr_util.rs @@ -1,16 +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 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; diff --git a/src/x2c/mod.rs b/src/x2c/mod.rs index 7f54f3ddb12fe96fd5281b162326f70763f4abec..7cd75f571344c3ac3c7108614f79f69840f1c109 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 = 137.03599967994; + let light_speed: f64 = LIGHT_SPEED; let (xmol, contr_coeff) = self.decontract_basis(); let n_contracted = self.num_basis; diff --git a/tests/analdrv/mod.rs b/tests/analdrv/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..9927b8629c61257ba36a5b2468f6029000e907e2 --- /dev/null +++ b/tests/analdrv/mod.rs @@ -0,0 +1,6 @@ +pub mod rhf; +pub mod rks_b3lyp; +pub mod rks_b3lyp_grid_levels; +pub mod uhf; +pub mod uks_tpss0_grid_levels; +pub mod uks_tpssh; diff --git a/tests/analdrv/rhf.rs b/tests/analdrv/rhf.rs new file mode 100644 index 0000000000000000000000000000000000000000..9531139086aad41c94f6c8c1c1c16797b807d2de --- /dev/null +++ b/tests/analdrv/rhf.rs @@ -0,0 +1,121 @@ +use pyrest::analdrv::config::AnalDrvConfig; +use pyrest::analdrv::rscf_interface::rscf_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 = "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 = 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]; + 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); + + let natm = 4; + let de = rt::asarray((&de, [3, 3, natm, natm])); + println!("Hessian:\n{:12.6}", de.t()); +} + +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 = AnalDrvConfig::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]); + } +} diff --git a/tests/analdrv/rks_b3lyp.rs b/tests/analdrv/rks_b3lyp.rs new file mode 100644 index 0000000000000000000000000000000000000000..d8c09a14345ae65901302e03ba3e36e2704e0020 --- /dev/null +++ b/tests/analdrv/rks_b3lyp.rs @@ -0,0 +1,53 @@ +use pyrest::analdrv::rscf_interface::rscf_hess_interface; +use pyrest::analdrv::config::AnalDrvConfig; + +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 = "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 = AnalDrvConfig::default(); + 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/analdrv/rks_b3lyp_grid_levels.rs b/tests/analdrv/rks_b3lyp_grid_levels.rs new file mode 100644 index 0000000000000000000000000000000000000000..cb181924c68ceaa1ee7ea20b10daa7163126df57 --- /dev/null +++ b/tests/analdrv/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::analdrv::config::AnalDrvConfig; +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}; +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: 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(); + 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 = AnalDrvConfig { 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 = AnalDrvConfig::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 = AnalDrvConfig { 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/analdrv/uhf.rs b/tests/analdrv/uhf.rs new file mode 100644 index 0000000000000000000000000000000000000000..b4ad6eb3ad4ece6d4d4157b3d92590095bd427da --- /dev/null +++ b/tests/analdrv/uhf.rs @@ -0,0 +1,56 @@ +use pyrest::analdrv::uscf_interface::uscf_hess_interface; +use pyrest::analdrv::config::AnalDrvConfig; + +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 + """ + +[analdrv] +atm_list = [0, 1, 3] +"##; + +#[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 = AnalDrvConfig::default(); + let (de, _vib, _th) = uscf_hess_interface(&mut scf_data, &config); + + 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 diff --git a/tests/analdrv/uks_tpss0_grid_levels.rs b/tests/analdrv/uks_tpss0_grid_levels.rs new file mode 100644 index 0000000000000000000000000000000000000000..6468bac0c5705ce5b2df019ee9a8faccd7a6dd1d --- /dev/null +++ b/tests/analdrv/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::analdrv::config::AnalDrvConfig; +use pyrest::analdrv::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: 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(); + 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 = AnalDrvConfig::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 = AnalDrvConfig { + 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/analdrv/uks_tpssh.rs b/tests/analdrv/uks_tpssh.rs new file mode 100644 index 0000000000000000000000000000000000000000..79a9cc32455f23194b9f79cf2149934c33ffc2fd --- /dev/null +++ b/tests/analdrv/uks_tpssh.rs @@ -0,0 +1,64 @@ +use pyrest::analdrv::config::AnalDrvConfig; +use pyrest::analdrv::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 = "TPSSh" + 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() { + // 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(); + let mut scf_data = scf_io::SCF::build(mol, &None); + scf_without_build(&mut scf_data, &None); + + let config = AnalDrvConfig::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()); + + // 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]); + } + } +} diff --git a/tests/test_hessian_opbase.rs b/tests/test_hessian_opbase.rs new file mode 100644 index 0000000000000000000000000000000000000000..ab0f9246bfb3f92e231829158f6302bc2bf0d56f --- /dev/null +++ b/tests/test_hessian_opbase.rs @@ -0,0 +1 @@ +pub mod analdrv; \ No newline at end of file