From 57d96849d8afccc3a4a76b46522fa493b009c92c Mon Sep 17 00:00:00 2001 From: Betty0621 <3309721592@qq.com> Date: Fri, 28 Aug 2026 21:25:23 +0800 Subject: [PATCH 01/16] Add analytic nuclear gradients for Dz0SCF --- src/nest/dz0scf/dz0scf.py | 5 + src/nest/grad/dz0scf.py | 562 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 567 insertions(+) create mode 100644 src/nest/grad/dz0scf.py diff --git a/src/nest/dz0scf/dz0scf.py b/src/nest/dz0scf/dz0scf.py index 943ab77..7d32fa3 100644 --- a/src/nest/dz0scf/dz0scf.py +++ b/src/nest/dz0scf/dz0scf.py @@ -84,6 +84,11 @@ def get_veff( def high_spin_energy(self): return evaluate_high_spin_energy(self) + def nuc_grad_method(self): + """Return the Dz0SCF analytic nuclear-gradient driver.""" + from nest.grad.dz0scf import Gradients + return Gradients(self) + class EnsembleROKS(_DZ0VeffMixin, dft.roks.ROKS): pass diff --git a/src/nest/grad/dz0scf.py b/src/nest/grad/dz0scf.py new file mode 100644 index 0000000..f730ad3 --- /dev/null +++ b/src/nest/grad/dz0scf.py @@ -0,0 +1,562 @@ +from __future__ import annotations + +import warnings +from dataclasses import dataclass + +import numpy as np +from scipy.sparse.linalg import LinearOperator, gmres + +from pyscf import dft, lib +from pyscf.grad import rhf as rhf_grad +from pyscf.hessian import rhf as rhf_hess +from pyscf.hessian import rks as rks_hess +from pyscf.lib import logger +from pyscf.scf import _response_functions + + +_OCC_TOL = 1e-8 + + +@dataclass(frozen=True) +class RotationSpace: + """Packed nonredundant rotations and their orbital occupations.""" + + p: np.ndarray + q: np.ndarray + f: np.ndarray + nalpha: np.ndarray + nbeta: np.ndarray + + @property + def size(self) -> int: + return int(self.p.size) + + @property + def occupation_gap(self) -> np.ndarray: + return self.f[self.q] - self.f[self.p] + + def unpack(self, vector: np.ndarray) -> np.ndarray: + """Map a packed vector to the full real anti-symmetric MO matrix.""" + vector = np.asarray(vector, dtype=float) + if vector.shape != (self.size,): + raise ValueError( + f"Expected a packed rotation of shape {(self.size,)}, " + f"got {vector.shape}." + ) + matrix = np.zeros((self.f.size, self.f.size)) + matrix[self.p, self.q] = vector + matrix[self.q, self.p] = -vector + return matrix + + def pack(self, matrix: np.ndarray) -> np.ndarray: + """Extract entries matching the packed (p,q) rotation order.""" + return np.asarray(matrix)[..., self.p, self.q] + + +def _block_pairs(rows: np.ndarray, cols: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Pairs matching ``matrix[np.ix_(rows, cols)].ravel()`` order.""" + rows = np.asarray(rows, dtype=int) + cols = np.asarray(cols, dtype=int) + if rows.size == 0 or cols.size == 0: + empty = np.empty(0, dtype=int) + return empty, empty + return np.repeat(rows, cols.size), np.tile(cols, rows.size) + + +def _rotation_space(mo_occ: np.ndarray) -> RotationSpace: + f = np.asarray(mo_occ, dtype=float) + if f.ndim != 1: + raise ValueError(f"Dz0 occupations must be a one-dimensional array; got {f.shape}.") + + is_c = np.isclose(f, 2.0, atol=_OCC_TOL, rtol=0.0) + is_o = np.isclose(f, 1.0, atol=_OCC_TOL, rtol=0.0) + is_v = np.isclose(f, 0.0, atol=_OCC_TOL, rtol=0.0) + if not np.all(is_c | is_o | is_v): + bad = np.where(~(is_c | is_o | is_v))[0] + raise NotImplementedError( + "This implementation requires Dz0 occupations 0, 1, or 2. " + f"Nonstandard occupations were found at MO indices {bad.tolist()}." + ) + + c = np.where(is_c)[0] + o = np.where(is_o)[0] + v = np.where(is_v)[0] + blocks = (_block_pairs(o, c), _block_pairs(v, c), _block_pairs(v, o)) + p = np.concatenate([block[0] for block in blocks]) + q = np.concatenate([block[1] for block in blocks]) + + nalpha = (f > 0.0).astype(float) + nbeta = np.isclose(f, 2.0, atol=_OCC_TOL, rtol=0.0).astype(float) + return RotationSpace(p=p, q=q, f=f, nalpha=nalpha, nbeta=nbeta) + + +def _copy_ks_settings(source, target) -> None: + """Copy numerical-integration settings without copying unsupported wrappers.""" + for name in ( + "xc", + "nlc", + "grids", + "nlcgrids", + "_numint", + "max_memory", + "verbose", + "stdout", + "direct_scf_tol", + "small_rho_cutoff", + ): + if hasattr(source, name): + setattr(target, name, getattr(source, name)) + + +def _ao_density(mo_coeff: np.ndarray, occupation: np.ndarray) -> np.ndarray: + """Return D_AO = C occupation C^dagger.""" + return (mo_coeff * occupation) @ mo_coeff.conj().T + + +def _transform_ao_to_mo(mo_coeff: np.ndarray, matrices: np.ndarray) -> np.ndarray: + """Transform one AO matrix or a leading batch of AO matrices to the MO basis.""" + return np.einsum("up,...uv,vq->...pq", mo_coeff.conj(), matrices, mo_coeff, optimize=True) + + +def _full_overlap_derivative( + one_sided_s1: np.ndarray, + atom: int, + aoslices: np.ndarray, +) -> np.ndarray: + """Build S^A from PySCF's one-sided overlap derivative matrix.""" + p0, p1 = aoslices[atom][2:] + s1 = np.zeros_like(one_sided_s1) + s1[:, p0:p1, :] += one_sided_s1[:, p0:p1, :] + s1[:, :, p0:p1] += one_sided_s1[:, p0:p1, :].transpose(0, 2, 1) + return s1 + + +def _fractional_rks_fock_skeleton( + charge_mf, + mo_coeff: np.ndarray, + mo_occ: np.ndarray, +) -> np.ndarray: + """Return F_AO^(0,[A]) for all atoms and Cartesian components. + + This is PySCF's RKS ``make_h1`` construction with the density corrected + from ``2 C_occ C_occ^T`` to the fractional-occupation density + ``C mo_occ C^T`` required by Dz0SCF. Range-separated exchange is included + through PySCF's ``(omega, alpha, hyb)`` decomposition. + """ + mol = charge_mf.mol + dm0 = _ao_density(mo_coeff, mo_occ) + hessobj = rks_hess.Hessian(charge_mf) + hcore_deriv = charge_mf.nuc_grad_method().hcore_generator(mol) + + ni = charge_mf._numint + ni.libxc.test_deriv_order(charge_mf.xc, 2, raise_error=True) + omega, alpha, hyb = ni.rsh_and_hybrid_coeff(charge_mf.xc, spin=mol.spin) + hybrid = ni.libxc.is_hybrid_xc(charge_mf.xc) + + memory_now = lib.current_memory()[0] + max_memory = max(2000, charge_mf.max_memory * 0.9 - memory_now) + h1ao = rks_hess._get_vxc_deriv1(hessobj, mo_coeff, mo_occ, max_memory) + + aoslices = mol.aoslice_by_atom() + for atom, (shl0, shl1, p0, p1) in enumerate(aoslices): + shls_slice = (shl0, shl1) + (0, mol.nbas) * 3 + if hybrid: + vj1, vj2, vk1, vk2 = rhf_hess._get_jk( + mol, + "int2e_ip1", + 3, + "s2kl", + [ + "ji->s2kl", + -dm0[:, p0:p1], + "lk->s1ij", + -dm0, + "li->s1kj", + -dm0[:, p0:p1], + "jk->s1il", + -dm0, + ], + shls_slice=shls_slice, + ) + veff = vj1 - 0.5 * hyb * vk1 + veff[:, p0:p1] += vj2 - 0.5 * hyb * vk2 + if omega != 0.0: + with mol.with_range_coulomb(omega): + vk1, vk2 = rhf_hess._get_jk( + mol, + "int2e_ip1", + 3, + "s2kl", + [ + "li->s1kj", + -dm0[:, p0:p1], + "jk->s1il", + -dm0, + ], + shls_slice=shls_slice, + ) + veff -= 0.5 * (alpha - hyb) * vk1 + veff[:, p0:p1] -= 0.5 * (alpha - hyb) * vk2 + else: + vj1, vj2 = rhf_hess._get_jk( + mol, + "int2e_ip1", + 3, + "s2kl", + ["ji->s2kl", -dm0[:, p0:p1], "lk->s1ij", -dm0], + shls_slice=shls_slice, + ) + veff = vj1 + veff[:, p0:p1] += vj2 + + h1ao[atom] += veff + veff.transpose(0, 2, 1) + h1ao[atom] += hcore_deriv(atom) + + return h1ao + + +class DZ0Gradients(lib.StreamObject): + """Analytic gradient driver for the Dz0SCF high-spin reference energy.""" + + _keys = { + "base", + "mol", + "max_memory", + "conv_tol", + "max_cycle", + "restart", + "grid_response", + "atmlst", + "de", + "z", + "g_hs", + "g_dz0", + "b", + "e_hs_unrelaxed", + } + + def __init__(self, mf): + self.base = mf + self.mol = mf.mol + self.verbose = mf.verbose + self.stdout = mf.stdout + self.max_memory = mf.max_memory + + self.conv_tol = 1e-9 + self.max_cycle = 80 + self.restart = 40 + self.grid_response = False + self.atmlst = None + + self.de = None + self.z = None + self.g_hs = None + self.g_dz0 = None + self.b = None + self.e_hs_unrelaxed = None + + self._space = None + self._charge_mf = None + self._hs_mf = None + self._charge_response = None + self._c0 = None + self._f0ao = None + self._f0mo = None + self._f_hs_ao = None + self._f_hs_mo = None + self._dm_hs = None + self._w_hs_mo = None + + def dump_flags(self, verbose=None): + log = logger.new_logger(self, verbose) + log.info("******** Dz0SCF high-spin-reference analytic gradient ********") + log.info("Z-vector tolerance = %.3g", self.conv_tol) + log.info("Z-vector max cycles = %d", self.max_cycle) + log.info("GMRES restart = %d", self.restart) + log.info("grid response = %s", self.grid_response) + return self + + def _validate(self) -> None: + mf = self.base + mol = self.mol + if getattr(mf, "mo_coeff", None) is None or getattr(mf, "mo_occ", None) is None: + raise RuntimeError("Run Dz0SCF before requesting its analytic gradient.") + if hasattr(mf, "converged") and not mf.converged: + warnings.warn("Dz0SCF is not converged; its analytic gradient is not stationary.") + if np.iscomplexobj(mf.mo_coeff) and np.max(np.abs(np.asarray(mf.mo_coeff).imag)) > 1e-12: + raise NotImplementedError("Complex-orbital Dz0SCF gradients are not implemented.") + if self.grid_response: + raise NotImplementedError( + "Moving-grid response is not implemented consistently in B^(0,A); " + "use grid_response=False." + ) + if getattr(mf, "with_df", None) is not None: + raise NotImplementedError("Density-fitted Dz0SCF gradients are not implemented.") + if getattr(mf, "with_x2c", None) is not None: + raise NotImplementedError("X2C Dz0SCF gradients are not implemented.") + if getattr(mf, "with_solvent", None) is not None: + raise NotImplementedError("Solvent-response Dz0SCF gradients are not implemented.") + if hasattr(mf, "do_nlc") and mf.do_nlc(): + raise NotImplementedError("Nonlocal-correlation (NLC/VV10) Dz0SCF gradients are not implemented.") + if hasattr(mf, "do_disp") and mf.do_disp(): + raise NotImplementedError("Dispersion-corrected Dz0SCF gradients are not implemented.") + if getattr(mol, "dimension", 3) != 3: + raise NotImplementedError("Only molecular (three-dimensional) calculations are supported.") + + def _build_intermediates(self) -> None: + mf = self.base + mol = self.mol + c0 = np.asarray(mf.mo_coeff).real + space = _rotation_space(np.asarray(mf.mo_occ)) + + # Do not call the dft.RKS factory here: for mol.spin != 0 it returns + # ROKS, whereas R^(0) is the spin-unpolarized, charge-only RKS kernel. + charge_mf = dft.rks.RKS(mol) + _copy_ks_settings(mf, charge_mf) + charge_mf.mo_coeff = c0 + charge_mf.mo_occ = space.f + + dm0 = _ao_density(c0, space.f) + hcore = charge_mf.get_hcore(mol) + f0ao = hcore + charge_mf.get_veff(mol, dm0) + f0mo = _transform_ao_to_mo(c0, f0ao) + + charge_response = _response_functions._gen_rhf_response( + charge_mf, + mo_coeff=c0, + mo_occ=space.f, + singlet=None, + hermi=1, + max_memory=self.max_memory, + with_nlc=False, + ) + + hs_mf = dft.ROKS(mol) + _copy_ks_settings(mf, hs_mf) + hs_mf.mo_coeff = c0 + hs_mf.mo_occ = space.f + dm_hs = hs_mf.make_rdm1(c0, space.f) + veff_hs = hs_mf.get_veff(mol, dm_hs) + f_hs_ao = np.asarray((hcore + veff_hs[0], hcore + veff_hs[1])) + f_hs_mo = _transform_ao_to_mo(c0, f_hs_ao) + + occ_spin = np.asarray((space.nalpha, space.nbeta)) + w_hs_mo = 0.5 * np.sum( + occ_spin[:, :, None] * f_hs_mo + f_hs_mo * occ_spin[:, None, :], + axis=0, + ) + + self._space = space + self._charge_mf = charge_mf + self._hs_mf = hs_mf + self._charge_response = charge_response + self._c0 = c0 + self._f0ao = f0ao + self._f0mo = f0mo + self._f_hs_ao = f_hs_ao + self._f_hs_mo = f_hs_mo + self._dm_hs = dm_hs + self._w_hs_mo = w_hs_mo + + gap = space.occupation_gap + self.g_dz0 = 2.0 * gap * space.pack(f0mo) + self.g_hs = 2.0 * ( + (space.nalpha[space.q] - space.nalpha[space.p]) * space.pack(f_hs_mo[0]) + + (space.nbeta[space.q] - space.nbeta[space.p]) * space.pack(f_hs_mo[1]) + ) + + def hessian_vector_product(self, vector: np.ndarray) -> np.ndarray: + """Evaluate A^(0) vector without constructing the orbital Hessian.""" + if self._space is None: + self._validate() + self._build_intermediates() + + space = self._space + kappa = space.unpack(vector) + delta_dm_mo = kappa * space.f[None, :] - space.f[:, None] * kappa + delta_dm_ao = self._c0 @ delta_dm_mo @ self._c0.T + delta_f_ao = self._charge_response(delta_dm_ao) + delta_f_mo = _transform_ao_to_mo(self._c0, delta_f_ao) + moving_mo = self._f0mo @ kappa - kappa @ self._f0mo + result = 2.0 * space.occupation_gap * space.pack(moving_mo + delta_f_mo) + return np.asarray(result).real + + def _solve_z(self) -> np.ndarray: + space = self._space + if space.size == 0: + return np.empty(0) + + operator = LinearOperator( + (space.size, space.size), + matvec=self.hessian_vector_product, + rmatvec=self.hessian_vector_product, + dtype=float, + ) + + # The exact real-orbital Dz0 Hessian is symmetric, so A^T z = g_HS + # is solved with the same matrix-free action. This diagonal contains + # the one-electron commutator part and is used only as a preconditioner. + diagonal = 2.0 * space.occupation_gap * ( + self._f0mo.diagonal()[space.p] - self._f0mo.diagonal()[space.q] + ) + floor = max(1e-8, 1e-6 * np.max(np.abs(diagonal))) + safe_diagonal = np.where( + np.abs(diagonal) > floor, + diagonal, + np.where(diagonal < 0.0, -floor, floor), + ) + preconditioner = LinearOperator( + operator.shape, + matvec=lambda x: np.asarray(x) / safe_diagonal, + dtype=float, + ) + + residuals = [] + common = dict( + A=operator, + b=np.asarray(self.g_hs).real, + M=preconditioner, + restart=self.restart, + maxiter=self.max_cycle, + callback=residuals.append, + ) + try: + z, info = gmres(rtol=self.conv_tol, atol=0.0, callback_type="pr_norm", **common) + except TypeError: # SciPy < 1.12 compatibility + z, info = gmres(tol=self.conv_tol, atol=0.0, **common) + + if info != 0: + last = residuals[-1] if residuals else np.nan + raise RuntimeError( + "Dz0 Z-vector GMRES did not converge: " + f"info={info}, last preconditioned residual={last:.3e}. " + "Increase max_cycle/restart or inspect a near-singular Dz0 orbital Hessian." + ) + return z + + def _high_spin_unrelaxed_gradient(self) -> np.ndarray: + """Evaluate E_HS^(A), including the nonstationary overlap term.""" + mol = self.mol + c0 = self._c0 + dm_hs = self._dm_hs + hs_grad = self._hs_mf.nuc_grad_method() + hs_grad.grid_response = False + hs_grad.max_memory = self.max_memory + + hcore_deriv = hs_grad.hcore_generator(mol) + one_sided_s1 = hs_grad.get_ovlp(mol) + + dm_hs_for_gradient = hs_grad._tag_rdm1( + np.asarray(dm_hs), c0, self._space.f + ) + vhf = hs_grad.get_veff(mol, dm_hs_for_gradient) + dm_total = dm_hs[0] + dm_hs[1] + w_hs_ao = c0 @ self._w_hs_mo @ c0.T + + if mol._pseudo: + from pyscf.gto.pp_int import vpploc_nuc_grad, vppnl_nuc_grad + + de = vpploc_nuc_grad(mol, dm_total) + de += vppnl_nuc_grad(mol, dm_total) + else: + de = np.zeros((mol.natm, 3)) + + aoslices = mol.aoslice_by_atom() + for atom, (_, _, p0, p1) in enumerate(aoslices): + de[atom] += np.einsum("xij,ij->x", hcore_deriv(atom), dm_total) + de[atom] += 2.0 * np.einsum( + "sxij,sij->x", vhf[:, :, p0:p1], dm_hs[:, p0:p1] + ) + # -Tr[W_HS^MO S_MO^A], written with PySCF's one-sided S derivative. + de[atom] -= 2.0 * np.einsum( + "xij,ij->x", one_sided_s1[:, p0:p1], w_hs_ao[p0:p1] + ) + + return de + hs_grad.grad_nuc(mol) + + def _build_b(self) -> np.ndarray: + """Build B_i^(0,A) for all atoms and Cartesian components.""" + mol = self.mol + c0 = self._c0 + space = self._space + aoslices = mol.aoslice_by_atom() + one_sided_s1 = rhf_grad.get_ovlp(mol) + fock_skeleton = _fractional_rks_fock_skeleton( + self._charge_mf, c0, space.f + ) + + b = np.empty((space.size, mol.natm, 3)) + for atom in range(mol.natm): + s1ao = _full_overlap_derivative(one_sided_s1, atom, aoslices) + s1mo = _transform_ao_to_mo(c0, s1ao) + + anticommutator_sf = s1mo * space.f[None, None, :] + anticommutator_sf += space.f[None, :, None] * s1mo + delta_dm_sym_mo = -0.5 * anticommutator_sf + delta_dm_sym_ao = np.einsum( + "up,xpq,vq->xuv", c0, delta_dm_sym_mo, c0.conj(), optimize=True + ) + response_ao = self._charge_response(delta_dm_sym_ao) + + explicit_mo = _transform_ao_to_mo(c0, fock_skeleton[atom]) + response_mo = _transform_ao_to_mo(c0, response_ao) + overlap_mo = -0.5 * ( + np.einsum("xpq,qr->xpr", s1mo, self._f0mo, optimize=True) + + np.einsum("pq,xqr->xpr", self._f0mo, s1mo, optimize=True) + ) + fock_fixed_k = explicit_mo + response_mo + overlap_mo + b[:, atom, :] = ( + 2.0 + * space.occupation_gap[:, None] + * space.pack(fock_fixed_k).T + ) + return b.real + + def kernel(self, atmlst=None, verbose=None) -> np.ndarray: + """Compute and return the Dz0SCF high-spin-reference nuclear gradient.""" + log = logger.new_logger(self, verbose) + self._validate() + self.dump_flags(verbose) + self._build_intermediates() + + max_g0 = float(np.max(np.abs(self.g_dz0))) if self.g_dz0.size else 0.0 + log.info("max |g_Dz0| = %.6g", max_g0) + scf_grad_tol = getattr(self.base, "conv_tol_grad", 0.0) or 0.0 + if max_g0 > max(1e-6, 100.0 * scf_grad_tol): + warnings.warn( + f"The packed Dz0 orbital gradient is not small (max={max_g0:.3e}); " + "the analytic-gradient stationarity equation may be inaccurate." + ) + + self.z = self._solve_z() + self.e_hs_unrelaxed = self._high_spin_unrelaxed_gradient() + self.b = self._build_b() + z_correction = np.einsum("i,iax->ax", self.z, self.b, optimize=True) + de = self.e_hs_unrelaxed - z_correction + + if self.mol.symmetry: + de = rhf_grad.symmetrize(self.mol, de) + + self.atmlst = atmlst + result = de if atmlst is None else de[np.asarray(atmlst, dtype=int)] + self.de = result + + if log.verbose >= logger.NOTE: + logger.note(self, "--------------- Dz0SCF reference gradients ---------------") + rhf_grad._write(log, self.mol, result, atmlst) + logger.note(self, "----------------------------------------------------------") + return result + + grad = kernel + + +Gradients = DZ0Gradients +Grad = DZ0Gradients + + +def nuc_grad_method(mf) -> DZ0Gradients: + """Functional constructor useful before binding the class as a method.""" + return DZ0Gradients(mf) + + +__all__ = ["DZ0Gradients", "Gradients", "Grad", "RotationSpace", "nuc_grad_method"] From a7abacabc39e90808171af40088b873fb4aec076 Mon Sep 17 00:00:00 2001 From: Betty0621 <3309721592@qq.com> Date: Fri, 28 Aug 2026 21:25:48 +0800 Subject: [PATCH 02/16] Add Dz0SCF gradient regression tests --- src/nest/grad/tests/test_dz0scf_grad.py | 187 ++++++++++++++++++++++++ 1 file changed, 187 insertions(+) create mode 100644 src/nest/grad/tests/test_dz0scf_grad.py diff --git a/src/nest/grad/tests/test_dz0scf_grad.py b/src/nest/grad/tests/test_dz0scf_grad.py new file mode 100644 index 0000000..6cef964 --- /dev/null +++ b/src/nest/grad/tests/test_dz0scf_grad.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Regression tests for the Dz0SCF high-spin-reference analytic gradient. + +The reference values were generated with PySCF 2.13.0. The analytic-gradient +implementation was independently checked against central finite differences. +NH2/PBE/6-31G exercises a 44-dimensional response space, while triplet +CH2/B3LYP/STO-3G explicitly exercises a singly occupied O space of dimension +two. +""" + +import unittest + +import numpy as np +from pyscf import gto + +from nest.dz0scf import DZ0SCF +from nest.grad.dz0scf import Gradients + + +GRADIENT_ATOL = 1.0e-6 +TRANSLATION_ATOL = 1.0e-6 +STATIONARITY_TOL = 1.0e-7 +Z_RESIDUAL_TOL = 1.0e-8 + + +NH2_PBE_GRAD = np.array( + [ + [-2.427124373439239e-15, -3.595373998342540e-02, 6.282335070410562e-02], + [1.181596094619805e-15, -1.750053024027649e-02, -2.075848019567249e-02], + [1.223560123578276e-15, 5.345440088313257e-02, -4.206455607347775e-02], + ] +) + + +CH2_B3LYP_OPEN2_GRAD = np.array( + [ + [-7.664563046880733e-04, -2.770013672468408e-02, 5.335666299824815e-02], + [2.403193394966376e-03, -6.315495031449209e-02, -2.395287280504814e-02], + [-1.636476595240756e-03, 9.085495314029288e-02, -2.940373960552241e-02], + ] +) + + +def run_dz0_gradient(atom, spin, xc, basis): + """Run one tightly converged and reproducible Dz0SCF gradient.""" + mol = gto.M( + atom=atom, + unit="Angstrom", + basis=basis, + charge=0, + spin=spin, + symmetry=False, + verbose=0, + output="/dev/null", + ) + + mf = DZ0SCF(mol, xc=xc) + mf.conv_tol = 1.0e-12 + mf.conv_tol_grad = 1.0e-9 + mf.max_cycle = 120 + mf.grids.level = 5 + mf.grids.prune = None + mf.small_rho_cutoff = 0.0 + mf.kernel() + if not mf.converged: + mol.stdout.close() + raise RuntimeError(f"{xc}/{basis} Dz0SCF did not converge") + + # Exercise the public API installed on the Dz0SCF class/mixin. + grad_obj = mf.nuc_grad_method() + if not isinstance(grad_obj, Gradients): + mol.stdout.close() + raise TypeError( + "DZ0SCF.nuc_grad_method() did not return " + "nest.grad.dz0scf.Gradients" + ) + + grad_obj.conv_tol = 1.0e-10 + grad_obj.max_cycle = 120 + grad_obj.restart = 50 + gradient = np.asarray(grad_obj.kernel()) + return mol, mf, grad_obj, gradient + + +class KnownValues(unittest.TestCase): + def assert_dz0_result( + self, + mf, + grad_obj, + gradient, + reference, + expected_space, + ): + """Check the orbital space, response equations, and nuclear gradient.""" + occupations = np.asarray(mf.mo_occ) + n_closed = int(np.count_nonzero(np.isclose(occupations, 2.0))) + n_open = int(np.count_nonzero(np.isclose(occupations, 1.0))) + n_virtual = int(np.count_nonzero(np.isclose(occupations, 0.0))) + observed_space = (n_closed, n_open, n_virtual, grad_obj._space.size) + + self.assertEqual(observed_space, expected_space) + self.assertLess( + float(np.max(np.abs(grad_obj.g_dz0))), + STATIONARITY_TOL, + ) + + z_residual = ( + grad_obj.hessian_vector_product(grad_obj.z) - grad_obj.g_hs + ) + self.assertLess(float(np.linalg.norm(z_residual)), Z_RESIDUAL_TOL) + + np.testing.assert_allclose( + np.sum(gradient, axis=0), + np.zeros(3), + rtol=0.0, + atol=TRANSLATION_ATOL, + ) + np.testing.assert_allclose( + gradient, + reference, + rtol=0.0, + atol=GRADIENT_ATOL, + ) + + def test_pbe_nh2_gradient(self): + """Check all NH2/PBE/6-31G components in a 44-D response space.""" + mol, mf, grad_obj, gradient = run_dz0_gradient( + atom=""" + N 0.000000 -0.040000 0.000000 + H 0.000000 0.780000 0.590000 + H 0.000000 -0.860000 0.520000 + """, + spin=1, + xc="PBE", + basis="6-31g", + ) + try: + self.assert_dz0_result( + mf, + grad_obj, + gradient, + NH2_PBE_GRAD, + expected_space=(4, 1, 8, 44), + ) + finally: + mol.stdout.close() + + def test_b3lyp_ch2_two_open_orbitals_gradient(self): + """Check all CH2 components with a two-dimensional open-shell space.""" + mol, mf, grad_obj, gradient = run_dz0_gradient( + atom=""" + C 0.020000 -0.030000 0.010000 + H -0.020000 0.800000 0.620000 + H 0.030000 -0.910000 0.500000 + """, + spin=2, + xc="B3LYP", + basis="sto-3g", + ) + try: + self.assert_dz0_result( + mf, + grad_obj, + gradient, + CH2_B3LYP_OPEN2_GRAD, + expected_space=(3, 2, 2, 16), + ) + finally: + mol.stdout.close() + + +if __name__ == "__main__": + unittest.main() From 954d6983eeeca976c6c90bb680153127839c50e5 Mon Sep 17 00:00:00 2001 From: Betty0621 <3309721592@qq.com> Date: Sat, 29 Aug 2026 11:47:08 +0800 Subject: [PATCH 03/16] Fix whitespace in Dz0SCF gradient --- src/nest/grad/dz0scf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nest/grad/dz0scf.py b/src/nest/grad/dz0scf.py index f730ad3..7033b69 100644 --- a/src/nest/grad/dz0scf.py +++ b/src/nest/grad/dz0scf.py @@ -445,7 +445,7 @@ def _high_spin_unrelaxed_gradient(self) -> np.ndarray: hcore_deriv = hs_grad.hcore_generator(mol) one_sided_s1 = hs_grad.get_ovlp(mol) - + dm_hs_for_gradient = hs_grad._tag_rdm1( np.asarray(dm_hs), c0, self._space.f ) From 958c1b089f2232816a0c20a79cb395933e700866 Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Mon, 31 Aug 2026 03:18:47 +0800 Subject: [PATCH 04/16] feat: add EnsembleRKS reference and reference-energy protocol Port the average/roks dual-semantics EnsembleRKS reference (fixed 2/1/0 occupations, Dz = 0) and align Dz0SCF with the shared reference_energy semantics so NTTDA can treat both reference types uniformly. --- src/nest/__init__.py | 4 +- src/nest/dz0scf/dz0scf.py | 12 +- src/nest/dz0scf/tests/test_dz0scf.py | 20 ++- src/nest/ensemble_rks/__init__.py | 3 + src/nest/ensemble_rks/ensemble_rks.py | 127 ++++++++++++++++++ .../ensemble_rks/tests/test_ensemble_rks.py | 71 ++++++++++ 6 files changed, 225 insertions(+), 12 deletions(-) create mode 100644 src/nest/ensemble_rks/__init__.py create mode 100644 src/nest/ensemble_rks/ensemble_rks.py create mode 100644 src/nest/ensemble_rks/tests/test_ensemble_rks.py diff --git a/src/nest/__init__.py b/src/nest/__init__.py index f90fbce..3ce6f9c 100644 --- a/src/nest/__init__.py +++ b/src/nest/__init__.py @@ -15,8 +15,8 @@ """nest package.""" -from nest import dz0scf, nttda, sftda +from nest import dz0scf, ensemble_rks, nttda, sftda __version__ = "0.1.0" -__all__ = ["__version__", "dz0scf", "nttda", "sftda"] +__all__ = ["__version__", "dz0scf", "ensemble_rks", "nttda", "sftda"] diff --git a/src/nest/dz0scf/dz0scf.py b/src/nest/dz0scf/dz0scf.py index 7d32fa3..6382416 100644 --- a/src/nest/dz0scf/dz0scf.py +++ b/src/nest/dz0scf/dz0scf.py @@ -54,7 +54,10 @@ def evaluate_high_spin_energy(mf): vhf=veff, ) -class _DZ0VeffMixin: +class _DZ0VeffMixin: + reference_energy_semantics = 'high_spin_roks_energy_on_dz0_orbitals' + reference_energy_stationary = False + def get_veff( self, mol=None, @@ -81,8 +84,11 @@ def get_veff( vhf_last, hermi, ) - def high_spin_energy(self): - return evaluate_high_spin_energy(self) + def high_spin_energy(self): + return evaluate_high_spin_energy(self) + + def reference_energy(self): + return self.high_spin_energy() def nuc_grad_method(self): """Return the Dz0SCF analytic nuclear-gradient driver.""" diff --git a/src/nest/dz0scf/tests/test_dz0scf.py b/src/nest/dz0scf/tests/test_dz0scf.py index fc88305..d5f57ee 100644 --- a/src/nest/dz0scf/tests/test_dz0scf.py +++ b/src/nest/dz0scf/tests/test_dz0scf.py @@ -79,12 +79,18 @@ def test_svwn_dz0scf(self): ]) self.assertTrue(np.all(np.asarray(td_s.converged))) - np.testing.assert_allclose( - np.asarray(omega_s), - omega_s_ref, - rtol=0.0, - atol=1e-6, - ) + np.testing.assert_allclose( + np.asarray(omega_s), + omega_s_ref, + rtol=0.0, + atol=1e-6, + ) + np.testing.assert_allclose( + td_s.total_energies(), + mf.high_spin_energy() + omega_s, + rtol=0.0, + atol=1e-12, + ) td_t = NTTDA(mf) td_t.deltaS = 0 @@ -176,4 +182,4 @@ def test_b3lyp_dz0scf(self): omega_t_ref, rtol=0.0, atol=1e-6, - ) \ No newline at end of file + ) diff --git a/src/nest/ensemble_rks/__init__.py b/src/nest/ensemble_rks/__init__.py new file mode 100644 index 0000000..dff4b69 --- /dev/null +++ b/src/nest/ensemble_rks/__init__.py @@ -0,0 +1,3 @@ +from .ensemble_rks import EnsembleRKS + +__all__ = ['EnsembleRKS'] diff --git a/src/nest/ensemble_rks/ensemble_rks.py b/src/nest/ensemble_rks/ensemble_rks.py new file mode 100644 index 0000000..0b0e595 --- /dev/null +++ b/src/nest/ensemble_rks/ensemble_rks.py @@ -0,0 +1,127 @@ +"""Average-occupation restricted ensemble Kohn--Sham references.""" + +import numpy as np + +from pyscf import dft +from pyscf.lib import logger +from pyscf.scf import hf + + +class EnsembleRKS(dft.rks.RKS): + """RKS with fixed ``2/1/0`` occupations and zero spin density.""" + + is_ensemble_rks = True + _keys = dft.rks.RKS._keys | {'nopen', 'reference_semantics'} + + def __init__(self, mol, xc='LDA,VWN', nopen=None, reference_semantics='average'): + super().__init__(mol, xc=xc) + if reference_semantics not in ('average', 'roks'): + raise ValueError( + "reference_semantics must be 'average' or 'roks', got %r" + % (reference_semantics,)) + self.reference_semantics = reference_semantics + if nopen is None: + nopen = mol.spin + if isinstance(nopen, bool) or int(nopen) != nopen: + raise ValueError('nopen must be a non-negative integer') + self.nopen = int(nopen) + self._validate_ensemble() + + @property + def reference_energy_semantics(self): + if self.reference_semantics == 'roks': + return 'roks_energy_on_ensemble_rks_orbitals' + return 'average_occupation_ensemble_rks_energy' + + @property + def reference_energy_stationary(self): + return self.reference_semantics == 'average' + + def _validate_ensemble(self): + if self.nopen < 0: + raise ValueError('nopen must be a non-negative integer') + if self.nopen != self.mol.spin: + raise ValueError('nopen must match mol.spin for NTTDA') + if self.nopen > self.mol.nelectron: + raise ValueError('nopen cannot exceed the electron count') + if (self.mol.nelectron - self.nopen) % 2: + raise ValueError('electron count and nopen have inconsistent parity') + + @property + def nclosed(self): + return (self.mol.nelectron - self.nopen) // 2 + + def check_sanity(self): + out = hf.SCF.check_sanity(self) + if self.do_nlc() and self.do_disp() and self._numint.libxc.is_nlc(self.xc): + import warnings + warnings.warn( + f'nlc-type xc {self.xc} and disp {self.disp} may lead to double counting in NLC.' + ) + return out + + def get_occ(self, mo_energy=None, mo_coeff=None): + self._validate_ensemble() + if mo_energy is None: + mo_energy = self.mo_energy + mo_energy = np.asarray(mo_energy) + if self.nclosed + self.nopen > mo_energy.size: + raise RuntimeError('not enough orbitals for the requested ensemble occupations') + + order = np.argsort(mo_energy, kind='stable') + mo_occ = np.zeros_like(mo_energy) + mo_occ[order[:self.nclosed]] = 2 + mo_occ[order[self.nclosed:self.nclosed + self.nopen]] = 1 + if self.verbose >= logger.INFO: + logger.info(self, 'EnsembleRKS occupations = %s', mo_occ) + return mo_occ + + def make_rdm1s(self, mo_coeff=None, mo_occ=None): + dm = self.make_rdm1(mo_coeff, mo_occ) + dm_spin = np.asarray(dm) * 0.5 + return dm_spin, dm_spin.copy() + + def reference_energy(self): + if self.reference_semantics == 'roks': + if self.mo_coeff is None or self.mo_occ is None: + raise RuntimeError( + 'run EnsembleRKS.kernel() before evaluating the reference energy') + evaluator = dft.ROKS(self.mol).set( + xc=self.xc, + nlc=self.nlc, + max_memory=self.max_memory, + verbose=0, + ) + evaluator.grids = self.grids + evaluator.nlcgrids = self.nlcgrids + dm = evaluator.make_rdm1(self.mo_coeff, self.mo_occ) + hcore = evaluator.get_hcore(self.mol) + veff = evaluator.get_veff(self.mol, dm) + return float(evaluator.energy_tot(dm=dm, h1e=hcore, vhf=veff)) + if self.e_tot is None: + raise RuntimeError('run EnsembleRKS.kernel() before evaluating the reference energy') + return float(self.e_tot) + + def nuc_grad_method(self): + """Return the analytic gradient driver of the selected reference energy. + + For the ``average`` semantics this is the gradient of the stationary + ensemble energy (the stock RKS driver). For the ``roks`` semantics it + differentiates the fixed-orbital ROKS energy returned by + :meth:`reference_energy`, including its Z-vector orbital-relaxation + contribution and the nuclear-repulsion gradient. + """ + if self.reference_semantics == 'roks': + from nest.grad.nttda.reference import ReferenceGradients + return ReferenceGradients(self) + return super().nuc_grad_method() + + def get_grad(self, mo_coeff, mo_occ, fock=None): + mo_occ = np.asarray(mo_occ) + if fock is None: + dm = self.make_rdm1(mo_coeff, mo_occ) + fock = self.get_hcore(self.mol) + self.get_veff(self.mol, dm) + fock_mo = mo_coeff.conj().T @ fock @ mo_coeff + unique = hf.uniq_var_indices(mo_occ) + occupation_difference = mo_occ[None, :] - mo_occ[:, None] + return (fock_mo * occupation_difference)[unique] diff --git a/src/nest/ensemble_rks/tests/test_ensemble_rks.py b/src/nest/ensemble_rks/tests/test_ensemble_rks.py new file mode 100644 index 0000000..4d4155f --- /dev/null +++ b/src/nest/ensemble_rks/tests/test_ensemble_rks.py @@ -0,0 +1,71 @@ +import unittest + +import numpy as np +from pyscf import gto + +from nest.ensemble_rks import EnsembleRKS +from nest.nttda import NTTDA + + +class EnsembleRKSReferenceTest(unittest.TestCase): + def make_reference(self): + mol = gto.M( + atom=''' + O 0.64372820 0.14077399 -0.04477253 + O -0.64862595 -0.12779073 -0.05445498 + H 1.16027512 -0.65947800 0.36730132 + H -1.12109306 0.55561188 0.42651873 + ''', + basis='6-31g', + spin=2, + unit='Angstrom', + verbose=0, + ) + mf = EnsembleRKS(mol, xc='SVWN') + mf.conv_tol = 1e-12 + mf.conv_tol_grad = 1e-9 + mf.max_cycle = 100 + mf.grids.level = 3 + mf.grids.prune = None + mf.small_rho_cutoff = 0.0 + mf.kernel() + self.assertTrue(mf.converged) + return mf + + def test_average_occupation_reference_energy_is_used_for_total_energies(self): + mf = self.make_reference() + self.assertEqual(np.count_nonzero(mf.mo_occ == 1), mf.nopen) + self.assertEqual(np.count_nonzero(mf.mo_occ == 2), mf.nclosed) + self.assertAlmostEqual(mf.reference_energy(), mf.e_tot, places=14) + + td = NTTDA(mf).set(deltaS=0, nstates=2, conv_tol=1e-5, max_cycle=200, verbose=0) + td.kernel() + + self.assertTrue(np.all(td.converged)) + np.testing.assert_allclose( + td.total_energies(), + mf.reference_energy() + td.e, + atol=1e-13, + rtol=0, + ) + + def test_nobeta_does_not_change_an_ensemble_reference(self): + mf = self.make_reference() + roots = [] + for nobeta in (False, True): + td = NTTDA(mf).set( + deltaS=0, + nobeta=nobeta, + nstates=2, + conv_tol=1e-5, + max_cycle=200, + verbose=0, + ) + td.kernel() + self.assertTrue(np.all(td.converged)) + roots.append(td.e) + np.testing.assert_allclose(roots[0], roots[1], atol=1e-12, rtol=0) + + +if __name__ == '__main__': + unittest.main() From ac8dbe662cad4b30202fd0998e56555dc1f2f4da Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Mon, 31 Aug 2026 03:18:47 +0800 Subject: [PATCH 05/16] feat: add NTTDA analytic nuclear gradients for ROKS and EnsembleRKS Port the NTTDA response/gradient stack (deltaS = -1/0 channels, XC ledger kernels, ensemble orbital-response backend) and wire NTTDA.nuc_grad_method to nest.grad.nttda.Gradients. The linear solver uses pyscf.tdscf._lr_eig: the nest._lr_eig fork does not converge on the ensemble orbital Hessian. --- src/nest/grad/nttda/__init__.py | 214 +++ src/nest/grad/nttda/delta_s_minus_one.py | 1230 +++++++++++++++++ src/nest/grad/nttda/delta_s_zero.py | 1236 +++++++++++++++++ src/nest/grad/nttda/ensemble.py | 199 +++ src/nest/grad/nttda/roks.py | 334 +++++ src/nest/grad/nttda/xc.py | 1546 ++++++++++++++++++++++ src/nest/nttda/nttda.py | 494 ++++++- 7 files changed, 5202 insertions(+), 51 deletions(-) create mode 100644 src/nest/grad/nttda/__init__.py create mode 100644 src/nest/grad/nttda/delta_s_minus_one.py create mode 100644 src/nest/grad/nttda/delta_s_zero.py create mode 100644 src/nest/grad/nttda/ensemble.py create mode 100644 src/nest/grad/nttda/roks.py create mode 100644 src/nest/grad/nttda/xc.py diff --git a/src/nest/grad/nttda/__init__.py b/src/nest/grad/nttda/__init__.py new file mode 100644 index 0000000..5d62961 --- /dev/null +++ b/src/nest/grad/nttda/__init__.py @@ -0,0 +1,214 @@ +"""Analytic nuclear gradients for :mod:`nest.nttda`.""" + +import numpy as np + +from pyscf import dft, lib +from pyscf.grad import rhf as rhf_grad +from pyscf.lib import logger +from nest.ensemble_rks import EnsembleRKS +from nest.nttda import NTTDA + +from . import delta_s_minus_one, delta_s_zero + + + +def _normalized_amplitude(xy): + vector = np.asarray(xy[0]).ravel() + return vector / np.linalg.norm(vector) + + +def _copy_td_settings(source, target): + for name in ( + "deltaS", "nobeta", "nstates", "conv_tol", "lindep", + "max_cycle", "max_memory"): + setattr(target, name, getattr(source, name)) + target.verbose = 0 + return target + + +def _displaced_reference(source, mol, fixed_grid): + if isinstance(source, EnsembleRKS): + reference = EnsembleRKS(mol, xc=source.xc, nopen=source.nopen) + elif isinstance(source, dft.KohnShamDFT): + reference = dft.ROKS(mol) + else: + reference = source.__class__(mol) + for name in ( + "conv_tol", "conv_tol_grad", "max_cycle", "max_memory", + "level_shift", "damp"): + if hasattr(source, name): + setattr(reference, name, getattr(source, name)) + reference.verbose = 0 + if isinstance(source, dft.KohnShamDFT): + reference.xc = source.xc + reference.nlc = source.nlc + reference.grids.level = source.grids.level + reference.grids.prune = source.grids.prune + if fixed_grid and source.grids.coords is not None: + reference.grids.coords = np.array(source.grids.coords, copy=True) + reference.grids.weights = np.array(source.grids.weights, copy=True) + reference.grids.non0tab = None + reference.grids.verbose = 0 + return reference + + +class Gradients(rhf_grad.GradientsBase): + """NTTDA gradients, including finite differences for EnsembleRKS.""" + + _keys = rhf_grad.GradientsBase._keys | { + "state", "method", "step", "fixed_grid", "root_overlap_tol", + "cphf_conv_tol", "cphf_max_cycle", + } + + def __init__(self, tdobj): + super().__init__(tdobj) + self.state = 1 + self.method = "analytic" + self.step = 1e-3 + self.fixed_grid = ( + isinstance(tdobj._scf, dft.KohnShamDFT) + and not isinstance(tdobj._scf, EnsembleRKS) + ) + self.root_overlap_tol = 0.5 + self.cphf_conv_tol = 1e-12 + self.cphf_max_cycle = None + self.nttda_details = None + + def dump_flags(self, verbose=None): + log = logger.new_logger(self, verbose) + log.info("******** NTTDA nuclear gradients ********") + log.info("state = %d", self.state) + log.info("deltaS = %d", self.base.deltaS) + log.info("nobeta = %s", self.base.nobeta) + log.info("method = %s", self.method) + log.info("fixed_grid = %s", self.fixed_grid) + if self.method == "finite_diff": + log.info("finite-difference step = %.6g Bohr", self.step) + return self + + def grad_nuc(self, atmlst=None): + """Ground-state reference gradient, including nuclear repulsion.""" + if atmlst is not None: + atmlst = list(atmlst) + return self.base._scf.nuc_grad_method().kernel(atmlst=atmlst) + + def _analytic_components(self, xy, atmlst): + tdobj = self.base + options = { + "atmlst": atmlst, + "tolerance": self.cphf_conv_tol, + "max_cycle": self.cphf_max_cycle, + } + if tdobj.deltaS == -1: + return delta_s_minus_one.grad_elec( + self, tdobj, xy, **options, + ) + if tdobj.deltaS == 0: + return delta_s_zero.grad_elec( + self, tdobj, xy, **options, + ) + if tdobj.deltaS == 1: + raise NotImplementedError( + "Analytic NTTDA gradients are not implemented for deltaS=1; " + "use method='finite_diff'." + ) + raise ValueError("deltaS must be -1, 0, or 1") + + def grad_elec(self, xy, atmlst=None): + """Return the analytic excitation-energy derivative ``d omega/dR``.""" + if atmlst is None: + atmlst = range(self.mol.natm) + components = self._analytic_components(xy, tuple(atmlst)) + self.nttda_details = components + return components.total + + def _energy_at(self, coords, reference_amplitude): + mol = self.mol.copy() + mol.set_geom_(coords, unit="Bohr") + mf = _displaced_reference(self.base._scf, mol, self.fixed_grid) + mf.kernel(dm0=self.base._scf.make_rdm1()) + if not mf.converged: + raise RuntimeError("displaced NTTDA reference did not converge") + tdobj = _copy_td_settings(self.base, NTTDA(mf)) + tdobj.kernel() + overlaps = np.asarray([ + abs(np.vdot(reference_amplitude, _normalized_amplitude(xy))) + for xy in tdobj.xy + ]) + root = int(np.argmax(overlaps)) + if overlaps[root] < self.root_overlap_tol: + raise RuntimeError( + "NTTDA state tracking overlap %.6f is below %.6f" % + (overlaps[root], self.root_overlap_tol) + ) + return mf.e_tot + tdobj.e[root] + + def _finite_difference(self, atmlst): + coords0 = self.mol.atom_coords() + reference_amplitude = _normalized_amplitude( + self.base.xy[self.state - 1], + ) + result = np.zeros((len(atmlst), 3)) + for index, atom in enumerate(atmlst): + for xyz in range(3): + coords_plus = coords0.copy() + coords_minus = coords0.copy() + coords_plus[atom, xyz] += self.step + coords_minus[atom, xyz] -= self.step + energy_plus = self._energy_at( + coords_plus, reference_amplitude, + ) + energy_minus = self._energy_at( + coords_minus, reference_amplitude, + ) + result[index, xyz] = ( + (energy_plus - energy_minus) / (2.0 * self.step) + ) + return result + + def kernel(self, state=None, atmlst=None, method=None, step=None): + """Return ``d(E_reference + omega_state)/dR`` in Eh/Bohr.""" + if state is not None: + self.state = state + if method is not None: + self.method = method + if step is not None: + self.step = step + if atmlst is None: + atmlst = self.atmlst + else: + self.atmlst = atmlst + if atmlst is None: + atmlst = range(self.mol.natm) + atmlst = tuple(atmlst) + + if self.state == 0: + return self.grad_nuc(atmlst=atmlst) + if self.base.xy is None: + self.base.run() + if not 1 <= self.state <= len(self.base.xy): + raise ValueError("state must be in [1, %d]" % len(self.base.xy)) + if self.verbose >= logger.INFO: + self.dump_flags() + + if self.method == "analytic": + excitation = self.grad_elec( + self.base.xy[self.state - 1], atmlst=atmlst, + ) + result = self.grad_nuc(atmlst=atmlst) + excitation + elif self.method == "finite_diff": + result = self._finite_difference(atmlst) + else: + raise ValueError("unknown NTTDA gradient method %s" % self.method) + self.de = result + if self.mol.symmetry: + self.de = self.symmetrize(self.de, atmlst) + self._finalize() + return self.de + + grad = lib.alias(kernel, alias_name="grad") + + +Grad = Gradients + +__all__ = ["Grad", "Gradients"] diff --git a/src/nest/grad/nttda/delta_s_minus_one.py b/src/nest/grad/nttda/delta_s_minus_one.py new file mode 100644 index 0000000..b6ef3ac --- /dev/null +++ b/src/nest/grad/nttda/delta_s_minus_one.py @@ -0,0 +1,1230 @@ +"""Analytic gradient for current NTTDA ``deltaS=-1``. + +This module owns the complete spin-lowering amplitude, Fock, response, and +AO derivative formulas. It does not import the ``deltaS=0`` channel. +""" + +from dataclasses import dataclass + +import numpy as np + +from pyscf import dft, lib +from nest.nttda import nttda as nttda_mod +from nest.nttda.nttda import gen_rohf_response_sfd + +from . import xc as xc_backend +from .roks import finish_gradient + + + +# Orbital spaces and native amplitudes + +@dataclass(frozen=True) +class OrbitalSpaces: + """Closed, open, and virtual spatial-orbital partitions.""" + + closed: np.ndarray + open: np.ndarray + virtual: np.ndarray + c_closed: np.ndarray + c_open: np.ndarray + c_virtual: np.ndarray + + @property + def spin(self): + return 0.5 * len(self.open) + + +@dataclass(frozen=True) +class SpinLoweringAmplitudes: + """Four native blocks used by ``NTTDA(deltaS=-1)``.""" + + co: np.ndarray + cv: np.ndarray + oo: np.ndarray + ov: np.ndarray + + +def orbital_spaces(tdobj): + """Return the ROKS ``C/O/V`` orbital partition used by NTTDA.""" + mf = tdobj._scf + occ = np.asarray(mf.mo_occ) + if occ.ndim != 1: + raise ValueError("NTTDA gradients require spatial ROKS orbitals") + closed = np.flatnonzero(occ == 2) + open_ = np.flatnonzero(occ == 1) + virtual = np.flatnonzero(occ == 0) + coeff = np.asarray(mf.mo_coeff) + return OrbitalSpaces( + closed=closed, + open=open_, + virtual=virtual, + c_closed=coeff[:, closed], + c_open=coeff[:, open_], + c_virtual=coeff[:, virtual], + ) + + +def pair_density(c_left, coefficient, c_right): + """Build ``C_left coefficient C_right^T`` without symmetrizing it.""" + return c_left @ np.asarray(coefficient) @ c_right.conj().T + + +def split_spin_lowering(tdobj, xy): + """Split a lowering-channel amplitude into ``CO/CV/OO/OV`` blocks.""" + spaces = orbital_spaces(tdobj) + if spaces.spin < 1.0: + raise ValueError("NTTDA deltaS=-1 requires reference spin Si >= 1") + vector = xy[0] if isinstance(xy, (tuple, list)) else xy + vector = np.asarray(vector) + nc = len(spaces.closed) + no = len(spaces.open) + nv = len(spaces.virtual) + expected = (nc + no, no + nv) + if vector.size != expected[0] * expected[1]: + raise ValueError( + "deltaS=-1 amplitude has size %d; expected %d" % + (vector.size, expected[0] * expected[1]) + ) + vector = vector.reshape(expected) + return spaces, SpinLoweringAmplitudes( + co=vector[:nc, :no], + cv=vector[:nc, no:], + oo=vector[nc:, :no], + ov=vector[nc:, no:], + ) + + +def spin_lowering_transition_densities(tdobj, xy): + """Directed alpha-occupied to beta-target transition densities.""" + spaces, amp = split_spin_lowering(tdobj, xy) + return spaces, amp, { + "CO": pair_density(spaces.c_open, amp.co.T, spaces.c_closed), + "CV": pair_density(spaces.c_virtual, amp.cv.T, spaces.c_closed), + "OO": pair_density(spaces.c_open, amp.oo.T, spaces.c_open), + "OV": pair_density(spaces.c_virtual, amp.ov.T, spaces.c_open), + } + + +def spin_lowering_block_data(spaces, amplitudes): + """MO index/factor map for variations of lowering transition densities.""" + return { + "CO": (spaces.open, spaces.closed, amplitudes.co.T), + "CV": (spaces.virtual, spaces.closed, amplitudes.cv.T), + "OO": (spaces.open, spaces.open, amplitudes.oo.T), + "OV": (spaces.virtual, spaces.open, amplitudes.ov.T), + } + + +# Channel-local immutable records + +@dataclass(frozen=True) +class FockProjection: + """One scalar term ``Tr[P (weight_f0 F0 + weight_fz Fz)]``.""" + + name: str + left_indices: np.ndarray + left_orbitals: np.ndarray + coefficient: np.ndarray + right_indices: np.ndarray + right_orbitals: np.ndarray + weight_f0: float + weight_fz: float + + def density(self): + return pair_density( + self.left_orbitals, self.coefficient, self.right_orbitals, + ) + + +@dataclass(frozen=True) +class ResponseTerm: + """Directed response term from one source density to one target block.""" + + target: str + source: str + vref0: float + vref1: float + + +# Shared-response evaluators copied into the lowering channel + +def _fxc_reference(tdobj): + mf = tdobj._scf + ni = mf._numint + fxc = ni.cache_xc_kernel( + mf.mol, mf.grids, mf.xc, mf.mo_coeff, mf.mo_occ, 1, + )[2] + return 0.5 * ( + fxc[0, :, 0] - fxc[0, :, 1] + - fxc[1, :, 0] + fxc[1, :, 1] + ) + + +def _apply_reference_responses(tdobj, densities, max_memory=None): + """Return separate ``vref0`` and ``vref1`` actions for each density.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if max_memory is None: + max_memory = tdobj.max_memory + labels = tuple(densities) + dms = np.asarray([densities[label] for label in labels]) + xctype = ni._xc_type(mf.xc) + if xctype == "HF": + vref0 = np.zeros_like(dms) + vref1 = np.zeros_like(dms) + else: + fxc_ref = _fxc_reference(tdobj) + vref0 = ni.nr_rks_fxc( + mol, mf.grids, mf.xc, None, dms, 0, 0, + None, None, fxc_ref, max_memory=max_memory, + ) + if xctype == "LDA": + vref1 = ni.nr_rks_fxc( + mol, mf.grids, mf.xc, None, dms, 0, 0, + None, None, fxc_ref, max_memory=max_memory, + ) + elif xctype == "GGA": + vref1 = nttda_mod.nr_rks_fxc1_gga( + ni, mol, mf.grids, mf.xc, dms, fxc_ref, + max_memory=max_memory, + ) + elif xctype == "MGGA": + vref1 = nttda_mod.nr_rks_fxc1_mgga( + ni, mol, mf.grids, mf.xc, dms, fxc_ref, + max_memory=max_memory, + ) + else: + raise NotImplementedError( + "NTTDA spin-lowering response does not support XC type %s" % + xctype + ) + + omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + if ni.libxc.is_hybrid_xc(mf.xc): + vref0 -= hyb * mf.get_k(mol, dms, hermi=0) + vref1 -= hyb * mf.get_j(mol, dms, hermi=0) + if omega != 0: + scale = alpha - hyb + vref0 -= scale * mf.get_k(mol, dms, hermi=0, omega=omega) + vref1 -= scale * mf.get_j(mol, dms, hermi=0, omega=omega) + return ( + {label: value for label, value in zip(labels, vref0)}, + {label: value for label, value in zip(labels, vref1)}, + ) + + +def _apply_hfx_responses(tdobj, densities): + """Return only the hybrid/RSH J/K portions of ``vref0/vref1``.""" + mf = tdobj._scf + labels = tuple(densities) + dms = np.asarray([densities[label] for label in labels]) + vref0 = np.zeros_like(dms) + vref1 = np.zeros_like(dms) + ni = mf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mf.mol.spin) + if ni.libxc.is_hybrid_xc(mf.xc): + vref0 -= hybrid * mf.get_k(mf.mol, dms, hermi=0) + vref1 -= hybrid * mf.get_j(mf.mol, dms, hermi=0) + if omega != 0: + scale = alpha - hybrid + vref0 -= scale * mf.get_k( + mf.mol, dms, hermi=0, omega=omega, + ) + vref1 -= scale * mf.get_j( + mf.mol, dms, hermi=0, omega=omega, + ) + return ( + {label: value for label, value in zip(labels, vref0)}, + {label: value for label, value in zip(labels, vref1)}, + ) + + +# Reference Fock response helper + +def _fock_response_q(tdobj, p_alpha, p_beta): + """Reference-density derivative of a spin-resolved Fock scalar.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + if getattr(mf, "is_ensemble_rks", False): + occupation = np.asarray(mf.mo_occ) + probe = np.asarray(p_alpha) + np.asarray(p_beta) + potential = mf.gen_response(hermi=0)(probe.T) + q_total = ( + mo.conj().T @ (potential + potential.T) @ mo + ) * occupation[None, :] + return 0.5 * q_total, 0.5 * q_total + occ_alpha = (np.asarray(mf.mo_occ) > 0).astype(float) + occ_beta = (np.asarray(mf.mo_occ) == 2).astype(float) + if (isinstance(mf, dft.KohnShamDFT) + and mf._numint._xc_type(mf.xc) != "HF"): + unrestricted = mf.to_uks() + unrestricted.verbose = 0 + v_alpha, v_beta = unrestricted.gen_response(hermi=0)( + np.asarray((p_alpha.T, p_beta.T)) + ) + else: + p_total = p_alpha + p_beta + coulomb = mf.get_j(mf.mol, p_total.T, hermi=0) + v_alpha = coulomb - mf.get_k(mf.mol, p_alpha.T, hermi=0) + v_beta = coulomb - mf.get_k(mf.mol, p_beta.T, hermi=0) + q_alpha = ( + mo.conj().T @ (v_alpha + v_alpha.T) @ mo + ) * occ_alpha[None, :] + q_beta = ( + mo.conj().T @ (v_beta + v_beta.T) @ mo + ) * occ_beta[None, :] + return q_alpha, q_beta + + +# Complete lowering scalar and M-matrix ledger + +def spin_lowering_response_terms(spin): + """Directed ``vref0/vref1`` coefficients in ``gen_rohf_response_sfd``.""" + denominator = 2.0 * spin - 1.0 + a = np.sqrt((2.0 * spin + 1.0) / (2.0 * spin)) + b = np.sqrt(2.0 * spin / denominator) + c = np.sqrt((2.0 * spin + 1.0) / denominator) + return ( + ResponseTerm("CO", "CO", 1.0, 1.0 / denominator), + ResponseTerm("CO", "CV", a, 0.0), + ResponseTerm("CO", "OO", b, 0.0), + ResponseTerm( + "CO", "OV", 2.0 * spin / denominator, + -1.0 / denominator, + ), + ResponseTerm("CV", "CO", a, 0.0), + ResponseTerm("CV", "CV", 1.0, 0.0), + ResponseTerm("CV", "OO", c, 0.0), + ResponseTerm("CV", "OV", a, 0.0), + ResponseTerm("OO", "CO", b, 0.0), + ResponseTerm("OO", "CV", c, 0.0), + ResponseTerm("OO", "OO", 1.0, 0.0), + ResponseTerm("OO", "OV", b, 0.0), + ResponseTerm( + "OV", "CO", 2.0 * spin / denominator, + -1.0 / denominator, + ), + ResponseTerm("OV", "CV", a, 0.0), + ResponseTerm("OV", "OO", b, 0.0), + ResponseTerm("OV", "OV", 1.0, 1.0 / denominator), + ) + + +def spin_lowering_fock0_fockz(tdobj, max_memory=None): + """Operators used by the current lowering-channel action.""" + mf = tdobj._scf + if max_memory is None: + max_memory = tdobj.max_memory + _response, fockz = gen_rohf_response_sfd( + mf, + mo_coeff=mf.mo_coeff, + mo_occ=mf.mo_occ, + hermi=0, + max_memory=max_memory, + ) + return nttda_mod._reference_fock0(mf, tdobj.nobeta), fockz + + +def spin_lowering_fock_projections(tdobj, xy): + """Complete explicit-Fock ledger of ``X.T A_sfd X``.""" + spaces, amplitudes = split_spin_lowering(tdobj, xy) + c = spaces.c_closed + o = spaces.c_open + v = spaces.c_virtual + block_data = ( + ("C", "O", amplitudes.co), + ("C", "V", amplitudes.cv), + ("O", "O", amplitudes.oo), + ("O", "V", amplitudes.ov), + ) + orbital_data = { + "C": (spaces.closed, c), + "O": (spaces.open, o), + "V": (spaces.virtual, v), + } + terms = [] + + def add(name, left_label, coefficient, right_label, f0, fz): + left_indices, left_orbitals = orbital_data[left_label] + right_indices, right_orbitals = orbital_data[right_label] + coefficient = np.asarray(coefficient) + if coefficient.size: + terms.append(FockProjection( + name=name, + left_indices=left_indices, + left_orbitals=left_orbitals, + coefficient=coefficient, + right_indices=right_indices, + right_orbitals=right_orbitals, + weight_f0=float(f0), + weight_fz=float(fz), + )) + + # Ordinary alpha-to-beta spin-flip Fock difference. + for row_left, column_left, x_left in block_data: + for row_right, column_right, x_right in block_data: + if row_left == row_right: + add( + "base-beta-%s%s-%s%s" % ( + row_left, column_left, row_right, column_right, + ), + column_left, + x_left.T @ x_right, + column_right, + 1.0, + -1.0, + ) + if column_left == column_right: + add( + "base-alpha-%s%s-%s%s" % ( + row_left, column_left, row_right, column_right, + ), + row_right, + -(x_right @ x_left.T), + row_left, + 1.0, + 1.0, + ) + + # Tensor spin-adaptation correction, expressed in the same F0/Fz basis. + spin = spaces.spin + trace_oo = float(np.trace(amplitudes.oo)) + eta = np.sqrt((2.0 * spin + 1.0) / (2.0 * spin)) - 1.0 + gamma = np.sqrt((2.0 * spin + 1.0) / (2.0 * spin - 1.0)) + zeta = np.sqrt(2.0 * spin / (2.0 * spin - 1.0)) - 1.0 + chi = 1.0 / np.sqrt(2.0 * spin * (2.0 * spin - 1.0)) + t_cc = ( + amplitudes.cv @ amplitudes.cv.T / spin + + amplitudes.co @ amplitudes.co.T * 2.0 / (2.0 * spin - 1.0) + ) + t_vv = ( + amplitudes.cv.T @ amplitudes.cv / spin + + amplitudes.ov.T @ amplitudes.ov * 2.0 / (2.0 * spin - 1.0) + ) + t_cv = gamma * (1.0 + 1.0 / spin) * trace_oo * amplitudes.cv + t_beta_vo = ( + 2.0 * eta * amplitudes.cv.T @ amplitudes.co + + 2.0 * zeta * amplitudes.ov.T @ amplitudes.oo + ) + t_beta_co = 2.0 * chi * trace_oo * amplitudes.co + t_alpha_oc = ( + -2.0 * eta * amplitudes.cv @ amplitudes.ov.T + - 2.0 * zeta * amplitudes.co @ amplitudes.oo.T + ).T + t_alpha_vo = -2.0 * chi * trace_oo * amplitudes.ov.T + + add("adapt-spin-cc", "C", t_cc, "C", 0.0, -1.0) + add("adapt-spin-vv", "V", t_vv, "V", 0.0, -1.0) + add("adapt-spin-cv", "C", t_cv, "V", 0.0, -1.0) + add("adapt-beta-vo", "V", t_beta_vo, "O", 1.0, -1.0) + add("adapt-beta-co", "C", t_beta_co, "O", 1.0, -1.0) + add("adapt-alpha-oc", "O", t_alpha_oc, "C", 1.0, 1.0) + add("adapt-alpha-vo", "V", t_alpha_vo, "O", 1.0, 1.0) + return tuple(terms) + + +def spin_lowering_fock_probes(tdobj, xy): + """Return AO probes ``P0,Pz`` for the lowering Fock ledger.""" + nao = tdobj.mol.nao_nr() + p0 = np.zeros((nao, nao)) + pz = np.zeros_like(p0) + for term in spin_lowering_fock_projections(tdobj, xy): + density = term.density() + p0 += term.weight_f0 * density + pz += term.weight_fz * density + return p0, pz + + +def spin_lowering_fock_scalar(tdobj, xy, max_memory=None): + fock0, fockz = spin_lowering_fock0_fockz( + tdobj, max_memory=max_memory, + ) + p0, pz = spin_lowering_fock_probes(tdobj, xy) + return float( + lib.einsum("pq,pq->", p0, fock0) + + lib.einsum("pq,pq->", pz, fockz) + ) + + +def spin_lowering_response_scalar(tdobj, xy, max_memory=None): + spaces, _amplitudes, densities = spin_lowering_transition_densities( + tdobj, xy, + ) + vref0, vref1 = _apply_reference_responses( + tdobj, densities, max_memory=max_memory, + ) + value = 0.0 + for term in spin_lowering_response_terms(spaces.spin): + target = densities[term.target] + if term.vref0: + value += term.vref0 * lib.einsum( + "pq,pq->", target, vref0[term.source], + ) + if term.vref1: + value += term.vref1 * lib.einsum( + "pq,pq->", target, vref1[term.source], + ) + return float(value) + + +def spin_lowering_ledger_scalar(tdobj, xy, max_memory=None): + """Independent reconstruction of ``X.T gen_vind_sfd(X)``.""" + return ( + spin_lowering_fock_scalar(tdobj, xy, max_memory=max_memory) + + spin_lowering_response_scalar(tdobj, xy, max_memory=max_memory) + ) + + +def spin_lowering_action_scalar(tdobj, xy): + vector = xy[0] if isinstance(xy, (tuple, list)) else xy + vector = np.asarray(vector) + vind, _diagonal = tdobj.gen_vind_sfd() + action = vind(vector.reshape(1, -1)).reshape(vector.shape) + return float(np.vdot(vector, action).real) + + +def _response_potentials(densities, vref0, vref1, terms): + potentials = {label: np.zeros_like(dm) for label, dm in densities.items()} + for term in terms: + if term.vref0: + potentials[term.target] += term.vref0 * vref0[term.source] + potentials[term.source] += term.vref0 * vref0[term.target] + if term.vref1: + potentials[term.target] += term.vref1 * vref1[term.source] + potentials[term.source] += term.vref1 * vref1[term.target] + return potentials + + +def _project_transition_potentials(tdobj, blocks, potentials): + mo = np.asarray(tdobj._scf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + for label, (target, source, coefficient) in blocks.items(): + potential = mo.conj().T @ potentials[label] @ mo + q_beta[:, target] += potential[:, source] @ coefficient.T + q_alpha[:, source] += potential[target, :].T @ coefficient + return q_alpha, q_beta + + +def spin_lowering_response_projection_q( + tdobj, xy, max_memory=None, hfx_only=False): + """Transition-factor derivative of the lowering response scalar.""" + spaces, amplitudes, densities = spin_lowering_transition_densities( + tdobj, xy, + ) + blocks = spin_lowering_block_data(spaces, amplitudes) + if hfx_only: + vref0, vref1 = _apply_hfx_responses(tdobj, densities) + else: + vref0, vref1 = _apply_reference_responses( + tdobj, densities, max_memory=max_memory, + ) + potentials = _response_potentials( + densities, + vref0, + vref1, + spin_lowering_response_terms(spaces.spin), + ) + return _project_transition_potentials(tdobj, blocks, potentials) + + +def spin_lowering_fock_q(tdobj, xy, max_memory=None): + """Explicit-Fock projection and reference-density response M matrices.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + nmo = mo.shape[1] + fock0, fockz = spin_lowering_fock0_fockz( + tdobj, max_memory=max_memory, + ) + fock0_mo = mo.conj().T @ fock0 @ mo + fockz_mo = mo.conj().T @ fockz @ mo + q_alpha = np.zeros((nmo, nmo)) + q_beta = np.zeros_like(q_alpha) + is_hf = mf._numint._xc_type(mf.xc) == "HF" + + for term in spin_lowering_fock_projections(tdobj, xy): + left = term.left_indices + right = term.right_indices + coefficient = term.coefficient + + def project(target, operator, scale): + if scale: + target[:, left] += ( + scale * operator[:, right] @ coefficient.T + ) + target[:, right] += ( + scale * operator[:, left] @ coefficient + ) + + project(q_alpha, fock0_mo, 0.5 * term.weight_f0) + project(q_beta, fock0_mo, 0.5 * term.weight_f0) + if is_hf: + project(q_alpha, fockz_mo, 0.5 * term.weight_fz) + project(q_beta, fockz_mo, 0.5 * term.weight_fz) + else: + project(q_alpha, fockz_mo, term.weight_fz) + + p0, pz = spin_lowering_fock_probes(tdobj, xy) + p_alpha = 0.5 * p0 + p_beta = 0.5 * p0 + if is_hf: + p_alpha = p_alpha + 0.5 * pz + p_beta = p_beta - 0.5 * pz + response_alpha, response_beta = _fock_response_q( + tdobj, p_alpha, p_beta, + ) + q_alpha += response_alpha + q_beta += response_beta + return q_alpha, q_beta + + +# AO J/K nuclear derivatives + +def _as_derivative_stack(array): + array = np.asarray(array) + if array.ndim == 3: + array = array[None] + return array + + +def _density_key(density): + density = np.asarray(density) + data = density.__array_interface__["data"][0] + return data, density.shape, density.strides, density.dtype.str + + +@dataclass(frozen=True) +class _JKDerivativeTerm: + """One fixed-AO bilinear derivative with a named output slot.""" + + left: np.ndarray + right: np.ndarray + scale: float + omega: float + slot: object + + +class _JKDerivativeLedger: + """Spin-lowering scheduler for fixed-AO J/K derivatives.""" + + def __init__(self): + self._terms = {"j": [], "k": []} + + def add(self, operator, slot, terms): + self._terms[operator].extend( + _JKDerivativeTerm(left, right, scale, omega, slot) + for left, right, scale, omega in terms + if scale != 0.0 + ) + + def contract(self, gradient_driver, mol, atoms, slots=()): + atoms = tuple(atoms) + shape = (len(atoms), 3) + gradients = {slot: np.zeros(shape) for slot in slots} + for operator in ("j", "k"): + for term in self._terms[operator]: + gradients.setdefault(term.slot, np.zeros(shape)) + _contract_derivative_terms( + gradients, + gradient_driver, + mol, + atoms, + mol.offset_nr_by_atom(), + self._terms[operator], + operator, + ) + return gradients + + +def _term_densities(term, exchange): + left, right = term.left, term.right + if exchange: + return left, right, left.T, right.T + return left, right + + +def _density_batches(terms, exchange, max_memory, nao): + """Group bilinear terms while bounding derivative-potential storage.""" + minimum = 4 if exchange else 2 + bytes_per_density = 4 * nao * nao * np.dtype(float).itemsize + batch_limit = max( + minimum, + int(0.2 * max_memory * 1e6 / bytes_per_density), + ) + batch = [] + keys = set() + for term in terms: + term_keys = { + _density_key(density) + for density in _term_densities(term, exchange) + } + if batch and len(keys | term_keys) > batch_limit: + yield batch + batch = [] + keys = set() + batch.append(term) + keys.update(term_keys) + if batch: + yield batch + + +def _jk_derivative_potentials( + gradient_driver, mol, terms, operator, omega): + exchange = operator == "k" + densities = {} + for term in terms: + for density in _term_densities(term, exchange): + density = np.asarray(density) + densities.setdefault(_density_key(density), density) + keys = tuple(densities) + stack = np.asarray([densities[key] for key in keys]) + if operator == "j": + if omega is None: + values = gradient_driver.get_j(mol, stack, hermi=0) + else: + values = gradient_driver.get_j( + mol, stack, hermi=0, omega=omega, + ) + else: + if omega is None: + values = gradient_driver.get_k(mol, stack, hermi=0) + else: + values = gradient_driver.get_k( + mol, stack, hermi=0, omega=omega, + ) + values = _as_derivative_stack(values) + return dict(zip(keys, values)) + + +def _contract_derivative_terms( + gradients, gradient_driver, mol, atoms, offsets, terms, + operator): + if not atoms: + return + terms_by_omega = {} + for term in terms: + terms_by_omega.setdefault(term.omega, []).append(term) + exchange = operator == "k" + for omega, omega_terms in terms_by_omega.items(): + for batch in _density_batches( + omega_terms, exchange, gradient_driver.max_memory, + mol.nao_nr()): + potentials = _jk_derivative_potentials( + gradient_driver, mol, batch, operator, omega, + ) + for term in batch: + left = np.asarray(term.left) + right = np.asarray(term.right) + right_derivative = potentials[_density_key(right)] + left_derivative = potentials[_density_key(left)] + if exchange: + right_t_derivative = potentials[ + _density_key(right.T) + ] + left_t_derivative = potentials[_density_key(left.T)] + for k, atom in enumerate(atoms): + p0, p1 = offsets[atom][2:] + if exchange: + value = lib.einsum( + "xpq,pq->x", + right_derivative[:, p0:p1, :], + left[p0:p1, :], + ) + value += lib.einsum( + "xqp,pq->x", + right_t_derivative[:, p0:p1, :], + left[:, p0:p1], + ) + value += lib.einsum( + "xpq,pq->x", + left_derivative[:, p0:p1, :], + right[p0:p1, :], + ) + value += lib.einsum( + "xqp,pq->x", + left_t_derivative[:, p0:p1, :], + right[:, p0:p1], + ) + else: + value = lib.einsum( + "xpq,pq->x", + right_derivative[:, p0:p1], + left[p0:p1], + ) + value += lib.einsum( + "xpq,qp->x", + right_derivative[:, p0:p1], + left[:, p0:p1], + ) + value += lib.einsum( + "xpq,pq->x", + left_derivative[:, p0:p1], + right[p0:p1], + ) + value += lib.einsum( + "xpq,qp->x", + left_derivative[:, p0:p1], + right[:, p0:p1], + ) + gradients[term.slot][k] += term.scale * value + + +def _reference_spin_densities(tdobj): + mf = tdobj._scf + if getattr(mf, "is_ensemble_rks", False): + return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) + mo = np.asarray(mf.mo_coeff) + return ( + mo[:, mf.mo_occ > 0] @ mo[:, mf.mo_occ > 0].T, + mo[:, mf.mo_occ == 2] @ mo[:, mf.mo_occ == 2].T, + ) + + +def _spin_probe_stacks(p_alpha, p_beta): + p_alpha = np.asarray(p_alpha) + p_beta = np.asarray(p_beta) + single_probe = p_alpha.ndim == 2 + if single_probe: + p_alpha = p_alpha[None] + p_beta = p_beta[None] + return p_alpha, p_beta, single_probe + + +def spin_fock_direct_dft( + gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, + nobeta_p0=None, jk_ledger=None, output_slots=None): + """Differentiate one or more ordinary UKS Fock scalar probes. + + The optional ``nobeta_p0`` correction belongs to the first, explicit-direct + probe in the batch. + """ + mf = tdobj._scf + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + p_alpha, p_beta, single_probe = _spin_probe_stacks( + p_alpha, p_beta, + ) + if output_slots is None: + output_slots = tuple(range(len(p_alpha))) + p_total = p_alpha + p_beta + density_alpha, density_beta = _reference_spin_densities(tdobj) + gradient = np.zeros((len(p_alpha), len(atmlst), 3)) + hcore_derivative = mf.nuc_grad_method().hcore_generator(mol) + for k, atom in enumerate(atmlst): + gradient[:, k] += lib.einsum( + "npq,xpq->nx", p_total, hcore_derivative(atom), + ) + ni = mf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + for probe in range(len(p_alpha)): + j_terms = [ + (p_total[probe], density_alpha, 1.0, None), + (p_total[probe], density_beta, 1.0, None), + ] + k_terms = [] + if ni.libxc.is_hybrid_xc(mf.xc): + k_terms.extend(( + (p_alpha[probe], density_alpha, -hybrid, None), + (p_beta[probe], density_beta, -hybrid, None), + )) + if omega != 0: + long_range = -(alpha - hybrid) + k_terms.extend(( + (p_alpha[probe], density_alpha, long_range, omega), + (p_beta[probe], density_beta, long_range, omega), + )) + ledger.add("j", output_slots[probe], j_terms) + ledger.add("k", output_slots[probe], k_terms) + xctype = ni._xc_type(mf.xc) + if xctype == "LDA": + derivative_contractor = xc_backend.contract_lda_vxc_derivative + elif xctype == "GGA": + derivative_contractor = xc_backend.contract_gga_vxc_derivative + elif xctype == "MGGA": + derivative_contractor = xc_backend.contract_mgga_vxc_derivative + else: + raise NotImplementedError( + "ordinary Fock direct derivative is not implemented for %s" % + xctype + ) + if (nobeta_p0 is not None and tdobj.nobeta + and not getattr(mf, "is_ensemble_rks", False)): + density0 = 0.5 * (density_alpha + density_beta) + actual_probe_alpha = np.array(p_alpha, copy=True) + actual_probe_beta = np.array(p_beta, copy=True) + actual_probe_alpha[0] -= 0.5 * nobeta_p0 + actual_probe_beta[0] -= 0.5 * nobeta_p0 + else: + density0 = None + actual_probe_alpha = p_alpha + actual_probe_beta = p_beta + gradient += derivative_contractor( + mf, + density_alpha, + density_beta, + actual_probe_alpha, + actual_probe_beta, + atmlst=atmlst, + max_memory=gradient_driver.max_memory, + ) + if density0 is not None: + gradient[0] += derivative_contractor( + mf, + density0, + density0, + 0.5 * nobeta_p0, + 0.5 * nobeta_p0, + atmlst=atmlst, + max_memory=gradient_driver.max_memory, + ) + if jk_ledger is None: + contractions = local_ledger.contract( + gradient_driver, mol, atmlst, slots=output_slots, + ) + for probe, slot in enumerate(output_slots): + gradient[probe] += contractions[slot] + return gradient[0] if single_probe else gradient + + +def spin_fock_direct_hf( + gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, + jk_ledger=None, output_slots=None): + """Differentiate one or more spin-resolved HF Fock scalar probes.""" + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + p_alpha, p_beta, single_probe = _spin_probe_stacks( + p_alpha, p_beta, + ) + if output_slots is None: + output_slots = tuple(range(len(p_alpha))) + p_total = p_alpha + p_beta + dm_alpha, dm_beta = _reference_spin_densities(tdobj) + gradient = np.zeros((len(p_alpha), len(atmlst), 3)) + + hcore_derivative = tdobj._scf.nuc_grad_method().hcore_generator(mol) + for k, atom in enumerate(atmlst): + gradient[:, k] += lib.einsum( + "npq,xpq->nx", p_total, hcore_derivative(atom), + ) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + for probe in range(len(p_alpha)): + ledger.add( + "j", output_slots[probe], ( + (p_total[probe], dm_alpha, 1.0, None), + (p_total[probe], dm_beta, 1.0, None), + ), + ) + ledger.add( + "k", output_slots[probe], ( + (p_alpha[probe], dm_alpha, -1.0, None), + (p_beta[probe], dm_beta, -1.0, None), + ), + ) + if jk_ledger is None: + contractions = local_ledger.contract( + gradient_driver, mol, atmlst, slots=output_slots, + ) + for probe, slot in enumerate(output_slots): + gradient[probe] += contractions[slot] + return gradient[0] if single_probe else gradient + + +def response_direct_hfx( + gradient_driver, tdobj, densities, response_terms, atmlst=None, + jk_ledger=None, output_slot=0): + """J/K skeleton derivative for a channel response-term ledger.""" + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + gradient = np.zeros((len(atmlst), 3)) + ni = tdobj._scf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff( + tdobj._scf.xc, mol.spin, + ) + if not ni.libxc.is_hybrid_xc(tdobj._scf.xc): + return gradient + + scales = [(hybrid, None)] + if omega != 0: + scales.append((alpha - hybrid, omega)) + j_terms = [] + k_terms = [] + for term in response_terms: + target = densities[term.target] + source = densities[term.source] + for coefficient, range_omega in scales: + if term.vref0: + k_terms.append(( + target, + source, + -coefficient * term.vref0, + range_omega, + )) + if term.vref1: + j_terms.append(( + target, + source, + -coefficient * term.vref1, + range_omega, + )) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + ledger.add("j", output_slot, j_terms) + ledger.add("k", output_slot, k_terms) + if jk_ledger is None: + gradient += local_ledger.contract( + gradient_driver, mol, atmlst, slots=(output_slot,), + )[output_slot] + return gradient + + +# Hybrid/RSH Fz correction + +def spin_lowering_fockz_hfx_terms( + gradient_driver, tdobj, pz, atmlst=None, with_direct=True, + jk_ledger=None, output_slot=0): + """Differentiate ``-1/2 Pz:K(D_OO)`` excluding the Pz projection.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + mo = np.asarray(mf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + direct = np.zeros((len(atmlst), 3)) + if not ni.libxc.is_hybrid_xc(mf.xc): + return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) + + spaces = orbital_spaces(tdobj) + density_open = spaces.c_open @ spaces.c_open.T + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + scales = [(hybrid, None)] + if omega != 0: + scales.append((alpha - hybrid, omega)) + k_terms = [] + for coefficient, range_omega in scales: + if coefficient == 0.0: + continue + if range_omega is None: + potential = mf.get_k(mol, pz, hermi=0) + else: + potential = mf.get_k( + mol, pz, hermi=0, omega=range_omega, + ) + q_alpha[:, spaces.open] -= 0.5 * coefficient * ( + mo.conj().T @ (potential + potential.T) @ spaces.c_open + ) + if with_direct: + k_terms.append(( + pz, + density_open, + -0.5 * coefficient, + range_omega, + )) + if with_direct: + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + ledger.add("k", output_slot, k_terms) + if jk_ledger is None: + direct += local_ledger.contract( + gradient_driver, mol, atmlst, slots=(output_slot,), + )[output_slot] + return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) + +# Channel assembly + +def grad_elec( + gradient_driver, tdobj, xy, atmlst=None, tolerance=1e-12, + max_cycle=None): + """Build the complete analytic excitation gradient for deltaS=-1.""" + if tdobj.deltaS != -1: + raise ValueError("deltaS=-1 gradient received a different spin channel") + if atmlst is None: + atmlst = range(tdobj.mol.natm) + atmlst = tuple(atmlst) + mf = tdobj._scf + xctype = mf._numint._xc_type(mf.xc) + + # 1. Native amplitudes, transition densities, and explicit Fock probes. + spaces, amplitudes, densities = spin_lowering_transition_densities( + tdobj, xy, + ) + blocks = spin_lowering_block_data(spaces, amplitudes) + response_terms = spin_lowering_response_terms(spaces.spin) + channel_data = (spaces, amplitudes, densities, blocks, response_terms) + p0, pz = spin_lowering_fock_probes(tdobj, xy) + jk_ledger = _JKDerivativeLedger() + direct_slot = "direct" + zvector_slot = "zvector" + + # 2. Explicit Fock contribution to the orbital-rotation M matrix. + fock_alpha, fock_beta = spin_lowering_fock_q(tdobj, xy) + hfx_alpha, hfx_beta = spin_lowering_response_projection_q( + tdobj, xy, hfx_only=True, + ) + + if xctype == "HF": + # 3a. HF response and fixed-orbital AO derivative. + m_matrix = ( + fock_alpha + fock_beta + hfx_alpha + hfx_beta + ) + direct = response_direct_hfx( + gradient_driver, + tdobj, + densities, + response_terms, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slot=direct_slot, + ) + if getattr(mf, "is_ensemble_rks", False): + fockz_hfx = spin_lowering_fockz_hfx_terms( + gradient_driver, + tdobj, + pz, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slot=direct_slot, + ) + m_matrix += fockz_hfx.q_alpha + fockz_hfx.q_beta + direct += fockz_hfx.direct + direct_fock_probes = (0.5 * p0, 0.5 * p0) + else: + direct_fock_probes = ( + 0.5 * (p0 + pz), + 0.5 * (p0 - pz), + ) + + def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): + local = spin_fock_direct_hf( + driver, + obj, + p_alpha, + p_beta, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slots=(direct_slot, zvector_slot), + ) + contractions = jk_ledger.contract( + driver, + obj.mol, + atmlst, + slots=(direct_slot, zvector_slot), + ) + local[0] += contractions[direct_slot] + local[1] += contractions[zvector_slot] + return local + else: + # 3b. Semilocal XC, hybrid/RSH, Fz, and nobeta contributions. + try: + response_builder, fockz_builder, nobeta_q_builder = { + "LDA": ( + xc_backend.lda_response_terms, + xc_backend.lda_fockz_terms, + xc_backend.lda_nobeta_reference_q, + ), + "GGA": ( + xc_backend.gga_response_terms, + xc_backend.gga_fockz_terms, + xc_backend.gga_nobeta_reference_q, + ), + "MGGA": ( + xc_backend.mgga_response_terms, + xc_backend.mgga_fockz_terms, + xc_backend.mgga_nobeta_reference_q, + ), + }[xctype] + except KeyError as error: + raise NotImplementedError( + "NTTDA deltaS=-1 gradient does not support XC type %s" % xctype + ) from error + + response_xc = response_builder( + gradient_driver, + tdobj, + channel_data, + atmlst=atmlst, + ) + fockz_xc = fockz_builder( + gradient_driver, + tdobj, + spaces, + pz, + atmlst=atmlst, + ) + fockz_hfx = spin_lowering_fockz_hfx_terms( + gradient_driver, + tdobj, + pz, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slot=direct_slot, + ) + common_alpha, common_beta = nobeta_q_builder(tdobj, p0) + m_matrix = ( + fock_alpha + fock_beta + + hfx_alpha + hfx_beta + + response_xc.q_alpha + response_xc.q_beta + + fockz_xc.q_alpha + fockz_xc.q_beta + + fockz_hfx.q_alpha + fockz_hfx.q_beta + + common_alpha + common_beta + ) + + direct_fock_probes = (0.5 * p0, 0.5 * p0) + direct = response_direct_hfx( + gradient_driver, + tdobj, + densities, + response_terms, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slot=direct_slot, + ) + direct += response_xc.direct + direct += fockz_xc.direct + direct += fockz_hfx.direct + def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): + local = spin_fock_direct_dft( + driver, + obj, + p_alpha, + p_beta, + atmlst=atmlst, + nobeta_p0=p0, + jk_ledger=jk_ledger, + output_slots=(direct_slot, zvector_slot), + ) + contractions = jk_ledger.contract( + driver, + obj.mol, + atmlst, + slots=(direct_slot, zvector_slot), + ) + local[0] += contractions[direct_slot] + local[1] += contractions[zvector_slot] + return local + + # 4-5. ROKS transpose-Hessian adjoint, Dz Fock derivative, and Pulay term. + return finish_gradient( + gradient_driver, + tdobj, + m_matrix, + direct, + atmlst, + tolerance, + max_cycle, + fock_direct, + direct_fock_probes=direct_fock_probes, + ) diff --git a/src/nest/grad/nttda/delta_s_zero.py b/src/nest/grad/nttda/delta_s_zero.py new file mode 100644 index 0000000..c807337 --- /dev/null +++ b/src/nest/grad/nttda/delta_s_zero.py @@ -0,0 +1,1236 @@ +"""Analytic gradient for current NTTDA ``deltaS=0``. + +The public ``grad_elec`` function exposes the complete scientific data flow. +All same-spin amplitudes, Fock projections, response coefficients, and direct +J/K contractions live in this module. Only XC quadrature and the ROKS +adjoint are delegated to sibling modules. +""" + +from dataclasses import dataclass + +import numpy as np + +from pyscf import dft, lib +from nest.nttda import nttda as nttda_mod +from nest.nttda.nttda import gen_rohf_response_sc + +from . import xc as xc_backend +from .roks import finish_gradient + + + +# Orbital spaces and native amplitudes + +@dataclass(frozen=True) +class OrbitalSpaces: + """Closed, open, and virtual spatial-orbital partitions.""" + + closed: np.ndarray + open: np.ndarray + virtual: np.ndarray + c_closed: np.ndarray + c_open: np.ndarray + c_virtual: np.ndarray + + @property + def spin(self): + return 0.5 * len(self.open) + + +@dataclass(frozen=True) +class SameSpinAmplitudes: + """Five amplitude blocks used by ``NTTDA(deltaS=0)``.""" + + co: np.ndarray + cv: np.ndarray + oo: float + ov: np.ndarray + cv0: np.ndarray + + +def orbital_spaces(tdobj): + """Return the ROKS ``C/O/V`` orbital partition used by NTTDA.""" + mf = tdobj._scf + occ = np.asarray(mf.mo_occ) + if occ.ndim != 1: + raise ValueError("NTTDA gradients require spatial ROKS orbitals") + closed = np.flatnonzero(occ == 2) + open_ = np.flatnonzero(occ == 1) + virtual = np.flatnonzero(occ == 0) + coeff = np.asarray(mf.mo_coeff) + return OrbitalSpaces( + closed=closed, + open=open_, + virtual=virtual, + c_closed=coeff[:, closed], + c_open=coeff[:, open_], + c_virtual=coeff[:, virtual], + ) + + +def same_spin_slices(spaces): + """Return canonical slices for ``CO/CV/OO/OV/CV0`` amplitudes.""" + nc = len(spaces.closed) + no = len(spaces.open) + nv = len(spaces.virtual) + nco = nc * no + ncv = nc * nv + nov = no * nv + i1 = nco + i2 = i1 + ncv + i3 = i2 + 1 + i4 = i3 + nov + return { + "CO": slice(0, i1), + "CV": slice(i1, i2), + "OO": slice(i2, i3), + "OV": slice(i3, i4), + "CV0": slice(i4, i4 + ncv), + } + + +def split_same_spin(tdobj, xy): + """Split one packed ``deltaS=0`` vector into its five native blocks.""" + spaces = orbital_spaces(tdobj) + if spaces.spin < 0.5: + raise ValueError("NTTDA deltaS=0 requires at least one open orbital") + vector = xy[0] if isinstance(xy, (tuple, list)) else xy + vector = np.asarray(vector).reshape(-1) + slices = same_spin_slices(spaces) + expected = slices["CV0"].stop + if vector.size != expected: + raise ValueError( + "deltaS=0 amplitude has size %d; expected %d" % + (vector.size, expected) + ) + nc = len(spaces.closed) + no = len(spaces.open) + nv = len(spaces.virtual) + return spaces, SameSpinAmplitudes( + co=vector[slices["CO"]].reshape(nc, no), + cv=vector[slices["CV"]].reshape(nc, nv), + oo=float(vector[slices["OO"]][0]), + ov=vector[slices["OV"]].reshape(no, nv), + cv0=vector[slices["CV0"]].reshape(nc, nv), + ) + + +def pair_density(c_left, coefficient, c_right): + """Build ``C_left coefficient C_right^T`` without symmetrizing it.""" + return c_left @ np.asarray(coefficient) @ c_right.conj().T + + +def same_spin_transition_densities(tdobj, xy): + """Return directed AO transition densities for the four response blocks.""" + spaces, amp = split_same_spin(tdobj, xy) + return spaces, amp, { + "CO": pair_density(spaces.c_open, amp.co.T, spaces.c_closed), + "CV": pair_density(spaces.c_virtual, amp.cv.T, spaces.c_closed), + "OV": pair_density(spaces.c_virtual, amp.ov.T, spaces.c_open), + "CV0": pair_density(spaces.c_virtual, amp.cv0.T, spaces.c_closed), + } + + +# Explicit F0/Fz ledger + +@dataclass(frozen=True) +class FockProjection: + """One scalar term ``Tr[P (weight_f0 F0 + weight_fz Fz)]``.""" + + name: str + left_indices: np.ndarray + left_orbitals: np.ndarray + coefficient: np.ndarray + right_indices: np.ndarray + right_orbitals: np.ndarray + weight_f0: float + weight_fz: float + + def density(self): + return pair_density( + self.left_orbitals, self.coefficient, self.right_orbitals, + ) + + +def fock0_fockz(tdobj, max_memory=None): + """Build exactly the ``F0`` and ``Fz`` matrices used by ``gen_vind_sc``.""" + mf = tdobj._scf + if max_memory is None: + max_memory = tdobj.max_memory + _response, fockz = gen_rohf_response_sc( + mf, + mo_coeff=mf.mo_coeff, + mo_occ=mf.mo_occ, + hermi=0, + max_memory=max_memory, + ) + fock0 = nttda_mod._reference_fock0(mf, tdobj.nobeta) + return fock0, fockz + + +def same_spin_fock_projections(tdobj, xy): + """Return the complete five-block explicit-Fock ledger.""" + spaces, x = split_same_spin(tdobj, xy) + spin = spaces.spin + c = spaces.c_closed + o = spaces.c_open + v = spaces.c_virtual + a = np.sqrt((spin + 1.0) / (2.0 * spin)) + b = np.sqrt(2.0 * (spin + 1.0) / spin) + d = np.sqrt((spin + 1.0) / spin) + h = np.sqrt(0.5) + + terms = [] + + def indices(orbitals): + if orbitals is c: + return spaces.closed + if orbitals is o: + return spaces.open + if orbitals is v: + return spaces.virtual + raise ValueError("Fock projection uses an unknown orbital space") + + def add(name, left, coefficient, right, f0, fz): + coefficient = np.asarray(coefficient) + if coefficient.size: + terms.append(FockProjection( + name, + indices(left), left, coefficient, + indices(right), right, + float(f0), float(fz), + )) + + # CO row/column and its couplings. + add("co-oo", o, x.co.T @ x.co, o, 1.0, -1.0) + add("co-cc", c, -x.co @ x.co.T, c, 1.0, -1.0) + add("co-cv", o, 2.0 * a * (x.co.T @ x.cv), v, 1.0, -1.0) + add("co-oo1", o, -2.0 * x.oo * x.co.T, c, 1.0, -1.0) + add("co-cv0", o, 2.0 * h * (x.co.T @ x.cv0), v, 1.0, -1.0) + + # CV block and its OO/OV/CV0 couplings. + add("cv-vv", v, x.cv.T @ x.cv, v, 1.0, -1.0 / spin) + add("cv-cc", c, -x.cv @ x.cv.T, c, 1.0, 1.0 / spin) + add("cv-oo1", v, 2.0 * b * x.oo * x.cv.T, c, 0.0, 1.0) + add("cv-ov", o, -2.0 * a * (x.ov @ x.cv.T), c, 1.0, 1.0) + add( + "cv-cv0-vv", v, + -d * (x.cv.T @ x.cv0 + x.cv0.T @ x.cv), v, 0.0, 1.0, + ) + add( + "cv-cv0-cc", c, + d * (x.cv0 @ x.cv.T + x.cv @ x.cv0.T), c, 0.0, 1.0, + ) + + # OV and CV0 diagonal/coupling terms. + add("ov-vv", v, x.ov.T @ x.ov, v, 1.0, 1.0) + add("ov-oo", o, -x.ov @ x.ov.T, o, 1.0, 1.0) + add("ov-oo1", v, 2.0 * x.oo * x.ov.T, o, 1.0, 1.0) + add("ov-cv0", c, 2.0 * h * (x.cv0 @ x.ov.T), o, 1.0, 1.0) + add("cv0-vv", v, x.cv0.T @ x.cv0, v, 1.0, 0.0) + add("cv0-cc", c, -x.cv0 @ x.cv0.T, c, 1.0, 0.0) + add("cv0-oo1", v, -2.0 * np.sqrt(2.0) * x.oo * x.cv0.T, c, 1.0, 0.0) + return tuple(terms) + + +def same_spin_fock_probes(tdobj, xy): + """Return AO probes ``(P0, Pz)`` generated from the Fock ledger.""" + nao = tdobj.mol.nao_nr() + p0 = np.zeros((nao, nao)) + pz = np.zeros_like(p0) + for term in same_spin_fock_projections(tdobj, xy): + density = term.density() + p0 += term.weight_f0 * density + pz += term.weight_fz * density + return p0, pz + + +def _fock_response_q(tdobj, p_alpha, p_beta): + """Reference-density derivative of a spin-resolved Fock scalar.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + if getattr(mf, "is_ensemble_rks", False): + occupation = np.asarray(mf.mo_occ) + probe = np.asarray(p_alpha) + np.asarray(p_beta) + potential = mf.gen_response(hermi=0)(probe.T) + q_total = ( + mo.conj().T @ (potential + potential.T) @ mo + ) * occupation[None, :] + return 0.5 * q_total, 0.5 * q_total + occ_alpha = (np.asarray(mf.mo_occ) > 0).astype(float) + occ_beta = (np.asarray(mf.mo_occ) == 2).astype(float) + if (isinstance(mf, dft.KohnShamDFT) + and mf._numint._xc_type(mf.xc) != "HF"): + unrestricted = mf.to_uks() + unrestricted.verbose = 0 + v_alpha, v_beta = unrestricted.gen_response(hermi=0)( + np.asarray((p_alpha.T, p_beta.T)) + ) + else: + p_total = p_alpha + p_beta + coulomb = mf.get_j(mf.mol, p_total.T, hermi=0) + v_alpha = coulomb - mf.get_k(mf.mol, p_alpha.T, hermi=0) + v_beta = coulomb - mf.get_k(mf.mol, p_beta.T, hermi=0) + q_alpha = ( + mo.conj().T @ (v_alpha + v_alpha.T) @ mo + ) * occ_alpha[None, :] + q_beta = ( + mo.conj().T @ (v_beta + v_beta.T) @ mo + ) * occ_beta[None, :] + return q_alpha, q_beta + + +def same_spin_fock_q(tdobj, xy, max_memory=None): + """Return the explicit-Fock contribution to ``(Q_alpha,Q_beta)``. + + For HF the complete ``Fz`` response is exactly represented by the + spin-resolved probes. DFT callers add the independent ``Fz`` and + ``nobeta`` response ledgers after this common ``F0`` contribution. + """ + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + nmo = mo.shape[1] + fock0, fockz = fock0_fockz(tdobj, max_memory=max_memory) + fock0_mo = mo.conj().T @ fock0 @ mo + fockz_mo = mo.conj().T @ fockz @ mo + q_alpha = np.zeros((nmo, nmo)) + q_beta = np.zeros_like(q_alpha) + + for term in same_spin_fock_projections(tdobj, xy): + left = term.left_indices + right = term.right_indices + coeff = term.coefficient + + def project(target, operator, scale): + if scale == 0.0: + return + target[:, left] += scale * operator[:, right] @ coeff.T + target[:, right] += scale * operator[:, left] @ coeff + + project(q_alpha, fock0_mo, 0.5 * term.weight_f0) + project(q_beta, fock0_mo, 0.5 * term.weight_f0) + if mf._numint._xc_type(mf.xc) == "HF": + project(q_alpha, fockz_mo, 0.5 * term.weight_fz) + project(q_beta, fockz_mo, 0.5 * term.weight_fz) + else: + project(q_alpha, fockz_mo, term.weight_fz) + + p0, pz = same_spin_fock_probes(tdobj, xy) + p_alpha = 0.5 * p0 + p_beta = 0.5 * p0 + if mf._numint._xc_type(mf.xc) == "HF": + p_alpha = p_alpha + 0.5 * pz + p_beta = p_beta - 0.5 * pz + response_alpha, response_beta = _fock_response_q( + tdobj, p_alpha, p_beta, + ) + q_alpha += response_alpha + q_beta += response_beta + return q_alpha, q_beta + + +def same_spin_fock_scalar(tdobj, xy, max_memory=None): + """Evaluate the complete explicit-Fock part of ``X.T A_sc X``.""" + fock0, fockz = fock0_fockz(tdobj, max_memory=max_memory) + return same_spin_fock_projection_scalar(tdobj, xy, fock0, fockz) + + +def same_spin_fock_projection_scalar(tdobj, xy, fock0, fockz): + """Evaluate the Fock ledger for caller-supplied frozen operators.""" + p0, pz = same_spin_fock_probes(tdobj, xy) + return float( + lib.einsum("pq,pq->", p0, fock0) + + lib.einsum("pq,pq->", pz, fockz) + ) + + +# vref0/vref1 response ledger + +@dataclass(frozen=True) +class ResponseTerm: + """Directed response term from one source density to one target block.""" + + target: str + source: str + vref0: float + vref1: float + + +def same_spin_response_terms(spin): + """Directed coefficients transcribed from ``gen_rohf_response_sc``.""" + a = np.sqrt((spin + 1.0) / (2.0 * spin)) + h = np.sqrt(0.5) + r2 = np.sqrt(2.0) + return ( + ResponseTerm("CO", "CO", 1.0, -1.0), + ResponseTerm("CO", "CV", a, 0.0), + ResponseTerm("CO", "OV", 0.0, 1.0), + ResponseTerm("CO", "CV0", h, -r2), + ResponseTerm("CV", "CO", a, 0.0), + ResponseTerm("CV", "CV", 1.0, 0.0), + ResponseTerm("CV", "OV", a, 0.0), + ResponseTerm("OV", "CO", 0.0, 1.0), + ResponseTerm("OV", "CV", a, 0.0), + ResponseTerm("OV", "OV", 1.0, -1.0), + ResponseTerm("OV", "CV0", -h, r2), + ResponseTerm("CV0", "CO", h, -r2), + ResponseTerm("CV0", "OV", -h, r2), + ResponseTerm("CV0", "CV0", 1.0, -2.0), + ) + + +def _fxc_reference(tdobj): + mf = tdobj._scf + ni = mf._numint + fxc = ni.cache_xc_kernel( + mf.mol, mf.grids, mf.xc, mf.mo_coeff, mf.mo_occ, 1, + )[2] + return 0.5 * ( + fxc[0, :, 0] - fxc[0, :, 1] + - fxc[1, :, 0] + fxc[1, :, 1] + ) + + +def _apply_reference_responses(tdobj, densities, max_memory=None): + """Return separate ``vref0`` and ``vref1`` actions for each density.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if max_memory is None: + max_memory = tdobj.max_memory + labels = tuple(densities) + dms = np.asarray([densities[label] for label in labels]) + xctype = ni._xc_type(mf.xc) + if xctype == "HF": + vref0 = np.zeros_like(dms) + vref1 = np.zeros_like(dms) + else: + fxc_ref = _fxc_reference(tdobj) + vref0 = ni.nr_rks_fxc( + mol, mf.grids, mf.xc, None, dms, 0, 0, + None, None, fxc_ref, max_memory=max_memory, + ) + if xctype == "LDA": + vref1 = ni.nr_rks_fxc( + mol, mf.grids, mf.xc, None, dms, 0, 0, + None, None, fxc_ref, max_memory=max_memory, + ) + elif xctype == "GGA": + vref1 = nttda_mod.nr_rks_fxc1_gga( + ni, mol, mf.grids, mf.xc, dms, fxc_ref, + max_memory=max_memory, + ) + elif xctype == "MGGA": + vref1 = nttda_mod.nr_rks_fxc1_mgga( + ni, mol, mf.grids, mf.xc, dms, fxc_ref, + max_memory=max_memory, + ) + else: + raise NotImplementedError( + "NTTDA same-spin response does not support XC type %s" % xctype + ) + + omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + if ni.libxc.is_hybrid_xc(mf.xc): + vref0 -= hyb * mf.get_k(mol, dms, hermi=0) + vref1 -= hyb * mf.get_j(mol, dms, hermi=0) + if omega != 0: + scale = alpha - hyb + vref0 -= scale * mf.get_k(mol, dms, hermi=0, omega=omega) + vref1 -= scale * mf.get_j(mol, dms, hermi=0, omega=omega) + return ( + {label: value for label, value in zip(labels, vref0)}, + {label: value for label, value in zip(labels, vref1)}, + ) + + +def _apply_hfx_responses(tdobj, densities): + """Return only the hybrid/RSH J/K portions of ``vref0/vref1``.""" + mf = tdobj._scf + labels = tuple(densities) + dms = np.asarray([densities[label] for label in labels]) + vref0 = np.zeros_like(dms) + vref1 = np.zeros_like(dms) + ni = mf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mf.mol.spin) + if ni.libxc.is_hybrid_xc(mf.xc): + vref0 -= hybrid * mf.get_k(mf.mol, dms, hermi=0) + vref1 -= hybrid * mf.get_j(mf.mol, dms, hermi=0) + if omega != 0: + scale = alpha - hybrid + vref0 -= scale * mf.get_k( + mf.mol, dms, hermi=0, omega=omega, + ) + vref1 -= scale * mf.get_j( + mf.mol, dms, hermi=0, omega=omega, + ) + return ( + {label: value for label, value in zip(labels, vref0)}, + {label: value for label, value in zip(labels, vref1)}, + ) + + +def _derivative_potentials(spaces, densities, vref0, vref1): + potentials = {label: np.zeros_like(dm) for label, dm in densities.items()} + for term in same_spin_response_terms(spaces.spin): + if term.vref0: + potentials[term.target] += term.vref0 * vref0[term.source] + potentials[term.source] += term.vref0 * vref0[term.target] + if term.vref1: + potentials[term.target] += term.vref1 * vref1[term.source] + potentials[term.source] += term.vref1 * vref1[term.target] + return potentials + + +def same_spin_response_scalar(tdobj, xy, max_memory=None): + """Evaluate all current-NTTDA response terms in ``X.T A_sc X``.""" + spaces, _amp, densities = same_spin_transition_densities(tdobj, xy) + vref0, vref1 = _apply_reference_responses( + tdobj, densities, max_memory=max_memory, + ) + value = 0.0 + for term in same_spin_response_terms(spaces.spin): + target = densities[term.target] + if term.vref0: + value += term.vref0 * lib.einsum( + "pq,pq->", target, vref0[term.source], + ) + if term.vref1: + value += term.vref1 * lib.einsum( + "pq,pq->", target, vref1[term.source], + ) + return float(value) + + +def same_spin_response_derivative_potentials(tdobj, xy, max_memory=None): + """AO potentials obtained by varying both sides of the response scalar.""" + spaces, _amp, densities = same_spin_transition_densities(tdobj, xy) + vref0, vref1 = _apply_reference_responses( + tdobj, densities, max_memory=max_memory, + ) + return _derivative_potentials(spaces, densities, vref0, vref1) + + +def same_spin_response_projection_q(tdobj, xy, max_memory=None): + """MO derivative from transition-density factors at frozen kernels.""" + spaces, amp, _densities = same_spin_transition_densities(tdobj, xy) + potentials = same_spin_response_derivative_potentials( + tdobj, xy, max_memory=max_memory, + ) + mo = np.asarray(tdobj._scf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + block_data = { + "CO": (spaces.open, spaces.closed, amp.co.T), + "CV": (spaces.virtual, spaces.closed, amp.cv.T), + "OV": (spaces.virtual, spaces.open, amp.ov.T), + "CV0": (spaces.virtual, spaces.closed, amp.cv0.T), + } + for label, (target, source, coefficient) in block_data.items(): + potential_mo = mo.conj().T @ potentials[label] @ mo + q_beta[:, target] += potential_mo[:, source] @ coefficient.T + q_alpha[:, source] += potential_mo[target, :].T @ coefficient + return q_alpha, q_beta + + +def same_spin_hfx_projection_q(tdobj, xy): + """Transition-factor derivative of only the hybrid/RSH response scalar.""" + spaces, amp, densities = same_spin_transition_densities(tdobj, xy) + vref0, vref1 = _apply_hfx_responses(tdobj, densities) + potentials = _derivative_potentials( + spaces, densities, vref0, vref1, + ) + mo = np.asarray(tdobj._scf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + block_data = { + "CO": (spaces.open, spaces.closed, amp.co.T), + "CV": (spaces.virtual, spaces.closed, amp.cv.T), + "OV": (spaces.virtual, spaces.open, amp.ov.T), + "CV0": (spaces.virtual, spaces.closed, amp.cv0.T), + } + for label, (target, source, coefficient) in block_data.items(): + potential_mo = mo.conj().T @ potentials[label] @ mo + q_beta[:, target] += potential_mo[:, source] @ coefficient.T + q_alpha[:, source] += potential_mo[target, :].T @ coefficient + return q_alpha, q_beta + + +# AO J/K nuclear derivatives + +def _as_derivative_stack(array): + array = np.asarray(array) + if array.ndim == 3: + array = array[None] + return array + + +def _density_key(density): + density = np.asarray(density) + data = density.__array_interface__["data"][0] + return data, density.shape, density.strides, density.dtype.str + + +@dataclass(frozen=True) +class _JKDerivativeTerm: + """One fixed-AO bilinear derivative with a named output slot.""" + + left: np.ndarray + right: np.ndarray + scale: float + omega: float + slot: object + + +class _JKDerivativeLedger: + """Channel-local scheduler for fixed-AO J/K derivative contractions.""" + + def __init__(self): + self._terms = {"j": [], "k": []} + + def add(self, operator, slot, terms): + self._terms[operator].extend( + _JKDerivativeTerm(left, right, scale, omega, slot) + for left, right, scale, omega in terms + if scale != 0.0 + ) + + def contract(self, gradient_driver, mol, atoms, slots=()): + atoms = tuple(atoms) + shape = (len(atoms), 3) + gradients = {slot: np.zeros(shape) for slot in slots} + for operator in ("j", "k"): + for term in self._terms[operator]: + gradients.setdefault(term.slot, np.zeros(shape)) + _contract_derivative_terms( + gradients, + gradient_driver, + mol, + atoms, + mol.offset_nr_by_atom(), + self._terms[operator], + operator, + ) + return gradients + + +def _term_densities(term, exchange): + left, right = term.left, term.right + if exchange: + return left, right, left.T, right.T + return left, right + + +def _density_batches(terms, exchange, max_memory, nao): + """Group bilinear terms while bounding derivative-potential storage.""" + minimum = 4 if exchange else 2 + bytes_per_density = 4 * nao * nao * np.dtype(float).itemsize + batch_limit = max( + minimum, + int(0.2 * max_memory * 1e6 / bytes_per_density), + ) + batch = [] + keys = set() + for term in terms: + term_keys = { + _density_key(density) + for density in _term_densities(term, exchange) + } + if batch and len(keys | term_keys) > batch_limit: + yield batch + batch = [] + keys = set() + batch.append(term) + keys.update(term_keys) + if batch: + yield batch + + +def _jk_derivative_potentials( + gradient_driver, mol, terms, operator, omega): + exchange = operator == "k" + densities = {} + for term in terms: + for density in _term_densities(term, exchange): + density = np.asarray(density) + densities.setdefault(_density_key(density), density) + keys = tuple(densities) + stack = np.asarray([densities[key] for key in keys]) + if operator == "j": + if omega is None: + values = gradient_driver.get_j(mol, stack, hermi=0) + else: + values = gradient_driver.get_j( + mol, stack, hermi=0, omega=omega, + ) + else: + if omega is None: + values = gradient_driver.get_k(mol, stack, hermi=0) + else: + values = gradient_driver.get_k( + mol, stack, hermi=0, omega=omega, + ) + values = _as_derivative_stack(values) + return dict(zip(keys, values)) + + +def _contract_derivative_terms( + gradients, gradient_driver, mol, atoms, offsets, terms, + operator): + if not atoms: + return + terms_by_omega = {} + for term in terms: + terms_by_omega.setdefault(term.omega, []).append(term) + exchange = operator == "k" + for omega, omega_terms in terms_by_omega.items(): + for batch in _density_batches( + omega_terms, exchange, gradient_driver.max_memory, + mol.nao_nr()): + potentials = _jk_derivative_potentials( + gradient_driver, mol, batch, operator, omega, + ) + for term in batch: + left = np.asarray(term.left) + right = np.asarray(term.right) + right_derivative = potentials[_density_key(right)] + left_derivative = potentials[_density_key(left)] + if exchange: + right_t_derivative = potentials[ + _density_key(right.T) + ] + left_t_derivative = potentials[_density_key(left.T)] + for k, atom in enumerate(atoms): + p0, p1 = offsets[atom][2:] + if exchange: + value = lib.einsum( + "xpq,pq->x", + right_derivative[:, p0:p1, :], + left[p0:p1, :], + ) + value += lib.einsum( + "xqp,pq->x", + right_t_derivative[:, p0:p1, :], + left[:, p0:p1], + ) + value += lib.einsum( + "xpq,pq->x", + left_derivative[:, p0:p1, :], + right[p0:p1, :], + ) + value += lib.einsum( + "xqp,pq->x", + left_t_derivative[:, p0:p1, :], + right[:, p0:p1], + ) + else: + value = lib.einsum( + "xpq,pq->x", + right_derivative[:, p0:p1], + left[p0:p1], + ) + value += lib.einsum( + "xpq,qp->x", + right_derivative[:, p0:p1], + left[:, p0:p1], + ) + value += lib.einsum( + "xpq,pq->x", + left_derivative[:, p0:p1], + right[p0:p1], + ) + value += lib.einsum( + "xpq,qp->x", + left_derivative[:, p0:p1], + right[:, p0:p1], + ) + gradients[term.slot][k] += term.scale * value + + +def _reference_spin_densities(tdobj): + mf = tdobj._scf + if getattr(mf, "is_ensemble_rks", False): + return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) + mo = np.asarray(mf.mo_coeff) + return ( + mo[:, mf.mo_occ > 0] @ mo[:, mf.mo_occ > 0].T, + mo[:, mf.mo_occ == 2] @ mo[:, mf.mo_occ == 2].T, + ) + + +def _spin_probe_stacks(p_alpha, p_beta): + p_alpha = np.asarray(p_alpha) + p_beta = np.asarray(p_beta) + single_probe = p_alpha.ndim == 2 + if single_probe: + p_alpha = p_alpha[None] + p_beta = p_beta[None] + return p_alpha, p_beta, single_probe + + +def spin_fock_direct_dft( + gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, + nobeta_p0=None, jk_ledger=None, output_slots=None): + """Differentiate one or more ordinary UKS Fock scalar probes. + + The optional ``nobeta_p0`` correction belongs to the first, explicit-direct + probe in the batch. + """ + mf = tdobj._scf + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + p_alpha, p_beta, single_probe = _spin_probe_stacks( + p_alpha, p_beta, + ) + if output_slots is None: + output_slots = tuple(range(len(p_alpha))) + p_total = p_alpha + p_beta + density_alpha, density_beta = _reference_spin_densities(tdobj) + gradient = np.zeros((len(p_alpha), len(atmlst), 3)) + hcore_derivative = mf.nuc_grad_method().hcore_generator(mol) + for k, atom in enumerate(atmlst): + gradient[:, k] += lib.einsum( + "npq,xpq->nx", p_total, hcore_derivative(atom), + ) + ni = mf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + for probe in range(len(p_alpha)): + j_terms = [ + (p_total[probe], density_alpha, 1.0, None), + (p_total[probe], density_beta, 1.0, None), + ] + k_terms = [] + if ni.libxc.is_hybrid_xc(mf.xc): + k_terms.extend(( + (p_alpha[probe], density_alpha, -hybrid, None), + (p_beta[probe], density_beta, -hybrid, None), + )) + if omega != 0: + long_range = -(alpha - hybrid) + k_terms.extend(( + (p_alpha[probe], density_alpha, long_range, omega), + (p_beta[probe], density_beta, long_range, omega), + )) + ledger.add("j", output_slots[probe], j_terms) + ledger.add("k", output_slots[probe], k_terms) + xctype = ni._xc_type(mf.xc) + if xctype == "LDA": + derivative_contractor = xc_backend.contract_lda_vxc_derivative + elif xctype == "GGA": + derivative_contractor = xc_backend.contract_gga_vxc_derivative + elif xctype == "MGGA": + derivative_contractor = xc_backend.contract_mgga_vxc_derivative + else: + raise NotImplementedError( + "ordinary Fock direct derivative is not implemented for %s" % + xctype + ) + if (nobeta_p0 is not None and tdobj.nobeta + and not getattr(mf, "is_ensemble_rks", False)): + density0 = 0.5 * (density_alpha + density_beta) + actual_probe_alpha = np.array(p_alpha, copy=True) + actual_probe_beta = np.array(p_beta, copy=True) + actual_probe_alpha[0] -= 0.5 * nobeta_p0 + actual_probe_beta[0] -= 0.5 * nobeta_p0 + else: + density0 = None + actual_probe_alpha = p_alpha + actual_probe_beta = p_beta + gradient += derivative_contractor( + mf, + density_alpha, + density_beta, + actual_probe_alpha, + actual_probe_beta, + atmlst=atmlst, + max_memory=gradient_driver.max_memory, + ) + if density0 is not None: + gradient[0] += derivative_contractor( + mf, + density0, + density0, + 0.5 * nobeta_p0, + 0.5 * nobeta_p0, + atmlst=atmlst, + max_memory=gradient_driver.max_memory, + ) + if jk_ledger is None: + contractions = local_ledger.contract( + gradient_driver, mol, atmlst, slots=output_slots, + ) + for probe, slot in enumerate(output_slots): + gradient[probe] += contractions[slot] + return gradient[0] if single_probe else gradient + + +def spin_fock_direct_hf( + gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, + jk_ledger=None, output_slots=None): + """Differentiate one or more spin-resolved HF Fock scalar probes.""" + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + p_alpha, p_beta, single_probe = _spin_probe_stacks( + p_alpha, p_beta, + ) + if output_slots is None: + output_slots = tuple(range(len(p_alpha))) + p_total = p_alpha + p_beta + dm_alpha, dm_beta = _reference_spin_densities(tdobj) + gradient = np.zeros((len(p_alpha), len(atmlst), 3)) + + hcore_derivative = tdobj._scf.nuc_grad_method().hcore_generator(mol) + for k, atom in enumerate(atmlst): + gradient[:, k] += lib.einsum( + "npq,xpq->nx", p_total, hcore_derivative(atom), + ) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + for probe in range(len(p_alpha)): + ledger.add( + "j", output_slots[probe], ( + (p_total[probe], dm_alpha, 1.0, None), + (p_total[probe], dm_beta, 1.0, None), + ), + ) + ledger.add( + "k", output_slots[probe], ( + (p_alpha[probe], dm_alpha, -1.0, None), + (p_beta[probe], dm_beta, -1.0, None), + ), + ) + if jk_ledger is None: + contractions = local_ledger.contract( + gradient_driver, mol, atmlst, slots=output_slots, + ) + for probe, slot in enumerate(output_slots): + gradient[probe] += contractions[slot] + return gradient[0] if single_probe else gradient + + +def response_direct_hfx( + gradient_driver, tdobj, densities, response_terms, atmlst=None, + jk_ledger=None, output_slot=0): + """J/K skeleton derivative for a channel response-term ledger.""" + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + gradient = np.zeros((len(atmlst), 3)) + ni = tdobj._scf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff( + tdobj._scf.xc, mol.spin, + ) + if not ni.libxc.is_hybrid_xc(tdobj._scf.xc): + return gradient + + scales = [(hybrid, None)] + if omega != 0: + scales.append((alpha - hybrid, omega)) + j_terms = [] + k_terms = [] + for term in response_terms: + target = densities[term.target] + source = densities[term.source] + for coefficient, range_omega in scales: + if term.vref0: + k_terms.append(( + target, + source, + -coefficient * term.vref0, + range_omega, + )) + if term.vref1: + j_terms.append(( + target, + source, + -coefficient * term.vref1, + range_omega, + )) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + ledger.add("j", output_slot, j_terms) + ledger.add("k", output_slot, k_terms) + if jk_ledger is None: + gradient += local_ledger.contract( + gradient_driver, mol, atmlst, slots=(output_slot,), + )[output_slot] + return gradient + + +# Hybrid/RSH Fz correction + +def same_spin_fockz_hfx_terms( + gradient_driver, tdobj, pz, atmlst=None, with_direct=True, + jk_ledger=None, output_slot=0): + """Differentiate ``-1/2 Pz:K(D_OO)`` excluding the Pz projection.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + mo = np.asarray(mf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + direct = np.zeros((len(atmlst), 3)) + if not ni.libxc.is_hybrid_xc(mf.xc): + return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) + + spaces = orbital_spaces(tdobj) + density_open = spaces.c_open @ spaces.c_open.T + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + scales = [(hybrid, None)] + if omega != 0: + scales.append((alpha - hybrid, omega)) + k_terms = [] + for coefficient, range_omega in scales: + if coefficient == 0.0: + continue + if range_omega is None: + potential = mf.get_k(mol, pz, hermi=0) + else: + potential = mf.get_k( + mol, pz, hermi=0, omega=range_omega, + ) + q_alpha[:, spaces.open] -= 0.5 * coefficient * ( + mo.conj().T @ (potential + potential.T) @ spaces.c_open + ) + if with_direct: + k_terms.append(( + pz, + density_open, + -0.5 * coefficient, + range_omega, + )) + if with_direct: + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + ledger.add("k", output_slot, k_terms) + if jk_ledger is None: + direct += local_ledger.contract( + gradient_driver, mol, atmlst, slots=(output_slot,), + )[output_slot] + return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) + + +# Scalar closure diagnostics (private to this channel) + +def same_spin_action_scalar(tdobj, xy): + """Evaluate ``X.T gen_vind_sc(X)`` using the production NTTDA action.""" + vector = xy[0] if isinstance(xy, (tuple, list)) else xy + vector = np.asarray(vector).reshape(-1) + vind, _hdiag = tdobj.gen_vind_sc() + action = vind(vector.reshape(1, -1))[0] + return float(np.vdot(vector, action).real) + + +def same_spin_ledger_scalar(tdobj, xy, max_memory=None, return_parts=False): + """Evaluate the independent ``F0/Fz + vref0/vref1`` scalar ledger.""" + fock = same_spin_fock_scalar(tdobj, xy, max_memory=max_memory) + response = same_spin_response_scalar(tdobj, xy, max_memory=max_memory) + total = fock + response + if return_parts: + return {"fock": fock, "response": response, "total": total} + return total + +# Channel assembly + +def grad_elec( + gradient_driver, tdobj, xy, atmlst=None, tolerance=1e-12, + max_cycle=None): + """Build the complete analytic excitation gradient for deltaS=0. + + The function follows the physical order of the Lagrangian: native + amplitudes and AO probes, XC/J/K contributions to the M matrix and direct + derivative, the ROKS adjoint, and the final overlap contraction. + """ + if tdobj.deltaS != 0: + raise ValueError("deltaS=0 gradient received a different spin channel") + if atmlst is None: + atmlst = range(tdobj.mol.natm) + atmlst = tuple(atmlst) + mf = tdobj._scf + xctype = mf._numint._xc_type(mf.xc) + + # 1. Native amplitudes, transition densities, and explicit Fock probes. + spaces, amplitudes, densities = same_spin_transition_densities(tdobj, xy) + blocks = { + "CO": (spaces.open, spaces.closed, amplitudes.co.T), + "CV": (spaces.virtual, spaces.closed, amplitudes.cv.T), + "OV": (spaces.virtual, spaces.open, amplitudes.ov.T), + "CV0": (spaces.virtual, spaces.closed, amplitudes.cv0.T), + } + response_terms = same_spin_response_terms(spaces.spin) + channel_data = (spaces, amplitudes, densities, blocks, response_terms) + p0, pz = same_spin_fock_probes(tdobj, xy) + jk_ledger = _JKDerivativeLedger() + direct_slot = "direct" + zvector_slot = "zvector" + + # 2. Explicit Fock contribution to the orbital-rotation M matrix. + fock_alpha, fock_beta = same_spin_fock_q(tdobj, xy) + + if xctype == "HF": + # 3a. HF response and fixed-orbital AO derivative. + response_alpha, response_beta = same_spin_response_projection_q( + tdobj, xy, + ) + m_matrix = ( + fock_alpha + fock_beta + response_alpha + response_beta + ) + direct = response_direct_hfx( + gradient_driver, + tdobj, + densities, + response_terms, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slot=direct_slot, + ) + if getattr(mf, "is_ensemble_rks", False): + fockz_hfx = same_spin_fockz_hfx_terms( + gradient_driver, + tdobj, + pz, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slot=direct_slot, + ) + m_matrix += fockz_hfx.q_alpha + fockz_hfx.q_beta + direct += fockz_hfx.direct + direct_fock_probes = (0.5 * p0, 0.5 * p0) + else: + direct_fock_probes = ( + 0.5 * (p0 + pz), + 0.5 * (p0 - pz), + ) + + def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): + local = spin_fock_direct_hf( + driver, + obj, + p_alpha, + p_beta, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slots=(direct_slot, zvector_slot), + ) + contractions = jk_ledger.contract( + driver, + obj.mol, + atmlst, + slots=(direct_slot, zvector_slot), + ) + local[0] += contractions[direct_slot] + local[1] += contractions[zvector_slot] + return local + else: + # 3b. Semilocal XC, hybrid/RSH, Fz, and nobeta contributions. + try: + response_builder, fockz_builder, nobeta_q_builder = { + "LDA": ( + xc_backend.lda_response_terms, + xc_backend.lda_fockz_terms, + xc_backend.lda_nobeta_reference_q, + ), + "GGA": ( + xc_backend.gga_response_terms, + xc_backend.gga_fockz_terms, + xc_backend.gga_nobeta_reference_q, + ), + "MGGA": ( + xc_backend.mgga_response_terms, + xc_backend.mgga_fockz_terms, + xc_backend.mgga_nobeta_reference_q, + ), + }[xctype] + except KeyError as error: + raise NotImplementedError( + "NTTDA deltaS=0 gradient does not support XC type %s" % xctype + ) from error + + hfx_alpha, hfx_beta = same_spin_hfx_projection_q(tdobj, xy) + response_xc = response_builder( + gradient_driver, + tdobj, + channel_data, + atmlst=atmlst, + ) + fockz_xc = fockz_builder( + gradient_driver, + tdobj, + spaces, + pz, + atmlst=atmlst, + ) + fockz_hfx = same_spin_fockz_hfx_terms( + gradient_driver, + tdobj, + pz, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slot=direct_slot, + ) + common_alpha, common_beta = nobeta_q_builder(tdobj, p0) + m_matrix = ( + fock_alpha + fock_beta + + hfx_alpha + hfx_beta + + response_xc.q_alpha + response_xc.q_beta + + fockz_xc.q_alpha + fockz_xc.q_beta + + fockz_hfx.q_alpha + fockz_hfx.q_beta + + common_alpha + common_beta + ) + + direct_fock_probes = (0.5 * p0, 0.5 * p0) + direct = response_direct_hfx( + gradient_driver, + tdobj, + densities, + response_terms, + atmlst=atmlst, + jk_ledger=jk_ledger, + output_slot=direct_slot, + ) + direct += response_xc.direct + direct += fockz_xc.direct + direct += fockz_hfx.direct + def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): + local = spin_fock_direct_dft( + driver, + obj, + p_alpha, + p_beta, + atmlst=atmlst, + nobeta_p0=p0, + jk_ledger=jk_ledger, + output_slots=(direct_slot, zvector_slot), + ) + contractions = jk_ledger.contract( + driver, + obj.mol, + atmlst, + slots=(direct_slot, zvector_slot), + ) + local[0] += contractions[direct_slot] + local[1] += contractions[zvector_slot] + return local + + # 4-5. ROKS transpose-Hessian adjoint, Dz Fock derivative, and Pulay term. + return finish_gradient( + gradient_driver, + tdobj, + m_matrix, + direct, + atmlst, + tolerance, + max_cycle, + fock_direct, + direct_fock_probes=direct_fock_probes, + ) diff --git a/src/nest/grad/nttda/ensemble.py b/src/nest/grad/nttda/ensemble.py new file mode 100644 index 0000000..afb1b90 --- /dev/null +++ b/src/nest/grad/nttda/ensemble.py @@ -0,0 +1,199 @@ +"""Average-occupation orbital response for EnsembleRKS NTTDA gradients.""" + +import numpy as np + +from pyscf import lib +from pyscf.scf import hf + +from .roks import GradientComponents, _orbital_gradient, pack_m_matrix + + +def canonical_pairs(tdobj): + """Independent rotations between unequal-occupation orbital spaces.""" + occupation = np.asarray(tdobj._scf.mo_occ) + labels = {2: "c", 1: "o", 0: "v"} + rows, columns = np.where(hf.uniq_var_indices(occupation)) + return tuple( + (int(p), int(q), labels[int(occupation[q])] + labels[int(occupation[p])]) + for p, q in zip(rows, columns) + ) + + +def _rotation_matrix(vector, pairs, nmo): + rotation = np.zeros((nmo, nmo)) + for value, (p, q, _name) in zip(vector, pairs): + rotation[p, q] += value + rotation[q, p] -= value + return rotation + + +def _weighted_source(tdobj, pairs, vector): + occupation = np.asarray(tdobj._scf.mo_occ) + nmo = occupation.size + source = np.zeros((nmo, nmo)) + for value, (p, q, _name) in zip(vector, pairs): + source[p, q] += value * (occupation[q] - occupation[p]) + return source + + +def _fock_mo(mf): + orbitals = np.asarray(mf.mo_coeff) + return orbitals.conj().T @ np.asarray(mf.get_fock()) @ orbitals + + +def make_hessian_transpose_action(tdobj, pairs=None): + """Return the symmetric average-occupation orbital Hessian action.""" + mf = tdobj._scf + orbitals = np.asarray(mf.mo_coeff) + occupation = np.asarray(mf.mo_occ) + occupation_difference = occupation[None, :] - occupation[:, None] + nmo = orbitals.shape[1] + if pairs is None: + pairs = canonical_pairs(tdobj) + fock = _fock_mo(mf) + response = mf.gen_response(hermi=1) + + def apply_one(vector): + rotation = _rotation_matrix(vector, pairs, nmo) + density_mo = rotation * occupation_difference + density_ao = orbitals @ density_mo @ orbitals.conj().T + potential = response(density_ao) + potential_mo = orbitals.conj().T @ potential @ orbitals + fock_derivative = ( + fock @ rotation - rotation @ fock + potential_mo + ) + return np.asarray([ + occupation_difference[p, q] * fock_derivative[p, q] + for p, q, _name in pairs + ]) + + def apply(vector): + vector = np.asarray(vector) + if vector.ndim == 1: + return apply_one(vector) + return np.asarray([apply_one(row) for row in vector]) + + return apply, pairs + + +def zvector_adjoint_matrix(tdobj, pairs, zvector): + """Return the full coefficient derivative of ``z . g_orbital``.""" + mf = tdobj._scf + orbitals = np.asarray(mf.mo_coeff) + occupation = np.asarray(mf.mo_occ) + source = _weighted_source(tdobj, pairs, zvector) + fock = _fock_mo(mf) + gradient = fock @ (source + source.T) + + density = orbitals @ source @ orbitals.conj().T + density = 0.5 * (density + density.conj().T) + potential = mf.gen_response(hermi=1)(density) + potential = orbitals.conj().T @ potential @ orbitals + gradient += potential * occupation[None, :] + gradient += potential.conj().T * occupation[None, :] + return gradient + + +def zvector_probe_densities(tdobj, pairs, zvector): + """Spin probes for the nuclear derivative of the common ensemble Fock.""" + orbitals = np.asarray(tdobj._scf.mo_coeff) + source = _weighted_source(tdobj, pairs, zvector) + total = orbitals @ source @ orbitals.conj().T + return 0.5 * total, 0.5 * total + + +def _preconditioner(tdobj, pairs): + occupation = np.asarray(tdobj._scf.mo_occ) + epsilon = np.diag(_fock_mo(tdobj._scf)) + diagonal = np.asarray([ + (occupation[q] - occupation[p]) * (epsilon[p] - epsilon[q]) + for p, q, _name in pairs + ]) + small = np.abs(diagonal) < 1e-8 + diagonal[small] = np.where(diagonal[small] < 0.0, -1e-8, 1e-8) + return diagonal + + +def solve_zvector(action, pairs, tdobj, rhs, tolerance=1e-12, + max_cycle=None): + """Solve the average-occupation orbital adjoint equation.""" + diagonal = _preconditioner(tdobj, pairs) + initial = rhs / diagonal + if max_cycle is None: + max_cycle = len(rhs) + + def operator(vector): + vector = np.asarray(vector) + if vector.ndim == 1: + return action(vector) / diagonal - vector + return np.asarray([action(row) / diagonal - row for row in vector]) + + solution = lib.krylov( + operator, + initial, + tol=tolerance, + max_cycle=max_cycle, + lindep=1e-22, + hermi=False, + verbose=0, + ) + return np.asarray(solution).reshape(-1) + + +def finish_gradient( + gradient_driver, tdobj, m_matrix, direct, atmlst, + tolerance, max_cycle, fock_direct, direct_fock_probes=None): + """Solve the EnsembleRKS adjoint and assemble ``d omega / dR``.""" + transpose_action, pairs = make_hessian_transpose_action(tdobj) + rhs = pack_m_matrix(m_matrix, pairs) + zvector = solve_zvector( + transpose_action, + pairs, + tdobj, + rhs, + tolerance=tolerance, + max_cycle=max_cycle, + ) + adjoint = zvector_adjoint_matrix(tdobj, pairs, zvector) + residual = float(np.max(np.abs(pack_m_matrix(adjoint, pairs) - rhs))) + probe_alpha, probe_beta = zvector_probe_densities( + tdobj, pairs, zvector, + ) + if direct_fock_probes is None: + fock_contraction = fock_direct( + gradient_driver, tdobj, probe_alpha, probe_beta, atmlst=atmlst, + ) + direct_total = direct + else: + direct_alpha, direct_beta = direct_fock_probes + fock_contractions = fock_direct( + gradient_driver, + tdobj, + np.asarray((direct_alpha, probe_alpha)), + np.asarray((direct_beta, probe_beta)), + atmlst=atmlst, + ) + direct_total = direct + fock_contractions[0] + fock_contraction = fock_contractions[1] + orbital = _orbital_gradient( + tdobj, m_matrix, adjoint, fock_contraction, atmlst=atmlst, + ) + return GradientComponents( + m_matrix=m_matrix, + direct=direct_total, + orbital=orbital, + total=direct_total + orbital, + zvector=zvector, + residual=residual, + ) + + +__all__ = [ + "canonical_pairs", + "finish_gradient", + "make_hessian_transpose_action", + "pack_m_matrix", + "solve_zvector", + "zvector_adjoint_matrix", + "zvector_probe_densities", +] diff --git a/src/nest/grad/nttda/roks.py b/src/nest/grad/nttda/roks.py new file mode 100644 index 0000000..63bb36c --- /dev/null +++ b/src/nest/grad/nttda/roks.py @@ -0,0 +1,334 @@ +"""ROKS transpose-Hessian adjoint and final NTTDA gradient assembly.""" + +from dataclasses import dataclass + +import numpy as np + +from pyscf import dft, lib + + + +@dataclass(frozen=True) +class GradientComponents: + """Excitation-gradient pieces and the solved ROKS adjoint.""" + + m_matrix: np.ndarray + direct: np.ndarray + orbital: np.ndarray + total: np.ndarray + zvector: np.ndarray + residual: float + + +def finish_gradient( + gradient_driver, tdobj, m_matrix, direct, atmlst, + tolerance, max_cycle, fock_direct, direct_fock_probes=None): + """Solve the common ROKS Z-vector equation and assemble ``d omega/dR``. + + ``direct_fock_probes`` enables one batched Fock-derivative evaluation: its + contraction is the first result and the Z-vector contraction is the second. + """ + from nest.ensemble_rks import EnsembleRKS + + if isinstance(tdobj._scf, EnsembleRKS): + from .ensemble import finish_gradient as finish_ensemble_gradient + return finish_ensemble_gradient( + gradient_driver, + tdobj, + m_matrix, + direct, + atmlst, + tolerance, + max_cycle, + fock_direct, + direct_fock_probes=direct_fock_probes, + ) + transpose_action, pairs = make_hessian_transpose_action(tdobj) + rhs = pack_m_matrix(m_matrix, pairs) + zvector = solve_zvector( + transpose_action, + pairs, + tdobj, + rhs, + tolerance=tolerance, + max_cycle=max_cycle, + ) + adjoint = zvector_adjoint_matrix(tdobj, pairs, zvector) + residual = float(np.max(np.abs(pack_m_matrix(adjoint, pairs) - rhs))) + probe_alpha, probe_beta = zvector_probe_densities( + tdobj, pairs, zvector, + ) + if direct_fock_probes is None: + fock_contraction = fock_direct( + gradient_driver, + tdobj, + probe_alpha, + probe_beta, + atmlst=atmlst, + ) + direct_total = direct + else: + direct_alpha, direct_beta = direct_fock_probes + fock_contractions = fock_direct( + gradient_driver, + tdobj, + np.asarray((direct_alpha, probe_alpha)), + np.asarray((direct_beta, probe_beta)), + atmlst=atmlst, + ) + direct_total = direct + fock_contractions[0] + fock_contraction = fock_contractions[1] + orbital = _orbital_gradient( + tdobj, + m_matrix, + adjoint, + fock_contraction, + atmlst=atmlst, + ) + return GradientComponents( + m_matrix=m_matrix, + direct=direct_total, + orbital=orbital, + total=direct_total + orbital, + zvector=zvector, + residual=residual, + ) + + +def canonical_pairs(tdobj, compact=True): + """Canonical spatial-orbital rotations and their ROKS residual type.""" + occ = np.asarray(tdobj._scf.mo_occ) + closed = np.flatnonzero(occ == 2) + open_ = np.flatnonzero(occ == 1) + virtual = np.flatnonzero(occ == 0) + pairs = [] + if not compact: + for indices, name in ( + (closed, "cc"), (open_, "oo"), (virtual, "vv")): + for p_local in range(1, len(indices)): + for q_local in range(p_local): + pairs.append((indices[p_local], indices[q_local], name)) + pairs.extend((o, c, "co") for o in open_ for c in closed) + pairs.extend((v, c, "cv") for v in virtual for c in closed) + pairs.extend((v, o, "ov") for v in virtual for o in open_) + return tuple(pairs) + + +def _spin_focks_mo(mf): + fock = mf.get_fock() + mo = np.asarray(mf.mo_coeff) + return mo.conj().T @ fock.focka @ mo, mo.conj().T @ fock.fockb @ mo + + +def _response_reference(mf): + if (isinstance(mf, dft.KohnShamDFT) + and mf._numint._xc_type(mf.xc) != "HF"): + reference = mf.to_uks() + reference.verbose = 0 + return reference + return mf + + +def make_hessian_transpose_action(tdobj, pairs=None): + """Return a matrix-free action for the transpose ROKS Hessian.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + occ = np.asarray(mf.mo_occ) + nmo = mo.shape[1] + if pairs is None: + pairs = canonical_pairs(tdobj, compact=True) + fock_alpha, fock_beta = _spin_focks_mo(mf) + occupation_alpha = (occ > 0).astype(float) + occupation_beta = (occ == 2).astype(float) + response = _response_reference(mf).gen_response(hermi=1) + + def unpack(vector): + source_alpha = np.zeros((nmo, nmo)) + source_beta = np.zeros_like(source_alpha) + for value, (p, q, name) in zip(vector, pairs): + if name in ("cc", "oo", "vv"): + source_alpha[p, q] += 0.5 * value + source_beta[p, q] += 0.5 * value + elif name == "co": + source_beta[p, q] += value + elif name == "cv": + source_alpha[p, q] += value + source_beta[p, q] += value + elif name == "ov": + source_alpha[p, q] += value + else: + raise ValueError("unknown ROKS pair type %s" % name) + return source_alpha, source_beta + + def apply_one(vector): + source_alpha, source_beta = unpack(vector) + gradient = fock_alpha @ (source_alpha + source_alpha.T) + gradient += fock_beta @ (source_beta + source_beta.T) + density_alpha = mo @ source_alpha @ mo.conj().T + density_beta = mo @ source_beta @ mo.conj().T + density_alpha = 0.5 * (density_alpha + density_alpha.T) + density_beta = 0.5 * (density_beta + density_beta.T) + potential_alpha, potential_beta = response( + np.asarray((density_alpha, density_beta)) + ) + potential_alpha = mo.conj().T @ potential_alpha @ mo + potential_beta = mo.conj().T @ potential_beta @ mo + gradient += potential_alpha * occupation_alpha[None, :] + gradient += potential_alpha.T * occupation_alpha[None, :] + gradient += potential_beta * occupation_beta[None, :] + gradient += potential_beta.T * occupation_beta[None, :] + return np.asarray([ + gradient[p, q] - gradient[q, p] + for p, q, _name in pairs + ]) + + def apply(vector): + vector = np.asarray(vector) + if vector.ndim == 1: + return apply_one(vector) + return np.asarray([apply_one(row) for row in vector]) + + return apply, pairs + + +def pack_m_matrix(matrix, pairs): + antisymmetric = matrix - matrix.T + return np.asarray([antisymmetric[p, q] for p, q, _name in pairs]) + + +def _preconditioner(tdobj, pairs): + fock_alpha, fock_beta = _spin_focks_mo(tdobj._scf) + epsilon_alpha = np.diag(fock_alpha) + epsilon_beta = np.diag(fock_beta) + epsilon_common = 0.5 * (epsilon_alpha + epsilon_beta) + diagonal = [] + for p, q, name in pairs: + if name in ("cc", "oo", "vv"): + value = epsilon_common[p] - epsilon_common[q] + elif name == "co": + value = epsilon_beta[p] - epsilon_beta[q] + elif name == "cv": + value = ( + epsilon_alpha[p] - epsilon_alpha[q] + + epsilon_beta[p] - epsilon_beta[q] + ) + elif name == "ov": + value = epsilon_alpha[p] - epsilon_alpha[q] + else: + raise ValueError("unknown ROKS pair type %s" % name) + diagonal.append(value) + diagonal = np.asarray(diagonal) + small = np.abs(diagonal) < 1e-8 + diagonal[small] = np.where(diagonal[small] < 0.0, -1e-8, 1e-8) + return diagonal + + +def solve_zvector(action, pairs, tdobj, rhs, tolerance=1e-12, max_cycle=None): + """Solve ``H.T z = rhs`` using the PySCF CPHF Krylov pattern.""" + diagonal = _preconditioner(tdobj, pairs) + initial = rhs / diagonal + if max_cycle is None: + max_cycle = len(rhs) + + def operator(vector): + vector = np.asarray(vector) + if vector.ndim == 1: + return action(vector) / diagonal - vector + return np.asarray([action(row) / diagonal - row for row in vector]) + + solution = lib.krylov( + operator, + initial, + tol=tolerance, + max_cycle=max_cycle, + lindep=1e-22, + hermi=False, + verbose=0, + ) + return np.asarray(solution).reshape(-1) + + +def _unpack_zvector_source(tdobj, pairs, zvector): + nmo = tdobj._scf.mo_coeff.shape[1] + source_alpha = np.zeros((nmo, nmo)) + source_beta = np.zeros_like(source_alpha) + for value, (p, q, name) in zip(zvector, pairs): + if name in ("cc", "oo", "vv"): + source_alpha[p, q] += 0.5 * value + source_beta[p, q] += 0.5 * value + elif name == "co": + source_beta[p, q] += value + elif name == "cv": + source_alpha[p, q] += value + source_beta[p, q] += value + elif name == "ov": + source_alpha[p, q] += value + else: + raise ValueError("unknown ROKS pair type %s" % name) + return source_alpha, source_beta + + +def zvector_adjoint_matrix(tdobj, pairs, zvector): + """Full MO adjoint matrix satisfying ``z.H(kappa)=Tr(G.T kappa)``.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + occ = np.asarray(mf.mo_occ) + fock_alpha, fock_beta = _spin_focks_mo(mf) + occupation_alpha = (occ > 0).astype(float) + occupation_beta = (occ == 2).astype(float) + source_alpha, source_beta = _unpack_zvector_source( + tdobj, pairs, zvector, + ) + gradient = fock_alpha @ (source_alpha + source_alpha.T) + gradient += fock_beta @ (source_beta + source_beta.T) + density_alpha = mo @ source_alpha @ mo.conj().T + density_beta = mo @ source_beta @ mo.conj().T + density_alpha = 0.5 * (density_alpha + density_alpha.T) + density_beta = 0.5 * (density_beta + density_beta.T) + potential_alpha, potential_beta = _response_reference(mf).gen_response( + hermi=1, + )(np.asarray((density_alpha, density_beta))) + potential_alpha = mo.conj().T @ potential_alpha @ mo + potential_beta = mo.conj().T @ potential_beta @ mo + gradient += potential_alpha * occupation_alpha[None, :] + gradient += potential_alpha.T * occupation_alpha[None, :] + gradient += potential_beta * occupation_beta[None, :] + gradient += potential_beta.T * occupation_beta[None, :] + return gradient + + +def zvector_probe_densities(tdobj, pairs, zvector): + mo = np.asarray(tdobj._scf.mo_coeff) + source_alpha, source_beta = _unpack_zvector_source( + tdobj, pairs, zvector, + ) + return ( + mo @ source_alpha @ mo.conj().T, + mo @ source_beta @ mo.conj().T, + ) + + +def _orbital_gradient( + tdobj, m_matrix, adjoint, fock_contraction, atmlst=None): + mol = tdobj.mol + mf = tdobj._scf + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + mo = np.asarray(mf.mo_coeff) + overlap_derivative = mf.nuc_grad_method().get_ovlp(mol) + offsets = mol.offset_nr_by_atom() + result = np.zeros((len(atmlst), 3)) + for k, atom in enumerate(atmlst): + p0, p1 = offsets[atom][2:] + for xyz in range(3): + overlap = np.zeros((mol.nao_nr(), mol.nao_nr())) + overlap[p0:p1] += overlap_derivative[xyz, p0:p1] + overlap[:, p0:p1] += overlap_derivative[xyz, p0:p1].T + symmetric_kappa = -0.5 * (mo.conj().T @ overlap @ mo) + result[k, xyz] = ( + -fock_contraction[k, xyz] + - np.trace(adjoint.T @ symmetric_kappa) + + np.trace(m_matrix @ symmetric_kappa) + ) + return result diff --git a/src/nest/grad/nttda/xc.py b/src/nest/grad/nttda/xc.py new file mode 100644 index 0000000..f7768d8 --- /dev/null +++ b/src/nest/grad/nttda/xc.py @@ -0,0 +1,1546 @@ +"""LDA, GGA, and meta-GGA quadrature for NTTDA gradients. + +This module is channel-neutral: callers provide orbital spaces, transition +densities, block projections, and response-term coefficients. +""" + +from dataclasses import dataclass + +import numpy as np + +from pyscf import lib +from pyscf.dft.gen_grid import NBINS +from pyscf.dft.numint import _dot_ao_ao_sparse, _scale_ao_sparse +from pyscf.grad import tdrks as tdrks_grad + + + +# Shared result and projection helpers + +@dataclass(frozen=True) +class XCGradientTerms: + q_alpha: np.ndarray + q_beta: np.ndarray + direct: np.ndarray + + +# AO feature algebra + +def sparse_context(mf): + cutoff = mf.grids.cutoff * 1e2 + nbins = NBINS * 2 - int(NBINS * np.log(cutoff) / np.log(mf.grids.cutoff)) + pair_mask = mf.mol.get_overlap_cond() < -np.log(mf._numint.cutoff) + return nbins, pair_mask, mf.mol.ao_loc_nr() + + +def add_gga_matrix(mol, output, ao, weights, mask, sparse): + nbins, pair_mask, ao_loc = sparse + weights = np.asarray(weights, order="C").copy() + weights[0] *= 0.5 + scaled = _scale_ao_sparse(ao[:4], weights, mask, ao_loc) + matrix = _dot_ao_ao_sparse( + ao[0], scaled, None, nbins, mask, pair_mask, ao_loc, + hermi=0, out=None, + ) + output += lib.hermi_sum(matrix) + + +def add_mgga_matrix(mol, output, ao, weights, mask, sparse=None): + """Accumulate one ordinary meta-GGA feature potential matrix.""" + del sparse + output += mgga_eval_matrix(mol, ao, weights, mask) + + +def pair_matrix(mol, ao, mask, tensor, sparse): + nbins, pair_mask, ao_loc = sparse + output = np.zeros((mol.nao_nr(), mol.nao_nr())) + for left in range(4): + scaled = _scale_ao_sparse( + ao[:4], np.asarray(tensor[left], order="C"), mask, ao_loc, + ) + output += _dot_ao_ao_sparse( + ao[left], scaled, None, nbins, mask, pair_mask, ao_loc, + hermi=0, out=None, + ) + return output + + +def second_derivative_index(first, second): + if first > second: + first, second = second, first + return { + (0, 0): 4, + (0, 1): 5, + (0, 2): 6, + (1, 1): 7, + (1, 2): 8, + (2, 2): 9, + }[(first, second)] + + +def _compact_ao_center_derivative(ao, p0, p1, xyz, xctype): + """AO-center derivative restricted to one atom's AO columns.""" + if xctype == "LDA": + return -ao[xyz + 1][:, p0:p1] + delta = np.empty((4, ao.shape[-2], p1 - p0)) + delta[0] = -ao[xyz + 1][:, p0:p1] + for feature in range(3): + delta[feature + 1] = -ao[ + second_derivative_index(xyz, feature) + ][:, p0:p1] + return delta + + +def _hermitian_density_derivative_batches( + ao, densities, p0, p1, xctype): + """Yield AO-center derivatives for a stack of real symmetric densities.""" + densities = np.asarray(densities) + feature_count = 1 if xctype == "LDA" else 4 + density_rows = densities[:, p0:p1] + packed_rows = density_rows.transpose(2, 0, 1).reshape( + density_rows.shape[-1], -1, + ) + contracted = (ao[:feature_count] @ packed_rows).reshape( + feature_count, ao.shape[-2], len(densities), p1 - p0, + ).transpose(2, 0, 1, 3) + + for xyz in range(3): + delta = _compact_ao_center_derivative( + ao, p0, p1, xyz, xctype, + ) + if xctype == "LDA": + yield 2.0 * lib.einsum( + "ga,nga->ng", delta, contracted[:, 0], + )[:, None] + continue + + derivative_count = 4 if xctype == "GGA" else 5 + output = np.empty(( + len(densities), derivative_count, ao.shape[-2], + )) + output[:, 0] = 2.0 * lib.einsum( + "ga,nga->ng", delta[0], contracted[:, 0], + ) + for feature in range(1, 4): + output[:, feature] = 2.0 * ( + lib.einsum( + "ga,nga->ng", delta[feature], contracted[:, 0], + ) + + lib.einsum( + "ga,nga->ng", delta[0], contracted[:, feature], + ) + ) + if xctype == "GGA": + yield output + continue + output[:, 4] = 0.0 + for feature in range(1, 4): + output[:, 4] += lib.einsum( + "ga,nga->ng", delta[feature], contracted[:, feature], + ) + yield output + + +def pair_feature_batches(ao, densities): + """Pair features and reusable ``AO @ D`` contractions by channel.""" + densities = np.asarray(densities) + grids = ao.shape[-2] + features = np.empty((len(densities), 4, 4, grids)) + contracted = np.asarray([ + ao[index] @ densities for index in range(4) + ]).transpose(1, 0, 2, 3) + for left in range(4): + for right in range(4): + features[:, left, right] = lib.einsum( + "ngu,gu->ng", contracted[:, left], ao[right], + ) + return features, contracted + + +def contract_pair_feature_derivatives( + ao, densities, delta, contracted_ao, p0, p1, + tensor_weights, grid_weights): + """Contract pair-feature derivatives without materializing ``dPair``.""" + densities = np.asarray(densities) + tensor_weights = np.asarray(tensor_weights) + value = 0.0 + for left in range(4): + contracted_delta = delta[left] @ densities[:, p0:p1] + value += lib.einsum( + "pbg,pgu,bgu,g->", + tensor_weights[:, left], contracted_delta, ao[:4], + grid_weights, optimize=True, + ) + contracted_atom = contracted_ao[:, :, :, p0:p1] + for right in range(4): + value += lib.einsum( + "pag,pagq,gq,g->", + tensor_weights[:, :, right], contracted_atom, + delta[right], grid_weights, optimize=True, + ) + return value + + +def gga_pair_potential(kernel, features): + output = np.zeros_like(features) + output[0, 0] = lib.einsum("abg,abg->g", kernel, features) + output[1:4, 0] = kernel[1:4, 0] * features[0, 0] + output[1:4, 0] += lib.einsum( + "ijg,jg->ig", kernel[1:4, 1:4], features[0, 1:4], + ) + output[0, 1:4] = kernel[0, 1:4] * features[0, 0] + output[0, 1:4] += lib.einsum( + "ijg,ig->jg", kernel[1:4, 1:4], features[1:4, 0], + ) + output[1:4, 1:4] = kernel[1:4, 1:4] * features[0, 0] + return output + + +def gga_pair_kernel_cross(left, right): + output = np.zeros_like(left) + output[0, 0] = left[0, 0] * right[0, 0] + output[1:4, 0] = ( + left[0, 0][None] * right[1:4, 0] + + left[1:4, 0] * right[0, 0][None] + ) + output[0, 1:4] = ( + left[0, 0][None] * right[0, 1:4] + + left[0, 1:4] * right[0, 0][None] + ) + output[1:4, 1:4] = ( + left[0, 0][None, None] * right[1:4, 1:4] + + left[1:4, 0][:, None] * right[0, 1:4][None] + + left[0, 1:4][None] * right[1:4, 0][:, None] + + left[1:4, 1:4] * right[0, 0][None, None] + ) + return output + + +def mgga_pair_potential(kernel, features): + output = np.zeros_like(features) + output[0, 0] = lib.einsum( + "abg,abg->g", kernel[:4, :4], features, + ) + output[1:4, 0] = kernel[1:4, 0] * features[0, 0] + output[1:4, 0] += lib.einsum( + "ijg,jg->ig", kernel[1:4, 1:4], features[0, 1:4], + ) + output[1:4, 0] += 0.5 * kernel[4, 0][None] * features[1:4, 0] + output[1:4, 0] += 0.5 * lib.einsum( + "jg,ijg->ig", kernel[4, 1:4], features[1:4, 1:4], + ) + output[0, 1:4] = kernel[0, 1:4] * features[0, 0] + output[0, 1:4] += lib.einsum( + "ijg,ig->jg", kernel[1:4, 1:4], features[1:4, 0], + ) + output[0, 1:4] += 0.5 * kernel[0, 4][None] * features[0, 1:4] + output[0, 1:4] += 0.5 * lib.einsum( + "ig,ijg->jg", kernel[1:4, 4], features[1:4, 1:4], + ) + output[1:4, 1:4] = kernel[1:4, 1:4] * features[0, 0] + output[1:4, 1:4] += 0.5 * lib.einsum( + "ig,jg->ijg", kernel[1:4, 4], features[0, 1:4], + ) + output[1:4, 1:4] += 0.5 * lib.einsum( + "jg,ig->ijg", kernel[4, 1:4], features[1:4, 0], + ) + output[1:4, 1:4] += 0.25 * kernel[4, 4][None, None] * features[1:4, 1:4] + return output + + +def mgga_pair_kernel_cross(left, right): + grids = left.shape[-1] + output = np.zeros((5, 5, grids)) + output[:4, :4] = gga_pair_kernel_cross(left, right) + output[4, 0] = 0.5 * lib.einsum( + "ig,ig->g", left[1:4, 0], right[1:4, 0], + ) + output[0, 4] = 0.5 * lib.einsum( + "jg,jg->g", left[0, 1:4], right[0, 1:4], + ) + output[4, 1:4] = 0.5 * lib.einsum( + "ig,ijg->jg", left[1:4, 0], right[1:4, 1:4], + ) + output[4, 1:4] += 0.5 * lib.einsum( + "ijg,ig->jg", left[1:4, 1:4], right[1:4, 0], + ) + output[1:4, 4] = 0.5 * lib.einsum( + "jg,ijg->ig", left[0, 1:4], right[1:4, 1:4], + ) + output[1:4, 4] += 0.5 * lib.einsum( + "ijg,jg->ig", left[1:4, 1:4], right[0, 1:4], + ) + output[4, 4] = 0.25 * lib.einsum( + "ijg,ijg->g", left[1:4, 1:4], right[1:4, 1:4], + ) + return output + + +def gga_eval_matrix(mol, ao, weights, mask): + output = np.zeros((4, mol.nao_nr(), mol.nao_nr())) + tdrks_grad._gga_eval_mat_( + mol, output, ao, np.array(weights, copy=True), mask, + (0, mol.nbas), mol.ao_loc_nr(), + ) + return output[0] + + +def mgga_eval_matrix(mol, ao, weights, mask): + output = np.zeros((4, mol.nao_nr(), mol.nao_nr())) + tdrks_grad._mgga_eval_mat_( + mol, output, ao, np.array(weights, copy=True), mask, + (0, mol.nbas), mol.ao_loc_nr(), + ) + return output[0] + + +# LDA quadrature + +def _lda_fref_kref(mf, ao0, mask): + ni = mf._numint + rho0 = ni.eval_rho2( + mf.mol, ao0, mf.mo_coeff, mf.mo_occ, mask, "LDA", + with_lapl=False, + ) * 0.5 + fxc, kxc = ni.eval_xc_eff( + mf.xc, (rho0, rho0), deriv=3, xctype="LDA", spin=1, + )[2:4] + fref = 0.5 * ( + fxc[0, 0, 0, 0] - fxc[0, 0, 1, 0] + - fxc[1, 0, 0, 0] + fxc[1, 0, 1, 0] + ) + kref_alpha = 0.5 * ( + kxc[0, 0, 0, 0, 0, 0] - kxc[0, 0, 1, 0, 0, 0] + - kxc[1, 0, 0, 0, 0, 0] + kxc[1, 0, 1, 0, 0, 0] + ) + kref_beta = 0.5 * ( + kxc[0, 0, 0, 0, 1, 0] - kxc[0, 0, 1, 0, 1, 0] + - kxc[1, 0, 0, 0, 1, 0] + kxc[1, 0, 1, 0, 1, 0] + ) + return fref, kref_alpha, kref_beta + + +def _lda_matrix(ao0, weights): + return ao0.T @ (ao0 * np.asarray(weights)[:, None]) + + +def _project_channel_potentials(tdobj, potentials, blocks): + """Project transition-factor potentials for any NTTDA spin channel.""" + mo = np.asarray(tdobj._scf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + for label, (target, source, coefficient) in blocks.items(): + potential = mo.conj().T @ potentials[label] @ mo + q_beta[:, target] += potential[:, source] @ coefficient.T + q_alpha[:, source] += potential[target, :].T @ coefficient + return q_alpha, q_beta + + +def _reference_spin_densities(tdobj): + """Spin densities of the variational reference used by the XC kernel.""" + mf = tdobj._scf + if getattr(mf, "is_ensemble_rks", False): + return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) + mo = np.asarray(mf.mo_coeff) + return ( + mo[:, mf.mo_occ > 0] @ mo[:, mf.mo_occ > 0].T, + mo[:, mf.mo_occ == 2] @ mo[:, mf.mo_occ == 2].T, + ) + + +def _reference_spin_occupations(tdobj): + """Per-orbital alpha/beta occupations of the reference density.""" + mf = tdobj._scf + occupation = np.asarray(mf.mo_occ) + if getattr(mf, "is_ensemble_rks", False): + return 0.5 * occupation, 0.5 * occupation + return (occupation > 0).astype(float), (occupation == 2).astype(float) + + +def _add_reference_q(tdobj, q_alpha, q_beta, matrix_alpha, matrix_beta): + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + occupation_alpha, occupation_beta = _reference_spin_occupations(tdobj) + q_alpha += ( + mo.conj().T @ (matrix_alpha + matrix_alpha.T) @ mo + ) * occupation_alpha[None, :] + q_beta += ( + mo.conj().T @ (matrix_beta + matrix_beta.T) @ mo + ) * occupation_beta[None, :] + + +def _spin_probe_stacks(probe_alpha, probe_beta): + probe_alpha = np.asarray(probe_alpha) + probe_beta = np.asarray(probe_beta) + single_probe = probe_alpha.ndim == 2 + if single_probe: + probe_alpha = probe_alpha[None] + probe_beta = probe_beta[None] + probe_alpha = 0.5 * ( + probe_alpha + probe_alpha.swapaxes(-1, -2) + ) + probe_beta = 0.5 * ( + probe_beta + probe_beta.swapaxes(-1, -2) + ) + return probe_alpha, probe_beta, single_probe + + +def _xc_density(ni, mol, ao, density, mask, xctype): + ao_values = ao[0] if xctype == "LDA" else ao + rho = ni.eval_rho( + mol, ao_values, density, mask, xctype, hermi=1, + with_lapl=False, + ) + return rho[None] if rho.ndim == 1 else rho + + +def _xc_ao_center_derivative(ao, p0, p1, xyz, xctype): + return _compact_ao_center_derivative( + ao, p0, p1, xyz, xctype, + ) + + +def _xc_density_derivatives( + ao, densities, p0, p1, xctype, ao_center_derivative): + """AO-center derivatives for a stack of probe/reference densities.""" + densities = np.asarray(densities) + if xctype == "LDA": + delta0 = ao_center_derivative + output = lib.einsum( + "ga,nau,gu->ng", + delta0, + densities[:, p0:p1], + ao[0], + ) + output += lib.einsum( + "gu,nua,ga->ng", + ao[0], + densities[:, :, p0:p1], + delta0, + ) + return output[:, None] + + delta = ao_center_derivative + feature_count = 4 if xctype == "GGA" else 5 + output = np.empty( + (len(densities), feature_count, ao.shape[-2]), + ) + output[:, 0] = lib.einsum( + "ga,nau,gu->ng", + delta[0], + densities[:, p0:p1], + ao[0], + ) + output[:, 0] += lib.einsum( + "gu,nua,ga->ng", + ao[0], + densities[:, :, p0:p1], + delta[0], + ) + for feature in range(1, 4): + output[:, feature] = lib.einsum( + "ga,nau,gu->ng", + delta[feature], + densities[:, p0:p1], + ao[0], + ) + output[:, feature] += lib.einsum( + "gu,nua,ga->ng", + ao[feature], + densities[:, :, p0:p1], + delta[0], + ) + output[:, feature] += lib.einsum( + "ga,nau,gu->ng", + delta[0], + densities[:, p0:p1], + ao[feature], + ) + output[:, feature] += lib.einsum( + "gu,nua,ga->ng", + ao[0], + densities[:, :, p0:p1], + delta[feature], + ) + if xctype == "GGA": + return output[:, :4] + + output[:, 4] = 0.0 + for feature in range(1, 4): + output[:, 4] += 0.5 * lib.einsum( + "ga,nau,gu->ng", + delta[feature], + densities[:, p0:p1], + ao[feature], + ) + output[:, 4] += 0.5 * lib.einsum( + "gu,nua,ga->ng", + ao[feature], + densities[:, :, p0:p1], + delta[feature], + ) + return output + + +def _response_density_stack( + densities, density_alpha, density_beta): + labels = tuple(densities) + stack = np.asarray( + [densities[label] for label in labels] + + [density_alpha, density_beta] + ) + return labels, stack + + +def _response_density_derivatives( + ao, density_stack, labels, p0, p1, xyz, xctype): + """Generate every channel/reference density derivative from one AO delta.""" + ao_center_derivative = _xc_ao_center_derivative( + ao, p0, p1, xyz, xctype, + ) + derivatives = _xc_density_derivatives( + ao, density_stack, p0, p1, xctype, ao_center_derivative, + ) + if xctype == "LDA": + derivatives = derivatives[:, 0] + channel_count = len(labels) + channel_derivatives = dict(zip( + labels, derivatives[:channel_count], + )) + return ( + channel_derivatives, + derivatives[channel_count], + derivatives[channel_count + 1], + ao_center_derivative, + ) + + +def _contract_vxc_derivative( + mf, density_alpha, density_beta, probe_alpha, probe_beta, + atmlst, xctype, max_memory): + """Contract all fixed-grid XC potential derivatives in one grid pass.""" + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + probe_alpha, probe_beta, single_probe = _spin_probe_stacks( + probe_alpha, probe_beta, + ) + output = np.zeros((len(probe_alpha), len(atmlst), 3)) + if not atmlst: + return output[0] if single_probe else output + + density_alpha = 0.5 * ( + np.asarray(density_alpha) + np.asarray(density_alpha).T + ) + density_beta = 0.5 * ( + np.asarray(density_beta) + np.asarray(density_beta).T + ) + probe_densities = np.stack( + (probe_alpha, probe_beta), axis=1, + ).reshape(-1, *probe_alpha.shape[1:]) + density_stack = np.concatenate(( + np.asarray((density_alpha, density_beta)), + probe_densities, + )) + offsets = mol.offset_nr_by_atom() + ao_deriv = 1 if xctype == "LDA" else 2 + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, mol.nao_nr(), ao_deriv, + max_memory=max_memory): + rho = np.asarray([ + _xc_density(ni, mol, ao, density, mask, xctype) + for density in density_stack + ]) + reference_rho = rho[:2] + probe_rho = rho[2:].reshape( + len(probe_alpha), 2, *rho.shape[1:], + ) + vxc, fxc = ni.eval_xc_eff( + mf.xc, reference_rho, deriv=2, xctype=xctype, spin=1, + )[1:3] + + for k, atom in enumerate(atmlst): + p0, p1 = offsets[atom][2:] + for xyz in range(3): + ao_center_derivative = _xc_ao_center_derivative( + ao, p0, p1, xyz, xctype, + ) + density_derivative = _xc_density_derivatives( + ao, density_stack, p0, p1, xctype, + ao_center_derivative, + ) + reference_derivative = density_derivative[:2] + probe_derivative = density_derivative[2:].reshape( + len(probe_alpha), 2, *density_derivative.shape[1:], + ) + output[:, k, xyz] += lib.einsum( + "nsxg,sxg,g->n", + probe_derivative, + vxc, + weights, + ) + response_weights = lib.einsum( + "axg,axbyg,g->byg", + reference_derivative, + fxc, + weights, + ) + output[:, k, xyz] += lib.einsum( + "nbyg,byg->n", probe_rho, response_weights, + ) + return output[0] if single_probe else output + + +def contract_lda_vxc_derivative( + mf, density_alpha, density_beta, probe_alpha, probe_beta, + atmlst=None, max_memory=2000): + """Contract all requested LDA XC potential nuclear derivatives.""" + return _contract_vxc_derivative( + mf, density_alpha, density_beta, probe_alpha, probe_beta, + atmlst, "LDA", max_memory, + ) + + +def contract_gga_vxc_derivative( + mf, density_alpha, density_beta, probe_alpha, probe_beta, + atmlst=None, max_memory=2000): + """Contract all requested GGA XC potential nuclear derivatives.""" + return _contract_vxc_derivative( + mf, density_alpha, density_beta, probe_alpha, probe_beta, + atmlst, "GGA", max_memory, + ) + + +def contract_mgga_vxc_derivative( + mf, density_alpha, density_beta, probe_alpha, probe_beta, + atmlst=None, max_memory=2000): + """Contract all requested MGGA XC potential nuclear derivatives.""" + return _contract_vxc_derivative( + mf, density_alpha, density_beta, probe_alpha, probe_beta, + atmlst, "MGGA", max_memory, + ) + + +def lda_response_terms( + gradient_driver, tdobj, channel_data, atmlst=None, + with_direct=True): + """Analytic LDA M/direct terms for the ``vref0/vref1`` scalar.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + _spaces, _amplitudes, densities, blocks, response_terms = channel_data + coefficients = tuple( + (term.target, term.source, term.vref0 + term.vref1) + for term in response_terms + if term.vref0 + term.vref1 + ) + nao = mol.nao_nr() + potentials = {label: np.zeros((nao, nao)) for label in densities} + reference_alpha = np.zeros((nao, nao)) + reference_beta = np.zeros_like(reference_alpha) + direct = np.zeros((len(atmlst), 3)) + mo = np.asarray(mf.mo_coeff) + density_alpha, density_beta = _reference_spin_densities(tdobj) + density_labels, density_stack = _response_density_stack( + densities, density_alpha, density_beta, + ) + offsets = mol.offset_nr_by_atom() + + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, nao, 1, max_memory=gradient_driver.max_memory): + ao0 = ao[0] + fref, kref_alpha, kref_beta = _lda_fref_kref(mf, ao0, mask) + rho = { + label: ni.eval_rho( + mol, ao0, density, mask, "LDA", hermi=0, + with_lapl=False, + ) + for label, density in densities.items() + } + potential_weights = { + label: np.zeros_like(weights) for label in densities + } + pair_alpha = np.zeros_like(weights) + pair_beta = np.zeros_like(weights) + for target, source, coefficient in coefficients: + potential_weights[target] += coefficient * fref * rho[source] + potential_weights[source] += coefficient * fref * rho[target] + pair = coefficient * rho[target] * rho[source] + pair_alpha += kref_alpha * pair + pair_beta += kref_beta * pair + potential_weight_stack = np.asarray([ + potential_weights[label] for label in density_labels + ]) + for label in potentials: + potentials[label] += _lda_matrix( + ao0, weights * potential_weights[label], + ) + reference_alpha += _lda_matrix(ao0, weights * pair_alpha) + reference_beta += _lda_matrix(ao0, weights * pair_beta) + + if not with_direct: + continue + for k, atom in enumerate(atmlst): + p0, p1 = offsets[atom][2:] + for xyz in range(3): + drho, drho_alpha, drho_beta, _ao_delta = ( + _response_density_derivatives( + ao, density_stack, density_labels, + p0, p1, xyz, "LDA", + ) + ) + drho_stack = np.asarray([ + drho[label] for label in density_labels + ]) + value = lib.einsum( + "ng,ng,g->", + potential_weight_stack, drho_stack, weights, + ) + value += lib.einsum( + "g,g,g->", pair_alpha, drho_alpha, weights, + ) + value += lib.einsum( + "g,g,g->", pair_beta, drho_beta, weights, + ) + direct[k, xyz] += value + + q_alpha, q_beta = _project_channel_potentials( + tdobj, potentials, blocks, + ) + _add_reference_q( + tdobj, q_alpha, q_beta, reference_alpha, reference_beta, + ) + return XCGradientTerms(q_alpha, q_beta, direct) + + +def lda_fockz_terms( + gradient_driver, tdobj, spaces, pz, atmlst=None, + with_direct=True): + """LDA response/direct derivative of ``Pz:Fz`` excluding Pz projection.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + density_open = spaces.c_open @ spaces.c_open.T + pz = np.asarray(pz) + pz_symmetric = 0.5 * (pz + pz.T) + nao = mol.nao_nr() + open_potential = np.zeros((nao, nao)) + reference_alpha = np.zeros((nao, nao)) + reference_beta = np.zeros_like(reference_alpha) + direct = np.zeros((len(atmlst), 3)) + mo = np.asarray(mf.mo_coeff) + density_alpha, density_beta = _reference_spin_densities(tdobj) + density_stack = np.asarray(( + pz_symmetric, density_open, density_alpha, density_beta, + )) + offsets = mol.offset_nr_by_atom() + + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, nao, 1, max_memory=gradient_driver.max_memory): + ao0 = ao[0] + fref, kref_alpha, kref_beta = _lda_fref_kref(mf, ao0, mask) + rho_pz = ni.eval_rho( + mol, ao0, pz_symmetric, mask, "LDA", hermi=1, + with_lapl=False, + ) + rho_open = ni.eval_rho( + mol, ao0, density_open, mask, "LDA", hermi=1, + with_lapl=False, + ) + open_potential += _lda_matrix( + ao0, 0.5 * weights * fref * rho_pz, + ) + pair = 0.5 * rho_pz * rho_open + reference_alpha += _lda_matrix( + ao0, weights * kref_alpha * pair, + ) + reference_beta += _lda_matrix( + ao0, weights * kref_beta * pair, + ) + if not with_direct: + continue + for k, atom in enumerate(atmlst): + p0, p1 = offsets[atom][2:] + derivative_batches = _hermitian_density_derivative_batches( + ao, density_stack, p0, p1, "LDA", + ) + for xyz, derivatives in enumerate(derivative_batches): + drho_pz, drho_open, drho_alpha, drho_beta = ( + derivatives[:, 0] + ) + direct[k, xyz] += 0.5 * np.dot( + weights, + fref * ( + drho_pz * rho_open + rho_pz * drho_open + ) + + rho_pz * rho_open * ( + kref_alpha * drho_alpha + + kref_beta * drho_beta + ), + ) + + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + q_alpha[:, spaces.open] += ( + mo.conj().T @ (open_potential + open_potential.T) + @ spaces.c_open + ) + _add_reference_q( + tdobj, q_alpha, q_beta, reference_alpha, reference_beta, + ) + return XCGradientTerms(q_alpha, q_beta, direct) + + +def lda_nobeta_reference_q(tdobj, p0, max_memory=None): + """Reference-density correction for the equal-spin ``nobeta`` Fock.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + nmo = mo.shape[1] + q_alpha = np.zeros((nmo, nmo)) + q_beta = np.zeros_like(q_alpha) + if not tdobj.nobeta or getattr(mf, "is_ensemble_rks", False): + return q_alpha, q_beta + if max_memory is None: + max_memory = tdobj.max_memory + ni = mf._numint + mol = mf.mol + nao = mol.nao_nr() + density_alpha, density_beta = _reference_spin_densities(tdobj) + density0 = 0.5 * (density_alpha + density_beta) + p0 = 0.5 * (np.asarray(p0) + np.asarray(p0).T) + matrix_alpha = np.zeros((nao, nao)) + matrix_beta = np.zeros_like(matrix_alpha) + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, nao, 1, max_memory=max_memory): + ao0 = ao[0] + rho_p = ni.eval_rho( + mol, ao0, p0, mask, "LDA", hermi=1, with_lapl=False, + ) + rho_alpha = ni.eval_rho( + mol, ao0, density_alpha, mask, "LDA", hermi=1, + with_lapl=False, + ) + rho_beta = ni.eval_rho( + mol, ao0, density_beta, mask, "LDA", hermi=1, + with_lapl=False, + ) + rho0 = ni.eval_rho( + mol, ao0, density0, mask, "LDA", hermi=1, + with_lapl=False, + ) + fxc_actual = ni.eval_xc_eff( + mf.xc, (rho_alpha, rho_beta), deriv=2, + xctype="LDA", spin=1, + )[2] + fxc_equal = ni.eval_xc_eff( + mf.xc, (rho0, rho0), deriv=2, + xctype="LDA", spin=1, + )[2] + equal_derivative = 0.25 * ( + fxc_equal[0, 0, 0, 0] + fxc_equal[0, 0, 1, 0] + + fxc_equal[1, 0, 0, 0] + fxc_equal[1, 0, 1, 0] + ) + actual_alpha = 0.5 * ( + fxc_actual[0, 0, 0, 0] + fxc_actual[1, 0, 0, 0] + ) + actual_beta = 0.5 * ( + fxc_actual[0, 0, 1, 0] + fxc_actual[1, 0, 1, 0] + ) + matrix_alpha += _lda_matrix( + ao0, weights * rho_p * (equal_derivative - actual_alpha), + ) + matrix_beta += _lda_matrix( + ao0, weights * rho_p * (equal_derivative - actual_beta), + ) + _add_reference_q( + tdobj, q_alpha, q_beta, matrix_alpha, matrix_beta, + ) + return q_alpha, q_beta + + +# GGA quadrature + +def _gga_fref_kref(mf, rho0): + fxc, kxc = mf._numint.eval_xc_eff( + mf.xc, (rho0, rho0), deriv=3, xctype="GGA", spin=1, + )[2:4] + fref = 0.5 * ( + fxc[0, :, 0] - fxc[0, :, 1] + - fxc[1, :, 0] + fxc[1, :, 1] + ) + kref_alpha = 0.5 * ( + kxc[0, :, 0, :, 0] - kxc[0, :, 1, :, 0] + - kxc[1, :, 0, :, 0] + kxc[1, :, 1, :, 0] + ) + kref_beta = 0.5 * ( + kxc[0, :, 0, :, 1] - kxc[0, :, 1, :, 1] + - kxc[1, :, 0, :, 1] + kxc[1, :, 1, :, 1] + ) + return fref, kref_alpha, kref_beta + + +def gga_response_terms( + gradient_driver, tdobj, channel_data, atmlst=None, + with_direct=True): + """GGA ``vref0/vref1`` M matrix and fixed-grid skeleton derivative.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + _spaces, _amplitudes, densities, blocks, terms = channel_data + pair_labels = tuple( + label for label in densities + if any( + term.vref1 and label in (term.target, term.source) + for term in terms + ) + ) + pair_density_stack = np.asarray([ + densities[label] for label in pair_labels + ]) + nao = mol.nao_nr() + potentials = {label: np.zeros((nao, nao)) for label in densities} + reference_alpha = np.zeros((nao, nao)) + reference_beta = np.zeros_like(reference_alpha) + direct = np.zeros((len(atmlst), 3)) + mo = np.asarray(mf.mo_coeff) + density_alpha, density_beta = _reference_spin_densities(tdobj) + density_labels, density_stack = _response_density_stack( + densities, density_alpha, density_beta, + ) + offsets = mol.offset_nr_by_atom() + sparse = sparse_context(mf) + + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): + rho0 = ni.eval_rho2( + mol, ao, mo, mf.mo_occ, mask, "GGA", with_lapl=False, + ) * 0.5 + fref, kref_alpha, kref_beta = _gga_fref_kref(mf, rho0) + rho = { + label: ni.eval_rho( + mol, ao, density, mask, "GGA", hermi=0, + with_lapl=False, + ) + for label, density in densities.items() + } + pair_values, contracted_pair_ao = pair_feature_batches( + ao, pair_density_stack, + ) + pairs = dict(zip(pair_labels, pair_values)) + pair_potentials = { + label: gga_pair_potential(fref, pairs[label]) + for label in pair_labels + } + ordinary_weights = { + label: np.zeros((4, weights.size)) for label in densities + } + special_weights = { + label: np.zeros((4, 4, weights.size)) for label in pair_labels + } + reference_weights_alpha = np.zeros((4, weights.size)) + reference_weights_beta = np.zeros_like(reference_weights_alpha) + + for term in terms: + if term.vref0: + ordinary_weights[term.target] += term.vref0 * lib.einsum( + "xyg,yg->xg", fref, rho[term.source], + ) + ordinary_weights[term.source] += term.vref0 * lib.einsum( + "xyg,xg->yg", fref, rho[term.target], + ) + pair = term.vref0 * lib.einsum( + "xg,yg->xyg", rho[term.target], rho[term.source], + ) + reference_weights_alpha += lib.einsum( + "xyg,xyzg->zg", pair, kref_alpha, + ) + reference_weights_beta += lib.einsum( + "xyg,xyzg->zg", pair, kref_beta, + ) + if term.vref1: + special_weights[term.target] += ( + term.vref1 * pair_potentials[term.source] + ) + special_weights[term.source] += ( + term.vref1 * pair_potentials[term.target] + ) + pair = term.vref1 * gga_pair_kernel_cross( + pairs[term.target], pairs[term.source], + ) + reference_weights_alpha += lib.einsum( + "xyg,xyzg->zg", pair, kref_alpha, + ) + reference_weights_beta += lib.einsum( + "xyg,xyzg->zg", pair, kref_beta, + ) + ordinary_weight_stack = np.asarray([ + ordinary_weights[label] for label in density_labels + ]) + special_weight_stack = np.asarray([ + special_weights[label] for label in pair_labels + ]) + + for label in potentials: + add_gga_matrix( + mol, potentials[label], ao, + ordinary_weights[label] * weights, mask, sparse, + ) + for label in pair_labels: + potentials[label] += pair_matrix( + mol, ao, mask, special_weights[label] * weights, sparse, + ) + reference_alpha += gga_eval_matrix( + mol, ao, reference_weights_alpha * weights, mask, + ) + reference_beta += gga_eval_matrix( + mol, ao, reference_weights_beta * weights, mask, + ) + + if not with_direct: + continue + for k, atom in enumerate(atmlst): + p0, p1 = offsets[atom][2:] + for xyz in range(3): + drho, drho_alpha, drho_beta, ao_delta = ( + _response_density_derivatives( + ao, density_stack, density_labels, + p0, p1, xyz, "GGA", + ) + ) + drho_stack = np.asarray([ + drho[label] for label in density_labels + ]) + value = lib.einsum( + "nfg,nfg,g->", + ordinary_weight_stack, drho_stack, weights, + ) + value += lib.einsum( + "fg,fg,g->", + reference_weights_alpha, drho_alpha, weights, + ) + value += lib.einsum( + "fg,fg,g->", + reference_weights_beta, drho_beta, weights, + ) + value += contract_pair_feature_derivatives( + ao, pair_density_stack, ao_delta, + contracted_pair_ao, p0, p1, + special_weight_stack, weights, + ) + direct[k, xyz] += value + + q_alpha, q_beta = _project_channel_potentials( + tdobj, potentials, blocks, + ) + _add_reference_q( + tdobj, q_alpha, q_beta, reference_alpha, reference_beta, + ) + return XCGradientTerms(q_alpha, q_beta, direct) + + +def gga_fockz_terms( + gradient_driver, tdobj, spaces, pz, atmlst=None, + with_direct=True): + """GGA derivative of ``Pz:Fz`` excluding its explicit Pz projection.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + density_open = spaces.c_open @ spaces.c_open.T + pz = 0.5 * (np.asarray(pz) + np.asarray(pz).T) + nao = mol.nao_nr() + open_potential = np.zeros((nao, nao)) + reference_alpha = np.zeros((nao, nao)) + reference_beta = np.zeros_like(reference_alpha) + direct = np.zeros((len(atmlst), 3)) + mo = np.asarray(mf.mo_coeff) + density_alpha, density_beta = _reference_spin_densities(tdobj) + density_stack = np.asarray(( + pz, density_open, density_alpha, density_beta, + )) + offsets = mol.offset_nr_by_atom() + sparse = sparse_context(mf) + + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): + rho0 = ni.eval_rho2( + mol, ao, mo, mf.mo_occ, mask, "GGA", with_lapl=False, + ) * 0.5 + fref, kref_alpha, kref_beta = _gga_fref_kref(mf, rho0) + rho_pz = ni.eval_rho( + mol, ao, pz, mask, "GGA", hermi=1, with_lapl=False, + ) + rho_open = ni.eval_rho( + mol, ao, density_open, mask, "GGA", hermi=1, + with_lapl=False, + ) + add_gga_matrix( + mol, + open_potential, + ao, + 0.5 * lib.einsum("xyg,yg->xg", fref, rho_pz) * weights, + mask, + sparse, + ) + pair = 0.5 * lib.einsum("xg,yg->xyg", rho_pz, rho_open) + reference_alpha += gga_eval_matrix( + mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_alpha) * weights, + mask, + ) + reference_beta += gga_eval_matrix( + mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_beta) * weights, + mask, + ) + if not with_direct: + continue + for k, atom in enumerate(atmlst): + p0, p1 = offsets[atom][2:] + derivative_batches = _hermitian_density_derivative_batches( + ao, density_stack, p0, p1, "GGA", + ) + for xyz, derivatives in enumerate(derivative_batches): + drho_pz, drho_open, drho_alpha, drho_beta = derivatives + direct[k, xyz] += 0.5 * lib.einsum( + "xg,xyg,yg,g->", drho_pz, fref, rho_open, weights, + ) + direct[k, xyz] += 0.5 * lib.einsum( + "xg,xyg,yg,g->", rho_pz, fref, drho_open, weights, + ) + direct[k, xyz] += lib.einsum( + "xyg,xyzg,zg,g->", + pair, kref_alpha, drho_alpha, weights, + ) + direct[k, xyz] += lib.einsum( + "xyg,xyzg,zg,g->", + pair, kref_beta, drho_beta, weights, + ) + + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + q_alpha[:, spaces.open] += ( + mo.conj().T @ (open_potential + open_potential.T) @ spaces.c_open + ) + _add_reference_q( + tdobj, q_alpha, q_beta, reference_alpha, reference_beta, + ) + return XCGradientTerms(q_alpha, q_beta, direct) + + +def gga_nobeta_reference_q(tdobj, p0, max_memory=None): + """Reference-density response of the GGA equal-spin common Fock.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + if not tdobj.nobeta or getattr(mf, "is_ensemble_rks", False): + return q_alpha, q_beta + if max_memory is None: + max_memory = tdobj.max_memory + ni = mf._numint + mol = mf.mol + density_alpha, density_beta = _reference_spin_densities(tdobj) + density0 = 0.5 * (density_alpha + density_beta) + p0 = 0.5 * (np.asarray(p0) + np.asarray(p0).T) + matrix_alpha = np.zeros((mol.nao_nr(), mol.nao_nr())) + matrix_beta = np.zeros_like(matrix_alpha) + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, mol.nao_nr(), 2, max_memory=max_memory): + rho_p = ni.eval_rho( + mol, ao, p0, mask, "GGA", hermi=1, with_lapl=False, + ) + rho_alpha = ni.eval_rho( + mol, ao, density_alpha, mask, "GGA", hermi=1, + with_lapl=False, + ) + rho_beta = ni.eval_rho( + mol, ao, density_beta, mask, "GGA", hermi=1, + with_lapl=False, + ) + rho_equal = ni.eval_rho( + mol, ao, density0, mask, "GGA", hermi=1, with_lapl=False, + ) + fxc_actual = ni.eval_xc_eff( + mf.xc, (rho_alpha, rho_beta), deriv=2, + xctype="GGA", spin=1, + )[2] + fxc_equal = ni.eval_xc_eff( + mf.xc, (rho_equal, rho_equal), deriv=2, + xctype="GGA", spin=1, + )[2] + equal = 0.25 * ( + fxc_equal[0, :, 0] + fxc_equal[0, :, 1] + + fxc_equal[1, :, 0] + fxc_equal[1, :, 1] + ) + actual_alpha = 0.5 * ( + fxc_actual[0, :, 0] + fxc_actual[1, :, 0] + ) + actual_beta = 0.5 * ( + fxc_actual[0, :, 1] + fxc_actual[1, :, 1] + ) + matrix_alpha += gga_eval_matrix( + mol, + ao, + lib.einsum("xg,xzg->zg", rho_p, equal - actual_alpha) * weights, + mask, + ) + matrix_beta += gga_eval_matrix( + mol, + ao, + lib.einsum("xg,xzg->zg", rho_p, equal - actual_beta) * weights, + mask, + ) + _add_reference_q(tdobj, q_alpha, q_beta, matrix_alpha, matrix_beta) + return q_alpha, q_beta + + +# meta-GGA quadrature + +def _mgga_fref_kref(mf, rho0): + fxc, kxc = mf._numint.eval_xc_eff( + mf.xc, (rho0, rho0), deriv=3, xctype="MGGA", spin=1, + )[2:4] + fref = 0.5 * ( + fxc[0, :, 0] - fxc[0, :, 1] + - fxc[1, :, 0] + fxc[1, :, 1] + ) + kref_alpha = 0.5 * ( + kxc[0, :, 0, :, 0] - kxc[0, :, 1, :, 0] + - kxc[1, :, 0, :, 0] + kxc[1, :, 1, :, 0] + ) + kref_beta = 0.5 * ( + kxc[0, :, 0, :, 1] - kxc[0, :, 1, :, 1] + - kxc[1, :, 0, :, 1] + kxc[1, :, 1, :, 1] + ) + return fref, kref_alpha, kref_beta + + +def mgga_response_terms( + gradient_driver, tdobj, channel_data, atmlst=None, + with_direct=True): + """MGGA ``vref0/vref1`` M matrix and fixed-grid skeleton derivative.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + _spaces, _amplitudes, densities, blocks, terms = channel_data + pair_labels = tuple( + label for label in densities + if any( + term.vref1 and label in (term.target, term.source) + for term in terms + ) + ) + pair_density_stack = np.asarray([ + densities[label] for label in pair_labels + ]) + nao = mol.nao_nr() + potentials = {label: np.zeros((nao, nao)) for label in densities} + reference_alpha = np.zeros((nao, nao)) + reference_beta = np.zeros_like(reference_alpha) + direct = np.zeros((len(atmlst), 3)) + mo = np.asarray(mf.mo_coeff) + density_alpha, density_beta = _reference_spin_densities(tdobj) + density_labels, density_stack = _response_density_stack( + densities, density_alpha, density_beta, + ) + offsets = mol.offset_nr_by_atom() + sparse = sparse_context(mf) + + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): + rho0 = ni.eval_rho2( + mol, ao, mo, mf.mo_occ, mask, "MGGA", with_lapl=False, + ) * 0.5 + fref, kref_alpha, kref_beta = _mgga_fref_kref(mf, rho0) + rho = { + label: ni.eval_rho( + mol, ao, density, mask, "MGGA", hermi=0, + with_lapl=False, + ) + for label, density in densities.items() + } + pair_values, contracted_pair_ao = pair_feature_batches( + ao, pair_density_stack, + ) + pairs = dict(zip(pair_labels, pair_values)) + pair_potentials = { + label: mgga_pair_potential(fref, pairs[label]) + for label in pair_labels + } + ordinary_weights = { + label: np.zeros((5, weights.size)) for label in densities + } + special_weights = { + label: np.zeros((4, 4, weights.size)) for label in pair_labels + } + reference_weights_alpha = np.zeros((5, weights.size)) + reference_weights_beta = np.zeros_like(reference_weights_alpha) + + for term in terms: + if term.vref0: + ordinary_weights[term.target] += term.vref0 * lib.einsum( + "xyg,yg->xg", fref, rho[term.source], + ) + ordinary_weights[term.source] += term.vref0 * lib.einsum( + "xyg,xg->yg", fref, rho[term.target], + ) + pair = term.vref0 * lib.einsum( + "xg,yg->xyg", rho[term.target], rho[term.source], + ) + reference_weights_alpha += lib.einsum( + "xyg,xyzg->zg", pair, kref_alpha, + ) + reference_weights_beta += lib.einsum( + "xyg,xyzg->zg", pair, kref_beta, + ) + if term.vref1: + special_weights[term.target] += ( + term.vref1 * pair_potentials[term.source] + ) + special_weights[term.source] += ( + term.vref1 * pair_potentials[term.target] + ) + pair = term.vref1 * mgga_pair_kernel_cross( + pairs[term.target], pairs[term.source], + ) + reference_weights_alpha += lib.einsum( + "xyg,xyzg->zg", pair, kref_alpha, + ) + reference_weights_beta += lib.einsum( + "xyg,xyzg->zg", pair, kref_beta, + ) + ordinary_weight_stack = np.asarray([ + ordinary_weights[label] for label in density_labels + ]) + special_weight_stack = np.asarray([ + special_weights[label] for label in pair_labels + ]) + + for label in potentials: + add_mgga_matrix( + mol, potentials[label], ao, + ordinary_weights[label] * weights, mask, sparse, + ) + for label in pair_labels: + potentials[label] += pair_matrix( + mol, ao, mask, special_weights[label] * weights, sparse, + ) + reference_alpha += mgga_eval_matrix( + mol, ao, reference_weights_alpha * weights, mask, + ) + reference_beta += mgga_eval_matrix( + mol, ao, reference_weights_beta * weights, mask, + ) + + if not with_direct: + continue + for k, atom in enumerate(atmlst): + p0, p1 = offsets[atom][2:] + for xyz in range(3): + drho, drho_alpha, drho_beta, ao_delta = ( + _response_density_derivatives( + ao, density_stack, density_labels, + p0, p1, xyz, "MGGA", + ) + ) + drho_stack = np.asarray([ + drho[label] for label in density_labels + ]) + value = lib.einsum( + "nfg,nfg,g->", + ordinary_weight_stack, drho_stack, weights, + ) + value += lib.einsum( + "fg,fg,g->", + reference_weights_alpha, drho_alpha, weights, + ) + value += lib.einsum( + "fg,fg,g->", + reference_weights_beta, drho_beta, weights, + ) + value += contract_pair_feature_derivatives( + ao, pair_density_stack, ao_delta, + contracted_pair_ao, p0, p1, + special_weight_stack, weights, + ) + direct[k, xyz] += value + + q_alpha, q_beta = _project_channel_potentials( + tdobj, potentials, blocks, + ) + _add_reference_q( + tdobj, q_alpha, q_beta, reference_alpha, reference_beta, + ) + return XCGradientTerms(q_alpha, q_beta, direct) + + +def mgga_fockz_terms( + gradient_driver, tdobj, spaces, pz, atmlst=None, + with_direct=True): + """MGGA derivative of ``Pz:Fz`` excluding its explicit Pz projection.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + density_open = spaces.c_open @ spaces.c_open.T + pz = 0.5 * (np.asarray(pz) + np.asarray(pz).T) + nao = mol.nao_nr() + open_potential = np.zeros((nao, nao)) + reference_alpha = np.zeros((nao, nao)) + reference_beta = np.zeros_like(reference_alpha) + direct = np.zeros((len(atmlst), 3)) + mo = np.asarray(mf.mo_coeff) + density_alpha, density_beta = _reference_spin_densities(tdobj) + density_stack = np.asarray(( + pz, density_open, density_alpha, density_beta, + )) + offsets = mol.offset_nr_by_atom() + sparse = sparse_context(mf) + + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): + rho0 = ni.eval_rho2( + mol, ao, mo, mf.mo_occ, mask, "MGGA", with_lapl=False, + ) * 0.5 + fref, kref_alpha, kref_beta = _mgga_fref_kref(mf, rho0) + rho_pz = ni.eval_rho( + mol, ao, pz, mask, "MGGA", hermi=1, with_lapl=False, + ) + rho_open = ni.eval_rho( + mol, ao, density_open, mask, "MGGA", hermi=1, + with_lapl=False, + ) + add_mgga_matrix( + mol, + open_potential, + ao, + 0.5 * lib.einsum("xyg,yg->xg", fref, rho_pz) * weights, + mask, + sparse, + ) + pair = 0.5 * lib.einsum("xg,yg->xyg", rho_pz, rho_open) + reference_alpha += mgga_eval_matrix( + mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_alpha) * weights, + mask, + ) + reference_beta += mgga_eval_matrix( + mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_beta) * weights, + mask, + ) + if not with_direct: + continue + for k, atom in enumerate(atmlst): + p0, p1 = offsets[atom][2:] + derivative_batches = _hermitian_density_derivative_batches( + ao, density_stack, p0, p1, "MGGA", + ) + for xyz, derivatives in enumerate(derivative_batches): + drho_pz, drho_open, drho_alpha, drho_beta = derivatives + direct[k, xyz] += 0.5 * lib.einsum( + "xg,xyg,yg,g->", drho_pz, fref, rho_open, weights, + ) + direct[k, xyz] += 0.5 * lib.einsum( + "xg,xyg,yg,g->", rho_pz, fref, drho_open, weights, + ) + direct[k, xyz] += lib.einsum( + "xyg,xyzg,zg,g->", + pair, kref_alpha, drho_alpha, weights, + ) + direct[k, xyz] += lib.einsum( + "xyg,xyzg,zg,g->", + pair, kref_beta, drho_beta, weights, + ) + + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + q_alpha[:, spaces.open] += ( + mo.conj().T @ (open_potential + open_potential.T) @ spaces.c_open + ) + _add_reference_q( + tdobj, q_alpha, q_beta, reference_alpha, reference_beta, + ) + return XCGradientTerms(q_alpha, q_beta, direct) + + +def mgga_nobeta_reference_q(tdobj, p0, max_memory=None): + """Reference-density response of the MGGA equal-spin common Fock.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + if not tdobj.nobeta or getattr(mf, "is_ensemble_rks", False): + return q_alpha, q_beta + if max_memory is None: + max_memory = tdobj.max_memory + ni = mf._numint + mol = mf.mol + density_alpha, density_beta = _reference_spin_densities(tdobj) + density0 = 0.5 * (density_alpha + density_beta) + p0 = 0.5 * (np.asarray(p0) + np.asarray(p0).T) + matrix_alpha = np.zeros((mol.nao_nr(), mol.nao_nr())) + matrix_beta = np.zeros_like(matrix_alpha) + for ao, mask, weights, _coords in ni.block_loop( + mol, mf.grids, mol.nao_nr(), 2, max_memory=max_memory): + rho_p = ni.eval_rho( + mol, ao, p0, mask, "MGGA", hermi=1, with_lapl=False, + ) + rho_alpha = ni.eval_rho( + mol, ao, density_alpha, mask, "MGGA", hermi=1, + with_lapl=False, + ) + rho_beta = ni.eval_rho( + mol, ao, density_beta, mask, "MGGA", hermi=1, + with_lapl=False, + ) + rho_equal = ni.eval_rho( + mol, ao, density0, mask, "MGGA", hermi=1, with_lapl=False, + ) + fxc_actual = ni.eval_xc_eff( + mf.xc, (rho_alpha, rho_beta), deriv=2, + xctype="MGGA", spin=1, + )[2] + fxc_equal = ni.eval_xc_eff( + mf.xc, (rho_equal, rho_equal), deriv=2, + xctype="MGGA", spin=1, + )[2] + equal = 0.25 * ( + fxc_equal[0, :, 0] + fxc_equal[0, :, 1] + + fxc_equal[1, :, 0] + fxc_equal[1, :, 1] + ) + actual_alpha = 0.5 * ( + fxc_actual[0, :, 0] + fxc_actual[1, :, 0] + ) + actual_beta = 0.5 * ( + fxc_actual[0, :, 1] + fxc_actual[1, :, 1] + ) + matrix_alpha += mgga_eval_matrix( + mol, + ao, + lib.einsum("xg,xzg->zg", rho_p, equal - actual_alpha) * weights, + mask, + ) + matrix_beta += mgga_eval_matrix( + mol, + ao, + lib.einsum("xg,xzg->zg", rho_p, equal - actual_beta) * weights, + mask, + ) + _add_reference_q(tdobj, q_alpha, q_beta, matrix_alpha, matrix_beta) + return q_alpha, q_beta diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index 994fcac..9568e6d 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -# Copyright 2026 The NEST Developers. All Rights Reserved. +# Copyright 2014-2024 The PySCF Developers. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,11 +27,175 @@ from pyscf.dft.gen_grid import NBINS from pyscf import __config__ from pyscf.dft.numint import _scale_ao_sparse, _dot_ao_ao_sparse, _dot_ao_dm_sparse, _contract_rho_sparse -from nest._lr_eig import eigh as lr_eigh +from pyscf.tdscf._lr_eig import eigh as lr_eigh from pyscf import symm from pyscf.data import nist +from nest.ensemble_rks import EnsembleRKS MO_BASE = getattr(__config__, 'MO_BASE', 1) +MO_GRID_FXC1 = True + + +def _is_ensemble_reference(mf): + return isinstance(mf, EnsembleRKS) + + +def _require_nttda_reference(mf): + supported = ( + dft.roks.ROKS, + dft.rks_symm.SymAdaptedROKS, + EnsembleRKS, + ) + if not isinstance(mf, supported): + raise TypeError("NTTDA response requires ROKS or EnsembleRKS reference") + + +def _reference_fock0(mf, nobeta): + """Return the common Fock used by the NTTDA orbital terms. + + EnsembleRKS is self-consistent in ``F0[D/2,D/2]``. For ROKS, retain the + two historical choices controlled by ``nobeta``. + """ + if _is_ensemble_reference(mf): + return np.asarray(mf.get_fock()) + if nobeta: + dma, dmb = mf.make_rdm1() + dm0 = 0.5 * (dma + dmb) + fock = mf.get_fock(dm=np.array([dm0, dm0])) + else: + fock = mf.get_fock() + return 0.5 * (fock.focka + fock.fockb) + +def _fxc1_gga_mo_wv(fxc, t, i): + nvec = t.shape[0] + ngrids = t.shape[-1] + wv = np.empty((nvec, 4, ngrids)) + t00 = t[:, 0, 0] + if i == 0: + wv[:, 0] = lib.einsum('ijg,xijg->xg', fxc[:4, :4], t) + wv[:, 1:4] = fxc[0, 1:4][None] * t00[:, None] + wv[:, 1:4] += lib.einsum('ijg,xig->xjg', fxc[1:4, 1:4], t[:, 1:4, 0]) + else: + wv[:, 0] = fxc[i, 0][None] * t00 + wv[:, 0] += lib.einsum('jg,xjg->xg', fxc[i, 1:4], t[:, 0, 1:4]) + wv[:, 1:4] = fxc[i, 1:4][None] * t00[:, None] + return wv + +def _fxc1_mgga_mo_wv(fxc, t, i): + nvec = t.shape[0] + ngrids = t.shape[-1] + wv = np.empty((nvec, 4, ngrids)) + t00 = t[:, 0, 0] + if i == 0: + wv[:, 0] = lib.einsum('ijg,xijg->xg', fxc[:4, :4], t) + wv[:, 1:4] = fxc[0, 1:4][None] * t00[:, None] + wv[:, 1:4] += lib.einsum('ijg,xig->xjg', fxc[1:4, 1:4], t[:, 1:4, 0]) + wv[:, 1:4] += 0.5 * fxc[0, 4][None, None] * t[:, 0, 1:4] + wv[:, 1:4] += 0.5 * lib.einsum('ig,xijg->xjg', fxc[1:4, 4], t[:, 1:4, 1:4]) + else: + wv[:, 0] = fxc[i, 0][None] * t00 + wv[:, 0] += lib.einsum('jg,xjg->xg', fxc[i, 1:4], t[:, 0, 1:4]) + wv[:, 0] += 0.5 * fxc[4, 0][None] * t[:, i, 0] + wv[:, 0] += 0.5 * lib.einsum('jg,xjg->xg', fxc[4, 1:4], t[:, i, 1:4]) + wv[:, 1:4] = fxc[i, 1:4][None] * t00[:, None] + wv[:, 1:4] += 0.5 * fxc[i, 4][None, None] * t[:, 0, 1:4] + wv[:, 1:4] += 0.5 * fxc[4, 1:4][None] * t[:, i, 0][:, None] + wv[:, 1:4] += 0.25 * fxc[4, 4][None, None] * t[:, i, 1:4] + return wv + +def _fxc1_mo_make_t(x, left_mo, right_mo): + '''Build T[x,i,j,g] = L[j,g,a] X[x,a,b] R[i,g,b]. + + ``left_mo`` and ``right_mo`` contain AO values and first derivatives + projected to the two MO spaces. The transpose of X lets the inner dot use + BLAS over the right-index dimension. + ''' + nvec = x.shape[0] + ngrids = left_mo.shape[1] + t = np.empty((nvec, 4, 4, ngrids)) + for num in range(nvec): + xt = np.asarray(x[num].T, order='C') + for i in range(4): + tmp = lib.dot(right_mo[i], xt) + for j in range(4): + t[num, i, j] = lib.einsum('go,go->g', tmp, left_mo[j]) + return t + +def _fxc1_mo_accumulate(out, left_mo, right_mo, wv, coef): + '''Accumulate coef * L[j].T @ diag(wv[x,i,j]) @ R[i] to a MO block.''' + if coef == 0: + return + nvec = out.shape[0] + for num in range(nvec): + for i in range(4): + ri = right_mo[i] + for j in range(4): + weighted_left = left_mo[j] * wv[num, i, j, :, None] + out[num] += coef * lib.dot(weighted_left.T, ri) + +def _nr_rks_fxc1_mo(ni, mol, grids, mo_blocks, in_blocks, out_blocks, + terms, fxc, xctype, max_memory=2000): + '''Contract the fxc1 kernel directly in selected MO spaces. + + ``in_blocks`` maps an input name to (X, left_mo_key, right_mo_key). + ``out_blocks`` maps an output name to its projection MO spaces. + ``terms`` is the existing NTTDA linear combination as + (input_name, output_name, coefficient). The function returns MO-basis + contributions only; the AO vref0 and hybrid JK paths stay outside. + ''' + if xctype == 'GGA': + fill_wv = _fxc1_gga_mo_wv + elif xctype == 'MGGA': + fill_wv = _fxc1_mgga_mo_wv + else: + raise ValueError(f'MO-grid fxc1 only supports GGA/MGGA, got {xctype}') + + nao = mol.nao_nr() + nvec = next(iter(in_blocks.values()))[0].shape[0] + out = { + name: np.zeros((nvec, mo_blocks[left_key].shape[1], + mo_blocks[right_key].shape[1])) + for name, (left_key, right_key) in out_blocks.items() + } + needed_mos = set() + for x, left_key, right_key in in_blocks.values(): + needed_mos.add(left_key) + needed_mos.add(right_key) + for left_key, right_key in out_blocks.values(): + needed_mos.add(left_key) + needed_mos.add(right_key) + + terms_by_input = {} + for in_name, out_name, coef in terms: + terms_by_input.setdefault(in_name, []).append((out_name, coef)) + + p1 = 0 + for ao, mask, weight, coords in ni.block_loop(mol, grids, nao, 1, max_memory=max_memory): + p0, p1 = p1, p1 + weight.size + ngrids = weight.size + _fxc = fxc[:, :, p0:p1] * weight + + mo_cache = {} + for key in needed_mos: + coeff = mo_blocks[key] + mo = np.empty((4, ngrids, coeff.shape[1])) + for i in range(4): + mo[i] = lib.dot(ao[i], coeff) + mo_cache[key] = mo + + for in_name, (x, left_key, right_key) in in_blocks.items(): + input_terms = terms_by_input.get(in_name) + if not input_terms: + continue + t = _fxc1_mo_make_t(x, mo_cache[left_key], mo_cache[right_key]) + wv = np.empty((nvec, 4, 4, ngrids)) + for i in range(4): + wv[:, i] = fill_wv(_fxc, t, i) + for out_name, coef in input_terms: + out_left_key, out_right_key = out_blocks[out_name] + _fxc1_mo_accumulate(out[out_name], mo_cache[out_left_key], + mo_cache[out_right_key], wv, coef) + return out def nr_rks_fxc1_gga(ni, mol, grids, xc_code, dms, fxc, max_memory=2000): nset = dms.shape[0] @@ -162,8 +326,7 @@ def gen_rohf_response_sfu(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=No mol = mf.mol if log is None: log = logger.new_logger(mf) - if not isinstance(mf, (dft.roks.ROKS, dft.rks_symm.SymAdaptedROKS)): - raise TypeError('NTTDA response requires ROKS reference') + _require_nttda_reference(mf) ni = mf._numint ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True) @@ -186,7 +349,7 @@ def vind(dms_cv): time_xc = (logger.process_clock(), logger.perf_counter()) v1ao_cv = ni.nr_rks_fxc(mol, mf.grids, mf.xc, None, dms_cv, 0, hermi, None, None, fxc_ref, max_memory=max_memory) - time_xc = log.timer('NTTDA response_sfu kernel v1ao_cv', *time_xc) + time_xc = log.timer('NTTDA response_sfu kernel xc response_cv', *time_xc) else: v1ao_cv = np.zeros_like(dms_cv) @@ -211,7 +374,8 @@ def vind(dms_cv): delta -= mf.get_k(mol, dmoo, 1, omega=omega) * (alpha - hyb) return vind, 0.5 * delta -def gen_rohf_response_sc(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=None, log=None): +def gen_rohf_response_sc(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=None, + log=None, fxc_ref=None, skip_xc_vref1=False): ''' response function for Sf=Si ''' @@ -223,8 +387,7 @@ def gen_rohf_response_sc(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=Non mol = mf.mol if log is None: log = logger.new_logger(mf) - if not isinstance(mf, (dft.roks.ROKS, dft.rks_symm.SymAdaptedROKS)): - raise TypeError('NTTDA response requires ROKS reference') + _require_nttda_reference(mf) s = (mol.nelec[0] - mol.nelec[1]) * 0.5 @@ -233,7 +396,7 @@ def gen_rohf_response_sc(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=Non omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) hybrid = ni.libxc.is_hybrid_xc(mf.xc) xctype = ni._xc_type(mf.xc) - if xctype != 'HF': + if xctype != 'HF' and fxc_ref is None: fxc_d0 = ni.cache_xc_kernel(mol, mf.grids, mf.xc, mo_coeff, mo_occ, 1)[2] fxc_ref = 0.5 * (fxc_d0[0, :, 0] - fxc_d0[0, :, 1] - fxc_d0[1, :, 0] + fxc_d0[1, :, 1]) @@ -266,7 +429,9 @@ def vind(dms_co, dms_cv, dms_ov, dms_cv0): vref0 = ni.nr_rks_fxc(mol, mf.grids, mf.xc, None, dms0, 0, hermi, None, None, fxc_ref, max_memory=max_memory) time_xc = log.timer('NTTDA response_sc kernel vref0', *time_xc) - if xctype == 'LDA': + if skip_xc_vref1 and xctype in ('GGA', 'MGGA'): + vref1 = np.zeros_like(dms1) + elif xctype == 'LDA': vref1 = ni.nr_rks_fxc(mol, mf.grids, mf.xc, None, dms1, 0, hermi, None, None, fxc_ref, max_memory=max_memory) elif xctype =='GGA': @@ -319,7 +484,8 @@ def vind(dms_co, dms_cv, dms_ov, dms_cv0): delta -= mf.get_k(mol, dmoo, 1, omega=omega) * (alpha - hyb) return vind, 0.5 * delta -def gen_rohf_response_sfd(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=None, log=None): +def gen_rohf_response_sfd(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=None, + log=None, fxc_ref=None, skip_xc_vref1=False): ''' response function for Sf=Si-1 ''' @@ -331,8 +497,7 @@ def gen_rohf_response_sfd(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=No mol = mf.mol if log is None: log = logger.new_logger(mf) - if not isinstance(mf, (dft.roks.ROKS, dft.rks_symm.SymAdaptedROKS)): - raise TypeError('NTTDA response requires ROKS reference') + _require_nttda_reference(mf) s = (mol.nelec[0] - mol.nelec[1]) * 0.5 @@ -342,7 +507,7 @@ def gen_rohf_response_sfd(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=No hybrid = ni.libxc.is_hybrid_xc(mf.xc) xctype = ni._xc_type(mf.xc) - if xctype != 'HF': + if xctype != 'HF' and fxc_ref is None: fxc_d0 = ni.cache_xc_kernel(mol, mf.grids, mf.xc, mo_coeff, mo_occ, 1)[2] fxc_ref = 0.5 * (fxc_d0[0, :, 0] - fxc_d0[0, :, 1] - fxc_d0[1, :, 0] + fxc_d0[1, :, 1]) @@ -374,7 +539,9 @@ def vind(dms_co, dms_cv, dms_oo, dms_ov): vref0 = ni.nr_rks_fxc(mol, mf.grids, mf.xc, None, dms0, 0, hermi, None, None, fxc_ref, max_memory=max_memory) time_xc = log.timer('NTTDA response_sf vref0', *time_xc) - if xctype == 'LDA': + if skip_xc_vref1 and xctype in ('GGA', 'MGGA'): + vref1 = np.zeros_like(dms1) + elif xctype == 'LDA': vref1 = ni.nr_rks_fxc(mol, mf.grids, mf.xc, None, dms1, 0, hermi, None, None, fxc_ref, max_memory=max_memory) elif xctype =='GGA': @@ -444,18 +611,9 @@ def gen_vind_sfu(td): vresp, fockz = gen_rohf_response_sfu(mf, mo_coeff=mo_coeff, mo_occ=mo_occ, hermi=0, max_memory=td.max_memory, log=log) - if td.nobeta: - dma, dmb = mf.make_rdm1() - dm0 = 0.5 * (dma + dmb) - fock = mf.get_fock(dm=np.array([dm0, dm0])) - fock0 = 0.5 * (fock.focka + fock.fockb) - focka = fock0 + fockz - fockb = fock0 - fockz - else: - fock = mf.get_fock() - fock0 = 0.5 * (fock.focka + fock.fockb) - focka = fock0 + fockz - fockb = fock0 - fockz + fock0 = _reference_fock0(mf, td.nobeta) + focka = fock0 + fockz + fockb = fock0 - fockz fock_v = orbvs.T @ focka @ orbvs fock_c = orbcs.T @ fockb @ orbcs @@ -492,6 +650,7 @@ def gen_vind_sc(td): orbcs = mo_coeff[:, csidx] orbos = mo_coeff[:, osidx] orbvs = mo_coeff[:, vsidx] + mo_blocks = {'c': orbcs, 'o': orbos, 'v': orbvs} ncs = orbcs.shape[1] nos = orbos.shape[1] nvs = orbvs.shape[1] @@ -502,21 +661,22 @@ def gen_vind_sc(td): assert s == (mf.mol.nelec[0] - mf.mol.nelec[1]) * 0.5 log = logger.new_logger(td) + xctype = mf._numint._xc_type(mf.xc) + use_mo_grid_fxc1 = MO_GRID_FXC1 and xctype in ('GGA', 'MGGA') + fxc_ref = None + if use_mo_grid_fxc1: + fxc_d0 = mf._numint.cache_xc_kernel(mf.mol, mf.grids, mf.xc, + mo_coeff, mo_occ, 1)[2] + fxc_ref = 0.5 * (fxc_d0[0, :, 0] - fxc_d0[0, :, 1] - + fxc_d0[1, :, 0] + fxc_d0[1, :, 1]) vresp, fockz = gen_rohf_response_sc(mf, mo_coeff=mo_coeff, mo_occ=mo_occ, hermi=0, - max_memory=td.max_memory, log=log) + max_memory=td.max_memory, log=log, + fxc_ref=fxc_ref, + skip_xc_vref1=use_mo_grid_fxc1) - if td.nobeta: - dma, dmb = mf.make_rdm1() - dm0 = 0.5 * (dma + dmb) - fock = mf.get_fock(dm=np.array([dm0, dm0])) - fock0 = 0.5 * (fock.focka + fock.fockb) - focka = fock0 + fockz - fockb = fock0 - fockz - else: - fock = mf.get_fock() - fock0 = 0.5 * (fock.focka + fock.fockb) - focka = fock0 + fockz - fockb = fock0 - fockz + fock0 = _reference_fock0(mf, td.nobeta) + focka = fock0 + fockz + fockb = fock0 - fockz fock_coco1 = orbos.T @ (fock0 - fockz) @ orbos fock_coco2 = orbcs.T @ (fock0 - fockz) @ orbcs @@ -566,6 +726,37 @@ def vind(zs): v1mo_cv0 = lib.einsum('xpq,qo,pv->xov', v1ao_cv0, orbcs, orbvs.conj()) time1 = log.timer('NTTDA gen_vind_sc AO->MO transform', *time1) + if use_mo_grid_fxc1: + time_mo = (logger.process_clock(), logger.perf_counter()) + in_blocks = { + 'co': (zs_co, 'c', 'o'), + 'ov': (zs_ov, 'o', 'v'), + 'cv0': (zs_cv0, 'c', 'v'), + } + out_blocks = { + 'co': ('c', 'o'), + 'ov': ('o', 'v'), + 'cv0': ('c', 'v'), + } + terms = ( + ('co', 'co', -1.0), + ('ov', 'co', 1.0), + ('cv0', 'co', -np.sqrt(2.0)), + ('co', 'ov', 1.0), + ('ov', 'ov', -1.0), + ('cv0', 'ov', np.sqrt(2.0)), + ('co', 'cv0', -np.sqrt(2.0)), + ('ov', 'cv0', np.sqrt(2.0)), + ('cv0', 'cv0', -2.0), + ) + vref1_mo = _nr_rks_fxc1_mo( + mf._numint, mf.mol, mf.grids, mo_blocks, in_blocks, + out_blocks, terms, fxc_ref, xctype, max_memory=td.max_memory) + v1mo_co += vref1_mo['co'] + v1mo_ov += vref1_mo['ov'] + v1mo_cv0 += vref1_mo['cv0'] + time1 = log.timer('NTTDA gen_vind_sc MO-grid vref1', *time_mo) + v1mo_co += lib.einsum('uv,xiv->xiu', fock_coco1, zs_co) v1mo_co -= lib.einsum('ji,xju->xiu', fock_coco2, zs_co) v1mo_co += lib.einsum('ub,xib->xiu', fock_cocv, zs_cv) * np.sqrt((s + 1) / 2 / s) @@ -622,6 +813,7 @@ def gen_vind_sfd(td): orbcs = mo_coeff[:, csidx] orbos = mo_coeff[:, osidx] orbvs = mo_coeff[:, vsidx] + mo_blocks = {'c': orbcs, 'o': orbos, 'v': orbvs} ncs = orbcs.shape[1] nos = orbos.shape[1] nvs = orbvs.shape[1] @@ -637,17 +829,20 @@ def gen_vind_sfd(td): assert s == (mf.mol.nelec[0] - mf.mol.nelec[1]) * 0.5 log = logger.new_logger(td) + xctype = mf._numint._xc_type(mf.xc) + use_mo_grid_fxc1 = MO_GRID_FXC1 and xctype in ('GGA', 'MGGA') + fxc_ref = None + if use_mo_grid_fxc1: + fxc_d0 = mf._numint.cache_xc_kernel(mf.mol, mf.grids, mf.xc, + mo_coeff, mo_occ, 1)[2] + fxc_ref = 0.5 * (fxc_d0[0, :, 0] - fxc_d0[0, :, 1] - + fxc_d0[1, :, 0] + fxc_d0[1, :, 1]) vresp, fockz = gen_rohf_response_sfd(mf, mo_coeff=mo_coeff, mo_occ=mo_occ, hermi=0, - max_memory=td.max_memory, log=log) + max_memory=td.max_memory, log=log, + fxc_ref=fxc_ref, + skip_xc_vref1=use_mo_grid_fxc1) - if td.nobeta: - dma, dmb = mf.make_rdm1() - dm0 = 0.5 * (dma + dmb) - fock = mf.get_fock(dm=np.array([dm0, dm0])) - fock0 = 0.5 * (fock.focka + fock.fockb) - else: - fock = mf.get_fock() - fock0 = 0.5 * (fock.focka + fock.fockb) + fock0 = _reference_fock0(mf, td.nobeta) fock_coco0 = orbos.T @ (fock0 - fockz) @ orbos fock_coco1 = orbcs.T @ (fock0 + fockz) @ orbcs @@ -700,6 +895,30 @@ def vind(zs): v1mo_ov = lib.einsum('xpq,qo,pv->xov', v1ao_ov, orbos, orbvs.conj()) time1 = log.timer('NTTDA gen_vind_sfd AO->MO transform', *time1) + if use_mo_grid_fxc1: + time_mo = (logger.process_clock(), logger.perf_counter()) + denom = 2 * s - 1 + in_blocks = { + 'co': (zs_co, 'c', 'o'), + 'ov': (zs_ov, 'o', 'v'), + } + out_blocks = { + 'co': ('c', 'o'), + 'ov': ('o', 'v'), + } + terms = ( + ('co', 'co', 1.0 / denom), + ('ov', 'co', -1.0 / denom), + ('co', 'ov', -1.0 / denom), + ('ov', 'ov', 1.0 / denom), + ) + vref1_mo = _nr_rks_fxc1_mo( + mf._numint, mf.mol, mf.grids, mo_blocks, in_blocks, + out_blocks, terms, fxc_ref, xctype, max_memory=td.max_memory) + v1mo_co += vref1_mo['co'] + v1mo_ov += vref1_mo['ov'] + time1 = log.timer('NTTDA gen_vind_sfd MO-grid vref1', *time_mo) + v1mo_co += lib.einsum('uv,xiv->xiu', fock_coco0, zs_co) v1mo_co -= lib.einsum('ji,xju->xiu', fock_coco1, zs_co) v1mo_co -= lib.einsum('ji,xju->xiu', fock_coco2, zs_co) * 2 / (2 * s - 1) @@ -759,11 +978,31 @@ class NTTDA(TDBase): nobeta: True for problemstic cases where there is no local beta electrons ''' - deltaS = -1 - nobeta = False + deltaS = getattr(__config__, 'NTTDA_delta_S', -1) + nobeta = getattr(__config__, 'NTTDA_nobeta', False) _keys = {'deltaS', 'nobeta'} + def reference_energy(self): + """Return the reference zero selected by the mean-field object.""" + selector = getattr(self._scf, 'reference_energy', None) + if selector is None: + return float(self._scf.e_tot) + return float(selector()) + + def total_energies(self): + """Return ``E_reference + omega`` for the converged NTTDA roots.""" + if self.e is None: + raise RuntimeError('run NTTDA.kernel() before requesting total energies') + return self.reference_energy() + np.asarray(self.e) + + def nuc_grad_method(self): + """Return the independent NTTDA nuclear-gradient driver.""" + from nest.grad.nttda import Gradients + return Gradients(self) + + Gradients = nuc_grad_method + def init_guess(self, hdiag, nstates=None): if nstates is None: nstates = self.nstates @@ -1073,5 +1312,158 @@ def analyze(tdobj, verbose=None): NTTDA.analyze = analyze + +def transition_dipole(tdobj, ref=1, state=None): + """ + Transition dipole moments between NTTDA DeltaS = -1 excited states. + + Parameters + ---------- + tdobj : NTTDA object + Need tdobj.deltaS == -1 and tdobj.xy. + ref : int + 1-based reference excited-state index. + state : int or array-like or None + 1-based target excited-state index/indices. If None, all states except ref. + + Returns + ------- + pol : ndarray, shape (nstates, 3) + in length gauge. + """ + mf = tdobj._scf + mol = mf.mol + + deltaS = getattr(tdobj, "deltaS", getattr(tdobj, "DeltaS", None)) + assert deltaS == -1 + + s = (mol.nelec[0] - mol.nelec[1]) * 0.5 + assert s >= 1 + + mo_coeff = mf.mo_coeff + assert mf.mo_occ.ndim == 1 + csidx, osidx, vsidx = _orbital_indices(tdobj) + + ncs = len(csidx) + nos = len(osidx) + nvs = len(vsidx) + nocc = ncs + nos + nvir = nos + nvs + nmo = mo_coeff.shape[1] + + if state is None: + states = np.arange(tdobj.nstates) + 1 + else: + states = np.atleast_1d(state).astype(int) + + states = states[states != ref] + ref0 = ref - 1 + states0 = states - 1 + + def get_x(i): + x = tdobj.xy[i] + if isinstance(x, (tuple, list)): + x = x[0] + return np.asarray(x).reshape(nocc, nvir) + + mx = get_x(ref0) + nxs = np.asarray([get_x(i) for i in states0]) + + # x_co: j v, x_cv: j b, x_oo: w v, x_ov: v b + m_co = mx[:ncs, :nos].conj() + m_cv = mx[:ncs, nos:].conj() + m_oo = mx[ncs:, :nos].conj() + m_ov = mx[ncs:, nos:].conj() + + n_co = nxs[:, :ncs, :nos] + n_cv = nxs[:, :ncs, nos:] + n_oo = nxs[:, ncs:, :nos] + n_ov = nxs[:, ncs:, nos:] + + nstate = len(states0) + gamma = np.zeros((nstate, nmo, nmo), dtype=np.result_type(mx, nxs, complex)) + ist = np.arange(nstate) + + def add(rows, cols, block): + gamma[np.ix_(ist, rows, cols)] += block + + a = np.sqrt(2 * s / (2 * s - 1)) + b = 1 / np.sqrt(2 * s * (2 * s - 1)) + # c = np.sqrt((2 * s - 1) / (2 * s)) + f = np.sqrt((2 * s + 1) / (2 * s)) + + tr_moo = np.einsum("tt->", m_oo) + tr_noo = np.einsum("ntt->n", n_oo) + + # OO-OO + add(osidx, osidx, lib.einsum("ut,nuv->ntv", m_oo, n_oo)) + add(osidx, osidx, -lib.einsum("ut,nwt->nwu", m_oo, n_oo)) + + # CO-CO + add(osidx, osidx, lib.einsum("iu,niv->nuv", m_co, n_co)) + add(csidx, csidx, -lib.einsum("iu,nju->nji", m_co, n_co)) + + # CV-CV + add(csidx, csidx, -lib.einsum("ia,nja->nji", m_cv, n_cv)) + add(vsidx, vsidx, lib.einsum("ia,nib->nab", m_cv, n_cv)) + + # OV-OV + add(osidx, osidx, -lib.einsum("ua,nva->nvu", m_ov, n_ov)) + add(vsidx, vsidx, lib.einsum("ua,nub->nab", m_ov, n_ov)) + + # OO-CO and CO-OO + add(csidx, osidx, + -a * lib.einsum("ut,njt->nju", m_oo, n_co) + + b * tr_moo * n_co) + + add(osidx, csidx, + -a * lib.einsum("iu,nwu->nwi", m_co, n_oo) + + b * lib.einsum("iu,n->nui", m_co, tr_noo)) + + # OO-OV and OV-OO + add(osidx, vsidx, + a * lib.einsum("ut,nub->ntb", m_oo, n_ov) + - b * tr_moo * n_ov) + + add(vsidx, osidx, + a * lib.einsum("ua,nuv->nav", m_ov, n_oo) + - b * lib.einsum("ua,n->nau", m_ov, tr_noo)) + + # CO-CV and CV-CO + add(osidx, vsidx, f * lib.einsum("iu,nib->nub", m_co, n_cv)) + add(vsidx, osidx, f * lib.einsum("ia,niv->nav", m_cv, n_co)) + + # OV-CV and CV-OV + add(csidx, osidx, -f * lib.einsum("ua,nja->nju", m_ov, n_cv)) + add(osidx, csidx, -f * lib.einsum("ia,nva->nvi", m_cv, n_ov)) + + dip_ao = mol.intor_symmetric("int1e_r", comp=3) + dip_mo = lib.einsum("up,xuv,vq->xpq", mo_coeff.conj(), dip_ao, mo_coeff) + + pol = lib.einsum("npq,xpq->nx", gamma, dip_mo) + return pol.real + +def oscillator_strength(tdobj, ref=1, state=None): + if state is None: + states = np.arange(tdobj.nstates) + 1 + else: + states = np.atleast_1d(state) + states = states[states != ref] + + trans_dip = transition_dipole(tdobj, ref, states) + + ref -= 1 + states -= 1 + es = tdobj.e[states] - tdobj.e[ref] + f = (2./3.) * lib.einsum('n,nx,nx->n', es, trans_dip.conj(), trans_dip).real + if isinstance(state, int): + return f[0] + else: + return f + +NTTDA.transition_dipole = transition_dipole +NTTDA.oscillator_strength = oscillator_strength + dft.roks.ROKS.NTTDA = lib.class_as_method(NTTDA) dft.rks_symm.SymAdaptedROKS.NTTDA = lib.class_as_method(NTTDA) +EnsembleRKS.NTTDA = lib.class_as_method(NTTDA) From 7b8d8aae3399461b868a1decadfaa6d2173e0139 Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Mon, 31 Aug 2026 03:18:47 +0800 Subject: [PATCH 06/16] test: cover NTTDA gradient and ensemble reference channels Add the ensemble FD/response, scalar-ledger, gradient-layer, and known-value NTTDA energy suites. --- src/nest/grad/tests/test_nttda_ensemble_fd.py | 222 ++++++++++++++++ .../tests/test_nttda_ensemble_response.py | 206 +++++++++++++++ src/nest/grad/tests/test_nttda_grad.py | 103 ++++++++ .../grad/tests/test_nttda_gradient_layers.py | 238 ++++++++++++++++++ .../grad/tests/test_nttda_scalar_ledger.py | 106 ++++++++ src/nest/nttda/tests/test_nttda.py | 222 ++++++++++------ src/nest/nttda/tests/test_nttda_ensemble.py | 169 +++++++++++++ 7 files changed, 1194 insertions(+), 72 deletions(-) create mode 100644 src/nest/grad/tests/test_nttda_ensemble_fd.py create mode 100644 src/nest/grad/tests/test_nttda_ensemble_response.py create mode 100644 src/nest/grad/tests/test_nttda_grad.py create mode 100644 src/nest/grad/tests/test_nttda_gradient_layers.py create mode 100644 src/nest/grad/tests/test_nttda_scalar_ledger.py create mode 100644 src/nest/nttda/tests/test_nttda_ensemble.py diff --git a/src/nest/grad/tests/test_nttda_ensemble_fd.py b/src/nest/grad/tests/test_nttda_ensemble_fd.py new file mode 100644 index 0000000..7ce56ed --- /dev/null +++ b/src/nest/grad/tests/test_nttda_ensemble_fd.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python +"""Finite-difference NTTDA gradients on an EnsembleRKS reference.""" + +import unittest + +import numpy as np + +from pyscf import gto +from nest.ensemble_rks import EnsembleRKS +from nest.nttda import NTTDA + + +class EnsembleRKSFiniteDifferenceGradient(unittest.TestCase): + @staticmethod + def make_td(): + mol = gto.M( + atom="Li 0 0 0; H 0 0 3.0", + basis="sto-3g", + charge=1, + spin=1, + unit="Bohr", + verbose=0, + ) + mf = EnsembleRKS(mol).set( + xc="SVWN", + conv_tol=1e-12, + conv_tol_grad=1e-9, + max_cycle=100, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + if not mf.converged: + raise RuntimeError("EnsembleRKS test reference did not converge") + tdobj = NTTDA(mf).set( + deltaS=0, + nstates=2, + conv_tol=1e-8, + max_cycle=100, + verbose=0, + ).run() + return tdobj + + def test_finite_difference_total_energy_gradient(self): + tdobj = self.make_td() + gradient = tdobj.Gradients().set( + verbose=0, + fixed_grid=False, + root_overlap_tol=0.5, + ) + self.assertFalse(tdobj.Gradients().fixed_grid) + from nest.grad.nttda import _displaced_reference + + displaced = _displaced_reference(tdobj._scf, tdobj.mol.copy(), False) + self.assertIsInstance(displaced, EnsembleRKS) + self.assertEqual(displaced.nopen, tdobj._scf.nopen) + result = gradient.kernel( + state=2, + method="finite_diff", + step=2e-3, + ) + self.assertEqual(result.shape, (2, 3)) + self.assertTrue(np.all(np.isfinite(result))) + np.testing.assert_allclose( + result, + [[0, 0, 0.09285896198862886], + [0, 0, -0.09285896198774068]], + atol=2e-5, + rtol=0, + ) + np.testing.assert_allclose(result.sum(axis=0), 0, atol=2e-5, rtol=0) + np.testing.assert_allclose(result[:, :2], 0, atol=2e-5, rtol=0) + self.assertGreater(abs(result[0, 2]), 1e-3) + + gradient.fixed_grid = True + analytic = gradient.kernel(state=2, method="analytic") + fixed_grid_difference = gradient.kernel( + state=2, + method="finite_diff", + step=2e-3, + ) + np.testing.assert_allclose( + analytic, fixed_grid_difference, atol=2e-5, rtol=0, + ) + + def test_all_spin_channels_have_a_finite_difference_path(self): + mol = gto.M( + atom="C 0 0 0; H 0 0 2.0; H 0 1.7 -0.5", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + mf = EnsembleRKS(mol).set( + xc="SVWN", + conv_tol=1e-11, + max_cycle=150, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + + for delta_s in (-1, 0, 1): + with self.subTest(deltaS=delta_s): + tdobj = NTTDA(mf).set( + deltaS=delta_s, + nstates=2, + conv_tol=1e-6, + max_cycle=200, + verbose=0, + ).run() + result = tdobj.Gradients().set( + verbose=0, + root_overlap_tol=0.5, + ).kernel( + state=1, + atmlst=[0], + method="finite_diff", + step=2e-3, + ) + self.assertEqual(result.shape, (1, 3)) + self.assertTrue(np.all(np.isfinite(result))) + + def test_spin_lowering_analytic_matches_fixed_grid_finite_difference(self): + mol = gto.M( + atom="C 0 0 0; H 0 0 2.0; H 0 1.7 -0.5", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + mf = EnsembleRKS(mol).set( + xc="SVWN", + conv_tol=1e-12, + conv_tol_grad=1e-9, + max_cycle=150, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + tdobj = NTTDA(mf).set( + deltaS=-1, + nstates=2, + conv_tol=1e-9, + max_cycle=200, + verbose=0, + ).run() + gradient = tdobj.Gradients().set( + verbose=0, + fixed_grid=True, + root_overlap_tol=0.5, + ) + analytic = gradient.kernel(state=1, atmlst=[1], method="analytic") + finite_difference = gradient.kernel( + state=1, + atmlst=[1], + method="finite_diff", + step=2e-3, + ) + np.testing.assert_allclose( + analytic, finite_difference, atol=3e-5, rtol=0, + ) + + def test_representative_functional_families(self): + cases = ( + (0, "HF", False), + (0, "PBE", False), + (0, "M06-2X", False), + (0, "CAM-B3LYP", False), + (-1, "PBE", False), + (-1, "M06-2X", True), + ) + for delta_s, xc, nobeta in cases: + with self.subTest(deltaS=delta_s, xc=xc, nobeta=nobeta): + mol = gto.M( + atom="C 0 0 0; H 0 0 2.0; H 0 1.7 -0.5", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + mf = EnsembleRKS(mol).set( + xc=xc, + conv_tol=1e-12, + conv_tol_grad=1e-9, + max_cycle=150, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + tdobj = NTTDA(mf).set( + deltaS=delta_s, + nobeta=nobeta, + nstates=2, + conv_tol=1e-9, + max_cycle=200, + verbose=0, + ).run() + gradient = tdobj.Gradients().set( + verbose=0, + fixed_grid=True, + root_overlap_tol=0.5, + ) + analytic = gradient.kernel( + state=1, atmlst=[1], method="analytic", + ) + finite_difference = gradient.kernel( + state=1, + atmlst=[1], + method="finite_diff", + step=2e-3, + ) + np.testing.assert_allclose( + analytic, finite_difference, atol=5e-5, rtol=0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/nest/grad/tests/test_nttda_ensemble_response.py b/src/nest/grad/tests/test_nttda_ensemble_response.py new file mode 100644 index 0000000..6a82860 --- /dev/null +++ b/src/nest/grad/tests/test_nttda_ensemble_response.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python +"""Orbital-response checks for average-occupation EnsembleRKS gradients.""" + +import unittest +from pathlib import Path + +import numpy as np +from scipy.linalg import expm + + + + +from pyscf import gto + + +from nest.grad.nttda.ensemble import ( # noqa: E402 + make_hessian_transpose_action, + pack_m_matrix, + zvector_adjoint_matrix, + zvector_probe_densities, +) +from nest.grad.nttda.delta_s_zero import ( # noqa: E402 + grad_elec, + same_spin_ledger_scalar, +) +from nest.grad.nttda.delta_s_minus_one import ( # noqa: E402 + grad_elec as spin_lowering_grad_elec, + spin_lowering_ledger_scalar, +) +from nest.ensemble_rks import EnsembleRKS # noqa: E402 +from nest.nttda import NTTDA # noqa: E402 + + +class EnsembleOrbitalResponse(unittest.TestCase): + @staticmethod + def make_reference(): + mol = gto.M( + atom="Li 0 0 0; H 0 0 3.0", + basis="sto-3g", + charge=1, + spin=1, + unit="Bohr", + verbose=0, + ) + mf = EnsembleRKS(mol).set( + xc="SVWN", + conv_tol=1e-13, + conv_tol_grad=1e-10, + max_cycle=100, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + if not mf.converged: + raise RuntimeError("EnsembleRKS reference did not converge") + return mf + + @staticmethod + def make_spin_one_reference(): + mol = gto.M( + atom="C 0 0 0; H 0 0 2.0; H 0 1.7 -0.5", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + mf = EnsembleRKS(mol).set( + xc="SVWN", + conv_tol=1e-12, + conv_tol_grad=1e-9, + max_cycle=150, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + if not mf.converged: + raise RuntimeError("spin-one EnsembleRKS reference did not converge") + return mf + + def test_hessian_action_matches_orbital_rotation_finite_difference(self): + mf = self.make_reference() + tdobj = NTTDA(mf) + action, pairs = make_hessian_transpose_action(tdobj) + rng = np.random.default_rng(19) + vector = rng.normal(size=len(pairs)) + + mo = np.asarray(mf.mo_coeff) + occ = np.asarray(mf.mo_occ) + kappa = np.zeros((mo.shape[1], mo.shape[1])) + for value, (p, q, _name) in zip(vector, pairs): + kappa[p, q] = value + kappa[q, p] = -value + + step = 1e-5 + gradients = [] + for sign in (1.0, -1.0): + displaced_mo = mo @ expm(sign * step * kappa) + density = mf.make_rdm1(displaced_mo, occ) + fock = mf.get_hcore() + mf.get_veff(mf.mol, density) + gradients.append(mf.get_grad(displaced_mo, occ, fock)) + finite_difference = (gradients[0] - gradients[1]) / (2.0 * step) + + np.testing.assert_allclose( + action(vector), finite_difference, atol=1e-8, rtol=0, + ) + + def test_adjoint_and_probe_are_consistent_with_explicit_hessian(self): + mf = self.make_reference() + tdobj = NTTDA(mf) + action, pairs = make_hessian_transpose_action(tdobj) + identity = np.eye(len(pairs)) + hessian = np.asarray(action(identity)).T + rng = np.random.default_rng(23) + zvector = rng.normal(size=len(pairs)) + + adjoint = zvector_adjoint_matrix(tdobj, pairs, zvector) + np.testing.assert_allclose( + pack_m_matrix(adjoint, pairs), + hessian.T @ zvector, + atol=1e-10, + rtol=0, + ) + + probe_alpha, probe_beta = zvector_probe_densities( + tdobj, pairs, zvector, + ) + perturbation = rng.normal(size=(mf.mol.nao_nr(),) * 2) + perturbation = perturbation + perturbation.T + mo = np.asarray(mf.mo_coeff) + occ = np.asarray(mf.mo_occ) + fock_mo = mo.T @ perturbation @ mo + expected = sum( + value * (occ[q] - occ[p]) * fock_mo[p, q] + for value, (p, q, _name) in zip(zvector, pairs) + ) + actual = np.einsum( + "ij,ji", probe_alpha + probe_beta, perturbation, + ) + self.assertAlmostEqual(actual, expected, places=11) + + def test_same_spin_m_matrix_is_the_orbital_derivative(self): + mf = self.make_reference() + tdobj = NTTDA(mf).set( + deltaS=0, + nstates=2, + conv_tol=1e-10, + max_cycle=200, + verbose=0, + ).run() + xy = tdobj.xy[1] + result = grad_elec( + mf.nuc_grad_method(), tdobj, xy, atmlst=(), + ) + + rng = np.random.default_rng(29) + perturbation = rng.normal(size=(mf.mo_coeff.shape[1],) * 2) + original = np.array(mf.mo_coeff, copy=True) + step = 1e-6 + values = [] + try: + for sign in (1.0, -1.0): + mf.mo_coeff = original @ ( + np.eye(original.shape[1]) + + sign * step * perturbation + ) + values.append(same_spin_ledger_scalar(tdobj, xy)) + finally: + mf.mo_coeff = original + finite_difference = (values[0] - values[1]) / (2.0 * step) + analytic = np.trace(result.m_matrix.T @ perturbation) + self.assertAlmostEqual(analytic, finite_difference, places=8) + + def test_spin_lowering_m_matrix_is_the_orbital_derivative(self): + mf = self.make_spin_one_reference() + tdobj = NTTDA(mf).set( + deltaS=-1, + nstates=2, + conv_tol=1e-9, + max_cycle=200, + verbose=0, + ).run() + xy = tdobj.xy[0] + result = spin_lowering_grad_elec( + mf.nuc_grad_method(), tdobj, xy, atmlst=(), + ) + + rng = np.random.default_rng(31) + perturbation = rng.normal(size=(mf.mo_coeff.shape[1],) * 2) + original = np.array(mf.mo_coeff, copy=True) + step = 1e-6 + values = [] + try: + for sign in (1.0, -1.0): + mf.mo_coeff = original @ ( + np.eye(original.shape[1]) + sign * step * perturbation + ) + values.append(spin_lowering_ledger_scalar(tdobj, xy)) + finally: + mf.mo_coeff = original + finite_difference = (values[0] - values[1]) / (2.0 * step) + analytic = np.trace(result.m_matrix.T @ perturbation) + self.assertAlmostEqual(analytic, finite_difference, places=8) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/nest/grad/tests/test_nttda_grad.py b/src/nest/grad/tests/test_nttda_grad.py new file mode 100644 index 0000000..c529a86 --- /dev/null +++ b/src/nest/grad/tests/test_nttda_grad.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python +"""Public NTTDA analytic-gradient acceptance tests.""" + +import unittest +from pathlib import Path + +import numpy as np + + + + +from pyscf import dft, gto + + +from nest.nttda import NTTDA # noqa: E402 + + +class NTTDAGradientAcceptance(unittest.TestCase): + @staticmethod + def molecule(): + return gto.M( + atom="N 0 0 0; O 0 0 1.20; H 0 0.90 -0.20", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + + def make_td(self, xc, delta_s, nobeta=False): + mf = dft.ROKS(self.molecule()).set( + xc=xc, + conv_tol=1e-14, + conv_tol_grad=1e-11, + max_cycle=200, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + tdobj = NTTDA(mf).set( + deltaS=delta_s, + nobeta=nobeta, + nstates=3, + conv_tol=1e-9, + max_cycle=200, + verbose=0, + ).run() + self.assertGreaterEqual(len(tdobj.xy), 2) + return tdobj + + def compare_public_gradient(self, xc, delta_s, nobeta, threshold): + tdobj = self.make_td(xc, delta_s, nobeta=nobeta) + gradient = tdobj.Gradients().set( + verbose=0, + fixed_grid=True, + root_overlap_tol=0.5, + ) + analytic = gradient.kernel(state=2, method="analytic") + finite_difference = gradient.kernel( + state=2, method="finite_diff", step=2e-4, + ) + error = np.max(np.abs(analytic - finite_difference)) + self.assertLess(error, threshold) + + def test_delta_s_zero_hf_lda_gga_mgga_hybrid_and_rsh(self): + cases = ( + ("HF", False, 3e-5), + ("SVWN", False, 1e-5), + ("PBE", False, 1e-5), + ("TPSS", False, 1e-5), + ("M06-2X", False, 1e-5), + ("M06-2X", True, 1e-5), + ("CAM-B3LYP", False, 1e-5), + ) + for xc, nobeta, threshold in cases: + with self.subTest(xc=xc, nobeta=nobeta): + self.compare_public_gradient( + xc, delta_s=0, nobeta=nobeta, threshold=threshold, + ) + + def test_delta_s_minus_one_shares_the_independent_driver(self): + for xc, nobeta in (("PBE", False), ("M06-2X", True)): + with self.subTest(xc=xc, nobeta=nobeta): + self.compare_public_gradient( + xc, delta_s=-1, nobeta=nobeta, threshold=1e-5, + ) + + def test_delta_s_plus_one_rejects_analytic_and_keeps_finite_difference(self): + tdobj = self.make_td("HF", delta_s=1) + gradient = tdobj.Gradients().set(verbose=0) + with self.assertRaisesRegex(NotImplementedError, "deltaS=1"): + gradient.kernel(state=1, method="analytic") + finite_difference = gradient.kernel( + state=1, + atmlst=[0], + method="finite_diff", + step=1e-3, + ) + self.assertTrue(np.all(np.isfinite(finite_difference))) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/nest/grad/tests/test_nttda_gradient_layers.py b/src/nest/grad/tests/test_nttda_gradient_layers.py new file mode 100644 index 0000000..0c88866 --- /dev/null +++ b/src/nest/grad/tests/test_nttda_gradient_layers.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python +"""Layered M-matrix and frozen-orbital checks for NTTDA gradients.""" + +import unittest +from pathlib import Path + +import numpy as np + + + + +from pyscf import dft, gto + + +from nest.grad.nttda.delta_s_minus_one import ( # noqa: E402 + grad_elec as lowering_grad_elec, + spin_lowering_ledger_scalar, +) +from nest.grad.nttda.delta_s_zero import ( # noqa: E402 + grad_elec as same_spin_grad_elec, + same_spin_ledger_scalar, +) +from nest.grad.nttda import xc as xc_backend # noqa: E402 +from nest.nttda import NTTDA # noqa: E402 + + +FUNCTIONALS = ("HF", "SVWN", "PBE", "TPSS", "M06-2X", "CAM-B3LYP") + + +def make_molecule(): + return gto.M( + atom="N 0 0 0; O 0 0 1.20; H 0 0.90 -0.20", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + + +def make_reference(mol, xc): + mf = dft.ROKS(mol).set( + xc=xc, + conv_tol=1e-14, + conv_tol_grad=1e-11, + max_cycle=200, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + if not mf.converged: + raise RuntimeError("ROKS reference did not converge") + return mf + + +def exact_state_two(mf, delta_s, nobeta): + tdobj = NTTDA(mf).set( + deltaS=delta_s, + nobeta=nobeta, + nstates=3, + max_memory=mf.max_memory, + verbose=0, + ) + if delta_s == 0: + vind, diagonal = tdobj.gen_vind_sc() + else: + vind, diagonal = tdobj.gen_vind_sfd() + size = diagonal.size + rows = np.asarray(vind(np.eye(size))).reshape(size, size) + if np.max(np.abs(rows - rows.T)) >= 1e-10: + raise AssertionError("NTTDA action is not symmetric") + energies, vectors = np.linalg.eigh(0.5 * (rows + rows.T)) + if delta_s == -1: + vectors = vectors[:, np.abs(energies) > 1e-8] + vector = vectors[:, 1] + if delta_s == -1: + nc = np.count_nonzero(mf.mo_occ == 2) + no = np.count_nonzero(mf.mo_occ == 1) + nv = np.count_nonzero(mf.mo_occ == 0) + vector = vector.reshape(nc + no, no + nv) + return tdobj, (vector, 0) + + +def gradient_components(mf, tdobj, xy, atmlst): + builder = lowering_grad_elec if tdobj.deltaS == -1 else same_spin_grad_elec + return builder(mf.nuc_grad_method(), tdobj, xy, atmlst=atmlst) + + +def channel_scalar(tdobj, xy): + if tdobj.deltaS == -1: + return spin_lowering_ledger_scalar(tdobj, xy) + return same_spin_ledger_scalar(tdobj, xy) + + +def frozen_scalar_at(base_mf, tdobj, xy, coords): + mol = base_mf.mol.copy() + mol.set_geom_(coords, unit="Bohr") + mf = dft.ROKS(mol).set(xc=base_mf.xc, verbose=0) + mf.grids.coords = np.array(base_mf.grids.coords, copy=True) + mf.grids.weights = np.array(base_mf.grids.weights, copy=True) + mf.grids.non0tab = None + mf.mo_coeff = np.array(base_mf.mo_coeff, copy=True) + mf.mo_occ = np.array(base_mf.mo_occ, copy=True) + mf.mo_energy = np.array(base_mf.mo_energy, copy=True) + displaced = NTTDA(mf).set( + deltaS=tdobj.deltaS, + nobeta=tdobj.nobeta, + max_memory=tdobj.max_memory, + verbose=0, + ) + return channel_scalar(displaced, xy) + + +class GradientLayerChecks(unittest.TestCase): + def test_batched_jk_and_shared_vxc_call_counts(self): + mf = make_reference(make_molecule(), "CAM-B3LYP") + for delta_s in (-1, 0): + for nobeta in (False, True): + tdobj, xy = exact_state_two(mf, delta_s, nobeta) + driver = mf.nuc_grad_method() + calls = {"j": 0, "k": 0, "vxc": 0, "j_batch": [], + "k_batch": []} + original_j = driver.get_j + original_k = driver.get_k + original_vxc = xc_backend.contract_gga_vxc_derivative + + def counted_j(mol=None, dm=None, **kwargs): + calls["j"] += 1 + calls["j_batch"].append(1 if dm.ndim == 2 else len(dm)) + return original_j(mol, dm, **kwargs) + + def counted_k(mol=None, dm=None, **kwargs): + calls["k"] += 1 + calls["k_batch"].append(1 if dm.ndim == 2 else len(dm)) + return original_k(mol, dm, **kwargs) + + def counted_vxc(*args, **kwargs): + calls["vxc"] += 1 + return original_vxc(*args, **kwargs) + + driver.get_j = counted_j + driver.get_k = counted_k + xc_backend.contract_gga_vxc_derivative = counted_vxc + try: + builder = ( + lowering_grad_elec if delta_s == -1 + else same_spin_grad_elec + ) + result = builder( + driver, + tdobj, + xy, + atmlst=range(mf.mol.natm), + ) + finally: + xc_backend.contract_gga_vxc_derivative = original_vxc + + with self.subTest(delta_s=delta_s, nobeta=nobeta): + self.assertTrue(np.all(np.isfinite(result.total))) + self.assertEqual(calls["j"], 2) + self.assertEqual(calls["k"], 2) + self.assertGreater(max(calls["j_batch"]), 1) + self.assertGreater(max(calls["k_batch"]), 1) + expected_vxc = 2 if nobeta else 1 + self.assertEqual(calls["vxc"], expected_vxc) + + def test_full_m_matrix_for_both_channels_and_fock_modes(self): + rng = np.random.default_rng(103) + for delta_s in (-1, 0): + for xc in FUNCTIONALS: + mf = make_reference(make_molecule(), xc) + for nobeta in (False, True): + tdobj, xy = exact_state_two(mf, delta_s, nobeta) + result = gradient_components(mf, tdobj, xy, atmlst=()) + nmo = mf.mo_coeff.shape[1] + perturbation = rng.normal(size=(nmo, nmo)) + original = np.array(mf.mo_coeff, copy=True) + step = 1e-6 + values = [] + try: + for sign in (1.0, -1.0): + mf.mo_coeff = original @ ( + np.eye(nmo) + sign * step * perturbation + ) + values.append(channel_scalar(tdobj, xy)) + finally: + mf.mo_coeff = original + finite_difference = (values[0] - values[1]) / (2 * step) + analytic = np.trace(result.m_matrix.T @ perturbation) + with self.subTest( + delta_s=delta_s, xc=xc, nobeta=nobeta): + self.assertLess(abs(analytic - finite_difference), 1e-8) + + def test_direct_derivative_for_representative_functionals(self): + cases = ( + (0, "HF", False), + (0, "SVWN", False), + (0, "PBE", False), + (0, "TPSS", False), + (0, "M06-2X", False), + (0, "M06-2X", True), + (0, "CAM-B3LYP", False), + (-1, "HF", False), + (-1, "PBE", False), + (-1, "TPSS", False), + (-1, "M06-2X", True), + ) + step = 1e-4 + for delta_s, xc, nobeta in cases: + mf = make_reference(make_molecule(), xc) + tdobj, xy = exact_state_two(mf, delta_s, nobeta) + result = gradient_components( + mf, tdobj, xy, atmlst=range(mf.mol.natm), + ) + coords0 = mf.mol.atom_coords() + finite_difference = np.zeros_like(coords0) + for atom in range(mf.mol.natm): + for xyz in range(3): + coords_plus = coords0.copy() + coords_minus = coords0.copy() + coords_plus[atom, xyz] += step + coords_minus[atom, xyz] -= step + value_plus = frozen_scalar_at( + mf, tdobj, xy, coords_plus, + ) + value_minus = frozen_scalar_at( + mf, tdobj, xy, coords_minus, + ) + finite_difference[atom, xyz] = ( + (value_plus - value_minus) / (2 * step) + ) + error = np.max(np.abs(result.direct - finite_difference)) + threshold = 1e-6 if xc == "HF" else 3e-6 + with self.subTest(delta_s=delta_s, xc=xc, nobeta=nobeta): + self.assertLess(error, threshold) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/nest/grad/tests/test_nttda_scalar_ledger.py b/src/nest/grad/tests/test_nttda_scalar_ledger.py new file mode 100644 index 0000000..4dabb0f --- /dev/null +++ b/src/nest/grad/tests/test_nttda_scalar_ledger.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python +"""Scalar-closure tests for both independent NTTDA analytic channels.""" + +import unittest +from pathlib import Path + +import numpy as np + +import nest +from pyscf import dft, gto +from nest.nttda import NTTDA + + + +from nest.grad.nttda.delta_s_zero import ( # noqa: E402 + same_spin_action_scalar, + same_spin_ledger_scalar, +) +from nest.grad.nttda.delta_s_minus_one import ( # noqa: E402 + spin_lowering_action_scalar, + spin_lowering_ledger_scalar, +) + + +class SameSpinScalarClosure(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.mol = gto.M( + atom="N 0 0 0; O 0 0 1.20; H 0 0.90 -0.20", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + + def test_selected_functionals_eigenvector_and_random_vector(self): + rng = np.random.default_rng(19) + for xc in ("HF", "SVWN", "PBE", "TPSS", "M06-2X", "CAM-B3LYP"): + mf = dft.ROKS(self.mol).set(xc=xc, conv_tol=1e-11, verbose=0) + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + for nobeta in (False, True): + td = NTTDA(mf).set( + deltaS=0, + nobeta=nobeta, + nstates=3, + conv_tol=1e-10, + max_cycle=200, + verbose=0, + ) + vind, hdiag = td.gen_vind_sc() + rows = np.asarray(vind(np.eye(hdiag.size))).reshape( + hdiag.size, hdiag.size, + ) + _energies, eigenvectors = np.linalg.eigh( + 0.5 * (rows + rows.T), + ) + vectors = { + "root2": eigenvectors[:, 1], + "random": rng.normal(size=hdiag.size), + } + for vector_kind, vector in vectors.items(): + with self.subTest( + xc=xc, nobeta=nobeta, vector=vector_kind): + action = same_spin_action_scalar(td, vector) + ledger = same_spin_ledger_scalar(td, vector) + self.assertLess(abs(action - ledger), 1e-11) + + def test_nttda_package_does_not_import_satda_gradient_modules(self): + imported = set(__import__("sys").modules) + self.assertNotIn("pyscf.grad.tdsatda_delta", imported) + self.assertNotIn("pyscf.grad.tdsatda_fast", imported) + self.assertTrue( + str(Path(nest.grad.nttda.__file__).resolve()).startswith(str(Path(nest.__file__).resolve().parent)) + ) + + def test_lowering_channel_uses_an_independent_closed_ledger(self): + rng = np.random.default_rng(31) + for xc in ("HF", "SVWN", "PBE", "TPSS", "M06-2X", "CAM-B3LYP"): + mf = dft.ROKS(self.mol).set(xc=xc, conv_tol=1e-11, verbose=0) + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + for nobeta in (False, True): + td = NTTDA(mf).set( + deltaS=-1, + nobeta=nobeta, + nstates=3, + conv_tol=1e-10, + verbose=0, + ).run() + _vind, diagonal = td.gen_vind_sfd() + vectors = ( + np.asarray(td.xy[1][0]), + rng.normal(size=diagonal.size), + ) + for vector in vectors: + with self.subTest(xc=xc, nobeta=nobeta): + action = spin_lowering_action_scalar(td, vector) + ledger = spin_lowering_ledger_scalar(td, vector) + self.assertLess(abs(action - ledger), 1e-11) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/nest/nttda/tests/test_nttda.py b/src/nest/nttda/tests/test_nttda.py index f807a32..bcb4394 100644 --- a/src/nest/nttda/tests/test_nttda.py +++ b/src/nest/nttda/tests/test_nttda.py @@ -1,4 +1,4 @@ -# Copyright 2026 The NEST Developers. All Rights Reserved. +# Copyright 2021-2024 The PySCF Developers. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,9 +13,110 @@ # limitations under the License. import unittest +from unittest import mock import numpy as np from pyscf import gto -from nest import nttda +from nest.nttda import nttda + + +REFS = { + 'HF': { + True: { + -1: np.array([-0.25588162251385949, 0.031791648059151634, + 0.082164577524099669, 0.10984172789557073, + 0.14436589862648663]), + 0: np.array([-0.021227306082552837, 0.036812245658307111, + 0.055828674198231877, 0.10820095281082762, + 0.13769598373589942]), + 1: np.array([0.26373033968267973, 0.32114587049263738, + 0.35767192060413755, 0.4089546816647468, + 0.48418822465436356]), + }, + False: { + -1: np.array([-0.25588162251385815, 0.03179164805915535, + 0.08216457752408901, 0.10984172789556618, + 0.14436589862650168]), + 0: np.array([-0.021227306082554027, 0.03681224565830669, + 0.05582867419822887, 0.10820095281083697, + 0.13769598373589134]), + 1: np.array([0.26373033968267307, 0.32114587049263676, + 0.35767192060412084, 0.4089546816647443, + 0.48418822465431444]), + }, + }, + 'SVWN': { + True: { + -1: np.array([-0.21136285952298853, 0.022829192982022128, + 0.04449709298041335, 0.070334528481137998, + 0.11581794978093166]), + 0: np.array([-0.0014224229333087768, 0.029907227771976085, + 0.042159504595931208, 0.087581948278004515, + 0.17827490921649417]), + 1: np.array([0.26097145556097057, 0.31399118616119204, + 0.40046031718535502, 0.44773897177486011, + 0.45191690809443946]), + }, + False: { + -1: np.array([-0.21170979048359168, 0.023046236405179832, + 0.04399445674703403, 0.07151698987123586, + 0.1149177949908444]), + 0: np.array([-0.0014102920144926014, 0.029534076286594643, + 0.043327478728623726, 0.08675602648118973, + 0.17868955257035488]), + 1: np.array([0.2621305574444208, 0.3146577468311684, + 0.400854855533031, 0.4488089982219909, + 0.45217145231139155]), + }, + }, + 'M062X': { + True: { + -1: np.array([-0.24666086824597583, 0.015820053409613927, + 0.050190722681826144, 0.071795073579681096, + 0.12358842176137239]), + 0: np.array([-0.0066422638316957381, 0.028055776231764321, + 0.034831792351000868, 0.097283193576694127, + 0.16108162207164683]), + 1: np.array([0.277635913239132, 0.33395796939250971, + 0.38888717645852439, 0.44267039406168623, + 0.49110742356819381]), + }, + False: { + -1: np.array([-0.24280053851498867, 0.011530030283297799, + 0.05005354330269396, 0.06762698114448712, + 0.12639763640154533]), + 0: np.array([-0.008184446338165025, 0.025150738879015422, + 0.032777031664227074, 0.09911876211938828, + 0.15953063092372488]), + 1: np.array([0.26880002289621757, 0.3280851476633962, + 0.3822461897656717, 0.4389979233432141, + 0.4818601324603382]), + }, + }, + 'CAM-B3LYP': { + True: { + -1: np.array([-0.22468903600466386, 0.022443972864282041, + 0.053034041517141139, 0.079464935567422901, + 0.12134548968286102]), + 0: np.array([-0.0044893465927124268, 0.035037117269294718, + 0.043274626762097285, 0.096035968020390092, + 0.17133618259284775]), + 1: np.array([0.27155932081326395, 0.32184531828332463, + 0.38819254300788419, 0.43485814799250122, + 0.47810311079140677]), + }, + False: { + -1: np.array([-0.22362676199942616, 0.02217598445976246, + 0.0521859034687622, 0.08054218557201234, + 0.12099781181846828]), + 0: np.array([-0.004488754315208394, 0.03433453230586325, + 0.044209793487172355, 0.09543419827896533, + 0.17172220940564042]), + 1: np.array([0.27064837890318744, 0.32121564431676974, + 0.3866085753082913, 0.43353862650338093, + 0.47757913242638816]), + }, + }, +} class KnownValues(unittest.TestCase): @@ -24,94 +125,71 @@ def setUpClass(cls): mol = gto.Mole() mol.verbose = 0 mol.output = '/dev/null' - mol.atom = """ + mol.atom = ''' O 0.64372820 0.14077399 -0.04477253 O -0.64862595 -0.12779073 -0.05445498 H 1.16027512 -0.65947800 0.36730132 H -1.12109306 0.55561188 0.42651873 - """ + ''' mol.charge = 0 mol.spin = 2 mol.basis = '631g' + mol.symmetry = True cls.mol = mol.build() @classmethod def tearDownClass(cls): cls.mol.stdout.close() - def test_hf_nttda(self): - mf = self.mol.ROKS(xc='HF').run() - - ref = np.array([0.26373033968267973, 0.32114587049263738]) - td = mf.NTTDA().set(nstates=2, deltaS=1, nobeta=True).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + def _check_functional(self, xc): + mf = self.mol.ROKS(xc=xc).run() + for nobeta, refs_by_delta_s in REFS[xc].items(): + for delta_s, ref in refs_by_delta_s.items(): + with self.subTest(xc=xc, nobeta=nobeta, deltaS=delta_s): + td = nttda.NTTDA(mf) + td.nstates = 5 + td.deltaS = delta_s + td.nobeta = nobeta + td.verbose = 0 + td.kernel() + self.assertTrue(np.all(td.converged)) + np.testing.assert_allclose(td.e, ref, atol=1e-6, rtol=0) - ref = np.array([-0.25588162251385815, 0.03179164805915535]) - td = mf.NTTDA().set(nstates=2, deltaS=-1, nobeta=False).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) - - ref = np.array([-0.021227306082554027, 0.03681224565830669]) - td = mf.NTTDA().set(nstates=2, deltaS=0, nobeta=False).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + def test_hf_nttda(self): + self._check_functional('HF') def test_svwn_nttda(self): - mf = self.mol.ROKS(xc='SVWN').run() - - ref = np.array([-0.21136285952298853, 0.022829192982022128]) - td = mf.NTTDA().set(nstates=2, deltaS=-1, nobeta=True).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) - - ref = np.array([-0.0014224229333087768, 0.029907227771976085]) - td = mf.NTTDA().set(nstates=2, deltaS=0, nobeta=True).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) - - ref = np.array([0.2621305574444208, 0.3146577468311684]) - td = mf.NTTDA().set(nstates=2, deltaS=1, nobeta=False).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + self._check_functional('SVWN') def test_m062x_nttda(self): - mf = self.mol.ROKS(xc='M062X').run() - - ref = np.array([-0.24666086824597583, 0.015820053409613927]) - td = mf.NTTDA().set(nstates=2, deltaS=-1, nobeta=True).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) - - ref = np.array([-0.008184446338165025, 0.025150738879015422]) - td = mf.NTTDA().set(nstates=2, deltaS=0, nobeta=False).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) - - ref = np.array([0.26880002289621757, 0.3280851476633962]) - td = mf.NTTDA().set(nstates=2, deltaS=1, nobeta=False).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + self._check_functional('M062X') def test_cam_b3lyp_nttda(self): - mf = self.mol.ROKS(xc='CAM-B3LYP').run() - - ref = np.array([-0.0044893465927124268, 0.035037117269294718]) - td = mf.NTTDA().set(nstates=2, deltaS=0, nobeta=True).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) - - ref = np.array([0.27155932081326395, 0.32184531828332463]) - td = mf.NTTDA().set(nstates=2, deltaS=1, nobeta=True).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) - - ref = np.array([-0.22362676199942616, 0.02217598445976246]) - td = mf.NTTDA().set(nstates=2, deltaS=-1, nobeta=False).run() - self.assertTrue(np.all(td.converged)) - self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) - - -if __name__ == '__main__': - print('Full tests for noncollinear tensor TDA based on ROKS reference') + self._check_functional('CAM-B3LYP') + + def test_mo_grid_fxc1_vind_matches_ao(self): + rng = np.random.default_rng(12) + for xc in ('BLYP', 'TPSS'): + mf = self.mol.ROKS(xc=xc).run() + for delta_s in (0, -1): + with self.subTest(xc=xc, deltaS=delta_s): + td0 = nttda.NTTDA(mf) + td0.deltaS = delta_s + td0.verbose = 0 + td1 = nttda.NTTDA(mf) + td1.deltaS = delta_s + td1.verbose = 0 + with mock.patch.object(nttda, 'MO_GRID_FXC1', False): + vind0, hdiag0 = (td0.gen_vind_sc() if delta_s == 0 + else td0.gen_vind_sfd()) + with mock.patch.object(nttda, 'MO_GRID_FXC1', True): + vind1, hdiag1 = (td1.gen_vind_sc() if delta_s == 0 + else td1.gen_vind_sfd()) + np.testing.assert_allclose(hdiag1, hdiag0, atol=1e-10, rtol=0) + zs = rng.standard_normal((2, hdiag0.size)) + np.testing.assert_allclose(vind1(zs), vind0(zs), atol=1e-9, rtol=0) + + +if __name__ == "__main__": + print("Full Tests for noncollinear tensor TDA based on ROKS reference") unittest.main() diff --git a/src/nest/nttda/tests/test_nttda_ensemble.py b/src/nest/nttda/tests/test_nttda_ensemble.py new file mode 100644 index 0000000..ea4f240 --- /dev/null +++ b/src/nest/nttda/tests/test_nttda_ensemble.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python +"""Acceptance tests for EnsembleRKS-based NTTDA energies.""" + +import unittest + +import numpy as np + +from pyscf import gto +from pyscf.scf import hf +from nest.ensemble_rks import EnsembleRKS +from nest.nttda import NTTDA + + +class EnsembleRKSReference(unittest.TestCase): + @staticmethod + def lithium_hydride_cation(): + return gto.M( + atom="Li 0 0 0; H 0 0 3.0", + basis="sto-3g", + charge=1, + spin=1, + unit="Bohr", + verbose=0, + ) + + def make_reference(self, xc="SVWN"): + mf = EnsembleRKS(self.lithium_hydride_cation()).set( + xc=xc, + conv_tol=1e-12, + conv_tol_grad=1e-9, + max_cycle=100, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + return mf + + def test_fixed_occupations_define_a_spin_unpolarized_ensemble(self): + mf = self.make_reference() + np.testing.assert_array_equal(mf.mo_occ, [2, 1, 0, 0, 0, 0]) + self.assertAlmostEqual(mf.mo_occ.sum(), mf.mol.nelectron) + + dm = mf.make_rdm1() + dma, dmb = mf.make_rdm1s() + np.testing.assert_allclose(dma, dmb, atol=0, rtol=0) + np.testing.assert_allclose(dma + dmb, dm, atol=1e-14, rtol=0) + self.assertAlmostEqual( + np.einsum("ij,ji", dm, mf.get_ovlp()), + mf.mol.nelectron, + places=10, + ) + + rng = np.random.default_rng(8) + fock = rng.standard_normal(dm.shape) + fock = fock + fock.T + fock_mo = mf.mo_coeff.T @ fock @ mf.mo_coeff + unique = hf.uniq_var_indices(mf.mo_occ) + occupation_difference = mf.mo_occ[None, :] - mf.mo_occ[:, None] + expected = (fock_mo * occupation_difference)[unique] + np.testing.assert_allclose( + mf.get_grad(mf.mo_coeff, mf.mo_occ, fock), + expected, + atol=1e-14, + rtol=0, + ) + + def test_explicit_open_shell_count_must_match_molecular_spin(self): + with self.assertRaisesRegex(ValueError, "nopen.*mol.spin"): + EnsembleRKS(self.lithium_hydride_cation(), nopen=3) + + def test_nttda_energy_uses_the_ensemble_fock_for_both_nobeta_values(self): + mf = self.make_reference() + energies = [] + for nobeta in (False, True): + tdobj = NTTDA(mf).set( + deltaS=0, + nobeta=nobeta, + nstates=2, + conv_tol=1e-8, + max_cycle=200, + verbose=0, + ).run() + self.assertTrue(np.all(tdobj.converged)) + self.assertTrue(np.all(np.isfinite(tdobj.e))) + energies.append(tdobj.e) + np.testing.assert_allclose(energies[0], energies[1], atol=1e-12, rtol=0) + + def test_reference_energy_is_the_stationary_average_occupation_energy(self): + mf = self.make_reference() + self.assertEqual( + mf.reference_energy_semantics, + "average_occupation_ensemble_rks_energy", + ) + self.assertTrue(mf.reference_energy_stationary) + self.assertAlmostEqual(mf.reference_energy(), mf.e_tot, places=14) + + def test_nttda_total_energies_use_the_selected_reference_energy(self): + mf = self.make_reference() + tdobj = NTTDA(mf).set( + deltaS=0, + nstates=2, + conv_tol=1e-8, + max_cycle=200, + verbose=0, + ).run() + + self.assertAlmostEqual(tdobj.reference_energy(), mf.e_tot) + np.testing.assert_allclose( + tdobj.total_energies(), + mf.e_tot + tdobj.e, + atol=1e-13, + rtol=0, + ) + + def test_nttda_supports_common_functional_families(self): + for xc in ("HF", "PBE", "TPSS", "M06-2X", "CAM-B3LYP"): + with self.subTest(xc=xc): + mf = self.make_reference(xc) + tdobj = NTTDA(mf).set( + deltaS=0, + nstates=2, + conv_tol=1e-6, + max_cycle=200, + verbose=0, + ).run() + self.assertTrue(np.all(tdobj.converged)) + self.assertTrue(np.all(np.isfinite(tdobj.e))) + + def test_all_three_spin_channels(self): + mol = gto.M( + atom="C 0 0 0; H 0 0 2.0; H 0 1.7 -0.5", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + mf = EnsembleRKS(mol).set( + xc="SVWN", + conv_tol=1e-12, + max_cycle=150, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + + references = { + -1: [0.01477165973875402, 0.06753726182098739], + 0: [-0.002622144798075737, 0.2393673138366825], + 1: [0.6745475947851435, 0.8387186686911036], + } + for delta_s, reference in references.items(): + with self.subTest(deltaS=delta_s): + tdobj = NTTDA(mf).set( + deltaS=delta_s, + nstates=2, + conv_tol=1e-7, + max_cycle=200, + verbose=0, + ).run() + self.assertTrue(np.all(tdobj.converged)) + np.testing.assert_allclose( + tdobj.e, reference, atol=2e-6, rtol=0, + ) + + +if __name__ == "__main__": + unittest.main() From d6808d08a5a57a3a22443bf8fad687c514581248 Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Wed, 9 Sep 2026 14:47:03 +0800 Subject: [PATCH 07/16] refactor: use Dz0SCF as the average-occupation NTTDA reference Drop the EnsembleRKS reference path and drive the NTTDA gradient from Dz0SCF (average-occupation) and ROKS only. Dz0SCF now exposes the average-occupation marker plus a charge-RKS companion for the spin-unpolarized response, rebuilt on every call so a reused mean-field object (reset/new geometry/xc change) never reuses a stale reference. Port the ensemble suites to Dz0SCF and add a cache-invalidation regression test. --- src/nest/__init__.py | 4 +- src/nest/dz0scf/dz0scf.py | 54 +++++++- src/nest/ensemble_rks/__init__.py | 3 - src/nest/ensemble_rks/ensemble_rks.py | 127 ------------------ .../ensemble_rks/tests/test_ensemble_rks.py | 71 ---------- src/nest/grad/nttda/__init__.py | 14 +- src/nest/grad/nttda/delta_s_minus_one.py | 13 +- src/nest/grad/nttda/delta_s_zero.py | 13 +- src/nest/grad/nttda/ensemble.py | 4 +- src/nest/grad/nttda/roks.py | 7 +- src/nest/grad/nttda/xc.py | 11 +- ...ensemble_fd.py => test_nttda_dz0scf_fd.py} | 30 ++--- ...ponse.py => test_nttda_dz0scf_response.py} | 78 +++++++++-- src/nest/nttda/nttda.py | 16 +-- ...nttda_ensemble.py => test_nttda_dz0scf.py} | 55 ++++---- 15 files changed, 189 insertions(+), 311 deletions(-) delete mode 100644 src/nest/ensemble_rks/__init__.py delete mode 100644 src/nest/ensemble_rks/ensemble_rks.py delete mode 100644 src/nest/ensemble_rks/tests/test_ensemble_rks.py rename src/nest/grad/tests/{test_nttda_ensemble_fd.py => test_nttda_dz0scf_fd.py} (89%) rename src/nest/grad/tests/{test_nttda_ensemble_response.py => test_nttda_dz0scf_response.py} (73%) rename src/nest/nttda/tests/{test_nttda_ensemble.py => test_nttda_dz0scf.py} (76%) diff --git a/src/nest/__init__.py b/src/nest/__init__.py index 3ce6f9c..f90fbce 100644 --- a/src/nest/__init__.py +++ b/src/nest/__init__.py @@ -15,8 +15,8 @@ """nest package.""" -from nest import dz0scf, ensemble_rks, nttda, sftda +from nest import dz0scf, nttda, sftda __version__ = "0.1.0" -__all__ = ["__version__", "dz0scf", "ensemble_rks", "nttda", "sftda"] +__all__ = ["__version__", "dz0scf", "nttda", "sftda"] diff --git a/src/nest/dz0scf/dz0scf.py b/src/nest/dz0scf/dz0scf.py index 6382416..7a4f1a0 100644 --- a/src/nest/dz0scf/dz0scf.py +++ b/src/nest/dz0scf/dz0scf.py @@ -54,11 +54,55 @@ def evaluate_high_spin_energy(mf): vhf=veff, ) -class _DZ0VeffMixin: - reference_energy_semantics = 'high_spin_roks_energy_on_dz0_orbitals' - reference_energy_stationary = False - - def get_veff( +class _DZ0VeffMixin: + reference_energy_semantics = 'high_spin_roks_energy_on_dz0_orbitals' + reference_energy_stationary = False + is_average_occupation_reference = True + + def _charge_rks(self): + """Return a fresh RKS view used for the spin-unpolarized charge response. + + The view is rebuilt on every call so that a reused mean-field object + (``reset(new_mol)``, ``xc`` change, new geometry) never feeds a stale + molecule or functional into the response. + """ + charge = dft.rks.RKS(self.mol) + for name in ( + 'xc', 'nlc', 'grids', 'nlcgrids', '_numint', + 'max_memory', 'small_rho_cutoff'): + if hasattr(self, name): + setattr(charge, name, getattr(self, name)) + charge.mo_coeff = np.asarray(self.mo_coeff) + charge.mo_occ = np.asarray(self.mo_occ) + charge.mo_energy = np.asarray(self.mo_energy) + charge.verbose = 0 + return charge + + def make_rdm1s(self, mo_coeff=None, mo_occ=None): + """Return equal spin densities ``D/2`` for the spin-unpolarized reference.""" + if mo_coeff is None: + mo_coeff = self.mo_coeff + if mo_occ is None: + mo_occ = self.mo_occ + mo_coeff = np.asarray(mo_coeff) + occupation = np.asarray(mo_occ) + dm0 = (mo_coeff * occupation) @ mo_coeff.conj().T + return 0.5 * dm0, 0.5 * dm0 + + def gen_response(self, mo_coeff=None, mo_occ=None, hermi=1, max_memory=None): + """Charge-only (spin-unpolarized) linear response of the reference.""" + if mo_coeff is None: + mo_coeff = self.mo_coeff + if mo_occ is None: + mo_occ = self.mo_occ + return self._charge_rks().gen_response( + mo_coeff=mo_coeff, + mo_occ=mo_occ, + hermi=hermi, + max_memory=max_memory, + ) + + def get_veff( self, mol=None, dm=None, diff --git a/src/nest/ensemble_rks/__init__.py b/src/nest/ensemble_rks/__init__.py deleted file mode 100644 index dff4b69..0000000 --- a/src/nest/ensemble_rks/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .ensemble_rks import EnsembleRKS - -__all__ = ['EnsembleRKS'] diff --git a/src/nest/ensemble_rks/ensemble_rks.py b/src/nest/ensemble_rks/ensemble_rks.py deleted file mode 100644 index 0b0e595..0000000 --- a/src/nest/ensemble_rks/ensemble_rks.py +++ /dev/null @@ -1,127 +0,0 @@ -"""Average-occupation restricted ensemble Kohn--Sham references.""" - -import numpy as np - -from pyscf import dft -from pyscf.lib import logger -from pyscf.scf import hf - - -class EnsembleRKS(dft.rks.RKS): - """RKS with fixed ``2/1/0`` occupations and zero spin density.""" - - is_ensemble_rks = True - _keys = dft.rks.RKS._keys | {'nopen', 'reference_semantics'} - - def __init__(self, mol, xc='LDA,VWN', nopen=None, reference_semantics='average'): - super().__init__(mol, xc=xc) - if reference_semantics not in ('average', 'roks'): - raise ValueError( - "reference_semantics must be 'average' or 'roks', got %r" - % (reference_semantics,)) - self.reference_semantics = reference_semantics - if nopen is None: - nopen = mol.spin - if isinstance(nopen, bool) or int(nopen) != nopen: - raise ValueError('nopen must be a non-negative integer') - self.nopen = int(nopen) - self._validate_ensemble() - - @property - def reference_energy_semantics(self): - if self.reference_semantics == 'roks': - return 'roks_energy_on_ensemble_rks_orbitals' - return 'average_occupation_ensemble_rks_energy' - - @property - def reference_energy_stationary(self): - return self.reference_semantics == 'average' - - def _validate_ensemble(self): - if self.nopen < 0: - raise ValueError('nopen must be a non-negative integer') - if self.nopen != self.mol.spin: - raise ValueError('nopen must match mol.spin for NTTDA') - if self.nopen > self.mol.nelectron: - raise ValueError('nopen cannot exceed the electron count') - if (self.mol.nelectron - self.nopen) % 2: - raise ValueError('electron count and nopen have inconsistent parity') - - @property - def nclosed(self): - return (self.mol.nelectron - self.nopen) // 2 - - def check_sanity(self): - out = hf.SCF.check_sanity(self) - if self.do_nlc() and self.do_disp() and self._numint.libxc.is_nlc(self.xc): - import warnings - warnings.warn( - f'nlc-type xc {self.xc} and disp {self.disp} may lead to double counting in NLC.' - ) - return out - - def get_occ(self, mo_energy=None, mo_coeff=None): - self._validate_ensemble() - if mo_energy is None: - mo_energy = self.mo_energy - mo_energy = np.asarray(mo_energy) - if self.nclosed + self.nopen > mo_energy.size: - raise RuntimeError('not enough orbitals for the requested ensemble occupations') - - order = np.argsort(mo_energy, kind='stable') - mo_occ = np.zeros_like(mo_energy) - mo_occ[order[:self.nclosed]] = 2 - mo_occ[order[self.nclosed:self.nclosed + self.nopen]] = 1 - if self.verbose >= logger.INFO: - logger.info(self, 'EnsembleRKS occupations = %s', mo_occ) - return mo_occ - - def make_rdm1s(self, mo_coeff=None, mo_occ=None): - dm = self.make_rdm1(mo_coeff, mo_occ) - dm_spin = np.asarray(dm) * 0.5 - return dm_spin, dm_spin.copy() - - def reference_energy(self): - if self.reference_semantics == 'roks': - if self.mo_coeff is None or self.mo_occ is None: - raise RuntimeError( - 'run EnsembleRKS.kernel() before evaluating the reference energy') - evaluator = dft.ROKS(self.mol).set( - xc=self.xc, - nlc=self.nlc, - max_memory=self.max_memory, - verbose=0, - ) - evaluator.grids = self.grids - evaluator.nlcgrids = self.nlcgrids - dm = evaluator.make_rdm1(self.mo_coeff, self.mo_occ) - hcore = evaluator.get_hcore(self.mol) - veff = evaluator.get_veff(self.mol, dm) - return float(evaluator.energy_tot(dm=dm, h1e=hcore, vhf=veff)) - if self.e_tot is None: - raise RuntimeError('run EnsembleRKS.kernel() before evaluating the reference energy') - return float(self.e_tot) - - def nuc_grad_method(self): - """Return the analytic gradient driver of the selected reference energy. - - For the ``average`` semantics this is the gradient of the stationary - ensemble energy (the stock RKS driver). For the ``roks`` semantics it - differentiates the fixed-orbital ROKS energy returned by - :meth:`reference_energy`, including its Z-vector orbital-relaxation - contribution and the nuclear-repulsion gradient. - """ - if self.reference_semantics == 'roks': - from nest.grad.nttda.reference import ReferenceGradients - return ReferenceGradients(self) - return super().nuc_grad_method() - - def get_grad(self, mo_coeff, mo_occ, fock=None): - mo_occ = np.asarray(mo_occ) - if fock is None: - dm = self.make_rdm1(mo_coeff, mo_occ) - fock = self.get_hcore(self.mol) + self.get_veff(self.mol, dm) - fock_mo = mo_coeff.conj().T @ fock @ mo_coeff - unique = hf.uniq_var_indices(mo_occ) - occupation_difference = mo_occ[None, :] - mo_occ[:, None] - return (fock_mo * occupation_difference)[unique] diff --git a/src/nest/ensemble_rks/tests/test_ensemble_rks.py b/src/nest/ensemble_rks/tests/test_ensemble_rks.py deleted file mode 100644 index 4d4155f..0000000 --- a/src/nest/ensemble_rks/tests/test_ensemble_rks.py +++ /dev/null @@ -1,71 +0,0 @@ -import unittest - -import numpy as np -from pyscf import gto - -from nest.ensemble_rks import EnsembleRKS -from nest.nttda import NTTDA - - -class EnsembleRKSReferenceTest(unittest.TestCase): - def make_reference(self): - mol = gto.M( - atom=''' - O 0.64372820 0.14077399 -0.04477253 - O -0.64862595 -0.12779073 -0.05445498 - H 1.16027512 -0.65947800 0.36730132 - H -1.12109306 0.55561188 0.42651873 - ''', - basis='6-31g', - spin=2, - unit='Angstrom', - verbose=0, - ) - mf = EnsembleRKS(mol, xc='SVWN') - mf.conv_tol = 1e-12 - mf.conv_tol_grad = 1e-9 - mf.max_cycle = 100 - mf.grids.level = 3 - mf.grids.prune = None - mf.small_rho_cutoff = 0.0 - mf.kernel() - self.assertTrue(mf.converged) - return mf - - def test_average_occupation_reference_energy_is_used_for_total_energies(self): - mf = self.make_reference() - self.assertEqual(np.count_nonzero(mf.mo_occ == 1), mf.nopen) - self.assertEqual(np.count_nonzero(mf.mo_occ == 2), mf.nclosed) - self.assertAlmostEqual(mf.reference_energy(), mf.e_tot, places=14) - - td = NTTDA(mf).set(deltaS=0, nstates=2, conv_tol=1e-5, max_cycle=200, verbose=0) - td.kernel() - - self.assertTrue(np.all(td.converged)) - np.testing.assert_allclose( - td.total_energies(), - mf.reference_energy() + td.e, - atol=1e-13, - rtol=0, - ) - - def test_nobeta_does_not_change_an_ensemble_reference(self): - mf = self.make_reference() - roots = [] - for nobeta in (False, True): - td = NTTDA(mf).set( - deltaS=0, - nobeta=nobeta, - nstates=2, - conv_tol=1e-5, - max_cycle=200, - verbose=0, - ) - td.kernel() - self.assertTrue(np.all(td.converged)) - roots.append(td.e) - np.testing.assert_allclose(roots[0], roots[1], atol=1e-12, rtol=0) - - -if __name__ == '__main__': - unittest.main() diff --git a/src/nest/grad/nttda/__init__.py b/src/nest/grad/nttda/__init__.py index 5d62961..b12f59e 100644 --- a/src/nest/grad/nttda/__init__.py +++ b/src/nest/grad/nttda/__init__.py @@ -5,7 +5,6 @@ from pyscf import dft, lib from pyscf.grad import rhf as rhf_grad from pyscf.lib import logger -from nest.ensemble_rks import EnsembleRKS from nest.nttda import NTTDA from . import delta_s_minus_one, delta_s_zero @@ -27,8 +26,8 @@ def _copy_td_settings(source, target): def _displaced_reference(source, mol, fixed_grid): - if isinstance(source, EnsembleRKS): - reference = EnsembleRKS(mol, xc=source.xc, nopen=source.nopen) + if getattr(source, "is_average_occupation_reference", False): + reference = source.__class__(mol) elif isinstance(source, dft.KohnShamDFT): reference = dft.ROKS(mol) else: @@ -53,7 +52,7 @@ def _displaced_reference(source, mol, fixed_grid): class Gradients(rhf_grad.GradientsBase): - """NTTDA gradients, including finite differences for EnsembleRKS.""" + """NTTDA gradients, including finite differences for Dz0SCF references.""" _keys = rhf_grad.GradientsBase._keys | { "state", "method", "step", "fixed_grid", "root_overlap_tol", @@ -65,10 +64,7 @@ def __init__(self, tdobj): self.state = 1 self.method = "analytic" self.step = 1e-3 - self.fixed_grid = ( - isinstance(tdobj._scf, dft.KohnShamDFT) - and not isinstance(tdobj._scf, EnsembleRKS) - ) + self.fixed_grid = isinstance(tdobj._scf, dft.KohnShamDFT) self.root_overlap_tol = 0.5 self.cphf_conv_tol = 1e-12 self.cphf_max_cycle = None @@ -141,7 +137,7 @@ def _energy_at(self, coords, reference_amplitude): "NTTDA state tracking overlap %.6f is below %.6f" % (overlaps[root], self.root_overlap_tol) ) - return mf.e_tot + tdobj.e[root] + return float(tdobj.total_energies()[root]) def _finite_difference(self, atmlst): coords0 = self.mol.atom_coords() diff --git a/src/nest/grad/nttda/delta_s_minus_one.py b/src/nest/grad/nttda/delta_s_minus_one.py index b6ef3ac..19ff17e 100644 --- a/src/nest/grad/nttda/delta_s_minus_one.py +++ b/src/nest/grad/nttda/delta_s_minus_one.py @@ -9,6 +9,7 @@ import numpy as np from pyscf import dft, lib +from pyscf.grad import rhf as rhf_grad from nest.nttda import nttda as nttda_mod from nest.nttda.nttda import gen_rohf_response_sfd @@ -247,7 +248,7 @@ def _fock_response_q(tdobj, p_alpha, p_beta): """Reference-density derivative of a spin-resolved Fock scalar.""" mf = tdobj._scf mo = np.asarray(mf.mo_coeff) - if getattr(mf, "is_ensemble_rks", False): + if getattr(mf, "is_average_occupation_reference", False): occupation = np.asarray(mf.mo_occ) probe = np.asarray(p_alpha) + np.asarray(p_beta) potential = mf.gen_response(hermi=0)(probe.T) @@ -773,7 +774,7 @@ def _contract_derivative_terms( def _reference_spin_densities(tdobj): mf = tdobj._scf - if getattr(mf, "is_ensemble_rks", False): + if getattr(mf, "is_average_occupation_reference", False): return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) mo = np.asarray(mf.mo_coeff) return ( @@ -813,7 +814,7 @@ def spin_fock_direct_dft( p_total = p_alpha + p_beta density_alpha, density_beta = _reference_spin_densities(tdobj) gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - hcore_derivative = mf.nuc_grad_method().hcore_generator(mol) + hcore_derivative = rhf_grad.Gradients(mf).hcore_generator(mol) for k, atom in enumerate(atmlst): gradient[:, k] += lib.einsum( "npq,xpq->nx", p_total, hcore_derivative(atom), @@ -854,7 +855,7 @@ def spin_fock_direct_dft( xctype ) if (nobeta_p0 is not None and tdobj.nobeta - and not getattr(mf, "is_ensemble_rks", False)): + and not getattr(mf, "is_average_occupation_reference", False)): density0 = 0.5 * (density_alpha + density_beta) actual_probe_alpha = np.array(p_alpha, copy=True) actual_probe_beta = np.array(p_beta, copy=True) @@ -909,7 +910,7 @@ def spin_fock_direct_hf( dm_alpha, dm_beta = _reference_spin_densities(tdobj) gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - hcore_derivative = tdobj._scf.nuc_grad_method().hcore_generator(mol) + hcore_derivative = rhf_grad.Gradients(tdobj._scf).hcore_generator(mol) for k, atom in enumerate(atmlst): gradient[:, k] += lib.einsum( "npq,xpq->nx", p_total, hcore_derivative(atom), @@ -1089,7 +1090,7 @@ def grad_elec( jk_ledger=jk_ledger, output_slot=direct_slot, ) - if getattr(mf, "is_ensemble_rks", False): + if getattr(mf, "is_average_occupation_reference", False): fockz_hfx = spin_lowering_fockz_hfx_terms( gradient_driver, tdobj, diff --git a/src/nest/grad/nttda/delta_s_zero.py b/src/nest/grad/nttda/delta_s_zero.py index c807337..18722b0 100644 --- a/src/nest/grad/nttda/delta_s_zero.py +++ b/src/nest/grad/nttda/delta_s_zero.py @@ -11,6 +11,7 @@ import numpy as np from pyscf import dft, lib +from pyscf.grad import rhf as rhf_grad from nest.nttda import nttda as nttda_mod from nest.nttda.nttda import gen_rohf_response_sc @@ -249,7 +250,7 @@ def _fock_response_q(tdobj, p_alpha, p_beta): """Reference-density derivative of a spin-resolved Fock scalar.""" mf = tdobj._scf mo = np.asarray(mf.mo_coeff) - if getattr(mf, "is_ensemble_rks", False): + if getattr(mf, "is_average_occupation_reference", False): occupation = np.asarray(mf.mo_occ) probe = np.asarray(p_alpha) + np.asarray(p_beta) potential = mf.gen_response(hermi=0)(probe.T) @@ -749,7 +750,7 @@ def _contract_derivative_terms( def _reference_spin_densities(tdobj): mf = tdobj._scf - if getattr(mf, "is_ensemble_rks", False): + if getattr(mf, "is_average_occupation_reference", False): return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) mo = np.asarray(mf.mo_coeff) return ( @@ -789,7 +790,7 @@ def spin_fock_direct_dft( p_total = p_alpha + p_beta density_alpha, density_beta = _reference_spin_densities(tdobj) gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - hcore_derivative = mf.nuc_grad_method().hcore_generator(mol) + hcore_derivative = rhf_grad.Gradients(mf).hcore_generator(mol) for k, atom in enumerate(atmlst): gradient[:, k] += lib.einsum( "npq,xpq->nx", p_total, hcore_derivative(atom), @@ -830,7 +831,7 @@ def spin_fock_direct_dft( xctype ) if (nobeta_p0 is not None and tdobj.nobeta - and not getattr(mf, "is_ensemble_rks", False)): + and not getattr(mf, "is_average_occupation_reference", False)): density0 = 0.5 * (density_alpha + density_beta) actual_probe_alpha = np.array(p_alpha, copy=True) actual_probe_beta = np.array(p_beta, copy=True) @@ -885,7 +886,7 @@ def spin_fock_direct_hf( dm_alpha, dm_beta = _reference_spin_densities(tdobj) gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - hcore_derivative = tdobj._scf.nuc_grad_method().hcore_generator(mol) + hcore_derivative = rhf_grad.Gradients(tdobj._scf).hcore_generator(mol) for k, atom in enumerate(atmlst): gradient[:, k] += lib.einsum( "npq,xpq->nx", p_total, hcore_derivative(atom), @@ -1094,7 +1095,7 @@ def grad_elec( jk_ledger=jk_ledger, output_slot=direct_slot, ) - if getattr(mf, "is_ensemble_rks", False): + if getattr(mf, "is_average_occupation_reference", False): fockz_hfx = same_spin_fockz_hfx_terms( gradient_driver, tdobj, diff --git a/src/nest/grad/nttda/ensemble.py b/src/nest/grad/nttda/ensemble.py index afb1b90..6129b8b 100644 --- a/src/nest/grad/nttda/ensemble.py +++ b/src/nest/grad/nttda/ensemble.py @@ -1,4 +1,4 @@ -"""Average-occupation orbital response for EnsembleRKS NTTDA gradients.""" +"""Average-occupation orbital response for Dz0SCF NTTDA gradients.""" import numpy as np @@ -143,7 +143,7 @@ def operator(vector): def finish_gradient( gradient_driver, tdobj, m_matrix, direct, atmlst, tolerance, max_cycle, fock_direct, direct_fock_probes=None): - """Solve the EnsembleRKS adjoint and assemble ``d omega / dR``.""" + """Solve the Dz0SCF average-occupation adjoint and assemble ``d omega / dR``.""" transpose_action, pairs = make_hessian_transpose_action(tdobj) rhs = pack_m_matrix(m_matrix, pairs) zvector = solve_zvector( diff --git a/src/nest/grad/nttda/roks.py b/src/nest/grad/nttda/roks.py index 63bb36c..82cb445 100644 --- a/src/nest/grad/nttda/roks.py +++ b/src/nest/grad/nttda/roks.py @@ -5,6 +5,7 @@ import numpy as np from pyscf import dft, lib +from pyscf.grad import rhf as rhf_grad @@ -28,9 +29,7 @@ def finish_gradient( ``direct_fock_probes`` enables one batched Fock-derivative evaluation: its contraction is the first result and the Z-vector contraction is the second. """ - from nest.ensemble_rks import EnsembleRKS - - if isinstance(tdobj._scf, EnsembleRKS): + if getattr(tdobj._scf, "is_average_occupation_reference", False): from .ensemble import finish_gradient as finish_ensemble_gradient return finish_ensemble_gradient( gradient_driver, @@ -316,7 +315,7 @@ def _orbital_gradient( atmlst = range(mol.natm) atmlst = tuple(atmlst) mo = np.asarray(mf.mo_coeff) - overlap_derivative = mf.nuc_grad_method().get_ovlp(mol) + overlap_derivative = rhf_grad.Gradients(mf).get_ovlp(mol) offsets = mol.offset_nr_by_atom() result = np.zeros((len(atmlst), 3)) for k, atom in enumerate(atmlst): diff --git a/src/nest/grad/nttda/xc.py b/src/nest/grad/nttda/xc.py index f7768d8..8d5c83e 100644 --- a/src/nest/grad/nttda/xc.py +++ b/src/nest/grad/nttda/xc.py @@ -339,7 +339,7 @@ def _project_channel_potentials(tdobj, potentials, blocks): def _reference_spin_densities(tdobj): """Spin densities of the variational reference used by the XC kernel.""" mf = tdobj._scf - if getattr(mf, "is_ensemble_rks", False): + if getattr(mf, "is_average_occupation_reference", False): return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) mo = np.asarray(mf.mo_coeff) return ( @@ -352,7 +352,7 @@ def _reference_spin_occupations(tdobj): """Per-orbital alpha/beta occupations of the reference density.""" mf = tdobj._scf occupation = np.asarray(mf.mo_occ) - if getattr(mf, "is_ensemble_rks", False): + if getattr(mf, "is_average_occupation_reference", False): return 0.5 * occupation, 0.5 * occupation return (occupation > 0).astype(float), (occupation == 2).astype(float) @@ -644,7 +644,6 @@ def lda_response_terms( reference_alpha = np.zeros((nao, nao)) reference_beta = np.zeros_like(reference_alpha) direct = np.zeros((len(atmlst), 3)) - mo = np.asarray(mf.mo_coeff) density_alpha, density_beta = _reference_spin_densities(tdobj) density_labels, density_stack = _response_density_stack( densities, density_alpha, density_beta, @@ -806,7 +805,7 @@ def lda_nobeta_reference_q(tdobj, p0, max_memory=None): nmo = mo.shape[1] q_alpha = np.zeros((nmo, nmo)) q_beta = np.zeros_like(q_alpha) - if not tdobj.nobeta or getattr(mf, "is_ensemble_rks", False): + if not tdobj.nobeta or getattr(mf, "is_average_occupation_reference", False): return q_alpha, q_beta if max_memory is None: max_memory = tdobj.max_memory @@ -1145,7 +1144,7 @@ def gga_nobeta_reference_q(tdobj, p0, max_memory=None): mo = np.asarray(mf.mo_coeff) q_alpha = np.zeros((mo.shape[1], mo.shape[1])) q_beta = np.zeros_like(q_alpha) - if not tdobj.nobeta or getattr(mf, "is_ensemble_rks", False): + if not tdobj.nobeta or getattr(mf, "is_average_occupation_reference", False): return q_alpha, q_beta if max_memory is None: max_memory = tdobj.max_memory @@ -1485,7 +1484,7 @@ def mgga_nobeta_reference_q(tdobj, p0, max_memory=None): mo = np.asarray(mf.mo_coeff) q_alpha = np.zeros((mo.shape[1], mo.shape[1])) q_beta = np.zeros_like(q_alpha) - if not tdobj.nobeta or getattr(mf, "is_ensemble_rks", False): + if not tdobj.nobeta or getattr(mf, "is_average_occupation_reference", False): return q_alpha, q_beta if max_memory is None: max_memory = tdobj.max_memory diff --git a/src/nest/grad/tests/test_nttda_ensemble_fd.py b/src/nest/grad/tests/test_nttda_dz0scf_fd.py similarity index 89% rename from src/nest/grad/tests/test_nttda_ensemble_fd.py rename to src/nest/grad/tests/test_nttda_dz0scf_fd.py index 7ce56ed..0736c5b 100644 --- a/src/nest/grad/tests/test_nttda_ensemble_fd.py +++ b/src/nest/grad/tests/test_nttda_dz0scf_fd.py @@ -1,16 +1,16 @@ #!/usr/bin/env python -"""Finite-difference NTTDA gradients on an EnsembleRKS reference.""" +"""Finite-difference NTTDA gradients on a Dz0SCF reference.""" import unittest import numpy as np from pyscf import gto -from nest.ensemble_rks import EnsembleRKS +from nest.dz0scf import DZ0SCF from nest.nttda import NTTDA -class EnsembleRKSFiniteDifferenceGradient(unittest.TestCase): +class Dz0SCFFiniteDifferenceGradient(unittest.TestCase): @staticmethod def make_td(): mol = gto.M( @@ -21,8 +21,7 @@ def make_td(): unit="Bohr", verbose=0, ) - mf = EnsembleRKS(mol).set( - xc="SVWN", + mf = DZ0SCF(mol, xc="SVWN").set( conv_tol=1e-12, conv_tol_grad=1e-9, max_cycle=100, @@ -31,7 +30,7 @@ def make_td(): mf.grids.level = 0 mf.kernel() if not mf.converged: - raise RuntimeError("EnsembleRKS test reference did not converge") + raise RuntimeError("Dz0SCF test reference did not converge") tdobj = NTTDA(mf).set( deltaS=0, nstates=2, @@ -48,12 +47,12 @@ def test_finite_difference_total_energy_gradient(self): fixed_grid=False, root_overlap_tol=0.5, ) - self.assertFalse(tdobj.Gradients().fixed_grid) + self.assertTrue(tdobj.Gradients().fixed_grid) from nest.grad.nttda import _displaced_reference displaced = _displaced_reference(tdobj._scf, tdobj.mol.copy(), False) - self.assertIsInstance(displaced, EnsembleRKS) - self.assertEqual(displaced.nopen, tdobj._scf.nopen) + self.assertTrue( + getattr(displaced, "is_average_occupation_reference", False)) result = gradient.kernel( state=2, method="finite_diff", @@ -63,8 +62,8 @@ def test_finite_difference_total_energy_gradient(self): self.assertTrue(np.all(np.isfinite(result))) np.testing.assert_allclose( result, - [[0, 0, 0.09285896198862886], - [0, 0, -0.09285896198774068]], + [[0, 0, 0.091524381309104896], + [0, 0, -0.091524381308882852]], atol=2e-5, rtol=0, ) @@ -91,8 +90,7 @@ def test_all_spin_channels_have_a_finite_difference_path(self): unit="Bohr", verbose=0, ) - mf = EnsembleRKS(mol).set( - xc="SVWN", + mf = DZ0SCF(mol, xc="SVWN").set( conv_tol=1e-11, max_cycle=150, verbose=0, @@ -130,8 +128,7 @@ def test_spin_lowering_analytic_matches_fixed_grid_finite_difference(self): unit="Bohr", verbose=0, ) - mf = EnsembleRKS(mol).set( - xc="SVWN", + mf = DZ0SCF(mol, xc="SVWN").set( conv_tol=1e-12, conv_tol_grad=1e-9, max_cycle=150, @@ -181,8 +178,7 @@ def test_representative_functional_families(self): unit="Bohr", verbose=0, ) - mf = EnsembleRKS(mol).set( - xc=xc, + mf = DZ0SCF(mol, xc=xc).set( conv_tol=1e-12, conv_tol_grad=1e-9, max_cycle=150, diff --git a/src/nest/grad/tests/test_nttda_ensemble_response.py b/src/nest/grad/tests/test_nttda_dz0scf_response.py similarity index 73% rename from src/nest/grad/tests/test_nttda_ensemble_response.py rename to src/nest/grad/tests/test_nttda_dz0scf_response.py index 6a82860..56cfe9b 100644 --- a/src/nest/grad/tests/test_nttda_ensemble_response.py +++ b/src/nest/grad/tests/test_nttda_dz0scf_response.py @@ -1,15 +1,12 @@ #!/usr/bin/env python -"""Orbital-response checks for average-occupation EnsembleRKS gradients.""" +"""Orbital-response checks for average-occupation Dz0SCF gradients.""" import unittest -from pathlib import Path import numpy as np from scipy.linalg import expm - - from pyscf import gto @@ -27,11 +24,11 @@ grad_elec as spin_lowering_grad_elec, spin_lowering_ledger_scalar, ) -from nest.ensemble_rks import EnsembleRKS # noqa: E402 +from nest.dz0scf import DZ0SCF # noqa: E402 from nest.nttda import NTTDA # noqa: E402 -class EnsembleOrbitalResponse(unittest.TestCase): +class Dz0SCFOrbitalResponse(unittest.TestCase): @staticmethod def make_reference(): mol = gto.M( @@ -42,8 +39,7 @@ def make_reference(): unit="Bohr", verbose=0, ) - mf = EnsembleRKS(mol).set( - xc="SVWN", + mf = DZ0SCF(mol, xc="SVWN").set( conv_tol=1e-13, conv_tol_grad=1e-10, max_cycle=100, @@ -52,7 +48,7 @@ def make_reference(): mf.grids.level = 0 mf.kernel() if not mf.converged: - raise RuntimeError("EnsembleRKS reference did not converge") + raise RuntimeError("Dz0SCF reference did not converge") return mf @staticmethod @@ -64,8 +60,7 @@ def make_spin_one_reference(): unit="Bohr", verbose=0, ) - mf = EnsembleRKS(mol).set( - xc="SVWN", + mf = DZ0SCF(mol, xc="SVWN").set( conv_tol=1e-12, conv_tol_grad=1e-9, max_cycle=150, @@ -74,9 +69,65 @@ def make_spin_one_reference(): mf.grids.level = 0 mf.kernel() if not mf.converged: - raise RuntimeError("spin-one EnsembleRKS reference did not converge") + raise RuntimeError("spin-one Dz0SCF reference did not converge") return mf + @staticmethod + def _common_fock(mf, mo_coeff, mo_occ): + density = mf.make_rdm1(mo_coeff, mo_occ) + veff = mf.get_veff(mf.mol, density) + return mf.get_hcore() + veff[0] + + def test_charge_response_tracks_a_reused_reference(self): + mol_a = gto.M( + atom="Li 0 0 0; H 0 0 3.0", + basis="sto-3g", + charge=1, + spin=1, + unit="Bohr", + verbose=0, + ) + mf = DZ0SCF(mol_a, xc="SVWN").set( + conv_tol=1e-12, + conv_tol_grad=1e-9, + max_cycle=100, + verbose=0, + ) + mf.grids.level = 0 + mf.kernel() + rng = np.random.default_rng(7) + density = rng.standard_normal((mol_a.nao_nr(),) * 2) + density = 0.5 * (density + density.T) + response_a = mf.gen_response(hermi=1)(density) + + mol_b = gto.M( + atom="Li 0 0 0; H 0 0 3.2", + basis="sto-3g", + charge=1, + spin=1, + unit="Bohr", + verbose=0, + ) + mf.reset(mol_b) + mf.grids.level = 0 + mf.kernel() + response_b = mf.gen_response(hermi=1)(density) + + fresh = DZ0SCF(mol_b, xc="SVWN").set( + conv_tol=1e-12, + conv_tol_grad=1e-9, + max_cycle=100, + verbose=0, + ) + fresh.grids.level = 0 + fresh.kernel() + response_fresh = fresh.gen_response(hermi=1)(density) + + np.testing.assert_allclose( + response_b, response_fresh, atol=1e-8, rtol=0, + ) + self.assertGreater(np.max(np.abs(response_b - response_a)), 1e-4) + def test_hessian_action_matches_orbital_rotation_finite_difference(self): mf = self.make_reference() tdobj = NTTDA(mf) @@ -95,8 +146,7 @@ def test_hessian_action_matches_orbital_rotation_finite_difference(self): gradients = [] for sign in (1.0, -1.0): displaced_mo = mo @ expm(sign * step * kappa) - density = mf.make_rdm1(displaced_mo, occ) - fock = mf.get_hcore() + mf.get_veff(mf.mol, density) + fock = self._common_fock(mf, displaced_mo, occ) gradients.append(mf.get_grad(displaced_mo, occ, fock)) finite_difference = (gradients[0] - gradients[1]) / (2.0 * step) diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index 9568e6d..d71c91b 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -30,33 +30,32 @@ from pyscf.tdscf._lr_eig import eigh as lr_eigh from pyscf import symm from pyscf.data import nist -from nest.ensemble_rks import EnsembleRKS MO_BASE = getattr(__config__, 'MO_BASE', 1) MO_GRID_FXC1 = True -def _is_ensemble_reference(mf): - return isinstance(mf, EnsembleRKS) +def _is_average_occupation_reference(mf): + return getattr(mf, "is_average_occupation_reference", False) def _require_nttda_reference(mf): supported = ( dft.roks.ROKS, dft.rks_symm.SymAdaptedROKS, - EnsembleRKS, ) if not isinstance(mf, supported): - raise TypeError("NTTDA response requires ROKS or EnsembleRKS reference") + raise TypeError("NTTDA response requires a ROKS or Dz0SCF reference") def _reference_fock0(mf, nobeta): """Return the common Fock used by the NTTDA orbital terms. - EnsembleRKS is self-consistent in ``F0[D/2,D/2]``. For ROKS, retain the - two historical choices controlled by ``nobeta``. + A Dz0SCF (average-occupation) reference is self-consistent in + ``F0[D/2,D/2]``. For ROKS, retain the two historical choices controlled + by ``nobeta``. """ - if _is_ensemble_reference(mf): + if _is_average_occupation_reference(mf): return np.asarray(mf.get_fock()) if nobeta: dma, dmb = mf.make_rdm1() @@ -1466,4 +1465,3 @@ def oscillator_strength(tdobj, ref=1, state=None): dft.roks.ROKS.NTTDA = lib.class_as_method(NTTDA) dft.rks_symm.SymAdaptedROKS.NTTDA = lib.class_as_method(NTTDA) -EnsembleRKS.NTTDA = lib.class_as_method(NTTDA) diff --git a/src/nest/nttda/tests/test_nttda_ensemble.py b/src/nest/nttda/tests/test_nttda_dz0scf.py similarity index 76% rename from src/nest/nttda/tests/test_nttda_ensemble.py rename to src/nest/nttda/tests/test_nttda_dz0scf.py index ea4f240..1fbb422 100644 --- a/src/nest/nttda/tests/test_nttda_ensemble.py +++ b/src/nest/nttda/tests/test_nttda_dz0scf.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -"""Acceptance tests for EnsembleRKS-based NTTDA energies.""" +"""Acceptance tests for Dz0SCF-based NTTDA energies.""" import unittest @@ -7,11 +7,11 @@ from pyscf import gto from pyscf.scf import hf -from nest.ensemble_rks import EnsembleRKS +from nest.dz0scf import DZ0SCF from nest.nttda import NTTDA -class EnsembleRKSReference(unittest.TestCase): +class Dz0SCFReference(unittest.TestCase): @staticmethod def lithium_hydride_cation(): return gto.M( @@ -24,24 +24,23 @@ def lithium_hydride_cation(): ) def make_reference(self, xc="SVWN"): - mf = EnsembleRKS(self.lithium_hydride_cation()).set( - xc=xc, - conv_tol=1e-12, - conv_tol_grad=1e-9, - max_cycle=100, - verbose=0, - ) + mf = DZ0SCF(self.lithium_hydride_cation(), xc=xc) + mf.conv_tol = 1e-12 + mf.conv_tol_grad = 1e-9 + mf.max_cycle = 100 + mf.verbose = 0 mf.grids.level = 0 mf.kernel() self.assertTrue(mf.converged) return mf - def test_fixed_occupations_define_a_spin_unpolarized_ensemble(self): + def test_fixed_occupations_define_a_spin_unpolarized_reference(self): mf = self.make_reference() np.testing.assert_array_equal(mf.mo_occ, [2, 1, 0, 0, 0, 0]) self.assertAlmostEqual(mf.mo_occ.sum(), mf.mol.nelectron) - dm = mf.make_rdm1() + mo = np.asarray(mf.mo_coeff) + dm = (mo * np.asarray(mf.mo_occ)) @ mo.conj().T dma, dmb = mf.make_rdm1s() np.testing.assert_allclose(dma, dmb, atol=0, rtol=0) np.testing.assert_allclose(dma + dmb, dm, atol=1e-14, rtol=0) @@ -65,11 +64,7 @@ def test_fixed_occupations_define_a_spin_unpolarized_ensemble(self): rtol=0, ) - def test_explicit_open_shell_count_must_match_molecular_spin(self): - with self.assertRaisesRegex(ValueError, "nopen.*mol.spin"): - EnsembleRKS(self.lithium_hydride_cation(), nopen=3) - - def test_nttda_energy_uses_the_ensemble_fock_for_both_nobeta_values(self): + def test_nttda_energy_is_independent_of_the_nobeta_flag(self): mf = self.make_reference() energies = [] for nobeta in (False, True): @@ -86,16 +81,18 @@ def test_nttda_energy_uses_the_ensemble_fock_for_both_nobeta_values(self): energies.append(tdobj.e) np.testing.assert_allclose(energies[0], energies[1], atol=1e-12, rtol=0) - def test_reference_energy_is_the_stationary_average_occupation_energy(self): + def test_reference_energy_is_the_high_spin_roks_energy(self): mf = self.make_reference() self.assertEqual( mf.reference_energy_semantics, - "average_occupation_ensemble_rks_energy", + "high_spin_roks_energy_on_dz0_orbitals", + ) + self.assertFalse(mf.reference_energy_stationary) + self.assertAlmostEqual( + mf.reference_energy(), mf.high_spin_energy(), places=14, ) - self.assertTrue(mf.reference_energy_stationary) - self.assertAlmostEqual(mf.reference_energy(), mf.e_tot, places=14) - def test_nttda_total_energies_use_the_selected_reference_energy(self): + def test_nttda_total_energies_use_the_reference_energy(self): mf = self.make_reference() tdobj = NTTDA(mf).set( deltaS=0, @@ -105,10 +102,10 @@ def test_nttda_total_energies_use_the_selected_reference_energy(self): verbose=0, ).run() - self.assertAlmostEqual(tdobj.reference_energy(), mf.e_tot) + self.assertAlmostEqual(tdobj.reference_energy(), mf.high_spin_energy()) np.testing.assert_allclose( tdobj.total_energies(), - mf.e_tot + tdobj.e, + mf.high_spin_energy() + tdobj.e, atol=1e-13, rtol=0, ) @@ -135,12 +132,10 @@ def test_all_three_spin_channels(self): unit="Bohr", verbose=0, ) - mf = EnsembleRKS(mol).set( - xc="SVWN", - conv_tol=1e-12, - max_cycle=150, - verbose=0, - ) + mf = DZ0SCF(mol, xc="SVWN") + mf.conv_tol = 1e-12 + mf.max_cycle = 150 + mf.verbose = 0 mf.grids.level = 0 mf.kernel() self.assertTrue(mf.converged) From ff1b136a6777a3988490855605eb0d5a8a1c0b30 Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Sat, 12 Sep 2026 10:50:07 +0800 Subject: [PATCH 08/16] docs: add Dz0SCF reference and NTTDA gradient examples --- examples/grad/02_dz0scf_grad.py | 58 +++++++++++++++++++++ examples/nttda/02_nttda_dz0scf_grad.py | 71 ++++++++++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100644 examples/grad/02_dz0scf_grad.py create mode 100644 examples/nttda/02_nttda_dz0scf_grad.py diff --git a/examples/grad/02_dz0scf_grad.py b/examples/grad/02_dz0scf_grad.py new file mode 100644 index 0000000..4e84998 --- /dev/null +++ b/examples/grad/02_dz0scf_grad.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +''' +Analytic nuclear gradient of the Dz0SCF (average-occupation) reference. + +Dz0SCF drives one set of orbitals with average occupations 2/1/0 for a +high-spin open-shell reference, and takes the high-spin ROKS energy evaluated +on those orbitals as the reference energy. ``nuc_grad_method()`` returns the +analytic gradient of that reference energy, which is the zero state used by +the NTTDA excited-state gradients (see examples/nttda/02_nttda_dz0scf_grad.py). + +The reference is non-stationary on the average-occupation orbitals, so the +driver solves a Z-vector equation for the orbital response; a diffuse enough +integration grid is required for the force sum to vanish. +''' + +from pyscf import gto +from nest import dz0scf # necessary import +from nest.dz0scf import DZ0SCF + +atom = ''' +N 0.000000 -0.040000 0.000000 +H 0.000000 0.780000 0.590000 +H 0.000000 -0.860000 0.520000 +''' +mol = gto.M(atom=atom, charge=0, spin=1, basis='6-31g', verbose=3) +fun = 'PBE' # try also 'SVWN', 'B3LYP', 'M06-2X', etc. +mf = DZ0SCF(mol, xc=fun) +mf.conv_tol = 1e-12 +mf.conv_tol_grad = 1e-9 +mf.max_cycle = 120 +mf.grids.level = 5 # dense grid: the force sum is grid-sensitive +mf.grids.prune = None +mf.small_rho_cutoff = 0.0 +mf.kernel() + +print('Dz0SCF reference energy: %.12f' % mf.high_spin_energy()) + +grad = mf.nuc_grad_method().kernel() +print('Analytic reference gradient (Eh/Bohr):\n', grad) +print('Force sum (should be ~0):\n', grad.sum(axis=0)) + +# Gradients can also be restricted to selected atoms: +grad_n = mf.nuc_grad_method().kernel(atmlst=[0]) +print('Gradient on the nitrogen atom only:\n', grad_n) diff --git a/examples/nttda/02_nttda_dz0scf_grad.py b/examples/nttda/02_nttda_dz0scf_grad.py new file mode 100644 index 0000000..1f6154c --- /dev/null +++ b/examples/nttda/02_nttda_dz0scf_grad.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +''' +NTTDA excited-state gradients on a Dz0SCF (average-occupation) reference. + +A Dz0SCF reference gives a common average-occupation orbital set for every +spin channel. NTTDA built on that reference can target the same-spin channel +(``deltaS=0``) and the spin-lowering channel (``deltaS=-1``); the total state +energy is the high-spin reference energy plus the NTTDA excitation energy, so +its gradient is the sum of the reference gradient and the excitation-energy +gradient. + +``td.Gradients().kernel(state=n)`` returns the analytic gradient of +``E_reference + omega_n`` for state ``n`` (1 for the lowest root); ``state=0`` +returns the reference gradient. Analytic gradients are available for +``deltaS = -1`` and ``0``; ``deltaS = +1`` is not implemented. +''' + +from pyscf import gto +from nest import dz0scf, nttda # necessary imports +from nest.dz0scf import DZ0SCF + +atom = ''' +C 0.020000 -0.030000 0.010000 +H -0.020000 0.800000 0.620000 +H 0.030000 -0.910000 0.500000 +''' +mol = gto.M(atom=atom, charge=0, spin=2, basis='sto-3g', verbose=3) +fun = 'B3LYP' +mf = DZ0SCF(mol, xc=fun) +mf.conv_tol = 1e-12 +mf.conv_tol_grad = 1e-9 +mf.max_cycle = 150 +mf.grids.level = 5 # dense grid: the force sum is grid-sensitive +mf.grids.prune = None +mf.small_rho_cutoff = 0.0 +mf.kernel() + +for delta_s in (-1, 0): + td = mf.NTTDA().set( + deltaS=delta_s, # Sf = Si + deltaS + nstates=3, + conv_tol=1e-9, + max_cycle=200, + verbose=0, + ).run() + print('deltaS = %+d' % delta_s) + print(' NTTDA excitation energies:', td.e) + print(' total energies (E_ref + omega):', td.total_energies()) + + grad = td.Gradients().kernel(state=1) + print(' state-1 analytic gradient (Eh/Bohr):\n', grad) + print(' force sum (should be ~0):\n', grad.sum(axis=0)) + + ref_grad = td.Gradients().kernel(state=0) + print(' reference (state-0) gradient:\n', ref_grad) + print(' excitation-only contribution (state 1 - state 0):\n', + grad - ref_grad) From 307cf7724a86dae48df03ffa49f9b83a6dc6f80f Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Tue, 22 Sep 2026 14:21:08 +0800 Subject: [PATCH 09/16] fix: retain upstream NTTDA transitions and response eigensolver Remove the stale spin-lowering-only transition methods that shadow the upstream all-channel implementation after merging main. Restore nest._lr_eig so oscillator strengths retain upstream convergence behavior. The existing all-channel oscillator-strength regression passes without changing its tolerance. --- .gitignore | 3 + src/nest/nttda/nttda.py | 153 +--------------------------------------- 2 files changed, 4 insertions(+), 152 deletions(-) diff --git a/.gitignore b/.gitignore index 2aa776c..093a151 100644 --- a/.gitignore +++ b/.gitignore @@ -220,3 +220,6 @@ __marimo__/ # Streamlit .streamlit/secrets.toml + +# Quarantined local research-code snapshots +.codex_trash/ diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index b8fc302..5fe1bbb 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -27,7 +27,7 @@ from pyscf.dft.gen_grid import NBINS from pyscf import __config__ from pyscf.dft.numint import _scale_ao_sparse, _dot_ao_ao_sparse, _dot_ao_dm_sparse, _contract_rho_sparse -from pyscf.tdscf._lr_eig import eigh as lr_eigh +from nest._lr_eig import eigh as lr_eigh from pyscf import symm from pyscf.data import nist @@ -1490,156 +1490,5 @@ def analyze(tdobj, verbose=None): NTTDA.analyze = analyze -def transition_dipole(tdobj, ref=1, state=None): - """ - Transition dipole moments between NTTDA DeltaS = -1 excited states. - - Parameters - ---------- - tdobj : NTTDA object - Need tdobj.deltaS == -1 and tdobj.xy. - ref : int - 1-based reference excited-state index. - state : int or array-like or None - 1-based target excited-state index/indices. If None, all states except ref. - - Returns - ------- - pol : ndarray, shape (nstates, 3) - in length gauge. - """ - mf = tdobj._scf - mol = mf.mol - - deltaS = getattr(tdobj, "deltaS", getattr(tdobj, "DeltaS", None)) - assert deltaS == -1 - - s = (mol.nelec[0] - mol.nelec[1]) * 0.5 - assert s >= 1 - - mo_coeff = mf.mo_coeff - assert mf.mo_occ.ndim == 1 - csidx, osidx, vsidx = _orbital_indices(tdobj) - - ncs = len(csidx) - nos = len(osidx) - nvs = len(vsidx) - nocc = ncs + nos - nvir = nos + nvs - nmo = mo_coeff.shape[1] - - if state is None: - states = np.arange(tdobj.nstates) + 1 - else: - states = np.atleast_1d(state).astype(int) - - states = states[states != ref] - ref0 = ref - 1 - states0 = states - 1 - - def get_x(i): - x = tdobj.xy[i] - if isinstance(x, (tuple, list)): - x = x[0] - return np.asarray(x).reshape(nocc, nvir) - - mx = get_x(ref0) - nxs = np.asarray([get_x(i) for i in states0]) - - # x_co: j v, x_cv: j b, x_oo: w v, x_ov: v b - m_co = mx[:ncs, :nos].conj() - m_cv = mx[:ncs, nos:].conj() - m_oo = mx[ncs:, :nos].conj() - m_ov = mx[ncs:, nos:].conj() - - n_co = nxs[:, :ncs, :nos] - n_cv = nxs[:, :ncs, nos:] - n_oo = nxs[:, ncs:, :nos] - n_ov = nxs[:, ncs:, nos:] - - nstate = len(states0) - gamma = np.zeros((nstate, nmo, nmo), dtype=np.result_type(mx, nxs, complex)) - ist = np.arange(nstate) - - def add(rows, cols, block): - gamma[np.ix_(ist, rows, cols)] += block - - a = np.sqrt(2 * s / (2 * s - 1)) - b = 1 / np.sqrt(2 * s * (2 * s - 1)) - # c = np.sqrt((2 * s - 1) / (2 * s)) - f = np.sqrt((2 * s + 1) / (2 * s)) - - tr_moo = np.einsum("tt->", m_oo) - tr_noo = np.einsum("ntt->n", n_oo) - - # OO-OO - add(osidx, osidx, lib.einsum("ut,nuv->ntv", m_oo, n_oo)) - add(osidx, osidx, -lib.einsum("ut,nwt->nwu", m_oo, n_oo)) - - # CO-CO - add(osidx, osidx, lib.einsum("iu,niv->nuv", m_co, n_co)) - add(csidx, csidx, -lib.einsum("iu,nju->nji", m_co, n_co)) - - # CV-CV - add(csidx, csidx, -lib.einsum("ia,nja->nji", m_cv, n_cv)) - add(vsidx, vsidx, lib.einsum("ia,nib->nab", m_cv, n_cv)) - - # OV-OV - add(osidx, osidx, -lib.einsum("ua,nva->nvu", m_ov, n_ov)) - add(vsidx, vsidx, lib.einsum("ua,nub->nab", m_ov, n_ov)) - - # OO-CO and CO-OO - add(csidx, osidx, - -a * lib.einsum("ut,njt->nju", m_oo, n_co) - + b * tr_moo * n_co) - - add(osidx, csidx, - -a * lib.einsum("iu,nwu->nwi", m_co, n_oo) - + b * lib.einsum("iu,n->nui", m_co, tr_noo)) - - # OO-OV and OV-OO - add(osidx, vsidx, - a * lib.einsum("ut,nub->ntb", m_oo, n_ov) - - b * tr_moo * n_ov) - - add(vsidx, osidx, - a * lib.einsum("ua,nuv->nav", m_ov, n_oo) - - b * lib.einsum("ua,n->nau", m_ov, tr_noo)) - - # CO-CV and CV-CO - add(osidx, vsidx, f * lib.einsum("iu,nib->nub", m_co, n_cv)) - add(vsidx, osidx, f * lib.einsum("ia,niv->nav", m_cv, n_co)) - - # OV-CV and CV-OV - add(csidx, osidx, -f * lib.einsum("ua,nja->nju", m_ov, n_cv)) - add(osidx, csidx, -f * lib.einsum("ia,nva->nvi", m_cv, n_ov)) - - dip_ao = mol.intor_symmetric("int1e_r", comp=3) - dip_mo = lib.einsum("up,xuv,vq->xpq", mo_coeff.conj(), dip_ao, mo_coeff) - - pol = lib.einsum("npq,xpq->nx", gamma, dip_mo) - return pol.real - -def oscillator_strength(tdobj, ref=1, state=None): - if state is None: - states = np.arange(tdobj.nstates) + 1 - else: - states = np.atleast_1d(state) - states = states[states != ref] - - trans_dip = transition_dipole(tdobj, ref, states) - - ref -= 1 - states -= 1 - es = tdobj.e[states] - tdobj.e[ref] - f = (2./3.) * lib.einsum('n,nx,nx->n', es, trans_dip.conj(), trans_dip).real - if isinstance(state, int): - return f[0] - else: - return f - -NTTDA.transition_dipole = transition_dipole -NTTDA.oscillator_strength = oscillator_strength - dft.roks.ROKS.NTTDA = lib.class_as_method(NTTDA) dft.rks_symm.SymAdaptedROKS.NTTDA = lib.class_as_method(NTTDA) From 79eacfe742553cb8125de77e12c85a296d18b871 Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Tue, 22 Sep 2026 14:23:43 +0800 Subject: [PATCH 10/16] refactor: share NTTDA gradient response and quadrature Keep spin-channel amplitudes and coefficients explicit while sharing orbital partitions, reference responses, J/K derivative batching, Fz exchange and adjoint assembly. Consolidate identical GGA and meta-GGA quadrature flows with explicit density dimensions and kernel dispatch. Remove 1216 net lines from the gradient package; retain existing numerical checks and tolerances. --- src/nest/grad/nttda/common.py | 747 ++++++++++++++++++++ src/nest/grad/nttda/delta_s_minus_one.py | 861 +---------------------- src/nest/grad/nttda/delta_s_zero.py | 860 +--------------------- src/nest/grad/nttda/xc.py | 436 ++---------- 4 files changed, 844 insertions(+), 2060 deletions(-) create mode 100644 src/nest/grad/nttda/common.py diff --git a/src/nest/grad/nttda/common.py b/src/nest/grad/nttda/common.py new file mode 100644 index 0000000..5f37211 --- /dev/null +++ b/src/nest/grad/nttda/common.py @@ -0,0 +1,747 @@ +"""Shared orbital, reference-response and J/K derivative operations for NTTDA. + +Spin-channel coefficients and amplitude projections stay in the channel modules. +""" + +from dataclasses import dataclass + +import numpy as np +from pyscf import dft, lib +from pyscf.grad import rhf as rhf_grad +from nest.nttda import nttda as nttda_mod + +from . import xc as xc_backend +from .roks import finish_gradient + + +def assemble_gradient( + gradient_driver, tdobj, channel_data, probes, fock_q, response_q, + atmlst=None, tolerance=1e-12, max_cycle=None): + """Assemble channel projections, XC/J/K derivatives and the adjoint. + + ``response_q`` contains only the hybrid/RSH part; semilocal response is + added here. The direct and Z-vector probes share one J/K derivative batch. + """ + mf = tdobj._scf + xctype = mf._numint._xc_type(mf.xc) + atmlst = tuple(range(tdobj.mol.natm) if atmlst is None else atmlst) + spaces, _amplitudes, densities, _blocks, response_terms = channel_data + p0, pz = probes + fock_alpha, fock_beta = fock_q + response_alpha, response_beta = response_q + m_matrix = fock_alpha + fock_beta + response_alpha + response_beta + ledger = _JKDerivativeLedger() + slots = ("direct", "zvector") + direct = response_direct_hfx( + gradient_driver, tdobj, densities, response_terms, + atmlst=atmlst, jk_ledger=ledger, output_slot=slots[0], + ) + + # ROKS/HF includes Fz in the spin-resolved Fock probes. All other + # references use charge-only probes and differentiate Fz separately. + spin_fock = xctype == "HF" and not getattr(mf, "is_average_occupation_reference", False) + direct_fock_probes = ( + (0.5 * (p0 + pz), 0.5 * (p0 - pz)) if spin_fock + else (0.5 * p0, 0.5 * p0) + ) + if xctype != "HF": + if xctype == "LDA": + response_builder = xc_backend.lda_response_terms + fockz_builder = xc_backend.lda_fockz_terms + nobeta_builder = xc_backend.lda_nobeta_reference_q + elif xctype in ("GGA", "MGGA"): + response_builder = xc_backend.semilocal_response_terms + fockz_builder = xc_backend.semilocal_fockz_terms + nobeta_builder = xc_backend.semilocal_nobeta_reference_q + else: + raise NotImplementedError("NTTDA gradient does not support XC type %s" % xctype) + for terms in ( + response_builder(gradient_driver, tdobj, channel_data, atmlst=atmlst), + fockz_builder(gradient_driver, tdobj, spaces, pz, atmlst=atmlst)): + m_matrix += terms.q_alpha + terms.q_beta + direct += terms.direct + common_alpha, common_beta = nobeta_builder(tdobj, p0) + m_matrix += common_alpha + common_beta + if not spin_fock: + terms = fockz_hfx_terms( + gradient_driver, tdobj, pz, atmlst=atmlst, + jk_ledger=ledger, output_slot=slots[0], + ) + m_matrix += terms.q_alpha + terms.q_beta + direct += terms.direct + + def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): + if xctype == "HF": + local = spin_fock_direct_hf( + driver, obj, p_alpha, p_beta, atmlst=atmlst, + jk_ledger=ledger, output_slots=slots, + ) + else: + local = spin_fock_direct_dft( + driver, obj, p_alpha, p_beta, atmlst=atmlst, nobeta_p0=p0, + jk_ledger=ledger, output_slots=slots, + ) + contractions = ledger.contract(driver, obj.mol, atmlst, slots=slots) + for index, slot in enumerate(slots): + local[index] += contractions[slot] + return local + + return finish_gradient( + gradient_driver, tdobj, m_matrix, direct, atmlst, tolerance, + max_cycle, fock_direct, direct_fock_probes=direct_fock_probes, + ) + +@dataclass(frozen=True) +class OrbitalSpaces: + """Closed, open, and virtual spatial-orbital partitions.""" + + closed: np.ndarray + open: np.ndarray + virtual: np.ndarray + c_closed: np.ndarray + c_open: np.ndarray + c_virtual: np.ndarray + + @property + def spin(self): + return 0.5 * len(self.open) + + +def orbital_spaces(tdobj): + """Return the ROKS ``C/O/V`` orbital partition used by NTTDA.""" + mf = tdobj._scf + occ = np.asarray(mf.mo_occ) + if occ.ndim != 1: + raise ValueError("NTTDA gradients require spatial ROKS orbitals") + closed = np.flatnonzero(occ == 2) + open_ = np.flatnonzero(occ == 1) + virtual = np.flatnonzero(occ == 0) + coeff = np.asarray(mf.mo_coeff) + return OrbitalSpaces( + closed=closed, + open=open_, + virtual=virtual, + c_closed=coeff[:, closed], + c_open=coeff[:, open_], + c_virtual=coeff[:, virtual], + ) + + +def pair_density(c_left, coefficient, c_right): + """Build ``C_left coefficient C_right^T`` without symmetrizing it.""" + return c_left @ np.asarray(coefficient) @ c_right.conj().T + + +@dataclass(frozen=True) +class FockProjection: + """One scalar term ``Tr[P (weight_f0 F0 + weight_fz Fz)]``.""" + + name: str + left_indices: np.ndarray + left_orbitals: np.ndarray + coefficient: np.ndarray + right_indices: np.ndarray + right_orbitals: np.ndarray + weight_f0: float + weight_fz: float + + def density(self): + return pair_density( + self.left_orbitals, self.coefficient, self.right_orbitals, + ) + + +def _fock_response_q(tdobj, p_alpha, p_beta): + """Reference-density derivative of a spin-resolved Fock scalar.""" + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + if getattr(mf, "is_average_occupation_reference", False): + occupation = np.asarray(mf.mo_occ) + probe = np.asarray(p_alpha) + np.asarray(p_beta) + potential = mf.gen_response(hermi=0)(probe.T) + q_total = ( + mo.conj().T @ (potential + potential.T) @ mo + ) * occupation[None, :] + return 0.5 * q_total, 0.5 * q_total + occ_alpha = (np.asarray(mf.mo_occ) > 0).astype(float) + occ_beta = (np.asarray(mf.mo_occ) == 2).astype(float) + if (isinstance(mf, dft.KohnShamDFT) + and mf._numint._xc_type(mf.xc) != "HF"): + unrestricted = mf.to_uks() + unrestricted.verbose = 0 + v_alpha, v_beta = unrestricted.gen_response(hermi=0)( + np.asarray((p_alpha.T, p_beta.T)) + ) + else: + p_total = p_alpha + p_beta + coulomb = mf.get_j(mf.mol, p_total.T, hermi=0) + v_alpha = coulomb - mf.get_k(mf.mol, p_alpha.T, hermi=0) + v_beta = coulomb - mf.get_k(mf.mol, p_beta.T, hermi=0) + q_alpha = ( + mo.conj().T @ (v_alpha + v_alpha.T) @ mo + ) * occ_alpha[None, :] + q_beta = ( + mo.conj().T @ (v_beta + v_beta.T) @ mo + ) * occ_beta[None, :] + return q_alpha, q_beta + + +@dataclass(frozen=True) +class ResponseTerm: + """Directed response term from one source density to one target block.""" + + target: str + source: str + vref0: float + vref1: float + + +def _fxc_reference(tdobj): + mf = tdobj._scf + ni = mf._numint + fxc = ni.cache_xc_kernel( + mf.mol, mf.grids, mf.xc, mf.mo_coeff, mf.mo_occ, 1, + )[2] + return 0.5 * ( + fxc[0, :, 0] - fxc[0, :, 1] + - fxc[1, :, 0] + fxc[1, :, 1] + ) + + +def _apply_reference_responses(tdobj, densities, max_memory=None): + """Return separate ``vref0`` and ``vref1`` actions for each density.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if max_memory is None: + max_memory = tdobj.max_memory + labels = tuple(densities) + dms = np.asarray([densities[label] for label in labels]) + xctype = ni._xc_type(mf.xc) + if xctype == "HF": + vref0 = np.zeros_like(dms) + vref1 = np.zeros_like(dms) + else: + fxc_ref = _fxc_reference(tdobj) + vref0 = ni.nr_rks_fxc( + mol, mf.grids, mf.xc, None, dms, 0, 0, + None, None, fxc_ref, max_memory=max_memory, + ) + if xctype == "LDA": + vref1 = ni.nr_rks_fxc( + mol, mf.grids, mf.xc, None, dms, 0, 0, + None, None, fxc_ref, max_memory=max_memory, + ) + elif xctype == "GGA": + vref1 = nttda_mod.nr_rks_fxc1_gga( + ni, mol, mf.grids, mf.xc, dms, fxc_ref, + max_memory=max_memory, + ) + elif xctype == "MGGA": + vref1 = nttda_mod.nr_rks_fxc1_mgga( + ni, mol, mf.grids, mf.xc, dms, fxc_ref, + max_memory=max_memory, + ) + else: + raise NotImplementedError( + "NTTDA response does not support XC type %s" % xctype + ) + + omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + if ni.libxc.is_hybrid_xc(mf.xc): + vref0 -= hyb * mf.get_k(mol, dms, hermi=0) + vref1 -= hyb * mf.get_j(mol, dms, hermi=0) + if omega != 0: + scale = alpha - hyb + vref0 -= scale * mf.get_k(mol, dms, hermi=0, omega=omega) + vref1 -= scale * mf.get_j(mol, dms, hermi=0, omega=omega) + return ( + {label: value for label, value in zip(labels, vref0)}, + {label: value for label, value in zip(labels, vref1)}, + ) + + +def _apply_hfx_responses(tdobj, densities): + """Return only the hybrid/RSH J/K portions of ``vref0/vref1``.""" + mf = tdobj._scf + labels = tuple(densities) + dms = np.asarray([densities[label] for label in labels]) + vref0 = np.zeros_like(dms) + vref1 = np.zeros_like(dms) + ni = mf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mf.mol.spin) + if ni.libxc.is_hybrid_xc(mf.xc): + vref0 -= hybrid * mf.get_k(mf.mol, dms, hermi=0) + vref1 -= hybrid * mf.get_j(mf.mol, dms, hermi=0) + if omega != 0: + scale = alpha - hybrid + vref0 -= scale * mf.get_k( + mf.mol, dms, hermi=0, omega=omega, + ) + vref1 -= scale * mf.get_j( + mf.mol, dms, hermi=0, omega=omega, + ) + return ( + {label: value for label, value in zip(labels, vref0)}, + {label: value for label, value in zip(labels, vref1)}, + ) + + +def _as_derivative_stack(array): + array = np.asarray(array) + if array.ndim == 3: + array = array[None] + return array + + +def _density_key(density): + density = np.asarray(density) + data = density.__array_interface__["data"][0] + return data, density.shape, density.strides, density.dtype.str + + +@dataclass(frozen=True) +class _JKDerivativeTerm: + """One fixed-AO bilinear derivative with a named output slot.""" + + left: np.ndarray + right: np.ndarray + scale: float + omega: float + slot: object + + +class _JKDerivativeLedger: + """Shared scheduler for fixed-AO J/K derivative contractions.""" + + def __init__(self): + self._terms = {"j": [], "k": []} + + def add(self, operator, slot, terms): + self._terms[operator].extend( + _JKDerivativeTerm(left, right, scale, omega, slot) + for left, right, scale, omega in terms + if scale != 0.0 + ) + + def contract(self, gradient_driver, mol, atoms, slots=()): + atoms = tuple(atoms) + shape = (len(atoms), 3) + gradients = {slot: np.zeros(shape) for slot in slots} + for operator in ("j", "k"): + for term in self._terms[operator]: + gradients.setdefault(term.slot, np.zeros(shape)) + _contract_derivative_terms( + gradients, + gradient_driver, + mol, + atoms, + mol.offset_nr_by_atom(), + self._terms[operator], + operator, + ) + return gradients + + +def _term_densities(term, exchange): + left, right = term.left, term.right + if exchange: + return left, right, left.T, right.T + return left, right + + +def _density_batches(terms, exchange, max_memory, nao): + """Group bilinear terms while bounding derivative-potential storage.""" + minimum = 4 if exchange else 2 + bytes_per_density = 4 * nao * nao * np.dtype(float).itemsize + batch_limit = max( + minimum, + int(0.2 * max_memory * 1e6 / bytes_per_density), + ) + batch = [] + keys = set() + for term in terms: + term_keys = { + _density_key(density) + for density in _term_densities(term, exchange) + } + if batch and len(keys | term_keys) > batch_limit: + yield batch + batch = [] + keys = set() + batch.append(term) + keys.update(term_keys) + if batch: + yield batch + + +def _jk_derivative_potentials( + gradient_driver, mol, terms, operator, omega): + exchange = operator == "k" + densities = {} + for term in terms: + for density in _term_densities(term, exchange): + density = np.asarray(density) + densities.setdefault(_density_key(density), density) + keys = tuple(densities) + stack = np.asarray([densities[key] for key in keys]) + if operator == "j": + if omega is None: + values = gradient_driver.get_j(mol, stack, hermi=0) + else: + values = gradient_driver.get_j( + mol, stack, hermi=0, omega=omega, + ) + else: + if omega is None: + values = gradient_driver.get_k(mol, stack, hermi=0) + else: + values = gradient_driver.get_k( + mol, stack, hermi=0, omega=omega, + ) + values = _as_derivative_stack(values) + return dict(zip(keys, values)) + + +def _contract_derivative_terms( + gradients, gradient_driver, mol, atoms, offsets, terms, + operator): + if not atoms: + return + terms_by_omega = {} + for term in terms: + terms_by_omega.setdefault(term.omega, []).append(term) + exchange = operator == "k" + for omega, omega_terms in terms_by_omega.items(): + for batch in _density_batches( + omega_terms, exchange, gradient_driver.max_memory, + mol.nao_nr()): + potentials = _jk_derivative_potentials( + gradient_driver, mol, batch, operator, omega, + ) + for term in batch: + left = np.asarray(term.left) + right = np.asarray(term.right) + right_derivative = potentials[_density_key(right)] + left_derivative = potentials[_density_key(left)] + if exchange: + right_t_derivative = potentials[ + _density_key(right.T) + ] + left_t_derivative = potentials[_density_key(left.T)] + for k, atom in enumerate(atoms): + p0, p1 = offsets[atom][2:] + if exchange: + value = lib.einsum( + "xpq,pq->x", + right_derivative[:, p0:p1, :], + left[p0:p1, :], + ) + value += lib.einsum( + "xqp,pq->x", + right_t_derivative[:, p0:p1, :], + left[:, p0:p1], + ) + value += lib.einsum( + "xpq,pq->x", + left_derivative[:, p0:p1, :], + right[p0:p1, :], + ) + value += lib.einsum( + "xqp,pq->x", + left_t_derivative[:, p0:p1, :], + right[:, p0:p1], + ) + else: + value = lib.einsum( + "xpq,pq->x", + right_derivative[:, p0:p1], + left[p0:p1], + ) + value += lib.einsum( + "xpq,qp->x", + right_derivative[:, p0:p1], + left[:, p0:p1], + ) + value += lib.einsum( + "xpq,pq->x", + left_derivative[:, p0:p1], + right[p0:p1], + ) + value += lib.einsum( + "xpq,qp->x", + left_derivative[:, p0:p1], + right[:, p0:p1], + ) + gradients[term.slot][k] += term.scale * value + + +def _reference_spin_densities(tdobj): + mf = tdobj._scf + if getattr(mf, "is_average_occupation_reference", False): + return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) + mo = np.asarray(mf.mo_coeff) + return ( + mo[:, mf.mo_occ > 0] @ mo[:, mf.mo_occ > 0].T, + mo[:, mf.mo_occ == 2] @ mo[:, mf.mo_occ == 2].T, + ) + + +def _spin_probe_stacks(p_alpha, p_beta): + p_alpha = np.asarray(p_alpha) + p_beta = np.asarray(p_beta) + single_probe = p_alpha.ndim == 2 + if single_probe: + p_alpha = p_alpha[None] + p_beta = p_beta[None] + return p_alpha, p_beta, single_probe + + +def spin_fock_direct_dft( + gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, + nobeta_p0=None, jk_ledger=None, output_slots=None): + """Differentiate one or more ordinary UKS Fock scalar probes. + + The optional ``nobeta_p0`` correction belongs to the first, explicit-direct + probe in the batch. + """ + mf = tdobj._scf + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + p_alpha, p_beta, single_probe = _spin_probe_stacks( + p_alpha, p_beta, + ) + if output_slots is None: + output_slots = tuple(range(len(p_alpha))) + p_total = p_alpha + p_beta + density_alpha, density_beta = _reference_spin_densities(tdobj) + gradient = np.zeros((len(p_alpha), len(atmlst), 3)) + hcore_derivative = rhf_grad.Gradients(mf).hcore_generator(mol) + for k, atom in enumerate(atmlst): + gradient[:, k] += lib.einsum( + "npq,xpq->nx", p_total, hcore_derivative(atom), + ) + ni = mf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + for probe in range(len(p_alpha)): + j_terms = [ + (p_total[probe], density_alpha, 1.0, None), + (p_total[probe], density_beta, 1.0, None), + ] + k_terms = [] + if ni.libxc.is_hybrid_xc(mf.xc): + k_terms.extend(( + (p_alpha[probe], density_alpha, -hybrid, None), + (p_beta[probe], density_beta, -hybrid, None), + )) + if omega != 0: + long_range = -(alpha - hybrid) + k_terms.extend(( + (p_alpha[probe], density_alpha, long_range, omega), + (p_beta[probe], density_beta, long_range, omega), + )) + ledger.add("j", output_slots[probe], j_terms) + ledger.add("k", output_slots[probe], k_terms) + xctype = ni._xc_type(mf.xc) + if xctype == "LDA": + derivative_contractor = xc_backend.contract_lda_vxc_derivative + elif xctype == "GGA": + derivative_contractor = xc_backend.contract_gga_vxc_derivative + elif xctype == "MGGA": + derivative_contractor = xc_backend.contract_mgga_vxc_derivative + else: + raise NotImplementedError( + "ordinary Fock direct derivative is not implemented for %s" % + xctype + ) + if (nobeta_p0 is not None and tdobj.nobeta + and not getattr(mf, "is_average_occupation_reference", False)): + density0 = 0.5 * (density_alpha + density_beta) + actual_probe_alpha = np.array(p_alpha, copy=True) + actual_probe_beta = np.array(p_beta, copy=True) + actual_probe_alpha[0] -= 0.5 * nobeta_p0 + actual_probe_beta[0] -= 0.5 * nobeta_p0 + else: + density0 = None + actual_probe_alpha = p_alpha + actual_probe_beta = p_beta + gradient += derivative_contractor( + mf, + density_alpha, + density_beta, + actual_probe_alpha, + actual_probe_beta, + atmlst=atmlst, + max_memory=gradient_driver.max_memory, + ) + if density0 is not None: + gradient[0] += derivative_contractor( + mf, + density0, + density0, + 0.5 * nobeta_p0, + 0.5 * nobeta_p0, + atmlst=atmlst, + max_memory=gradient_driver.max_memory, + ) + if jk_ledger is None: + contractions = local_ledger.contract( + gradient_driver, mol, atmlst, slots=output_slots, + ) + for probe, slot in enumerate(output_slots): + gradient[probe] += contractions[slot] + return gradient[0] if single_probe else gradient + + +def spin_fock_direct_hf( + gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, + jk_ledger=None, output_slots=None): + """Differentiate one or more spin-resolved HF Fock scalar probes.""" + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + p_alpha, p_beta, single_probe = _spin_probe_stacks( + p_alpha, p_beta, + ) + if output_slots is None: + output_slots = tuple(range(len(p_alpha))) + p_total = p_alpha + p_beta + dm_alpha, dm_beta = _reference_spin_densities(tdobj) + gradient = np.zeros((len(p_alpha), len(atmlst), 3)) + + hcore_derivative = rhf_grad.Gradients(tdobj._scf).hcore_generator(mol) + for k, atom in enumerate(atmlst): + gradient[:, k] += lib.einsum( + "npq,xpq->nx", p_total, hcore_derivative(atom), + ) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + for probe in range(len(p_alpha)): + ledger.add( + "j", output_slots[probe], ( + (p_total[probe], dm_alpha, 1.0, None), + (p_total[probe], dm_beta, 1.0, None), + ), + ) + ledger.add( + "k", output_slots[probe], ( + (p_alpha[probe], dm_alpha, -1.0, None), + (p_beta[probe], dm_beta, -1.0, None), + ), + ) + if jk_ledger is None: + contractions = local_ledger.contract( + gradient_driver, mol, atmlst, slots=output_slots, + ) + for probe, slot in enumerate(output_slots): + gradient[probe] += contractions[slot] + return gradient[0] if single_probe else gradient + + +def response_direct_hfx( + gradient_driver, tdobj, densities, response_terms, atmlst=None, + jk_ledger=None, output_slot=0): + """J/K skeleton derivative for a channel response-term ledger.""" + mol = tdobj.mol + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + gradient = np.zeros((len(atmlst), 3)) + ni = tdobj._scf._numint + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff( + tdobj._scf.xc, mol.spin, + ) + if not ni.libxc.is_hybrid_xc(tdobj._scf.xc): + return gradient + + scales = [(hybrid, None)] + if omega != 0: + scales.append((alpha - hybrid, omega)) + j_terms = [] + k_terms = [] + for term in response_terms: + target = densities[term.target] + source = densities[term.source] + for coefficient, range_omega in scales: + if term.vref0: + k_terms.append(( + target, + source, + -coefficient * term.vref0, + range_omega, + )) + if term.vref1: + j_terms.append(( + target, + source, + -coefficient * term.vref1, + range_omega, + )) + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + ledger.add("j", output_slot, j_terms) + ledger.add("k", output_slot, k_terms) + if jk_ledger is None: + gradient += local_ledger.contract( + gradient_driver, mol, atmlst, slots=(output_slot,), + )[output_slot] + return gradient + + +def fockz_hfx_terms( + gradient_driver, tdobj, pz, atmlst=None, with_direct=True, + jk_ledger=None, output_slot=0): + """Differentiate ``-1/2 Pz:K(D_OO)`` excluding the Pz projection.""" + mf = tdobj._scf + mol = mf.mol + ni = mf._numint + if atmlst is None: + atmlst = range(mol.natm) + atmlst = tuple(atmlst) + mo = np.asarray(mf.mo_coeff) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + direct = np.zeros((len(atmlst), 3)) + if not ni.libxc.is_hybrid_xc(mf.xc): + return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) + + spaces = orbital_spaces(tdobj) + density_open = spaces.c_open @ spaces.c_open.T + omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + scales = [(hybrid, None)] + if omega != 0: + scales.append((alpha - hybrid, omega)) + k_terms = [] + for coefficient, range_omega in scales: + if coefficient == 0.0: + continue + if range_omega is None: + potential = mf.get_k(mol, pz, hermi=0) + else: + potential = mf.get_k( + mol, pz, hermi=0, omega=range_omega, + ) + q_alpha[:, spaces.open] -= 0.5 * coefficient * ( + mo.conj().T @ (potential + potential.T) @ spaces.c_open + ) + if with_direct: + k_terms.append(( + pz, + density_open, + -0.5 * coefficient, + range_omega, + )) + if with_direct: + local_ledger = _JKDerivativeLedger() + ledger = jk_ledger if jk_ledger is not None else local_ledger + ledger.add("k", output_slot, k_terms) + if jk_ledger is None: + direct += local_ledger.contract( + gradient_driver, mol, atmlst, slots=(output_slot,), + )[output_slot] + return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) diff --git a/src/nest/grad/nttda/delta_s_minus_one.py b/src/nest/grad/nttda/delta_s_minus_one.py index 19ff17e..c0fcc84 100644 --- a/src/nest/grad/nttda/delta_s_minus_one.py +++ b/src/nest/grad/nttda/delta_s_minus_one.py @@ -1,40 +1,31 @@ """Analytic gradient for current NTTDA ``deltaS=-1``. -This module owns the complete spin-lowering amplitude, Fock, response, and -AO derivative formulas. It does not import the ``deltaS=0`` channel. +This module owns the spin-lowering amplitudes, Fock projections and response +coefficients. Reference response and AO derivatives are shared between channels. """ from dataclasses import dataclass import numpy as np -from pyscf import dft, lib -from pyscf.grad import rhf as rhf_grad +from pyscf import lib from nest.nttda import nttda as nttda_mod from nest.nttda.nttda import gen_rohf_response_sfd -from . import xc as xc_backend -from .roks import finish_gradient - +from .common import ( + assemble_gradient, + orbital_spaces, + pair_density, + FockProjection, + _fock_response_q, + ResponseTerm, + _apply_reference_responses, + _apply_hfx_responses, +) # Orbital spaces and native amplitudes -@dataclass(frozen=True) -class OrbitalSpaces: - """Closed, open, and virtual spatial-orbital partitions.""" - - closed: np.ndarray - open: np.ndarray - virtual: np.ndarray - c_closed: np.ndarray - c_open: np.ndarray - c_virtual: np.ndarray - - @property - def spin(self): - return 0.5 * len(self.open) - @dataclass(frozen=True) class SpinLoweringAmplitudes: @@ -46,31 +37,6 @@ class SpinLoweringAmplitudes: ov: np.ndarray -def orbital_spaces(tdobj): - """Return the ROKS ``C/O/V`` orbital partition used by NTTDA.""" - mf = tdobj._scf - occ = np.asarray(mf.mo_occ) - if occ.ndim != 1: - raise ValueError("NTTDA gradients require spatial ROKS orbitals") - closed = np.flatnonzero(occ == 2) - open_ = np.flatnonzero(occ == 1) - virtual = np.flatnonzero(occ == 0) - coeff = np.asarray(mf.mo_coeff) - return OrbitalSpaces( - closed=closed, - open=open_, - virtual=virtual, - c_closed=coeff[:, closed], - c_open=coeff[:, open_], - c_virtual=coeff[:, virtual], - ) - - -def pair_density(c_left, coefficient, c_right): - """Build ``C_left coefficient C_right^T`` without symmetrizing it.""" - return c_left @ np.asarray(coefficient) @ c_right.conj().T - - def split_spin_lowering(tdobj, xy): """Split a lowering-channel amplitude into ``CO/CV/OO/OV`` blocks.""" spaces = orbital_spaces(tdobj) @@ -117,168 +83,6 @@ def spin_lowering_block_data(spaces, amplitudes): } -# Channel-local immutable records - -@dataclass(frozen=True) -class FockProjection: - """One scalar term ``Tr[P (weight_f0 F0 + weight_fz Fz)]``.""" - - name: str - left_indices: np.ndarray - left_orbitals: np.ndarray - coefficient: np.ndarray - right_indices: np.ndarray - right_orbitals: np.ndarray - weight_f0: float - weight_fz: float - - def density(self): - return pair_density( - self.left_orbitals, self.coefficient, self.right_orbitals, - ) - - -@dataclass(frozen=True) -class ResponseTerm: - """Directed response term from one source density to one target block.""" - - target: str - source: str - vref0: float - vref1: float - - -# Shared-response evaluators copied into the lowering channel - -def _fxc_reference(tdobj): - mf = tdobj._scf - ni = mf._numint - fxc = ni.cache_xc_kernel( - mf.mol, mf.grids, mf.xc, mf.mo_coeff, mf.mo_occ, 1, - )[2] - return 0.5 * ( - fxc[0, :, 0] - fxc[0, :, 1] - - fxc[1, :, 0] + fxc[1, :, 1] - ) - - -def _apply_reference_responses(tdobj, densities, max_memory=None): - """Return separate ``vref0`` and ``vref1`` actions for each density.""" - mf = tdobj._scf - mol = mf.mol - ni = mf._numint - if max_memory is None: - max_memory = tdobj.max_memory - labels = tuple(densities) - dms = np.asarray([densities[label] for label in labels]) - xctype = ni._xc_type(mf.xc) - if xctype == "HF": - vref0 = np.zeros_like(dms) - vref1 = np.zeros_like(dms) - else: - fxc_ref = _fxc_reference(tdobj) - vref0 = ni.nr_rks_fxc( - mol, mf.grids, mf.xc, None, dms, 0, 0, - None, None, fxc_ref, max_memory=max_memory, - ) - if xctype == "LDA": - vref1 = ni.nr_rks_fxc( - mol, mf.grids, mf.xc, None, dms, 0, 0, - None, None, fxc_ref, max_memory=max_memory, - ) - elif xctype == "GGA": - vref1 = nttda_mod.nr_rks_fxc1_gga( - ni, mol, mf.grids, mf.xc, dms, fxc_ref, - max_memory=max_memory, - ) - elif xctype == "MGGA": - vref1 = nttda_mod.nr_rks_fxc1_mgga( - ni, mol, mf.grids, mf.xc, dms, fxc_ref, - max_memory=max_memory, - ) - else: - raise NotImplementedError( - "NTTDA spin-lowering response does not support XC type %s" % - xctype - ) - - omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) - if ni.libxc.is_hybrid_xc(mf.xc): - vref0 -= hyb * mf.get_k(mol, dms, hermi=0) - vref1 -= hyb * mf.get_j(mol, dms, hermi=0) - if omega != 0: - scale = alpha - hyb - vref0 -= scale * mf.get_k(mol, dms, hermi=0, omega=omega) - vref1 -= scale * mf.get_j(mol, dms, hermi=0, omega=omega) - return ( - {label: value for label, value in zip(labels, vref0)}, - {label: value for label, value in zip(labels, vref1)}, - ) - - -def _apply_hfx_responses(tdobj, densities): - """Return only the hybrid/RSH J/K portions of ``vref0/vref1``.""" - mf = tdobj._scf - labels = tuple(densities) - dms = np.asarray([densities[label] for label in labels]) - vref0 = np.zeros_like(dms) - vref1 = np.zeros_like(dms) - ni = mf._numint - omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mf.mol.spin) - if ni.libxc.is_hybrid_xc(mf.xc): - vref0 -= hybrid * mf.get_k(mf.mol, dms, hermi=0) - vref1 -= hybrid * mf.get_j(mf.mol, dms, hermi=0) - if omega != 0: - scale = alpha - hybrid - vref0 -= scale * mf.get_k( - mf.mol, dms, hermi=0, omega=omega, - ) - vref1 -= scale * mf.get_j( - mf.mol, dms, hermi=0, omega=omega, - ) - return ( - {label: value for label, value in zip(labels, vref0)}, - {label: value for label, value in zip(labels, vref1)}, - ) - - -# Reference Fock response helper - -def _fock_response_q(tdobj, p_alpha, p_beta): - """Reference-density derivative of a spin-resolved Fock scalar.""" - mf = tdobj._scf - mo = np.asarray(mf.mo_coeff) - if getattr(mf, "is_average_occupation_reference", False): - occupation = np.asarray(mf.mo_occ) - probe = np.asarray(p_alpha) + np.asarray(p_beta) - potential = mf.gen_response(hermi=0)(probe.T) - q_total = ( - mo.conj().T @ (potential + potential.T) @ mo - ) * occupation[None, :] - return 0.5 * q_total, 0.5 * q_total - occ_alpha = (np.asarray(mf.mo_occ) > 0).astype(float) - occ_beta = (np.asarray(mf.mo_occ) == 2).astype(float) - if (isinstance(mf, dft.KohnShamDFT) - and mf._numint._xc_type(mf.xc) != "HF"): - unrestricted = mf.to_uks() - unrestricted.verbose = 0 - v_alpha, v_beta = unrestricted.gen_response(hermi=0)( - np.asarray((p_alpha.T, p_beta.T)) - ) - else: - p_total = p_alpha + p_beta - coulomb = mf.get_j(mf.mol, p_total.T, hermi=0) - v_alpha = coulomb - mf.get_k(mf.mol, p_alpha.T, hermi=0) - v_beta = coulomb - mf.get_k(mf.mol, p_beta.T, hermi=0) - q_alpha = ( - mo.conj().T @ (v_alpha + v_alpha.T) @ mo - ) * occ_alpha[None, :] - q_beta = ( - mo.conj().T @ (v_beta + v_beta.T) @ mo - ) * occ_beta[None, :] - return q_alpha, q_beta - - # Complete lowering scalar and M-matrix ledger def spin_lowering_response_terms(spin): @@ -581,484 +385,14 @@ def project(target, operator, scale): return q_alpha, q_beta -# AO J/K nuclear derivatives - -def _as_derivative_stack(array): - array = np.asarray(array) - if array.ndim == 3: - array = array[None] - return array - - -def _density_key(density): - density = np.asarray(density) - data = density.__array_interface__["data"][0] - return data, density.shape, density.strides, density.dtype.str - - -@dataclass(frozen=True) -class _JKDerivativeTerm: - """One fixed-AO bilinear derivative with a named output slot.""" - - left: np.ndarray - right: np.ndarray - scale: float - omega: float - slot: object - - -class _JKDerivativeLedger: - """Spin-lowering scheduler for fixed-AO J/K derivatives.""" - - def __init__(self): - self._terms = {"j": [], "k": []} - - def add(self, operator, slot, terms): - self._terms[operator].extend( - _JKDerivativeTerm(left, right, scale, omega, slot) - for left, right, scale, omega in terms - if scale != 0.0 - ) - - def contract(self, gradient_driver, mol, atoms, slots=()): - atoms = tuple(atoms) - shape = (len(atoms), 3) - gradients = {slot: np.zeros(shape) for slot in slots} - for operator in ("j", "k"): - for term in self._terms[operator]: - gradients.setdefault(term.slot, np.zeros(shape)) - _contract_derivative_terms( - gradients, - gradient_driver, - mol, - atoms, - mol.offset_nr_by_atom(), - self._terms[operator], - operator, - ) - return gradients - - -def _term_densities(term, exchange): - left, right = term.left, term.right - if exchange: - return left, right, left.T, right.T - return left, right - - -def _density_batches(terms, exchange, max_memory, nao): - """Group bilinear terms while bounding derivative-potential storage.""" - minimum = 4 if exchange else 2 - bytes_per_density = 4 * nao * nao * np.dtype(float).itemsize - batch_limit = max( - minimum, - int(0.2 * max_memory * 1e6 / bytes_per_density), - ) - batch = [] - keys = set() - for term in terms: - term_keys = { - _density_key(density) - for density in _term_densities(term, exchange) - } - if batch and len(keys | term_keys) > batch_limit: - yield batch - batch = [] - keys = set() - batch.append(term) - keys.update(term_keys) - if batch: - yield batch - - -def _jk_derivative_potentials( - gradient_driver, mol, terms, operator, omega): - exchange = operator == "k" - densities = {} - for term in terms: - for density in _term_densities(term, exchange): - density = np.asarray(density) - densities.setdefault(_density_key(density), density) - keys = tuple(densities) - stack = np.asarray([densities[key] for key in keys]) - if operator == "j": - if omega is None: - values = gradient_driver.get_j(mol, stack, hermi=0) - else: - values = gradient_driver.get_j( - mol, stack, hermi=0, omega=omega, - ) - else: - if omega is None: - values = gradient_driver.get_k(mol, stack, hermi=0) - else: - values = gradient_driver.get_k( - mol, stack, hermi=0, omega=omega, - ) - values = _as_derivative_stack(values) - return dict(zip(keys, values)) - - -def _contract_derivative_terms( - gradients, gradient_driver, mol, atoms, offsets, terms, - operator): - if not atoms: - return - terms_by_omega = {} - for term in terms: - terms_by_omega.setdefault(term.omega, []).append(term) - exchange = operator == "k" - for omega, omega_terms in terms_by_omega.items(): - for batch in _density_batches( - omega_terms, exchange, gradient_driver.max_memory, - mol.nao_nr()): - potentials = _jk_derivative_potentials( - gradient_driver, mol, batch, operator, omega, - ) - for term in batch: - left = np.asarray(term.left) - right = np.asarray(term.right) - right_derivative = potentials[_density_key(right)] - left_derivative = potentials[_density_key(left)] - if exchange: - right_t_derivative = potentials[ - _density_key(right.T) - ] - left_t_derivative = potentials[_density_key(left.T)] - for k, atom in enumerate(atoms): - p0, p1 = offsets[atom][2:] - if exchange: - value = lib.einsum( - "xpq,pq->x", - right_derivative[:, p0:p1, :], - left[p0:p1, :], - ) - value += lib.einsum( - "xqp,pq->x", - right_t_derivative[:, p0:p1, :], - left[:, p0:p1], - ) - value += lib.einsum( - "xpq,pq->x", - left_derivative[:, p0:p1, :], - right[p0:p1, :], - ) - value += lib.einsum( - "xqp,pq->x", - left_t_derivative[:, p0:p1, :], - right[:, p0:p1], - ) - else: - value = lib.einsum( - "xpq,pq->x", - right_derivative[:, p0:p1], - left[p0:p1], - ) - value += lib.einsum( - "xpq,qp->x", - right_derivative[:, p0:p1], - left[:, p0:p1], - ) - value += lib.einsum( - "xpq,pq->x", - left_derivative[:, p0:p1], - right[p0:p1], - ) - value += lib.einsum( - "xpq,qp->x", - left_derivative[:, p0:p1], - right[:, p0:p1], - ) - gradients[term.slot][k] += term.scale * value - - -def _reference_spin_densities(tdobj): - mf = tdobj._scf - if getattr(mf, "is_average_occupation_reference", False): - return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) - mo = np.asarray(mf.mo_coeff) - return ( - mo[:, mf.mo_occ > 0] @ mo[:, mf.mo_occ > 0].T, - mo[:, mf.mo_occ == 2] @ mo[:, mf.mo_occ == 2].T, - ) - - -def _spin_probe_stacks(p_alpha, p_beta): - p_alpha = np.asarray(p_alpha) - p_beta = np.asarray(p_beta) - single_probe = p_alpha.ndim == 2 - if single_probe: - p_alpha = p_alpha[None] - p_beta = p_beta[None] - return p_alpha, p_beta, single_probe - - -def spin_fock_direct_dft( - gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, - nobeta_p0=None, jk_ledger=None, output_slots=None): - """Differentiate one or more ordinary UKS Fock scalar probes. - - The optional ``nobeta_p0`` correction belongs to the first, explicit-direct - probe in the batch. - """ - mf = tdobj._scf - mol = tdobj.mol - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - p_alpha, p_beta, single_probe = _spin_probe_stacks( - p_alpha, p_beta, - ) - if output_slots is None: - output_slots = tuple(range(len(p_alpha))) - p_total = p_alpha + p_beta - density_alpha, density_beta = _reference_spin_densities(tdobj) - gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - hcore_derivative = rhf_grad.Gradients(mf).hcore_generator(mol) - for k, atom in enumerate(atmlst): - gradient[:, k] += lib.einsum( - "npq,xpq->nx", p_total, hcore_derivative(atom), - ) - ni = mf._numint - omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - for probe in range(len(p_alpha)): - j_terms = [ - (p_total[probe], density_alpha, 1.0, None), - (p_total[probe], density_beta, 1.0, None), - ] - k_terms = [] - if ni.libxc.is_hybrid_xc(mf.xc): - k_terms.extend(( - (p_alpha[probe], density_alpha, -hybrid, None), - (p_beta[probe], density_beta, -hybrid, None), - )) - if omega != 0: - long_range = -(alpha - hybrid) - k_terms.extend(( - (p_alpha[probe], density_alpha, long_range, omega), - (p_beta[probe], density_beta, long_range, omega), - )) - ledger.add("j", output_slots[probe], j_terms) - ledger.add("k", output_slots[probe], k_terms) - xctype = ni._xc_type(mf.xc) - if xctype == "LDA": - derivative_contractor = xc_backend.contract_lda_vxc_derivative - elif xctype == "GGA": - derivative_contractor = xc_backend.contract_gga_vxc_derivative - elif xctype == "MGGA": - derivative_contractor = xc_backend.contract_mgga_vxc_derivative - else: - raise NotImplementedError( - "ordinary Fock direct derivative is not implemented for %s" % - xctype - ) - if (nobeta_p0 is not None and tdobj.nobeta - and not getattr(mf, "is_average_occupation_reference", False)): - density0 = 0.5 * (density_alpha + density_beta) - actual_probe_alpha = np.array(p_alpha, copy=True) - actual_probe_beta = np.array(p_beta, copy=True) - actual_probe_alpha[0] -= 0.5 * nobeta_p0 - actual_probe_beta[0] -= 0.5 * nobeta_p0 - else: - density0 = None - actual_probe_alpha = p_alpha - actual_probe_beta = p_beta - gradient += derivative_contractor( - mf, - density_alpha, - density_beta, - actual_probe_alpha, - actual_probe_beta, - atmlst=atmlst, - max_memory=gradient_driver.max_memory, - ) - if density0 is not None: - gradient[0] += derivative_contractor( - mf, - density0, - density0, - 0.5 * nobeta_p0, - 0.5 * nobeta_p0, - atmlst=atmlst, - max_memory=gradient_driver.max_memory, - ) - if jk_ledger is None: - contractions = local_ledger.contract( - gradient_driver, mol, atmlst, slots=output_slots, - ) - for probe, slot in enumerate(output_slots): - gradient[probe] += contractions[slot] - return gradient[0] if single_probe else gradient - - -def spin_fock_direct_hf( - gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, - jk_ledger=None, output_slots=None): - """Differentiate one or more spin-resolved HF Fock scalar probes.""" - mol = tdobj.mol - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - p_alpha, p_beta, single_probe = _spin_probe_stacks( - p_alpha, p_beta, - ) - if output_slots is None: - output_slots = tuple(range(len(p_alpha))) - p_total = p_alpha + p_beta - dm_alpha, dm_beta = _reference_spin_densities(tdobj) - gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - - hcore_derivative = rhf_grad.Gradients(tdobj._scf).hcore_generator(mol) - for k, atom in enumerate(atmlst): - gradient[:, k] += lib.einsum( - "npq,xpq->nx", p_total, hcore_derivative(atom), - ) - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - for probe in range(len(p_alpha)): - ledger.add( - "j", output_slots[probe], ( - (p_total[probe], dm_alpha, 1.0, None), - (p_total[probe], dm_beta, 1.0, None), - ), - ) - ledger.add( - "k", output_slots[probe], ( - (p_alpha[probe], dm_alpha, -1.0, None), - (p_beta[probe], dm_beta, -1.0, None), - ), - ) - if jk_ledger is None: - contractions = local_ledger.contract( - gradient_driver, mol, atmlst, slots=output_slots, - ) - for probe, slot in enumerate(output_slots): - gradient[probe] += contractions[slot] - return gradient[0] if single_probe else gradient - - -def response_direct_hfx( - gradient_driver, tdobj, densities, response_terms, atmlst=None, - jk_ledger=None, output_slot=0): - """J/K skeleton derivative for a channel response-term ledger.""" - mol = tdobj.mol - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - gradient = np.zeros((len(atmlst), 3)) - ni = tdobj._scf._numint - omega, alpha, hybrid = ni.rsh_and_hybrid_coeff( - tdobj._scf.xc, mol.spin, - ) - if not ni.libxc.is_hybrid_xc(tdobj._scf.xc): - return gradient - - scales = [(hybrid, None)] - if omega != 0: - scales.append((alpha - hybrid, omega)) - j_terms = [] - k_terms = [] - for term in response_terms: - target = densities[term.target] - source = densities[term.source] - for coefficient, range_omega in scales: - if term.vref0: - k_terms.append(( - target, - source, - -coefficient * term.vref0, - range_omega, - )) - if term.vref1: - j_terms.append(( - target, - source, - -coefficient * term.vref1, - range_omega, - )) - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - ledger.add("j", output_slot, j_terms) - ledger.add("k", output_slot, k_terms) - if jk_ledger is None: - gradient += local_ledger.contract( - gradient_driver, mol, atmlst, slots=(output_slot,), - )[output_slot] - return gradient - - -# Hybrid/RSH Fz correction - -def spin_lowering_fockz_hfx_terms( - gradient_driver, tdobj, pz, atmlst=None, with_direct=True, - jk_ledger=None, output_slot=0): - """Differentiate ``-1/2 Pz:K(D_OO)`` excluding the Pz projection.""" - mf = tdobj._scf - mol = mf.mol - ni = mf._numint - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - mo = np.asarray(mf.mo_coeff) - q_alpha = np.zeros((mo.shape[1], mo.shape[1])) - q_beta = np.zeros_like(q_alpha) - direct = np.zeros((len(atmlst), 3)) - if not ni.libxc.is_hybrid_xc(mf.xc): - return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) - - spaces = orbital_spaces(tdobj) - density_open = spaces.c_open @ spaces.c_open.T - omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) - scales = [(hybrid, None)] - if omega != 0: - scales.append((alpha - hybrid, omega)) - k_terms = [] - for coefficient, range_omega in scales: - if coefficient == 0.0: - continue - if range_omega is None: - potential = mf.get_k(mol, pz, hermi=0) - else: - potential = mf.get_k( - mol, pz, hermi=0, omega=range_omega, - ) - q_alpha[:, spaces.open] -= 0.5 * coefficient * ( - mo.conj().T @ (potential + potential.T) @ spaces.c_open - ) - if with_direct: - k_terms.append(( - pz, - density_open, - -0.5 * coefficient, - range_omega, - )) - if with_direct: - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - ledger.add("k", output_slot, k_terms) - if jk_ledger is None: - direct += local_ledger.contract( - gradient_driver, mol, atmlst, slots=(output_slot,), - )[output_slot] - return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) - # Channel assembly def grad_elec( gradient_driver, tdobj, xy, atmlst=None, tolerance=1e-12, max_cycle=None): - """Build the complete analytic excitation gradient for deltaS=-1.""" + """Build the analytic excitation gradient for deltaS=-1.""" if tdobj.deltaS != -1: raise ValueError("deltaS=-1 gradient received a different spin channel") - if atmlst is None: - atmlst = range(tdobj.mol.natm) - atmlst = tuple(atmlst) - mf = tdobj._scf - xctype = mf._numint._xc_type(mf.xc) - - # 1. Native amplitudes, transition densities, and explicit Fock probes. spaces, amplitudes, densities = spin_lowering_transition_densities( tdobj, xy, ) @@ -1066,166 +400,9 @@ def grad_elec( response_terms = spin_lowering_response_terms(spaces.spin) channel_data = (spaces, amplitudes, densities, blocks, response_terms) p0, pz = spin_lowering_fock_probes(tdobj, xy) - jk_ledger = _JKDerivativeLedger() - direct_slot = "direct" - zvector_slot = "zvector" - - # 2. Explicit Fock contribution to the orbital-rotation M matrix. - fock_alpha, fock_beta = spin_lowering_fock_q(tdobj, xy) - hfx_alpha, hfx_beta = spin_lowering_response_projection_q( - tdobj, xy, hfx_only=True, - ) - - if xctype == "HF": - # 3a. HF response and fixed-orbital AO derivative. - m_matrix = ( - fock_alpha + fock_beta + hfx_alpha + hfx_beta - ) - direct = response_direct_hfx( - gradient_driver, - tdobj, - densities, - response_terms, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slot=direct_slot, - ) - if getattr(mf, "is_average_occupation_reference", False): - fockz_hfx = spin_lowering_fockz_hfx_terms( - gradient_driver, - tdobj, - pz, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slot=direct_slot, - ) - m_matrix += fockz_hfx.q_alpha + fockz_hfx.q_beta - direct += fockz_hfx.direct - direct_fock_probes = (0.5 * p0, 0.5 * p0) - else: - direct_fock_probes = ( - 0.5 * (p0 + pz), - 0.5 * (p0 - pz), - ) - - def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): - local = spin_fock_direct_hf( - driver, - obj, - p_alpha, - p_beta, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slots=(direct_slot, zvector_slot), - ) - contractions = jk_ledger.contract( - driver, - obj.mol, - atmlst, - slots=(direct_slot, zvector_slot), - ) - local[0] += contractions[direct_slot] - local[1] += contractions[zvector_slot] - return local - else: - # 3b. Semilocal XC, hybrid/RSH, Fz, and nobeta contributions. - try: - response_builder, fockz_builder, nobeta_q_builder = { - "LDA": ( - xc_backend.lda_response_terms, - xc_backend.lda_fockz_terms, - xc_backend.lda_nobeta_reference_q, - ), - "GGA": ( - xc_backend.gga_response_terms, - xc_backend.gga_fockz_terms, - xc_backend.gga_nobeta_reference_q, - ), - "MGGA": ( - xc_backend.mgga_response_terms, - xc_backend.mgga_fockz_terms, - xc_backend.mgga_nobeta_reference_q, - ), - }[xctype] - except KeyError as error: - raise NotImplementedError( - "NTTDA deltaS=-1 gradient does not support XC type %s" % xctype - ) from error - - response_xc = response_builder( - gradient_driver, - tdobj, - channel_data, - atmlst=atmlst, - ) - fockz_xc = fockz_builder( - gradient_driver, - tdobj, - spaces, - pz, - atmlst=atmlst, - ) - fockz_hfx = spin_lowering_fockz_hfx_terms( - gradient_driver, - tdobj, - pz, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slot=direct_slot, - ) - common_alpha, common_beta = nobeta_q_builder(tdobj, p0) - m_matrix = ( - fock_alpha + fock_beta - + hfx_alpha + hfx_beta - + response_xc.q_alpha + response_xc.q_beta - + fockz_xc.q_alpha + fockz_xc.q_beta - + fockz_hfx.q_alpha + fockz_hfx.q_beta - + common_alpha + common_beta - ) - - direct_fock_probes = (0.5 * p0, 0.5 * p0) - direct = response_direct_hfx( - gradient_driver, - tdobj, - densities, - response_terms, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slot=direct_slot, - ) - direct += response_xc.direct - direct += fockz_xc.direct - direct += fockz_hfx.direct - def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): - local = spin_fock_direct_dft( - driver, - obj, - p_alpha, - p_beta, - atmlst=atmlst, - nobeta_p0=p0, - jk_ledger=jk_ledger, - output_slots=(direct_slot, zvector_slot), - ) - contractions = jk_ledger.contract( - driver, - obj.mol, - atmlst, - slots=(direct_slot, zvector_slot), - ) - local[0] += contractions[direct_slot] - local[1] += contractions[zvector_slot] - return local - - # 4-5. ROKS transpose-Hessian adjoint, Dz Fock derivative, and Pulay term. - return finish_gradient( - gradient_driver, - tdobj, - m_matrix, - direct, - atmlst, - tolerance, - max_cycle, - fock_direct, - direct_fock_probes=direct_fock_probes, + return assemble_gradient( + gradient_driver, tdobj, channel_data, (p0, pz), + spin_lowering_fock_q(tdobj, xy), + spin_lowering_response_projection_q(tdobj, xy, hfx_only=True), + atmlst=atmlst, tolerance=tolerance, max_cycle=max_cycle, ) diff --git a/src/nest/grad/nttda/delta_s_zero.py b/src/nest/grad/nttda/delta_s_zero.py index 18722b0..60d4998 100644 --- a/src/nest/grad/nttda/delta_s_zero.py +++ b/src/nest/grad/nttda/delta_s_zero.py @@ -1,42 +1,32 @@ """Analytic gradient for current NTTDA ``deltaS=0``. The public ``grad_elec`` function exposes the complete scientific data flow. -All same-spin amplitudes, Fock projections, response coefficients, and direct -J/K contractions live in this module. Only XC quadrature and the ROKS -adjoint are delegated to sibling modules. +Same-spin amplitudes, Fock projections and response coefficients live here. +Reference response, J/K derivatives, XC quadrature and the adjoint are shared. """ from dataclasses import dataclass import numpy as np -from pyscf import dft, lib -from pyscf.grad import rhf as rhf_grad +from pyscf import lib from nest.nttda import nttda as nttda_mod from nest.nttda.nttda import gen_rohf_response_sc -from . import xc as xc_backend -from .roks import finish_gradient - +from .common import ( + assemble_gradient, + orbital_spaces, + pair_density, + FockProjection, + _fock_response_q, + ResponseTerm, + _apply_reference_responses, + _apply_hfx_responses, +) # Orbital spaces and native amplitudes -@dataclass(frozen=True) -class OrbitalSpaces: - """Closed, open, and virtual spatial-orbital partitions.""" - - closed: np.ndarray - open: np.ndarray - virtual: np.ndarray - c_closed: np.ndarray - c_open: np.ndarray - c_virtual: np.ndarray - - @property - def spin(self): - return 0.5 * len(self.open) - @dataclass(frozen=True) class SameSpinAmplitudes: @@ -49,26 +39,6 @@ class SameSpinAmplitudes: cv0: np.ndarray -def orbital_spaces(tdobj): - """Return the ROKS ``C/O/V`` orbital partition used by NTTDA.""" - mf = tdobj._scf - occ = np.asarray(mf.mo_occ) - if occ.ndim != 1: - raise ValueError("NTTDA gradients require spatial ROKS orbitals") - closed = np.flatnonzero(occ == 2) - open_ = np.flatnonzero(occ == 1) - virtual = np.flatnonzero(occ == 0) - coeff = np.asarray(mf.mo_coeff) - return OrbitalSpaces( - closed=closed, - open=open_, - virtual=virtual, - c_closed=coeff[:, closed], - c_open=coeff[:, open_], - c_virtual=coeff[:, virtual], - ) - - def same_spin_slices(spaces): """Return canonical slices for ``CO/CV/OO/OV/CV0`` amplitudes.""" nc = len(spaces.closed) @@ -116,11 +86,6 @@ def split_same_spin(tdobj, xy): ) -def pair_density(c_left, coefficient, c_right): - """Build ``C_left coefficient C_right^T`` without symmetrizing it.""" - return c_left @ np.asarray(coefficient) @ c_right.conj().T - - def same_spin_transition_densities(tdobj, xy): """Return directed AO transition densities for the four response blocks.""" spaces, amp = split_same_spin(tdobj, xy) @@ -134,24 +99,6 @@ def same_spin_transition_densities(tdobj, xy): # Explicit F0/Fz ledger -@dataclass(frozen=True) -class FockProjection: - """One scalar term ``Tr[P (weight_f0 F0 + weight_fz Fz)]``.""" - - name: str - left_indices: np.ndarray - left_orbitals: np.ndarray - coefficient: np.ndarray - right_indices: np.ndarray - right_orbitals: np.ndarray - weight_f0: float - weight_fz: float - - def density(self): - return pair_density( - self.left_orbitals, self.coefficient, self.right_orbitals, - ) - def fock0_fockz(tdobj, max_memory=None): """Build exactly the ``F0`` and ``Fz`` matrices used by ``gen_vind_sc``.""" @@ -246,41 +193,6 @@ def same_spin_fock_probes(tdobj, xy): return p0, pz -def _fock_response_q(tdobj, p_alpha, p_beta): - """Reference-density derivative of a spin-resolved Fock scalar.""" - mf = tdobj._scf - mo = np.asarray(mf.mo_coeff) - if getattr(mf, "is_average_occupation_reference", False): - occupation = np.asarray(mf.mo_occ) - probe = np.asarray(p_alpha) + np.asarray(p_beta) - potential = mf.gen_response(hermi=0)(probe.T) - q_total = ( - mo.conj().T @ (potential + potential.T) @ mo - ) * occupation[None, :] - return 0.5 * q_total, 0.5 * q_total - occ_alpha = (np.asarray(mf.mo_occ) > 0).astype(float) - occ_beta = (np.asarray(mf.mo_occ) == 2).astype(float) - if (isinstance(mf, dft.KohnShamDFT) - and mf._numint._xc_type(mf.xc) != "HF"): - unrestricted = mf.to_uks() - unrestricted.verbose = 0 - v_alpha, v_beta = unrestricted.gen_response(hermi=0)( - np.asarray((p_alpha.T, p_beta.T)) - ) - else: - p_total = p_alpha + p_beta - coulomb = mf.get_j(mf.mol, p_total.T, hermi=0) - v_alpha = coulomb - mf.get_k(mf.mol, p_alpha.T, hermi=0) - v_beta = coulomb - mf.get_k(mf.mol, p_beta.T, hermi=0) - q_alpha = ( - mo.conj().T @ (v_alpha + v_alpha.T) @ mo - ) * occ_alpha[None, :] - q_beta = ( - mo.conj().T @ (v_beta + v_beta.T) @ mo - ) * occ_beta[None, :] - return q_alpha, q_beta - - def same_spin_fock_q(tdobj, xy, max_memory=None): """Return the explicit-Fock contribution to ``(Q_alpha,Q_beta)``. @@ -347,15 +259,6 @@ def same_spin_fock_projection_scalar(tdobj, xy, fock0, fockz): # vref0/vref1 response ledger -@dataclass(frozen=True) -class ResponseTerm: - """Directed response term from one source density to one target block.""" - - target: str - source: str - vref0: float - vref1: float - def same_spin_response_terms(spin): """Directed coefficients transcribed from ``gen_rohf_response_sc``.""" @@ -380,97 +283,6 @@ def same_spin_response_terms(spin): ) -def _fxc_reference(tdobj): - mf = tdobj._scf - ni = mf._numint - fxc = ni.cache_xc_kernel( - mf.mol, mf.grids, mf.xc, mf.mo_coeff, mf.mo_occ, 1, - )[2] - return 0.5 * ( - fxc[0, :, 0] - fxc[0, :, 1] - - fxc[1, :, 0] + fxc[1, :, 1] - ) - - -def _apply_reference_responses(tdobj, densities, max_memory=None): - """Return separate ``vref0`` and ``vref1`` actions for each density.""" - mf = tdobj._scf - mol = mf.mol - ni = mf._numint - if max_memory is None: - max_memory = tdobj.max_memory - labels = tuple(densities) - dms = np.asarray([densities[label] for label in labels]) - xctype = ni._xc_type(mf.xc) - if xctype == "HF": - vref0 = np.zeros_like(dms) - vref1 = np.zeros_like(dms) - else: - fxc_ref = _fxc_reference(tdobj) - vref0 = ni.nr_rks_fxc( - mol, mf.grids, mf.xc, None, dms, 0, 0, - None, None, fxc_ref, max_memory=max_memory, - ) - if xctype == "LDA": - vref1 = ni.nr_rks_fxc( - mol, mf.grids, mf.xc, None, dms, 0, 0, - None, None, fxc_ref, max_memory=max_memory, - ) - elif xctype == "GGA": - vref1 = nttda_mod.nr_rks_fxc1_gga( - ni, mol, mf.grids, mf.xc, dms, fxc_ref, - max_memory=max_memory, - ) - elif xctype == "MGGA": - vref1 = nttda_mod.nr_rks_fxc1_mgga( - ni, mol, mf.grids, mf.xc, dms, fxc_ref, - max_memory=max_memory, - ) - else: - raise NotImplementedError( - "NTTDA same-spin response does not support XC type %s" % xctype - ) - - omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) - if ni.libxc.is_hybrid_xc(mf.xc): - vref0 -= hyb * mf.get_k(mol, dms, hermi=0) - vref1 -= hyb * mf.get_j(mol, dms, hermi=0) - if omega != 0: - scale = alpha - hyb - vref0 -= scale * mf.get_k(mol, dms, hermi=0, omega=omega) - vref1 -= scale * mf.get_j(mol, dms, hermi=0, omega=omega) - return ( - {label: value for label, value in zip(labels, vref0)}, - {label: value for label, value in zip(labels, vref1)}, - ) - - -def _apply_hfx_responses(tdobj, densities): - """Return only the hybrid/RSH J/K portions of ``vref0/vref1``.""" - mf = tdobj._scf - labels = tuple(densities) - dms = np.asarray([densities[label] for label in labels]) - vref0 = np.zeros_like(dms) - vref1 = np.zeros_like(dms) - ni = mf._numint - omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mf.mol.spin) - if ni.libxc.is_hybrid_xc(mf.xc): - vref0 -= hybrid * mf.get_k(mf.mol, dms, hermi=0) - vref1 -= hybrid * mf.get_j(mf.mol, dms, hermi=0) - if omega != 0: - scale = alpha - hybrid - vref0 -= scale * mf.get_k( - mf.mol, dms, hermi=0, omega=omega, - ) - vref1 -= scale * mf.get_j( - mf.mol, dms, hermi=0, omega=omega, - ) - return ( - {label: value for label, value in zip(labels, vref0)}, - {label: value for label, value in zip(labels, vref1)}, - ) - - def _derivative_potentials(spaces, densities, vref0, vref1): potentials = {label: np.zeros_like(dm) for label, dm in densities.items()} for term in same_spin_response_terms(spaces.spin): @@ -557,470 +369,6 @@ def same_spin_hfx_projection_q(tdobj, xy): return q_alpha, q_beta -# AO J/K nuclear derivatives - -def _as_derivative_stack(array): - array = np.asarray(array) - if array.ndim == 3: - array = array[None] - return array - - -def _density_key(density): - density = np.asarray(density) - data = density.__array_interface__["data"][0] - return data, density.shape, density.strides, density.dtype.str - - -@dataclass(frozen=True) -class _JKDerivativeTerm: - """One fixed-AO bilinear derivative with a named output slot.""" - - left: np.ndarray - right: np.ndarray - scale: float - omega: float - slot: object - - -class _JKDerivativeLedger: - """Channel-local scheduler for fixed-AO J/K derivative contractions.""" - - def __init__(self): - self._terms = {"j": [], "k": []} - - def add(self, operator, slot, terms): - self._terms[operator].extend( - _JKDerivativeTerm(left, right, scale, omega, slot) - for left, right, scale, omega in terms - if scale != 0.0 - ) - - def contract(self, gradient_driver, mol, atoms, slots=()): - atoms = tuple(atoms) - shape = (len(atoms), 3) - gradients = {slot: np.zeros(shape) for slot in slots} - for operator in ("j", "k"): - for term in self._terms[operator]: - gradients.setdefault(term.slot, np.zeros(shape)) - _contract_derivative_terms( - gradients, - gradient_driver, - mol, - atoms, - mol.offset_nr_by_atom(), - self._terms[operator], - operator, - ) - return gradients - - -def _term_densities(term, exchange): - left, right = term.left, term.right - if exchange: - return left, right, left.T, right.T - return left, right - - -def _density_batches(terms, exchange, max_memory, nao): - """Group bilinear terms while bounding derivative-potential storage.""" - minimum = 4 if exchange else 2 - bytes_per_density = 4 * nao * nao * np.dtype(float).itemsize - batch_limit = max( - minimum, - int(0.2 * max_memory * 1e6 / bytes_per_density), - ) - batch = [] - keys = set() - for term in terms: - term_keys = { - _density_key(density) - for density in _term_densities(term, exchange) - } - if batch and len(keys | term_keys) > batch_limit: - yield batch - batch = [] - keys = set() - batch.append(term) - keys.update(term_keys) - if batch: - yield batch - - -def _jk_derivative_potentials( - gradient_driver, mol, terms, operator, omega): - exchange = operator == "k" - densities = {} - for term in terms: - for density in _term_densities(term, exchange): - density = np.asarray(density) - densities.setdefault(_density_key(density), density) - keys = tuple(densities) - stack = np.asarray([densities[key] for key in keys]) - if operator == "j": - if omega is None: - values = gradient_driver.get_j(mol, stack, hermi=0) - else: - values = gradient_driver.get_j( - mol, stack, hermi=0, omega=omega, - ) - else: - if omega is None: - values = gradient_driver.get_k(mol, stack, hermi=0) - else: - values = gradient_driver.get_k( - mol, stack, hermi=0, omega=omega, - ) - values = _as_derivative_stack(values) - return dict(zip(keys, values)) - - -def _contract_derivative_terms( - gradients, gradient_driver, mol, atoms, offsets, terms, - operator): - if not atoms: - return - terms_by_omega = {} - for term in terms: - terms_by_omega.setdefault(term.omega, []).append(term) - exchange = operator == "k" - for omega, omega_terms in terms_by_omega.items(): - for batch in _density_batches( - omega_terms, exchange, gradient_driver.max_memory, - mol.nao_nr()): - potentials = _jk_derivative_potentials( - gradient_driver, mol, batch, operator, omega, - ) - for term in batch: - left = np.asarray(term.left) - right = np.asarray(term.right) - right_derivative = potentials[_density_key(right)] - left_derivative = potentials[_density_key(left)] - if exchange: - right_t_derivative = potentials[ - _density_key(right.T) - ] - left_t_derivative = potentials[_density_key(left.T)] - for k, atom in enumerate(atoms): - p0, p1 = offsets[atom][2:] - if exchange: - value = lib.einsum( - "xpq,pq->x", - right_derivative[:, p0:p1, :], - left[p0:p1, :], - ) - value += lib.einsum( - "xqp,pq->x", - right_t_derivative[:, p0:p1, :], - left[:, p0:p1], - ) - value += lib.einsum( - "xpq,pq->x", - left_derivative[:, p0:p1, :], - right[p0:p1, :], - ) - value += lib.einsum( - "xqp,pq->x", - left_t_derivative[:, p0:p1, :], - right[:, p0:p1], - ) - else: - value = lib.einsum( - "xpq,pq->x", - right_derivative[:, p0:p1], - left[p0:p1], - ) - value += lib.einsum( - "xpq,qp->x", - right_derivative[:, p0:p1], - left[:, p0:p1], - ) - value += lib.einsum( - "xpq,pq->x", - left_derivative[:, p0:p1], - right[p0:p1], - ) - value += lib.einsum( - "xpq,qp->x", - left_derivative[:, p0:p1], - right[:, p0:p1], - ) - gradients[term.slot][k] += term.scale * value - - -def _reference_spin_densities(tdobj): - mf = tdobj._scf - if getattr(mf, "is_average_occupation_reference", False): - return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) - mo = np.asarray(mf.mo_coeff) - return ( - mo[:, mf.mo_occ > 0] @ mo[:, mf.mo_occ > 0].T, - mo[:, mf.mo_occ == 2] @ mo[:, mf.mo_occ == 2].T, - ) - - -def _spin_probe_stacks(p_alpha, p_beta): - p_alpha = np.asarray(p_alpha) - p_beta = np.asarray(p_beta) - single_probe = p_alpha.ndim == 2 - if single_probe: - p_alpha = p_alpha[None] - p_beta = p_beta[None] - return p_alpha, p_beta, single_probe - - -def spin_fock_direct_dft( - gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, - nobeta_p0=None, jk_ledger=None, output_slots=None): - """Differentiate one or more ordinary UKS Fock scalar probes. - - The optional ``nobeta_p0`` correction belongs to the first, explicit-direct - probe in the batch. - """ - mf = tdobj._scf - mol = tdobj.mol - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - p_alpha, p_beta, single_probe = _spin_probe_stacks( - p_alpha, p_beta, - ) - if output_slots is None: - output_slots = tuple(range(len(p_alpha))) - p_total = p_alpha + p_beta - density_alpha, density_beta = _reference_spin_densities(tdobj) - gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - hcore_derivative = rhf_grad.Gradients(mf).hcore_generator(mol) - for k, atom in enumerate(atmlst): - gradient[:, k] += lib.einsum( - "npq,xpq->nx", p_total, hcore_derivative(atom), - ) - ni = mf._numint - omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - for probe in range(len(p_alpha)): - j_terms = [ - (p_total[probe], density_alpha, 1.0, None), - (p_total[probe], density_beta, 1.0, None), - ] - k_terms = [] - if ni.libxc.is_hybrid_xc(mf.xc): - k_terms.extend(( - (p_alpha[probe], density_alpha, -hybrid, None), - (p_beta[probe], density_beta, -hybrid, None), - )) - if omega != 0: - long_range = -(alpha - hybrid) - k_terms.extend(( - (p_alpha[probe], density_alpha, long_range, omega), - (p_beta[probe], density_beta, long_range, omega), - )) - ledger.add("j", output_slots[probe], j_terms) - ledger.add("k", output_slots[probe], k_terms) - xctype = ni._xc_type(mf.xc) - if xctype == "LDA": - derivative_contractor = xc_backend.contract_lda_vxc_derivative - elif xctype == "GGA": - derivative_contractor = xc_backend.contract_gga_vxc_derivative - elif xctype == "MGGA": - derivative_contractor = xc_backend.contract_mgga_vxc_derivative - else: - raise NotImplementedError( - "ordinary Fock direct derivative is not implemented for %s" % - xctype - ) - if (nobeta_p0 is not None and tdobj.nobeta - and not getattr(mf, "is_average_occupation_reference", False)): - density0 = 0.5 * (density_alpha + density_beta) - actual_probe_alpha = np.array(p_alpha, copy=True) - actual_probe_beta = np.array(p_beta, copy=True) - actual_probe_alpha[0] -= 0.5 * nobeta_p0 - actual_probe_beta[0] -= 0.5 * nobeta_p0 - else: - density0 = None - actual_probe_alpha = p_alpha - actual_probe_beta = p_beta - gradient += derivative_contractor( - mf, - density_alpha, - density_beta, - actual_probe_alpha, - actual_probe_beta, - atmlst=atmlst, - max_memory=gradient_driver.max_memory, - ) - if density0 is not None: - gradient[0] += derivative_contractor( - mf, - density0, - density0, - 0.5 * nobeta_p0, - 0.5 * nobeta_p0, - atmlst=atmlst, - max_memory=gradient_driver.max_memory, - ) - if jk_ledger is None: - contractions = local_ledger.contract( - gradient_driver, mol, atmlst, slots=output_slots, - ) - for probe, slot in enumerate(output_slots): - gradient[probe] += contractions[slot] - return gradient[0] if single_probe else gradient - - -def spin_fock_direct_hf( - gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, - jk_ledger=None, output_slots=None): - """Differentiate one or more spin-resolved HF Fock scalar probes.""" - mol = tdobj.mol - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - p_alpha, p_beta, single_probe = _spin_probe_stacks( - p_alpha, p_beta, - ) - if output_slots is None: - output_slots = tuple(range(len(p_alpha))) - p_total = p_alpha + p_beta - dm_alpha, dm_beta = _reference_spin_densities(tdobj) - gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - - hcore_derivative = rhf_grad.Gradients(tdobj._scf).hcore_generator(mol) - for k, atom in enumerate(atmlst): - gradient[:, k] += lib.einsum( - "npq,xpq->nx", p_total, hcore_derivative(atom), - ) - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - for probe in range(len(p_alpha)): - ledger.add( - "j", output_slots[probe], ( - (p_total[probe], dm_alpha, 1.0, None), - (p_total[probe], dm_beta, 1.0, None), - ), - ) - ledger.add( - "k", output_slots[probe], ( - (p_alpha[probe], dm_alpha, -1.0, None), - (p_beta[probe], dm_beta, -1.0, None), - ), - ) - if jk_ledger is None: - contractions = local_ledger.contract( - gradient_driver, mol, atmlst, slots=output_slots, - ) - for probe, slot in enumerate(output_slots): - gradient[probe] += contractions[slot] - return gradient[0] if single_probe else gradient - - -def response_direct_hfx( - gradient_driver, tdobj, densities, response_terms, atmlst=None, - jk_ledger=None, output_slot=0): - """J/K skeleton derivative for a channel response-term ledger.""" - mol = tdobj.mol - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - gradient = np.zeros((len(atmlst), 3)) - ni = tdobj._scf._numint - omega, alpha, hybrid = ni.rsh_and_hybrid_coeff( - tdobj._scf.xc, mol.spin, - ) - if not ni.libxc.is_hybrid_xc(tdobj._scf.xc): - return gradient - - scales = [(hybrid, None)] - if omega != 0: - scales.append((alpha - hybrid, omega)) - j_terms = [] - k_terms = [] - for term in response_terms: - target = densities[term.target] - source = densities[term.source] - for coefficient, range_omega in scales: - if term.vref0: - k_terms.append(( - target, - source, - -coefficient * term.vref0, - range_omega, - )) - if term.vref1: - j_terms.append(( - target, - source, - -coefficient * term.vref1, - range_omega, - )) - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - ledger.add("j", output_slot, j_terms) - ledger.add("k", output_slot, k_terms) - if jk_ledger is None: - gradient += local_ledger.contract( - gradient_driver, mol, atmlst, slots=(output_slot,), - )[output_slot] - return gradient - - -# Hybrid/RSH Fz correction - -def same_spin_fockz_hfx_terms( - gradient_driver, tdobj, pz, atmlst=None, with_direct=True, - jk_ledger=None, output_slot=0): - """Differentiate ``-1/2 Pz:K(D_OO)`` excluding the Pz projection.""" - mf = tdobj._scf - mol = mf.mol - ni = mf._numint - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - mo = np.asarray(mf.mo_coeff) - q_alpha = np.zeros((mo.shape[1], mo.shape[1])) - q_beta = np.zeros_like(q_alpha) - direct = np.zeros((len(atmlst), 3)) - if not ni.libxc.is_hybrid_xc(mf.xc): - return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) - - spaces = orbital_spaces(tdobj) - density_open = spaces.c_open @ spaces.c_open.T - omega, alpha, hybrid = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) - scales = [(hybrid, None)] - if omega != 0: - scales.append((alpha - hybrid, omega)) - k_terms = [] - for coefficient, range_omega in scales: - if coefficient == 0.0: - continue - if range_omega is None: - potential = mf.get_k(mol, pz, hermi=0) - else: - potential = mf.get_k( - mol, pz, hermi=0, omega=range_omega, - ) - q_alpha[:, spaces.open] -= 0.5 * coefficient * ( - mo.conj().T @ (potential + potential.T) @ spaces.c_open - ) - if with_direct: - k_terms.append(( - pz, - density_open, - -0.5 * coefficient, - range_omega, - )) - if with_direct: - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - ledger.add("k", output_slot, k_terms) - if jk_ledger is None: - direct += local_ledger.contract( - gradient_driver, mol, atmlst, slots=(output_slot,), - )[output_slot] - return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) - - # Scalar closure diagnostics (private to this channel) def same_spin_action_scalar(tdobj, xy): @@ -1046,21 +394,9 @@ def same_spin_ledger_scalar(tdobj, xy, max_memory=None, return_parts=False): def grad_elec( gradient_driver, tdobj, xy, atmlst=None, tolerance=1e-12, max_cycle=None): - """Build the complete analytic excitation gradient for deltaS=0. - - The function follows the physical order of the Lagrangian: native - amplitudes and AO probes, XC/J/K contributions to the M matrix and direct - derivative, the ROKS adjoint, and the final overlap contraction. - """ + """Build the analytic excitation gradient for deltaS=0.""" if tdobj.deltaS != 0: raise ValueError("deltaS=0 gradient received a different spin channel") - if atmlst is None: - atmlst = range(tdobj.mol.natm) - atmlst = tuple(atmlst) - mf = tdobj._scf - xctype = mf._numint._xc_type(mf.xc) - - # 1. Native amplitudes, transition densities, and explicit Fock probes. spaces, amplitudes, densities = same_spin_transition_densities(tdobj, xy) blocks = { "CO": (spaces.open, spaces.closed, amplitudes.co.T), @@ -1071,167 +407,9 @@ def grad_elec( response_terms = same_spin_response_terms(spaces.spin) channel_data = (spaces, amplitudes, densities, blocks, response_terms) p0, pz = same_spin_fock_probes(tdobj, xy) - jk_ledger = _JKDerivativeLedger() - direct_slot = "direct" - zvector_slot = "zvector" - - # 2. Explicit Fock contribution to the orbital-rotation M matrix. - fock_alpha, fock_beta = same_spin_fock_q(tdobj, xy) - - if xctype == "HF": - # 3a. HF response and fixed-orbital AO derivative. - response_alpha, response_beta = same_spin_response_projection_q( - tdobj, xy, - ) - m_matrix = ( - fock_alpha + fock_beta + response_alpha + response_beta - ) - direct = response_direct_hfx( - gradient_driver, - tdobj, - densities, - response_terms, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slot=direct_slot, - ) - if getattr(mf, "is_average_occupation_reference", False): - fockz_hfx = same_spin_fockz_hfx_terms( - gradient_driver, - tdobj, - pz, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slot=direct_slot, - ) - m_matrix += fockz_hfx.q_alpha + fockz_hfx.q_beta - direct += fockz_hfx.direct - direct_fock_probes = (0.5 * p0, 0.5 * p0) - else: - direct_fock_probes = ( - 0.5 * (p0 + pz), - 0.5 * (p0 - pz), - ) - - def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): - local = spin_fock_direct_hf( - driver, - obj, - p_alpha, - p_beta, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slots=(direct_slot, zvector_slot), - ) - contractions = jk_ledger.contract( - driver, - obj.mol, - atmlst, - slots=(direct_slot, zvector_slot), - ) - local[0] += contractions[direct_slot] - local[1] += contractions[zvector_slot] - return local - else: - # 3b. Semilocal XC, hybrid/RSH, Fz, and nobeta contributions. - try: - response_builder, fockz_builder, nobeta_q_builder = { - "LDA": ( - xc_backend.lda_response_terms, - xc_backend.lda_fockz_terms, - xc_backend.lda_nobeta_reference_q, - ), - "GGA": ( - xc_backend.gga_response_terms, - xc_backend.gga_fockz_terms, - xc_backend.gga_nobeta_reference_q, - ), - "MGGA": ( - xc_backend.mgga_response_terms, - xc_backend.mgga_fockz_terms, - xc_backend.mgga_nobeta_reference_q, - ), - }[xctype] - except KeyError as error: - raise NotImplementedError( - "NTTDA deltaS=0 gradient does not support XC type %s" % xctype - ) from error - - hfx_alpha, hfx_beta = same_spin_hfx_projection_q(tdobj, xy) - response_xc = response_builder( - gradient_driver, - tdobj, - channel_data, - atmlst=atmlst, - ) - fockz_xc = fockz_builder( - gradient_driver, - tdobj, - spaces, - pz, - atmlst=atmlst, - ) - fockz_hfx = same_spin_fockz_hfx_terms( - gradient_driver, - tdobj, - pz, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slot=direct_slot, - ) - common_alpha, common_beta = nobeta_q_builder(tdobj, p0) - m_matrix = ( - fock_alpha + fock_beta - + hfx_alpha + hfx_beta - + response_xc.q_alpha + response_xc.q_beta - + fockz_xc.q_alpha + fockz_xc.q_beta - + fockz_hfx.q_alpha + fockz_hfx.q_beta - + common_alpha + common_beta - ) - - direct_fock_probes = (0.5 * p0, 0.5 * p0) - direct = response_direct_hfx( - gradient_driver, - tdobj, - densities, - response_terms, - atmlst=atmlst, - jk_ledger=jk_ledger, - output_slot=direct_slot, - ) - direct += response_xc.direct - direct += fockz_xc.direct - direct += fockz_hfx.direct - def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): - local = spin_fock_direct_dft( - driver, - obj, - p_alpha, - p_beta, - atmlst=atmlst, - nobeta_p0=p0, - jk_ledger=jk_ledger, - output_slots=(direct_slot, zvector_slot), - ) - contractions = jk_ledger.contract( - driver, - obj.mol, - atmlst, - slots=(direct_slot, zvector_slot), - ) - local[0] += contractions[direct_slot] - local[1] += contractions[zvector_slot] - return local - - # 4-5. ROKS transpose-Hessian adjoint, Dz Fock derivative, and Pulay term. - return finish_gradient( - gradient_driver, - tdobj, - m_matrix, - direct, - atmlst, - tolerance, - max_cycle, - fock_direct, - direct_fock_probes=direct_fock_probes, + return assemble_gradient( + gradient_driver, tdobj, channel_data, (p0, pz), + same_spin_fock_q(tdobj, xy), + same_spin_hfx_projection_q(tdobj, xy), + atmlst=atmlst, tolerance=tolerance, max_cycle=max_cycle, ) diff --git a/src/nest/grad/nttda/xc.py b/src/nest/grad/nttda/xc.py index 8d5c83e..82f39e3 100644 --- a/src/nest/grad/nttda/xc.py +++ b/src/nest/grad/nttda/xc.py @@ -865,351 +865,10 @@ def lda_nobeta_reference_q(tdobj, p0, max_memory=None): return q_alpha, q_beta -# GGA quadrature - -def _gga_fref_kref(mf, rho0): - fxc, kxc = mf._numint.eval_xc_eff( - mf.xc, (rho0, rho0), deriv=3, xctype="GGA", spin=1, - )[2:4] - fref = 0.5 * ( - fxc[0, :, 0] - fxc[0, :, 1] - - fxc[1, :, 0] + fxc[1, :, 1] - ) - kref_alpha = 0.5 * ( - kxc[0, :, 0, :, 0] - kxc[0, :, 1, :, 0] - - kxc[1, :, 0, :, 0] + kxc[1, :, 1, :, 0] - ) - kref_beta = 0.5 * ( - kxc[0, :, 0, :, 1] - kxc[0, :, 1, :, 1] - - kxc[1, :, 0, :, 1] + kxc[1, :, 1, :, 1] - ) - return fref, kref_alpha, kref_beta - - -def gga_response_terms( - gradient_driver, tdobj, channel_data, atmlst=None, - with_direct=True): - """GGA ``vref0/vref1`` M matrix and fixed-grid skeleton derivative.""" - mf = tdobj._scf - mol = mf.mol - ni = mf._numint - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - _spaces, _amplitudes, densities, blocks, terms = channel_data - pair_labels = tuple( - label for label in densities - if any( - term.vref1 and label in (term.target, term.source) - for term in terms - ) - ) - pair_density_stack = np.asarray([ - densities[label] for label in pair_labels - ]) - nao = mol.nao_nr() - potentials = {label: np.zeros((nao, nao)) for label in densities} - reference_alpha = np.zeros((nao, nao)) - reference_beta = np.zeros_like(reference_alpha) - direct = np.zeros((len(atmlst), 3)) - mo = np.asarray(mf.mo_coeff) - density_alpha, density_beta = _reference_spin_densities(tdobj) - density_labels, density_stack = _response_density_stack( - densities, density_alpha, density_beta, - ) - offsets = mol.offset_nr_by_atom() - sparse = sparse_context(mf) - - for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): - rho0 = ni.eval_rho2( - mol, ao, mo, mf.mo_occ, mask, "GGA", with_lapl=False, - ) * 0.5 - fref, kref_alpha, kref_beta = _gga_fref_kref(mf, rho0) - rho = { - label: ni.eval_rho( - mol, ao, density, mask, "GGA", hermi=0, - with_lapl=False, - ) - for label, density in densities.items() - } - pair_values, contracted_pair_ao = pair_feature_batches( - ao, pair_density_stack, - ) - pairs = dict(zip(pair_labels, pair_values)) - pair_potentials = { - label: gga_pair_potential(fref, pairs[label]) - for label in pair_labels - } - ordinary_weights = { - label: np.zeros((4, weights.size)) for label in densities - } - special_weights = { - label: np.zeros((4, 4, weights.size)) for label in pair_labels - } - reference_weights_alpha = np.zeros((4, weights.size)) - reference_weights_beta = np.zeros_like(reference_weights_alpha) - - for term in terms: - if term.vref0: - ordinary_weights[term.target] += term.vref0 * lib.einsum( - "xyg,yg->xg", fref, rho[term.source], - ) - ordinary_weights[term.source] += term.vref0 * lib.einsum( - "xyg,xg->yg", fref, rho[term.target], - ) - pair = term.vref0 * lib.einsum( - "xg,yg->xyg", rho[term.target], rho[term.source], - ) - reference_weights_alpha += lib.einsum( - "xyg,xyzg->zg", pair, kref_alpha, - ) - reference_weights_beta += lib.einsum( - "xyg,xyzg->zg", pair, kref_beta, - ) - if term.vref1: - special_weights[term.target] += ( - term.vref1 * pair_potentials[term.source] - ) - special_weights[term.source] += ( - term.vref1 * pair_potentials[term.target] - ) - pair = term.vref1 * gga_pair_kernel_cross( - pairs[term.target], pairs[term.source], - ) - reference_weights_alpha += lib.einsum( - "xyg,xyzg->zg", pair, kref_alpha, - ) - reference_weights_beta += lib.einsum( - "xyg,xyzg->zg", pair, kref_beta, - ) - ordinary_weight_stack = np.asarray([ - ordinary_weights[label] for label in density_labels - ]) - special_weight_stack = np.asarray([ - special_weights[label] for label in pair_labels - ]) - - for label in potentials: - add_gga_matrix( - mol, potentials[label], ao, - ordinary_weights[label] * weights, mask, sparse, - ) - for label in pair_labels: - potentials[label] += pair_matrix( - mol, ao, mask, special_weights[label] * weights, sparse, - ) - reference_alpha += gga_eval_matrix( - mol, ao, reference_weights_alpha * weights, mask, - ) - reference_beta += gga_eval_matrix( - mol, ao, reference_weights_beta * weights, mask, - ) - - if not with_direct: - continue - for k, atom in enumerate(atmlst): - p0, p1 = offsets[atom][2:] - for xyz in range(3): - drho, drho_alpha, drho_beta, ao_delta = ( - _response_density_derivatives( - ao, density_stack, density_labels, - p0, p1, xyz, "GGA", - ) - ) - drho_stack = np.asarray([ - drho[label] for label in density_labels - ]) - value = lib.einsum( - "nfg,nfg,g->", - ordinary_weight_stack, drho_stack, weights, - ) - value += lib.einsum( - "fg,fg,g->", - reference_weights_alpha, drho_alpha, weights, - ) - value += lib.einsum( - "fg,fg,g->", - reference_weights_beta, drho_beta, weights, - ) - value += contract_pair_feature_derivatives( - ao, pair_density_stack, ao_delta, - contracted_pair_ao, p0, p1, - special_weight_stack, weights, - ) - direct[k, xyz] += value - - q_alpha, q_beta = _project_channel_potentials( - tdobj, potentials, blocks, - ) - _add_reference_q( - tdobj, q_alpha, q_beta, reference_alpha, reference_beta, - ) - return XCGradientTerms(q_alpha, q_beta, direct) - - -def gga_fockz_terms( - gradient_driver, tdobj, spaces, pz, atmlst=None, - with_direct=True): - """GGA derivative of ``Pz:Fz`` excluding its explicit Pz projection.""" - mf = tdobj._scf - mol = mf.mol - ni = mf._numint - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - density_open = spaces.c_open @ spaces.c_open.T - pz = 0.5 * (np.asarray(pz) + np.asarray(pz).T) - nao = mol.nao_nr() - open_potential = np.zeros((nao, nao)) - reference_alpha = np.zeros((nao, nao)) - reference_beta = np.zeros_like(reference_alpha) - direct = np.zeros((len(atmlst), 3)) - mo = np.asarray(mf.mo_coeff) - density_alpha, density_beta = _reference_spin_densities(tdobj) - density_stack = np.asarray(( - pz, density_open, density_alpha, density_beta, - )) - offsets = mol.offset_nr_by_atom() - sparse = sparse_context(mf) - - for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): - rho0 = ni.eval_rho2( - mol, ao, mo, mf.mo_occ, mask, "GGA", with_lapl=False, - ) * 0.5 - fref, kref_alpha, kref_beta = _gga_fref_kref(mf, rho0) - rho_pz = ni.eval_rho( - mol, ao, pz, mask, "GGA", hermi=1, with_lapl=False, - ) - rho_open = ni.eval_rho( - mol, ao, density_open, mask, "GGA", hermi=1, - with_lapl=False, - ) - add_gga_matrix( - mol, - open_potential, - ao, - 0.5 * lib.einsum("xyg,yg->xg", fref, rho_pz) * weights, - mask, - sparse, - ) - pair = 0.5 * lib.einsum("xg,yg->xyg", rho_pz, rho_open) - reference_alpha += gga_eval_matrix( - mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_alpha) * weights, - mask, - ) - reference_beta += gga_eval_matrix( - mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_beta) * weights, - mask, - ) - if not with_direct: - continue - for k, atom in enumerate(atmlst): - p0, p1 = offsets[atom][2:] - derivative_batches = _hermitian_density_derivative_batches( - ao, density_stack, p0, p1, "GGA", - ) - for xyz, derivatives in enumerate(derivative_batches): - drho_pz, drho_open, drho_alpha, drho_beta = derivatives - direct[k, xyz] += 0.5 * lib.einsum( - "xg,xyg,yg,g->", drho_pz, fref, rho_open, weights, - ) - direct[k, xyz] += 0.5 * lib.einsum( - "xg,xyg,yg,g->", rho_pz, fref, drho_open, weights, - ) - direct[k, xyz] += lib.einsum( - "xyg,xyzg,zg,g->", - pair, kref_alpha, drho_alpha, weights, - ) - direct[k, xyz] += lib.einsum( - "xyg,xyzg,zg,g->", - pair, kref_beta, drho_beta, weights, - ) - - q_alpha = np.zeros((mo.shape[1], mo.shape[1])) - q_beta = np.zeros_like(q_alpha) - q_alpha[:, spaces.open] += ( - mo.conj().T @ (open_potential + open_potential.T) @ spaces.c_open - ) - _add_reference_q( - tdobj, q_alpha, q_beta, reference_alpha, reference_beta, - ) - return XCGradientTerms(q_alpha, q_beta, direct) - - -def gga_nobeta_reference_q(tdobj, p0, max_memory=None): - """Reference-density response of the GGA equal-spin common Fock.""" - mf = tdobj._scf - mo = np.asarray(mf.mo_coeff) - q_alpha = np.zeros((mo.shape[1], mo.shape[1])) - q_beta = np.zeros_like(q_alpha) - if not tdobj.nobeta or getattr(mf, "is_average_occupation_reference", False): - return q_alpha, q_beta - if max_memory is None: - max_memory = tdobj.max_memory - ni = mf._numint - mol = mf.mol - density_alpha, density_beta = _reference_spin_densities(tdobj) - density0 = 0.5 * (density_alpha + density_beta) - p0 = 0.5 * (np.asarray(p0) + np.asarray(p0).T) - matrix_alpha = np.zeros((mol.nao_nr(), mol.nao_nr())) - matrix_beta = np.zeros_like(matrix_alpha) - for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, mol.nao_nr(), 2, max_memory=max_memory): - rho_p = ni.eval_rho( - mol, ao, p0, mask, "GGA", hermi=1, with_lapl=False, - ) - rho_alpha = ni.eval_rho( - mol, ao, density_alpha, mask, "GGA", hermi=1, - with_lapl=False, - ) - rho_beta = ni.eval_rho( - mol, ao, density_beta, mask, "GGA", hermi=1, - with_lapl=False, - ) - rho_equal = ni.eval_rho( - mol, ao, density0, mask, "GGA", hermi=1, with_lapl=False, - ) - fxc_actual = ni.eval_xc_eff( - mf.xc, (rho_alpha, rho_beta), deriv=2, - xctype="GGA", spin=1, - )[2] - fxc_equal = ni.eval_xc_eff( - mf.xc, (rho_equal, rho_equal), deriv=2, - xctype="GGA", spin=1, - )[2] - equal = 0.25 * ( - fxc_equal[0, :, 0] + fxc_equal[0, :, 1] - + fxc_equal[1, :, 0] + fxc_equal[1, :, 1] - ) - actual_alpha = 0.5 * ( - fxc_actual[0, :, 0] + fxc_actual[1, :, 0] - ) - actual_beta = 0.5 * ( - fxc_actual[0, :, 1] + fxc_actual[1, :, 1] - ) - matrix_alpha += gga_eval_matrix( - mol, - ao, - lib.einsum("xg,xzg->zg", rho_p, equal - actual_alpha) * weights, - mask, - ) - matrix_beta += gga_eval_matrix( - mol, - ao, - lib.einsum("xg,xzg->zg", rho_p, equal - actual_beta) * weights, - mask, - ) - _add_reference_q(tdobj, q_alpha, q_beta, matrix_alpha, matrix_beta) - return q_alpha, q_beta - - -# meta-GGA quadrature - -def _mgga_fref_kref(mf, rho0): +# GGA/meta-GGA quadrature +def _semilocal_fref_kref(mf, rho0, xctype): fxc, kxc = mf._numint.eval_xc_eff( - mf.xc, (rho0, rho0), deriv=3, xctype="MGGA", spin=1, + mf.xc, (rho0, rho0), deriv=3, xctype=xctype, spin=1, )[2:4] fref = 0.5 * ( fxc[0, :, 0] - fxc[0, :, 1] @@ -1226,11 +885,20 @@ def _mgga_fref_kref(mf, rho0): return fref, kref_alpha, kref_beta -def mgga_response_terms( +def semilocal_response_terms( gradient_driver, tdobj, channel_data, atmlst=None, with_direct=True): - """MGGA ``vref0/vref1`` M matrix and fixed-grid skeleton derivative.""" + """GGA/meta-GGA ``vref0/vref1`` M matrix and fixed-grid skeleton derivative.""" mf = tdobj._scf + xctype = mf._numint._xc_type(mf.xc) + if xctype == "GGA": + nvar, pair_potential, pair_cross = 4, gga_pair_potential, gga_pair_kernel_cross + add_matrix, eval_matrix = add_gga_matrix, gga_eval_matrix + elif xctype == "MGGA": + nvar, pair_potential, pair_cross = 5, mgga_pair_potential, mgga_pair_kernel_cross + add_matrix, eval_matrix = add_mgga_matrix, mgga_eval_matrix + else: + raise NotImplementedError("Semilocal response requires GGA or MGGA") mol = mf.mol ni = mf._numint if atmlst is None: @@ -1263,12 +931,12 @@ def mgga_response_terms( for ao, mask, weights, _coords in ni.block_loop( mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): rho0 = ni.eval_rho2( - mol, ao, mo, mf.mo_occ, mask, "MGGA", with_lapl=False, + mol, ao, mo, mf.mo_occ, mask, xctype, with_lapl=False, ) * 0.5 - fref, kref_alpha, kref_beta = _mgga_fref_kref(mf, rho0) + fref, kref_alpha, kref_beta = _semilocal_fref_kref(mf, rho0, xctype) rho = { label: ni.eval_rho( - mol, ao, density, mask, "MGGA", hermi=0, + mol, ao, density, mask, xctype, hermi=0, with_lapl=False, ) for label, density in densities.items() @@ -1278,16 +946,16 @@ def mgga_response_terms( ) pairs = dict(zip(pair_labels, pair_values)) pair_potentials = { - label: mgga_pair_potential(fref, pairs[label]) + label: pair_potential(fref, pairs[label]) for label in pair_labels } ordinary_weights = { - label: np.zeros((5, weights.size)) for label in densities + label: np.zeros((nvar, weights.size)) for label in densities } special_weights = { label: np.zeros((4, 4, weights.size)) for label in pair_labels } - reference_weights_alpha = np.zeros((5, weights.size)) + reference_weights_alpha = np.zeros((nvar, weights.size)) reference_weights_beta = np.zeros_like(reference_weights_alpha) for term in terms: @@ -1314,7 +982,7 @@ def mgga_response_terms( special_weights[term.source] += ( term.vref1 * pair_potentials[term.target] ) - pair = term.vref1 * mgga_pair_kernel_cross( + pair = term.vref1 * pair_cross( pairs[term.target], pairs[term.source], ) reference_weights_alpha += lib.einsum( @@ -1331,7 +999,7 @@ def mgga_response_terms( ]) for label in potentials: - add_mgga_matrix( + add_matrix( mol, potentials[label], ao, ordinary_weights[label] * weights, mask, sparse, ) @@ -1339,10 +1007,10 @@ def mgga_response_terms( potentials[label] += pair_matrix( mol, ao, mask, special_weights[label] * weights, sparse, ) - reference_alpha += mgga_eval_matrix( + reference_alpha += eval_matrix( mol, ao, reference_weights_alpha * weights, mask, ) - reference_beta += mgga_eval_matrix( + reference_beta += eval_matrix( mol, ao, reference_weights_beta * weights, mask, ) @@ -1354,7 +1022,7 @@ def mgga_response_terms( drho, drho_alpha, drho_beta, ao_delta = ( _response_density_derivatives( ao, density_stack, density_labels, - p0, p1, xyz, "MGGA", + p0, p1, xyz, xctype, ) ) drho_stack = np.asarray([ @@ -1388,11 +1056,18 @@ def mgga_response_terms( return XCGradientTerms(q_alpha, q_beta, direct) -def mgga_fockz_terms( +def semilocal_fockz_terms( gradient_driver, tdobj, spaces, pz, atmlst=None, with_direct=True): - """MGGA derivative of ``Pz:Fz`` excluding its explicit Pz projection.""" + """GGA/meta-GGA derivative of ``Pz:Fz`` excluding its explicit Pz projection.""" mf = tdobj._scf + xctype = mf._numint._xc_type(mf.xc) + if xctype == "GGA": + add_matrix, eval_matrix = add_gga_matrix, gga_eval_matrix + elif xctype == "MGGA": + add_matrix, eval_matrix = add_mgga_matrix, mgga_eval_matrix + else: + raise NotImplementedError("Semilocal response requires GGA or MGGA") mol = mf.mol ni = mf._numint if atmlst is None: @@ -1416,17 +1091,17 @@ def mgga_fockz_terms( for ao, mask, weights, _coords in ni.block_loop( mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): rho0 = ni.eval_rho2( - mol, ao, mo, mf.mo_occ, mask, "MGGA", with_lapl=False, + mol, ao, mo, mf.mo_occ, mask, xctype, with_lapl=False, ) * 0.5 - fref, kref_alpha, kref_beta = _mgga_fref_kref(mf, rho0) + fref, kref_alpha, kref_beta = _semilocal_fref_kref(mf, rho0, xctype) rho_pz = ni.eval_rho( - mol, ao, pz, mask, "MGGA", hermi=1, with_lapl=False, + mol, ao, pz, mask, xctype, hermi=1, with_lapl=False, ) rho_open = ni.eval_rho( - mol, ao, density_open, mask, "MGGA", hermi=1, + mol, ao, density_open, mask, xctype, hermi=1, with_lapl=False, ) - add_mgga_matrix( + add_matrix( mol, open_potential, ao, @@ -1435,11 +1110,11 @@ def mgga_fockz_terms( sparse, ) pair = 0.5 * lib.einsum("xg,yg->xyg", rho_pz, rho_open) - reference_alpha += mgga_eval_matrix( + reference_alpha += eval_matrix( mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_alpha) * weights, mask, ) - reference_beta += mgga_eval_matrix( + reference_beta += eval_matrix( mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_beta) * weights, mask, ) @@ -1448,7 +1123,7 @@ def mgga_fockz_terms( for k, atom in enumerate(atmlst): p0, p1 = offsets[atom][2:] derivative_batches = _hermitian_density_derivative_batches( - ao, density_stack, p0, p1, "MGGA", + ao, density_stack, p0, p1, xctype, ) for xyz, derivatives in enumerate(derivative_batches): drho_pz, drho_open, drho_alpha, drho_beta = derivatives @@ -1478,9 +1153,16 @@ def mgga_fockz_terms( return XCGradientTerms(q_alpha, q_beta, direct) -def mgga_nobeta_reference_q(tdobj, p0, max_memory=None): - """Reference-density response of the MGGA equal-spin common Fock.""" +def semilocal_nobeta_reference_q(tdobj, p0, max_memory=None): + """Reference-density response of the GGA/meta-GGA equal-spin common Fock.""" mf = tdobj._scf + xctype = mf._numint._xc_type(mf.xc) + if xctype == "GGA": + eval_matrix = gga_eval_matrix + elif xctype == "MGGA": + eval_matrix = mgga_eval_matrix + else: + raise NotImplementedError("Semilocal response requires GGA or MGGA") mo = np.asarray(mf.mo_coeff) q_alpha = np.zeros((mo.shape[1], mo.shape[1])) q_beta = np.zeros_like(q_alpha) @@ -1498,26 +1180,26 @@ def mgga_nobeta_reference_q(tdobj, p0, max_memory=None): for ao, mask, weights, _coords in ni.block_loop( mol, mf.grids, mol.nao_nr(), 2, max_memory=max_memory): rho_p = ni.eval_rho( - mol, ao, p0, mask, "MGGA", hermi=1, with_lapl=False, + mol, ao, p0, mask, xctype, hermi=1, with_lapl=False, ) rho_alpha = ni.eval_rho( - mol, ao, density_alpha, mask, "MGGA", hermi=1, + mol, ao, density_alpha, mask, xctype, hermi=1, with_lapl=False, ) rho_beta = ni.eval_rho( - mol, ao, density_beta, mask, "MGGA", hermi=1, + mol, ao, density_beta, mask, xctype, hermi=1, with_lapl=False, ) rho_equal = ni.eval_rho( - mol, ao, density0, mask, "MGGA", hermi=1, with_lapl=False, + mol, ao, density0, mask, xctype, hermi=1, with_lapl=False, ) fxc_actual = ni.eval_xc_eff( mf.xc, (rho_alpha, rho_beta), deriv=2, - xctype="MGGA", spin=1, + xctype=xctype, spin=1, )[2] fxc_equal = ni.eval_xc_eff( mf.xc, (rho_equal, rho_equal), deriv=2, - xctype="MGGA", spin=1, + xctype=xctype, spin=1, )[2] equal = 0.25 * ( fxc_equal[0, :, 0] + fxc_equal[0, :, 1] @@ -1529,13 +1211,13 @@ def mgga_nobeta_reference_q(tdobj, p0, max_memory=None): actual_beta = 0.5 * ( fxc_actual[0, :, 1] + fxc_actual[1, :, 1] ) - matrix_alpha += mgga_eval_matrix( + matrix_alpha += eval_matrix( mol, ao, lib.einsum("xg,xzg->zg", rho_p, equal - actual_alpha) * weights, mask, ) - matrix_beta += mgga_eval_matrix( + matrix_beta += eval_matrix( mol, ao, lib.einsum("xg,xzg->zg", rho_p, equal - actual_beta) * weights, From 56d007a58fc846c5f450fb58784ea68fce03083f Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Tue, 22 Sep 2026 14:28:51 +0800 Subject: [PATCH 11/16] fix: preserve Dz0SCF convergence and tighten transition regression Retain the branch's PySCF response solver: the shared NEST solver fails the existing small-system Dz0SCF convergence checks. Explicitly converge oscillator-strength amplitudes to 1e-9 with lindep=1e-18 rather than relying on the default 1e-5 residual; reference values and assertion tolerances are unchanged. Seven targeted tests and eight subtests pass. --- src/nest/nttda/nttda.py | 2 +- src/nest/nttda/tests/test_nttda_oscillator_strength.py | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index 5fe1bbb..fa06e72 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -27,7 +27,7 @@ from pyscf.dft.gen_grid import NBINS from pyscf import __config__ from pyscf.dft.numint import _scale_ao_sparse, _dot_ao_ao_sparse, _dot_ao_dm_sparse, _contract_rho_sparse -from nest._lr_eig import eigh as lr_eigh +from pyscf.tdscf._lr_eig import eigh as lr_eigh from pyscf import symm from pyscf.data import nist diff --git a/src/nest/nttda/tests/test_nttda_oscillator_strength.py b/src/nest/nttda/tests/test_nttda_oscillator_strength.py index d9983da..81f6be8 100644 --- a/src/nest/nttda/tests/test_nttda_oscillator_strength.py +++ b/src/nest/nttda/tests/test_nttda_oscillator_strength.py @@ -50,7 +50,11 @@ def test_svwn_nttda_oscillator_strength(self): [-0.068297131072, -0.365336383706, 0.031212437760], ]) ref_f = np.array([0.007681628385, 0.037787591904, 0.026266335987]) - td = mf.NTTDA().set(deltaS=-1, nstates=4).run() + # Converge amplitudes more tightly than the observable assertions. + # The default residual tolerance (1e-5) leaves dipole outer products + # sensitive to the eigensolver's trial-space ordering. Retain small + # independent corrections when requesting the tighter residual. + td = mf.NTTDA().set(deltaS=-1, nstates=4, conv_tol=1e-9, lindep=1e-18).run() self.assertTrue(np.all(td.converged)) dip = td.transition_dipole() dip_outer = np.einsum('nx,ny->nxy', dip.conj(), dip) @@ -64,7 +68,7 @@ def test_svwn_nttda_oscillator_strength(self): [-0.056333525366, 0.071326673082, -0.457181696110], ]) ref_f = np.array([0.067496526009, 0.003926438694, 0.012770953752]) - td = mf.NTTDA().set(deltaS=0, nstates=4).run() + td = mf.NTTDA().set(deltaS=0, nstates=4, conv_tol=1e-9, lindep=1e-18).run() self.assertTrue(np.all(td.converged)) dip = td.transition_dipole() dip_outer = np.einsum('nx,ny->nxy', dip.conj(), dip) @@ -78,7 +82,7 @@ def test_svwn_nttda_oscillator_strength(self): [-1.078796933744, -0.328635627721, -0.013042401171], ]) ref_f = np.array([0.109988279396, 0.000658252112, 0.158300119889]) - td = mf.NTTDA().set(deltaS=1, nstates=4).run() + td = mf.NTTDA().set(deltaS=1, nstates=4, conv_tol=1e-9, lindep=1e-18).run() self.assertTrue(np.all(td.converged)) dip = td.transition_dipole() dip_outer = np.einsum('nx,ny->nxy', dip.conj(), dip) From 1295a50487db92e764a897527540b2028f84a37e Mon Sep 17 00:00:00 2001 From: jingjjkk <2501110377@stu.pku.edu.cn> Date: Tue, 22 Sep 2026 15:01:59 +0800 Subject: [PATCH 12/16] refactor: consolidate NTTDA projections and XC gradient loops Build channel projections once per gradient and reuse the shared transition-potential projection. Share adjoint assembly and Krylov mechanics while retaining reference-specific Hessians and preconditioners. Use a singleton density-feature axis for LDA to share the GGA/MGGA quadrature loops, retaining distinct pair kernels. Consolidate HF/DFT Fock derivatives and remove unreferenced channel projection wrappers. Net reduction: 604 implementation lines; existing numerical tests and tolerances are unchanged. --- src/nest/grad/nttda/common.py | 242 +++++++------ src/nest/grad/nttda/delta_s_minus_one.py | 121 +------ src/nest/grad/nttda/delta_s_zero.py | 157 ++------- src/nest/grad/nttda/ensemble.py | 79 +---- src/nest/grad/nttda/roks.py | 104 ++---- src/nest/grad/nttda/xc.py | 429 +++++------------------ 6 files changed, 264 insertions(+), 868 deletions(-) diff --git a/src/nest/grad/nttda/common.py b/src/nest/grad/nttda/common.py index 5f37211..0a51154 100644 --- a/src/nest/grad/nttda/common.py +++ b/src/nest/grad/nttda/common.py @@ -11,6 +11,7 @@ from nest.nttda import nttda as nttda_mod from . import xc as xc_backend +from .xc import _reference_spin_densities from .roks import finish_gradient @@ -45,22 +46,12 @@ def assemble_gradient( else (0.5 * p0, 0.5 * p0) ) if xctype != "HF": - if xctype == "LDA": - response_builder = xc_backend.lda_response_terms - fockz_builder = xc_backend.lda_fockz_terms - nobeta_builder = xc_backend.lda_nobeta_reference_q - elif xctype in ("GGA", "MGGA"): - response_builder = xc_backend.semilocal_response_terms - fockz_builder = xc_backend.semilocal_fockz_terms - nobeta_builder = xc_backend.semilocal_nobeta_reference_q - else: - raise NotImplementedError("NTTDA gradient does not support XC type %s" % xctype) for terms in ( - response_builder(gradient_driver, tdobj, channel_data, atmlst=atmlst), - fockz_builder(gradient_driver, tdobj, spaces, pz, atmlst=atmlst)): + xc_backend.response_terms(gradient_driver, tdobj, channel_data, atmlst=atmlst), + xc_backend.fockz_terms(gradient_driver, tdobj, spaces, pz, atmlst=atmlst)): m_matrix += terms.q_alpha + terms.q_beta direct += terms.direct - common_alpha, common_beta = nobeta_builder(tdobj, p0) + common_alpha, common_beta = xc_backend.nobeta_reference_q(tdobj, p0) m_matrix += common_alpha + common_beta if not spin_fock: terms = fockz_hfx_terms( @@ -71,16 +62,10 @@ def assemble_gradient( direct += terms.direct def fock_direct(driver, obj, p_alpha, p_beta, atmlst=None): - if xctype == "HF": - local = spin_fock_direct_hf( - driver, obj, p_alpha, p_beta, atmlst=atmlst, - jk_ledger=ledger, output_slots=slots, - ) - else: - local = spin_fock_direct_dft( - driver, obj, p_alpha, p_beta, atmlst=atmlst, nobeta_p0=p0, - jk_ledger=ledger, output_slots=slots, - ) + local = spin_fock_direct( + driver, obj, p_alpha, p_beta, atmlst=atmlst, nobeta_p0=p0, + jk_ledger=ledger, output_slots=slots, + ) contractions = ledger.contract(driver, obj.mol, atmlst, slots=slots) for index, slot in enumerate(slots): local[index] += contractions[slot] @@ -228,10 +213,7 @@ def _apply_reference_responses(tdobj, densities, max_memory=None): None, None, fxc_ref, max_memory=max_memory, ) if xctype == "LDA": - vref1 = ni.nr_rks_fxc( - mol, mf.grids, mf.xc, None, dms, 0, 0, - None, None, fxc_ref, max_memory=max_memory, - ) + vref1 = vref0.copy() elif xctype == "GGA": vref1 = nttda_mod.nr_rks_fxc1_gga( ni, mol, mf.grids, mf.xc, dms, fxc_ref, @@ -476,17 +458,6 @@ def _contract_derivative_terms( gradients[term.slot][k] += term.scale * value -def _reference_spin_densities(tdobj): - mf = tdobj._scf - if getattr(mf, "is_average_occupation_reference", False): - return tuple(np.asarray(dm) for dm in mf.make_rdm1s()) - mo = np.asarray(mf.mo_coeff) - return ( - mo[:, mf.mo_occ > 0] @ mo[:, mf.mo_occ > 0].T, - mo[:, mf.mo_occ == 2] @ mo[:, mf.mo_occ == 2].T, - ) - - def _spin_probe_stacks(p_alpha, p_beta): p_alpha = np.asarray(p_alpha) p_beta = np.asarray(p_beta) @@ -497,10 +468,10 @@ def _spin_probe_stacks(p_alpha, p_beta): return p_alpha, p_beta, single_probe -def spin_fock_direct_dft( +def spin_fock_direct( gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, nobeta_p0=None, jk_ledger=None, output_slots=None): - """Differentiate one or more ordinary UKS Fock scalar probes. + """Differentiate one or more HF/DFT spin Fock scalar probes. The optional ``nobeta_p0`` correction belongs to the first, explicit-direct probe in the batch. @@ -547,93 +518,48 @@ def spin_fock_direct_dft( ledger.add("j", output_slots[probe], j_terms) ledger.add("k", output_slots[probe], k_terms) xctype = ni._xc_type(mf.xc) - if xctype == "LDA": - derivative_contractor = xc_backend.contract_lda_vxc_derivative - elif xctype == "GGA": - derivative_contractor = xc_backend.contract_gga_vxc_derivative - elif xctype == "MGGA": - derivative_contractor = xc_backend.contract_mgga_vxc_derivative - else: - raise NotImplementedError( - "ordinary Fock direct derivative is not implemented for %s" % - xctype - ) - if (nobeta_p0 is not None and tdobj.nobeta - and not getattr(mf, "is_average_occupation_reference", False)): - density0 = 0.5 * (density_alpha + density_beta) - actual_probe_alpha = np.array(p_alpha, copy=True) - actual_probe_beta = np.array(p_beta, copy=True) - actual_probe_alpha[0] -= 0.5 * nobeta_p0 - actual_probe_beta[0] -= 0.5 * nobeta_p0 - else: - density0 = None - actual_probe_alpha = p_alpha - actual_probe_beta = p_beta - gradient += derivative_contractor( - mf, - density_alpha, - density_beta, - actual_probe_alpha, - actual_probe_beta, - atmlst=atmlst, - max_memory=gradient_driver.max_memory, - ) - if density0 is not None: - gradient[0] += derivative_contractor( + if xctype != "HF": + if xctype == "LDA": + derivative_contractor = xc_backend.contract_lda_vxc_derivative + elif xctype == "GGA": + derivative_contractor = xc_backend.contract_gga_vxc_derivative + elif xctype == "MGGA": + derivative_contractor = xc_backend.contract_mgga_vxc_derivative + else: + raise NotImplementedError( + "ordinary Fock direct derivative is not implemented for %s" % + xctype + ) + if (nobeta_p0 is not None and tdobj.nobeta + and not getattr(mf, "is_average_occupation_reference", False)): + density0 = 0.5 * (density_alpha + density_beta) + actual_probe_alpha = np.array(p_alpha, copy=True) + actual_probe_beta = np.array(p_beta, copy=True) + actual_probe_alpha[0] -= 0.5 * nobeta_p0 + actual_probe_beta[0] -= 0.5 * nobeta_p0 + else: + density0 = None + actual_probe_alpha = p_alpha + actual_probe_beta = p_beta + gradient += derivative_contractor( mf, - density0, - density0, - 0.5 * nobeta_p0, - 0.5 * nobeta_p0, + density_alpha, + density_beta, + actual_probe_alpha, + actual_probe_beta, atmlst=atmlst, max_memory=gradient_driver.max_memory, ) - if jk_ledger is None: - contractions = local_ledger.contract( - gradient_driver, mol, atmlst, slots=output_slots, - ) - for probe, slot in enumerate(output_slots): - gradient[probe] += contractions[slot] - return gradient[0] if single_probe else gradient - - -def spin_fock_direct_hf( - gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, - jk_ledger=None, output_slots=None): - """Differentiate one or more spin-resolved HF Fock scalar probes.""" - mol = tdobj.mol - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - p_alpha, p_beta, single_probe = _spin_probe_stacks( - p_alpha, p_beta, - ) - if output_slots is None: - output_slots = tuple(range(len(p_alpha))) - p_total = p_alpha + p_beta - dm_alpha, dm_beta = _reference_spin_densities(tdobj) - gradient = np.zeros((len(p_alpha), len(atmlst), 3)) - - hcore_derivative = rhf_grad.Gradients(tdobj._scf).hcore_generator(mol) - for k, atom in enumerate(atmlst): - gradient[:, k] += lib.einsum( - "npq,xpq->nx", p_total, hcore_derivative(atom), - ) - local_ledger = _JKDerivativeLedger() - ledger = jk_ledger if jk_ledger is not None else local_ledger - for probe in range(len(p_alpha)): - ledger.add( - "j", output_slots[probe], ( - (p_total[probe], dm_alpha, 1.0, None), - (p_total[probe], dm_beta, 1.0, None), - ), - ) - ledger.add( - "k", output_slots[probe], ( - (p_alpha[probe], dm_alpha, -1.0, None), - (p_beta[probe], dm_beta, -1.0, None), - ), - ) + if density0 is not None: + gradient[0] += derivative_contractor( + mf, + density0, + density0, + 0.5 * nobeta_p0, + 0.5 * nobeta_p0, + atmlst=atmlst, + max_memory=gradient_driver.max_memory, + ) if jk_ledger is None: contractions = local_ledger.contract( gradient_driver, mol, atmlst, slots=output_slots, @@ -745,3 +671,73 @@ def fockz_hfx_terms( gradient_driver, mol, atmlst, slots=(output_slot,), )[output_slot] return xc_backend.XCGradientTerms(q_alpha, q_beta, direct) + + +def fock_probes(tdobj, projections): + """AO probes of the explicit F0/Fz scalar for either spin channel.""" + p0 = np.zeros((tdobj.mol.nao_nr(), tdobj.mol.nao_nr())) + pz = np.zeros_like(p0) + for term in projections: + density = term.density() + p0 += term.weight_f0 * density + pz += term.weight_fz * density + return p0, pz + + +def fock_projection_q(tdobj, projections, operators, probes): + """Differentiate the Fock projections and the reference density. + + HF spin probes include the full Fz response. DFT adds Fz and nobeta + corrections during gradient assembly, after the common F0 response. + """ + mf = tdobj._scf + mo = np.asarray(mf.mo_coeff) + fock0, fockz = (mo.conj().T @ operator @ mo for operator in operators) + q_alpha = np.zeros((mo.shape[1], mo.shape[1])) + q_beta = np.zeros_like(q_alpha) + is_hf = mf._numint._xc_type(mf.xc) == "HF" + for term in projections: + left, right = term.left_indices, term.right_indices + coefficient = term.coefficient + def project(target, operator, scale): + if scale: + target[:, left] += scale * operator[:, right] @ coefficient.T + target[:, right] += scale * operator[:, left] @ coefficient + project(q_alpha, fock0, 0.5 * term.weight_f0) + project(q_beta, fock0, 0.5 * term.weight_f0) + if is_hf: + project(q_alpha, fockz, 0.5 * term.weight_fz) + project(q_beta, fockz, 0.5 * term.weight_fz) + else: + project(q_alpha, fockz, term.weight_fz) + p0, pz = probes + p_alpha, p_beta = 0.5 * p0, 0.5 * p0 + if is_hf: + p_alpha = p_alpha + 0.5 * pz + p_beta = p_beta - 0.5 * pz + response_alpha, response_beta = _fock_response_q(tdobj, p_alpha, p_beta) + return q_alpha + response_alpha, q_beta + response_beta + + +def response_potentials(densities, vref0, vref1, terms): + """Vary both transition-density factors of the response scalar.""" + potentials = {label: np.zeros_like(dm) for label, dm in densities.items()} + for term in terms: + if term.vref0: + potentials[term.target] += term.vref0 * vref0[term.source] + potentials[term.source] += term.vref0 * vref0[term.target] + if term.vref1: + potentials[term.target] += term.vref1 * vref1[term.source] + potentials[term.source] += term.vref1 * vref1[term.target] + return potentials + + +def response_projection_q(tdobj, channel_data, max_memory=None, hfx_only=False): + """MO derivative of the response at fixed kernels, for either channel.""" + _spaces, _amplitudes, densities, blocks, terms = channel_data + if hfx_only: + vref0, vref1 = _apply_hfx_responses(tdobj, densities) + else: + vref0, vref1 = _apply_reference_responses(tdobj, densities, max_memory) + potentials = response_potentials(densities, vref0, vref1, terms) + return xc_backend._project_channel_potentials(tdobj, potentials, blocks) diff --git a/src/nest/grad/nttda/delta_s_minus_one.py b/src/nest/grad/nttda/delta_s_minus_one.py index c0fcc84..526491d 100644 --- a/src/nest/grad/nttda/delta_s_minus_one.py +++ b/src/nest/grad/nttda/delta_s_minus_one.py @@ -17,10 +17,12 @@ orbital_spaces, pair_density, FockProjection, - _fock_response_q, + fock_probes, + fock_projection_q, + response_projection_q, ResponseTerm, _apply_reference_responses, - _apply_hfx_responses, + ) @@ -231,15 +233,8 @@ def add(name, left_label, coefficient, right_label, f0, fz): def spin_lowering_fock_probes(tdobj, xy): - """Return AO probes ``P0,Pz`` for the lowering Fock ledger.""" - nao = tdobj.mol.nao_nr() - p0 = np.zeros((nao, nao)) - pz = np.zeros_like(p0) - for term in spin_lowering_fock_projections(tdobj, xy): - density = term.density() - p0 += term.weight_f0 * density - pz += term.weight_fz * density - return p0, pz + """AO probes of the channel's explicit Fock scalar.""" + return fock_probes(tdobj, spin_lowering_fock_projections(tdobj, xy)) def spin_lowering_fock_scalar(tdobj, xy, max_memory=None): @@ -290,101 +285,6 @@ def spin_lowering_action_scalar(tdobj, xy): return float(np.vdot(vector, action).real) -def _response_potentials(densities, vref0, vref1, terms): - potentials = {label: np.zeros_like(dm) for label, dm in densities.items()} - for term in terms: - if term.vref0: - potentials[term.target] += term.vref0 * vref0[term.source] - potentials[term.source] += term.vref0 * vref0[term.target] - if term.vref1: - potentials[term.target] += term.vref1 * vref1[term.source] - potentials[term.source] += term.vref1 * vref1[term.target] - return potentials - - -def _project_transition_potentials(tdobj, blocks, potentials): - mo = np.asarray(tdobj._scf.mo_coeff) - q_alpha = np.zeros((mo.shape[1], mo.shape[1])) - q_beta = np.zeros_like(q_alpha) - for label, (target, source, coefficient) in blocks.items(): - potential = mo.conj().T @ potentials[label] @ mo - q_beta[:, target] += potential[:, source] @ coefficient.T - q_alpha[:, source] += potential[target, :].T @ coefficient - return q_alpha, q_beta - - -def spin_lowering_response_projection_q( - tdobj, xy, max_memory=None, hfx_only=False): - """Transition-factor derivative of the lowering response scalar.""" - spaces, amplitudes, densities = spin_lowering_transition_densities( - tdobj, xy, - ) - blocks = spin_lowering_block_data(spaces, amplitudes) - if hfx_only: - vref0, vref1 = _apply_hfx_responses(tdobj, densities) - else: - vref0, vref1 = _apply_reference_responses( - tdobj, densities, max_memory=max_memory, - ) - potentials = _response_potentials( - densities, - vref0, - vref1, - spin_lowering_response_terms(spaces.spin), - ) - return _project_transition_potentials(tdobj, blocks, potentials) - - -def spin_lowering_fock_q(tdobj, xy, max_memory=None): - """Explicit-Fock projection and reference-density response M matrices.""" - mf = tdobj._scf - mo = np.asarray(mf.mo_coeff) - nmo = mo.shape[1] - fock0, fockz = spin_lowering_fock0_fockz( - tdobj, max_memory=max_memory, - ) - fock0_mo = mo.conj().T @ fock0 @ mo - fockz_mo = mo.conj().T @ fockz @ mo - q_alpha = np.zeros((nmo, nmo)) - q_beta = np.zeros_like(q_alpha) - is_hf = mf._numint._xc_type(mf.xc) == "HF" - - for term in spin_lowering_fock_projections(tdobj, xy): - left = term.left_indices - right = term.right_indices - coefficient = term.coefficient - - def project(target, operator, scale): - if scale: - target[:, left] += ( - scale * operator[:, right] @ coefficient.T - ) - target[:, right] += ( - scale * operator[:, left] @ coefficient - ) - - project(q_alpha, fock0_mo, 0.5 * term.weight_f0) - project(q_beta, fock0_mo, 0.5 * term.weight_f0) - if is_hf: - project(q_alpha, fockz_mo, 0.5 * term.weight_fz) - project(q_beta, fockz_mo, 0.5 * term.weight_fz) - else: - project(q_alpha, fockz_mo, term.weight_fz) - - p0, pz = spin_lowering_fock_probes(tdobj, xy) - p_alpha = 0.5 * p0 - p_beta = 0.5 * p0 - if is_hf: - p_alpha = p_alpha + 0.5 * pz - p_beta = p_beta - 0.5 * pz - response_alpha, response_beta = _fock_response_q( - tdobj, p_alpha, p_beta, - ) - q_alpha += response_alpha - q_beta += response_beta - return q_alpha, q_beta - - # Channel assembly def grad_elec( @@ -399,10 +299,11 @@ def grad_elec( blocks = spin_lowering_block_data(spaces, amplitudes) response_terms = spin_lowering_response_terms(spaces.spin) channel_data = (spaces, amplitudes, densities, blocks, response_terms) - p0, pz = spin_lowering_fock_probes(tdobj, xy) + projections = spin_lowering_fock_projections(tdobj, xy) + probes = fock_probes(tdobj, projections) return assemble_gradient( - gradient_driver, tdobj, channel_data, (p0, pz), - spin_lowering_fock_q(tdobj, xy), - spin_lowering_response_projection_q(tdobj, xy, hfx_only=True), + gradient_driver, tdobj, channel_data, probes, + fock_projection_q(tdobj, projections, spin_lowering_fock0_fockz(tdobj), probes), + response_projection_q(tdobj, channel_data, hfx_only=True), atmlst=atmlst, tolerance=tolerance, max_cycle=max_cycle, ) diff --git a/src/nest/grad/nttda/delta_s_zero.py b/src/nest/grad/nttda/delta_s_zero.py index 60d4998..764f918 100644 --- a/src/nest/grad/nttda/delta_s_zero.py +++ b/src/nest/grad/nttda/delta_s_zero.py @@ -18,10 +18,11 @@ orbital_spaces, pair_density, FockProjection, - _fock_response_q, + fock_probes, + fock_projection_q, + response_projection_q, ResponseTerm, _apply_reference_responses, - _apply_hfx_responses, ) @@ -86,6 +87,16 @@ def split_same_spin(tdobj, xy): ) +def same_spin_block_data(spaces, amplitudes): + """MO index/factor map for variations of same-spin transition densities.""" + return { + "CO": (spaces.open, spaces.closed, amplitudes.co.T), + "CV": (spaces.virtual, spaces.closed, amplitudes.cv.T), + "OV": (spaces.virtual, spaces.open, amplitudes.ov.T), + "CV0": (spaces.virtual, spaces.closed, amplitudes.cv0.T), + } + + def same_spin_transition_densities(tdobj, xy): """Return directed AO transition densities for the four response blocks.""" spaces, amp = split_same_spin(tdobj, xy) @@ -182,64 +193,8 @@ def add(name, left, coefficient, right, f0, fz): def same_spin_fock_probes(tdobj, xy): - """Return AO probes ``(P0, Pz)`` generated from the Fock ledger.""" - nao = tdobj.mol.nao_nr() - p0 = np.zeros((nao, nao)) - pz = np.zeros_like(p0) - for term in same_spin_fock_projections(tdobj, xy): - density = term.density() - p0 += term.weight_f0 * density - pz += term.weight_fz * density - return p0, pz - - -def same_spin_fock_q(tdobj, xy, max_memory=None): - """Return the explicit-Fock contribution to ``(Q_alpha,Q_beta)``. - - For HF the complete ``Fz`` response is exactly represented by the - spin-resolved probes. DFT callers add the independent ``Fz`` and - ``nobeta`` response ledgers after this common ``F0`` contribution. - """ - mf = tdobj._scf - mo = np.asarray(mf.mo_coeff) - nmo = mo.shape[1] - fock0, fockz = fock0_fockz(tdobj, max_memory=max_memory) - fock0_mo = mo.conj().T @ fock0 @ mo - fockz_mo = mo.conj().T @ fockz @ mo - q_alpha = np.zeros((nmo, nmo)) - q_beta = np.zeros_like(q_alpha) - - for term in same_spin_fock_projections(tdobj, xy): - left = term.left_indices - right = term.right_indices - coeff = term.coefficient - - def project(target, operator, scale): - if scale == 0.0: - return - target[:, left] += scale * operator[:, right] @ coeff.T - target[:, right] += scale * operator[:, left] @ coeff - - project(q_alpha, fock0_mo, 0.5 * term.weight_f0) - project(q_beta, fock0_mo, 0.5 * term.weight_f0) - if mf._numint._xc_type(mf.xc) == "HF": - project(q_alpha, fockz_mo, 0.5 * term.weight_fz) - project(q_beta, fockz_mo, 0.5 * term.weight_fz) - else: - project(q_alpha, fockz_mo, term.weight_fz) - - p0, pz = same_spin_fock_probes(tdobj, xy) - p_alpha = 0.5 * p0 - p_beta = 0.5 * p0 - if mf._numint._xc_type(mf.xc) == "HF": - p_alpha = p_alpha + 0.5 * pz - p_beta = p_beta - 0.5 * pz - response_alpha, response_beta = _fock_response_q( - tdobj, p_alpha, p_beta, - ) - q_alpha += response_alpha - q_beta += response_beta - return q_alpha, q_beta + """AO probes of the channel's explicit Fock scalar.""" + return fock_probes(tdobj, same_spin_fock_projections(tdobj, xy)) def same_spin_fock_scalar(tdobj, xy, max_memory=None): @@ -283,18 +238,6 @@ def same_spin_response_terms(spin): ) -def _derivative_potentials(spaces, densities, vref0, vref1): - potentials = {label: np.zeros_like(dm) for label, dm in densities.items()} - for term in same_spin_response_terms(spaces.spin): - if term.vref0: - potentials[term.target] += term.vref0 * vref0[term.source] - potentials[term.source] += term.vref0 * vref0[term.target] - if term.vref1: - potentials[term.target] += term.vref1 * vref1[term.source] - potentials[term.source] += term.vref1 * vref1[term.target] - return potentials - - def same_spin_response_scalar(tdobj, xy, max_memory=None): """Evaluate all current-NTTDA response terms in ``X.T A_sc X``.""" spaces, _amp, densities = same_spin_transition_densities(tdobj, xy) @@ -315,60 +258,6 @@ def same_spin_response_scalar(tdobj, xy, max_memory=None): return float(value) -def same_spin_response_derivative_potentials(tdobj, xy, max_memory=None): - """AO potentials obtained by varying both sides of the response scalar.""" - spaces, _amp, densities = same_spin_transition_densities(tdobj, xy) - vref0, vref1 = _apply_reference_responses( - tdobj, densities, max_memory=max_memory, - ) - return _derivative_potentials(spaces, densities, vref0, vref1) - - -def same_spin_response_projection_q(tdobj, xy, max_memory=None): - """MO derivative from transition-density factors at frozen kernels.""" - spaces, amp, _densities = same_spin_transition_densities(tdobj, xy) - potentials = same_spin_response_derivative_potentials( - tdobj, xy, max_memory=max_memory, - ) - mo = np.asarray(tdobj._scf.mo_coeff) - q_alpha = np.zeros((mo.shape[1], mo.shape[1])) - q_beta = np.zeros_like(q_alpha) - block_data = { - "CO": (spaces.open, spaces.closed, amp.co.T), - "CV": (spaces.virtual, spaces.closed, amp.cv.T), - "OV": (spaces.virtual, spaces.open, amp.ov.T), - "CV0": (spaces.virtual, spaces.closed, amp.cv0.T), - } - for label, (target, source, coefficient) in block_data.items(): - potential_mo = mo.conj().T @ potentials[label] @ mo - q_beta[:, target] += potential_mo[:, source] @ coefficient.T - q_alpha[:, source] += potential_mo[target, :].T @ coefficient - return q_alpha, q_beta - - -def same_spin_hfx_projection_q(tdobj, xy): - """Transition-factor derivative of only the hybrid/RSH response scalar.""" - spaces, amp, densities = same_spin_transition_densities(tdobj, xy) - vref0, vref1 = _apply_hfx_responses(tdobj, densities) - potentials = _derivative_potentials( - spaces, densities, vref0, vref1, - ) - mo = np.asarray(tdobj._scf.mo_coeff) - q_alpha = np.zeros((mo.shape[1], mo.shape[1])) - q_beta = np.zeros_like(q_alpha) - block_data = { - "CO": (spaces.open, spaces.closed, amp.co.T), - "CV": (spaces.virtual, spaces.closed, amp.cv.T), - "OV": (spaces.virtual, spaces.open, amp.ov.T), - "CV0": (spaces.virtual, spaces.closed, amp.cv0.T), - } - for label, (target, source, coefficient) in block_data.items(): - potential_mo = mo.conj().T @ potentials[label] @ mo - q_beta[:, target] += potential_mo[:, source] @ coefficient.T - q_alpha[:, source] += potential_mo[target, :].T @ coefficient - return q_alpha, q_beta - - # Scalar closure diagnostics (private to this channel) def same_spin_action_scalar(tdobj, xy): @@ -398,18 +287,14 @@ def grad_elec( if tdobj.deltaS != 0: raise ValueError("deltaS=0 gradient received a different spin channel") spaces, amplitudes, densities = same_spin_transition_densities(tdobj, xy) - blocks = { - "CO": (spaces.open, spaces.closed, amplitudes.co.T), - "CV": (spaces.virtual, spaces.closed, amplitudes.cv.T), - "OV": (spaces.virtual, spaces.open, amplitudes.ov.T), - "CV0": (spaces.virtual, spaces.closed, amplitudes.cv0.T), - } + blocks = same_spin_block_data(spaces, amplitudes) response_terms = same_spin_response_terms(spaces.spin) channel_data = (spaces, amplitudes, densities, blocks, response_terms) - p0, pz = same_spin_fock_probes(tdobj, xy) + projections = same_spin_fock_projections(tdobj, xy) + probes = fock_probes(tdobj, projections) return assemble_gradient( - gradient_driver, tdobj, channel_data, (p0, pz), - same_spin_fock_q(tdobj, xy), - same_spin_hfx_projection_q(tdobj, xy), + gradient_driver, tdobj, channel_data, probes, + fock_projection_q(tdobj, projections, fock0_fockz(tdobj), probes), + response_projection_q(tdobj, channel_data, hfx_only=True), atmlst=atmlst, tolerance=tolerance, max_cycle=max_cycle, ) diff --git a/src/nest/grad/nttda/ensemble.py b/src/nest/grad/nttda/ensemble.py index 6129b8b..fe62e98 100644 --- a/src/nest/grad/nttda/ensemble.py +++ b/src/nest/grad/nttda/ensemble.py @@ -2,10 +2,9 @@ import numpy as np -from pyscf import lib from pyscf.scf import hf -from .roks import GradientComponents, _orbital_gradient, pack_m_matrix +from .roks import finish_gradient, pack_m_matrix, _solve_zvector def canonical_pairs(tdobj): @@ -114,78 +113,10 @@ def _preconditioner(tdobj, pairs): return diagonal -def solve_zvector(action, pairs, tdobj, rhs, tolerance=1e-12, - max_cycle=None): - """Solve the average-occupation orbital adjoint equation.""" - diagonal = _preconditioner(tdobj, pairs) - initial = rhs / diagonal - if max_cycle is None: - max_cycle = len(rhs) - - def operator(vector): - vector = np.asarray(vector) - if vector.ndim == 1: - return action(vector) / diagonal - vector - return np.asarray([action(row) / diagonal - row for row in vector]) - - solution = lib.krylov( - operator, - initial, - tol=tolerance, - max_cycle=max_cycle, - lindep=1e-22, - hermi=False, - verbose=0, - ) - return np.asarray(solution).reshape(-1) - - -def finish_gradient( - gradient_driver, tdobj, m_matrix, direct, atmlst, - tolerance, max_cycle, fock_direct, direct_fock_probes=None): - """Solve the Dz0SCF average-occupation adjoint and assemble ``d omega / dR``.""" - transpose_action, pairs = make_hessian_transpose_action(tdobj) - rhs = pack_m_matrix(m_matrix, pairs) - zvector = solve_zvector( - transpose_action, - pairs, - tdobj, - rhs, - tolerance=tolerance, - max_cycle=max_cycle, - ) - adjoint = zvector_adjoint_matrix(tdobj, pairs, zvector) - residual = float(np.max(np.abs(pack_m_matrix(adjoint, pairs) - rhs))) - probe_alpha, probe_beta = zvector_probe_densities( - tdobj, pairs, zvector, - ) - if direct_fock_probes is None: - fock_contraction = fock_direct( - gradient_driver, tdobj, probe_alpha, probe_beta, atmlst=atmlst, - ) - direct_total = direct - else: - direct_alpha, direct_beta = direct_fock_probes - fock_contractions = fock_direct( - gradient_driver, - tdobj, - np.asarray((direct_alpha, probe_alpha)), - np.asarray((direct_beta, probe_beta)), - atmlst=atmlst, - ) - direct_total = direct + fock_contractions[0] - fock_contraction = fock_contractions[1] - orbital = _orbital_gradient( - tdobj, m_matrix, adjoint, fock_contraction, atmlst=atmlst, - ) - return GradientComponents( - m_matrix=m_matrix, - direct=direct_total, - orbital=orbital, - total=direct_total + orbital, - zvector=zvector, - residual=residual, - ) +def solve_zvector(action, pairs, tdobj, rhs, tolerance=1e-12, max_cycle=None): + """Solve the average-occupation adjoint with its occupation-weighted diagonal.""" + return _solve_zvector(action, _preconditioner(tdobj, pairs), rhs, + tolerance, max_cycle) __all__ = [ diff --git a/src/nest/grad/nttda/roks.py b/src/nest/grad/nttda/roks.py index 82cb445..d485317 100644 --- a/src/nest/grad/nttda/roks.py +++ b/src/nest/grad/nttda/roks.py @@ -30,21 +30,12 @@ def finish_gradient( contraction is the first result and the Z-vector contraction is the second. """ if getattr(tdobj._scf, "is_average_occupation_reference", False): - from .ensemble import finish_gradient as finish_ensemble_gradient - return finish_ensemble_gradient( - gradient_driver, - tdobj, - m_matrix, - direct, - atmlst, - tolerance, - max_cycle, - fock_direct, - direct_fock_probes=direct_fock_probes, - ) - transpose_action, pairs = make_hessian_transpose_action(tdobj) + from . import ensemble as response + else: + from . import roks as response + transpose_action, pairs = response.make_hessian_transpose_action(tdobj) rhs = pack_m_matrix(m_matrix, pairs) - zvector = solve_zvector( + zvector = response.solve_zvector( transpose_action, pairs, tdobj, @@ -52,9 +43,9 @@ def finish_gradient( tolerance=tolerance, max_cycle=max_cycle, ) - adjoint = zvector_adjoint_matrix(tdobj, pairs, zvector) + adjoint = response.zvector_adjoint_matrix(tdobj, pairs, zvector) residual = float(np.max(np.abs(pack_m_matrix(adjoint, pairs) - rhs))) - probe_alpha, probe_beta = zvector_probe_densities( + probe_alpha, probe_beta = response.zvector_probe_densities( tdobj, pairs, zvector, ) if direct_fock_probes is None: @@ -128,39 +119,18 @@ def _response_reference(mf): return mf -def make_hessian_transpose_action(tdobj, pairs=None): - """Return a matrix-free action for the transpose ROKS Hessian.""" +def _make_adjoint_action(tdobj, pairs): + """Cache reference operators for the full ROKS adjoint matrix.""" mf = tdobj._scf mo = np.asarray(mf.mo_coeff) occ = np.asarray(mf.mo_occ) - nmo = mo.shape[1] - if pairs is None: - pairs = canonical_pairs(tdobj, compact=True) fock_alpha, fock_beta = _spin_focks_mo(mf) occupation_alpha = (occ > 0).astype(float) occupation_beta = (occ == 2).astype(float) response = _response_reference(mf).gen_response(hermi=1) - def unpack(vector): - source_alpha = np.zeros((nmo, nmo)) - source_beta = np.zeros_like(source_alpha) - for value, (p, q, name) in zip(vector, pairs): - if name in ("cc", "oo", "vv"): - source_alpha[p, q] += 0.5 * value - source_beta[p, q] += 0.5 * value - elif name == "co": - source_beta[p, q] += value - elif name == "cv": - source_alpha[p, q] += value - source_beta[p, q] += value - elif name == "ov": - source_alpha[p, q] += value - else: - raise ValueError("unknown ROKS pair type %s" % name) - return source_alpha, source_beta - def apply_one(vector): - source_alpha, source_beta = unpack(vector) + source_alpha, source_beta = _unpack_zvector_source(tdobj, pairs, vector) gradient = fock_alpha @ (source_alpha + source_alpha.T) gradient += fock_beta @ (source_beta + source_beta.T) density_alpha = mo @ source_alpha @ mo.conj().T @@ -176,16 +146,22 @@ def apply_one(vector): gradient += potential_alpha.T * occupation_alpha[None, :] gradient += potential_beta * occupation_beta[None, :] gradient += potential_beta.T * occupation_beta[None, :] - return np.asarray([ - gradient[p, q] - gradient[q, p] - for p, q, _name in pairs - ]) + return gradient + + return apply_one + + +def make_hessian_transpose_action(tdobj, pairs=None): + """Return a matrix-free action for the transpose ROKS Hessian.""" + if pairs is None: + pairs = canonical_pairs(tdobj, compact=True) + adjoint = _make_adjoint_action(tdobj, pairs) def apply(vector): vector = np.asarray(vector) if vector.ndim == 1: - return apply_one(vector) - return np.asarray([apply_one(row) for row in vector]) + return pack_m_matrix(adjoint(vector), pairs) + return np.asarray([pack_m_matrix(adjoint(row), pairs) for row in vector]) return apply, pairs @@ -222,9 +198,8 @@ def _preconditioner(tdobj, pairs): return diagonal -def solve_zvector(action, pairs, tdobj, rhs, tolerance=1e-12, max_cycle=None): - """Solve ``H.T z = rhs`` using the PySCF CPHF Krylov pattern.""" - diagonal = _preconditioner(tdobj, pairs) +def _solve_zvector(action, diagonal, rhs, tolerance=1e-12, max_cycle=None): + """Shared preconditioned Krylov solve, with a reference-specific diagonal.""" initial = rhs / diagonal if max_cycle is None: max_cycle = len(rhs) @@ -246,6 +221,11 @@ def operator(vector): ) return np.asarray(solution).reshape(-1) +def solve_zvector(action, pairs, tdobj, rhs, tolerance=1e-12, max_cycle=None): + """Solve the ROKS adjoint with its spin-resolved preconditioner.""" + return _solve_zvector(action, _preconditioner(tdobj, pairs), rhs, + tolerance, max_cycle) + def _unpack_zvector_source(tdobj, pairs, zvector): nmo = tdobj._scf.mo_coeff.shape[1] @@ -269,31 +249,7 @@ def _unpack_zvector_source(tdobj, pairs, zvector): def zvector_adjoint_matrix(tdobj, pairs, zvector): """Full MO adjoint matrix satisfying ``z.H(kappa)=Tr(G.T kappa)``.""" - mf = tdobj._scf - mo = np.asarray(mf.mo_coeff) - occ = np.asarray(mf.mo_occ) - fock_alpha, fock_beta = _spin_focks_mo(mf) - occupation_alpha = (occ > 0).astype(float) - occupation_beta = (occ == 2).astype(float) - source_alpha, source_beta = _unpack_zvector_source( - tdobj, pairs, zvector, - ) - gradient = fock_alpha @ (source_alpha + source_alpha.T) - gradient += fock_beta @ (source_beta + source_beta.T) - density_alpha = mo @ source_alpha @ mo.conj().T - density_beta = mo @ source_beta @ mo.conj().T - density_alpha = 0.5 * (density_alpha + density_alpha.T) - density_beta = 0.5 * (density_beta + density_beta.T) - potential_alpha, potential_beta = _response_reference(mf).gen_response( - hermi=1, - )(np.asarray((density_alpha, density_beta))) - potential_alpha = mo.conj().T @ potential_alpha @ mo - potential_beta = mo.conj().T @ potential_beta @ mo - gradient += potential_alpha * occupation_alpha[None, :] - gradient += potential_alpha.T * occupation_alpha[None, :] - gradient += potential_beta * occupation_beta[None, :] - gradient += potential_beta.T * occupation_beta[None, :] - return gradient + return _make_adjoint_action(tdobj, pairs)(zvector) def zvector_probe_densities(tdobj, pairs, zvector): diff --git a/src/nest/grad/nttda/xc.py b/src/nest/grad/nttda/xc.py index 82f39e3..97bf844 100644 --- a/src/nest/grad/nttda/xc.py +++ b/src/nest/grad/nttda/xc.py @@ -14,7 +14,6 @@ from pyscf.grad import tdrks as tdrks_grad - # Shared result and projection helpers @dataclass(frozen=True) @@ -294,32 +293,6 @@ def mgga_eval_matrix(mol, ao, weights, mask): return output[0] -# LDA quadrature - -def _lda_fref_kref(mf, ao0, mask): - ni = mf._numint - rho0 = ni.eval_rho2( - mf.mol, ao0, mf.mo_coeff, mf.mo_occ, mask, "LDA", - with_lapl=False, - ) * 0.5 - fxc, kxc = ni.eval_xc_eff( - mf.xc, (rho0, rho0), deriv=3, xctype="LDA", spin=1, - )[2:4] - fref = 0.5 * ( - fxc[0, 0, 0, 0] - fxc[0, 0, 1, 0] - - fxc[1, 0, 0, 0] + fxc[1, 0, 1, 0] - ) - kref_alpha = 0.5 * ( - kxc[0, 0, 0, 0, 0, 0] - kxc[0, 0, 1, 0, 0, 0] - - kxc[1, 0, 0, 0, 0, 0] + kxc[1, 0, 1, 0, 0, 0] - ) - kref_beta = 0.5 * ( - kxc[0, 0, 0, 0, 1, 0] - kxc[0, 0, 1, 0, 1, 0] - - kxc[1, 0, 0, 0, 1, 0] + kxc[1, 0, 1, 0, 1, 0] - ) - return fref, kref_alpha, kref_beta - - def _lda_matrix(ao0, weights): return ao0.T @ (ao0 * np.asarray(weights)[:, None]) @@ -385,10 +358,10 @@ def _spin_probe_stacks(probe_alpha, probe_beta): return probe_alpha, probe_beta, single_probe -def _xc_density(ni, mol, ao, density, mask, xctype): +def _xc_density(ni, mol, ao, density, mask, xctype, hermi=1): ao_values = ao[0] if xctype == "LDA" else ao rho = ni.eval_rho( - mol, ao_values, density, mask, xctype, hermi=1, + mol, ao_values, density, mask, xctype, hermi=hermi, with_lapl=False, ) return rho[None] if rho.ndim == 1 else rho @@ -501,8 +474,6 @@ def _response_density_derivatives( derivatives = _xc_density_derivatives( ao, density_stack, p0, p1, xctype, ao_center_derivative, ) - if xctype == "LDA": - derivatives = derivatives[:, 0] channel_count = len(labels) channel_derivatives = dict(zip( labels, derivatives[:channel_count], @@ -623,250 +594,8 @@ def contract_mgga_vxc_derivative( ) -def lda_response_terms( - gradient_driver, tdobj, channel_data, atmlst=None, - with_direct=True): - """Analytic LDA M/direct terms for the ``vref0/vref1`` scalar.""" - mf = tdobj._scf - mol = mf.mol - ni = mf._numint - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - _spaces, _amplitudes, densities, blocks, response_terms = channel_data - coefficients = tuple( - (term.target, term.source, term.vref0 + term.vref1) - for term in response_terms - if term.vref0 + term.vref1 - ) - nao = mol.nao_nr() - potentials = {label: np.zeros((nao, nao)) for label in densities} - reference_alpha = np.zeros((nao, nao)) - reference_beta = np.zeros_like(reference_alpha) - direct = np.zeros((len(atmlst), 3)) - density_alpha, density_beta = _reference_spin_densities(tdobj) - density_labels, density_stack = _response_density_stack( - densities, density_alpha, density_beta, - ) - offsets = mol.offset_nr_by_atom() - - for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, nao, 1, max_memory=gradient_driver.max_memory): - ao0 = ao[0] - fref, kref_alpha, kref_beta = _lda_fref_kref(mf, ao0, mask) - rho = { - label: ni.eval_rho( - mol, ao0, density, mask, "LDA", hermi=0, - with_lapl=False, - ) - for label, density in densities.items() - } - potential_weights = { - label: np.zeros_like(weights) for label in densities - } - pair_alpha = np.zeros_like(weights) - pair_beta = np.zeros_like(weights) - for target, source, coefficient in coefficients: - potential_weights[target] += coefficient * fref * rho[source] - potential_weights[source] += coefficient * fref * rho[target] - pair = coefficient * rho[target] * rho[source] - pair_alpha += kref_alpha * pair - pair_beta += kref_beta * pair - potential_weight_stack = np.asarray([ - potential_weights[label] for label in density_labels - ]) - for label in potentials: - potentials[label] += _lda_matrix( - ao0, weights * potential_weights[label], - ) - reference_alpha += _lda_matrix(ao0, weights * pair_alpha) - reference_beta += _lda_matrix(ao0, weights * pair_beta) - - if not with_direct: - continue - for k, atom in enumerate(atmlst): - p0, p1 = offsets[atom][2:] - for xyz in range(3): - drho, drho_alpha, drho_beta, _ao_delta = ( - _response_density_derivatives( - ao, density_stack, density_labels, - p0, p1, xyz, "LDA", - ) - ) - drho_stack = np.asarray([ - drho[label] for label in density_labels - ]) - value = lib.einsum( - "ng,ng,g->", - potential_weight_stack, drho_stack, weights, - ) - value += lib.einsum( - "g,g,g->", pair_alpha, drho_alpha, weights, - ) - value += lib.einsum( - "g,g,g->", pair_beta, drho_beta, weights, - ) - direct[k, xyz] += value - - q_alpha, q_beta = _project_channel_potentials( - tdobj, potentials, blocks, - ) - _add_reference_q( - tdobj, q_alpha, q_beta, reference_alpha, reference_beta, - ) - return XCGradientTerms(q_alpha, q_beta, direct) - - -def lda_fockz_terms( - gradient_driver, tdobj, spaces, pz, atmlst=None, - with_direct=True): - """LDA response/direct derivative of ``Pz:Fz`` excluding Pz projection.""" - mf = tdobj._scf - mol = mf.mol - ni = mf._numint - if atmlst is None: - atmlst = range(mol.natm) - atmlst = tuple(atmlst) - density_open = spaces.c_open @ spaces.c_open.T - pz = np.asarray(pz) - pz_symmetric = 0.5 * (pz + pz.T) - nao = mol.nao_nr() - open_potential = np.zeros((nao, nao)) - reference_alpha = np.zeros((nao, nao)) - reference_beta = np.zeros_like(reference_alpha) - direct = np.zeros((len(atmlst), 3)) - mo = np.asarray(mf.mo_coeff) - density_alpha, density_beta = _reference_spin_densities(tdobj) - density_stack = np.asarray(( - pz_symmetric, density_open, density_alpha, density_beta, - )) - offsets = mol.offset_nr_by_atom() - - for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, nao, 1, max_memory=gradient_driver.max_memory): - ao0 = ao[0] - fref, kref_alpha, kref_beta = _lda_fref_kref(mf, ao0, mask) - rho_pz = ni.eval_rho( - mol, ao0, pz_symmetric, mask, "LDA", hermi=1, - with_lapl=False, - ) - rho_open = ni.eval_rho( - mol, ao0, density_open, mask, "LDA", hermi=1, - with_lapl=False, - ) - open_potential += _lda_matrix( - ao0, 0.5 * weights * fref * rho_pz, - ) - pair = 0.5 * rho_pz * rho_open - reference_alpha += _lda_matrix( - ao0, weights * kref_alpha * pair, - ) - reference_beta += _lda_matrix( - ao0, weights * kref_beta * pair, - ) - if not with_direct: - continue - for k, atom in enumerate(atmlst): - p0, p1 = offsets[atom][2:] - derivative_batches = _hermitian_density_derivative_batches( - ao, density_stack, p0, p1, "LDA", - ) - for xyz, derivatives in enumerate(derivative_batches): - drho_pz, drho_open, drho_alpha, drho_beta = ( - derivatives[:, 0] - ) - direct[k, xyz] += 0.5 * np.dot( - weights, - fref * ( - drho_pz * rho_open + rho_pz * drho_open - ) - + rho_pz * rho_open * ( - kref_alpha * drho_alpha - + kref_beta * drho_beta - ), - ) - - q_alpha = np.zeros((mo.shape[1], mo.shape[1])) - q_beta = np.zeros_like(q_alpha) - q_alpha[:, spaces.open] += ( - mo.conj().T @ (open_potential + open_potential.T) - @ spaces.c_open - ) - _add_reference_q( - tdobj, q_alpha, q_beta, reference_alpha, reference_beta, - ) - return XCGradientTerms(q_alpha, q_beta, direct) - - -def lda_nobeta_reference_q(tdobj, p0, max_memory=None): - """Reference-density correction for the equal-spin ``nobeta`` Fock.""" - mf = tdobj._scf - mo = np.asarray(mf.mo_coeff) - nmo = mo.shape[1] - q_alpha = np.zeros((nmo, nmo)) - q_beta = np.zeros_like(q_alpha) - if not tdobj.nobeta or getattr(mf, "is_average_occupation_reference", False): - return q_alpha, q_beta - if max_memory is None: - max_memory = tdobj.max_memory - ni = mf._numint - mol = mf.mol - nao = mol.nao_nr() - density_alpha, density_beta = _reference_spin_densities(tdobj) - density0 = 0.5 * (density_alpha + density_beta) - p0 = 0.5 * (np.asarray(p0) + np.asarray(p0).T) - matrix_alpha = np.zeros((nao, nao)) - matrix_beta = np.zeros_like(matrix_alpha) - for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, nao, 1, max_memory=max_memory): - ao0 = ao[0] - rho_p = ni.eval_rho( - mol, ao0, p0, mask, "LDA", hermi=1, with_lapl=False, - ) - rho_alpha = ni.eval_rho( - mol, ao0, density_alpha, mask, "LDA", hermi=1, - with_lapl=False, - ) - rho_beta = ni.eval_rho( - mol, ao0, density_beta, mask, "LDA", hermi=1, - with_lapl=False, - ) - rho0 = ni.eval_rho( - mol, ao0, density0, mask, "LDA", hermi=1, - with_lapl=False, - ) - fxc_actual = ni.eval_xc_eff( - mf.xc, (rho_alpha, rho_beta), deriv=2, - xctype="LDA", spin=1, - )[2] - fxc_equal = ni.eval_xc_eff( - mf.xc, (rho0, rho0), deriv=2, - xctype="LDA", spin=1, - )[2] - equal_derivative = 0.25 * ( - fxc_equal[0, 0, 0, 0] + fxc_equal[0, 0, 1, 0] - + fxc_equal[1, 0, 0, 0] + fxc_equal[1, 0, 1, 0] - ) - actual_alpha = 0.5 * ( - fxc_actual[0, 0, 0, 0] + fxc_actual[1, 0, 0, 0] - ) - actual_beta = 0.5 * ( - fxc_actual[0, 0, 1, 0] + fxc_actual[1, 0, 1, 0] - ) - matrix_alpha += _lda_matrix( - ao0, weights * rho_p * (equal_derivative - actual_alpha), - ) - matrix_beta += _lda_matrix( - ao0, weights * rho_p * (equal_derivative - actual_beta), - ) - _add_reference_q( - tdobj, q_alpha, q_beta, matrix_alpha, matrix_beta, - ) - return q_alpha, q_beta - - -# GGA/meta-GGA quadrature -def _semilocal_fref_kref(mf, rho0, xctype): +# XC quadrature +def _reference_fref_kref(mf, rho0, xctype): fxc, kxc = mf._numint.eval_xc_eff( mf.xc, (rho0, rho0), deriv=3, xctype=xctype, spin=1, )[2:4] @@ -885,20 +614,19 @@ def _semilocal_fref_kref(mf, rho0, xctype): return fref, kref_alpha, kref_beta -def semilocal_response_terms( +def response_terms( gradient_driver, tdobj, channel_data, atmlst=None, with_direct=True): - """GGA/meta-GGA ``vref0/vref1`` M matrix and fixed-grid skeleton derivative.""" + """LDA/GGA/meta-GGA ``vref0/vref1`` M matrix and fixed-grid skeleton derivative.""" mf = tdobj._scf xctype = mf._numint._xc_type(mf.xc) - if xctype == "GGA": + add_matrix, eval_matrix = _xc_matrix_builders(xctype) + if xctype == "LDA": + nvar = 1 + elif xctype == "GGA": nvar, pair_potential, pair_cross = 4, gga_pair_potential, gga_pair_kernel_cross - add_matrix, eval_matrix = add_gga_matrix, gga_eval_matrix - elif xctype == "MGGA": - nvar, pair_potential, pair_cross = 5, mgga_pair_potential, mgga_pair_kernel_cross - add_matrix, eval_matrix = add_mgga_matrix, mgga_eval_matrix else: - raise NotImplementedError("Semilocal response requires GGA or MGGA") + nvar, pair_potential, pair_cross = 5, mgga_pair_potential, mgga_pair_kernel_cross mol = mf.mol ni = mf._numint if atmlst is None: @@ -907,7 +635,7 @@ def semilocal_response_terms( _spaces, _amplitudes, densities, blocks, terms = channel_data pair_labels = tuple( label for label in densities - if any( + if xctype != "LDA" and any( term.vref1 and label in (term.target, term.source) for term in terms ) @@ -926,29 +654,27 @@ def semilocal_response_terms( densities, density_alpha, density_beta, ) offsets = mol.offset_nr_by_atom() - sparse = sparse_context(mf) + sparse = sparse_context(mf) if xctype != "LDA" else None + ao_deriv = 1 if xctype == "LDA" else 2 for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): + mol, mf.grids, nao, ao_deriv, max_memory=gradient_driver.max_memory): rho0 = ni.eval_rho2( - mol, ao, mo, mf.mo_occ, mask, xctype, with_lapl=False, + mol, ao[0] if xctype == "LDA" else ao, mo, mf.mo_occ, + mask, xctype, with_lapl=False, ) * 0.5 - fref, kref_alpha, kref_beta = _semilocal_fref_kref(mf, rho0, xctype) + fref, kref_alpha, kref_beta = _reference_fref_kref(mf, rho0, xctype) rho = { - label: ni.eval_rho( - mol, ao, density, mask, xctype, hermi=0, - with_lapl=False, - ) + label: _xc_density(ni, mol, ao, density, mask, xctype, hermi=0) for label, density in densities.items() } - pair_values, contracted_pair_ao = pair_feature_batches( - ao, pair_density_stack, - ) - pairs = dict(zip(pair_labels, pair_values)) - pair_potentials = { - label: pair_potential(fref, pairs[label]) - for label in pair_labels - } + pairs, pair_potentials = {}, {} + if pair_labels: + pair_values, contracted_pair_ao = pair_feature_batches(ao, pair_density_stack) + pairs = dict(zip(pair_labels, pair_values)) + pair_potentials = { + label: pair_potential(fref, pairs[label]) for label in pair_labels + } ordinary_weights = { label: np.zeros((nvar, weights.size)) for label in densities } @@ -959,14 +685,16 @@ def semilocal_response_terms( reference_weights_beta = np.zeros_like(reference_weights_alpha) for term in terms: - if term.vref0: - ordinary_weights[term.target] += term.vref0 * lib.einsum( + # In LDA the two kernels coincide; no pair-feature correction remains. + ordinary_coefficient = term.vref0 + term.vref1 if xctype == "LDA" else term.vref0 + if ordinary_coefficient: + ordinary_weights[term.target] += ordinary_coefficient * lib.einsum( "xyg,yg->xg", fref, rho[term.source], ) - ordinary_weights[term.source] += term.vref0 * lib.einsum( + ordinary_weights[term.source] += ordinary_coefficient * lib.einsum( "xyg,xg->yg", fref, rho[term.target], ) - pair = term.vref0 * lib.einsum( + pair = ordinary_coefficient * lib.einsum( "xg,yg->xyg", rho[term.target], rho[term.source], ) reference_weights_alpha += lib.einsum( @@ -975,7 +703,7 @@ def semilocal_response_terms( reference_weights_beta += lib.einsum( "xyg,xyzg->zg", pair, kref_beta, ) - if term.vref1: + if xctype != "LDA" and term.vref1: special_weights[term.target] += ( term.vref1 * pair_potentials[term.source] ) @@ -1040,11 +768,12 @@ def semilocal_response_terms( "fg,fg,g->", reference_weights_beta, drho_beta, weights, ) - value += contract_pair_feature_derivatives( - ao, pair_density_stack, ao_delta, - contracted_pair_ao, p0, p1, - special_weight_stack, weights, - ) + if pair_labels: + value += contract_pair_feature_derivatives( + ao, pair_density_stack, ao_delta, + contracted_pair_ao, p0, p1, + special_weight_stack, weights, + ) direct[k, xyz] += value q_alpha, q_beta = _project_channel_potentials( @@ -1056,18 +785,33 @@ def semilocal_response_terms( return XCGradientTerms(q_alpha, q_beta, direct) -def semilocal_fockz_terms( +def _lda_eval_matrix(mol, ao, weights, mask): + """LDA potential with the same feature axis as GGA/meta-GGA.""" + return _lda_matrix(ao[0], weights[0]) + + +def _add_lda_matrix(mol, output, ao, weights, mask, sparse): + output += _lda_eval_matrix(mol, ao, weights, mask) + + +def _xc_matrix_builders(xctype): + """Feature-potential builders; LDA retains one density feature.""" + if xctype == "LDA": + return _add_lda_matrix, _lda_eval_matrix + if xctype == "GGA": + return add_gga_matrix, gga_eval_matrix + if xctype == "MGGA": + return add_mgga_matrix, mgga_eval_matrix + raise NotImplementedError("Unsupported XC type %s" % xctype) + + +def fockz_terms( gradient_driver, tdobj, spaces, pz, atmlst=None, with_direct=True): - """GGA/meta-GGA derivative of ``Pz:Fz`` excluding its explicit Pz projection.""" + """LDA/GGA/meta-GGA derivative of ``Pz:Fz`` excluding its explicit Pz projection.""" mf = tdobj._scf xctype = mf._numint._xc_type(mf.xc) - if xctype == "GGA": - add_matrix, eval_matrix = add_gga_matrix, gga_eval_matrix - elif xctype == "MGGA": - add_matrix, eval_matrix = add_mgga_matrix, mgga_eval_matrix - else: - raise NotImplementedError("Semilocal response requires GGA or MGGA") + add_matrix, eval_matrix = _xc_matrix_builders(xctype) mol = mf.mol ni = mf._numint if atmlst is None: @@ -1086,21 +830,18 @@ def semilocal_fockz_terms( pz, density_open, density_alpha, density_beta, )) offsets = mol.offset_nr_by_atom() - sparse = sparse_context(mf) + sparse = sparse_context(mf) if xctype != "LDA" else None + ao_deriv = 1 if xctype == "LDA" else 2 for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, nao, 2, max_memory=gradient_driver.max_memory): + mol, mf.grids, nao, ao_deriv, max_memory=gradient_driver.max_memory): rho0 = ni.eval_rho2( - mol, ao, mo, mf.mo_occ, mask, xctype, with_lapl=False, + mol, ao[0] if xctype == "LDA" else ao, mo, mf.mo_occ, + mask, xctype, with_lapl=False, ) * 0.5 - fref, kref_alpha, kref_beta = _semilocal_fref_kref(mf, rho0, xctype) - rho_pz = ni.eval_rho( - mol, ao, pz, mask, xctype, hermi=1, with_lapl=False, - ) - rho_open = ni.eval_rho( - mol, ao, density_open, mask, xctype, hermi=1, - with_lapl=False, - ) + fref, kref_alpha, kref_beta = _reference_fref_kref(mf, rho0, xctype) + rho_pz = _xc_density(ni, mol, ao, pz, mask, xctype) + rho_open = _xc_density(ni, mol, ao, density_open, mask, xctype) add_matrix( mol, open_potential, @@ -1153,16 +894,11 @@ def semilocal_fockz_terms( return XCGradientTerms(q_alpha, q_beta, direct) -def semilocal_nobeta_reference_q(tdobj, p0, max_memory=None): - """Reference-density response of the GGA/meta-GGA equal-spin common Fock.""" +def nobeta_reference_q(tdobj, p0, max_memory=None): + """Reference-density response of the LDA/GGA/meta-GGA equal-spin common Fock.""" mf = tdobj._scf xctype = mf._numint._xc_type(mf.xc) - if xctype == "GGA": - eval_matrix = gga_eval_matrix - elif xctype == "MGGA": - eval_matrix = mgga_eval_matrix - else: - raise NotImplementedError("Semilocal response requires GGA or MGGA") + _add_matrix, eval_matrix = _xc_matrix_builders(xctype) mo = np.asarray(mf.mo_coeff) q_alpha = np.zeros((mo.shape[1], mo.shape[1])) q_beta = np.zeros_like(q_alpha) @@ -1177,21 +913,12 @@ def semilocal_nobeta_reference_q(tdobj, p0, max_memory=None): p0 = 0.5 * (np.asarray(p0) + np.asarray(p0).T) matrix_alpha = np.zeros((mol.nao_nr(), mol.nao_nr())) matrix_beta = np.zeros_like(matrix_alpha) + ao_deriv = 1 if xctype == "LDA" else 2 for ao, mask, weights, _coords in ni.block_loop( - mol, mf.grids, mol.nao_nr(), 2, max_memory=max_memory): - rho_p = ni.eval_rho( - mol, ao, p0, mask, xctype, hermi=1, with_lapl=False, - ) - rho_alpha = ni.eval_rho( - mol, ao, density_alpha, mask, xctype, hermi=1, - with_lapl=False, - ) - rho_beta = ni.eval_rho( - mol, ao, density_beta, mask, xctype, hermi=1, - with_lapl=False, - ) - rho_equal = ni.eval_rho( - mol, ao, density0, mask, xctype, hermi=1, with_lapl=False, + mol, mf.grids, mol.nao_nr(), ao_deriv, max_memory=max_memory): + rho_p, rho_alpha, rho_beta, rho_equal = ( + _xc_density(ni, mol, ao, density, mask, xctype) + for density in (p0, density_alpha, density_beta, density0) ) fxc_actual = ni.eval_xc_eff( mf.xc, (rho_alpha, rho_beta), deriv=2, From 8d5cc716bfa634895d2e00479f983462925785b6 Mon Sep 17 00:00:00 2001 From: wtpeter Date: Sat, 26 Sep 2026 18:27:19 +0800 Subject: [PATCH 13/16] fix(nttda): restore NEST eigensolver and correct subspace expansion --- src/nest/_lr_eig.py | 2 +- src/nest/nttda/nttda.py | 2 +- .../tests/test_nttda_aocscf_reference.py | 8 ++++++++ src/nest/soc/tests/test_sftda_soc.py | 20 +++++++++++-------- 4 files changed, 22 insertions(+), 10 deletions(-) diff --git a/src/nest/_lr_eig.py b/src/nest/_lr_eig.py index 088dd9f..bb74718 100644 --- a/src/nest/_lr_eig.py +++ b/src/nest/_lr_eig.py @@ -189,7 +189,7 @@ def eigh(aop, x0, precond, tol_residual=1e-5, lindep=1e-12, nroots=1, space_inc = nroots else: # Adding too many trial bases in each iteration may cause larger errors - space_inc = min(nroots, min(MAX_SPACE_INC, x0_size//2)) + space_inc = max(nroots, min(MAX_SPACE_INC, x0_size//2)) max_space = int(max_memory*1e6/8/x0_size / 2 - nroots - space_inc) if max_space < nroots * 4 < x0_size: diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index 88bdfd8..1008cbc 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -27,7 +27,7 @@ from pyscf.dft.gen_grid import NBINS from pyscf import __config__ from pyscf.dft.numint import _scale_ao_sparse, _dot_ao_ao_sparse, _dot_ao_dm_sparse, _contract_rho_sparse -from pyscf.tdscf._lr_eig import eigh as lr_eigh +from nest._lr_eig import eigh as lr_eigh from pyscf import symm from pyscf.data import nist from nest.aocscf import AverageOccupationROKS, SymAdaptedAverageOccupationROKS diff --git a/src/nest/nttda/tests/test_nttda_aocscf_reference.py b/src/nest/nttda/tests/test_nttda_aocscf_reference.py index 4510d05..6ee2222 100644 --- a/src/nest/nttda/tests/test_nttda_aocscf_reference.py +++ b/src/nest/nttda/tests/test_nttda_aocscf_reference.py @@ -93,6 +93,13 @@ def test_nttda_energy_is_independent_of_the_nobeta_flag(self): ).run() self.assertTrue(np.all(tdobj.converged)) self.assertTrue(np.all(np.isfinite(tdobj.e))) + matrix = tdobj.get_ab() + vectors = np.asarray([xy[0].ravel() for xy in tdobj.xy]) + residuals = np.linalg.norm(matrix @ vectors.T - vectors.T * tdobj.e, axis=0) + self.assertLess(np.max(residuals), tdobj.conv_tol) + np.testing.assert_allclose( + tdobj.e, np.linalg.eigvalsh(matrix)[:2], atol=1e-10, rtol=0, + ) energies.append(tdobj.e) np.testing.assert_allclose(energies[0], energies[1], atol=1e-12, rtol=0) @@ -162,6 +169,7 @@ def test_all_three_spin_channels(self): deltaS=delta_s, nstates=2, conv_tol=1e-7, + lindep=1e-18, max_cycle=200, verbose=0, ).run() diff --git a/src/nest/soc/tests/test_sftda_soc.py b/src/nest/soc/tests/test_sftda_soc.py index 94259c9..2b5b6d2 100644 --- a/src/nest/soc/tests/test_sftda_soc.py +++ b/src/nest/soc/tests/test_sftda_soc.py @@ -47,8 +47,11 @@ def test_rank_one_clebsch_gordan(self): def test_roks_sftda_soc(self): mf = self.mol.ROKS(xc='SVWN').run() + # Spin and SOC references use a dense diagonalization of get_ab_sf(). + # Converge amplitudes tightly enough for the observable assertions. td = sftda.TDA_SF(mf).set( extype=1, collinear='mcol', collinear_samples=50, nstates=3, + conv_tol=1e-10, lindep=1e-18, ).run() driver = td.SOC(soctype='SOMF') driver.kernel() @@ -60,7 +63,7 @@ def test_roks_sftda_soc(self): -0.2104295981711506, -0.0007174487394460, 0.0251523165536107, ], atol=1e-8, rtol=0) np.testing.assert_allclose(td.spin_square(), [ - 0.0010848962999618905, 1.9999490812871423, 0.031289468589663194, + 0.0010849068616214375, 1.999949080770187, 0.03128946766229479, ], atol=1e-8, rtol=0) self.assertEqual(driver.h_soc.shape, (5, 5)) np.testing.assert_allclose(driver.h_soc, driver.h_soc.conj().T, atol=1e-12) @@ -68,15 +71,16 @@ def test_roks_sftda_soc(self): 0.0, 46026.50192719676, 46026.502933180134, 46026.50860068575, 51704.26176254972, ], atol=1e-5, rtol=0) assert_allclose_up_to_sign(self, driver.get_block(1, 0) * HARTREE2WAVENUMBER, np.array([ - [0.4674642361078794 - 6.685128880147436j], - [0.0 - 14.141141940607179j], - [0.4674642361078794 + 6.685128880147436j], + [0.4674528693952154 - 6.685125063070073j], + [0.0 - 14.141112265887136j], + [0.4674528693952154 + 6.685125063070073j], ]), 1e-8) def test_uks_sftda_soc(self): mf = self.mol.UKS(xc='SVWN').run() td = sftda.TDA_SF(mf).set( extype=1, collinear='mcol', collinear_samples=50, nstates=3, + conv_tol=1e-10, lindep=1e-18, ).run() driver = td.SOC(soctype='SOMF') driver.kernel() @@ -88,7 +92,7 @@ def test_uks_sftda_soc(self): -0.2087681123003969, 0.0008054142507056, 0.0266315014304553, ], atol=1e-8, rtol=0) np.testing.assert_allclose(td.spin_square(), [ - 0.0026539116310360, 2.0039691808012283, 0.0372918876839510, + 0.0026539051945104752, 2.0039691867433698, 0.03729208427888597, ], atol=1e-8, rtol=0) self.assertEqual(driver.h_soc.shape, (5, 5)) np.testing.assert_allclose(driver.h_soc, driver.h_soc.conj().T, atol=1e-12) @@ -96,9 +100,9 @@ def test_uks_sftda_soc(self): 0.0, 45996.07760059907, 45996.07872056033, 45996.08436902181, 51664.25143819543, ], atol=1e-5, rtol=0) assert_allclose_up_to_sign(self, driver.get_block(1, 0) * HARTREE2WAVENUMBER, np.array([ - [-0.4695042701337793 + 6.699202067333119j], - [0.0 + 14.109500425297364j], - [-0.4695042701337793 - 6.699202067333119j], + [-0.4695297473851865 + 6.6992258754796055j], + [0.0 + 14.109663673762803j], + [-0.4695297473851865 - 6.6992258754796055j], ]), 1e-8) From dbbba7120fd1faaaa851876a8ee72056a8abaa2f Mon Sep 17 00:00:00 2001 From: wtpeter Date: Sat, 26 Sep 2026 19:50:29 +0800 Subject: [PATCH 14/16] fix(nttda): preserve original solver strategy and consolidate tests --- src/nest/_lr_eig.py | 2 +- src/nest/grad/tests/test_nttda_aocscf_fd.py | 2 + src/nest/grad/tests/test_nttda_grad.py | 71 +------ src/nest/nttda/tests/test_nttda_aocscf.py | 156 +++++++++++++- .../tests/test_nttda_aocscf_reference.py | 190 ------------------ src/nest/soc/tests/test_sftda_soc.py | 20 +- 6 files changed, 167 insertions(+), 274 deletions(-) delete mode 100644 src/nest/nttda/tests/test_nttda_aocscf_reference.py diff --git a/src/nest/_lr_eig.py b/src/nest/_lr_eig.py index bb74718..088dd9f 100644 --- a/src/nest/_lr_eig.py +++ b/src/nest/_lr_eig.py @@ -189,7 +189,7 @@ def eigh(aop, x0, precond, tol_residual=1e-5, lindep=1e-12, nroots=1, space_inc = nroots else: # Adding too many trial bases in each iteration may cause larger errors - space_inc = max(nroots, min(MAX_SPACE_INC, x0_size//2)) + space_inc = min(nroots, min(MAX_SPACE_INC, x0_size//2)) max_space = int(max_memory*1e6/8/x0_size / 2 - nroots - space_inc) if max_space < nroots * 4 < x0_size: diff --git a/src/nest/grad/tests/test_nttda_aocscf_fd.py b/src/nest/grad/tests/test_nttda_aocscf_fd.py index 03899db..eff90b0 100644 --- a/src/nest/grad/tests/test_nttda_aocscf_fd.py +++ b/src/nest/grad/tests/test_nttda_aocscf_fd.py @@ -49,6 +49,7 @@ def make_td(): deltaS=0, nstates=2, conv_tol=1e-8, + lindep=1e-18, max_cycle=100, verbose=0, ).run() @@ -205,6 +206,7 @@ def test_representative_functional_families(self): nobeta=nobeta, nstates=2, conv_tol=1e-9, + lindep=1e-20, max_cycle=200, verbose=0, ).run() diff --git a/src/nest/grad/tests/test_nttda_grad.py b/src/nest/grad/tests/test_nttda_grad.py index 6d416ec..37d8c4f 100644 --- a/src/nest/grad/tests/test_nttda_grad.py +++ b/src/nest/grad/tests/test_nttda_grad.py @@ -56,6 +56,7 @@ def make_td(self, xc, delta_s, nobeta=False): nobeta=nobeta, nstates=3, conv_tol=1e-9, + lindep=1e-20, max_cycle=200, verbose=0, ).run() @@ -70,7 +71,6 @@ def compare_public_gradient(self, xc, delta_s, nobeta, threshold): root_overlap_tol=0.5, ) analytic = gradient.kernel(state=2, method="analytic") - self.assertLess(gradient.nttda_details.residual, 1e-8) finite_difference = gradient.kernel( state=2, method="finite_diff", step=2e-4, ) @@ -113,75 +113,6 @@ def test_delta_s_plus_one_rejects_analytic_and_keeps_finite_difference(self): ) self.assertTrue(np.all(np.isfinite(finite_difference))) - def test_scanner_returns_the_selected_energy_and_gradient(self): - for average_occ in (False, True): - for delta_s in (-1, 0): - with self.subTest(average_occ=average_occ, deltaS=delta_s): - mol = gto.M( - atom="C 0 0 0; H 0 0 2.0; H 0 1.7 -0.5", - basis="sto-3g", spin=2, unit="Bohr", verbose=0, - ) - mf = mol.ROKS(xc="HF") - if average_occ: - mf = mf.average_occ() - mf.set(conv_tol=1e-13, conv_tol_grad=1e-10).run() - td = NTTDA(mf).set( - deltaS=delta_s, nstates=2, conv_tol=1e-9, lindep=1e-18, - ).run() - grad = td.Gradients().set(state=2) - expected = grad.kernel() - scanner = grad.as_scanner() - energy, derivative = scanner(mol) - self.assertTrue(scanner.converged) - self.assertEqual(np.ndim(energy), 0) - self.assertAlmostEqual(energy, td.e_tot[1], delta=1e-9) - np.testing.assert_allclose(derivative, expected, atol=1e-7, rtol=0) - - coords = mol.atom_coords() - coords[1, 2] += 0.02 - moved = mol.set_geom_(coords, unit="Bohr", inplace=False) - energy, derivative = scanner(moved) - self.assertTrue(scanner.converged) - self.assertAlmostEqual(energy, scanner.base.e_tot[1], delta=1e-12) - np.testing.assert_allclose( - derivative, scanner.base.Gradients().kernel(state=2), atol=1e-9, rtol=0, - ) - - energy, derivative = scanner(moved, state=0) - self.assertTrue(scanner.converged) - self.assertAlmostEqual(energy, scanner.base._scf.e_tot, delta=1e-12) - np.testing.assert_allclose( - derivative, scanner.base._scf.nuc_grad_method().kernel(), atol=1e-9, rtol=0, - ) - - def test_symmetry_and_atom_selection(self): - mol = gto.M( - atom="H 0 .934473 -.588078; H 0 -.934473 -.588078; C 0 0 0; O 0 0 1.221104", - basis="sto-3g", spin=2, symmetry=True, verbose=0, - ) - atoms = [3, 0] - for average_occ in (False, True): - mf = mol.ROKS(xc="HF") - if average_occ: - mf = mf.average_occ() - mf.set(conv_tol=1e-12, conv_tol_grad=1e-9, max_cycle=150).run() - for delta_s in (-1, 0): - with self.subTest(average_occ=average_occ, deltaS=delta_s): - td = mf.NTTDA().set( - deltaS=delta_s, nstates=2, conv_tol=1e-9, lindep=1e-18, - ).run() - grad = td.Gradients() - full = grad.kernel(state=1) - self.assertEqual(full.shape, (4, 3)) - np.testing.assert_allclose(full.sum(axis=0), 0, atol=1e-7, rtol=0) - np.testing.assert_allclose( - grad.kernel(atmlst=atoms), full[atoms], atol=1e-10, rtol=0, - ) - np.testing.assert_allclose( - grad.kernel(state=0, atmlst=atoms), - mf.nuc_grad_method().kernel()[atoms], atol=1e-10, rtol=0, - ) - if __name__ == "__main__": unittest.main() diff --git a/src/nest/nttda/tests/test_nttda_aocscf.py b/src/nest/nttda/tests/test_nttda_aocscf.py index 413bdcb..0946b96 100644 --- a/src/nest/nttda/tests/test_nttda_aocscf.py +++ b/src/nest/nttda/tests/test_nttda_aocscf.py @@ -15,8 +15,10 @@ import unittest import numpy as np -from pyscf import gto +from pyscf import dft, gto, lib +from pyscf.scf import hf from nest import aocscf, nttda +from nest.nttda import NTTDA class KnownValues(unittest.TestCase): @@ -78,3 +80,155 @@ def test_b3lyp_nttda_aocscf(self): td = mf.NTTDA().set(nstates=2, deltaS=0, nobeta=True, conv_tol=1e-5, max_cycle=200).run() self.assertTrue(np.all(td.converged)) self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + + +class AOCSCFReference(unittest.TestCase): + @staticmethod + def lithium_hydride_cation(): + return gto.M( + atom="Li 0 0 0; H 0 0 3.0", + basis="sto-3g", + charge=1, + spin=1, + unit="Bohr", + verbose=0, + ) + + def make_reference(self, xc="SVWN"): + mf = self.lithium_hydride_cation().ROKS(xc=xc).average_occ() + mf.conv_tol = 1e-12 + mf.conv_tol_grad = 1e-9 + mf.max_cycle = 100 + mf.verbose = 0 + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + return mf + + def test_fixed_occupations_define_a_spin_unpolarized_reference(self): + mf = self.make_reference() + np.testing.assert_array_equal(mf.mo_occ, [2, 1, 0, 0, 0, 0]) + self.assertAlmostEqual(mf.mo_occ.sum(), mf.mol.nelectron) + + mo = np.asarray(mf.mo_coeff) + dm = (mo * np.asarray(mf.mo_occ)) @ mo.conj().T + veff = mf.get_veff() + charge_veff = lib.view(mf, dft.rks.RKS).get_veff(dm=dm) + np.testing.assert_allclose(veff[0], veff[1], atol=0, rtol=0) + np.testing.assert_allclose(veff[0], charge_veff, atol=1e-12, rtol=0) + self.assertAlmostEqual( + np.einsum("ij,ji", dm, mf.get_ovlp()), + mf.mol.nelectron, + places=10, + ) + + rng = np.random.default_rng(8) + fock = rng.standard_normal(dm.shape) + fock = fock + fock.T + fock_mo = mf.mo_coeff.T @ fock @ mf.mo_coeff + unique = hf.uniq_var_indices(mf.mo_occ) + occupation_difference = mf.mo_occ[None, :] - mf.mo_occ[:, None] + expected = (fock_mo * occupation_difference)[unique] + np.testing.assert_allclose( + mf.get_grad(mf.mo_coeff, mf.mo_occ, fock), + expected, + atol=1e-14, + rtol=0, + ) + + def test_nttda_energy_is_independent_of_the_nobeta_flag(self): + mf = self.make_reference() + energies = [] + for nobeta in (False, True): + tdobj = NTTDA(mf).set( + deltaS=0, + nobeta=nobeta, + nstates=2, + conv_tol=1e-8, + lindep=1e-18, + max_cycle=200, + verbose=0, + ).run() + self.assertTrue(np.all(tdobj.converged)) + self.assertTrue(np.all(np.isfinite(tdobj.e))) + energies.append(tdobj.e) + np.testing.assert_allclose(energies[0], energies[1], atol=1e-12, rtol=0) + + def test_reference_energy_is_the_high_spin_roks_energy(self): + mf = self.make_reference() + high_spin = lib.view(mf, dft.roks.ROKS) + self.assertAlmostEqual(mf.e_tot, high_spin.energy_tot(), places=12) + self.assertAlmostEqual(mf.e_avg_occ, mf.energy_tot(), places=12) + + def test_nttda_total_energies_use_the_reference_energy(self): + mf = self.make_reference() + tdobj = NTTDA(mf).set( + deltaS=0, + nstates=2, + conv_tol=1e-8, + lindep=1e-18, + max_cycle=200, + verbose=0, + ).run() + + self.assertAlmostEqual(tdobj.reference_energy(), mf.e_tot) + np.testing.assert_allclose( + tdobj.total_energies(), + mf.e_tot + tdobj.e, + atol=1e-13, + rtol=0, + ) + + def test_nttda_supports_common_functional_families(self): + for xc in ("HF", "PBE", "TPSS", "M06-2X", "CAM-B3LYP"): + with self.subTest(xc=xc): + mf = self.make_reference(xc) + tdobj = NTTDA(mf).set( + deltaS=0, + nstates=2, + conv_tol=1e-6, + max_cycle=200, + verbose=0, + ).run() + self.assertTrue(np.all(tdobj.converged)) + self.assertTrue(np.all(np.isfinite(tdobj.e))) + + def test_all_three_spin_channels(self): + mol = gto.M( + atom="C 0 0 0; H 0 0 2.0; H 0 1.7 -0.5", + basis="sto-3g", + spin=2, + unit="Bohr", + verbose=0, + ) + mf = mol.ROKS(xc="SVWN").average_occ() + mf.conv_tol = 1e-12 + mf.max_cycle = 150 + mf.verbose = 0 + mf.grids.level = 0 + mf.kernel() + self.assertTrue(mf.converged) + + references = { + -1: [0.01477165973875402, 0.06753726182098739], + 0: [-0.002622144798075737, 0.2393673138366825], + 1: [0.6745475947851435, 0.8387186686911036], + } + for delta_s, reference in references.items(): + with self.subTest(deltaS=delta_s): + tdobj = NTTDA(mf).set( + deltaS=delta_s, + nstates=2, + conv_tol=1e-7, + lindep=1e-18, + max_cycle=200, + verbose=0, + ).run() + self.assertTrue(np.all(tdobj.converged)) + np.testing.assert_allclose( + tdobj.e, reference, atol=2e-6, rtol=0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/nest/nttda/tests/test_nttda_aocscf_reference.py b/src/nest/nttda/tests/test_nttda_aocscf_reference.py deleted file mode 100644 index 6ee2222..0000000 --- a/src/nest/nttda/tests/test_nttda_aocscf_reference.py +++ /dev/null @@ -1,190 +0,0 @@ -#!/usr/bin/env python -# Copyright 2026 The NEST Developers. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Acceptance tests for AOCSCF-based NTTDA energies.""" - -import unittest - -import numpy as np - -from pyscf import dft, gto, lib -from pyscf.scf import hf -from nest import aocscf -from nest.nttda import NTTDA - - -class AOCSCFReference(unittest.TestCase): - @staticmethod - def lithium_hydride_cation(): - return gto.M( - atom="Li 0 0 0; H 0 0 3.0", - basis="sto-3g", - charge=1, - spin=1, - unit="Bohr", - verbose=0, - ) - - def make_reference(self, xc="SVWN"): - mf = self.lithium_hydride_cation().ROKS(xc=xc).average_occ() - mf.conv_tol = 1e-12 - mf.conv_tol_grad = 1e-9 - mf.max_cycle = 100 - mf.verbose = 0 - mf.grids.level = 0 - mf.kernel() - self.assertTrue(mf.converged) - return mf - - def test_fixed_occupations_define_a_spin_unpolarized_reference(self): - mf = self.make_reference() - np.testing.assert_array_equal(mf.mo_occ, [2, 1, 0, 0, 0, 0]) - self.assertAlmostEqual(mf.mo_occ.sum(), mf.mol.nelectron) - - mo = np.asarray(mf.mo_coeff) - dm = (mo * np.asarray(mf.mo_occ)) @ mo.conj().T - veff = mf.get_veff() - charge_veff = lib.view(mf, dft.rks.RKS).get_veff(dm=dm) - np.testing.assert_allclose(veff[0], veff[1], atol=0, rtol=0) - np.testing.assert_allclose(veff[0], charge_veff, atol=1e-12, rtol=0) - self.assertAlmostEqual( - np.einsum("ij,ji", dm, mf.get_ovlp()), - mf.mol.nelectron, - places=10, - ) - - rng = np.random.default_rng(8) - fock = rng.standard_normal(dm.shape) - fock = fock + fock.T - fock_mo = mf.mo_coeff.T @ fock @ mf.mo_coeff - unique = hf.uniq_var_indices(mf.mo_occ) - occupation_difference = mf.mo_occ[None, :] - mf.mo_occ[:, None] - expected = (fock_mo * occupation_difference)[unique] - np.testing.assert_allclose( - mf.get_grad(mf.mo_coeff, mf.mo_occ, fock), - expected, - atol=1e-14, - rtol=0, - ) - - def test_nttda_energy_is_independent_of_the_nobeta_flag(self): - mf = self.make_reference() - energies = [] - for nobeta in (False, True): - tdobj = NTTDA(mf).set( - deltaS=0, - nobeta=nobeta, - nstates=2, - conv_tol=1e-8, - max_cycle=200, - verbose=0, - ).run() - self.assertTrue(np.all(tdobj.converged)) - self.assertTrue(np.all(np.isfinite(tdobj.e))) - matrix = tdobj.get_ab() - vectors = np.asarray([xy[0].ravel() for xy in tdobj.xy]) - residuals = np.linalg.norm(matrix @ vectors.T - vectors.T * tdobj.e, axis=0) - self.assertLess(np.max(residuals), tdobj.conv_tol) - np.testing.assert_allclose( - tdobj.e, np.linalg.eigvalsh(matrix)[:2], atol=1e-10, rtol=0, - ) - energies.append(tdobj.e) - np.testing.assert_allclose(energies[0], energies[1], atol=1e-12, rtol=0) - - def test_reference_energy_is_the_high_spin_roks_energy(self): - mf = self.make_reference() - high_spin = lib.view(mf, dft.roks.ROKS) - self.assertAlmostEqual(mf.e_tot, high_spin.energy_tot(), places=12) - self.assertAlmostEqual(mf.e_avg_occ, mf.energy_tot(), places=12) - self.assertGreater(abs(mf.e_tot - mf.e_avg_occ), 1e-5) - - def test_nttda_total_energies_use_the_reference_energy(self): - mf = self.make_reference() - tdobj = NTTDA(mf).set( - deltaS=0, - nstates=2, - conv_tol=1e-8, - max_cycle=200, - verbose=0, - ).run() - - self.assertAlmostEqual(tdobj.reference_energy(), mf.e_tot) - np.testing.assert_allclose( - tdobj.total_energies(), - tdobj.e_tot, - atol=1e-13, - rtol=0, - ) - - def test_nttda_supports_common_functional_families(self): - for xc in ("HF", "PBE", "TPSS", "M06-2X", "CAM-B3LYP"): - with self.subTest(xc=xc): - mf = self.make_reference(xc) - tdobj = NTTDA(mf).set( - deltaS=0, - nstates=2, - conv_tol=1e-6, - max_cycle=200, - verbose=0, - ).run() - self.assertTrue(np.all(tdobj.converged)) - self.assertTrue(np.all(np.isfinite(tdobj.e))) - - def test_all_three_spin_channels(self): - mol = gto.M( - atom="C 0 0 0; H 0 0 2.0; H 0 1.7 -0.5", - basis="sto-3g", - spin=2, - unit="Bohr", - verbose=0, - ) - mf = mol.ROKS(xc="SVWN").average_occ() - mf.conv_tol = 1e-12 - mf.max_cycle = 150 - mf.verbose = 0 - mf.grids.level = 0 - mf.kernel() - self.assertTrue(mf.converged) - - references = { - -1: [0.01477165973875402, 0.06753726182098739], - 0: [-0.002622144798075737, 0.2393673138366825], - 1: [0.6745475947851435, 0.8387186686911036], - } - for delta_s, reference in references.items(): - with self.subTest(deltaS=delta_s): - tdobj = NTTDA(mf).set( - deltaS=delta_s, - nstates=2, - conv_tol=1e-7, - lindep=1e-18, - max_cycle=200, - verbose=0, - ).run() - self.assertTrue(np.all(tdobj.converged)) - np.testing.assert_allclose( - tdobj.e, reference, atol=2e-6, rtol=0, - ) - # The independent explicit matrix must describe the same physical roots. - amplitudes = np.asarray([xy[0].ravel() for xy in tdobj.xy]) - np.testing.assert_allclose( - tdobj.get_ab() @ amplitudes.T, - (tdobj.e[:, None] * amplitudes).T, - atol=2e-7, rtol=0, - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/src/nest/soc/tests/test_sftda_soc.py b/src/nest/soc/tests/test_sftda_soc.py index 2b5b6d2..94259c9 100644 --- a/src/nest/soc/tests/test_sftda_soc.py +++ b/src/nest/soc/tests/test_sftda_soc.py @@ -47,11 +47,8 @@ def test_rank_one_clebsch_gordan(self): def test_roks_sftda_soc(self): mf = self.mol.ROKS(xc='SVWN').run() - # Spin and SOC references use a dense diagonalization of get_ab_sf(). - # Converge amplitudes tightly enough for the observable assertions. td = sftda.TDA_SF(mf).set( extype=1, collinear='mcol', collinear_samples=50, nstates=3, - conv_tol=1e-10, lindep=1e-18, ).run() driver = td.SOC(soctype='SOMF') driver.kernel() @@ -63,7 +60,7 @@ def test_roks_sftda_soc(self): -0.2104295981711506, -0.0007174487394460, 0.0251523165536107, ], atol=1e-8, rtol=0) np.testing.assert_allclose(td.spin_square(), [ - 0.0010849068616214375, 1.999949080770187, 0.03128946766229479, + 0.0010848962999618905, 1.9999490812871423, 0.031289468589663194, ], atol=1e-8, rtol=0) self.assertEqual(driver.h_soc.shape, (5, 5)) np.testing.assert_allclose(driver.h_soc, driver.h_soc.conj().T, atol=1e-12) @@ -71,16 +68,15 @@ def test_roks_sftda_soc(self): 0.0, 46026.50192719676, 46026.502933180134, 46026.50860068575, 51704.26176254972, ], atol=1e-5, rtol=0) assert_allclose_up_to_sign(self, driver.get_block(1, 0) * HARTREE2WAVENUMBER, np.array([ - [0.4674528693952154 - 6.685125063070073j], - [0.0 - 14.141112265887136j], - [0.4674528693952154 + 6.685125063070073j], + [0.4674642361078794 - 6.685128880147436j], + [0.0 - 14.141141940607179j], + [0.4674642361078794 + 6.685128880147436j], ]), 1e-8) def test_uks_sftda_soc(self): mf = self.mol.UKS(xc='SVWN').run() td = sftda.TDA_SF(mf).set( extype=1, collinear='mcol', collinear_samples=50, nstates=3, - conv_tol=1e-10, lindep=1e-18, ).run() driver = td.SOC(soctype='SOMF') driver.kernel() @@ -92,7 +88,7 @@ def test_uks_sftda_soc(self): -0.2087681123003969, 0.0008054142507056, 0.0266315014304553, ], atol=1e-8, rtol=0) np.testing.assert_allclose(td.spin_square(), [ - 0.0026539051945104752, 2.0039691867433698, 0.03729208427888597, + 0.0026539116310360, 2.0039691808012283, 0.0372918876839510, ], atol=1e-8, rtol=0) self.assertEqual(driver.h_soc.shape, (5, 5)) np.testing.assert_allclose(driver.h_soc, driver.h_soc.conj().T, atol=1e-12) @@ -100,9 +96,9 @@ def test_uks_sftda_soc(self): 0.0, 45996.07760059907, 45996.07872056033, 45996.08436902181, 51664.25143819543, ], atol=1e-5, rtol=0) assert_allclose_up_to_sign(self, driver.get_block(1, 0) * HARTREE2WAVENUMBER, np.array([ - [-0.4695297473851865 + 6.6992258754796055j], - [0.0 + 14.109663673762803j], - [-0.4695297473851865 - 6.6992258754796055j], + [-0.4695042701337793 + 6.699202067333119j], + [0.0 + 14.109500425297364j], + [-0.4695042701337793 - 6.699202067333119j], ]), 1e-8) From 19f97440a89bd932e09f19671a66c60eebbc8a4b Mon Sep 17 00:00:00 2001 From: wtpeter Date: Sat, 26 Sep 2026 20:18:43 +0800 Subject: [PATCH 15/16] refactor(nttda): use existing energy and Fock interfaces --- examples/nttda/02_nttda_aocscf_grad.py | 2 +- src/nest/grad/nttda/__init__.py | 6 +- src/nest/grad/nttda/common.py | 7 ++- src/nest/grad/nttda/delta_s_minus_one.py | 10 +++- src/nest/grad/nttda/delta_s_zero.py | 9 ++- src/nest/grad/nttda/roks.py | 4 +- src/nest/grad/nttda/xc.py | 8 +-- src/nest/nttda/nttda.py | 70 +++++++---------------- src/nest/nttda/tests/test_nttda.py | 4 -- src/nest/nttda/tests/test_nttda_aocscf.py | 19 ------ 10 files changed, 50 insertions(+), 89 deletions(-) diff --git a/examples/nttda/02_nttda_aocscf_grad.py b/examples/nttda/02_nttda_aocscf_grad.py index 053889c..bd8f718 100644 --- a/examples/nttda/02_nttda_aocscf_grad.py +++ b/examples/nttda/02_nttda_aocscf_grad.py @@ -58,7 +58,7 @@ ).run() print('deltaS = %+d' % delta_s) print(' NTTDA excitation energies:', td.e) - print(' total energies (E_ref + omega):', td.total_energies()) + print(' total energies (E_ref + omega):', td.e_tot) grad = td.Gradients().kernel(state=1) print(' state-1 analytic gradient (Eh/Bohr):\n', grad) diff --git a/src/nest/grad/nttda/__init__.py b/src/nest/grad/nttda/__init__.py index 0c501dc..2e4d7b6 100644 --- a/src/nest/grad/nttda/__init__.py +++ b/src/nest/grad/nttda/__init__.py @@ -21,7 +21,7 @@ from pyscf.grad import tdrhf as tdrhf_grad from pyscf.lib import logger from nest.nttda import NTTDA -from nest.nttda.nttda import _is_average_occupation_reference +from nest.aocscf import AverageOccupationROKS, SymAdaptedAverageOccupationROKS from . import delta_s_minus_one, delta_s_zero @@ -42,7 +42,7 @@ def _copy_td_settings(source, target): def _displaced_reference(source, mol, fixed_grid): - if _is_average_occupation_reference(source): + if isinstance(source, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)): reference = dft.ROKS(mol).average_occ() elif isinstance(source, dft.KohnShamDFT): reference = dft.ROKS(mol) @@ -196,7 +196,7 @@ def _energy_at(self, coords, reference_amplitude): "NTTDA state tracking overlap %.6f is below %.6f" % (overlaps[root], self.root_overlap_tol) ) - return float(tdobj.total_energies()[root]) + return float(tdobj.e_tot[root]) def _finite_difference(self, atmlst): coords0 = self.mol.atom_coords() diff --git a/src/nest/grad/nttda/common.py b/src/nest/grad/nttda/common.py index 653db29..2c1b0ec 100644 --- a/src/nest/grad/nttda/common.py +++ b/src/nest/grad/nttda/common.py @@ -23,6 +23,7 @@ from pyscf import dft, lib from pyscf.grad import rhf as rhf_grad from nest.nttda import nttda as nttda_mod +from nest.aocscf import AverageOccupationROKS, SymAdaptedAverageOccupationROKS from . import xc as xc_backend from .xc import _reference_spin_densities @@ -54,7 +55,7 @@ def assemble_gradient( # ROKS/HF includes Fz in the spin-resolved Fock probes. All other # references use charge-only probes and differentiate Fz separately. - spin_fock = xctype == "HF" and not nttda_mod._is_average_occupation_reference(mf) + spin_fock = xctype == "HF" and not isinstance(mf, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)) direct_fock_probes = ( (0.5 * (p0 + pz), 0.5 * (p0 - pz)) if spin_fock else (0.5 * p0, 0.5 * p0) @@ -154,7 +155,7 @@ def _fock_response_q(tdobj, p_alpha, p_beta): """Reference-density derivative of a spin-resolved Fock scalar.""" mf = tdobj._scf mo = np.asarray(mf.mo_coeff) - if nttda_mod._is_average_occupation_reference(mf): + if isinstance(mf, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)): occupation = np.asarray(mf.mo_occ) probe = np.asarray(p_alpha) + np.asarray(p_beta) potential = lib.view(mf, dft.rks.RKS).gen_response(hermi=0)(probe.T) @@ -545,7 +546,7 @@ def spin_fock_direct( xctype ) if (nobeta_p0 is not None and tdobj.nobeta - and not nttda_mod._is_average_occupation_reference(mf)): + and not isinstance(mf, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS))): density0 = 0.5 * (density_alpha + density_beta) actual_probe_alpha = np.array(p_alpha, copy=True) actual_probe_beta = np.array(p_beta, copy=True) diff --git a/src/nest/grad/nttda/delta_s_minus_one.py b/src/nest/grad/nttda/delta_s_minus_one.py index f9cda2b..49430ad 100644 --- a/src/nest/grad/nttda/delta_s_minus_one.py +++ b/src/nest/grad/nttda/delta_s_minus_one.py @@ -23,7 +23,6 @@ import numpy as np from pyscf import lib -from nest.nttda import nttda as nttda_mod from nest.nttda.nttda import gen_rohf_response_sfd from .common import ( @@ -145,7 +144,14 @@ def spin_lowering_fock0_fockz(tdobj, max_memory=None): hermi=0, max_memory=max_memory, ) - return nttda_mod._reference_fock0(mf, tdobj.nobeta), fockz + if tdobj.nobeta: + dma, dmb = mf.make_rdm1() + dm0 = 0.5 * (dma + dmb) + fock = mf.get_fock(dm=np.array([dm0, dm0])) + else: + fock = mf.get_fock() + fock0 = 0.5 * (fock.focka + fock.fockb) + return fock0, fockz def spin_lowering_fock_projections(tdobj, xy): diff --git a/src/nest/grad/nttda/delta_s_zero.py b/src/nest/grad/nttda/delta_s_zero.py index 6ae0617..f8a43e0 100644 --- a/src/nest/grad/nttda/delta_s_zero.py +++ b/src/nest/grad/nttda/delta_s_zero.py @@ -24,7 +24,6 @@ import numpy as np from pyscf import lib -from nest.nttda import nttda as nttda_mod from nest.nttda.nttda import gen_rohf_response_sc from .common import ( @@ -137,7 +136,13 @@ def fock0_fockz(tdobj, max_memory=None): hermi=0, max_memory=max_memory, ) - fock0 = nttda_mod._reference_fock0(mf, tdobj.nobeta) + if tdobj.nobeta: + dma, dmb = mf.make_rdm1() + dm0 = 0.5 * (dma + dmb) + fock = mf.get_fock(dm=np.array([dm0, dm0])) + else: + fock = mf.get_fock() + fock0 = 0.5 * (fock.focka + fock.fockb) return fock0, fockz diff --git a/src/nest/grad/nttda/roks.py b/src/nest/grad/nttda/roks.py index 1c3e10c..c73c70f 100644 --- a/src/nest/grad/nttda/roks.py +++ b/src/nest/grad/nttda/roks.py @@ -20,7 +20,7 @@ from pyscf import dft, lib from pyscf.grad import rhf as rhf_grad -from nest.nttda.nttda import _is_average_occupation_reference +from nest.aocscf import AverageOccupationROKS, SymAdaptedAverageOccupationROKS @@ -44,7 +44,7 @@ def finish_gradient( ``direct_fock_probes`` enables one batched Fock-derivative evaluation: its contraction is the first result and the Z-vector contraction is the second. """ - if _is_average_occupation_reference(tdobj._scf): + if isinstance(tdobj._scf, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)): from . import ensemble as response else: from . import roks as response diff --git a/src/nest/grad/nttda/xc.py b/src/nest/grad/nttda/xc.py index 257e02d..3901816 100644 --- a/src/nest/grad/nttda/xc.py +++ b/src/nest/grad/nttda/xc.py @@ -26,7 +26,7 @@ from pyscf.dft.gen_grid import NBINS from pyscf.dft.numint import _dot_ao_ao_sparse, _scale_ao_sparse from pyscf.grad import tdrks as tdrks_grad -from nest.nttda.nttda import _is_average_occupation_reference +from nest.aocscf import AverageOccupationROKS, SymAdaptedAverageOccupationROKS # Shared result and projection helpers @@ -328,7 +328,7 @@ def _reference_spin_densities(tdobj): """Spin densities of the variational reference used by the XC kernel.""" mf = tdobj._scf mo = np.asarray(mf.mo_coeff) - if _is_average_occupation_reference(mf): + if isinstance(mf, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)): density = (mo * mf.mo_occ) @ mo.conj().T return 0.5 * density, 0.5 * density return ( @@ -341,7 +341,7 @@ def _reference_spin_occupations(tdobj): """Per-orbital alpha/beta occupations of the reference density.""" mf = tdobj._scf occupation = np.asarray(mf.mo_occ) - if _is_average_occupation_reference(mf): + if isinstance(mf, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)): return 0.5 * occupation, 0.5 * occupation return (occupation > 0).astype(float), (occupation == 2).astype(float) @@ -918,7 +918,7 @@ def nobeta_reference_q(tdobj, p0, max_memory=None): mo = np.asarray(mf.mo_coeff) q_alpha = np.zeros((mo.shape[1], mo.shape[1])) q_beta = np.zeros_like(q_alpha) - if not tdobj.nobeta or _is_average_occupation_reference(mf): + if not tdobj.nobeta or isinstance(mf, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)): return q_alpha, q_beta if max_memory is None: max_memory = tdobj.max_memory diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index 1008cbc..11ba4c4 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -30,42 +30,11 @@ from nest._lr_eig import eigh as lr_eigh from pyscf import symm from pyscf.data import nist -from nest.aocscf import AverageOccupationROKS, SymAdaptedAverageOccupationROKS MO_BASE = getattr(__config__, 'MO_BASE', 1) MO_GRID_FXC1 = True -def _is_average_occupation_reference(mf): - return isinstance(mf, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)) - - -def _require_nttda_reference(mf): - supported = ( - dft.roks.ROKS, - dft.rks_symm.SymAdaptedROKS, - ) - if not isinstance(mf, supported): - raise TypeError("NTTDA response requires a ROKS or AOCSCF reference") - - -def _reference_fock0(mf, nobeta): - """Return the common Fock used by the NTTDA orbital terms. - - An AOCSCF (average-occupation) reference is self-consistent in - ``F0[D/2,D/2]``. For ROKS, retain the two historical choices controlled - by ``nobeta``. - """ - if _is_average_occupation_reference(mf): - return np.asarray(mf.get_fock()) - if nobeta: - dma, dmb = mf.make_rdm1() - dm0 = 0.5 * (dma + dmb) - fock = mf.get_fock(dm=np.array([dm0, dm0])) - else: - fock = mf.get_fock() - return 0.5 * (fock.focka + fock.fockb) - def _fxc1_gga_mo_wv(fxc, t, i): nvec = t.shape[0] ngrids = t.shape[-1] @@ -326,7 +295,6 @@ def gen_rohf_response_sfu(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=No mol = mf.mol if log is None: log = logger.new_logger(mf) - _require_nttda_reference(mf) ni = mf._numint ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True) @@ -387,7 +355,6 @@ def gen_rohf_response_sc(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=Non mol = mf.mol if log is None: log = logger.new_logger(mf) - _require_nttda_reference(mf) s = (mol.nelec[0] - mol.nelec[1]) * 0.5 @@ -497,7 +464,6 @@ def gen_rohf_response_sfd(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=No mol = mf.mol if log is None: log = logger.new_logger(mf) - _require_nttda_reference(mf) s = (mol.nelec[0] - mol.nelec[1]) * 0.5 @@ -611,7 +577,13 @@ def gen_vind_sfu(td): vresp, fockz = gen_rohf_response_sfu(mf, mo_coeff=mo_coeff, mo_occ=mo_occ, hermi=0, max_memory=td.max_memory, log=log) - fock0 = _reference_fock0(mf, td.nobeta) + if td.nobeta: + dma, dmb = mf.make_rdm1() + dm0 = 0.5 * (dma + dmb) + fock = mf.get_fock(dm=np.array([dm0, dm0])) + else: + fock = mf.get_fock() + fock0 = 0.5 * (fock.focka + fock.fockb) focka = fock0 + fockz fockb = fock0 - fockz @@ -674,7 +646,13 @@ def gen_vind_sc(td): fxc_ref=fxc_ref, skip_xc_vref1=use_mo_grid_fxc1) - fock0 = _reference_fock0(mf, td.nobeta) + if td.nobeta: + dma, dmb = mf.make_rdm1() + dm0 = 0.5 * (dma + dmb) + fock = mf.get_fock(dm=np.array([dm0, dm0])) + else: + fock = mf.get_fock() + fock0 = 0.5 * (fock.focka + fock.fockb) focka = fock0 + fockz fockb = fock0 - fockz @@ -842,7 +820,13 @@ def gen_vind_sfd(td): fxc_ref=fxc_ref, skip_xc_vref1=use_mo_grid_fxc1) - fock0 = _reference_fock0(mf, td.nobeta) + if td.nobeta: + dma, dmb = mf.make_rdm1() + dm0 = 0.5 * (dma + dmb) + fock = mf.get_fock(dm=np.array([dm0, dm0])) + else: + fock = mf.get_fock() + fock0 = 0.5 * (fock.focka + fock.fockb) fock_coco0 = orbos.T @ (fock0 - fockz) @ orbos fock_coco1 = orbcs.T @ (fock0 + fockz) @ orbcs @@ -984,16 +968,6 @@ class NTTDA(TDBase): _keys = {'deltaS', 'nobeta'} - def reference_energy(self): - """Return the reported reference energy, including AOCSCF's high-spin energy.""" - return float(self._scf.e_tot) - - def total_energies(self): - """Return ``E_reference + omega`` for the converged NTTDA roots.""" - if self.e is None: - raise RuntimeError('run NTTDA.kernel() before requesting total energies') - return self.reference_energy() + np.asarray(self.e) - def nuc_grad_method(self): """Return the independent NTTDA nuclear-gradient driver.""" from nest.grad.nttda import Gradients @@ -1016,8 +990,6 @@ def dump_flags(self, verbose=None): def check_sanity(self): if self.deltaS not in (-1, 0, 1): raise ValueError('deltaS must be -1, 0, or 1') - if not isinstance(self.nstates, (int, np.integer)) or self.nstates <= 0: - raise ValueError('nstates must be a positive integer') TDBase.check_sanity(self) return self diff --git a/src/nest/nttda/tests/test_nttda.py b/src/nest/nttda/tests/test_nttda.py index 7a03c80..cb8a7a3 100644 --- a/src/nest/nttda/tests/test_nttda.py +++ b/src/nest/nttda/tests/test_nttda.py @@ -231,10 +231,6 @@ def test_nttda_check_sanity(self): with self.assertRaisesRegex(ValueError, 'deltaS must be'): td.kernel() td.deltaS = -1 - with self.assertRaisesRegex(ValueError, 'nstates must be a positive integer'): - td.kernel(nstates=0) - with self.assertRaisesRegex(ValueError, 'nstates must be a positive integer'): - td.kernel(nstates=1.5) mol = gto.M(atom='H 0 0 0', spin=1, basis='sto-3g', verbose=0) td = mol.ROKS(xc='HF').run().NTTDA() diff --git a/src/nest/nttda/tests/test_nttda_aocscf.py b/src/nest/nttda/tests/test_nttda_aocscf.py index 0946b96..24bd7a7 100644 --- a/src/nest/nttda/tests/test_nttda_aocscf.py +++ b/src/nest/nttda/tests/test_nttda_aocscf.py @@ -160,25 +160,6 @@ def test_reference_energy_is_the_high_spin_roks_energy(self): self.assertAlmostEqual(mf.e_tot, high_spin.energy_tot(), places=12) self.assertAlmostEqual(mf.e_avg_occ, mf.energy_tot(), places=12) - def test_nttda_total_energies_use_the_reference_energy(self): - mf = self.make_reference() - tdobj = NTTDA(mf).set( - deltaS=0, - nstates=2, - conv_tol=1e-8, - lindep=1e-18, - max_cycle=200, - verbose=0, - ).run() - - self.assertAlmostEqual(tdobj.reference_energy(), mf.e_tot) - np.testing.assert_allclose( - tdobj.total_energies(), - mf.e_tot + tdobj.e, - atol=1e-13, - rtol=0, - ) - def test_nttda_supports_common_functional_families(self): for xc in ("HF", "PBE", "TPSS", "M06-2X", "CAM-B3LYP"): with self.subTest(xc=xc): From 29e1b60761f9436037c334755a23a33ee4936736 Mon Sep 17 00:00:00 2001 From: wtpeter Date: Sun, 27 Sep 2026 11:20:44 +0800 Subject: [PATCH 16/16] fix(nttda): preserve displaced settings and align state amplitudes --- src/nest/grad/nttda/__init__.py | 63 +++++++++++++++++++++++---------- 1 file changed, 45 insertions(+), 18 deletions(-) diff --git a/src/nest/grad/nttda/__init__.py b/src/nest/grad/nttda/__init__.py index 2e4d7b6..3e3dba8 100644 --- a/src/nest/grad/nttda/__init__.py +++ b/src/nest/grad/nttda/__init__.py @@ -21,7 +21,7 @@ from pyscf.grad import tdrhf as tdrhf_grad from pyscf.lib import logger from nest.nttda import NTTDA -from nest.aocscf import AverageOccupationROKS, SymAdaptedAverageOccupationROKS +from nest.nttda.nttda import _orbital_indices, _unpack_amplitudes from . import delta_s_minus_one, delta_s_zero @@ -32,6 +32,40 @@ def _normalized_amplitude(xy): return vector / np.linalg.norm(vector) +def _aligned_amplitudes(reference, displaced): + """Express displaced amplitudes in the reference's C/O/V orbital frames.""" + overlap = gto.intor_cross("int1e_ovlp", reference.mol, displaced.mol) + rotations = [] + for ref_idx, new_idx in zip(_orbital_indices(reference), _orbital_indices(displaced)): + ref_mo = reference._scf.mo_coeff[:, ref_idx] + new_mo = displaced._scf.mo_coeff[:, new_idx] + # The polar factor aligns each occupation space, including signs, + # permutations and rotations, without mixing distinct CSF sectors. + u, _, vh = np.linalg.svd(ref_mo.conj().T @ overlap @ new_mo, full_matrices=False) + rotations.append(u @ vh) + c, o, v = rotations + nc, no, nv = (rotation.shape[1] for rotation in rotations) + amplitudes = [] + for state in range(len(displaced.xy)): + co, cv, oo, ov, cv0 = _unpack_amplitudes(displaced, state, nc, no, nv) + # X[i,a] carries a conjugated hole and an unconjugated particle index. + cv = c.conj() @ cv @ v.T + if displaced.deltaS == 1: + x = cv + else: + co = c.conj() @ co @ o.T + ov = o.conj() @ ov @ v.T + if displaced.deltaS == 0: + # The spin-conserving OO amplitude is a scalar, not an O×O block. + cv0 = c.conj() @ cv0 @ v.T + x = np.concatenate((co.ravel(), cv.ravel(), [oo], ov.ravel(), cv0.ravel())) + else: + oo = o.conj() @ oo @ o.T + x = np.block([[co, cv], [oo, ov]]) + amplitudes.append(_normalized_amplitude((x, 0))) + return np.asarray(amplitudes) + + def _copy_td_settings(source, target): for name in ( "deltaS", "nobeta", "nstates", "conv_tol", "lindep", @@ -42,23 +76,16 @@ def _copy_td_settings(source, target): def _displaced_reference(source, mol, fixed_grid): - if isinstance(source, (AverageOccupationROKS, SymAdaptedAverageOccupationROKS)): - reference = dft.ROKS(mol).average_occ() - elif isinstance(source, dft.KohnShamDFT): - reference = dft.ROKS(mol) - else: - reference = source.__class__(mol) - for name in ( - "conv_tol", "conv_tol_grad", "max_cycle", "max_memory", - "level_shift", "damp"): - if hasattr(source, name): - setattr(reference, name, getattr(source, name)) + reference = source.copy() + if isinstance(source, dft.KohnShamDFT): + # reset() mutates grids; detach them before clearing geometry caches. + reference.grids = source.grids.copy() + reference.nlcgrids = source.nlcgrids.copy() + reference._numint = source._numint.copy() + reference.reset(mol) + reference.chkfile = None reference.verbose = 0 if isinstance(source, dft.KohnShamDFT): - reference.xc = source.xc - reference.nlc = source.nlc - reference.grids.level = source.grids.level - reference.grids.prune = source.grids.prune if fixed_grid and source.grids.coords is not None: reference.grids.coords = np.array(source.grids.coords, copy=True) reference.grids.weights = np.array(source.grids.weights, copy=True) @@ -185,8 +212,8 @@ def _energy_at(self, coords, reference_amplitude): tdobj = _copy_td_settings(self.base, NTTDA(mf)) tdobj.kernel() overlaps = np.asarray([ - abs(np.vdot(reference_amplitude, _normalized_amplitude(xy))) - for xy in tdobj.xy + abs(np.vdot(reference_amplitude, amplitude)) + for amplitude in _aligned_amplitudes(self.base, tdobj) ]) root = int(np.argmax(overlaps)) if not tdobj.converged[root]: