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/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) diff --git a/src/nest/dz0scf/dz0scf.py b/src/nest/dz0scf/dz0scf.py index 7d32fa3..7a4f1a0 100644 --- a/src/nest/dz0scf/dz0scf.py +++ b/src/nest/dz0scf/dz0scf.py @@ -55,7 +55,54 @@ def evaluate_high_spin_energy(mf): ) class _DZ0VeffMixin: - def get_veff( + 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, @@ -81,8 +128,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/grad/nttda/__init__.py b/src/nest/grad/nttda/__init__.py new file mode 100644 index 0000000..b12f59e --- /dev/null +++ b/src/nest/grad/nttda/__init__.py @@ -0,0 +1,210 @@ +"""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.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 getattr(source, "is_average_occupation_reference", False): + reference = source.__class__(mol) + 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 Dz0SCF references.""" + + _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) + 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 float(tdobj.total_energies()[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/common.py b/src/nest/grad/nttda/common.py new file mode 100644 index 0000000..0a51154 --- /dev/null +++ b/src/nest/grad/nttda/common.py @@ -0,0 +1,743 @@ +"""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 .xc import _reference_spin_densities +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": + for terms in ( + 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 = xc_backend.nobeta_reference_q(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): + 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] + 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 = vref0.copy() + 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 _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( + gradient_driver, tdobj, p_alpha, p_beta, atmlst=None, + nobeta_p0=None, jk_ledger=None, output_slots=None): + """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. + """ + 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 != "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, + 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 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) + + +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 new file mode 100644 index 0000000..526491d --- /dev/null +++ b/src/nest/grad/nttda/delta_s_minus_one.py @@ -0,0 +1,309 @@ +"""Analytic gradient for current NTTDA ``deltaS=-1``. + +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 lib +from nest.nttda import nttda as nttda_mod +from nest.nttda.nttda import gen_rohf_response_sfd + +from .common import ( + assemble_gradient, + orbital_spaces, + pair_density, + FockProjection, + fock_probes, + fock_projection_q, + response_projection_q, + ResponseTerm, + _apply_reference_responses, + +) + + +# Orbital spaces and native amplitudes + + +@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 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), + } + + +# 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): + """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): + 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) + + +# Channel assembly + +def grad_elec( + gradient_driver, tdobj, xy, atmlst=None, tolerance=1e-12, + max_cycle=None): + """Build the analytic excitation gradient for deltaS=-1.""" + if tdobj.deltaS != -1: + raise ValueError("deltaS=-1 gradient received a different spin channel") + 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) + projections = spin_lowering_fock_projections(tdobj, xy) + probes = fock_probes(tdobj, projections) + return assemble_gradient( + 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 new file mode 100644 index 0000000..764f918 --- /dev/null +++ b/src/nest/grad/nttda/delta_s_zero.py @@ -0,0 +1,300 @@ +"""Analytic gradient for current NTTDA ``deltaS=0``. + +The public ``grad_elec`` function exposes the complete scientific data flow. +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 lib +from nest.nttda import nttda as nttda_mod +from nest.nttda.nttda import gen_rohf_response_sc + +from .common import ( + assemble_gradient, + orbital_spaces, + pair_density, + FockProjection, + fock_probes, + fock_projection_q, + response_projection_q, + ResponseTerm, + _apply_reference_responses, +) + + +# Orbital spaces and native amplitudes + + +@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 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 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) + 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 + + +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): + """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): + """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 + + +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 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) + + +# 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 analytic excitation gradient for deltaS=0.""" + if tdobj.deltaS != 0: + raise ValueError("deltaS=0 gradient received a different spin channel") + spaces, amplitudes, densities = same_spin_transition_densities(tdobj, xy) + blocks = same_spin_block_data(spaces, amplitudes) + response_terms = same_spin_response_terms(spaces.spin) + channel_data = (spaces, amplitudes, densities, blocks, response_terms) + projections = same_spin_fock_projections(tdobj, xy) + probes = fock_probes(tdobj, projections) + return assemble_gradient( + 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 new file mode 100644 index 0000000..fe62e98 --- /dev/null +++ b/src/nest/grad/nttda/ensemble.py @@ -0,0 +1,130 @@ +"""Average-occupation orbital response for Dz0SCF NTTDA gradients.""" + +import numpy as np + +from pyscf.scf import hf + +from .roks import finish_gradient, pack_m_matrix, _solve_zvector + + +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 adjoint with its occupation-weighted diagonal.""" + return _solve_zvector(action, _preconditioner(tdobj, pairs), rhs, + tolerance, max_cycle) + + +__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..d485317 --- /dev/null +++ b/src/nest/grad/nttda/roks.py @@ -0,0 +1,289 @@ +"""ROKS transpose-Hessian adjoint and final NTTDA gradient assembly.""" + +from dataclasses import dataclass + +import numpy as np + +from pyscf import dft, lib +from pyscf.grad import rhf as rhf_grad + + + +@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. + """ + if getattr(tdobj._scf, "is_average_occupation_reference", False): + 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 = response.solve_zvector( + transpose_action, + pairs, + tdobj, + rhs, + tolerance=tolerance, + max_cycle=max_cycle, + ) + adjoint = response.zvector_adjoint_matrix(tdobj, pairs, zvector) + residual = float(np.max(np.abs(pack_m_matrix(adjoint, pairs) - rhs))) + probe_alpha, probe_beta = response.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_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) + 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 apply_one(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 + 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 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 pack_m_matrix(adjoint(vector), pairs) + return np.asarray([pack_m_matrix(adjoint(row), pairs) 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, 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) + + 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 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] + 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)``.""" + return _make_adjoint_action(tdobj, pairs)(zvector) + + +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 = 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): + 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..97bf844 --- /dev/null +++ b/src/nest/grad/nttda/xc.py @@ -0,0 +1,954 @@ +"""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] + + +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_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 _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_average_occupation_reference", 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, hermi=1): + ao_values = ao[0] if xctype == "LDA" else ao + rho = ni.eval_rho( + mol, ao_values, density, mask, xctype, hermi=hermi, + 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, + ) + 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, + ) + + +# 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] + 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 response_terms( + gradient_driver, tdobj, channel_data, atmlst=None, + with_direct=True): + """LDA/GGA/meta-GGA ``vref0/vref1`` M matrix and fixed-grid skeleton derivative.""" + mf = tdobj._scf + xctype = mf._numint._xc_type(mf.xc) + 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 + else: + nvar, pair_potential, pair_cross = 5, mgga_pair_potential, mgga_pair_kernel_cross + 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 xctype != "LDA" and 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) 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, ao_deriv, max_memory=gradient_driver.max_memory): + rho0 = ni.eval_rho2( + mol, ao[0] if xctype == "LDA" else ao, mo, mf.mo_occ, + mask, xctype, with_lapl=False, + ) * 0.5 + fref, kref_alpha, kref_beta = _reference_fref_kref(mf, rho0, xctype) + rho = { + label: _xc_density(ni, mol, ao, density, mask, xctype, hermi=0) + for label, density in densities.items() + } + 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 + } + special_weights = { + label: np.zeros((4, 4, weights.size)) for label in pair_labels + } + reference_weights_alpha = np.zeros((nvar, weights.size)) + reference_weights_beta = np.zeros_like(reference_weights_alpha) + + for term in terms: + # 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] += ordinary_coefficient * lib.einsum( + "xyg,xg->yg", fref, rho[term.target], + ) + pair = ordinary_coefficient * 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 xctype != "LDA" and 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 * pair_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_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 += eval_matrix( + mol, ao, reference_weights_alpha * weights, mask, + ) + reference_beta += 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, xctype, + ) + ) + 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, + ) + 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( + tdobj, potentials, blocks, + ) + _add_reference_q( + tdobj, q_alpha, q_beta, reference_alpha, reference_beta, + ) + return XCGradientTerms(q_alpha, q_beta, direct) + + +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): + """LDA/GGA/meta-GGA derivative of ``Pz:Fz`` excluding its explicit Pz projection.""" + mf = tdobj._scf + xctype = mf._numint._xc_type(mf.xc) + add_matrix, eval_matrix = _xc_matrix_builders(xctype) + 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) 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, ao_deriv, max_memory=gradient_driver.max_memory): + rho0 = ni.eval_rho2( + mol, ao[0] if xctype == "LDA" else ao, mo, mf.mo_occ, + mask, xctype, with_lapl=False, + ) * 0.5 + 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, + 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 += eval_matrix( + mol, ao, lib.einsum("xyg,xyzg->zg", pair, kref_alpha) * weights, + mask, + ) + reference_beta += 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, xctype, + ) + 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 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) + _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) + 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) + 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_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, + xctype=xctype, spin=1, + )[2] + fxc_equal = ni.eval_xc_eff( + mf.xc, (rho_equal, rho_equal), deriv=2, + xctype=xctype, 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 += eval_matrix( + mol, + ao, + lib.einsum("xg,xzg->zg", rho_p, equal - actual_alpha) * weights, + mask, + ) + matrix_beta += 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/grad/tests/test_nttda_dz0scf_fd.py b/src/nest/grad/tests/test_nttda_dz0scf_fd.py new file mode 100644 index 0000000..0736c5b --- /dev/null +++ b/src/nest/grad/tests/test_nttda_dz0scf_fd.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python +"""Finite-difference NTTDA gradients on a Dz0SCF reference.""" + +import unittest + +import numpy as np + +from pyscf import gto +from nest.dz0scf import DZ0SCF +from nest.nttda import NTTDA + + +class Dz0SCFFiniteDifferenceGradient(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 = DZ0SCF(mol, xc="SVWN").set( + 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("Dz0SCF 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.assertTrue(tdobj.Gradients().fixed_grid) + from nest.grad.nttda import _displaced_reference + + displaced = _displaced_reference(tdobj._scf, tdobj.mol.copy(), False) + self.assertTrue( + getattr(displaced, "is_average_occupation_reference", False)) + 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.091524381309104896], + [0, 0, -0.091524381308882852]], + 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 = DZ0SCF(mol, xc="SVWN").set( + 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 = DZ0SCF(mol, xc="SVWN").set( + 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 = DZ0SCF(mol, xc=xc).set( + 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_dz0scf_response.py b/src/nest/grad/tests/test_nttda_dz0scf_response.py new file mode 100644 index 0000000..56cfe9b --- /dev/null +++ b/src/nest/grad/tests/test_nttda_dz0scf_response.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python +"""Orbital-response checks for average-occupation Dz0SCF gradients.""" + +import unittest + +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.dz0scf import DZ0SCF # noqa: E402 +from nest.nttda import NTTDA # noqa: E402 + + +class Dz0SCFOrbitalResponse(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 = DZ0SCF(mol, xc="SVWN").set( + 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("Dz0SCF 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 = DZ0SCF(mol, xc="SVWN").set( + 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 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) + 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) + 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) + + 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/nttda.py b/src/nest/nttda/nttda.py index 7724a13..fa06e72 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,174 @@ 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 MO_BASE = getattr(__config__, 'MO_BASE', 1) +MO_GRID_FXC1 = True + + +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, + ) + if not isinstance(mf, supported): + 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. + + 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_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] + 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 +325,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 +348,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 +373,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 +386,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 +395,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 +428,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 +483,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 +496,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 +506,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 +538,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 +610,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 +649,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 +660,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 +725,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 +812,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 +828,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 +894,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 +977,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 @@ -1251,5 +1489,6 @@ def analyze(tdobj, verbose=None): NTTDA.analyze = analyze + dft.roks.ROKS.NTTDA = lib.class_as_method(NTTDA) dft.rks_symm.SymAdaptedROKS.NTTDA = lib.class_as_method(NTTDA) 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_dz0scf.py b/src/nest/nttda/tests/test_nttda_dz0scf.py new file mode 100644 index 0000000..1fbb422 --- /dev/null +++ b/src/nest/nttda/tests/test_nttda_dz0scf.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python +"""Acceptance tests for Dz0SCF-based NTTDA energies.""" + +import unittest + +import numpy as np + +from pyscf import gto +from pyscf.scf import hf +from nest.dz0scf import DZ0SCF +from nest.nttda import NTTDA + + +class Dz0SCFReference(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 = 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_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 + 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_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))) + 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() + self.assertEqual( + mf.reference_energy_semantics, + "high_spin_roks_energy_on_dz0_orbitals", + ) + self.assertFalse(mf.reference_energy_stationary) + self.assertAlmostEqual( + mf.reference_energy(), mf.high_spin_energy(), places=14, + ) + + 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.high_spin_energy()) + np.testing.assert_allclose( + tdobj.total_energies(), + mf.high_spin_energy() + 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 = 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) + + 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() 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)