diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 49eeede3..0db3d8ea 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -15,6 +15,7 @@ from ebcc.core.dump import Dump from ebcc.core.logging import ANSI from ebcc.core.precision import astype, types +from ebcc.ham.base import BaseERIs from ebcc.util import _BaseOptions if TYPE_CHECKING: @@ -26,7 +27,7 @@ from ebcc.core.damping import BaseDamping from ebcc.core.logging import Logger - from ebcc.ham.base import BaseElectronBoson, BaseERIs, BaseFock + from ebcc.ham.base import BaseElectronBoson, BaseFock from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace @@ -891,16 +892,21 @@ def init_space(self) -> SpaceType: """ pass - @abstractmethod def get_fock(self) -> BaseFock: """Get the Fock matrix. Returns: Fock matrix. """ - pass + return self.Fock( + self.mf, + space=(self.space, self.space), + mo_coeff=(self.mo_coeff, self.mo_coeff), + g=self.g, + shift=self.options.shift, + xi=self.xi if self.boson_ansatz else None, + ) - @abstractmethod def get_eris(self, eris: Optional[ERIsInputType] = None) -> BaseERIs: """Get the electron repulsion integrals. @@ -910,7 +916,23 @@ def get_eris(self, eris: Optional[ERIsInputType] = None) -> BaseERIs: Returns: Electron repulsion integrals. """ - pass + use_df = getattr(self.mf, "with_df", None) is not None + if isinstance(eris, BaseERIs): + return eris + elif eris is not None: + raise TypeError(f"`eris` must be an `BaseERIs` object, got {eris.__class__.__name__}.") + elif use_df: + return self.CDERIs( + self.mf, + space=(self.space, self.space, self.space, self.space), + mo_coeff=(self.mo_coeff, self.mo_coeff, self.mo_coeff, self.mo_coeff), + ) + else: + return self.ERIs( + self.mf, + space=(self.space, self.space, self.space, self.space), + mo_coeff=(self.mo_coeff, self.mo_coeff, self.mo_coeff, self.mo_coeff), + ) def get_g(self) -> BaseElectronBoson: """Get the blocks of the electron-boson coupling matrix. @@ -920,7 +942,14 @@ def get_g(self) -> BaseElectronBoson: Returns: Electron-boson coupling matrix. """ - return self.ElectronBoson(self, array=self.bare_g) + if self.bare_g is None: + raise ValueError("Bare electron-boson coupling matrix not provided.") + return self.ElectronBoson( + self.mf, + self.bare_g, + (self.space, self.space), + (self.mo_coeff, self.mo_coeff), + ) @abstractmethod def get_mean_field_G(self) -> Any: @@ -931,18 +960,6 @@ def get_mean_field_G(self) -> Any: """ pass - @property - @abstractmethod - def bare_fock(self) -> Any: - """Get the mean-field Fock matrix in the MO basis, including frozen parts. - - Returns an array and not a `BaseFock` object. - - Returns: - Mean-field Fock matrix. - """ - pass - @property @abstractmethod def xi(self) -> NDArray[T]: diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 0803ccc2..123826bd 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -856,19 +856,6 @@ def get_mean_field_G(self) -> NDArray[T]: val += self.bare_G return val - @property - def bare_fock(self) -> NDArray[T]: - """Get the mean-field Fock matrix in the MO basis, including frozen parts. - - Returns an array and not a `BaseFock` object. - - Returns: - Mean-field Fock matrix. - """ - fock_ao: NDArray[T] = np.asarray(self.mf.get_fock(), dtype=types[float]) - fock = util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff, self.mo_coeff) - return fock - @property def xi(self) -> NDArray[T]: """Get the shift in the bosonic operators to diagonalise the photon Hamiltonian. @@ -888,28 +875,6 @@ def xi(self) -> NDArray[T]: xi = np.zeros(self.omega.shape) return xi - def get_fock(self) -> GFock: - """Get the Fock matrix. - - Returns: - Fock matrix. - """ - return self.Fock(self, array=self.bare_fock, g=self.g) - - def get_eris(self, eris: Optional[ERIsInputType] = None) -> GERIs: - """Get the electron repulsion integrals. - - Args: - eris: Input electron repulsion integrals. - - Returns: - Electron repulsion integrals. - """ - if isinstance(eris, GERIs): - return eris - else: - return self.ERIs(self, array=eris) - @property def nmo(self) -> int: """Get the number of molecular orbitals. diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index b16a0e88..1d2e0c39 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -614,19 +614,6 @@ def get_mean_field_G(self) -> NDArray[T]: val += self.bare_G return val - @property - def bare_fock(self) -> NDArray[T]: - """Get the mean-field Fock matrix in the MO basis, including frozen parts. - - Returns an array and not a `BaseFock` object. - - Returns: - Mean-field Fock matrix. - """ - fock_ao: NDArray[T] = np.asarray(self.mf.get_fock(), dtype=types[float]) - fock = util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff, self.mo_coeff) - return fock - @property def xi(self) -> NDArray[T]: """Get the shift in the bosonic operators to diagonalise the photon Hamiltonian. @@ -646,31 +633,6 @@ def xi(self) -> NDArray[T]: xi = np.zeros(self.omega.shape, dtype=types[float]) return xi - def get_fock(self) -> RFock: - """Get the Fock matrix. - - Returns: - Fock matrix. - """ - return self.Fock(self, array=self.bare_fock, g=self.g) - - def get_eris(self, eris: Optional[ERIsInputType] = None) -> Union[RERIs, RCDERIs]: - """Get the electron repulsion integrals. - - Args: - eris: Input electron repulsion integrals. - - Returns: - Electron repulsion integrals. - """ - use_df = getattr(self.mf, "with_df", None) is not None - if isinstance(eris, (RERIs, RCDERIs)): - return eris - elif (isinstance(eris, np.ndarray) and eris.ndim == 3) or use_df: - return self.CDERIs(self, array=eris) - else: - return self.ERIs(self, array=eris) - @property def nmo(self) -> int: """Get the number of molecular orbitals. diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 32c1c968..a6164318 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -809,24 +809,6 @@ def get_mean_field_G(self) -> NDArray[T]: val += self.bare_G return val - @property - def bare_fock(self) -> Namespace[NDArray[T]]: - """Get the mean-field Fock matrix in the MO basis, including frozen parts. - - Returns an array and not a `BaseFock` object. - - Returns: - Mean-field Fock matrix. - """ - fock_array = util.einsum( - "npq,npi,nqj->nij", - np.asarray(self.mf.get_fock(), dtype=types[float]), - self.mo_coeff, - self.mo_coeff, - ) - fock = util.Namespace(aa=fock_array[0], bb=fock_array[1]) - return fock - @property def xi(self) -> NDArray[T]: """Get the shift in the bosonic operators to diagonalise the photon Hamiltonian. @@ -847,33 +829,6 @@ def xi(self) -> NDArray[T]: xi = np.zeros(self.omega.shape) return xi - def get_fock(self) -> UFock: - """Get the Fock matrix. - - Returns: - Fock matrix. - """ - return self.Fock(self, array=(self.bare_fock.aa, self.bare_fock.bb), g=self.g) - - def get_eris(self, eris: Optional[ERIsInputType] = None) -> Union[UERIs, UCDERIs]: - """Get the electron repulsion integrals. - - Args: - eris: Input electron repulsion integrals. - - Returns: - Electron repulsion integrals. - """ - use_df = getattr(self.mf, "with_df", None) is not None - if isinstance(eris, (UERIs, UCDERIs)): - return eris - elif ( - isinstance(eris, tuple) and isinstance(eris[0], np.ndarray) and eris[0].ndim == 3 - ) or use_df: - return self.CDERIs(self, array=eris) - else: - return self.ERIs(self, array=eris) - @property def nmo(self) -> int: """Get the number of molecular orbitals. diff --git a/ebcc/ham/base.py b/ebcc/ham/base.py index 50b9f0f2..353d9e25 100644 --- a/ebcc/ham/base.py +++ b/ebcc/ham/base.py @@ -2,24 +2,62 @@ from __future__ import annotations +import sys from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any +import numpy + +from ebcc import BACKEND +from ebcc import numpy as np +from ebcc.backend import to_numpy +from ebcc.core.precision import types from ebcc.util import Namespace if TYPE_CHECKING: from typing import Optional - from ebcc.cc.base import BaseEBCC - from ebcc.cc.gebcc import GEBCC - from ebcc.cc.rebcc import REBCC - from ebcc.cc.uebcc import UEBCC + from numpy import float64 + from numpy.typing import NDArray + from pyscf.scf.hf import SCF + + from ebcc.cc.base import SpaceType + from ebcc.cc.gebcc import SpaceType as GSpaceType + from ebcc.cc.rebcc import SpaceType as RSpaceType + from ebcc.cc.uebcc import SpaceType as USpaceType + + CoeffType = Any + RCoeffType = NDArray[float64] + UCoeffType = tuple[NDArray[float64], NDArray[float64]] + GCoeffType = NDArray[float64] class BaseHamiltonian(Namespace[Any], ABC): """Base class for Hamiltonians.""" - cc: BaseEBCC + mf: SCF + space: tuple[SpaceType, ...] + mo_coeff: tuple[CoeffType, ...] + + def __init__( + self, + mf: SCF, + space: tuple[SpaceType, ...], + mo_coeff: Optional[tuple[CoeffType, ...]] = None, + ) -> None: + """Initialise the Hamiltonian. + + Args: + mf: Mean-field object. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. + """ + Namespace.__init__(self) + + # Parameters: + self.__dict__["mf"] = mf + self.__dict__["space"] = space + self.__dict__["mo_coeff"] = mo_coeff if mo_coeff is not None else (mf.mo_coeff,) * 4 @abstractmethod def __getitem__(self, key: str) -> Any: @@ -33,23 +71,68 @@ def __getitem__(self, key: str) -> Any: """ pass + def _to_pyscf_backend(self, array: NDArray[Any]) -> NDArray[Any]: + """Convert an array to the NumPy backend used by PySCF.""" + if BACKEND == "jax" and "pyscfad" in sys.modules: + return array + else: + return to_numpy(array, dtype=numpy.float64) + + def _to_ebcc_backend(self, array: NDArray[Any]) -> NDArray[Any]: + """Convert an array to the NumPy backend used by `ebcc`.""" + return np.asarray(array, dtype=types[float]) + + def _get_slices(self, key: str) -> tuple[slice, ...]: + """Get the slices for the given key. + + Args: + key: Key to get. + + Returns: + Slices for the given key. + """ + slices = tuple(s.slice(k) for s, k in zip(self.space, key)) + return slices + + def _get_coeffs(self, key: str, offset: int = 0) -> tuple[NDArray[Any], ...]: + """Get the coefficients for the given key. + + Args: + key: Key to get. + + Returns: + Coefficients for the given key. + """ + coeffs = tuple( + ( + self.mo_coeff[i + offset][:, self.space[i + offset].slice(k)] + if k != "p" + else np.eye(self.mo_coeff[i + offset].shape[0]) + ) + for i, k in enumerate(key) + ) + return coeffs + class BaseRHamiltonian(BaseHamiltonian): """Base class for restricted Hamiltonians.""" - cc: REBCC + space: tuple[RSpaceType, ...] + mo_coeff: tuple[RCoeffType, ...] class BaseUHamiltonian(BaseHamiltonian): """Base class for unrestricted Hamiltonians.""" - cc: UEBCC + space: tuple[USpaceType, ...] + mo_coeff: tuple[UCoeffType, ...] class BaseGHamiltonian(BaseHamiltonian): """Base class for general Hamiltonians.""" - cc: GEBCC + space: tuple[GSpaceType, ...] + mo_coeff: tuple[GCoeffType, ...] class BaseFock(BaseHamiltonian): @@ -57,65 +140,37 @@ class BaseFock(BaseHamiltonian): def __init__( self, - cc: BaseEBCC, - array: Optional[Any] = None, - space: Optional[tuple[Any, ...]] = None, - mo_coeff: Optional[tuple[Any, ...]] = None, + mf: SCF, + space: tuple[SpaceType, ...], + mo_coeff: Optional[tuple[CoeffType, ...]] = None, g: Optional[Namespace[Any]] = None, + shift: Optional[bool] = None, + xi: Optional[NDArray[float64]] = None, ) -> None: - """Initialise the Fock matrix. + """Initialise the Hamiltonian. Args: - cc: Coupled cluster object. - array: Fock matrix in the MO basis. + mf: Mean-field object. space: Space object for each index. mo_coeff: Molecular orbital coefficients for each index. g: Namespace containing blocks of the electron-boson coupling matrix. + shift: Shift the boson operators such that the Hamiltonian is normal-ordered with + respect to a coherent state. This removes the bosonic coupling to the static + mean-field density, introducing a constant energy shift. + xi: Shift in the bosonic operators to diagonalise the photon Hamiltonian. """ - Namespace.__init__(self) - - # Parameters: - self.__dict__["cc"] = cc - self.__dict__["space"] = space if space is not None else (cc.space,) * 2 - self.__dict__["mo_coeff"] = mo_coeff if mo_coeff is not None else (cc.mo_coeff,) * 2 - self.__dict__["array"] = array if array is not None else self._get_fock() - self.__dict__["g"] = g if g is not None else cc.g + super().__init__(mf, space, mo_coeff=mo_coeff) # Boson parameters: - self.__dict__["shift"] = cc.options.shift if g is not None else None - self.__dict__["xi"] = cc.xi if g is not None else None - - @abstractmethod - def _get_fock(self) -> Any: - """Get the Fock matrix.""" - pass + self.__dict__["g"] = g + self.__dict__["shift"] = shift + self.__dict__["xi"] = xi class BaseERIs(BaseHamiltonian): """Base class for electronic repulsion integrals.""" - def __init__( - self, - cc: BaseEBCC, - array: Optional[Any] = None, - space: Optional[tuple[Any, ...]] = None, - mo_coeff: Optional[tuple[Any, ...]] = None, - ) -> None: - """Initialise the ERIs. - - Args: - cc: Coupled cluster object. - array: ERIs in the MO basis. - space: Space object for each index. - mo_coeff: Molecular orbital coefficients for each index. - """ - Namespace.__init__(self) - - # Parameters: - self.__dict__["cc"] = cc - self.__dict__["space"] = space if space is not None else (cc.space,) * 4 - self.__dict__["mo_coeff"] = mo_coeff if mo_coeff is not None else (cc.mo_coeff,) * 4 - self.__dict__["array"] = array if array is not None else None + pass class BaseElectronBoson(BaseHamiltonian): @@ -123,20 +178,20 @@ class BaseElectronBoson(BaseHamiltonian): def __init__( self, - cc: BaseEBCC, - array: Optional[Any] = None, - space: Optional[tuple[Any, ...]] = None, + mf: SCF, + g: NDArray[float64], + space: tuple[SpaceType, ...], + mo_coeff: Optional[tuple[CoeffType, ...]] = None, ) -> None: - """Initialise the electron-boson coupling matrix. + """Initialise the Hamiltonian. Args: - cc: Coupled cluster object. - array: Electron-boson coupling matrix in the MO basis. + mf: Mean-field object. + g: The electron-boson coupling matrix array. space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. """ - Namespace.__init__(self) + super().__init__(mf, space, mo_coeff=mo_coeff) - # Parameters: - self.__dict__["cc"] = cc - self.__dict__["space"] = space if space is not None else (cc.space,) * 2 - self.__dict__["array"] = array if array is not None else self._get_g() + # Boson parameters: + self.__dict__["g"] = g diff --git a/ebcc/ham/cderis.py b/ebcc/ham/cderis.py index d6420784..72902fb6 100644 --- a/ebcc/ham/cderis.py +++ b/ebcc/ham/cderis.py @@ -4,13 +4,10 @@ from typing import TYPE_CHECKING -import numpy # PySCF uses true numpy, no backend stuff here from pyscf import lib from ebcc import numpy as np from ebcc import pyscf, util -from ebcc.backend import to_numpy -from ebcc.core.precision import types from ebcc.ham.base import BaseERIs, BaseRHamiltonian, BaseUHamiltonian if TYPE_CHECKING: @@ -37,9 +34,6 @@ def __getitem__(self, key: str, e2: Optional[bool] = False) -> NDArray[T]: Returns: CDERIs for the given spaces. """ - if self.array is not None: - raise NotImplementedError("`array` is not supported for CDERIs.") - if len(key) == 4: v1 = self.__getitem__("Q" + key[:2]) v2 = self.__getitem__("Q" + key[2:], e2=True) # type: ignore @@ -51,28 +45,30 @@ def __getitem__(self, key: str, e2: Optional[bool] = False) -> NDArray[T]: key_e2 = f"{key}_{'e1' if not e2 else 'e2'}" # Check the DF is built incore - if not isinstance(self.cc.mf.with_df._cderi, np.ndarray): - with lib.temporary_env(self.cc.mf.with_df, max_memory=1e6): - self.cc.mf.with_df.build() + if not isinstance(self.mf.with_df._cderi, np.ndarray): + with lib.temporary_env(self.mf.with_df, max_memory=1e6): + self.mf.with_df.build() if key_e2 not in self._members: s = 0 if not e2 else 2 - coeffs = [ - to_numpy(self.mo_coeff[i + s][:, self.space[i + s].slice(k)], dtype=numpy.float64) - for i, k in enumerate(key) - ] - ijslice = ( - 0, - coeffs[0].shape[-1], - coeffs[0].shape[-1], - coeffs[0].shape[-1] + coeffs[1].shape[-1], - ) - coeffs = numpy.concatenate(coeffs, axis=1) + + # Get the coefficients and shape + coeffs = self._get_coeffs(key, offset=s) + shape = tuple(c.shape[-1] for c in coeffs) + ijslice = (0, shape[0], shape[0], shape[0] + shape[1]) + coeffs = np.concatenate(coeffs, axis=1) + coeffs = self._to_pyscf_backend(coeffs) + + # Transform the block + # TODO: Optimise for (L|pp) block = pyscf.ao2mo._ao2mo.nr_e2( - self.cc.mf.with_df._cderi, coeffs, ijslice, aosym="s2", mosym="s1" + self.mf.with_df._cderi, coeffs, ijslice, aosym="s2", mosym="s1" ) - block = np.reshape(block, (-1, ijslice[1] - ijslice[0], ijslice[3] - ijslice[2])) - self._members[key_e2] = np.asarray(block, dtype=types[float]) + block = self._to_ebcc_backend(block) + block = np.reshape(block, (-1, *shape)) + + # Store the block + self._members[key_e2] = block return self._members[key_e2] @@ -102,7 +98,7 @@ def __getitem__(self, key: str) -> RCDERIs: if key not in self._members: self._members[key] = RCDERIs( - self.cc, + self.mf, space=(self.space[0][i], self.space[1][i], self.space[2][j], self.space[3][j]), mo_coeff=( self.mo_coeff[0][i], diff --git a/ebcc/ham/elbos.py b/ebcc/ham/elbos.py index 5a4ec4e4..d98b00ff 100644 --- a/ebcc/ham/elbos.py +++ b/ebcc/ham/elbos.py @@ -30,15 +30,16 @@ def __getitem__(self, key: str) -> NDArray[T]: """ if key not in self._members: assert key[0] == "b" - i = self.space[0].slice(key[1]) - j = self.space[1].slice(key[2]) - self._members[key] = np.copy(self.array[:, i, j]) - return self._members[key] + if "p" in key: + raise NotImplementedError(f"AO basis not supported in {self.__class__.__name__}.") + + # Get the slices + slices = (slice(None),) + self._get_slices(key[1:]) + + # Store the block + self._members[key] = np.copy(self.g[slices]) - def _get_g(self) -> NDArray[T]: - """Get the electron-boson coupling matrix.""" - assert self.cc.bare_g is not None - return self.cc.bare_g + return self._members[key] class UElectronBoson(BaseElectronBoson, BaseUHamiltonian): @@ -57,19 +58,17 @@ def __getitem__(self, key: str) -> RElectronBoson: """ if key not in ("aa", "bb"): raise KeyError(f"Invalid key: {key}") + if key not in self._members: i = "ab".index(key[0]) self._members[key] = RElectronBoson( - self.cc, - array=self.array[i] if self.array.ndim == 4 else self.array, + self.mf, + self.g[i] if self.g.ndim == 4 else self.g, space=(self.space[0][i], self.space[1][i]), ) - return self._members[key] + self._members[key]._spin_index = i - def _get_g(self) -> NDArray[T]: - """Get the electron-boson coupling matrix.""" - assert self.cc.bare_g is not None - return self.cc.bare_g + return self._members[key] class GElectronBoson(BaseElectronBoson, BaseGHamiltonian): @@ -88,12 +87,13 @@ def __getitem__(self, key: str) -> NDArray[T]: """ if key not in self._members: assert key[0] == "b" - i = self.space[0].slice(key[1]) - j = self.space[1].slice(key[2]) - self._members[key] = np.copy(self.array[:, i, j]) - return self._members[key] + if "p" in key: + raise NotImplementedError(f"AO basis not supported in {self.__class__.__name__}.") + + # Get the slices + slices = (slice(None),) + self._get_slices(key[1:]) - def _get_g(self) -> NDArray[T]: - """Get the electron-boson coupling matrix.""" - assert self.cc.bare_g is not None - return self.cc.bare_g + # Store the block + self._members[key] = np.copy(self.g[slices]) + + return self._members[key] diff --git a/ebcc/ham/eris.py b/ebcc/ham/eris.py index 5ec8712f..fa781b29 100644 --- a/ebcc/ham/eris.py +++ b/ebcc/ham/eris.py @@ -4,16 +4,13 @@ from typing import TYPE_CHECKING -import numpy # PySCF uses true numpy, no backend stuff here - from ebcc import numpy as np from ebcc import pyscf -from ebcc.backend import to_numpy from ebcc.core.precision import types from ebcc.ham.base import BaseERIs, BaseGHamiltonian, BaseRHamiltonian, BaseUHamiltonian if TYPE_CHECKING: - from typing import Any, Optional + from typing import Any from numpy import float64 from numpy.typing import NDArray @@ -25,7 +22,6 @@ class RERIs(BaseERIs, BaseRHamiltonian): """Restricted ERIs container class.""" _members: dict[str, NDArray[T]] - array: Optional[NDArray[T]] def __getitem__(self, key: str) -> NDArray[T]: """Just-in-time getter. @@ -36,29 +32,30 @@ def __getitem__(self, key: str) -> NDArray[T]: Returns: ERIs for the given spaces. """ - if self.array is None: - if key not in self._members.keys(): - coeffs = [ - to_numpy(self.mo_coeff[i][:, self.space[i].slice(k)], dtype=numpy.float64) - for i, k in enumerate(key) - ] - if getattr(self.cc.mf, "_eri", None) is not None: - block = pyscf.ao2mo.incore.general(self.cc.mf._eri, coeffs, compact=False) - else: - block = pyscf.ao2mo.kernel(self.cc.mf.mol, coeffs, compact=False) - block = np.reshape(block, [c.shape[-1] for c in coeffs]) - self._members[key] = np.asarray(block, dtype=types[float]) - return self._members[key] - else: - ijkl = tuple(self.space[i].slice(k) for i, k in enumerate(key)) - return self.array[ijkl] # type: ignore + if key not in self._members.keys(): + # Get the coefficients and shape + coeffs = self._get_coeffs(key) + coeffs = tuple(self._to_pyscf_backend(c) for c in coeffs) + shape = tuple(c.shape[-1] for c in coeffs) + + # Transform the block + # TODO: Optimise for patially AO + if getattr(self.mf, "_eri", None) is not None: + block = pyscf.ao2mo.incore.general(self.mf._eri, coeffs, compact=False) + else: + block = pyscf.ao2mo.kernel(self.mf.mol, coeffs, compact=False) + block = np.reshape(block, shape) + + # Store the block + self._members[key] = np.asarray(block, dtype=types[float]) + + return self._members[key] class UERIs(BaseERIs, BaseUHamiltonian): """Unrestricted ERIs container class.""" _members: dict[str, RERIs] - array: Optional[tuple[NDArray[T], ...]] def __getitem__(self, key: str) -> RERIs: """Just-in-time getter. @@ -71,34 +68,13 @@ def __getitem__(self, key: str) -> RERIs: """ if key not in ("aaaa", "aabb", "bbaa", "bbbb"): raise KeyError(f"Invalid key: {key}") + if key not in self._members: i = "ab".index(key[0]) j = "ab".index(key[2]) - ij = i * (i + 1) // 2 + j - - if self.array is not None: - array = self.array[ij] - if key == "bbaa": - array = np.transpose(array, (2, 3, 0, 1)) - elif isinstance(self.cc.mf._eri, tuple): - # Support spin-dependent integrals in the mean-field - coeffs = [ - to_numpy(self.mo_coeff[x][y], dtype=numpy.float64) - for y, x in enumerate(sorted((i, i, j, j))) - ] - if getattr(self.cc.mf, "_eri", None) is not None: - array = pyscf.ao2mo.incore.general(self.cc.mf.mol, coeffs, compact=False) - else: - array = pyscf.ao2mo.kernel(self.cc.mf.mol, coeffs, compact=False) - if key == "bbaa": - array = np.transpose(array, (2, 3, 0, 1)) - array = np.asarray(array, dtype=types[float]) - else: - array = None self._members[key] = RERIs( - self.cc, - array=array, + self.mf, space=(self.space[0][i], self.space[1][i], self.space[2][j], self.space[3][j]), mo_coeff=( self.mo_coeff[0][i], @@ -107,6 +83,7 @@ def __getitem__(self, key: str) -> RERIs: self.mo_coeff[3][j], ), ) + return self._members[key] @@ -114,30 +91,42 @@ class GERIs(BaseERIs, BaseGHamiltonian): """Generalised ERIs container class.""" _members: dict[str, NDArray[T]] - array: NDArray[T] + _array: NDArray[T] def __init__(self, *args: Any, **kwargs: Any) -> None: """Initialise the class.""" super().__init__(*args, **kwargs) - if self.array is None: - mo_a = [to_numpy(mo[: self.cc.mf.mol.nao], dtype=numpy.float64) for mo in self.mo_coeff] - mo_b = [to_numpy(mo[self.cc.mf.mol.nao :], dtype=numpy.float64) for mo in self.mo_coeff] - if getattr(self.cc.mf, "_eri", None) is not None: - array = pyscf.ao2mo.incore.general(self.cc.mf._eri, mo_a) - array += pyscf.ao2mo.incore.general(self.cc.mf._eri, mo_b) - array += pyscf.ao2mo.incore.general(self.cc.mf._eri, mo_a[:2] + mo_b[2:]) - array += pyscf.ao2mo.incore.general(self.cc.mf._eri, mo_b[:2] + mo_a[2:]) - else: - array = pyscf.ao2mo.kernel(self.cc.mf.mol, mo_a) - array += pyscf.ao2mo.kernel(self.cc.mf.mol, mo_b) - array += pyscf.ao2mo.kernel(self.cc.mf.mol, mo_a[:2] + mo_b[2:]) - array += pyscf.ao2mo.kernel(self.cc.mf.mol, mo_b[:2] + mo_a[2:]) - array = np.reshape( - pyscf.ao2mo.addons.restore(1, array, self.cc.nmo), (self.cc.nmo,) * 4 + + # Get the coefficients and shape + mo_a = [self._to_pyscf_backend(mo[: self.mf.mol.nao]) for mo in self.mo_coeff] + mo_b = [self._to_pyscf_backend(mo[self.mf.mol.nao :]) for mo in self.mo_coeff] + shape = tuple(mo.shape[-1] for mo in self.mo_coeff) + if len(set(shape)) != 1: + raise ValueError( + "MO coefficients must have the same number of basis functions for " + f"{self.__class__.__name__}." ) - array = np.asarray(array, dtype=types[float]) - array = np.transpose(array, (0, 2, 1, 3)) - np.transpose(array, (0, 2, 3, 1)) - self.__dict__["array"] = array + nmo = shape[0] + + if getattr(self.mf, "_eri", None) is not None: + array = pyscf.ao2mo.incore.general(self.mf._eri, mo_a) + array += pyscf.ao2mo.incore.general(self.mf._eri, mo_b) + array += pyscf.ao2mo.incore.general(self.mf._eri, mo_a[:2] + mo_b[2:]) + array += pyscf.ao2mo.incore.general(self.mf._eri, mo_b[:2] + mo_a[2:]) + else: + array = pyscf.ao2mo.kernel(self.mf.mol, mo_a) + array += pyscf.ao2mo.kernel(self.mf.mol, mo_b) + array += pyscf.ao2mo.kernel(self.mf.mol, mo_a[:2] + mo_b[2:]) + array += pyscf.ao2mo.kernel(self.mf.mol, mo_b[:2] + mo_a[2:]) + array = pyscf.ao2mo.addons.restore(1, array, nmo) + array = np.reshape(array, shape) + array = self._to_ebcc_backend(array) + + # Transform to antisymmetric Physicist's notation + array = np.transpose(array, (0, 2, 1, 3)) - np.transpose(array, (0, 2, 3, 1)) + + # Store the array + self.__dict__["_array"] = array def __getitem__(self, key: str) -> NDArray[T]: """Just-in-time getter. @@ -148,5 +137,6 @@ def __getitem__(self, key: str) -> NDArray[T]: Returns: ERIs for the given spaces. """ - ijkl = tuple(self.space[i].slice(k) for i, k in enumerate(key)) - return self.array[ijkl] # type: ignore + if "p" in key: + raise NotImplementedError(f"AO basis not supported in {self.__class__.__name__}.") + return self._array[self._get_slices(key)] diff --git a/ebcc/ham/fock.py b/ebcc/ham/fock.py index 7691a15d..fd5af0ec 100644 --- a/ebcc/ham/fock.py +++ b/ebcc/ham/fock.py @@ -7,9 +7,11 @@ from ebcc import numpy as np from ebcc import util from ebcc.core.precision import types -from ebcc.ham.base import BaseFock, BaseGHamiltonian, BaseRHamiltonian, BaseUHamiltonian +from ebcc.ham.base import BaseFock, BaseRHamiltonian, BaseUHamiltonian if TYPE_CHECKING: + from typing import Optional + from numpy import float64 from numpy.typing import NDArray @@ -20,10 +22,7 @@ class RFock(BaseFock, BaseRHamiltonian): """Restricted Fock matrix container class.""" _members: dict[str, NDArray[T]] - - def _get_fock(self) -> NDArray[T]: - fock_ao: NDArray[T] = np.asarray(self.cc.mf.get_fock(), dtype=types[float]) - return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff[0], self.mo_coeff[1]) + _spin_index: Optional[int] = None def __getitem__(self, key: str) -> NDArray[T]: """Just-in-time getter. @@ -35,11 +34,20 @@ def __getitem__(self, key: str) -> NDArray[T]: Fock matrix for the given spaces. """ if key not in self._members: - i = self.space[0].slice(key[0]) - j = self.space[1].slice(key[1]) - self._members[key] = np.copy(self.array[i, j]) + # Get the coefficients + coeffs = self._get_coeffs(key) + + # Transform the block + fock_ao: NDArray[T] = np.asarray(self.mf.get_fock(), dtype=types[float]) + if self._spin_index is not None: + fock_ao = fock_ao[self._spin_index] + block = util.einsum("pq,pi,qj->ij", fock_ao, *coeffs) + + # Store the block + self._members[key] = block if self.shift: + # Shift for bosons xi = self.xi g = np.copy(self.g.__getattr__(f"b{key}")) g += np.transpose(self.g.__getattr__(f"b{key[::-1]}"), (0, 2, 1)) @@ -53,13 +61,6 @@ class UFock(BaseFock, BaseUHamiltonian): _members: dict[str, RFock] - def _get_fock(self) -> tuple[NDArray[T], NDArray[T]]: - fock_ao: NDArray[T] = np.asarray(self.cc.mf.get_fock(), dtype=types[float]) - return ( - util.einsum("pq,pi,qj->ij", fock_ao[0], self.mo_coeff[0][0], self.mo_coeff[1][0]), - util.einsum("pq,pi,qj->ij", fock_ao[1], self.mo_coeff[0][1], self.mo_coeff[1][1]), - ) - def __getitem__(self, key: str) -> RFock: """Just-in-time getter. @@ -71,45 +72,23 @@ def __getitem__(self, key: str) -> RFock: """ if key not in ("aa", "bb"): raise KeyError(f"Invalid key: {key}") + if key not in self._members: i = "ab".index(key[0]) self._members[key] = RFock( - self.cc, - array=self.array[i], + self.mf, space=(self.space[0][i], self.space[1][i]), mo_coeff=(self.mo_coeff[0][i], self.mo_coeff[1][i]), g=self.g[key] if self.g is not None else None, + shift=self.shift, + xi=self.xi, ) + self._members[key].__dict__["_spin_index"] = i + return self._members[key] -class GFock(BaseFock, BaseGHamiltonian): +class GFock(RFock): """Generalised Fock matrix container class.""" - _members: dict[str, NDArray[T]] - - def _get_fock(self) -> NDArray[T]: - fock_ao: NDArray[T] = np.asarray(self.cc.mf.get_fock(), dtype=types[float]) - return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff[0], self.mo_coeff[1]) - - def __getitem__(self, key: str) -> NDArray[T]: - """Just-in-time getter. - - Args: - key: Key to get. - - Returns: - Fock matrix for the given spin. - """ - if key not in self._members: - i = self.space[0].slice(key[0]) - j = self.space[1].slice(key[1]) - self._members[key] = np.copy(self.array[i, j]) - - if self.shift: - xi = self.xi - g = np.copy(self.g.__getattr__(f"b{key}")) - g += np.transpose(self.g.__getattr__(f"b{key[::-1]}"), (0, 2, 1)) - self._members[key] -= util.einsum("I,Ipq->pq", xi, g) - - return self._members[key] + pass diff --git a/tests/test_GCC2.py b/tests/test_GCC2.py index 0539838d..b8472f1c 100644 --- a/tests/test_GCC2.py +++ b/tests/test_GCC2.py @@ -83,7 +83,7 @@ def test_t1_amplitudes(self): # c = self.mf.to_ghf().mo_coeff # h = self.mf.to_ghf().get_hcore() # h = np.linalg.multi_dot((c.T, h, c)) - # v = self.ccsd.get_eris().array + # v = self.ccsd.get_eris().xxxx # e_rdm = util.einsum("pq,pq->", h, dm1) # e_rdm += util.einsum("pqrs,pqrs->", v, dm2) * 0.5 # e_rdm += self.mf.mol.energy_nuc() diff --git a/tests/test_GCCSD.py b/tests/test_GCCSD.py index ec3394bd..239ecafd 100644 --- a/tests/test_GCCSD.py +++ b/tests/test_GCCSD.py @@ -397,7 +397,7 @@ def test_rdm_energy(self): c = self.mf.to_ghf().mo_coeff h = self.mf.to_ghf().get_hcore() h = np.linalg.multi_dot((c.T, h, c)) - v = self.ccsd.get_eris().array + v = self.ccsd.get_eris().xxxx e_rdm = np.einsum("pq,pq->", h, dm1) e_rdm += np.einsum("pqrs,pqrs->", v, dm2) * 0.5 e_rdm += self.mf.mol.energy_nuc() diff --git a/tests/test_GCCSDT.py b/tests/test_GCCSDT.py index 12221662..87ad75f3 100644 --- a/tests/test_GCCSDT.py +++ b/tests/test_GCCSDT.py @@ -69,7 +69,7 @@ def test_rdm_energy(self): c = mf.to_ghf().mo_coeff h = mf.to_ghf().get_hcore() h = np.linalg.multi_dot((c.T, h, c)) - v = ccsdt.get_eris().array + v = ccsdt.get_eris().xxxx e_rdm = util.einsum("pq,pq->", h, dm1) e_rdm += util.einsum("pqrs,pqrs->", v, dm2) * 0.5 e_rdm += mol.energy_nuc()