diff --git a/adcc/ElectronicStates.py b/adcc/ElectronicStates.py index 417daf886..7df50d2b7 100644 --- a/adcc/ElectronicStates.py +++ b/adcc/ElectronicStates.py @@ -160,6 +160,19 @@ def excitation_energy(self): """Excitation energies including all corrections in atomic units""" return self._excitation_energy + @property + def total_energy(self): + """Total energies (ground state + excitation) of all computed states""" + return np.array([self._total_energy(i) for i in range(self.size)]) + + def _total_energy(self, state_n: int) -> float: + """Computes the total energy for a single state""" + if self.method.level == MethodLevel.ZERO: + return (self._excitation_energy[state_n] + + self.reference_state.energy_scf) + return (self._excitation_energy[state_n] + + self.ground_state.energy(self.method.level.to_int())) + @property def excitation_energy_uncorrected(self): """Excitation energies without any corrections in atomic units""" diff --git a/adcc/Excitation.py b/adcc/Excitation.py index 71337181c..517dfb959 100644 --- a/adcc/Excitation.py +++ b/adcc/Excitation.py @@ -94,6 +94,21 @@ def transition_quadrupole_moment_velocity(self, state_n=self.index, gauge_origin=gauge_origin ) + @property + def ground_state(self): + """The ground state (LazyMp) of the parent ElectronicStates""" + return self._parent_state.ground_state + + @property + def reference_state(self): + """The HF reference state of the parent ElectronicStates""" + return self._parent_state.reference_state + + @property + def method(self): + """The ADC method of the parent ElectronicStates""" + return self._parent_state.method + @property def oscillator_strength(self) -> np.float64: """The oscillator strength""" diff --git a/adcc/ReferenceState.py b/adcc/ReferenceState.py index d6d718c1d..1d58fb99d 100644 --- a/adcc/ReferenceState.py +++ b/adcc/ReferenceState.py @@ -159,7 +159,9 @@ def __init__(self, hfdata, core_orbitals=None, frozen_core=None, ) self.environment = None # no environment attached by default - for name in ["excitation_energy_corrections", "environment"]: + for name in ["excitation_energy_corrections", + "environment", + "gradient_provider"]: if hasattr(hfdata, name): setattr(self, name, getattr(hfdata, name)) diff --git a/adcc/StateView.py b/adcc/StateView.py index 186d113de..639496742 100644 --- a/adcc/StateView.py +++ b/adcc/StateView.py @@ -39,6 +39,11 @@ def excitation_energy(self): """Excitation energy including all corrections in atomic units""" return self._parent_state.excitation_energy[self.index] + @property + def total_energy(self) -> float: + """Total energy (ground state + excitation) of this state""" + return self._parent_state._total_energy(self.index) + @property def excitation_energy_uncorrected(self): """Excitation energy without any corrections in atomic units""" diff --git a/adcc/__init__.py b/adcc/__init__.py index 672089217..c91d31397 100644 --- a/adcc/__init__.py +++ b/adcc/__init__.py @@ -47,6 +47,7 @@ from .TwoParticleOperator import TwoParticleOperator from .TwoParticleDensity import TwoParticleDensity from .opt_einsum_integration import register_with_opt_einsum +from .gradients import NuclearGradientScanner, nuclear_gradient # This has to be the last set of import from .guess import (guess_symmetries, guess_zero, guesses_any, guesses_singlet, @@ -68,6 +69,7 @@ "guess_symmetries", "guesses_spin_flip", "guess_zero", "LazyMp", "adc0", "cis", "adc1", "adc2", "adc2x", "adc3", "cvs_adc0", "cvs_adc1", "cvs_adc2", "cvs_adc2x", "cvs_adc3", + "nuclear_gradient", "NuclearGradientScanner", "banner"] __version__ = "0.18.0" diff --git a/adcc/backends/psi4.py b/adcc/backends/psi4.py index 36bb146fa..9356c08dc 100644 --- a/adcc/backends/psi4.py +++ b/adcc/backends/psi4.py @@ -24,6 +24,8 @@ from libadcc import HartreeFockProvider +from adcc.gradients import GradientComponents + import psi4 from .EriBuilder import EriBuilder @@ -32,6 +34,72 @@ from ..OneParticleOperator import OneParticleOperator +class Psi4GradientProvider: + def __init__(self, wfn): + self.wfn = wfn + self.mol = self.wfn.molecule() + self.backend = "psi4" + self.mints = psi4.core.MintsHelper(self.wfn) + + def correlated_gradient(self, g1_ao, w_ao, g2_ao_1, g2_ao_2): + """ + g1_ao: relaxed one-particle density matrix + w_ao: energy-weighted density matrix + g2_ao_1, g2_ao_2: relaxed two-particle density matrices + """ + natoms = self.mol.natom() + Gradient = {} + Gradient["N"] = np.zeros((natoms, 3)) + Gradient["S"] = np.zeros((natoms, 3)) + Gradient["T"] = np.zeros((natoms, 3)) + Gradient["V"] = np.zeros((natoms, 3)) + Gradient["OEI"] = np.zeros((natoms, 3)) + Gradient["ERI"] = np.zeros((natoms, 3)) + Gradient["Total"] = np.zeros((natoms, 3)) + + # 1st Derivative of Nuclear Repulsion + Gradient["N"] = psi4.core.Matrix.to_array( + self.mol.nuclear_repulsion_energy_deriv1([0, 0, 0]) + ) + # Build Integral Derivatives + cart = ['_X', '_Y', '_Z'] + oei_dict = {"S": "OVERLAP", "T": "KINETIC", "V": "POTENTIAL"} + + deriv1_mat = {} + deriv1_np = {} + # 1st Derivative of OEIs + for atom in range(natoms): + for key in oei_dict: + string = key + str(atom) + deriv1_mat[string] = self.mints.ao_oei_deriv1(oei_dict[key], atom) + for p in range(3): + map_key = string + cart[p] + deriv1_np[map_key] = np.asarray(deriv1_mat[string][p]) + if key == "S": + Gradient["S"][atom, p] = np.sum(w_ao * deriv1_np[map_key]) + else: + Gradient[key][atom, p] = np.sum(g1_ao * deriv1_np[map_key]) + # Build Total OEI Gradient + Gradient["OEI"] = Gradient["T"] + Gradient["V"] + Gradient["S"] + + # Build TE contributions + for atom in range(natoms): + deriv2 = self.mints.ao_tei_deriv1(atom) + for p in range(3): + deriv2_np = np.asarray(deriv2[p]) + Gradient["ERI"][atom, p] += np.einsum( + 'pqrs,prqs->', g2_ao_1, deriv2_np, optimize=True + ) + Gradient["ERI"][atom, p] -= np.einsum( + 'pqrs,psqr->', g2_ao_2, deriv2_np, optimize=True + ) + ret = GradientComponents( + natoms, Gradient["N"], Gradient["S"], + Gradient["T"] + Gradient["V"], Gradient["ERI"] + ) + return ret + + class Psi4OperatorIntegralProvider: available: tuple[str, ...] = ( "overlap", "electric_dipole", "electric_dipole_velocity", "magnetic_dipole", @@ -169,6 +237,7 @@ def __init__(self, wfn): wfn.nalpha(), wfn.nbeta(), self.restricted) self.operator_integral_provider = Psi4OperatorIntegralProvider(self.wfn) + self.gradient_provider = Psi4GradientProvider(self.wfn) self.environment = None self.environment_implementation = None diff --git a/adcc/backends/pyscf.py b/adcc/backends/pyscf.py index c00c6da2c..c692542c1 100644 --- a/adcc/backends/pyscf.py +++ b/adcc/backends/pyscf.py @@ -20,19 +20,247 @@ ## along with adcc. If not, see . ## ## --------------------------------------------------------------------- +import contextlib +import os +import tempfile + +import h5py import numpy as np from libadcc import HartreeFockProvider +from adcc.gradients import GradientComponents +from adcc.gradients.TwoParticleDensityMatrix import ao_pair_index, ao_pair_indices + from .EriBuilder import EriBuilder from ..exceptions import InvalidReference from ..ElectronicStates import EnergyCorrection from ..OneParticleOperator import OneParticleOperator -from pyscf import ao2mo, gto, scf +from pyscf import ao2mo, gto, scf, grad from pyscf.solvent import ddcosmo +class PyScfGradientProvider: + def __init__(self, scfres): + self.scfres = scfres + self.mol = self.scfres.mol + self.backend = "pyscf" + + def _empty_gradient(self): + natoms = self.mol.natm + return { + "S": np.zeros((natoms, 3)), + "T+V": np.zeros((natoms, 3)), + "ERI": np.zeros((natoms, 3)), + } + + def _add_one_electron_terms(self, gradient_components, g1_ao, w_ao, + pyscf_gradient): + hcore_deriv = pyscf_gradient.hcore_generator() + Sx = -1.0 * self.mol.intor('int1e_ipovlp', aosym='s1') + ao_slices = self.mol.aoslice_by_atom() + for ia in range(self.mol.natm): + k0, k1 = ao_slices[ia, 2:] + + # derivative of the overlap matrix + Sx_a = np.zeros_like(Sx) + Sx_a[:, k0:k1] = Sx[:, k0:k1] + Sx_a += Sx_a.transpose(0, 2, 1) + gradient_components["S"][ia] += np.einsum( + "xpq,pq->x", Sx_a, w_ao + ) + + # derivative of the core Hamiltonian + Hx_a = hcore_deriv(ia) + gradient_components["T+V"][ia] += np.einsum( + "xpq,pq->x", Hx_a, g1_ao + ) + + def _final_gradient_components(self, pyscf_gradient, gradient_components): + return GradientComponents( + self.mol.natm, pyscf_gradient.grad_nuc(), gradient_components["S"], + gradient_components["T+V"], gradient_components["ERI"] + ) + + def correlated_gradient(self, g1_ao, w_ao, g2_ao_1, g2_ao_2): + """Full-AO reference path. Allocates the full derivative ERI tensor.""" + Gradient = self._empty_gradient() + + # TODO: does RHF/UHF matter here? + gradient = grad.RHF(self.scfres) + self._add_one_electron_terms(Gradient, g1_ao, w_ao, gradient) + ERIx = -1.0 * self.mol.intor('int2e_ip1', aosym='s1') + + ao_slices = self.mol.aoslice_by_atom() + for ia in range(self.mol.natm): + # TODO: only contract/compute with specific slices + # of density matrices (especially TPDM) + # this requires a lot of work however... + k0, k1 = ao_slices[ia, 2:] + + # derivatives of the ERIs + ERIx_a = np.zeros_like(ERIx) + ERIx_a[:, k0:k1] = ERIx[:, k0:k1] + ERIx_a += ( + + ERIx_a.transpose(0, 2, 1, 4, 3) + + ERIx_a.transpose(0, 3, 4, 1, 2) + + ERIx_a.transpose(0, 4, 3, 2, 1) + ) + Gradient["ERI"][ia] += np.einsum( + "pqrs,xprqs->x", g2_ao_1, ERIx_a, optimize=True + ) + Gradient["ERI"][ia] -= np.einsum( + "pqrs,xpsqr->x", g2_ao_2, ERIx_a, optimize=True + ) + return self._final_gradient_components(gradient, Gradient) + + def _symmetrized_density_slice(self, pair_density, a0, a1): + """ + Build the four-center derivative-symmetrized packed density slice. + + ``pair_density`` stores ``D[p,r,q,s]`` with packed ket pair ``q,s``. + For a derivative integral block ``E[a,b,q,s]`` this returns the packed + equivalent of + + ``D[a,b,q,s] + D[b,a,s,q] + D[q,s,a,b] + D[s,q,b,a]``. + """ + nao = pair_density.shape[0] + qidx, sidx = ao_pair_indices(nao) + offdiag = qidx != sidx + any_offdiag = bool(np.any(offdiag)) + density_slice = np.asarray(pair_density[a0:a1, :, :]).copy() + density_slice += np.asarray(pair_density[:, a0:a1, :]).transpose(1, 0, 2) + + for local_a, a in enumerate(range(a0, a1)): + for b in range(nao): + pair_ab = int(ao_pair_index(a, b)) + ket_as_bra = np.asarray(pair_density[:, :, pair_ab]) + term = ket_as_bra[qidx, sidx] + if any_offdiag: + term[offdiag] += ket_as_bra[sidx[offdiag], qidx[offdiag]] + if a == b: + term *= 2.0 + density_slice[local_a, b, :] += term + return density_slice + + def _contract_eri_with_packed_density(self, pair_density, + shell_chunk_size=1): + """ + Contract derivative ERIs with packed AO-pair effective density. + + PySCF only packs the ket pair for ``aosym=\"s2kl\"``. Off-diagonal + packed density entries therefore contain the sum of both ket orders. + Shell slices are kept atom-local so atom AO and shell ranges cannot be + mixed accidentally. + """ + if shell_chunk_size <= 0: + raise ValueError("shell_chunk_size needs to be positive.") + eri_grad = np.zeros((self.mol.natm, 3)) + ao_slices = self.mol.aoslice_by_atom() + ao_loc = self.mol.ao_loc_nr() + nbas = self.mol.nbas + + for ia in range(self.mol.natm): + sh0, sh1 = ao_slices[ia, :2] + for p0 in range(sh0, sh1, shell_chunk_size): + p1 = min(p0 + shell_chunk_size, sh1) + a0, a1 = ao_loc[p0], ao_loc[p1] + shls_slice = (p0, p1, 0, nbas, 0, nbas, 0, nbas) + erix = -1.0 * self.mol.intor( + 'int2e_ip1', comp=3, aosym='s2kl', shls_slice=shls_slice + ) + density_slice = self._symmetrized_density_slice( + pair_density, a0, a1 + ) + eri_grad[ia] += np.einsum( + "xabk,abk->x", erix, density_slice, optimize=True + ) + return eri_grad + + def correlated_gradient_direct(self, g1_ao, w_ao, g2, refstate=None, + shell_chunk_size=1, pair_chunk_size=None, + pair_density_storage="memory", + scratch_directory=None): + """ + Direct PySCF gradient path from the MO/block-sparse TPDM. + + The TPDM is transformed block-by-block to the packed AO-pair effective + density. It never forms ``g2_ao_1``/``g2_ao_2`` or the full derivative + ERI tensor. + + ``pair_density_storage`` selects where the packed AO-pair density is + kept: + + - ``"memory"`` (default): an in-memory ``(nao, nao, npair)`` array. + - ``"hdf5"``: a temporary HDF5 file under ``scratch_directory`` (or the + ``PYSCF_TMPDIR`` / system temp directory), removed after contraction. + + Note that out-of-core storage only bounds the packed *output* density. + The transform's in-memory working buffer is sized by ``pair_chunk_size`` + (the number of AO pairs processed at once). When ``pair_chunk_size`` is + left unset, the ``"hdf5"`` path therefore picks a bounded default so the + working buffer does not grow to the full ``(nao, nao, npair)`` size; the + ``"memory"`` path keeps the full ``npair`` chunk. + """ + valid_storage = {"memory", "hdf5"} + if pair_density_storage not in valid_storage: + raise ValueError( + f"pair_density_storage needs to be one of " + f"{sorted(valid_storage)}." + ) + + Gradient = self._empty_gradient() + gradient = grad.RHF(self.scfres) + self._add_one_electron_terms(Gradient, g1_ao, w_ao, gradient) + + nao = self.mol.nao_nr() + npair = nao * (nao + 1) // 2 + if pair_chunk_size is None and pair_density_storage == "hdf5": + # Bound the in-memory working buffer for the out-of-core path + # instead of silently allocating the full (nao, nao, npair) chunk. + pair_chunk_size = min(npair, 256) + + h5file = None + h5path = None + pair_density = None + try: + if pair_density_storage == "memory": + pair_density = np.zeros((nao, nao, npair)) + else: # "hdf5" + scratch_directory = scratch_directory or os.environ.get( + "PYSCF_TMPDIR", tempfile.gettempdir() + ) + tmp = tempfile.NamedTemporaryFile( + prefix="adcc_pyscf_gradient_", suffix=".h5", + dir=scratch_directory, delete=False + ) + h5path = tmp.name + tmp.close() + h5file = h5py.File(h5path, "w") + pair_density = h5file.create_dataset( + "pair_density", shape=(nao, nao, npair), dtype="f8", + chunks=(min(nao, 16), min(nao, 16), min(npair, 256)) + ) + + g2.to_ao_pair_density( + refstate, pair_chunk_size=pair_chunk_size, out=pair_density + ) + Gradient["ERI"] += self._contract_eri_with_packed_density( + pair_density, shell_chunk_size=shell_chunk_size + ) + finally: + # Cleanup must not mask an exception raised inside the try block + # (e.g. an out-of-space write error followed by a failing flush). + if h5file is not None: + with contextlib.suppress(Exception): + h5file.close() + if h5path is not None: + with contextlib.suppress(OSError): + os.unlink(h5path) + return self._final_gradient_components(gradient, Gradient) + + class PyScfOperatorIntegralProvider: available: tuple[str, ...] = ( "overlap", "electric_dipole", "electric_dipole_velocity", "magnetic_dipole", @@ -218,6 +446,7 @@ def __init__(self, scfres): self.operator_integral_provider = PyScfOperatorIntegralProvider( self.scfres ) + self.gradient_provider = PyScfGradientProvider(self.scfres) if not self.restricted: assert self.scfres.mo_coeff[0].shape[1] == \ self.scfres.mo_coeff[1].shape[1] diff --git a/adcc/gradients/TwoParticleDensityMatrix.py b/adcc/gradients/TwoParticleDensityMatrix.py new file mode 100644 index 000000000..a5d6de0ca --- /dev/null +++ b/adcc/gradients/TwoParticleDensityMatrix.py @@ -0,0 +1,553 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2021 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +import numpy as np + +import libadcc +import adcc.block as b +from adcc.MoSpaces import split_spaces +from adcc.Tensor import Tensor +from adcc.functions import einsum + + +def ao_pair_index(p, q): + """ + Return the packed index of the symmetric AO pair ``(p, q)``. + + AO pairs are stored in lower-triangular order: the pair is identified by + ``(max(p, q), min(p, q))`` and mapped to + ``max(p, q) * (max(p, q) + 1) // 2 + min(p, q)``. This is the conventional + symmetric-pair packing of a derivative-ERI backend that exposes only one + entry per symmetric ket pair (matching, e.g., PySCF ``aosym=\"s2kl\"``). + """ + p, q = np.maximum(p, q), np.minimum(p, q) + return p * (p + 1) // 2 + q + + +def ao_pair_indices(nao): + """ + Return arrays ``q, s`` enumerating the lower-triangular AO pairs. + """ + q, s = np.tril_indices(nao) + return q.astype(np.intp, copy=False), s.astype(np.intp, copy=False) + + +class TwoParticleDensityMatrix: + """ + Two-particle density matrix (TPDM) used for gradient evaluations + """ + def __init__(self, spaces): + if hasattr(spaces, "mospaces"): + self.mospaces = spaces.mospaces + else: + self.mospaces = spaces + # Set reference_state if possible + if isinstance(spaces, libadcc.ReferenceState): + self.reference_state = spaces + elif hasattr(spaces, "reference_state"): + assert isinstance(spaces.reference_state, libadcc.ReferenceState) + self.reference_state = spaces.reference_state + + occs = sorted(self.mospaces.subspaces_occupied, reverse=True) + virts = sorted(self.mospaces.subspaces_virtual, reverse=True) + self.orbital_subspaces = occs + virts + # check that orbital subspaces are correct + assert sum(self.mospaces.n_orbs(ss) for ss in self.orbital_subspaces) \ + == self.mospaces.n_orbs("f") + # set the canonical blocks explicitly + self.blocks = [ + b.oooo, b.ooov, b.oovv, + b.ovov, b.ovvv, b.vvvv, + ] + if self.mospaces.has_core_occupied_space: + self.blocks += [ + b.cccc, b.ococ, b.cvcv, + b.ocov, b.cccv, b.cocv, b.ocoo, + b.ccco, b.occv, b.ccvv, b.ocvv, + ] + # make sure we didn't add any block twice! + assert len(list(set(self.blocks))) == len(self.blocks) + self._tensors = {} + + @property + def shape(self): + """ + Returns the shape tuple of the TwoParticleDensityMatrix + """ + size = self.mospaces.n_orbs("f") + return 4 * (size,) + + @property + def size(self): + """ + Returns the number of elements of the TwoParticleDensityMatrix + """ + return np.prod(self.shape) + + def __setitem__(self, block, tensor): + """ + Assigns a tensor to the specified block + """ + if block not in self.blocks: + raise KeyError(f"Invalid block {block} assigned. " + f"Available blocks are: {self.blocks}.") + s1, s2, s3, s4 = split_spaces(block) + expected_shape = (self.mospaces.n_orbs(s1), + self.mospaces.n_orbs(s2), + self.mospaces.n_orbs(s3), + self.mospaces.n_orbs(s4)) + if expected_shape != tensor.shape: + raise ValueError("Invalid shape of incoming tensor. " + f"Expected shape {expected_shape}, but " + f"got shape {tensor.shape} instead.") + self._tensors[block] = tensor + + def __getitem__(self, block): + if block not in self.blocks: + raise KeyError(f"Invalid block {block} requested. " + f"Available blocks are: {self.blocks}.") + if block not in self._tensors: + sym = libadcc.make_symmetry_eri(self.mospaces, block) + self._tensors[block] = Tensor(sym) + return self._tensors[block] + + def to_ndarray(self): + raise NotImplementedError("ndarray export not implemented for TPDM.") + + def copy(self): + """ + Return a deep copy of the TwoParticleDensityMatrix + """ + ret = TwoParticleDensityMatrix(self.mospaces) + for bl in self.blocks_nonzero: + ret[bl] = self.block(bl).copy() + if hasattr(self, "reference_state"): + ret.reference_state = self.reference_state + return ret + + @property + def blocks_nonzero(self): + """ + Returns a list of the non-zero block labels + """ + return [b for b in self.blocks if b in self._tensors] + + def is_zero_block(self, block): + """ + Checks if block is explicitly marked as zero block. + Returns False if the block does not exist. + """ + if block not in self.blocks: + return False + return block not in self.blocks_nonzero + + def block(self, block): + """ + Returns tensor of the given block. + Does not create a block in case it is marked as a zero block. + Use __getitem__ for that purpose. + """ + if block not in self.blocks_nonzero: + raise KeyError("The block function does not support " + "access to zero-blocks. Available non-zero " + f"blocks are: {self.blocks_nonzero}.") + return self._tensors[block] + + def __getattr__(self, attr): + return self.__getitem__(b.__getattr__(attr)) + + def __setattr__(self, attr, value): + try: + self.__setitem__(b.__getattr__(attr), value) + except AttributeError: + super().__setattr__(attr, value) + + def set_zero_block(self, block): + """ + Set a given block as zero block + """ + if block not in self.blocks: + raise KeyError(f"Invalid block {block} set as zero block. " + f"Available blocks are: {self.blocks}.") + self._tensors.pop(block) + + def __transform_to_ao(self, refstate_or_coefficients): + if not len(self.blocks_nonzero): + raise ValueError("At least one non-zero block is needed to " + "transform the TwoParticleDensityMatrix.") + if isinstance(refstate_or_coefficients, libadcc.ReferenceState): + hf = refstate_or_coefficients + coeff_map = {} + for sp in self.orbital_subspaces: + coeff_map[sp + "_a"] = hf.orbital_coefficients_alpha(sp + "b") + coeff_map[sp + "_b"] = hf.orbital_coefficients_beta(sp + "b") + else: + coeff_map = refstate_or_coefficients + + g2_ao_1 = 0 + g2_ao_2 = 0 + transf = "ip,jq,ijkl,kr,ls->pqrs" + cc = coeff_map + for block in self.blocks_nonzero: + s1, s2, s3, s4 = split_spaces(block) + ten = self[block] + aaaa = einsum(transf, cc[f"{s1}_a"], cc[f"{s2}_a"], + ten, cc[f"{s3}_a"], cc[f"{s4}_a"]) + bbbb = einsum(transf, cc[f"{s1}_b"], cc[f"{s2}_b"], + ten, cc[f"{s3}_b"], cc[f"{s4}_b"]) + g2_ao_1 += ( + + aaaa + + bbbb + + einsum(transf, cc[f"{s1}_a"], cc[f"{s2}_b"], + ten, cc[f"{s3}_a"], cc[f"{s4}_b"]) # abab + + einsum(transf, cc[f"{s1}_b"], cc[f"{s2}_a"], + ten, cc[f"{s3}_b"], cc[f"{s4}_a"]) # baba + ) + g2_ao_2 += ( + + aaaa + + bbbb + + einsum(transf, cc[f"{s1}_a"], cc[f"{s2}_b"], + ten, cc[f"{s3}_b"], cc[f"{s4}_a"]) # abba + + einsum(transf, cc[f"{s1}_b"], cc[f"{s2}_a"], + ten, cc[f"{s3}_a"], cc[f"{s4}_b"]) # baab + ) + return (g2_ao_1.evaluate(), g2_ao_2.evaluate()) + + def to_ao_basis(self, refstate_or_coefficients=None): + """ + Transform the density to the AO basis for contraction + with two-electron integrals. + ALL coefficients are already accounted for in the density matrix. + Two blocks are returned, the first one needs to be contracted with + prqs, the second one with -psqr (in Chemists' notation). + """ + if isinstance(refstate_or_coefficients, (dict, libadcc.ReferenceState)): + return self.__transform_to_ao(refstate_or_coefficients) + elif refstate_or_coefficients is None: + if not hasattr(self, "reference_state"): + raise ValueError("Argument reference_state is required if no " + "reference_state is stored in the " + "TwoParticleDensityMatrix") + return self.__transform_to_ao(self.reference_state) + else: + raise TypeError("Argument type not supported.") + + def _ao_coefficient_map(self, refstate_or_coefficients=None): + """Return AO coefficient matrices as NumPy arrays.""" + if refstate_or_coefficients is None: + if not hasattr(self, "reference_state"): + raise ValueError("Argument reference_state is required if no " + "reference_state is stored in the " + "TwoParticleDensityMatrix") + refstate_or_coefficients = self.reference_state + + if isinstance(refstate_or_coefficients, libadcc.ReferenceState): + hf = refstate_or_coefficients + coeff_map = {} + for sp in self.orbital_subspaces: + coeff_map[sp + "_a"] = hf.orbital_coefficients_alpha( + sp + "b" + ).to_ndarray() + coeff_map[sp + "_b"] = hf.orbital_coefficients_beta( + sp + "b" + ).to_ndarray() + elif isinstance(refstate_or_coefficients, dict): + coeff_map = {} + for key, coeff in refstate_or_coefficients.items(): + if hasattr(coeff, "to_ndarray"): + coeff = coeff.to_ndarray() + coeff_map[key] = np.asarray(coeff) + else: + raise TypeError("Argument type not supported.") + return coeff_map + + @staticmethod + def ao_pair_density_from_dense(g2_ao_1, g2_ao_2, out=None): + """ + Pack the effective AO density for a derivative-ERI contraction. + + The effective density is stored in the order + ``D[p,r,q,s] = g2_ao_1[p,q,r,s] - g2_ao_2[p,q,s,r]``. The last two + AO indices are packed into the lower-triangular ket pair ``q,s`` (see + :func:`ao_pair_index`). For an off-diagonal ket pair ``q != s`` the + packed entry holds the sum of both full-density orderings, since a + symmetric-ket-pair integral layout stores only one entry per pair (the + layout used, e.g., by PySCF ``aosym=\"s2kl\"``). + """ + nao = g2_ao_1.shape[0] + npair = nao * (nao + 1) // 2 + if out is None: + out = np.zeros( + (nao, nao, npair), dtype=np.result_type(g2_ao_1, g2_ao_2) + ) + else: + out[...] = 0 + qidx, sidx = ao_pair_indices(nao) + for pair, (q, s) in enumerate(zip(qidx, sidx)): + out[:, :, pair] += g2_ao_1[:, q, :, s] + out[:, :, pair] -= g2_ao_2[:, q, s, :] + if q != s: + out[:, :, pair] += g2_ao_1[:, s, :, q] + out[:, :, pair] -= g2_ao_2[:, s, q, :] + return out + + @staticmethod + def _add_direct_pair_transform(out, tensor, c1, c2, c3, c4, + qidx, sidx, sign, exchange): + """Accumulate one spin case into a packed AO-pair density chunk.""" + if exchange: + right = c2[:, qidx][:, None, :] * c3[:, sidx][None, :, :] + out += sign * np.einsum( + "ip,lr,ijkl,jkm->prm", c1, c4, tensor, right, optimize=True + ) + else: + right = c2[:, qidx][:, None, :] * c4[:, sidx][None, :, :] + out += sign * np.einsum( + "ip,kr,ijkl,jlm->prm", c1, c3, tensor, right, optimize=True + ) + + def to_ao_pair_density(self, refstate_or_coefficients=None, + pair_chunk_size=None, out=None): + """ + Transform directly to packed AO-pair effective density. + + This is the memory-bounded counterpart of ``to_ao_basis()`` for the + gradient two-electron contraction. It reproduces the spin cases from + ``__transform_to_ao`` without forming the two full AO rank-4 TPDMs: + + - ``g2_ao_1``: ``aaaa``, ``bbbb``, ``abab``, ``baba`` + - ``g2_ao_2``: ``aaaa``, ``bbbb``, ``abba``, ``baab`` + + For restricted references, equivalent spin cases are merged and + accumulated with their spin multiplicity (``aaaa``/``abab`` for + ``g2_ao_1`` and ``aaaa``/``abba`` for ``g2_ao_2``, each with a factor + of two). Raw coefficient dictionaries do not carry a trustworthy + restricted/unrestricted flag and therefore use the full spin expansion. + + The returned/filled array has shape ``(nao, nao, nao * (nao + 1) // 2)`` + and contains ``D[p,r,q,s] = g2_ao_1[p,q,r,s] - g2_ao_2[p,q,s,r]`` with + the ``q,s`` ket pair packed in lower-triangular order (see + :func:`ao_pair_index`; this matches a symmetric-ket-pair integral + layout such as PySCF ``aosym=\"s2kl\"``). Existing block prefactors in + this ``TwoParticleDensityMatrix`` are assumed to have been applied + upstream and are not changed here. + + Each MO axis of the spin-orbital TPDM splits into a leading alpha range + ``[0, n_alpha)`` and a trailing beta range ``[n_alpha, n_orbs)``; the + spin-blocked MO coefficients are non-zero only on the matching range. + For every spin case we therefore extract just the relevant rank-4 spin + sub-block via :meth:`Tensor.export_block` and contract it with the + compacted (non-zero-row) coefficients, instead of densifying the full + zero-padded block. + """ + if not len(self.blocks_nonzero): + raise ValueError("At least one non-zero block is needed to " + "transform the TwoParticleDensityMatrix.") + restricted_refstate = None + if isinstance(refstate_or_coefficients, libadcc.ReferenceState): + restricted_refstate = refstate_or_coefficients + elif (refstate_or_coefficients is None + and hasattr(self, "reference_state")): + restricted_refstate = self.reference_state + use_restricted_fast_path = bool( + getattr(restricted_refstate, "restricted", False) + ) + + coeff_map = self._ao_coefficient_map(refstate_or_coefficients) + nao = next(iter(coeff_map.values())).shape[1] + npair = nao * (nao + 1) // 2 + if pair_chunk_size is None: + pair_chunk_size = npair + if pair_chunk_size <= 0: + raise ValueError("pair_chunk_size needs to be positive.") + + if out is None: + out = np.zeros((nao, nao, npair), dtype=float) + else: + if out.shape != (nao, nao, npair): + raise ValueError("Invalid output shape for packed AO-pair density.") + out[...] = 0 + + # Compact each spin-blocked coefficient matrix to its non-zero MO rows. + # For spin-blocked coefficients these rows are contiguous and coincide + # with the alpha (``[0, n_alpha)``) or beta (``[n_alpha, n_orbs)``) + # range of the corresponding MO axis, so the same range slices the TPDM + # spin sub-block. ``compact[(space, spin)] = (coeff_rows, lo, hi)``. + compact = {} + for sp in self.orbital_subspaces: + for spin in ("a", "b"): + c = np.asarray(coeff_map[f"{sp}_{spin}"]) + nz = np.nonzero(np.any(c != 0.0, axis=1))[0] + if nz.size == 0: + compact[(sp, spin)] = (None, 0, 0) + else: + lo, hi = int(nz[0]), int(nz[-1]) + 1 + compact[(sp, spin)] = (c[lo:hi], lo, hi) + + if use_restricted_fast_path: + direct_spin_cases = [ + (("a", "a", "a", "a"), 2.0), + (("a", "b", "a", "b"), 2.0), + ] + exchange_spin_cases = [ + (("a", "a", "a", "a"), 2.0), + (("a", "b", "b", "a"), 2.0), + ] + else: + direct_spin_cases = [ + (("a", "a", "a", "a"), 1.0), + (("b", "b", "b", "b"), 1.0), + (("a", "b", "a", "b"), 1.0), + (("b", "a", "b", "a"), 1.0), + ] + exchange_spin_cases = [ + (("a", "a", "a", "a"), 1.0), + (("b", "b", "b", "b"), 1.0), + (("a", "b", "b", "a"), 1.0), + (("b", "a", "a", "b"), 1.0), + ] + required_spin_cases = { + spins for spin_cases in (direct_spin_cases, exchange_spin_cases) + for spins, _ in spin_cases + } + + # Pre-extract every required rank-4 spin sub-block exactly once (each is + # only the non-zero-row portion of a block, i.e. a small fraction of the + # zero-padded block and far smaller than a dense AO tensor). Caching + # them keeps the transform memory-bounded by the sub-block sizes plus a + # single pair chunk, while avoiding repeated extraction across chunks. + subblocks = {} # (block, spins) -> rank-4 ndarray or None + coeff_sets = {} # (block, spins) -> [coeff_rows, ...] + for block in self.blocks_nonzero: + spaces = split_spaces(block) + ten_obj = self.block(block) + for spins in required_spin_cases: + ranges = [compact[(sp, spin)] for sp, spin in zip(spaces, spins)] + coeff_sets[(block, spins)] = [r[0] for r in ranges] + if any(r[0] is None for r in ranges): + subblocks[(block, spins)] = None + continue + starts = [r[1] for r in ranges] + ends = [r[2] for r in ranges] + subblocks[(block, spins)] = ten_obj.export_block(starts, ends) + + qall, sall = ao_pair_indices(nao) + for start in range(0, npair, pair_chunk_size): + stop = min(start + pair_chunk_size, npair) + qidx = qall[start:stop] + sidx = sall[start:stop] + chunk = np.zeros((nao, nao, stop - start), dtype=out.dtype) + + offdiag = qidx != sidx + any_offdiag = bool(np.any(offdiag)) + if any_offdiag: + qswap = sidx[offdiag] + sswap = qidx[offdiag] + swapped = np.zeros( + (nao, nao, np.count_nonzero(offdiag)), dtype=out.dtype + ) + + for block in self.blocks_nonzero: + for spin_cases, sign, exchange in ( + (direct_spin_cases, +1.0, False), + (exchange_spin_cases, -1.0, True)): + for spins, prefactor in spin_cases: + sub = subblocks[(block, spins)] + if sub is None: + continue + coeffs = coeff_sets[(block, spins)] + scaled_sign = sign * prefactor + self._add_direct_pair_transform( + chunk, sub, *coeffs, qidx, sidx, scaled_sign, + exchange + ) + if any_offdiag: + self._add_direct_pair_transform( + swapped, sub, *coeffs, qswap, sswap, + scaled_sign, exchange + ) + + if any_offdiag: + chunk[:, :, offdiag] += swapped + + out[:, :, start:stop] += chunk + return out + + def __iadd__(self, other): + if self.mospaces != other.mospaces: + raise ValueError("Cannot add TwoParticleDensityMatrices with " + "differing mospaces.") + + for bl in other.blocks_nonzero: + if self.is_zero_block(bl): + self[bl] = other.block(bl).copy() + else: + self[bl] = self.block(bl) + other.block(bl) + + # Update ReferenceState pointer + if hasattr(self, "reference_state"): + if hasattr(other, "reference_state") \ + and self.reference_state != other.reference_state: + delattr(self, "reference_state") + return self + + def __isub__(self, other): + if self.mospaces != other.mospaces: + raise ValueError("Cannot subtract TwoParticleDensityMatrix with " + "differing mospaces.") + + for bl in other.blocks_nonzero: + if self.is_zero_block(bl): + self[bl] = -1.0 * other.block(bl) # The copy is implicit + else: + self[bl] = self.block(bl) - other.block(bl) + + # Update ReferenceState pointer + if hasattr(self, "reference_state"): + if hasattr(other, "reference_state") \ + and self.reference_state != other.reference_state: + delattr(self, "reference_state") + return self + + def __imul__(self, other): + if not isinstance(other, (float, int)): + return NotImplemented + for bl in self.blocks_nonzero: + self[bl] = self.block(bl) * other + return self + + def __add__(self, other): + return self.copy().__iadd__(other) + + def __sub__(self, other): + return self.copy().__isub__(other) + + def __mul__(self, other): + return self.copy().__imul__(other) + + def __rmul__(self, other): + return self.copy().__imul__(other) + + def evaluate(self): + for bl in self.blocks_nonzero: + self.block(bl).evaluate() + return self diff --git a/adcc/gradients/__init__.py b/adcc/gradients/__init__.py new file mode 100644 index 000000000..8465a1503 --- /dev/null +++ b/adcc/gradients/__init__.py @@ -0,0 +1,263 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2021 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +from dataclasses import dataclass +from typing import Dict, Union, Optional +import numpy as np + +from adcc.LazyMp import LazyMp +from adcc.Excitation import Excitation +from adcc.timings import Timer +from adcc.functions import einsum, evaluate + +from adcc.OneParticleDensity import OneParticleDensity +from adcc.NParticleOperator import OperatorSymmetry, product_trace +from .TwoParticleDensityMatrix import TwoParticleDensityMatrix +from .orbital_response import ( + orbital_response, orbital_response_rhs, energy_weighted_density_matrix +) +from .amplitude_response import amplitude_relaxed_densities +from .scanner import (ExcitedStateTarget, GroundStateTarget, + NuclearGradientScanner, density_overlap_score) + +__all__ = [ + "nuclear_gradient", "NuclearGradientScanner", + "GroundStateTarget", "ExcitedStateTarget", "density_overlap_score", +] + + +@dataclass(frozen=True) +class GradientComponents: + natoms: int + nuc: np.ndarray + overlap: np.ndarray + hcore: np.ndarray + two_electron: np.ndarray + custom: Optional[Dict[str, np.ndarray]] = None + + @property + def total(self): + """Returns the total gradient""" + ret = sum([self.nuc, self.overlap, self.hcore, self.two_electron]) + if self.custom is None: + return ret + for c in self.custom: + ret += self.custom[c] + return ret + + @property + def one_electron(self): + """Returns the one-electron gradient""" + return sum([self.nuc, self.overlap, self.hcore]) + + +@dataclass(frozen=True) +class GradientResult: + excitation_or_mp: Union[LazyMp, Excitation] + components: GradientComponents + g1: OneParticleDensity + g2: TwoParticleDensityMatrix + timer: Timer + g1a: Optional[OneParticleDensity] = None + g2a: Optional[TwoParticleDensityMatrix] = None + + @property + def reference_state(self): + return self.excitation_or_mp.reference_state + + @property + def _energy(self): + """Compute energy based on density matrices + for testing purposes""" + if self.g1a is None: + raise ValueError("No unrelaxed one-particle " + "density available.") + if self.g2a is None: + raise ValueError("No unrelaxed two-particle " + "density available.") + ret = 0.0 + hf = self.reference_state + for b in self.g1a.blocks_nonzero: + ret += self.g1a[b].dot(hf.fock(b)) + for b in self.g2a.blocks_nonzero: + ret += self.g2a[b].dot(hf.eri(b)) + return ret + + @property + def dipole_moment_relaxed(self): + """Returns the orbital-relaxed electric dipole moment""" + return self.__dipole_moment_electric(self.g1) + + @property + def dipole_moment_unrelaxed(self): + """Returns the unrelaxed electric dipole moment""" + if self.g1a is None: + raise ValueError("No unrelaxed one-particle " + "density available.") + hf = self.reference_state + return self.__dipole_moment_electric(self.g1a + hf.density) + + @property + def total(self): + """Returns the total gradient""" + return self.components.total + + def __dipole_moment_electric(self, dm): + dips = self.reference_state.operators.electric_dipole + elec_dip = -1.0 * np.array( + [product_trace(dm, dip) for dip in dips] + ) + return elec_dip + self.reference_state.nuclear_dipole + + +def nuclear_gradient(excitation_or_mp, conv_tol=1e-9, eri_contraction=None, + eri_shell_chunk_size=1, eri_pair_chunk_size=None, + eri_pair_density_storage="memory"): + if isinstance(excitation_or_mp, LazyMp): + mp = excitation_or_mp + elif isinstance(excitation_or_mp, Excitation): + mp = excitation_or_mp.ground_state + else: + raise TypeError("Gradient can only be computed for " + "Excitation or LazyMp object.") + + timer = Timer() + hf = mp.reference_state + with timer.record("amplitude_response"): + g1a, g2a = amplitude_relaxed_densities(excitation_or_mp) + + with timer.record("orbital_response"): + rhs = orbital_response_rhs(hf, g1a, g2a).evaluate() + lam = orbital_response(hf, rhs, conv_tol=conv_tol) + + # orbital-relaxed OPDM (without reference state) + g1o = g1a.copy() + g1o.ov = 0.5 * lam.ov + if hf.has_core_occupied_space: + g1o.cv = 0.5 * lam.cv + # For NOSYMMETRY operators (CVS-ADC1+), explicitly set transpose blocks + if g1o.symmetry == OperatorSymmetry.NOSYMMETRY: + from adcc.functions import transpose + g1o.vo = transpose(g1o.ov) + g1o.vc = transpose(g1o.cv) + if not g1o.is_zero_block("o2o1"): + g1o.oc = transpose(g1o.co) + # orbital-relaxed OPDM (including reference state) + g1 = g1o.copy() + g1 += hf.density + + with timer.record("energy_weighted_density_matrix"): + w = energy_weighted_density_matrix(hf, g1o, g2a) + + # build two-particle density matrices for contraction with ERIs + # prefactors see eqs 17 and A4 in 10.1063/1.5085117 + with timer.record("form_tpdm"): + g2_hf = TwoParticleDensityMatrix(hf) + g2_oresp = TwoParticleDensityMatrix(hf) + delta_ij = hf.density.oo + if hf.has_core_occupied_space: + delta_IJ = hf.density.cc + + g2_hf.oooo = 0.25 * (- einsum("li,jk->ijkl", delta_ij, delta_ij) + + einsum("ki,jl->ijkl", delta_ij, delta_ij)) + g2_hf.cccc = -0.5 * einsum("IK,JL->IJKL", delta_IJ, delta_IJ) + g2_hf.ococ = -1.0 * einsum("ik,JL->iJkL", delta_ij, delta_IJ) + + g2_oresp.cccc = einsum("IK,JL->IJKL", delta_IJ, g1o.cc + delta_IJ) + g2_oresp.ococ = ( + + einsum("ik,JL->iJkL", delta_ij, g1o.cc + 2.0 * delta_IJ) + + einsum("ik,JL->iJkL", g1o.oo, delta_IJ) + ) + g2_oresp.oooo = einsum("ij,kl->kilj", delta_ij, g1o.oo) + g2_oresp.ovov = einsum("ij,ab->iajb", delta_ij, g1o.vv) + g2_oresp.cvcv = einsum("IJ,ab->IaJb", delta_IJ, g1o.vv) + g2_oresp.ocov = 2 * einsum("ik,Ja->iJka", delta_ij, g1o.cv) + g2_oresp.cccv = 2 * einsum("IK,Ja->IJKa", delta_IJ, g1o.cv) + g2_oresp.ooov = 2 * einsum("ik,ja->ijka", delta_ij, g1o.ov) + g2_oresp.cocv = 2 * einsum("IK,ja->IjKa", delta_IJ, g1o.ov) + g2_oresp.ocoo = 2 * einsum("ik,Jl->iJkl", delta_ij, g1o.co) + g2_oresp.ccco = 2 * einsum("IK,Jl->IJKl", delta_IJ, g1o.co) + + # scale for contraction with integrals + g2a.oovv *= 0.5 + g2a.ccvv *= 0.5 + g2a.occv *= 2.0 + g2a.vvvv *= 0.25 + + g2_total = evaluate(g2_hf + g2a + g2_oresp) + else: + g2_hf.oooo = 0.25 * (- einsum("li,jk->ijkl", delta_ij, delta_ij) + + einsum("ki,jl->ijkl", delta_ij, delta_ij)) + + g2_oresp.oooo = einsum("ij,kl->kilj", delta_ij, g1o.oo) + g2_oresp.ovov = einsum("ij,ab->iajb", delta_ij, g1o.vv) + g2_oresp.ooov = (- einsum("kj,ia->ijka", delta_ij, g1o.ov) + + einsum("ki,ja->ijka", delta_ij, g1o.ov)) + + # scale for contraction with integrals + g2a.oovv *= 0.5 + g2a.oooo *= 0.25 + g2a.vvvv *= 0.25 + g2_total = evaluate(g2_hf + g2a + g2_oresp) + + provider = hf.gradient_provider + if eri_contraction is None: + if getattr(provider, "backend", None) == "pyscf" \ + and hasattr(provider, "correlated_gradient_direct"): + eri_contraction = "direct" + else: + eri_contraction = "full_ao" + valid_eri_contractions = {"direct", "full_ao"} + if eri_contraction not in valid_eri_contractions: + raise ValueError( + f"Invalid eri_contraction '{eri_contraction}'. Valid values are " + f"{sorted(valid_eri_contractions)}." + ) + if eri_contraction == "direct" and not hasattr( + provider, "correlated_gradient_direct"): + raise NotImplementedError( + "eri_contraction='direct' is currently only available for PySCF." + ) + + with timer.record("transform_ao"): + g1_ao = sum(g1.to_ao_basis(hf)).to_ndarray() + w_ao = sum(w.to_ao_basis(hf)).to_ndarray() + if eri_contraction == "full_ao": + g2_ao_1, g2_ao_2 = g2_total.to_ao_basis() + g2_ao_1, g2_ao_2 = g2_ao_1.to_ndarray(), g2_ao_2.to_ndarray() + + with timer.record("contract_integral_derivatives"): + if eri_contraction == "direct": + grad = provider.correlated_gradient_direct( + g1_ao, w_ao, g2_total, refstate=hf, + shell_chunk_size=eri_shell_chunk_size, + pair_chunk_size=eri_pair_chunk_size, + pair_density_storage=eri_pair_density_storage, + ) + else: + grad = provider.correlated_gradient( + g1_ao, w_ao, g2_ao_1, g2_ao_2 + ) + + ret = GradientResult(excitation_or_mp, grad, g1, g2_total, + timer, g1a=g1a, g2a=g2a) + return ret diff --git a/adcc/gradients/amplitude_response.py b/adcc/gradients/amplitude_response.py new file mode 100644 index 000000000..8620dc636 --- /dev/null +++ b/adcc/gradients/amplitude_response.py @@ -0,0 +1,713 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2021 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +from math import sqrt + +from adcc.functions import einsum, direct_sum, evaluate +from adcc.NParticleOperator import OperatorSymmetry + +from .TwoParticleDensityMatrix import TwoParticleDensityMatrix +from adcc.OneParticleDensity import OneParticleDensity +from adcc.LazyMp import LazyMp +from adcc.Excitation import Excitation +import adcc.block as b + + +def t2bar_oovv_adc2(exci, g1a_adc0): + mp = exci.ground_state + hf = mp.reference_state + u = exci.excitation_vector + df_ia = mp.df(b.ov) + t2bar = 0.5 * ( + hf.oovv + - 2.0 * einsum( + "ijcb,ac->ijab", hf.oovv, g1a_adc0.vv + ).antisymmetrise(2, 3) + + 2.0 * einsum( + "kjab,ik->ijab", hf.oovv, g1a_adc0.oo + ).antisymmetrise(0, 1) + + 4.0 * einsum( + "ia,jkbc,kc->ijab", u.ph, hf.oovv, u.ph + ).antisymmetrise(2, 3).antisymmetrise(0, 1) + ) / ( + 2.0 * direct_sum("ia+jb->ijab", df_ia, df_ia).symmetrise(0, 1) + ) + return t2bar + + +def tbarD_oovv_adc3(exci, g1a_adc0): + mp = exci.ground_state + hf = mp.reference_state + u = exci.excitation_vector + df_ia = mp.df(b.ov) + # Table XI (10.1063/1.5085117) + tbarD = 0.5 * ( + hf.oovv + - 2.0 * einsum( + "ijcb,ac->ijab", hf.oovv, g1a_adc0.vv + ).antisymmetrise(2, 3) + + 2.0 * einsum( + "kjab,ik->ijab", hf.oovv, g1a_adc0.oo + ).antisymmetrise(0, 1) + + 4.0 * einsum( + "ia,jkbc,kc->ijab", u.ph, hf.oovv, u.ph + ).antisymmetrise(2, 3).antisymmetrise(0, 1) + ) / ( + 2.0 * direct_sum("ia+jb->ijab", df_ia, df_ia).symmetrise(0, 1) + ) + return tbarD + + +def tbar_TD_oovv_adc3(exci, g1a_adc0): + mp = exci.ground_state + hf = mp.reference_state + # Table XI (10.1063/1.5085117) + tbarD = tbarD_oovv_adc3(exci, g1a_adc0) + tbarD.evaluate() + ret = 2.0 * 4.0 * ( + + 2.0 * einsum("ikac,jckb->ijab", tbarD, hf.ovov) + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ret += 2.0 * ( + - 1.0 * einsum("ijcd,cdab->ijab", tbarD, hf.vvvv) + - 1.0 * einsum("klab,ijkl->ijab", tbarD, hf.oooo) + ) + return ret + + +def rho_bar_adc3(exci, g1a_adc0): + mp = exci.ground_state + hf = mp.reference_state + u = exci.excitation_vector + df_ia = mp.df(b.ov) + rho_bar = 2.0 * ( + + 1.0 * einsum("ijka,jk->ia", hf.ooov, g1a_adc0.oo) + + 1.0 * einsum("ijkb,jb,ka->ia", hf.ooov, u.ph, u.ph) + - 1.0 * einsum("icab,bc->ia", hf.ovvv, g1a_adc0.vv) + + 1.0 * einsum("jcab,jb,ic->ia", hf.ovvv, u.ph, u.ph) + ) / df_ia + return rho_bar + + +def tbar_rho_oovv_adc3(exci, g1a_adc0): + mp = exci.ground_state + hf = mp.reference_state + # Table XI (10.1063/1.5085117) + rho_bar = rho_bar_adc3(exci, g1a_adc0) + ret = ( + - 1.0 * 2.0 * einsum("jcab,ic->ijab", hf.ovvv, rho_bar).antisymmetrise(0, 1) + - 1.0 * 2.0 * einsum("ijkb,ka->ijab", hf.ooov, rho_bar).antisymmetrise(2, 3) + ) + return ret + + +def t2bar_oovv_adc3(exci, g1a_adc0, g2a_adc1): + mp = exci.ground_state + hf = mp.reference_state + u = exci.excitation_vector + tbar_TD = tbar_TD_oovv_adc3(exci, g1a_adc0) + tbar_TD.evaluate() + tbar_rho = tbar_rho_oovv_adc3(exci, g1a_adc0) + tbar_rho.evaluate() + t2 = mp.t2(b.oovv).evaluate() + df_ia = mp.df(b.ov) + rx = einsum("jb,ijab->ia", u.ph, t2).evaluate() + + ttilde1 = ( + - 1.0 * einsum("klab,ijkm,lm->ijab", t2, hf.oooo, g1a_adc0.oo) + - 1.0 * 2.0 * einsum( + "ka,ijkl,lb->ijab", u.ph, hf.oooo, rx + ).antisymmetrise(2, 3) + + 1.0 * 2.0 * ( + + 0.5 * einsum("ik,jklm,lmab->ijab", g1a_adc0.oo, hf.oooo, t2) + - 2.0 * einsum("jkab,lm,ilkm->ijab", t2, g1a_adc0.oo, hf.oooo) + ).antisymmetrise(0, 1) + - 1.0 * 4.0 * ( + + 0.5 * einsum("ia,kc,jklm,lmbc->ijab", u.ph, u.ph, hf.oooo, t2) + - 2.0 * einsum("ka,likm,jmbc,lc->ijab", u.ph, hf.oooo, t2, u.ph) + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ) + ttilde1.evaluate() + + ttilde2 = 4.0 * ( + - 1.0 * einsum("ijkc,lc,klab->ijab", hf.ooov, u.ph, u.pphh) + + 1.0 * 2.0 * einsum( + "ikab,jlkc,lc->ijab", u.pphh, hf.ooov, u.ph + ).antisymmetrise(0, 1) + - 1.0 * 4.0 * einsum( + "jklb,kc,ilac->ijab", hf.ooov, u.ph, u.pphh + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ) + ttilde2.evaluate() + + ttilde3 = ( + + 1.0 * 2.0 * einsum( + "ijbc,ac->ijab", hf.oovv, g1a_adc0.vv + ).antisymmetrise(2, 3) + - 1.0 * 2.0 * einsum( + "jkab,ik->ijab", hf.oovv, g1a_adc0.oo + ).antisymmetrise(0, 1) + + 1.0 * 4.0 * einsum( + "ia,jkbc,kc->ijab", u.ph, hf.oovv, u.ph + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ) + ttilde3.evaluate() + + # TODO: intermediate x_ka ? + # TODO: intermediate x_lc t_jlbd? + ttilde4 = ( + + 1.0 * 2.0 * ( + - 2.0 * einsum("jkab,ickd,cd->ijab", t2, hf.ovov, g1a_adc0.vv) + - 2.0 * einsum("ic,kd,jcld,klab->ijab", u.ph, u.ph, hf.ovov, t2) + + 1.0 * einsum("ic,jkab,ld,lckd->ijab", u.ph, t2, u.ph, hf.ovov) + + 1.0 * einsum("jkab,kc,ld,lcid->ijab", t2, u.ph, u.ph, hf.ovov) + ).antisymmetrise(0, 1) + + 1.0 * 2.0 * ( + - 2.0 * einsum("ijcb,kalc,kl->ijab", t2, hf.ovov, g1a_adc0.oo) + - 2.0 * einsum("ka,lc,kbld,ijcd->ijab", u.ph, u.ph, hf.ovov, t2) + + 1.0 * einsum("ka,ijbc,ld,kdlc->ijab", u.ph, t2, u.ph, hf.ovov) + + 1.0 * einsum("ijbc,kc,ld,kdla->ijab", t2, u.ph, u.ph, hf.ovov) + ).antisymmetrise(2, 3) + + 1.0 * 4.0 * ( + + 1.0 * einsum("ac,jdkc,ikbd->ijab", g1a_adc0.vv, hf.ovov, t2) + - 1.0 * einsum("jckb,ikad,cd->ijab", hf.ovov, t2, g1a_adc0.vv) + - 1.0 * einsum("ik,kclb,jlac->ijab", g1a_adc0.oo, hf.ovov, t2) + + 1.0 * einsum("jckb,ilac,lk->ijab", hf.ovov, t2, g1a_adc0.oo) + - 1.0 * einsum("ic,jckb,ka->ijab", u.ph, hf.ovov, rx) + - 1.0 * einsum("jb,kc,lcid,klad->ijab", u.ph, u.ph, hf.ovov, t2) + - 1.0 * einsum("ka,jckb,ic->ijab", u.ph, hf.ovov, rx) + - 1.0 * einsum("jb,kc,lakd,ilcd->ijab", u.ph, u.ph, hf.ovov, t2) + + 2.0 * einsum("ka,ickd,ld,jlbc->ijab", u.ph, hf.ovov, u.ph, t2) + + 2.0 * einsum("ic,kcla,kd,jlbd->ijab", u.ph, hf.ovov, u.ph, t2) + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ) + ttilde4.evaluate() + + ttilde5 = - 4.0 * ( + + 1.0 * einsum("kcab,kd,ijcd->ijab", hf.ovvv, u.ph, u.pphh) + + 1.0 * 2.0 * einsum( + "ijbc,kcad,kd->ijab", u.pphh, hf.ovvv, u.ph + ).antisymmetrise(2, 3) + + 1.0 * 4.0 * einsum( + "jcbd,kd,ikac->ijab", hf.ovvv, u.ph, u.pphh + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ) + ttilde5.evaluate() + + ttilde6 = ( + - 1.0 * einsum("ijcd,abde,ce->ijab", t2, hf.vvvv, g1a_adc0.vv) + - 1.0 * 2.0 * einsum( + "ic,abcd,jkde,ke->ijab", u.ph, hf.vvvv, t2, u.ph + ).antisymmetrise(0, 1) + - 1.0 * 2.0 * ( + + 0.5 * einsum("ac,bcde,ijde->ijab", g1a_adc0.vv, hf.vvvv, t2) + - 2.0 * einsum("ijbc,de,adce->ijab", t2, g1a_adc0.vv, hf.vvvv) + ).antisymmetrise(2, 3) + - 1.0 * 4.0 * ( + + 0.5 * einsum("ia,kc,bcde,jkde->ijab", u.ph, u.ph, hf.vvvv, t2) + - 2.0 * einsum("ic,adce,jkeb,kd->ijab", u.ph, hf.vvvv, t2, u.ph) + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ) + ttilde6.evaluate() + + ret = 0.5 * ( + hf.oovv + + ttilde1 + ttilde2 + ttilde3 + ttilde4 + ttilde5 + ttilde6 + + tbar_TD + tbar_rho + ) / (2.0 * direct_sum("ia+jb->ijab", df_ia, df_ia).symmetrise(0, 1)) + return ret + + +def t2bar_oovv_cvs_adc2(exci, g1a_adc0): + mp = exci.ground_state + hf = mp.reference_state + df_ia = mp.df(b.ov) + t2bar = 0.5 * ( + - einsum("ijcb,ac->ijab", hf.oovv, g1a_adc0.vv).antisymmetrise(2, 3) + ) / direct_sum("ia+jb->ijab", df_ia, df_ia).symmetrise(0, 1) + return t2bar + + +def ampl_relaxed_dms_mp2(mp): + hf = mp.reference_state + t2 = mp.t2(b.oovv) + g1a = OneParticleDensity(hf) + g2a = TwoParticleDensityMatrix(hf) + g1a.oo = -0.5 * einsum('ikab,jkab->ij', t2, t2) + g1a.vv = 0.5 * einsum('ijac,ijbc->ab', t2, t2) + g2a.oovv = -1.0 * mp.t2(b.oovv) + return g1a, g2a + + +def ampl_relaxed_dms_adc0(exci): + hf = exci.reference_state + u = exci.excitation_vector + g1a = OneParticleDensity(hf) + g2a = TwoParticleDensityMatrix(hf) + # g2a is not required for the adc0 gradient, + # but expected by amplitude_relaxed_densities + g1a.oo = - 1.0 * einsum("ia,ja->ij", u.ph, u.ph) + g1a.vv = + 1.0 * einsum("ia,ib->ab", u.ph, u.ph) + return g1a, g2a + + +def ampl_relaxed_dms_adc1(exci): + hf = exci.reference_state + u = exci.excitation_vector + g1a = OneParticleDensity(hf) + g2a = TwoParticleDensityMatrix(hf) + g1a.oo = - 1.0 * einsum("ia,ja->ij", u.ph, u.ph) + g1a.vv = + 1.0 * einsum("ia,ib->ab", u.ph, u.ph) + g2a.ovov = - 1.0 * einsum("ja,ib->iajb", u.ph, u.ph) + return g1a, g2a + + +def ampl_relaxed_dms_adc2(exci): + u = exci.excitation_vector + mp = exci.ground_state + g1a_adc1, g2a_adc1 = ampl_relaxed_dms_adc1(exci) + t2 = mp.t2(b.oovv) + t2bar = t2bar_oovv_adc2(exci, g1a_adc1).evaluate() + + g1a = g1a_adc1.copy() + g1a.oo += ( + - 2.0 * einsum('jkab,ikab->ij', u.pphh, u.pphh) + - 2.0 * einsum('jkab,ikab->ij', t2bar, t2).symmetrise((0, 1)) + ) + g1a.vv += ( + + 2.0 * einsum("ijac,ijbc->ab", u.pphh, u.pphh) + + 2.0 * einsum("ijac,ijbc->ab", t2bar, t2).symmetrise((0, 1)) + ) + + g2a = g2a_adc1.copy() + ru_ov = einsum("ijab,jb->ia", t2, u.ph) + g2a.oovv = ( + 0.5 * ( + - 1.0 * t2 + + 2.0 * einsum("ijcb,ca->ijab", t2, g1a_adc1.vv).antisymmetrise(2, 3) + - 2.0 * einsum("kjab,ki->ijab", t2, g1a_adc1.oo).antisymmetrise(0, 1) + - 4.0 * einsum( + "ia,jb->ijab", u.ph, ru_ov + ).antisymmetrise((0, 1)).antisymmetrise((2, 3)) + ) + - 2.0 * t2bar + ) + g2a.ooov = -2.0 * einsum("kb,ijab->ijka", u.ph, u.pphh) + g2a.ovvv = -2.0 * einsum("ja,ijbc->iabc", u.ph, u.pphh) + return g1a, g2a + + +def ampl_relaxed_dms_adc2x(exci): + u = exci.excitation_vector + g1a, g2a = ampl_relaxed_dms_adc2(exci) + + g2a.ovov += -4.0 * einsum("ikbc,jkac->iajb", u.pphh, u.pphh) + g2a.oooo = 2.0 * einsum('ijab,klab->ijkl', u.pphh, u.pphh) + g2a.vvvv = 2.0 * einsum('ijcd,ijab->abcd', u.pphh, u.pphh) + + return g1a, g2a + + +def ampl_relaxed_dms_adc3(exci): + u = exci.excitation_vector + mp = exci.ground_state + hf = mp.reference_state + g1a_adc1, g2a_adc1 = ampl_relaxed_dms_adc1(exci) + t2bar = t2bar_oovv_adc3(exci, g1a_adc1, g2a_adc1).evaluate() + tbarD = tbarD_oovv_adc3(exci, g1a_adc1).evaluate() + rho_bar = rho_bar_adc3(exci, g1a_adc1).evaluate() + t2 = mp.t2(b.oovv) + tD = mp.td2(b.oovv) + rho = mp.mp2_diffdm + + # Table IX (10.1063/1.5085117) + g1a = g1a_adc1.copy() + g1a.oo += ( + - 2.0 * einsum("jkab,ikab->ij", u.pphh, u.pphh) + - 1.0 * 2.0 * ( + + 1.0 * einsum("ikab,jkab->ij", t2, t2bar) + + 1.0 * einsum("ikab,jkab->ij", tD, tbarD) + + 0.5 * einsum("ia,ja->ij", rho_bar, rho.ov) + ).symmetrise(0, 1) + ) + g1a.vv += ( + + 2.0 * einsum("ijac,ijbc->ab", u.pphh, u.pphh) + + 1.0 * 2.0 * ( + + 1.0 * einsum("ijac,ijbc->ab", t2bar, t2) + + 1.0 * einsum("ijac,ijbc->ab", tbarD, tD) + + 0.5 * einsum("ia,ib->ab", rho_bar, rho.ov) + ).symmetrise(0, 1) + ) + + tsq_ovov = einsum("ikac,jkbc->iajb", t2, t2).evaluate() + tsq_vvvv = einsum("klab,klcd->abcd", t2, t2).evaluate() + tsq_oooo = einsum("ijcd,klcd->ijkl", t2, t2).evaluate() + rx = einsum("ijab,jb->ia", t2, u.ph) + + g2a = TwoParticleDensityMatrix(hf) + g2a.oooo = ( + + 2.0 * einsum("ijab,klab->ijkl", u.pphh, u.pphh) + + 0.5 * 2.0 * ( + + 2.0 * einsum("ijab,klab->ijkl", tbarD, t2) + + 0.5 * 2.0 * einsum( + "jm,imab,klab->ijkl", g1a_adc1.oo, t2, t2 + ).antisymmetrise(0, 1) + + 1.0 * 2.0 * einsum( + "kc,ijbc,ma,lmab->ijkl", u.ph, t2, u.ph, t2 + ).antisymmetrise(2, 3) + + 1.0 * 4.0 * ( + + 1.0 * einsum("ik,jl->ijkl", rho.oo, g1a_adc1.oo) + - 1.0 * einsum("lajb,ia,kb->ijkl", tsq_ovov, u.ph, u.ph) + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ).symmetrise([(0, 2), (1, 3)]) + ) + g2a.ooov = ( + - 2.0 * einsum("kb,ijab->ijka", u.ph, u.pphh) + + 1.0 * einsum("la,ijbc,klbc->ijka", u.ph, t2, u.pphh) + + 1.0 * 2.0 * ( + + 2.0 * einsum("ic,jlab,klbc->ijka", u.ph, t2, u.pphh) + - 1.0 * einsum("ja,ilbc,klbc->ijka", u.ph, t2, u.pphh) + + 1.0 * einsum("ja,ik->ijka", rho.ov, g1a_adc1.oo) + + 1.0 * einsum("ia,jb,kb->ijka", u.ph, rho.ov, u.ph) + ).antisymmetrise(0, 1) + - 0.5 * einsum("ijab,kb->ijka", t2, rho_bar) + ) + g2a.oovv = 0.5 * ( + - 1.0 * t2 - 1.0 * tD - 4.0 * t2bar + - 1.0 * 2.0 * ( + + 1.0 * einsum("ijbc,ac->ijab", tD + t2, g1a_adc1.vv) + ).antisymmetrise(2, 3) + + 1.0 * 2.0 * ( + + 1.0 * einsum("jkab,ik->ijab", tD + t2, g1a_adc1.oo) + ).antisymmetrise(0, 1) + - 1.0 * 4.0 * ( + + 1.0 * einsum("ia,jb->ijab", u.ph, + einsum("jkbc,kc->jb", tD, u.ph) + rx) + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ) + g2a.ovov = ( + + 1.0 * g2a_adc1.ovov + - 4.0 * einsum("ikbc,jkac->iajb", u.pphh, u.pphh) + + 1.0 * einsum("ij,ab->iajb", rho.oo, g1a_adc1.vv) + + 1.0 * einsum("ab,ij->iajb", rho.vv, g1a_adc1.oo) + + 0.5 * 2.0 * ( + - 4.0 * einsum("ikbc,jkac->iajb", t2, tbarD) + + 1.0 * einsum("ibjc,ac->iajb", tsq_ovov, g1a_adc1.vv) + - 1.0 * einsum("ibka,jk->iajb", tsq_ovov, g1a_adc1.oo) + + 1.0 * einsum("ka,ikbc,jc->iajb", u.ph, t2, rx) + + 1.0 * einsum("jc,ikbc,ka->iajb", u.ph, t2, rx) + + 1.0 * einsum("ja,cb,ic->iajb", u.ph, rho.vv, u.ph) + - 1.0 * einsum("ib,kj,ka->iajb", u.ph, rho.oo, u.ph) + - 2.0 * einsum("jckb,ic,ka->iajb", tsq_ovov, u.ph, u.ph) + + 0.5 * einsum("acbd,ic,jd->iajb", tsq_vvvv, u.ph, u.ph) + + 0.5 * einsum("ikjl,ka,lb->iajb", tsq_oooo, u.ph, u.ph) + ).symmetrise([(0, 2), (1, 3)]) + ) + g2a.ovvv = ( + - 2.0 * einsum("ja,ijbc->iabc", u.ph, u.pphh) + + 1.0 * einsum("id,jkbc,jkad->iabc", u.ph, t2, u.pphh) + + 1.0 * 2.0 * ( + - 2.0 * einsum("jb,ikcd,jkad->iabc", u.ph, t2, u.pphh) + + 1.0 * einsum("ib,jkcd,jkad->iabc", u.ph, t2, u.pphh) + - 1.0 * einsum("ic,ab->iabc", rho.ov, g1a_adc1.vv) + + 1.0 * einsum("ib,jc,ja->iabc", u.ph, rho.ov, u.ph) + ).antisymmetrise(2, 3) + - 0.5 * einsum("ijbc,ja->iabc", t2, rho_bar) + ) + g2a.vvvv = ( + + 2.0 * einsum("ijcd,ijab->abcd", u.pphh, u.pphh) + + 0.5 * 2.0 * ( + + 2.0 * einsum("ijab,ijcd->abcd", tbarD, t2) + + 0.5 * 2.0 * ( + - 1.0 * einsum("be,aecd->abcd", g1a_adc1.vv, tsq_vvvv) + ).antisymmetrise(0, 1) + + 1.0 * 2.0 * ( + + 1.0 * einsum("ia,ijcd,jkbe,ke->abcd", u.ph, t2, t2, u.ph) + ).antisymmetrise(0, 1) + + 1.0 * 4.0 * ( + + 1.0 * einsum("bd,ac->abcd", rho.vv, g1a_adc1.vv) + - 1.0 * einsum("idjb,ia,jc->abcd", tsq_ovov, u.ph, u.ph) + ).antisymmetrise(0, 1).antisymmetrise(2, 3) + ).symmetrise([(0, 2), (1, 3)]) + ) + return g1a, g2a + + +### CVS ### + +def ampl_relaxed_dms_cvs_adc0(exci): + hf = exci.reference_state + u = exci.excitation_vector + g1a = OneParticleDensity(hf) + g2a = TwoParticleDensityMatrix(hf) + # g2a is not required for cvs-adc0 gradient, + # but expected by amplitude_relaxed_densities + g1a.cc = - 1.0 * einsum("Ia,Ja->IJ", u.ph, u.ph) + g1a.vv = + 1.0 * einsum("Ia,Ib->ab", u.ph, u.ph) + return g1a, g2a + + +def ampl_relaxed_dms_cvs_adc1(exci): + hf = exci.reference_state + u = exci.excitation_vector + g1a = OneParticleDensity(hf, symmetry=OperatorSymmetry.NOSYMMETRY) + g2a = TwoParticleDensityMatrix(hf) + g2a.cvcv = - 1.0 * einsum("Ja,Ib->IaJb", u.ph, u.ph) + g1a.cc = - 1.0 * einsum("Ia,Ja->IJ", u.ph, u.ph) + g1a.vv = + 1.0 * einsum("Ia,Ib->ab", u.ph, u.ph) + + # Prerequisites for the OC block of the + # orbital response Lagrange multipliers: + fc = hf.fock(b.cc).diagonal() + fo = hf.fock(b.oo).diagonal() + fco = direct_sum("-j+I->jI", fc, fo).evaluate() + g1a.co = - 1.0 * einsum('JbKc,ibKc->Ji', g2a.cvcv, hf.ovcv) / fco + from adcc.functions import transpose + g1a.oc = transpose(g1a.co) + return g1a, g2a + + +def ampl_relaxed_dms_cvs_adc2(exci): + hf = exci.reference_state + mp = exci.ground_state + u = exci.excitation_vector + g1a = OneParticleDensity(hf, symmetry=OperatorSymmetry.NOSYMMETRY) + g2a = TwoParticleDensityMatrix(hf) + + # Determine the t-amplitudes and multipliers: + t2oovv = mp.t2(b.oovv) + t2ccvv = mp.t2(b.ccvv) + t2ocvv = mp.t2(b.ocvv) + g1a_cvs0, g2a_cvs0 = ampl_relaxed_dms_cvs_adc0(exci) + t2bar = t2bar_oovv_cvs_adc2(exci, g1a_cvs0).evaluate() + + g1a.cc = ( + - 1.0 * einsum("Ia,Ja->IJ", u.ph, u.ph) + - 1.0 * einsum("kJba,kIba->IJ", u.pphh, u.pphh) + - 0.5 * einsum('IKab,JKab->IJ', t2ccvv, t2ccvv) + - 0.5 * einsum('kIab,kJab->IJ', t2ocvv, t2ocvv) + ) + + g1a.oo = ( + - 1.0 * einsum("jKba,iKba->ij", u.pphh, u.pphh) + - 2.0 * einsum("ikab,jkab->ij", t2bar, t2oovv).symmetrise((0, 1)) + - 0.5 * einsum('iKab,jKab->ij', t2ocvv, t2ocvv) + - 0.5 * einsum('ikab,jkab->ij', t2oovv, t2oovv) + ) + + # Pre-requisites for the OC block of the + # orbital response Lagrange multipliers: + fc = hf.fock(b.cc).diagonal() + fo = hf.fock(b.oo).diagonal() + fco = direct_sum("-j+I->jI", fc, fo).evaluate() + + g1a.vv = ( + + 1.0 * einsum("Ia,Ib->ab", u.ph, u.ph) + + 2.0 * einsum('jIcb,jIca->ab', u.pphh, u.pphh) + + 2.0 * einsum('ijac,ijbc->ab', t2bar, t2oovv).symmetrise((0, 1)) + + 0.5 * einsum('IJac,IJbc->ab', t2ccvv, t2ccvv) + + 0.5 * einsum('ijac,ijbc->ab', t2oovv, t2oovv) + + 1.0 * einsum('iJac,iJbc->ab', t2ocvv, t2ocvv) + ) + + g2a.cvcv = ( + - einsum("Ja,Ib->IaJb", u.ph, u.ph) + ) + + # The factor 1/sqrt(2) is needed because of the scaling used in adcc + # for the ph-pphh blocks. + g2a.occv = (1 / sqrt(2)) * ( + 2.0 * einsum('Ib,kJba->kJIa', u.ph, u.pphh) + ) + + g2a.oovv = ( + + 1.0 * einsum('ijcb,ca->ijab', t2oovv, g1a_cvs0.vv).antisymmetrise((2, 3)) + - 1.0 * t2oovv + - 2.0 * t2bar + ) + + # The factor 2/sqrt(2) is necessary because of the way + # that the ph-pphh is scaled. + g2a.ovvv = (2 / sqrt(2)) * ( + einsum('Ja,iJcb->iabc', u.ph, u.pphh) + ) + + g2a.ccvv = - 1.0 * t2ccvv + g2a.ocvv = - 1.0 * t2ocvv + + # This is the OC block of the orbital response + # Lagrange multipliers (lambda): + g1a.co = ( + - 1.0 * einsum('JbKc,ibKc->Ji', g2a.cvcv, hf.ovcv) + - 0.5 * einsum('JKab,iKab->Ji', g2a.ccvv, hf.ocvv) + + 1.0 * einsum('kJLa,ikLa->Ji', g2a.occv, hf.oocv) + + 0.5 * einsum('kJab,ikab->Ji', g2a.ocvv, hf.oovv) + - 1.0 * einsum('kLJa,kLia->Ji', g2a.occv, hf.ocov) + + 1.0 * einsum('iKLa,JKLa->Ji', g2a.occv, hf.cccv) + + 0.5 * einsum('iKab,JKab->Ji', g2a.ocvv, hf.ccvv) + - 0.5 * einsum('ikab,kJab->Ji', g2a.oovv, hf.ocvv) + + 0.5 * einsum('iabc,Jabc->Ji', g2a.ovvv, hf.cvvv) + ) / fco + from adcc.functions import transpose + g1a.oc = transpose(g1a.co) + + return g1a, g2a + + +def ampl_relaxed_dms_cvs_adc2x(exci): + hf = exci.reference_state + mp = exci.ground_state + u = exci.excitation_vector + g1a = OneParticleDensity(hf, symmetry=OperatorSymmetry.NOSYMMETRY) + g2a = TwoParticleDensityMatrix(hf) + + # Determine the t-amplitudes and multipliers: + t2oovv = mp.t2(b.oovv) + t2ccvv = mp.t2(b.ccvv) + t2ocvv = mp.t2(b.ocvv) + g1a_cvs0, _ = ampl_relaxed_dms_cvs_adc0(exci) + t2bar = t2bar_oovv_cvs_adc2(exci, g1a_cvs0).evaluate() + + g1a.cc = ( + - 1.0 * einsum("Ia,Ja->IJ", u.ph, u.ph) + - 1.0 * einsum("kJba,kIba->IJ", u.pphh, u.pphh) + - 0.5 * einsum('IKab,JKab->IJ', t2ccvv, t2ccvv) + - 0.5 * einsum('kIab,kJab->IJ', t2ocvv, t2ocvv) + ) + + g1a.oo = ( + - 1.0 * einsum("jKba,iKba->ij", u.pphh, u.pphh) + - 2.0 * einsum("ikab,jkab->ij", t2bar, t2oovv).symmetrise((0, 1)) + - 0.5 * einsum('iKab,jKab->ij', t2ocvv, t2ocvv) + - 0.5 * einsum('ikab,jkab->ij', t2oovv, t2oovv) + ) + + # Pre-requisites for the OC block of the + # orbital response Lagrange multipliers: + fc = hf.fock(b.cc).diagonal() + fo = hf.fock(b.oo).diagonal() + fco = direct_sum("-j+I->jI", fc, fo).evaluate() + + g1a.vv = ( + + 1.0 * einsum("Ia,Ib->ab", u.ph, u.ph) + + 2.0 * einsum('jIcb,jIca->ab', u.pphh, u.pphh) + + 2.0 * einsum('ijac,ijbc->ab', t2bar, t2oovv).symmetrise((0, 1)) + + 0.5 * einsum('IJac,IJbc->ab', t2ccvv, t2ccvv) + + 0.5 * einsum('ijac,ijbc->ab', t2oovv, t2oovv) + + 1.0 * einsum('iJac,iJbc->ab', t2ocvv, t2ocvv) + ) + + g2a.cvcv = ( + - 1.0 * einsum("Ja,Ib->IaJb", u.ph, u.ph) + - 1.0 * einsum('kIbc,kJac->IaJb', u.pphh, u.pphh) + + 1.0 * einsum('kIcb,kJac->IaJb', u.pphh, u.pphh) + ) + + # The factor 1/sqrt(2) is needed because of the scaling used in adcc + # for the ph-pphh blocks. + g2a.occv = (1 / sqrt(2)) * ( + 2.0 * einsum('Ib,kJba->kJIa', u.ph, u.pphh) + ) + + g2a.oovv = ( + + 1.0 * einsum('ijcb,ca->ijab', t2oovv, g1a_cvs0.vv).antisymmetrise((2, 3)) + - 1.0 * t2oovv + - 2.0 * t2bar + ) + + # The factor 2/sqrt(2) is necessary because of + # the way that the ph-pphh is scaled + g2a.ovvv = (2 / sqrt(2)) * ( + einsum('Ja,iJcb->iabc', u.ph, u.pphh) + ) + + g2a.ovov = 1.0 * ( + - einsum("iKbc,jKac->iajb", u.pphh, u.pphh) + + einsum("iKcb,jKac->iajb", u.pphh, u.pphh) + ) + + g2a.ccvv = - 1.0 * t2ccvv + g2a.ocvv = - 1.0 * t2ocvv + g2a.ococ = 1.0 * einsum("iJab,kLab->iJkL", u.pphh, u.pphh) + g2a.vvvv = 2.0 * einsum("iJcd,iJab->abcd", u.pphh, u.pphh) + + g1a.co = ( + - 1.0 * einsum('JbKc,ibKc->Ji', g2a.cvcv, hf.ovcv) + - 0.5 * einsum('JKab,iKab->Ji', g2a.ccvv, hf.ocvv) + + 1.0 * einsum('kJLa,ikLa->Ji', g2a.occv, hf.oocv) + + 0.5 * einsum('kJab,ikab->Ji', g2a.ocvv, hf.oovv) + - 1.0 * einsum('kLJa,kLia->Ji', g2a.occv, hf.ocov) + + 1.0 * einsum('iKLa,JKLa->Ji', g2a.occv, hf.cccv) + + 0.5 * einsum('iKab,JKab->Ji', g2a.ocvv, hf.ccvv) + - 0.5 * einsum('ikab,kJab->Ji', g2a.oovv, hf.ocvv) + + 0.5 * einsum('iabc,Jabc->Ji', g2a.ovvv, hf.cvvv) + + 1.0 * einsum('kJmL,ikmL->Ji', g2a.ococ, hf.oooc) + - 1.0 * einsum('iKlM,JKMl->Ji', g2a.ococ, hf.ccco) + + 1.0 * einsum('iakb,kbJa->Ji', g2a.ovov, hf.ovcv) + ) / fco + from adcc.functions import transpose + g1a.oc = transpose(g1a.co) + + return g1a, g2a + + +DISPATCH = { + "mp2": ampl_relaxed_dms_mp2, + "adc0": ampl_relaxed_dms_adc0, + "adc1": ampl_relaxed_dms_adc1, + "adc2": ampl_relaxed_dms_adc2, + "adc2x": ampl_relaxed_dms_adc2x, + "adc3": ampl_relaxed_dms_adc3, + "cvs-adc0": ampl_relaxed_dms_cvs_adc0, + "cvs-adc1": ampl_relaxed_dms_cvs_adc1, + "cvs-adc2": ampl_relaxed_dms_cvs_adc2, + "cvs-adc2x": ampl_relaxed_dms_cvs_adc2x, +} + + +def amplitude_relaxed_densities(excitation_or_mp): + """Computation of amplitude-relaxed one- and two-particle density matrices + + Parameters + ---------- + excitation_or_mp : LazyMp, Excitation + Data for which the densities are requested, either LazyMp for ground + state densities or Excitation for excited state densities + + Returns + ------- + (OneParticleOperator, TwoParticleDensityMatrix) + Tuple of amplitude-relaxed one- and two-particle density matrices + + Raises + ------ + NotImplementedError + if density matrices are not implemented for a given method + """ + if isinstance(excitation_or_mp, LazyMp): + method_name = "mp2" + elif isinstance(excitation_or_mp, Excitation): + method_name = excitation_or_mp.method.name + if method_name not in DISPATCH: + raise NotImplementedError("Amplitude response is not " + f"implemented for {method_name}.") + g1a, g2a = DISPATCH[method_name](excitation_or_mp) + return evaluate(g1a), evaluate(g2a) diff --git a/adcc/gradients/orbital_response.py b/adcc/gradients/orbital_response.py new file mode 100644 index 000000000..97eec2092 --- /dev/null +++ b/adcc/gradients/orbital_response.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2021 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +import adcc.block as b +from adcc.OneParticleDensity import OneParticleDensity +from adcc.NParticleOperator import OperatorSymmetry +from adcc.functions import direct_sum, einsum, evaluate +from adcc.AmplitudeVector import AmplitudeVector +from adcc.solver.conjugate_gradient import conjugate_gradient, default_print + + +def orbital_response_rhs(hf, g1a, g2a): + """ + Build the right-hand side for solving the orbital + response given amplitude-relaxed density matrices (method-specific) + """ + # TODO: only add non-zero blocks to equations! + + if hf.has_core_occupied_space: + ret_ov = -1.0 * ( + + 2.0 * einsum("JiKa,JK->ia", hf.cocv, g1a.cc) + + 2.0 * einsum("icab,bc->ia", hf.ovvv, g1a.vv) + + 2.0 * einsum('kija,jk->ia', hf.ooov, g1a.oo) + + 2.0 * einsum("kiJa,Jk->ia", hf.oocv, g1a.co) + - 2.0 * einsum("iJka,Jk->ia", hf.ocov, g1a.co) + + 2.0 * einsum("iJKb,JaKb->ia", hf.occv, g2a.cvcv) # 2PDMs start + + 2.0 * einsum('lKJa,lKiJ->ia', g2a.occv, hf.ococ) + + 1.0 * einsum('jabc,ijbc->ia', g2a.ovvv, hf.oovv) + - 1.0 * einsum('JKab,JKib->ia', g2a.ccvv, hf.ccov) + - 2.0 * einsum('jKab,jKib->ia', g2a.ocvv, hf.ocov) + - 1.0 * einsum('jkab,jkib->ia', g2a.oovv, hf.ooov) + - 2.0 * einsum('jcab,ibjc->ia', g2a.ovvv, hf.ovov) + - 2.0 * einsum('iJKb,JaKb->ia', g2a.occv, hf.cvcv) + - 1.0 * einsum('iJbc,Jabc->ia', g2a.ocvv, hf.cvvv) + - 1.0 * einsum('ijbc,jabc->ia', g2a.oovv, hf.ovvv) + + 1.0 * einsum('ibcd,abcd->ia', g2a.ovvv, hf.vvvv) + - 1.0 * einsum('abcd,ibcd->ia', g2a.vvvv, hf.ovvv) # cvs-adc2x + + 2.0 * einsum('jakb,ijkb->ia', g2a.ovov, hf.ooov) # cvs-adc2x + + 2.0 * einsum('ibjc,jcab->ia', g2a.ovov, hf.ovvv) # cvs-adc2x + - 2.0 * einsum('iJkL,kLJa->ia', g2a.ococ, hf.occv) # cvs-adc2x + ) + + ret_cv = -1.0 * ( + + 2.0 * einsum("JIKa,JK->Ia", hf.cccv, g1a.cc) + + 2.0 * einsum("Icab,bc->Ia", hf.cvvv, g1a.vv) + + 2.0 * einsum('kIja,jk->Ia', hf.ocov, g1a.oo) + + 2.0 * einsum("kIJa,Jk->Ia", hf.occv, g1a.co) + + 2.0 * einsum("JIka,Jk->Ia", hf.ccov, g1a.co) + + 2.0 * einsum("IJKb,JaKb->Ia", hf.cccv, g2a.cvcv) # 2PDMs start + + 2.0 * einsum("Jcab,IbJc->Ia", hf.cvvv, g2a.cvcv) + + 2.0 * einsum('lKJa,lKIJ->Ia', g2a.occv, hf.occc) + - 1.0 * einsum('jabc,jIbc->Ia', g2a.ovvv, hf.ocvv) + - 1.0 * einsum('JKab,JKIb->Ia', g2a.ccvv, hf.cccv) + - 2.0 * einsum('jKab,jKIb->Ia', g2a.ocvv, hf.occv) + - 1.0 * einsum('jkab,jkIb->Ia', g2a.oovv, hf.oocv) + - 2.0 * einsum('jcab,jcIb->Ia', g2a.ovvv, hf.ovcv) + - 1.0 * einsum('IJbc,Jabc->Ia', g2a.ccvv, hf.cvvv) + + 2.0 * einsum('jIKb,jaKb->Ia', g2a.occv, hf.ovcv) + + 1.0 * einsum('jIbc,jabc->Ia', g2a.ocvv, hf.ovvv) + + 2.0 * einsum('kJIb,kJab->Ia', g2a.occv, hf.ocvv) + - 1.0 * einsum('abcd,Ibcd->Ia', g2a.vvvv, hf.cvvv) # cvs-adc2x + - 2.0 * einsum('jakb,jIkb->Ia', g2a.ovov, hf.ocov) # cvs-adc2x + + 2.0 * einsum('jIlK,lKja->Ia', g2a.ococ, hf.ocov) # cvs-adc2x + ) + ret = AmplitudeVector(cv=ret_cv, ov=ret_ov) + else: + # equal to the ov block of the energy-weighted density + # matrix when lambda_ov multipliers are zero + w_ov = 0.5 * ( + + 1.0 * einsum("ijkl,klja->ia", hf.oooo, g2a.ooov) # not in cvs-adc + - 1.0 * einsum("ibcd,abcd->ia", hf.ovvv, g2a.vvvv) + - 1.0 * einsum("jkib,jkab->ia", hf.ooov, g2a.oovv) + + 2.0 * einsum("ijkb,jakb->ia", hf.ooov, g2a.ovov) + + 1.0 * einsum("ijbc,jabc->ia", hf.oovv, g2a.ovvv) + - 2.0 * einsum("ibjc,jcab->ia", hf.ovov, g2a.ovvv) + ) + w_ov = w_ov.evaluate() + + ret_ov = -1.0 * ( + 2.0 * w_ov + - 1.0 * einsum("klja,ijkl->ia", hf.ooov, g2a.oooo) + + 1.0 * einsum("abcd,ibcd->ia", hf.vvvv, g2a.ovvv) + - 2.0 * einsum("jakb,ijkb->ia", hf.ovov, g2a.ooov) # not in cvs-adc + + 1.0 * einsum("jkab,jkib->ia", hf.oovv, g2a.ooov) # not in cvs-adc + + 2.0 * einsum("jcab,ibjc->ia", hf.ovvv, g2a.ovov) + - 1.0 * einsum("jabc,ijbc->ia", hf.ovvv, g2a.oovv) + + 2.0 * einsum("jika,jk->ia", hf.ooov, g1a.oo) + + 2.0 * einsum("icab,bc->ia", hf.ovvv, g1a.vv) + ) + ret = AmplitudeVector(ov=ret_ov) + + return ret + + +def energy_weighted_density_matrix(hf, g1o, g2a): + if hf.has_core_occupied_space: + w = OneParticleDensity(hf, symmetry=OperatorSymmetry.NOSYMMETRY) + w.cc = - 0.5 * ( + + einsum("JKab,IKab->IJ", g2a.ccvv, hf.ccvv) + + einsum("kJab,kIab->IJ", g2a.ocvv, hf.ocvv) + ) + w.cc += ( + - hf.fcc + - einsum("IK,JK->IJ", g1o.cc, hf.fcc) + - einsum("ab,IaJb->IJ", g1o.vv, hf.cvcv) + - einsum("KL,IKJL->IJ", g1o.cc, hf.cccc) + - einsum("kl,kIlJ->IJ", g1o.oo, hf.ococ) + + einsum("ka,kJIa->IJ", g1o.ov, hf.occv) + + einsum("ka,kIJa->IJ", g1o.ov, hf.occv) + + einsum("Ka,KJIa->IJ", g1o.cv, hf.cccv) + + einsum("Ka,KIJa->IJ", g1o.cv, hf.cccv) + - einsum("Lk,kILJ->IJ", g1o.co, hf.occc) + - einsum("Lk,kJLI->IJ", g1o.co, hf.occc) + - einsum("JbKc,IbKc->IJ", g2a.cvcv, hf.cvcv) + - einsum("kJLa,kILa->IJ", g2a.occv, hf.occv) + - einsum("kLJa,kLIa->IJ", g2a.occv, hf.occv) + - einsum("kJmL,kImL->IJ", g2a.ococ, hf.ococ) # cvs-adc2x + ) + w.oo = - 0.5 * ( + + einsum("jKab,iKab->ij", g2a.ocvv, hf.ocvv) + + einsum("jkab,ikab->ij", g2a.oovv, hf.oovv) + + einsum("jabc,iabc->ij", g2a.ovvv, hf.ovvv) + ) + w.oo += ( + - hf.foo + - einsum("ij,ii->ij", g1o.oo, hf.foo) + - einsum("KL,iKjL->ij", g1o.cc, hf.ococ) + - einsum("ab,iajb->ij", g1o.vv, hf.ovov) + - einsum("kl,ikjl->ij", g1o.oo, hf.oooo) + - einsum("ka,ikja->ij", g1o.ov, hf.ooov) + - einsum("ka,jkia->ij", g1o.ov, hf.ooov) + - einsum("Ka,iKja->ij", g1o.cv, hf.ocov) + - einsum("Ka,jKia->ij", g1o.cv, hf.ocov) + + einsum("Lk,kijL->ij", g1o.co, hf.oooc) + + einsum("Lk,kjiL->ij", g1o.co, hf.oooc) + - einsum("jKLa,iKLa->ij", g2a.occv, hf.occv) + - einsum("jKlM,iKlM->ij", g2a.ococ, hf.ococ) # cvs-adc2x + - einsum("jakb,iakb->ij", g2a.ovov, hf.ovov) # cvs-adc2x + ) + w.vv = - 0.5 * ( + + einsum("ibcd,iacd->ab", g2a.ovvv, hf.ovvv) + + einsum("IJbc,IJac->ab", g2a.ccvv, hf.ccvv) + + einsum("ijbc,ijac->ab", g2a.oovv, hf.oovv) + + einsum("bcde,acde->ab", g2a.vvvv, hf.vvvv) # cvs-adc2x + ) + w.vv += ( + - einsum("ac,cb->ab", g1o.vv, hf.fvv) + - einsum('JbKc,JaKc->ab', g2a.cvcv, hf.cvcv) + - einsum("jKIb,jKIa->ab", g2a.occv, hf.occv) + - einsum("idbc,idac->ab", g2a.ovvv, hf.ovvv) + - einsum("iJbc,iJac->ab", g2a.ocvv, hf.ocvv) + - einsum("ibjc,iajc->ab", g2a.ovov, hf.ovov) # cvs-adc2x + ) + w.co = 0.5 * ( + - einsum("JKab,iKab->Ji", g2a.ccvv, hf.ocvv) + + einsum("kJab,ikab->Ji", g2a.ocvv, hf.oovv) + ) + w.co += ( + + einsum("KL,iKLJ->Ji", g1o.cc, hf.occc) + - einsum("ab,iaJb->Ji", g1o.vv, hf.ovcv) + - einsum("kl,kilJ->Ji", g1o.oo, hf.oooc) + - einsum("Ka,iKJa->Ji", g1o.cv, hf.occv) + - einsum("Ka,JKia->Ji", g1o.cv, hf.ccov) + - einsum("ka,ikJa->Ji", g1o.ov, hf.oocv) + - einsum("ka,Jkia->Ji", g1o.ov, hf.coov) + - einsum("Lk,kiLJ->Ji", g1o.co, hf.oocc) + + einsum("Lk,iLkJ->Ji", g1o.co, hf.ococ) + - einsum("Ik,jk->Ij", g1o.co, hf.foo) + - einsum("JbKc,ibKc->Ji", g2a.cvcv, hf.ovcv) + + einsum("kJLa,ikLa->Ji", g2a.occv, hf.oocv) + - einsum("kLJa,kLia->Ji", g2a.occv, hf.ocov) + + einsum("kJmL,ikmL->Ji", g2a.ococ, hf.oooc) # cvs-adc2x + ) + w.ov = 0.5 * ( + + einsum("jabc,ijbc->ia", g2a.ovvv, hf.oovv) + - einsum("JKab,JKib->ia", g2a.ccvv, hf.ccov) + - einsum("jkab,jkib->ia", g2a.oovv, hf.ooov) + - einsum("abcd,ibcd->ia", g2a.vvvv, hf.ovvv) # cvs-adc2x + ) + w.ov += ( + - einsum("ij,ja->ia", hf.foo, g1o.ov) + + einsum("JaKc,iJKc->ia", g2a.cvcv, hf.occv) + + einsum("kLJa,iJkL->ia", g2a.occv, hf.ococ) + - einsum("jKab,jKib->ia", g2a.ocvv, hf.ocov) + - einsum("jcab,ibjc->ia", g2a.ovvv, hf.ovov) + + einsum("jakb,ijkb->ia", g2a.ovov, hf.ooov) # cvs-adc2x + ) + w.cv = - 0.5 * ( + + einsum("jabc,jIbc->Ia", g2a.ovvv, hf.ocvv) + + einsum("JKab,JKIb->Ia", g2a.ccvv, hf.cccv) + + einsum("jkab,jkIb->Ia", g2a.oovv, hf.oocv) + + einsum("abcd,Ibcd->Ia", g2a.vvvv, hf.cvvv) # cvs-adc2x + ) + w.cv += ( + - einsum("IJ,Ja->Ia", hf.fcc, g1o.cv) + + einsum("JaKc,IJKc->Ia", g2a.cvcv, hf.cccv) + + einsum("lKJa,lKIJ->Ia", g2a.occv, hf.occc) + - einsum("kJab,kJIb->Ia", g2a.ocvv, hf.occv) + - einsum("jcab,jcIb->Ia", g2a.ovvv, hf.ovcv) + - einsum("jakb,jIkb->Ia", g2a.ovov, hf.ocov) # cvs-adc2x + ) + # For NOSYMMETRY operators, explicitly set transpose blocks + # to maintain the same AO transform behaviour as the old HERMITIAN code + from adcc.functions import transpose + w.vo = transpose(w.ov) + w.vc = transpose(w.cv) + w.oc = transpose(w.co) + else: + gi_oo = -0.5 * ( + + 1.0 * einsum("jklm,iklm->ij", hf.oooo, g2a.oooo) + + 1.0 * einsum("jabc,iabc->ij", hf.ovvv, g2a.ovvv) + + 1.0 * einsum("klja,klia->ij", hf.ooov, g2a.ooov) + + 2.0 * einsum("jkla,ikla->ij", hf.ooov, g2a.ooov) + + 1.0 * einsum("jkab,ikab->ij", hf.oovv, g2a.oovv) + + 2.0 * einsum("jakb,iakb->ij", hf.ovov, g2a.ovov) + ) + gi_vv = -0.5 * ( + + 1.0 * einsum("kjib,kjia->ab", hf.ooov, g2a.ooov) + + einsum("bcde,acde->ab", hf.vvvv, g2a.vvvv) + + 1.0 * einsum("ijcb,ijca->ab", hf.oovv, g2a.oovv) + + 2.0 * einsum("jcib,jcia->ab", hf.ovov, g2a.ovov) + + 1.0 * einsum("ibcd,iacd->ab", hf.ovvv, g2a.ovvv) + + 2.0 * einsum("idcb,idca->ab", hf.ovvv, g2a.ovvv) + ) + gi_oo = gi_oo.evaluate() + gi_vv = gi_vv.evaluate() + w = OneParticleDensity(hf) + w.ov = 0.5 * ( + - 2.0 * einsum("ij,ja->ia", hf.foo, g1o.ov) + + 1.0 * einsum("ijkl,klja->ia", hf.oooo, g2a.ooov) + - 1.0 * einsum("ibcd,abcd->ia", hf.ovvv, g2a.vvvv) + - 1.0 * einsum("jkib,jkab->ia", hf.ooov, g2a.oovv) + + 2.0 * einsum("ijkb,jakb->ia", hf.ooov, g2a.ovov) + + 1.0 * einsum("ijbc,jabc->ia", hf.oovv, g2a.ovvv) + - 2.0 * einsum("ibjc,jcab->ia", hf.ovov, g2a.ovvv) + ) + w.oo = ( + + gi_oo - hf.foo + - einsum("ik,jk->ij", g1o.oo, hf.foo) + - einsum("ikjl,kl->ij", hf.oooo, g1o.oo) + - einsum("ikja,ka->ij", hf.ooov, g1o.ov) + - einsum("jkia,ka->ij", hf.ooov, g1o.ov) + - einsum("jaib,ab->ij", hf.ovov, g1o.vv) + ) + w.vv = gi_vv - einsum("ac,cb->ab", g1o.vv, hf.fvv) + return evaluate(w) + + +class OrbitalResponseMatrix: + def __init__(self, hf): + self.hf = hf + + @property + def shape(self): + no1 = self.hf.mospaces.n_orbs(b.o) + if self.hf.has_core_occupied_space: + no1 += self.hf.mospaces.n_orbs(b.c) + nv1 = self.hf.mospaces.n_orbs(b.v) + size = no1 * nv1 + return (size, size) + + def __matmul__(self, lam): + ret_ov = ( + + einsum("ab,ib->ia", self.hf.fvv, lam.ov) + - einsum("ij,ja->ia", self.hf.foo, lam.ov) + + einsum("ijab,jb->ia", self.hf.oovv, lam.ov) + - einsum("ibja,jb->ia", self.hf.ovov, lam.ov) + ) + if self.hf.has_core_occupied_space: + ret_ov += ( + + einsum("iJab,Jb->ia", self.hf.ocvv, lam.cv) + - einsum("ibJa,Jb->ia", self.hf.ovcv, lam.cv) + ) + ret_cv = ( + + einsum("ab,Ib->Ia", self.hf.fvv, lam.cv) + - einsum("IJ,Ja->Ia", self.hf.fcc, lam.cv) + + einsum("IJab,Jb->Ia", self.hf.ccvv, lam.cv) + - einsum("IbJa,Jb->Ia", self.hf.cvcv, lam.cv) + + einsum("Ijab,jb->Ia", self.hf.covv, lam.ov) + - einsum("Ibja,jb->Ia", self.hf.cvov, lam.ov) + ) + ret = AmplitudeVector(cv=ret_cv, ov=ret_ov) + else: + ret = AmplitudeVector(ov=ret_ov) + return evaluate(ret) + + +class OrbitalResponsePinv: + def __init__(self, hf): + self.hf = hf + # Terms common to adc and cvs-adc + fo = hf.fock(b.oo).diagonal() + fv = hf.fock(b.vv).diagonal() + fov = direct_sum("-i+a->ia", fo, fv) + + if hf.has_core_occupied_space: + fc = hf.fock(b.cc).diagonal() + fcv = direct_sum("-I+a->Ia", fc, fv) + self.df = AmplitudeVector(cv=fcv, ov=fov) + else: + self.df = AmplitudeVector(ov=fov) + self.df.evaluate() + + @property + def shape(self): + no1 = self.hf.mospaces.n_orbs(b.o) + if self.hf.has_core_occupied_space: + no1 += self.hf.mospaces.n_orbs(b.c) + nv1 = self.hf.mospaces.n_orbs(b.v) + size = no1 * nv1 + return (size, size) + + def __matmul__(self, invec): + return invec / self.df + + +def orbital_response(hf, rhs, conv_tol=1e-9, **solver_kwargs): + """ + Solves the orbital response equations + for a given reference state and right-hand side + """ + A = OrbitalResponseMatrix(hf) + Pinv = OrbitalResponsePinv(hf) + x0 = (Pinv @ rhs).evaluate() + lam = conjugate_gradient(A, rhs=rhs, x0=x0, Pinv=Pinv, + conv_tol=conv_tol, + explicit_symmetrisation=None, + callback=default_print, **solver_kwargs) + return lam.solution diff --git a/adcc/gradients/scanner.py b/adcc/gradients/scanner.py new file mode 100644 index 000000000..451853d7f --- /dev/null +++ b/adcc/gradients/scanner.py @@ -0,0 +1,405 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2026 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +"""Geometry-optimisation scanner for adcc nuclear gradients. + +The scanner owns the per-geometry PySCF -> adcc -> gradient loop. Users pass a +configured PySCF SCF object, whose settings define how SCF is performed at every +geometry. adcc-specific keyword arguments are forwarded unchanged to +:func:`adcc.run_adc` for excited-state targets. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Optional +import warnings + +import numpy as np + +from adcc.Excitation import Excitation +from adcc.LazyMp import LazyMp + + +@dataclass(frozen=True) +class GroundStateTarget: + """Ground-state MP target for a nuclear-gradient scanner.""" + + level: int = 2 + + +@dataclass(frozen=True) +class ExcitedStateTarget: + """Excited-state ADC target for a nuclear-gradient scanner.""" + + method: str + state_index: int = 0 + run_adc_kwargs: dict[str, Any] = field(default_factory=dict) + + def kwargs(self) -> dict[str, Any]: + """Return keyword arguments forwarded to :func:`adcc.run_adc`.""" + ret = dict(self.run_adc_kwargs) + ret["method"] = self.method + return ret + + +@dataclass +class TrackingResult: + """Diagnostic information for the most recent root-tracking decision.""" + + index: int + scores: np.ndarray + best_score: float = 0.0 + gap: float = float("inf") + previous_index: Optional[int] = None + switched: bool = False + unavailable_channels: tuple[str, ...] = () + + +@dataclass +class _TrackingDescriptor: + mol: Any + transition_dm: Optional[np.ndarray] + state_diffdm: Optional[np.ndarray] + + +class NuclearGradientScanner: + """Callable scanner returning adcc energies and nuclear gradients. + + ``scfres`` is a configured PySCF SCF object. Its molecule and SCF settings + are used as the template for every scanner call; only the nuclear geometry + changes. Coordinates are Cartesian in Bohr, energies are Hartree and + gradients are Hartree/Bohr. + """ + + def __init__(self, scfres, *, + target: GroundStateTarget | ExcitedStateTarget | str | None = None, + method: Optional[str] = None, state_index: int = 0, + mp_level: int = 2, follow: str = "overlap", + tracking_min_score: float = 0.0, tracking_min_gap: float = 0.0, + gradient_kwargs: Optional[dict[str, Any]] = None, + **run_adc_kwargs): + self._pyscf = _import_pyscf() + if not _is_pyscf_scf(self._pyscf, scfres): + raise TypeError( + "NuclearGradientScanner expects a PySCF SCF object, e.g. " + "scf.RHF(mol), as its first argument." + ) + if not hasattr(scfres, "as_scanner"): + raise TypeError( + "The provided PySCF SCF object has no as_scanner method." + ) + + self.scf_template = scfres + self.scf_scanner = scfres.as_scanner() + self.base_mol = scfres.mol.copy() + self.atom_symbols = [self.base_mol.atom_symbol(i) + for i in range(self.base_mol.natm)] + self.initial_coords = np.asarray( + self.base_mol.atom_coords(unit="Bohr"), dtype=float + ) + + self.follow = follow + if self.follow not in ("overlap", "index"): + raise ValueError("follow needs to be 'overlap' or 'index'.") + self.tracking_min_score = tracking_min_score + self.tracking_min_gap = tracking_min_gap + self.gradient_kwargs = dict(gradient_kwargs or {}) + self.target = self._normalise_target( + target, method, state_index, mp_level, run_adc_kwargs + ) + + self.previous_scf = None + self.previous_descriptor: Optional[_TrackingDescriptor] = None + self.previous_index: Optional[int] = None + self.last_scf = None + self.last_states = None + self.last_excitation = None + self.last_gradient = None + self.last_tracking: Optional[TrackingResult] = None + + @property + def natoms(self) -> int: + return len(self.atom_symbols) + + def __call__(self, coords): + """Return ``(energy, gradient)`` for Cartesian coordinates in Bohr.""" + coords = self._coords_array(coords) + scfres = self._run_scf(coords) + target = self._build_target(scfres) + + from adcc.gradients import nuclear_gradient + grad = nuclear_gradient(target, **self.gradient_kwargs) + if isinstance(target, LazyMp): + energy = target.energy(self.target.level) + else: + energy = target.total_energy + + self.last_scf = scfres + self.last_gradient = grad + self.previous_scf = scfres + if isinstance(target, Excitation): + self.previous_descriptor = self._descriptor(target, scfres.mol) + if self.last_tracking is not None: + self.previous_index = self.last_tracking.index + else: + self.previous_index = self.target.state_index + return float(energy), np.asarray(grad.total) + + def calc_new(self, coords): + """geomeTRIC custom-engine entry point. + + ``coords`` is a flattened Cartesian coordinate array in Bohr. The + returned gradient is flattened in Hartree/Bohr. + """ + energy, gradient = self(coords) + return {"energy": energy, "gradient": np.asarray(gradient).ravel()} + + def _normalise_target(self, target, method, state_index, mp_level, + run_adc_kwargs): + if isinstance(target, (GroundStateTarget, ExcitedStateTarget)): + if isinstance(target, GroundStateTarget): + _check_ground_state_level(target.level) + return target + if isinstance(target, str): + method = target + if method is None or method.lower().startswith("mp"): + if method is not None: + mp_level = _mp_level(method, mp_level) + if run_adc_kwargs: + warnings.warn( + "Ignoring run_adc keyword arguments for ground-state MP " + "scanner target.", RuntimeWarning, + ) + _check_ground_state_level(mp_level) + return GroundStateTarget(level=mp_level) + return ExcitedStateTarget( + method=method, state_index=state_index, + run_adc_kwargs=dict(run_adc_kwargs), + ) + + def _coords_array(self, coords): + coords = np.asarray(coords, dtype=float) + if coords.shape == (3 * self.natoms,): + coords = coords.reshape(self.natoms, 3) + if coords.shape != (self.natoms, 3): + raise ValueError( + f"Expected coordinates with shape {(self.natoms, 3)} or " + f"{(3 * self.natoms,)}, got {coords.shape}." + ) + return coords + + def _mol_at(self, coords): + # PySCF's set_geom_ preserves the original Mole settings (basis, charge, + # spin, symmetry setting, unit conventions, etc.) and changes only atom + # coordinates. This keeps the scanner interface anchored on the user's + # configured PySCF object rather than mirroring PySCF keyword arguments. + return self.base_mol.set_geom_(coords, unit="Bohr", inplace=False) + + def _run_scf(self, coords): + mol = self._mol_at(coords) + self.scf_scanner(mol) + if not self.scf_scanner.converged: + raise RuntimeError("PySCF SCF did not converge at scanner geometry.") + # Snapshot the current scanner state for adcc. The scanner itself is + # reused on the next geometry to keep PySCF's orbital/density continuity. + return self.scf_scanner.undo_scanner() + + def _build_target(self, scfres): + if isinstance(self.target, GroundStateTarget): + from adcc import ReferenceState + return LazyMp(ReferenceState(scfres)) + from adcc import run_adc + states = run_adc(scfres, **self.target.kwargs()) + self.last_states = states + excitation = self._select_excitation(states, scfres.mol) + self.last_excitation = excitation + return excitation + + def _select_excitation(self, states, mol): + excitations = states.excitations + if not excitations: + raise RuntimeError("ADC calculation did not return any excited states.") + if self.follow == "index" or self.previous_descriptor is None: + index = self.target.state_index + if index >= len(excitations) or index < -len(excitations): + raise ValueError( + f"state_index {index} is out of range for " + f"{len(excitations)} computed states." + ) + self.last_tracking = None + return excitations[index] + unavailable: set[str] = set() + scores = np.array([ + self._tracking_score(self.previous_descriptor, + self._descriptor(excitation, mol), mol, + unavailable) + for excitation in excitations + ]) + order = np.argsort(scores)[::-1] + best = int(order[0]) + gap = (scores[best] - scores[int(order[1])] + if len(order) > 1 else float("inf")) + # Always record a diagnostic so silent state switching is observable on + # every call via ``last_tracking`` (best score, gap to second-best and + # whether the followed positional index changed). + switched = (self.previous_index is not None + and best != self.previous_index) + self.last_tracking = TrackingResult( + index=best, scores=scores, best_score=float(scores[best]), + gap=float(gap), previous_index=self.previous_index, + switched=bool(switched), + unavailable_channels=tuple(sorted(unavailable)), + ) + if scores[best] < self.tracking_min_score: + raise RuntimeError( + "Root tracking failed: best state-character overlap " + f"{scores[best]:.6g} is below threshold " + f"{self.tracking_min_score:.6g}." + ) + if len(order) > 1 and gap < self.tracking_min_gap: + raise RuntimeError( + "Root tracking is ambiguous: best and second-best overlaps " + f"differ by {gap:.6g}." + ) + return excitations[best] + + def _descriptor(self, excitation, mol): + # PySCF SCF scanner objects mutate their internal Mole in-place between + # geometries. Keep an independent copy for cross-overlap root tracking; + # otherwise the "previous" AO basis silently turns into the current one. + return _TrackingDescriptor( + mol=mol.copy(), + transition_dm=_safe_ao_ndarray(excitation, "transition_dm_ao"), + state_diffdm=_safe_ao_ndarray(excitation, "state_diffdm_ao"), + ) + + def _tracking_score(self, previous, current, current_mol, unavailable=None): + s_cross = self._pyscf.gto.intor_cross( + "int1e_ovlp", previous.mol, current_mol + ) + s_old = previous.mol.intor_symmetric("int1e_ovlp") + s_new = current_mol.intor_symmetric("int1e_ovlp") + score = 0.0 + weight = 0 + for name, old, new, use_abs in [ + ("transition_dm_ao", previous.transition_dm, + current.transition_dm, True), + ("state_diffdm_ao", previous.state_diffdm, + current.state_diffdm, False), + ]: + if old is None or new is None: + if unavailable is not None: + unavailable.add(name) + continue + channel = density_overlap_score( + old, new, s_cross, s_old=s_old, s_new=s_new, use_abs=use_abs, + ) + if np.isfinite(channel): + score += channel + weight += 1 + if weight == 0: + raise RuntimeError( + "Root tracking requested, but neither transition_dm_ao nor " + "state_diffdm_ao could be computed." + ) + return score / weight + + +def density_overlap_score(old_dm, new_dm, s_cross, *, s_old=None, s_new=None, + use_abs=False) -> float: + """Cosine-like overlap of AO-basis density matrices at two geometries. + + ``s_cross`` is the rectangular AO overlap between the old and new PySCF + molecules, i.e. ``pyscf.gto.intor_cross('int1e_ovlp', old_mol, new_mol)``. + ``s_old`` and ``s_new`` are the AO overlap matrices within each geometry. + They default to identity matrices for simple unit tests. + """ + old_dm = np.asarray(old_dm) + new_dm = np.asarray(new_dm) + s_cross = np.asarray(s_cross) + if s_old is None: + s_old = np.eye(old_dm.shape[0]) + if s_new is None: + s_new = np.eye(new_dm.shape[0]) + s_old = np.asarray(s_old) + s_new = np.asarray(s_new) + numerator = np.einsum("pq,pr,rs,qs->", old_dm, s_cross, new_dm, s_cross) + old_norm = np.einsum("pq,pr,rs,qs->", old_dm, s_old, old_dm, s_old) + new_norm = np.einsum("pq,pr,rs,qs->", new_dm, s_new, new_dm, s_new) + denom = np.sqrt(abs(old_norm) * abs(new_norm)) + if denom == 0.0: + return np.nan + score = numerator / denom + if use_abs: + score = abs(score) + return float(score) + + +def _safe_ao_ndarray(excitation, attr): + try: + value = getattr(excitation, attr) + return value.to_ndarray() + except (AttributeError, NotImplementedError): + # The operator/attribute is genuinely unavailable for this method. The + # degradation is recorded per call on ``last_tracking`` (see + # ``_tracking_score``), so it stays observable on every call rather than + # being deduplicated by Python's warning filter. Any other exception + # signals a real bug and is allowed to propagate. + return None + + +def _check_ground_state_level(level): + if level != 2: + raise NotImplementedError( + "adcc only provides MP2 ground-state nuclear gradients, so the " + f"scanner cannot pair an MP{level} energy with an MP2 gradient. " + "Use level=2 for a ground-state MP target." + ) + + +def _mp_level(method, default): + try: + return int(method.lower().removeprefix("mp")) + except ValueError: + return default + + +def _import_pyscf(): + try: + from pyscf import gto, scf + except ImportError as exc: # pragma: no cover - depends on optional dep + raise ModuleNotFoundError( + "NuclearGradientScanner currently requires PySCF. Install PySCF " + "and pass a configured PySCF SCF object." + ) from exc + return _PyscfModules(gto=gto, scf=scf) + + +@dataclass(frozen=True) +class _PyscfModules: + gto: Any + scf: Any + + +def _is_pyscf_scf(pyscf, obj) -> bool: + return isinstance(obj, pyscf.scf.hf.SCF) diff --git a/adcc/tests/ExcitedStates_test.py b/adcc/tests/ExcitedStates_test.py index cfd91274b..af8a33b9a 100644 --- a/adcc/tests/ExcitedStates_test.py +++ b/adcc/tests/ExcitedStates_test.py @@ -44,7 +44,8 @@ def test_excitation_view(self): if key.startswith("_"): continue blacklist = ["__", "index", "_ao", "excitation_vector", - "method", "parent_state"] + "method", "parent_state", "ground_state", + "reference_state"] if any(b in key for b in blacklist): continue try: diff --git a/adcc/tests/data/gradient_data.json b/adcc/tests/data/gradient_data.json new file mode 100644 index 000000000..49506f780 --- /dev/null +++ b/adcc/tests/data/gradient_data.json @@ -0,0 +1 @@ +{"basissets": ["sto3g", "ccpvdz"], "methods": ["mp2", "adc0", "adc1", "adc2", "adc2x", "adc3", "cvs-adc0", "cvs-adc1", "cvs-adc2", "cvs-adc2x"], "molecules": ["h2o"], "h2o": {"sto3g": {"mp2": {"energy": -74.99357808910112, "gradient": [[0.09648089344409527, -6.548361852765083e-11, 0.06886449033481767], [-0.026183461952314246, 1.1641532182693481e-10, -0.06597386243811343], [-0.07029743152088486, 1.0913936421275139e-10, -0.0028906277802889235]], "config": {"conv_tol": 1e-08, "core_orbitals": null}}, "adc0": {"energy": [-73.96883482574164, -73.91481187925336, -73.80575589453754, -73.75173294804925, -73.7272127969465], "gradient": [[[0.2150052078141016, 1.025910023599863e-09, 0.15122539438016247], [0.07489849092962686, 1.1641532182693481e-09, -0.3332042335314327], [-0.2899036992675974, 5.093170329928398e-11, 0.18197884024993982]], [[0.15869260571344057, 6.184563972055912e-10, 0.1113579279917758], [0.12612935056677088, 5.748006515204906e-10, -0.34585730692197103], [-0.28482195638207486, 1.4551915228366852e-11, 0.23449938005796866]], [[0.4883978162833955, 7.130438461899757e-10, 0.3478099758794997], [-0.09521430447784951, 6.111804395914078e-10, -0.3859635299522779], [-0.39318351215479197, 0.0, 0.038153555309691]], [[0.4320852141318028, 3.128661774098873e-10, 0.3079425094401813], [-0.04398344477522187, -2.9103830456733704e-11, -0.39861660321912495], [-0.38810176937113283, -1.4551915228366852e-10, 0.09067409506678814]], [[0.4312687490310054, 5.529727786779404e-10, 0.30399765187030425], [0.0012444711828720756, 4.802132025361061e-10, -0.4583379852992948], [-0.43251322035939666, -1.7462298274040222e-10, 0.1543403345203842]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "adc1": {"energy": [-74.47588319567681, -74.38511885022422, -74.35718229188835, -74.24904200347017, -74.13327522700904], "gradient": [[[0.2519565792172216, 9.458744898438454e-11, 0.17643765192042338], [0.043837964309204835, -2.9103830456733704e-11, -0.3275681610321044], [-0.2957945424786885, 1.0913936421275139e-10, 0.1511305107269436]], [[0.4420402010146063, 1.0186340659856796e-10, 0.315973125914752], [-0.08205664000706747, 1.0186340659856796e-10, -0.35633136319665937], [-0.3599835601489758, 2.9103830456733704e-11, 0.04035823874437483]], [[0.10246273293887498, 0.0, 0.07140888599678874], [0.11964530166005716, 2.0372681319713593e-10, -0.27684558628243394], [-0.22210803395864787, -2.9103830456733704e-11, 0.20543670164624928]], [[0.3338767148015904, -3.128661774098873e-10, 0.23875348411093], [-2.251266414532438e-05, 7.275957614183426e-11, -0.3568454429623671], [-0.3338542017590953, -4.3655745685100555e-11, 0.11809195968817221]], [[0.4270152220633463, -2.764863893389702e-10, 0.3020655645086663], [-0.05104077036958188, 1.8189894035458565e-10, -0.38096642339223763], [-0.3759744515264174, -2.9103830456733704e-11, 0.07890085967665073]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "adc2": {"energy": [-74.5230649472871, -74.4210231425989, -74.39990473399217, -74.28060927279226, -74.15388076822933], "gradient": [[[0.27296630615455797, -1.2150849215686321e-09, 0.19153599689161638], [0.043751712088123895, 4.729372449219227e-10, -0.3499776703538373], [-0.31671801855554804, -6.039044819772243e-10, 0.1584416729165241]], [[0.46611144659254933, -9.38598532229662e-10, 0.3327183593864902], [-0.08526564063504338, 4.43833414465189e-10, -0.3770551477718982], [-0.3808458058992983, -4.220055416226387e-10, 0.044336788152577356]], [[0.124540212797001, -6.548361852765083e-10, 0.0870925094714039], [0.11524763608758803, 2.6921043172478676e-10, -0.29412153096927796], [-0.23978784900100436, -4.511093720793724e-10, 0.20702902116318]], [[0.3449249533659895, -5.165929906070232e-10, 0.24646419959753985], [-0.0014078256208449602, 2.9103830456733704e-10, -0.3665060474377242], [-0.3435171277815243, -4.511093720793724e-10, 0.120041847898392]], [[0.4314259999373462, -4.001776687800884e-10, 0.30517971327208215], [-0.050288449245272204, 2.0372681319713593e-10, -0.3867049110485823], [-0.3811375506775221, -4.5838532969355583e-10, 0.08152519780560397]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "adc2x": {"energy": [-74.54846626855937, -74.44218300963513, -74.42210922347253, -74.30300532544527, -74.16527252516222], "gradient": [[[0.28895747945352923, 4.874891601502895e-10, 0.20262688676302787], [0.04506526638579089, -2.6921043172478676e-10, -0.3685835975629743], [-0.3340227447260986, -8.440110832452774e-10, 0.1659567122333101]], [[0.4846122064409428, 2.1100277081131935e-10, 0.3459748480745475], [-0.08685884059377713, -2.473825588822365e-10, -0.39460422116098925], [-0.3977533652941929, -5.966285243630409e-10, 0.04862937382858945]], [[0.13661488314392045, 2.6921043172478676e-10, 0.09541062157950364], [0.115302309350227, -8.003553375601768e-11, -0.3067889071535319], [-0.25191719177382765, -5.602487362921238e-10, 0.21137828633072786]], [[0.3605390388838714, 2.9103830456733704e-11, 0.25766010260849725], [-0.0009222042790497653, 1.2369127944111824e-10, -0.38391282907832647], [-0.3596168342119199, -3.055902197957039e-10, 0.12625272648438113]], [[0.43669924524147063, 1.382431946694851e-10, 0.30894342893589055], [-0.050859652561484836, 2.4010660126805305e-10, -0.3915265554387588], [-0.3858395925781224, -1.8189894035458565e-10, 0.08258312664111145]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "adc3": {"energy": [-74.55090439628677, -74.44714801103761, -74.4242627369054, -74.3062014889106, -74.166943866151], "gradient": [[[0.29086252811248414, -6.039044819772243e-10, 0.20394485098222503], [0.04500433925568359, 1.4551915228366852e-11, -0.3704894038455677], [-0.33586686704074964, 1.899024937301874e-09, 0.16654455348179908]], [[0.4865842863218859, -4.874891601502895e-10, 0.3473889869055711], [-0.08571265536738792, -2.1827872842550278e-11, -0.39833684905897826], [-0.40087163069983944, 1.3023964129388332e-09, 0.050947862553584855]], [[0.13570507876283955, -4.5838532969355583e-10, 0.09476317655207822], [0.11679784373700386, 3.637978807091713e-11, -0.30793435761734145], [-0.25250292226701276, 1.2878444977104664e-09, 0.21317118139995728]], [[0.3612387815810507, -4.220055416226387e-10, 0.2581562035193201], [0.0004597214501700364, -1.0913936421275139e-10, -0.38661062418395886], [-0.36169850303122075, 6.330083124339581e-10, 0.12845442099205684]], [[0.440449568668555, -2.837623469531536e-10, 0.31157616744894767], [-0.05225038550270256, -7.275957614183426e-12, -0.3935195546873729], [-0.3881991831076448, 3.7834979593753815e-10, 0.08194338760949904]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "cvs-adc0": {"energy": [-54.12308247364575, -53.96000354244165], "gradient": [[[0.13670423539588228, -9.094947017729282e-10, 0.09585740224429173], [0.09903087812563172, -2.5393092073500156e-09, -0.28425848513143137], [-0.23573511977883754, 2.5684130378067493e-09, 0.18840108285075985]], [[0.41009684355231, -7.421476766467094e-10, 0.2924419834816945], [-0.07108191672887187, -2.0590960048139095e-09, -0.3370177815158968], [-0.33901493183657294, 1.9936123862862587e-09, 0.04457579792506294]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": 1}}, "cvs-adc1": {"energy": [-54.853544135891156, -54.794194262470164], "gradient": [[[0.20441433423548006, -9.240466170012951e-10, 0.14198894493893022], [0.033883337360748556, 7.203198038041592e-10, -0.26222279854118824], [-0.23829767041024752, -2.371962182223797e-09, 0.12023385263455566]], [[0.3112917303442373, -8.87666828930378e-10, 0.22435430077894125], [-0.030362786164914723, 8.003553375601768e-10, -0.291552770700946], [-0.28092894334258744, -2.342858351767063e-09, 0.0671984688815428]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": 1}}, "cvs-adc2": {"energy": [-54.989035890092495, -54.90586009741678, -53.166856814459365, -53.10110478594787, -53.003777883255246], "gradient": [[[0.2387200986413518, -5.165929906070232e-10, 0.16710098999465117], [0.033488719614979345, 3.346940502524376e-10, -0.29891402351495344], [-0.27220882266556146, 1.3023964129388332e-09, 0.13181303529563593]], [[0.3619300255377311, -5.966285243630409e-10, 0.2592484486376634], [-0.04305157013004646, 1.8917489796876907e-10, -0.3264197408570908], [-0.3188784597659833, 1.3606040738523006e-09, 0.06717129400931299]], [[0.30556881017400883, -7.203198038041592e-10, 0.21378264685336035], [0.19212377052463125, 5.311449058353901e-10, -0.5935662723713904], [-0.4976925876835594, 2.495653461664915e-09, 0.3797836279409239]], [[0.2346877919335384, -6.111804395914078e-10, 0.16384246804955183], [0.24505627128382912, 4.220055416226387e-10, -0.5933986337622628], [-0.47974406894354615, 1.9572325982153416e-09, 0.4295561676772195]], [[0.5789614186214749, -6.257323548197746e-10, 0.4103672272322001], [0.022010975182638504, 3.637978807091713e-10, -0.6463255684648175], [-0.6009723997267429, 1.9063008949160576e-09, 0.2359583434081287]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": 1}}, "cvs-adc2x": {"energy": [-55.0882980383054, -54.994531290945545, -54.26690371515961, -54.25805543504194, -54.20500159628774], "gradient": [[[0.27884997289220337, -2.750311978161335e-09, 0.19537103317998117], [0.030986498801212292, -1.025910023599863e-09, -0.3378435056147282], [-0.30983647142420523, -7.712515071034431e-10, 0.14247247339517344]], [[0.4121253871926456, -2.473825588822365e-09, 0.29478439864033135], [-0.053034817494335584, -9.240466170012951e-10, -0.3655976063673734], [-0.35909056951641105, -6.621121428906918e-10, 0.07081320842553396]], [[0.6789990261531784, -2.684828359633684e-09, 0.4736056826513959], [-0.019374828414584044, -1.0477378964424133e-09, -0.6864346818911145], [-0.6596241970837582, -7.8580342233181e-10, 0.21282900028018048]], [[0.6081797281003674, -2.6702764444053173e-09, 0.4384273022369598], [0.010582070404780097, -6.693881005048752e-10, -0.6685617406692472], [-0.6187617980031064, -5.748006515204906e-10, 0.2301344396400964]], [[0.4674347668842529, -3.572495188564062e-09, 0.33145696960855275], [0.022216291108634323, -1.127773430198431e-09, -0.5282481098765857], [-0.48965105717797996, -1.2078089639544487e-09, 0.1967911415413255]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": 1}}}, "ccpvdz": {"mp2": {"energy": -76.22940338807781, "gradient": [[0.024395445980189834, 6.912159733474255e-10, 0.017738173199177254], [-0.010762450154288672, 1.1350493878126144e-09, -0.011150374659337103], [-0.013632994901854545, -1.4551915228366852e-10, -0.006587799092812929]], "config": {"conv_tol": 1e-08, "core_orbitals": null}}, "adc0": {"energy": [-75.34687593232839, -75.28099897874145, -75.27790714329215, -75.2120301897052, -75.12638003476145], "gradient": [[[0.05781756929354742, -7.566995918750763e-10, 0.04124683982809074], [-0.0010629517419147305, -1.4551915228366852e-10, -0.06019930707407184], [-0.05675461760984035, -2.546585164964199e-10, 0.01895246607455192]], [[0.013116556481691077, -7.930793799459934e-10, 0.00958195509883808], [0.05172139628120931, -2.1827872842550278e-11, -0.08735981061909115], [-0.06483795282110805, -2.3283064365386963e-10, 0.07777785430516815]], [[0.03587014273944078, -6.548361852765083e-10, 0.025914298770658206], [0.007818847821909003, -2.9103830456733704e-11, -0.04966130383400014], [-0.0436889906268334, -1.5279510989785194e-10, 0.02374700368818594]], [[-0.008830870065139607, -5.602487362921238e-10, -0.005750585907662753], [0.06060319573589368, -2.9103830456733704e-11, -0.07682180725532817], [-0.05177232583082514, -2.473825588822365e-10, 0.08257239179511089]], [[0.2704835088807158, -7.203198038041592e-10, 0.1915996252646437], [-0.0700494504199014, -7.275957614183426e-11, -0.18824094198498642], [-0.2004340584462625, -2.6921043172478676e-10, -0.0033586846766411327]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "adc1": {"energy": [-75.68543998223063, -75.61931015978162, -75.59868950684974, -75.53275299230333, -75.4555818736797], "gradient": [[[0.09273998593562283, -3.637978807091713e-11, 0.06558003614190966], [-0.0026301230100216344, -2.764863893389702e-10, -0.09467178914928809], [-0.09010986285284162, -2.255546860396862e-10, 0.029091751930536702]], [[0.10931789746246068, 2.1827872842550278e-11, 0.07809548162913416], [-0.0066606216569198295, -2.255546860396862e-10, -0.10735303382534767], [-0.10265727566729765, -1.8189894035458565e-10, 0.02925755091564497]], [[0.011670344370941166, 5.093170329928398e-11, 0.008451255693216808], [0.05684476407623151, -1.8189894035458565e-10, -0.09296221136173699], [-0.0685151085126563, -8.731149137020111e-11, 0.08451095448253909]], [[0.026288290071533993, 2.9103830456733704e-11, 0.019256163730460685], [0.057597592058300506, -1.4551915228366852e-10, -0.11000341594626661], [-0.08388588184607215, -2.255546860396862e-10, 0.09074725099344505]], [[0.28210162517643766, -1.4551915228366852e-11, 0.19977476378699066], [-0.07840580360061722, -3.8562575355172157e-10, -0.18871010188013315], [-0.20369582153944066, -1.382431946694851e-10, -0.011064662874559872]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "adc2": {"energy": [-75.92967540577142, -75.85499790072596, -75.84309170584581, -75.76675241368349, -75.66953798762675], "gradient": [[[0.11448491387272952, -5.456968210637569e-10, 0.08113860062439926], [-0.008900549270038027, 3.2741809263825417e-10, -0.10905681972508319], [-0.10558436510473257, 4.0745362639427185e-10, 0.027918217885599006]], [[0.11318143692915328, -4.3655745685100555e-10, 0.08065580731636146], [-0.004472568667551968, 2.546585164964199e-10, -0.11437430455407593], [-0.10870886873453856, 4.0745362639427185e-10, 0.033718496284564026]], [[0.022055673631257378, -6.330083124339581e-10, 0.015903145344054792], [0.05098410337086534, 2.473825588822365e-10, -0.09580103585903998], [-0.07303977738047251, 3.2014213502407074e-10, 0.07989788945997134]], [[0.020305685422499664, -5.456968210637569e-10, 0.014960420950956177], [0.058050360916240606, 3.8562575355172157e-10, -0.10423130493290955], [-0.07835604677529773, 4.220055416226387e-10, 0.08927088286873186]], [[0.26928583782864735, -5.165929906070232e-10, 0.1907385332233389], [-0.07746144225529861, 4.147295840084553e-10, -0.17647510213282658], [-0.19182439600263024, 4.3655745685100555e-10, -0.01426343232014915]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "adc2x": {"energy": [-75.94713027411659, -75.87072120063236, -75.86069046788475, -75.78212932723245, -75.683392163303], "gradient": [[[0.12108458548755152, -2.1827872842550278e-10, 0.08575616330199409], [-0.00893010448635323, -3.8562575355172157e-10, -0.11596751655451953], [-0.11215448191796895, -4.802132025361061e-10, 0.030211352779588196]], [[0.12104126124904724, -7.275957614183426e-11, 0.08622114626632538], [-0.005538550809433218, -4.874891601502895e-10, -0.12121304059110116], [-0.11550271141459234, -3.2741809263825417e-10, 0.03499189356807619]], [[0.027651805263303686, -2.1100277081131935e-10, 0.019813952560070902], [0.05061759612726746, -4.2928149923682213e-10, -0.10117345918115461], [-0.07826940238010138, -3.2741809263825417e-10, 0.08135950614814647]], [[0.02727153090381762, -2.4010660126805305e-10, 0.019892480071575847], [0.05681764219480101, -4.2928149923682213e-10, -0.10988470019947272], [-0.08408917413180461, -4.3655745685100555e-10, 0.08999221950944047]], [[0.27164540941885207, -2.0372681319713593e-10, 0.1923906145602814], [-0.0771593524113996, -4.0745362639427185e-10, -0.17938915952981915], [-0.19448605793877505, -3.710738383233547e-10, -0.01300145529967267]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "adc3": {"energy": [-75.92857834780224, -75.85463334358096, -75.84167672062061, -75.76664280204564, -75.67470300248684], "gradient": [[[0.11627264085109346, 5.529727786779404e-10, 0.08230351519887336], [-0.0062954875465948135, -1.5279510989785194e-10, -0.11453793959663017], [-0.1099771520748618, -4.43833414465189e-10, 0.03223442495072959]], [[0.12384287333406974, 4.874891601502895e-10, 0.08825098589295521], [-0.006875561492051929, -2.546585164964199e-10, -0.12234345744218444], [-0.11696731067058863, -3.346940502524376e-10, 0.034092471811163705]], [[0.02580865211348282, 4.802132025361061e-10, 0.018469874739821535], [0.053307208931073546, -1.5279510989785194e-10, -0.10298056269675726], [-0.07911585993861081, -4.001776687800884e-10, 0.08451068834983744]], [[0.033784865772759076, 4.0745362639427185e-10, 0.024526490720745642], [0.056045929311949294, -1.5279510989785194e-10, -0.1157318419936928], [-0.08983079368044855, -4.0745362639427185e-10, 0.09120535184047185]], [[0.2793041700133472, 5.165929906070232e-10, 0.19778781363129383], [-0.07754871377983363, -1.8917489796876907e-10, -0.18694540472642984], [-0.20175545511301607, -3.5652192309498787e-10, -0.01084240862110164]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": null}}, "cvs-adc0": {"energy": [-55.29234316066312, -55.22337437162689, -54.66716621079568, -54.64254226522229, -54.31941455398455], "gradient": [[[-0.013648041604028549, 2.1100277081131935e-10, -0.00933605365571566], [0.005706590804038569, -2.4010660126805305e-10, 0.006095180557167623], [0.007941450814541895, 3.128661774098873e-10, 0.0032408732440671884]], [[-0.03559546809992753, 1.4551915228366852e-10, -0.02466859494597884], [0.01458839019323932, -1.2369127944111824e-10, 0.016633183913654648], [0.02100707790668821, 3.4924596548080444e-10, 0.008035410843149293]], [[0.29560473508172436, 2.4010660126805305e-10, 0.20220448746113107], [-0.11422406550263986, -2.546585164964199e-10, -0.1452678925197688], [-0.18138066952087684, 5.165929906070232e-10, -0.05693659470125567]], [[-0.04259348352206871, 6.548361852765083e-11, -0.02192712562828092], [0.16518394282320514, -2.3283064365386963e-10, -0.19658394833822967], [-0.12259045930113643, 3.5652192309498787e-10, 0.21851107419206528]], [[-0.09435172603116371, 1.5279510989785194e-10, -0.06654396395606454], [0.01681859556265408, -2.6193447411060333e-10, 0.07614255287626293], [0.07753313046123367, 3.41970007866621e-10, -0.009598588767403271]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": 1}}, "cvs-adc1": {"energy": [-55.762358022721386, -55.74423118010104, -55.238855752322834, -55.200484693742794, -55.17897922566969], "gradient": [[[0.11375782875256846, 6.548361852765083e-10, 0.07506364177243086], [-0.006232537889445666, 4.511093720793724e-10, -0.10649735498736845], [-0.10752529080491513, -8.076312951743603e-10, 0.031433712792932056]], [[0.13675098276144126, 7.057678885757923e-10, 0.10287786688422784], [-0.010133649768249597, 3.710738383233547e-10, -0.13693001883802935], [-0.12661733281856868, -7.930793799459934e-10, 0.0340521515608998]], [[-0.0358682753139874, 8.294591680169106e-10, -0.02505332922009984], [0.028073582740034908, 4.729372449219227e-10, -0.00195446856378112], [0.0077946928577148356, -7.203198038041592e-10, 0.027007797441910952]], [[0.06371899313671747, 7.712515071034431e-10, 0.046058176092628855], [-0.044796131522161886, 5.384208634495735e-10, -0.0052570854386431165], [-0.018922861039754935, -7.930793799459934e-10, -0.0408010909523]], [[-0.052627472738095094, 6.83940015733242e-10, -0.03674976604816038], [0.0053149469385971315, 4.729372449219227e-10, 0.047853415919234976], [0.04731252626515925, -8.149072527885437e-10, -0.011103650183940772]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": 1}}, "cvs-adc2": {"energy": [-56.46332440410863, -56.391989749250264, -55.93190966804171, -55.913659032092795, -55.63489101855794], "gradient": [[[0.06424999397131614, -5.238689482212067e-10, 0.04552041432179976], [-0.0016863320706761442, -3.41970007866621e-10, -0.06586713928118115], [-0.06256366208253894, -5.966285243630409e-10, 0.02034672508307267]], [[0.05899440823122859, 3.7834979593753815e-10, 0.04236195599514758], [0.0059706283354898915, -4.220055416226387e-10, -0.07167670349736], [-0.06496503714151913, -1.0913936421275139e-10, 0.029314748200704344]], [[0.26677654623199487, -5.020410753786564e-10, 0.18210553026437992], [-0.08747906567441532, -4.3655745685100555e-11, -0.15278877816308523], [-0.17929748162714532, -3.4924596548080444e-10, -0.02931675294530578]], [[0.01987398319033673, -6.693881005048752e-10, 0.022062283314880915], [0.11682128573011141, 3.2014213502407074e-10, -0.19428619428072125], [-0.13669526929152198, -5.384208634495735e-10, 0.17222390823008027]], [[-0.08171111583214952, -7.712515071034431e-10, -0.057449760992312804], [0.049808511350420304, 7.057678885757923e-10, 0.01592636768327793], [0.03190260559495073, -8.512870408594608e-10, 0.04152339239954017]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": 1}}, "cvs-adc2x": {"energy": [-56.52769099730406, -56.45532005465755, -56.007565680366255, -55.9924144865235, -55.72589636871257], "gradient": [[[0.08822462362149963, 3.41970007866621e-10, 0.062298513439600356], [-0.0035465051696519367, 1.0913936421275139e-10, -0.08849693409138126], [-0.08467811883747345, 8.949427865445614e-10, 0.026198420448054094]], [[0.09281926582480082, -4.874891601502895e-10, 0.06639221474324586], [-0.0013965219841338694, -1.0186340659856796e-10, -0.09725640671240399], [-0.09142274381156312, 9.167706593871117e-10, 0.030864191117871087]], [[0.26090828989981674, -1.0186340659856796e-10, 0.17750934670766583], [-0.08874587644822896, -6.111804395914078e-10, -0.14432508649042575], [-0.17216241290589096, 2.255546860396862e-10, -0.03318426080659265]], [[0.020961150199582335, 3.41970007866621e-10, 0.023280395165784284], [0.11284214541228721, -1.382431946694851e-10, -0.19026218561339192], [-0.13380329590290785, 2.1100277081131935e-10, 0.16698178952356102]], [[-0.005848991961102001, 1.7462298274040222e-10, -0.0038023634187993594], [0.0007165410861489363, 3.055902197957039e-10, 0.004858471031184308], [0.0051324504529475234, 8.003553375601768e-11, -0.0010561081508058123]]], "config": {"conv_tol": 1e-08, "n_singlets": 5, "core_orbitals": 1}}}, "xyz": "\n O 0 0 0\n H 0 0 1.795239827225189\n H 1.693194615993441 0 -0.599043184453037\n "}} \ No newline at end of file diff --git a/adcc/tests/functionality_geomopt_test.py b/adcc/tests/functionality_geomopt_test.py new file mode 100644 index 000000000..2e65738d9 --- /dev/null +++ b/adcc/tests/functionality_geomopt_test.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2026 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +"""End-to-end geometry-optimisation smoke tests. + +These drive the :class:`adcc.NuclearGradientScanner` through geomeTRIC (via +PySCF's geomopt bridge) and are skipped cleanly when either optional dependency +is unavailable. +""" +import importlib.util + +import numpy as np +import pytest + +import adcc +import adcc.backends + + +def _missing(*modules): + return [m for m in modules if importlib.util.find_spec(m) is None] + + +_required = ["pyscf", "geometric"] +pytestmark = pytest.mark.skipif( + "pyscf" not in adcc.backends.available() or _missing(*_required), + reason="PySCF and geomeTRIC are required for end-to-end geomopt tests.", +) + + +def _h2o_scf(): + from pyscf import gto, scf + mol = gto.M( + atom=""" + O 0 0 0 + H 0 0 1.795239827225189 + H 1.693194615993441 0 -0.599043184453037 + """, + basis="sto-3g", unit="Bohr", symmetry=False, + verbose=0, parse_arg=False, + ) + mf = scf.RHF(mol) + mf.conv_tol = 1e-11 + mf.conv_tol_grad = 1e-9 + return mf + + +def test_ground_state_optimization_reaches_stationary_point(): + from pyscf.geomopt import as_pyscf_method, geometric_solver + + scfres = _h2o_scf() + scanner = adcc.NuclearGradientScanner(scfres, method="mp2") + + def energy_and_gradient(mol_at_step): + return scanner(mol_at_step.atom_coords(unit="Bohr")) + + method = as_pyscf_method(scfres.mol, energy_and_gradient) + mol_eq = geometric_solver.optimize( + method, maxsteps=50, + convergence_grms=1e-4, convergence_gmax=2e-4, + ) + + # The gradient at the optimised geometry must be (near) zero. + _, gradient = scanner(mol_eq.atom_coords(unit="Bohr")) + gmax = np.abs(gradient).max() + assert gmax < 5e-4 + + +def test_calc_new_drives_geometric_internal_engine(): + # Exercise the geomeTRIC custom-engine calc_new contract directly: flattened + # Bohr coordinates in, energy plus flattened Hartree/Bohr gradient out. + scanner = adcc.NuclearGradientScanner(_h2o_scf(), method="mp2") + result = scanner.calc_new(scanner.initial_coords.ravel()) + assert set(result) == {"energy", "gradient"} + assert np.isfinite(result["energy"]) + assert result["gradient"].shape == (3 * scanner.natoms,) + + +def test_excited_state_optimization_tracks_state_across_steps(): + from pyscf.geomopt import as_pyscf_method, geometric_solver + + scfres = _h2o_scf() + scanner = adcc.NuclearGradientScanner( + scfres, method="adc2", n_singlets=3, state_index=0, + follow="overlap", conv_tol=1e-8, + gradient_kwargs={"eri_contraction": "full_ao"}, + ) + + seen_energies = [] + + def energy_and_gradient(mol_at_step): + energy, gradient = scanner(mol_at_step.atom_coords(unit="Bohr")) + seen_energies.append(scanner.last_excitation.excitation_energy) + return energy, gradient + + method = as_pyscf_method(scfres.mol, energy_and_gradient) + + # The optimisation does not need to fully converge here; a few steps are + # enough to exercise root tracking across geometries. + geometric_solver.optimize( + method, maxsteps=3, + convergence_grms=1e-4, convergence_gmax=2e-4, + ) + + # Multiple scanner calls were made and state tracking stayed active. + assert len(seen_energies) >= 2 + assert scanner.last_tracking is not None + # The followed excitation energy must stay finite and positive every step. + assert all(np.isfinite(omega) and omega > 0.0 for omega in seen_energies) diff --git a/adcc/tests/functionality_gradients_test.py b/adcc/tests/functionality_gradients_test.py new file mode 100644 index 000000000..e7fe766fb --- /dev/null +++ b/adcc/tests/functionality_gradients_test.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2020 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +import adcc +import adcc.backends + +from numpy.testing import assert_allclose + +import pytest + +from .backends.testing import cached_backend_hf +from .testdata_cache import gradient_data + + +backends = [b for b in adcc.backends.available() + if b not in ["molsturm", "veloxchem"]] + +molecules = gradient_data["molecules"] +basissets = gradient_data["basissets"] +methods = gradient_data["methods"] + + +@pytest.mark.skipif(len(backends) == 0, reason="No backend found.") +@pytest.mark.parametrize("backend", backends) +@pytest.mark.parametrize("method", methods) +@pytest.mark.parametrize("basis", basissets) +@pytest.mark.parametrize("molecule", molecules) +def test_nuclear_gradient(molecule, basis, method, backend): + grad_ref = gradient_data[molecule][basis][method] + + energy_ref = grad_ref["energy"] + grad_fdiff = grad_ref["gradient"] + kwargs = dict(grad_ref["config"]) + conv_tol = kwargs["conv_tol"] + + scfres = cached_backend_hf(backend, f"{molecule}_{basis}", conv_tol=1e-11) + print(kwargs) + + if "adc" in method: + # Request exactly as many states as reference data is available for. + # Some CVS cases (e.g. h2o/sto3g cvs-adc0/adc1) support fewer states + # than the default n_singlets, so cap the request accordingly. + if "n_singlets" in kwargs: + kwargs["n_singlets"] = len(energy_ref) + state = adcc.run_adc(scfres, method=method, **kwargs) + for ee in state.excitations: + grad = adcc.nuclear_gradient(ee) + assert_allclose(energy_ref[ee.index], ee.total_energy, + atol=conv_tol) + # check energy computed with unrelaxed densities + gs_corr = 0.0 + if ee.method.level.to_int() > 0: + # compute the ground state contribution + # to the correlation energy + gs_energy = ee.ground_state.energy(ee.method.level.to_int()) + gs_corr = gs_energy - ee.reference_state.energy_scf + assert_allclose( + gs_corr + ee.excitation_energy, + grad._energy, atol=1e-10 + ) + assert_allclose( + grad_fdiff[ee.index], grad.total, atol=2e-7, rtol=0, + err_msg=f'Gradient for state {ee.index} wrong.' + ) + else: + # MP2 gradients + refstate = adcc.ReferenceState(scfres) + mp = adcc.LazyMp(refstate) + grad = adcc.nuclear_gradient(mp) + assert_allclose(energy_ref, mp.energy(2), atol=1e-8) + # check energy computed with unrelaxed densities + assert_allclose( + mp.energy_correction(2), grad._energy, atol=5e-8, rtol=0 + ) + assert_allclose( + grad_fdiff, grad.total, atol=1e-8, rtol=0 + ) diff --git a/adcc/tests/geomopt_scanner_test.py b/adcc/tests/geomopt_scanner_test.py new file mode 100644 index 000000000..76603e628 --- /dev/null +++ b/adcc/tests/geomopt_scanner_test.py @@ -0,0 +1,418 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2026 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +import types + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +import adcc +import adcc.backends +from adcc.gradients.scanner import ( + GroundStateTarget, + _TrackingDescriptor, + _safe_ao_ndarray, + density_overlap_score, +) + + +pytestmark = pytest.mark.skipif( + "pyscf" not in adcc.backends.available(), reason="PySCF not found." +) + + +def _h2o_scf(): + from pyscf import gto, scf + mol = gto.M( + atom=""" + O 0 0 0 + H 0 0 1.795239827225189 + H 1.693194615993441 0 -0.599043184453037 + """, + basis="sto-3g", unit="Bohr", symmetry=False, + verbose=0, parse_arg=False, + ) + mf = scf.RHF(mol) + mf.conv_tol = 1e-11 + mf.conv_tol_grad = 1e-9 + return mf + + +def test_density_overlap_score_uses_absolute_transition_phase(): + dm = np.eye(2) + s = np.eye(2) + assert density_overlap_score(dm, dm, s) == pytest.approx(1.0) + assert density_overlap_score(dm, -dm, s) == pytest.approx(-1.0) + assert density_overlap_score(dm, -dm, s, use_abs=True) == pytest.approx(1.0) + + +def test_scanner_requires_pyscf_scf_object(): + with pytest.raises(TypeError, match="PySCF SCF object"): + adcc.NuclearGradientScanner(_h2o_scf().mol, method="mp2") + + +def test_scanner_validates_coordinate_shape(): + scanner = adcc.NuclearGradientScanner(_h2o_scf(), method="mp2") + with pytest.raises(ValueError, match="Expected coordinates"): + scanner(np.zeros((2, 3))) + + +def test_ground_state_scanner_matches_explicit_gradient_loop(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="mp2", gradient_kwargs={"eri_contraction": "full_ao"}, + ) + coords = scanner.initial_coords + + energy, gradient = scanner(coords) + + mp = adcc.LazyMp(adcc.ReferenceState(scanner.last_scf)) + explicit_gradient = adcc.nuclear_gradient(mp, eri_contraction="full_ao") + assert energy == pytest.approx(mp.energy(2)) + assert gradient.shape == (3, 3) + assert_allclose(gradient, explicit_gradient.total) + + +def test_calc_new_returns_geometric_engine_shape(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="mp2", gradient_kwargs={"eri_contraction": "full_ao"}, + ) + result = scanner.calc_new(scanner.initial_coords.ravel()) + assert set(result) == {"energy", "gradient"} + assert isinstance(result["energy"], float) + assert result["gradient"].shape == (9,) + + +def test_run_adc_kwargs_are_forwarded_with_native_names(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=2, conv_tol=1e-7, + follow="index", gradient_kwargs={"eri_contraction": "full_ao"}, + ) + assert scanner.target.kwargs()["conv_tol"] == pytest.approx(1e-7) + assert scanner.target.kwargs()["n_singlets"] == 2 + assert "output" not in scanner.target.kwargs() + + +def test_tracking_descriptor_keeps_independent_mol_snapshot(): + class FakeAoOperator: + def to_ndarray(self): + return np.eye(7) + + class FakeExcitation: + transition_dm_ao = FakeAoOperator() + state_diffdm_ao = FakeAoOperator() + + scanner = adcc.NuclearGradientScanner(_h2o_scf(), method="adc2") + mol = scanner.base_mol.copy() + descriptor = scanner._descriptor(FakeExcitation(), mol) + old_coords = descriptor.mol.atom_coords(unit="Bohr").copy() + + mol.set_geom_(old_coords + 0.1, unit="Bohr") + + assert_allclose(descriptor.mol.atom_coords(unit="Bohr"), old_coords) + + +def test_unconverged_scf_raises(): + from pyscf import scf + mf = scf.RHF(_h2o_scf().mol) + mf.max_cycle = 1 + mf.conv_tol = 1e-12 + mf.conv_tol_grad = 1e-10 + scanner = adcc.NuclearGradientScanner(mf, method="mp2") + with pytest.raises(RuntimeError, match="did not converge"): + scanner(scanner.initial_coords) + + +def test_excited_state_scanner_matches_explicit_gradient_loop(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=3, state_index=0, + follow="index", conv_tol=1e-9, + gradient_kwargs={"eri_contraction": "full_ao"}, + ) + + energy, gradient = scanner(scanner.initial_coords) + + # Fixed-index selection must pick exactly the requested positional state. + assert scanner.last_excitation.index == 0 + assert scanner.last_tracking is None + + states = adcc.run_adc(scanner.last_scf, method="adc2", n_singlets=3, + conv_tol=1e-9) + explicit = adcc.nuclear_gradient(states.excitations[0], + eri_contraction="full_ao") + assert energy == pytest.approx(states.excitations[0].total_energy, abs=1e-7) + assert gradient.shape == (3, 3) + assert_allclose(gradient, explicit.total, atol=1e-7) + + +def test_overlap_tracking_follows_same_state_character(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=3, state_index=1, + follow="overlap", conv_tol=1e-9, + gradient_kwargs={"eri_contraction": "full_ao"}, + ) + + # First call seeds the tracker from the positional index. + scanner(scanner.initial_coords) + assert scanner.last_excitation.index == 1 + assert scanner.last_tracking is None + + # Second call at the same geometry must follow the seeded state via density + # overlap. At identical geometry the matching state has self-overlap ~1. + scanner(scanner.initial_coords) + assert scanner.last_tracking is not None + assert scanner.last_tracking.index == 1 + scores = scanner.last_tracking.scores + assert np.argmax(scores) == 1 + assert scores[1] == pytest.approx(1.0, abs=1e-3) + + +def test_overlap_tracking_continuous_under_small_displacement(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=3, state_index=0, + follow="overlap", conv_tol=1e-9, + gradient_kwargs={"eri_contraction": "full_ao"}, + ) + + scanner(scanner.initial_coords) + omega_initial = scanner.last_excitation.excitation_energy + + displaced = scanner.initial_coords.copy() + displaced[0, 2] += 0.02 + scanner(displaced) + + # The followed state should remain the same physical state, i.e. its + # excitation energy must not jump discontinuously between candidates. + assert scanner.last_tracking is not None + assert abs(scanner.last_excitation.excitation_energy - omega_initial) < 0.05 + + +def test_overlap_tracking_below_min_score_raises(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=3, state_index=0, + follow="overlap", tracking_min_score=2.0, conv_tol=1e-9, + gradient_kwargs={"eri_contraction": "full_ao"}, + ) + scanner(scanner.initial_coords) # seed + with pytest.raises(RuntimeError, match="below threshold"): + scanner(scanner.initial_coords) + + +def test_overlap_tracking_ambiguous_gap_raises(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=3, state_index=0, + follow="overlap", tracking_min_gap=10.0, conv_tol=1e-9, + gradient_kwargs={"eri_contraction": "full_ao"}, + ) + scanner(scanner.initial_coords) # seed + with pytest.raises(RuntimeError, match="ambiguous"): + scanner(scanner.initial_coords) + + +def test_scf_guess_continuity_updates_previous_state(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="mp2", gradient_kwargs={"eri_contraction": "full_ao"}, + ) + + energy_first, _ = scanner(scanner.initial_coords) + assert scanner.previous_scf is scanner.last_scf + + # Move away and return: the scanner reuses one PySCF scanner object across + # geometries (orbital/guess continuity), and the surface stays consistent so + # the energy at the original geometry is reproduced. + displaced = scanner.initial_coords.copy() + displaced[0, 2] += 0.05 + scanner(displaced) + energy_back, _ = scanner(scanner.initial_coords) + assert energy_back == pytest.approx(energy_first, abs=1e-9) + + +# --------------------------------------------------------------------------- +# Root-reorder tracking (FINDING 4) +# --------------------------------------------------------------------------- + +class _FakeAoOperator: + def __init__(self, matrix): + self._matrix = matrix + + def to_ndarray(self): + return self._matrix + + +class _FakeExcitation: + def __init__(self, index, transition_dm, state_diffdm): + self.index = index + self.transition_dm_ao = _FakeAoOperator(transition_dm) + self.state_diffdm_ao = _FakeAoOperator(state_diffdm) + + +def test_overlap_tracking_follows_reordered_root_by_overlap(): + # Prove that overlap tracking follows a physical state whose *positional* + # index differs from the originally requested ``state_index``: the seeded + # descriptor matches candidate #1, while the requested state_index is 0. + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=3, state_index=0, + follow="overlap", + ) + nao = scanner.base_mol.nao + + def projector(i): + m = np.zeros((nao, nao)) + m[i, i] = 1.0 + return m + + # The previously-followed physical state has this character. + seed_transition = projector(0) + seed_diffdm = projector(1) + scanner.previous_descriptor = _TrackingDescriptor( + mol=scanner.base_mol.copy(), + transition_dm=seed_transition, state_diffdm=seed_diffdm, + ) + scanner.previous_index = 0 + + # Candidate roots: index 1 carries the seeded character (energy ordering + # swapped so it is no longer at positional index 0). + candidates = [ + _FakeExcitation(0, projector(2), projector(3)), + _FakeExcitation(1, seed_transition.copy(), seed_diffdm.copy()), + _FakeExcitation(2, projector(4), projector(5)), + ] + states = types.SimpleNamespace(excitations=candidates) + + selected = scanner._select_excitation(states, scanner.base_mol.copy()) + + # Tracking must pick the reordered root (positional index 1) by overlap, + # which differs from the requested state_index (0). + assert selected.index == 1 + assert scanner.target.state_index == 0 + assert scanner.last_tracking is not None + assert scanner.last_tracking.index == 1 + assert scanner.last_tracking.index != scanner.target.state_index + assert np.argmax(scanner.last_tracking.scores) == 1 + # New diagnostic fields are populated and observable on every call. + assert scanner.last_tracking.best_score == pytest.approx(1.0) + assert scanner.last_tracking.gap > 0.0 + assert np.isfinite(scanner.last_tracking.gap) + assert scanner.last_tracking.previous_index == 0 + assert scanner.last_tracking.switched is True + + +# --------------------------------------------------------------------------- +# Validation / construction-time branches (FINDING 10) +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("kwargs", [ + {"target": "mp3"}, + {"method": "mp3"}, + {"target": GroundStateTarget(level=3)}, +]) +def test_ground_state_level_other_than_two_is_rejected(kwargs): + with pytest.raises(NotImplementedError, match="MP2 ground-state"): + adcc.NuclearGradientScanner(_h2o_scf(), **kwargs) + + +def test_invalid_follow_raises_eagerly_at_construction(): + with pytest.raises(ValueError, match="overlap.*index"): + adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=2, follow="bogus", + ) + + +def test_state_index_out_of_range_raises(): + scanner = adcc.NuclearGradientScanner( + _h2o_scf(), method="adc2", n_singlets=2, state_index=5, + follow="index", conv_tol=1e-9, + gradient_kwargs={"eri_contraction": "full_ao"}, + ) + with pytest.raises(ValueError, match="out of range"): + scanner(scanner.initial_coords) + + +def test_empty_excitations_raises_runtime_error(): + scanner = adcc.NuclearGradientScanner(_h2o_scf(), method="adc2") + states = types.SimpleNamespace(excitations=[]) + with pytest.raises(RuntimeError, match="did not return any excited states"): + scanner._select_excitation(states, scanner.base_mol.copy()) + + +@pytest.mark.parametrize("kwargs", [ + {"method": "mp2", "n_singlets": 3}, + {"target": "mp2", "n_singlets": 3}, +]) +def test_ground_state_run_adc_kwargs_emit_runtime_warning(kwargs): + with pytest.warns(RuntimeWarning, match="Ignoring run_adc"): + adcc.NuclearGradientScanner(_h2o_scf(), **kwargs) + + +# --------------------------------------------------------------------------- +# Density-overlap scoring branches (FINDING 10) +# --------------------------------------------------------------------------- + +def test_safe_ao_ndarray_swallows_unavailable_channels(): + class FailingOperator: + def to_ndarray(self): + raise NotImplementedError("not available") + + class Excitation: + transition_dm_ao = FailingOperator() + + @property + def state_diffdm_ao(self): + raise AttributeError("no diffdm") + + excitation = Excitation() + assert _safe_ao_ndarray(excitation, "transition_dm_ao") is None + assert _safe_ao_ndarray(excitation, "state_diffdm_ao") is None + + +def test_tracking_score_records_unavailable_and_raises_when_both_missing(): + scanner = adcc.NuclearGradientScanner(_h2o_scf(), method="adc2") + mol = scanner.base_mol.copy() + previous = _TrackingDescriptor(mol=mol.copy(), + transition_dm=None, state_diffdm=None) + current = _TrackingDescriptor(mol=mol.copy(), + transition_dm=None, state_diffdm=None) + unavailable = set() + with pytest.raises(RuntimeError, match="neither transition_dm_ao"): + scanner._tracking_score(previous, current, mol, unavailable) + # Both channels are surfaced as unavailable before the weight==0 failure. + assert unavailable == {"transition_dm_ao", "state_diffdm_ao"} + + +def test_density_overlap_score_zero_density_returns_nan(): + zero = np.zeros((2, 2)) + s = np.eye(2) + assert np.isnan(density_overlap_score(zero, np.eye(2), s)) + + +def test_density_overlap_score_nonidentity_metrics_matches_hand_value(): + old = np.array([[1.0, 0.0], [0.0, 0.0]]) + new = np.array([[1.0, 0.0], [0.0, 0.0]]) + s_cross = np.array([[0.9, 0.1], [0.1, 0.9]]) + s_old = np.array([[1.0, 0.2], [0.2, 1.0]]) + s_new = np.array([[1.0, 0.3], [0.3, 1.0]]) + + # numerator = (D_old . s_cross . D_new . s_cross), norms use s_old/s_new. + # With single populated diagonal element: numerator = s_cross[0,0]**2 = 0.81, + # old_norm = s_old[0,0]**2 = 1.0, new_norm = s_new[0,0]**2 = 1.0. + score = density_overlap_score(old, new, s_cross, s_old=s_old, s_new=s_new) + assert score == pytest.approx(0.81) diff --git a/adcc/tests/gradient_contraction_test.py b/adcc/tests/gradient_contraction_test.py new file mode 100644 index 000000000..53b5b1a09 --- /dev/null +++ b/adcc/tests/gradient_contraction_test.py @@ -0,0 +1,620 @@ +#!/usr/bin/env python3 +## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab +## --------------------------------------------------------------------- +## +## Copyright (C) 2026 by the adcc authors +## +## This file is part of adcc. +## +## adcc is free software: you can redistribute it and/or modify +## it under the terms of the GNU General Public License as published +## by the Free Software Foundation, either version 3 of the License, or +## (at your option) any later version. +## +## adcc is distributed in the hope that it will be useful, +## but WITHOUT ANY WARRANTY; without even the implied warranty of +## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +## GNU General Public License for more details. +## +## You should have received a copy of the GNU General Public License +## along with adcc. If not, see . +## +## --------------------------------------------------------------------- +import glob +import os + +import numpy as np +from numpy.testing import assert_allclose +import pytest + +import adcc +import adcc.backends +import adcc.block as b +from adcc.backends import have_backend +from adcc.MoSpaces import split_spaces +from adcc.Tensor import Tensor +from adcc.gradients.TwoParticleDensityMatrix import ( + TwoParticleDensityMatrix, ao_pair_indices, +) + +from .backends.testing import cached_backend_hf + + +pytestmark = pytest.mark.skipif( + not have_backend("pyscf"), reason="pyscf not found." +) + +# Spin cases reproduced by ``to_ao_pair_density`` / ``to_ao_basis``. The +# direct (``g2_ao_1``) cases enter with ``+`` sign, the exchange (``g2_ao_2``) +# cases with ``-`` sign and a swapped ket pair. +_DIRECT_SPIN_CASES = [("a", "a", "a", "a"), ("b", "b", "b", "b"), + ("a", "b", "a", "b"), ("b", "a", "b", "a")] +_EXCHANGE_SPIN_CASES = [("a", "a", "a", "a"), ("b", "b", "b", "b"), + ("a", "b", "b", "a"), ("b", "a", "a", "b")] + + +def _random_ao_inputs(nao): + rng = np.random.default_rng(20240611) + g1_ao = rng.standard_normal((nao, nao)) + g1_ao = 0.5 * (g1_ao + g1_ao.T) + w_ao = rng.standard_normal((nao, nao)) + w_ao = 0.5 * (w_ao + w_ao.T) + g2_ao_1 = rng.standard_normal((nao, nao, nao, nao)) + g2_ao_2 = rng.standard_normal((nao, nao, nao, nao)) + return g1_ao, w_ao, g2_ao_1, g2_ao_2 + + +def _random_tpdm(refstate, blocks=(b.oooo, b.oovv, b.ovov)): + """ + Build a ``TwoParticleDensityMatrix`` with deterministically filled blocks. + + Using an explicit ``TwoParticleDensityMatrix`` (instead of a physical MP2 + density) gives dense, non-degenerate data in every block so the full spin + expansion (``aaaa``/``bbbb`` and the mixed ``abab``/``baba``/``abba``/ + ``baab`` cases) is exercised with generic numbers rather than values that + might accidentally cancel. + """ + g2 = TwoParticleDensityMatrix(refstate) + g2.reference_state = refstate + for blk in blocks: + tensor = g2[blk] + tensor.set_random() + g2[blk] = tensor + return g2 + + +def _random_tpdm_with_distinct_spin_blocks( + refstate, blocks=(b.oooo, b.oovv, b.ovov) +): + """ + Build a ``TwoParticleDensityMatrix`` whose ``aaaa`` and ``bbbb`` spin + sub-blocks are independent random data. + + This intentionally bypasses the spin-block maps that a restricted + ``ReferenceState`` would impose, so ``aaaa != bbbb`` and the full + eight-spin-case transform cannot be silently replaced by the restricted + fast path. + """ + g2 = TwoParticleDensityMatrix(refstate) + g2.reference_state = refstate + rng = np.random.default_rng(20240612) + for blk in blocks: + tensor = Tensor(refstate.mospaces, blk) + tensor.set_from_ndarray(rng.standard_normal(tensor.shape), 1e-14) + g2[blk] = tensor + return g2 + + +def _perturbed_restricted_coefficient_map( + g2, refstate, magnitude=1e-2, seed=20240613 +): + """ + Return a coefficient map based on a restricted reference but with a + deliberately different beta spatial coefficient matrix. + + This makes a raw coefficient dict distinguishable from the reference-state + fast path, which assumes alpha and beta spatial coefficients are identical. + """ + cc = g2._ao_coefficient_map(refstate) + rng = np.random.default_rng(seed) + perturbed = {} + for sp in g2.orbital_subspaces: + perturbed[f"{sp}_a"] = cc[f"{sp}_a"].copy() + beta = cc[f"{sp}_b"].copy() + beta += magnitude * rng.standard_normal(beta.shape) + perturbed[f"{sp}_b"] = beta + return perturbed + + +def _direct_gradient_inputs(mp): + """Inputs for ``correlated_gradient_direct`` from an ``LazyMp``. + + Returns ``(g1_ao, w_ao, g2_total)`` where the one-electron AO matrices are + zero (the storage/validation tests only exercise the two-electron packed + density path) and ``g2_total`` is the real MP2 two-particle density matrix + with its block prefactors applied, as built by ``nuclear_gradient``. + """ + grad = adcc.nuclear_gradient(mp, eri_contraction="full_ao") + nao = mp.reference_state.gradient_provider.mol.nao_nr() + g1_ao = np.zeros((nao, nao)) + w_ao = np.zeros((nao, nao)) + return g1_ao, w_ao, grad.g2 + + +def _two_stage_packed_density(g2, refstate, coeff_map=None): + """ + Dense two-stage reference for the packed AO-pair effective density. + + This is intentionally written as an explicit + ``(p, r) -> (mu, nu)`` half-transform followed by a + ``(q, s) -> (lambda, sigma)`` transform + ``s2kl`` pack, so it pins the + *staging contract* that the libtensor-based refactor (milestones m2/m3) + must reproduce, independently of the current fused implementation. + + An optional ``coeff_map`` can be supplied to exercise the transform with + non-default (e.g. deliberately alpha/beta-asymmetric) coefficients. + """ + if coeff_map is None: + cc = g2._ao_coefficient_map(refstate) + else: + cc = { + key: np.asarray(coeff) for key, coeff in coeff_map.items() + } + nao = next(iter(cc.values())).shape[1] + npair = nao * (nao + 1) // 2 + qall, sall = ao_pair_indices(nao) + + def accumulate(out, qidx, sidx): + for block in g2.blocks_nonzero: + spaces = split_spaces(block) + tensor = np.asarray(g2[block].to_ndarray()) + for spins in _DIRECT_SPIN_CASES: + c1, c2, c3, c4 = (cc[f"{sp}_{spin}"] + for sp, spin in zip(spaces, spins)) + # stage 1: bra pair (p, r) -> (mu, nu); keep ket pair in MO + half = np.einsum("ip,kr,ijkl->prjl", c1, c3, tensor, + optimize=True) + # stage 2: ket pair (q, s) -> packed (lambda, sigma) + right = c2[:, qidx][:, None, :] * c4[:, sidx][None, :, :] + out += np.einsum("prjl,jlm->prm", half, right, optimize=True) + for spins in _EXCHANGE_SPIN_CASES: + c1, c2, c3, c4 = (cc[f"{sp}_{spin}"] + for sp, spin in zip(spaces, spins)) + half = np.einsum("ip,lr,ijkl->prjk", c1, c4, tensor, + optimize=True) + right = c2[:, qidx][:, None, :] * c3[:, sidx][None, :, :] + out -= np.einsum("prjk,jkm->prm", half, right, optimize=True) + + out = np.zeros((nao, nao, npair)) + accumulate(out, qall, sall) + # complete off-diagonal ket pairs (s2kl stores both orders summed) + offdiag = qall != sall + if np.any(offdiag): + swapped = np.zeros((nao, nao, int(np.count_nonzero(offdiag)))) + accumulate(swapped, sall[offdiag], qall[offdiag]) + out[:, :, offdiag] += swapped + return out + + +@pytest.mark.parametrize("shell_chunk_size", [1, 2, 1000]) +def test_packed_ao_pair_contraction_matches_full_pyscf_reference( + shell_chunk_size, +): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + provider = hf.gradient_provider + inputs = _random_ao_inputs(hf.n_bas) + _, _, g2_ao_1, g2_ao_2 = inputs + + full = provider.correlated_gradient(*inputs) + pair_density = TwoParticleDensityMatrix.ao_pair_density_from_dense( + g2_ao_1, g2_ao_2 + ) + packed_eri = provider._contract_eri_with_packed_density( + pair_density, shell_chunk_size=shell_chunk_size + ) + + assert_allclose(packed_eri, full.two_electron, atol=1e-10) + + +def test_direct_mp2_pair_density_matches_dense_ao_transform(): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + mp = adcc.LazyMp(adcc.ReferenceState(hf)) + gradient = adcc.nuclear_gradient(mp, eri_contraction="full_ao") + + g2_ao_1, g2_ao_2 = gradient.g2.to_ao_basis() + dense_pair_density = TwoParticleDensityMatrix.ao_pair_density_from_dense( + g2_ao_1.to_ndarray(), g2_ao_2.to_ndarray() + ) + direct_pair_density = gradient.g2.to_ao_pair_density( + gradient.reference_state, pair_chunk_size=5 + ) + + assert_allclose(direct_pair_density, dense_pair_density, atol=1e-10) + + +def test_direct_mp2_gradient_matches_full_ao_fallback(): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + mp = adcc.LazyMp(adcc.ReferenceState(hf)) + + full = adcc.nuclear_gradient(mp, eri_contraction="full_ao") + direct = adcc.nuclear_gradient( + mp, eri_contraction="direct", eri_shell_chunk_size=2, + eri_pair_chunk_size=5 + ) + + assert_allclose(direct.total, full.total, atol=1e-10) + assert_allclose(direct.components.two_electron, + full.components.two_electron, atol=1e-10) + + +def test_restricted_fast_path_matches_eight_case_transform(): + """RHF fast path must reproduce the legacy eight-spin-case transform.""" + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + assert refstate.restricted + mp = adcc.LazyMp(refstate) + gradient = adcc.nuclear_gradient(mp, eri_contraction="full_ao") + g2 = gradient.g2 + + # A raw coefficient dict has no restricted flag and therefore exercises the + # legacy eight-case spin expansion with the same fused contraction order. + coeffs = g2._ao_coefficient_map(refstate) + eight_case = g2.to_ao_pair_density(coeffs, pair_chunk_size=5) + fast_path = g2.to_ao_pair_density(refstate, pair_chunk_size=5) + + assert_allclose(fast_path, eight_case, atol=1e-12) + + +def test_restricted_coefficient_dict_uses_full_spin_expansion(): + """Coefficient dictionaries must not enable the RHF-only shortcut. + + A raw dict has no reliable restricted flag, so the implementation must + use the full eight-spin-case expansion even when the underlying + reference is restricted. The test uses different alpha/beta spatial + coefficients and a TPDM with distinct ``aaaa``/``bbbb`` sub-blocks so + that an accidental fast-path activation would produce a numerically + different result. + """ + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + assert refstate.restricted + + g2 = _random_tpdm_with_distinct_spin_blocks(refstate) + coeffs = _perturbed_restricted_coefficient_map(g2, refstate) + reference = _two_stage_packed_density(g2, refstate, coeff_map=coeffs) + produced = g2.to_ao_pair_density(coeffs, pair_chunk_size=3) + + assert_allclose(produced, reference, atol=1e-10) + + +# --------------------------------------------------------------------------- +# Guardrail tests (milestone m0) +# +# These pin the libtensor -> dense handover *at the intermediate boundary* +# that the upcoming refactor will move, so a correct relocation of that +# boundary can be told apart from a subtly wrong one. They include an +# open-shell reference and an explicit two-stage staging reference. +# --------------------------------------------------------------------------- + + +def test_open_shell_pair_density_matches_dense_ao_transform(): + """Direct packed density must match the dense AO transform for UHF. + + An open-shell reference (CN doublet, n_alpha != n_beta, distinct alpha and + beta spatial orbitals) is the strongest available check that the mixed-spin + cases of the transform are handled correctly when the refactor moves the + handover boundary. + """ + hf = cached_backend_hf("pyscf", "cn_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + assert not refstate.restricted + + g2 = _random_tpdm(refstate) + g2_ao_1, g2_ao_2 = g2.to_ao_basis() + dense_pair_density = TwoParticleDensityMatrix.ao_pair_density_from_dense( + g2_ao_1.to_ndarray(), g2_ao_2.to_ndarray() + ) + direct_pair_density = g2.to_ao_pair_density(refstate, pair_chunk_size=3) + + assert_allclose(direct_pair_density, dense_pair_density, atol=1e-10) + + +@pytest.mark.parametrize("system,use_coeff_dict", [ + ("h2o_sto3g", True), + ("h2o_sto3g", False), + ("cn_sto3g", False), +]) +def test_packed_density_matches_two_stage_staging_contract(system, use_coeff_dict): + """Packed density must equal an explicit bra/ket two-stage transform. + + This pins the staging contract (bra half-transform ``(p,r)->(mu,nu)`` then + ket transform + ``s2kl`` pack) that milestones m2/m3 must reproduce when + the transform is lifted into libtensor, independently of the current fused + einsum implementation. + + For restricted references the test runs both via a raw coefficient dict + (full eight-spin-case path) and via the reference state (restricted fast + path), so the fast path is checked against an independent reference on + generic random TPDM data. + """ + hf = cached_backend_hf("pyscf", system, conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + + g2 = _random_tpdm(refstate) + reference = _two_stage_packed_density(g2, refstate) + if use_coeff_dict: + coeffs_or_refstate = g2._ao_coefficient_map(refstate) + else: + coeffs_or_refstate = refstate + produced = g2.to_ao_pair_density(coeffs_or_refstate, pair_chunk_size=3) + + assert_allclose(produced, reference, atol=1e-10) + + +@pytest.mark.parametrize("pair_chunk_size", [1, 2, 3, 1000]) +def test_packed_density_independent_of_pair_chunk_size(pair_chunk_size): + """The packed density must not depend on the pair-chunk batching. + + The refactor changes how AO pairs are batched; this guards against a + chunk-boundary bug that would only appear for particular chunk sizes. + """ + hf = cached_backend_hf("pyscf", "cn_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + + g2 = _random_tpdm(refstate) + reference = g2.to_ao_pair_density(refstate, pair_chunk_size=None) + chunked = g2.to_ao_pair_density(refstate, pair_chunk_size=pair_chunk_size) + + assert_allclose(chunked, reference, atol=1e-12) + + +@pytest.mark.parametrize("shell_chunk_size", [1, 2, 1000]) +def test_open_shell_packed_contraction_matches_dense_reference( + shell_chunk_size, +): + """Full packed-density ERI contraction must match the dense path (UHF). + + This anchors the end-to-end two-electron gradient for an open-shell + reference, downstream of the packed density, so the ERI-side contraction + is pinned together with the spin expansion. + """ + hf = cached_backend_hf("pyscf", "cn_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + provider = hf.gradient_provider + + g2 = _random_tpdm(refstate) + g2_ao_1, g2_ao_2 = g2.to_ao_basis() + g2_ao_1, g2_ao_2 = g2_ao_1.to_ndarray(), g2_ao_2.to_ndarray() + + nao = g2_ao_1.shape[0] + dummy = np.zeros((nao, nao)) + full = provider.correlated_gradient(dummy, dummy, g2_ao_1, g2_ao_2) + + pair_density = TwoParticleDensityMatrix.ao_pair_density_from_dense( + g2_ao_1, g2_ao_2 + ) + packed_eri = provider._contract_eri_with_packed_density( + pair_density, shell_chunk_size=shell_chunk_size + ) + + assert_allclose(packed_eri, full.two_electron, atol=1e-10) + + +# --------------------------------------------------------------------------- +# export_block: dense sub-block / spin-subblock tensor view (C++) +# +# These pin the new libadcc Tensor.export_block primitive that the direct +# gradient transform will consume. export_block must agree with slicing the +# full dense to_ndarray() export for arbitrary ranges and, in particular, for +# spin sub-blocks (the alpha/beta partition along each axis). +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("system", ["h2o_sto3g", "cn_sto3g"]) +@pytest.mark.parametrize("block", [b.oooo, b.oovv, b.ovov]) +def test_export_block_matches_to_ndarray_slices(system, block): + hf = cached_backend_hf("pyscf", system, conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + + g2 = _random_tpdm(refstate, blocks=(block,)) + tensor = g2[block] + full = tensor.to_ndarray() + shape = full.shape + + # Full range must reproduce to_ndarray exactly. + eb_full = tensor.export_block([0] * len(shape), list(shape)) + assert_allclose(eb_full, full, atol=1e-12) + + # A handful of random sub-ranges. + rng = np.random.default_rng(42) + for _ in range(8): + start = [int(rng.integers(0, s)) for s in shape] + end = [int(rng.integers(st + 1, s + 1)) for st, s in zip(start, shape)] + eb = tensor.export_block(start, end) + ref_slice = full[tuple(slice(a, b_) for a, b_ in zip(start, end))] + assert_allclose(eb, ref_slice, atol=1e-12) + + +def test_export_block_spin_subblocks_match(): + """Every (alpha/beta)^4 spin sub-block must match the to_ndarray slice. + + This is the key use case for the direct gradient transform: extracting a + single spin sub-block (e.g. abab) of a 4-index TPDM block without + materialising the full dense tensor. Uses an open-shell reference so the + alpha and beta ranges differ in size. + """ + from itertools import product + from adcc.MoSpaces import split_spaces + + hf = cached_backend_hf("pyscf", "cn_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + ms = refstate.mospaces + + for block in (b.oooo, b.oovv, b.ovov): + g2 = _random_tpdm(refstate, blocks=(block,)) + tensor = g2[block] + full = tensor.to_ndarray() + spaces = split_spaces(block) + splits = [(ms.n_orbs_alpha(sp), ms.n_orbs(sp)) for sp in spaces] + + for combo in product([0, 1], repeat=4): + start, end = [], [] + for (na, ntot), spin in zip(splits, combo): + start.append(0 if spin == 0 else na) + end.append(na if spin == 0 else ntot) + eb = tensor.export_block(start, end) + ref_slice = full[tuple(slice(a, b_) for a, b_ in zip(start, end))] + assert_allclose(eb, ref_slice, atol=1e-12, + err_msg=f"{block} spin {combo} mismatch") + + +def test_export_block_validates_arguments(): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + g2 = _random_tpdm(refstate, blocks=(b.oovv,)) + tensor = g2[b.oovv] + shape = tensor.to_ndarray().shape + + with pytest.raises(Exception): + tensor.export_block([0, 0, 0], list(shape)) # wrong ndim + with pytest.raises(Exception): + # end exceeds shape + tensor.export_block([0, 0, 0, 0], + [shape[0] + 1, shape[1], shape[2], shape[3]]) + + # Zero-size ranges must return correctly shaped empty arrays. + empty_all = tensor.export_block([0, 0, 0, 0], [0, 0, 0, 0]) + assert empty_all.shape == (0, 0, 0, 0) + empty_one = tensor.export_block([0, 0, 0, 0], + [0, shape[1], shape[2], shape[3]]) + assert empty_one.shape == (0, shape[1], shape[2], shape[3]) + assert_allclose(empty_one, np.empty((0, shape[1], shape[2], shape[3])), + atol=1e-12) + + +# --------------------------------------------------------------------------- +# Out-of-core (HDF5) packed-density storage and API validation +# +# These pin the new eri_pair_density_storage="hdf5" path (equivalence with the +# in-memory path and scratch-file cleanup) and the user-facing validation +# contracts of the new keyword arguments. +# --------------------------------------------------------------------------- + + +def test_direct_gradient_hdf5_matches_memory(tmp_path): + """The out-of-core HDF5 storage path must match the in-memory path. + + Exercises the h5py-dataset writes inside ``to_ao_pair_density`` (the + ``out[...] = 0`` reset and the ``out[:, :, start:stop] += chunk`` + read-modify-write), which behave differently for an HDF5 dataset than for + a plain ndarray. + """ + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + mp = adcc.LazyMp(adcc.ReferenceState(hf)) + + memory = adcc.nuclear_gradient( + mp, eri_contraction="direct", eri_pair_density_storage="memory" + ) + hdf5 = adcc.nuclear_gradient( + mp, eri_contraction="direct", eri_pair_density_storage="hdf5", + eri_pair_chunk_size=3 + ) + + assert_allclose(hdf5.total, memory.total, atol=1e-10) + assert_allclose(hdf5.components.two_electron, + memory.components.two_electron, atol=1e-10) + + +def test_direct_gradient_hdf5_removes_scratch_file(tmp_path): + """The temporary HDF5 scratch file must be removed after contraction.""" + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + mp = adcc.LazyMp(adcc.ReferenceState(hf)) + provider = hf.gradient_provider + + g1_ao, w_ao, g2_total = _direct_gradient_inputs(mp) + before = set(glob.glob(os.path.join(str(tmp_path), "adcc_pyscf_gradient_*"))) + provider.correlated_gradient_direct( + g1_ao, w_ao, g2_total, refstate=adcc.ReferenceState(hf), + pair_density_storage="hdf5", scratch_directory=str(tmp_path) + ) + after = set(glob.glob(os.path.join(str(tmp_path), "adcc_pyscf_gradient_*"))) + assert before == after, "scratch HDF5 file was not cleaned up" + + +def test_direct_gradient_hdf5_default_pair_chunk_is_bounded(): + """The hdf5 path must not silently use the full npair working buffer. + + Out-of-core storage only offloads the output density; if the in-memory + working chunk defaulted to the full ``npair`` it would defeat the memory + bound. The hdf5 path therefore selects a bounded default chunk size. + """ + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + mp = adcc.LazyMp(adcc.ReferenceState(hf)) + provider = hf.gradient_provider + g1_ao, w_ao, g2_total = _direct_gradient_inputs(mp) + + nao = provider.mol.nao_nr() + npair = nao * (nao + 1) // 2 + + captured = {} + original = g2_total.to_ao_pair_density + + def spy(refstate_or_coefficients=None, pair_chunk_size=None, out=None): + captured["pair_chunk_size"] = pair_chunk_size + return original(refstate_or_coefficients, + pair_chunk_size=pair_chunk_size, out=out) + + g2_total.to_ao_pair_density = spy + provider.correlated_gradient_direct( + g1_ao, w_ao, g2_total, refstate=adcc.ReferenceState(hf), + pair_density_storage="hdf5" + ) + assert captured["pair_chunk_size"] is not None + assert captured["pair_chunk_size"] <= npair + + +def test_nuclear_gradient_rejects_invalid_eri_contraction(): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + mp = adcc.LazyMp(adcc.ReferenceState(hf)) + with pytest.raises(ValueError, match="Invalid eri_contraction"): + adcc.nuclear_gradient(mp, eri_contraction="bogus") + + +def test_correlated_gradient_direct_rejects_invalid_storage(): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + mp = adcc.LazyMp(adcc.ReferenceState(hf)) + provider = hf.gradient_provider + g1_ao, w_ao, g2_total = _direct_gradient_inputs(mp) + with pytest.raises(ValueError, match="pair_density_storage"): + provider.correlated_gradient_direct( + g1_ao, w_ao, g2_total, refstate=adcc.ReferenceState(hf), + pair_density_storage="bogus" + ) + + +def test_contract_eri_rejects_non_positive_shell_chunk(): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + provider = hf.gradient_provider + nao = provider.mol.nao_nr() + npair = nao * (nao + 1) // 2 + pair_density = np.zeros((nao, nao, npair)) + with pytest.raises(ValueError, match="shell_chunk_size"): + provider._contract_eri_with_packed_density( + pair_density, shell_chunk_size=0 + ) + + +def test_to_ao_pair_density_rejects_non_positive_pair_chunk(): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + g2 = _random_tpdm(refstate) + with pytest.raises(ValueError, match="pair_chunk_size"): + g2.to_ao_pair_density(refstate, pair_chunk_size=0) + + +def test_to_ao_pair_density_rejects_bad_output_shape(): + hf = cached_backend_hf("pyscf", "h2o_sto3g", conv_tol=1e-11) + refstate = adcc.ReferenceState(hf) + g2 = _random_tpdm(refstate) + with pytest.raises(ValueError, match="Invalid output shape"): + g2.to_ao_pair_density(refstate, out=np.zeros((1, 1, 1))) diff --git a/adcc/tests/testdata_cache.py b/adcc/tests/testdata_cache.py index a1effde56..167f39696 100644 --- a/adcc/tests/testdata_cache.py +++ b/adcc/tests/testdata_cache.py @@ -307,3 +307,4 @@ def _import_hook(data: dict): psi4_data = read_json_data("psi4_data.json") pyscf_data = read_json_data("pyscf_data.json") tmole_data = read_json_data("tmole_data.json") +gradient_data = read_json_data("gradient_data.json") diff --git a/docs/gradients.rst b/docs/gradients.rst new file mode 100644 index 000000000..6acb59241 --- /dev/null +++ b/docs/gradients.rst @@ -0,0 +1,324 @@ +:github_url: https://github.com/adc-connect/adcc/blob/master/docs/gradients.rst + +.. _nuclear-gradients: + +Nuclear gradients +================= + +.. note:: + Analytic nuclear gradients in adcc are currently experimental and + available for the PySCF backend. + +adcc can evaluate analytic nuclear gradients of the ground-state (MP2) and +excited-state (ADC) energies. This page first shows how to request a gradient, +then documents the supported methods and finally describes the underlying +theory and the memory-bounded contraction algorithm. + +Usage +----- + +A nuclear gradient is requested via :py:func:`adcc.nuclear_gradient`. For an +excited state, pass the corresponding :py:class:`adcc.Excitation` object (e.g. +one element of ``state.excitations``): + +.. code-block:: python + + from pyscf import gto, scf + import adcc + + mol = gto.M(atom="O 0 0 0; H 0 0 1.795; H 1.693 0 -0.599", + basis="cc-pvdz", unit="Bohr") + scfres = scf.RHF(mol) + scfres.kernel() + + state = adcc.adc2(scfres, n_singlets=3) + grad = adcc.nuclear_gradient(state.excitations[0]) + print(grad.total) + +For a ground-state (MP2) gradient, pass the :py:class:`adcc.LazyMp` object +instead: + +.. code-block:: python + + mp = adcc.LazyMp(adcc.ReferenceState(scfres)) + grad = adcc.nuclear_gradient(mp) + +The returned object exposes the total gradient via ``grad.total`` as well as +the individual one- and two-electron contributions. + +Geometry optimisation scanners +------------------------------ + +For geometry optimisers it is useful to expose the whole PySCF/adcc/gradient +loop as a single callable. :py:class:`adcc.NuclearGradientScanner` takes a +configured PySCF SCF object as its template and accepts Cartesian coordinates in +Bohr. It returns ``(energy, gradient)`` in Hartree and Hartree/Bohr: + +.. code-block:: python + + scfres = scf.RHF(mol) + scfres.conv_tol = 1e-11 + scfres.conv_tol_grad = 1e-9 + + scanner = adcc.NuclearGradientScanner( + scfres, + method="adc2", + state_index=0, + n_singlets=3, + conv_tol=1e-7, # forwarded unchanged to adcc.run_adc + ) + + energy, gradient = scanner(mol.atom_coords()) + +The scanner uses PySCF's scanner machinery to rerun the SCF calculation at each +new geometry with the original SCF object settings. This keeps the SCF +interface on the PySCF object, where users already configure basis, charge, +spin/reference type, symmetry behaviour and convergence options. For geomeTRIC +optimisations it is usually safest to construct the PySCF molecule with +``symmetry=False``. PySCF's scanner also provides the SCF guess continuity +between geometry steps. For excited states the scanner stores the previous +:class:`adcc.ExcitedStates` and selected :class:`adcc.Excitation`; by default it +follows the state by comparing AO-basis transition and state-difference densities +across geometries using PySCF AO cross-overlap integrals. A fixed-index mode is +available via ``follow="index"`` for debugging or well-separated states. + +The same object also implements geomeTRIC's custom-engine ``calc_new`` protocol: + +.. code-block:: python + + result = scanner.calc_new(mol.atom_coords().ravel()) + # result == {"energy": energy, "gradient": gradient.ravel()} + +geomeTRIC remains an optional dependency; the plain scanner only requires the +PySCF backend. Minimum-energy crossing point workflows can be built from two +scanner targets because geomeTRIC's penalty-constrained formulation only needs +the two state energies and gradients, not derivative couplings. + +The two-electron term can be evaluated with two different strategies, selected +through the ``eri_contraction`` keyword (see +:ref:`gradients-eri-contraction` below). For the PySCF backend the +memory-bounded ``"direct"`` strategy is chosen automatically; it can be +requested or overridden explicitly together with its tuning knobs: + +.. code-block:: python + + grad = adcc.nuclear_gradient( + state.excitations[0], + eri_contraction="direct", # or "full_ao" + eri_pair_density_storage="hdf5", # out-of-core packed density + eri_pair_chunk_size=256, # bound the working buffer + ) + +.. _gradients-supported-methods: + +Supported methods +----------------- + +Analytic gradients are available for the following methods (for restricted and +unrestricted references, unless noted otherwise): + +.. list-table:: + :header-rows: 1 + :widths: 30 20 50 + + * - Method + - Gradient + - Notes + * - MP2 (ground state) + - ✓ + - via :py:class:`adcc.LazyMp` + * - ADC(0), ADC(1) + - ✓ + - + * - ADC(2), ADC(2)-x + - ✓ + - + * - ADC(3) + - ✓ + - + * - CVS-ADC(0), CVS-ADC(1) + - ✓ + - core-valence separation :cite:`Brumboiu2021` + * - CVS-ADC(2), CVS-ADC(2)-x + - ✓ + - core-valence separation :cite:`Brumboiu2021` + * - CVS-ADC(3) + - ✗ + - not available + +Requesting a gradient for an unsupported method raises a +:py:exc:`NotImplementedError`. + +.. _gradients-theory: + +Theory and algorithm +-------------------- + +Gradient Lagrangian +~~~~~~~~~~~~~~~~~~~~ + +Analytic energy gradients for non-variational methods such as MP and ADC are +conveniently formulated within a Lagrangian formalism, a well-established +approach in quantum chemistry. The working equations implemented in adcc follow +the derivation of Rehn and Dreuw :cite:`Rehn2019`; the extension to +core-excited states within the core-valence separation is due to +Brumboiu *et al.* :cite:`Brumboiu2021`. + +The gradient of the total energy with respect to a nuclear coordinate :math:`x` +separates into a one-electron and a two-electron contribution, + +.. math:: + :label: eqn:gradient_total + + \frac{\mathrm{d} E}{\mathrm{d} x} + = \sum_{\mu\nu} \gamma_{\mu\nu}\, + \frac{\partial h_{\mu\nu}}{\partial x} + - \sum_{\mu\nu} W_{\mu\nu}\, + \frac{\partial S_{\mu\nu}}{\partial x} + + \sum_{\mu\nu\kappa\lambda} \Gamma_{\mu\nu\kappa\lambda}\, + \frac{\partial (\mu\nu|\kappa\lambda)}{\partial x}, + +where :math:`\mu,\nu,\kappa,\lambda` are atomic-orbital (AO) indices, +:math:`\gamma` is the relaxed one-particle density matrix (OPDM), +:math:`W` the energy-weighted density matrix, +:math:`S` the overlap matrix, +:math:`h` the one-electron Hamiltonian, and +:math:`\Gamma` the relaxed two-particle density matrix (TPDM). +The one-electron terms are cheap; the bottleneck is the contraction of the +TPDM with the derivative electron-repulsion integrals (ERIs) +:math:`\partial (\mu\nu|\kappa\lambda) / \partial x`. + +.. _gradients-eri-contraction: + +Two-electron contraction strategies +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +adcc offers two strategies for the two-electron term, selected through the +``eri_contraction`` keyword of :py:func:`adcc.nuclear_gradient`. + +``full_ao`` + The reference implementation. The MO-basis TPDM is transformed to two full + AO-basis rank-4 tensors :math:`g^{(1)}` and :math:`g^{(2)}`, which are then + contracted with the full derivative ERI tensor. Both the AO TPDMs and the + derivative integrals scale as :math:`\mathcal{O}(N_\text{AO}^4)` in memory + (the derivative ERIs even as :math:`3\,N_\text{AO}^4`), so this path becomes + impractical already for moderate basis sets. It is kept as a + validation/debug fallback. + +``direct`` + The memory-bounded production path (default for PySCF). It never forms the + full AO TPDMs nor the full derivative ERIs. Instead, the TPDM is transformed + block-by-block directly into a *packed* AO-pair effective density, which is + streamed against shell-batched derivative integrals. Peak memory is bounded + by a single AO-pair chunk rather than by :math:`N_\text{AO}^4`. + +The remaining subsections describe the ``direct`` algorithm. + +The packed AO-pair effective density +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``full_ao`` path produces two AO TPDMs that are contracted with the ERIs in +chemists' notation as :math:`+(pr|qs)` and :math:`-(ps|qr)`, i.e. + +.. math:: + :label: eqn:ao_tpdm + + g^{(1)}_{pqrs} &= \sum_{ijkl} C_{ip} C_{jq} \Gamma_{ijkl} C_{kr} C_{ls}, \\ + g^{(2)}_{pqrs} &= \sum_{ijkl} C_{ip} C_{jq} \Gamma_{ijkl} C_{ks} C_{lr}, + +summed over the spin cases that survive for a restricted reference +(:math:`g^{(1)}`: ``aaaa``, ``bbbb``, ``abab``, ``baba``; +:math:`g^{(2)}`: ``aaaa``, ``bbbb``, ``abba``, ``baab``). +Here :math:`i,j,k,l` are spin-orbital MO indices and :math:`C` are the +spin-blocked MO coefficients. + +The ``direct`` path instead builds the combined effective density only in the +layout actually required for the integral contraction. Defining the effective +density as the difference of the two orderings, + +.. math:: + :label: eqn:effective_density + + D_{pr,qs} = g^{(1)}_{pqrs} - g^{(2)}_{pqrs}, + +the two leading AO indices :math:`p,r` are kept dense while the trailing pair +:math:`q,s` is stored in lower-triangular packed form. With the symmetric pair +identified by :math:`(\max(q,s), \min(q,s))` and mapped to the packed index + +.. math:: + :label: eqn:pair_index + + m(q,s) = \frac{\max(q,s)\,[\max(q,s)+1]}{2} + \min(q,s), + +the packed density has shape +:math:`(N_\text{AO},\, N_\text{AO},\, N_\text{pair})` with +:math:`N_\text{pair} = N_\text{AO}(N_\text{AO}+1)/2`. +This is the conventional symmetric-pair packing used by derivative-ERI +backends that expose only one integral per symmetric ket pair (e.g. PySCF +``aosym="s2kl"``). Because only one entry is stored for an off-diagonal pair +:math:`q \neq s`, the packed entry accumulates both contributing orderings. + +Block-wise streaming transform +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The key observation is that each MO axis of the spin-orbital TPDM splits into a +leading alpha range :math:`[0, n_\alpha)` and a trailing beta range +:math:`[n_\alpha, n_\text{orb})`, and the spin-blocked coefficients are +non-zero only on the matching range. For every required spin case adcc +therefore extracts just the relevant rank-4 spin sub-block of a TPDM block +(via the ``export_block`` method of :py:class:`adcc.Tensor`) and contracts it +with the compacted (non-zero-row) coefficients, never densifying the full +zero-padded block. + +For a chunk of AO pairs the central contraction reads, for the direct ordering, + +.. math:: + :label: eqn:direct_transform + + D_{pr,m} \mathrel{+}= \sum_{ijkl} + C_{ip}\, C_{kr}\, \Gamma_{ijkl}\, + \left( C_{j\,q(m)}\, C_{l\,s(m)} \right), + +and analogously for the exchange ordering, where :math:`q(m)` and :math:`s(m)` +are the AO indices of packed pair :math:`m`. The intermediate +:math:`C_{j\,q(m)} C_{l\,s(m)}` is built once per chunk and contracted with the +sub-block, so the working memory is bounded by the chunk size +(the number of AO pairs processed at once) rather than by the full +:math:`N_\text{pair}`. + +Contraction with derivative integrals +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The packed density is then contracted with the shell-batched derivative ERIs. +PySCF only packs the ket pair (``aosym="s2kl"``), so for a derivative integral +block :math:`E_{ab,qs} = \partial (ab|qs)/\partial x` restricted to the AO +shells of one atom, the symmetrised contraction effectively evaluates + +.. math:: + :label: eqn:eri_contraction + + \frac{\partial E_\text{2e}}{\partial x_A} = \sum_{ab}\sum_{q\le s} + E^{A}_{ab,qs}\, + \left( D_{ab,qs} + D_{ba,sq} + D_{qs,ab} + D_{sq,ba} \right), + +with the four terms accounting for the bra/ket permutational symmetry of the +ERIs. Keeping the shell slices atom-local ensures that atom AO ranges and shell +ranges are never mixed accidentally. + +Memory and storage options +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``direct`` path exposes two knobs on :py:func:`adcc.nuclear_gradient`: + +* ``eri_pair_chunk_size`` -- the number of AO pairs transformed at once. This + bounds the in-memory working buffer of the transform. When left unset it + defaults to the full :math:`N_\text{pair}` for in-memory storage, and to a + smaller bounded value for out-of-core storage. +* ``eri_pair_density_storage`` -- where the packed AO-pair density is kept: + + - ``"memory"`` (default): a single in-memory + :math:`(N_\text{AO}, N_\text{AO}, N_\text{pair})` array. + - ``"hdf5"``: a temporary HDF5 file under the scratch directory (or + ``PYSCF_TMPDIR`` / the system temporary directory), removed after the + contraction. This bounds the packed *output* density; the transform's + working buffer is bounded separately by ``eri_pair_chunk_size``. diff --git a/docs/index.rst b/docs/index.rst index c0e6930da..941f28498 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -90,6 +90,7 @@ Contents installation calculations theory + gradients benchmarks topics diff --git a/docs/ref.bib b/docs/ref.bib index ded2396c5..d23e02594 100644 --- a/docs/ref.bib +++ b/docs/ref.bib @@ -330,6 +330,30 @@ @article{Marefat2018 year = {2018} } +@article{Brumboiu2021, + author = {Brumboiu, Iulia Emilia and Rehn, Dirk R. and Dreuw, Andreas and Rhee, Young Min and Norman, Patrick}, + doi = {10.1063/5.0058221}, + issn = {0021-9606}, + journal = {J. Chem. Phys.}, + number = {4}, + pages = {044106}, + title = {{Analytical gradients for core-excited states in the algebraic diagrammatic construction (ADC) framework}}, + volume = {155}, + year = {2021} +} + +@article{Rehn2019, + author = {Rehn, Dirk R. and Dreuw, Andreas}, + doi = {10.1063/1.5085117}, + issn = {0021-9606}, + journal = {J. Chem. Phys.}, + number = {17}, + pages = {174110}, + title = {{Analytic nuclear gradients of the algebraic-diagrammatic construction scheme for the polarization propagator up to third order of perturbation theory}}, + volume = {150}, + year = {2019} +} + @article{Scheurer2019, title={CPPE: An open-source C++ and Python library for polarizable embedding}, author={Scheurer, Maximilian and Reinholdt, Peter and Kjellgren, Erik Rosendahl and Olsen, Jógvan Magnus Haugaard and Dreuw, Andreas and Kongsted, Jacob}, diff --git a/examples/optimization/pyscf_adcc_geoopt.py b/examples/optimization/pyscf_adcc_geoopt.py new file mode 100644 index 000000000..7a9a477c9 --- /dev/null +++ b/examples/optimization/pyscf_adcc_geoopt.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +"""Ground-state MP2 geometry optimisation with adcc gradients. + +The scanner is built from a configured PySCF SCF object. PySCF owns all SCF +settings and adcc runs MP2 plus the analytic gradient at each geometry. +""" + +import adcc +from pyscf import gto, scf +from pyscf.geomopt import as_pyscf_method, geometric_solver + + +mol = gto.M( + atom=""" + O 0.000000000000 0.000000000000 0.000000000000 + H 0.000000000000 0.000000000000 1.795239827225 + H 1.693194615993 0.000000000000 -0.599043184453 + """, + basis="sto-3g", + unit="Bohr", + symmetry=False, + verbose=0, +) + +mf = scf.RHF(mol) +mf.conv_tol = 1e-11 +mf.conv_tol_grad = 1e-9 + +scanner = adcc.NuclearGradientScanner(mf, method="mp2") + + +def energy_and_gradient(mol_at_step): + """PySCF geomopt callback: return energy and gradient for a Mole.""" + return scanner(mol_at_step.atom_coords(unit="Bohr")) + + +if __name__ == "__main__": + method = as_pyscf_method(mol, energy_and_gradient) + mol_eq = geometric_solver.optimize(method, maxsteps=20) + + print("\nOptimised geometry (Bohr):") + for symbol, xyz in zip(scanner.atom_symbols, mol_eq.atom_coords(unit="Bohr")): + print(f"{symbol:2s} {xyz[0]:16.10f} {xyz[1]:16.10f} {xyz[2]:16.10f}") diff --git a/examples/optimization/pyscf_adcc_geoopt_excited.py b/examples/optimization/pyscf_adcc_geoopt_excited.py new file mode 100644 index 000000000..d187f96eb --- /dev/null +++ b/examples/optimization/pyscf_adcc_geoopt_excited.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +"""Excited-state geometry optimisation with adcc gradients. + +This example reproduces the s-trans-butadiene setup used as a geometry +optimisation showcase by Rehn & Dreuw (*J. Chem. Phys.* 150, 174110, 2019). +It first relaxes the molecule on the MP2 ground-state surface and then uses that +geometry as the starting point for an ADC(2) excited-state optimisation. The +:class:`adcc.NuclearGradientScanner` is built from a configured PySCF SCF +object, so PySCF owns all SCF settings while adcc keyword arguments are passed +unchanged to :func:`adcc.run_adc`. + +geomeTRIC is optional; install it separately, e.g. ``pip install geometric`` or +``pip install adcc[geomopt]`` once the optional extra is available. +""" + +import adcc +from pyscf import gto, scf +from pyscf.geomopt import as_pyscf_method, geometric_solver + + +def make_scf(mol): + """Build the PySCF reference used as scanner template.""" + mf = scf.RHF(mol) + mf.conv_tol = 1e-10 + mf.conv_tol_grad = 1e-6 + return mf + + +def print_geometry(title, mol): + print(f"\n{title} (Bohr):") + for i, xyz in enumerate(mol.atom_coords(unit="Bohr")): + print(f"{mol.atom_symbol(i):2s} " + f"{xyz[0]:16.10f} {xyz[1]:16.10f} {xyz[2]:16.10f}") + + +def make_energy_gradient(scanner): + """PySCF geomopt callback: return energy and gradient for a Mole.""" + def energy_and_gradient(mol_at_step): + energy, gradient = scanner(mol_at_step.atom_coords(unit="Bohr")) + excitation = scanner.last_excitation + state_info = "" + if excitation is not None: + state_info = ( + f", state = {excitation.index}, " + f"omega = {excitation.excitation_energy:.8f} Eh" + ) + print( + f"E = {energy:.12f} Eh, |g| = " + f"{float((gradient ** 2).sum() ** 0.5):.6e} Eh/Bohr" + f"{state_info}" + ) + return energy, gradient + return energy_and_gradient + + +# Use Bohr coordinates and keep molecular symmetry disabled so the optimizer's +# Cartesian frame and the gradient frame stay identical during the scan. +# s-trans-butadiene (C4H6), planar all-trans conformation. +mol = gto.M( + atom=""" +C -3.4705405625 0.4297233578 -0.0000000000 +C -1.1911517402 -0.6904425316 0.0000000000 +C 1.1911517402 0.6904425315 0.0000000000 +C 3.4705405625 -0.4297233578 -0.0000000000 +H -3.6579517613 2.4744185046 -0.0000000000 +H -5.2070544325 -0.6589625683 0.0000000000 +H -1.0698622893 -2.7464869632 0.0000000000 +H 1.0698622893 2.7464869632 0.0000000000 +H 3.6579517613 -2.4744185046 -0.0000000000 +H 5.2070544325 0.6589625682 0.0000000000 + """, + basis="6-31G*", + unit="Bohr", + symmetry=False, + verbose=0, +) + + +if __name__ == "__main__": + print_geometry("Initial geometry", mol) + + # Stage 1: relax the ground-state structure at MP2 level. This provides a + # much better starting point for the subsequent state-specific optimisation + # than the arbitrary input geometry. + mp2_scanner = adcc.NuclearGradientScanner( + make_scf(mol), + method="mp2", + gradient_kwargs={"conv_tol": 1e-7}, + ) + mp2_method = as_pyscf_method(mol, make_energy_gradient(mp2_scanner)) + mol_mp2 = geometric_solver.optimize( + mp2_method, + maxsteps=20, + convergence_grms=1e-3, + convergence_gmax=1.5e-3, + ) + print_geometry("MP2-optimised geometry", mol_mp2) + + # Stage 2: optimise the lowest singlet ADC(2) state. Requesting a few + # singlet roots lets the scanner compare candidates and follow the same + # physical state if root ordering changes along the optimisation. + adc_scanner = adcc.NuclearGradientScanner( + make_scf(mol_mp2), + method="adc2", + state_index=0, + n_singlets=3, + follow="overlap", + conv_tol=1e-6, + gradient_kwargs={"conv_tol": 1e-7}, + ) + adc_method = as_pyscf_method(mol_mp2, make_energy_gradient(adc_scanner)) + mol_adc = geometric_solver.optimize( + adc_method, + maxsteps=20, + convergence_grms=1e-3, + convergence_gmax=1.5e-3, + ) + print_geometry("ADC(2)-optimised demonstration geometry", mol_adc) diff --git a/libadcc.pyi b/libadcc.pyi index 9383c04a2..2cc591824 100644 --- a/libadcc.pyi +++ b/libadcc.pyi @@ -1,5 +1,7 @@ from __future__ import annotations +import collections.abc import numpy +import numpy.typing import pybind11_stubgen.typing_ext import typing @@ -713,6 +715,14 @@ class Tensor: """ Ensure the tensor to be fully evaluated and resilient in memory. Usually happens automatically when needed. Might be useful for fine-tuning, however. """ + def export_block( + self, + start: collections.abc.Sequence[typing.SupportsInt | typing.SupportsIndex], + end: collections.abc.Sequence[typing.SupportsInt | typing.SupportsIndex], + ) -> numpy.typing.NDArray[numpy.float64]: + """ + Export a dense sub-block of the tensor (half-open per-axis range [start, end)) to a np::ndarray, respecting the full tensor symmetry, without materialising the full dense tensor. + """ def is_allowed(self, arg0: tuple) -> bool: """ Is a particular index allowed by symmetry diff --git a/libadcc_src/Tensor.hh b/libadcc_src/Tensor.hh index 75b0e0cea..8b1d9bba4 100644 --- a/libadcc_src/Tensor.hh +++ b/libadcc_src/Tensor.hh @@ -322,6 +322,31 @@ class Tensor { export_to(output.data(), output.size()); } + /** Extract a dense sub-block of the tensor into plain memory. + * + * The sub-block is described by a half-open per-axis range + * ``[start[i], end[i])``. This allows extracting e.g. a single spin + * sub-block (such as the alpha-beta-alpha-beta block of a 4-index tensor) + * without materialising the full dense tensor via ``export_to``. + * + * The full tensor symmetry is respected: requested elements that map onto a + * non-canonical or zero block are filled with the symmetry-implied value + * (including sign/scalar transformations), exactly as ``get_element`` would + * return them. + * + * \param start Per-axis start index (inclusive), one entry per dimension. + * \param end Per-axis end index (exclusive), one entry per dimension. + * \param memptr Dense memory location to fill. At least + * ``prod(end[i] - start[i])`` elements are assumed. + * \param size Size of the dense memory provided. + * + * \note The data is stored in row-major (C-like) format with respect to the + * requested sub-block extents ``end[i] - start[i]``. + */ + virtual void export_block(const std::vector& start, + const std::vector& end, scalar_type* memptr, + size_t size) const = 0; + /** Import the tensor from plain memory provided by the given * pointer. The memory will be copied and all existing data overwritten. * If symmetry_check is true, the process will check that the data has the diff --git a/libadcc_src/TensorImpl.cc b/libadcc_src/TensorImpl.cc index 6fde3f27a..c04368a29 100644 --- a/libadcc_src/TensorImpl.cc +++ b/libadcc_src/TensorImpl.cc @@ -1375,6 +1375,133 @@ void TensorImpl::export_to(scalar_type* memptr, size_t size) const { lt::btod_export(*libtensor_ptr()).perform(memptr); } +template +void TensorImpl::export_block(const std::vector& start, + const std::vector& end, scalar_type* memptr, + size_t size) const { + if (start.size() != N || end.size() != N) { + throw invalid_argument("start and end must each have exactly " + std::to_string(N) + + " entries."); + } + + // Compute requested sub-block extents and the row-major output strides. + std::array ext; + std::array out_strides; + size_t total = 1; + for (size_t i = 0; i < N; ++i) { + if (end[i] < start[i]) { + throw invalid_argument("end must not be smaller than start along any axis."); + } + if (end[i] > shape()[i]) { + throw invalid_argument("end exceeds the tensor shape along axis " + + std::to_string(i) + "."); + } + ext[i] = end[i] - start[i]; + total *= ext[i]; + } + if (size != total) { + throw invalid_argument("The memory provided (== " + std::to_string(size) + + ") does not agree with the requested sub-block size (== " + + std::to_string(total) + ")."); + } + if (total == 0) return; + + out_strides[N - 1] = 1; + for (size_t i = N - 1; i > 0; --i) { + out_strides[i - 1] = out_strides[i] * ext[i]; + } + + // Default to zero: requested elements falling onto symmetry-forbidden or zero + // blocks must come out as zero. + std::fill(memptr, memptr + total, scalar_type(0)); + + lt::block_tensor_ctrl ctrl(*libtensor_ptr()); + const lt::block_index_space& bis = libtensor_ptr()->get_bis(); + lt::dimensions bidims(bis.get_block_index_dims()); + + // This mirrors btod_export (the routine behind export_to / to_ndarray): for + // every non-zero canonical block, scatter the block data to all blocks of its + // symmetry orbit, applying the orbit permutation and scalar transform. The + // difference here is that we only write elements that fall inside the + // requested [start, end) window, and into a compact sub-block buffer. + lt::orbit_list orbitlist(ctrl.req_const_symmetry()); + lt::index canon_idx; + for (auto oit = orbitlist.begin(); oit != orbitlist.end(); ++oit) { + orbitlist.get_index(oit, canon_idx); + if (ctrl.req_is_zero_block(canon_idx)) continue; + lt::orbit obit(ctrl.req_const_symmetry(), canon_idx); + + auto& cblock = ctrl.req_const_block(canon_idx); + { + lt::dense_tensor_rd_ctrl blkctrl(cblock); + const scalar_type* sptr = blkctrl.req_const_dataptr(); + const lt::dimensions cblk_dims = cblock.get_dims(); + + // Walk every block of the orbit (each is a destination block in the + // full dense tensor). + for (auto j = obit.begin(); j != obit.end(); ++j) { + lt::abs_index aidx(obit.get_abs_index(j), bidims); + const lt::tensor_transf& tr = obit.get_transf(j); + lt::index dst_start(bis.get_block_start(aidx.get_index())); + lt::dimensions dst_dims(bis.get_block_dims(aidx.get_index())); + + // Does this destination block intersect the requested window? + std::array lo, hi; + bool empty = false; + for (size_t i = 0; i < N; ++i) { + const size_t blo = dst_start[i]; + const size_t bhi = dst_start[i] + dst_dims[i]; + lo[i] = std::max(blo, start[i]); + hi[i] = std::min(bhi, end[i]); + if (lo[i] >= hi[i]) { + empty = true; + break; + } + } + if (empty) continue; + + // Inverse permutation maps a destination (output) index back to the + // source (canonical block) index, exactly as btod_export::copy_block. + lt::permutation inv_perm(tr.get_perm()); + inv_perm.invert(); + const scalar_type coeff = tr.get_scalar_tr().get_coeff(); + + std::array isect_ext; + size_t isect_total = 1; + for (size_t i = 0; i < N; ++i) { + isect_ext[i] = hi[i] - lo[i]; + isect_total *= isect_ext[i]; + } + + for (size_t flat = 0; flat < isect_total; ++flat) { + // Decode flat -> per-axis destination index within the intersection. + lt::index dst_in_block; // index within the destination block + size_t out_abs = 0; + for (size_t i = 0; i < N; ++i) { + size_t stride = 1; + for (size_t k = i + 1; k < N; ++k) stride *= isect_ext[k]; + const size_t off = (flat / stride) % isect_ext[i]; + const size_t abs_ix = lo[i] + off; // absolute tensor index on axis i + dst_in_block[i] = abs_ix - dst_start[i]; + out_abs += (abs_ix - start[i]) * out_strides[i]; + } + + // Map destination in-block index to the source (canonical) in-block + // index via the inverse permutation. + lt::index src_in_block(dst_in_block); + src_in_block.permute(inv_perm); + const size_t src_abs = + lt::abs_index(src_in_block, cblk_dims).get_abs_index(); + memptr[out_abs] = coeff * sptr[src_abs]; + } + } + + blkctrl.ret_const_dataptr(sptr); + } + ctrl.ret_const_block(canon_idx); + } +} + template void TensorImpl::import_from(const scalar_type* memptr, size_t size, scalar_type tolerance, bool symmetry_check) { diff --git a/libadcc_src/TensorImpl.hh b/libadcc_src/TensorImpl.hh index 6d3391b67..a53624700 100644 --- a/libadcc_src/TensorImpl.hh +++ b/libadcc_src/TensorImpl.hh @@ -100,6 +100,8 @@ class TensorImpl : public Tensor { size_t n, bool unique_by_symmetry) const override; void export_to(scalar_type* memptr, size_t size) const override; + void export_block(const std::vector& start, const std::vector& end, + scalar_type* memptr, size_t size) const override; void import_from(const scalar_type* memptr, size_t size, scalar_type tolerance, bool symmetry_check) override; virtual void import_from( diff --git a/libadcc_src/pyiface/export_Tensor.cc b/libadcc_src/pyiface/export_Tensor.cc index ecff8ac73..422e8f48d 100644 --- a/libadcc_src/pyiface/export_Tensor.cc +++ b/libadcc_src/pyiface/export_Tensor.cc @@ -112,6 +112,27 @@ static py::array_t Tensor_to_ndarray(const Tensor& self) { return res; } +static py::array_t Tensor_export_block(const Tensor& self, + std::vector start, + std::vector end) { + const size_t ndim = self.ndim(); + if (start.size() != ndim || end.size() != ndim) { + throw invalid_argument("start and end must each have one entry per tensor axis."); + } + std::vector ext(ndim); + size_t total = 1; + for (size_t i = 0; i < ndim; ++i) { + if (end[i] < start[i]) { + throw invalid_argument("end must not be smaller than start along any axis."); + } + ext[i] = static_cast(end[i] - start[i]); + total *= (end[i] - start[i]); + } + py::array_t res(ext); + self.export_block(start, end, res.mutable_data(), total); + return res; +} + static ten_ptr Tensor_from_ndarray_tol( ten_ptr self, py::array_t in_array, @@ -471,6 +492,10 @@ void export_Tensor(py::module& m) { .def("antisymmetrise", &Tensor_antisymmetrise_2) .def("to_ndarray", &Tensor_to_ndarray, "Export the tensor data to a standard np::ndarray by making a copy.") + .def("export_block", &Tensor_export_block, "start"_a, "end"_a, + "Export a dense sub-block of the tensor (half-open per-axis range " + "[start, end)) to a np::ndarray, respecting the full tensor " + "symmetry, without materialising the full dense tensor.") .def("set_from_ndarray", &Tensor_from_ndarray, "Set all tensor elements from a standard np::ndarray by making a copy. " "Provide an optional tolerance argument to increase the tolerance for the " diff --git a/pyproject.toml b/pyproject.toml index feb379251..635daee78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "adcc" description = "adcc: Seamlessly connect your host program to ADC" -version = "0.18.0" +version = "0.17.1" dynamic = ["readme"] # handled in setup.py keywords = [ "ADC", "algebraic-diagrammatic", "construction", "excited", "states", @@ -54,6 +54,7 @@ build_docs = [ "sphinx-automodapi", "sphinx-rtd-theme" ] analysis = ["matplotlib >= 3.0", "pandas >= 0.25.0"] +geomopt = ["geometric"] build_stub = ["ruff", "pybind11-stubgen"] [project.urls] @@ -74,7 +75,7 @@ exclude = [ ] [tool.bumpversion] -current_version = "0.18.0" +current_version = "0.17.1" commit = true tag = true files = [