diff --git a/.gitignore b/.gitignore index 093a151..2aa776c 100644 --- a/.gitignore +++ b/.gitignore @@ -220,6 +220,3 @@ __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 deleted file mode 100644 index 4e84998..0000000 --- a/examples/grad/02_dz0scf_grad.py +++ /dev/null @@ -1,58 +0,0 @@ -#!/usr/bin/env python -# Copyright 2026 The NEST Developers. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -''' -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 deleted file mode 100644 index 1f6154c..0000000 --- a/examples/nttda/02_nttda_dz0scf_grad.py +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env python -# Copyright 2026 The NEST Developers. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -''' -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 7a4f1a0..7d32fa3 100644 --- a/src/nest/dz0scf/dz0scf.py +++ b/src/nest/dz0scf/dz0scf.py @@ -55,54 +55,7 @@ def evaluate_high_spin_energy(mf): ) class _DZ0VeffMixin: - reference_energy_semantics = 'high_spin_roks_energy_on_dz0_orbitals' - reference_energy_stationary = False - is_average_occupation_reference = True - - def _charge_rks(self): - """Return a fresh RKS view used for the spin-unpolarized charge response. - - The view is rebuilt on every call so that a reused mean-field object - (``reset(new_mol)``, ``xc`` change, new geometry) never feeds a stale - molecule or functional into the response. - """ - charge = dft.rks.RKS(self.mol) - for name in ( - 'xc', 'nlc', 'grids', 'nlcgrids', '_numint', - 'max_memory', 'small_rho_cutoff'): - if hasattr(self, name): - setattr(charge, name, getattr(self, name)) - charge.mo_coeff = np.asarray(self.mo_coeff) - charge.mo_occ = np.asarray(self.mo_occ) - charge.mo_energy = np.asarray(self.mo_energy) - charge.verbose = 0 - return charge - - def make_rdm1s(self, mo_coeff=None, mo_occ=None): - """Return equal spin densities ``D/2`` for the spin-unpolarized reference.""" - if mo_coeff is None: - mo_coeff = self.mo_coeff - if mo_occ is None: - mo_occ = self.mo_occ - mo_coeff = np.asarray(mo_coeff) - occupation = np.asarray(mo_occ) - dm0 = (mo_coeff * occupation) @ mo_coeff.conj().T - return 0.5 * dm0, 0.5 * dm0 - - def gen_response(self, mo_coeff=None, mo_occ=None, hermi=1, max_memory=None): - """Charge-only (spin-unpolarized) linear response of the reference.""" - if mo_coeff is None: - mo_coeff = self.mo_coeff - if mo_occ is None: - mo_occ = self.mo_occ - return self._charge_rks().gen_response( - mo_coeff=mo_coeff, - mo_occ=mo_occ, - hermi=hermi, - max_memory=max_memory, - ) - - def get_veff( + def get_veff( self, mol=None, dm=None, @@ -128,11 +81,8 @@ def get_veff( vhf_last, hermi, ) - def high_spin_energy(self): - return evaluate_high_spin_energy(self) - - def reference_energy(self): - return self.high_spin_energy() + def high_spin_energy(self): + return evaluate_high_spin_energy(self) 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 d5f57ee..fc88305 100644 --- a/src/nest/dz0scf/tests/test_dz0scf.py +++ b/src/nest/dz0scf/tests/test_dz0scf.py @@ -79,18 +79,12 @@ 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( - td_s.total_energies(), - mf.high_spin_energy() + omega_s, - rtol=0.0, - atol=1e-12, - ) + np.testing.assert_allclose( + np.asarray(omega_s), + omega_s_ref, + rtol=0.0, + atol=1e-6, + ) td_t = NTTDA(mf) td_t.deltaS = 0 @@ -182,4 +176,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 deleted file mode 100644 index b12f59e..0000000 --- a/src/nest/grad/nttda/__init__.py +++ /dev/null @@ -1,210 +0,0 @@ -"""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 deleted file mode 100644 index 0a51154..0000000 --- a/src/nest/grad/nttda/common.py +++ /dev/null @@ -1,743 +0,0 @@ -"""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 deleted file mode 100644 index 526491d..0000000 --- a/src/nest/grad/nttda/delta_s_minus_one.py +++ /dev/null @@ -1,309 +0,0 @@ -"""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 deleted file mode 100644 index 764f918..0000000 --- a/src/nest/grad/nttda/delta_s_zero.py +++ /dev/null @@ -1,300 +0,0 @@ -"""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 deleted file mode 100644 index fe62e98..0000000 --- a/src/nest/grad/nttda/ensemble.py +++ /dev/null @@ -1,130 +0,0 @@ -"""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 deleted file mode 100644 index d485317..0000000 --- a/src/nest/grad/nttda/roks.py +++ /dev/null @@ -1,289 +0,0 @@ -"""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 deleted file mode 100644 index 97bf844..0000000 --- a/src/nest/grad/nttda/xc.py +++ /dev/null @@ -1,954 +0,0 @@ -"""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 deleted file mode 100644 index 0736c5b..0000000 --- a/src/nest/grad/tests/test_nttda_dz0scf_fd.py +++ /dev/null @@ -1,218 +0,0 @@ -#!/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 deleted file mode 100644 index 56cfe9b..0000000 --- a/src/nest/grad/tests/test_nttda_dz0scf_response.py +++ /dev/null @@ -1,256 +0,0 @@ -#!/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 deleted file mode 100644 index c529a86..0000000 --- a/src/nest/grad/tests/test_nttda_grad.py +++ /dev/null @@ -1,103 +0,0 @@ -#!/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 deleted file mode 100644 index 0c88866..0000000 --- a/src/nest/grad/tests/test_nttda_gradient_layers.py +++ /dev/null @@ -1,238 +0,0 @@ -#!/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 deleted file mode 100644 index 4dabb0f..0000000 --- a/src/nest/grad/tests/test_nttda_scalar_ledger.py +++ /dev/null @@ -1,106 +0,0 @@ -#!/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 fa06e72..7724a13 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -# Copyright 2014-2024 The PySCF Developers. All Rights Reserved. +# 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. @@ -27,174 +27,11 @@ from pyscf.dft.gen_grid import NBINS from pyscf import __config__ from pyscf.dft.numint import _scale_ao_sparse, _dot_ao_ao_sparse, _dot_ao_dm_sparse, _contract_rho_sparse -from pyscf.tdscf._lr_eig import eigh as lr_eigh +from nest._lr_eig import eigh as lr_eigh from pyscf import symm from pyscf.data import nist 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] @@ -325,7 +162,8 @@ def gen_rohf_response_sfu(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=No mol = mf.mol if log is None: log = logger.new_logger(mf) - _require_nttda_reference(mf) + if not isinstance(mf, (dft.roks.ROKS, dft.rks_symm.SymAdaptedROKS)): + raise TypeError('NTTDA response requires ROKS reference') ni = mf._numint ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True) @@ -348,7 +186,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 xc response_cv', *time_xc) + time_xc = log.timer('NTTDA response_sfu kernel v1ao_cv', *time_xc) else: v1ao_cv = np.zeros_like(dms_cv) @@ -373,8 +211,7 @@ 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, fxc_ref=None, skip_xc_vref1=False): +def gen_rohf_response_sc(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=None, log=None): ''' response function for Sf=Si ''' @@ -386,7 +223,8 @@ def gen_rohf_response_sc(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=Non mol = mf.mol if log is None: log = logger.new_logger(mf) - _require_nttda_reference(mf) + if not isinstance(mf, (dft.roks.ROKS, dft.rks_symm.SymAdaptedROKS)): + raise TypeError('NTTDA response requires ROKS reference') s = (mol.nelec[0] - mol.nelec[1]) * 0.5 @@ -395,7 +233,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' and fxc_ref is None: + if xctype != 'HF': 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]) @@ -428,9 +266,7 @@ 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 skip_xc_vref1 and xctype in ('GGA', 'MGGA'): - vref1 = np.zeros_like(dms1) - elif xctype == 'LDA': + if 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': @@ -483,8 +319,7 @@ 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, fxc_ref=None, skip_xc_vref1=False): +def gen_rohf_response_sfd(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=None, log=None): ''' response function for Sf=Si-1 ''' @@ -496,7 +331,8 @@ def gen_rohf_response_sfd(mf, mo_coeff=None, mo_occ=None, hermi=0, max_memory=No mol = mf.mol if log is None: log = logger.new_logger(mf) - _require_nttda_reference(mf) + if not isinstance(mf, (dft.roks.ROKS, dft.rks_symm.SymAdaptedROKS)): + raise TypeError('NTTDA response requires ROKS reference') s = (mol.nelec[0] - mol.nelec[1]) * 0.5 @@ -506,7 +342,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' and fxc_ref is None: + if xctype != 'HF': 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]) @@ -538,9 +374,7 @@ 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 skip_xc_vref1 and xctype in ('GGA', 'MGGA'): - vref1 = np.zeros_like(dms1) - elif xctype == 'LDA': + if 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': @@ -610,9 +444,18 @@ def gen_vind_sfu(td): vresp, fockz = gen_rohf_response_sfu(mf, mo_coeff=mo_coeff, mo_occ=mo_occ, hermi=0, max_memory=td.max_memory, log=log) - fock0 = _reference_fock0(mf, td.nobeta) - focka = fock0 + fockz - fockb = fock0 - fockz + 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 fock_v = orbvs.T @ focka @ orbvs fock_c = orbcs.T @ fockb @ orbcs @@ -649,7 +492,6 @@ 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] @@ -660,22 +502,21 @@ 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, - fxc_ref=fxc_ref, - skip_xc_vref1=use_mo_grid_fxc1) + max_memory=td.max_memory, log=log) - fock0 = _reference_fock0(mf, td.nobeta) - focka = fock0 + fockz - fockb = fock0 - fockz + 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 fock_coco1 = orbos.T @ (fock0 - fockz) @ orbos fock_coco2 = orbcs.T @ (fock0 - fockz) @ orbcs @@ -725,37 +566,6 @@ 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) @@ -812,7 +622,6 @@ 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] @@ -828,20 +637,17 @@ 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, - fxc_ref=fxc_ref, - skip_xc_vref1=use_mo_grid_fxc1) + max_memory=td.max_memory, log=log) - fock0 = _reference_fock0(mf, td.nobeta) + if td.nobeta: + dma, dmb = mf.make_rdm1() + dm0 = 0.5 * (dma + dmb) + fock = mf.get_fock(dm=np.array([dm0, dm0])) + fock0 = 0.5 * (fock.focka + fock.fockb) + else: + fock = mf.get_fock() + fock0 = 0.5 * (fock.focka + fock.fockb) fock_coco0 = orbos.T @ (fock0 - fockz) @ orbos fock_coco1 = orbcs.T @ (fock0 + fockz) @ orbcs @@ -894,30 +700,6 @@ 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) @@ -977,31 +759,11 @@ class NTTDA(TDBase): nobeta: True for problemstic cases where there is no local beta electrons ''' - deltaS = getattr(__config__, 'NTTDA_delta_S', -1) - nobeta = getattr(__config__, 'NTTDA_nobeta', False) + deltaS = -1 + 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 @@ -1489,6 +1251,5 @@ 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 bcb4394..f807a32 100644 --- a/src/nest/nttda/tests/test_nttda.py +++ b/src/nest/nttda/tests/test_nttda.py @@ -1,4 +1,4 @@ -# Copyright 2021-2024 The PySCF Developers. All Rights Reserved. +# 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. @@ -13,110 +13,9 @@ # limitations under the License. import unittest -from unittest import mock import numpy as np from pyscf import gto -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]), - }, - }, -} +from nest import nttda class KnownValues(unittest.TestCase): @@ -125,71 +24,94 @@ 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 _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) - def test_hf_nttda(self): - self._check_functional('HF') + 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) + + 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_svwn_nttda(self): - self._check_functional('SVWN') + 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) def test_m062x_nttda(self): - self._check_functional('M062X') + 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) def test_cam_b3lyp_nttda(self): - 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") + 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') unittest.main() diff --git a/src/nest/nttda/tests/test_nttda_dz0scf.py b/src/nest/nttda/tests/test_nttda_dz0scf.py deleted file mode 100644 index 1fbb422..0000000 --- a/src/nest/nttda/tests/test_nttda_dz0scf.py +++ /dev/null @@ -1,164 +0,0 @@ -#!/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 81f6be8..d9983da 100644 --- a/src/nest/nttda/tests/test_nttda_oscillator_strength.py +++ b/src/nest/nttda/tests/test_nttda_oscillator_strength.py @@ -50,11 +50,7 @@ def test_svwn_nttda_oscillator_strength(self): [-0.068297131072, -0.365336383706, 0.031212437760], ]) ref_f = np.array([0.007681628385, 0.037787591904, 0.026266335987]) - # 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() + td = mf.NTTDA().set(deltaS=-1, nstates=4).run() self.assertTrue(np.all(td.converged)) dip = td.transition_dipole() dip_outer = np.einsum('nx,ny->nxy', dip.conj(), dip) @@ -68,7 +64,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, conv_tol=1e-9, lindep=1e-18).run() + td = mf.NTTDA().set(deltaS=0, nstates=4).run() self.assertTrue(np.all(td.converged)) dip = td.transition_dipole() dip_outer = np.einsum('nx,ny->nxy', dip.conj(), dip) @@ -82,7 +78,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, conv_tol=1e-9, lindep=1e-18).run() + td = mf.NTTDA().set(deltaS=1, nstates=4).run() self.assertTrue(np.all(td.converged)) dip = td.transition_dipole() dip_outer = np.einsum('nx,ny->nxy', dip.conj(), dip)