From 5ba55d14966afcbb5a8f5a01cdc007d97bb18e74 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sat, 27 Jul 2024 23:36:28 +0100 Subject: [PATCH 01/37] Moving base classes --- ebcc/base.py | 35 ++++++++++++++++++ ebcc/brueckner.py | 13 ++----- ebcc/cderis.py | 12 ++---- ebcc/eris.py | 9 +---- ebcc/fock.py | 10 +---- ebcc/gebcc.py | 16 +------- ebcc/geom.py | 3 +- ebcc/rebcc.py | 7 +--- ebcc/reom.py | 35 ++---------------- ebcc/uebcc.py | 33 +---------------- ebcc/ueom.py | 19 ++-------- ebcc/util/__init__.py | 3 +- ebcc/util/inheritance.py | 80 ---------------------------------------- ebcc/util/misc.py | 9 +++++ 14 files changed, 66 insertions(+), 218 deletions(-) create mode 100644 ebcc/base.py delete mode 100644 ebcc/util/inheritance.py diff --git a/ebcc/base.py b/ebcc/base.py new file mode 100644 index 00000000..f8e879b6 --- /dev/null +++ b/ebcc/base.py @@ -0,0 +1,35 @@ +"""Base classes.""" + +from abc import ABC, abstractmethod + +from ebcc import util + + +class EBCC(ABC): + """Base class for electron-boson coupled cluster.""" + + pass + + +class EOM(ABC): + """Base class for equation-of-motion methods.""" + + pass + + +class BruecknerEBCC(ABC): + """Base class for Brueckner orbital methods.""" + + pass + + +class ERIs(ABC, util.Namespace): + """Base class for electronic repulsion integrals.""" + + pass + + +class Fock(ABC, util.Namespace): + """Base class for Fock matrices.""" + + pass diff --git a/ebcc/brueckner.py b/ebcc/brueckner.py index 1bd1f0db..7f232f11 100644 --- a/ebcc/brueckner.py +++ b/ebcc/brueckner.py @@ -11,6 +11,7 @@ from ebcc.damping import DIIS from ebcc.logging import ANSI from ebcc.precision import types +from ebcc.base import BruecknerEBCC @dataclasses.dataclass @@ -42,7 +43,7 @@ class Options: damping: float = 0.0 -class BruecknerREBCC: +class BruecknerREBCC(BruecknerEBCC): """ Brueckner orbital self-consistency for coupled cluster calculations. Iteratively solve for a new mean-field that presents a vanishing T1 @@ -376,9 +377,7 @@ def spin_type(self): return self.cc.spin_type -@util.has_docstring -class BruecknerUEBCC(BruecknerREBCC, metaclass=util.InheritDocstrings): - @util.has_docstring +class BruecknerUEBCC(BruecknerREBCC): def get_rotation_matrix(self, u_tot=None, diis=None, t1=None): if t1 is None: t1 = self.cc.t1 @@ -426,7 +425,6 @@ def get_rotation_matrix(self, u_tot=None, diis=None, t1=None): return u, u_tot - @util.has_docstring def transform_amplitudes(self, u, amplitudes=None): if amplitudes is None: amplitudes = self.cc.amplitudes @@ -456,7 +454,6 @@ def transform_amplitudes(self, u, amplitudes=None): return self.cc.amplitudes - @util.has_docstring def get_t1_norm(self, amplitudes=None): if amplitudes is None: amplitudes = self.cc.amplitudes @@ -466,21 +463,18 @@ def get_t1_norm(self, amplitudes=None): return np.linalg.norm([norm_a, norm_b]) - @util.has_docstring def mo_to_correlated(self, mo_coeff): return ( mo_coeff[0][:, self.cc.space[0].correlated], mo_coeff[1][:, self.cc.space[1].correlated], ) - @util.has_docstring def mo_update_correlated(self, mo_coeff, mo_coeff_corr): mo_coeff[0][:, self.cc.space[0].correlated] = mo_coeff_corr[0] mo_coeff[1][:, self.cc.space[1].correlated] = mo_coeff_corr[1] return mo_coeff - @util.has_docstring def update_coefficients(self, u_tot, mo_coeff, mo_coeff_ref): mo_coeff_new_corr = ( util.einsum("pi,ij->pj", mo_coeff_ref[0], u_tot.aa), @@ -490,6 +484,5 @@ def update_coefficients(self, u_tot, mo_coeff, mo_coeff_ref): return mo_coeff_new -@util.has_docstring class BruecknerGEBCC(BruecknerREBCC): pass diff --git a/ebcc/cderis.py b/ebcc/cderis.py index 466e104d..1787ec40 100644 --- a/ebcc/cderis.py +++ b/ebcc/cderis.py @@ -4,15 +4,10 @@ from pyscf import ao2mo from ebcc import precision, util +from ebcc.base import ERIs -class CDERIs(util.Namespace): - """Base class for Cholesky decomposed ERIs.""" - - pass - - -class RCDERIs(CDERIs): +class RCDERIs(ERIs): """ Cholesky decomposed ERI container class for `REBCC`. Consists of a just-in-time namespace containing blocks of the integrals. @@ -123,8 +118,7 @@ def __getattr__(self, key): __getitem__ = __getattr__ -@util.has_docstring -class UCDERIs(CDERIs): +class UCDERIs(ERIs): """ Cholesky decomposed ERI container class for `UEBCC`. Consists of a just-in-time namespace containing blocks of the integrals. diff --git a/ebcc/eris.py b/ebcc/eris.py index 6105ea06..0fc34ec3 100644 --- a/ebcc/eris.py +++ b/ebcc/eris.py @@ -5,12 +5,7 @@ from ebcc import numpy as np from ebcc import util from ebcc.precision import types - - -class ERIs(util.Namespace): - """Base class for electronic repulsion integrals.""" - - pass +from ebcc.base import ERIs class RERIs(ERIs): @@ -97,7 +92,6 @@ def __getattr__(self, key): __getitem__ = __getattr__ -@util.has_docstring class UERIs(ERIs): """ Electronic repulsion integral container class for `UEBCC`. Consists @@ -198,7 +192,6 @@ def __init__(self, ebcc, array=None, slices=None, mo_coeff=None): ) -@util.has_docstring class GERIs(RERIs): __doc__ = __doc__.replace("REBCC", "GEBCC") diff --git a/ebcc/fock.py b/ebcc/fock.py index 39e5856d..7a47217a 100644 --- a/ebcc/fock.py +++ b/ebcc/fock.py @@ -3,12 +3,7 @@ from ebcc import numpy as np from ebcc import util from ebcc.precision import types - - -class Fock(util.Namespace): - """Base class for Fock matrices.""" - - pass +from ebcc.base import Fock class RFock(Fock): @@ -100,7 +95,7 @@ def __getattr__(self, key): __getitem__ = __getattr__ -class UFock(Fock, metaclass=util.InheritDocstrings): +class UFock(Fock): """ Fock matrix container class for `UEBCC`. Consists of a namespace of `RFock` objects, on for each spin signature. @@ -174,6 +169,5 @@ def __init__(self, ebcc, array=None, slices=None, mo_coeff=None): ) -@util.has_docstring class GFock(RFock): __doc__ = RFock.__doc__.replace("REBCC", "GEBCC") diff --git a/ebcc/gebcc.py b/ebcc/gebcc.py index c7c4ef9f..da6da863 100644 --- a/ebcc/gebcc.py +++ b/ebcc/gebcc.py @@ -13,8 +13,7 @@ from ebcc.space import Space -@util.has_docstring -class GEBCC(REBCC, metaclass=util.InheritDocstrings): +class GEBCC(REBCC): __doc__ = __doc__.replace("Restricted", "Generalised", 1) ERIs = GERIs @@ -261,7 +260,6 @@ def from_rebcc(cls, rcc): return gcc - @util.has_docstring def init_amps(self, eris=None): eris = self.get_eris(eris) amplitudes = util.Namespace() @@ -301,7 +299,6 @@ def init_amps(self, eris=None): return amplitudes - @util.has_docstring def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): func, kwargs = self._load_function( "make_rdm2_f", @@ -318,7 +315,6 @@ def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): return dm - @util.has_docstring def excitations_to_vector_ip(self, *excitations): vectors = [] m = 0 @@ -336,7 +332,6 @@ def excitations_to_vector_ip(self, *excitations): return np.concatenate(vectors) - @util.has_docstring def excitations_to_vector_ee(self, *excitations): vectors = [] m = 0 @@ -353,7 +348,6 @@ def excitations_to_vector_ee(self, *excitations): return np.concatenate(vectors) - @util.has_docstring def vector_to_excitations_ip(self, vector): excitations = [] i0 = 0 @@ -375,7 +369,6 @@ def vector_to_excitations_ip(self, vector): return tuple(excitations) - @util.has_docstring def vector_to_excitations_ea(self, vector): excitations = [] i0 = 0 @@ -397,7 +390,6 @@ def vector_to_excitations_ea(self, vector): return tuple(excitations) - @util.has_docstring def vector_to_excitations_ee(self, vector): excitations = [] i0 = 0 @@ -418,7 +410,6 @@ def vector_to_excitations_ee(self, vector): return tuple(excitations) - @util.has_docstring def get_mean_field_G(self): val = lib.einsum("Ipp->I", self.g.boo) val -= self.xi * self.omega @@ -449,20 +440,16 @@ def get_eris(self, eris=None): else: return eris - @util.has_docstring def ip_eom(self, options=None, **kwargs): return geom.IP_GEOM(self, options=options, **kwargs) - @util.has_docstring def ea_eom(self, options=None, **kwargs): return geom.EA_GEOM(self, options=options, **kwargs) - @util.has_docstring def ee_eom(self, options=None, **kwargs): return geom.EE_GEOM(self, options=options, **kwargs) @property - @util.has_docstring def xi(self): if self.options.shift: xi = lib.einsum("Iii->I", self.g.boo) @@ -474,6 +461,5 @@ def xi(self): return xi @property - @util.has_docstring def spin_type(self): return "G" diff --git a/ebcc/geom.py b/ebcc/geom.py index c3960ff1..263a6425 100644 --- a/ebcc/geom.py +++ b/ebcc/geom.py @@ -3,11 +3,10 @@ from ebcc import reom, util -class GEOM(reom.REOM, metaclass=util.InheritDocstrings): +class GEOM(reom.REOM): """Generalised equation-of-motion base class.""" @property - @util.has_docstring def name(self): return self.excitation_type.upper() + "-GEOM-" + self.ebcc.name diff --git a/ebcc/rebcc.py b/ebcc/rebcc.py index dff7ee72..8aee28dd 100644 --- a/ebcc/rebcc.py +++ b/ebcc/rebcc.py @@ -17,12 +17,7 @@ from ebcc.logging import ANSI from ebcc.precision import types from ebcc.space import Space - - -class EBCC: - """Base class for EBCC.""" - - pass +from ebcc.base import EBCC @dataclasses.dataclass diff --git a/ebcc/reom.py b/ebcc/reom.py index 1c7ce752..fd34a2ec 100644 --- a/ebcc/reom.py +++ b/ebcc/reom.py @@ -9,12 +9,7 @@ from ebcc import util from ebcc.logging import ANSI from ebcc.precision import types - - -class EOM: - """Base class for equation-of-motion methods.""" - - pass +from ebcc.base import EOM @dataclasses.dataclass @@ -312,24 +307,20 @@ def nvir(self): return self.ebcc.nvir -class IP_REOM(REOM, metaclass=util.InheritDocstrings): +class IP_REOM(REOM): """Restricted equation-of-motion class for ionisation potentials.""" - @util.has_docstring def amplitudes_to_vector(self, *amplitudes): return self.ebcc.excitations_to_vector_ip(*amplitudes) - @util.has_docstring def vector_to_amplitudes(self, vector): return self.ebcc.vector_to_excitations_ip(vector) - @util.has_docstring def matvec(self, vector, eris=None): amplitudes = self.vector_to_amplitudes(vector) amplitudes = self.ebcc.hbar_matvec_ip(*amplitudes, eris=eris) return self.amplitudes_to_vector(*amplitudes) - @util.has_docstring def diag(self, eris=None): parts = [] @@ -348,7 +339,6 @@ def _bras(self, eris=None): def _kets(self, eris=None): return self.ebcc.make_ip_mom_kets(eris=eris) - @util.has_docstring def bras(self, eris=None): bras_raw = list(self._bras(eris=eris)) bras = np.array( @@ -356,7 +346,6 @@ def bras(self, eris=None): ) return bras - @util.has_docstring def kets(self, eris=None): kets_raw = list(self._kets(eris=eris)) kets = np.array( @@ -365,29 +354,24 @@ def kets(self, eris=None): return kets @property - @util.has_docstring def excitation_type(self): return "ip" -class EA_REOM(REOM, metaclass=util.InheritDocstrings): +class EA_REOM(REOM): """Equation-of-motion class for electron affinities.""" - @util.has_docstring def amplitudes_to_vector(self, *amplitudes): return self.ebcc.excitations_to_vector_ea(*amplitudes) - @util.has_docstring def vector_to_amplitudes(self, vector): return self.ebcc.vector_to_excitations_ea(vector) - @util.has_docstring def matvec(self, vector, eris=None): amplitudes = self.vector_to_amplitudes(vector) amplitudes = self.ebcc.hbar_matvec_ea(*amplitudes, eris=eris) return self.amplitudes_to_vector(*amplitudes) - @util.has_docstring def diag(self, eris=None): parts = [] @@ -406,7 +390,6 @@ def _bras(self, eris=None): def _kets(self, eris=None): return self.ebcc.make_ea_mom_kets(eris=eris) - @util.has_docstring def bras(self, eris=None): bras_raw = list(self._bras(eris=eris)) bras = np.array( @@ -414,7 +397,6 @@ def bras(self, eris=None): ) return bras - @util.has_docstring def kets(self, eris=None): kets_raw = list(self._kets(eris=eris)) kets = np.array( @@ -423,29 +405,24 @@ def kets(self, eris=None): return kets @property - @util.has_docstring def excitation_type(self): return "ea" -class EE_REOM(REOM, metaclass=util.InheritDocstrings): +class EE_REOM(REOM): """Equation-of-motion class for neutral excitations.""" - @util.has_docstring def amplitudes_to_vector(self, *amplitudes): return self.ebcc.excitations_to_vector_ee(*amplitudes) - @util.has_docstring def vector_to_amplitudes(self, vector): return self.ebcc.vector_to_excitations_ee(vector) - @util.has_docstring def matvec(self, vector, eris=None): amplitudes = self.vector_to_amplitudes(vector) amplitudes = self.ebcc.hbar_matvec_ee(*amplitudes, eris=eris) return self.amplitudes_to_vector(*amplitudes) - @util.has_docstring def diag(self, eris=None): parts = [] @@ -457,7 +434,6 @@ def diag(self, eris=None): return self.amplitudes_to_vector(*parts) - @util.has_docstring def bras(self, eris=None): bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) bras = np.array( @@ -468,7 +444,6 @@ def bras(self, eris=None): ) return bras - @util.has_docstring def kets(self, eris=None): kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) kets = np.array( @@ -482,7 +457,6 @@ def kets(self, eris=None): ) return kets - @util.has_docstring def moments(self, nmom, eris=None, amplitudes=None, hermitise=True, diagonal_only=True): """Construct the moments of the EOM Hamiltonian.""" @@ -519,6 +493,5 @@ def moments(self, nmom, eris=None, amplitudes=None, hermitise=True, diagonal_onl return moments @property - @util.has_docstring def excitation_type(self): return "ee" diff --git a/ebcc/uebcc.py b/ebcc/uebcc.py index 19b2f2d5..352adccc 100644 --- a/ebcc/uebcc.py +++ b/ebcc/uebcc.py @@ -12,8 +12,7 @@ from ebcc.space import Space -@util.has_docstring -class UEBCC(rebcc.REBCC, metaclass=util.InheritDocstrings): +class UEBCC(rebcc.REBCC): ERIs = UERIs Fock = UFock CDERIs = UCDERIs @@ -118,7 +117,6 @@ def _pack_codegen_kwargs(self, *extra_kwargs, eris=None): return kwargs - @util.has_docstring def init_space(self): space = ( Space( @@ -135,7 +133,6 @@ def init_space(self): return space - @util.has_docstring def init_amps(self, eris=None): eris = self.get_eris(eris) amplitudes = util.Namespace() @@ -196,7 +193,6 @@ def init_amps(self, eris=None): return amplitudes - @util.has_docstring def init_lams(self, amplitudes=None): if amplitudes is None: amplitudes = self.amplitudes @@ -228,7 +224,6 @@ def init_lams(self, amplitudes=None): return lambdas - @util.has_docstring def update_amps(self, eris=None, amplitudes=None): func, kwargs = self._load_function( "update_amps", @@ -268,7 +263,6 @@ def update_amps(self, eris=None, amplitudes=None): return res - @util.has_docstring def update_lams(self, eris=None, amplitudes=None, lambdas=None, lambdas_pert=None): func, kwargs = self._load_function( "update_lams", @@ -314,7 +308,6 @@ def update_lams(self, eris=None, amplitudes=None, lambdas=None, lambdas_pert=Non return res - @util.has_docstring def make_rdm1_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): func, kwargs = self._load_function( "make_rdm1_f", @@ -331,7 +324,6 @@ def make_rdm1_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): return dm - @util.has_docstring def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): func, kwargs = self._load_function( "make_rdm2_f", @@ -358,7 +350,6 @@ def transpose2(dm): return dm - @util.has_docstring def make_eb_coup_rdm( self, eris=None, @@ -391,7 +382,6 @@ def make_eb_coup_rdm( return dm_eb - @util.has_docstring def get_mean_field_G(self): val = lib.einsum("Ipp->I", self.g.aa.boo) val += lib.einsum("Ipp->I", self.g.bb.boo) @@ -404,7 +394,6 @@ def get_mean_field_G(self): return val - @util.has_docstring def get_g(self, g): if np.array(g).ndim != 4: g = np.array([g, g]) @@ -441,7 +430,6 @@ def __getitem__(selffer, key): return gs @property - @util.has_docstring def bare_fock(self): fock = lib.einsum( "npq,npi,nqj->nij", @@ -453,7 +441,6 @@ def bare_fock(self): return fock @property - @util.has_docstring def xi(self): if self.options.shift: xi = lib.einsum("Iii->I", self.g.aa.boo) @@ -466,7 +453,6 @@ def xi(self): return xi - @util.has_docstring def get_fock(self): return self.Fock(self, array=(self.bare_fock.aa, self.bare_fock.bb)) @@ -496,19 +482,15 @@ def get_eris(self, eris=None): else: return eris - @util.has_docstring def ip_eom(self, options=None, **kwargs): return ueom.IP_UEOM(self, options=options, **kwargs) - @util.has_docstring def ea_eom(self, options=None, **kwargs): return ueom.EA_UEOM(self, options=options, **kwargs) - @util.has_docstring def ee_eom(self, options=None, **kwargs): return ueom.EE_UEOM(self, options=options, **kwargs) - @util.has_docstring def amplitudes_to_vector(self, amplitudes): vectors = [] @@ -529,7 +511,6 @@ def amplitudes_to_vector(self, amplitudes): return np.concatenate(vectors) - @util.has_docstring def vector_to_amplitudes(self, vector): amplitudes = util.Namespace() i0 = 0 @@ -569,7 +550,6 @@ def vector_to_amplitudes(self, vector): return amplitudes - @util.has_docstring def lambdas_to_vector(self, lambdas): vectors = [] @@ -592,7 +572,6 @@ def lambdas_to_vector(self, lambdas): return np.concatenate(vectors) - @util.has_docstring def vector_to_lambdas(self, vector): lambdas = util.Namespace() i0 = 0 @@ -636,7 +615,6 @@ def vector_to_lambdas(self, vector): return lambdas - @util.has_docstring def excitations_to_vector_ip(self, *excitations): vectors = [] m = 0 @@ -656,7 +634,6 @@ def excitations_to_vector_ip(self, *excitations): return np.concatenate(vectors) - @util.has_docstring def excitations_to_vector_ea(self, *excitations): vectors = [] m = 0 @@ -677,7 +654,6 @@ def excitations_to_vector_ea(self, *excitations): return np.concatenate(vectors) - @util.has_docstring def excitations_to_vector_ee(self, *excitations): vectors = [] m = 0 @@ -697,7 +673,6 @@ def excitations_to_vector_ee(self, *excitations): return np.concatenate(vectors) - @util.has_docstring def vector_to_excitations_ip(self, vector): excitations = [] i0 = 0 @@ -730,7 +705,6 @@ def vector_to_excitations_ip(self, vector): return tuple(excitations) - @util.has_docstring def vector_to_excitations_ea(self, vector): excitations = [] i0 = 0 @@ -763,7 +737,6 @@ def vector_to_excitations_ea(self, vector): return tuple(excitations) - @util.has_docstring def vector_to_excitations_ee(self, vector): excitations = [] i0 = 0 @@ -796,23 +769,19 @@ def vector_to_excitations_ee(self, vector): return tuple(excitations) @property - @util.has_docstring def spin_type(self): return "U" @property - @util.has_docstring def nmo(self): assert self.mo_occ[0].size == self.mo_occ[1].size return self.mo_occ[0].size @property - @util.has_docstring def nocc(self): return tuple(np.sum(mo_occ > 0) for mo_occ in self.mo_occ) @property - @util.has_docstring def nvir(self): return tuple(self.nmo - nocc for nocc in self.nocc) diff --git a/ebcc/ueom.py b/ebcc/ueom.py index e67319cf..c95adec2 100644 --- a/ebcc/ueom.py +++ b/ebcc/ueom.py @@ -7,7 +7,7 @@ from ebcc.precision import types -class UEOM(reom.REOM, metaclass=util.InheritDocstrings): +class UEOM(reom.REOM): """Unrestricted equation-of-motion base class.""" def _argsort_guess(self, diag): @@ -23,7 +23,6 @@ def _argsort_guess(self, diag): def _quasiparticle_weight(self, r1): return np.linalg.norm(r1.a) ** 2 + np.linalg.norm(r1.b) ** 2 - @util.has_docstring def moments(self, nmom, eris=None, amplitudes=None, hermitise=True): if eris is None: eris = self.ebcc.get_eris() @@ -55,10 +54,9 @@ def moments(self, nmom, eris=None, amplitudes=None, hermitise=True): return moments -class IP_UEOM(UEOM, reom.IP_REOM, metaclass=util.InheritDocstrings): +class IP_UEOM(UEOM, reom.IP_REOM): """Unrestricted equation-of-motion class for ionisation potentials.""" - @util.has_docstring def diag(self, eris=None): parts = [] @@ -74,7 +72,6 @@ def diag(self, eris=None): return self.amplitudes_to_vector(*parts) - @util.has_docstring def bras(self, eris=None): bras_raw = list(self._bras(eris=eris)) bras = util.Namespace(a=[], b=[]) @@ -123,7 +120,6 @@ def bras(self, eris=None): return bras - @util.has_docstring def kets(self, eris=None): kets_raw = list(self._kets(eris=eris)) kets = util.Namespace(a=[], b=[]) @@ -174,10 +170,9 @@ def kets(self, eris=None): return kets -class EA_UEOM(UEOM, reom.EA_REOM, metaclass=util.InheritDocstrings): +class EA_UEOM(UEOM, reom.EA_REOM): """Unrestricted equation-of-motion class for electron affinities.""" - @util.has_docstring def diag(self, eris=None): parts = [] @@ -193,7 +188,6 @@ def diag(self, eris=None): return self.amplitudes_to_vector(*parts) - @util.has_docstring def bras(self, eris=None): bras_raw = list(self._bras(eris=eris)) bras = util.Namespace(a=[], b=[]) @@ -242,7 +236,6 @@ def bras(self, eris=None): return bras - @util.has_docstring def kets(self, eris=None): kets_raw = list(self._kets(eris=eris)) kets = util.Namespace(a=[], b=[]) @@ -293,13 +286,12 @@ def kets(self, eris=None): return kets -class EE_UEOM(UEOM, reom.EE_REOM, metaclass=util.InheritDocstrings): +class EE_UEOM(UEOM, reom.EE_REOM): """Unrestricted equation-of-motion class for neutral excitations.""" def _quasiparticle_weight(self, r1): return np.linalg.norm(r1.aa) ** 2 + np.linalg.norm(r1.bb) ** 2 - @util.has_docstring def diag(self, eris=None): parts = [] @@ -314,7 +306,6 @@ def diag(self, eris=None): return self.amplitudes_to_vector(*parts) - @util.has_docstring def bras(self, eris=None): # pragma: no cover raise util.ModelNotImplemented("EE moments for UEBCC not working.") @@ -378,7 +369,6 @@ def bras(self, eris=None): # pragma: no cover return bras - @util.has_docstring def kets(self, eris=None): # pragma: no cover raise util.ModelNotImplemented("EE moments for UEBCC not working.") @@ -439,7 +429,6 @@ def kets(self, eris=None): # pragma: no cover return kets - @util.has_docstring def moments( self, nmom, diff --git a/ebcc/util/__init__.py b/ebcc/util/__init__.py index d8477f67..94d89101 100644 --- a/ebcc/util/__init__.py +++ b/ebcc/util/__init__.py @@ -1,8 +1,7 @@ """Utilities.""" from ebcc.util.einsumfunc import direct_sum, dot, einsum -from ebcc.util.inheritance import InheritDocstrings, Inherited, InheritedType, _mro, has_docstring -from ebcc.util.misc import ModelNotImplemented, Namespace, Timer +from ebcc.util.misc import Inherited, ModelNotImplemented, Namespace, Timer from ebcc.util.permutations import ( antisymmetrise_array, combine_subscripts, diff --git a/ebcc/util/inheritance.py b/ebcc/util/inheritance.py deleted file mode 100644 index baa53f5a..00000000 --- a/ebcc/util/inheritance.py +++ /dev/null @@ -1,80 +0,0 @@ -"""Tools for class inheritance.""" - - -class InheritedType: - """Type for an inherited variable in `Options` classes.""" - - pass - - -Inherited = InheritedType() - - -class InheritDocstrings(type): - """ - Metaclass to inherit docstrings from superclasses. All attributes which - are public (no underscore prefix) are updated with the docstring of the - first superclass in the MRO containing that attribute with a docstring. - - Additionally checks that all methods are documented at runtime. - """ - - def __new__(cls, name, bases, attrs): - """Create an instance of the class with inherited docstrings.""" - - for key, val in attrs.items(): - if key.startswith("_") or val.__doc__ is not None: - continue - - for supcls in _mro(*bases): - supval = getattr(supcls, key, None) - if supval is None: - continue - val.__doc__ = supval.__doc__ - break - else: - raise RuntimeError("Method {} does not exist in superclass".format(key)) - - if val.__doc__ is None: - raise RuntimeError("Could not find docstring for {}".format(key)) - - attrs[key] = val - - return super().__new__(cls, name, bases, attrs) - - -def has_docstring(obj): - """ - Decorate a function or class to inform a static analyser that it has a - docstring even if one is not visible, for example via inheritance. - """ - return obj - - -def _mro(*bases): - """Find the method resolution order of bases using the C3 algorithm.""" - - seqs = [list(x.__mro__) for x in bases] + [list(bases)] - res = [] - - while True: - non_empty = list(filter(None, seqs)) - if not non_empty: - return tuple(res) - - for seq in non_empty: - candidate = seq[0] - not_head = [s for s in non_empty if candidate in s[1:]] - if not_head: - candidate = None - else: - break - - if not candidate: - raise TypeError("Inconsistent hierarchy") - - res.append(candidate) - - for seq in non_empty: - if seq[0] == candidate: - del seq[0] diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index d9ec2951..3fdbcaf2 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -3,6 +3,15 @@ import time +class InheritedType: + """Type for an inherited variable.""" + + pass + + +Inherited = InheritedType() + + class ModelNotImplemented(NotImplementedError): """Error for unsupported models.""" From 32d4610ff4bca19acf3496cd5bdc0b7a4ea7e812 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sun, 28 Jul 2024 00:05:54 +0100 Subject: [PATCH 02/37] Progress on base stuff --- ebcc/base.py | 83 ++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 81 insertions(+), 2 deletions(-) diff --git a/ebcc/base.py b/ebcc/base.py index f8e879b6..253e117c 100644 --- a/ebcc/base.py +++ b/ebcc/base.py @@ -1,14 +1,93 @@ """Base classes.""" +from __future__ import annotations + from abc import ABC, abstractmethod +from typing import TYPE_CHECKING from ebcc import util +if TYPE_CHECKING: + from logging import Logger + from typing import Optional, Union, Any + from dataclasses import dataclass + + from pyscf.scf import SCF + + from ebcc.ansatz import Ansatz + from ebcc.space import Space + class EBCC(ABC): - """Base class for electron-boson coupled cluster.""" + """Base class for electron-boson coupled cluster. + + """ + + # Types + Options: dataclass + ERIs: ERIs + Fock: Fock + CDERIs: ERIs + Brueckner: BruecknerEBCC + + # Metadata + log: Logger + options: dataclass + + # Mean-field + mf: SCF + _mo_coeff: Any + _mo_occ: Any + fock: Any + + # Ansatz + ansatz: Ansatz + _eqns: ModuleType + space: Any + + # Bosons + omega: Any + bare_g: Any + bare_G: Any + + # Results + e_corr: float + amplitudes: Any + converged: bool + + def __init__( + self, + mf: SCF, + log: Optional[Logger] = None, + ansatz: Optional[Union[Ansatz, str]] = "CCSD", + space: Optional[Any] = None, + omega: Optional[Any] = None, + g: Optional[Any] = None, + G: Optional[Any] = None, + mo_coeff: Optional[Any] = None, + mo_occ: Optional[Any] = None, + fock: Optional[Any] = None, + **kwargs, + ): + """Initialize the EBCC object. + + Args: + mf: PySCF mean-field object. + log: Log to write output to. Default is the global logger, outputting to `stderr`. + ansatz: Overall ansatz. + space: Space containing the frozen, correlated, and active fermionic spaces. Default + assumes all electrons are correlated. + omega: Bosonic frequencies. + g: Electron-boson coupling matrix corresponding to the bosonic annihilation operator + :math:`g_{bpq} p^\dagger q b`. The creation part is assumed to be the fermionic + conjugate transpose to retain Hermiticity in the Hamiltonian. + G: Boson non-conserving term :math:`G_{b} (b^\dagger + b)`. + mo_coeff: Molecular orbital coefficients. Default is the mean-field coefficients. + mo_occ: Molecular orbital occupation numbers. Default is the mean-field occupation. + fock: Fock matrix. Default is the mean-field Fock matrix. + **kwargs: Additional keyword arguments used to update `options`. + """ - pass class EOM(ABC): From 61df59ce852f48ddc7a289dca3ca5f7cef3a70d5 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sun, 28 Jul 2024 01:10:17 +0100 Subject: [PATCH 03/37] Moving to submodule --- ebcc/base.py | 72 ----------- ebcc/cc/__init__.py | 1 + ebcc/cc/base.py | 154 ++++++++++++++++++++++++ ebcc/rebcc.py | 284 +++----------------------------------------- 4 files changed, 169 insertions(+), 342 deletions(-) create mode 100644 ebcc/cc/__init__.py create mode 100644 ebcc/cc/base.py diff --git a/ebcc/base.py b/ebcc/base.py index 253e117c..1f62681d 100644 --- a/ebcc/base.py +++ b/ebcc/base.py @@ -18,78 +18,6 @@ from ebcc.space import Space -class EBCC(ABC): - """Base class for electron-boson coupled cluster. - - """ - - # Types - Options: dataclass - ERIs: ERIs - Fock: Fock - CDERIs: ERIs - Brueckner: BruecknerEBCC - - # Metadata - log: Logger - options: dataclass - - # Mean-field - mf: SCF - _mo_coeff: Any - _mo_occ: Any - fock: Any - - # Ansatz - ansatz: Ansatz - _eqns: ModuleType - space: Any - - # Bosons - omega: Any - bare_g: Any - bare_G: Any - - # Results - e_corr: float - amplitudes: Any - converged: bool - - def __init__( - self, - mf: SCF, - log: Optional[Logger] = None, - ansatz: Optional[Union[Ansatz, str]] = "CCSD", - space: Optional[Any] = None, - omega: Optional[Any] = None, - g: Optional[Any] = None, - G: Optional[Any] = None, - mo_coeff: Optional[Any] = None, - mo_occ: Optional[Any] = None, - fock: Optional[Any] = None, - **kwargs, - ): - """Initialize the EBCC object. - - Args: - mf: PySCF mean-field object. - log: Log to write output to. Default is the global logger, outputting to `stderr`. - ansatz: Overall ansatz. - space: Space containing the frozen, correlated, and active fermionic spaces. Default - assumes all electrons are correlated. - omega: Bosonic frequencies. - g: Electron-boson coupling matrix corresponding to the bosonic annihilation operator - :math:`g_{bpq} p^\dagger q b`. The creation part is assumed to be the fermionic - conjugate transpose to retain Hermiticity in the Hamiltonian. - G: Boson non-conserving term :math:`G_{b} (b^\dagger + b)`. - mo_coeff: Molecular orbital coefficients. Default is the mean-field coefficients. - mo_occ: Molecular orbital occupation numbers. Default is the mean-field occupation. - fock: Fock matrix. Default is the mean-field Fock matrix. - **kwargs: Additional keyword arguments used to update `options`. - """ - - - class EOM(ABC): """Base class for equation-of-motion methods.""" diff --git a/ebcc/cc/__init__.py b/ebcc/cc/__init__.py new file mode 100644 index 00000000..f6b49bd0 --- /dev/null +++ b/ebcc/cc/__init__.py @@ -0,0 +1 @@ +"""Coupled cluster solvers.""" diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py new file mode 100644 index 00000000..7eb15968 --- /dev/null +++ b/ebcc/cc/base.py @@ -0,0 +1,154 @@ +"""Base classes for `ebcc.cc`.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from ebcc import default_log, init_logging +from ebcc import numpy as np +from ebcc import util +from ebcc.ansatz import Ansatz +from ebcc.space import Space +from ebcc.logging import ANSI +from ebcc.precision import types + +if TYPE_CHECKING: + from logging import Logger + from typing import Optional, Union, Any + from dataclasses import dataclass + + from pyscf.scf import SCF + + +class EBCC(ABC): + """Base class for electron-boson coupled cluster. + + Attributes: + mf: PySCF mean-field object. + log: Log to write output to. + options: Options for the EBCC calculation. + e_corr: Correlation energy. + amplitudes: Cluster amplitudes. + converged: Convergence flag. + lambdas: Cluster lambda amplitudes. + converged_lambda: Lambda convergence flag. + name: Name of the method. + """ + + # Types + Options: dataclass + ERIs: ERIs + Fock: Fock + CDERIs: ERIs + Brueckner: BruecknerEBCC + + def __init__( + self, + mf: SCF, + log: Optional[Logger] = None, + ansatz: Optional[Union[Ansatz, str]] = "CCSD", + options: Optional[dataclass] = None, + space: Optional[Any] = None, + omega: Optional[Any] = None, + g: Optional[Any] = None, + G: Optional[Any] = None, + mo_coeff: Optional[Any] = None, + mo_occ: Optional[Any] = None, + fock: Optional[Any] = None, + **kwargs, + ): + """Initialize the EBCC object. + + Args: + mf: PySCF mean-field object. + log: Log to write output to. Default is the global logger, outputting to `stderr`. + ansatz: Overall ansatz. + space: Space containing the frozen, correlated, and active fermionic spaces. Default + assumes all electrons are correlated. + omega: Bosonic frequencies. + g: Electron-boson coupling matrix corresponding to the bosonic annihilation operator + :math:`g_{bpq} p^\dagger q b`. The creation part is assumed to be the fermionic + conjugate transpose to retain Hermiticity in the Hamiltonian. + G: Boson non-conserving term :math:`G_{b} (b^\dagger + b)`. + mo_coeff: Molecular orbital coefficients. Default is the mean-field coefficients. + mo_occ: Molecular orbital occupation numbers. Default is the mean-field occupation. + fock: Fock matrix. Default is the mean-field Fock matrix. + **kwargs: Additional keyword arguments used to update `options`. + """ + # Options: + if options is None: + options = self.Options() + self.options = options + for key, val in kwargs.items(): + setattr(self.options, key, val) + + # Parameters: + self.log = default_log if log is None else log + self.mf = self._convert_mf(mf) + self._mo_coeff = np.asarray(mo_coeff).astype(types[float]) if mo_coeff is not None else None + self._mo_occ = np.asarray(mo_occ).astype(types[float]) if mo_occ is not None else None + + # Ansatz: + if isinstance(ansatz, Ansatz): + self.ansatz = ansatz + else: + self.ansatz = Ansatz.from_string( + ansatz, density_fitting=getattr(self.mf, "with_df", None) is not None + ) + self._eqns = self.ansatz._get_eqns(self.spin_type) + + # Space: + if space is not None: + self.space = space + else: + self.space = self.init_space() + + # Boson parameters: + if bool(self.fermion_coupling_rank) != bool(self.boson_coupling_rank): + raise ValueError( + "Fermionic and bosonic coupling ranks must both be zero, or both non-zero." + ) + self.omega = omega.astype(types[float]) if omega is not None else None + self.bare_g = g.astype(types[float]) if g is not None else None + self.bare_G = G.astype(types[float]) if G is not None else None + if self.boson_ansatz != "": + self.g = self.get_g(g) + self.G = self.get_mean_field_G() + if self.options.shift: + self.log.info(" > Energy shift due to polaritonic basis: %.10f", self.const) + else: + assert self.nbos == 0 + self.options.shift = False + self.g = None + self.G = None + + # Fock matrix: + if fock is None: + self.fock = self.get_fock() + else: + self.fock = fock + + # Attributes: + self.e_corr = None + self.amplitudes = None + self.converged = False + self.lambdas = None + self.converged_lambda = False + + # Logging: + init_logging(self.log) + self.log.info(f"\n{ANSI.B}{ANSI.U}{self.name}{ANSI.R}") + self.log.debug(f"{ANSI.B}{'*' * len(self.name)}{ANSI.R}") + self.log.debug("") + self.log.info(f"{ANSI.B}Options{ANSI.R}:") + self.log.info(f" > e_tol: {ANSI.y}{self.options.e_tol}{ANSI.R}") + self.log.info(f" > t_tol: {ANSI.y}{self.options.t_tol}{ANSI.R}") + self.log.info(f" > max_iter: {ANSI.y}{self.options.max_iter}{ANSI.R}") + self.log.info(f" > diis_space: {ANSI.y}{self.options.diis_space}{ANSI.R}") + self.log.info(f" > damping: {ANSI.y}{self.options.damping}{ANSI.R}") + self.log.debug("") + self.log.info(f"{ANSI.B}Ansatz{ANSI.R}: {ANSI.m}{self.ansatz}{ANSI.R}") + self.log.debug("") + self.log.info(f"{ANSI.B}Space{ANSI.R}: {ANSI.m}{self.space}{ANSI.R}") + self.log.debug("") diff --git a/ebcc/rebcc.py b/ebcc/rebcc.py index 8aee28dd..79db4d4f 100644 --- a/ebcc/rebcc.py +++ b/ebcc/rebcc.py @@ -17,7 +17,7 @@ from ebcc.logging import ANSI from ebcc.precision import types from ebcc.space import Space -from ebcc.base import EBCC +from ebcc.cc.base import EBCC @dataclasses.dataclass @@ -56,283 +56,27 @@ class Options: class REBCC(EBCC): - r""" - Restricted electron-boson coupled cluster class. - - Parameters - ---------- - mf : pyscf.scf.hf.SCF - PySCF mean-field object. - log : logging.Logger, optional - Log to print output to. Default value is the global logger which - outputs to `sys.stderr`. - ansatz : str or Ansatz, optional - Overall ansatz, as a string representation or an `Ansatz` object. - If `None`, construct from the individual ansatz parameters, - otherwise override them. Default value is `None`. - space : Space, optional - Space object defining the size of frozen, correlated and active - fermionic spaces. If `None`, all fermionic degrees of freedom are - assumed correlated. Default value is `None`. - omega : numpy.ndarray (nbos,), optional - Bosonic frequencies. Default value is `None`. - g : numpy.ndarray (nbos, nmo, nmo), optional - Electron-boson coupling matrix corresponding to the bosonic - annihilation operator i.e. - - .. math:: g_{xpq} p^\dagger q b - - The creation part is assume to be the fermionic transpose of this - tensor to retain hermiticity in the overall Hamiltonian. Default - value is `None`. - G : numpy.ndarray (nbos,), optional - Boson non-conserving term of the Hamiltonian i.e. - - .. math:: G_x (b^\dagger + b) - - Default value is `None`. - mo_coeff : numpy.ndarray, optional - Molecular orbital coefficients. Default value is `mf.mo_coeff`. - mo_occ : numpy.ndarray, optional - Molecular orbital occupancies. Default value is `mf.mo_occ`. - fock : util.Namespace, optional - Fock input. Default value is calculated using `get_fock()`. - options : dataclasses.dataclass, optional - Object containing the options. Default value is `Options()`. - **kwargs : dict - Additional keyword arguments used to update `options`. - - Attributes - ---------- - mf : pyscf.scf.hf.SCF - PySCF mean-field object. - log : logging.Logger - Log to print output to. - options : dataclasses.dataclass - Object containing the options. - e_corr : float - Correlation energy. - amplitudes : Namespace - Cluster amplitudes. - lambdas : Namespace - Cluster lambda amplitudes. - converged : bool - Whether the coupled cluster equations converged. - converged_lambda : bool - Whether the lambda coupled cluster equations converged. - omega : numpy.ndarray (nbos,) - Bosonic frequencies. - g : util.Namespace - Namespace containing blocks of the electron-boson coupling matrix. - Each attribute should be a length-3 string of `b`, `o` or `v` - signifying whether the corresponding axis is bosonic, occupied, or - virtual. - G : numpy.ndarray (nbos,) - Mean-field boson non-conserving term of the Hamiltonian. - bare_G : numpy.ndarray (nbos,) - Boson non-conserving term of the Hamiltonian. - fock : util.Namespace - Namespace containing blocks of the Fock matrix. Each attribute - should be a length-2 string of `o` or `v` signifying whether the - corresponding axis is occupied or virtual. - bare_fock : numpy.ndarray (nmo, nmo) - The mean-field Fock matrix in the MO basis. - xi : numpy.ndarray (nbos,) - Shift in bosonic operators to diagonalise the phononic Hamiltonian. - const : float - Shift in the energy from moving to polaritonic basis. - name : str - Name of the method. - - Methods - ------- - init_amps(eris=None) - Initialise the amplitudes. - init_lams(amplitudes=None) - Initialise the lambda amplitudes. - kernel(eris=None) - Run the coupled cluster calculation. - solve_lambda(amplitudes=None, eris=None) - Solve the lambda coupled cluster equations. - energy(eris=None, amplitudes=None) - Compute the correlation energy. - update_amps(eris=None, amplitudes=None) - Update the amplitudes. - update_lams(eris=None, amplitudes=None, lambdas=None) - Update the lambda amplitudes. - make_sing_b_dm(eris=None, amplitudes=None, lambdas=None) - Build the single boson density matrix. - make_rdm1_b(eris=None, amplitudes=None, lambdas=None, unshifted=None, - hermitise=None) - Build the bosonic one-particle reduced density matrix. - make_rdm1_f(eris=None, amplitudes=None, lambdas=None, hermitise=None) - Build the fermionic one-particle reduced density matrix. - make_rdm2_f(eris=None, amplitudes=None, lambdas=None, hermitise=None) - Build the fermionic two-particle reduced density matrix. - make_eb_coup_rdm(eris=None, amplitudes=None, lambdas=None, - unshifted=True, hermitise=True) - Build the electron-boson coupling reduced density matrices. - hbar_matvec_ip(r1, r2, eris=None, amplitudes=None) - Compute the product between a state vector and the EOM Hamiltonian - for the IP. - hbar_matvec_ea(r1, r2, eris=None, amplitudes=None) - Compute the product between a state vector and the EOM Hamiltonian - for the EA. - hbar_matvec_ee(r1, r2, eris=None, amplitudes=None) - Compute the product between a state vector and the EOM Hamiltonian - for the EE. - make_ip_mom_bras(eris=None, amplitudes=None, lambdas=None) - Get the bra IP vectors to construct EOM moments. - make_ea_mom_bras(eris=None, amplitudes=None, lambdas=None) - Get the bra EA vectors to construct EOM moments. - make_ee_mom_bras(eris=None, amplitudes=None, lambdas=None) - Get the bra EE vectors to construct EOM moments. - make_ip_mom_kets(eris=None, amplitudes=None, lambdas=None) - Get the ket IP vectors to construct EOM moments. - make_ea_mom_kets(eris=None, amplitudes=None, lambdas=None) - Get the ket EA vectors to construct EOM moments. - make_ee_mom_kets(eris=None, amplitudes=None, lambdas=None) - Get the ket EE vectors to construct EOM moments. - amplitudes_to_vector(amplitudes) - Construct a vector containing all of the amplitudes used in the - given ansatz. - vector_to_amplitudes(vector) - Construct all of the amplitudes used in the given ansatz from a - vector. - lambdas_to_vector(lambdas) - Construct a vector containing all of the lambda amplitudes used in - the given ansatz. - vector_to_lambdas(vector) - Construct all of the lambdas used in the given ansatz from a - vector. - excitations_to_vector_ip(*excitations) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the IP. - excitations_to_vector_ea(*excitations) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EA. - excitations_to_vector_ee(*excitations) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EE. - vector_to_excitations_ip(vector) - Construct all of the excitation amplitudes used in the given - ansatz from a vector for the IP. - vector_to_excitations_ea(vector) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EA. - vector_to_excitations_ee(vector) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EE. - get_mean_field_G() - Get the mean-field boson non-conserving term of the Hamiltonian. - get_g(g) - Get the blocks of the electron-boson coupling matrix corresponding - to the bosonic annihilation operator. - get_fock() - Get the blocks of the Fock matrix, shifted due to bosons where the - ansatz requires. - get_eris() - Get blocks of the ERIs. + """Restricted electron-boson coupled cluster. + + Attributes: + mf: PySCF mean-field object. + log: Log to write output to. + options: Options for the EBCC calculation. + e_corr: Correlation energy. + amplitudes: Cluster amplitudes. + converged: Convergence flag. + lambdas: Cluster lambda amplitudes. + converged_lambda: Lambda convergence flag. + name: Name of the method. """ + # Types Options = Options ERIs = RERIs Fock = RFock CDERIs = RCDERIs Brueckner = BruecknerREBCC - def __init__( - self, - mf, - log=None, - ansatz="CCSD", - space=None, - omega=None, - g=None, - G=None, - mo_coeff=None, - mo_occ=None, - fock=None, - options=None, - **kwargs, - ): - # Options: - if options is None: - options = self.Options() - self.options = options - for key, val in kwargs.items(): - setattr(self.options, key, val) - - # Parameters: - self.log = default_log if log is None else log - self.mf = self._convert_mf(mf) - self._mo_coeff = np.asarray(mo_coeff).astype(types[float]) if mo_coeff is not None else None - self._mo_occ = np.asarray(mo_occ).astype(types[float]) if mo_occ is not None else None - - # Ansatz: - if isinstance(ansatz, Ansatz): - self.ansatz = ansatz - else: - self.ansatz = Ansatz.from_string( - ansatz, density_fitting=getattr(self.mf, "with_df", None) is not None - ) - self._eqns = self.ansatz._get_eqns(self.spin_type) - - # Space: - if space is not None: - self.space = space - else: - self.space = self.init_space() - - # Boson parameters: - if bool(self.fermion_coupling_rank) != bool(self.boson_coupling_rank): - raise ValueError( - "Fermionic and bosonic coupling ranks must both be zero, or both non-zero." - ) - self.omega = omega.astype(types[float]) if omega is not None else None - self.bare_g = g.astype(types[float]) if g is not None else None - self.bare_G = G.astype(types[float]) if G is not None else None - if self.boson_ansatz != "": - self.g = self.get_g(g) - self.G = self.get_mean_field_G() - if self.options.shift: - self.log.info(" > Energy shift due to polaritonic basis: %.10f", self.const) - else: - assert self.nbos == 0 - self.options.shift = False - self.g = None - self.G = None - - # Fock matrix: - if fock is None: - self.fock = self.get_fock() - else: - self.fock = fock - - # Attributes: - self.e_corr = None - self.amplitudes = None - self.converged = False - self.lambdas = None - self.converged_lambda = False - - # Logging: - init_logging(self.log) - self.log.info(f"\n{ANSI.B}{ANSI.U}{self.name}{ANSI.R}") - self.log.debug(f"{ANSI.B}{'*' * len(self.name)}{ANSI.R}") - self.log.debug("") - self.log.info(f"{ANSI.B}Options{ANSI.R}:") - self.log.info(f" > e_tol: {ANSI.y}{self.options.e_tol}{ANSI.R}") - self.log.info(f" > t_tol: {ANSI.y}{self.options.t_tol}{ANSI.R}") - self.log.info(f" > max_iter: {ANSI.y}{self.options.max_iter}{ANSI.R}") - self.log.info(f" > diis_space: {ANSI.y}{self.options.diis_space}{ANSI.R}") - self.log.info(f" > damping: {ANSI.y}{self.options.damping}{ANSI.R}") - self.log.debug("") - self.log.info(f"{ANSI.B}Ansatz{ANSI.R}: {ANSI.m}{self.ansatz}{ANSI.R}") - self.log.debug("") - self.log.info(f"{ANSI.B}Space{ANSI.R}: {ANSI.m}{self.space}{ANSI.R}") - self.log.debug("") - def kernel(self, eris=None): """ Run the coupled cluster calculation. From 5b14bd1034bd60f085f808613c32930b55f42f02 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Wed, 31 Jul 2024 09:10:01 +0100 Subject: [PATCH 04/37] More base --- ebcc/cc/base.py | 258 ++++++++++++++++++++++++++++++++++++++++++++-- ebcc/rebcc.py | 130 +---------------------- ebcc/util/misc.py | 92 ++++++++++------- 3 files changed, 305 insertions(+), 175 deletions(-) diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 7eb15968..d6f9b4fe 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -4,6 +4,7 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING +from dataclasses import dataclass from ebcc import default_log, init_logging from ebcc import numpy as np @@ -16,12 +17,38 @@ if TYPE_CHECKING: from logging import Logger from typing import Optional, Union, Any - from dataclasses import dataclass - from pyscf.scf import SCF + from pyscf.scf import SCF # type: ignore + from ebcc.numpy.typing import NDArray # type: ignore + from ebcc.base import ERIs as BaseERIs + from ebcc.base import Fock as BaseFock + from ebcc.base import BruecknerEBCC as BaseBrueckner -class EBCC(ABC): + +@dataclass +class BaseOptions: + """Options for EBCC calculations. + + Args: + 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. + e_tol: Threshold for convergence in the correlation energy. + t_tol: Threshold for convergence in the amplitude norm. + max_iter: Maximum number of iterations. + diis_space: Number of amplitudes to use in DIIS extrapolation. + damping: Damping factor for DIIS extrapolation. + """ + shift: bool = True + e_tol: float = 1e-8 + t_tol: float = 1e-8 + max_iter: int = 200 + diis_space: int = 12 + damping: float = 0.0 + + +class BaseEBCC(ABC): """Base class for electron-boson coupled cluster. Attributes: @@ -37,18 +64,18 @@ class EBCC(ABC): """ # Types - Options: dataclass - ERIs: ERIs - Fock: Fock - CDERIs: ERIs - Brueckner: BruecknerEBCC + Options: type[BaseOptions] = BaseOptions + ERIs: type[BaseERIs] + Fock: type[BaseFock] + CDERIs: type[BaseERIs] + Brueckner: type[BaseBrueckner] def __init__( self, mf: SCF, log: Optional[Logger] = None, ansatz: Optional[Union[Ansatz, str]] = "CCSD", - options: Optional[dataclass] = None, + options: Optional[BaseOptions] = None, space: Optional[Any] = None, omega: Optional[Any] = None, g: Optional[Any] = None, @@ -152,3 +179,216 @@ def __init__( self.log.debug("") self.log.info(f"{ANSI.B}Space{ANSI.R}: {ANSI.m}{self.space}{ANSI.R}") self.log.debug("") + + @abstractmethod + @staticmethod + def _convert_mf(mf: SCF) -> SCF: + """Convert the mean-field object to the appropriate type.""" + pass + + @abstractmethod + @property + def spin_type(self) -> str: + """Get a string representation of the spin type.""" + pass + + @abstractmethod + @property + def name(self) -> str: + """Get the name of the method.""" + pass + + @property + def fermion_ansatz(self) -> str: + """Get a string representation of the fermion ansatz.""" + return self.ansatz.fermion_ansatz + + @property + def boson_ansatz(self) -> str: + """Get a string representation of the boson ansatz.""" + return self.ansatz.boson_ansatz + + @property + def fermion_coupling_rank(self) -> int: + """Get an integer representation of the fermion coupling rank.""" + return self.ansatz.fermion_coupling_rank + + @property + def boson_coupling_rank(self) -> int: + """Get an integer representation of the boson coupling rank.""" + return self.ansatz.boson_coupling_rank + + @abstractmethod + def init_space(self) -> Any: + """Initialise the fermionic space. + + Returns: + Fermionic space. + """ + pass + + @abstractmethod + def get_fock(self) -> Any: + """Get the Fock matrix. + + Returns: + Fock matrix. + """ + pass + + @abstractmethod + def get_eris(self, eris: Optional[Union[type[ERIs], NDArray[float]]]) -> Any: + """Get the electron repulsion integrals. + + Args: + eris: Input electron repulsion integrals. + + Returns: + Electron repulsion integrals. + """ + pass + + @abstractmethod + def get_g(self, g: NDArray[float]) -> Any: + """Get the blocks of the electron-boson coupling matrix. + + This matrix corresponds to the bosonic annihilation operator. + + Args: + g: Electron-boson coupling matrix. + + Returns: + Blocks of the electron-boson coupling matrix. + """ + pass + + @abstractmethod + def get_mean_field_G(self) -> Any: + """Get the mean-field boson non-conserving term. + + Returns: + Mean-field boson non-conserving term. + """ + pass + + def const(self) -> float: + """Get the shift in energy from moving to the polaritonic basis. + + Returns: + Constant energy shift due to the polaritonic basis. + """ + if self.options.shift: + return util.einsum("I,I->", self.omega, self.xi ** 2) + return 0.0 + + @abstractmethod + @property + def xi(self) -> NDArray[float]: + """Get the shift in the bosonic operators to diagonalise the photon Hamiltonian. + + Returns: + Shift in the bosonic operators. + """ + pass + + @property + def mo_coeff(self) -> NDArray[float]: + """Get the molecular orbital coefficients. + + Returns: + Molecular orbital coefficients. + """ + if self._mo_coeff is None: + return np.asarray(self.mf.mo_coeff).astype(types[float]) + return self._mo_coeff + + @property + def mo_occ(self) -> NDArray[float]: + """Get the molecular orbital occupation numbers. + + Returns: + Molecular orbital occupation numbers. + """ + if self._mo_occ is None: + return np.asarray(self.mf.mo_occ).astype(types[float]) + return self._mo_occ + + @abstractmethod + @property + def nmo(self) -> Any: + """Get the number of molecular orbitals. + + Returns: + Number of molecular orbitals. + """ + pass + + @abstractmethod + @property + def nocc(self) -> Any: + """Get the number of occupied molecular orbitals. + + Returns: + Number of occupied molecular orbitals. + """ + pass + + @abstractmethod + @property + def nvir(self) -> Any: + """Get the number of virtual molecular orbitals. + + Returns: + Number of virtual molecular orbitals. + """ + pass + + @property + def nbos(self) -> int: + """Get the number of bosonic modes. + + Returns: + Number of bosonic modes. + """ + if self.omega is None: + return 0 + return self.omega.shape[0] + + @property + def e_tot(self) -> float: + """Get the total energy (mean-field plus correlation). + + Returns: + Total energy. + """ + return types[float](self.mf.e_tot) + self.e_corr + + @property + def t1(self) -> Any: + """Get the T1 amplitudes.""" + return self.amplitudes["t1"] + + @property + def t2(self) -> Any: + """Get the T2 amplitudes.""" + return self.amplitudes["t2"] + + @property + def t3(self) -> Any: + """Get the T3 amplitudes.""" + return self.amplitudes["t3"] + + @property + def l1(self) -> Any: + """Get the L1 amplitudes.""" + return self.lambdas["l1"] + + @property + def l2(self) -> Any: + """Get the L2 amplitudes.""" + return self.lambdas["l2"] + + @property + def l3(self) -> Any: + """Get the L3 amplitudes.""" + return self.lambdas["l3"] diff --git a/ebcc/rebcc.py b/ebcc/rebcc.py index 79db4d4f..f9bf76ce 100644 --- a/ebcc/rebcc.py +++ b/ebcc/rebcc.py @@ -20,41 +20,6 @@ from ebcc.cc.base import EBCC -@dataclasses.dataclass -class Options: - """ - Options for EBCC calculations. - - Attributes - ---------- - shift : bool, optional - If `True`, 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. Default value is `True`. - e_tol : float, optional - Threshold for convergence in the correlation energy. Default value - is `1e-8`. - t_tol : float, optional - Threshold for convergence in the amplitude norm. Default value is - `1e-8`. - max_iter : int, optional - Maximum number of iterations. Default value is `200`. - diis_space : int, optional - Number of amplitudes to use in DIIS extrapolation. Default value is - `12`. - damping : float, optional - Damping factor for DIIS extrapolation. Default value is `0.0`. - """ - - shift: bool = True - e_tol: float = 1e-8 - t_tol: float = 1e-8 - max_iter: int = 200 - diis_space: int = 12 - damping: float = 0.0 - - class REBCC(EBCC): """Restricted electron-boson coupled cluster. @@ -71,7 +36,6 @@ class REBCC(EBCC): """ # Types - Options = Options ERIs = RERIs Fock = RFock CDERIs = RCDERIs @@ -1778,64 +1742,16 @@ def const(self): else: return 0.0 - @property - def fermion_ansatz(self): - """Get a string representation of the fermion ansatz.""" - return self.ansatz.fermion_ansatz - - @property - def boson_ansatz(self): - """Get a string representation of the boson ansatz.""" - return self.ansatz.boson_ansatz - - @property - def fermion_coupling_rank(self): - """Get an integer representation of the fermion coupling rank.""" - return self.ansatz.fermion_coupling_rank - - @property - def boson_coupling_rank(self): - """Get an integer representation of the boson coupling rank.""" - return self.ansatz.boson_coupling_rank - @property def name(self): - """Get a string representation of the method name.""" + """Get the name of the method.""" return self.spin_type + self.ansatz.name @property def spin_type(self): - """Get a string represent of the spin channel.""" + """Get a string representation of the spin type.""" return "R" - @property - def mo_coeff(self): - """ - Get the molecular orbital coefficients. - - Returns - ------- - mo_coeff : numpy.ndarray - Molecular orbital coefficients. - """ - if self._mo_coeff is None: - return np.asarray(self.mf.mo_coeff).astype(types[float]) - return self._mo_coeff - - @property - def mo_occ(self): - """ - Get the molecular orbital occupancies. - - Returns - ------- - mo_occ : numpy.ndarray - Molecular orbital occupancies. - """ - if self._mo_occ is None: - return np.asarray(self.mf.mo_occ).astype(types[float]) - return self._mo_occ - @property def nmo(self): """Get the number of MOs.""" @@ -1916,45 +1832,3 @@ def next_char(): energy_sum = lib.direct_sum(subscript, *energies) return energy_sum - - @property - def e_tot(self): - """ - Return the total energy (mean-field plus correlation). - - Returns - ------- - e_tot : float - Total energy. - """ - return types[float](self.mf.e_tot) + self.e_corr - - @property - def t1(self): - """Get the T1 amplitudes.""" - return self.amplitudes["t1"] - - @property - def t2(self): - """Get the T2 amplitudes.""" - return self.amplitudes["t2"] - - @property - def t3(self): - """Get the T3 amplitudes.""" - return self.amplitudes["t3"] - - @property - def l1(self): - """Get the L1 amplitudes.""" - return self.lambdas["l1"] - - @property - def l2(self): - """Get the L2 amplitudes.""" - return self.lambdas["l2"] - - @property - def l3(self): - """Get the L3 amplitudes.""" - return self.lambdas["l3"] diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index 3fdbcaf2..5b40d04e 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -1,7 +1,16 @@ """Miscellaneous utilities.""" +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import TYPE_CHECKING, Generic, TypeVar import time +if TYPE_CHECKING: + from typing import Any, Callable, Iterable, Iterator + +T = TypeVar("T") + class InheritedType: """Type for an inherited variable.""" @@ -18,7 +27,7 @@ class ModelNotImplemented(NotImplementedError): pass -class Namespace: +class Namespace(MutableMapping[str, T], Generic[T]): """ Replacement for SimpleNamespace, which does not trivially allow conversion to a dict for heterogenously nested objects. @@ -27,66 +36,73 @@ class Namespace: accessing the attribute directly. """ - def __init__(self, **kwargs): - self.__dict__["_keys"] = set() + _members: dict[str, T] + + def __init__(self, **kwargs: T): + """Initialise the namespace.""" + self.__dict__["_members"] = {} for key, val in kwargs.items(): - self[key] = val + self.__dict__["_members"][key] = val + + def __setitem__(self, key: str, val: T) -> None: + """Set an item.""" + self.__dict__["_members"][key] = value - def __setitem__(self, key, val): + def __setattr__(self, key: str, val: T) -> None: """Set an attribute.""" - self._keys.add(key) - self.__dict__[key] = val + return self.__setitem__(key, value) - def __getitem__(self, key): + def __getitem__(self, key: str) -> T: + """Get an item.""" + value: T = self.__dict__["_members"][key] + return value + + def __getattr__(self, key: str) -> T: """Get an attribute.""" - if key not in self._keys: - raise IndexError(key) - return self.__dict__[key] + return self.__getitem__(key) - def __delitem__(self, key): - """Delete an attribute.""" - if key not in self._keys: - raise IndexError(key) - del self.__dict__[key] + def __delitem__(self, key: str) -> None: + """Delete an item.""" + self._members.pop(key) - __setattr__ = __setitem__ + def __delattr__(self, key: str) -> None: + """Delete an attribute.""" + return self.__delitem__(key) - def __iter__(self): + def __iter__(self) -> Iterator[str]: """Iterate over the namespace as a dictionary.""" - yield from {key: self[key] for key in self._keys} + yield from self._members - def __eq__(self, other): + def __eq__(self, other: Any) -> bool: """Check equality.""" + if not isinstance(other, Namespace): + return False return dict(self) == dict(other) - def __ne__(self, other): + def __ne__(self, other: Any) -> bool: """Check inequality.""" - return dict(self) != dict(other) + return not self == other - def __contains__(self, key): + def __contains__(self, key: str) -> bool: """Check if an attribute exists.""" - return key in self._keys + return key in self._members - def __len__(self): - """Return the number of attributes.""" - return len(self._keys) + def __len__(self) -> int: + """Get the number of attributes.""" + return len(self._members) - def keys(self): - """Return keys of the namespace as a dictionary.""" - return {k: None for k in self._keys}.keys() + def keys(self) -> Iterable[str]: + """Get keys of the namespace as a dictionary.""" + return dict(self).keys() - def values(self): - """Return values of the namespace as a dictionary.""" + def values(self) -> Iterable: + """Get values of the namespace as a dictionary.""" return dict(self).values() - def items(self): - """Return items of the namespace as a dictionary.""" + def items(self) -> Iterable[Tuple[str, T]]: + """Get items of the namespace as a dictionary.""" return dict(self).items() - def get(self, *args, **kwargs): - """Get an item of the namespace as a dictionary.""" - return dict(self).get(*args, **kwargs) - class Timer: """Timer class.""" From e76f2abf744a4c1aa8d091877f0cf2866af1aeef Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Thu, 1 Aug 2024 19:04:37 +0100 Subject: [PATCH 05/37] Properly typed ebccc base class --- ebcc/cc/base.py | 29 ++++++++++++++++++----------- pyproject.toml | 5 ++--- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index d6f9b4fe..49a5eee1 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -3,27 +3,27 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import TYPE_CHECKING from dataclasses import dataclass +from typing import TYPE_CHECKING from ebcc import default_log, init_logging from ebcc import numpy as np from ebcc import util from ebcc.ansatz import Ansatz -from ebcc.space import Space from ebcc.logging import ANSI from ebcc.precision import types if TYPE_CHECKING: from logging import Logger - from typing import Optional, Union, Any + from typing import Any, Optional, Union from pyscf.scf import SCF # type: ignore - from ebcc.numpy.typing import NDArray # type: ignore + from ebcc.base import BruecknerEBCC as BaseBrueckner from ebcc.base import ERIs as BaseERIs from ebcc.base import Fock as BaseFock - from ebcc.base import BruecknerEBCC as BaseBrueckner + from ebcc.numpy.typing import NDArray # type: ignore + from ebcc.util import Namespace @dataclass @@ -40,6 +40,7 @@ class BaseOptions: diis_space: Number of amplitudes to use in DIIS extrapolation. damping: Damping factor for DIIS extrapolation. """ + shift: bool = True e_tol: float = 1e-8 t_tol: float = 1e-8 @@ -70,6 +71,11 @@ class BaseEBCC(ABC): CDERIs: type[BaseERIs] Brueckner: type[BaseBrueckner] + # Attributes + space: Any + amplitudes: Namespace[Any] + lambdas: Namespace[Any] + def __init__( self, mf: SCF, @@ -85,12 +91,13 @@ def __init__( fock: Optional[Any] = None, **kwargs, ): - """Initialize the EBCC object. + r"""Initialize the EBCC object. Args: mf: PySCF mean-field object. log: Log to write output to. Default is the global logger, outputting to `stderr`. ansatz: Overall ansatz. + options: Options for the EBCC calculation. space: Space containing the frozen, correlated, and active fermionic spaces. Default assumes all electrons are correlated. omega: Bosonic frequencies. @@ -157,10 +164,10 @@ def __init__( self.fock = fock # Attributes: - self.e_corr = None - self.amplitudes = None + self.e_corr = types[float](0.0) + self.amplitudes = util.Namespace() self.converged = False - self.lambdas = None + self.lambdas = util.Namespace() self.converged_lambda = False # Logging: @@ -237,7 +244,7 @@ def get_fock(self) -> Any: pass @abstractmethod - def get_eris(self, eris: Optional[Union[type[ERIs], NDArray[float]]]) -> Any: + def get_eris(self, eris: Optional[Union[type[BaseERIs], NDArray[float]]]) -> Any: """Get the electron repulsion integrals. Args: @@ -278,7 +285,7 @@ def const(self) -> float: Constant energy shift due to the polaritonic basis. """ if self.options.shift: - return util.einsum("I,I->", self.omega, self.xi ** 2) + return util.einsum("I,I->", self.omega, self.xi**2) return 0.0 @abstractmethod diff --git a/pyproject.toml b/pyproject.toml index b1b6a065..967d5db9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,7 +81,7 @@ skip_glob = [ [tool.flake8] max-line-length = 100 -max-doc-length = 75 +max-doc-length = 100 ignore = [ "E203", # Whitespace before ':' "E731", # Do not assign a lambda expression, use a def @@ -95,8 +95,7 @@ ignore = [ per-file-ignores = [ "__init__.py:E402,W605,F401,F811,D103", ] -docstring-convention = "numpy" -ignore-decorators = "has_docstring" +docstring-convention = "google" count = true include = "ebcc" exclude = """ From 120d2b7e6d874fa58ca359333d7325bcbba26ac2 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Thu, 1 Aug 2024 20:30:29 +0100 Subject: [PATCH 06/37] Weak typing for einsum --- ebcc/util/einsumfunc.py | 124 +++++++++++++++++----------------------- 1 file changed, 54 insertions(+), 70 deletions(-) diff --git a/ebcc/util/einsumfunc.py b/ebcc/util/einsumfunc.py index 70e14777..dde1d2d7 100644 --- a/ebcc/util/einsumfunc.py +++ b/ebcc/util/einsumfunc.py @@ -1,18 +1,28 @@ """Einstein summation convention.""" +from __future__ import annotations + import ctypes +from typing import TYPE_CHECKING -from pyscf.lib import direct_sum, dot # noqa: F401 -from pyscf.lib import einsum as pyscf_einsum # noqa: F401 +from pyscf.lib import direct_sum, dot # type: ignore # noqa: F401 +from pyscf.lib import einsum as pyscf_einsum # type: ignore # noqa: F401 from ebcc import numpy as np +if TYPE_CHECKING: + from typing import Any, TypeVar + + from ebcc.numpy.typing import NDArray # type: ignore + + T = TypeVar("T") + # Try to import TBLIS try: try: - import tblis_einsum + import tblis_einsum # type: ignore except ImportError: - from pyscf.tblis_einsum import tblis_einsum + from pyscf.tblis_einsum import tblis_einsum # type: ignore FOUND_TBLIS = True except ImportError: FOUND_TBLIS = False @@ -27,9 +37,8 @@ class EinsumOperandError(ValueError): pass -def _fallback_einsum(*operands, **kwargs): +def _fallback_einsum(*operands: Any, **kwargs: Any) -> NDArray[T]: """Handle the fallback to `numpy.einsum`.""" - # Parse the kwargs kwargs = kwargs.copy() alpha = kwargs.pop("alpha", 1.0) @@ -47,35 +56,24 @@ def _fallback_einsum(*operands, **kwargs): return res -def contract(subscript, *args, **kwargs): - """ - Contract a pair of terms in an einsum. Supports additional keyword - arguments `alpha` and `beta` which operate as `pyscf.lib.dot`. In some - cases this will still require copying, but it minimises the memory - overhead in simple cases. - - Parameters - ---------- - subscript : str - Subscript representation of the contraction. - args : tuple of numpy.ndarray - Arrays to contract. - alpha : float, optional - Scaling factor for contraction. Default value is `1.0`. - beta : float, optional - Scaling factor for the output. Default value is `0.0`. - out : numpy.ndarray, optional - Output array. If `None`, a new array is created. Default value is - `None`. - **kwargs : dict - Additional keyword arguments to `numpy.einsum`. - - Returns - ------- - array : numpy.ndarray +def contract(subscript: str, *args: Any, **kwargs: Any) -> NDArray[T]: + """Contraction a pair of terms in an Einstein summation. + + Supports additional keyword arguments `alpha` and `beta` which operate as `pyscf.lib.dot`. + In some cases this will still require copying, but it minimises the memory overhead in simple + cases. + + Args: + subscript: Subscript representation of the contraction. + args: Arrays to contract. + alpha: Scaling factor for contraction. + beta: Scaling factor for the output. + out: Output array. If `None`, a new array is created. + kwargs: Additional keyword arguments to `numpy.einsum`. + + Returns: Result of the contraction. """ - alpha = kwargs.get("alpha", 1.0) beta = kwargs.get("beta", 0.0) buf = kwargs.get("out", None) @@ -99,7 +97,7 @@ def contract(subscript, *args, **kwargs): return _fallback_einsum(subscript, *args, **kwargs) # Get the characters for each input and output - inp, out, args = np.core.einsumfunc._parse_einsum_input((subscript, a, b)) + inp, out, args = np.core.einsumfunc._parse_einsum_input((subscript, a, b)) # type: ignore inp_a, inp_b = inps = inp.split(",") assert len(inps) == len(args) == 2 assert all(len(inp) == arg.ndim for inp, arg in zip(inps, args)) @@ -115,7 +113,7 @@ def contract(subscript, *args, **kwargs): return _fallback_einsum(subscript, *args, **kwargs) # Find the sizes of the indices - ranges = {} + ranges: dict[int, int] = {} for inp, arg in zip(inps, args): for i, s in zip(inp, arg.shape): if i in ranges: @@ -201,7 +199,7 @@ def contract(subscript, *args, **kwargs): # Get the shapes shape_a = a.shape shape_b = b.shape - shape_c = tuple(ranges[x] for x in out) + shape_c = [ranges[x] for x in out] # If any dimension has size zero, return here if a.size == 0 or b.size == 0: @@ -252,9 +250,8 @@ def contract(subscript, *args, **kwargs): return c -def einsum(*operands, **kwargs): - """ - Evaluate an Einstein summation convention on the operands. +def einsum(*operands: Any, **kwargs: Any) -> NDArray[T]: + """Evaluate an Einstein summation convention on the operands. Using the Einstein summation convention, many common multi-dimensional, linear algebraic array operations can be @@ -268,37 +265,23 @@ def einsum(*operands, **kwargs): See the `numpy.einsum` documentation for clarification. - Parameters - ---------- - operands : list - Any valid input to `numpy.einsum`. - alpha : float, optional - Scaling factor for the contraction. Default value is `1.0`. - beta : float, optional - Scaling factor for the output. Default value is `0.0`. - out : ndarray, optional - If provided, the calculation is done into this array. - contract : callable, optional - The function to use for contraction. Default value is - `contract`. - optimize : bool, optional - If `True`, use the `numpy.einsum_path` to optimize the - contraction. Default value is `True`. - - Returns - ------- - output : ndarray + Args: + operands: Any valid input to `numpy.einsum`. + alpha: Scaling factor for the contraction. + beta: Scaling factor for the output. + out: If provided, the calculation is done into this array. + contract: The function to use for contraction. + optimize: If `True`, use the `numpy.einsum_path` to optimize the contraction. + + Returns: The calculation based on the Einstein summation convention. - Notes - ----- - This function may use `numpy.einsum`, `pyscf.lib.einsum`, or - `tblis_einsum` as a backend, depending on the problem size and the - modules available. + Notes: + This function may use `numpy.einsum`, `pyscf.lib.einsum`, or `tblis_einsum` as a backend, + depending on the problem size and the modules available. """ - # Parse the kwargs - inp, out, args = np.core.einsumfunc._parse_einsum_input(operands) + inp, out, args = np.core.einsumfunc._parse_einsum_input(operands) # type: ignore subscript = "%s->%s" % (inp, out) _contract = kwargs.get("contract", contract) @@ -313,11 +296,12 @@ def einsum(*operands, **kwargs): # If it's a chain of contractions, use the path optimizer optimize = kwargs.pop("optimize", True) args = list(args) - contractions = np.einsum_path(subscript, *args, optimize=optimize, einsum_call=True)[1] + kwargs = dict(optimize=optimize, einsum_call=True) + _, contractions = np.einsum_path(subscript, *args, **kwargs) # type: ignore for contraction in contractions: - inds, idx_rm, einsum_str, remain = contraction[:4] - operands = [args.pop(x) for x in inds] - out = _contract(einsum_str, *operands) + inds, idx_rm, einsum_str, remain = list(contraction[:4]) + contraction_args = [args.pop(x) for x in inds] + out = _contract(einsum_str, *contraction_args, **kwargs) args.append(out) return out From bc8bc3f8eb52de911553d31281e9117518cb3097 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Thu, 1 Aug 2024 20:36:25 +0100 Subject: [PATCH 07/37] Timer typing --- ebcc/util/misc.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index 5b40d04e..9cda5239 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -2,12 +2,12 @@ from __future__ import annotations +import time from collections.abc import MutableMapping from typing import TYPE_CHECKING, Generic, TypeVar -import time if TYPE_CHECKING: - from typing import Any, Callable, Iterable, Iterator + from typing import Any, ItemsView, Iterator, KeysView, ValuesView T = TypeVar("T") @@ -28,7 +28,8 @@ class ModelNotImplemented(NotImplementedError): class Namespace(MutableMapping[str, T], Generic[T]): - """ + """Namespace class. + Replacement for SimpleNamespace, which does not trivially allow conversion to a dict for heterogenously nested objects. @@ -46,11 +47,11 @@ def __init__(self, **kwargs: T): def __setitem__(self, key: str, val: T) -> None: """Set an item.""" - self.__dict__["_members"][key] = value + self.__dict__["_members"][key] = val def __setattr__(self, key: str, val: T) -> None: """Set an attribute.""" - return self.__setitem__(key, value) + return self.__setitem__(key, val) def __getitem__(self, key: str) -> T: """Get an item.""" @@ -83,7 +84,7 @@ def __ne__(self, other: Any) -> bool: """Check inequality.""" return not self == other - def __contains__(self, key: str) -> bool: + def __contains__(self, key: Any) -> bool: """Check if an attribute exists.""" return key in self._members @@ -91,15 +92,15 @@ def __len__(self) -> int: """Get the number of attributes.""" return len(self._members) - def keys(self) -> Iterable[str]: + def keys(self) -> KeysView[str]: """Get keys of the namespace as a dictionary.""" return dict(self).keys() - def values(self) -> Iterable: + def values(self) -> ValuesView[T]: """Get values of the namespace as a dictionary.""" return dict(self).values() - def items(self) -> Iterable[Tuple[str, T]]: + def items(self) -> ItemsView[str, T]: """Get items of the namespace as a dictionary.""" return dict(self).items() @@ -107,24 +108,25 @@ def items(self) -> Iterable[Tuple[str, T]]: class Timer: """Timer class.""" - def __init__(self): + def __init__(self) -> None: + """Initialise the timer.""" self.t_init = time.perf_counter() self.t_prev = time.perf_counter() self.t_curr = time.perf_counter() - def lap(self): + def lap(self) -> float: """Return the time since the last call to `lap`.""" self.t_prev, self.t_curr = self.t_curr, time.perf_counter() return self.t_curr - self.t_prev __call__ = lap - def total(self): + def total(self) -> float: """Return the total time since initialization.""" return time.perf_counter() - self.t_init @staticmethod - def format_time(seconds, precision=2): + def format_time(seconds: float, precision: int = 2) -> str: """Return a formatted time.""" seconds, milliseconds = divmod(seconds, 1) From a6169354662d5bc221c489ff277e40cd778e8753 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Fri, 2 Aug 2024 21:17:49 +0100 Subject: [PATCH 08/37] Improve CC class structure, better base class --- ebcc/__init__.py | 12 +- ebcc/base.py | 16 +- ebcc/cc/__init__.py | 4 + ebcc/cc/base.py | 912 +++++++++++++++++- ebcc/cc/gebcc.py | 1035 +++++++++++++++++++++ ebcc/cc/rebcc.py | 812 ++++++++++++++++ ebcc/{ => cc}/uebcc.py | 642 +++++++++---- ebcc/fock.py | 4 +- ebcc/gebcc.py | 465 ---------- ebcc/logging.py | 22 +- ebcc/precision.py | 67 +- ebcc/rebcc.py | 1834 ------------------------------------- ebcc/util/einsumfunc.py | 11 +- ebcc/util/misc.py | 17 +- ebcc/util/permutations.py | 457 ++++----- pyproject.toml | 38 +- 16 files changed, 3476 insertions(+), 2872 deletions(-) create mode 100644 ebcc/cc/gebcc.py create mode 100644 ebcc/cc/rebcc.py rename ebcc/{ => cc}/uebcc.py (70%) delete mode 100644 ebcc/gebcc.py delete mode 100644 ebcc/rebcc.py diff --git a/ebcc/__init__.py b/ebcc/__init__.py index ce8e119d..bc009fb6 100644 --- a/ebcc/__init__.py +++ b/ebcc/__init__.py @@ -39,15 +39,16 @@ import os import sys +import numpy + +from ebcc.logging import NullLogger, default_log, init_logging # --- Import NumPy here to allow drop-in replacements -import numpy # --- Logging: -from ebcc.logging import default_log, init_logging, NullLogger # --- Types of ansatz supporting by the EBCC solvers: @@ -57,9 +58,9 @@ # --- General constructor: -from ebcc.gebcc import GEBCC -from ebcc.rebcc import REBCC -from ebcc.uebcc import UEBCC +from ebcc.cc.gebcc import GEBCC +from ebcc.cc.rebcc import REBCC +from ebcc.cc.uebcc import UEBCC def EBCC(mf, *args, **kwargs): @@ -113,7 +114,6 @@ def constructor(mf, *args, **kwargs): from ebcc.brueckner import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC from ebcc.space import Space - # --- List available methods: diff --git a/ebcc/base.py b/ebcc/base.py index 1f62681d..0b581851 100644 --- a/ebcc/base.py +++ b/ebcc/base.py @@ -8,9 +8,9 @@ from ebcc import util if TYPE_CHECKING: - from logging import Logger - from typing import Optional, Union, Any from dataclasses import dataclass + from logging import Logger + from typing import Any, Optional, Union from pyscf.scf import SCF @@ -30,13 +30,17 @@ class BruecknerEBCC(ABC): pass -class ERIs(ABC, util.Namespace): +class ERIs(ABC): """Base class for electronic repulsion integrals.""" - pass + def __getitem__(self, key: str) -> Any: + """Get an item.""" + return self.__dict__[key] -class Fock(ABC, util.Namespace): +class Fock(ABC): """Base class for Fock matrices.""" - pass + def __getitem__(self, key: str) -> Any: + """Get an item.""" + return self.__dict__[key] diff --git a/ebcc/cc/__init__.py b/ebcc/cc/__init__.py index f6b49bd0..4e4c5999 100644 --- a/ebcc/cc/__init__.py +++ b/ebcc/cc/__init__.py @@ -1 +1,5 @@ """Coupled cluster solvers.""" + +from ebcc.cc.gebcc import GEBCC +from ebcc.cc.rebcc import REBCC +from ebcc.cc.uebcc import UEBCC diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 49a5eee1..e7148acd 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -10,21 +10,28 @@ from ebcc import numpy as np from ebcc import util from ebcc.ansatz import Ansatz +from ebcc.damping import DIIS +from ebcc.dump import Dump from ebcc.logging import ANSI -from ebcc.precision import types +from ebcc.precision import cast, types if TYPE_CHECKING: - from logging import Logger - from typing import Any, Optional, Union + from typing import Any, Callable, Literal, Mapping, Optional, Type, TypeVar, Union - from pyscf.scf import SCF # type: ignore + from pyscf.scf.hf import SCF # type: ignore from ebcc.base import BruecknerEBCC as BaseBrueckner from ebcc.base import ERIs as BaseERIs from ebcc.base import Fock as BaseFock + from ebcc.logging import Logger from ebcc.numpy.typing import NDArray # type: ignore + from ebcc.space import Space from ebcc.util import Namespace + ERIsInputType = Union[type[BaseERIs], NDArray[float]] + AmplitudeType = TypeVar("AmplitudeType") + SpaceType = TypeVar("SpaceType") + @dataclass class BaseOptions: @@ -72,9 +79,9 @@ class BaseEBCC(ABC): Brueckner: type[BaseBrueckner] # Attributes - space: Any - amplitudes: Namespace[Any] - lambdas: Namespace[Any] + space: SpaceType + amplitudes: Namespace[AmplitudeType] + lambdas: Namespace[AmplitudeType] def __init__( self, @@ -82,14 +89,14 @@ def __init__( log: Optional[Logger] = None, ansatz: Optional[Union[Ansatz, str]] = "CCSD", options: Optional[BaseOptions] = None, - space: Optional[Any] = None, - omega: Optional[Any] = None, - g: Optional[Any] = None, - G: Optional[Any] = None, - mo_coeff: Optional[Any] = None, - mo_occ: Optional[Any] = None, - fock: Optional[Any] = None, - **kwargs, + space: Optional[SpaceType] = None, + omega: Optional[NDArray[float]] = None, + g: Optional[NDArray[float]] = None, + G: Optional[NDArray[float]] = None, + mo_coeff: Optional[NDArray[float]] = None, + mo_occ: Optional[NDArray[float]] = None, + fock: Optional[BaseFock] = None, + **kwargs: Any, ): r"""Initialize the EBCC object. @@ -164,7 +171,7 @@ def __init__( self.fock = fock # Attributes: - self.e_corr = types[float](0.0) + self.e_corr = 0.0 self.amplitudes = util.Namespace() self.converged = False self.lambdas = util.Namespace() @@ -187,22 +194,839 @@ def __init__( self.log.info(f"{ANSI.B}Space{ANSI.R}: {ANSI.m}{self.space}{ANSI.R}") self.log.debug("") + @property @abstractmethod + def spin_type(self) -> str: + """Get a string representation of the spin type.""" + pass + + @property + def name(self) -> str: + """Get the name of the method.""" + return self.spin_type + self.ansatz.name + + def kernel(self, eris: Optional[ERIsInputType] = None) -> float: + """Run the coupled cluster calculation. + + Args: + eris: Electron repulsion integrals. + + Returns: + Correlation energy. + """ + timer = util.Timer() + + # Get the ERIs: + eris = self.get_eris(eris) + + # Get the amplitude guesses: + amplitudes = self.amplitudes + if not amplitudes: + amplitudes = self.init_amps(eris=eris) + + # Get the initial energy: + e_cc = self.energy(amplitudes=amplitudes, eris=eris) + + self.log.output("Solving for excitation amplitudes.") + self.log.debug("") + self.log.info( + f"{ANSI.B}{'Iter':>4s} {'Energy (corr.)':>16s} {'Energy (tot.)':>18s} " + f"{'Δ(Energy)':>13s} {'Δ(Ampl.)':>13s}{ANSI.R}" + ) + self.log.info(f"{0:4d} {e_cc:16.10f} {e_cc + self.mf.e_tot:18.10f}") + + if not self.ansatz.is_one_shot: + # Set up DIIS: + diis = DIIS() + diis.space = self.options.diis_space + diis.damping = self.options.damping + + converged = False + for niter in range(1, self.options.max_iter + 1): + # Update the amplitudes, extrapolate with DIIS and calculate change: + amplitudes_prev = amplitudes + amplitudes = self.update_amps(amplitudes=amplitudes, eris=eris) + vector = self.amplitudes_to_vector(amplitudes) + vector = diis.update(vector) + amplitudes = self.vector_to_amplitudes(vector) + dt = np.linalg.norm(vector - self.amplitudes_to_vector(amplitudes_prev), ord=np.inf) + + # Update the energy and calculate change: + e_prev = e_cc + e_cc = self.energy(amplitudes=amplitudes, eris=eris) + de = abs(e_prev - e_cc) + + # Log the iteration: + converged_e = bool(de < self.options.e_tol) + converged_t = bool(dt < self.options.t_tol) + self.log.info( + f"{niter:4d} {e_cc:16.10f} {e_cc + self.mf.e_tot:18.10f}" + f" {[ANSI.r, ANSI.g][bool(converged_e)]}{de:13.3e}{ANSI.R}" + f" {[ANSI.r, ANSI.g][bool(converged_t)]}{dt:13.3e}{ANSI.R}" + ) + + # Check for convergence: + converged = converged_e and converged_t + if converged: + self.log.debug("") + self.log.output(f"{ANSI.g}Converged.{ANSI.R}") + break + else: + self.log.debug("") + self.log.warning(f"{ANSI.r}Failed to converge.{ANSI.R}") + + # Include perturbative correction if required: + if self.ansatz.has_perturbative_correction: + self.log.debug("") + self.log.info("Computing perturbative energy correction.") + e_pert = self.energy_perturbative(amplitudes=amplitudes, eris=eris) + e_cc += e_pert + self.log.info(f"E(pert) = {e_pert:.10f}") + + else: + converged = True + + # Update attributes: + self.e_corr = e_cc + self.amplitudes = amplitudes + self.converged = converged + + self.log.debug("") + self.log.output(f"E(corr) = {self.e_corr:.10f}") + self.log.output(f"E(tot) = {self.e_tot:.10f}") + self.log.debug("") + self.log.debug("Time elapsed: %s", timer.format_time(timer())) + self.log.debug("") + + return e_cc + + def solve_lambda(self, amplitudes: Optional[Namespace[AmplitudeType]] = None, eris: Optional[ERIsInputType] = None) -> None: + """Solve for the lambda amplitudes. + + Args: + amplitudes: Cluster amplitudes. + eris: Electron repulsion integrals. + """ + timer = util.Timer() + + # Get the ERIs: + eris = self.get_eris(eris) + + # Get the amplitudes: + amplitudes = self.amplitudes + if not amplitudes: + amplitudes = self.init_amps(eris=eris) + + # If needed, get the perturbative part of the lambda amplitudes: + lambdas_pert = None + if self.ansatz.has_perturbative_correction: + lambdas_pert = self.update_lams(eris=eris, amplitudes=amplitudes, perturbative=True) + + # Get the initial lambda amplitudes: + lambdas = self.lambdas + if not lambdas: + lambdas = self.init_lams(amplitudes=amplitudes) + + self.log.output("Solving for de-excitation (lambda) amplitudes.") + self.log.debug("") + self.log.info(f"{ANSI.B}{'Iter':>4s} {'Δ(Ampl.)':>13s}{ANSI.R}") + + # Set up DIIS: + diis = DIIS() + diis.space = self.options.diis_space + diis.damping = self.options.damping + + converged = False + for niter in range(1, self.options.max_iter + 1): + # Update the lambda amplitudes, extrapolate with DIIS and calculate change: + lambdas_prev = lambdas + lambdas = self.update_lams( + amplitudes=amplitudes, + lambdas=lambdas, + lambdas_pert=lambdas_pert, + eris=eris, + ) + vector = self.lambdas_to_vector(lambdas) + vector = diis.update(vector) + lambdas = self.vector_to_lambdas(vector) + dl = np.linalg.norm(vector - self.lambdas_to_vector(lambdas_prev), ord=np.inf) + + # Log the iteration: + converged = bool(dl < self.options.t_tol) + self.log.info(f"{niter:4d} {[ANSI.r, ANSI.g][converged]}{dl:13.3e}{ANSI.R}") + + # Check for convergence: + if converged: + self.log.debug("") + self.log.output(f"{ANSI.g}Converged.{ANSI.R}") + break + else: + self.log.debug("") + self.log.warning(f"{ANSI.r}Failed to converge.{ANSI.R}") + + self.log.debug("") + self.log.debug("Time elapsed: %s", timer.format_time(timer())) + self.log.debug("") + self.log.debug("") + + # Update attributes: + self.lambdas = lambdas + self.converged_lambda = converged + + @abstractmethod + def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> Any: + """Get the IP-EOM object. + + Args: + options: Options for the IP-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + IP-EOM object. + """ + pass + + @abstractmethod + def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> Any: + """Get the EA-EOM object. + + Args: + options: Options for the EA-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + EA-EOM object. + """ + pass + + @abstractmethod + def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> Any: + """Get the EE-EOM object. + + Args: + options: Options for the EE-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + EE-EOM object. + """ + pass + + def brueckner(self, *args: Any, **kwargs: Any) -> float: + """Run a Brueckner orbital coupled cluster calculation. + + The coupled cluster object will be update in-place. + + Args: + *args: Arguments to pass to the Brueckner object. + **kwargs: Keyword arguments to pass to the Brueckner object. + + Returns: + Correlation energy. + """ + bcc = self.Brueckner(self, *args, **kwargs) + return bcc.kernel() + + def write(self, file: str) -> None: + """Write the EBCC object to a file. + + Args: + file: File to write the object to. + """ + writer = Dump(file) + writer.write(self) + + @classmethod + def read(cls, file: str, log: Optional[Logger] = None) -> BaseEBCC: + """Read the EBCC object from a file. + + Args: + file: File to read the object from. + log: Logger to use for new object. + + Returns: + EBCC object. + """ + reader = Dump(file) + return reader.read(cls=cls, log=log) + @staticmethod + @abstractmethod def _convert_mf(mf: SCF) -> SCF: """Convert the mean-field object to the appropriate type.""" pass + def _load_function( + self, + name: str, + eris: Optional[Union[ERIsInputType, Literal[False]]] = False, + amplitudes: Optional[Union[Namespace[AmplitudeType], Literal[False]]] = False, + lambdas: Optional[Union[Namespace[AmplitudeType], Literal[False]]] = False, + **kwargs: Any, + ) -> tuple[Callable[..., Any], dict[str, Any]]: + """Load a function from the generated code, and return the arguments.""" + dicts = [] + + # Get the ERIs: + if not (eris is False): + eris = self.get_eris(eris) + else: + eris = None + + # Get the amplitudes: + if not (amplitudes is False): + if amplitudes is None: + amplitudes = self.amplitudes + if amplitudes is None: + amplitudes = self.init_amps(eris=eris) + dicts.append(dict(amplitudes)) + + # Get the lambda amplitudes: + if not (lambdas is False): + if lambdas is None: + lambdas = self.lambdas + if lambdas is None: + lambdas = self.init_lams(amplitudes=amplitudes if amplitudes else None) + dicts.append(dict(lambdas)) + + # Get the function: + func = getattr(self._eqns, name, None) + if func is None: + raise util.ModelNotImplemented("%s for %s" % (name, self.name)) + + # Get the arguments: + if kwargs: + dicts.append(kwargs) + all_kwargs = self._pack_codegen_kwargs(*dicts, eris=eris) + + return func, all_kwargs + @abstractmethod - @property - def spin_type(self) -> str: - """Get a string representation of the spin type.""" + def _pack_codegen_kwargs( + self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None + ) -> dict[str, Any]: + """Pack all the keyword arguments for the generated code.""" pass @abstractmethod - @property - def name(self) -> str: - """Get the name of the method.""" + def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + """Initialise the cluster amplitudes. + + Args: + eris: Electron repulsion integrals. + + Returns: + Initial cluster amplitudes. + """ + pass + + @abstractmethod + def init_lams(self, amplitude: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + """Initialise the cluster lambda amplitudes. + + Args: + amplitude: Cluster amplitudes. + + Returns: + Initial cluster lambda amplitudes. + """ + pass + + def energy( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> float: + """Calculate the correlation energy. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + + Returns: + Correlation energy. + """ + func, kwargs = self._load_function( + "energy", + eris=eris, + amplitudes=amplitudes, + ) + return cast(func(**kwargs).real, float) + + def energy_perturbative( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + ) -> float: + """Calculate the perturbative correction to the correlation energy. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + + Returns: + Perturbative energy correction. + """ + func, kwargs = self._load_function( + "energy_perturbative", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + return cast(func(**kwargs).real, float) + + @abstractmethod + def update_amps( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> Namespace[AmplitudeType]: + """Update the cluster amplitudes. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + + Returns: + Updated cluster amplitudes. + """ + pass + + @abstractmethod + def update_lams( + self, + eris: ERIsInputType = None, + amplitude: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + lambdas_pert: Optional[Namespace[AmplitudeType]] = None, + perturbative: bool = False, + ) -> Namespace[AmplitudeType]: + """Update the cluster lambda amplitudes. + + Args: + eris: Electron repulsion integrals. + amplitude: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + lambdas_pert: Perturbative cluster lambda amplitudes. + perturbative: Flag to include perturbative correction. + + Returns: + Updated cluster lambda amplitudes. + """ + pass + + def make_sing_b_dm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> Any: + r"""Make the single boson density matrix :math:`\langle b \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + + Returns: + Single boson density matrix. + """ + func, kwargs = self._load_function( + "make_sing_b_dm", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + return func(**kwargs) + + def make_rdm1_b(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + r"""Make the one-particle boson reduced density matrix :math:`\langle b^+ c \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + unshifted: If `self.options.shift` is `True`, return the unshifted density matrix. Has + no effect if `self.options.shift` is `False`. + hermitise: Hermitise the density matrix. + + Returns: + One-particle boson reduced density matrix. + """ + func, kwargs = self._load_function( + "make_rdm1_b", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + dm = func(**kwargs) + + if hermitise: + dm = 0.5 * (dm + dm.T) + + if unshifted and self.options.shift: + dm_cre, dm_ann = self.make_sing_b_dm() + xi = self.xi + dm[np.diag_indices_from(dm)] -= xi * (dm_cre + dm_ann) - xi**2 + + return dm + + @abstractmethod + def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + hermitise: Hermitise the density matrix. + + Returns: + One-particle fermion reduced density matrix. + """ + pass + + @abstractmethod + def make_rdm2_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + hermitise: Hermitise the density matrix. + + Returns: + Two-particle fermion reduced density matrix. + """ + pass + + @abstractmethod + def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + unshifted: If `self.options.shift` is `True`, return the unshifted density matrix. Has + no effect if `self.options.shift` is `False`. + hermitise: Hermitise the density matrix. + + Returns: + Electron-boson coupling reduced density matrix. + """ + pass + + def hbar_matvec_ip(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType]: + """Compute the product between a state vector and the IP-EOM Hamiltonian. + + Args: + r1: State vector (single excitations). + r2: State vector (double excitations). + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + + Returns: + Products between the state vectors and the IP-EOM Hamiltonian for the singles and + doubles. + """ + func, kwargs = self._load_function( + "hbar_matvec_ip", + eris=eris, + amplitudes=amplitudes, + r1=r1, + r2=r2, + ) + return func(**kwargs) + + def hbar_matvec_ea(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, AmplitudeType]: + """Compute the product between a state vector and the EA-EOM Hamiltonian. + + Args: + r1: State vector (single excitations). + r2: State vector (double excitations). + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + + Returns: + Products between the state vectors and the EA-EOM Hamiltonian for the singles and + doubles. + """ + func, kwargs = self._load_function( + "hbar_matvec_ea", + eris=eris, + amplitudes=amplitudes, + r1=r1, + r2=r2, + ) + return func(**kwargs) + + def hbar_matvec_ee(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, AmplitudeType]: + """Compute the product between a state vector and the EE-EOM Hamiltonian. + + Args: + r1: State vector (single excitations). + r2: State vector (double excitations). + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + + Returns: + Products between the state vectors and the EE-EOM Hamiltonian for the singles and + doubles. + """ + func, kwargs = self._load_function( + "hbar_matvec_ee", + eris=eris, + amplitudes=amplitudes, + r1=r1, + r2=r2, + ) + return func(**kwargs) + + def make_ip_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + """Get the bra vectors to construct IP-EOM moments. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + + Returns: + Bra vectors for IP-EOM moments. + """ + func, kwargs = self._load_function( + "make_ip_mom_bras", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + return func(**kwargs) + + def make_ea_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + """Get the bra vectors to construct EA-EOM moments. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + + Returns: + Bra vectors for EA-EOM moments. + """ + func, kwargs = self._load_function( + "make_ea_mom_bras", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + return func(**kwargs) + + def make_ee_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + """Get the bra vectors to construct EE-EOM moments. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + + Returns: + Bra vectors for EE-EOM moments. + """ + func, kwargs = self._load_function( + "make_ee_mom_bras", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + return func(**kwargs) + + def make_ip_mom_kets(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + """Get the ket vectors to construct IP-EOM moments. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + + Returns: + Ket vectors for IP-EOM moments. + """ + func, kwargs = self._load_function( + "make_ip_mom_kets", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + return func(**kwargs) + + def make_ea_mom_kets(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + """Get the ket vectors to construct EA-EOM moments. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + + Returns: + Ket vectors for EA-EOM moments. + """ + func, kwargs = self._load_function( + "make_ea_mom_kets", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + return func(**kwargs) + + def make_ee_mom_kets(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + """Get the ket vectors to construct EE-EOM moments. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + + Returns: + Ket vectors for EE-EOM moments. + """ + func, kwargs = self._load_function( + "make_ee_mom_kets", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + return func(**kwargs) + + @abstractmethod + def energy_sum(self, *args: str, signs_dict: Optional[dict[str, int]] = None) -> NDArray[float]: + """Get a direct sum of energies. + + Args: + *args: Energies to sum. Subclass should specify a subscript, and optionally spins. + signs_dict: Signs of the energies in the sum. Default sets `("o", "O", "i")` to be + positive, and `("v", "V", "a", "b")` to be negative. + + Returns: + Sum of energies. + """ + pass + + @abstractmethod + def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the amplitudes used in the given ansatz. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Cluster amplitudes as a vector. + """ + pass + + @abstractmethod + def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + """Construct a namespace of amplitudes from a vector. + + Args: + vector: Cluster amplitudes as a vector. + + Returns: + Cluster amplitudes. + """ + pass + + @abstractmethod + def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the lambda amplitudes used in the given ansatz. + + Args: + lambdas: Cluster lambda amplitudes. + + Returns: + Cluster lambda amplitudes as a vector. + """ + pass + + @abstractmethod + def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + """Construct a namespace of lambda amplitudes from a vector. + + Args: + vector: Cluster lambda amplitudes as a vector. + + Returns: + Cluster lambda amplitudes. + """ + pass + + @abstractmethod + def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the IP-EOM excitations. + + Args: + excitations: IP-EOM excitations. + + Returns: + IP-EOM excitations as a vector. + """ + pass + + @abstractmethod + def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the EA-EOM excitations. + + Args: + excitations: EA-EOM excitations. + + Returns: + EA-EOM excitations as a vector. + """ + pass + + @abstractmethod + def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the EE-EOM excitations. + + Args: + excitations: EE-EOM excitations. + + Returns: + EE-EOM excitations as a vector. + """ + pass + + @abstractmethod + def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of IP-EOM excitations from a vector. + + Args: + vector: IP-EOM excitations as a vector. + + Returns: + IP-EOM excitations. + """ + pass + + @abstractmethod + def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of EA-EOM excitations from a vector. + + Args: + vector: EA-EOM excitations as a vector. + + Returns: + EA-EOM excitations. + """ + pass + + @abstractmethod + def vector_to_excitations_ee(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of EE-EOM excitations from a vector. + + Args: + vector: EE-EOM excitations as a vector. + + Returns: + EE-EOM excitations. + """ pass @property @@ -230,12 +1054,12 @@ def init_space(self) -> Any: """Initialise the fermionic space. Returns: - Fermionic space. + Fermionic space. All fermionic degrees of freedom are assumed to be correlated. """ pass @abstractmethod - def get_fock(self) -> Any: + def get_fock(self) -> BaseFock: """Get the Fock matrix. Returns: @@ -244,7 +1068,7 @@ def get_fock(self) -> Any: pass @abstractmethod - def get_eris(self, eris: Optional[Union[type[BaseERIs], NDArray[float]]]) -> Any: + def get_eris(self, eris: Optional[ERIsInputType] = None) -> BaseERIs: """Get the electron repulsion integrals. Args: @@ -256,7 +1080,7 @@ def get_eris(self, eris: Optional[Union[type[BaseERIs], NDArray[float]]]) -> Any pass @abstractmethod - def get_g(self, g: NDArray[float]) -> Any: + def get_g(self, g: NDArray[float]) -> Namespace[Any]: """Get the blocks of the electron-boson coupling matrix. This matrix corresponds to the bosonic annihilation operator. @@ -278,18 +1102,19 @@ def get_mean_field_G(self) -> Any: """ pass - def const(self) -> float: - """Get the shift in energy from moving to the polaritonic basis. + @property + 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: - Constant energy shift due to the polaritonic basis. + Mean-field Fock matrix. """ - if self.options.shift: - return util.einsum("I,I->", self.omega, self.xi**2) - return 0.0 + pass - @abstractmethod @property + @abstractmethod def xi(self) -> NDArray[float]: """Get the shift in the bosonic operators to diagonalise the photon Hamiltonian. @@ -298,6 +1123,17 @@ def xi(self) -> NDArray[float]: """ pass + @property + def const(self) -> float: + """Get the shift in energy from moving to the polaritonic basis. + + Returns: + Constant energy shift due to the polaritonic basis. + """ + if self.options.shift: + return util.einsum("I,I->", self.omega, self.xi**2) + return 0.0 + @property def mo_coeff(self) -> NDArray[float]: """Get the molecular orbital coefficients. @@ -320,8 +1156,8 @@ def mo_occ(self) -> NDArray[float]: return np.asarray(self.mf.mo_occ).astype(types[float]) return self._mo_occ - @abstractmethod @property + @abstractmethod def nmo(self) -> Any: """Get the number of molecular orbitals. @@ -330,8 +1166,8 @@ def nmo(self) -> Any: """ pass - @abstractmethod @property + @abstractmethod def nocc(self) -> Any: """Get the number of occupied molecular orbitals. @@ -340,8 +1176,8 @@ def nocc(self) -> Any: """ pass - @abstractmethod @property + @abstractmethod def nvir(self) -> Any: """Get the number of virtual molecular orbitals. @@ -359,7 +1195,7 @@ def nbos(self) -> int: """ if self.omega is None: return 0 - return self.omega.shape[0] + return int(self.omega.shape[0]) @property def e_tot(self) -> float: @@ -368,7 +1204,7 @@ def e_tot(self) -> float: Returns: Total energy. """ - return types[float](self.mf.e_tot) + self.e_corr + return cast(self.mf.e_tot + self.e_corr, float) @property def t1(self) -> Any: diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py new file mode 100644 index 00000000..4196b71e --- /dev/null +++ b/ebcc/cc/gebcc.py @@ -0,0 +1,1035 @@ +"""General electron-boson coupled cluster.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyscf import lib, scf + +from ebcc import geom +from ebcc import numpy as np +from ebcc import uebcc, util +from ebcc.brueckner import BruecknerGEBCC +from ebcc.cc.base import BaseEBCC +from ebcc.eris import GERIs +from ebcc.fock import GFock +from ebcc.precision import types +from ebcc.space import Space + +if TYPE_CHECKING: + from typing import Optional + + from pyscf.scf.ghf import GHF + from pyscf.scf.hf import SCF + + from ebcc.cc.base import ERIsInputType + from ebcc.cc.rebcc import REBCC + from ebcc.cc.uebcc import UEBCC + from ebcc.numpy.typing import NDArray + from ebcc.util import Namespace + + AmplitudeType = NDArray[float] + + +class GEBCC(BaseEBCC): + """Restricted electron-boson coupled cluster. + + Attributes: + mf: PySCF mean-field object. + log: Log to write output to. + options: Options for the EBCC calculation. + e_corr: Correlation energy. + amplitudes: Cluster amplitudes. + converged: Convergence flag. + lambdas: Cluster lambda amplitudes. + converged_lambda: Lambda convergence flag. + name: Name of the method. + """ + + # Types + ERIs = GERIs + Fock = GFock + CDERIs = None + Brueckner = BruecknerGEBCC + + @property + def spin_type(self): + return "G" + + def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.IP_GEOM: + """Get the IP-EOM object. + + Args: + options: Options for the IP-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + IP-EOM object. + """ + return geom.IP_GEOM(self, options=options, **kwargs) + + def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.EA_GEOM: + """Get the EA-EOM object. + + Args: + options: Options for the EA-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + EA-EOM object. + """ + return geom.EA_GEOM(self, options=options, **kwargs) + + def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.EE_GEOM: + """Get the EE-EOM object. + + Args: + options: Options for the EE-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + EE-EOM object. + """ + return geom.EE_GEOM(self, options=options, **kwargs) + + @staticmethod + def _convert_mf(mf: SCF) -> GHF: + """Convert a mean-field object to a GHF object. + + Note: + Converts to UHF first to ensure consistent ordering. + """ + if isinstance(mf, scf.ghf.GHF): + return mf + return mf.to_uhf().to_ghf() + + @classmethod + def from_uebcc(cls, ucc: UEBCC) -> GEBCC: + """Initialise a `GEBCC` object from an `UEBCC` object. + + Args: + ucc: Unrestricted electron-boson coupled cluster object. + + Returns: + GEBCC object. + """ + orbspin = scf.addons.get_ghf_orbspin(ucc.mf.mo_energy, ucc.mf.mo_occ, False) + nocc = ucc.space[0].nocc + ucc.space[1].nocc + nvir = ucc.space[0].nvir + ucc.space[1].nvir + nbos = ucc.nbos + sa = np.where(orbspin == 0)[0] + sb = np.where(orbspin == 1)[0] + + occupied = np.zeros((nocc + nvir,), dtype=bool) + occupied[sa] = ucc.space[0]._occupied.copy() + occupied[sb] = ucc.space[1]._occupied.copy() + frozen = np.zeros((nocc + nvir,), dtype=bool) + frozen[sa] = ucc.space[0]._frozen.copy() + frozen[sb] = ucc.space[1]._frozen.copy() + active = np.zeros((nocc + nvir,), dtype=bool) + active[sa] = ucc.space[0]._active.copy() + active[sb] = ucc.space[1]._active.copy() + space = Space(occupied, frozen, active) + + slices = util.Namespace( + a=util.Namespace(**{k: np.where(orbspin[space.mask(k)] == 0)[0] for k in "oOivVa"}), + b=util.Namespace(**{k: np.where(orbspin[space.mask(k)] == 1)[0] for k in "oOivVa"}), + ) + + if ucc.bare_g is not None: + if np.asarray(ucc.bare_g).ndim == 3: + bare_g_a = bare_g_b = ucc.bare_g + else: + bare_g_a, bare_g_b = ucc.bare_g + g = np.zeros((ucc.nbos, ucc.nmo * 2, ucc.nmo * 2)) + g[np.ix_(range(ucc.nbos), sa, sa)] = bare_g_a.copy() + g[np.ix_(range(ucc.nbos), sb, sb)] = bare_g_b.copy() + else: + g = None + + gcc = cls( + ucc.mf, + log=ucc.log, + ansatz=ucc.ansatz, + space=space, + omega=ucc.omega, + g=g, + G=ucc.bare_G, + options=ucc.options, + ) + + gcc.e_corr = ucc.e_corr + gcc.converged = ucc.converged + gcc.converged_lambda = ucc.converged_lambda + + has_amps = ucc.amplitudes is not None + has_lams = ucc.lambdas is not None + + if has_amps: + amplitudes = util.Namespace() + + for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): + shape = tuple(space.size(k) for k in key) + amplitudes[name] = np.zeros(shape, dtype=types[float]) + for comb in util.generate_spin_combinations(n, unique=True): + done = set() + for lperm, lsign in util.permutations_with_signs(tuple(range(n))): + for uperm, usign in util.permutations_with_signs(tuple(range(n))): + combn = util.permute_string(comb[:n], lperm) + combn += util.permute_string(comb[n:], uperm) + if combn in done: + continue + mask = np.ix_(*[slices[s][k] for s, k in zip(combn, key)]) + transpose = tuple(lperm) + tuple(p + n for p in uperm) + amp = ( + getattr(ucc.amplitudes[name], comb).transpose(transpose) + * lsign + * usign + ) + for perm, sign in util.permutations_with_signs(tuple(range(n))): + transpose = tuple(perm) + tuple(range(n, 2 * n)) + if util.permute_string(comb[:n], perm) == comb[:n]: + amplitudes[name][mask] += amp.transpose(transpose).copy() * sign + done.add(combn) + + for name, key, n in ucc.ansatz.bosonic_cluster_ranks(spin_type=ucc.spin_type): + amplitudes[name] = ucc.amplitudes[name].copy() + + for name, key, nf, nb in ucc.ansatz.coupling_cluster_ranks(spin_type=ucc.spin_type): + shape = (nbos,) * nb + tuple(space.size(k) for k in key[nb:]) + amplitudes[name] = np.zeros(shape, dtype=types[float]) + for comb in util.generate_spin_combinations(nf): + done = set() + for lperm, lsign in util.permutations_with_signs(tuple(range(nf))): + for uperm, usign in util.permutations_with_signs(tuple(range(nf))): + combn = util.permute_string(comb[:nf], lperm) + combn += util.permute_string(comb[nf:], uperm) + if combn in done: + continue + mask = np.ix_( + *([range(nbos)] * nb), + *[slices[s][k] for s, k in zip(combn, key[nb:])], + ) + transpose = ( + tuple(range(nb)) + + tuple(p + nb for p in lperm) + + tuple(p + nb + nf for p in uperm) + ) + amp = ( + getattr(ucc.amplitudes[name], comb).transpose(transpose) + * lsign + * usign + ) + for perm, sign in util.permutations_with_signs(tuple(range(nf))): + transpose = ( + tuple(range(nb)) + + tuple(p + nb for p in perm) + + tuple(range(nb + nf, nb + 2 * nf)) + ) + if util.permute_string(comb[:nf], perm) == comb[:nf]: + amplitudes[name][mask] += amp.transpose(transpose).copy() * sign + done.add(combn) + + gcc.amplitudes = amplitudes + + if has_lams: + lambdas = gcc.init_lams() # Easier this way - but have to build ERIs... + + for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): + lname = name.replace("t", "l") + shape = tuple(space.size(k) for k in key[n:] + key[:n]) + lambdas[lname] = np.zeros(shape, dtype=types[float]) + for comb in util.generate_spin_combinations(n, unique=True): + done = set() + for lperm, lsign in util.permutations_with_signs(tuple(range(n))): + for uperm, usign in util.permutations_with_signs(tuple(range(n))): + combn = util.permute_string(comb[:n], lperm) + combn += util.permute_string(comb[n:], uperm) + if combn in done: + continue + mask = np.ix_(*[slices[s][k] for s, k in zip(combn, key[n:] + key[:n])]) + transpose = tuple(lperm) + tuple(p + n for p in uperm) + amp = ( + getattr(ucc.lambdas[lname], comb).transpose(transpose) + * lsign + * usign + ) + for perm, sign in util.permutations_with_signs(tuple(range(n))): + transpose = tuple(perm) + tuple(range(n, 2 * n)) + if util.permute_string(comb[:n], perm) == comb[:n]: + lambdas[lname][mask] += amp.transpose(transpose).copy() * sign + done.add(combn) + + for name, key, n in ucc.ansatz.bosonic_cluster_ranks(spin_type=ucc.spin_type): + lname = "l" + name + lambdas[lname] = ucc.lambdas[lname].copy() + + for name, key, nf, nb in ucc.ansatz.coupling_cluster_ranks(spin_type=ucc.spin_type): + lname = "l" + name + shape = (nbos,) * nb + tuple( + space.size(k) for k in key[nb + nf :] + key[nb : nb + nf] + ) + lambdas[lname] = np.zeros(shape, dtype=types[float]) + for comb in util.generate_spin_combinations(nf, unique=True): + done = set() + for lperm, lsign in util.permutations_with_signs(tuple(range(nf))): + for uperm, usign in util.permutations_with_signs(tuple(range(nf))): + combn = util.permute_string(comb[:nf], lperm) + combn += util.permute_string(comb[nf:], uperm) + if combn in done: + continue + mask = np.ix_( + *([range(nbos)] * nb), + *[ + slices[s][k] + for s, k in zip(combn, key[nb + nf :] + key[nb : nb + nf]) + ], + ) + transpose = ( + tuple(range(nb)) + + tuple(p + nb for p in lperm) + + tuple(p + nb + nf for p in uperm) + ) + amp = ( + getattr(ucc.lambdas[lname], comb).transpose(transpose) + * lsign + * usign + ) + for perm, sign in util.permutations_with_signs(tuple(range(nf))): + transpose = ( + tuple(range(nb)) + + tuple(p + nb for p in perm) + + tuple(range(nb + nf, nb + 2 * nf)) + ) + if util.permute_string(comb[:nf], perm) == comb[:nf]: + lambdas[lname][mask] += amp.transpose(transpose).copy() * sign + done.add(combn) + + gcc.lambdas = lambdas + + return gcc + + @classmethod + def from_rebcc(cls, rcc: REBCC) -> GEBCC: + """Initialise a `GEBCC` object from an `REBCC` object. + + Args: + rcc: Restricted electron-boson coupled cluster object. + + Returns: + GEBCC object. + """ + ucc = uebcc.UEBCC.from_rebcc(rcc) + gcc = cls.from_uebcc(ucc) + return gcc + + def init_space(self) -> Space: + """Initialise the fermionic space. + + Returns: + Fermionic space. All fermionic degrees of freedom are assumed to be correlated. + """ + space = Space( + self.mo_occ > 0, + np.zeros_like(self.mo_occ, dtype=bool), + np.zeros_like(self.mo_occ, dtype=bool), + ) + return space + + def _pack_codegen_kwargs(self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None) -> dict[str, Any]: + """Pack all the keyword arguments for the generated code.""" + kwargs = dict( + f=self.fock, + v=self.get_eris(eris), + g=self.g, + G=self.G, + w=np.diag(self.omega) if self.omega is not None else None, + space=self.space, + nocc=self.space.ncocc, + nvir=self.space.ncvir, + nbos=self.nbos, + ) + + for kw in extra_kwargs: + if kw is not None: + kwargs.update(kw) + + return kwargs + + def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + """Initialise the cluster amplitudes. + + Args: + eris: Electron repulsion integrals. + + Returns: + Initial cluster amplitudes. + """ + eris = self.get_eris(eris) + amplitudes = util.Namespace() + + # Build T amplitudes: + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + if n == 1: + amplitudes[name] = getattr(self.fock, key) / self.energy_sum(key) + elif n == 2: + amplitudes[name] = getattr(eris, key) / self.energy_sum(key) + else: + shape = tuple(self.space.size(k) for k in key) + amplitudes[name] = np.zeros(shape, dtype=types[float]) + + if self.boson_ansatz: + # Only true for real-valued couplings: + h = self.g + H = self.G + + # Build S amplitudes: + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + if n == 1: + amplitudes[name] = -H / self.omega + else: + shape = (self.nbos,) * n + amplitudes[name] = np.zeros(shape, dtype=types[float]) + + # Build U amplitudes: + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + if nf != 1: + raise util.ModelNotImplemented + if n == 1: + amplitudes[name] = h[key] / self.energy_sum(key) + else: + shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) + amplitudes[name] = np.zeros(shape, dtype=types[float]) + + return amplitudes + + def init_lams(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + """Initialise the cluster lambda amplitudes. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Initial cluster lambda amplitudes. + """ + if amplitudes is None: + amplitudes = self.amplitudes + lambdas = util.Namespace() + + # Build L amplitudes: + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + lname = name.replace("t", "l") + perm = list(range(n, 2 * n)) + list(range(n)) + lambdas[lname] = amplitudes[name].transpose(perm) + + # Build LS amplitudes: + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + lname = "l" + name + lambdas[lname] = amplitudes[name] + + # Build LU amplitudes: + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + if nf != 1: + raise util.ModelNotImplemented + lname = "l" + name + perm = list(range(nb)) + [nb + 1, nb] + lambdas[lname] = amplitudes[name].transpose(perm) + + return lambdas + + def update_amps( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> Namespace[AmplitudeType]: + """Update the cluster amplitudes. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + + Returns: + Updated cluster amplitudes. + """ + func, kwargs = self._load_function( + "update_amps", + eris=eris, + amplitudes=amplitudes, + ) + res = func(**kwargs) + res = {key.rstrip("new"): val for key, val in res.items()} + + # Divide T amplitudes: + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + res[name] /= self.energy_sum(key) + res[name] += amplitudes[name] + + # Divide S amplitudes: + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + res[name] /= self.energy_sum(key) + res[name] += amplitudes[name] + + # Divide U amplitudes: + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + if nf != 1: + raise util.ModelNotImplemented + res[name] /= self.energy_sum(key) + res[name] += amplitudes[name] + + return res + + def update_lams( + self, + eris: ERIsInputType = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + lambdas_pert: Optional[Namespace[AmplitudeType]] = None, + perturbative: bool = False, + ) -> Namespace[AmplitudeType]: + """Update the cluster lambda amplitudes. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + lambdas_pert: Perturbative cluster lambda amplitudes. + perturbative: Flag to include perturbative correction. + + Returns: + Updated cluster lambda amplitudes. + """ + # TODO active + if lambdas_pert is not None: + lambdas.update(lambdas_pert) + + func, kwargs = self._load_function( + "update_lams%s" % ("_perturbative" if perturbative else ""), + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + res = func(**kwargs) + res = {key.rstrip("new"): val for key, val in res.items()} + + # Divide T amplitudes: + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + lname = name.replace("t", "l") + res[lname] /= self.energy_sum(key[n:] + key[:n]) + if not perturbative: + res[lname] += lambdas[lname] + + # Divide S amplitudes: + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + lname = "l" + name + res[lname] /= self.energy_sum(key[n:] + key[:n]) + if not perturbative: + res[lname] += lambdas[lname] + + # Divide U amplitudes: + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + if nf != 1: + raise util.ModelNotImplemented + lname = "l" + name + res[lname] /= self.energy_sum(key[:nb] + key[nb + nf :] + key[nb : nb + nf]) + if not perturbative: + res[lname] += lambdas[lname] + + if perturbative: + res = {key + "pert": val for key, val in res.items()} + + return res + + def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + hermitise: Hermitise the density matrix. + + Returns: + One-particle fermion reduced density matrix. + """ + func, kwargs = self._load_function( + "make_rdm1_f", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + dm = func(**kwargs) + + if hermitise: + dm = 0.5 * (dm + dm.T) + + return dm + + def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): + func, kwargs = self._load_function( + "make_rdm2_f", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + + dm = func(**kwargs) + + if hermitise: + dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) + dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(1, 0, 3, 2)) + + return dm + + def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + unshifted: If `self.options.shift` is `True`, return the unshifted density matrix. Has + no effect if `self.options.shift` is `False`. + hermitise: Hermitise the density matrix. + + Returns: + Electron-boson coupling reduced density matrix. + """ + func, kwargs = self._load_function( + "make_eb_coup_rdm", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + dm_eb = func(**kwargs) + + if hermitise: + dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) + dm_eb[1] = dm_eb[0].transpose(0, 2, 1).copy() + + if unshifted and self.options.shift: + rdm1_f = self.make_rdm1_f(hermitise=hermitise) + shift = util.einsum("x,ij->xij", self.xi, rdm1_f) + dm_eb -= shift[None] + + return dm_eb + + def energy_sum(self, subscript: str, signs_dict: dict[str, int] = None) -> NDArray[float]: + """Get a direct sum of energies. + + Args: + subscript: Subscript for the direct sum. + signs_dict: Signs of the energies in the sum. Default sets `("o", "O", "i")` to be + positive, and `("v", "V", "a", "b")` to be negative. + + Returns: + Sum of energies. + """ + n = 0 + + def next_char() -> str: + nonlocal n + if n < 26: + char = chr(ord("a") + n) + else: + char = chr(ord("A") + n) + n += 1 + return char + + if signs_dict is None: + signs_dict = {} + for k, s in zip("vVaoOib", "---+++-"): + if k not in signs_dict: + signs_dict[k] = s + + energies = [] + for key in subscript: + if key == "b": + energies.append(self.omega) + else: + energies.append(np.diag(self.fock[key + key])) + + subscript = "".join([signs_dict[k] + next_char() for k in subscript]) + energy_sum = lib.direct_sum(subscript, *energies) + + return energy_sum + + def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the amplitudes used in the given ansatz. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Cluster amplitudes as a vector. + """ + vectors = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + vectors.append(amplitudes[name].ravel()) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + vectors.append(amplitudes[name].ravel()) + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + vectors.append(amplitudes[name].ravel()) + + return np.concatenate(vectors) + + def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + """Construct a namespace of amplitudes from a vector. + + Args: + vector: Cluster amplitudes as a vector. + + Returns: + Cluster amplitudes. + """ + amplitudes = util.Namespace() + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + shape = tuple(self.space.size(k) for k in key) + size = np.prod(shape) + amplitudes[name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + shape = (self.nbos,) * n + size = np.prod(shape) + amplitudes[name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) + size = np.prod(shape) + amplitudes[name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + return amplitudes + + def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the lambda amplitudes used in the given ansatz. + + Args: + lambdas: Cluster lambda amplitudes. + + Returns: + Cluster lambda amplitudes as a vector. + """ + vectors = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + vectors.append(lambdas[name.replace("t", "l")].ravel()) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + vectors.append(lambdas["l" + name].ravel()) + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + vectors.append(lambdas["l" + name].ravel()) + + return np.concatenate(vectors) + + def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + """Construct a namespace of lambda amplitudes from a vector. + + Args: + vector: Cluster lambda amplitudes as a vector. + + Returns: + Cluster lambda amplitudes. + """ + lambdas = util.Namespace() + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + lname = name.replace("t", "l") + key = key[n:] + key[:n] + shape = tuple(self.space.size(k) for k in key) + size = np.prod(shape) + lambdas[lname] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + shape = (self.nbos,) * n + size = np.prod(shape) + lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + key = key[:nb] + key[nb + nf :] + key[nb : nb + nf] + shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) + size = np.prod(shape) + lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + return lambdas + + def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the IP-EOM excitations. + + Args: + excitations: IP-EOM excitations. + + Returns: + IP-EOM excitations as a vector. + """ + vectors = [] + m = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[:-1] + vectors.append(util.compress_axes(key, excitations[m]).ravel()) + m += 1 + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return np.concatenate(vectors) + + def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the EA-EOM excitations. + + Args: + excitations: EA-EOM excitations. + + Returns: + EA-EOM excitations as a vector. + """ + return self.excitations_to_vector_ip(*excitations) + + def excitations_to_vector_ee(self, *excitations): + """Construct a vector containing all of the EE-EOM excitations. + + Args: + excitations: EE-EOM excitations. + + Returns: + EE-EOM excitations as a vector. + """ + vectors = [] + m = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + vectors.append(util.compress_axes(key, excitations[m]).ravel()) + m += 1 + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return np.concatenate(vectors) + + def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of IP-EOM excitations from a vector. + + Args: + vector: IP-EOM excitations as a vector. + + Returns: + IP-EOM excitations. + """ + excitations = [] + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[:-1] + size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) + shape = tuple(self.space.size(k) for k in key) + vn_tril = vector[i0 : i0 + size] + vn = util.decompress_axes(key, vn_tril, shape=shape) + excitations.append(vn) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return tuple(excitations) + + def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of EA-EOM excitations from a vector. + + Args: + vector: EA-EOM excitations as a vector. + + Returns: + EA-EOM excitations. + """ + excitations = [] + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[n:] + key[: n - 1] + size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) + shape = tuple(self.space.size(k) for k in key) + vn_tril = vector[i0 : i0 + size] + vn = util.decompress_axes(key, vn_tril, shape=shape) + excitations.append(vn) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return tuple(excitations) + + def vector_to_excitations_ee(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of EE-EOM excitations from a vector. + + Args: + vector: EE-EOM excitations as a vector. + + Returns: + EE-EOM excitations. + """ + excitations = [] + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) + shape = tuple(self.space.size(k) for k in key) + vn_tril = vector[i0 : i0 + size] + vn = util.decompress_axes(key, vn_tril, shape=shape) + excitations.append(vn) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return tuple(excitations) + + def get_mean_field_G(self) -> NDArray[float]: + """Get the mean-field boson non-conserving term. + + Returns: + Mean-field boson non-conserving term. + """ + # FIXME should this also sum in frozen orbitals? + val = lib.einsum("Ipp->I", self.g.boo) + val -= self.xi * self.omega + if self.bare_G is not None: + val += self.bare_G + return val + + def get_g(self, g: NDArray[float]) -> Namespace[NDArray[float]]: + """Get the blocks of the electron-boson coupling matrix. + + This matrix corresponds to the bosonic annihilation operator. + + Args: + g: Electron-boson coupling matrix. + + Returns: + Blocks of the electron-boson coupling matrix. + """ + # TODO make a proper class for this + slices = { + "x": self.space.correlated, + "o": self.space.correlated_occupied, + "v": self.space.correlated_virtual, + "O": self.space.active_occupied, + "V": self.space.active_virtual, + "i": self.space.inactive_occupied, + "a": self.space.inactive_virtual, + } + + class Blocks(util.Namespace): + def __getitem__(selffer, key): + assert key[0] == "b" + i = slices[key[1]] + j = slices[key[2]] + return g[:, i][:, :, j].copy() + + __getattr__ = __getitem__ + + return Blocks() + + @property + def bare_fock(self) -> NDArray[float]: + """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 = self.mf.get_fock().astype(types[float]) + fock = util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff, self.mo_coeff) + return fock + + @property + def xi(self) -> NDArray[float]: + """Get the shift in the bosonic operators to diagonalise the photon Hamiltonian. + + Returns: + Shift in the bosonic operators. + """ + if self.options.shift: + xi = lib.einsum("Iii->I", self.g.boo) + xi /= self.omega + if self.bare_G is not None: + xi += self.bare_G / self.omega + else: + xi = np.zeros_like(self.omega) + return xi + + def get_fock(self) -> GFock: + """Get the Fock matrix. + + Returns: + Fock matrix. + """ + return self.Fock(self, array=self.bare_fock) + + 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 (eris is None) or isinstance(eris, np.ndarray): + return self.ERIs(self, array=eris) + else: + return eris + + @property + def nmo(self) -> int: + """Get the number of molecular orbitals. + + Returns: + Number of molecular orbitals. + """ + return self.space.nmo + + @property + def nocc(self) -> int: + """Get the number of occupied molecular orbitals. + + Returns: + Number of occupied molecular orbitals. + """ + return self.space.nocc + + @property + def nvir(self) -> int: + """Get the number of virtual molecular orbitals. + + Returns: + Number of virtual molecular orbitals. + """ + return self.space.nvir diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py new file mode 100644 index 00000000..09c3c469 --- /dev/null +++ b/ebcc/cc/rebcc.py @@ -0,0 +1,812 @@ +"""Restricted electron-boson coupled cluster.""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +from pyscf import lib + +from ebcc import default_log, init_logging +from ebcc import numpy as np +from ebcc import reom, util +from ebcc.ansatz import Ansatz +from ebcc.brueckner import BruecknerREBCC +from ebcc.cc.base import BaseEBCC +from ebcc.cderis import RCDERIs +from ebcc.damping import DIIS +from ebcc.dump import Dump +from ebcc.eris import RERIs +from ebcc.fock import RFock +from ebcc.logging import ANSI +from ebcc.precision import types +from ebcc.space import Space + +if TYPE_CHECKING: + from typing import Optional + + from pyscf.scf.hf import RHF, SCF + + from ebcc.cc.base import ERIsInputType + from ebcc.numpy.typing import NDArray + from ebcc.util import Namespace + + AmplitudeType = NDArray[float] + + +class REBCC(BaseEBCC): + """Restricted electron-boson coupled cluster. + + Attributes: + mf: PySCF mean-field object. + log: Log to write output to. + options: Options for the EBCC calculation. + e_corr: Correlation energy. + amplitudes: Cluster amplitudes. + converged: Convergence flag. + lambdas: Cluster lambda amplitudes. + converged_lambda: Lambda convergence flag. + name: Name of the method. + """ + + # Types + ERIs = RERIs + Fock = RFock + CDERIs = RCDERIs + Brueckner = BruecknerREBCC + + @property + def spin_type(self): + """Get a string representation of the spin type.""" + return "R" + + def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.IP_REOM: + """Get the IP-EOM object. + + Args: + options: Options for the IP-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + IP-EOM object. + """ + return reom.IP_REOM(self, options=options, **kwargs) + + def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.EA_REOM: + """Get the EA-EOM object. + + Args: + options: Options for the EA-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + EA-EOM object. + """ + return reom.EA_REOM(self, options=options, **kwargs) + + def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.EE_REOM: + """Get the EE-EOM object. + + Args: + options: Options for the EE-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + EE-EOM object. + """ + return reom.EE_REOM(self, options=options, **kwargs) + + @staticmethod + def _convert_mf(mf: SCF) -> RHF: + """Convert the mean-field object to the appropriate type.""" + return mf.to_rhf() + + def init_space(self) -> Space: + """Initialise the fermionic space. + + Returns: + Fermionic space. All fermionic degrees of freedom are assumed to be correlated. + """ + space = Space( + self.mo_occ > 0, + np.zeros_like(self.mo_occ, dtype=bool), + np.zeros_like(self.mo_occ, dtype=bool), + ) + return space + + def _pack_codegen_kwargs(self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None) -> dict[str, Any]: + """Pack all the keyword arguments for the generated code.""" + kwargs = dict( + f=self.fock, + v=self.get_eris(eris), + g=self.g, + G=self.G, + w=np.diag(self.omega) if self.omega is not None else None, + space=self.space, + nocc=self.space.ncocc, + nvir=self.space.ncvir, + nbos=self.nbos, + ) + + if isinstance(kwargs["v"], self.CDERIs): + kwargs["naux"] = self.mf.with_df.get_naoaux() + + for kw in extra_kwargs: + if kw is not None: + kwargs.update(kw) + + return kwargs + + def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + """Initialise the cluster amplitudes. + + Args: + eris: Electron repulsion integrals. + + Returns: + Initial cluster amplitudes. + """ + eris = self.get_eris(eris) + amplitudes = util.Namespace() + + # Build T amplitudes: + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + if n == 1: + amplitudes[name] = self.fock[key] / self.energy_sum(key) + elif n == 2: + key_t = key[0] + key[2] + key[1] + key[3] + amplitudes[name] = eris[key_t].swapaxes(1, 2) / self.energy_sum(key) + else: + shape = tuple(self.space.size(k) for k in key) + amplitudes[name] = np.zeros(shape, dtype=types[float]) + + if self.boson_ansatz: + # Only true for real-valued couplings: + h = self.g + H = self.G + + # Build S amplitudes: + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + if n == 1: + amplitudes[name] = -H / self.omega + else: + shape = (self.nbos,) * n + amplitudes[name] = np.zeros(shape, dtype=types[float]) + + # Build U amplitudes: + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + if nf != 1: + raise util.ModelNotImplemented + if nb == 1: + amplitudes[name] = h[key] / self.energy_sum(key) + else: + shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) + amplitudes[name] = np.zeros(shape, dtype=types[float]) + + return amplitudes + + def init_lams(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + """Initialise the cluster lambda amplitudes. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Initial cluster lambda amplitudes. + """ + if amplitudes is None: + amplitudes = self.amplitudes + lambdas = util.Namespace() + + # Build L amplitudes: + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + lname = name.replace("t", "l") + perm = list(range(n, 2 * n)) + list(range(n)) + lambdas[lname] = amplitudes[name].transpose(perm) + + # Build LS amplitudes: + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + lname = "l" + name + lambdas[lname] = amplitudes[name] + + # Build LU amplitudes: + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + if nf != 1: + raise util.ModelNotImplemented + lname = "l" + name + perm = list(range(nb)) + [nb + 1, nb] + lambdas[lname] = amplitudes[name].transpose(perm) + + return lambdas + + def update_amps( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> Namespace[AmplitudeType]: + """Update the cluster amplitudes. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + + Returns: + Updated cluster amplitudes. + """ + func, kwargs = self._load_function( + "update_amps", + eris=eris, + amplitudes=amplitudes, + ) + res = func(**kwargs) + res = {key.rstrip("new"): val for key, val in res.items()} + + # Divide T amplitudes: + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + res[name] /= self.energy_sum(key) + res[name] += amplitudes[name] + + # Divide S amplitudes: + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + res[name] /= self.energy_sum(key) + res[name] += amplitudes[name] + + # Divide U amplitudes: + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + if nf != 1: + raise util.ModelNotImplemented + res[name] /= self.energy_sum(key) + res[name] += amplitudes[name] + + return res + + def update_lams( + self, + eris: ERIsInputType = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + lambdas_pert: Optional[Namespace[AmplitudeType]] = None, + perturbative: bool = False, + ) -> Namespace[AmplitudeType]: + """Update the cluster lambda amplitudes. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + lambdas_pert: Perturbative cluster lambda amplitudes. + perturbative: Flag to include perturbative correction. + + Returns: + Updated cluster lambda amplitudes. + """ + # TODO active + if lambdas_pert is not None: + lambdas.update(lambdas_pert) + + func, kwargs = self._load_function( + "update_lams%s" % ("_perturbative" if perturbative else ""), + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + res = func(**kwargs) + res = {key.rstrip("new"): val for key, val in res.items()} + + # Divide T amplitudes: + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + lname = name.replace("t", "l") + res[lname] /= self.energy_sum(key[n:] + key[:n]) + if not perturbative: + res[lname] += lambdas[lname] + + # Divide S amplitudes: + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + lname = "l" + name + res[lname] /= self.energy_sum(key[n:] + key[:n]) + if not perturbative: + res[lname] += lambdas[lname] + + # Divide U amplitudes: + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + if nf != 1: + raise util.ModelNotImplemented + lname = "l" + name + res[lname] /= self.energy_sum(key[:nb] + key[nb + nf :] + key[nb : nb + nf]) + if not perturbative: + res[lname] += lambdas[lname] + + if perturbative: + res = {key + "pert": val for key, val in res.items()} + + return res + + def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + hermitise: Hermitise the density matrix. + + Returns: + One-particle fermion reduced density matrix. + """ + func, kwargs = self._load_function( + "make_rdm1_f", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + dm = func(**kwargs) + + if hermitise: + dm = 0.5 * (dm + dm.T) + + return dm + + def make_rdm2_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + hermitise: Hermitise the density matrix. + + Returns: + Two-particle fermion reduced density matrix. + """ + func, kwargs = self._load_function( + "make_rdm2_f", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + dm = func(**kwargs) + + if hermitise: + dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) + dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(1, 0, 3, 2)) + + return dm + + def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + unshifted: If `self.options.shift` is `True`, return the unshifted density matrix. Has + no effect if `self.options.shift` is `False`. + hermitise: Hermitise the density matrix. + + Returns: + Electron-boson coupling reduced density matrix. + """ + func, kwargs = self._load_function( + "make_eb_coup_rdm", + eris=eris, + amplitudes=amplitudes, + lambdas=lambdas, + ) + dm_eb = func(**kwargs) + + if hermitise: + dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) + dm_eb[1] = dm_eb[0].transpose(0, 2, 1).copy() + + if unshifted and self.options.shift: + rdm1_f = self.make_rdm1_f(hermitise=hermitise) + shift = util.einsum("x,ij->xij", self.xi, rdm1_f) + dm_eb -= shift[None] + + return dm_eb + + def energy_sum(self, subscript: str, signs_dict: dict[str, int] = None) -> NDArray[float]: + """Get a direct sum of energies. + + Args: + subscript: Subscript for the direct sum. + signs_dict: Signs of the energies in the sum. Default sets `("o", "O", "i")` to be + positive, and `("v", "V", "a", "b")` to be negative. + + Returns: + Sum of energies. + """ + n = 0 + + def next_char() -> str: + nonlocal n + if n < 26: + char = chr(ord("a") + n) + else: + char = chr(ord("A") + n) + n += 1 + return char + + if signs_dict is None: + signs_dict = {} + for k, s in zip("vVaoOib", "---+++-"): + if k not in signs_dict: + signs_dict[k] = s + + energies = [] + for key in subscript: + if key == "b": + energies.append(self.omega) + else: + energies.append(np.diag(self.fock[key + key])) + + subscript = "".join([signs_dict[k] + next_char() for k in subscript]) + energy_sum = lib.direct_sum(subscript, *energies) + + return energy_sum + + def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the amplitudes used in the given ansatz. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Cluster amplitudes as a vector. + """ + vectors = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + vectors.append(amplitudes[name].ravel()) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + vectors.append(amplitudes[name].ravel()) + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + vectors.append(amplitudes[name].ravel()) + + return np.concatenate(vectors) + + def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + """Construct a namespace of amplitudes from a vector. + + Args: + vector: Cluster amplitudes as a vector. + + Returns: + Cluster amplitudes. + """ + amplitudes = util.Namespace() + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + shape = tuple(self.space.size(k) for k in key) + size = np.prod(shape) + amplitudes[name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + shape = (self.nbos,) * n + size = np.prod(shape) + amplitudes[name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) + size = np.prod(shape) + amplitudes[name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + return amplitudes + + def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the lambda amplitudes used in the given ansatz. + + Args: + lambdas: Cluster lambda amplitudes. + + Returns: + Cluster lambda amplitudes as a vector. + """ + vectors = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + vectors.append(lambdas[name.replace("t", "l")].ravel()) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + vectors.append(lambdas["l" + name].ravel()) + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + vectors.append(lambdas["l" + name].ravel()) + + return np.concatenate(vectors) + + def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + """Construct a namespace of lambda amplitudes from a vector. + + Args: + vector: Cluster lambda amplitudes as a vector. + + Returns: + Cluster lambda amplitudes. + """ + lambdas = util.Namespace() + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + lname = name.replace("t", "l") + key = key[n:] + key[:n] + shape = tuple(self.space.size(k) for k in key) + size = np.prod(shape) + lambdas[lname] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + shape = (self.nbos,) * n + size = np.prod(shape) + lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + key = key[:nb] + key[nb + nf :] + key[nb : nb + nf] + shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) + size = np.prod(shape) + lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) + i0 += size + + return lambdas + + def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the IP-EOM excitations. + + Args: + excitations: IP-EOM excitations. + + Returns: + IP-EOM excitations as a vector. + """ + vectors = [] + m = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + vectors.append(excitations[m].ravel()) + m += 1 + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return np.concatenate(vectors) + + def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the EA-EOM excitations. + + Args: + excitations: EA-EOM excitations. + + Returns: + EA-EOM excitations as a vector. + """ + return self.excitations_to_vector_ip(*excitations) + + def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the EE-EOM excitations. + + Args: + excitations: EE-EOM excitations. + + Returns: + EE-EOM excitations as a vector. + """ + return self.excitations_to_vector_ip(*excitations) + + def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of IP-EOM excitations from a vector. + + Args: + vector: IP-EOM excitations as a vector. + + Returns: + IP-EOM excitations. + """ + excitations = [] + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[:-1] + shape = tuple(self.space.size(k) for k in key) + size = np.prod(shape) + excitations.append(vector[i0 : i0 + size].reshape(shape)) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return tuple(excitations) + + def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of EA-EOM excitations from a vector. + + Args: + vector: EA-EOM excitations as a vector. + + Returns: + EA-EOM excitations. + """ + excitations = [] + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[n:] + key[: n - 1] + shape = tuple(self.space.size(k) for k in key) + size = np.prod(shape) + excitations.append(vector[i0 : i0 + size].reshape(shape)) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return tuple(excitations) + + def vector_to_excitations_ee(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of EE-EOM excitations from a vector. + + Args: + vector: EE-EOM excitations as a vector. + + Returns: + EE-EOM excitations. + """ + excitations = [] + i0 = 0 + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + shape = tuple(self.space.size(k) for k in key) + size = np.prod(shape) + excitations.append(vector[i0 : i0 + size].reshape(shape)) + i0 += size + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return tuple(excitations) + + def get_mean_field_G(self) -> NDArray[float]: + """Get the mean-field boson non-conserving term. + + Returns: + Mean-field boson non-conserving term. + """ + # FIXME should this also sum in frozen orbitals? + val = lib.einsum("Ipp->I", self.g.boo) * 2.0 + val -= self.xi * self.omega + if self.bare_G is not None: + val += self.bare_G + return val + + def get_g(self, g: NDArray[float]) -> Namespace[NDArray[float]]: + """Get the blocks of the electron-boson coupling matrix. + + This matrix corresponds to the bosonic annihilation operator. + + Args: + g: Electron-boson coupling matrix. + + Returns: + Blocks of the electron-boson coupling matrix. + """ + # TODO make a proper class for this + slices = { + "x": self.space.correlated, + "o": self.space.correlated_occupied, + "v": self.space.correlated_virtual, + "O": self.space.active_occupied, + "V": self.space.active_virtual, + "i": self.space.inactive_occupied, + "a": self.space.inactive_virtual, + } + + class Blocks(util.Namespace): + def __getitem__(selffer, key): + assert key[0] == "b" + i = slices[key[1]] + j = slices[key[2]] + return g[:, i][:, :, j].copy() + + __getattr__ = __getitem__ + + return Blocks() + + @property + def bare_fock(self) -> NDArray[float]: + """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 = self.mf.get_fock().astype(types[float]) + fock = util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff, self.mo_coeff) + return fock + + @property + def xi(self) -> NDArray[float]: + """Get the shift in the bosonic operators to diagonalise the photon Hamiltonian. + + Returns: + Shift in the bosonic operators. + """ + if self.options.shift: + xi = lib.einsum("Iii->I", self.g.boo) * 2.0 + xi /= self.omega + if self.bare_G is not None: + xi += self.bare_G / self.omega + else: + xi = np.zeros_like(self.omega, 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) + + 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. + """ + if (eris is None) or isinstance(eris, np.ndarray): + if (isinstance(eris, np.ndarray) and eris.ndim == 3) or getattr( + self.mf, "with_df", None + ): + return self.CDERIs(self, array=eris) + else: + return self.ERIs(self, array=eris) + else: + return eris + + @property + def nmo(self) -> int: + """Get the number of molecular orbitals. + + Returns: + Number of molecular orbitals. + """ + return self.space.nmo + + @property + def nocc(self) -> int: + """Get the number of occupied molecular orbitals. + + Returns: + Number of occupied molecular orbitals. + """ + return self.space.nocc + + @property + def nvir(self) -> int: + """Get the number of virtual molecular orbitals. + + Returns: + Number of virtual molecular orbitals. + """ + return self.space.nvir diff --git a/ebcc/uebcc.py b/ebcc/cc/uebcc.py similarity index 70% rename from ebcc/uebcc.py rename to ebcc/cc/uebcc.py index 352adccc..5166159e 100644 --- a/ebcc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -1,31 +1,103 @@ """Unrestricted electron-boson coupled cluster.""" +from __future__ import annotations + +from typing import TYPE_CHECKING + from pyscf import lib from ebcc import numpy as np from ebcc import rebcc, ueom, util from ebcc.brueckner import BruecknerUEBCC +from ebcc.cc.base import BaseEBCC from ebcc.cderis import UCDERIs from ebcc.eris import UERIs from ebcc.fock import UFock from ebcc.precision import types from ebcc.space import Space +if TYPE_CHECKING: + from ebcc.cc.rebcc import REBCC + from ebcc.numpy.typing import NDArray + -class UEBCC(rebcc.REBCC): +class UEBCC(BaseEBCC): + """Unrestricted electron-boson coupled cluster. + + Attributes: + mf: PySCF mean-field object. + log: Log to write output to. + options: Options for the EBCC calculation. + e_corr: Correlation energy. + amplitudes: Cluster amplitudes. + converged: Convergence flag. + lambdas: Cluster lambda amplitudes. + converged_lambda: Lambda convergence flag. + name: Name of the method. + """ + + # Types ERIs = UERIs Fock = UFock CDERIs = UCDERIs Brueckner = BruecknerUEBCC + @property + def spin_type(self): + """Get a string representation of the spin type.""" + return "U" + + def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.IP_REOM: + """Get the IP-EOM object. + + Args: + options: Options for the IP-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + IP-EOM object. + """ + return ueom.IP_UEOM(self, options=options, **kwargs) + + def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.EA_REOM: + """Get the EA-EOM object. + + Args: + options: Options for the EA-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + EA-EOM object. + """ + return ueom.EA_UEOM(self, options=options, **kwargs) + + def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.EE_REOM: + """Get the EE-EOM object. + + Args: + options: Options for the EE-EOM calculation. + **kwargs: Additional keyword arguments. + + Returns: + EE-EOM object. + """ + return ueom.EE_UEOM(self, options=options, **kwargs) + @staticmethod - def _convert_mf(mf): + def _convert_mf(mf: SCF) -> UHF: + """Convert the mean-field object to the appropriate type.""" return mf.to_uhf() @classmethod - def from_rebcc(cls, rcc): - """Initialise an UEBCC object from an REBCC object.""" + def from_rebcc(cls, rcc: REBCC) -> UEBCC: + """Initialise an `UEBCC` object from an `REBCC` object. + Args: + rcc: Restricted electron-boson coupled cluster object. + + Returns: + UEBCC object. + """ ucc = cls( rcc.mf, log=rcc.log, @@ -93,47 +165,58 @@ def from_rebcc(cls, rcc): return ucc - def _pack_codegen_kwargs(self, *extra_kwargs, eris=None): - eris = self.get_eris(eris) + def init_space(self) -> Namespace[Space]: + """Initialise the fermionic space. - omega = np.diag(self.omega) if self.omega is not None else None + Returns: + Fermionic space. All fermionic degrees of freedom are assumed to be correlated. + """ + space = ( + Space( + self.mo_occ[0] > 0, + np.zeros_like(self.mo_occ[0], dtype=bool), + np.zeros_like(self.mo_occ[0], dtype=bool), + ), + Space( + self.mo_occ[1] > 0, + np.zeros_like(self.mo_occ[1], dtype=bool), + np.zeros_like(self.mo_occ[1], dtype=bool), + ), + ) + return space + def _pack_codegen_kwargs(self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None) -> dict[str, Any]: + """Pack all the keyword arguments for the generated code.""" kwargs = dict( f=self.fock, - v=eris, + v=self.get_eris(eris), g=self.g, G=self.G, - w=omega, + w=np.diag(self.omega) if self.omega is not None else None, space=self.space, nocc=(self.space[0].ncocc, self.space[1].ncocc), # FIXME rename? nvir=(self.space[0].ncvir, self.space[1].ncvir), # FIXME rename? nbos=self.nbos, ) + if isinstance(eris, self.CDERIs): kwargs["naux"] = self.mf.with_df.get_naoaux() + for kw in extra_kwargs: if kw is not None: kwargs.update(kw) return kwargs - def init_space(self): - space = ( - Space( - self.mo_occ[0] > 0, - np.zeros_like(self.mo_occ[0], dtype=bool), - np.zeros_like(self.mo_occ[0], dtype=bool), - ), - Space( - self.mo_occ[1] > 0, - np.zeros_like(self.mo_occ[1], dtype=bool), - np.zeros_like(self.mo_occ[1], dtype=bool), - ), - ) + def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + """Initialise the cluster amplitudes. - return space + Args: + eris: Electron repulsion integrals. - def init_amps(self, eris=None): + Returns: + Initial cluster amplitudes. + """ eris = self.get_eris(eris) amplitudes = util.Namespace() @@ -193,10 +276,17 @@ def init_amps(self, eris=None): return amplitudes - def init_lams(self, amplitudes=None): + def init_lams(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + """Initialise the cluster lambda amplitudes. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Initial cluster lambda amplitudes. + """ if amplitudes is None: amplitudes = self.amplitudes - lambdas = util.Namespace() # Build L amplitudes: @@ -224,7 +314,20 @@ def init_lams(self, amplitudes=None): return lambdas - def update_amps(self, eris=None, amplitudes=None): + def update_amps( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> Namespace[AmplitudeType]: + """Update the cluster amplitudes. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + + Returns: + Updated cluster amplitudes. + """ func, kwargs = self._load_function( "update_amps", eris=eris, @@ -263,7 +366,27 @@ def update_amps(self, eris=None, amplitudes=None): return res - def update_lams(self, eris=None, amplitudes=None, lambdas=None, lambdas_pert=None): + def update_lams( + self, + eris: ERIsInputType = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + lambdas_pert: Optional[Namespace[AmplitudeType]] = None, + perturbative: bool = False, + ) -> Namespace[AmplitudeType]: + """Update the cluster lambda amplitudes. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + lambdas_pert: Perturbative cluster lambda amplitudes. + perturbative: Flag to include perturbative correction. + + Returns: + Updated cluster lambda amplitudes. + """ + # TODO active func, kwargs = self._load_function( "update_lams", eris=eris, @@ -308,7 +431,18 @@ def update_lams(self, eris=None, amplitudes=None, lambdas=None, lambdas_pert=Non return res - def make_rdm1_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): + def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + hermitise: Hermitise the density matrix. + + Returns: + One-particle fermion reduced density matrix. + """ func, kwargs = self._load_function( "make_rdm1_f", eris=eris, @@ -324,7 +458,18 @@ def make_rdm1_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): return dm - def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): + def make_rdm2_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + hermitise: Hermitise the density matrix. + + Returns: + Two-particle fermion reduced density matrix. + """ func, kwargs = self._load_function( "make_rdm2_f", eris=eris, @@ -350,14 +495,20 @@ def transpose2(dm): return dm - def make_eb_coup_rdm( - self, - eris=None, - amplitudes=None, - lambdas=None, - unshifted=True, - hermitise=True, - ): + def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + unshifted: If `self.options.shift` is `True`, return the unshifted density matrix. Has + no effect if `self.options.shift` is `False`. + hermitise: Hermitise the density matrix. + + Returns: + Electron-boson coupling reduced density matrix. + """ func, kwargs = self._load_function( "make_eb_coup_rdm", eris=eris, @@ -382,116 +533,56 @@ def make_eb_coup_rdm( return dm_eb - def get_mean_field_G(self): - val = lib.einsum("Ipp->I", self.g.aa.boo) - val += lib.einsum("Ipp->I", self.g.bb.boo) - val -= self.xi * self.omega + def energy_sum(self, subscript: str, spins: str, signs_dict: dict[str, int] = None) -> NDArray[float]: + """Get a direct sum of energies. - if self.bare_G is not None: - # Require bare_G to have a spin index for now: - assert np.shape(self.bare_G) == val.shape - val += self.bare_G - - return val - - def get_g(self, g): - if np.array(g).ndim != 4: - g = np.array([g, g]) - - slices = [ - { - "x": space.correlated, - "o": space.correlated_occupied, - "v": space.correlated_virtual, - "O": space.active_occupied, - "V": space.active_virtual, - "i": space.inactive_occupied, - "a": space.inactive_virtual, - } - for space in self.space - ] - - def constructor(s): - class Blocks(util.Namespace): - def __getitem__(selffer, key): - assert key[0] == "b" - i = slices[s][key[1]] - j = slices[s][key[2]] - return g[s][:, i][:, :, j].copy() - - __getattr__ = __getitem__ - - return Blocks() + Args: + subscript: Subscript for the direct sum. + spins: Spins for energies. + signs_dict: Signs of the energies in the sum. Default sets `("o", "O", "i")` to be + positive, and `("v", "V", "a", "b")` to be negative. - gs = util.Namespace() - gs.aa = constructor(0) - gs.bb = constructor(1) + Returns: + Sum of energies. + """ + n = 0 - return gs + def next_char() -> str: + nonlocal n + if n < 26: + char = chr(ord("a") + n) + else: + char = chr(ord("A") + n) + n += 1 + return char - @property - def bare_fock(self): - fock = lib.einsum( - "npq,npi,nqj->nij", - self.mf.get_fock().astype(types[float]), - self.mo_coeff, - self.mo_coeff, - ) - fock = util.Namespace(aa=fock[0], bb=fock[1]) - return fock + if signs_dict is None: + signs_dict = {} + for k, s in zip("vVaoOib", "---+++-"): + if k not in signs_dict: + signs_dict[k] = s - @property - def xi(self): - if self.options.shift: - xi = lib.einsum("Iii->I", self.g.aa.boo) - xi += lib.einsum("Iii->I", self.g.bb.boo) - xi /= self.omega - if self.bare_G is not None: - xi += self.bare_G / self.omega - else: - xi = np.zeros_like(self.omega) + energies = [] + for key, spin in zip(subscript, spins): + if key == "b": + energies.append(self.omega) + else: + energies.append(np.diag(self.fock[spin + spin][key + key])) - return xi + subscript = "".join([signs_dict[k] + next_char() for k in subscript]) + energy_sum = lib.direct_sum(subscript, *energies) - def get_fock(self): - return self.Fock(self, array=(self.bare_fock.aa, self.bare_fock.bb)) + return energy_sum - def get_eris(self, eris=None): - """Get blocks of the ERIs. + def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the amplitudes used in the given ansatz. - Parameters - ---------- - eris : tuple of np.ndarray or ERIs, optional. - Electronic repulsion integrals, either in the form of a - dense array for each spin channel or an ERIs object. - Default value is `None`. + Args: + amplitudes: Cluster amplitudes. - Returns - ------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.ERIs()`. + Returns: + Cluster amplitudes as a vector. """ - if (eris is None) or isinstance(eris, tuple): - if ( - isinstance(eris, tuple) and isinstance(eris[0], np.ndarray) and eris[0].ndim == 3 - ) or getattr(self.mf, "with_df", None): - return self.CDERIs(self, array=eris) - else: - return self.ERIs(self, array=eris) - else: - return eris - - def ip_eom(self, options=None, **kwargs): - return ueom.IP_UEOM(self, options=options, **kwargs) - - def ea_eom(self, options=None, **kwargs): - return ueom.EA_UEOM(self, options=options, **kwargs) - - def ee_eom(self, options=None, **kwargs): - return ueom.EE_UEOM(self, options=options, **kwargs) - - def amplitudes_to_vector(self, amplitudes): vectors = [] for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -511,7 +602,15 @@ def amplitudes_to_vector(self, amplitudes): return np.concatenate(vectors) - def vector_to_amplitudes(self, vector): + def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + """Construct a namespace of amplitudes from a vector. + + Args: + vector: Cluster amplitudes as a vector. + + Returns: + Cluster amplitudes. + """ amplitudes = util.Namespace() i0 = 0 sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} @@ -550,7 +649,15 @@ def vector_to_amplitudes(self, vector): return amplitudes - def lambdas_to_vector(self, lambdas): + def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the lambda amplitudes used in the given ansatz. + + Args: + lambdas: Cluster lambda amplitudes. + + Returns: + Cluster lambda amplitudes as a vector. + """ vectors = [] for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -572,7 +679,15 @@ def lambdas_to_vector(self, lambdas): return np.concatenate(vectors) - def vector_to_lambdas(self, vector): + def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + """Construct a namespace of lambda amplitudes from a vector. + + Args: + vector: Cluster lambda amplitudes as a vector. + + Returns: + Cluster lambda amplitudes. + """ lambdas = util.Namespace() i0 = 0 sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} @@ -615,7 +730,15 @@ def vector_to_lambdas(self, vector): return lambdas - def excitations_to_vector_ip(self, *excitations): + def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the IP-EOM excitations. + + Args: + excitations: IP-EOM excitations. + + Returns: + IP-EOM excitations as a vector. + """ vectors = [] m = 0 @@ -634,7 +757,15 @@ def excitations_to_vector_ip(self, *excitations): return np.concatenate(vectors) - def excitations_to_vector_ea(self, *excitations): + def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the EA-EOM excitations. + + Args: + excitations: EA-EOM excitations. + + Returns: + EA-EOM excitations as a vector. + """ vectors = [] m = 0 @@ -654,7 +785,15 @@ def excitations_to_vector_ea(self, *excitations): return np.concatenate(vectors) - def excitations_to_vector_ee(self, *excitations): + def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + """Construct a vector containing all of the EE-EOM excitations. + + Args: + excitations: EE-EOM excitations. + + Returns: + EE-EOM excitations as a vector. + """ vectors = [] m = 0 @@ -673,7 +812,15 @@ def excitations_to_vector_ee(self, *excitations): return np.concatenate(vectors) - def vector_to_excitations_ip(self, vector): + def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of IP-EOM excitations from a vector. + + Args: + vector: IP-EOM excitations as a vector. + + Returns: + IP-EOM excitations. + """ excitations = [] i0 = 0 sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} @@ -705,7 +852,15 @@ def vector_to_excitations_ip(self, vector): return tuple(excitations) - def vector_to_excitations_ea(self, vector): + def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of EA-EOM excitations from a vector. + + Args: + vector: EA-EOM excitations as a vector. + + Returns: + EA-EOM excitations. + """ excitations = [] i0 = 0 sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} @@ -737,7 +892,15 @@ def vector_to_excitations_ea(self, vector): return tuple(excitations) - def vector_to_excitations_ee(self, vector): + def vector_to_excitations_ee(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + """Construct a namespace of EE-EOM excitations from a vector. + + Args: + vector: EE-EOM excitations as a vector. + + Returns: + EE-EOM excitations. + """ excitations = [] i0 = 0 sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} @@ -768,74 +931,153 @@ def vector_to_excitations_ee(self, vector): return tuple(excitations) - @property - def spin_type(self): - return "U" + def get_mean_field_G(self) -> NDArray[float]: + """Get the mean-field boson non-conserving term. - @property - def nmo(self): - assert self.mo_occ[0].size == self.mo_occ[1].size - return self.mo_occ[0].size + Returns: + Mean-field boson non-conserving term. + """ + # FIXME should this also sum in frozen orbitals? + val = lib.einsum("Ipp->I", self.g.aa.boo) + val += lib.einsum("Ipp->I", self.g.bb.boo) + val -= self.xi * self.omega + if self.bare_G is not None: + # Require bare_G to have a spin index for now: + assert np.shape(self.bare_G) == val.shape + val += self.bare_G + return val + + def get_g(self, g: NDArray[float]) -> Namespace[Namespace[NDArray[float]]]: + """Get the blocks of the electron-boson coupling matrix. + + This matrix corresponds to the bosonic annihilation operator. + + Args: + g: Electron-boson coupling matrix. + + Returns: + Blocks of the electron-boson coupling matrix. + """ + # TODO make a proper class for this + if np.array(g).ndim != 4: + g = np.array([g, g]) + slices = [ + { + "x": space.correlated, + "o": space.correlated_occupied, + "v": space.correlated_virtual, + "O": space.active_occupied, + "V": space.active_virtual, + "i": space.inactive_occupied, + "a": space.inactive_virtual, + } + for space in self.space + ] + + def constructor(s): + class Blocks(util.Namespace): + def __getitem__(selffer, key): + assert key[0] == "b" + i = slices[s][key[1]] + j = slices[s][key[2]] + return g[s][:, i][:, :, j].copy() + + __getattr__ = __getitem__ + + return Blocks() + + gs = util.Namespace() + gs.aa = constructor(0) + gs.bb = constructor(1) + + return gs @property - def nocc(self): - return tuple(np.sum(mo_occ > 0) for mo_occ in self.mo_occ) + def bare_fock(self) -> Namespace[NDArray[float]]: + """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 = lib.einsum( + "npq,npi,nqj->nij", + self.mf.get_fock().astype(types[float]), + self.mo_coeff, + self.mo_coeff, + ) + fock = util.Namespace(aa=fock[0], bb=fock[1]) + return fock @property - def nvir(self): - return tuple(self.nmo - nocc for nocc in self.nocc) + def xi(self) -> NDArray[float]: + """Get the shift in the bosonic operators to diagonalise the photon Hamiltonian. - def energy_sum(self, subscript, spins, signs_dict=None): + Returns: + Shift in the bosonic operators. """ - Get a direct sum of energies. + if self.options.shift: + xi = lib.einsum("Iii->I", self.g.aa.boo) + xi += lib.einsum("Iii->I", self.g.bb.boo) + xi /= self.omega + if self.bare_G is not None: + xi += self.bare_G / self.omega + else: + xi = np.zeros_like(self.omega) + return xi - Parameters - ---------- - subscript : str - The direct sum subscript, where each character indicates the - sector for each energy. For the default slice characters, see - `Space`. Occupied degrees of freedom are assumed to be - positive, virtual and bosonic negative (the signs can be - changed via the `signs_dict` keyword argument). - spins : str - String of spins, length must be the same as `subscript` with - each character being one of `"a"` or `"b"`. - signs_dict : dict, optional - Dictionary defining custom signs for each sector. If `None`, - initialised such that `["o", "O", "i"]` are positive, and - `["v", "V", "a", "b"]` negative. Default value is `None`. + def get_fock(self) -> UFock: + """Get the Fock matrix. - Returns - ------- - energy_sum : numpy.ndarray - Array of energy sums. + Returns: + Fock matrix. """ + return self.Fock(self, array=(self.bare_fock.aa, self.bare_fock.bb)) - n = 0 + def get_eris(self, eris: Optional[ERIsInputType] = None) -> Union[UERIs, UCDERIs]: + """Get the electron repulsion integrals. - def next_char(): - nonlocal n - if n < 26: - char = chr(ord("a") + n) + Args: + eris: Input electron repulsion integrals. + + Returns: + Electron repulsion integrals. + """ + if (eris is None) or isinstance(eris, tuple): + if ( + isinstance(eris, tuple) and isinstance(eris[0], np.ndarray) and eris[0].ndim == 3 + ) or getattr(self.mf, "with_df", None): + return self.CDERIs(self, array=eris) else: - char = chr(ord("A") + n) - n += 1 - return char + return self.ERIs(self, array=eris) + else: + return eris - if signs_dict is None: - signs_dict = {} - for k, s in zip("vVaoOib", "---+++-"): - if k not in signs_dict: - signs_dict[k] = s + @property + def nmo(self) -> int: + """Get the number of molecular orbitals. - energies = [] - for key, spin in zip(subscript, spins): - if key == "b": - energies.append(self.omega) - else: - energies.append(np.diag(self.fock[spin + spin][key + key])) + Returns: + Number of molecular orbitals. + """ + assert self.mo_occ[0].size == self.mo_occ[1].size + return self.mo_occ[0].size - subscript = "".join([signs_dict[k] + next_char() for k in subscript]) - energy_sum = lib.direct_sum(subscript, *energies) + @property + def nocc(self) -> tuple[int, int]: + """Get the number of occupied molecular orbitals. - return energy_sum + Returns: + Number of occupied molecular orbitals for each spin. + """ + return tuple(np.sum(mo_occ > 0) for mo_occ in self.mo_occ) + + @property + def nvir(self) -> tuple[int, int]: + """Get the number of virtual molecular orbitals. + + Returns: + Number of virtual molecular orbitals for each spin. + """ + return tuple(self.nmo - nocc for nocc in self.nocc) diff --git a/ebcc/fock.py b/ebcc/fock.py index 7a47217a..ff9360c6 100644 --- a/ebcc/fock.py +++ b/ebcc/fock.py @@ -2,8 +2,8 @@ from ebcc import numpy as np from ebcc import util -from ebcc.precision import types from ebcc.base import Fock +from ebcc.precision import types class RFock(Fock): @@ -78,7 +78,7 @@ def __init__(self, ebcc, array=None, slices=None, mo_coeff=None, g=None): def __getattr__(self, key): """Just-in-time attribute getter.""" - if key not in self.__dict__.keys(): + if key not in self.__dict__: ki, kj = key i = self.slices[0][ki] j = self.slices[1][kj] diff --git a/ebcc/gebcc.py b/ebcc/gebcc.py deleted file mode 100644 index da6da863..00000000 --- a/ebcc/gebcc.py +++ /dev/null @@ -1,465 +0,0 @@ -"""General electron-boson coupled cluster.""" - -from pyscf import lib, scf - -from ebcc import geom -from ebcc import numpy as np -from ebcc import uebcc, util -from ebcc.brueckner import BruecknerGEBCC -from ebcc.eris import GERIs -from ebcc.fock import GFock -from ebcc.precision import types -from ebcc.rebcc import REBCC -from ebcc.space import Space - - -class GEBCC(REBCC): - __doc__ = __doc__.replace("Restricted", "Generalised", 1) - - ERIs = GERIs - Fock = GFock - Brueckner = BruecknerGEBCC - - @staticmethod - def _convert_mf(mf): - if isinstance(mf, scf.ghf.GHF): - return mf - # NOTE: First convert to UHF - otherwise conversions from - # RHF->GHF and UHF->GHF may have inconsistent ordering - return mf.to_uhf().to_ghf() - - @classmethod - def from_uebcc(cls, ucc): - """Initialise a `GEBCC` object from an `UEBCC` object. - - Parameters - ---------- - ucc : UEBCC - The UEBCC object to initialise from. - - Returns - ------- - gcc : GEBCC - The GEBCC object. - """ - - orbspin = scf.addons.get_ghf_orbspin(ucc.mf.mo_energy, ucc.mf.mo_occ, False) - nocc = ucc.space[0].nocc + ucc.space[1].nocc - nvir = ucc.space[0].nvir + ucc.space[1].nvir - nbos = ucc.nbos - sa = np.where(orbspin == 0)[0] - sb = np.where(orbspin == 1)[0] - - occupied = np.zeros((nocc + nvir,), dtype=bool) - occupied[sa] = ucc.space[0]._occupied.copy() - occupied[sb] = ucc.space[1]._occupied.copy() - frozen = np.zeros((nocc + nvir,), dtype=bool) - frozen[sa] = ucc.space[0]._frozen.copy() - frozen[sb] = ucc.space[1]._frozen.copy() - active = np.zeros((nocc + nvir,), dtype=bool) - active[sa] = ucc.space[0]._active.copy() - active[sb] = ucc.space[1]._active.copy() - space = Space(occupied, frozen, active) - - slices = util.Namespace( - a=util.Namespace(**{k: np.where(orbspin[space.mask(k)] == 0)[0] for k in "oOivVa"}), - b=util.Namespace(**{k: np.where(orbspin[space.mask(k)] == 1)[0] for k in "oOivVa"}), - ) - - if ucc.bare_g is not None: - if np.asarray(ucc.bare_g).ndim == 3: - bare_g_a = bare_g_b = ucc.bare_g - else: - bare_g_a, bare_g_b = ucc.bare_g - g = np.zeros((ucc.nbos, ucc.nmo * 2, ucc.nmo * 2)) - g[np.ix_(range(ucc.nbos), sa, sa)] = bare_g_a.copy() - g[np.ix_(range(ucc.nbos), sb, sb)] = bare_g_b.copy() - else: - g = None - - gcc = cls( - ucc.mf, - log=ucc.log, - ansatz=ucc.ansatz, - space=space, - omega=ucc.omega, - g=g, - G=ucc.bare_G, - options=ucc.options, - ) - - gcc.e_corr = ucc.e_corr - gcc.converged = ucc.converged - gcc.converged_lambda = ucc.converged_lambda - - has_amps = ucc.amplitudes is not None - has_lams = ucc.lambdas is not None - - if has_amps: - amplitudes = util.Namespace() - - for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): - shape = tuple(space.size(k) for k in key) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - for comb in util.generate_spin_combinations(n, unique=True): - done = set() - for lperm, lsign in util.permutations_with_signs(tuple(range(n))): - for uperm, usign in util.permutations_with_signs(tuple(range(n))): - combn = util.permute_string(comb[:n], lperm) - combn += util.permute_string(comb[n:], uperm) - if combn in done: - continue - mask = np.ix_(*[slices[s][k] for s, k in zip(combn, key)]) - transpose = tuple(lperm) + tuple(p + n for p in uperm) - amp = ( - getattr(ucc.amplitudes[name], comb).transpose(transpose) - * lsign - * usign - ) - for perm, sign in util.permutations_with_signs(tuple(range(n))): - transpose = tuple(perm) + tuple(range(n, 2 * n)) - if util.permute_string(comb[:n], perm) == comb[:n]: - amplitudes[name][mask] += amp.transpose(transpose).copy() * sign - done.add(combn) - - for name, key, n in ucc.ansatz.bosonic_cluster_ranks(spin_type=ucc.spin_type): - amplitudes[name] = ucc.amplitudes[name].copy() - - for name, key, nf, nb in ucc.ansatz.coupling_cluster_ranks(spin_type=ucc.spin_type): - shape = (nbos,) * nb + tuple(space.size(k) for k in key[nb:]) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - for comb in util.generate_spin_combinations(nf): - done = set() - for lperm, lsign in util.permutations_with_signs(tuple(range(nf))): - for uperm, usign in util.permutations_with_signs(tuple(range(nf))): - combn = util.permute_string(comb[:nf], lperm) - combn += util.permute_string(comb[nf:], uperm) - if combn in done: - continue - mask = np.ix_( - *([range(nbos)] * nb), - *[slices[s][k] for s, k in zip(combn, key[nb:])], - ) - transpose = ( - tuple(range(nb)) - + tuple(p + nb for p in lperm) - + tuple(p + nb + nf for p in uperm) - ) - amp = ( - getattr(ucc.amplitudes[name], comb).transpose(transpose) - * lsign - * usign - ) - for perm, sign in util.permutations_with_signs(tuple(range(nf))): - transpose = ( - tuple(range(nb)) - + tuple(p + nb for p in perm) - + tuple(range(nb + nf, nb + 2 * nf)) - ) - if util.permute_string(comb[:nf], perm) == comb[:nf]: - amplitudes[name][mask] += amp.transpose(transpose).copy() * sign - done.add(combn) - - gcc.amplitudes = amplitudes - - if has_lams: - lambdas = gcc.init_lams() # Easier this way - but have to build ERIs... - - for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): - lname = name.replace("t", "l") - shape = tuple(space.size(k) for k in key[n:] + key[:n]) - lambdas[lname] = np.zeros(shape, dtype=types[float]) - for comb in util.generate_spin_combinations(n, unique=True): - done = set() - for lperm, lsign in util.permutations_with_signs(tuple(range(n))): - for uperm, usign in util.permutations_with_signs(tuple(range(n))): - combn = util.permute_string(comb[:n], lperm) - combn += util.permute_string(comb[n:], uperm) - if combn in done: - continue - mask = np.ix_(*[slices[s][k] for s, k in zip(combn, key[n:] + key[:n])]) - transpose = tuple(lperm) + tuple(p + n for p in uperm) - amp = ( - getattr(ucc.lambdas[lname], comb).transpose(transpose) - * lsign - * usign - ) - for perm, sign in util.permutations_with_signs(tuple(range(n))): - transpose = tuple(perm) + tuple(range(n, 2 * n)) - if util.permute_string(comb[:n], perm) == comb[:n]: - lambdas[lname][mask] += amp.transpose(transpose).copy() * sign - done.add(combn) - - for name, key, n in ucc.ansatz.bosonic_cluster_ranks(spin_type=ucc.spin_type): - lname = "l" + name - lambdas[lname] = ucc.lambdas[lname].copy() - - for name, key, nf, nb in ucc.ansatz.coupling_cluster_ranks(spin_type=ucc.spin_type): - lname = "l" + name - shape = (nbos,) * nb + tuple( - space.size(k) for k in key[nb + nf :] + key[nb : nb + nf] - ) - lambdas[lname] = np.zeros(shape, dtype=types[float]) - for comb in util.generate_spin_combinations(nf, unique=True): - done = set() - for lperm, lsign in util.permutations_with_signs(tuple(range(nf))): - for uperm, usign in util.permutations_with_signs(tuple(range(nf))): - combn = util.permute_string(comb[:nf], lperm) - combn += util.permute_string(comb[nf:], uperm) - if combn in done: - continue - mask = np.ix_( - *([range(nbos)] * nb), - *[ - slices[s][k] - for s, k in zip(combn, key[nb + nf :] + key[nb : nb + nf]) - ], - ) - transpose = ( - tuple(range(nb)) - + tuple(p + nb for p in lperm) - + tuple(p + nb + nf for p in uperm) - ) - amp = ( - getattr(ucc.lambdas[lname], comb).transpose(transpose) - * lsign - * usign - ) - for perm, sign in util.permutations_with_signs(tuple(range(nf))): - transpose = ( - tuple(range(nb)) - + tuple(p + nb for p in perm) - + tuple(range(nb + nf, nb + 2 * nf)) - ) - if util.permute_string(comb[:nf], perm) == comb[:nf]: - lambdas[lname][mask] += amp.transpose(transpose).copy() * sign - done.add(combn) - - gcc.lambdas = lambdas - - return gcc - - @classmethod - def from_rebcc(cls, rcc): - """ - Initialise a `GEBCC` object from an `REBCC` object. - - Parameters - ---------- - rcc : REBCC - The REBCC object to initialise from. - - Returns - ------- - gcc : GEBCC - The GEBCC object. - """ - - ucc = uebcc.UEBCC.from_rebcc(rcc) - gcc = cls.from_uebcc(ucc) - - return gcc - - def init_amps(self, eris=None): - eris = self.get_eris(eris) - amplitudes = util.Namespace() - - # Build T amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - if n == 1: - amplitudes[name] = getattr(self.fock, key) / self.energy_sum(key) - elif n == 2: - amplitudes[name] = getattr(eris, key) / self.energy_sum(key) - else: - shape = tuple(self.space.size(k) for k in key) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - if self.boson_ansatz: - # Only true for real-valued couplings: - h = self.g - H = self.G - - # Build S amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - if n == 1: - amplitudes[name] = -H / self.omega - else: - shape = (self.nbos,) * n - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - # Build U amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - if n == 1: - amplitudes[name] = h[key] / self.energy_sum(key) - else: - shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - return amplitudes - - def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): - func, kwargs = self._load_function( - "make_rdm2_f", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm = func(**kwargs) - - if hermitise: - dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) - dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(1, 0, 3, 2)) - - return dm - - def excitations_to_vector_ip(self, *excitations): - vectors = [] - m = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[:-1] - vectors.append(util.compress_axes(key, excitations[m]).ravel()) - m += 1 - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return np.concatenate(vectors) - - def excitations_to_vector_ee(self, *excitations): - vectors = [] - m = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - vectors.append(util.compress_axes(key, excitations[m]).ravel()) - m += 1 - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return np.concatenate(vectors) - - def vector_to_excitations_ip(self, vector): - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[:-1] - size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) - shape = tuple(self.space.size(k) for k in key) - vn_tril = vector[i0 : i0 + size] - vn = util.decompress_axes(key, vn_tril, shape=shape) - excitations.append(vn) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def vector_to_excitations_ea(self, vector): - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[n:] + key[: n - 1] - size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) - shape = tuple(self.space.size(k) for k in key) - vn_tril = vector[i0 : i0 + size] - vn = util.decompress_axes(key, vn_tril, shape=shape) - excitations.append(vn) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def vector_to_excitations_ee(self, vector): - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) - shape = tuple(self.space.size(k) for k in key) - vn_tril = vector[i0 : i0 + size] - vn = util.decompress_axes(key, vn_tril, shape=shape) - excitations.append(vn) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def get_mean_field_G(self): - val = lib.einsum("Ipp->I", self.g.boo) - val -= self.xi * self.omega - - if self.bare_G is not None: - val += self.bare_G - - return val - - def get_eris(self, eris=None): - """ - Get blocks of the ERIs. - - Parameters - ---------- - eris : np.ndarray or ERIs, optional. - Electronic repulsion integrals, either in the form of a - dense array or an ERIs object. Default value is `None`. - - Returns - ------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.ERIs()`. - """ - if (eris is None) or isinstance(eris, np.ndarray): - return self.ERIs(self, array=eris) - else: - return eris - - def ip_eom(self, options=None, **kwargs): - return geom.IP_GEOM(self, options=options, **kwargs) - - def ea_eom(self, options=None, **kwargs): - return geom.EA_GEOM(self, options=options, **kwargs) - - def ee_eom(self, options=None, **kwargs): - return geom.EE_GEOM(self, options=options, **kwargs) - - @property - def xi(self): - if self.options.shift: - xi = lib.einsum("Iii->I", self.g.boo) - xi /= self.omega - if self.bare_G is not None: - xi += self.bare_G / self.omega - else: - xi = np.zeros_like(self.omega) - return xi - - @property - def spin_type(self): - return "G" diff --git a/ebcc/logging.py b/ebcc/logging.py index b3561132..415b3ebd 100644 --- a/ebcc/logging.py +++ b/ebcc/logging.py @@ -17,20 +17,28 @@ %s""" # noqa: W605 -def output(self, msg, *args, **kwargs): - """Output a message at the `"OUTPUT"` level.""" - if self.isEnabledFor(25): - self._log(25, msg, args, **kwargs) +class Logger(logging.Logger): + """Logger with a custom output level.""" + + def __init__(self, name, level=logging.INFO): + super().__init__(name, level) + + def output(self, msg, *args, **kwargs): + """Output a message at the `"OUTPUT"` level.""" + if self.isEnabledFor(25): + self._log(25, msg, args, **kwargs) + + +logging.setLoggerClass(Logger) +logging.addLevelName(25, "OUTPUT") default_log = logging.getLogger(__name__) default_log.setLevel(logging.INFO) default_log.addHandler(logging.StreamHandler(sys.stderr)) -logging.addLevelName(25, "OUTPUT") -logging.Logger.output = output -class NullLogger(logging.Logger): +class NullLogger(Logger): """A logger that does nothing.""" def __init__(self, *args, **kwargs): diff --git a/ebcc/precision.py b/ebcc/precision.py index 7b98aed2..3646ca62 100644 --- a/ebcc/precision.py +++ b/ebcc/precision.py @@ -1,39 +1,63 @@ """Floating point precision control.""" +from __future__ import annotations + from contextlib import contextmanager +from typing import TYPE_CHECKING from ebcc import numpy as np -types = { - float: np.float64, - complex: np.complex128, -} +if TYPE_CHECKING: + from collections.abc import Iterator + from typing import Literal, Type, TypeVar, Union + + T = TypeVar("T", float, complex) + + +types: dict[type, type] +if TYPE_CHECKING: + types = { + float: float, + complex: complex, + } +else: + types = { + float: np.float64, + complex: np.complex128, + } + +def cast(value: Any, dtype: Type[T]) -> T: + """Cast a value to the current floating point type. -def set_precision(**kwargs): + Args: + value: The value to cast. + dtype: The type to cast to. + + Returns: + The value cast to the current floating point type. + """ + return types[dtype](value) + + +def set_precision(**kwargs: type) -> None: """Set the floating point type. - Parameters - ---------- - float : type, optional - The floating point type to use. - complex : type, optional - The complex type to use. + Args: + float: The floating point type to use. + complex: The complex type to use. """ types[float] = kwargs.get("float", types[float]) types[complex] = kwargs.get("complex", types[complex]) @contextmanager -def precision(**kwargs): +def precision(**kwargs: type) -> Iterator[None]: """Context manager for setting the floating point precision. - Parameters - ---------- - float : type, optional - The floating point type to use. - complex : type, optional - The complex type to use. + Args: + float: The floating point type to use. + complex: The complex type to use. """ old = { "float": types[float], @@ -45,10 +69,7 @@ def precision(**kwargs): @contextmanager -def single_precision(): - """ - Context manager for setting the floating point precision to single - precision. - """ +def single_precision() -> Iterator[None]: + """Context manager for setting the floating point precision to single precision.""" with precision(float=np.float32, complex=np.complex64): yield diff --git a/ebcc/rebcc.py b/ebcc/rebcc.py deleted file mode 100644 index f9bf76ce..00000000 --- a/ebcc/rebcc.py +++ /dev/null @@ -1,1834 +0,0 @@ -"""Restricted electron-boson coupled cluster.""" - -import dataclasses - -from pyscf import lib - -from ebcc import default_log, init_logging -from ebcc import numpy as np -from ebcc import reom, util -from ebcc.ansatz import Ansatz -from ebcc.brueckner import BruecknerREBCC -from ebcc.cderis import RCDERIs -from ebcc.damping import DIIS -from ebcc.dump import Dump -from ebcc.eris import RERIs -from ebcc.fock import RFock -from ebcc.logging import ANSI -from ebcc.precision import types -from ebcc.space import Space -from ebcc.cc.base import EBCC - - -class REBCC(EBCC): - """Restricted electron-boson coupled cluster. - - Attributes: - mf: PySCF mean-field object. - log: Log to write output to. - options: Options for the EBCC calculation. - e_corr: Correlation energy. - amplitudes: Cluster amplitudes. - converged: Convergence flag. - lambdas: Cluster lambda amplitudes. - converged_lambda: Lambda convergence flag. - name: Name of the method. - """ - - # Types - ERIs = RERIs - Fock = RFock - CDERIs = RCDERIs - Brueckner = BruecknerREBCC - - def kernel(self, eris=None): - """ - Run the coupled cluster calculation. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - - Returns - ------- - e_cc : float - Correlation energy. - """ - - # Start a timer: - timer = util.Timer() - - # Get the ERIs: - eris = self.get_eris(eris) - - # Get the amplitude guesses: - if self.amplitudes is None: - amplitudes = self.init_amps(eris=eris) - else: - amplitudes = self.amplitudes - - # Get the initial energy: - e_cc = self.energy(amplitudes=amplitudes, eris=eris) - - self.log.output("Solving for excitation amplitudes.") - self.log.debug("") - self.log.info( - f"{ANSI.B}{'Iter':>4s} {'Energy (corr.)':>16s} {'Energy (tot.)':>18s} " - f"{'Δ(Energy)':>13s} {'Δ(Ampl.)':>13s}{ANSI.R}" - ) - self.log.info(f"{0:4d} {e_cc:16.10f} {e_cc + self.mf.e_tot:18.10f}") - - if not self.ansatz.is_one_shot: - # Set up DIIS: - diis = DIIS() - diis.space = self.options.diis_space - diis.damping = self.options.damping - - converged = False - for niter in range(1, self.options.max_iter + 1): - # Update the amplitudes, extrapolate with DIIS and - # calculate change: - amplitudes_prev = amplitudes - amplitudes = self.update_amps(amplitudes=amplitudes, eris=eris) - vector = self.amplitudes_to_vector(amplitudes) - vector = diis.update(vector) - amplitudes = self.vector_to_amplitudes(vector) - dt = np.linalg.norm(vector - self.amplitudes_to_vector(amplitudes_prev), ord=np.inf) - - # Update the energy and calculate change: - e_prev = e_cc - e_cc = self.energy(amplitudes=amplitudes, eris=eris) - de = abs(e_prev - e_cc) - - # Log the iteration: - converged_e = de < self.options.e_tol - converged_t = dt < self.options.t_tol - self.log.info( - f"{niter:4d} {e_cc:16.10f} {e_cc + self.mf.e_tot:18.10f}" - f" {[ANSI.r, ANSI.g][converged_e]}{de:13.3e}{ANSI.R}" - f" {[ANSI.r, ANSI.g][converged_t]}{dt:13.3e}{ANSI.R}" - ) - - # Check for convergence: - converged = converged_e and converged_t - if converged: - self.log.debug("") - self.log.output(f"{ANSI.g}Converged.{ANSI.R}") - break - else: - self.log.debug("") - self.log.warning(f"{ANSI.r}Failed to converge.{ANSI.R}") - - # Include perturbative correction if required: - if self.ansatz.has_perturbative_correction: - self.log.debug("") - self.log.info("Computing perturbative energy correction.") - e_pert = self.energy_perturbative(amplitudes=amplitudes, eris=eris) - e_cc += e_pert - self.log.info(f"E(pert) = {e_pert:.10f}") - - else: - converged = True - - # Update attributes: - self.e_corr = e_cc - self.amplitudes = amplitudes - self.converged = converged - - self.log.debug("") - self.log.output(f"E(corr) = {self.e_corr:.10f}") - self.log.output(f"E(tot) = {self.e_tot:.10f}") - self.log.debug("") - self.log.debug("Time elapsed: %s", timer.format_time(timer())) - self.log.debug("") - - return e_cc - - def solve_lambda(self, amplitudes=None, eris=None): - """ - Solve the lambda coupled cluster equations. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - """ - - # Start a timer: - timer = util.Timer() - - # Get the ERIs: - eris = self.get_eris(eris) - - # Get the amplitudes: - if amplitudes is None: - amplitudes = self.amplitudes - if amplitudes is None: - amplitudes = self.init_amps(eris=eris) - - # If needed, precompute the perturbative part of the lambda - # amplitudes: - if self.ansatz.has_perturbative_correction: - lambdas_pert = self.update_lams(eris=eris, amplitudes=amplitudes, perturbative=True) - else: - lambdas_pert = None - - # Get the lambda amplitude guesses: - if self.lambdas is None: - lambdas = self.init_lams(amplitudes=amplitudes) - else: - lambdas = self.lambdas - - # Set up DIIS: - diis = DIIS() - diis.space = self.options.diis_space - diis.damping = self.options.damping - - self.log.output("Solving for de-excitation (lambda) amplitudes.") - self.log.debug("") - self.log.info(f"{ANSI.B}{'Iter':>4s} {'Δ(Ampl.)':>13s}{ANSI.R}") - - converged = False - for niter in range(1, self.options.max_iter + 1): - # Update the lambda amplitudes, extrapolate with DIIS and - # calculate change: - lambdas_prev = lambdas - lambdas = self.update_lams( - amplitudes=amplitudes, - lambdas=lambdas, - lambdas_pert=lambdas_pert, - eris=eris, - ) - vector = self.lambdas_to_vector(lambdas) - vector = diis.update(vector) - lambdas = self.vector_to_lambdas(vector) - dl = np.linalg.norm(vector - self.lambdas_to_vector(lambdas_prev), ord=np.inf) - - # Log the iteration: - converged = dl < self.options.t_tol - self.log.info(f"{niter:4d} {[ANSI.r, ANSI.g][converged]}{dl:13.3e}{ANSI.R}") - - # Check for convergence: - if converged: - self.log.debug("") - self.log.output(f"{ANSI.g}Converged.{ANSI.R}") - break - else: - self.log.debug("") - self.log.warning(f"{ANSI.r}Failed to converge.{ANSI.R}") - - self.log.debug("") - self.log.debug("Time elapsed: %s", timer.format_time(timer())) - self.log.debug("") - self.log.debug("") - - # Update attributes: - self.lambdas = lambdas - self.converged_lambda = converged - - def brueckner(self, *args, **kwargs): - """ - Run a Brueckner orbital coupled cluster calculation. - - Returns - ------- - e_cc : float - Correlation energy. - """ - - bcc = self.Brueckner(self, *args, **kwargs) - - return bcc.kernel() - - def write(self, file): - """ - Write the EBCC data to a file. - - Parameters - ---------- - file : str - Path of file to write to. - """ - - writer = Dump(file) - writer.write(self) - - @classmethod - def read(cls, file, log=None): - """ - Read the data from a file. - - Parameters - ---------- - file : str - Path of file to read from. - log : Logger, optional - Logger to assign to the EBCC object. - - Returns - ------- - ebcc : EBCC - The EBCC object loaded from the file. - """ - - reader = Dump(file) - cc = reader.read(cls=cls, log=log) - - return cc - - @staticmethod - def _convert_mf(mf): - """ - Convert the input PySCF mean-field object to the one required for - the current class. - """ - return mf.to_rhf() - - def _load_function(self, name, eris=False, amplitudes=False, lambdas=False, **kwargs): - """ - Load a function from the generated code, and return a dict of - arguments. - """ - - if not (eris is False): - eris = self.get_eris(eris) - else: - eris = None - - dicts = [] - - if not (amplitudes is False): - if amplitudes is None: - amplitudes = self.amplitudes - if amplitudes is None: - amplitudes = self.init_amps(eris=eris) - dicts.append(amplitudes) - - if not (lambdas is False): - if lambdas is None: - lambdas = self.lambdas - if lambdas is None: - self.log.warning("Using Λ = T* for %s", name) - lambdas = self.init_lams(amplitudes=amplitudes) - dicts.append(lambdas) - - if kwargs: - dicts.append(kwargs) - - func = getattr(self._eqns, name, None) - - if func is None: - raise util.ModelNotImplemented("%s for rank = %s" % (name, self.name)) - - kwargs = self._pack_codegen_kwargs(*dicts, eris=eris) - - return func, kwargs - - def _pack_codegen_kwargs(self, *extra_kwargs, eris=None): - """ - Pack all the possible keyword arguments for generated code - into a dictionary. - """ - # TODO change all APIs to take the space object instead of - # nocc, nvir, nbos, etc. - - eris = self.get_eris(eris) - - omega = np.diag(self.omega) if self.omega is not None else None - - kwargs = dict( - f=self.fock, - v=eris, - g=self.g, - G=self.G, - w=omega, - space=self.space, - nocc=self.space.ncocc, # FIXME rename? - nvir=self.space.ncvir, # FIXME rename? - nbos=self.nbos, - ) - if isinstance(eris, self.CDERIs): - kwargs["naux"] = self.mf.with_df.get_naoaux() - for kw in extra_kwargs: - if kw is not None: - kwargs.update(kw) - - return kwargs - - def init_space(self): - """ - Initialise the default `Space` object. - - Returns - ------- - space : Space - Space object in which all fermionic degrees of freedom are - considered inactive. - """ - - space = Space( - self.mo_occ > 0, - np.zeros_like(self.mo_occ, dtype=bool), - np.zeros_like(self.mo_occ, dtype=bool), - ) - - return space - - def init_amps(self, eris=None): - """ - Initialise the amplitudes. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - - Returns - ------- - amplitudes : Namespace - Cluster amplitudes. - """ - - eris = self.get_eris(eris) - amplitudes = util.Namespace() - - # Build T amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - if n == 1: - amplitudes[name] = self.fock[key] / self.energy_sum(key) - elif n == 2: - key_t = key[0] + key[2] + key[1] + key[3] - amplitudes[name] = eris[key_t].swapaxes(1, 2) / self.energy_sum(key) - else: - shape = tuple(self.space.size(k) for k in key) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - if self.boson_ansatz: - # Only true for real-valued couplings: - h = self.g - H = self.G - - # Build S amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - if n == 1: - amplitudes[name] = -H / self.omega - else: - shape = (self.nbos,) * n - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - # Build U amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - if nb == 1: - amplitudes[name] = h[key] / self.energy_sum(key) - else: - shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - return amplitudes - - def init_lams(self, amplitudes=None): - """ - Initialise the lambda amplitudes. - - Parameters - ---------- - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - lambdas : Namespace - Cluster lambda amplitudes. - """ - - if amplitudes is None: - amplitudes = self.amplitudes - lambdas = util.Namespace() - - # Build L amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - lname = name.replace("t", "l") - perm = list(range(n, 2 * n)) + list(range(n)) - lambdas[lname] = amplitudes[name].transpose(perm) - - # Build LS amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - lname = "l" + name - lambdas[lname] = amplitudes[name] - - # Build LU amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - lname = "l" + name - perm = list(range(nb)) + [nb + 1, nb] - lambdas[lname] = amplitudes[name].transpose(perm) - - return lambdas - - def energy(self, eris=None, amplitudes=None): - """ - Compute the correlation energy. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - e_cc : float - Correlation energy. - """ - - func, kwargs = self._load_function( - "energy", - eris=eris, - amplitudes=amplitudes, - ) - - return types[float](func(**kwargs).real) - - def energy_perturbative(self, eris=None, amplitudes=None, lambdas=None): - """ - Compute the perturbative part to the correlation energy. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - e_pert : float - Perturbative correction to the correlation energy. - """ - - func, kwargs = self._load_function( - "energy_perturbative", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return types[float](func(**kwargs).real) - - def update_amps(self, eris=None, amplitudes=None): - """ - Update the amplitudes. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - amplitudes : Namespace - Updated cluster amplitudes. - """ - - func, kwargs = self._load_function( - "update_amps", - eris=eris, - amplitudes=amplitudes, - ) - res = func(**kwargs) - res = {key.rstrip("new"): val for key, val in res.items()} - - # Divide T amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - res[name] /= self.energy_sum(key) - res[name] += amplitudes[name] - - # Divide S amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - res[name] /= self.energy_sum(key) - res[name] += amplitudes[name] - - # Divide U amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - res[name] /= self.energy_sum(key) - res[name] += amplitudes[name] - - return res - - def update_lams( - self, - eris=None, - amplitudes=None, - lambdas=None, - lambdas_pert=None, - perturbative=False, - ): - """ - Update the lambda amplitudes. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - perturbative : bool, optional - Whether to compute the perturbative part of the lambda - amplitudes. Default value is `False`. - - Returns - ------- - lambdas : Namespace - Updated cluster lambda amplitudes. - """ - # TODO active - - if lambdas_pert is not None: - lambdas.update(lambdas_pert) - - func, kwargs = self._load_function( - "update_lams%s" % ("_perturbative" if perturbative else ""), - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - res = func(**kwargs) - res = {key.rstrip("new"): val for key, val in res.items()} - - # Divide T amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - lname = name.replace("t", "l") - res[lname] /= self.energy_sum(key[n:] + key[:n]) - if not perturbative: - res[lname] += lambdas[lname] - - # Divide S amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - lname = "l" + name - res[lname] /= self.energy_sum(key[n:] + key[:n]) - if not perturbative: - res[lname] += lambdas[lname] - - # Divide U amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - lname = "l" + name - res[lname] /= self.energy_sum(key[:nb] + key[nb + nf :] + key[nb : nb + nf]) - if not perturbative: - res[lname] += lambdas[lname] - - if perturbative: - res = {key + "pert": val for key, val in res.items()} - - return res - - def make_sing_b_dm(self, eris=None, amplitudes=None, lambdas=None): - r""" - Build the single boson density matrix: - - ..math :: \langle b^+ \rangle - - and - - ..math :: \langle b \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - dm_b : numpy.ndarray (nbos,) - Single boson density matrix. - """ - - func, kwargs = self._load_function( - "make_sing_b_dm", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_rdm1_b(self, eris=None, amplitudes=None, lambdas=None, unshifted=True, hermitise=True): - r""" - Build the bosonic one-particle reduced density matrix: - - ..math :: \langle b^+ b \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - unshifted : bool, optional - If `self.shift` is `True`, then `unshifted=True` applies the - reverse transformation such that the bosonic operators are - defined with respect to the unshifted bosons. Default value is - `True`. Has no effect if `self.shift` is `False`. - hermitise : bool, optional - Force Hermiticity in the output. Default value is `True`. - - Returns - ------- - rdm1_b : numpy.ndarray (nbos, nbos) - Bosonic one-particle reduced density matrix. - """ - - func, kwargs = self._load_function( - "make_rdm1_b", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm = func(**kwargs) - - if hermitise: - dm = 0.5 * (dm + dm.T) - - if unshifted and self.options.shift: - dm_cre, dm_ann = self.make_sing_b_dm() - xi = self.xi - dm[np.diag_indices_from(dm)] -= xi * (dm_cre + dm_ann) - xi**2 - - return dm - - def make_rdm1_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): - r""" - Build the fermionic one-particle reduced density matrix: - - ..math :: \langle i^+ j \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - hermitise : bool, optional - Force Hermiticity in the output. Default value is `True`. - - Returns - ------- - rdm1_f : numpy.ndarray (nmo, nmo) - Fermionic one-particle reduced density matrix. - """ - - func, kwargs = self._load_function( - "make_rdm1_f", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm = func(**kwargs) - - if hermitise: - dm = 0.5 * (dm + dm.T) - - return dm - - def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): - r""" - Build the fermionic two-particle reduced density matrix: - - ..math :: \Gamma_{ijkl} = \langle i^+ j^+ l k \rangle - - which is stored in Chemist's notation. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - hermitise : bool, optional - Force Hermiticity in the output. Default value is `True`. - - Returns - ------- - rdm2_f : numpy.ndarray (nmo, nmo, nmo, nmo) - Fermionic two-particle reduced density matrix. - """ - - func, kwargs = self._load_function( - "make_rdm2_f", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm = func(**kwargs) - - if hermitise: - dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) - dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(1, 0, 3, 2)) - - return dm - - def make_eb_coup_rdm( - self, - eris=None, - amplitudes=None, - lambdas=None, - unshifted=True, - hermitise=True, - ): - r""" - Build the electron-boson coupling reduced density matrices: - - ..math :: \langle b^+ i^+ j \rangle - - and - - ..math :: \langle b i^+ j \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - unshifted : bool, optional - If `self.shift` is `True`, then `unshifted=True` applies the - reverse transformation such that the bosonic operators are - defined with respect to the unshifted bosons. Default value is - `True`. Has no effect if `self.shift` is `False`. - hermitise : bool, optional - Force Hermiticity in the output. Default value is `True`. - - Returns - ------- - dm_eb : numpy.ndarray (2, nbos, nmo, nmo) - Electron-boson coupling reduce density matrices. First - index corresponds to creation and second to annihilation - of the bosonic index. - """ - - func, kwargs = self._load_function( - "make_eb_coup_rdm", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm_eb = func(**kwargs) - - if hermitise: - dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) - dm_eb[1] = dm_eb[0].transpose(0, 2, 1).copy() - - if unshifted and self.options.shift: - rdm1_f = self.make_rdm1_f(hermitise=hermitise) - shift = util.einsum("x,ij->xij", self.xi, rdm1_f) - dm_eb -= shift[None] - - return dm_eb - - def hbar_matvec_ip(self, r1, r2, eris=None, amplitudes=None): - """ - Compute the product between a state vector and the EOM Hamiltonian - for the IP. - - Parameters - ---------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector. Keys are - strings of the name of each vector, and values are arrays - whose dimension depends on the particular sector. - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector resulting from - the matrix-vector product with the input vectors. Keys are - strings of the name of each vector, and values are arrays whose - dimension depends on the particular sector. - """ - # TODO generalise vectors input - - func, kwargs = self._load_function( - "hbar_matvec_ip", - eris=eris, - amplitudes=amplitudes, - r1=r1, - r2=r2, - ) - - return func(**kwargs) - - def hbar_matvec_ea(self, r1, r2, eris=None, amplitudes=None): - """ - Compute the product between a state vector and the EOM Hamiltonian - for the EA. - - Parameters - ---------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector. Keys are - strings of the name of each vector, and values are arrays - whose dimension depends on the particular sector. - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector resulting from - the matrix-vector product with the input vectors. Keys are - strings of the name of each vector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "hbar_matvec_ea", - eris=eris, - amplitudes=amplitudes, - r1=r1, - r2=r2, - ) - - return func(**kwargs) - - def hbar_matvec_ee(self, r1, r2, eris=None, amplitudes=None): - """ - Compute the product between a state vector and the EOM Hamiltonian - for the EE. - - Parameters - ---------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector. Keys are - strings of the name of each vector, and values are arrays - whose dimension depends on the particular sector. - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector resulting from - the matrix-vector product with the input vectors. Keys are - strings of the name of each vector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "hbar_matvec_ee", - eris=eris, - amplitudes=amplitudes, - r1=r1, - r2=r2, - ) - - return func(**kwargs) - - def make_ip_mom_bras(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the bra IP vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - bras : dict of (str, numpy.ndarray) - Dictionary containing the bra vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ip_mom_bras", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ea_mom_bras(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the bra EA vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - bras : dict of (str, numpy.ndarray) - Dictionary containing the bra vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ea_mom_bras", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ee_mom_bras(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the bra EE vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - bras : dict of (str, numpy.ndarray) - Dictionary containing the bra vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ee_mom_bras", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ip_mom_kets(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the ket IP vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - kets : dict of (str, numpy.ndarray) - Dictionary containing the ket vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ip_mom_kets", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ea_mom_kets(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the ket IP vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - kets : dict of (str, numpy.ndarray) - Dictionary containing the ket vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ea_mom_kets", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ee_mom_kets(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the ket EE vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - kets : dict of (str, numpy.ndarray) - Dictionary containing the ket vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ee_mom_kets", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ip_1mom(self, eris=None, amplitudes=None, lambdas=None): - r""" - Build the first fermionic hole single-particle moment. - - .. math:: T_{pq} = \langle c_p^+ (H - E_0) c_q \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - mom : numpy.ndarray (nmo, nmo) - Array of the first moment. - """ - - raise util.ModelNotImplemented # TODO - - def make_ea_1mom(self, eris=None, amplitudes=None, lambdas=None): - r""" - Build the first fermionic particle single-particle moment. - - .. math:: T_{pq} = \langle c_p (H - E_0) c_q^+ \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - mom : numpy.ndarray (nmo, nmo) - Array of the first moment. - """ - - raise util.ModelNotImplemented # TODO - - def ip_eom(self, options=None, **kwargs): - """Get the IP EOM object.""" - return reom.IP_REOM(self, options=options, **kwargs) - - def ea_eom(self, options=None, **kwargs): - """Get the EA EOM object.""" - return reom.EA_REOM(self, options=options, **kwargs) - - def ee_eom(self, options=None, **kwargs): - """Get the EE EOM object.""" - return reom.EE_REOM(self, options=options, **kwargs) - - def amplitudes_to_vector(self, amplitudes): - """ - Construct a vector containing all of the amplitudes used in the - given ansatz. - - Parameters - ---------- - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the amplitudes flattened and - concatenated. Size depends on the ansatz. - """ - - vectors = [] - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - vectors.append(amplitudes[name].ravel()) - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - vectors.append(amplitudes[name].ravel()) - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - vectors.append(amplitudes[name].ravel()) - - return np.concatenate(vectors) - - def vector_to_amplitudes(self, vector): - """ - Construct all of the amplitudes used in the given ansatz from a - vector. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the amplitudes flattened - and concatenated. Size depends on the ansatz. - - Returns - ------- - amplitudes : Namespace - Cluster amplitudes. - """ - - amplitudes = util.Namespace() - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - amplitudes[name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - shape = (self.nbos,) * n - size = np.prod(shape) - amplitudes[name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) - size = np.prod(shape) - amplitudes[name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - return amplitudes - - def lambdas_to_vector(self, lambdas): - """ - Construct a vector containing all of the lambda amplitudes used in - the given ansatz. - - Parameters - ---------- - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the lambdas flattened and - concatenated. Size depends on the ansatz. - """ - - vectors = [] - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - vectors.append(lambdas[name.replace("t", "l")].ravel()) - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - vectors.append(lambdas["l" + name].ravel()) - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - vectors.append(lambdas["l" + name].ravel()) - - return np.concatenate(vectors) - - def vector_to_lambdas(self, vector): - """ - Construct all of the lambdas used in the given ansatz from a - vector. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the lambdas flattened and - concatenated. Size depends on the ansatz. - - Returns - ------- - lambdas : Namespace - Cluster lambda amplitudes. - """ - - lambdas = util.Namespace() - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - lname = name.replace("t", "l") - key = key[n:] + key[:n] - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - lambdas[lname] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - shape = (self.nbos,) * n - size = np.prod(shape) - lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - key = key[:nb] + key[nb + nf :] + key[nb : nb + nf] - shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) - size = np.prod(shape) - lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - return lambdas - - def excitations_to_vector_ip(self, *excitations): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the IP. - - Parameters - ---------- - *excitations : iterable of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - """ - - vectors = [] - m = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - vectors.append(excitations[m].ravel()) - m += 1 - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return np.concatenate(vectors) - - def excitations_to_vector_ea(self, *excitations): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EA. - - Parameters - ---------- - *excitations : iterable of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - """ - return self.excitations_to_vector_ip(*excitations) - - def excitations_to_vector_ee(self, *excitations): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EE. - - Parameters - ---------- - *excitations : iterable of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - """ - return self.excitations_to_vector_ip(*excitations) - - def vector_to_excitations_ip(self, vector): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the IP. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - - Returns - ------- - excitations : tuple of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - """ - - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[:-1] - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - excitations.append(vector[i0 : i0 + size].reshape(shape)) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def vector_to_excitations_ea(self, vector): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EA. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - - Returns - ------- - excitations : tuple of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - """ - - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[n:] + key[: n - 1] - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - excitations.append(vector[i0 : i0 + size].reshape(shape)) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def vector_to_excitations_ee(self, vector): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EE. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - - Returns - ------- - excitations : tuple of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - """ - - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - excitations.append(vector[i0 : i0 + size].reshape(shape)) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def get_mean_field_G(self): - """ - Get the mean-field boson non-conserving term of the Hamiltonian. - - Returns - ------- - G_mf : numpy.ndarray (nbos,) - Mean-field boson non-conserving term of the Hamiltonian. - """ - - # FIXME should this also sum in frozen orbitals? - val = lib.einsum("Ipp->I", self.g.boo) * 2.0 - val -= self.xi * self.omega - - if self.bare_G is not None: - val += self.bare_G - - return val - - def get_g(self, g): - """ - Get blocks of the electron-boson coupling matrix corresponding to - the bosonic annihilation operator. - - Parameters - ---------- - g : numpy.ndarray (nbos, nmo, nmo) - Array of the electron-boson coupling matrix. - - Returns - ------- - g : Namespace - Namespace containing blocks of the electron-boson coupling - matrix. Each attribute should be a length-3 string of - `b`, `o` or `v` signifying whether the corresponding axis - is bosonic, occupied or virtual. - """ - - slices = { - "x": self.space.correlated, - "o": self.space.correlated_occupied, - "v": self.space.correlated_virtual, - "O": self.space.active_occupied, - "V": self.space.active_virtual, - "i": self.space.inactive_occupied, - "a": self.space.inactive_virtual, - } - - class Blocks(util.Namespace): - def __getitem__(selffer, key): - assert key[0] == "b" - i = slices[key[1]] - j = slices[key[2]] - return g[:, i][:, :, j].copy() - - __getattr__ = __getitem__ - - return Blocks() - - def get_fock(self): - """ - Get blocks of the Fock matrix, shifted due to bosons where the - ansatz requires. - - Returns - ------- - fock : Namespace - Namespace containing blocks of the Fock matrix. Each attribute - should be a length-2 string of `o` or `v` signifying whether - the corresponding axis is occupied or virtual. - """ - return self.Fock(self, array=self.bare_fock) - - def get_eris(self, eris=None): - """Get blocks of the ERIs. - - Parameters - ---------- - eris : np.ndarray or ERIs, optional. - Electronic repulsion integrals, either in the form of a dense - array or an `ERIs` object. Default value is `None`. - - Returns - ------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.ERIs()`. - """ - if (eris is None) or isinstance(eris, np.ndarray): - if (isinstance(eris, np.ndarray) and eris.ndim == 3) or getattr( - self.mf, "with_df", None - ): - return self.CDERIs(self, array=eris) - else: - return self.ERIs(self, array=eris) - else: - return eris - - @property - def bare_fock(self): - """ - Get the mean-field Fock matrix in the MO basis, including frozen - parts. - - Returns - ------- - bare_fock : numpy.ndarray (nmo, nmo) - The mean-field Fock matrix in the MO basis. - """ - - fock_ao = self.mf.get_fock().astype(types[float]) - mo_coeff = self.mo_coeff - - fock = util.einsum("pq,pi,qj->ij", fock_ao, mo_coeff, mo_coeff) - - return fock - - @property - def xi(self): - """ - Get the shift in bosonic operators to diagonalise the photonic - Hamiltonian. - - Returns - ------- - xi : numpy.ndarray (nbos,) - Shift in bosonic operators to diagonalise the phononic - Hamiltonian. - """ - - if self.options.shift: - xi = lib.einsum("Iii->I", self.g.boo) * 2.0 - xi /= self.omega - if self.bare_G is not None: - xi += self.bare_G / self.omega - else: - xi = np.zeros_like(self.omega, dtype=types[float]) - - return xi - - @property - def const(self): - """ - Get the shift in the energy from moving to polaritonic basis. - - Returns - ------- - const : float - Shift in the energy from moving to polaritonic basis. - """ - if self.options.shift: - return lib.einsum("I,I->", self.omega, self.xi**2) - else: - return 0.0 - - @property - def name(self): - """Get the name of the method.""" - return self.spin_type + self.ansatz.name - - @property - def spin_type(self): - """Get a string representation of the spin type.""" - return "R" - - @property - def nmo(self): - """Get the number of MOs.""" - return self.space.nmo - - @property - def nocc(self): - """Get the number of occupied MOs.""" - return self.space.nocc - - @property - def nvir(self): - """Get the number of virtual MOs.""" - return self.space.nvir - - @property - def nbos(self): - """ - Get the number of bosonic degrees of freedom. - - Returns - ------- - nbos : int - Number of bosonic degrees of freedom. - """ - if self.omega is None: - return 0 - return self.omega.shape[0] - - def energy_sum(self, subscript, signs_dict=None): - """ - Get a direct sum of energies. - - Parameters - ---------- - subscript : str - The direct sum subscript, where each character indicates the - sector for each energy. For the default slice characters, see - `Space`. Occupied degrees of freedom are assumed to be - positive, virtual and bosonic negative (the signs can be - changed via the `signs_dict` keyword argument). - signs_dict : dict, optional - Dictionary defining custom signs for each sector. If `None`, - initialised such that `["o", "O", "i"]` are positive, and - `["v", "V", "a", "b"]` negative. Default value is `None`. - - Returns - ------- - energy_sum : numpy.ndarray - Array of energy sums. - """ - - n = 0 - - def next_char(): - nonlocal n - if n < 26: - char = chr(ord("a") + n) - else: - char = chr(ord("A") + n) - n += 1 - return char - - if signs_dict is None: - signs_dict = {} - for k, s in zip("vVaoOib", "---+++-"): - if k not in signs_dict: - signs_dict[k] = s - - energies = [] - for key in subscript: - if key == "b": - energies.append(self.omega) - else: - energies.append(np.diag(self.fock[key + key])) - - subscript = "".join([signs_dict[k] + next_char() for k in subscript]) - energy_sum = lib.direct_sum(subscript, *energies) - - return energy_sum diff --git a/ebcc/util/einsumfunc.py b/ebcc/util/einsumfunc.py index dde1d2d7..527ae211 100644 --- a/ebcc/util/einsumfunc.py +++ b/ebcc/util/einsumfunc.py @@ -11,7 +11,7 @@ from ebcc import numpy as np if TYPE_CHECKING: - from typing import Any, TypeVar + from typing import Any, TypeVar, Union from ebcc.numpy.typing import NDArray # type: ignore @@ -44,9 +44,10 @@ def _fallback_einsum(*operands: Any, **kwargs: Any) -> NDArray[T]: alpha = kwargs.pop("alpha", 1.0) beta = kwargs.pop("beta", 0.0) out = kwargs.pop("out", None) + kwargs["optimize"] = True # Perform the contraction - res = np.einsum(*operands, **kwargs, optimize=True) + res = np.einsum(*operands, **kwargs) res *= alpha # Scale the output @@ -250,7 +251,7 @@ def contract(subscript: str, *args: Any, **kwargs: Any) -> NDArray[T]: return c -def einsum(*operands: Any, **kwargs: Any) -> NDArray[T]: +def einsum(*operands: Any, **kwargs: Any) -> Union[T, NDArray[T]]: """Evaluate an Einstein summation convention on the operands. Using the Einstein summation convention, many common @@ -296,8 +297,8 @@ def einsum(*operands: Any, **kwargs: Any) -> NDArray[T]: # If it's a chain of contractions, use the path optimizer optimize = kwargs.pop("optimize", True) args = list(args) - kwargs = dict(optimize=optimize, einsum_call=True) - _, contractions = np.einsum_path(subscript, *args, **kwargs) # type: ignore + path_kwargs = dict(optimize=optimize, einsum_call=True) + contractions = np.einsum_path(subscript, *args, **path_kwargs)[1] # type: ignore for contraction in contractions: inds, idx_rm, einsum_str, remain = list(contraction[:4]) contraction_args = [args.pop(x) for x in inds] diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index 9cda5239..8533d5c8 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -60,7 +60,12 @@ def __getitem__(self, key: str) -> T: def __getattr__(self, key: str) -> T: """Get an attribute.""" - return self.__getitem__(key) + if key in self.__dict__: + return self.__dict__[key] + try: + return self.__getitem__(key) + except KeyError: + raise AttributeError(f"Namespace object has no attribute {key}") def __delitem__(self, key: str) -> None: """Delete an item.""" @@ -94,15 +99,19 @@ def __len__(self) -> int: def keys(self) -> KeysView[str]: """Get keys of the namespace as a dictionary.""" - return dict(self).keys() + return self._members.keys() def values(self) -> ValuesView[T]: """Get values of the namespace as a dictionary.""" - return dict(self).values() + return self._members.values def items(self) -> ItemsView[str, T]: """Get items of the namespace as a dictionary.""" - return dict(self).items() + return self._members.items() + + def __repr__(self) -> str: + """Return a string representation.""" + return f"Namespace({self._members})" class Timer: diff --git a/ebcc/util/permutations.py b/ebcc/util/permutations.py index fc6f3725..ed02e2b3 100644 --- a/ebcc/util/permutations.py +++ b/ebcc/util/permutations.py @@ -1,12 +1,22 @@ """Symmetry and permutational utilities.""" +from __future__ import annotations + import functools import itertools +from typing import TYPE_CHECKING import numpy as np +if TYPE_CHECKING: + from typing import Any, Callable, Generator, Hashable, Iterable, Optional, TypeVar, Union + + from ebcc.numpy.typing import NDArray + + T = TypeVar("T") + -def factorial(n): +def factorial(n: int) -> int: """Return the factorial of `n`.""" if n in (0, 1): return 1 @@ -14,67 +24,48 @@ def factorial(n): return n * factorial(n - 1) -def permute_string(string, permutation): - """ - Permute a string. - - Parameters - ---------- - string : str - String to permute. - permutation : list of int - Permutation to apply. - - Returns - ------- - permuted : str +def permute_string(string: str, permutation: tuple[int, ...]) -> str: + """Permute a string. + + Args: + string: String to permute. + permutation: Permutation to apply. + + Returns: Permuted string. - Examples - -------- - >>> permute_string("abcd", [2, 0, 3, 1]) - "cbda" + Examples: + >>> permute_string("abcd", (2, 0, 3, 1)) + "cbda" """ return "".join([string[i] for i in permutation]) -def tril_indices_ndim(n, dims, include_diagonal=False): - """ - Return lower triangular indices for a multidimensional array. - - Parameters - ---------- - n : int - Size of each dimension. - dims : int - Number of dimensions. - include_diagonal : bool, optional - If `True`, include diagonal elements. Default value is `False`. - - Returns - ------- - tril : tuple of ndarray +def tril_indices_ndim(n: int, dims: int, include_diagonal: Optional[bool] = False) -> tuple[NDArray[int]]: + """Return lower triangular indices for a multidimensional array. + + Args: + n: Size of each dimension. + dims: Number of dimensions. + include_diagonal: If True, include diagonal elements. + + Returns: Lower triangular indices for each dimension. """ - ranges = [np.arange(n)] * dims - - if dims == 0: - return tuple() - elif dims == 1: + if dims == 1: return (ranges[0],) - - if include_diagonal: - func = np.greater_equal - else: - func = np.greater + #func: Callable[[Any, ...], Any] = np.greater_equal if include_diagonal else np.greater slices = [ tuple(slice(None) if i == j else np.newaxis for i in range(dims)) for j in range(dims) ] casted = [rng[ind] for rng, ind in zip(ranges, slices)] - mask = functools.reduce(np.logical_and, [func(a, b) for a, b in zip(casted[:-1], casted[1:])]) + if include_diagonal: + mask = functools.reduce(np.logical_and, map(np.greater_equal, casted[:-1], casted[1:])) + else: + mask = functools.reduce(np.logical_and, map(np.greater, casted[:-1], casted[1:])) tril = tuple( np.broadcast_to(inds, mask.shape)[mask] for inds in np.indices(mask.shape, sparse=True) @@ -83,7 +74,7 @@ def tril_indices_ndim(n, dims, include_diagonal=False): return tril -def ntril_ndim(n, dims, include_diagonal=False): +def ntril_ndim(n: int, dims: int, include_diagonal: Optional[bool] = False) -> int: """Return `len(tril_indices_ndim(n, dims, include_diagonal))`.""" # FIXME hack until this function is fixed: @@ -104,38 +95,27 @@ def ntril_ndim(n, dims, include_diagonal=False): return out -def generate_spin_combinations(n, excited=False, unique=False): - """ - Generate combinations of spin components for a given number of - occupied and virtual axes. - - Parameters - ---------- - n : int - Order of cluster amplitude. - excited : bool, optional - If True, treat the amplitudes as excited. Default value is - `False`. - unique : bool, optional - If True, return only unique combinations. - - Yields - ------ - spin : str - String of spin combination. - - Examples - -------- - >>> generate_spin_combinations(1) - ['aa', 'bb'] - >>> generate_spin_combinations(2) - ['aaaa', 'abab', 'baba', 'bbbb'] - >>> generate_spin_combinations(2, excited=True) - ['aaa', 'aba', 'bab', 'bbb'] - >>> generate_spin_combinations(2, unique=True) - ['aaaa', 'abab', 'bbbb'] - """ +def generate_spin_combinations(n: int, excited: Optional[bool] = False, unique: Optional[bool] = False) -> Generator[str, None, None]: + """Generate combinations of spin components for a given number of occupied and virtual axes. + Args: + n: Order of cluster amplitude. + excited: If True, treat the amplitudes as excited. + unique: If True, return only unique combinations. + + Returns: + List of spin combinations. + + Examples: + >>> generate_spin_combinations(1) + ['aa', 'bb'] + >>> generate_spin_combinations(2) + ['aaaa', 'abab', 'baba', 'bbbb'] + >>> generate_spin_combinations(2, excited=True) + ['aaa', 'aba', 'bab', 'bbb'] + >>> generate_spin_combinations(2, unique=True) + ['aaaa', 'abab', 'bbbb'] + """ if unique: check = set() @@ -160,31 +140,28 @@ def generate_spin_combinations(n, excited=False, unique=False): yield comb -def permutations_with_signs(seq): - """ - Return permutations of seq, yielding also a sign which is equal to +1 - for an even number of swaps, and -1 for an odd number of swaps. +def permutations_with_signs(seq: Iterable[Any]) -> list[tuple[Any, int]]: + """Return permutations of a sequence with a sign indicating the number of swaps. + + The sign is equal to +1 for an even number of swaps, and -1 for an odd number of swaps. - Parameters - ---------- - seq : iterable - Sequence to permute. + Args: + seq: Sequence to permute. - Returns - ------- - permuted : list of tuple + Returns: List of tuples of the form (permuted, sign). """ - def _permutations(seq): + def _permutations(seq: list[Any]) -> list[list[Any]]: if not seq: return [[]] items = [] for i, item in enumerate(_permutations(seq[:-1])): - inds = range(len(item) + 1) - if i % 2 == 0: - inds = reversed(inds) + if i % 2 == 1: + inds = range(len(item) + 1) + else: + inds = range(len(item), -1, -1) items += [item[:i] + seq[-1:] + item[i:] for i in inds] return items @@ -192,55 +169,41 @@ def _permutations(seq): return [(item, -1 if i % 2 else 1) for i, item in enumerate(_permutations(list(seq)))] -def get_symmetry_factor(*numbers): - """ - Get a floating point value corresponding to the factor from the - neglection of symmetry in repeated indices. +def get_symmetry_factor(*numbers: int) -> float: + """Get a value corresponding to the factor from the neglection of symmetry in repeated indices. - Parameters - ---------- - numbers : tuple of int - Multiplicity of each distinct degree of freedom. + Args: + numbers: Multiplicity of each distinct degree of freedom. - Returns - ------- - factor : float + Returns: Symmetry factor. - Examples - -------- - >>> build_symmetry_factor(1, 1) - 1.0 - >>> build_symmetry_factor(2, 2) - 0.25 - >>> build_symmetry_factor(3, 2, 1) - 0.125 + Examples: + >>> get_symmetry_factor(1, 1) + 1.0 + >>> get_symmetry_factor(2, 2) + 0.25 + >>> get_symmetry_factor(3, 2, 1) + 0.125 """ - ntot = 0 for n in numbers: ntot += max(0, n - 1) - return 1.0 / (2.0**ntot) -def antisymmetrise_array(v, axes=(0, 1)): - """ - Antisymmetrise an array. - - Parameters - ---------- - v : ndarray - Array to antisymmetrise. - axes : tuple of int, optional - Axes to antisymmetrise over. Default value is `(0, 1)`. - - Returns - ------- - v_as : ndarray +def antisymmetrise_array(v: NDArray[T], axes: Optional[tuple[int, ...]] = None) -> NDArray[T]: + """Antisymmetrise an array. + + Args: + v: Array to antisymmetrise. + axes: Axes to antisymmetrise over. + + Returns: Antisymmetrised array. """ - + if axes is None: + axes = tuple(range(v.ndim)) v_as = np.zeros_like(v) for perm, sign in permutations_with_signs(axes): @@ -254,15 +217,16 @@ def antisymmetrise_array(v, axes=(0, 1)): return v_as -def is_mixed_spin(spin): +def is_mixed_spin(spin: Iterable[Hashable]) -> bool: """Return a boolean indicating if a list of spins is mixed.""" return len(set(spin)) != 1 -def combine_subscripts(*subscripts, sizes=None): - """ - Combine subscripts into new unique subscripts for functions such as - `compress_axes`. +def combine_subscripts( + *subscripts: str, + sizes: Optional[dict[tuple[str, ...], int]] = None, +) -> Union[str, tuple[str, dict[str, int]]]: + """Combine subscripts into new unique subscripts for functions such as `compress_axes`. For example, one may wish to compress an amplitude according to both occupancy and spin signatures. @@ -277,30 +241,19 @@ def combine_subscripts(*subscripts, sizes=None): the size of the corresponding original character in the dictionary `sizes`. - Parameters - ---------- - subscripts : tuple of str - Subscripts to combine. Each subscript must be a string of the same - length. - sizes : dict, optional - Dictionary of sizes for each index. Keys should be - `tuple(s[i] for s in subscripts)`. Default value is `None`. - - Returns - ------- - new_subscript : str - Output subscript. - new_sizes : dict, optional - Dictionary of the sizes of each new index. Only returned if the - `sizes` keyword argument is provided. - """ + Args: + subscripts: Subscripts to combine. + sizes: Dictionary of sizes for each index. + Returns: + New subscript, with a dictionary of sizes of each new index if `sizes` is passed. + """ if len(set(len(s) for s in subscripts)) != 1: raise ValueError("Subscripts must be of the same length.") - char_map = {} + char_map: dict[tuple[str, ...], str] = {} new_subscript = "" - new_sizes = {} + new_sizes: dict[str, int] = {} j = 0 for i in range(len(subscripts[0])): key = tuple(s[i] for s in subscripts) @@ -321,34 +274,21 @@ def combine_subscripts(*subscripts, sizes=None): return new_subscript, new_sizes -def compress_axes(subscript, array, include_diagonal=False): - """ - Compress an array into lower-triangular representations using an - einsum-like input. - - Parameters - ---------- - subscript : str - Subscript for the input array. The output array will have a - compressed representation of the input array according to this - subscript, where repeated characters indicate symmetrically - equivalent axes. - array : numpy.ndarray - Array to compress. - include_diagonal : bool, optional - Whether to include the diagonal elements of the input array in the - output array. Default value is `False`. - - Returns - ------- - compressed_array : numpy.ndarray +def compress_axes(subscript: str, array: NDArray[T], include_diagonal: Optional[bool] = False) -> NDArray[T]: + """Compress an array into lower-triangular representations using an einsum-like input. + + Args: + subscript: Subscript for the input array. + array: Array to compress. + include_diagonal: Whether to include the diagonal elements of the input array in the output array. + + Returns: Compressed array. - Examples - -------- - >>> t2 = np.zeros((4, 4, 10, 10)) - >>> compress_axes("iiaa", t2).shape - (6, 45) + Examples: + >>> t2 = np.zeros((4, 4, 10, 10)) + >>> compress_axes("iiaa", t2).shape + (6, 45) """ # TODO out # TODO can this be OpenMP parallel? @@ -365,12 +305,12 @@ def compress_axes(subscript, array, include_diagonal=False): subscript = "".join([subs[s] for s in subscript]) # Reshape array so that all axes of the same character are adjacent: - arg = np.argsort(list(subscript)) + arg = tuple(np.argsort(list(subscript))) array = array.transpose(arg) subscript = permute_string(subscript, arg) # Reshape array so that all axes of the same character are flattened: - sizes = {} + sizes: dict[str, int] = {} for char, n in zip(subscript, array.shape): if char in sizes: assert sizes[char] == n @@ -399,40 +339,29 @@ def compress_axes(subscript, array, include_diagonal=False): def decompress_axes( - subscript, - array_flat, - shape=None, - include_diagonal=False, - symmetry=None, - out=None, -): - """ - Reverse operation of `compress_axes`, subscript input is the same. One - of `shape` or `out` must be passed. - - Parameters - ---------- - subscript : str - Subscript for the output array. The input array will have a - compressed representation of the output array according to this - subscript, where repeated characters indicate symmetrically - equivalent axes. - array_flat : numpy.ndarray - Array to decompress. - shape : tuple of int, optional - Shape of the output array. Must be passed if `out` is `None`. - Default value is `None`. - include_diagonal : bool, optional - Whether to include the diagonal elements of the output array in the - input array. Default value is `False`. - symmetry : str, optional - Symmetry of the output array, with a `"+"` indicating symmetry and - `"-"` indicating antisymmetry for each dimension in the - decompressed array. If `None`, defaults to fully antisymmetric - (i.e. all characters are `"-"`). Default value is `None`. - out : numpy.ndarray, optional - Output array. If `None`, a new array is created, and `shape` must - be passed. Default value is `None`. + subscript: str, + array_flat: NDArray[T], + shape: Optional[tuple[int, ...]] = None, + include_diagonal: Optional[bool] = False, + symmetry: Optional[str] = None, + out: Optional[NDArray[T]] = None, +) -> NDArray[T]: + """Reverse operation of `compress_axes`, subscript input is the same. + + One of `shape` or `out` must be passed. + + Args: + subscript: Subscript for the output array. + array_flat: Array to decompress. + shape: Shape of the output array. Must be passed if `out` is `None`. + include_diagonal: Whether to include the diagonal elements of the output array in the input + array. + symmetry: Symmetry of the output array, with a "+" indicating symmetry and "-" indicating + antisymmetry for each dimension in the decompressed array. + out: Output array. If `None`, a new array is created, and `shape` must be passed. + + Returns: + Decompressed array. """ assert "->" not in subscript @@ -444,6 +373,8 @@ def decompress_axes( # Initialise decompressed array if out is None: + if shape is None: + raise ValueError("One of `shape` or `out` must be passed.") array = np.zeros(shape, dtype=array_flat.dtype) else: array = out @@ -459,12 +390,12 @@ def decompress_axes( subscript = "".join([subs[s] for s in subscript]) # Reshape array so that all axes of the same character are adjacent: - arg = np.argsort(list(subscript)) + arg = tuple(np.argsort(list(subscript))) array = array.transpose(arg) subscript = permute_string(subscript, arg) # Reshape array so that all axes of the same character are flattened: - sizes = {} + sizes: dict[str, int] = {} for char, n in zip(subscript, array.shape): if char in sizes: assert sizes[char] == n @@ -489,19 +420,19 @@ def decompress_axes( # Iterate over permutations with signs: for tup in itertools.product(*[permutations_with_signs(ind) for ind in indices]): indices_perm, signs = zip(*tup) - signs = [s if symm == "-" else 1 for s, symm in zip(signs, symmetry_compressed)] + signs = tuple(s if symm == "-" else 1 for s, symm in zip(signs, symmetry_compressed)) # Apply the indices: - indices_perm = [ + indices_perm = tuple( np.ravel_multi_index(ind, (sizes[char],) * subscript.count(char)) for ind, char in zip(indices_perm, sorted(set(subscript))) - ] - indices_perm = [ + ) + indices_perm = tuple( ind[tuple(np.newaxis if i != j else slice(None) for i in range(len(indices_perm)))] for j, ind in enumerate(indices_perm) - ] - shape = array[tuple(indices_perm)].shape - array[tuple(indices_perm)] = array_flat.reshape(shape) * np.prod(signs) + ) + shape = array[indices_perm].shape + array[indices_perm] = array_flat.reshape(shape) * np.prod(signs) # Reshape array to non-flattened format array = array.reshape( @@ -509,71 +440,46 @@ def decompress_axes( ) # Undo transpose: - arg = np.argsort(arg) + arg = tuple(np.argsort(arg)) array = array.transpose(arg) return array -def get_compressed_size(subscript, **sizes): - """ - Get the size of a compressed representation of a matrix based on the - subscript input to `compressed_axes` and the sizes of each character. - - Parameters - ---------- - subscript : str - Subscript for the output array. See `compressed_axes` for details. - **sizes : int - Sizes of each character in the subscript. - - Returns - ------- - n : int +def get_compressed_size(subscript: str, **sizes: int) -> int: + """Get the size of a compressed representation of a matrix based on the subscript input. + + Args: + subscript: Subscript for the output array. See `compressed_axes` for details. + **sizes: Sizes of each character in the subscript. + + Returns: Size of the compressed representation of the array. - Examples - -------- - >>> get_compressed_shape("iiaa", i=5, a=3) - 30 + Examples: + >>> get_compressed_size("iiaa", i=5, a=3) + 30 """ - n = 1 for char in set(subscript): dims = subscript.count(char) n *= ntril_ndim(sizes[char], dims) - return n -def symmetrise(subscript, array, symmetry=None, apply_factor=True): - """ - Enforce a symmetry in an array. - - Parameters - ---------- - subscript : str - Subscript for the input array. The output array will have a - compressed representation of the input array according to this - subscript, where repeated characters indicate symmetrically - equivalent axes. - array : numpy.ndarray - Array to compress. - symmetry : str, optional - Symmetry of the output array, with a `"+"` indicating symmetry and - `"-"` indicating antisymmetry for each dimension in the - decompressed array. If `None`, defaults to fully symmetric (i.e. - all characters are `"+"`). Default value is `None`. - apply_factor : bool, optional - Whether to apply a factor of to the output array, to account for - the symmetry. Default value is `True`. - - Returns - ------- - array : numpy.ndarray +def symmetrise(subscript: str, array: NDArray[T], symmetry: Optional[str] = None, apply_factor: Optional[bool] = True) -> NDArray[T]: + """Enforce a symmetry in an array. + + Args: + subscript: Subscript for the input array. + array: Array to symmetrise. + symmetry: Symmetry of the output array, with a "+" indicating symmetry and "-" indicating + antisymmetry for each dimension in the decompressed array. + apply_factor: Whether to apply a factor to the output array, to account for the symmetry. + + Returns: Symmetrised array. """ - # Substitute the input characters so that they are ordered: subs = {} i = 0 @@ -603,7 +509,6 @@ def symmetrise(subscript, array, symmetry=None, apply_factor=True): for inds_part, perms_part in zip(inds, perms): for i, p in zip(inds_part, perms_part): perm[i] = p - perm = tuple(perm) sign = np.prod(signs) if symmetry[perm[0]] == "-" else 1 array_as = array_as + sign * array.transpose(perm) @@ -615,7 +520,7 @@ def symmetrise(subscript, array, symmetry=None, apply_factor=True): return array_as -def pack_2e(*args): # noqa +def pack_2e(*args): # type: ignore # noqa # args should be in the order of ov_2e # TODO remove @@ -653,14 +558,12 @@ def pack_2e(*args): # noqa return out -def unique(lst): +def unique(lst: list[Hashable]) -> list[Hashable]: """Get unique elements of a list.""" - done = set() out = [] for el in lst: if el not in done: out.append(el) done.add(el) - return out diff --git a/pyproject.toml b/pyproject.toml index 967d5db9..4ccba447 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,10 +18,10 @@ classifiers = [ "Topic :: Scientific/Engineering", "Topic :: Software Development", "Programming Language :: Python", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Operating System :: MacOS :: MacOS X", "Operating System :: POSIX :: Linux", ] @@ -46,6 +46,7 @@ dev = [ "flake8-pyproject>=1.2.3", "flake8-bugbear>=23.0.0", "flake8-docstrings>=1.6.0", + "mypy>=1.8.0", "coverage[toml]>=5.5.0", "pytest>=6.2.4", "pytest-cov>=4.0.0", @@ -61,7 +62,7 @@ include = "ebcc" exclude = """ /( | ebcc/codegen - | ebcc/__pycache__ + | __pycache__ | .git )/ """ @@ -74,7 +75,7 @@ src_paths = [ "ebcc", ] skip_glob = [ - "ebcc/__pycache__/*", + "*/__pycache__/*", "ebcc/codegen/*", "ebcc/__init__.py", ] @@ -101,11 +102,38 @@ include = "ebcc" exclude = """ /( | ebcc/codegen - | ebcc/__pycache__ + | __pycache__ | .git )/ """ +[tool.mypy] +python_version = "3.10" +exclude = """ +/( + | ebcc/codegen + | __pycache__ + | .git +)/ +""" + +[[tool.mypy.overrides]] +module = "ebcc.*" +strict = true +ignore_missing_imports = true +#implicit_optional = true +disallow_subclassing_any = false +implicit_reexport = true +warn_unused_ignores = false + +[[tool.mypy.overrides]] +module = "scipy.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "h5py.*" +ignore_missing_imports = true + [tool.coverage.run] branch = true source = [ From ff248e00053ce0cc2c88857d2ac197bc3117b6c7 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Fri, 2 Aug 2024 22:16:39 +0100 Subject: [PATCH 09/37] Start on EOM base class --- ebcc/cc/base.py | 4 +- ebcc/eom/__init__.py | 1 + ebcc/eom/base.py | 599 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 602 insertions(+), 2 deletions(-) create mode 100644 ebcc/eom/__init__.py create mode 100644 ebcc/eom/base.py diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index e7148acd..921c5f2a 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -97,8 +97,8 @@ def __init__( mo_occ: Optional[NDArray[float]] = None, fock: Optional[BaseFock] = None, **kwargs: Any, - ): - r"""Initialize the EBCC object. + ) -> None: + r"""Initialise the EBCC object. Args: mf: PySCF mean-field object. diff --git a/ebcc/eom/__init__.py b/ebcc/eom/__init__.py new file mode 100644 index 00000000..71ebcacd --- /dev/null +++ b/ebcc/eom/__init__.py @@ -0,0 +1 @@ +"""Equation-of-motion coupled cluster solvers.""" diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py new file mode 100644 index 00000000..c499388e --- /dev/null +++ b/ebcc/eom/base.py @@ -0,0 +1,599 @@ +"""Base classes for `ebcc.eom`.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from pyscf import lib + +from ebcc import numpy as np +from ebcc import util +from ebcc.logging import ANSI +from ebcc.precision import types + +if TYPE_CHECKING: + from typing import Any, Optional, TypeVar, Union, Callable + + from ebcc.util import Namespace + from ebcc.numpy.typing import NDArray + from ebcc.cc.base import BaseEBCC + + ERIsInputType = Union[type[BaseERIs], NDArray[float]] + AmplitudeType = TypeVar("AmplitudeType") + PickFunctionType = Callable[ + [NDArray[float], NDArray[float], int, dict[str, Any]], + tuple[NDArray[float], NDArray[float], int], + ] + + +@dataclass +class Options: + """Options for EOM calculations. + + Args: + nroots: Number of roots to solver for. + e_tol: Threshold for convergence in the eigenvalues. + max_iter: Maximum number of iterations. + max_space: Maximum size of the Lanczos vector space. + koopmans: Whether to use a Koopmans'-like guess. + """ + + nroots: int = 5 + e_tol: float = util.Inherited * 1e2 + max_iter: int = util.Inherited + max_space: int = 12 + + +class BaseEOM(ABC): + """Base class for equation-of-motion coupled cluster. + + Attributes: + ebcc: Parent `EBCC` object. + options: Options for the EBCC calculation. + """ + + # Types + Options = Options + + def __init__( + self, + ebcc: BaseEBCC, + options: Optional[Options] = None, + **kwargs: Any, + ) -> None: + """Initialise the EOM object. + + Args: + ebcc: Parent `EBCC` object. + options: Options for the EOM calculation. + **kwargs: Additional keyword arguments used to update `options`. + """ + # Options: + if options is None: + options = self.Options() + self.options = options + for key, val in kwargs.items(): + setattr(self.options, key, val) + + # Parameters: + self.ebcc = ebcc + self.space = ebcc.space + self.ansatz = ebcc.anzatz + self.log = ebcc.log + + # Attributes: + self.converged = False + self.e = None + self.v = None + + # Logging: + self.log.info(f"\n{ANSI.B}{ANSI.U}{self.name}{ANSI.R}") + self.log.debug(f"{ANSI.B}{'*' * len(self.name)}{ANSI.R}") + self.log.debug("") + self.log.info(f"{ANSI.B}Options{ANSI.R}:") + self.log.info(f" > nroots: {ANSI.y}{self.options.nroots}{ANSI.R}") + self.log.info(f" > e_tol: {ANSI.y}{self.options.e_tol}{ANSI.R}") + self.log.info(f" > max_iter: {ANSI.y}{self.options.max_iter}{ANSI.R}") + self.log.info(f" > max_space: {ANSI.y}{self.options.max_space}{ANSI.R}") + self.log.debug("") + + @property + @abstractmethod + def excitation_type(self) -> str: + """Get the type of excitation.""" + pass + + @property + def name(self) -> str: + """Get the name of the method.""" + return f"{self.excitation_type.upper()}-EOM-{self.spin_type}{self.ansatz.name}" + + @abstractmethod + def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: + """Construct a vector containing all of the amplitudes used in the given ansatz. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Cluster amplitudes as a vector. + """ + pass + + @abstractmethod + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, ...]: + """Construct amplitudes from a vector. + + Args: + vector: Cluster amplitudes as a vector. + + Returns: + Cluster amplitudes. + """ + pass + + @abstractmethod + def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Apply the Hamiltonian to a vector. + + Args: + vector: State vector to apply the Hamiltonian to. + eris: Electronic repulsion integrals. + + Returns: + Resulting vector. + """ + pass + + @abstractmethod + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ + pass + + @abstractmethod + def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + pass + + @abstractmethod + def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + pass + + def dot_braket(self, bra: NDArray[float], ket: NDArray[float]) -> float: + """Compute the dot product of a bra and ket.""" + return np.dot(bra, ket) + + def get_pick(self, guesses: Optional[NDArray[float]] = None, real_system: bool = True) -> PickFunctionType: + """Get the function to pick the eigenvalues matching the criteria. + + Args: + guesses: Initial guesses for the roots. + real_system: Whether the system is real-valued. + + Returns: + Function to pick the eigenvalues. + """ + if self.options.koopmans: + assert guesses is None + guesses_array = np.asarray(guesses) + + def pick(w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any]) -> tuple[NDArray[float], NDArray[float], int]: + """Pick the eigenvalues.""" + x0 = lib.linalg_helper._gen_x0(envs["v"], envs["xs"]) + x0 = np.asarray(x0) + s = np.dot(g.conj(), x0.T) + s = util.einsum("pi,qi->i", s.conj(), s) + idx = np.argsort(-s)[:nroots] + return lib.linalg_helper._eigs_cmplx2real(w, v, idx, real_system) + + else: + + def pick(w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any]) -> tuple[NDArray[float], NDArray[float], int]: + """Pick the eigenvalues.""" + real_idx = np.where(abs(w.imag) < 1e-3)[0] + w, v, idx = lib.linalg_helper._eigs_cmplx2real(w, v, real_idx, real_system) + mask = np.argsort(np.abs(w)) + w, v = w[mask], v[:, mask] + return w, v, 0 + + return pick + + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(diag[:r1.size]) + else: + arg = np.argsort(diag) + return arg + + def _quasiparticle_weights(self, r1: NDArray[float]) -> float: + """Get the quasiparticle weight.""" + return np.linalg.norm(r1) ** 2 + + def get_guesses(self, diag: NDArray[float]) -> list[NDArray[float]]: + """Get the initial guesses vectors. + + Args: + diag: Diagonal of the Hamiltonian. + + Returns: + Initial guesses. + """ + if diag is None: + diag = self.diag() + arg = self._argsort_guesses(diag) + + nroots = min(self.options.nroots, diag.size) + guesses = np.zeros((nroots, diag.size), dtype=diag.dtype) + for root, guess in enumerate(arg[:nroots]): + guesses[root, guess] = 1.0 + + return list(guesses) + + def callback(self, envs: dict[str, Any]) -> None: + """Callback function for the eigensolver.""" # noqa: D401 + pass + + def davidson(self, eris: Optional[ERIsInputType] = None, guesses: Optional[list[NDArray[float]]] = None) -> NDArray[float]: + """Solve the EOM Hamiltonian using the Davidson solver. + + Args: + eris: Electronic repulsion integrals. + guesses: Initial guesses for the roots. + + Returns: + Energies of the roots. + """ + timer = util.Timer() + + # Get the ERIs: + eris = self.ebcc.get_eris(eris) + + self.log.output( + "Solving for %s excitations using the Davidson solver.", self.excitation_type.upper() + ) + + # Get the matrix-vector products and the diagonal: + matvecs = lambda vs: [self.matvec(v, eris=eris) for v in vs] + diag = self.diag(eris=eris) + + # Get the guesses: + if guesses is None: + guesses = self.get_guesses(diag=diag) + + # Solve the EOM Hamiltonian: + nroots = min(len(guesses), self.options.nroots) + pick = self.get_pick(guesses=guesses) + converged, e, v = lib.davidson_nosym1( + matvecs, + guesses, + diag, + tol=self.options.e_tol, + nroots=nroots, + pick=pick, + max_cycle=self.options.max_iter, + max_space=self.options.max_space, + callback=self.callback, + verbose=0, + ) + + # Check for convergence: + if all(converged): + self.log.debug("") + self.log.output(f"{ANSI.g}Converged.{ANSI.R}") + else: + self.log.debug("") + self.log.warning( + f"{ANSI.r}Failed to converge {sum(not c for c in converged)} roots.{ANSI.R}" + ) + + # Update attributes: + self.converged = converged + self.e = e + self.v = v + + self.log.debug("") + self.log.output( + f"{ANSI.B}{'Root':>4s} {'Energy':>16s} {'Weight':>13s} {'Conv.':>8s}{ANSI.R}" + ) + for n, (en, vn, cn) in enumerate(zip(e, v, converged)): + r1n = self.vector_to_amplitudes(vn)[0] + qpwt = self._quasiparticle_weight(r1n) + self.log.output( + f"{n:>4d} {en:>16.10f} {qpwt:>13.5g} " f"{[ANSI.r, ANSI.g][cn]}{cn!r:>8s}{ANSI.R}" + ) + + self.log.debug("") + self.log.debug("Time elapsed: %s", timer.format_time(timer())) + self.log.debug("") + + return e + + kernel = davidson + + def moments(self, nmom: int, eris: Optional[ERIsInputType] = None, amplitudes: Namespace[AmplitudeType] = None, hermitise: bool = True) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian.""" + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) + + for j in range(self.nmo): + ket = kets[j] + for n in range(nmom): + for i in range(self.nmo): + bra = bras[i] + moments[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.swapaxes(1, 2)) + + return moments + + @property + def nmo(self) -> Any: + """Get the number of MOs.""" + return self.ebcc.nmo + + @property + def nocc(self) -> Any: + """Get the number of occupied MOs.""" + return self.ebcc.nocc + + @property + def nvir(self) -> Any: + """Get the number of virtual MOs.""" + return self.ebcc.nvir + + +class BaseIPEOM(BaseEOM): + """Base class for ionization-potential EOM-CC.""" + + @property + def excitation_type(self) -> str: + """Get the type of excitation.""" + return "ip" + + def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: + """Construct a vector containing all of the amplitudes used in the given ansatz. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Cluster amplitudes as a vector. + """ + return self.ebcc.excitations_to_vector_ip(*amplitudes) + + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, ...]: + """Construct amplitudes from a vector. + + Args: + vector: Cluster amplitudes as a vector. + + Returns: + Cluster amplitudes. + """ + return self.ebcc.vector_to_excitations_ip(vector) + + def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Apply the Hamiltonian to a vector. + + Args: + vector: State vector to apply the Hamiltonian to. + eris: Electronic repulsion integrals. + + Returns: + Resulting vector. + """ + amplitudes = self.vector_to_amplitudes(vector) + result = self.ebcc.hbar_matvec_ip(*amplitudes, eris=eris) + return self.amplitudes_to_vector(*result) + + def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ip_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ip_mom_kets(eris=eris)) + kets = np.array( + [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + ) + return kets + + +class BaseEAEOM(BaseEOM): + """Base class for electron-affinity EOM-CC.""" + + @property + def excitation_type(self) -> str: + """Get the type of excitation.""" + return "ea" + + def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: + """Construct a vector containing all of the amplitudes used in the given ansatz. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Cluster amplitudes as a vector. + """ + return self.ebcc.excitations_to_vector_ea(*amplitudes) + + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, ...]: + """Construct amplitudes from a vector. + + Args: + vector: Cluster amplitudes as a vector. + + Returns: + Cluster amplitudes. + """ + return self.ebcc.vector_to_excitations_ea(vector) + + def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Apply the Hamiltonian to a vector. + + Args: + vector: State vector to apply the Hamiltonian to. + eris: Electronic repulsion integrals. + + Returns: + Resulting vector. + """ + amplitudes = self.vector_to_amplitudes(vector) + result = self.ebcc.hbar_matvec_ea(*amplitudes, eris=eris) + return self.amplitudes_to_vector(*result) + + def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ea_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ea_mom_kets(eris=eris)) + kets = np.array( + [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + ) + return kets + + +class BaseEEEOM(BaseEOM): + """Base class for electron-electron EOM-CC.""" + + @property + def excitation_type(self) -> str: + """Get the type of excitation.""" + return "ee" + + def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: + """Construct a vector containing all of the amplitudes used in the given ansatz. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Cluster amplitudes as a vector. + """ + return self.ebcc.excitations_to_vector_ee(*amplitudes) + + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, ...]: + """Construct amplitudes from a vector. + + Args: + vector: Cluster amplitudes as a vector. + + Returns: + Cluster amplitudes. + """ + return self.ebcc.vector_to_excitations_ee(vector) + + def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Apply the Hamiltonian to a vector. + + Args: + vector: State vector to apply the Hamiltonian to. + eris: Electronic repulsion integrals. + + Returns: + Resulting vector. + """ + amplitudes = self.vector_to_amplitudes(vector) + result = self.ebcc.hbar_matvec_ee(*amplitudes, eris=eris) + return self.amplitudes_to_vector(*result) + + def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) From d078683797be21ffe610493bcced954a67706805 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sat, 3 Aug 2024 11:34:33 +0100 Subject: [PATCH 10/37] EOM base classes --- ebcc/__init__.py | 2 - ebcc/cc/base.py | 129 ++++++++-- ebcc/cc/gebcc.py | 59 +++-- ebcc/cc/rebcc.py | 60 +++-- ebcc/cc/uebcc.py | 64 +++-- ebcc/eom/__init__.py | 4 + ebcc/eom/base.py | 146 +++++------ ebcc/eom/geom.py | 227 +++++++++++++++++ ebcc/eom/reom.py | 227 +++++++++++++++++ ebcc/{ => eom}/ueom.py | 315 +++++++++--------------- ebcc/geom.py | 29 --- ebcc/reom.py | 497 -------------------------------------- ebcc/util/permutations.py | 21 +- 13 files changed, 897 insertions(+), 883 deletions(-) create mode 100644 ebcc/eom/geom.py create mode 100644 ebcc/eom/reom.py rename ebcc/{ => eom}/ueom.py (55%) delete mode 100644 ebcc/geom.py delete mode 100644 ebcc/reom.py diff --git a/ebcc/__init__.py b/ebcc/__init__.py index bc009fb6..027ce828 100644 --- a/ebcc/__init__.py +++ b/ebcc/__init__.py @@ -46,11 +46,9 @@ # --- Import NumPy here to allow drop-in replacements - # --- Logging: - # --- Types of ansatz supporting by the EBCC solvers: METHOD_TYPES = ["MP", "CC", "LCC", "QCI", "QCC", "DC"] diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 921c5f2a..ceae2f99 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -300,7 +300,11 @@ def kernel(self, eris: Optional[ERIsInputType] = None) -> float: return e_cc - def solve_lambda(self, amplitudes: Optional[Namespace[AmplitudeType]] = None, eris: Optional[ERIsInputType] = None) -> None: + def solve_lambda( + self, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + eris: Optional[ERIsInputType] = None, + ) -> None: """Solve for the lambda amplitudes. Args: @@ -521,7 +525,9 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude pass @abstractmethod - def init_lams(self, amplitude: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + def init_lams( + self, amplitude: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: """Initialise the cluster lambda amplitudes. Args: @@ -617,7 +623,12 @@ def update_lams( """ pass - def make_sing_b_dm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> Any: + def make_sing_b_dm( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + ) -> Any: r"""Make the single boson density matrix :math:`\langle b \rangle`. Args: @@ -636,7 +647,14 @@ def make_sing_b_dm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optio ) return func(**kwargs) - def make_rdm1_b(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + def make_rdm1_b( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + unshifted: bool = True, + hermitise: bool = True, + ) -> Any: r"""Make the one-particle boson reduced density matrix :math:`\langle b^+ c \rangle`. Args: @@ -669,7 +687,13 @@ def make_rdm1_b(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional return dm @abstractmethod - def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + def make_rdm1_f( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + ) -> Any: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -684,7 +708,13 @@ def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional pass @abstractmethod - def make_rdm2_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + def make_rdm2_f( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + ) -> Any: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -699,7 +729,14 @@ def make_rdm2_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional pass @abstractmethod - def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + def make_eb_coup_rdm( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + unshifted: bool = True, + hermitise: bool = True, + ) -> Any: r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. Args: @@ -715,7 +752,13 @@ def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Opt """ pass - def hbar_matvec_ip(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType]: + def hbar_matvec_ip( + self, + r1: AmplitudeType, + r2: AmplitudeType, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType]: """Compute the product between a state vector and the IP-EOM Hamiltonian. Args: @@ -737,7 +780,13 @@ def hbar_matvec_ip(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ER ) return func(**kwargs) - def hbar_matvec_ea(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, AmplitudeType]: + def hbar_matvec_ea( + self, + r1: AmplitudeType, + r2: AmplitudeType, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType, AmplitudeType]: """Compute the product between a state vector and the EA-EOM Hamiltonian. Args: @@ -759,7 +808,13 @@ def hbar_matvec_ea(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ER ) return func(**kwargs) - def hbar_matvec_ee(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, AmplitudeType]: + def hbar_matvec_ee( + self, + r1: AmplitudeType, + r2: AmplitudeType, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType, AmplitudeType]: """Compute the product between a state vector and the EE-EOM Hamiltonian. Args: @@ -781,7 +836,12 @@ def hbar_matvec_ee(self, r1: AmplitudeType, r2: AmplitudeType, eris: Optional[ER ) return func(**kwargs) - def make_ip_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + def make_ip_mom_bras( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType, ...]: """Get the bra vectors to construct IP-EOM moments. Args: @@ -800,7 +860,12 @@ def make_ip_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Opt ) return func(**kwargs) - def make_ea_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + def make_ea_mom_bras( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType, ...]: """Get the bra vectors to construct EA-EOM moments. Args: @@ -819,7 +884,12 @@ def make_ea_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Opt ) return func(**kwargs) - def make_ee_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + def make_ee_mom_bras( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType, ...]: """Get the bra vectors to construct EE-EOM moments. Args: @@ -838,7 +908,12 @@ def make_ee_mom_bras(self, eris: Optional[ERIsInputType] = None, amplitudes: Opt ) return func(**kwargs) - def make_ip_mom_kets(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + def make_ip_mom_kets( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType, ...]: """Get the ket vectors to construct IP-EOM moments. Args: @@ -857,7 +932,12 @@ def make_ip_mom_kets(self, eris: Optional[ERIsInputType] = None, amplitudes: Opt ) return func(**kwargs) - def make_ea_mom_kets(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + def make_ea_mom_kets( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType, ...]: """Get the ket vectors to construct EA-EOM moments. Args: @@ -876,7 +956,12 @@ def make_ea_mom_kets(self, eris: Optional[ERIsInputType] = None, amplitudes: Opt ) return func(**kwargs) - def make_ee_mom_kets(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None) -> tuple[AmplitudeType, ...]: + def make_ee_mom_kets( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + ) -> tuple[AmplitudeType, ...]: """Get the ket vectors to construct EE-EOM moments. Args: @@ -994,7 +1079,9 @@ def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> ND pass @abstractmethod - def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ip( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of IP-EOM excitations from a vector. Args: @@ -1006,7 +1093,9 @@ def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[Am pass @abstractmethod - def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ea( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of EA-EOM excitations from a vector. Args: @@ -1018,7 +1107,9 @@ def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[Am pass @abstractmethod - def vector_to_excitations_ee(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ee( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of EE-EOM excitations from a vector. Args: diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 4196b71e..00ca8318 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -1,4 +1,4 @@ -"""General electron-boson coupled cluster.""" +"""Generalised electron-boson coupled cluster.""" from __future__ import annotations @@ -6,11 +6,11 @@ from pyscf import lib, scf -from ebcc import geom from ebcc import numpy as np -from ebcc import uebcc, util +from ebcc import util from ebcc.brueckner import BruecknerGEBCC from ebcc.cc.base import BaseEBCC +from ebcc.eom import EA_GEOM, EE_GEOM, IP_GEOM from ebcc.eris import GERIs from ebcc.fock import GFock from ebcc.precision import types @@ -56,7 +56,7 @@ class GEBCC(BaseEBCC): def spin_type(self): return "G" - def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.IP_GEOM: + def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_GEOM: """Get the IP-EOM object. Args: @@ -66,9 +66,9 @@ def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.I Returns: IP-EOM object. """ - return geom.IP_GEOM(self, options=options, **kwargs) + return IP_GEOM(self, options=options, **kwargs) - def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.EA_GEOM: + def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EA_GEOM: """Get the EA-EOM object. Args: @@ -78,9 +78,9 @@ def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.E Returns: EA-EOM object. """ - return geom.EA_GEOM(self, options=options, **kwargs) + return EA_GEOM(self, options=options, **kwargs) - def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.EE_GEOM: + def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EE_GEOM: """Get the EE-EOM object. Args: @@ -90,7 +90,7 @@ def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> geom.E Returns: EE-EOM object. """ - return geom.EE_GEOM(self, options=options, **kwargs) + return EE_GEOM(self, options=options, **kwargs) @staticmethod def _convert_mf(mf: SCF) -> GHF: @@ -319,7 +319,9 @@ def from_rebcc(cls, rcc: REBCC) -> GEBCC: Returns: GEBCC object. """ - ucc = uebcc.UEBCC.from_rebcc(rcc) + from ebcc.cc.uebcc import UEBCC + + ucc = UEBCC.from_rebcc(rcc) gcc = cls.from_uebcc(ucc) return gcc @@ -336,7 +338,9 @@ def init_space(self) -> Space: ) return space - def _pack_codegen_kwargs(self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None) -> dict[str, Any]: + def _pack_codegen_kwargs( + self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None + ) -> dict[str, Any]: """Pack all the keyword arguments for the generated code.""" kwargs = dict( f=self.fock, @@ -403,7 +407,9 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude return amplitudes - def init_lams(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + def init_lams( + self, amplitudes: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: """Initialise the cluster lambda amplitudes. Args: @@ -539,7 +545,13 @@ def update_lams( return res - def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + def make_rdm1_f( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + ) -> Any: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -580,7 +592,14 @@ def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): return dm - def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + def make_eb_coup_rdm( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + unshifted: bool = True, + hermitise: bool = True, + ) -> Any: r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. Args: @@ -824,7 +843,9 @@ def excitations_to_vector_ee(self, *excitations): return np.concatenate(vectors) - def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ip( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of IP-EOM excitations from a vector. Args: @@ -853,7 +874,9 @@ def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[Am return tuple(excitations) - def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ea( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of EA-EOM excitations from a vector. Args: @@ -882,7 +905,9 @@ def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[Am return tuple(excitations) - def vector_to_excitations_ee(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ee( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of EE-EOM excitations from a vector. Args: diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 09c3c469..8dc89faf 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -9,13 +9,14 @@ from ebcc import default_log, init_logging from ebcc import numpy as np -from ebcc import reom, util +from ebcc import util from ebcc.ansatz import Ansatz from ebcc.brueckner import BruecknerREBCC from ebcc.cc.base import BaseEBCC from ebcc.cderis import RCDERIs from ebcc.damping import DIIS from ebcc.dump import Dump +from ebcc.eom import EA_REOM, EE_REOM, IP_REOM from ebcc.eris import RERIs from ebcc.fock import RFock from ebcc.logging import ANSI @@ -60,7 +61,7 @@ def spin_type(self): """Get a string representation of the spin type.""" return "R" - def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.IP_REOM: + def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_REOM: """Get the IP-EOM object. Args: @@ -70,9 +71,9 @@ def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.I Returns: IP-EOM object. """ - return reom.IP_REOM(self, options=options, **kwargs) + return IP_REOM(self, options=options, **kwargs) - def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.EA_REOM: + def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EA_REOM: """Get the EA-EOM object. Args: @@ -82,9 +83,9 @@ def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.E Returns: EA-EOM object. """ - return reom.EA_REOM(self, options=options, **kwargs) + return EA_REOM(self, options=options, **kwargs) - def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.EE_REOM: + def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EE_REOM: """Get the EE-EOM object. Args: @@ -94,7 +95,7 @@ def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> reom.E Returns: EE-EOM object. """ - return reom.EE_REOM(self, options=options, **kwargs) + return EE_REOM(self, options=options, **kwargs) @staticmethod def _convert_mf(mf: SCF) -> RHF: @@ -114,7 +115,9 @@ def init_space(self) -> Space: ) return space - def _pack_codegen_kwargs(self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None) -> dict[str, Any]: + def _pack_codegen_kwargs( + self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None + ) -> dict[str, Any]: """Pack all the keyword arguments for the generated code.""" kwargs = dict( f=self.fock, @@ -185,7 +188,9 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude return amplitudes - def init_lams(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + def init_lams( + self, amplitudes: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: """Initialise the cluster lambda amplitudes. Args: @@ -321,7 +326,13 @@ def update_lams( return res - def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + def make_rdm1_f( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + ) -> Any: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -346,7 +357,13 @@ def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional return dm - def make_rdm2_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + def make_rdm2_f( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + ) -> Any: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -372,7 +389,14 @@ def make_rdm2_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional return dm - def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + def make_eb_coup_rdm( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + unshifted: bool = True, + hermitise: bool = True, + ) -> Any: r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. Args: @@ -602,7 +626,9 @@ def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> ND """ return self.excitations_to_vector_ip(*excitations) - def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ip( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of IP-EOM excitations from a vector. Args: @@ -629,7 +655,9 @@ def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[Am return tuple(excitations) - def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ea( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of EA-EOM excitations from a vector. Args: @@ -656,7 +684,9 @@ def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[Am return tuple(excitations) - def vector_to_excitations_ee(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ee( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of EE-EOM excitations from a vector. Args: diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 5166159e..5e760f5d 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -7,10 +7,11 @@ from pyscf import lib from ebcc import numpy as np -from ebcc import rebcc, ueom, util +from ebcc import util from ebcc.brueckner import BruecknerUEBCC from ebcc.cc.base import BaseEBCC from ebcc.cderis import UCDERIs +from ebcc.eom import EA_UEOM, EE_UEOM, IP_UEOM from ebcc.eris import UERIs from ebcc.fock import UFock from ebcc.precision import types @@ -47,7 +48,7 @@ def spin_type(self): """Get a string representation of the spin type.""" return "U" - def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.IP_REOM: + def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_UEOM: """Get the IP-EOM object. Args: @@ -57,9 +58,9 @@ def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.I Returns: IP-EOM object. """ - return ueom.IP_UEOM(self, options=options, **kwargs) + return IP_UEOM(self, options=options, **kwargs) - def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.EA_REOM: + def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EA_UEOM: """Get the EA-EOM object. Args: @@ -69,9 +70,9 @@ def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.E Returns: EA-EOM object. """ - return ueom.EA_UEOM(self, options=options, **kwargs) + return EA_UEOM(self, options=options, **kwargs) - def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.EE_REOM: + def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EE_UEOM: """Get the EE-EOM object. Args: @@ -81,7 +82,7 @@ def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> ueom.E Returns: EE-EOM object. """ - return ueom.EE_UEOM(self, options=options, **kwargs) + return EE_UEOM(self, options=options, **kwargs) @staticmethod def _convert_mf(mf: SCF) -> UHF: @@ -185,7 +186,9 @@ def init_space(self) -> Namespace[Space]: ) return space - def _pack_codegen_kwargs(self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None) -> dict[str, Any]: + def _pack_codegen_kwargs( + self, *extra_kwargs: dict[str, Any], eris: Optional[ERIsInputType] = None + ) -> dict[str, Any]: """Pack all the keyword arguments for the generated code.""" kwargs = dict( f=self.fock, @@ -276,7 +279,9 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude return amplitudes - def init_lams(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + def init_lams( + self, amplitudes: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: """Initialise the cluster lambda amplitudes. Args: @@ -431,7 +436,13 @@ def update_lams( return res - def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + def make_rdm1_f( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + ) -> Any: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -458,7 +469,13 @@ def make_rdm1_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional return dm - def make_rdm2_f(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True) -> Any: + def make_rdm2_f( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + ) -> Any: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -495,7 +512,14 @@ def transpose2(dm): return dm - def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True) -> Any: + def make_eb_coup_rdm( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + unshifted: bool = True, + hermitise: bool = True, + ) -> Any: r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. Args: @@ -533,7 +557,9 @@ def make_eb_coup_rdm(self, eris: Optional[ERIsInputType] = None, amplitudes: Opt return dm_eb - def energy_sum(self, subscript: str, spins: str, signs_dict: dict[str, int] = None) -> NDArray[float]: + def energy_sum( + self, subscript: str, spins: str, signs_dict: dict[str, int] = None + ) -> NDArray[float]: """Get a direct sum of energies. Args: @@ -812,7 +838,9 @@ def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> ND return np.concatenate(vectors) - def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ip( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of IP-EOM excitations from a vector. Args: @@ -852,7 +880,9 @@ def vector_to_excitations_ip(self, vector: NDArray[float]) -> tuple[Namespace[Am return tuple(excitations) - def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ea( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of EA-EOM excitations from a vector. Args: @@ -892,7 +922,9 @@ def vector_to_excitations_ea(self, vector: NDArray[float]) -> tuple[Namespace[Am return tuple(excitations) - def vector_to_excitations_ee(self, vector: NDArray[float]) -> tuple[Namespace[AmplitudeType], ...]: + def vector_to_excitations_ee( + self, vector: NDArray[float] + ) -> tuple[Namespace[AmplitudeType], ...]: """Construct a namespace of EE-EOM excitations from a vector. Args: diff --git a/ebcc/eom/__init__.py b/ebcc/eom/__init__.py index 71ebcacd..08783d00 100644 --- a/ebcc/eom/__init__.py +++ b/ebcc/eom/__init__.py @@ -1 +1,5 @@ """Equation-of-motion coupled cluster solvers.""" + +from ebcc.eom.geom import EA_GEOM, EE_GEOM, IP_GEOM +from ebcc.eom.reom import EA_REOM, EE_REOM, IP_REOM +from ebcc.eom.ueom import EA_UEOM, EE_UEOM, IP_UEOM diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py index c499388e..25c1a71a 100644 --- a/ebcc/eom/base.py +++ b/ebcc/eom/base.py @@ -14,14 +14,12 @@ from ebcc.precision import types if TYPE_CHECKING: - from typing import Any, Optional, TypeVar, Union, Callable + from typing import Any, Callable, Optional, TypeVar, Union - from ebcc.util import Namespace + from ebcc.cc.base import AmplitudeType, BaseEBCC, ERIsInputType from ebcc.numpy.typing import NDArray - from ebcc.cc.base import BaseEBCC + from ebcc.util import Namespace - ERIsInputType = Union[type[BaseERIs], NDArray[float]] - AmplitudeType = TypeVar("AmplitudeType") PickFunctionType = Callable[ [NDArray[float], NDArray[float], int, dict[str, Any]], tuple[NDArray[float], NDArray[float], int], @@ -41,9 +39,10 @@ class Options: """ nroots: int = 5 - e_tol: float = util.Inherited * 1e2 - max_iter: int = util.Inherited + e_tol: float = 1e-6 + max_iter: int = 100 max_space: int = 12 + koopmans: bool = False class BaseEOM(ABC): @@ -80,13 +79,13 @@ def __init__( # Parameters: self.ebcc = ebcc self.space = ebcc.space - self.ansatz = ebcc.anzatz + self.ansatz = ebcc.ansatz self.log = ebcc.log # Attributes: self.converged = False - self.e = None - self.v = None + self.e: NDArray[float] = np.empty((0), dtype=types[float]) + self.v: NDArray[float] = np.empty((0, 0), dtype=types[float]) # Logging: self.log.info(f"\n{ANSI.B}{ANSI.U}{self.name}{ANSI.R}") @@ -105,6 +104,11 @@ def excitation_type(self) -> str: """Get the type of excitation.""" pass + @property + def spin_type(self) -> str: + """Get the spin type.""" + return self.ebcc.spin_type + @property def name(self) -> str: """Get the name of the method.""" @@ -135,7 +139,9 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, pass @abstractmethod - def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) -> NDArray[float]: + def matvec( + self, vector: NDArray[float], eris: Optional[ERIsInputType] = None + ) -> NDArray[float]: """Apply the Hamiltonian to a vector. Args: @@ -160,7 +166,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: pass @abstractmethod - def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + def bras(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: """Get the bra vectors. Args: @@ -172,7 +178,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: pass @abstractmethod - def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + def kets(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: """Get the ket vectors. Args: @@ -187,7 +193,9 @@ def dot_braket(self, bra: NDArray[float], ket: NDArray[float]) -> float: """Compute the dot product of a bra and ket.""" return np.dot(bra, ket) - def get_pick(self, guesses: Optional[NDArray[float]] = None, real_system: bool = True) -> PickFunctionType: + def get_pick( + self, guesses: Optional[NDArray[float]] = None, real_system: bool = True + ) -> PickFunctionType: """Get the function to pick the eigenvalues matching the criteria. Args: @@ -196,12 +204,14 @@ def get_pick(self, guesses: Optional[NDArray[float]] = None, real_system: bool = Returns: Function to pick the eigenvalues. - """ + """ if self.options.koopmans: assert guesses is None guesses_array = np.asarray(guesses) - def pick(w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any]) -> tuple[NDArray[float], NDArray[float], int]: + def pick( + w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any] + ) -> tuple[NDArray[float], NDArray[float], int]: """Pick the eigenvalues.""" x0 = lib.linalg_helper._gen_x0(envs["v"], envs["xs"]) x0 = np.asarray(x0) @@ -212,7 +222,9 @@ def pick(w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any]) else: - def pick(w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any]) -> tuple[NDArray[float], NDArray[float], int]: + def pick( + w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any] + ) -> tuple[NDArray[float], NDArray[float], int]: """Pick the eigenvalues.""" real_idx = np.where(abs(w.imag) < 1e-3)[0] w, v, idx = lib.linalg_helper._eigs_cmplx2real(w, v, real_idx, real_system) @@ -222,18 +234,15 @@ def pick(w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any]) return pick + @abstractmethod def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: """Sort the diagonal to inform the initial guesses.""" - if self.options.koopmans: - r1 = self.vector_to_amplitudes(diag)[0] - arg = np.argsort(diag[:r1.size]) - else: - arg = np.argsort(diag) - return arg + pass - def _quasiparticle_weights(self, r1: NDArray[float]) -> float: + @abstractmethod + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1) ** 2 + pass def get_guesses(self, diag: NDArray[float]) -> list[NDArray[float]]: """Get the initial guesses vectors. @@ -259,7 +268,9 @@ def callback(self, envs: dict[str, Any]) -> None: """Callback function for the eigensolver.""" # noqa: D401 pass - def davidson(self, eris: Optional[ERIsInputType] = None, guesses: Optional[list[NDArray[float]]] = None) -> NDArray[float]: + def davidson( + self, eris: Optional[ERIsInputType] = None, guesses: Optional[list[NDArray[float]]] = None + ) -> NDArray[float]: """Solve the EOM Hamiltonian using the Davidson solver. Args: @@ -336,31 +347,16 @@ def davidson(self, eris: Optional[ERIsInputType] = None, guesses: Optional[list[ kernel = davidson - def moments(self, nmom: int, eris: Optional[ERIsInputType] = None, amplitudes: Namespace[AmplitudeType] = None, hermitise: bool = True) -> NDArray[float]: + @abstractmethod + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian.""" - if eris is None: - eris = self.ebcc.get_eris() - if not amplitudes: - amplitudes = self.ebcc.amplitudes - - bras = self.bras(eris=eris) - kets = self.kets(eris=eris) - - moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) - - for j in range(self.nmo): - ket = kets[j] - for n in range(nmom): - for i in range(self.nmo): - bra = bras[i] - moments[n, i, j] = self.dot_braket(bra, ket) - if n != (nmom - 1): - ket = self.matvec(ket, eris=eris) - - if hermitise: - moments = 0.5 * (moments + moments.swapaxes(1, 2)) - - return moments + pass @property def nmo(self) -> Any: @@ -378,8 +374,8 @@ def nvir(self) -> Any: return self.ebcc.nvir -class BaseIPEOM(BaseEOM): - """Base class for ionization-potential EOM-CC.""" +class BaseIP_EOM(BaseEOM): + """Base class for ionisation-potential EOM-CC.""" @property def excitation_type(self) -> str: @@ -408,7 +404,9 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, """ return self.ebcc.vector_to_excitations_ip(vector) - def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) -> NDArray[float]: + def matvec( + self, vector: NDArray[float], eris: Optional[ERIsInputType] = None + ) -> NDArray[float]: """Apply the Hamiltonian to a vector. Args: @@ -422,38 +420,8 @@ def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) - result = self.ebcc.hbar_matvec_ip(*amplitudes, eris=eris) return self.amplitudes_to_vector(*result) - def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: - """Get the bra vectors. - - Args: - eris: Electronic repulsion integrals. - - Returns: - Bra vectors. - """ - bras_raw = list(self.ebcc.make_ip_mom_bras(eris=eris)) - bras = np.array( - [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] - ) - return bras - - def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: - """Get the ket vectors. - - Args: - eris: Electronic repulsion integrals. - - Returns: - Ket vectors. - """ - kets_raw = list(self.ebcc.make_ip_mom_kets(eris=eris)) - kets = np.array( - [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] - ) - return kets - -class BaseEAEOM(BaseEOM): +class BaseEA_EOM(BaseEOM): """Base class for electron-affinity EOM-CC.""" @property @@ -483,7 +451,9 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, """ return self.ebcc.vector_to_excitations_ea(vector) - def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) -> NDArray[float]: + def matvec( + self, vector: NDArray[float], eris: Optional[ERIsInputType] = None + ) -> NDArray[float]: """Apply the Hamiltonian to a vector. Args: @@ -528,7 +498,7 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return kets -class BaseEEEOM(BaseEOM): +class BaseEE_EOM(BaseEOM): """Base class for electron-electron EOM-CC.""" @property @@ -558,7 +528,9 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, """ return self.ebcc.vector_to_excitations_ee(vector) - def matvec(self, vector: NDArray[float], eris: Optional[ERIsInputType] = None) -> NDArray[float]: + def matvec( + self, vector: NDArray[float], eris: Optional[ERIsInputType] = None + ) -> NDArray[float]: """Apply the Hamiltonian to a vector. Args: diff --git a/ebcc/eom/geom.py b/ebcc/eom/geom.py new file mode 100644 index 00000000..a3132649 --- /dev/null +++ b/ebcc/eom/geom.py @@ -0,0 +1,227 @@ +"""Generalised equation-of-motion coupled cluster.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ebcc import numpy as np +from ebcc import util +from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM, Options +from ebcc.precision import types + +if TYPE_CHECKING: + from typing import Optional, Tuple, Union + + from ebcc.cc.gebcc import AmplitudeType, ERIsInputType + from ebcc.numpy.typing import NDArray + from ebcc.util import Namespace + + +class GEOM(BaseEOM): + """Generalised equation-of-motion coupled cluster.""" + + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(diag[: r1.size]) + else: + arg = np.argsort(diag) + return arg + + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" + return np.linalg.norm(r1) ** 2 + + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian.""" + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) + + for j in range(self.nmo): + ket = kets[j] + for n in range(nmom): + for i in range(self.nmo): + bra = bras[i] + moments[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.swapaxes(1, 2)) + + return moments + + +class IP_GEOM(BaseIP_EOM, BaseEOM): + """Generalised ionisation potential equation-of-motion coupled cluster.""" + + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ + parts = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[:-1] + parts.append(self.ebcc.energy_sum(key)) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return self.amplitudes_to_vector(*parts) + + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ip_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ip_mom_kets(eris=eris)) + kets = np.array( + [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + ) + return kets + + +class EA_GEOM(BaseEA_EOM, BaseEOM): + """Generalised electron affinity equation-of-motion coupled cluster.""" + + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ + parts = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[n:] + key[: n - 1] + parts.append(-self.ebcc.energy_sum(key)) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return self.amplitudes_to_vector(*parts) + + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ea_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ea_mom_kets(eris=eris)) + kets = np.array( + [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + ) + return kets + + +class EE_GEOM(BaseEE_EOM, BaseEOM): + """Generalised electron-electron equation-of-motion coupled cluster.""" + + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ + parts = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + parts.append(-self.ebcc.energy_sum(key)) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return self.amplitudes_to_vector(*parts) + + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) + kets = np.array( + [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + ) + return kets diff --git a/ebcc/eom/reom.py b/ebcc/eom/reom.py new file mode 100644 index 00000000..8547e13c --- /dev/null +++ b/ebcc/eom/reom.py @@ -0,0 +1,227 @@ +"""Restricted equation-of-motion coupled cluster.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ebcc import numpy as np +from ebcc import util +from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM, Options +from ebcc.precision import types + +if TYPE_CHECKING: + from typing import Optional, Tuple, Union + + from ebcc.cc.rebcc import AmplitudeType, ERIsInputType + from ebcc.numpy.typing import NDArray + from ebcc.util import Namespace + + +class REOM(BaseEOM): + """Restricted equation-of-motion coupled cluster.""" + + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(diag[: r1.size]) + else: + arg = np.argsort(diag) + return arg + + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" + return np.linalg.norm(r1) ** 2 + + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian.""" + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) + + for j in range(self.nmo): + ket = kets[j] + for n in range(nmom): + for i in range(self.nmo): + bra = bras[i] + moments[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.swapaxes(1, 2)) + + return moments + + +class IP_REOM(REOM, BaseIP_EOM): + """Restricted ionisation potential equation-of-motion coupled cluster.""" + + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ + parts = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[:-1] + parts.append(self.ebcc.energy_sum(key)) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return self.amplitudes_to_vector(*parts) + + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ip_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ip_mom_kets(eris=eris)) + kets = np.array( + [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + ) + return kets + + +class EA_REOM(REOM, BaseEA_EOM): + """Restricted electron affinity equation-of-motion coupled cluster.""" + + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ + parts = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + key = key[n:] + key[: n - 1] + parts.append(-self.ebcc.energy_sum(key)) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return self.amplitudes_to_vector(*parts) + + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ea_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ea_mom_kets(eris=eris)) + kets = np.array( + [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + ) + return kets + + +class EE_REOM(REOM, BaseEE_EOM): + """Restricted electron-electron equation-of-motion coupled cluster.""" + + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ + parts = [] + + for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + parts.append(-self.ebcc.energy_sum(key)) + + for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented + + return self.amplitudes_to_vector(*parts) + + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) + bras = np.array( + [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + ) + return bras + + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) + kets = np.array( + [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + ) + return kets diff --git a/ebcc/ueom.py b/ebcc/eom/ueom.py similarity index 55% rename from ebcc/ueom.py rename to ebcc/eom/ueom.py index c95adec2..861367a2 100644 --- a/ebcc/ueom.py +++ b/ebcc/eom/ueom.py @@ -1,32 +1,49 @@ -"""Unrestricted equation-of-motion solver.""" +"""Unrestricted equation-of-motion coupled cluster.""" -import warnings +from __future__ import annotations + +from typing import TYPE_CHECKING from ebcc import numpy as np -from ebcc import reom, util +from ebcc import util +from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM, Options from ebcc.precision import types +if TYPE_CHECKING: + from typing import Optional, Tuple, Union + + from ebcc.cc.base import AmplitudeType, ERIsInputType + from ebcc.numpy.typing import NDArray + from ebcc.util import Namespace + -class UEOM(reom.REOM): - """Unrestricted equation-of-motion base class.""" +class UEOM(BaseEOM): + """Unrestricted equation-of-motion coupled cluster.""" - def _argsort_guess(self, diag): + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" if self.options.koopmans: - r_mf = self.vector_to_amplitudes(diag)[0] - size = r_mf.a.size + r_mf.b.size - arg = np.argsort(np.diag(diag[:size])) + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(diag[: r1.a.size + r1.b.size]) else: arg = np.argsort(diag) - return arg - def _quasiparticle_weight(self, r1): + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" return np.linalg.norm(r1.a) ** 2 + np.linalg.norm(r1.b) ** 2 - def moments(self, nmom, eris=None, amplitudes=None, hermitise=True): + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian.""" if eris is None: eris = self.ebcc.get_eris() - if amplitudes is None: + if not amplitudes: amplitudes = self.ebcc.amplitudes bras = self.bras(eris=eris) @@ -54,10 +71,18 @@ def moments(self, nmom, eris=None, amplitudes=None, hermitise=True): return moments -class IP_UEOM(UEOM, reom.IP_REOM): - """Unrestricted equation-of-motion class for ionisation potentials.""" +class IP_UEOM(UEOM, BaseIP_EOM): + """Unrestricted ionisation potential equation-of-motion coupled cluster.""" - def diag(self, eris=None): + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ parts = [] for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -72,8 +97,16 @@ def diag(self, eris=None): return self.amplitudes_to_vector(*parts) - def bras(self, eris=None): - bras_raw = list(self._bras(eris=eris)) + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ip_mom_bras(eris=eris)) bras = util.Namespace(a=[], b=[]) for i in range(self.nmo): @@ -120,8 +153,16 @@ def bras(self, eris=None): return bras - def kets(self, eris=None): - kets_raw = list(self._kets(eris=eris)) + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ip_mom_kets(eris=eris)) kets = util.Namespace(a=[], b=[]) for i in range(self.nmo): @@ -170,10 +211,18 @@ def kets(self, eris=None): return kets -class EA_UEOM(UEOM, reom.EA_REOM): - """Unrestricted equation-of-motion class for electron affinities.""" +class EA_UEOM(UEOM, BaseEA_EOM): + """Unrestricted electron affinity equation-of-motion coupled cluster.""" - def diag(self, eris=None): + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ parts = [] for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -188,8 +237,16 @@ def diag(self, eris=None): return self.amplitudes_to_vector(*parts) - def bras(self, eris=None): - bras_raw = list(self._bras(eris=eris)) + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Bra vectors. + """ + bras_raw = list(self.ebcc.make_ea_mom_bras(eris=eris)) bras = util.Namespace(a=[], b=[]) for i in range(self.nmo): @@ -236,8 +293,16 @@ def bras(self, eris=None): return bras - def kets(self, eris=None): - kets_raw = list(self._kets(eris=eris)) + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Ket vectors. + """ + kets_raw = list(self.ebcc.make_ea_mom_kets(eris=eris)) kets = util.Namespace(a=[], b=[]) for i in range(self.nmo): @@ -286,13 +351,22 @@ def kets(self, eris=None): return kets -class EE_UEOM(UEOM, reom.EE_REOM): - """Unrestricted equation-of-motion class for neutral excitations.""" +class EE_UEOM(UEOM, BaseEE_EOM): + """Unrestricted electron-electron equation-of-motion coupled cluster.""" - def _quasiparticle_weight(self, r1): + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" return np.linalg.norm(r1.aa) ** 2 + np.linalg.norm(r1.bb) ** 2 - def diag(self, eris=None): + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: + """Get the diagonal of the Hamiltonian. + + Args: + eris: Electronic repulsion integrals. + + Returns: + Diagonal of the Hamiltonian. + """ parts = [] for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -306,175 +380,24 @@ def diag(self, eris=None): return self.amplitudes_to_vector(*parts) - def bras(self, eris=None): # pragma: no cover - raise util.ModelNotImplemented("EE moments for UEBCC not working.") - - bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) - bras = util.Namespace(aa=[], bb=[]) - - for i in range(self.nmo): - for j in range(self.nmo): - amps_aa = [] - amps_bb = [] - - m = 0 - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_aa = util.Namespace() - amp_bb = util.Namespace() - for spin in util.generate_spin_combinations(n): - shape = tuple( - [ - *[self.space["ab".index(s)].ncocc for s in spin[:n]], - *[self.space["ab".index(s)].ncvir for s in spin[n:]], - ] - ) - setattr( - amp_aa, - spin, - getattr( - bras_raw[m], - "aa" + spin, - {(i, j): np.zeros(shape, dtype=types[float])}, - )[i, j], - ) - setattr( - amp_bb, - spin, - getattr( - bras_raw[m], - "bb" + spin, - {(i, j): np.zeros(shape, dtype=types[float])}, - )[i, j], - ) - amps_aa.append(amp_aa) - amps_bb.append(amp_bb) - m += 1 - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks( - spin_type=self.spin_type - ): - raise util.ModelNotImplemented - - bras.aa.append(self.amplitudes_to_vector(*amps_aa)) - bras.bb.append(self.amplitudes_to_vector(*amps_bb)) - - bras.aa = np.array(bras.aa) - bras.bb = np.array(bras.bb) - - bras.aa = bras.aa.reshape(self.nmo, self.nmo, *bras.aa.shape[1:]) - bras.bb = bras.bb.reshape(self.nmo, self.nmo, *bras.bb.shape[1:]) + def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the bra vectors. - return bras + Args: + eris: Electronic repulsion integrals. - def kets(self, eris=None): # pragma: no cover + Returns: + Bra vectors. + """ raise util.ModelNotImplemented("EE moments for UEBCC not working.") - kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) - kets = util.Namespace(aa=[], bb=[]) - - for i in range(self.nmo): - for j in range(self.nmo): - k = (Ellipsis, i, j) - amps_aa = [] - amps_bb = [] - - m = 0 - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_aa = util.Namespace() - amp_bb = util.Namespace() - for spin in util.generate_spin_combinations(n): - shape = tuple( - [ - *[self.space["ab".index(s)].ncocc for s in spin[:n]], - *[self.space["ab".index(s)].ncvir for s in spin[n:]], - ] - ) - setattr( - amp_aa, - spin, - getattr( - kets_raw[m], spin + "aa", {k: np.zeros(shape, dtype=types[float])} - )[k], - ) - setattr( - amp_bb, - spin, - getattr( - kets_raw[m], spin + "bb", {k: np.zeros(shape, dtype=types[float])} - )[k], - ) - amps_aa.append(amp_aa) - amps_bb.append(amp_bb) - m += 1 - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks( - spin_type=self.spin_type - ): - raise util.ModelNotImplemented - - kets.aa.append(self.amplitudes_to_vector(*amps_aa)) - kets.bb.append(self.amplitudes_to_vector(*amps_bb)) - - kets.aa = np.array(kets.aa) - kets.bb = np.array(kets.bb) - - kets.aa = kets.aa.reshape(self.nmo, self.nmo, *kets.aa.shape[1:]) - kets.bb = kets.bb.reshape(self.nmo, self.nmo, *kets.bb.shape[1:]) + def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + """Get the ket vectors. - return kets + Args: + eris: Electronic repulsion integrals. - def moments( - self, - nmom, - eris=None, - amplitudes=None, - hermitise=True, - diagonal_only=True, - ): # pragma: no cover + Returns: + Ket vectors. + """ raise util.ModelNotImplemented("EE moments for UEBCC not working.") - - if not diagonal_only: - warnings.warn( - "Constructing EE moments with `diagonal_only=False` will be very slow.", - stacklevel=2, - ) - - if eris is None: - eris = self.ebcc.get_eris() - if amplitudes is None: - amplitudes = self.ebcc.amplitudes - - bras = self.bras(eris=eris) - kets = self.kets(eris=eris) - - moments = util.Namespace( - aaaa=np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]), - aabb=np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]), - bbaa=np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]), - bbbb=np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]), - ) - - for spin in util.generate_spin_combinations(2): - spin = util.permute_string(spin, (0, 2, 1, 3)) - for k in range(self.nmo): - for l in [k] if diagonal_only else range(self.nmo): - ket = getattr(kets, spin[2:])[k, l] - for n in range(nmom): - for i in range(self.nmo): - for j in [i] if diagonal_only else range(self.nmo): - bra = getattr(bras, spin[:2])[i, j] - getattr(moments, spin)[n, i, j, k, l] = self.dot_braket(bra, ket) - if n != (nmom - 1): - ket = self.matvec(ket, eris=eris) - - if hermitise: - m = getattr(moments, spin) - setattr(moments, spin, 0.5 * (m + m.transpose(0, 3, 4, 1, 2))) - - return moments diff --git a/ebcc/geom.py b/ebcc/geom.py deleted file mode 100644 index 263a6425..00000000 --- a/ebcc/geom.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Generalised equation-of-motion solver.""" - -from ebcc import reom, util - - -class GEOM(reom.REOM): - """Generalised equation-of-motion base class.""" - - @property - def name(self): - return self.excitation_type.upper() + "-GEOM-" + self.ebcc.name - - -class IP_GEOM(GEOM, reom.IP_REOM): - """Generalised equation-of-motion class for ionisation potentials.""" - - pass - - -class EA_GEOM(GEOM, reom.EA_REOM): - """Generalised equation-of-motion class for electron affinities.""" - - pass - - -class EE_GEOM(GEOM, reom.EE_REOM): - """Generalised equation-of-motion class for neutral excitations.""" - - pass diff --git a/ebcc/reom.py b/ebcc/reom.py deleted file mode 100644 index fd34a2ec..00000000 --- a/ebcc/reom.py +++ /dev/null @@ -1,497 +0,0 @@ -"""Restricted equation-of-motion solver.""" - -import dataclasses -import warnings - -from pyscf import lib - -from ebcc import numpy as np -from ebcc import util -from ebcc.logging import ANSI -from ebcc.precision import types -from ebcc.base import EOM - - -@dataclasses.dataclass -class Options: - """ - Options for EOM calculations. - - Attributes - ---------- - nroots : int, optional - Number of roots to solve for. Default value is `5`. - e_tol : float, optional - Threshold for convergence in the correlation energy. Default value - is inherited from `self.ebcc.options`. - max_iter : int, optional - Maximum number of iterations. Default value is inherited from - `self.ebcc.options`. - max_space : int, optional - Maximum size of Lanczos vector space. Default value is `12`. - """ - - nroots: int = 5 - koopmans: bool = False - e_tol: float = util.Inherited - max_iter: int = util.Inherited - max_space: int = 12 - - -class REOM(EOM): - """Restricted equation-of-motion base class.""" - - Options = Options - - def __init__(self, ebcc, options=None, **kwargs): - # Options: - if options is None: - options = self.Options() - self.options = options - for key, val in kwargs.items(): - setattr(self.options, key, val) - for key, val in self.options.__dict__.items(): - if val is util.Inherited: - setattr(self.options, key, getattr(ebcc.options, key)) - - # Parameters: - self.ebcc = ebcc - self.space = ebcc.space - self.ansatz = ebcc.ansatz - self.log = ebcc.log - - # Attributes: - self.converged = False - self.e = None - self.v = None - - # Logging: - self.log.info(f"\n{ANSI.B}{ANSI.U}{self.name}{ANSI.R}") - self.log.debug(f"{ANSI.B}{'*' * len(self.name)}{ANSI.R}") - self.log.debug("") - self.log.info(f"{ANSI.B}Options{ANSI.R}:") - self.log.info(f" > nroots: {ANSI.y}{self.options.nroots}{ANSI.R}") - self.log.info(f" > e_tol: {ANSI.y}{self.options.e_tol}{ANSI.R}") - self.log.info(f" > max_iter: {ANSI.y}{self.options.max_iter}{ANSI.R}") - self.log.info(f" > max_space: {ANSI.y}{self.options.max_space}{ANSI.R}") - self.log.debug("") - - def amplitudes_to_vector(self, *amplitudes): - """Convert the amplitudes to a vector.""" - raise NotImplementedError - - def vector_to_amplitudes(self, vector): - """Convert the vector to amplitudes.""" - raise NotImplementedError - - def matvec(self, vector, eris=None): - """Apply the EOM Hamiltonian to a vector.""" - raise NotImplementedError - - def diag(self, eris=None): - """Find the diagonal of the EOM Hamiltonian.""" - raise NotImplementedError - - def bras(self, eris=None): - """Construct the bra vectors.""" - raise NotImplementedError - - def kets(self, eris=None): - """Construct the ket vectors.""" - raise NotImplementedError - - def dot_braket(self, bra, ket): - """Find the dot-product between the bra and the ket.""" - return np.dot(bra, ket) - - def get_pick(self, guesses=None, real_system=True): - """Pick eigenvalues which match the criterion.""" - - if self.options.koopmans: - assert guesses is not None - g = np.asarray(guesses) - - def pick(w, v, nroots, envs): - x0 = lib.linalg_helper._gen_x0(envs["v"], envs["xs"]) - x0 = np.asarray(x0) - s = np.dot(g.conj(), x0.T) - s = np.einsum("pi,qi->i", s.conj(), s) - idx = np.argsort(-s)[:nroots] - return lib.linalg_helper._eigs_cmplx2real(w, v, idx, real_system) - - else: - - def pick(w, v, nroots, envs): - real_idx = np.where(abs(w.imag) < 1e-3)[0] - w, v, idx = lib.linalg_helper._eigs_cmplx2real(w, v, real_idx, real_system) - mask = np.argsort(np.abs(w)) - w, v = w[mask], v[:, mask] - return w, v, 0 - - return pick - - def _argsort_guess(self, diag): - if self.options.koopmans: - r_mf = self.vector_to_amplitudes(diag)[0] - size = r_mf.size - arg = np.argsort(np.diag(diag[:size])) - else: - arg = np.argsort(diag) - - return arg - - def get_guesses(self, diag=None): - """Generate guess vectors.""" - - if diag is None: - diag = self.diag() - - arg = self._argsort_guess(diag) - - nroots = min(self.options.nroots, diag.size) - guesses = np.zeros((nroots, diag.size), dtype=diag.dtype) - for root, guess in enumerate(arg[:nroots]): - guesses[root, guess] = 1.0 - - return list(guesses) - - def callback(self, envs): - """Callback function for the Davidson solver.""" # noqa: D401 - - pass - - def _quasiparticle_weight(self, r1): - return np.linalg.norm(r1) ** 2 - - def davidson(self, eris=None, guesses=None): - """Solve the EOM Hamiltonian using the Davidson solver. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.ebcc.get_eris()`. - guesses : list of np.ndarray, optional - Initial guesses for the roots. Default value is generated - using `self.get_guesses()`. - - Returns - ------- - e : np.ndarray - The energies of the roots. - """ - - # Start a timer: - timer = util.Timer() - - # Get the ERIs: - eris = self.ebcc.get_eris(eris) - - self.log.output( - "Solving for %s excitations using the Davidson solver.", self.excitation_type.upper() - ) - - # Get the matrix-vector products and the diagonal: - matvecs = lambda vs: [self.matvec(v, eris=eris) for v in vs] - diag = self.diag(eris=eris) - - # Get the guesses: - if guesses is None: - guesses = self.get_guesses(diag=diag) - - # Solve the EOM Hamiltonian: - nroots = min(len(guesses), self.options.nroots) - pick = self.get_pick(guesses=guesses) - converged, e, v = lib.davidson_nosym1( - matvecs, - guesses, - diag, - tol=self.options.e_tol, - nroots=nroots, - pick=pick, - max_cycle=self.options.max_iter, - max_space=self.options.max_space, - callback=self.callback, - verbose=0, - ) - - # Check for convergence: - if all(converged): - self.log.debug("") - self.log.output(f"{ANSI.g}Converged.{ANSI.R}") - else: - self.log.debug("") - self.log.warning( - f"{ANSI.r}Failed to converge {sum(not c for c in converged)} roots.{ANSI.R}" - ) - - # Update attributes: - self.converged = converged - self.e = e - self.v = v - - self.log.debug("") - self.log.output( - f"{ANSI.B}{'Root':>4s} {'Energy':>16s} {'Weight':>13s} {'Conv.':>8s}{ANSI.R}" - ) - for n, (en, vn, cn) in enumerate(zip(e, v, converged)): - r1n = self.vector_to_amplitudes(vn)[0] - qpwt = self._quasiparticle_weight(r1n) - self.log.output( - f"{n:>4d} {en:>16.10f} {qpwt:>13.5g} " f"{[ANSI.r, ANSI.g][cn]}{cn!r:>8s}{ANSI.R}" - ) - - self.log.debug("") - self.log.debug("Time elapsed: %s", timer.format_time(timer())) - self.log.debug("") - - return e - - kernel = davidson - - def moments(self, nmom, eris=None, amplitudes=None, hermitise=True): - """Construct the moments of the EOM Hamiltonian.""" - - if eris is None: - eris = self.ebcc.get_eris() - if amplitudes is None: - amplitudes = self.ebcc.amplitudes - - bras = self.bras(eris=eris) - kets = self.kets(eris=eris) - - moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) - - for j in range(self.nmo): - ket = kets[j] - for n in range(nmom): - for i in range(self.nmo): - bra = bras[i] - moments[n, i, j] = self.dot_braket(bra, ket) - if n != (nmom - 1): - ket = self.matvec(ket, eris=eris) - - if hermitise: - moments = 0.5 * (moments + moments.swapaxes(1, 2)) - - return moments - - @property - def name(self): - """Get a string representation of the method name.""" - return f"{self.excitation_type.upper()}-{self.spin_type}EOM-{self.ebcc.name}" - - @property - def spin_type(self): - """Get a string representation of the spin channel.""" - return self.ebcc.spin_type - - @property - def excitation_type(self): - """Get a string representation of the excitation type.""" - raise NotImplementedError - - @property - def nmo(self): - """Get the number of MOs.""" - return self.ebcc.nmo - - @property - def nocc(self): - """Get the number of occupied MOs.""" - return self.ebcc.nocc - - @property - def nvir(self): - """Get the number of virtual MOs.""" - return self.ebcc.nvir - - -class IP_REOM(REOM): - """Restricted equation-of-motion class for ionisation potentials.""" - - def amplitudes_to_vector(self, *amplitudes): - return self.ebcc.excitations_to_vector_ip(*amplitudes) - - def vector_to_amplitudes(self, vector): - return self.ebcc.vector_to_excitations_ip(vector) - - def matvec(self, vector, eris=None): - amplitudes = self.vector_to_amplitudes(vector) - amplitudes = self.ebcc.hbar_matvec_ip(*amplitudes, eris=eris) - return self.amplitudes_to_vector(*amplitudes) - - def diag(self, eris=None): - parts = [] - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[:-1] - parts.append(self.ebcc.energy_sum(key)) - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return self.amplitudes_to_vector(*parts) - - def _bras(self, eris=None): - return self.ebcc.make_ip_mom_bras(eris=eris) - - def _kets(self, eris=None): - return self.ebcc.make_ip_mom_kets(eris=eris) - - def bras(self, eris=None): - bras_raw = list(self._bras(eris=eris)) - bras = np.array( - [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] - ) - return bras - - def kets(self, eris=None): - kets_raw = list(self._kets(eris=eris)) - kets = np.array( - [self.amplitudes_to_vector(*[b[..., i] for b in kets_raw]) for i in range(self.nmo)] - ) - return kets - - @property - def excitation_type(self): - return "ip" - - -class EA_REOM(REOM): - """Equation-of-motion class for electron affinities.""" - - def amplitudes_to_vector(self, *amplitudes): - return self.ebcc.excitations_to_vector_ea(*amplitudes) - - def vector_to_amplitudes(self, vector): - return self.ebcc.vector_to_excitations_ea(vector) - - def matvec(self, vector, eris=None): - amplitudes = self.vector_to_amplitudes(vector) - amplitudes = self.ebcc.hbar_matvec_ea(*amplitudes, eris=eris) - return self.amplitudes_to_vector(*amplitudes) - - def diag(self, eris=None): - parts = [] - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[n:] + key[: n - 1] - parts.append(-self.ebcc.energy_sum(key)) - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return self.amplitudes_to_vector(*parts) - - def _bras(self, eris=None): - return self.ebcc.make_ea_mom_bras(eris=eris) - - def _kets(self, eris=None): - return self.ebcc.make_ea_mom_kets(eris=eris) - - def bras(self, eris=None): - bras_raw = list(self._bras(eris=eris)) - bras = np.array( - [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] - ) - return bras - - def kets(self, eris=None): - kets_raw = list(self._kets(eris=eris)) - kets = np.array( - [self.amplitudes_to_vector(*[b[..., i] for b in kets_raw]) for i in range(self.nmo)] - ) - return kets - - @property - def excitation_type(self): - return "ea" - - -class EE_REOM(REOM): - """Equation-of-motion class for neutral excitations.""" - - def amplitudes_to_vector(self, *amplitudes): - return self.ebcc.excitations_to_vector_ee(*amplitudes) - - def vector_to_amplitudes(self, vector): - return self.ebcc.vector_to_excitations_ee(vector) - - def matvec(self, vector, eris=None): - amplitudes = self.vector_to_amplitudes(vector) - amplitudes = self.ebcc.hbar_matvec_ee(*amplitudes, eris=eris) - return self.amplitudes_to_vector(*amplitudes) - - def diag(self, eris=None): - parts = [] - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - parts.append(-self.ebcc.energy_sum(key)) - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return self.amplitudes_to_vector(*parts) - - def bras(self, eris=None): - bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) - bras = np.array( - [ - [self.amplitudes_to_vector(*[b[i, j] for b in bras_raw]) for j in range(self.nmo)] - for i in range(self.nmo) - ] - ) - return bras - - def kets(self, eris=None): - kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) - kets = np.array( - [ - [ - self.amplitudes_to_vector(*[k[..., i, j] for k in kets_raw]) - for j in range(self.nmo) - ] - for i in range(self.nmo) - ] - ) - return kets - - def moments(self, nmom, eris=None, amplitudes=None, hermitise=True, diagonal_only=True): - """Construct the moments of the EOM Hamiltonian.""" - - if not diagonal_only: - warnings.warn( - "Constructing EE moments with `diagonal_only=False` will be very slow.", - stacklevel=2, - ) - - if eris is None: - eris = self.ebcc.get_eris() - if amplitudes is None: - amplitudes = self.ebcc.amplitudes - - bras = self.bras(eris=eris) - kets = self.kets(eris=eris) - - moments = np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]) - - for k in range(self.nmo): - for l in [k] if diagonal_only else range(self.nmo): - ket = kets[k, l] - for n in range(nmom): - for i in range(self.nmo): - for j in [i] if diagonal_only else range(self.nmo): - bra = bras[i, j] - moments[n, i, j, k, l] = self.dot_braket(bra, ket) - if n != (nmom - 1): - ket = self.matvec(ket, eris=eris) - - if hermitise: - moments = 0.5 * (moments + moments.transpose(0, 3, 4, 1, 2)) - - return moments - - @property - def excitation_type(self): - return "ee" diff --git a/ebcc/util/permutations.py b/ebcc/util/permutations.py index ed02e2b3..8a7a03e0 100644 --- a/ebcc/util/permutations.py +++ b/ebcc/util/permutations.py @@ -41,7 +41,9 @@ def permute_string(string: str, permutation: tuple[int, ...]) -> str: return "".join([string[i] for i in permutation]) -def tril_indices_ndim(n: int, dims: int, include_diagonal: Optional[bool] = False) -> tuple[NDArray[int]]: +def tril_indices_ndim( + n: int, dims: int, include_diagonal: Optional[bool] = False +) -> tuple[NDArray[int]]: """Return lower triangular indices for a multidimensional array. Args: @@ -55,7 +57,7 @@ def tril_indices_ndim(n: int, dims: int, include_diagonal: Optional[bool] = Fals ranges = [np.arange(n)] * dims if dims == 1: return (ranges[0],) - #func: Callable[[Any, ...], Any] = np.greater_equal if include_diagonal else np.greater + # func: Callable[[Any, ...], Any] = np.greater_equal if include_diagonal else np.greater slices = [ tuple(slice(None) if i == j else np.newaxis for i in range(dims)) for j in range(dims) @@ -95,7 +97,9 @@ def ntril_ndim(n: int, dims: int, include_diagonal: Optional[bool] = False) -> i return out -def generate_spin_combinations(n: int, excited: Optional[bool] = False, unique: Optional[bool] = False) -> Generator[str, None, None]: +def generate_spin_combinations( + n: int, excited: Optional[bool] = False, unique: Optional[bool] = False +) -> Generator[str, None, None]: """Generate combinations of spin components for a given number of occupied and virtual axes. Args: @@ -274,7 +278,9 @@ def combine_subscripts( return new_subscript, new_sizes -def compress_axes(subscript: str, array: NDArray[T], include_diagonal: Optional[bool] = False) -> NDArray[T]: +def compress_axes( + subscript: str, array: NDArray[T], include_diagonal: Optional[bool] = False +) -> NDArray[T]: """Compress an array into lower-triangular representations using an einsum-like input. Args: @@ -467,7 +473,12 @@ def get_compressed_size(subscript: str, **sizes: int) -> int: return n -def symmetrise(subscript: str, array: NDArray[T], symmetry: Optional[str] = None, apply_factor: Optional[bool] = True) -> NDArray[T]: +def symmetrise( + subscript: str, + array: NDArray[T], + symmetry: Optional[str] = None, + apply_factor: Optional[bool] = True, +) -> NDArray[T]: """Enforce a symmetry in an array. Args: From da00ac06298ceb1fdbdc54a5700a8a279bb7ba62 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sat, 3 Aug 2024 12:40:17 +0100 Subject: [PATCH 11/37] Fixes for EOM --- ebcc/eom/base.py | 22 ++++-- ebcc/eom/geom.py | 191 +++++++++++++++++++++++++++++++++++++---------- ebcc/eom/reom.py | 185 ++++++++++++++++++++++++++++++++++++--------- ebcc/eom/ueom.py | 157 ++++++++++++++++++++++++++++---------- 4 files changed, 435 insertions(+), 120 deletions(-) diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py index 25c1a71a..e7d23261 100644 --- a/ebcc/eom/base.py +++ b/ebcc/eom/base.py @@ -14,7 +14,7 @@ from ebcc.precision import types if TYPE_CHECKING: - from typing import Any, Callable, Optional, TypeVar, Union + from typing import Any, Callable, Optional from ebcc.cc.base import AmplitudeType, BaseEBCC, ERIsInputType from ebcc.numpy.typing import NDArray @@ -206,16 +206,16 @@ def get_pick( Function to pick the eigenvalues. """ if self.options.koopmans: - assert guesses is None + assert guesses is not None guesses_array = np.asarray(guesses) def pick( w: NDArray[float], v: NDArray[float], nroots: int, env: dict[str, Any] ) -> tuple[NDArray[float], NDArray[float], int]: """Pick the eigenvalues.""" - x0 = lib.linalg_helper._gen_x0(envs["v"], envs["xs"]) + x0 = lib.linalg_helper._gen_x0(env["v"], env["xs"]) x0 = np.asarray(x0) - s = np.dot(g.conj(), x0.T) + s = np.dot(guesses_array.conj(), x0.T) s = util.einsum("pi,qi->i", s.conj(), s) idx = np.argsort(-s)[:nroots] return lib.linalg_helper._eigs_cmplx2real(w, v, idx, real_system) @@ -354,8 +354,18 @@ def moments( eris: Optional[ERIsInputType] = None, amplitudes: Namespace[AmplitudeType] = None, hermitise: bool = True, - ) -> NDArray[float]: - """Construct the moments of the EOM Hamiltonian.""" + ) -> AmplitudeType: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + + Returns: + Moments of the EOM Hamiltonian. + """ pass @property diff --git a/ebcc/eom/geom.py b/ebcc/eom/geom.py index a3132649..c9e65e85 100644 --- a/ebcc/eom/geom.py +++ b/ebcc/eom/geom.py @@ -6,11 +6,11 @@ from ebcc import numpy as np from ebcc import util -from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM, Options +from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM from ebcc.precision import types if TYPE_CHECKING: - from typing import Optional, Tuple, Union + from typing import Optional from ebcc.cc.gebcc import AmplitudeType, ERIsInputType from ebcc.numpy.typing import NDArray @@ -33,40 +33,8 @@ def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" return np.linalg.norm(r1) ** 2 - def moments( - self, - nmom: int, - eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, - hermitise: bool = True, - ) -> NDArray[float]: - """Construct the moments of the EOM Hamiltonian.""" - if eris is None: - eris = self.ebcc.get_eris() - if not amplitudes: - amplitudes = self.ebcc.amplitudes - - bras = self.bras(eris=eris) - kets = self.kets(eris=eris) - - moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) - - for j in range(self.nmo): - ket = kets[j] - for n in range(nmom): - for i in range(self.nmo): - bra = bras[i] - moments[n, i, j] = self.dot_braket(bra, ket) - if n != (nmom - 1): - ket = self.matvec(ket, eris=eris) - if hermitise: - moments = 0.5 * (moments + moments.swapaxes(1, 2)) - - return moments - - -class IP_GEOM(BaseIP_EOM, BaseEOM): +class IP_GEOM(GEOM, BaseIP_EOM): """Generalised ionisation potential equation-of-motion coupled cluster.""" def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: @@ -119,8 +87,50 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return kets + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + + Returns: + Moments of the EOM Hamiltonian. + """ + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) + + for j in range(self.nmo): + ket = kets[j] + for n in range(nmom): + for i in range(self.nmo): + bra = bras[i] + moments[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.swapaxes(1, 2)) + + return moments + -class EA_GEOM(BaseEA_EOM, BaseEOM): +class EA_GEOM(GEOM, BaseEA_EOM): """Generalised electron affinity equation-of-motion coupled cluster.""" def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: @@ -173,8 +183,50 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return kets + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + + Returns: + Moments of the EOM Hamiltonian. + """ + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) -class EE_GEOM(BaseEE_EOM, BaseEOM): + moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) + + for j in range(self.nmo): + ket = kets[j] + for n in range(nmom): + for i in range(self.nmo): + bra = bras[i] + moments[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.swapaxes(1, 2)) + + return moments + + +class EE_GEOM(GEOM, BaseEE_EOM): """Generalised electron-electron equation-of-motion coupled cluster.""" def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: @@ -207,7 +259,10 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: """ bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) bras = np.array( - [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + [ + [self.amplitudes_to_vector(*[b[i, j] for b in bras_raw]) for j in range(self.nmo)] + for i in range(self.nmo) + ] ) return bras @@ -222,6 +277,64 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: """ kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) kets = np.array( - [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + [ + [ + self.amplitudes_to_vector(*[k[..., i, j] for k in kets_raw]) + for j in range(self.nmo) + ] + for i in range(self.nmo) + ] ) return kets + + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + diagonal_only: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + diagonal_only: Only compute the diagonal elements. + + Returns: + Moments of the EOM Hamiltonian. + """ + if not diagonal_only: + warnings.warn( + "Constructing EE moments with `diagonal_only=False` will be very slow.", + stacklevel=2, + ) + + if eris is None: + eris = self.ebcc.get_eris() + if amplitudes is None: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]) + + for k in range(self.nmo): + for l in [k] if diagonal_only else range(self.nmo): + ket = kets[k, l] + for n in range(nmom): + for i in range(self.nmo): + for j in [i] if diagonal_only else range(self.nmo): + bra = bras[i, j] + moments[n, i, j, k, l] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.transpose(0, 3, 4, 1, 2)) + + return moments diff --git a/ebcc/eom/reom.py b/ebcc/eom/reom.py index 8547e13c..450480e6 100644 --- a/ebcc/eom/reom.py +++ b/ebcc/eom/reom.py @@ -6,11 +6,11 @@ from ebcc import numpy as np from ebcc import util -from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM, Options +from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM from ebcc.precision import types if TYPE_CHECKING: - from typing import Optional, Tuple, Union + from typing import Optional from ebcc.cc.rebcc import AmplitudeType, ERIsInputType from ebcc.numpy.typing import NDArray @@ -33,38 +33,6 @@ def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" return np.linalg.norm(r1) ** 2 - def moments( - self, - nmom: int, - eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, - hermitise: bool = True, - ) -> NDArray[float]: - """Construct the moments of the EOM Hamiltonian.""" - if eris is None: - eris = self.ebcc.get_eris() - if not amplitudes: - amplitudes = self.ebcc.amplitudes - - bras = self.bras(eris=eris) - kets = self.kets(eris=eris) - - moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) - - for j in range(self.nmo): - ket = kets[j] - for n in range(nmom): - for i in range(self.nmo): - bra = bras[i] - moments[n, i, j] = self.dot_braket(bra, ket) - if n != (nmom - 1): - ket = self.matvec(ket, eris=eris) - - if hermitise: - moments = 0.5 * (moments + moments.swapaxes(1, 2)) - - return moments - class IP_REOM(REOM, BaseIP_EOM): """Restricted ionisation potential equation-of-motion coupled cluster.""" @@ -119,6 +87,48 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return kets + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> AmplitudeType: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + + Returns: + Moments of the EOM Hamiltonian. + """ + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) + + for j in range(self.nmo): + ket = kets[j] + for n in range(nmom): + for i in range(self.nmo): + bra = bras[i] + moments[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.swapaxes(1, 2)) + + return moments + class EA_REOM(REOM, BaseEA_EOM): """Restricted electron affinity equation-of-motion coupled cluster.""" @@ -173,6 +183,48 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return kets + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> AmplitudeType: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + + Returns: + Moments of the EOM Hamiltonian. + """ + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) + + for j in range(self.nmo): + ket = kets[j] + for n in range(nmom): + for i in range(self.nmo): + bra = bras[i] + moments[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.swapaxes(1, 2)) + + return moments + class EE_REOM(REOM, BaseEE_EOM): """Restricted electron-electron equation-of-motion coupled cluster.""" @@ -207,7 +259,10 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: """ bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) bras = np.array( - [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] + [ + [self.amplitudes_to_vector(*[b[i, j] for b in bras_raw]) for j in range(self.nmo)] + for i in range(self.nmo) + ] ) return bras @@ -222,6 +277,64 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: """ kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) kets = np.array( - [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] + [ + [ + self.amplitudes_to_vector(*[k[..., i, j] for k in kets_raw]) + for j in range(self.nmo) + ] + for i in range(self.nmo) + ] ) return kets + + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + diagonal_only: bool = True, + ) -> AmplitudeType: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + diagonal_only: Only compute the diagonal elements. + + Returns: + Moments of the EOM Hamiltonian. + """ + if not diagonal_only: + warnings.warn( + "Constructing EE moments with `diagonal_only=False` will be very slow.", + stacklevel=2, + ) + + if eris is None: + eris = self.ebcc.get_eris() + if amplitudes is None: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]) + + for k in range(self.nmo): + for l in [k] if diagonal_only else range(self.nmo): + ket = kets[k, l] + for n in range(nmom): + for i in range(self.nmo): + for j in [i] if diagonal_only else range(self.nmo): + bra = bras[i, j] + moments[n, i, j, k, l] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + moments = 0.5 * (moments + moments.transpose(0, 3, 4, 1, 2)) + + return moments diff --git a/ebcc/eom/ueom.py b/ebcc/eom/ueom.py index 861367a2..3c6ce157 100644 --- a/ebcc/eom/ueom.py +++ b/ebcc/eom/ueom.py @@ -6,11 +6,11 @@ from ebcc import numpy as np from ebcc import util -from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM, Options +from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM from ebcc.precision import types if TYPE_CHECKING: - from typing import Optional, Tuple, Union + from typing import Optional from ebcc.cc.base import AmplitudeType, ERIsInputType from ebcc.numpy.typing import NDArray @@ -33,43 +33,6 @@ def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" return np.linalg.norm(r1.a) ** 2 + np.linalg.norm(r1.b) ** 2 - def moments( - self, - nmom: int, - eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, - hermitise: bool = True, - ) -> NDArray[float]: - """Construct the moments of the EOM Hamiltonian.""" - if eris is None: - eris = self.ebcc.get_eris() - if not amplitudes: - amplitudes = self.ebcc.amplitudes - - bras = self.bras(eris=eris) - kets = self.kets(eris=eris) - - moments = util.Namespace( - aa=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), - bb=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), - ) - - for spin in util.generate_spin_combinations(1): - for j in range(self.nmo): - ket = getattr(kets, spin[1])[j] - for n in range(nmom): - for i in range(self.nmo): - bra = getattr(bras, spin[0])[i] - getattr(moments, spin)[n, i, j] = self.dot_braket(bra, ket) - if n != (nmom - 1): - ket = self.matvec(ket, eris=eris) - - if hermitise: - t = getattr(moments, spin) - setattr(moments, spin, 0.5 * (t + t.swapaxes(1, 2))) - - return moments - class IP_UEOM(UEOM, BaseIP_EOM): """Unrestricted ionisation potential equation-of-motion coupled cluster.""" @@ -210,6 +173,53 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: return kets + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + + Returns: + Moments of the EOM Hamiltonian. + """ + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = util.Namespace( + aa=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), + bb=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), + ) + + for spin in util.generate_spin_combinations(1): + for j in range(self.nmo): + ket = getattr(kets, spin[1])[j] + for n in range(nmom): + for i in range(self.nmo): + bra = getattr(bras, spin[0])[i] + getattr(moments, spin)[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + t = getattr(moments, spin) + setattr(moments, spin, 0.5 * (t + t.swapaxes(1, 2))) + + return moments + class EA_UEOM(UEOM, BaseEA_EOM): """Unrestricted electron affinity equation-of-motion coupled cluster.""" @@ -350,6 +360,53 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: return kets + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + + Returns: + Moments of the EOM Hamiltonian. + """ + if eris is None: + eris = self.ebcc.get_eris() + if not amplitudes: + amplitudes = self.ebcc.amplitudes + + bras = self.bras(eris=eris) + kets = self.kets(eris=eris) + + moments = util.Namespace( + aa=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), + bb=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), + ) + + for spin in util.generate_spin_combinations(1): + for j in range(self.nmo): + ket = getattr(kets, spin[1])[j] + for n in range(nmom): + for i in range(self.nmo): + bra = getattr(bras, spin[0])[i] + getattr(moments, spin)[n, i, j] = self.dot_braket(bra, ket) + if n != (nmom - 1): + ket = self.matvec(ket, eris=eris) + + if hermitise: + t = getattr(moments, spin) + setattr(moments, spin, 0.5 * (t + t.swapaxes(1, 2))) + + return moments + class EE_UEOM(UEOM, BaseEE_EOM): """Unrestricted electron-electron equation-of-motion coupled cluster.""" @@ -401,3 +458,25 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: Ket vectors. """ raise util.ModelNotImplemented("EE moments for UEBCC not working.") + + def moments( + self, + nmom: int, + eris: Optional[ERIsInputType] = None, + amplitudes: Namespace[AmplitudeType] = None, + hermitise: bool = True, + diagonal_only: bool = True, + ) -> NDArray[float]: + """Construct the moments of the EOM Hamiltonian. + + Args: + nmom: Number of moments. + eris: Electronic repulsion integrals. + amplitudes: Cluster amplitudes. + hermitise: Hermitise the moments. + diagonal_only: Only compute the diagonal elements. + + Returns: + Moments of the EOM Hamiltonian. + """ + raise util.ModelNotImplemented("EE moments for UEBCC not working.") From 4906c3c5e4c944d2a8109516a05bb2460f56bad8 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sat, 3 Aug 2024 15:33:12 +0100 Subject: [PATCH 12/37] Cleans up Brueckner methods --- ebcc/base.py | 25 +- ebcc/brueckner.py | 488 ---------------------------------------- ebcc/cc/base.py | 11 +- ebcc/cc/gebcc.py | 6 +- ebcc/cc/rebcc.py | 8 +- ebcc/cc/uebcc.py | 6 +- ebcc/eom/base.py | 6 +- ebcc/eom/geom.py | 42 +++- ebcc/eom/reom.py | 42 +++- ebcc/eom/ueom.py | 36 ++- ebcc/opt/__init__.py | 5 + ebcc/opt/base.py | 286 +++++++++++++++++++++++ ebcc/opt/gbrueckner.py | 165 ++++++++++++++ ebcc/opt/rbrueckner.py | 165 ++++++++++++++ ebcc/opt/ubrueckner.py | 200 ++++++++++++++++ ebcc/util/einsumfunc.py | 2 + pyproject.toml | 12 +- tests/test_UCCSD.py | 97 ++++---- tests/test_util.py | 73 +++--- 19 files changed, 1045 insertions(+), 630 deletions(-) delete mode 100644 ebcc/brueckner.py create mode 100644 ebcc/opt/__init__.py create mode 100644 ebcc/opt/base.py create mode 100644 ebcc/opt/gbrueckner.py create mode 100644 ebcc/opt/rbrueckner.py create mode 100644 ebcc/opt/ubrueckner.py diff --git a/ebcc/base.py b/ebcc/base.py index 0b581851..e8e64fe6 100644 --- a/ebcc/base.py +++ b/ebcc/base.py @@ -2,32 +2,11 @@ from __future__ import annotations -from abc import ABC, abstractmethod +from abc import ABC from typing import TYPE_CHECKING -from ebcc import util - if TYPE_CHECKING: - from dataclasses import dataclass - from logging import Logger - from typing import Any, Optional, Union - - from pyscf.scf import SCF - - from ebcc.ansatz import Ansatz - from ebcc.space import Space - - -class EOM(ABC): - """Base class for equation-of-motion methods.""" - - pass - - -class BruecknerEBCC(ABC): - """Base class for Brueckner orbital methods.""" - - pass + from typing import Any class ERIs(ABC): diff --git a/ebcc/brueckner.py b/ebcc/brueckner.py deleted file mode 100644 index 7f232f11..00000000 --- a/ebcc/brueckner.py +++ /dev/null @@ -1,488 +0,0 @@ -"""Brueckner orbital self-consistency.""" - -import dataclasses - -import scipy.linalg -from pyscf import lib - -from ebcc import NullLogger, init_logging -from ebcc import numpy as np -from ebcc import util -from ebcc.damping import DIIS -from ebcc.logging import ANSI -from ebcc.precision import types -from ebcc.base import BruecknerEBCC - - -@dataclasses.dataclass -class Options: - """ - Options for Brueckner CC calculations. - - Attributes - ---------- - e_tol : float, optional - Threshold for convergence in the correlation energy. Default value - is `1e-8`. - t_tol : float, optional - Threshold for convergence in the amplitude norm. Default value is - `1e-8`. - max_iter : int, optional - Maximum number of iterations. Default value is `20`. - diis_space : int, optional - Number of amplitudes to use in DIIS extrapolation. Default value is - `12`. - damping : float, optional - Damping factor for DIIS extrapolation. Default value is `0.0`. - """ - - e_tol: float = 1e-8 - t_tol: float = 1e-8 - max_iter: int = 20 - diis_space: int = 12 - damping: float = 0.0 - - -class BruecknerREBCC(BruecknerEBCC): - """ - Brueckner orbital self-consistency for coupled cluster calculations. - Iteratively solve for a new mean-field that presents a vanishing T1 - under the given ansatz. - - Parameters - ---------- - cc : EBCC - EBCC coupled cluster object. - log : logging.Logger, optional - Log to print output to. Default value is `cc.log`. - options : dataclasses.dataclass, optional - Object containing the options. Default value is `Options`. - **kwargs : dict - Additional keyword arguments used to update `options`. - """ - - Options = Options - - def __init__(self, cc, log=None, options=None, **kwargs): - # Options: - if options is None: - options = self.Options() - self.options = options - for key, val in kwargs.items(): - setattr(self.options, key, val) - - # Parameters: - self.log = cc.log if log is None else log - self.mf = cc.mf - self.cc = cc - - # Attributes: - self.converged = False - - # Logging: - init_logging(cc.log) - cc.log.info(f"\n{ANSI.B}{ANSI.U}{self.name}{ANSI.R}") - cc.log.info(f"{ANSI.B}Options{ANSI.R}:") - cc.log.info(f" > e_tol: {ANSI.y}{self.options.e_tol}{ANSI.R}") - cc.log.info(f" > t_tol: {ANSI.y}{self.options.t_tol}{ANSI.R}") - cc.log.info(f" > max_iter: {ANSI.y}{self.options.max_iter}{ANSI.R}") - cc.log.info(f" > diis_space: {ANSI.y}{self.options.diis_space}{ANSI.R}") - cc.log.info(f" > damping: {ANSI.y}{self.options.damping}{ANSI.R}") - cc.log.debug("") - - def get_rotation_matrix(self, u_tot=None, diis=None, t1=None): - """ - Update the rotation matrix, and also return the total rotation - matrix. - - Parameters - ---------- - u_tot : np.ndarray, optional - Total rotation matrix. If `None`, then it is assumed to be the - identity matrix. Default value is `None`. - diis : DIIS, optional - DIIS object. If `None`, then DIIS is not used. Default value is - `None`. - t1 : np.ndarray, optional - T1 amplitudes. If `None`, then `cc.t1` is used. Default value - is `None`. - - Returns - ------- - u : np.ndarray - Rotation matrix. - u_tot : np.ndarray - Total rotation matrix. - """ - - if t1 is None: - t1 = self.cc.t1 - if u_tot is None: - u_tot = np.eye(self.cc.space.ncorr, dtype=types[float]) - - t1_block = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) - t1_block[: self.cc.space.ncocc, self.cc.space.ncocc :] = -t1 - t1_block[self.cc.space.ncocc :, : self.cc.space.ncocc] = t1.T - - u = scipy.linalg.expm(t1_block) - - u_tot = np.dot(u_tot, u) - if scipy.linalg.det(u_tot) < 0: - u_tot[:, 0] *= -1 - - a = scipy.linalg.logm(u_tot) - a = a.real.astype(types[float]) - if diis is not None: - a = diis.update(a, xerr=t1) - - u_tot = scipy.linalg.expm(a) - - return u, u_tot - - def transform_amplitudes(self, u, amplitudes=None): - """ - Transform the amplitudes into the Brueckner orbital basis. - - Parameters - ---------- - u : np.ndarray - Rotation matrix. - amplitudes : Namespace, optional - Amplitudes. If `None`, then `cc.amplitudes` is used. Default - value is `None`. - - Returns - ------- - amplitudes : Namespace - Rotated amplitudes. - """ - - if amplitudes is None: - amplitudes = self.cc.amplitudes - - nocc = self.cc.space.ncocc - ci = u[:nocc, :nocc] - ca = u[nocc:, nocc:] - - # Transform T amplitudes: - for name, key, n in self.cc.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - args = [self.cc.amplitudes[name], tuple(range(n * 2))] - for i in range(n): - args += [ci, (i, i + n * 2)] - for i in range(n): - args += [ca, (i + n, i + n * 3)] - args += [tuple(range(n * 2, n * 4))] - self.cc.amplitudes[name] = util.einsum(*args) - - # Transform S amplitudes: - for name, key, n in self.cc.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented # TODO - - # Transform U amplitudes: - for name, key, nf, nb in self.cc.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented # TODO - - return self.cc.amplitudes - - def get_t1_norm(self, amplitudes=None): - """ - Get the norm of the T1 amplitude. - - Parameters - ---------- - amplitudes : Namespace, optional - Amplitudes. If `None`, then `cc.amplitudes` is used. Default - value is `None`. - - Returns - ------- - t1_norm : float - Norm of the T1 amplitude. - """ - - if amplitudes is None: - amplitudes = self.cc.amplitudes - - return np.linalg.norm(amplitudes["t1"]) - - def mo_to_correlated(self, mo_coeff): - """ - For a given set of MO coefficients, return the correlated slice. - - Parameters - ---------- - mo_coeff : np.ndarray - MO coefficients. - - Returns - ------- - mo_coeff_corr : np.ndarray - Correlated slice of the MO coefficients. - """ - - return mo_coeff[:, self.cc.space.correlated] - - def mo_update_correlated(self, mo_coeff, mo_coeff_corr): - """ - Update the correlated slice of a set of MO coefficients. - - Parameters - ---------- - mo_coeff : np.ndarray - MO coefficients. - mo_coeff_corr : np.ndarray - Correlated slice of the MO coefficients. - - Returns - ------- - mo_coeff : np.ndarray - Updated MO coefficients. - """ - - mo_coeff[:, self.cc.space.correlated] = mo_coeff_corr - - return mo_coeff - - def update_coefficients(self, u_tot, mo_coeff, mo_coeff_ref): - """Get the updated coefficients. - - Parameters - ---------- - u_tot : np.ndarray - Total rotation matrix. - mo_coeff : np.ndarray - New MO coefficients. - mo_coeff_ref : np.ndarray - Reference MO coefficients. - - Returns - ------- - mo_coeff_new : np.ndarray - Updated MO coefficients. - u : np.ndarray - Rotation matrix. - """ - mo_coeff_new_corr = util.einsum("pi,ij->pj", mo_coeff_ref, u_tot) - mo_coeff_new = self.mo_update_correlated(mo_coeff, mo_coeff_new_corr) - return mo_coeff_new - - def kernel(self): - """ - Run the Bruckner orbital coupled cluster calculation. - - Returns - ------- - e_cc : float - Correlation energy. - """ - - # Start a timer: - timer = util.Timer() - - # Make sure the initial CC calculation is converged: - if not self.cc.converged: - with lib.temporary_env(self.cc, log=NullLogger()): - self.cc.kernel() - - # Set up DIIS: - diis = DIIS() - diis.space = self.options.diis_space - diis.damping = self.options.damping - - # Initialise coefficients: - mo_coeff_new = np.array(self.cc.mo_coeff, copy=True, dtype=types[float]) - mo_coeff_ref = np.array(self.cc.mo_coeff, copy=True, dtype=types[float]) - mo_coeff_ref = self.mo_to_correlated(mo_coeff_ref) - u_tot = None - - self.cc.log.output("Solving for Brueckner orbitals.") - self.cc.log.debug("") - self.log.info( - f"{ANSI.B}{'Iter':>4s} {'Energy (corr.)':>16s} {'Energy (tot.)':>18s} " - f"{'Conv.':>8s} {'Δ(Energy)':>13s} {'|T1|':>13s}{ANSI.R}" - ) - self.log.info( - f"{0:4d} {self.cc.e_corr:16.10f} {self.cc.e_tot:18.10f} " - f"{[ANSI.r, ANSI.g][self.cc.converged]}{self.cc.converged!r:>8}{ANSI.R}" - ) - - converged = False - for niter in range(1, self.options.max_iter + 1): - # Update rotation matrix: - u, u_tot = self.get_rotation_matrix(u_tot=u_tot, diis=diis) - - # Update MO coefficients: - mo_coeff_new = self.update_coefficients(u_tot, mo_coeff_new, mo_coeff_ref) - - # Transform mean-field and amplitudes: - self.mf.mo_coeff = mo_coeff_new - self.mf.e_tot = self.mf.energy_tot() - amplitudes = self.transform_amplitudes(u) - - # Run CC calculation: - e_prev = self.cc.e_tot - with lib.temporary_env(self.cc, log=NullLogger()): - self.cc.__init__( - self.mf, - log=self.cc.log, - ansatz=self.cc.ansatz, - space=self.cc.space, - omega=self.cc.omega, - g=self.cc.bare_g, - G=self.cc.bare_G, - options=self.cc.options, - ) - self.cc.amplitudes = amplitudes - self.cc.kernel() - de = abs(e_prev - self.cc.e_tot) - dt = self.get_t1_norm() - - # Log the iteration: - converged_e = de < self.options.e_tol - converged_t = dt < self.options.t_tol - self.log.info( - f"{niter:4d} {self.cc.e_corr:16.10f} {self.cc.e_tot:18.10f}" - f" {[ANSI.r, ANSI.g][self.cc.converged]}{self.cc.converged!r:>8}{ANSI.R}" - f" {[ANSI.r, ANSI.g][converged_e]}{de:13.3e}{ANSI.R}" - f" {[ANSI.r, ANSI.g][converged_t]}{dt:13.3e}{ANSI.R}" - ) - - # Check for convergence: - converged = converged_e and converged_t - if converged: - self.log.debug("") - self.log.output(f"{ANSI.g}Converged.{ANSI.R}") - break - else: - self.log.debug("") - self.log.warning(f"{ANSI.r}Failed to converge.{ANSI.R}") - - self.cc.log.debug("") - self.cc.log.output("E(corr) = %.10f", self.cc.e_corr) - self.cc.log.output("E(tot) = %.10f", self.cc.e_tot) - self.cc.log.debug("") - self.cc.log.debug("Time elapsed: %s", timer.format_time(timer())) - self.cc.log.debug("") - - return self.cc.e_corr - - @property - def name(self): - """Get a string representation of the method name.""" - return self.spin_type + "B" + self.cc.ansatz.name - - @property - def spin_type(self): - """Return the spin type.""" - return self.cc.spin_type - - -class BruecknerUEBCC(BruecknerREBCC): - def get_rotation_matrix(self, u_tot=None, diis=None, t1=None): - if t1 is None: - t1 = self.cc.t1 - if u_tot is None: - u_tot = util.Namespace( - aa=np.eye(self.cc.space[0].ncorr, dtype=types[float]), - bb=np.eye(self.cc.space[1].ncorr, dtype=types[float]), - ) - - t1_block = util.Namespace( - aa=np.zeros((self.cc.space[0].ncorr, self.cc.space[0].ncorr), dtype=types[float]), - bb=np.zeros((self.cc.space[1].ncorr, self.cc.space[1].ncorr), dtype=types[float]), - ) - t1_block.aa[: self.cc.space[0].ncocc, self.cc.space[0].ncocc :] = -t1.aa - t1_block.aa[self.cc.space[0].ncocc :, : self.cc.space[0].ncocc] = t1.aa.T - t1_block.bb[: self.cc.space[1].ncocc, self.cc.space[1].ncocc :] = -t1.bb - t1_block.bb[self.cc.space[1].ncocc :, : self.cc.space[1].ncocc] = t1.bb.T - - u = util.Namespace( - aa=scipy.linalg.expm(t1_block.aa), - bb=scipy.linalg.expm(t1_block.bb), - ) - - u_tot.aa = np.dot(u_tot.aa, u.aa) - u_tot.bb = np.dot(u_tot.bb, u.bb) - if scipy.linalg.det(u_tot.aa) < 0: - u_tot.aa[:, 0] *= -1 - if scipy.linalg.det(u_tot.bb) < 0: - u_tot.bb[:, 0] *= -1 - - a = np.concatenate( - [ - scipy.linalg.logm(u_tot.aa).ravel(), - scipy.linalg.logm(u_tot.bb).ravel(), - ], - axis=0, - ) - a = a.real.astype(types[float]) - if diis is not None: - xerr = np.concatenate([t1.aa.ravel(), t1.bb.ravel()]) - a = diis.update(a, xerr=xerr) - - u_tot.aa = scipy.linalg.expm(a[: u_tot.aa.size].reshape(u_tot.aa.shape)) - u_tot.bb = scipy.linalg.expm(a[u_tot.aa.size :].reshape(u_tot.bb.shape)) - - return u, u_tot - - def transform_amplitudes(self, u, amplitudes=None): - if amplitudes is None: - amplitudes = self.cc.amplitudes - - nocc = (self.cc.space[0].ncocc, self.cc.space[1].ncocc) - ci = {"a": u.aa[: nocc[0], : nocc[0]], "b": u.bb[: nocc[1], : nocc[1]]} - ca = {"a": u.aa[nocc[0] :, nocc[0] :], "b": u.bb[nocc[1] :, nocc[1] :]} - - # Transform T amplitudes: - for name, key, n in self.cc.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - for comb in util.generate_spin_combinations(n, unique=True): - args = [getattr(self.cc.amplitudes[name], comb), tuple(range(n * 2))] - for i in range(n): - args += [ci[comb[i]], (i, i + n * 2)] - for i in range(n): - args += [ca[comb[i + n]], (i + n, i + n * 3)] - args += [tuple(range(n * 2, n * 4))] - setattr(self.cc.amplitudes[name], comb, util.einsum(*args)) - - # Transform S amplitudes: - for name, key, n in self.cc.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented # TODO - - # Transform U amplitudes: - for name, key, nf, nb in self.cc.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented # TODO - - return self.cc.amplitudes - - def get_t1_norm(self, amplitudes=None): - if amplitudes is None: - amplitudes = self.cc.amplitudes - - norm_a = np.linalg.norm(amplitudes["t1"].aa) - norm_b = np.linalg.norm(amplitudes["t1"].bb) - - return np.linalg.norm([norm_a, norm_b]) - - def mo_to_correlated(self, mo_coeff): - return ( - mo_coeff[0][:, self.cc.space[0].correlated], - mo_coeff[1][:, self.cc.space[1].correlated], - ) - - def mo_update_correlated(self, mo_coeff, mo_coeff_corr): - mo_coeff[0][:, self.cc.space[0].correlated] = mo_coeff_corr[0] - mo_coeff[1][:, self.cc.space[1].correlated] = mo_coeff_corr[1] - - return mo_coeff - - def update_coefficients(self, u_tot, mo_coeff, mo_coeff_ref): - mo_coeff_new_corr = ( - util.einsum("pi,ij->pj", mo_coeff_ref[0], u_tot.aa), - util.einsum("pi,ij->pj", mo_coeff_ref[1], u_tot.bb), - ) - mo_coeff_new = self.mo_update_correlated(mo_coeff, mo_coeff_new_corr) - return mo_coeff_new - - -class BruecknerGEBCC(BruecknerREBCC): - pass diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index ceae2f99..6385caff 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -16,7 +16,7 @@ from ebcc.precision import cast, types if TYPE_CHECKING: - from typing import Any, Callable, Literal, Mapping, Optional, Type, TypeVar, Union + from typing import Any, Callable, Literal, Optional, TypeVar, Union from pyscf.scf.hf import SCF # type: ignore @@ -25,7 +25,6 @@ from ebcc.base import Fock as BaseFock from ebcc.logging import Logger from ebcc.numpy.typing import NDArray # type: ignore - from ebcc.space import Space from ebcc.util import Namespace ERIsInputType = Union[type[BaseERIs], NDArray[float]] @@ -479,17 +478,17 @@ def _load_function( # Get the amplitudes: if not (amplitudes is False): - if amplitudes is None: + if not amplitudes: amplitudes = self.amplitudes - if amplitudes is None: + if not amplitudes: amplitudes = self.init_amps(eris=eris) dicts.append(dict(amplitudes)) # Get the lambda amplitudes: if not (lambdas is False): - if lambdas is None: + if not lambdas: lambdas = self.lambdas - if lambdas is None: + if not lambdas: lambdas = self.init_lams(amplitudes=amplitudes if amplitudes else None) dicts.append(dict(lambdas)) diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 00ca8318..dfc229ff 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -8,11 +8,11 @@ from ebcc import numpy as np from ebcc import util -from ebcc.brueckner import BruecknerGEBCC from ebcc.cc.base import BaseEBCC from ebcc.eom import EA_GEOM, EE_GEOM, IP_GEOM from ebcc.eris import GERIs from ebcc.fock import GFock +from ebcc.opt.gbrueckner import BruecknerGEBCC from ebcc.precision import types from ebcc.space import Space @@ -162,8 +162,8 @@ def from_uebcc(cls, ucc: UEBCC) -> GEBCC: gcc.converged = ucc.converged gcc.converged_lambda = ucc.converged_lambda - has_amps = ucc.amplitudes is not None - has_lams = ucc.lambdas is not None + has_amps = bool(ucc.amplitudes) + has_lams = bool(ucc.lambdas) if has_amps: amplitudes = util.Namespace() diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 8dc89faf..66413f74 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -2,24 +2,18 @@ from __future__ import annotations -import dataclasses from typing import TYPE_CHECKING from pyscf import lib -from ebcc import default_log, init_logging from ebcc import numpy as np from ebcc import util -from ebcc.ansatz import Ansatz -from ebcc.brueckner import BruecknerREBCC from ebcc.cc.base import BaseEBCC from ebcc.cderis import RCDERIs -from ebcc.damping import DIIS -from ebcc.dump import Dump from ebcc.eom import EA_REOM, EE_REOM, IP_REOM from ebcc.eris import RERIs from ebcc.fock import RFock -from ebcc.logging import ANSI +from ebcc.opt.rbrueckner import BruecknerREBCC from ebcc.precision import types from ebcc.space import Space diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 5e760f5d..e99bccee 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -8,12 +8,12 @@ from ebcc import numpy as np from ebcc import util -from ebcc.brueckner import BruecknerUEBCC from ebcc.cc.base import BaseEBCC from ebcc.cderis import UCDERIs from ebcc.eom import EA_UEOM, EE_UEOM, IP_UEOM from ebcc.eris import UERIs from ebcc.fock import UFock +from ebcc.opt.ubrueckner import BruecknerUEBCC from ebcc.precision import types from ebcc.space import Space @@ -114,8 +114,8 @@ def from_rebcc(cls, rcc: REBCC) -> UEBCC: ucc.converged = rcc.converged ucc.converged_lambda = rcc.converged_lambda - has_amps = rcc.amplitudes is not None - has_lams = rcc.lambdas is not None + has_amps = bool(rcc.amplitudes) + has_lams = bool(rcc.lambdas) if has_amps: amplitudes = util.Namespace() diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py index e7d23261..815fb845 100644 --- a/ebcc/eom/base.py +++ b/ebcc/eom/base.py @@ -27,7 +27,7 @@ @dataclass -class Options: +class BaseOptions: """Options for EOM calculations. Args: @@ -54,7 +54,7 @@ class BaseEOM(ABC): """ # Types - Options = Options + Options = BaseOptions def __init__( self, @@ -244,7 +244,7 @@ def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" pass - def get_guesses(self, diag: NDArray[float]) -> list[NDArray[float]]: + def get_guesses(self, diag: Optional[NDArray[float]] = None) -> list[NDArray[float]]: """Get the initial guesses vectors. Args: diff --git a/ebcc/eom/geom.py b/ebcc/eom/geom.py index c9e65e85..32b25ea3 100644 --- a/ebcc/eom/geom.py +++ b/ebcc/eom/geom.py @@ -20,23 +20,25 @@ class GEOM(BaseEOM): """Generalised equation-of-motion coupled cluster.""" + pass + + +class IP_GEOM(GEOM, BaseIP_EOM): + """Generalised ionisation potential equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: """Sort the diagonal to inform the initial guesses.""" if self.options.koopmans: r1 = self.vector_to_amplitudes(diag)[0] - arg = np.argsort(diag[: r1.size]) + arg = np.argsort(np.abs(diag[: r1.size])) else: - arg = np.argsort(diag) + arg = np.argsort(np.abs(diag)) return arg def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" return np.linalg.norm(r1) ** 2 - -class IP_GEOM(GEOM, BaseIP_EOM): - """Generalised ionisation potential equation-of-motion coupled cluster.""" - def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -133,6 +135,19 @@ def moments( class EA_GEOM(GEOM, BaseEA_EOM): """Generalised electron affinity equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(np.abs(diag[: r1.size])) + else: + arg = np.argsort(np.abs(diag)) + return arg + + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" + return np.linalg.norm(r1) ** 2 + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -229,6 +244,19 @@ def moments( class EE_GEOM(GEOM, BaseEE_EOM): """Generalised electron-electron equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(diag[: r1.size]) + else: + arg = np.argsort(diag) + return arg + + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" + return np.linalg.norm(r1) ** 2 + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -315,7 +343,7 @@ def moments( if eris is None: eris = self.ebcc.get_eris() - if amplitudes is None: + if not amplitudes: amplitudes = self.ebcc.amplitudes bras = self.bras(eris=eris) diff --git a/ebcc/eom/reom.py b/ebcc/eom/reom.py index 450480e6..e69c1379 100644 --- a/ebcc/eom/reom.py +++ b/ebcc/eom/reom.py @@ -20,23 +20,25 @@ class REOM(BaseEOM): """Restricted equation-of-motion coupled cluster.""" + pass + + +class IP_REOM(REOM, BaseIP_EOM): + """Restricted ionisation potential equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: """Sort the diagonal to inform the initial guesses.""" if self.options.koopmans: r1 = self.vector_to_amplitudes(diag)[0] - arg = np.argsort(diag[: r1.size]) + arg = np.argsort(np.abs(diag[: r1.size])) else: - arg = np.argsort(diag) + arg = np.argsort(np.abs(diag)) return arg def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" return np.linalg.norm(r1) ** 2 - -class IP_REOM(REOM, BaseIP_EOM): - """Restricted ionisation potential equation-of-motion coupled cluster.""" - def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -133,6 +135,19 @@ def moments( class EA_REOM(REOM, BaseEA_EOM): """Restricted electron affinity equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(np.abs(diag[: r1.size])) + else: + arg = np.argsort(np.abs(diag)) + return arg + + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" + return np.linalg.norm(r1) ** 2 + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -229,6 +244,19 @@ def moments( class EE_REOM(REOM, BaseEE_EOM): """Restricted electron-electron equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(diag[: r1.size]) + else: + arg = np.argsort(diag) + return arg + + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" + return np.linalg.norm(r1) ** 2 + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -315,7 +343,7 @@ def moments( if eris is None: eris = self.ebcc.get_eris() - if amplitudes is None: + if not amplitudes: amplitudes = self.ebcc.amplitudes bras = self.bras(eris=eris) diff --git a/ebcc/eom/ueom.py b/ebcc/eom/ueom.py index 3c6ce157..841f6e6b 100644 --- a/ebcc/eom/ueom.py +++ b/ebcc/eom/ueom.py @@ -20,23 +20,25 @@ class UEOM(BaseEOM): """Unrestricted equation-of-motion coupled cluster.""" + pass + + +class IP_UEOM(UEOM, BaseIP_EOM): + """Unrestricted ionisation potential equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: """Sort the diagonal to inform the initial guesses.""" if self.options.koopmans: r1 = self.vector_to_amplitudes(diag)[0] - arg = np.argsort(diag[: r1.a.size + r1.b.size]) + arg = np.argsort(np.abs(diag[: r1.a.size + r1.b.size])) else: - arg = np.argsort(diag) + arg = np.argsort(np.abs(diag)) return arg def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" return np.linalg.norm(r1.a) ** 2 + np.linalg.norm(r1.b) ** 2 - -class IP_UEOM(UEOM, BaseIP_EOM): - """Unrestricted ionisation potential equation-of-motion coupled cluster.""" - def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -224,6 +226,19 @@ def moments( class EA_UEOM(UEOM, BaseEA_EOM): """Unrestricted electron affinity equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(np.abs(diag[: r1.a.size + r1.b.size])) + else: + arg = np.argsort(np.abs(diag)) + return arg + + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + """Get the quasiparticle weight.""" + return np.linalg.norm(r1.a) ** 2 + np.linalg.norm(r1.b) ** 2 + def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -411,6 +426,15 @@ def moments( class EE_UEOM(UEOM, BaseEE_EOM): """Unrestricted electron-electron equation-of-motion coupled cluster.""" + def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: + """Sort the diagonal to inform the initial guesses.""" + if self.options.koopmans: + r1 = self.vector_to_amplitudes(diag)[0] + arg = np.argsort(diag[: r1.aa.size + r1.bb.size]) + else: + arg = np.argsort(diag) + return arg + def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" return np.linalg.norm(r1.aa) ** 2 + np.linalg.norm(r1.bb) ** 2 diff --git a/ebcc/opt/__init__.py b/ebcc/opt/__init__.py new file mode 100644 index 00000000..bfd9a84d --- /dev/null +++ b/ebcc/opt/__init__.py @@ -0,0 +1,5 @@ +"""Orbital-optimised coupled cluster approaches.""" + +from ebcc.opt.gbrueckner import BruecknerGEBCC +from ebcc.opt.rbrueckner import BruecknerREBCC +from ebcc.opt.ubrueckner import BruecknerUEBCC diff --git a/ebcc/opt/base.py b/ebcc/opt/base.py new file mode 100644 index 00000000..cdf4d52e --- /dev/null +++ b/ebcc/opt/base.py @@ -0,0 +1,286 @@ +"""Base classes for `ebcc.opt`.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from pyscf import lib + +from ebcc import numpy as np +from ebcc import util +from ebcc.damping import DIIS +from ebcc.logging import ANSI, NullLogger, init_logging +from ebcc.precision import types + +if TYPE_CHECKING: + from typing import Any, Optional, TypeVar + + from ebcc.cc.base import AmplitudeType, BaseEBCC + from ebcc.util import Namespace + + T = TypeVar("T") + + +@dataclass +class BaseOptions: + """Options for Brueckner-orbital calculations. + + Args: + e_tol: Threshold for converged in the correlation energy. + t_tol: Threshold for converged in the amplitude norm. + max_iter: Maximum number of iterations. + diis_space: Number of amplitudes to use in DIIS extrapolation. + damping: Damping factor for DIIS extrapolation. + """ + + e_tol: float = 1e-8 + t_tol: float = 1e-8 + max_iter: int = 20 + diis_space: int = 12 + damping: float = 0.0 + + +class BaseBruecknerEBCC(ABC): + """Base class for Brueckner-orbital coupled cluster. + + Attributes: + cc: Parent `BaseEBCC` object. + options: Options for the EOM calculation. + """ + + # Types + Options: type[BaseOptions] = BaseOptions + + def __init__( + self, + cc: BaseEBCC, + options: Optional[BaseOptions] = None, + **kwargs: Any, + ) -> None: + r"""Initialise the Brueckner EBCC object. + + Args: + cc: Parent `EBCC` object. + options: Options for the EOM calculation. + **kwargs: Additional keyword arguments used to update `options`. + """ + # Options: + if options is None: + options = self.Options() + self.options = options + for key, val in kwargs.items(): + setattr(self.options, key, val) + + # Parameters: + self.cc = cc + self.mf = cc.mf + self.space = cc.space + self.log = cc.log + + # Attributes: + self.converged = False + + # Logging: + init_logging(cc.log) + cc.log.info(f"\n{ANSI.B}{ANSI.U}{self.name}{ANSI.R}") + cc.log.info(f"{ANSI.B}Options{ANSI.R}:") + cc.log.info(f" > e_tol: {ANSI.y}{self.options.e_tol}{ANSI.R}") + cc.log.info(f" > t_tol: {ANSI.y}{self.options.t_tol}{ANSI.R}") + cc.log.info(f" > max_iter: {ANSI.y}{self.options.max_iter}{ANSI.R}") + cc.log.info(f" > diis_space: {ANSI.y}{self.options.diis_space}{ANSI.R}") + cc.log.info(f" > damping: {ANSI.y}{self.options.damping}{ANSI.R}") + cc.log.debug("") + + @property + def spin_type(self) -> str: + """Get the spin type.""" + return self.cc.spin_type + + @property + def name(self) -> str: + """Get the name of the method.""" + return f"{self.spin_type}B{self.cc.ansatz.name}" + + def kernel(self) -> float: + """Run the Bruckner-orbital coupled cluster calculation. + + Returns: + Correlation energy. + """ + timer = util.Timer() + + # Make sure the initial CC calculation is converged: + if not self.cc.converged: + with lib.temporary_env(self.cc, log=NullLogger()): + self.cc.kernel() + + # Set up DIIS: + diis = DIIS() + diis.space = self.options.diis_space + diis.damping = self.options.damping + + # Initialise coefficients: + mo_coeff_new = np.array(self.cc.mo_coeff, copy=True, dtype=types[float]) + mo_coeff_ref = np.array(self.cc.mo_coeff, copy=True, dtype=types[float]) + mo_coeff_ref = self.mo_to_correlated(mo_coeff_ref) + u_tot = None + + self.cc.log.output("Solving for Brueckner orbitals.") + self.cc.log.debug("") + self.log.info( + f"{ANSI.B}{'Iter':>4s} {'Energy (corr.)':>16s} {'Energy (tot.)':>18s} " + f"{'Conv.':>8s} {'Δ(Energy)':>13s} {'|T1|':>13s}{ANSI.R}" + ) + self.log.info( + f"{0:4d} {self.cc.e_corr:16.10f} {self.cc.e_tot:18.10f} " + f"{[ANSI.r, ANSI.g][self.cc.converged]}{self.cc.converged!r:>8}{ANSI.R}" + ) + + converged = False + for niter in range(1, self.options.max_iter + 1): + # Update rotation matrix: + u, u_tot = self.get_rotation_matrix(u_tot=u_tot, diis=diis) + + # Update MO coefficients: + mo_coeff_new = self.update_coefficients(u_tot, mo_coeff_new, mo_coeff_ref) + + # Transform mean-field and amplitudes: + self.mf.mo_coeff = mo_coeff_new + self.mf.e_tot = self.mf.energy_tot() + amplitudes = self.transform_amplitudes(u) + + # Run CC calculation: + e_prev = self.cc.e_tot + with lib.temporary_env(self.cc, log=NullLogger()): + self.cc.__init__( + self.mf, + log=self.cc.log, + ansatz=self.cc.ansatz, + space=self.cc.space, + omega=self.cc.omega, + g=self.cc.bare_g, + G=self.cc.bare_G, + options=self.cc.options, + ) + self.cc.amplitudes = amplitudes + self.cc.kernel() + de = abs(e_prev - self.cc.e_tot) + dt = self.get_t1_norm() + + # Log the iteration: + converged_e = de < self.options.e_tol + converged_t = dt < self.options.t_tol + self.log.info( + f"{niter:4d} {self.cc.e_corr:16.10f} {self.cc.e_tot:18.10f}" + f" {[ANSI.r, ANSI.g][self.cc.converged]}{self.cc.converged!r:>8}{ANSI.R}" + f" {[ANSI.r, ANSI.g][converged_e]}{de:13.3e}{ANSI.R}" + f" {[ANSI.r, ANSI.g][converged_t]}{dt:13.3e}{ANSI.R}" + ) + + # Check for convergence: + converged = converged_e and converged_t + if converged: + self.log.debug("") + self.log.output(f"{ANSI.g}Converged.{ANSI.R}") + break + else: + self.log.debug("") + self.log.warning(f"{ANSI.r}Failed to converge.{ANSI.R}") + + self.cc.log.debug("") + self.cc.log.output("E(corr) = %.10f", self.cc.e_corr) + self.cc.log.output("E(tot) = %.10f", self.cc.e_tot) + self.cc.log.debug("") + self.cc.log.debug("Time elapsed: %s", timer.format_time(timer())) + self.cc.log.debug("") + + return self.cc.e_corr + + @abstractmethod + def get_rotation_matrix( + self, + u_tot: Optional[AmplitudeType] = None, + diis: Optional[DIIS] = None, + t1: Optional[AmplitudeType] = None, + ) -> tuple[AmplitudeType, AmplitudeType]: + """Update the rotation matrix. + + Also returns the total rotation matrix. + + Args: + u_tot: Total rotation matrix. + diis: DIIS object. + t1: T1 amplitude. + + Returns: + Rotation matrix and total rotation matrix. + """ + pass + + @abstractmethod + def transform_amplitudes( + self, u: AmplitudeType, amplitudes: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: + """Transform the amplitudes into the Brueckner orbital basis. + + Args: + u: Rotation matrix. + amplitudes: Cluster amplitudes. + + Returns: + Transformed cluster amplitudes. + """ + pass + + @abstractmethod + def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> float: + """Get the norm of the T1 amplitude. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Norm of the T1 amplitude. + """ + pass + + @abstractmethod + def mo_to_correlated(self, mo_coeff: T) -> T: + """Transform the MO coefficients into the correlated basis. + + Args: + mo_coeff: MO coefficients. + + Returns: + Correlated slice of MO coefficients. + """ + pass + + @abstractmethod + def mo_update_correlated(self, mo_coeff: T, mo_coeff_corr: T) -> T: + """Update the correlated slice of a set of MO coefficients. + + Args: + mo_coeff: MO coefficients. + mo_coeff_corr: Correlated slice of MO coefficients. + + Returns: + Updated MO coefficients. + """ + pass + + @abstractmethod + def update_coefficients(self, u_tot: AmplitudeType, mo_coeff_new: T, mo_coeff_ref: T) -> T: + """Update the MO coefficients. + + Args: + u_tot: Total rotation matrix. + mo_coeff_new: New MO coefficients. + mo_coeff_ref: Reference MO coefficients. + + Returns: + Updated MO coefficients. + """ + pass diff --git a/ebcc/opt/gbrueckner.py b/ebcc/opt/gbrueckner.py new file mode 100644 index 00000000..a9a74183 --- /dev/null +++ b/ebcc/opt/gbrueckner.py @@ -0,0 +1,165 @@ +"""Generalised Brueckner-orbital coupled cluster.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import scipy.linalg + +from ebcc import numpy as np +from ebcc import util +from ebcc.opt.base import BaseBruecknerEBCC +from ebcc.precision import types + +if TYPE_CHECKING: + from typing import Optional + + from ebcc.cc.gebcc import AmplitudeType + from ebcc.damping import DIIS + from ebcc.util import Namespace + + +class BruecknerGEBCC(BaseBruecknerEBCC): + """Generalised Brueckner-orbital coupled cluster. + + Attributes: + cc: Parent `BaseEBCC` object. + options: Options for the EOM calculation. + """ + + def get_rotation_matrix( + self, + u_tot: Optional[AmplitudeType] = None, + diis: Optional[DIIS] = None, + t1: Optional[AmplitudeType] = None, + ) -> tuple[AmplitudeType, AmplitudeType]: + """Update the rotation matrix. + + Also returns the total rotation matrix. + + Args: + u_tot: Total rotation matrix. + diis: DIIS object. + t1: T1 amplitude. + + Returns: + Rotation matrix and total rotation matrix. + """ + if t1 is None: + t1 = self.cc.t1 + if u_tot is None: + u_tot = np.eye(self.cc.space.ncorr, dtype=types[float]) + + t1_block = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) + t1_block[: self.cc.space.ncocc, self.cc.space.ncocc :] = -t1 + t1_block[self.cc.space.ncocc :, : self.cc.space.ncocc] = t1.T + + u = scipy.linalg.expm(t1_block) + + u_tot = np.dot(u_tot, u) + if scipy.linalg.det(u_tot) < 0: + u_tot[:, 0] *= -1 + + a = scipy.linalg.logm(u_tot) + a = a.real.astype(types[float]) + if diis is not None: + a = diis.update(a, xerr=t1) + + u_tot = scipy.linalg.expm(a) + + return u, u_tot + + def transform_amplitudes( + self, u: AmplitudeType, amplitudes: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: + """Transform the amplitudes into the Brueckner orbital basis. + + Args: + u: Rotation matrix. + amplitudes: Cluster amplitudes. + + Returns: + Transformed cluster amplitudes. + """ + if not amplitudes: + amplitudes = self.cc.amplitudes + + nocc = self.cc.space.ncocc + ci = u[:nocc, :nocc] + ca = u[nocc:, nocc:] + + # Transform T amplitudes: + for name, key, n in self.cc.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + args = [self.cc.amplitudes[name], tuple(range(n * 2))] + for i in range(n): + args += [ci, (i, i + n * 2)] + for i in range(n): + args += [ca, (i + n, i + n * 3)] + args += [tuple(range(n * 2, n * 4))] + self.cc.amplitudes[name] = util.einsum(*args) + + # Transform S amplitudes: + for name, key, n in self.cc.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented # TODO + + # Transform U amplitudes: + for name, key, nf, nb in self.cc.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented # TODO + + return self.cc.amplitudes + + def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> float: + """Get the norm of the T1 amplitude. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Norm of the T1 amplitude. + """ + if not amplitudes: + amplitudes = self.cc.amplitudes + return np.linalg.norm(amplitudes["t1"]) + + def mo_to_correlated(self, mo_coeff: NDArray[float]) -> NDArray[float]: + """Transform the MO coefficients into the correlated basis. + + Args: + mo_coeff: MO coefficients. + + Returns: + Correlated slice of MO coefficients. + """ + return mo_coeff[:, self.cc.space.correlated] + + def mo_update_correlated( + self, mo_coeff: NDArray[float], mo_coeff_corr: NDArray[float] + ) -> NDArray[float]: + """Update the correlated slice of a set of MO coefficients. + + Args: + mo_coeff: MO coefficients. + mo_coeff_corr: Correlated slice of MO coefficients. + + Returns: + Updated MO coefficients. + """ + mo_coeff[:, self.cc.space.correlated] = mo_coeff_corr + return mo_coeff + + def update_coefficients( + self, u_tot: AmplitudeType, mo_coeff: NDArray[float], mo_coeff_ref: NDArray[float] + ) -> NDArray[float]: + """Update the MO coefficients. + + Args: + u_tot: Total rotation matrix. + mo_coeff: New MO coefficients. + mo_coeff_ref: Reference MO coefficients. + + Returns: + Updated MO coefficients. + """ + mo_coeff_new_corr = util.einsum("pi,ij->pj", mo_coeff_ref, u_tot) + mo_coeff_new = self.mo_update_correlated(mo_coeff, mo_coeff_new_corr) + return mo_coeff_new diff --git a/ebcc/opt/rbrueckner.py b/ebcc/opt/rbrueckner.py new file mode 100644 index 00000000..e0d678df --- /dev/null +++ b/ebcc/opt/rbrueckner.py @@ -0,0 +1,165 @@ +"""Restricted Brueckner-orbital coupled cluster.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import scipy.linalg + +from ebcc import numpy as np +from ebcc import util +from ebcc.opt.base import BaseBruecknerEBCC +from ebcc.precision import types + +if TYPE_CHECKING: + from typing import Optional + + from ebcc.cc.rebcc import AmplitudeType + from ebcc.damping import DIIS + from ebcc.util import Namespace + + +class BruecknerREBCC(BaseBruecknerEBCC): + """Restricted Brueckner-orbital coupled cluster. + + Attributes: + cc: Parent `BaseEBCC` object. + options: Options for the EOM calculation. + """ + + def get_rotation_matrix( + self, + u_tot: Optional[AmplitudeType] = None, + diis: Optional[DIIS] = None, + t1: Optional[AmplitudeType] = None, + ) -> tuple[AmplitudeType, AmplitudeType]: + """Update the rotation matrix. + + Also returns the total rotation matrix. + + Args: + u_tot: Total rotation matrix. + diis: DIIS object. + t1: T1 amplitude. + + Returns: + Rotation matrix and total rotation matrix. + """ + if t1 is None: + t1 = self.cc.t1 + if u_tot is None: + u_tot = np.eye(self.cc.space.ncorr, dtype=types[float]) + + t1_block = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) + t1_block[: self.cc.space.ncocc, self.cc.space.ncocc :] = -t1 + t1_block[self.cc.space.ncocc :, : self.cc.space.ncocc] = t1.T + + u = scipy.linalg.expm(t1_block) + + u_tot = np.dot(u_tot, u) + if scipy.linalg.det(u_tot) < 0: + u_tot[:, 0] *= -1 + + a = scipy.linalg.logm(u_tot) + a = a.real.astype(types[float]) + if diis is not None: + a = diis.update(a, xerr=t1) + + u_tot = scipy.linalg.expm(a) + + return u, u_tot + + def transform_amplitudes( + self, u: AmplitudeType, amplitudes: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: + """Transform the amplitudes into the Brueckner orbital basis. + + Args: + u: Rotation matrix. + amplitudes: Cluster amplitudes. + + Returns: + Transformed cluster amplitudes. + """ + if not amplitudes: + amplitudes = self.cc.amplitudes + + nocc = self.cc.space.ncocc + ci = u[:nocc, :nocc] + ca = u[nocc:, nocc:] + + # Transform T amplitudes: + for name, key, n in self.cc.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + args = [self.cc.amplitudes[name], tuple(range(n * 2))] + for i in range(n): + args += [ci, (i, i + n * 2)] + for i in range(n): + args += [ca, (i + n, i + n * 3)] + args += [tuple(range(n * 2, n * 4))] + self.cc.amplitudes[name] = util.einsum(*args) + + # Transform S amplitudes: + for name, key, n in self.cc.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented # TODO + + # Transform U amplitudes: + for name, key, nf, nb in self.cc.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented # TODO + + return self.cc.amplitudes + + def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> float: + """Get the norm of the T1 amplitude. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Norm of the T1 amplitude. + """ + if not amplitudes: + amplitudes = self.cc.amplitudes + return np.linalg.norm(amplitudes["t1"]) + + def mo_to_correlated(self, mo_coeff: NDArray[float]) -> NDArray[float]: + """Transform the MO coefficients into the correlated basis. + + Args: + mo_coeff: MO coefficients. + + Returns: + Correlated slice of MO coefficients. + """ + return mo_coeff[:, self.cc.space.correlated] + + def mo_update_correlated( + self, mo_coeff: NDArray[float], mo_coeff_corr: NDArray[float] + ) -> NDArray[float]: + """Update the correlated slice of a set of MO coefficients. + + Args: + mo_coeff: MO coefficients. + mo_coeff_corr: Correlated slice of MO coefficients. + + Returns: + Updated MO coefficients. + """ + mo_coeff[:, self.cc.space.correlated] = mo_coeff_corr + return mo_coeff + + def update_coefficients( + self, u_tot: AmplitudeType, mo_coeff: NDArray[float], mo_coeff_ref: NDArray[float] + ) -> NDArray[float]: + """Update the MO coefficients. + + Args: + u_tot: Total rotation matrix. + mo_coeff: New MO coefficients. + mo_coeff_ref: Reference MO coefficients. + + Returns: + Updated MO coefficients. + """ + mo_coeff_new_corr = util.einsum("pi,ij->pj", mo_coeff_ref, u_tot) + mo_coeff_new = self.mo_update_correlated(mo_coeff, mo_coeff_new_corr) + return mo_coeff_new diff --git a/ebcc/opt/ubrueckner.py b/ebcc/opt/ubrueckner.py new file mode 100644 index 00000000..649e2ab3 --- /dev/null +++ b/ebcc/opt/ubrueckner.py @@ -0,0 +1,200 @@ +"""Unrestricted Brueckner-orbital coupled cluster.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import scipy.linalg + +from ebcc import numpy as np +from ebcc import util +from ebcc.opt.base import BaseBruecknerEBCC +from ebcc.precision import types + +if TYPE_CHECKING: + from typing import Optional + + from ebcc.cc.uebcc import AmplitudeType + from ebcc.damping import DIIS + from ebcc.util import Namespace + + +class BruecknerUEBCC(BaseBruecknerEBCC): + """Unrestricted Brueckner-orbital coupled cluster. + + Attributes: + cc: Parent `BaseEBCC` object. + options: Options for the EOM calculation. + """ + + def get_rotation_matrix( + self, + u_tot: Optional[AmplitudeType] = None, + diis: Optional[DIIS] = None, + t1: Optional[AmplitudeType] = None, + ) -> tuple[AmplitudeType, AmplitudeType]: + """Update the rotation matrix. + + Also returns the total rotation matrix. + + Args: + u_tot: Total rotation matrix. + diis: DIIS object. + t1: T1 amplitude. + + Returns: + Rotation matrix and total rotation matrix. + """ + if t1 is None: + t1 = self.cc.t1 + if u_tot is None: + u_tot = util.Namespace( + aa=np.eye(self.cc.space[0].ncorr, dtype=types[float]), + bb=np.eye(self.cc.space[1].ncorr, dtype=types[float]), + ) + + t1_block = util.Namespace( + aa=np.zeros((self.cc.space[0].ncorr, self.cc.space[0].ncorr), dtype=types[float]), + bb=np.zeros((self.cc.space[1].ncorr, self.cc.space[1].ncorr), dtype=types[float]), + ) + t1_block.aa[: self.cc.space[0].ncocc, self.cc.space[0].ncocc :] = -t1.aa + t1_block.aa[self.cc.space[0].ncocc :, : self.cc.space[0].ncocc] = t1.aa.T + t1_block.bb[: self.cc.space[1].ncocc, self.cc.space[1].ncocc :] = -t1.bb + t1_block.bb[self.cc.space[1].ncocc :, : self.cc.space[1].ncocc] = t1.bb.T + + u = util.Namespace( + aa=scipy.linalg.expm(t1_block.aa), + bb=scipy.linalg.expm(t1_block.bb), + ) + + u_tot.aa = np.dot(u_tot.aa, u.aa) + u_tot.bb = np.dot(u_tot.bb, u.bb) + if scipy.linalg.det(u_tot.aa) < 0: + u_tot.aa[:, 0] *= -1 + if scipy.linalg.det(u_tot.bb) < 0: + u_tot.bb[:, 0] *= -1 + + a = np.concatenate( + [ + scipy.linalg.logm(u_tot.aa).ravel(), + scipy.linalg.logm(u_tot.bb).ravel(), + ], + axis=0, + ) + a = a.real.astype(types[float]) + if diis is not None: + xerr = np.concatenate([t1.aa.ravel(), t1.bb.ravel()]) + a = diis.update(a, xerr=xerr) + + u_tot.aa = scipy.linalg.expm(a[: u_tot.aa.size].reshape(u_tot.aa.shape)) + u_tot.bb = scipy.linalg.expm(a[u_tot.aa.size :].reshape(u_tot.bb.shape)) + + return u, u_tot + + def transform_amplitudes( + self, u: AmplitudeType, amplitudes: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: + """Transform the amplitudes into the Brueckner orbital basis. + + Args: + u: Rotation matrix. + amplitudes: Cluster amplitudes. + + Returns: + Transformed cluster amplitudes. + """ + if not amplitudes: + amplitudes = self.cc.amplitudes + + nocc = (self.cc.space[0].ncocc, self.cc.space[1].ncocc) + ci = {"a": u.aa[: nocc[0], : nocc[0]], "b": u.bb[: nocc[1], : nocc[1]]} + ca = {"a": u.aa[nocc[0] :, nocc[0] :], "b": u.bb[nocc[1] :, nocc[1] :]} + + # Transform T amplitudes: + for name, key, n in self.cc.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): + for comb in util.generate_spin_combinations(n, unique=True): + args = [getattr(self.cc.amplitudes[name], comb), tuple(range(n * 2))] + for i in range(n): + args += [ci[comb[i]], (i, i + n * 2)] + for i in range(n): + args += [ca[comb[i + n]], (i + n, i + n * 3)] + args += [tuple(range(n * 2, n * 4))] + setattr(self.cc.amplitudes[name], comb, util.einsum(*args)) + + # Transform S amplitudes: + for name, key, n in self.cc.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented # TODO + + # Transform U amplitudes: + for name, key, nf, nb in self.cc.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): + raise util.ModelNotImplemented # TODO + + return self.cc.amplitudes + + def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> float: + """Get the norm of the T1 amplitude. + + Args: + amplitudes: Cluster amplitudes. + + Returns: + Norm of the T1 amplitude. + """ + if not amplitudes: + amplitudes = self.cc.amplitudes + norm_a = np.linalg.norm(amplitudes["t1"].aa) + norm_b = np.linalg.norm(amplitudes["t1"].bb) + return np.linalg.norm([norm_a, norm_b]) + + def mo_to_correlated(self, mo_coeff: tuple[NDArray[float]]) -> tuple[NDArray[float]]: + """Transform the MO coefficients into the correlated basis. + + Args: + mo_coeff: MO coefficients. + + Returns: + Correlated slice of MO coefficients. + """ + return ( + mo_coeff[0][:, self.cc.space[0].correlated], + mo_coeff[1][:, self.cc.space[1].correlated], + ) + + def mo_update_correlated( + self, mo_coeff: tuple[NDArray[float]], mo_coeff_corr: tuple[NDArray[float]] + ) -> tuple[NDArray[float]]: + """Update the correlated slice of a set of MO coefficients. + + Args: + mo_coeff: MO coefficients. + mo_coeff_corr: Correlated slice of MO coefficients. + + Returns: + Updated MO coefficients. + """ + mo_coeff[0][:, self.cc.space[0].correlated] = mo_coeff_corr[0] + mo_coeff[1][:, self.cc.space[1].correlated] = mo_coeff_corr[1] + return mo_coeff + + def update_coefficients( + self, + u_tot: AmplitudeType, + mo_coeff: tuple[NDArray[float]], + mo_coeff_ref: tuple[NDArray[float]], + ) -> tuple[NDArray[float]]: + """Update the MO coefficients. + + Args: + u_tot: Total rotation matrix. + mo_coeff: New MO coefficients. + mo_coeff_ref: Reference MO coefficients. + + Returns: + Updated MO coefficients. + """ + mo_coeff_new_corr = ( + util.einsum("pi,ij->pj", mo_coeff_ref[0], u_tot.aa), + util.einsum("pi,ij->pj", mo_coeff_ref[1], u_tot.bb), + ) + mo_coeff_new = self.mo_update_correlated(mo_coeff, mo_coeff_new_corr) + return mo_coeff_new diff --git a/ebcc/util/einsumfunc.py b/ebcc/util/einsumfunc.py index 527ae211..4b5ed4ff 100644 --- a/ebcc/util/einsumfunc.py +++ b/ebcc/util/einsumfunc.py @@ -302,6 +302,8 @@ def einsum(*operands: Any, **kwargs: Any) -> Union[T, NDArray[T]]: for contraction in contractions: inds, idx_rm, einsum_str, remain = list(contraction[:4]) contraction_args = [args.pop(x) for x in inds] + if kwargs.get("alpha", 1.0) != 1.0 or kwargs.get("beta", 0.0) != 0.0: + raise NotImplementedError("Scaling factors not supported for >2 arguments") out = _contract(einsum_str, *contraction_args, **kwargs) args.append(out) diff --git a/pyproject.toml b/pyproject.toml index 4ccba447..18d0621c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,9 +77,19 @@ src_paths = [ skip_glob = [ "*/__pycache__/*", "ebcc/codegen/*", - "ebcc/__init__.py", + "*/__init__.py", ] +[tool.unimport] +include_star_import = true +ignore_init = true +include = '\.pyi?$' +exclude = """ +/( + | __init__.py +)/ +""" + [tool.flake8] max-line-length = 100 max-doc-length = 100 diff --git a/tests/test_UCCSD.py b/tests/test_UCCSD.py index 2407b121..0682df2b 100644 --- a/tests/test_UCCSD.py +++ b/tests/test_UCCSD.py @@ -4,15 +4,15 @@ import itertools import os import pickle -import unittest import tempfile +import unittest import numpy as np import pytest import scipy.linalg -from pyscf import cc, gto, lib, scf +from pyscf import cc, gto, scf -from ebcc import REBCC, UEBCC, GEBCC, NullLogger, Space +from ebcc import REBCC, UEBCC, NullLogger, Space @pytest.mark.reference @@ -42,18 +42,26 @@ def setUpClass(cls): mf = mf.to_uhf() ccsd = UEBCC( - mf, - ansatz="CCSD", - log=NullLogger(), + mf, + ansatz="CCSD", + log=NullLogger(), ) ccsd.options.e_tol = 1e-12 eris = ccsd.get_eris() ccsd.kernel(eris=eris) ccsd.solve_lambda(eris=eris) - osort = list(itertools.chain(*zip(range(ccsd.nocc[0]), range(ccsd.nocc[0], ccsd.nocc[0]+ccsd.nocc[1])))) - vsort = list(itertools.chain(*zip(range(ccsd.nvir[1]), range(ccsd.nvir[0], ccsd.nvir[0]+ccsd.nvir[1])))) - fsort = list(itertools.chain(*zip(range(ccsd.nmo), range(ccsd.nmo, 2*ccsd.nmo)))) + osort = list( + itertools.chain( + *zip(range(ccsd.nocc[0]), range(ccsd.nocc[0], ccsd.nocc[0] + ccsd.nocc[1])) + ) + ) + vsort = list( + itertools.chain( + *zip(range(ccsd.nvir[1]), range(ccsd.nvir[0], ccsd.nvir[0] + ccsd.nvir[1])) + ) + ) + fsort = list(itertools.chain(*zip(range(ccsd.nmo), range(ccsd.nmo, 2 * ccsd.nmo)))) cls.mf, cls.ccsd, cls.eris, cls.data = mf, ccsd, eris, data cls.osort, cls.vsort, cls.fsort = osort, vsort, fsort @@ -87,9 +95,9 @@ def test_l1_amplitudes(self): def test_from_rebcc(self): rebcc = REBCC( - self.mf, - ansatz="CCSD", - log=NullLogger(), + self.mf, + ansatz="CCSD", + log=NullLogger(), ) rebcc.options.e_tol = 1e-12 rebcc.kernel() @@ -116,10 +124,10 @@ def test_from_rebcc_frozen(self): space = Space(occupied, frozen, active) rebcc = REBCC( - mf, - ansatz="CCSD", - space=space, - log=NullLogger(), + mf, + ansatz="CCSD", + space=space, + log=NullLogger(), ) rebcc.options.e_tol = 1e-12 rebcc.kernel() @@ -140,10 +148,10 @@ def test_from_rebcc_frozen(self): space_b = Space(occupied, frozen, active) uebcc_2 = UEBCC( - mf, - ansatz="CCSD", - space=(space_a, space_b), - log=NullLogger(), + mf, + ansatz="CCSD", + space=(space_a, space_b), + log=NullLogger(), ) uebcc_2.options.e_tol = 1e-12 uebcc_2.kernel() @@ -188,7 +196,7 @@ def _test_ee_moments_diag(self): eom = self.ccsd.ee_eom() nmo = self.ccsd.nmo a = self.data[True]["dd_moms"].transpose(4, 0, 1, 2, 3) - a = np.einsum("npqrs,pq,rs->npqrs", a, np.eye(nmo*2), np.eye(nmo*2)) + a = np.einsum("npqrs,pq,rs->npqrs", a, np.eye(nmo * 2), np.eye(nmo * 2)) t = eom.moments(4, diagonal_only=True) b = np.zeros_like(a) for i in range(a.shape[0]): @@ -205,8 +213,7 @@ def _test_ee_moments_diag(self): @pytest.mark.reference class UCCSD_PySCF_Tests(unittest.TestCase): - """Test UCCSD against the PySCF values. - """ + """Test UCCSD against the PySCF values.""" @classmethod def setUpClass(cls): @@ -229,9 +236,9 @@ def setUpClass(cls): ccsd_ref.solve_lambda() ccsd = UEBCC( - mf, - ansatz="CCSD", - log=NullLogger(), + mf, + ansatz="CCSD", + log=NullLogger(), ) ccsd.options.e_tol = 1e-12 eris = ccsd.get_eris() @@ -277,6 +284,8 @@ def test_rdm2(self): def test_eom_ip(self): e1 = self.ccsd.ip_eom(nroot=5).kernel() e2, v2 = self.ccsd_ref.ipccsd(nroots=5) + print(e1) + print(e2) self.assertAlmostEqual(e1[0], e2[0], 5) def test_eom_ea(self): @@ -291,8 +300,8 @@ def test_eom_ee(self): # Disabled until PySCF fix bug # TODO -#@pytest.mark.reference -#class FNOUCCSD_PySCF_Tests(UCCSD_PySCF_Tests): +# @pytest.mark.reference +# class FNOUCCSD_PySCF_Tests(UCCSD_PySCF_Tests): # """Test FNO-UCCSD against the PySCF values. # """ # @@ -349,8 +358,7 @@ def test_eom_ee(self): @pytest.mark.reference class UCCSD_PySCF_Frozen_Tests(unittest.TestCase): - """Test UCCSD against the PySCF values with frozen orbitals. - """ + """Test UCCSD against the PySCF values with frozen orbitals.""" @classmethod def setUpClass(cls): @@ -384,22 +392,22 @@ def setUpClass(cls): ccsd_ref.solve_lambda() space_a = Space( - mf.mo_occ[0] > 0, - frozen_a, - np.zeros_like(mf.mo_occ[0]), + mf.mo_occ[0] > 0, + frozen_a, + np.zeros_like(mf.mo_occ[0]), ) space_b = Space( - mf.mo_occ[1] > 0, - frozen_b, - np.zeros_like(mf.mo_occ[1]), + mf.mo_occ[1] > 0, + frozen_b, + np.zeros_like(mf.mo_occ[1]), ) ccsd = UEBCC( - mf, - ansatz="CCSD", - space=(space_a, space_b), - log=NullLogger(), + mf, + ansatz="CCSD", + space=(space_a, space_b), + log=NullLogger(), ) ccsd.options.e_tol = 1e-12 eris = ccsd.get_eris() @@ -460,8 +468,7 @@ def test_eom_ee(self): @pytest.mark.reference class UCCSD_Dump_Tests(UCCSD_PySCF_Tests): - """Test UCCSD against the PySCF after dumping and loading. - """ + """Test UCCSD against the PySCF after dumping and loading.""" @classmethod def setUpClass(cls): @@ -484,9 +491,9 @@ def setUpClass(cls): ccsd_ref.solve_lambda() ccsd = UEBCC( - mf, - ansatz="CCSD", - log=NullLogger(), + mf, + ansatz="CCSD", + log=NullLogger(), ) ccsd.options.e_tol = 1e-12 eris = ccsd.get_eris() diff --git a/tests/test_util.py b/tests/test_util.py index 7d90077f..2f309753 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -1,12 +1,12 @@ """Tests for util module. """ -import unittest import itertools +import unittest -from pyscf import gto, scf import numpy as np import pytest +from pyscf import gto, scf import ebcc from ebcc import util @@ -14,8 +14,7 @@ @pytest.mark.regression class Util_Tests(unittest.TestCase): - """Test util module against known values. - """ + """Test util module against known values.""" @classmethod def setUpClass(cls): @@ -41,8 +40,8 @@ def test_tril_indices_ndim(self): for n in (1, 2, 3, 4): for ndim in (1, 2, 3, 4): for combinations, include_diagonal in [ - (itertools.combinations, False), - (itertools.combinations_with_replacement, True), + (itertools.combinations, False), + (itertools.combinations_with_replacement, True), ]: x = np.zeros((n,) * ndim) y = np.zeros((n,) * ndim) @@ -55,8 +54,8 @@ def test_ntril_ndim(self): for n in (1, 2, 3, 4): for ndim in (1, 2, 3, 4): for combinations, include_diagonal in [ - (itertools.combinations, False), - (itertools.combinations_with_replacement, True), + (itertools.combinations, False), + (itertools.combinations_with_replacement, True), ]: a = sum(1 for tup in combinations(range(n), r=ndim)) b = util.ntril_ndim(n, ndim, include_diagonal=include_diagonal) @@ -64,31 +63,42 @@ def test_ntril_ndim(self): def test_generate_spin_combinations(self): for n, combs in [ - (1, {"aa", "bb"}), - (2, {"aaaa", "abab", "baba", "bbbb"}), - (3, {"aaaaaa", "aabaab", "abaaba", "baabaa", "abbabb", "babbab", "bbabba", "bbbbbb"}), + (1, {"aa", "bb"}), + (2, {"aaaa", "abab", "baba", "bbbb"}), + (3, {"aaaaaa", "aabaab", "abaaba", "baabaa", "abbabb", "babbab", "bbabba", "bbbbbb"}), ]: self.assertEqual(set(util.generate_spin_combinations(n)), combs) for n, combs in [ - (1, {"aa", "bb"}), - (2, {"aaaa", "abab", "bbbb"}), - (3, {"aaaaaa", "abaaba", "babbab", "bbbbbb"}), + (1, {"aa", "bb"}), + (2, {"aaaa", "abab", "bbbb"}), + (3, {"aaaaaa", "abaaba", "babbab", "bbbbbb"}), ]: self.assertEqual(set(util.generate_spin_combinations(n, unique=True)), combs) for n, combs in [ - (1, {"a", "b"}), - (2, {"aaa", "aba", "bab", "bbb"}), - (3, {"aaaaa", "aabaa", "abaab", "baaba", "abbab", "babba", "bbabb", "bbbbb"}), + (1, {"a", "b"}), + (2, {"aaa", "aba", "bab", "bbb"}), + (3, {"aaaaa", "aabaa", "abaab", "baaba", "abbab", "babba", "bbabb", "bbbbb"}), ]: self.assertEqual(set(util.generate_spin_combinations(n, excited=True)), combs) def test_permutations_with_signs(self): for seq, res in [ - ([0, 1], (([0, 1], 1), ([1, 0], -1))), - ([0, 1, 2], tuple(sorted([ - ([0, 1, 2], 1), ([0, 2, 1], -1), ([1, 0, 2], -1), - ([1, 2, 0], 1), ([2, 0, 1], 1), ([2, 1, 0], -1), - ]))), + ([0, 1], (([0, 1], 1), ([1, 0], -1))), + ( + [0, 1, 2], + tuple( + sorted( + [ + ([0, 1, 2], 1), + ([0, 2, 1], -1), + ([1, 0, 2], -1), + ([1, 2, 0], 1), + ([2, 0, 1], 1), + ([2, 1, 0], -1), + ] + ) + ), + ), ]: self.assertEqual(tuple(sorted(util.permutations_with_signs(seq))), res) @@ -100,7 +110,7 @@ def test_symmetry_factor(self): def test_antisymmetrise_array(self): for n in (1, 2, 3, 4): for ndim in (1, 2, 3, 4, 5, 6): - array = np.cos(np.arange(1, n**ndim+1).reshape((n,) * ndim)) + array = np.cos(np.arange(1, n**ndim + 1).reshape((n,) * ndim)) array = util.antisymmetrise_array(array, axes=range(ndim)) for perm, sign in util.permutations_with_signs(range(ndim)): np.testing.assert_almost_equal(array, sign * array.transpose(perm)) @@ -125,11 +135,11 @@ def test_get_compressed_size(self): def test_symmetrise(self): for n in (1, 2, 3, 4): for ndim in (2, 4, 6): - subscript = "i" * (ndim//2) + "a" * (ndim//2) + subscript = "i" * (ndim // 2) + "a" * (ndim // 2) array = np.cos(np.arange(n**ndim).reshape((n,) * ndim)) array = util.symmetrise(subscript, array) - for p1, s1 in util.permutations_with_signs(range(ndim//2)): - for p2, s2 in util.permutations_with_signs(range(ndim//2, ndim)): + for p1, s1 in util.permutations_with_signs(range(ndim // 2)): + for p2, s2 in util.permutations_with_signs(range(ndim // 2, ndim)): perm = tuple(p1) + tuple(p2) sign = s1 * s2 np.testing.assert_almost_equal(array, sign * array.transpose(perm)) @@ -175,10 +185,10 @@ def test_einsum(self): np.testing.assert_almost_equal(a, b) b = util.einsum("ijkl,lkab->ijab", x, y, out=b, alpha=1.0, beta=1.0) - np.testing.assert_almost_equal(a*2, b) + np.testing.assert_almost_equal(a * 2, b) b = util.einsum("ijkl,lkab->ijab", x, y, out=b, alpha=2.0, beta=0.0) - np.testing.assert_almost_equal(a*2, b) + np.testing.assert_almost_equal(a * 2, b) a = np.einsum("ijkl,lkab,bacd->ijcd", x, y, z) b = util.einsum("ijkl,lkab,bacd->ijcd", x, y, z) @@ -192,9 +202,10 @@ def test_einsum(self): b = util.einsum("ijk,jki,kji->ik", x, y, z) np.testing.assert_almost_equal(a, b) - a = np.einsum("ijk,jki,kji->ik", x, y, z) - b = util.einsum("ijk,jki,kji->ik", x, y, z, alpha=0.5, beta=0.5) - np.testing.assert_almost_equal(a, b) + with pytest.raises(NotImplementedError): + b = util.einsum("ijk,jki,kji->ik", x, y, z, alpha=0.5) + b = util.einsum("ijk,jki,kji->ik", x, y, z, beta=0.5) + b = util.einsum("ijk,jki,kji->ik", x, y, z, alpha=0.5, beta=0.5) a = np.einsum("iik,kjj->ij", x, y) b = util.einsum("iik,kjj->ij", x, y) From 19856798bd34821c39108dfbe234183b5c7f27c0 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sat, 3 Aug 2024 19:36:21 +0100 Subject: [PATCH 13/37] Moving hamiltonians stuff --- ebcc/__init__.py | 4 +- ebcc/cc/base.py | 7 +- ebcc/cc/gebcc.py | 4 +- ebcc/cc/rebcc.py | 4 +- ebcc/cc/uebcc.py | 4 +- ebcc/dump.py | 2 +- ebcc/eom/ueom.py | 2 +- ebcc/fock.py | 173 -------------------- ebcc/ham/__init__.py | 3 + ebcc/ham/base.py | 85 ++++++++++ ebcc/ham/fock.py | 126 +++++++++++++++ ebcc/{ => ham}/space.py | 339 +++++++++++++++++++--------------------- examples/13-fno_ccsd.py | 3 +- tests/test_RCCSD.py | 89 +++++------ 14 files changed, 433 insertions(+), 412 deletions(-) delete mode 100644 ebcc/fock.py create mode 100644 ebcc/ham/__init__.py create mode 100644 ebcc/ham/base.py create mode 100644 ebcc/ham/fock.py rename ebcc/{ => ham}/space.py (52%) diff --git a/ebcc/__init__.py b/ebcc/__init__.py index 027ce828..0f0f4b17 100644 --- a/ebcc/__init__.py +++ b/ebcc/__init__.py @@ -109,8 +109,8 @@ def constructor(mf, *args, **kwargs): # --- Other imports: from ebcc.ansatz import Ansatz -from ebcc.brueckner import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC -from ebcc.space import Space +from ebcc.ham import Space +from ebcc.opt import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC # --- List available methods: diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 6385caff..d0d16ac2 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -20,11 +20,10 @@ from pyscf.scf.hf import SCF # type: ignore - from ebcc.base import BruecknerEBCC as BaseBrueckner - from ebcc.base import ERIs as BaseERIs - from ebcc.base import Fock as BaseFock + from ebcc.ham.base import BaseERIs, BaseFock from ebcc.logging import Logger from ebcc.numpy.typing import NDArray # type: ignore + from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace ERIsInputType = Union[type[BaseERIs], NDArray[float]] @@ -75,7 +74,7 @@ class BaseEBCC(ABC): ERIs: type[BaseERIs] Fock: type[BaseFock] CDERIs: type[BaseERIs] - Brueckner: type[BaseBrueckner] + Brueckner: type[BaseBruecknerEBCC] # Attributes space: SpaceType diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index dfc229ff..8238efa1 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -11,10 +11,10 @@ from ebcc.cc.base import BaseEBCC from ebcc.eom import EA_GEOM, EE_GEOM, IP_GEOM from ebcc.eris import GERIs -from ebcc.fock import GFock +from ebcc.ham import Space +from ebcc.ham.fock import GFock from ebcc.opt.gbrueckner import BruecknerGEBCC from ebcc.precision import types -from ebcc.space import Space if TYPE_CHECKING: from typing import Optional diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 66413f74..dea68f4a 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -12,10 +12,10 @@ from ebcc.cderis import RCDERIs from ebcc.eom import EA_REOM, EE_REOM, IP_REOM from ebcc.eris import RERIs -from ebcc.fock import RFock +from ebcc.ham import Space +from ebcc.ham.fock import RFock from ebcc.opt.rbrueckner import BruecknerREBCC from ebcc.precision import types -from ebcc.space import Space if TYPE_CHECKING: from typing import Optional diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index e99bccee..2e930ed2 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -12,10 +12,10 @@ from ebcc.cderis import UCDERIs from ebcc.eom import EA_UEOM, EE_UEOM, IP_UEOM from ebcc.eris import UERIs -from ebcc.fock import UFock +from ebcc.ham import Space +from ebcc.ham.fock import UFock from ebcc.opt.ubrueckner import BruecknerUEBCC from ebcc.precision import types -from ebcc.space import Space if TYPE_CHECKING: from ebcc.cc.rebcc import REBCC diff --git a/ebcc/dump.py b/ebcc/dump.py index 2f65628c..c514b1b6 100644 --- a/ebcc/dump.py +++ b/ebcc/dump.py @@ -5,7 +5,7 @@ from ebcc import util from ebcc.ansatz import Ansatz -from ebcc.space import Space +from ebcc.ham import Space class Dump: diff --git a/ebcc/eom/ueom.py b/ebcc/eom/ueom.py index 841f6e6b..c031a2af 100644 --- a/ebcc/eom/ueom.py +++ b/ebcc/eom/ueom.py @@ -12,7 +12,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.base import AmplitudeType, ERIsInputType + from ebcc.cc.uebcc import AmplitudeType, ERIsInputType from ebcc.numpy.typing import NDArray from ebcc.util import Namespace diff --git a/ebcc/fock.py b/ebcc/fock.py deleted file mode 100644 index ff9360c6..00000000 --- a/ebcc/fock.py +++ /dev/null @@ -1,173 +0,0 @@ -"""Fock matrix containers.""" - -from ebcc import numpy as np -from ebcc import util -from ebcc.base import Fock -from ebcc.precision import types - - -class RFock(Fock): - """ - Fock matrix container class for `REBCC`. - - The default slices are: - * `"x"`: correlated - * `"o"`: correlated occupied - * `"v"`: correlated virtual - * `"O"`: active occupied - * `"V"`: active virtual - * `"i"`: inactive occupied - * `"a"`: inactive virtual - - Parameters - ---------- - ebcc : REBCC - The EBCC object. - array : np.ndarray, optional - The array of the Fock matrix in the MO basis. If provided, do - not perform just-in-time transformations but instead slice the - array. Default value is `None`. - slices : iterable of slice, optional - The slices to use for each dimension. If provided, the default - slices outlined above are used. - mo_coeff : np.ndarray, optional - The MO coefficients. If not provided, the MO coefficients from - `ebcc` are used. Default value is `None`. - g : Namespace, optional - Namespace containing blocks of the electron-boson coupling - matrix. Default value is `None`. - """ - - def __init__(self, ebcc, array=None, slices=None, mo_coeff=None, g=None): - util.Namespace.__init__(self) - - self.mf = ebcc.mf - self.space = ebcc.space - self.slices = slices - self.mo_coeff = mo_coeff - self.array = array - - self.shift = ebcc.options.shift - self.xi = ebcc.xi - self.g = g - if self.g is None: - self.g = ebcc.g - - if self.mo_coeff is None: - self.mo_coeff = ebcc.mo_coeff - if not (isinstance(self.mo_coeff, (tuple, list)) or self.mo_coeff.ndim == 3): - self.mo_coeff = [self.mo_coeff] * 2 - - if self.array is None: - fock_ao = self.mf.get_fock().astype(types[float]) - self.array = util.einsum("pq,pi,qj->ij", fock_ao, *self.mo_coeff) - - if self.slices is None: - self.slices = { - "x": self.space.correlated, - "o": self.space.correlated_occupied, - "v": self.space.correlated_virtual, - "O": self.space.active_occupied, - "V": self.space.active_virtual, - "i": self.space.inactive_occupied, - "a": self.space.inactive_virtual, - } - if not isinstance(self.slices, (tuple, list)): - self.slices = [self.slices] * 2 - - def __getattr__(self, key): - """Just-in-time attribute getter.""" - - if key not in self.__dict__: - ki, kj = key - i = self.slices[0][ki] - j = self.slices[1][kj] - self.__dict__[key] = self.array[i][:, j].copy() - - if self.shift: - xi = self.xi - g = self.g.__getattr__(f"b{ki}{kj}") - g += self.g.__getattr__(f"b{kj}{ki}").transpose(0, 2, 1) - self.__dict__[key] -= util.einsum("I,Ipq->pq", xi, g) - - return self.__dict__[key] - - __getitem__ = __getattr__ - - -class UFock(Fock): - """ - Fock matrix container class for `UEBCC`. Consists of a namespace of - `RFock` objects, on for each spin signature. - - Parameters - ---------- - ebcc : UEBCC - The EBCC object. - array : iterable of np.ndarray, optional - The array of the Fock matrix in the MO basis. If provided, do - not perform just-in-time transformations but instead slice the - array. Default value is `None`. - slices : iterable of iterable of slice, optional - The slices to use for each dimension. If provided, the default - slices outlined above are used. - mo_coeff : iterable of np.ndarray, optional - The MO coefficients. If not provided, the MO coefficients from - `ebcc` are used. Default value is `None`. - """ - - def __init__(self, ebcc, array=None, slices=None, mo_coeff=None): - util.Namespace.__init__(self) - - self.mf = ebcc.mf - self.space = ebcc.space - self.slices = slices - self.mo_coeff = mo_coeff - self.array = array - - self.shift = ebcc.options.shift - self.xi = ebcc.xi - self.g = ebcc.g - - if self.mo_coeff is None: - self.mo_coeff = ebcc.mo_coeff - - if self.slices is None: - self.slices = [ - { - "x": space.correlated, - "o": space.correlated_occupied, - "v": space.correlated_virtual, - "O": space.active_occupied, - "V": space.active_virtual, - "i": space.inactive_occupied, - "a": space.inactive_virtual, - } - for space in self.space - ] - - if self.array is None: - fock_ao = np.asarray(self.mf.get_fock()).astype(types[float]) - self.array = ( - util.einsum("pq,pi,qj->ij", fock_ao[0], self.mo_coeff[0], self.mo_coeff[0]), - util.einsum("pq,pi,qj->ij", fock_ao[1], self.mo_coeff[1], self.mo_coeff[1]), - ) - - self.aa = RFock( - ebcc, - array=self.array[0], - slices=[self.slices[0], self.slices[0]], - mo_coeff=[self.mo_coeff[0], self.mo_coeff[0]], - g=self.g.aa if self.g is not None else None, - ) - self.bb = RFock( - ebcc, - array=self.array[1], - slices=[self.slices[1], self.slices[1]], - mo_coeff=[self.mo_coeff[1], self.mo_coeff[1]], - g=self.g.bb if self.g is not None else None, - ) - - -class GFock(RFock): - __doc__ = RFock.__doc__.replace("REBCC", "GEBCC") diff --git a/ebcc/ham/__init__.py b/ebcc/ham/__init__.py new file mode 100644 index 00000000..1a8a38d0 --- /dev/null +++ b/ebcc/ham/__init__.py @@ -0,0 +1,3 @@ +"""Hamiltonian objects.""" + +from ebcc.ham.space import Space diff --git a/ebcc/ham/base.py b/ebcc/ham/base.py new file mode 100644 index 00000000..f1fe93f7 --- /dev/null +++ b/ebcc/ham/base.py @@ -0,0 +1,85 @@ +"""Base classes for `ebcc.ham`.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, TypeVar + + from ebcc.cc.base import BaseEBCC + from ebcc.util import Namespace + + T = TypeVar("T") + + +class BaseFock(ABC): + """Base class for Fock matrices. + + Attributes: + cc: Coupled cluster object. + space: Space object. + mo_coeff: Molecular orbital coefficients. + array: Fock matrix in the MO basis. + g: Namespace containing blocks of the electron-boson coupling matrix. + shift: Shift parameter. + xi: Boson parameters. + """ + + def __init__( + self, + cc: BaseEBCC, + array: T = None, + space: Any = None, + mo_coeff: T = None, + g: Namespace[Any] = None, + ) -> None: + """Initialise the Fock matrix. + + Args: + cc: Coupled cluster object. + array: Fock matrix in the MO basis. + space: Space object. + mo_coeff: Molecular orbital coefficients. + g: Namespace containing blocks of the electron-boson coupling matrix. + """ + # Parameters: + self.cc = cc + self.space = space if space is not None else cc.space + self.mo_coeff = mo_coeff if mo_coeff is not None else cc.mo_coeff + self.array = array if array is not None else self._get_fock() + self.g = g if g is not None else cc.g + + # Boson parameters: + self.shift = cc.options.shift + self.xi = cc.xi + + @abstractmethod + def _get_fock(self) -> T: + """Get the Fock matrix.""" + pass + + @abstractmethod + def __getattr__(self, key: str) -> Any: + """Just-in-time attribute getter. + + Args: + key: Key to get. + + Returns: + Slice of the Fock matrix. + """ + pass + + def __getitem__(self, key: str) -> Any: + """Get an item.""" + return self.__getattr__(key) + + +class BaseERIs(ABC): + """Base class for electronic repulsion integrals.""" + + def __getitem__(self, key: str) -> Any: + """Get an item.""" + return self.__dict__[key] diff --git a/ebcc/ham/fock.py b/ebcc/ham/fock.py new file mode 100644 index 00000000..ac188d43 --- /dev/null +++ b/ebcc/ham/fock.py @@ -0,0 +1,126 @@ +"""Fock matrix containers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ebcc import numpy as np +from ebcc import util +from ebcc.ham.base import BaseFock +from ebcc.precision import types + +if TYPE_CHECKING: + from ebcc.numpy.typing import NDArray + + +class RFock(BaseFock): + """Restricted Fock matrix container class. + + Attributes: + cc: Coupled cluster object. + space: Space object. + mo_coeff: Molecular orbital coefficients. + array: Fock matrix in the MO basis. + g: Namespace containing blocks of the electron-boson coupling matrix. + shift: Shift parameter. + xi: Boson parameters. + """ + + def _get_fock(self) -> NDArray[float]: + fock_ao = self.mf.get_fock().astype(types[float]) + return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff, self.mo_coeff) + + def __getattr__(self, key: str) -> NDArray[float]: + """Just-in-time attribute getter. + + Args: + key: Key to get. + + Returns: + Fock matrix for the given spin. + """ + if key not in self.__dict__: + i = self.space.mask(key[0]) + j = self.space.mask(key[1]) + self.__dict__[key] = self.array[i][:, j].copy() + + if self.shift: + xi = self.xi + g = self.g.__getattr__(f"b{key}") + g += self.g.__getattr__(f"b{key[::-1]}").transpose(0, 2, 1) + self.__dict__[key] -= np.einsum("I,Ipq->pq", xi, g) + + return self.__dict__[key] + + +class UFock(BaseFock): + """Unrestricted Fock matrix container class.""" + + def _get_fock(self) -> tuple[NDArray[float]]: + fock_ao = self.mf.get_fock().astype(types[float]) + return ( + util.einsum("pq,pi,qj->ij", fock_ao[0], self.mo_coeff[0], self.mo_coeff[0]), + util.einsum("pq,pi,qj->ij", fock_ao[1], self.mo_coeff[1], self.mo_coeff[1]), + ) + + def __getattr__(self, key: str) -> RFock: + """Just-in-time attribute getter. + + Args: + key: Key to get. + + Returns: + Slice of the Fock matrix. + """ + if key not in ("aa", "bb"): + raise KeyError(f"Invalid key: {key}") + if key not in self.__dict__: + i = "ab".index(key[0]) + self.__dict__[key] = RFock( + self.cc, + array=self.array[i], + space=self.space[i], + mo_coeff=self.mo_coeff[i], + g=self.g[key] if self.g is not None else None, + ) + return self.__dict__[key] + + +class GFock(BaseFock): + """Generalised Fock matrix container class. + + Attributes: + cc: Coupled cluster object. + space: Space object. + mo_coeff: Molecular orbital coefficients. + array: Fock matrix in the MO basis. + g: Namespace containing blocks of the electron-boson coupling matrix. + shift: Shift parameter. + xi: Boson parameters. + """ + + def _get_fock(self) -> NDArray[float]: + fock_ao = self.mf.get_fock().astype(types[float]) + return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff, self.mo_coeff) + + def __getattr__(self, key: str) -> NDArray[float]: + """Just-in-time attribute getter. + + Args: + key: Key to get. + + Returns: + Fock matrix for the given spin. + """ + if key not in self.__dict__: + i = self.space.mask(key[0]) + j = self.space.mask(key[1]) + self.__dict__[key] = self.array[i][:, j].copy() + + if self.shift: + xi = self.xi + g = self.g.__getattr__(f"b{key}") + g += self.g.__getattr__(f"b{key[::-1]}").transpose(0, 2, 1) + self.__dict__[key] -= np.einsum("I,Ipq->pq", xi, g) + + return self.__dict__[key] diff --git a/ebcc/space.py b/ebcc/ham/space.py similarity index 52% rename from ebcc/space.py rename to ebcc/ham/space.py index ad0e6250..512e7eea 100644 --- a/ebcc/space.py +++ b/ebcc/ham/space.py @@ -1,58 +1,83 @@ """Space definition.""" +from __future__ import annotations + +from typing import TYPE_CHECKING + from pyscf.mp import MP2 from ebcc import numpy as np from ebcc import util from ebcc.precision import types +if TYPE_CHECKING: + from typing import Optional, Union + + from pyscf.scf.hf import SCF + + from ebcc.cc.base import AmplitudeType + from ebcc.numpy.typing import NDArray + + ConstructSpaceReturnType = Union[ + tuple[NDArray[float], NDArray[float], Space], + tuple[ + tuple[NDArray[float], NDArray[float]], + tuple[NDArray[float], NDArray[float]], + tuple[Space, Space], + ], + ] + class Space: """ Space class. - - +----------+ - | | frozen | - | +----------+ - - occupied | | active | | - | +----------+ | correlated - | | inactive | | - - #==========# - - | | inactive | | - | +----------+ | correlated - virtual | | active | | - | +----------+ - - | | frozen | - - +----------+ - - Parameters - ---------- - occupied : np.ndarray - Array containing boolean flags indicating whether or not each - orbital is occupied. - frozen : np.ndarray - Array containing boolean flags indicating whether or not each - orbital is frozen. - active : np.ndarray - Array containing boolean flags indicating whether or not each - orbital is active. + ─┬─ ┌──────────┐ + │ │ frozen │ + │ ├──────────┤ ─┬─ + virtual │ │ active │ │ + │ ├──────────┤ │ correlated + │ │ inactive │ │ + ─┼─ ├══════════┤ ─┼─ + │ │ inactive │ │ + │ ├──────────┤ │ correlated + occupied │ │ active │ │ + │ ├──────────┤ ─┴─ + │ │ frozen │ + ─┴─ └──────────┘ + + Args: + occupied: Array containing boolean flags indicating whether or not each orbital is occupied. + frozen: Array containing boolean flags indicating whether or not each orbital is frozen. + active: Array containing boolean flags indicating whether or not each orbital is active. """ def __init__( self, - occupied: np.ndarray, - frozen: np.ndarray, - active: np.ndarray, - ): + occupied: NDArray[bool], + frozen: NDArray[bool], + active: NDArray[bool], + ) -> None: + """Initialise the space. + + Args: + occupied: Array containing boolean flags indicating whether or not each orbital is + occupied. + frozen: Array containing boolean flags indicating whether or not each orbital is frozen. + active: Array containing boolean flags indicating whether or not each orbital is active. + """ self._occupied = np.asarray(occupied, dtype=bool) self._frozen = np.asarray(frozen, dtype=bool) self._active = np.asarray(active, dtype=bool) - assert self._occupied.size == self._frozen.size == self._active.size - assert not np.any(np.logical_and(self._frozen, self._active)) + # Checks: + if not (self._occupied.size == self._frozen.size == self._active.size): + raise ValueError("The sizes of the space arrays must match.") + if np.any(np.logical_and(self._frozen, self._active)): + raise ValueError("Frozen and active orbitals must be mutually exclusive.") - def __repr__(self): - """Return a string representation of the space.""" + def __repr__(self) -> str: + """Get a string representation of the space.""" out = "(%do, %dv)" % (self.nocc, self.nvir) parts = [] if self.nfroz: @@ -63,23 +88,17 @@ def __repr__(self): out += " [" + ", ".join(parts) + "]" return out - def size(self, char): - """ - Convert a character in the standard `ebcc` notation to the size - corresponding to this space. See `ebcc.eris` for details on the - default slices. - - Parameters - ---------- - char : str - The character to convert. - - Returns - ------- - n : int - The size of the space. + def size(self, char: str) -> int: + """Convert a character corresponding to a space to the size of that space. + + Args: + char: Character to convert. + + Returns: + Size of the space. """ return { + "x": self.ncorr, "o": self.ncocc, "O": self.naocc, "i": self.niocc, @@ -88,23 +107,17 @@ def size(self, char): "a": self.nivir, }[char] - def mask(self, char): - """ - Convert a character in the standard `ebcc` notation to the mask - corresponding to this space. See `ebcc.eris` for details on the - default slices. - - Parameters - ---------- - char : str - The character to convert. - - Returns - ------- - mask : np.ndarray - The mask corresponding to the space. + def mask(self, char: str) -> NDArray[bool]: + """Convert a character corresponding to a space to a mask of that space. + + Args: + char: Character to convert. + + Returns: + Mask of the space. """ return { + "x": self.correlated, "o": self.correlated_occupied, "O": self.active_occupied, "i": self.inactive_occupied, @@ -113,215 +126,197 @@ def mask(self, char): "a": self.inactive_virtual, }[char] - def omask(self, char): - """ - Like `mask`, but returns only a mask into only the occupied sector. + def omask(self, char: str) -> NDArray[bool]: + """Like `mask`, but returns only a mask into only the occupied sector. - Parameters - ---------- - char : str - The character to convert. + Args: + char: Character to convert. - Returns - ------- - mask : np.ndarray - The mask corresponding to the space. + Returns: + Mask of the space. """ return self.mask(char)[self.occupied] - def vmask(self, char): - """ - Like `mask`, but returns only a mask into only the virtual sector. + def vmask(self, char: str) -> NDArray[bool]: + """Like `mask`, but returns only a mask into only the virtual sector. - Parameters - ---------- - char : str - The character to convert. + Args: + char: Character to convert. - Returns - ------- - mask : np.ndarray - The mask corresponding to the space. + Returns: + Mask of the space. """ return self.mask(char)[self.virtual] # Full space: @property - def occupied(self): + def occupied(self) -> NDArray[bool]: """Get a boolean mask of occupied orbitals.""" return self._occupied @property - def virtual(self): + def virtual(self) -> NDArray[bool]: """Get a boolean mask of virtual orbitals.""" return ~self.occupied @property - def nmo(self): + def nmo(self) -> int: """Get the number of orbitals.""" return self.occupied.size @property - def nocc(self): + def nocc(self) -> int: """Get the number of occupied orbitals.""" return np.sum(self.occupied) @property - def nvir(self): + def nvir(self) -> int: """Get the number of virtual orbitals.""" return np.sum(self.virtual) # Correlated space: @property - def correlated(self): + def correlated(self) -> NDArray[bool]: """Get a boolean mask of correlated orbitals.""" return ~self.frozen @property - def correlated_occupied(self): + def correlated_occupied(self) -> NDArray[bool]: """Get a boolean mask of occupied correlated orbitals.""" return np.logical_and(self.correlated, self.occupied) @property - def correlated_virtual(self): + def correlated_virtual(self) -> NDArray[bool]: """Get a boolean mask of virtual correlated orbitals.""" return np.logical_and(self.correlated, self.virtual) @property - def ncorr(self): + def ncorr(self) -> int: """Get the number of correlated orbitals.""" return np.sum(self.correlated) @property - def ncocc(self): + def ncocc(self) -> int: """Get the number of occupied correlated orbitals.""" return np.sum(self.correlated_occupied) @property - def ncvir(self): + def ncvir(self) -> int: """Get the number of virtual correlated orbitals.""" return np.sum(self.correlated_virtual) # Inactive space: @property - def inactive(self): + def inactive(self) -> NDArray[bool]: """Get a boolean mask of inactive orbitals.""" return ~self.active @property - def inactive_occupied(self): + def inactive_occupied(self) -> NDArray[bool]: """Get a boolean mask of occupied inactive orbitals.""" return np.logical_and(self.inactive, self.occupied) @property - def inactive_virtual(self): + def inactive_virtual(self) -> NDArray[bool]: """Get a boolean mask of virtual inactive orbitals.""" return np.logical_and(self.inactive, self.virtual) @property - def ninact(self): + def ninact(self) -> int: """Get the number of inactive orbitals.""" return np.sum(self.inactive) @property - def niocc(self): + def niocc(self) -> int: """Get the number of occupied inactive orbitals.""" return np.sum(self.inactive_occupied) @property - def nivir(self): + def nivir(self) -> int: """Get the number of virtual inactive orbitals.""" return np.sum(self.inactive_virtual) # Frozen space: @property - def frozen(self): + def frozen(self) -> NDArray[bool]: """Get a boolean mask of frozen orbitals.""" return self._frozen @property - def frozen_occupied(self): + def frozen_occupied(self) -> NDArray[bool]: """Get a boolean mask of occupied frozen orbitals.""" return np.logical_and(self.frozen, self.occupied) @property - def frozen_virtual(self): + def frozen_virtual(self) -> NDArray[bool]: """Get a boolean mask of virtual frozen orbitals.""" return np.logical_and(self.frozen, self.virtual) @property - def nfroz(self): + def nfroz(self) -> int: """Get the number of frozen orbitals.""" return np.sum(self.frozen) @property - def nfocc(self): + def nfocc(self) -> int: """Get the number of occupied frozen orbitals.""" return np.sum(self.frozen_occupied) @property - def nfvir(self): + def nfvir(self) -> int: """Get the number of virtual frozen orbitals.""" return np.sum(self.frozen_virtual) # Active space: @property - def active(self): + def active(self) -> NDArray[bool]: """Get a boolean mask of active orbitals.""" return self._active @property - def active_occupied(self): + def active_occupied(self) -> NDArray[bool]: """Get a boolean mask of occupied active orbitals.""" return np.logical_and(self.active, self.occupied) @property - def active_virtual(self): + def active_virtual(self) -> NDArray[bool]: """Get a boolean mask of virtual active orbitals.""" return np.logical_and(self.active, self.virtual) @property - def nact(self): + def nact(self) -> int: """Get the number of active orbitals.""" return np.sum(self.active) @property - def naocc(self): + def naocc(self) -> int: """Get the number of occupied active orbitals.""" return np.sum(self.active_occupied) @property - def navir(self): + def navir(self) -> int: """Get the number of virtual active orbitals.""" return np.sum(self.active_virtual) -def construct_default_space(mf): - """ - Construct a default space. - - Parameters - ---------- - mf : pyscf.scf.hf.SCF - PySCF mean-field object. - - Returns - ------- - mo_coeff : np.ndarray - The molecular orbital coefficients. - mo_occ : np.ndarray - The molecular orbital occupation numbers. - space : Space - The default space. +def construct_default_space(mf: SCF) -> ConstructSpaceReturnType: + """Construct a default space. + + Args: + mf: PySCF mean-field object. + + Returns: + The molecular orbital coefficients, the molecular orbital occupation numbers, and the + default space. """ - def _construct(mo_occ): - # Build the default space + def _construct(mo_occ: NDArray[float]) -> Space: + """Build the default space.""" frozen = np.zeros_like(mo_occ, dtype=bool) active = np.zeros_like(mo_occ, dtype=bool) space = Space( @@ -335,42 +330,31 @@ def _construct(mo_occ): if np.ndim(mf.mo_occ) == 2: space_a = _construct(mf.mo_occ[0]) space_b = _construct(mf.mo_occ[1]) - space = (space_a, space_b) + return mf.mo_coeff, mf.mo_occ, (space_a, space_b) else: - space = _construct(mf.mo_occ) - - return mf.mo_coeff, mf.mo_occ, space - - -def construct_fno_space(mf, occ_tol=1e-5, occ_frac=None, amplitudes=None): + return mf.mo_coeff, mf.mo_occ, _construct(mf.mo_occ) + + +def construct_fno_space( + mf: SCF, + occ_tol: Optional[float] = 1e-5, + occ_frac: Optional[float] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, +) -> ConstructSpaceReturnType: + """Construct a frozen natural orbital space. + + Args: + mf: PySCF mean-field object. + occ_tol: Threshold in the natural orbital occupation numbers. + occ_frac: Fraction of the natural orbital occupation numbers to be retained. Overrides + `occ_tol` if both are specified. + amplitudes: Cluster amplitudes. If provided, use these amplitudes when calculating the MP2 + 1RDM. + + Returns: + The natural orbital coefficients, the natural orbital occupation numbers, and the frozen + natural orbital space. """ - Construct a frozen natural orbital space. - - Parameters - ---------- - mf : pyscf.scf.hf.SCF - PySCF mean-field object. - occ_tol : float, optional - Threshold in the natural orbital occupation numbers. Default - value is `1e-5`. - occ_frac : float, optional - Fraction of the natural orbital occupation numbers to be - retained. Overrides `occ_tol` if both are specified. Default - value is `None`. - amplitudes : Namespace, optional - Cluster amplitudes. If provided, use these amplitudes when - calculating the MP2 1RDM. Default value is `None`. - - Returns - ------- - no_coeff : np.ndarray - The natural orbital coefficients. - no_occ : np.ndarray - The natural orbital occupation numbers. - no_space : Space - The frozen natural orbital space. - """ - # Get the MP2 1RDM solver = MP2(mf) if amplitudes is None: @@ -383,7 +367,13 @@ def construct_fno_space(mf, occ_tol=1e-5, occ_frac=None, amplitudes=None): else: dm1 = solver.make_rdm1(t2=amplitudes.t2) - def _construct(dm1, mo_energy, mo_coeff, mo_occ): + # def _construct(dm1, mo_energy, mo_coeff, mo_occ): + def _construct( + dm1: NDArray[float], + mo_energy: NDArray[float], + mo_coeff: NDArray[float], + mo_occ: NDArray[float], + ) -> ConstructSpaceReturnType: # Get the number of occupied orbitals nocc = np.sum(mo_occ > 0) @@ -418,15 +408,12 @@ def _construct(dm1, mo_energy, mo_coeff, mo_occ): active=np.zeros_like(mo_occ, dtype=bool), ) - return no_coeff, no_space + return no_coeff, mo_occ, no_space # Construct the natural orbitals if np.ndim(mf.mo_occ) == 2: - no_coeff_a, no_space_a = _construct(dm1[0], mf.mo_energy[0], mf.mo_coeff[0], mf.mo_occ[0]) - no_coeff_b, no_space_b = _construct(dm1[1], mf.mo_energy[1], mf.mo_coeff[1], mf.mo_occ[1]) - no_coeff = (no_coeff_a, no_coeff_b) - no_space = (no_space_a, no_space_b) + coeff_a, occ_a, space_a = _construct(dm1[0], mf.mo_energy[0], mf.mo_coeff[0], mf.mo_occ[0]) + coeff_b, occ_b, space_b = _construct(dm1[1], mf.mo_energy[1], mf.mo_coeff[1], mf.mo_occ[1]) + return (coeff_a, coeff_b), (occ_a, occ_b), (space_a, space_b) else: - no_coeff, no_space = _construct(dm1, mf.mo_energy, mf.mo_coeff, mf.mo_occ) - - return no_coeff, mf.mo_occ, no_space + return _construct(dm1, mf.mo_energy, mf.mo_coeff, mf.mo_occ) diff --git a/examples/13-fno_ccsd.py b/examples/13-fno_ccsd.py index 799794b7..8469f170 100644 --- a/examples/13-fno_ccsd.py +++ b/examples/13-fno_ccsd.py @@ -2,11 +2,10 @@ Example of a simple FNO-CCSD calculation. """ -import numpy as np from pyscf import gto, scf from ebcc import EBCC -from ebcc.space import construct_fno_space +from ebcc.ham.space import construct_fno_space # Define the molecule using PySCF mol = gto.Mole() diff --git a/tests/test_RCCSD.py b/tests/test_RCCSD.py index fc1b7da7..fc668036 100644 --- a/tests/test_RCCSD.py +++ b/tests/test_RCCSD.py @@ -4,22 +4,21 @@ import itertools import os import pickle -import unittest import tempfile +import unittest import numpy as np import pytest import scipy.linalg -from pyscf import cc, gto, lib, scf +from pyscf import cc, gto, scf from ebcc import REBCC, NullLogger, Space -from ebcc.space import construct_fno_space +from ebcc.ham.space import construct_fno_space @pytest.mark.reference class RCCSD_Tests(unittest.TestCase): - """Test RCCSD against the legacy GCCSD values. - """ + """Test RCCSD against the legacy GCCSD values.""" @classmethod def setUpClass(cls): @@ -41,18 +40,18 @@ def setUpClass(cls): mf.mo_coeff = mo_coeff ccsd = REBCC( - mf, - ansatz="CCSD", - log=NullLogger(), + mf, + ansatz="CCSD", + log=NullLogger(), ) ccsd.options.e_tol = 1e-10 eris = ccsd.get_eris() ccsd.kernel(eris=eris) ccsd.solve_lambda(eris=eris) - osort = list(itertools.chain(*zip(range(ccsd.nocc), range(ccsd.nocc, 2*ccsd.nocc)))) - vsort = list(itertools.chain(*zip(range(ccsd.nvir), range(ccsd.nvir, 2*ccsd.nvir)))) - fsort = list(itertools.chain(*zip(range(ccsd.nmo), range(ccsd.nmo, 2*ccsd.nmo)))) + osort = list(itertools.chain(*zip(range(ccsd.nocc), range(ccsd.nocc, 2 * ccsd.nocc)))) + vsort = list(itertools.chain(*zip(range(ccsd.nvir), range(ccsd.nvir, 2 * ccsd.nvir)))) + fsort = list(itertools.chain(*zip(range(ccsd.nmo), range(ccsd.nmo, 2 * ccsd.nmo)))) cls.mf, cls.ccsd, cls.eris, cls.data = mf, ccsd, eris, data cls.osort, cls.vsort, cls.fsort = osort, vsort, fsort @@ -84,7 +83,7 @@ def test_l1_amplitudes(self): b = scipy.linalg.block_diag(self.ccsd.l1, self.ccsd.l1)[self.vsort][:, self.osort] np.testing.assert_almost_equal(a, b, 6) - #def test_ip_moments(self): + # def test_ip_moments(self): # eom = self.ccsd.ip_eom() # ip_moms = eom.moments(4) # a = self.data[True]["ip_moms"].transpose(2, 0, 1) @@ -110,7 +109,7 @@ def test_l1_amplitudes(self): # y /= np.max(np.abs(y)) # np.testing.assert_almost_equal(x, y, 6) - #def test_ea_moments(self): + # def test_ea_moments(self): # eom = self.ccsd.ea_eom() # ea_moms = eom.moments(4) # a = self.data[True]["ea_moms"].transpose(2, 0, 1) @@ -127,13 +126,13 @@ def test_l1_amplitudes(self): # y /= np.max(np.abs(y)) # np.testing.assert_almost_equal(x, y, 6) - #def test_ip_1mom(self): + # def test_ip_1mom(self): # ip_1mom = self.ccsd.make_ip_1mom() # a = self.data[True]["ip_1mom"] # b = scipy.linalg.block_diag(ip_1mom, ip_1mom) # np.testing.assert_almost_equal(a, b, 6) - #def test_ea_1mom(self): + # def test_ea_1mom(self): # ea_1mom = self.ccsd.make_ea_1mom() # a = self.data[True]["ea_1mom"] # b = sceay.linalg.block_diag(ea_1mom, ea_1mom) @@ -142,14 +141,13 @@ def test_l1_amplitudes(self): @pytest.mark.reference class RCCSD_PySCF_Tests(unittest.TestCase): - """Test RCCSD against the PySCF values. - """ + """Test RCCSD against the PySCF values.""" @classmethod def setUpClass(cls): mol = gto.Mole() mol.atom = "O 0.0 0.0 0.11779; H 0.0 0.755453 -0.471161; H 0.0 -0.755453 -0.471161" - #mol.atom = "Li 0 0 0; H 0 0 1.4" + # mol.atom = "Li 0 0 0; H 0 0 1.4" mol.basis = "cc-pvdz" mol.verbose = 0 mol.build() @@ -166,9 +164,9 @@ def setUpClass(cls): ccsd_ref.solve_lambda() ccsd = REBCC( - mf, - ansatz="CCSD", - log=NullLogger(), + mf, + ansatz="CCSD", + log=NullLogger(), ) ccsd.options.e_tol = 1e-10 eris = ccsd.get_eris() @@ -222,13 +220,13 @@ def test_rdm2(self): b = self.ccsd.make_rdm2_f(eris=self.eris) np.testing.assert_almost_equal(a, b, 6, verbose=True) - #def test_eom_ip(self): + # def test_eom_ip(self): # eom = self.ccsd.ip_eom(nroots=5) # e1 = eom.kernel() # e2, v2 = self.ccsd_ref.ipccsd(nroots=5) # self.assertAlmostEqual(e1[0], e2[0], 6) - #def test_eom_ea(self): + # def test_eom_ea(self): # eom = self.ccsd.ea_eom(nroots=5) # e1 = eom.kernel() # e2, v2 = self.ccsd_ref.eaccsd(nroots=5) @@ -237,14 +235,13 @@ def test_rdm2(self): @pytest.mark.reference class FNORCCSD_PySCF_Tests(RCCSD_PySCF_Tests): - """Test FNO-RCCSD against the PySCF values. - """ + """Test FNO-RCCSD against the PySCF values.""" @classmethod def setUpClass(cls): mol = gto.Mole() mol.atom = "O 0.0 0.0 0.11779; H 0.0 0.755453 -0.471161; H 0.0 -0.755453 -0.471161" - #mol.atom = "Li 0 0 0; H 0 0 1.4" + # mol.atom = "Li 0 0 0; H 0 0 1.4" mol.basis = "cc-pvdz" mol.verbose = 0 mol.build() @@ -265,12 +262,12 @@ def setUpClass(cls): no_coeff = ccsd_ref.mo_coeff ccsd = REBCC( - mf, - mo_coeff=no_coeff, - mo_occ=no_occ, - space=no_space, - ansatz="CCSD", - log=NullLogger(), + mf, + mo_coeff=no_coeff, + mo_occ=no_occ, + space=no_space, + ansatz="CCSD", + log=NullLogger(), ) ccsd.options.e_tol = 1e-10 eris = ccsd.get_eris() @@ -292,8 +289,7 @@ def test_rdm2(self): @pytest.mark.reference class RCCSD_PySCF_Frozen_Tests(unittest.TestCase): - """Test RCCSD against the PySCF values with frozen orbitals. - """ + """Test RCCSD against the PySCF values with frozen orbitals.""" @classmethod def setUpClass(cls): @@ -321,16 +317,16 @@ def setUpClass(cls): ccsd_ref.solve_lambda() space = Space( - mf.mo_occ > 0, - frozen, - np.zeros_like(mf.mo_occ), + mf.mo_occ > 0, + frozen, + np.zeros_like(mf.mo_occ), ) ccsd = REBCC( - mf, - ansatz="CCSD", - space=space, - log=NullLogger(), + mf, + ansatz="CCSD", + space=space, + log=NullLogger(), ) ccsd.options.e_tol = 1e-13 eris = ccsd.get_eris() @@ -387,14 +383,13 @@ def test_rdm2(self): @pytest.mark.reference class RCCSD_Dump_Tests(RCCSD_PySCF_Tests): - """Test RCCSD against PySCF after dumping and loading. - """ + """Test RCCSD against PySCF after dumping and loading.""" @classmethod def setUpClass(cls): mol = gto.Mole() mol.atom = "O 0.0 0.0 0.11779; H 0.0 0.755453 -0.471161; H 0.0 -0.755453 -0.471161" - #mol.atom = "Li 0 0 0; H 0 0 1.4" + # mol.atom = "Li 0 0 0; H 0 0 1.4" mol.basis = "cc-pvdz" mol.verbose = 0 mol.build() @@ -411,9 +406,9 @@ def setUpClass(cls): ccsd_ref.solve_lambda() ccsd = REBCC( - mf, - ansatz="CCSD", - log=NullLogger(), + mf, + ansatz="CCSD", + log=NullLogger(), ) ccsd.options.e_tol = 1e-10 eris = ccsd.get_eris() From 391d2d4f9b19a8111b0668599b073965a9b23a84 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sat, 3 Aug 2024 22:11:43 +0100 Subject: [PATCH 14/37] Move ERIs to ham module --- ebcc/cc/gebcc.py | 2 +- ebcc/cc/rebcc.py | 2 +- ebcc/cc/uebcc.py | 2 +- ebcc/eom/base.py | 2 +- ebcc/eris.py | 220 ----------------------------------------------- ebcc/ham/base.py | 76 +++++++++++----- ebcc/ham/eris.py | 151 ++++++++++++++++++++++++++++++++ ebcc/ham/fock.py | 82 ++++++++++-------- 8 files changed, 253 insertions(+), 284 deletions(-) delete mode 100644 ebcc/eris.py create mode 100644 ebcc/ham/eris.py diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 8238efa1..c9845b88 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -10,8 +10,8 @@ from ebcc import util from ebcc.cc.base import BaseEBCC from ebcc.eom import EA_GEOM, EE_GEOM, IP_GEOM -from ebcc.eris import GERIs from ebcc.ham import Space +from ebcc.ham.eris import GERIs from ebcc.ham.fock import GFock from ebcc.opt.gbrueckner import BruecknerGEBCC from ebcc.precision import types diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index dea68f4a..06b5fb41 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -11,8 +11,8 @@ from ebcc.cc.base import BaseEBCC from ebcc.cderis import RCDERIs from ebcc.eom import EA_REOM, EE_REOM, IP_REOM -from ebcc.eris import RERIs from ebcc.ham import Space +from ebcc.ham.eris import RERIs from ebcc.ham.fock import RFock from ebcc.opt.rbrueckner import BruecknerREBCC from ebcc.precision import types diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 2e930ed2..efd388b4 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -11,8 +11,8 @@ from ebcc.cc.base import BaseEBCC from ebcc.cderis import UCDERIs from ebcc.eom import EA_UEOM, EE_UEOM, IP_UEOM -from ebcc.eris import UERIs from ebcc.ham import Space +from ebcc.ham.eris import UERIs from ebcc.ham.fock import UFock from ebcc.opt.ubrueckner import BruecknerUEBCC from ebcc.precision import types diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py index 815fb845..b489f2c7 100644 --- a/ebcc/eom/base.py +++ b/ebcc/eom/base.py @@ -336,7 +336,7 @@ def davidson( r1n = self.vector_to_amplitudes(vn)[0] qpwt = self._quasiparticle_weight(r1n) self.log.output( - f"{n:>4d} {en:>16.10f} {qpwt:>13.5g} " f"{[ANSI.r, ANSI.g][cn]}{cn!r:>8s}{ANSI.R}" + f"{n:>4d} {en:>16.10f} {qpwt:>13.5g} {[ANSI.r, ANSI.g][bool(cn)]}{cn!r:>8s}{ANSI.R}" ) self.log.debug("") diff --git a/ebcc/eris.py b/ebcc/eris.py deleted file mode 100644 index 0fc34ec3..00000000 --- a/ebcc/eris.py +++ /dev/null @@ -1,220 +0,0 @@ -"""Electronic repulsion integral containers.""" - -from pyscf import ao2mo - -from ebcc import numpy as np -from ebcc import util -from ebcc.precision import types -from ebcc.base import ERIs - - -class RERIs(ERIs): - """ - Electronic repulsion integral container class for `REBCC`. Consists - of a just-in-time namespace containing blocks of the integrals. - - The default slices are: - * `"x"`: correlated - * `"o"`: correlated occupied - * `"v"`: correlated virtual - * `"O"`: active occupied - * `"V"`: active virtual - * `"i"`: inactive occupied - * `"a"`: inactive virtual - - Attributes - ---------- - ebcc : REBCC - The EBCC object. - array : np.ndarray, optional - The array of integrals in the MO basis. If provided, do not - perform just-in-time transformations but instead slice the - array. Default value is `None`. - slices : iterable of slice, optional - The slices to use for each dimension. If provided, the default - slices outlined above are used. - mo_coeff : np.ndarray, optional - The MO coefficients. If not provided, the MO coefficients from - `ebcc` are used. Default value is `None`. - """ - - def __init__(self, ebcc, array=None, slices=None, mo_coeff=None): - util.Namespace.__init__(self) - - self.mf = ebcc.mf - self.space = ebcc.space - self.slices = slices - self.mo_coeff = mo_coeff - self.array = array - - if self.mo_coeff is None: - self.mo_coeff = ebcc.mo_coeff - if not (isinstance(self.mo_coeff, (tuple, list)) or self.mo_coeff.ndim == 3): - self.mo_coeff = [self.mo_coeff] * 4 - - if self.slices is None: - self.slices = { - "x": self.space.correlated, - "o": self.space.correlated_occupied, - "v": self.space.correlated_virtual, - "O": self.space.active_occupied, - "V": self.space.active_virtual, - "i": self.space.inactive_occupied, - "a": self.space.inactive_virtual, - } - if not isinstance(self.slices, (tuple, list)): - self.slices = [self.slices] * 4 - - def __getattr__(self, key): - """Just-in-time attribute getter.""" - - if self.array is None: - if key not in self.__dict__.keys(): - coeffs = [] - for i, k in enumerate(key): - coeffs.append(self.mo_coeff[i][:, self.slices[i][k]].astype(np.float64)) - block = ao2mo.incore.general( - self.mf._eri, - coeffs, - compact=False, - ) - block = block.reshape([c.shape[-1] for c in coeffs]) - self.__dict__[key] = block.astype(types[float]) - return self.__dict__[key] - else: - slices = [] - for i, k in enumerate(key): - slices.append(self.slices[i][k]) - si, sj, sk, sl = slices - block = self.array[si][:, sj][:, :, sk][:, :, :, sl] - return block - - __getitem__ = __getattr__ - - -class UERIs(ERIs): - """ - Electronic repulsion integral container class for `UEBCC`. Consists - of a namespace of `REBCC` objects, one for each spin signature. - - Attributes - ---------- - ebcc : UEBCC - The EBCC object. - array : iterable of np.ndarray, optional - The array of integrals in the MO basis. If provided, do not - perform just-in-time transformations but instead slice the - array. Default value is `None`. - slices : iterable of iterable of slice, optional - The slices to use for each spin and each dimension therein. - If provided, the default slices outlined above are used. - mo_coeff : iterable of np.ndarray, optional - The MO coefficients for each spin. If not provided, the MO - coefficients from `ebcc` are used. Default value is `None`. - """ - - def __init__(self, ebcc, array=None, slices=None, mo_coeff=None): - util.Namespace.__init__(self) - - self.mf = ebcc.mf - self.space = ebcc.space - self.slices = slices - self.mo_coeff = mo_coeff - - if self.slices is None: - self.slices = [ - { - "x": space.correlated, - "o": space.correlated_occupied, - "v": space.correlated_virtual, - "O": space.active_occupied, - "V": space.active_virtual, - "i": space.inactive_occupied, - "a": space.inactive_virtual, - } - for space in self.space - ] - - if self.mo_coeff is None: - self.mo_coeff = ebcc.mo_coeff - - if array is not None: - arrays = (array[0], array[1], array[1].transpose((2, 3, 0, 1)), array[2]) - elif isinstance(self.mf._eri, tuple): - # Have spin-dependent coulomb interaction; precalculate - # required arrays for simplicity. - arrays_aabb = ao2mo.incore.general( - self.mf._eri[1], - [self.mo_coeff[i].astype(np.float64) for i in (0, 0, 1, 1)], - compact=False, - ) - arrays = ( - ao2mo.incore.general( - self.mf._eri[0], - [self.mo_coeff[i].astype(np.float64) for i in (0, 0, 0, 0)], - compact=False, - ), - arrays_aabb, - arrays_aabb.transpose(2, 3, 0, 1), - ao2mo.incore.general( - self.mf._eri[2], - [self.mo_coeff[i].astype(np.float64) for i in (1, 1, 1, 1)], - compact=False, - ), - ) - arrays = tuple(array.astype(types[float]) for array in arrays) - else: - arrays = (None, None, None, None) - - self.aaaa = RERIs( - ebcc, - arrays[0], - slices=[self.slices[i] for i in (0, 0, 0, 0)], - mo_coeff=[self.mo_coeff[i] for i in (0, 0, 0, 0)], - ) - self.aabb = RERIs( - ebcc, - arrays[1], - slices=[self.slices[i] for i in (0, 0, 1, 1)], - mo_coeff=[self.mo_coeff[i] for i in (0, 0, 1, 1)], - ) - self.bbaa = RERIs( - ebcc, - arrays[2], - slices=[self.slices[i] for i in (1, 1, 0, 0)], - mo_coeff=[self.mo_coeff[i] for i in (1, 1, 0, 0)], - ) - self.bbbb = RERIs( - ebcc, - arrays[3], - slices=[self.slices[i] for i in (1, 1, 1, 1)], - mo_coeff=[self.mo_coeff[i] for i in (1, 1, 1, 1)], - ) - - -class GERIs(RERIs): - __doc__ = __doc__.replace("REBCC", "GEBCC") - - def __init__(self, ebcc, array=None, slices=None, mo_coeff=None): - util.Namespace.__init__(self) - - if mo_coeff is None: - mo_coeff = ebcc.mo_coeff - if not (isinstance(mo_coeff, (tuple, list)) or mo_coeff.ndim == 3): - mo_coeff = [mo_coeff] * 4 - - if array is None: - mo_a = [mo[: ebcc.mf.mol.nao].astype(np.float64) for mo in mo_coeff] - mo_b = [mo[ebcc.mf.mol.nao :].astype(np.float64) for mo in mo_coeff] - - array = ao2mo.kernel(ebcc.mf._eri, mo_a) - array += ao2mo.kernel(ebcc.mf._eri, mo_b) - array += ao2mo.kernel(ebcc.mf._eri, mo_a[:2] + mo_b[2:]) - array += ao2mo.kernel(ebcc.mf._eri, mo_b[:2] + mo_a[2:]) - - array = ao2mo.addons.restore(1, array, ebcc.nmo) - array = array.astype(types[float]) - array = array.reshape((ebcc.nmo,) * 4) - array = array.transpose(0, 2, 1, 3) - array.transpose(0, 2, 3, 1) - - RERIs.__init__(self, ebcc, slices=slices, mo_coeff=mo_coeff, array=array) diff --git a/ebcc/ham/base.py b/ebcc/ham/base.py index f1fe93f7..1a7ba4e5 100644 --- a/ebcc/ham/base.py +++ b/ebcc/ham/base.py @@ -5,16 +5,17 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING +from ebcc.util import Namespace + if TYPE_CHECKING: from typing import Any, TypeVar from ebcc.cc.base import BaseEBCC - from ebcc.util import Namespace T = TypeVar("T") -class BaseFock(ABC): +class BaseFock(Namespace, ABC): """Base class for Fock matrices. Attributes: @@ -30,9 +31,9 @@ class BaseFock(ABC): def __init__( self, cc: BaseEBCC, - array: T = None, - space: Any = None, - mo_coeff: T = None, + array: Any = None, + space: tuple[Any] = None, + mo_coeff: tuple[Any] = None, g: Namespace[Any] = None, ) -> None: """Initialise the Fock matrix. @@ -40,20 +41,22 @@ def __init__( Args: cc: Coupled cluster object. array: Fock matrix in the MO basis. - space: Space object. - mo_coeff: Molecular orbital coefficients. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. g: Namespace containing blocks of the electron-boson coupling matrix. """ + Namespace.__init__(self) + # Parameters: - self.cc = cc - self.space = space if space is not None else cc.space - self.mo_coeff = mo_coeff if mo_coeff is not None else cc.mo_coeff - self.array = array if array is not None else self._get_fock() - self.g = g if g is not None else cc.g + 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 # Boson parameters: - self.shift = cc.options.shift - self.xi = cc.xi + self.__dict__["shift"] = cc.options.shift + self.__dict__["xi"] = cc.xi @abstractmethod def _get_fock(self) -> T: @@ -61,8 +64,8 @@ def _get_fock(self) -> T: pass @abstractmethod - def __getattr__(self, key: str) -> Any: - """Just-in-time attribute getter. + def __getitem__(self, key: str) -> Any: + """Just-in-time getter. Args: key: Key to get. @@ -72,14 +75,41 @@ def __getattr__(self, key: str) -> Any: """ pass - def __getitem__(self, key: str) -> Any: - """Get an item.""" - return self.__getattr__(key) - -class BaseERIs(ABC): +class BaseERIs(Namespace, ABC): """Base class for electronic repulsion integrals.""" + def __init__( + self, + cc: BaseEBCC, + array: Any = None, + space: tuple[Any] = None, + mo_coeff: 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 + + @abstractmethod def __getitem__(self, key: str) -> Any: - """Get an item.""" - return self.__dict__[key] + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + Slice of the ERIs. + """ + pass diff --git a/ebcc/ham/eris.py b/ebcc/ham/eris.py new file mode 100644 index 00000000..c9b5b2b1 --- /dev/null +++ b/ebcc/ham/eris.py @@ -0,0 +1,151 @@ +"""Electronic repulsion integral containers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyscf import ao2mo + +from ebcc import numpy as np +from ebcc.ham.base import BaseERIs +from ebcc.precision import types + +if TYPE_CHECKING: + from ebcc.numpy.typing import NDArray + + +class RERIs(BaseERIs): + """Restricted ERIs container class. + + Attributes: + cc: Coupled cluster object. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. + array: ERIs in the MO basis. + """ + + def __getitem__(self, key: str) -> NDArray[float]: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + ERIs for the given spaces. + """ + if self.array is None: + if key not in self._members.keys(): + coeffs = [ + self.mo_coeff[i][:, self.space[i].mask(k)].astype(np.float64) + for i, k in enumerate(key) + ] + block = ao2mo.incore.general( + self.cc.mf._eri, + coeffs, + compact=False, + ) + block = block.reshape([c.shape[-1] for c in coeffs]) + self._members[key] = block.astype(types[float]) + return self._members[key] + else: + i, j, k, l = [self.space[i].mask(k) for i, k in enumerate(key)] + block = self.array[i][:, j][:, :, k][:, :, :, l] + return block + + +class UERIs(BaseERIs): + """Unrestricted ERIs container class. + + Attributes: + cc: Coupled cluster object. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. + array: ERIs in the MO basis. + """ + + def __getitem__(self, key: str) -> RERIs: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + ERIs for the given spins. + """ + 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 + + array: Optional[NDArray[float]] = None + if self.array is not None: + array = self.array[ij] + if key == "bbaa": + array = array.transpose(2, 3, 0, 1) + elif isinstance(self.cc.mf._eri, tuple): + # Support spin-dependent integrals in the mean-field + array = ao2mo.incore.general( + self.cc.mf._eri[ij], + [ + self.mo_coeff[x][y].astype(np.float64) + for y, x in enumerate(sorted((i, i, j, j))) + ], + compact=False, + ) + if key == "bbaa": + array = array.transpose(2, 3, 0, 1) + array = array.astype(types[float]) + + self._members[key] = RERIs( + self.cc, + array=array, + space=(self.space[0][i], self.space[1][i], self.space[2][j], self.space[3][j]), + mo_coeff=( + self.mo_coeff[0][i], + self.mo_coeff[1][i], + self.mo_coeff[2][j], + self.mo_coeff[3][j], + ), + ) + return self._members[key] + + +class GERIs(BaseERIs): + """Generalised ERIs container class. + + Attributes: + cc: Coupled cluster object. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. + array: ERIs in the MO basis. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Initialise the class.""" + super().__init__(*args, **kwargs) + if self.array is None: + mo_a = [mo[: self.cc.mf.mol.nao].astype(np.float64) for mo in self.mo_coeff] + mo_b = [mo[self.cc.mf.mol.nao :].astype(np.float64) for mo in self.mo_coeff] + array = ao2mo.kernel(self.cc.mf._eri, mo_a) + array += ao2mo.kernel(self.cc.mf._eri, mo_b) + array += ao2mo.kernel(self.cc.mf._eri, mo_a[:2] + mo_b[2:]) + array += ao2mo.kernel(self.cc.mf._eri, mo_b[:2] + mo_a[2:]) + array = ao2mo.addons.restore(1, array, self.cc.nmo).reshape((self.cc.nmo,) * 4) + array = array.astype(types[float]) + array = array.transpose(0, 2, 1, 3) - array.transpose(0, 2, 3, 1) + self.__dict__["array"] = array + + def __getitem__(self, key: str) -> NDArray[float]: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + ERIs for the given spaces. + """ + i, j, k, l = [self.space[i].mask(k) for i, k in enumerate(key)] + block = self.array[i][:, j][:, :, k][:, :, :, l] + return block diff --git a/ebcc/ham/fock.py b/ebcc/ham/fock.py index ac188d43..974bb914 100644 --- a/ebcc/ham/fock.py +++ b/ebcc/ham/fock.py @@ -18,8 +18,8 @@ class RFock(BaseFock): Attributes: cc: Coupled cluster object. - space: Space object. - mo_coeff: Molecular orbital coefficients. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. array: Fock matrix in the MO basis. g: Namespace containing blocks of the electron-boson coupling matrix. shift: Shift parameter. @@ -27,63 +27,71 @@ class RFock(BaseFock): """ def _get_fock(self) -> NDArray[float]: - fock_ao = self.mf.get_fock().astype(types[float]) - return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff, self.mo_coeff) + fock_ao = self.cc.mf.get_fock().astype(types[float]) + return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff[0], self.mo_coeff[1]) - def __getattr__(self, key: str) -> NDArray[float]: - """Just-in-time attribute getter. + def __getitem__(self, key: str) -> NDArray[float]: + """Just-in-time getter. Args: key: Key to get. Returns: - Fock matrix for the given spin. + Fock matrix for the given spaces. """ - if key not in self.__dict__: - i = self.space.mask(key[0]) - j = self.space.mask(key[1]) - self.__dict__[key] = self.array[i][:, j].copy() + if key not in self._members: + i = self.space[0].mask(key[0]) + j = self.space[1].mask(key[1]) + self._members[key] = self.array[i][:, j].copy() if self.shift: xi = self.xi g = self.g.__getattr__(f"b{key}") g += self.g.__getattr__(f"b{key[::-1]}").transpose(0, 2, 1) - self.__dict__[key] -= np.einsum("I,Ipq->pq", xi, g) + self._members[key] -= np.einsum("I,Ipq->pq", xi, g) - return self.__dict__[key] + return self._members[key] class UFock(BaseFock): - """Unrestricted Fock matrix container class.""" + """Unrestricted Fock matrix container class. + + Attributes: + cc: Coupled cluster object. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. + array: Fock matrix in the MO basis. + g: Namespace containing blocks of the electron-boson coupling matrix + """ def _get_fock(self) -> tuple[NDArray[float]]: - fock_ao = self.mf.get_fock().astype(types[float]) + fock_ao = self.cc.mf.get_fock().astype(types[float]) return ( - util.einsum("pq,pi,qj->ij", fock_ao[0], self.mo_coeff[0], self.mo_coeff[0]), - util.einsum("pq,pi,qj->ij", fock_ao[1], self.mo_coeff[1], self.mo_coeff[1]), + 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 __getattr__(self, key: str) -> RFock: - """Just-in-time attribute getter. + def __getitem__(self, key: str) -> RFock: + """Just-in-time getter. Args: key: Key to get. Returns: - Slice of the Fock matrix. + Fock matrix for the given spin. """ if key not in ("aa", "bb"): raise KeyError(f"Invalid key: {key}") - if key not in self.__dict__: + if key not in self._members: i = "ab".index(key[0]) - self.__dict__[key] = RFock( + self._members[key] = RFock( self.cc, array=self.array[i], - space=self.space[i], - mo_coeff=self.mo_coeff[i], + 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, ) - return self.__dict__[key] + return self._members[key] class GFock(BaseFock): @@ -91,8 +99,8 @@ class GFock(BaseFock): Attributes: cc: Coupled cluster object. - space: Space object. - mo_coeff: Molecular orbital coefficients. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. array: Fock matrix in the MO basis. g: Namespace containing blocks of the electron-boson coupling matrix. shift: Shift parameter. @@ -100,11 +108,11 @@ class GFock(BaseFock): """ def _get_fock(self) -> NDArray[float]: - fock_ao = self.mf.get_fock().astype(types[float]) - return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff, self.mo_coeff) + fock_ao = self.cc.mf.get_fock().astype(types[float]) + return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff[0], self.mo_coeff[1]) - def __getattr__(self, key: str) -> NDArray[float]: - """Just-in-time attribute getter. + def __getitem__(self, key: str) -> NDArray[float]: + """Just-in-time getter. Args: key: Key to get. @@ -112,15 +120,15 @@ def __getattr__(self, key: str) -> NDArray[float]: Returns: Fock matrix for the given spin. """ - if key not in self.__dict__: - i = self.space.mask(key[0]) - j = self.space.mask(key[1]) - self.__dict__[key] = self.array[i][:, j].copy() + if key not in self._members: + i = self.space[0].mask(key[0]) + j = self.space[1].mask(key[1]) + self._members[key] = self.array[i][:, j].copy() if self.shift: xi = self.xi g = self.g.__getattr__(f"b{key}") g += self.g.__getattr__(f"b{key[::-1]}").transpose(0, 2, 1) - self.__dict__[key] -= np.einsum("I,Ipq->pq", xi, g) + self._members[key] -= np.einsum("I,Ipq->pq", xi, g) - return self.__dict__[key] + return self._members[key] From 0e4e435777fc2e71b5c468b158e941438d428ef4 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sun, 4 Aug 2024 11:48:11 +0100 Subject: [PATCH 15/37] Move CDERIs --- ebcc/cc/rebcc.py | 2 +- ebcc/cc/uebcc.py | 2 +- ebcc/cderis.py | 207 --------------------------------------------- ebcc/ham/cderis.py | 128 ++++++++++++++++++++++++++++ 4 files changed, 130 insertions(+), 209 deletions(-) delete mode 100644 ebcc/cderis.py create mode 100644 ebcc/ham/cderis.py diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 06b5fb41..05d83938 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -9,9 +9,9 @@ from ebcc import numpy as np from ebcc import util from ebcc.cc.base import BaseEBCC -from ebcc.cderis import RCDERIs from ebcc.eom import EA_REOM, EE_REOM, IP_REOM from ebcc.ham import Space +from ebcc.ham.cderis import RCDERIs from ebcc.ham.eris import RERIs from ebcc.ham.fock import RFock from ebcc.opt.rbrueckner import BruecknerREBCC diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index efd388b4..bb86c21f 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -9,9 +9,9 @@ from ebcc import numpy as np from ebcc import util from ebcc.cc.base import BaseEBCC -from ebcc.cderis import UCDERIs from ebcc.eom import EA_UEOM, EE_UEOM, IP_UEOM from ebcc.ham import Space +from ebcc.ham.cderis import UCDERIs from ebcc.ham.eris import UERIs from ebcc.ham.fock import UFock from ebcc.opt.ubrueckner import BruecknerUEBCC diff --git a/ebcc/cderis.py b/ebcc/cderis.py deleted file mode 100644 index 1787ec40..00000000 --- a/ebcc/cderis.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Cholesky decomposed ERI containers.""" - -import numpy as np -from pyscf import ao2mo - -from ebcc import precision, util -from ebcc.base import ERIs - - -class RCDERIs(ERIs): - """ - Cholesky decomposed ERI container class for `REBCC`. Consists of a - just-in-time namespace containing blocks of the integrals. - - The default slices are: - * `"x"`: correlated - * `"o"`: correlated occupied - * `"v"`: correlated virtual - * `"O"`: active occupied - * `"V"`: active virtual - * `"i"`: inactive occupied - * `"a"`: inactive virtual - - Attributes - ---------- - ebcc : REBCC - The EBCC object. - array : np.ndarray, optional - The array of integrals in the MO basis. If provided, do not - perform just-in-time transformations but instead slice the - array. Default value is `None`. - slices : iterable of slice, optional - The slices to use for each dimension. If provided, the default - slices outlined above are used. - mo_coeff : np.ndarray, optional - The MO coefficients. If not provided, the MO coefficients from - `ebcc` are used. Default value is `None`. - """ - - def __init__(self, ebcc, array=None, slices=None, mo_coeff=None): - util.Namespace.__init__(self) - - self.mf = ebcc.mf - self.space = ebcc.space - self.slices = slices - self.mo_coeff = mo_coeff - self.array = array - - if self.mo_coeff is None: - self.mo_coeff = ebcc.mo_coeff - if not (isinstance(self.mo_coeff, (tuple, list)) or self.mo_coeff.ndim == 3): - self.mo_coeff = [self.mo_coeff] * 2 - - if self.slices is None: - self.slices = { - "x": self.space.correlated, - "o": self.space.correlated_occupied, - "v": self.space.correlated_virtual, - "O": self.space.active_occupied, - "V": self.space.active_virtual, - "i": self.space.inactive_occupied, - "a": self.space.inactive_virtual, - } - if not isinstance(self.slices, (tuple, list)): - self.slices = [self.slices] * 2 - - def __getattr__(self, key): - """Just-in-time attribute getter.""" - - if len(key) == 4: - e1 = getattr(self, key[:2]) - e2 = getattr(self, key[2:]) - return util.einsum("Qij,Qkl->ijkl", e1, e2) - elif len(key) == 3: - key = key[1:] - - if self.array is None: - if key not in self.__dict__.keys(): - coeffs = [] - for i, k in enumerate(key): - coeffs.append(self.mo_coeff[i][:, self.slices[i][k]]) - if precision.types[float] == np.float64: - ijslice = ( - 0, - coeffs[0].shape[-1], - coeffs[0].shape[-1], - coeffs[0].shape[-1] + coeffs[1].shape[-1], - ) - coeffs = np.concatenate(coeffs, axis=1) - block = ao2mo._ao2mo.nr_e2( - self.mf.with_df._cderi, coeffs, ijslice, aosym="s2", mosym="s1" - ) - block = block.reshape(-1, ijslice[1] - ijslice[0], ijslice[3] - ijslice[2]) - else: - shape = ( - self.mf.with_df.get_naoaux(), - self.mf.with_df.mol.nao_nr(), - self.mf.with_df.mol.nao_nr(), - ) - block = util.decompress_axes( - "Qpp", - self.mf.with_df._cderi.astype(precision.types[float]), - include_diagonal=True, - symmetry="+++", - shape=shape, - ) - block = util.einsum("Qpq,pi,qj->Qij", block, coeffs[0].conj(), coeffs[1]) - self.__dict__[key] = block - return self.__dict__[key] - else: - slices = [] - for i, k in enumerate(key): - slices.append(self.slices[i][k]) - si, sj = slices - block = self.array[:, si][:, :, sj] - return block - - __getitem__ = __getattr__ - - -class UCDERIs(ERIs): - """ - Cholesky decomposed ERI container class for `UEBCC`. Consists of a - just-in-time namespace containing blocks of the integrals. - - Attributes - ---------- - ebcc : UEBCC - The EBCC object. - array : iterable of np.ndarray, optional - The array of integrals in the MO basis for each spin. If - provided, do not perform just-in-time transformations but - instead slice the array. Default value is `None`. - slices : iterable of iterable of slice, optional - The slices to use for each spin and each dimension therein. - If provided, the default slices outlined above are used. - mo_coeff : iterable of np.ndarray, optional - The MO coefficients for each spin. If not provided, the MO - coefficients from `ebcc` are used. Default value is `None`. - """ - - def __init__(self, ebcc, array=None, slices=None, mo_coeff=None): - util.Namespace.__init__(self) - - self.mf = ebcc.mf - self.space = ebcc.space - self.slices = slices - self.mo_coeff = mo_coeff - - if self.mo_coeff is None: - self.mo_coeff = ebcc.mo_coeff - - if self.slices is None: - self.slices = [ - { - "x": space.correlated, - "o": space.correlated_occupied, - "v": space.correlated_virtual, - "O": space.active_occupied, - "V": space.active_virtual, - "i": space.inactive_occupied, - "a": space.inactive_virtual, - } - for space in self.space - ] - - if array is not None: - arrays = array - else: - arrays = (None, None) - - self.aa = RCDERIs( - ebcc, - arrays[0], - slices=[self.slices[0], self.slices[0]], - mo_coeff=[self.mo_coeff[0], self.mo_coeff[0]], - ) - self.bb = RCDERIs( - ebcc, - arrays[1], - slices=[self.slices[1], self.slices[1]], - mo_coeff=[self.mo_coeff[1], self.mo_coeff[1]], - ) - - def __getattr__(self, key): - """Just-in-time attribute getter.""" - - if len(key) == 4: - # Hacks in support for i.e. `UCDERIs.aaaa.ovov` - v1 = getattr(self, key[:2]) - v2 = getattr(self, key[2:]) - - class FakeCDERIs: - def __getattr__(self, key): - e1 = getattr(v1, key[:2]) - e2 = getattr(v2, key[2:]) - return util.einsum("Qij,Qkl->ijkl", e1, e2) - - __getitem__ = __getattr__ - - return FakeCDERIs() - - elif len(key) == 3: - key = key[1:] - return getattr(self, key) - - __getitem__ = __getattr__ diff --git a/ebcc/ham/cderis.py b/ebcc/ham/cderis.py new file mode 100644 index 00000000..5df3de67 --- /dev/null +++ b/ebcc/ham/cderis.py @@ -0,0 +1,128 @@ +"""Cholesky-decomposed electron repulsion integral containers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyscf import ao2mo + +from ebcc import numpy as np +from ebcc import util +from ebcc.ham.base import BaseERIs +from ebcc.precision import types + +if TYPE_CHECKING: + from ebcc.numpy.typing import NDArray + + +class RCDERIs(BaseERIs): + """Restricted Cholesky-decomposed ERIs container class. + + Attributes: + cc: Coupled cluster object. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. + array: ERIs in the MO basis. + """ + + def __getitem__(self, key: str, e2: bool = False) -> NDArray[float]: + """Just-in-time getter. + + Args: + key: Key to get. + e2: Whether the key is for the second electron. + + Returns: + CDERIs for the given spaces. + """ + if self.array is not None: + raise NotImplementedError("`array` is not supported for CDERIs.") + + if len(key) == 4: + e1 = self.__getitem__("Q" + key[:2]) + e2 = self.__getitem__("Q" + key[2:], e2=True) + return util.einsum("Qij,Qkl->ijkl", e1, e2) + elif len(key) == 3: + key = key[1:] + else: + raise KeyError("Key must be of length 3 or 4.") + key_e2 = f"{key}_{'e1' if not e2 else 'e2'}" + + if key_e2 not in self._members: + s = 0 if not e2 else 2 + coeffs = [ + self.mo_coeff[i + s][:, self.space[i + s].mask(k)].astype(np.float64) + for i, k in enumerate(key) + ] + if types[float] == np.float64: + ijslice = ( + 0, + coeffs[0].shape[-1], + coeffs[0].shape[-1], + coeffs[0].shape[-1] + coeffs[1].shape[-1], + ) + coeffs = np.concatenate(coeffs, axis=1) + block = ao2mo._ao2mo.nr_e2( + self.cc.mf.with_df._cderi, coeffs, ijslice, aosym="s2", mosym="s1" + ) + block = block.reshape(-1, ijslice[1] - ijslice[0], ijslice[3] - ijslice[2]) + else: + shape = ( + self.cc.mf.with_df.get_naoaux(), + self.cc.mf.with_df.mol.nao_nr(), + self.cc.mf.with_df.mol.nao_nr(), + ) + block = util.decompress_axes( + "Qpp", + self.cc.mf.with_df._cderi.astype(precesion.types[float]), + include_diagonal=True, + symmetry="+++", + shape=shape, + ) + block = util.einsum("Qpq,pi,qj->Qij", block, coeffs[s].conj(), coeffs[s + 1]) + self._members[key_e2] = block + + return self._members[key_e2] + + +class UCDERIs(BaseERIs): + """Unrestricted Cholesky-decomposed ERIs container class. + + Attributes: + cc: Coupled cluster object. + space: Space object for each index. + mo_coeff: Molecular orbital coefficients for each index. + array: ERIs in the MO basis. + """ + + def __getitem__(self, key: str) -> RCDERIs: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + CDERIs for the given spins. + """ + if len(key) == 3: + key = key[1:] + if key not in ("aa", "bb", "aaaa", "aabb", "bbaa", "bbbb"): + raise KeyError(f"Invalid key: {key}") + if len(key) == 2: + key = key + key + i = "ab".index(key[0]) + j = "ab".index(key[2]) + + if key not in self._members: + self._members[key] = RCDERIs( + self.cc, + space=(self.space[0][i], self.space[1][i], self.space[2][j], self.space[3][j]), + mo_coeff=( + self.mo_coeff[0][i], + self.mo_coeff[1][i], + self.mo_coeff[2][j], + self.mo_coeff[3][j], + ), + ) + + return self._members[key] From 66353a4f373e6a6767529fdad110a2a96ec6f48f Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sun, 4 Aug 2024 13:05:08 +0100 Subject: [PATCH 16/37] Move rest to core module --- ebcc/__init__.py | 13 +- ebcc/ansatz.py | 372 ---------------------------------------------- ebcc/base.py | 25 ---- ebcc/damping.py | 29 ---- ebcc/dump.py | 258 -------------------------------- ebcc/logging.py | 126 ---------------- ebcc/precision.py | 75 ---------- 7 files changed, 5 insertions(+), 893 deletions(-) delete mode 100644 ebcc/ansatz.py delete mode 100644 ebcc/base.py delete mode 100644 ebcc/damping.py delete mode 100644 ebcc/dump.py delete mode 100644 ebcc/logging.py delete mode 100644 ebcc/precision.py diff --git a/ebcc/__init__.py b/ebcc/__init__.py index 0f0f4b17..bbe305ce 100644 --- a/ebcc/__init__.py +++ b/ebcc/__init__.py @@ -61,7 +61,7 @@ from ebcc.cc.uebcc import UEBCC -def EBCC(mf, *args, **kwargs): +def EBCC(mf, *args, **kwargs): # type: ignore from pyscf import scf if isinstance(mf, scf.uhf.UHF): @@ -78,8 +78,8 @@ def EBCC(mf, *args, **kwargs): # --- Constructors for boson-free calculations: -def _factory(ansatz): - def constructor(mf, *args, **kwargs): +def _factory(ansatz): # type: ignore + def constructor(mf, *args, **kwargs): # type: ignore from pyscf import scf kwargs["ansatz"] = ansatz @@ -115,11 +115,8 @@ def constructor(mf, *args, **kwargs): # --- List available methods: -def available_models(verbose=True): # pragma: no cover - """List available coupled-cluster models for each of general (G), - restricted (R) and unrestricted (U) Hartree--Fock references. - """ - +def available_models(verbose=True): # type: ignore # pragma: no cover + """List available coupled-cluster models for each spin type.""" cd = os.path.dirname(os.path.realpath(__file__)) path = os.path.join(cd, "codegen") _, _, files = list(os.walk(path))[0] diff --git a/ebcc/ansatz.py b/ebcc/ansatz.py deleted file mode 100644 index bed095de..00000000 --- a/ebcc/ansatz.py +++ /dev/null @@ -1,372 +0,0 @@ -"""Ansatz definition.""" - -import importlib - -from ebcc import METHOD_TYPES, util - -named_ansatzes = { - "MP2": ("MP2", "", 0, 0), - "MP3": ("MP3", "", 0, 0), - "CCD": ("CCD", "", 0, 0), - "CCSD": ("CCSD", "", 0, 0), - "CCSDT": ("CCSDT", "", 0, 0), - "CCSDTQ": ("CCSDTQ", "", 0, 0), - "CCSD(T)": ("CCSD(T)", "", 0, 0), - "CC2": ("CC2", "", 0, 0), - "CC3": ("CC3", "", 0, 0), - "QCISD": ("QCISD", "", 0, 0), - "DCD": ("DCD", "", 0, 0), - "DCSD": ("DCSD", "", 0, 0), - "CCSDt'": ("CCSDt'", "", 0, 0), - "CCSD-S-1-1": ("CCSD", "S", 1, 1), - "CCSD-SD-1-1": ("CCSD", "SD", 1, 1), - "CCSD-SD-1-2": ("CCSD", "SD", 1, 2), -} - - -def name_to_identifier(name): - """ - Convert an ansatz name to an identifier. The identifier is used as for - the filename of the module containing the generated equations, where - the name may contain illegal characters. - - Parameters - ---------- - name : str - Name of the ansatz. - - Returns - ------- - iden : str - Identifier for the ansatz. - - Examples - -------- - >>> identifier_to_name("CCSD(T)") - CCSDxTx - >>> identifier_to_name("CCSD-SD-1-2") - CCSD_SD_1_2 - """ - - iden = name.replace("(", "x").replace(")", "x") - iden = iden.replace("[", "y").replace("]", "y") - iden = iden.replace("-", "_") - iden = iden.replace("'", "p") - - return iden - - -def identifity_to_name(iden): - """ - Convert an ansatz identifier to a name. Inverse operation of - `name_to_identifier`. - - Parameters - ---------- - iden : str - Identifier for the ansatz. - - Returns - ------- - name : str - Name of the ansatz. - - Examples - -------- - >>> identifier_to_name("CCSDxTx") - CCSD(T) - >>> identifier_to_name("CCSD_SD_1_2") - CCSD-SD-1-2 - """ - - name = iden.replace("-", "_") - while "x" in name: - name = name.replace("x", "(", 1).replace("x", ")", 1) - while "y" in name: - name = name.replace("y", "(", 1).replace("y", ")", 1) - name = name.replace("p", "'") - - return name - - -class Ansatz: - """ - Ansatz class. - - Parameters - ---------- - fermion_ansatz : str, optional - Fermionic ansatz. Default value is `"CCSD"`. - boson_ansatz : str, optional - Rank of bosonic excitations. Default value is `""`. - fermion_coupling_rank : int, optional - Rank of fermionic term in coupling. Default value is `0`. - boson_coupling_rank : int, optional - Rank of bosonic term in coupling. Default value is `0`. - density_fitting : bool, optional - Use density fitting. Default value is `False`. - module_name : str, optional - Name of the module containing the generated equations. If `None`, - the module name is generated from the ansatz name. Default value is - `None`. - """ - - def __init__( - self, - fermion_ansatz: str = "CCSD", - boson_ansatz: str = "", - fermion_coupling_rank: int = 0, - boson_coupling_rank: int = 0, - density_fitting: bool = False, - module_name: str = None, - ): - self.fermion_ansatz = fermion_ansatz - self.boson_ansatz = boson_ansatz - self.fermion_coupling_rank = fermion_coupling_rank - self.boson_coupling_rank = boson_coupling_rank - self.density_fitting = density_fitting - self.module_name = module_name - - def _get_eqns(self, prefix): - """Get the module containing the generated equations.""" - - if self.module_name is None: - name = prefix + name_to_identifier(self.name) - else: - name = self.module_name - - eqns = importlib.import_module("ebcc.codegen.%s" % name) - - return eqns - - @classmethod - def from_string(cls, string, density_fitting=False): - """ - Build an Ansatz from a string for the default ansatzes. - - Parameters - ---------- - input : str - Input string - density_fitting : bool, optional - Use density fitting. Default value is `False`. - - Returns - ------- - ansatz : Ansatz - Ansatz object - """ - - if string not in named_ansatzes: - raise util.ModelNotImplemented(string) - - return cls(*named_ansatzes[string], density_fitting=density_fitting) - - def __repr__(self): - """ - Get a string with the name of the method. - - Returns - ------- - name : str - Name of the method. - """ - name = "" - if self.density_fitting: - name += "DF" - name += self.fermion_ansatz - if self.boson_ansatz: - name += "-%s" % self.boson_ansatz - if self.fermion_coupling_rank or self.boson_coupling_rank: - name += "-%d" % self.fermion_coupling_rank - name += "-%d" % self.boson_coupling_rank - return name - - @property - def name(self): - """Get the name of the ansatz.""" - return repr(self) - - @property - def has_perturbative_correction(self): - """ - Return a boolean indicating whether the ansatz includes a - perturbative correction e.g. CCSD(T). - - Returns - ------- - perturbative : bool - Boolean indicating if the ansatz is perturbatively - corrected. - """ - return any( - "(" in ansatz and ")" in ansatz for ansatz in (self.fermion_ansatz, self.boson_ansatz) - ) - - @property - def is_one_shot(self): - """ - Return a boolean indicating whether the ansatz is simply a one-shot - energy calculation e.g. MP2. - - Returns - ------- - one_shot : bool - Boolean indicating if the ansatz is a one-shot energy - calculation. - """ - return all( - ansatz.startswith("MP") or ansatz == "" - for ansatz in (self.fermion_ansatz, self.boson_ansatz) - ) - - def fermionic_cluster_ranks(self, spin_type="G"): - """ - Get a list of cluster operator ranks for the fermionic space. - - Parameters - ---------- - spin_type : str, optional - Spin type of the cluster operator. Default value is `"G"`. - - Returns - ------- - ranks : list of tuples - List of cluster operator ranks, each element is a tuple - containing the name, the slices and the rank. - """ - - ranks = [] - if not self.fermion_ansatz: - return ranks - - notations = { - "S": [("t1", "ov", 1)], - "D": [("t2", "oovv", 2)], - "T": [("t3", "ooovvv", 3)], - "t": [("t3", "ooOvvV", 3)], - "t'": [("t3", "OOOVVV", 3)], - } - if spin_type == "R": - notations["Q"] = [("t4a", "oooovvvv", 4), ("t4b", "oooovvvv", 4)] - else: - notations["Q"] = [("t4", "oooovvvv", 4)] - notations["2"] = notations["S"] + notations["D"] - notations["3"] = notations["2"] + notations["T"] - notations["4"] = notations["3"] + notations["Q"] - - # Remove any perturbative corrections - op = self.fermion_ansatz - while "(" in op: - start = op.index("(") - end = op.index(")") - op = op[:start] - if (end + 1) < len(op): - op += op[end + 1 :] - - # Check in order of longest to shortest string in case one - # method name starts with a substring equal to the name of - # another method - for method_type in sorted(METHOD_TYPES, key=len)[::-1]: - if op.startswith(method_type): - op = op.replace(method_type, "", 1) - break - - # If it's MP we only ever need to initialise second-order - # amplitudes - if method_type == "MP": - op = "D" - - # Determine the ranks - for key in sorted(notations.keys(), key=len)[::-1]: - if key in op: - ranks += notations[key] - op = op.replace(key, "") - - # Check there are no duplicates - if len(ranks) != len(set(ranks)): - raise util.ModelNotImplemented("Duplicate ranks in %s" % self.fermion_ansatz) - - # Sort the ranks by the cluster operator dimension - ranks = sorted(ranks, key=lambda x: x[2]) - - return ranks - - def bosonic_cluster_ranks(self, spin_type="G"): - """ - Get a list of cluster operator ranks for the bosonic space. - - Parameters - ---------- - spin_type : str, optional - Spin type of the cluster operator. Default value is `"G"`. - - Returns - ------- - ranks : list of tuples - List of cluster operator ranks, each element is a tuple - containing the name, the slices and the rank. - """ - - ranks = [] - if not self.boson_ansatz: - return ranks - - notations = { - "S": [("s1", "b", 1)], - "D": [("s2", "bb", 2)], - "T": [("s3", "bbb", 3)], - } - notations["2"] = notations["S"] + notations["D"] - notations["3"] = notations["2"] + notations["T"] - - # Remove any perturbative corrections - op = self.boson_ansatz - while "(" in op: - start = op.index("(") - end = op.index(")") - op = op[:start] - if (end + 1) < len(op): - op += op[end + 1 :] - - # Determine the ranks - for key in sorted(notations.keys(), key=len)[::-1]: - if key in op: - ranks += notations[key] - op = op.replace(key, "") - - # Check there are no duplicates - if len(ranks) != len(set(ranks)): - raise util.ModelNotImplemented("Duplicate ranks in %s" % self.boson_ansatz) - - # Sort the ranks by the cluster operator dimension - ranks = sorted(ranks, key=lambda x: x[2]) - - return ranks - - def coupling_cluster_ranks(self, spin_type="G"): - """ - Get a list of cluster operator ranks for the coupling between - fermionic and bosonic spaces. - - Parameters - ---------- - spin_type : str, optional - Spin type of the cluster operator. Default value is `"G"`. - - Returns - ------- - ranks : list of tuple - List of cluster operator ranks, each element is a tuple - containing the name, slice, fermionic rank and bosonic rank. - """ - - ranks = [] - - for fermion_rank in range(1, self.fermion_coupling_rank + 1): - for boson_rank in range(1, self.boson_coupling_rank + 1): - name = f"u{fermion_rank}{boson_rank}" - key = "b" * boson_rank + "o" * fermion_rank + "v" * fermion_rank - ranks.append((name, key, fermion_rank, boson_rank)) - - return ranks diff --git a/ebcc/base.py b/ebcc/base.py deleted file mode 100644 index e8e64fe6..00000000 --- a/ebcc/base.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Base classes.""" - -from __future__ import annotations - -from abc import ABC -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from typing import Any - - -class ERIs(ABC): - """Base class for electronic repulsion integrals.""" - - def __getitem__(self, key: str) -> Any: - """Get an item.""" - return self.__dict__[key] - - -class Fock(ABC): - """Base class for Fock matrices.""" - - def __getitem__(self, key: str) -> Any: - """Get an item.""" - return self.__dict__[key] diff --git a/ebcc/damping.py b/ebcc/damping.py deleted file mode 100644 index a8805a41..00000000 --- a/ebcc/damping.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Damping and DIIS control.""" - -from pyscf.lib import diis - - -class DIIS(diis.DIIS): - """Direct inversion in the iterative subspace.""" - - def __init__(self, space=6, min_space=1, damping=0.0): - super().__init__(incore=True) - self.verbose = 0 - self.space = space - self.min_space = min_space - self.damping = damping - - def update(self, x, xerr=None): - """Extrapolate a vector.""" - - # Extrapolate the vector - x = super().update(x, xerr=xerr) - - # Apply damping - if self.damping: - nd = self.get_num_vec() - if nd > 1: - xprev = self.get_vec(self.get_num_vec() - 1) - x = (1.0 - self.damping) * x + self.damping * xprev - - return x diff --git a/ebcc/dump.py b/ebcc/dump.py deleted file mode 100644 index c514b1b6..00000000 --- a/ebcc/dump.py +++ /dev/null @@ -1,258 +0,0 @@ -"""File dumping and reading functionality.""" - -from pyscf import scf -from pyscf.lib.chkfile import dump, dump_mol, load, load_mol - -from ebcc import util -from ebcc.ansatz import Ansatz -from ebcc.ham import Space - - -class Dump: - """ - File handler for reading and writing EBCC calculations. - - Attributes - ---------- - name : str - The name of the file. - """ - - def __init__(self, name): - self.name = name - - def write(self, ebcc): - """ - Write the EBCC object to the file. - - Parameters - ---------- - ebcc : EBCC - The EBCC object to write. - """ - - # Write the options - dic = {} - for key, val in ebcc.options.__dict__.items(): - if val is not None: - dic[key] = val - dump(self.name, "options", dic) - - # Write the mean-field data - dic = { - "e_tot": ebcc.mf.e_tot, - "mo_energy": ebcc.mf.mo_energy, - "mo_coeff": ebcc.mf.mo_coeff, - "mo_occ": ebcc.mf.mo_occ, - } - dump_mol(ebcc.mf.mol, self.name) - dump(self.name, "mean-field", dic) - - # Write the MOs used - dic = { - "mo_coeff": ebcc.mo_coeff, - "mo_occ": ebcc.mo_occ, - } - dump(self.name, "mo", dic) - - # Write the ansatz - dic = { - "fermion_ansatz": ebcc.ansatz.fermion_ansatz, - "boson_ansatz": ebcc.ansatz.boson_ansatz, - "fermion_coupling_rank": ebcc.ansatz.fermion_coupling_rank, - "boson_coupling_rank": ebcc.ansatz.boson_coupling_rank, - } - if ebcc.ansatz.module_name is not None: - dic["module_name"] = ebcc.ansatz.module_name - dump(self.name, "ansatz", dic) - - # Write the space - if ebcc.spin_type == "U": - dic = { - "occupied": (ebcc.space[0]._occupied, ebcc.space[1]._occupied), - "frozen": (ebcc.space[0]._frozen, ebcc.space[1]._frozen), - "active": (ebcc.space[0]._active, ebcc.space[1]._active), - } - else: - dic = { - "occupied": ebcc.space._occupied, - "frozen": ebcc.space._frozen, - "active": ebcc.space._active, - } - dump(self.name, "space", dic) - - # Write the bosonic parameters - dic = {} - if ebcc.omega is not None: - dic["omega"] = ebcc.omega - if ebcc.bare_g is not None: - dic["bare_g"] = ebcc.bare_g - if ebcc.bare_G is not None: - dic["bare_G"] = ebcc.bare_G - dump(self.name, "bosons", dic) - - # Write the Fock matrix - # TODO write the Fock matrix class instead - - # Write miscellaneous data - kwargs = { - "spin_type": ebcc.spin_type, - } - if ebcc.e_corr is not None: - kwargs["e_corr"] = ebcc.e_corr - if ebcc.converged is not None: - kwargs["converged"] = ebcc.converged - if ebcc.converged_lambda is not None: - kwargs["converged_lambda"] = ebcc.converged_lambda - dump(self.name, "misc", kwargs) - - # Write the amplitudes - if ebcc.spin_type == "U": - if ebcc.amplitudes is not None: - dump( - self.name, - "amplitudes", - { - key: ({**val} if isinstance(val, (util.Namespace, dict)) else val) - for key, val in ebcc.amplitudes.items() - }, - ) - if ebcc.lambdas is not None: - dump( - self.name, - "lambdas", - { - key: ({**val} if isinstance(val, (util.Namespace, dict)) else val) - for key, val in ebcc.lambdas.items() - }, - ) - else: - if ebcc.amplitudes is not None: - dump(self.name, "amplitudes", {**ebcc.amplitudes}) - if ebcc.lambdas is not None: - dump(self.name, "lambdas", {**ebcc.lambdas}) - - def read(self, cls, log=None): - """ - Load the file to an EBCC object. - - Parameters - ---------- - cls : type - EBCC class to load the file to. - log : Logger, optional - Logger to assign to the EBCC object. - - Returns - ------- - ebcc : EBCC - The EBCC object loaded from the file. - """ - - # Load the options - dic = load(self.name, "options") - options = cls.Options() - for key, val in dic.items(): - setattr(options, key, val) - - # Load the miscellaneous data - misc = load(self.name, "misc") - spin_type = misc.pop("spin_type").decode("ascii") - - # Load the mean-field data - mf_cls = {"G": scf.GHF, "U": scf.UHF, "R": scf.RHF}[spin_type] - mol = load_mol(self.name) - dic = load(self.name, "mean-field") - mf = mf_cls(mol) - mf.__dict__.update(dic) - - # Load the MOs used - dic = load(self.name, "mo") - mo_coeff = dic.get("mo_coeff", None) - mo_occ = dic.get("mo_occ", None) - - # Load the ansatz - dic = load(self.name, "ansatz") - module_name = dic.get("module_name", None) - if isinstance(module_name, str): - module_name = module_name.encode("ascii") - ansatz = Ansatz( - dic.get("fermion_ansatz", b"CCSD").decode("ascii"), - dic.get("boson_ansatz", b"").decode("ascii"), - dic.get("fermion_coupling_rank", 0), - dic.get("boson_coupling_rank", 0), - module_name, - ) - - # Load the space - dic = load(self.name, "space") - if spin_type == "U": - space = ( - Space( - dic.get("occupied", None)[0], - dic.get("frozen", None)[0], - dic.get("active", None)[0], - ), - Space( - dic.get("occupied", None)[1], - dic.get("frozen", None)[1], - dic.get("active", None)[1], - ), - ) - else: - space = Space( - dic.get("occupied", None), - dic.get("frozen", None), - dic.get("active", None), - ) - - # Load the bosonic parameters - dic = load(self.name, "bosons") - omega = dic.get("omega", None) - bare_g = dic.get("bare_g", None) - bare_G = dic.get("bare_G", None) - - # Load the Fock matrix - # TODO load the Fock matrix class instead - - # Load the amplitudes - amplitudes = load(self.name, "amplitudes") - lambdas = load(self.name, "lambdas") - if spin_type == "U": - if amplitudes is not None: - amplitudes = { - key: (util.Namespace(**val) if isinstance(val, dict) else val) - for key, val in amplitudes.items() - } - amplitudes = util.Namespace(**amplitudes) - if lambdas is not None: - lambdas = { - key: (util.Namespace(**val) if isinstance(val, dict) else val) - for key, val in lambdas.items() - } - lambdas = util.Namespace(**lambdas) - else: - if amplitudes is not None: - amplitudes = util.Namespace(**amplitudes) - if lambdas is not None: - lambdas = util.Namespace(**lambdas) - - # Initialise the EBCC object - cc = cls( - mf, - log=log, - ansatz=ansatz, - space=space, - omega=omega, - g=bare_g, - G=bare_G, - mo_coeff=mo_coeff, - mo_occ=mo_occ, - # fock=fock, - options=options, - ) - cc.__dict__.update(misc) - cc.amplitudes = amplitudes - cc.lambdas = lambdas - - return cc diff --git a/ebcc/logging.py b/ebcc/logging.py deleted file mode 100644 index 415b3ebd..00000000 --- a/ebcc/logging.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Logging.""" - -import logging -import os -import subprocess -import sys - -from ebcc import __version__ -from ebcc.util import Namespace - -HEADER = """ _ - | | - ___ | |__ ___ ___ - / _ \| '_ \ / __| / __| - | __/| |_) || (__ | (__ - \___||_.__/ \___| \___| -%s""" # noqa: W605 - - -class Logger(logging.Logger): - """Logger with a custom output level.""" - - def __init__(self, name, level=logging.INFO): - super().__init__(name, level) - - def output(self, msg, *args, **kwargs): - """Output a message at the `"OUTPUT"` level.""" - if self.isEnabledFor(25): - self._log(25, msg, args, **kwargs) - - -logging.setLoggerClass(Logger) -logging.addLevelName(25, "OUTPUT") - - -default_log = logging.getLogger(__name__) -default_log.setLevel(logging.INFO) -default_log.addHandler(logging.StreamHandler(sys.stderr)) - - -class NullLogger(Logger): - """A logger that does nothing.""" - - def __init__(self, *args, **kwargs): - super().__init__("null") - - def _log(self, level, msg, args, **kwargs): - pass - - -def init_logging(log): - """Initialise the logging with a header.""" - - if globals().get("_EBCC_LOG_INITIALISED", False): - return - - # Print header - header_size = max([len(line) for line in HEADER.split("\n")]) - space = " " * (header_size - len(__version__)) - log.info(f"{ANSI.B}{HEADER}{ANSI.R}" % f"{space}{ANSI.B}{__version__}{ANSI.R}") - - # Print versions of dependencies and ebcc - def get_git_hash(directory): - git_directory = os.path.join(directory, ".git") - cmd = ["git", "--git-dir=%s" % git_directory, "rev-parse", "--short", "HEAD"] - try: - git_hash = subprocess.check_output( - cmd, universal_newlines=True, stderr=subprocess.STDOUT - ).rstrip() - except subprocess.CalledProcessError: - git_hash = "N/A" - return git_hash - - import numpy - import pyscf - - log.info("numpy:") - log.info(" > Version: %s" % numpy.__version__) - log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(numpy.__file__), ".."))) - - log.info("pyscf:") - log.info(" > Version: %s" % pyscf.__version__) - log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(pyscf.__file__), ".."))) - - log.info("ebcc:") - log.info(" > Version: %s" % __version__) - log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(__file__), ".."))) - - # Environment variables - log.info("OMP_NUM_THREADS = %s" % os.environ.get("OMP_NUM_THREADS", "")) - - log.debug("") - - globals()["_EBCC_LOG_INITIALISED"] = True - - -def _check_output(*args, **kwargs): - """ - Call a command. If the return code is non-zero, an empty `bytes` - object is returned. - """ - try: - return subprocess.check_output(*args, **kwargs) - except subprocess.CalledProcessError: - return bytes() - - -ANSI = Namespace( - B="\x1b[1m", - H="\x1b[3m", - R="\x1b[m\x0f", - U="\x1b[4m", - b="\x1b[34m", - c="\x1b[36m", - g="\x1b[32m", - k="\x1b[30m", - m="\x1b[35m", - r="\x1b[31m", - w="\x1b[37m", - y="\x1b[33m", -) - - -def colour(text, *cs): - """Colour a string.""" - return f"{''.join([ANSI[c] for c in cs])}{text}{ANSI[None]}" diff --git a/ebcc/precision.py b/ebcc/precision.py deleted file mode 100644 index 3646ca62..00000000 --- a/ebcc/precision.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Floating point precision control.""" - -from __future__ import annotations - -from contextlib import contextmanager -from typing import TYPE_CHECKING - -from ebcc import numpy as np - -if TYPE_CHECKING: - from collections.abc import Iterator - from typing import Literal, Type, TypeVar, Union - - T = TypeVar("T", float, complex) - - -types: dict[type, type] -if TYPE_CHECKING: - types = { - float: float, - complex: complex, - } -else: - types = { - float: np.float64, - complex: np.complex128, - } - - -def cast(value: Any, dtype: Type[T]) -> T: - """Cast a value to the current floating point type. - - Args: - value: The value to cast. - dtype: The type to cast to. - - Returns: - The value cast to the current floating point type. - """ - return types[dtype](value) - - -def set_precision(**kwargs: type) -> None: - """Set the floating point type. - - Args: - float: The floating point type to use. - complex: The complex type to use. - """ - types[float] = kwargs.get("float", types[float]) - types[complex] = kwargs.get("complex", types[complex]) - - -@contextmanager -def precision(**kwargs: type) -> Iterator[None]: - """Context manager for setting the floating point precision. - - Args: - float: The floating point type to use. - complex: The complex type to use. - """ - old = { - "float": types[float], - "complex": types[complex], - } - set_precision(**kwargs) - yield - set_precision(**old) - - -@contextmanager -def single_precision() -> Iterator[None]: - """Context manager for setting the floating point precision to single precision.""" - with precision(float=np.float32, complex=np.complex64): - yield From 09cd0f8e2d0d51f9ad30b959b81ebc21376775bd Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sun, 4 Aug 2024 13:38:56 +0100 Subject: [PATCH 17/37] Finalising package structure --- ebcc/__init__.py | 12 +++++------- ebcc/cc/base.py | 12 ++++++------ ebcc/cc/gebcc.py | 4 ++-- ebcc/cc/rebcc.py | 4 ++-- ebcc/cc/uebcc.py | 4 ++-- ebcc/eom/base.py | 4 ++-- ebcc/eom/geom.py | 2 +- ebcc/eom/reom.py | 2 +- ebcc/eom/ueom.py | 2 +- ebcc/ham/cderis.py | 2 +- ebcc/ham/eris.py | 2 +- ebcc/ham/fock.py | 2 +- ebcc/ham/space.py | 2 +- ebcc/opt/base.py | 16 ++++++++-------- ebcc/opt/gbrueckner.py | 4 ++-- ebcc/opt/rbrueckner.py | 4 ++-- ebcc/opt/ubrueckner.py | 4 ++-- 17 files changed, 40 insertions(+), 42 deletions(-) diff --git a/ebcc/__init__.py b/ebcc/__init__.py index bbe305ce..05447777 100644 --- a/ebcc/__init__.py +++ b/ebcc/__init__.py @@ -41,12 +41,10 @@ import numpy -from ebcc.logging import NullLogger, default_log, init_logging +from ebcc.core import precision +from ebcc.core.logging import NullLogger, default_log, init_logging -# --- Import NumPy here to allow drop-in replacements - - -# --- Logging: +sys.modules["ebcc.precision"] = precision # Compatibility # --- Types of ansatz supporting by the EBCC solvers: @@ -108,8 +106,8 @@ def constructor(mf, *args, **kwargs): # type: ignore # --- Other imports: -from ebcc.ansatz import Ansatz -from ebcc.ham import Space +from ebcc.core.ansatz import Ansatz +from ebcc.ham.space import Space from ebcc.opt import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC # --- List available methods: diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index d0d16ac2..406a345b 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -9,19 +9,19 @@ from ebcc import default_log, init_logging from ebcc import numpy as np from ebcc import util -from ebcc.ansatz import Ansatz -from ebcc.damping import DIIS -from ebcc.dump import Dump -from ebcc.logging import ANSI -from ebcc.precision import cast, types +from ebcc.core.ansatz import Ansatz +from ebcc.core.damping import DIIS +from ebcc.core.dump import Dump +from ebcc.core.logging import ANSI +from ebcc.core.precision import cast, types if TYPE_CHECKING: from typing import Any, Callable, Literal, Optional, TypeVar, Union from pyscf.scf.hf import SCF # type: ignore + from ebcc.core.logging import Logger from ebcc.ham.base import BaseERIs, BaseFock - from ebcc.logging import Logger from ebcc.numpy.typing import NDArray # type: ignore from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index c9845b88..0784290e 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -9,12 +9,12 @@ from ebcc import numpy as np from ebcc import util from ebcc.cc.base import BaseEBCC +from ebcc.core.precision import types from ebcc.eom import EA_GEOM, EE_GEOM, IP_GEOM -from ebcc.ham import Space from ebcc.ham.eris import GERIs from ebcc.ham.fock import GFock +from ebcc.ham.space import Space from ebcc.opt.gbrueckner import BruecknerGEBCC -from ebcc.precision import types if TYPE_CHECKING: from typing import Optional diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 05d83938..71509f6f 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -9,13 +9,13 @@ from ebcc import numpy as np from ebcc import util from ebcc.cc.base import BaseEBCC +from ebcc.core.precision import types from ebcc.eom import EA_REOM, EE_REOM, IP_REOM -from ebcc.ham import Space from ebcc.ham.cderis import RCDERIs from ebcc.ham.eris import RERIs from ebcc.ham.fock import RFock +from ebcc.ham.space import Space from ebcc.opt.rbrueckner import BruecknerREBCC -from ebcc.precision import types if TYPE_CHECKING: from typing import Optional diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index bb86c21f..1f2d6d80 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -9,13 +9,13 @@ from ebcc import numpy as np from ebcc import util from ebcc.cc.base import BaseEBCC +from ebcc.core.precision import types from ebcc.eom import EA_UEOM, EE_UEOM, IP_UEOM -from ebcc.ham import Space from ebcc.ham.cderis import UCDERIs from ebcc.ham.eris import UERIs from ebcc.ham.fock import UFock +from ebcc.ham.space import Space from ebcc.opt.ubrueckner import BruecknerUEBCC -from ebcc.precision import types if TYPE_CHECKING: from ebcc.cc.rebcc import REBCC diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py index b489f2c7..7af4a0d7 100644 --- a/ebcc/eom/base.py +++ b/ebcc/eom/base.py @@ -10,8 +10,8 @@ from ebcc import numpy as np from ebcc import util -from ebcc.logging import ANSI -from ebcc.precision import types +from ebcc.core.logging import ANSI +from ebcc.core.precision import types if TYPE_CHECKING: from typing import Any, Callable, Optional diff --git a/ebcc/eom/geom.py b/ebcc/eom/geom.py index 32b25ea3..e5df0d17 100644 --- a/ebcc/eom/geom.py +++ b/ebcc/eom/geom.py @@ -6,8 +6,8 @@ from ebcc import numpy as np from ebcc import util +from ebcc.core.precision import types from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM -from ebcc.precision import types if TYPE_CHECKING: from typing import Optional diff --git a/ebcc/eom/reom.py b/ebcc/eom/reom.py index e69c1379..dc1687d9 100644 --- a/ebcc/eom/reom.py +++ b/ebcc/eom/reom.py @@ -6,8 +6,8 @@ from ebcc import numpy as np from ebcc import util +from ebcc.core.precision import types from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM -from ebcc.precision import types if TYPE_CHECKING: from typing import Optional diff --git a/ebcc/eom/ueom.py b/ebcc/eom/ueom.py index c031a2af..f1637a4b 100644 --- a/ebcc/eom/ueom.py +++ b/ebcc/eom/ueom.py @@ -6,8 +6,8 @@ from ebcc import numpy as np from ebcc import util +from ebcc.core.precision import types from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM -from ebcc.precision import types if TYPE_CHECKING: from typing import Optional diff --git a/ebcc/ham/cderis.py b/ebcc/ham/cderis.py index 5df3de67..7c767f03 100644 --- a/ebcc/ham/cderis.py +++ b/ebcc/ham/cderis.py @@ -8,8 +8,8 @@ from ebcc import numpy as np from ebcc import util +from ebcc.core.precision import types from ebcc.ham.base import BaseERIs -from ebcc.precision import types if TYPE_CHECKING: from ebcc.numpy.typing import NDArray diff --git a/ebcc/ham/eris.py b/ebcc/ham/eris.py index c9b5b2b1..3c199fd2 100644 --- a/ebcc/ham/eris.py +++ b/ebcc/ham/eris.py @@ -7,8 +7,8 @@ from pyscf import ao2mo from ebcc import numpy as np +from ebcc.core.precision import types from ebcc.ham.base import BaseERIs -from ebcc.precision import types if TYPE_CHECKING: from ebcc.numpy.typing import NDArray diff --git a/ebcc/ham/fock.py b/ebcc/ham/fock.py index 974bb914..8deb6f94 100644 --- a/ebcc/ham/fock.py +++ b/ebcc/ham/fock.py @@ -6,8 +6,8 @@ from ebcc import numpy as np from ebcc import util +from ebcc.core.precision import types from ebcc.ham.base import BaseFock -from ebcc.precision import types if TYPE_CHECKING: from ebcc.numpy.typing import NDArray diff --git a/ebcc/ham/space.py b/ebcc/ham/space.py index 512e7eea..906f7623 100644 --- a/ebcc/ham/space.py +++ b/ebcc/ham/space.py @@ -8,7 +8,7 @@ from ebcc import numpy as np from ebcc import util -from ebcc.precision import types +from ebcc.core.precision import types if TYPE_CHECKING: from typing import Optional, Union diff --git a/ebcc/opt/base.py b/ebcc/opt/base.py index cdf4d52e..49faabb1 100644 --- a/ebcc/opt/base.py +++ b/ebcc/opt/base.py @@ -10,9 +10,9 @@ from ebcc import numpy as np from ebcc import util -from ebcc.damping import DIIS -from ebcc.logging import ANSI, NullLogger, init_logging -from ebcc.precision import types +from ebcc.core.damping import DIIS +from ebcc.core.logging import ANSI, NullLogger, init_logging +from ebcc.core.precision import types if TYPE_CHECKING: from typing import Any, Optional, TypeVar @@ -170,13 +170,13 @@ def kernel(self) -> float: dt = self.get_t1_norm() # Log the iteration: - converged_e = de < self.options.e_tol - converged_t = dt < self.options.t_tol + converged_e = bool(de < self.options.e_tol) + converged_t = bool(dt < self.options.t_tol) self.log.info( f"{niter:4d} {self.cc.e_corr:16.10f} {self.cc.e_tot:18.10f}" - f" {[ANSI.r, ANSI.g][self.cc.converged]}{self.cc.converged!r:>8}{ANSI.R}" - f" {[ANSI.r, ANSI.g][converged_e]}{de:13.3e}{ANSI.R}" - f" {[ANSI.r, ANSI.g][converged_t]}{dt:13.3e}{ANSI.R}" + f" {[ANSI.r, ANSI.g][int(self.cc.converged)]}{self.cc.converged!r:>8}{ANSI.R}" + f" {[ANSI.r, ANSI.g][int(converged_e)]}{de:13.3e}{ANSI.R}" + f" {[ANSI.r, ANSI.g][int(converged_t)]}{dt:13.3e}{ANSI.R}" ) # Check for convergence: diff --git a/ebcc/opt/gbrueckner.py b/ebcc/opt/gbrueckner.py index a9a74183..87611cfd 100644 --- a/ebcc/opt/gbrueckner.py +++ b/ebcc/opt/gbrueckner.py @@ -8,14 +8,14 @@ from ebcc import numpy as np from ebcc import util +from ebcc.core.precision import types from ebcc.opt.base import BaseBruecknerEBCC -from ebcc.precision import types if TYPE_CHECKING: from typing import Optional from ebcc.cc.gebcc import AmplitudeType - from ebcc.damping import DIIS + from ebcc.core.damping import DIIS from ebcc.util import Namespace diff --git a/ebcc/opt/rbrueckner.py b/ebcc/opt/rbrueckner.py index e0d678df..196b2125 100644 --- a/ebcc/opt/rbrueckner.py +++ b/ebcc/opt/rbrueckner.py @@ -8,14 +8,14 @@ from ebcc import numpy as np from ebcc import util +from ebcc.core.precision import types from ebcc.opt.base import BaseBruecknerEBCC -from ebcc.precision import types if TYPE_CHECKING: from typing import Optional from ebcc.cc.rebcc import AmplitudeType - from ebcc.damping import DIIS + from ebcc.core.damping import DIIS from ebcc.util import Namespace diff --git a/ebcc/opt/ubrueckner.py b/ebcc/opt/ubrueckner.py index 649e2ab3..6ae487ee 100644 --- a/ebcc/opt/ubrueckner.py +++ b/ebcc/opt/ubrueckner.py @@ -8,14 +8,14 @@ from ebcc import numpy as np from ebcc import util +from ebcc.core.precision import types from ebcc.opt.base import BaseBruecknerEBCC -from ebcc.precision import types if TYPE_CHECKING: from typing import Optional from ebcc.cc.uebcc import AmplitudeType - from ebcc.damping import DIIS + from ebcc.core.damping import DIIS from ebcc.util import Namespace From bcdd02256b76694e833793d3b2d37fba49111079 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sun, 4 Aug 2024 14:56:00 +0100 Subject: [PATCH 18/37] Casting codegen outputs --- ebcc/__init__.py | 64 ++++++++++++++++----------------- ebcc/cc/base.py | 51 ++++++++++++++++----------- ebcc/cc/gebcc.py | 49 +++++++++++++++++++------- ebcc/cc/rebcc.py | 30 ++++++++++------ ebcc/cc/uebcc.py | 40 ++++++++++++++------- ebcc/eom/base.py | 74 +++++---------------------------------- ebcc/eom/geom.py | 1 + ebcc/eom/reom.py | 1 + ebcc/ham/cderis.py | 4 ++- ebcc/ham/eris.py | 2 ++ ebcc/ham/space.py | 25 +++++++------ ebcc/opt/gbrueckner.py | 1 + ebcc/opt/rbrueckner.py | 1 + ebcc/opt/ubrueckner.py | 1 + ebcc/util/permutations.py | 5 +-- pyproject.toml | 6 +++- 16 files changed, 184 insertions(+), 171 deletions(-) diff --git a/ebcc/__init__.py b/ebcc/__init__.py index 05447777..ec53ddec 100644 --- a/ebcc/__init__.py +++ b/ebcc/__init__.py @@ -34,33 +34,41 @@ >>> ccsd.kernel() """ +from __future__ import annotations + +"""Version of the package.""" __version__ = "1.4.3" +"""List of supported ansatz types.""" +METHOD_TYPES = ["MP", "CC", "LCC", "QCI", "QCC", "DC"] + import os import sys +from typing import TYPE_CHECKING import numpy +from ebcc.cc import GEBCC, REBCC, UEBCC from ebcc.core import precision +from ebcc.core.ansatz import Ansatz from ebcc.core.logging import NullLogger, default_log, init_logging +from ebcc.ham.space import Space +from ebcc.opt import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC -sys.modules["ebcc.precision"] = precision # Compatibility - +if TYPE_CHECKING: + from typing import Any, Callable -# --- Types of ansatz supporting by the EBCC solvers: + from pyscf.scf.hf import SCF # type: ignore -METHOD_TYPES = ["MP", "CC", "LCC", "QCI", "QCC", "DC"] + from ebcc.cc.base import BaseEBCC -# --- General constructor: +sys.modules["ebcc.precision"] = precision # Compatibility with older codegen versions -from ebcc.cc.gebcc import GEBCC -from ebcc.cc.rebcc import REBCC -from ebcc.cc.uebcc import UEBCC - -def EBCC(mf, *args, **kwargs): # type: ignore - from pyscf import scf +def EBCC(mf: SCF, *args: Any, **kwargs: Any) -> BaseEBCC: + """Construct an EBCC object for the given mean-field object.""" + from pyscf import scf # type: ignore if isinstance(mf, scf.uhf.UHF): return UEBCC(mf, *args, **kwargs) @@ -73,25 +81,20 @@ def EBCC(mf, *args, **kwargs): # type: ignore EBCC.__doc__ = REBCC.__doc__ -# --- Constructors for boson-free calculations: - - -def _factory(ansatz): # type: ignore - def constructor(mf, *args, **kwargs): # type: ignore - from pyscf import scf +def _factory(ansatz: str) -> Callable[[SCF, Any, Any], BaseEBCC]: + """Constructor for some specific ansatz.""" + from pyscf import scf # type: ignore + def constructor(mf: SCF, *args: Any, **kwargs: Any) -> BaseEBCC: + """Construct an EBCC object for the given mean-field object.""" kwargs["ansatz"] = ansatz if isinstance(mf, scf.uhf.UHF): - cc = UEBCC(mf, *args, **kwargs) + return UEBCC(mf, *args, **kwargs) elif isinstance(mf, scf.ghf.GHF): - cc = GEBCC(mf, *args, **kwargs) + return GEBCC(mf, *args, **kwargs) else: - cc = REBCC(mf, *args, **kwargs) - - cc.__doc__ = REBCC.__doc__ - - return cc + return REBCC(mf, *args, **kwargs) return constructor @@ -104,16 +107,9 @@ def constructor(mf, *args, **kwargs): # type: ignore del _factory -# --- Other imports: - -from ebcc.core.ansatz import Ansatz -from ebcc.ham.space import Space -from ebcc.opt import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC - -# --- List available methods: - - -def available_models(verbose=True): # type: ignore # pragma: no cover +def available_models( + verbose: bool = True, +) -> tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]]: # pragma: no cover """List available coupled-cluster models for each spin type.""" cd = os.path.dirname(os.path.realpath(__file__)) path = os.path.join(cd, "codegen") diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 406a345b..a0e8e202 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from ebcc import default_log, init_logging from ebcc import numpy as np @@ -13,7 +13,7 @@ from ebcc.core.damping import DIIS from ebcc.core.dump import Dump from ebcc.core.logging import ANSI -from ebcc.core.precision import cast, types +from ebcc.core.precision import astype, types if TYPE_CHECKING: from typing import Any, Callable, Literal, Optional, TypeVar, Union @@ -22,7 +22,7 @@ from ebcc.core.logging import Logger from ebcc.ham.base import BaseERIs, BaseFock - from ebcc.numpy.typing import NDArray # type: ignore + from ebcc.numpy.typing import NDArray from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace @@ -555,7 +555,7 @@ def energy( eris=eris, amplitudes=amplitudes, ) - return cast(func(**kwargs).real, float) + return astype(func(**kwargs).real, float) def energy_perturbative( self, @@ -579,7 +579,7 @@ def energy_perturbative( amplitudes=amplitudes, lambdas=lambdas, ) - return cast(func(**kwargs).real, float) + return astype(func(**kwargs).real, float) @abstractmethod def update_amps( @@ -626,7 +626,7 @@ def make_sing_b_dm( eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, - ) -> Any: + ) -> NDArray[float]: r"""Make the single boson density matrix :math:`\langle b \rangle`. Args: @@ -643,7 +643,7 @@ def make_sing_b_dm( amplitudes=amplitudes, lambdas=lambdas, ) - return func(**kwargs) + return cast(NDArray[float], func(**kwargs)) def make_rdm1_b( self, @@ -652,7 +652,7 @@ def make_rdm1_b( lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True, - ) -> Any: + ) -> NDArray[float]: r"""Make the one-particle boson reduced density matrix :math:`\langle b^+ c \rangle`. Args: @@ -672,7 +672,7 @@ def make_rdm1_b( amplitudes=amplitudes, lambdas=lambdas, ) - dm = func(**kwargs) + dm = cast(NDArray[float], func(**kwargs)) if hermitise: dm = 0.5 * (dm + dm.T) @@ -735,7 +735,15 @@ def make_eb_coup_rdm( unshifted: bool = True, hermitise: bool = True, ) -> Any: - r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. + r"""Make the electron-boson coupling reduced density matrix. + + .. math:: + \langle b^+ i^+ j \rangle + + and + + .. math:: + \langle b i^+ j \rangle Args: eris: Electron repulsion integrals. @@ -756,7 +764,7 @@ def hbar_matvec_ip( r2: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType]: + ) -> tuple[AmplitudeType, AmplitudeType]: """Compute the product between a state vector and the IP-EOM Hamiltonian. Args: @@ -776,7 +784,7 @@ def hbar_matvec_ip( r1=r1, r2=r2, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, AmplitudeType], func(**kwargs)) def hbar_matvec_ea( self, @@ -804,7 +812,7 @@ def hbar_matvec_ea( r1=r1, r2=r2, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, AmplitudeType], func(**kwargs)) def hbar_matvec_ee( self, @@ -832,7 +840,7 @@ def hbar_matvec_ee( r1=r1, r2=r2, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, AmplitudeType], func(**kwargs)) def make_ip_mom_bras( self, @@ -856,7 +864,7 @@ def make_ip_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, ...], func(**kwargs)) def make_ea_mom_bras( self, @@ -880,7 +888,7 @@ def make_ea_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, ...], func(**kwargs)) def make_ee_mom_bras( self, @@ -904,7 +912,7 @@ def make_ee_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, ...], func(**kwargs)) def make_ip_mom_kets( self, @@ -928,7 +936,7 @@ def make_ip_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, ...], func(**kwargs)) def make_ea_mom_kets( self, @@ -952,7 +960,7 @@ def make_ea_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, ...], func(**kwargs)) def make_ee_mom_kets( self, @@ -976,7 +984,7 @@ def make_ee_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - return func(**kwargs) + return cast(tuple[AmplitudeType, ...], func(**kwargs)) @abstractmethod def energy_sum(self, *args: str, signs_dict: Optional[dict[str, int]] = None) -> NDArray[float]: @@ -1192,6 +1200,7 @@ 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. @@ -1293,7 +1302,7 @@ def e_tot(self) -> float: Returns: Total energy. """ - return cast(self.mf.e_tot + self.e_corr, float) + return astype(self.mf.e_tot + self.e_corr, float) @property def t1(self) -> Any: diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 0784290e..29180109 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -17,12 +17,12 @@ from ebcc.opt.gbrueckner import BruecknerGEBCC if TYPE_CHECKING: - from typing import Optional + from typing import Any, Optional, cast from pyscf.scf.ghf import GHF from pyscf.scf.hf import SCF - from ebcc.cc.base import ERIsInputType + from ebcc.cc.base import BaseOptions, ERIsInputType from ebcc.cc.rebcc import REBCC from ebcc.cc.uebcc import UEBCC from ebcc.numpy.typing import NDArray @@ -54,6 +54,7 @@ class GEBCC(BaseEBCC): @property def spin_type(self): + """Get a string representation of the spin type.""" return "G" def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_GEOM: @@ -462,7 +463,7 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res = func(**kwargs) + res = cast(Namespace[AmplitudeType], func(**kwargs)) res = {key.rstrip("new"): val for key, val in res.items()} # Divide T amplitudes: @@ -514,7 +515,7 @@ def update_lams( amplitudes=amplitudes, lambdas=lambdas, ) - res = func(**kwargs) + res = cast(Namespace[AmplitudeType], func(**kwargs)) res = {key.rstrip("new"): val for key, val in res.items()} # Divide T amplitudes: @@ -551,7 +552,7 @@ def make_rdm1_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> Any: + ) -> NDArray[float]: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -569,22 +570,38 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = func(**kwargs) + dm = cast(NDArray[float], func(**kwargs)) if hermitise: dm = 0.5 * (dm + dm.T) return dm - def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): + def make_rdm2_f( + self, + eris: Optional[ERIsInputType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + lambdas: Optional[Namespace[AmplitudeType]] = None, + hermitise: bool = True, + ) -> NDArray[float]: + r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. + + Args: + eris: Electron repulsion integrals. + amplitudes: Cluster amplitudes. + lambdas: Cluster lambda amplitudes. + hermitise: Hermitise the density matrix. + + Returns: + Two-particle fermion reduced density matrix. + """ func, kwargs = self._load_function( "make_rdm2_f", eris=eris, amplitudes=amplitudes, lambdas=lambdas, ) - - dm = func(**kwargs) + dm = cast(NDArray[float], func(**kwargs)) if hermitise: dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) @@ -599,8 +616,16 @@ def make_eb_coup_rdm( lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True, - ) -> Any: - r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. + ) -> NDArray[float]: + r"""Make the electron-boson coupling reduced density matrix. + + .. math:: + \langle b^+ i^+ j \rangle + + and + + .. math:: + \langle b i^+ j \rangle Args: eris: Electron repulsion integrals. @@ -619,7 +644,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - dm_eb = func(**kwargs) + dm_eb = cast(NDArray[float], func(**kwargs)) if hermitise: dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 71509f6f..93227212 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -18,11 +18,11 @@ from ebcc.opt.rbrueckner import BruecknerREBCC if TYPE_CHECKING: - from typing import Optional + from typing import Any, Optional, Union, cast from pyscf.scf.hf import RHF, SCF - from ebcc.cc.base import ERIsInputType + from ebcc.cc.base import BaseOptions, ERIsInputType from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -237,7 +237,7 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res = func(**kwargs) + res = cast(Namespace[AmplitudeType], func(**kwargs)) res = {key.rstrip("new"): val for key, val in res.items()} # Divide T amplitudes: @@ -289,7 +289,7 @@ def update_lams( amplitudes=amplitudes, lambdas=lambdas, ) - res = func(**kwargs) + res = cast(Namespace[AmplitudeType], func(**kwargs)) res = {key.rstrip("new"): val for key, val in res.items()} # Divide T amplitudes: @@ -326,7 +326,7 @@ def make_rdm1_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> Any: + ) -> NDArray[float]: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -344,7 +344,7 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = func(**kwargs) + dm = cast(NDArray[float], func(**kwargs)) if hermitise: dm = 0.5 * (dm + dm.T) @@ -357,7 +357,7 @@ def make_rdm2_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> Any: + ) -> NDArray[float]: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -375,7 +375,7 @@ def make_rdm2_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = func(**kwargs) + dm = cast(NDArray[float], func(**kwargs)) if hermitise: dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) @@ -390,8 +390,16 @@ def make_eb_coup_rdm( lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True, - ) -> Any: - r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. + ) -> NDArray[float]: + r"""Make the electron-boson coupling reduced density matrix. + + .. math:: + \langle b^+ i^+ j \rangle + + and + + .. math:: + \langle b i^+ j \rangle Args: eris: Electron repulsion integrals. @@ -410,7 +418,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - dm_eb = func(**kwargs) + dm_eb = cast(NDArray[float], func(**kwargs)) if hermitise: dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 1f2d6d80..8cd29e1e 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -18,8 +18,17 @@ from ebcc.opt.ubrueckner import BruecknerUEBCC if TYPE_CHECKING: - from ebcc.cc.rebcc import REBCC + from typing import Any, Optional, Union, cast + + from pyscf.scf.hf import SCF + from pyscf.scf.uhf import UHF + + from ebcc.cc.base import ERIsInputType + from ebcc.cc.rebcc import REBCC, BaseOptions from ebcc.numpy.typing import NDArray + from ebcc.util import Namespace + + AmplitudeType = Namespace[NDArray[float]] class UEBCC(BaseEBCC): @@ -338,7 +347,7 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res = func(**kwargs) + res = cast(Namespace[AmplitudeType], func(**kwargs)) res = {key.rstrip("new"): val for key, val in res.items()} # Divide T amplitudes: @@ -399,7 +408,7 @@ def update_lams( lambdas=lambdas, lambdas_pert=lambdas_pert, ) - res = func(**kwargs) + res = cast(Namespace[AmplitudeType], func(**kwargs)) res = {key.rstrip("new"): val for key, val in res.items()} # Divide T amplitudes: @@ -442,7 +451,7 @@ def make_rdm1_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> Any: + ) -> Namespace[NDArray[float]]: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -460,8 +469,7 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - - dm = func(**kwargs) + dm = cast(Namespace[NDArray[float]], func(**kwargs)) if hermitise: dm.aa = 0.5 * (dm.aa + dm.aa.T) @@ -475,7 +483,7 @@ def make_rdm2_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> Any: + ) -> Namespace[NDArray[float]]: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -493,8 +501,7 @@ def make_rdm2_f( amplitudes=amplitudes, lambdas=lambdas, ) - - dm = func(**kwargs) + dm = cast(Namespace[NDArray[float]], func(**kwargs)) if hermitise: @@ -519,8 +526,16 @@ def make_eb_coup_rdm( lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True, - ) -> Any: - r"""Make the electron-boson coupling reduced density matrix :math:`\langle b^+ c^+ c b \rangle`. + ) -> Namespace[NDArray[float]]: + r"""Make the electron-boson coupling reduced density matrix. + + .. math:: + \langle b^+ i^+ j \rangle + + and + + .. math:: + \langle b i^+ j \rangle Args: eris: Electron repulsion integrals. @@ -539,8 +554,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - - dm_eb = func(**kwargs) + dm_eb = cast(Namespace[NDArray[float]], func(**kwargs)) if hermitise: dm_eb.aa[0] = 0.5 * (dm_eb.aa[0] + dm_eb.aa[1].transpose(0, 2, 1)) diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py index 7af4a0d7..50c82ebd 100644 --- a/ebcc/eom/base.py +++ b/ebcc/eom/base.py @@ -115,7 +115,7 @@ def name(self) -> str: return f"{self.excitation_type.upper()}-EOM-{self.spin_type}{self.ansatz.name}" @abstractmethod - def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: + def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -127,7 +127,7 @@ def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: pass @abstractmethod - def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, ...]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudeType, ...]: """Construct amplitudes from a vector. Args: @@ -264,7 +264,7 @@ def get_guesses(self, diag: Optional[NDArray[float]] = None) -> list[NDArray[flo return list(guesses) - def callback(self, envs: dict[str, Any]) -> None: + def callback(self, envs: dict[str, Any]) -> None: # noqa: B027 """Callback function for the eigensolver.""" # noqa: D401 pass @@ -392,7 +392,7 @@ def excitation_type(self) -> str: """Get the type of excitation.""" return "ip" - def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: + def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -403,7 +403,7 @@ def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: """ return self.ebcc.excitations_to_vector_ip(*amplitudes) - def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, ...]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudeType, ...]: """Construct amplitudes from a vector. Args: @@ -439,7 +439,7 @@ def excitation_type(self) -> str: """Get the type of excitation.""" return "ea" - def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: + def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -450,7 +450,7 @@ def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: """ return self.ebcc.excitations_to_vector_ea(*amplitudes) - def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, ...]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudeType, ...]: """Construct amplitudes from a vector. Args: @@ -477,36 +477,6 @@ def matvec( result = self.ebcc.hbar_matvec_ea(*amplitudes, eris=eris) return self.amplitudes_to_vector(*result) - def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: - """Get the bra vectors. - - Args: - eris: Electronic repulsion integrals. - - Returns: - Bra vectors. - """ - bras_raw = list(self.ebcc.make_ea_mom_bras(eris=eris)) - bras = np.array( - [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] - ) - return bras - - def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: - """Get the ket vectors. - - Args: - eris: Electronic repulsion integrals. - - Returns: - Ket vectors. - """ - kets_raw = list(self.ebcc.make_ea_mom_kets(eris=eris)) - kets = np.array( - [self.amplitudes_to_vector(*[k[..., i] for k in kets_raw]) for i in range(self.nmo)] - ) - return kets - class BaseEE_EOM(BaseEOM): """Base class for electron-electron EOM-CC.""" @@ -516,7 +486,7 @@ def excitation_type(self) -> str: """Get the type of excitation.""" return "ee" - def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: + def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -527,7 +497,7 @@ def amplitudes_to_vector(self, *amplitudes: AmplitudesType) -> NDArray[float]: """ return self.ebcc.excitations_to_vector_ee(*amplitudes) - def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudesType, ...]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudeType, ...]: """Construct amplitudes from a vector. Args: @@ -553,29 +523,3 @@ def matvec( amplitudes = self.vector_to_amplitudes(vector) result = self.ebcc.hbar_matvec_ee(*amplitudes, eris=eris) return self.amplitudes_to_vector(*result) - - def bras(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: - """Get the bra vectors. - - Args: - eris: Electronic repulsion integrals. - - Returns: - Bra vectors. - """ - bras_raw = list(self.ebcc.make_ee_mom_bras(eris=eris)) - bras = np.array( - [self.amplitudes_to_vector(*[b[i] for b in bras_raw]) for i in range(self.nmo)] - ) - return bras - - def kets(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: - """Get the ket vectors. - - Args: - eris: Electronic repulsion integrals. - - Returns: - Ket vectors. - """ - kets_raw = list(self.ebcc.make_ee_mom_kets(eris=eris)) diff --git a/ebcc/eom/geom.py b/ebcc/eom/geom.py index e5df0d17..5d42d64a 100644 --- a/ebcc/eom/geom.py +++ b/ebcc/eom/geom.py @@ -2,6 +2,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING from ebcc import numpy as np diff --git a/ebcc/eom/reom.py b/ebcc/eom/reom.py index dc1687d9..57752656 100644 --- a/ebcc/eom/reom.py +++ b/ebcc/eom/reom.py @@ -2,6 +2,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING from ebcc import numpy as np diff --git a/ebcc/ham/cderis.py b/ebcc/ham/cderis.py index 7c767f03..d73101ce 100644 --- a/ebcc/ham/cderis.py +++ b/ebcc/ham/cderis.py @@ -74,7 +74,7 @@ def __getitem__(self, key: str, e2: bool = False) -> NDArray[float]: ) block = util.decompress_axes( "Qpp", - self.cc.mf.with_df._cderi.astype(precesion.types[float]), + self.cc.mf.with_df._cderi.astype(types[float]), include_diagonal=True, symmetry="+++", shape=shape, @@ -95,6 +95,8 @@ class UCDERIs(BaseERIs): array: ERIs in the MO basis. """ + _members: dict[str, RCDERIs] + def __getitem__(self, key: str) -> RCDERIs: """Just-in-time getter. diff --git a/ebcc/ham/eris.py b/ebcc/ham/eris.py index 3c199fd2..726b50b8 100644 --- a/ebcc/ham/eris.py +++ b/ebcc/ham/eris.py @@ -11,6 +11,8 @@ from ebcc.ham.base import BaseERIs if TYPE_CHECKING: + from typing import Any, Optional + from ebcc.numpy.typing import NDArray diff --git a/ebcc/ham/space.py b/ebcc/ham/space.py index 906f7623..2bb245a0 100644 --- a/ebcc/ham/space.py +++ b/ebcc/ham/space.py @@ -17,20 +17,11 @@ from ebcc.cc.base import AmplitudeType from ebcc.numpy.typing import NDArray - - ConstructSpaceReturnType = Union[ - tuple[NDArray[float], NDArray[float], Space], - tuple[ - tuple[NDArray[float], NDArray[float]], - tuple[NDArray[float], NDArray[float]], - tuple[Space, Space], - ], - ] + from ebcc.util import Namespace class Space: - """ - Space class. + """Space class. ─┬─ ┌──────────┐ │ │ frozen │ @@ -304,6 +295,18 @@ def navir(self) -> int: return np.sum(self.active_virtual) +if TYPE_CHECKING: + # Needs to be defined after Space + ConstructSpaceReturnType = Union[ + tuple[NDArray[float], NDArray[float], Space], + tuple[ + tuple[NDArray[float], NDArray[float]], + tuple[NDArray[float], NDArray[float]], + tuple[Space, Space], + ], + ] + + def construct_default_space(mf: SCF) -> ConstructSpaceReturnType: """Construct a default space. diff --git a/ebcc/opt/gbrueckner.py b/ebcc/opt/gbrueckner.py index 87611cfd..761c4292 100644 --- a/ebcc/opt/gbrueckner.py +++ b/ebcc/opt/gbrueckner.py @@ -16,6 +16,7 @@ from ebcc.cc.gebcc import AmplitudeType from ebcc.core.damping import DIIS + from ebcc.numpy.typing import NDArray from ebcc.util import Namespace diff --git a/ebcc/opt/rbrueckner.py b/ebcc/opt/rbrueckner.py index 196b2125..f36f2864 100644 --- a/ebcc/opt/rbrueckner.py +++ b/ebcc/opt/rbrueckner.py @@ -16,6 +16,7 @@ from ebcc.cc.rebcc import AmplitudeType from ebcc.core.damping import DIIS + from ebcc.numpy.typing import NDArray from ebcc.util import Namespace diff --git a/ebcc/opt/ubrueckner.py b/ebcc/opt/ubrueckner.py index 6ae487ee..5f1fa272 100644 --- a/ebcc/opt/ubrueckner.py +++ b/ebcc/opt/ubrueckner.py @@ -16,6 +16,7 @@ from ebcc.cc.uebcc import AmplitudeType from ebcc.core.damping import DIIS + from ebcc.numpy.typing import NDArray from ebcc.util import Namespace diff --git a/ebcc/util/permutations.py b/ebcc/util/permutations.py index 8a7a03e0..46e2b565 100644 --- a/ebcc/util/permutations.py +++ b/ebcc/util/permutations.py @@ -9,7 +9,7 @@ import numpy as np if TYPE_CHECKING: - from typing import Any, Callable, Generator, Hashable, Iterable, Optional, TypeVar, Union + from typing import Any, Generator, Hashable, Iterable, Optional, TypeVar, Union from ebcc.numpy.typing import NDArray @@ -286,7 +286,8 @@ def compress_axes( Args: subscript: Subscript for the input array. array: Array to compress. - include_diagonal: Whether to include the diagonal elements of the input array in the output array. + include_diagonal: Whether to include the diagonal elements of the input array in the output + array. Returns: Compressed array. diff --git a/pyproject.toml b/pyproject.toml index 18d0621c..97a22e36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,7 +104,7 @@ ignore = [ "B007", # Loop control variable not used within the loop body ] per-file-ignores = [ - "__init__.py:E402,W605,F401,F811,D103", + "__init__.py:E402,W605,F401,F811,D103,D212,D415", ] docstring-convention = "google" count = true @@ -136,6 +136,10 @@ disallow_subclassing_any = false implicit_reexport = true warn_unused_ignores = false +[[tool.mypy.overrides]] +module = "ebcc.codegen.*" +ignore_errors = true + [[tool.mypy.overrides]] module = "scipy.*" ignore_missing_imports = true From e988c23fb00069b84b645b075ef78fdb629fb9e5 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sun, 4 Aug 2024 16:18:00 +0100 Subject: [PATCH 19/37] Some typing fixes --- ebcc/__init__.py | 2 +- ebcc/cc/base.py | 26 +++++++++++++++----------- ebcc/cc/gebcc.py | 23 ++++++++++++----------- ebcc/cc/rebcc.py | 10 +++++----- ebcc/cc/uebcc.py | 10 +++++----- 5 files changed, 38 insertions(+), 33 deletions(-) diff --git a/ebcc/__init__.py b/ebcc/__init__.py index ec53ddec..b034a3aa 100644 --- a/ebcc/__init__.py +++ b/ebcc/__init__.py @@ -48,10 +48,10 @@ import numpy +from ebcc.core.logging import NullLogger, default_log, init_logging from ebcc.cc import GEBCC, REBCC, UEBCC from ebcc.core import precision from ebcc.core.ansatz import Ansatz -from ebcc.core.logging import NullLogger, default_log, init_logging from ebcc.ham.space import Space from ebcc.opt import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index a0e8e202..19e92ee4 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -16,7 +16,7 @@ from ebcc.core.precision import astype, types if TYPE_CHECKING: - from typing import Any, Callable, Literal, Optional, TypeVar, Union + from typing import Any, Callable, Literal, Optional, TypeVar, Union, Generic, cast from pyscf.scf.hf import SCF # type: ignore @@ -26,9 +26,9 @@ from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace - ERIsInputType = Union[type[BaseERIs], NDArray[float]] - AmplitudeType = TypeVar("AmplitudeType") - SpaceType = TypeVar("SpaceType") + ERIsInputType = Union[BaseERIs, NDArray[float]] + AmplitudeType = Any + SpaceType = Any @dataclass @@ -80,6 +80,8 @@ class BaseEBCC(ABC): space: SpaceType amplitudes: Namespace[AmplitudeType] lambdas: Namespace[AmplitudeType] + g: Optional[NDArray[float]] + G: Optional[NDArray[float]] def __init__( self, @@ -125,16 +127,18 @@ def __init__( # Parameters: self.log = default_log if log is None else log self.mf = self._convert_mf(mf) - self._mo_coeff = np.asarray(mo_coeff).astype(types[float]) if mo_coeff is not None else None - self._mo_occ = np.asarray(mo_occ).astype(types[float]) if mo_occ is not None else None + self._mo_coeff: Optional[NDArray[float]] = np.asarray(mo_coeff).astype(types[float]) if mo_coeff is not None else None + self._mo_occ: Optional[NDArray[float]] = np.asarray(mo_occ).astype(types[float]) if mo_occ is not None else None # Ansatz: if isinstance(ansatz, Ansatz): self.ansatz = ansatz - else: + elif isinstance(ansatz, str): self.ansatz = Ansatz.from_string( ansatz, density_fitting=getattr(self.mf, "with_df", None) is not None ) + else: + raise TypeError("ansatz must be an Ansatz object or a string.") self._eqns = self.ansatz._get_eqns(self.spin_type) # Space: @@ -524,12 +528,12 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude @abstractmethod def init_lams( - self, amplitude: Optional[Namespace[AmplitudeType]] = None + self, amplitudes: Optional[Namespace[AmplitudeType]] = None ) -> Namespace[AmplitudeType]: """Initialise the cluster lambda amplitudes. Args: - amplitude: Cluster amplitudes. + amplitudes: Cluster amplitudes. Returns: Initial cluster lambda amplitudes. @@ -602,7 +606,7 @@ def update_amps( def update_lams( self, eris: ERIsInputType = None, - amplitude: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, lambdas_pert: Optional[Namespace[AmplitudeType]] = None, perturbative: bool = False, @@ -611,7 +615,7 @@ def update_lams( Args: eris: Electron repulsion integrals. - amplitude: Cluster amplitudes. + amplitudes: Cluster amplitudes. lambdas: Cluster lambda amplitudes. lambdas_pert: Perturbative cluster lambda amplitudes. perturbative: Flag to include perturbative correction. diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 29180109..aa690308 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -2,9 +2,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast -from pyscf import lib, scf +from pyscf import scf from ebcc import numpy as np from ebcc import util @@ -17,18 +17,19 @@ from ebcc.opt.gbrueckner import BruecknerGEBCC if TYPE_CHECKING: - from typing import Any, Optional, cast + from typing import Any, Optional, TypeAlias, Union from pyscf.scf.ghf import GHF from pyscf.scf.hf import SCF - from ebcc.cc.base import BaseOptions, ERIsInputType + from ebcc.cc.base import BaseOptions from ebcc.cc.rebcc import REBCC from ebcc.cc.uebcc import UEBCC from ebcc.numpy.typing import NDArray from ebcc.util import Namespace - AmplitudeType = NDArray[float] + ERIsInputType: TypeAlias = Union[GERIs, NDArray[float]] + AmplitudeType: TypeAlias = NDArray[float] class GEBCC(BaseEBCC): @@ -693,7 +694,7 @@ def next_char() -> str: energies.append(np.diag(self.fock[key + key])) subscript = "".join([signs_dict[k] + next_char() for k in subscript]) - energy_sum = lib.direct_sum(subscript, *energies) + energy_sum = util.direct_sum(subscript, *energies) return energy_sum @@ -967,7 +968,7 @@ def get_mean_field_G(self) -> NDArray[float]: Mean-field boson non-conserving term. """ # FIXME should this also sum in frozen orbitals? - val = lib.einsum("Ipp->I", self.g.boo) + val = util.einsum("Ipp->I", self.g.boo) val -= self.xi * self.omega if self.bare_G is not None: val += self.bare_G @@ -1027,7 +1028,7 @@ def xi(self) -> NDArray[float]: Shift in the bosonic operators. """ if self.options.shift: - xi = lib.einsum("Iii->I", self.g.boo) + xi = util.einsum("Iii->I", self.g.boo) xi /= self.omega if self.bare_G is not None: xi += self.bare_G / self.omega @@ -1064,7 +1065,7 @@ def nmo(self) -> int: Returns: Number of molecular orbitals. """ - return self.space.nmo + return cast(int, self.space.nmo) @property def nocc(self) -> int: @@ -1073,7 +1074,7 @@ def nocc(self) -> int: Returns: Number of occupied molecular orbitals. """ - return self.space.nocc + return cast(int, self.space.nocc) @property def nvir(self) -> int: @@ -1082,4 +1083,4 @@ def nvir(self) -> int: Returns: Number of virtual molecular orbitals. """ - return self.space.nvir + return cast(int, self.space.nvir) diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 93227212..76681907 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from pyscf import lib @@ -18,7 +18,7 @@ from ebcc.opt.rbrueckner import BruecknerREBCC if TYPE_CHECKING: - from typing import Any, Optional, Union, cast + from typing import Any, Optional, Union from pyscf.scf.hf import RHF, SCF @@ -823,7 +823,7 @@ def nmo(self) -> int: Returns: Number of molecular orbitals. """ - return self.space.nmo + return cast(int, self.space.nmo) @property def nocc(self) -> int: @@ -832,7 +832,7 @@ def nocc(self) -> int: Returns: Number of occupied molecular orbitals. """ - return self.space.nocc + return cast(int, self.space.nocc) @property def nvir(self) -> int: @@ -841,4 +841,4 @@ def nvir(self) -> int: Returns: Number of virtual molecular orbitals. """ - return self.space.nvir + return cast(int, self.space.nvir) diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 8cd29e1e..04856507 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from pyscf import lib @@ -18,7 +18,7 @@ from ebcc.opt.ubrueckner import BruecknerUEBCC if TYPE_CHECKING: - from typing import Any, Optional, Union, cast + from typing import Any, Optional, Union from pyscf.scf.hf import SCF from pyscf.scf.uhf import UHF @@ -1108,7 +1108,7 @@ def nmo(self) -> int: Number of molecular orbitals. """ assert self.mo_occ[0].size == self.mo_occ[1].size - return self.mo_occ[0].size + return cast(int, self.mo_occ[0].size) @property def nocc(self) -> tuple[int, int]: @@ -1117,7 +1117,7 @@ def nocc(self) -> tuple[int, int]: Returns: Number of occupied molecular orbitals for each spin. """ - return tuple(np.sum(mo_occ > 0) for mo_occ in self.mo_occ) + return cast(tuple[int], tuple(np.sum(mo_occ > 0) for mo_occ in self.mo_occ)) @property def nvir(self) -> tuple[int, int]: @@ -1126,4 +1126,4 @@ def nvir(self) -> tuple[int, int]: Returns: Number of virtual molecular orbitals for each spin. """ - return tuple(self.nmo - nocc for nocc in self.nocc) + return cast(tuple[int], tuple(self.nmo - nocc for nocc in self.nocc)) From b4035bcc183692b85e3d7d8f98f012dde07cfb6f Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Sun, 4 Aug 2024 17:46:36 +0100 Subject: [PATCH 20/37] More typing --- ebcc/cc/base.py | 59 ++++++++++++++++++--------------- ebcc/cc/gebcc.py | 55 +++++++------------------------ ebcc/cc/rebcc.py | 57 +++++++------------------------- ebcc/cc/uebcc.py | 85 ++++++++++++------------------------------------ ebcc/ham/base.py | 46 ++++++++++++++++++++++++++ 5 files changed, 123 insertions(+), 179 deletions(-) diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 19e92ee4..42cd2c54 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING from ebcc import default_log, init_logging from ebcc import numpy as np @@ -16,12 +16,12 @@ from ebcc.core.precision import astype, types if TYPE_CHECKING: - from typing import Any, Callable, Literal, Optional, TypeVar, Union, Generic, cast + from typing import Any, Callable, Literal, Optional, TypeVar, Union, Generic from pyscf.scf.hf import SCF # type: ignore from ebcc.core.logging import Logger - from ebcc.ham.base import BaseERIs, BaseFock + from ebcc.ham.base import BaseERIs, BaseFock, BaseElectronBoson from ebcc.numpy.typing import NDArray from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace @@ -36,22 +36,22 @@ class BaseOptions: """Options for EBCC calculations. Args: - 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. e_tol: Threshold for convergence in the correlation energy. t_tol: Threshold for convergence in the amplitude norm. max_iter: Maximum number of iterations. diis_space: Number of amplitudes to use in DIIS extrapolation. damping: Damping factor for DIIS extrapolation. + 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. """ - shift: bool = True e_tol: float = 1e-8 t_tol: float = 1e-8 max_iter: int = 200 diis_space: int = 12 damping: float = 0.0 + shift: bool = True class BaseEBCC(ABC): @@ -74,6 +74,7 @@ class BaseEBCC(ABC): ERIs: type[BaseERIs] Fock: type[BaseFock] CDERIs: type[BaseERIs] + ElectronBoson: type[BaseElectronBoson] Brueckner: type[BaseBruecknerEBCC] # Attributes @@ -156,7 +157,7 @@ def __init__( self.bare_g = g.astype(types[float]) if g is not None else None self.bare_G = G.astype(types[float]) if G is not None else None if self.boson_ansatz != "": - self.g = self.get_g(g) + self.g = self.get_g() self.G = self.get_mean_field_G() if self.options.shift: self.log.info(" > Energy shift due to polaritonic basis: %.10f", self.const) @@ -647,7 +648,8 @@ def make_sing_b_dm( amplitudes=amplitudes, lambdas=lambdas, ) - return cast(NDArray[float], func(**kwargs)) + res: NDArray[float] = func(**kwargs) + return res def make_rdm1_b( self, @@ -676,7 +678,7 @@ def make_rdm1_b( amplitudes=amplitudes, lambdas=lambdas, ) - dm = cast(NDArray[float], func(**kwargs)) + dm: NDArray[float] = func(**kwargs) if hermitise: dm = 0.5 * (dm + dm.T) @@ -788,7 +790,8 @@ def hbar_matvec_ip( r1=r1, r2=r2, ) - return cast(tuple[AmplitudeType, AmplitudeType], func(**kwargs)) + res: tuple[AmplitudeType, AmplitudeType] = func(**kwargs) + return res def hbar_matvec_ea( self, @@ -816,7 +819,8 @@ def hbar_matvec_ea( r1=r1, r2=r2, ) - return cast(tuple[AmplitudeType, AmplitudeType], func(**kwargs)) + res: tuple[AmplitudeType, AmplitudeType] = func(**kwargs) + return res def hbar_matvec_ee( self, @@ -844,7 +848,8 @@ def hbar_matvec_ee( r1=r1, r2=r2, ) - return cast(tuple[AmplitudeType, AmplitudeType], func(**kwargs)) + res: tuple[AmplitudeType, AmplitudeType] = func(**kwargs) + return res def make_ip_mom_bras( self, @@ -868,7 +873,8 @@ def make_ip_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - return cast(tuple[AmplitudeType, ...], func(**kwargs)) + res: tuple[AmplitudeType, ...] = func(**kwargs) + return res def make_ea_mom_bras( self, @@ -892,7 +898,8 @@ def make_ea_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - return cast(tuple[AmplitudeType, ...], func(**kwargs)) + res: tuple[AmplitudeType, ...] = func(**kwargs) + return res def make_ee_mom_bras( self, @@ -916,7 +923,8 @@ def make_ee_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - return cast(tuple[AmplitudeType, ...], func(**kwargs)) + res: tuple[AmplitudeType, ...] = func(**kwargs) + return res def make_ip_mom_kets( self, @@ -940,7 +948,8 @@ def make_ip_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - return cast(tuple[AmplitudeType, ...], func(**kwargs)) + res: tuple[AmplitudeType, ...] = func(**kwargs) + return res def make_ea_mom_kets( self, @@ -964,7 +973,8 @@ def make_ea_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - return cast(tuple[AmplitudeType, ...], func(**kwargs)) + res: tuple[AmplitudeType, ...] = func(**kwargs) + return res def make_ee_mom_kets( self, @@ -988,7 +998,8 @@ def make_ee_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - return cast(tuple[AmplitudeType, ...], func(**kwargs)) + res: tuple[AmplitudeType, ...] = func(**kwargs) + return res @abstractmethod def energy_sum(self, *args: str, signs_dict: Optional[dict[str, int]] = None) -> NDArray[float]: @@ -1180,19 +1191,15 @@ def get_eris(self, eris: Optional[ERIsInputType] = None) -> BaseERIs: """ pass - @abstractmethod - def get_g(self, g: NDArray[float]) -> Namespace[Any]: + def get_g(self) -> BaseElectronBoson: """Get the blocks of the electron-boson coupling matrix. This matrix corresponds to the bosonic annihilation operator. - Args: - g: Electron-boson coupling matrix. - Returns: - Blocks of the electron-boson coupling matrix. + Electron-boson coupling matrix. """ - pass + return self.ElectronBoson(self, array=self.bare_g) @abstractmethod def get_mean_field_G(self) -> Any: diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index aa690308..2ee30a4d 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -14,6 +14,7 @@ from ebcc.ham.eris import GERIs from ebcc.ham.fock import GFock from ebcc.ham.space import Space +from ebcc.ham.elbos import GElectronBoson from ebcc.opt.gbrueckner import BruecknerGEBCC if TYPE_CHECKING: @@ -51,6 +52,7 @@ class GEBCC(BaseEBCC): ERIs = GERIs Fock = GFock CDERIs = None + ElectronBoson = GElectronBoson Brueckner = BruecknerGEBCC @property @@ -464,8 +466,8 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res = cast(Namespace[AmplitudeType], func(**kwargs)) - res = {key.rstrip("new"): val for key, val in res.items()} + res: Namespace[AmplitudeType] = func(**kwargs) + res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -516,8 +518,8 @@ def update_lams( amplitudes=amplitudes, lambdas=lambdas, ) - res = cast(Namespace[AmplitudeType], func(**kwargs)) - res = {key.rstrip("new"): val for key, val in res.items()} + res: Namespace[AmplitudeType] = func(**kwargs) + res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -553,7 +555,7 @@ def make_rdm1_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> NDArray[float]: + ) -> AmplitudeType: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -571,7 +573,7 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = cast(NDArray[float], func(**kwargs)) + dm: AmpliduteType = func(**kwargs) if hermitise: dm = 0.5 * (dm + dm.T) @@ -584,7 +586,7 @@ def make_rdm2_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> NDArray[float]: + ) -> AmplitudeType: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -602,7 +604,7 @@ def make_rdm2_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = cast(NDArray[float], func(**kwargs)) + dm: AmpliduteType = func(**kwargs) if hermitise: dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) @@ -617,7 +619,7 @@ def make_eb_coup_rdm( lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True, - ) -> NDArray[float]: + ) -> AmplitudeType: r"""Make the electron-boson coupling reduced density matrix. .. math:: @@ -645,7 +647,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - dm_eb = cast(NDArray[float], func(**kwargs)) + dm_eb: AmpliduteType = func(**kwargs) if hermitise: dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) @@ -974,39 +976,6 @@ def get_mean_field_G(self) -> NDArray[float]: val += self.bare_G return val - def get_g(self, g: NDArray[float]) -> Namespace[NDArray[float]]: - """Get the blocks of the electron-boson coupling matrix. - - This matrix corresponds to the bosonic annihilation operator. - - Args: - g: Electron-boson coupling matrix. - - Returns: - Blocks of the electron-boson coupling matrix. - """ - # TODO make a proper class for this - slices = { - "x": self.space.correlated, - "o": self.space.correlated_occupied, - "v": self.space.correlated_virtual, - "O": self.space.active_occupied, - "V": self.space.active_virtual, - "i": self.space.inactive_occupied, - "a": self.space.inactive_virtual, - } - - class Blocks(util.Namespace): - def __getitem__(selffer, key): - assert key[0] == "b" - i = slices[key[1]] - j = slices[key[2]] - return g[:, i][:, :, j].copy() - - __getattr__ = __getitem__ - - return Blocks() - @property def bare_fock(self) -> NDArray[float]: """Get the mean-field Fock matrix in the MO basis, including frozen parts. diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 76681907..1dacc738 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -4,8 +4,6 @@ from typing import TYPE_CHECKING, cast -from pyscf import lib - from ebcc import numpy as np from ebcc import util from ebcc.cc.base import BaseEBCC @@ -15,6 +13,7 @@ from ebcc.ham.eris import RERIs from ebcc.ham.fock import RFock from ebcc.ham.space import Space +from ebcc.ham.elbos import RElectronBoson from ebcc.opt.rbrueckner import BruecknerREBCC if TYPE_CHECKING: @@ -48,6 +47,7 @@ class REBCC(BaseEBCC): ERIs = RERIs Fock = RFock CDERIs = RCDERIs + ElectronBoson = RElectronBoson Brueckner = BruecknerREBCC @property @@ -237,8 +237,8 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res = cast(Namespace[AmplitudeType], func(**kwargs)) - res = {key.rstrip("new"): val for key, val in res.items()} + res: Namespace[AmplitudeType] = func(**kwargs) + res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -289,8 +289,8 @@ def update_lams( amplitudes=amplitudes, lambdas=lambdas, ) - res = cast(Namespace[AmplitudeType], func(**kwargs)) - res = {key.rstrip("new"): val for key, val in res.items()} + res: Namespace[AmplitudeType] = func(**kwargs) + res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -344,7 +344,7 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = cast(NDArray[float], func(**kwargs)) + dm: NDArray[float] = func(**kwargs) if hermitise: dm = 0.5 * (dm + dm.T) @@ -375,7 +375,7 @@ def make_rdm2_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = cast(NDArray[float], func(**kwargs)) + dm: NDArray[float] = func(**kwargs) if hermitise: dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) @@ -418,7 +418,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - dm_eb = cast(NDArray[float], func(**kwargs)) + dm_eb: NDArray[float] = func(**kwargs) if hermitise: dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) @@ -467,7 +467,7 @@ def next_char() -> str: energies.append(np.diag(self.fock[key + key])) subscript = "".join([signs_dict[k] + next_char() for k in subscript]) - energy_sum = lib.direct_sum(subscript, *energies) + energy_sum = util.direct_sum(subscript, *energies) return energy_sum @@ -721,45 +721,12 @@ def get_mean_field_G(self) -> NDArray[float]: Mean-field boson non-conserving term. """ # FIXME should this also sum in frozen orbitals? - val = lib.einsum("Ipp->I", self.g.boo) * 2.0 + val = util.einsum("Ipp->I", self.g.boo) * 2.0 val -= self.xi * self.omega if self.bare_G is not None: val += self.bare_G return val - def get_g(self, g: NDArray[float]) -> Namespace[NDArray[float]]: - """Get the blocks of the electron-boson coupling matrix. - - This matrix corresponds to the bosonic annihilation operator. - - Args: - g: Electron-boson coupling matrix. - - Returns: - Blocks of the electron-boson coupling matrix. - """ - # TODO make a proper class for this - slices = { - "x": self.space.correlated, - "o": self.space.correlated_occupied, - "v": self.space.correlated_virtual, - "O": self.space.active_occupied, - "V": self.space.active_virtual, - "i": self.space.inactive_occupied, - "a": self.space.inactive_virtual, - } - - class Blocks(util.Namespace): - def __getitem__(selffer, key): - assert key[0] == "b" - i = slices[key[1]] - j = slices[key[2]] - return g[:, i][:, :, j].copy() - - __getattr__ = __getitem__ - - return Blocks() - @property def bare_fock(self) -> NDArray[float]: """Get the mean-field Fock matrix in the MO basis, including frozen parts. @@ -781,7 +748,7 @@ def xi(self) -> NDArray[float]: Shift in the bosonic operators. """ if self.options.shift: - xi = lib.einsum("Iii->I", self.g.boo) * 2.0 + xi = util.einsum("Iii->I", self.g.boo) * 2.0 xi /= self.omega if self.bare_G is not None: xi += self.bare_G / self.omega diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 04856507..336c8e86 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -4,8 +4,6 @@ from typing import TYPE_CHECKING, cast -from pyscf import lib - from ebcc import numpy as np from ebcc import util from ebcc.cc.base import BaseEBCC @@ -15,6 +13,7 @@ from ebcc.ham.eris import UERIs from ebcc.ham.fock import UFock from ebcc.ham.space import Space +from ebcc.ham.elbos import UElectronBoson from ebcc.opt.ubrueckner import BruecknerUEBCC if TYPE_CHECKING: @@ -50,6 +49,7 @@ class UEBCC(BaseEBCC): ERIs = UERIs Fock = UFock CDERIs = UCDERIs + ElectronBoson = UElectronBoson Brueckner = BruecknerUEBCC @property @@ -347,8 +347,8 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res = cast(Namespace[AmplitudeType], func(**kwargs)) - res = {key.rstrip("new"): val for key, val in res.items()} + res: Namespace[AmplitudeType] = func(**kwargs) + res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -408,8 +408,8 @@ def update_lams( lambdas=lambdas, lambdas_pert=lambdas_pert, ) - res = cast(Namespace[AmplitudeType], func(**kwargs)) - res = {key.rstrip("new"): val for key, val in res.items()} + res: Namespace[AmplitudeType] = func(**kwargs) + res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -451,7 +451,7 @@ def make_rdm1_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> Namespace[NDArray[float]]: + ) -> AmplitudeType: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -469,7 +469,7 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = cast(Namespace[NDArray[float]], func(**kwargs)) + dm: AmplitudeType = func(**kwargs) if hermitise: dm.aa = 0.5 * (dm.aa + dm.aa.T) @@ -483,7 +483,7 @@ def make_rdm2_f( amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, - ) -> Namespace[NDArray[float]]: + ) -> AmplitudeType: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -501,7 +501,7 @@ def make_rdm2_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm = cast(Namespace[NDArray[float]], func(**kwargs)) + dm: AmplitudeType = func(**kwargs) if hermitise: @@ -526,7 +526,7 @@ def make_eb_coup_rdm( lambdas: Optional[Namespace[AmplitudeType]] = None, unshifted: bool = True, hermitise: bool = True, - ) -> Namespace[NDArray[float]]: + ) -> AmplitudeType: r"""Make the electron-boson coupling reduced density matrix. .. math:: @@ -554,7 +554,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - dm_eb = cast(Namespace[NDArray[float]], func(**kwargs)) + dm_eb: AmplitudeType = func(**kwargs) if hermitise: dm_eb.aa[0] = 0.5 * (dm_eb.aa[0] + dm_eb.aa[1].transpose(0, 2, 1)) @@ -564,9 +564,9 @@ def make_eb_coup_rdm( if unshifted and self.options.shift: rdm1_f = self.make_rdm1_f(hermitise=hermitise) - shift = lib.einsum("x,ij->xij", self.xi, rdm1_f.aa) + shift = util.einsum("x,ij->xij", self.xi, rdm1_f.aa) dm_eb.aa -= shift[None] - shift = lib.einsum("x,ij->xij", self.xi, rdm1_f.bb) + shift = util.einsum("x,ij->xij", self.xi, rdm1_f.bb) dm_eb.bb -= shift[None] return dm_eb @@ -610,7 +610,7 @@ def next_char() -> str: energies.append(np.diag(self.fock[spin + spin][key + key])) subscript = "".join([signs_dict[k] + next_char() for k in subscript]) - energy_sum = lib.direct_sum(subscript, *energies) + energy_sum = util.direct_sum(subscript, *energies) return energy_sum @@ -984,8 +984,8 @@ def get_mean_field_G(self) -> NDArray[float]: Mean-field boson non-conserving term. """ # FIXME should this also sum in frozen orbitals? - val = lib.einsum("Ipp->I", self.g.aa.boo) - val += lib.einsum("Ipp->I", self.g.bb.boo) + val = util.einsum("Ipp->I", self.g.aa.boo) + val += util.einsum("Ipp->I", self.g.bb.boo) val -= self.xi * self.omega if self.bare_G is not None: # Require bare_G to have a spin index for now: @@ -993,51 +993,6 @@ def get_mean_field_G(self) -> NDArray[float]: val += self.bare_G return val - def get_g(self, g: NDArray[float]) -> Namespace[Namespace[NDArray[float]]]: - """Get the blocks of the electron-boson coupling matrix. - - This matrix corresponds to the bosonic annihilation operator. - - Args: - g: Electron-boson coupling matrix. - - Returns: - Blocks of the electron-boson coupling matrix. - """ - # TODO make a proper class for this - if np.array(g).ndim != 4: - g = np.array([g, g]) - slices = [ - { - "x": space.correlated, - "o": space.correlated_occupied, - "v": space.correlated_virtual, - "O": space.active_occupied, - "V": space.active_virtual, - "i": space.inactive_occupied, - "a": space.inactive_virtual, - } - for space in self.space - ] - - def constructor(s): - class Blocks(util.Namespace): - def __getitem__(selffer, key): - assert key[0] == "b" - i = slices[s][key[1]] - j = slices[s][key[2]] - return g[s][:, i][:, :, j].copy() - - __getattr__ = __getitem__ - - return Blocks() - - gs = util.Namespace() - gs.aa = constructor(0) - gs.bb = constructor(1) - - return gs - @property def bare_fock(self) -> Namespace[NDArray[float]]: """Get the mean-field Fock matrix in the MO basis, including frozen parts. @@ -1047,7 +1002,7 @@ def bare_fock(self) -> Namespace[NDArray[float]]: Returns: Mean-field Fock matrix. """ - fock = lib.einsum( + fock = util.einsum( "npq,npi,nqj->nij", self.mf.get_fock().astype(types[float]), self.mo_coeff, @@ -1064,8 +1019,8 @@ def xi(self) -> NDArray[float]: Shift in the bosonic operators. """ if self.options.shift: - xi = lib.einsum("Iii->I", self.g.aa.boo) - xi += lib.einsum("Iii->I", self.g.bb.boo) + xi = util.einsum("Iii->I", self.g.aa.boo) + xi += util.einsum("Iii->I", self.g.bb.boo) xi /= self.omega if self.bare_G is not None: xi += self.bare_G / self.omega diff --git a/ebcc/ham/base.py b/ebcc/ham/base.py index 1a7ba4e5..c9c2b939 100644 --- a/ebcc/ham/base.py +++ b/ebcc/ham/base.py @@ -113,3 +113,49 @@ def __getitem__(self, key: str) -> Any: Slice of the ERIs. """ pass + + +class BaseElectronBoson(Namespace, ABC): + """Base class for electron-boson coupling matrices. + + Attributes: + cc: Coupled cluster object. + space: Space object. + array: Electron-boson coupling matrix in the MO basis. + """ + + def __init__( + self, + cc: BaseEBCC, + array: Any = None, + space: tuple[Any] = None, + ) -> None: + """Initialise the electron-boson coupling matrix. + + Args: + cc: Coupled cluster object. + array: Electron-boson coupling matrix in the MO basis. + space: Space object for each index. + """ + Namespace.__init__(self) + + # 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() + + def _get_g(self) -> NDArray[float]: + """Get the electron-boson coupling matrix.""" + return self.cc.bare_g + + @abstractmethod + def __getitem__(self, key: str) -> Any: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + Slice of the electron-boson coupling matrix. + """ + pass From a01f87fff7b9765c43b70293ebe692bf49eb4868 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 19:13:25 +0100 Subject: [PATCH 21/37] More type hinting --- ebcc/cc/base.py | 39 +++++++++----- ebcc/cc/gebcc.py | 51 ++++++++++-------- ebcc/cc/rebcc.py | 39 +++++++++----- ebcc/cc/uebcc.py | 107 ++++++++++++++++++++++---------------- ebcc/util/permutations.py | 18 +++---- pyproject.toml | 4 ++ 6 files changed, 154 insertions(+), 104 deletions(-) diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 42cd2c54..8ffadcf0 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -18,7 +18,7 @@ if TYPE_CHECKING: from typing import Any, Callable, Literal, Optional, TypeVar, Union, Generic - from pyscf.scf.hf import SCF # type: ignore + from pyscf.scf.hf import SCF from ebcc.core.logging import Logger from ebcc.ham.base import BaseERIs, BaseFock, BaseElectronBoson @@ -26,7 +26,7 @@ from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace - ERIsInputType = Union[BaseERIs, NDArray[float]] + ERIsInputType = Any AmplitudeType = Any SpaceType = Any @@ -81,7 +81,7 @@ class BaseEBCC(ABC): space: SpaceType amplitudes: Namespace[AmplitudeType] lambdas: Namespace[AmplitudeType] - g: Optional[NDArray[float]] + g: Optional[BaseElectronBoson] G: Optional[NDArray[float]] def __init__( @@ -482,19 +482,14 @@ def _load_function( # Get the amplitudes: if not (amplitudes is False): - if not amplitudes: - amplitudes = self.amplitudes - if not amplitudes: - amplitudes = self.init_amps(eris=eris) - dicts.append(dict(amplitudes)) + dicts.append(dict(self._get_amps(amplitudes=amplitudes))) # Get the lambda amplitudes: if not (lambdas is False): - if not lambdas: - lambdas = self.lambdas - if not lambdas: - lambdas = self.init_lams(amplitudes=amplitudes if amplitudes else None) - dicts.append(dict(lambdas)) + if not (amplitudes is False): + dicts.append(dict(self._get_lams(lambdas=lambdas, amplitudes=amplitudes))) + else: + dicts.append(dict(self._get_lams(lambdas=lambdas))) # Get the function: func = getattr(self._eqns, name, None) @@ -541,6 +536,22 @@ def init_lams( """ pass + def _get_amps(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + """Get the cluster amplitudes, initialising if required.""" + if not amplitudes: + amplitudes = self.amplitudes + if not amplitudes: + amplitudes = self.init_amps() + return amplitudes + + def _get_lams(self, lambdas: Optional[Namespace[AmplitudeType]] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + """Get the cluster lambda amplitudes, initialising if required.""" + if not lambdas: + lambdas = self.lambdas + if not lambdas: + lambdas = self.init_lams(amplitudes=self._get_amps(amplitudes=amplitudes)) + return lambdas + def energy( self, eris: Optional[ERIsInputType] = None, @@ -1002,7 +1013,7 @@ def make_ee_mom_kets( return res @abstractmethod - def energy_sum(self, *args: str, signs_dict: Optional[dict[str, int]] = None) -> NDArray[float]: + def energy_sum(self, *args: str, signs_dict: Optional[dict[str, str]] = None) -> NDArray[float]: """Get a direct sum of energies. Args: diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 2ee30a4d..cfcb9145 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -51,12 +51,11 @@ class GEBCC(BaseEBCC): # Types ERIs = GERIs Fock = GFock - CDERIs = None ElectronBoson = GElectronBoson Brueckner = BruecknerGEBCC @property - def spin_type(self): + def spin_type(self) -> str: """Get a string representation of the spin type.""" return "G" @@ -146,8 +145,8 @@ def from_uebcc(cls, ucc: UEBCC) -> GEBCC: else: bare_g_a, bare_g_b = ucc.bare_g g = np.zeros((ucc.nbos, ucc.nmo * 2, ucc.nmo * 2)) - g[np.ix_(range(ucc.nbos), sa, sa)] = bare_g_a.copy() - g[np.ix_(range(ucc.nbos), sb, sb)] = bare_g_b.copy() + g[np.ix_(np.arange(ucc.nbos), sa, sa)] = bare_g_a.copy() + g[np.ix_(np.arange(ucc.nbos), sb, sb)] = bare_g_b.copy() else: g = None @@ -170,7 +169,7 @@ def from_uebcc(cls, ucc: UEBCC) -> GEBCC: has_lams = bool(ucc.lambdas) if has_amps: - amplitudes = util.Namespace() + amplitudes: Namespace[AmplitudeType] = util.Namespace() for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): shape = tuple(space.size(k) for k in key) @@ -211,7 +210,7 @@ def from_uebcc(cls, ucc: UEBCC) -> GEBCC: if combn in done: continue mask = np.ix_( - *([range(nbos)] * nb), + *([np.arange(nbos)] * nb), *[slices[s][k] for s, k in zip(combn, key[nb:])], ) transpose = ( @@ -283,7 +282,7 @@ def from_uebcc(cls, ucc: UEBCC) -> GEBCC: if combn in done: continue mask = np.ix_( - *([range(nbos)] * nb), + *([np.arange(nbos)] * nb), *[ slices[s][k] for s, k in zip(combn, key[nb + nf :] + key[nb : nb + nf]) @@ -374,7 +373,7 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude Initial cluster amplitudes. """ eris = self.get_eris(eris) - amplitudes = util.Namespace() + amplitudes: Namespace[AmplitudeType] = util.Namespace() # Build T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -388,6 +387,8 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude if self.boson_ansatz: # Only true for real-valued couplings: + assert self.g is not None + assert self.G is not None h = self.g H = self.G @@ -424,7 +425,7 @@ def init_lams( """ if amplitudes is None: amplitudes = self.amplitudes - lambdas = util.Namespace() + lambdas: Namespace[AmplitudeType] = util.Namespace() # Build L amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -461,6 +462,7 @@ def update_amps( Returns: Updated cluster amplitudes. """ + amplitudes = self._get_amps(amplitudes=amplitudes) func, kwargs = self._load_function( "update_amps", eris=eris, @@ -509,6 +511,8 @@ def update_lams( Updated cluster lambda amplitudes. """ # TODO active + amplitudes = self._get_amps(amplitudes=amplitudes) + lambdas = self._get_lams(lambdas=lambdas, amplitudes=amplitudes) if lambdas_pert is not None: lambdas.update(lambdas_pert) @@ -545,7 +549,7 @@ def update_lams( res[lname] += lambdas[lname] if perturbative: - res = {key + "pert": val for key, val in res.items()} + res = Namespace(**{key + "pert": val for key, val in res.items()}) return res @@ -573,7 +577,7 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm: AmpliduteType = func(**kwargs) + dm: AmplitudeType = func(**kwargs) if hermitise: dm = 0.5 * (dm + dm.T) @@ -604,7 +608,7 @@ def make_rdm2_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm: AmpliduteType = func(**kwargs) + dm: AmplitudeType = func(**kwargs) if hermitise: dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) @@ -647,7 +651,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - dm_eb: AmpliduteType = func(**kwargs) + dm_eb: AmplitudeType = func(**kwargs) if hermitise: dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) @@ -660,17 +664,18 @@ def make_eb_coup_rdm( return dm_eb - def energy_sum(self, subscript: str, signs_dict: dict[str, int] = None) -> NDArray[float]: + def energy_sum(self, *args: str, signs_dict: Optional[dict[str, str]] = None) -> NDArray[float]: """Get a direct sum of energies. Args: - subscript: Subscript for the direct sum. + *args: Energies to sum. Should specify a subscript only. signs_dict: Signs of the energies in the sum. Default sets `("o", "O", "i")` to be positive, and `("v", "V", "a", "b")` to be negative. Returns: Sum of energies. """ + subscript, = args n = 0 def next_char() -> str: @@ -731,7 +736,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp Returns: Cluster amplitudes. """ - amplitudes = util.Namespace() + amplitudes: Namespace[AmplitudeType] = util.Namespace() i0 = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -785,7 +790,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: Returns: Cluster lambda amplitudes. """ - lambdas = util.Namespace() + lambdas: Namespace[AmplitudeType] = util.Namespace() i0 = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -847,7 +852,7 @@ def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> ND """ return self.excitations_to_vector_ip(*excitations) - def excitations_to_vector_ee(self, *excitations): + def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: """Construct a vector containing all of the EE-EOM excitations. Args: @@ -969,6 +974,8 @@ def get_mean_field_G(self) -> NDArray[float]: Returns: Mean-field boson non-conserving term. """ + assert self.g is not None + assert self.omega is not None # FIXME should this also sum in frozen orbitals? val = util.einsum("Ipp->I", self.g.boo) val -= self.xi * self.omega @@ -996,7 +1003,9 @@ def xi(self) -> NDArray[float]: Returns: Shift in the bosonic operators. """ + assert self.omega is not None if self.options.shift: + assert self.g is not None xi = util.einsum("Iii->I", self.g.boo) xi /= self.omega if self.bare_G is not None: @@ -1022,10 +1031,10 @@ def get_eris(self, eris: Optional[ERIsInputType] = None) -> GERIs: Returns: Electron repulsion integrals. """ - if (eris is None) or isinstance(eris, np.ndarray): - return self.ERIs(self, array=eris) - else: + if isinstance(eris, GERIs): return eris + else: + return self.ERIs(self, array=eris) @property def nmo(self) -> int: diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 1dacc738..cec69224 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -17,15 +17,16 @@ from ebcc.opt.rbrueckner import BruecknerREBCC if TYPE_CHECKING: - from typing import Any, Optional, Union + from typing import Any, Optional, Union, TypeAlias from pyscf.scf.hf import RHF, SCF - from ebcc.cc.base import BaseOptions, ERIsInputType + from ebcc.cc.base import BaseOptions from ebcc.numpy.typing import NDArray from ebcc.util import Namespace - AmplitudeType = NDArray[float] + ERIsInputType: TypeAlias = Union[RERIs, RCDERIs, NDArray[float]] + AmplitudeType: TypeAlias = NDArray[float] class REBCC(BaseEBCC): @@ -51,7 +52,7 @@ class REBCC(BaseEBCC): Brueckner = BruecknerREBCC @property - def spin_type(self): + def spin_type(self) -> str: """Get a string representation of the spin type.""" return "R" @@ -144,7 +145,7 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude Initial cluster amplitudes. """ eris = self.get_eris(eris) - amplitudes = util.Namespace() + amplitudes: Namespace[AmplitudeType] = util.Namespace() # Build T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -159,6 +160,8 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude if self.boson_ansatz: # Only true for real-valued couplings: + assert self.g is not None + assert self.G is not None h = self.g H = self.G @@ -195,7 +198,7 @@ def init_lams( """ if amplitudes is None: amplitudes = self.amplitudes - lambdas = util.Namespace() + lambdas: Namespace[AmplitudeType] = util.Namespace() # Build L amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -232,6 +235,7 @@ def update_amps( Returns: Updated cluster amplitudes. """ + amplitudes = self._get_amps(amplitudes=amplitudes) func, kwargs = self._load_function( "update_amps", eris=eris, @@ -280,6 +284,8 @@ def update_lams( Updated cluster lambda amplitudes. """ # TODO active + amplitudes = self._get_amps(amplitudes=amplitudes) + lambdas = self._get_lams(lambdas=lambdas, amplitudes=amplitudes) if lambdas_pert is not None: lambdas.update(lambdas_pert) @@ -316,7 +322,7 @@ def update_lams( res[lname] += lambdas[lname] if perturbative: - res = {key + "pert": val for key, val in res.items()} + res = util.Namespace(**{key + "pert": val for key, val in res.items()}) return res @@ -431,17 +437,18 @@ def make_eb_coup_rdm( return dm_eb - def energy_sum(self, subscript: str, signs_dict: dict[str, int] = None) -> NDArray[float]: + def energy_sum(self, *args: str, signs_dict: Optional[dict[str, str]] = None) -> NDArray[float]: """Get a direct sum of energies. Args: - subscript: Subscript for the direct sum. + *args: Energies to sum. Should specify a subscript only. signs_dict: Signs of the energies in the sum. Default sets `("o", "O", "i")` to be positive, and `("v", "V", "a", "b")` to be negative. Returns: Sum of energies. """ + subscript, = args n = 0 def next_char() -> str: @@ -502,7 +509,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp Returns: Cluster amplitudes. """ - amplitudes = util.Namespace() + amplitudes: Namespace[AmplitudeType] = util.Namespace() i0 = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -556,7 +563,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: Returns: Cluster lambda amplitudes. """ - lambdas = util.Namespace() + lambdas: Namespace[AmplitudeType] = util.Namespace() i0 = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -721,6 +728,8 @@ def get_mean_field_G(self) -> NDArray[float]: Mean-field boson non-conserving term. """ # FIXME should this also sum in frozen orbitals? + assert self.omega is not None + assert self.g is not None val = util.einsum("Ipp->I", self.g.boo) * 2.0 val -= self.xi * self.omega if self.bare_G is not None: @@ -747,7 +756,9 @@ def xi(self) -> NDArray[float]: Returns: Shift in the bosonic operators. """ + assert self.omega is not None if self.options.shift: + assert self.g is not None xi = util.einsum("Iii->I", self.g.boo) * 2.0 xi /= self.omega if self.bare_G is not None: @@ -773,15 +784,15 @@ def get_eris(self, eris: Optional[ERIsInputType] = None) -> Union[RERIs, RCDERIs Returns: Electron repulsion integrals. """ - if (eris is None) or isinstance(eris, np.ndarray): + if isinstance(eris, (RERIs, RCDERIs)): + return eris + else: if (isinstance(eris, np.ndarray) and eris.ndim == 3) or getattr( self.mf, "with_df", None ): return self.CDERIs(self, array=eris) else: return self.ERIs(self, array=eris) - else: - return eris @property def nmo(self) -> int: diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 336c8e86..1f618eb5 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -17,17 +17,17 @@ from ebcc.opt.ubrueckner import BruecknerUEBCC if TYPE_CHECKING: - from typing import Any, Optional, Union + from typing import Any, Optional, Union, TypeAlias from pyscf.scf.hf import SCF from pyscf.scf.uhf import UHF - from ebcc.cc.base import ERIsInputType from ebcc.cc.rebcc import REBCC, BaseOptions from ebcc.numpy.typing import NDArray from ebcc.util import Namespace - AmplitudeType = Namespace[NDArray[float]] + ERIsInputType: TypeAlias = Union[UERIs, UCDERIs, tuple[NDArray[float], ...]] + AmplitudeType: TypeAlias = Union[NDArray[float], Namespace[NDArray[float]]] # S_{n} has no spin class UEBCC(BaseEBCC): @@ -53,7 +53,7 @@ class UEBCC(BaseEBCC): Brueckner = BruecknerUEBCC @property - def spin_type(self): + def spin_type(self) -> str: """Get a string representation of the spin type.""" return "U" @@ -127,12 +127,12 @@ def from_rebcc(cls, rcc: REBCC) -> UEBCC: has_lams = bool(rcc.lambdas) if has_amps: - amplitudes = util.Namespace() + amplitudes: Namespace[AmplitudeType] = util.Namespace() for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): amplitudes[name] = util.Namespace() for comb in util.generate_spin_combinations(n, unique=True): - subscript = util.combine_subscripts(key, comb) + subscript, _ = util.combine_subscripts(key, comb) tn = rcc.amplitudes[name] tn = util.symmetrise(subscript, tn, symmetry="-" * 2 * n) amplitudes[name][comb] = tn @@ -149,13 +149,13 @@ def from_rebcc(cls, rcc: REBCC) -> UEBCC: ucc.amplitudes = amplitudes if has_lams: - lambdas = util.Namespace() + lambdas: AmplitudeType = util.Namespace() for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): lname = name.replace("t", "l") lambdas[lname] = util.Namespace() for comb in util.generate_spin_combinations(n, unique=True): - subscript = util.combine_subscripts(key, comb) + subscript, _ = util.combine_subscripts(key, comb) tn = rcc.lambdas[lname] tn = util.symmetrise(subscript, tn, symmetry="-" * 2 * n) lambdas[lname][comb] = tn @@ -175,7 +175,7 @@ def from_rebcc(cls, rcc: REBCC) -> UEBCC: return ucc - def init_space(self) -> Namespace[Space]: + def init_space(self) -> tuple[Space, Space]: """Initialise the fermionic space. Returns: @@ -230,11 +230,11 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude Initial cluster amplitudes. """ eris = self.get_eris(eris) - amplitudes = util.Namespace() + amplitudes: Namespace[AmplitudeType] = util.Namespace() # Build T amplitudes for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - tn = util.Namespace() + tn: AmplitudeType = util.Namespace() for comb in util.generate_spin_combinations(n, unique=True): if n == 1: tn[comb] = self.fock[comb][key] / self.energy_sum(key, comb) @@ -252,6 +252,8 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude if self.boson_ansatz: # Only tue for real-valued couplings: + assert self.g is not None + assert self.G is not None h = self.g H = self.G @@ -301,7 +303,7 @@ def init_lams( """ if amplitudes is None: amplitudes = self.amplitudes - lambdas = util.Namespace() + lambdas: Namespace[AmplitudeType] = util.Namespace() # Build L amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -342,6 +344,7 @@ def update_amps( Returns: Updated cluster amplitudes. """ + amplitudes = self._get_amps(amplitudes=amplitudes) func, kwargs = self._load_function( "update_amps", eris=eris, @@ -353,7 +356,7 @@ def update_amps( # Divide T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): for comb in util.generate_spin_combinations(n, unique=True): - subscript = util.combine_subscripts(key, comb) + subscript, _ = util.combine_subscripts(key, comb) tn = res[name][comb] tn /= self.energy_sum(key, comb) tn += amplitudes[name][comb] @@ -382,7 +385,7 @@ def update_amps( def update_lams( self, - eris: ERIsInputType = None, + eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, lambdas: Optional[Namespace[AmplitudeType]] = None, lambdas_pert: Optional[Namespace[AmplitudeType]] = None, @@ -401,6 +404,8 @@ def update_lams( Updated cluster lambda amplitudes. """ # TODO active + amplitudes = self._get_amps(amplitudes=amplitudes) + lambdas = self._get_lams(lambdas=lambdas, amplitudes=amplitudes) func, kwargs = self._load_function( "update_lams", eris=eris, @@ -415,7 +420,7 @@ def update_lams( for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): lname = name.replace("t", "l") for comb in util.generate_spin_combinations(n, unique=True): - subscript = util.combine_subscripts(key, comb) + subscript, _ = util.combine_subscripts(key, comb) tn = res[lname][comb] tn /= self.energy_sum(key[n:] + key[:n], comb) tn += lambdas[lname][comb] @@ -505,11 +510,11 @@ def make_rdm2_f( if hermitise: - def transpose1(dm): + def transpose1(dm: NDArray[float]) -> NDArray[float]: dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) return dm - def transpose2(dm): + def transpose2(dm: NDArray[float]) -> NDArray[float]: dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(1, 0, 3, 2)) return dm @@ -572,19 +577,19 @@ def make_eb_coup_rdm( return dm_eb def energy_sum( - self, subscript: str, spins: str, signs_dict: dict[str, int] = None + self, *args: str, signs_dict: Optional[dict[str, str]] = None ) -> NDArray[float]: """Get a direct sum of energies. Args: - subscript: Subscript for the direct sum. - spins: Spins for energies. + *args: Energies to sum. Should specify a subscript and spins. signs_dict: Signs of the energies in the sum. Default sets `("o", "O", "i")` to be positive, and `("v", "V", "a", "b")` to be negative. Returns: Sum of energies. """ + subscript, spins = args n = 0 def next_char() -> str: @@ -628,7 +633,7 @@ def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[ for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): for spin in util.generate_spin_combinations(n, unique=True): tn = amplitudes[name][spin] - subscript = util.combine_subscripts(key, spin) + subscript, _ = util.combine_subscripts(key, spin) vectors.append(util.compress_axes(subscript, tn).ravel()) for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): @@ -651,9 +656,11 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp Returns: Cluster amplitudes. """ - amplitudes = util.Namespace() + amplitudes: Namespace[AmplitudeType] = util.Namespace() i0 = 0 - sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} + sizes: dict[tuple[str, ...], int] = { + (o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab") + } for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): amplitudes[name] = util.Namespace() @@ -668,7 +675,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): shape = (self.nbos,) * n - size = np.prod(shape) + size = self.nbos ** n amplitudes[name] = vector[i0 : i0 + size].reshape(shape) i0 += size @@ -705,7 +712,7 @@ def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float] key = key[n:] + key[:n] for spin in util.generate_spin_combinations(n, unique=True): tn = lambdas[lname][spin] - subscript = util.combine_subscripts(key, spin) + subscript, _ = util.combine_subscripts(key, spin) vectors.append(util.compress_axes(subscript, tn).ravel()) for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): @@ -728,9 +735,11 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: Returns: Cluster lambda amplitudes. """ - lambdas = util.Namespace() + lambdas: AmplitudeType = util.Namespace() i0 = 0 - sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} + sizes: dict[tuple[str, ...], int] = { + (o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab") + } for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): lname = name.replace("t", "l") @@ -747,7 +756,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): shape = (self.nbos,) * n - size = np.prod(shape) + size = self.nbos ** n lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) i0 += size @@ -785,7 +794,7 @@ def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> ND for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): for spin in util.generate_spin_combinations(n, excited=True, unique=True): vn = excitations[m][spin] - subscript = util.combine_subscripts(key[:-1], spin) + subscript, _ = util.combine_subscripts(key[:-1], spin) vectors.append(util.compress_axes(subscript, vn).ravel()) m += 1 @@ -813,7 +822,7 @@ def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> ND key = key[n:] + key[:n] for spin in util.generate_spin_combinations(n, excited=True, unique=True): vn = excitations[m][spin] - subscript = util.combine_subscripts(key[:-1], spin) + subscript, _ = util.combine_subscripts(key[:-1], spin) vectors.append(util.compress_axes(subscript, vn).ravel()) m += 1 @@ -840,7 +849,7 @@ def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> ND for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): for spin in util.generate_spin_combinations(n): vn = excitations[m][spin] - subscript = util.combine_subscripts(key, spin) + subscript, _ = util.combine_subscripts(key, spin) vectors.append(util.compress_axes(subscript, vn).ravel()) m += 1 @@ -865,11 +874,13 @@ def vector_to_excitations_ip( """ excitations = [] i0 = 0 - sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} + sizes: dict[tuple[str, ...], int] = { + (o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab") + } for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): key = key[:-1] - amp = util.Namespace() + amp: AmplitudeType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True, unique=True): subscript, csizes = util.combine_subscripts(key, spin, sizes=sizes) size = util.get_compressed_size(subscript, **csizes) @@ -907,11 +918,13 @@ def vector_to_excitations_ea( """ excitations = [] i0 = 0 - sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} + sizes: dict[tuple[str, ...], int] = { + (o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab") + } for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): key = key[n:] + key[: n - 1] - amp = util.Namespace() + amp: AmplitudeType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True, unique=True): subscript, csizes = util.combine_subscripts(key, spin, sizes=sizes) size = util.get_compressed_size(subscript, **csizes) @@ -949,10 +962,12 @@ def vector_to_excitations_ee( """ excitations = [] i0 = 0 - sizes = {(o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab")} + sizes: dict[tuple[str, ...], int] = { + (o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab") + } for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp = util.Namespace() + amp: AmplitudeType = util.Namespace() for spin in util.generate_spin_combinations(n): subscript, csizes = util.combine_subscripts(key, spin, sizes=sizes) size = util.get_compressed_size(subscript, **csizes) @@ -984,6 +999,8 @@ def get_mean_field_G(self) -> NDArray[float]: Mean-field boson non-conserving term. """ # FIXME should this also sum in frozen orbitals? + assert self.omega is not None + assert self.g is not None val = util.einsum("Ipp->I", self.g.aa.boo) val += util.einsum("Ipp->I", self.g.bb.boo) val -= self.xi * self.omega @@ -1002,13 +1019,13 @@ def bare_fock(self) -> Namespace[NDArray[float]]: Returns: Mean-field Fock matrix. """ - fock = util.einsum( + fock_array = util.einsum( "npq,npi,nqj->nij", self.mf.get_fock().astype(types[float]), self.mo_coeff, self.mo_coeff, ) - fock = util.Namespace(aa=fock[0], bb=fock[1]) + fock = util.Namespace(aa=fock_array[0], bb=fock_array[1]) return fock @property @@ -1018,7 +1035,9 @@ def xi(self) -> NDArray[float]: Returns: Shift in the bosonic operators. """ + assert self.omega is not None if self.options.shift: + assert self.g is not None xi = util.einsum("Iii->I", self.g.aa.boo) xi += util.einsum("Iii->I", self.g.bb.boo) xi /= self.omega @@ -1045,15 +1064,15 @@ def get_eris(self, eris: Optional[ERIsInputType] = None) -> Union[UERIs, UCDERIs Returns: Electron repulsion integrals. """ - if (eris is None) or isinstance(eris, tuple): + if isinstance(eris, (UERIs, UCDERIs)): + return eris + else: if ( isinstance(eris, tuple) and isinstance(eris[0], np.ndarray) and eris[0].ndim == 3 ) or getattr(self.mf, "with_df", None): return self.CDERIs(self, array=eris) else: return self.ERIs(self, array=eris) - else: - return eris @property def nmo(self) -> int: @@ -1072,7 +1091,7 @@ def nocc(self) -> tuple[int, int]: Returns: Number of occupied molecular orbitals for each spin. """ - return cast(tuple[int], tuple(np.sum(mo_occ > 0) for mo_occ in self.mo_occ)) + return cast(tuple[int, int], tuple(np.sum(mo_occ > 0) for mo_occ in self.mo_occ)) @property def nvir(self) -> tuple[int, int]: @@ -1081,4 +1100,4 @@ def nvir(self) -> tuple[int, int]: Returns: Number of virtual molecular orbitals for each spin. """ - return cast(tuple[int], tuple(self.nmo - nocc for nocc in self.nocc)) + return cast(tuple[int, int], tuple(self.nmo - nocc for nocc in self.nocc)) diff --git a/ebcc/util/permutations.py b/ebcc/util/permutations.py index 46e2b565..d4b77135 100644 --- a/ebcc/util/permutations.py +++ b/ebcc/util/permutations.py @@ -228,8 +228,8 @@ def is_mixed_spin(spin: Iterable[Hashable]) -> bool: def combine_subscripts( *subscripts: str, - sizes: Optional[dict[tuple[str, ...], int]] = None, -) -> Union[str, tuple[str, dict[str, int]]]: + sizes: dict[tuple[str, ...], int] = {}, +) -> tuple[str, dict[str, int]]: """Combine subscripts into new unique subscripts for functions such as `compress_axes`. For example, one may wish to compress an amplitude according to both @@ -240,9 +240,8 @@ def combine_subscripts( such that it is unique for a unique value of `tuple(s[i] for s in subscripts)` among other values of `i`. - If `sizes` is passed, this function also returns a dictionary - indicating the size of each new character in the subscript according to - the size of the corresponding original character in the dictionary + This function also returns a dictionary indicating the size of each new character in the + subscript according to the size of the corresponding original character in the dictionary `sizes`. Args: @@ -250,7 +249,7 @@ def combine_subscripts( sizes: Dictionary of sizes for each index. Returns: - New subscript, with a dictionary of sizes of each new index if `sizes` is passed. + New subscript, with a dictionary of sizes of each new index. """ if len(set(len(s) for s in subscripts)) != 1: raise ValueError("Subscripts must be of the same length.") @@ -269,13 +268,10 @@ def combine_subscripts( if j == 123: j = 65 new_subscript += char_map[key] - if sizes is not None: + if sizes: new_sizes[char_map[key]] = sizes[key] - if sizes is None: - return new_subscript - else: - return new_subscript, new_sizes + return new_subscript, new_sizes def compress_axes( diff --git a/pyproject.toml b/pyproject.toml index 97a22e36..d9bd44db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -140,6 +140,10 @@ warn_unused_ignores = false module = "ebcc.codegen.*" ignore_errors = true +[[tool.mypy.overrides]] +module = "pyscf.*" +ignore_missing_imports = true + [[tool.mypy.overrides]] module = "scipy.*" ignore_missing_imports = true From b30eb42f80d10abb4512f123e7aee06376481762 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 21:36:38 +0100 Subject: [PATCH 22/37] Getting there --- ebcc/cc/base.py | 30 ++++++------- ebcc/cc/gebcc.py | 15 +++---- ebcc/cc/rebcc.py | 12 +++--- ebcc/cc/uebcc.py | 15 ++++--- ebcc/eom/base.py | 20 ++++++--- ebcc/eom/geom.py | 27 ++++++++---- ebcc/eom/reom.py | 25 +++++++---- ebcc/eom/ueom.py | 106 +++++++++++++++++++++++----------------------- ebcc/ham/base.py | 81 ++++++++++++++--------------------- ebcc/util/misc.py | 1 + 10 files changed, 169 insertions(+), 163 deletions(-) diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index 8ffadcf0..a9dff88d 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -381,11 +381,10 @@ def solve_lambda( self.converged_lambda = converged @abstractmethod - def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> Any: + def ip_eom(self, **kwargs: Any) -> Any: """Get the IP-EOM object. Args: - options: Options for the IP-EOM calculation. **kwargs: Additional keyword arguments. Returns: @@ -394,11 +393,10 @@ def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> Any: pass @abstractmethod - def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> Any: + def ea_eom(self, **kwargs: Any) -> Any: """Get the EA-EOM object. Args: - options: Options for the EA-EOM calculation. **kwargs: Additional keyword arguments. Returns: @@ -407,11 +405,10 @@ def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> Any: pass @abstractmethod - def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> Any: + def ee_eom(self, **kwargs: Any) -> Any: """Get the EE-EOM object. Args: - options: Options for the EE-EOM calculation. **kwargs: Additional keyword arguments. Returns: @@ -571,7 +568,8 @@ def energy( eris=eris, amplitudes=amplitudes, ) - return astype(func(**kwargs).real, float) + res: float = func(**kwargs).real + return astype(res, float) def energy_perturbative( self, @@ -595,7 +593,8 @@ def energy_perturbative( amplitudes=amplitudes, lambdas=lambdas, ) - return astype(func(**kwargs).real, float) + res: float = func(**kwargs).real + return astype(res, float) @abstractmethod def update_amps( @@ -777,8 +776,7 @@ def make_eb_coup_rdm( def hbar_matvec_ip( self, - r1: AmplitudeType, - r2: AmplitudeType, + *excitations: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, ) -> tuple[AmplitudeType, AmplitudeType]: @@ -794,6 +792,7 @@ def hbar_matvec_ip( Products between the state vectors and the IP-EOM Hamiltonian for the singles and doubles. """ + r1, r2 = excitations # FIXME func, kwargs = self._load_function( "hbar_matvec_ip", eris=eris, @@ -806,8 +805,7 @@ def hbar_matvec_ip( def hbar_matvec_ea( self, - r1: AmplitudeType, - r2: AmplitudeType, + *excitations: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, ) -> tuple[AmplitudeType, AmplitudeType]: @@ -823,6 +821,7 @@ def hbar_matvec_ea( Products between the state vectors and the EA-EOM Hamiltonian for the singles and doubles. """ + r1, r2 = excitations # FIXME func, kwargs = self._load_function( "hbar_matvec_ea", eris=eris, @@ -835,8 +834,7 @@ def hbar_matvec_ea( def hbar_matvec_ee( self, - r1: AmplitudeType, - r2: AmplitudeType, + *excitations: AmplitudeType, eris: Optional[ERIsInputType] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, ) -> tuple[AmplitudeType, AmplitudeType]: @@ -852,6 +850,7 @@ def hbar_matvec_ee( Products between the state vectors and the EE-EOM Hamiltonian for the singles and doubles. """ + r1, r2 = excitations # FIXME func, kwargs = self._load_function( "hbar_matvec_ee", eris=eris, @@ -1324,7 +1323,8 @@ def e_tot(self) -> float: Returns: Total energy. """ - return astype(self.mf.e_tot + self.e_corr, float) + e_tot: float = self.mf.e_tot + self.e_corr + return astype(e_tot, float) @property def t1(self) -> Any: diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index cfcb9145..e273b5b1 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -59,41 +59,38 @@ def spin_type(self) -> str: """Get a string representation of the spin type.""" return "G" - def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_GEOM: + def ip_eom(self, **kwargs: Any) -> IP_GEOM: """Get the IP-EOM object. Args: - options: Options for the IP-EOM calculation. **kwargs: Additional keyword arguments. Returns: IP-EOM object. """ - return IP_GEOM(self, options=options, **kwargs) + return IP_GEOM(self, **kwargs) - def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EA_GEOM: + def ea_eom(self, **kwargs: Any) -> EA_GEOM: """Get the EA-EOM object. Args: - options: Options for the EA-EOM calculation. **kwargs: Additional keyword arguments. Returns: EA-EOM object. """ - return EA_GEOM(self, options=options, **kwargs) + return EA_GEOM(self, **kwargs) - def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EE_GEOM: + def ee_eom(self, **kwargs: Any) -> EE_GEOM: """Get the EE-EOM object. Args: - options: Options for the EE-EOM calculation. **kwargs: Additional keyword arguments. Returns: EE-EOM object. """ - return EE_GEOM(self, options=options, **kwargs) + return EE_GEOM(self, **kwargs) @staticmethod def _convert_mf(mf: SCF) -> GHF: diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index cec69224..482834c5 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -56,7 +56,7 @@ def spin_type(self) -> str: """Get a string representation of the spin type.""" return "R" - def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_REOM: + def ip_eom(self, **kwargs: Any) -> IP_REOM: """Get the IP-EOM object. Args: @@ -66,9 +66,9 @@ def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_REO Returns: IP-EOM object. """ - return IP_REOM(self, options=options, **kwargs) + return IP_REOM(self, **kwargs) - def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EA_REOM: + def ea_eom(self, **kwargs: Any) -> EA_REOM: """Get the EA-EOM object. Args: @@ -78,9 +78,9 @@ def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EA_REO Returns: EA-EOM object. """ - return EA_REOM(self, options=options, **kwargs) + return EA_REOM(self, **kwargs) - def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EE_REOM: + def ee_eom(self, **kwargs: Any) -> EE_REOM: """Get the EE-EOM object. Args: @@ -90,7 +90,7 @@ def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EE_REO Returns: EE-EOM object. """ - return EE_REOM(self, options=options, **kwargs) + return EE_REOM(self, **kwargs) @staticmethod def _convert_mf(mf: SCF) -> RHF: diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 1f618eb5..cb23d5e1 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -22,7 +22,8 @@ from pyscf.scf.hf import SCF from pyscf.scf.uhf import UHF - from ebcc.cc.rebcc import REBCC, BaseOptions + from ebcc.cc.base import BaseOptions + from ebcc.cc.rebcc import REBCC from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -57,7 +58,7 @@ def spin_type(self) -> str: """Get a string representation of the spin type.""" return "U" - def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_UEOM: + def ip_eom(self, **kwargs: Any) -> IP_UEOM: """Get the IP-EOM object. Args: @@ -67,9 +68,9 @@ def ip_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> IP_UEO Returns: IP-EOM object. """ - return IP_UEOM(self, options=options, **kwargs) + return IP_UEOM(self, **kwargs) - def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EA_UEOM: + def ea_eom(self, **kwargs: Any) -> EA_UEOM: """Get the EA-EOM object. Args: @@ -79,9 +80,9 @@ def ea_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EA_UEO Returns: EA-EOM object. """ - return EA_UEOM(self, options=options, **kwargs) + return EA_UEOM(self, **kwargs) - def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EE_UEOM: + def ee_eom(self, **kwargs: Any) -> EE_UEOM: """Get the EE-EOM object. Args: @@ -91,7 +92,7 @@ def ee_eom(self, options: Optional[BaseOptions] = None, **kwargs: Any) -> EE_UEO Returns: EE-EOM object. """ - return EE_UEOM(self, options=options, **kwargs) + return EE_UEOM(self, **kwargs) @staticmethod def _convert_mf(mf: SCF) -> UHF: diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py index 50c82ebd..4adcb1e6 100644 --- a/ebcc/eom/base.py +++ b/ebcc/eom/base.py @@ -14,9 +14,10 @@ from ebcc.core.precision import types if TYPE_CHECKING: - from typing import Any, Callable, Optional + from typing import Any, Callable, Optional, Union - from ebcc.cc.base import AmplitudeType, BaseEBCC, ERIsInputType + from ebcc.cc.base import AmplitudeType, BaseEBCC, ERIsInputType, SpaceType + from ebcc.core.ansatz import Ansatz from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -56,10 +57,15 @@ class BaseEOM(ABC): # Types Options = BaseOptions + # Attrbutes + ebcc: BaseEBCC + space: SpaceType + ansatz: Ansatz + def __init__( self, ebcc: BaseEBCC, - options: Optional[Options] = None, + options: Optional[BaseOptions] = None, **kwargs: Any, ) -> None: """Initialise the EOM object. @@ -189,7 +195,7 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType] """ pass - def dot_braket(self, bra: NDArray[float], ket: NDArray[float]) -> float: + def dot_braket(self, bra: NDArray[float], ket: NDArray[float]) -> Union[float, NDArray[float]]: """Compute the dot product of a bra and ket.""" return np.dot(bra, ket) @@ -217,8 +223,8 @@ def pick( x0 = np.asarray(x0) s = np.dot(guesses_array.conj(), x0.T) s = util.einsum("pi,qi->i", s.conj(), s) - idx = np.argsort(-s)[:nroots] - return lib.linalg_helper._eigs_cmplx2real(w, v, idx, real_system) + arg = np.argsort(-s)[:nroots] + return lib.linalg_helper._eigs_cmplx2real(w, v, arg, real_system) # type: ignore else: @@ -352,7 +358,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, ) -> AmplitudeType: """Construct the moments of the EOM Hamiltonian. diff --git a/ebcc/eom/geom.py b/ebcc/eom/geom.py index 5d42d64a..0fe417c1 100644 --- a/ebcc/eom/geom.py +++ b/ebcc/eom/geom.py @@ -7,13 +7,14 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types +from ebcc.core.precision import types, astype from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM if TYPE_CHECKING: from typing import Optional - from ebcc.cc.gebcc import AmplitudeType, ERIsInputType + from ebcc.ham.space import Space + from ebcc.cc.gebcc import GEBCC, AmplitudeType, ERIsInputType from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -21,7 +22,9 @@ class GEOM(BaseEOM): """Generalised equation-of-motion coupled cluster.""" - pass + # Attributes + ebcc: GEBCC + space: Space class IP_GEOM(GEOM, BaseIP_EOM): @@ -38,7 +41,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1) ** 2 + weight: float = np.dot(r1.ravel(), r1.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -94,7 +98,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian. @@ -116,6 +120,7 @@ def moments( bras = self.bras(eris=eris) kets = self.kets(eris=eris) + moments: NDArray[float] moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) for j in range(self.nmo): @@ -147,7 +152,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1) ** 2 + weight: float = np.dot(r1.ravel(), r1.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -203,7 +209,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian. @@ -225,6 +231,7 @@ def moments( bras = self.bras(eris=eris) kets = self.kets(eris=eris) + moments: NDArray[float] moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) for j in range(self.nmo): @@ -256,7 +263,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1) ** 2 + weight: float = np.dot(r1.ravel(), r1.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -320,7 +328,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, diagonal_only: bool = True, ) -> NDArray[float]: @@ -350,6 +358,7 @@ def moments( bras = self.bras(eris=eris) kets = self.kets(eris=eris) + moments: NDArray[float] moments = np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]) for k in range(self.nmo): diff --git a/ebcc/eom/reom.py b/ebcc/eom/reom.py index 57752656..b10ecff9 100644 --- a/ebcc/eom/reom.py +++ b/ebcc/eom/reom.py @@ -7,13 +7,14 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types +from ebcc.core.precision import types, astype from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM if TYPE_CHECKING: from typing import Optional - from ebcc.cc.rebcc import AmplitudeType, ERIsInputType + from ebcc.ham.space import Space + from ebcc.cc.rebcc import REBCC, AmplitudeType, ERIsInputType from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -21,7 +22,9 @@ class REOM(BaseEOM): """Restricted equation-of-motion coupled cluster.""" - pass + # Attributes + ebcc: REBCC + space: Space class IP_REOM(REOM, BaseIP_EOM): @@ -38,7 +41,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1) ** 2 + weight: float = np.dot(r1.ravel(), r1.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -94,7 +98,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, ) -> AmplitudeType: """Construct the moments of the EOM Hamiltonian. @@ -116,6 +120,7 @@ def moments( bras = self.bras(eris=eris) kets = self.kets(eris=eris) + moments: NDArray[float] moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) for j in range(self.nmo): @@ -147,7 +152,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1) ** 2 + weight: float = np.dot(r1.ravel(), r1.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -203,7 +209,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, ) -> AmplitudeType: """Construct the moments of the EOM Hamiltonian. @@ -225,6 +231,7 @@ def moments( bras = self.bras(eris=eris) kets = self.kets(eris=eris) + moments: NDArray[float] moments = np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]) for j in range(self.nmo): @@ -256,7 +263,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1) ** 2 + weight: float = np.dot(r1.ravel(), r1.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -350,6 +358,7 @@ def moments( bras = self.bras(eris=eris) kets = self.kets(eris=eris) + moments: NDArray[float] moments = np.zeros((nmom, self.nmo, self.nmo, self.nmo, self.nmo), dtype=types[float]) for k in range(self.nmo): diff --git a/ebcc/eom/ueom.py b/ebcc/eom/ueom.py index f1637a4b..74ed633b 100644 --- a/ebcc/eom/ueom.py +++ b/ebcc/eom/ueom.py @@ -6,13 +6,14 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types +from ebcc.core.precision import types, astype from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM if TYPE_CHECKING: from typing import Optional - from ebcc.cc.uebcc import AmplitudeType, ERIsInputType + from ebcc.ham.space import Space + from ebcc.cc.uebcc import UEBCC, AmplitudeType, ERIsInputType from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -20,7 +21,9 @@ class UEOM(BaseEOM): """Unrestricted equation-of-motion coupled cluster.""" - pass + # Attributes + ebcc: UEBCC + space: tuple[Space, Space] class IP_UEOM(UEOM, BaseIP_EOM): @@ -37,7 +40,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1.a) ** 2 + np.linalg.norm(r1.b) ** 2 + weight: float = np.dot(r1.a.ravel(), r1.a.ravel()) + np.dot(r1.b.ravel(), r1.b.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -52,7 +56,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): key = key[:-1] - spin_part = util.Namespace() + spin_part: AmplitudeType = util.Namespace() for comb in util.generate_spin_combinations(n, excited=True): spin_part[comb] = self.ebcc.energy_sum(key, comb) parts.append(spin_part) @@ -72,16 +76,16 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: Bra vectors. """ bras_raw = list(self.ebcc.make_ip_mom_bras(eris=eris)) - bras = util.Namespace(a=[], b=[]) + bras_tmp: Namespace[list[NDArray[float]]] = util.Namespace(a=[], b=[]) for i in range(self.nmo): - amps_a = [] - amps_b = [] + amps_a: list[AmplitudeType] = [] + amps_b: list[AmplitudeType] = [] m = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_a = util.Namespace() - amp_b = util.Namespace() + amp_a: AmplitudeType = util.Namespace() + amp_b: AmplitudeType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True): shape = tuple(self.space["ab".index(s)].ncocc for s in spin[:n]) + tuple( self.space["ab".index(s)].ncvir for s in spin[n:] @@ -110,11 +114,10 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): raise util.ModelNotImplemented - bras.a.append(self.amplitudes_to_vector(*amps_a)) - bras.b.append(self.amplitudes_to_vector(*amps_b)) + bras_tmp.a.append(self.amplitudes_to_vector(*amps_a)) + bras_tmp.b.append(self.amplitudes_to_vector(*amps_b)) - bras.a = np.array(bras.a) - bras.b = np.array(bras.b) + bras: Namespace[NDArray[float]] = util.Namespace(a=np.array(bras_tmp.a), b=np.array(bras_tmp.b)) return bras @@ -128,17 +131,17 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: Ket vectors. """ kets_raw = list(self.ebcc.make_ip_mom_kets(eris=eris)) - kets = util.Namespace(a=[], b=[]) + kets_tmp: Namespace[list[NDArray[float]]] = util.Namespace(a=[], b=[]) for i in range(self.nmo): j = (Ellipsis, i) - amps_a = [] - amps_b = [] + amps_a: list[AmplitudeType] = [] + amps_b: list[AmplitudeType] = [] m = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_a = util.Namespace() - amp_b = util.Namespace() + amp_a: AmplitudeType = util.Namespace() + amp_b: AmplitudeType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True): shape = tuple(self.space["ab".index(s)].ncocc for s in spin[:n]) + tuple( self.space["ab".index(s)].ncvir for s in spin[n:] @@ -167,11 +170,10 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): raise util.ModelNotImplemented - kets.a.append(self.amplitudes_to_vector(*amps_a)) - kets.b.append(self.amplitudes_to_vector(*amps_b)) + kets_tmp.a.append(self.amplitudes_to_vector(*amps_a)) + kets_tmp.b.append(self.amplitudes_to_vector(*amps_b)) - kets.a = np.array(kets.a) - kets.b = np.array(kets.b) + kets: Namespace[NDArray[float]] = util.Namespace(a=np.array(kets_tmp.a), b=np.array(kets_tmp.b)) return kets @@ -179,7 +181,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian. @@ -193,14 +195,14 @@ def moments( Returns: Moments of the EOM Hamiltonian. """ - if eris is None: - eris = self.ebcc.get_eris() + eris = self.ebcc.get_eris(eris) if not amplitudes: amplitudes = self.ebcc.amplitudes bras = self.bras(eris=eris) kets = self.kets(eris=eris) + moments: Namespace[NDArray[float]] moments = util.Namespace( aa=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), bb=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), @@ -237,7 +239,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1.a) ** 2 + np.linalg.norm(r1.b) ** 2 + weight: float = np.dot(r1.a.ravel(), r1.a.ravel()) + np.dot(r1.b.ravel(), r1.b.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -252,7 +255,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): key = key[n:] + key[: n - 1] - spin_part = util.Namespace() + spin_part: AmplitudeType = util.Namespace() for comb in util.generate_spin_combinations(n, excited=True): spin_part[comb] = -self.ebcc.energy_sum(key, comb) parts.append(spin_part) @@ -272,16 +275,16 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: Bra vectors. """ bras_raw = list(self.ebcc.make_ea_mom_bras(eris=eris)) - bras = util.Namespace(a=[], b=[]) + bras_tmp: Namespace[list[NDArray[float]]] = util.Namespace(a=[], b=[]) for i in range(self.nmo): - amps_a = [] - amps_b = [] + amps_a: list[AmplitudeType] = [] + amps_b: list[AmplitudeType] = [] m = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_a = util.Namespace() - amp_b = util.Namespace() + amp_a: AmplitudeType = util.Namespace() + amp_b: AmplitudeType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True): shape = tuple(self.space["ab".index(s)].ncvir for s in spin[:n]) + tuple( self.space["ab".index(s)].ncocc for s in spin[n:] @@ -310,11 +313,10 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): raise util.ModelNotImplemented - bras.a.append(self.amplitudes_to_vector(*amps_a)) - bras.b.append(self.amplitudes_to_vector(*amps_b)) + bras_tmp.a.append(self.amplitudes_to_vector(*amps_a)) + bras_tmp.b.append(self.amplitudes_to_vector(*amps_b)) - bras.a = np.array(bras.a) - bras.b = np.array(bras.b) + bras: Namespace[NDArray[float]] = util.Namespace(a=np.array(bras_tmp.a), b=np.array(bras_tmp.b)) return bras @@ -328,17 +330,17 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: Ket vectors. """ kets_raw = list(self.ebcc.make_ea_mom_kets(eris=eris)) - kets = util.Namespace(a=[], b=[]) + kets_tmp: Namespace[list[NDArray[float]]] = util.Namespace(a=[], b=[]) for i in range(self.nmo): j = (Ellipsis, i) - amps_a = [] - amps_b = [] + amps_a: list[AmplitudeType] = [] + amps_b: list[AmplitudeType] = [] m = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_a = util.Namespace() - amp_b = util.Namespace() + amp_a: AmplitudeType = util.Namespace() + amp_b: AmplitudeType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True): shape = tuple(self.space["ab".index(s)].ncvir for s in spin[:n]) + tuple( self.space["ab".index(s)].ncocc for s in spin[n:] @@ -367,11 +369,10 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): raise util.ModelNotImplemented - kets.a.append(self.amplitudes_to_vector(*amps_a)) - kets.b.append(self.amplitudes_to_vector(*amps_b)) + kets_tmp.a.append(self.amplitudes_to_vector(*amps_a)) + kets_tmp.b.append(self.amplitudes_to_vector(*amps_b)) - kets.a = np.array(kets.a) - kets.b = np.array(kets.b) + kets: Namespace[NDArray[float]] = util.Namespace(a=np.array(kets_tmp.a), b=np.array(kets_tmp.b)) return kets @@ -379,7 +380,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian. @@ -393,14 +394,14 @@ def moments( Returns: Moments of the EOM Hamiltonian. """ - if eris is None: - eris = self.ebcc.get_eris() + eris = self.ebcc.get_eris(eris) if not amplitudes: amplitudes = self.ebcc.amplitudes bras = self.bras(eris=eris) kets = self.kets(eris=eris) + moments: Namespace[NDArray[float]] moments = util.Namespace( aa=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), bb=np.zeros((nmom, self.nmo, self.nmo), dtype=types[float]), @@ -437,7 +438,8 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: def _quasiparticle_weight(self, r1: AmplitudeType) -> float: """Get the quasiparticle weight.""" - return np.linalg.norm(r1.aa) ** 2 + np.linalg.norm(r1.bb) ** 2 + weight: float = np.dot(r1.aa.ravel(), r1.aa.ravel()) + np.dot(r1.bb.ravel(), r1.bb.ravel()) + return astype(weight, float) def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: """Get the diagonal of the Hamiltonian. @@ -451,7 +453,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: parts = [] for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - spin_part = util.Namespace() + spin_part: AmplitudeType = util.Namespace() for comb in util.generate_spin_combinations(n): spin_part[comb] = self.ebcc.energy_sum(key, comb) parts.append(spin_part) @@ -487,7 +489,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Namespace[AmplitudeType] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, hermitise: bool = True, diagonal_only: bool = True, ) -> NDArray[float]: diff --git a/ebcc/ham/base.py b/ebcc/ham/base.py index c9c2b939..04f4a5c4 100644 --- a/ebcc/ham/base.py +++ b/ebcc/ham/base.py @@ -8,14 +8,31 @@ from ebcc.util import Namespace if TYPE_CHECKING: - from typing import Any, TypeVar + from typing import Any, TypeVar, Optional from ebcc.cc.base import BaseEBCC + from ebcc.numpy.typing import NDArray T = TypeVar("T") -class BaseFock(Namespace, ABC): +class BaseHamiltonian(Namespace[Any], ABC): + """Base class for Hamiltonians.""" + + @abstractmethod + def __getitem__(self, key: str) -> Any: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + Slice of the Hamiltonian. + """ + pass + + +class BaseFock(BaseHamiltonian): """Base class for Fock matrices. Attributes: @@ -31,10 +48,10 @@ class BaseFock(Namespace, ABC): def __init__( self, cc: BaseEBCC, - array: Any = None, - space: tuple[Any] = None, - mo_coeff: tuple[Any] = None, - g: Namespace[Any] = None, + array: Optional[Any] = None, + space: Optional[tuple[Any]] = None, + mo_coeff: Optional[tuple[Any]] = None, + g: Optional[Namespace[Any]] = None, ) -> None: """Initialise the Fock matrix. @@ -59,32 +76,20 @@ def __init__( self.__dict__["xi"] = cc.xi @abstractmethod - def _get_fock(self) -> T: + def _get_fock(self) -> Any: """Get the Fock matrix.""" pass - @abstractmethod - def __getitem__(self, key: str) -> Any: - """Just-in-time getter. - - Args: - key: Key to get. - - Returns: - Slice of the Fock matrix. - """ - pass - -class BaseERIs(Namespace, ABC): +class BaseERIs(BaseHamiltonian): """Base class for electronic repulsion integrals.""" def __init__( self, cc: BaseEBCC, - array: Any = None, - space: tuple[Any] = None, - mo_coeff: tuple[Any] = None, + array: Optional[Any] = None, + space: Optional[tuple[Any]] = None, + mo_coeff: Optional[tuple[Any]] = None, ) -> None: """Initialise the ERIs. @@ -102,20 +107,8 @@ def __init__( 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 - @abstractmethod - def __getitem__(self, key: str) -> Any: - """Just-in-time getter. - - Args: - key: Key to get. - - Returns: - Slice of the ERIs. - """ - pass - -class BaseElectronBoson(Namespace, ABC): +class BaseElectronBoson(BaseHamiltonian): """Base class for electron-boson coupling matrices. Attributes: @@ -127,8 +120,8 @@ class BaseElectronBoson(Namespace, ABC): def __init__( self, cc: BaseEBCC, - array: Any = None, - space: tuple[Any] = None, + array: Optional[Any] = None, + space: Optional[tuple[Any]] = None, ) -> None: """Initialise the electron-boson coupling matrix. @@ -147,15 +140,3 @@ def __init__( def _get_g(self) -> NDArray[float]: """Get the electron-boson coupling matrix.""" return self.cc.bare_g - - @abstractmethod - def __getitem__(self, key: str) -> Any: - """Just-in-time getter. - - Args: - key: Key to get. - - Returns: - Slice of the electron-boson coupling matrix. - """ - pass diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index 8533d5c8..50a15624 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -4,6 +4,7 @@ import time from collections.abc import MutableMapping +from dataclasses import dataclass from typing import TYPE_CHECKING, Generic, TypeVar if TYPE_CHECKING: From 01a4bc140afb5581e39f12f6aaffde7c00a7eade Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 22:33:39 +0100 Subject: [PATCH 23/37] mypy done --- ebcc/ham/base.py | 18 ++++----- ebcc/ham/cderis.py | 6 ++- ebcc/ham/eris.py | 6 +++ ebcc/ham/fock.py | 8 +++- ebcc/ham/space.py | 79 +++++++++++++++++++++++---------------- ebcc/opt/base.py | 17 ++++++--- ebcc/opt/gbrueckner.py | 12 ++++-- ebcc/opt/rbrueckner.py | 12 ++++-- ebcc/opt/ubrueckner.py | 26 +++++++------ ebcc/util/misc.py | 8 ++-- ebcc/util/permutations.py | 4 +- 11 files changed, 120 insertions(+), 76 deletions(-) diff --git a/ebcc/ham/base.py b/ebcc/ham/base.py index 04f4a5c4..36c1fbbb 100644 --- a/ebcc/ham/base.py +++ b/ebcc/ham/base.py @@ -3,12 +3,12 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from ebcc.util import Namespace if TYPE_CHECKING: - from typing import Any, TypeVar, Optional + from typing import TypeVar, Optional from ebcc.cc.base import BaseEBCC from ebcc.numpy.typing import NDArray @@ -49,8 +49,8 @@ def __init__( self, cc: BaseEBCC, array: Optional[Any] = None, - space: Optional[tuple[Any]] = None, - mo_coeff: Optional[tuple[Any]] = None, + space: Optional[tuple[Any, ...]] = None, + mo_coeff: Optional[tuple[Any, ...]] = None, g: Optional[Namespace[Any]] = None, ) -> None: """Initialise the Fock matrix. @@ -72,8 +72,8 @@ def __init__( self.__dict__["g"] = g if g is not None else cc.g # Boson parameters: - self.__dict__["shift"] = cc.options.shift - self.__dict__["xi"] = cc.xi + 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: @@ -88,8 +88,8 @@ def __init__( self, cc: BaseEBCC, array: Optional[Any] = None, - space: Optional[tuple[Any]] = None, - mo_coeff: Optional[tuple[Any]] = None, + space: Optional[tuple[Any, ...]] = None, + mo_coeff: Optional[tuple[Any, ...]] = None, ) -> None: """Initialise the ERIs. @@ -121,7 +121,7 @@ def __init__( self, cc: BaseEBCC, array: Optional[Any] = None, - space: Optional[tuple[Any]] = None, + space: Optional[tuple[Any, ...]] = None, ) -> None: """Initialise the electron-boson coupling matrix. diff --git a/ebcc/ham/cderis.py b/ebcc/ham/cderis.py index d73101ce..8dfb0381 100644 --- a/ebcc/ham/cderis.py +++ b/ebcc/ham/cderis.py @@ -12,6 +12,8 @@ from ebcc.ham.base import BaseERIs if TYPE_CHECKING: + from typing import Optional + from ebcc.numpy.typing import NDArray @@ -25,7 +27,7 @@ class RCDERIs(BaseERIs): array: ERIs in the MO basis. """ - def __getitem__(self, key: str, e2: bool = False) -> NDArray[float]: + def __getitem__(self, key: str, e2: Optional[bool] = False) -> NDArray[float]: """Just-in-time getter. Args: @@ -40,7 +42,7 @@ def __getitem__(self, key: str, e2: bool = False) -> NDArray[float]: if len(key) == 4: e1 = self.__getitem__("Q" + key[:2]) - e2 = self.__getitem__("Q" + key[2:], e2=True) + e2 = self.__getitem__("Q" + key[2:], e2=True) # type: ignore return util.einsum("Qij,Qkl->ijkl", e1, e2) elif len(key) == 3: key = key[1:] diff --git a/ebcc/ham/eris.py b/ebcc/ham/eris.py index 726b50b8..bd7571ad 100644 --- a/ebcc/ham/eris.py +++ b/ebcc/ham/eris.py @@ -26,6 +26,8 @@ class RERIs(BaseERIs): array: ERIs in the MO basis. """ + _members: dict[str, NDArray[float]] + def __getitem__(self, key: str) -> NDArray[float]: """Just-in-time getter. @@ -65,6 +67,8 @@ class UERIs(BaseERIs): array: ERIs in the MO basis. """ + _members: dict[str, RERIs] + def __getitem__(self, key: str) -> RERIs: """Just-in-time getter. @@ -124,6 +128,8 @@ class GERIs(BaseERIs): array: ERIs in the MO basis. """ + _members: dict[str, UERIs] + def __init__(self, *args: Any, **kwargs: Any) -> None: """Initialise the class.""" super().__init__(*args, **kwargs) diff --git a/ebcc/ham/fock.py b/ebcc/ham/fock.py index 8deb6f94..cf0f22c9 100644 --- a/ebcc/ham/fock.py +++ b/ebcc/ham/fock.py @@ -26,6 +26,8 @@ class RFock(BaseFock): xi: Boson parameters. """ + _members: dict[str, NDArray[float]] + def _get_fock(self) -> NDArray[float]: fock_ao = self.cc.mf.get_fock().astype(types[float]) return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff[0], self.mo_coeff[1]) @@ -64,7 +66,9 @@ class UFock(BaseFock): g: Namespace containing blocks of the electron-boson coupling matrix """ - def _get_fock(self) -> tuple[NDArray[float]]: + _members: dict[str, RFock] + + def _get_fock(self) -> tuple[NDArray[float], NDArray[float]]: fock_ao = self.cc.mf.get_fock().astype(types[float]) return ( util.einsum("pq,pi,qj->ij", fock_ao[0], self.mo_coeff[0][0], self.mo_coeff[1][0]), @@ -107,6 +111,8 @@ class GFock(BaseFock): xi: Boson parameters. """ + _members: dict[str, NDArray[float]] + def _get_fock(self) -> NDArray[float]: fock_ao = self.cc.mf.get_fock().astype(types[float]) return util.einsum("pq,pi,qj->ij", fock_ao, self.mo_coeff[0], self.mo_coeff[1]) diff --git a/ebcc/ham/space.py b/ebcc/ham/space.py index 2bb245a0..314b594c 100644 --- a/ebcc/ham/space.py +++ b/ebcc/ham/space.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from pyscf.mp import MP2 @@ -154,17 +154,17 @@ def virtual(self) -> NDArray[bool]: @property def nmo(self) -> int: """Get the number of orbitals.""" - return self.occupied.size + return cast(int, self.occupied.size) @property def nocc(self) -> int: """Get the number of occupied orbitals.""" - return np.sum(self.occupied) + return cast(int, np.sum(self.occupied)) @property def nvir(self) -> int: """Get the number of virtual orbitals.""" - return np.sum(self.virtual) + return cast(int, np.sum(self.virtual)) # Correlated space: @@ -186,17 +186,17 @@ def correlated_virtual(self) -> NDArray[bool]: @property def ncorr(self) -> int: """Get the number of correlated orbitals.""" - return np.sum(self.correlated) + return cast(int, np.sum(self.correlated)) @property def ncocc(self) -> int: """Get the number of occupied correlated orbitals.""" - return np.sum(self.correlated_occupied) + return cast(int, np.sum(self.correlated_occupied)) @property def ncvir(self) -> int: """Get the number of virtual correlated orbitals.""" - return np.sum(self.correlated_virtual) + return cast(int, np.sum(self.correlated_virtual)) # Inactive space: @@ -218,17 +218,17 @@ def inactive_virtual(self) -> NDArray[bool]: @property def ninact(self) -> int: """Get the number of inactive orbitals.""" - return np.sum(self.inactive) + return cast(int, np.sum(self.inactive)) @property def niocc(self) -> int: """Get the number of occupied inactive orbitals.""" - return np.sum(self.inactive_occupied) + return cast(int, np.sum(self.inactive_occupied)) @property def nivir(self) -> int: """Get the number of virtual inactive orbitals.""" - return np.sum(self.inactive_virtual) + return cast(int, np.sum(self.inactive_virtual)) # Frozen space: @@ -250,17 +250,17 @@ def frozen_virtual(self) -> NDArray[bool]: @property def nfroz(self) -> int: """Get the number of frozen orbitals.""" - return np.sum(self.frozen) + return cast(int, np.sum(self.frozen)) @property def nfocc(self) -> int: """Get the number of occupied frozen orbitals.""" - return np.sum(self.frozen_occupied) + return cast(int, np.sum(self.frozen_occupied)) @property def nfvir(self) -> int: """Get the number of virtual frozen orbitals.""" - return np.sum(self.frozen_virtual) + return cast(int, np.sum(self.frozen_virtual)) # Active space: @@ -282,32 +282,30 @@ def active_virtual(self) -> NDArray[bool]: @property def nact(self) -> int: """Get the number of active orbitals.""" - return np.sum(self.active) + return cast(int, np.sum(self.active)) @property def naocc(self) -> int: """Get the number of occupied active orbitals.""" - return np.sum(self.active_occupied) + return cast(int, np.sum(self.active_occupied)) @property def navir(self) -> int: """Get the number of virtual active orbitals.""" - return np.sum(self.active_virtual) + return cast(int, np.sum(self.active_virtual)) if TYPE_CHECKING: # Needs to be defined after Space - ConstructSpaceReturnType = Union[ - tuple[NDArray[float], NDArray[float], Space], - tuple[ - tuple[NDArray[float], NDArray[float]], - tuple[NDArray[float], NDArray[float]], - tuple[Space, Space], - ], + RConstructSpaceReturnType = tuple[NDArray[float], NDArray[float], Space] + UConstructSpaceReturnType = tuple[ + tuple[NDArray[float], NDArray[float]], + tuple[NDArray[float], NDArray[float]], + tuple[Space, Space], ] -def construct_default_space(mf: SCF) -> ConstructSpaceReturnType: +def construct_default_space(mf: SCF) -> Union[RConstructSpaceReturnType, UConstructSpaceReturnType]: """Construct a default space. Args: @@ -343,7 +341,7 @@ def construct_fno_space( occ_tol: Optional[float] = 1e-5, occ_frac: Optional[float] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None, -) -> ConstructSpaceReturnType: +) -> Union[RConstructSpaceReturnType, UConstructSpaceReturnType]: """Construct a frozen natural orbital space. Args: @@ -376,9 +374,9 @@ def _construct( mo_energy: NDArray[float], mo_coeff: NDArray[float], mo_occ: NDArray[float], - ) -> ConstructSpaceReturnType: + ) -> RConstructSpaceReturnType: # Get the number of occupied orbitals - nocc = np.sum(mo_occ > 0) + nocc = int(np.sum(mo_occ > 0)) # Calculate the natural orbitals n, c = np.linalg.eigh(dm1[nocc:, nocc:]) @@ -389,10 +387,10 @@ def _construct( active_vir = n > occ_tol else: active_vir = np.cumsum(n / np.sum(n)) <= occ_frac - num_active_vir = np.sum(active_vir) + num_active_vir = int(np.sum(active_vir)) # Canonicalise the natural orbitals - fock_vv = np.diag(mo_energy[nocc:]).astype(types[float]) + fock_vv = np.diag(mo_energy[nocc:]) fock_vv = util.einsum("ab,au,bv->uv", fock_vv, c, c) _, c_can = np.linalg.eigh(fock_vv[active_vir][:, active_vir]) @@ -400,7 +398,7 @@ def _construct( no_coeff_avir = np.linalg.multi_dot((mo_coeff[:, nocc:], c[:, :num_active_vir], c_can)) no_coeff_fvir = np.dot(mo_coeff[:, nocc:], c[:, num_active_vir:]) no_coeff_occ = mo_coeff[:, :nocc] - no_coeff = np.hstack((no_coeff_occ, no_coeff_avir, no_coeff_fvir)).astype(types[float]) + no_coeff: NDArray[float] = np.hstack((no_coeff_occ, no_coeff_avir, no_coeff_fvir)) # Build the natural orbital space frozen = np.zeros_like(mo_occ, dtype=bool) @@ -415,8 +413,23 @@ def _construct( # Construct the natural orbitals if np.ndim(mf.mo_occ) == 2: - coeff_a, occ_a, space_a = _construct(dm1[0], mf.mo_energy[0], mf.mo_coeff[0], mf.mo_occ[0]) - coeff_b, occ_b, space_b = _construct(dm1[1], mf.mo_energy[1], mf.mo_coeff[1], mf.mo_occ[1]) + coeff_a, occ_a, space_a = _construct( + dm1[0].astype(types[float]), + mf.mo_energy[0].astype(types[float]), + mf.mo_coeff[0].astype(types[float]), + mf.mo_occ[0].astype(types[float]), + ) + coeff_b, occ_b, space_b = _construct( + dm1[1].astype(types[float]), + mf.mo_energy[1].astype(types[float]), + mf.mo_coeff[1].astype(types[float]), + mf.mo_occ[1].astype(types[float]), + ) return (coeff_a, coeff_b), (occ_a, occ_b), (space_a, space_b) else: - return _construct(dm1, mf.mo_energy, mf.mo_coeff, mf.mo_occ) + return _construct( + dm1.astype(types[float]), + mf.mo_energy.astype(types[float]), + mf.mo_coeff.astype(types[float]), + mf.mo_occ.astype(types[float]), + ) diff --git a/ebcc/opt/base.py b/ebcc/opt/base.py index 49faabb1..e1d6cb06 100644 --- a/ebcc/opt/base.py +++ b/ebcc/opt/base.py @@ -17,6 +17,7 @@ if TYPE_CHECKING: from typing import Any, Optional, TypeVar + from ebcc.numpy.typing import NDArray from ebcc.cc.base import AmplitudeType, BaseEBCC from ebcc.util import Namespace @@ -53,6 +54,9 @@ class BaseBruecknerEBCC(ABC): # Types Options: type[BaseOptions] = BaseOptions + # Attributes + cc: BaseEBCC + def __init__( self, cc: BaseEBCC, @@ -122,8 +126,8 @@ def kernel(self) -> float: diis.damping = self.options.damping # Initialise coefficients: - mo_coeff_new = np.array(self.cc.mo_coeff, copy=True, dtype=types[float]) - mo_coeff_ref = np.array(self.cc.mo_coeff, copy=True, dtype=types[float]) + mo_coeff_new: NDArray[float] = np.array(self.cc.mo_coeff, copy=True, dtype=types[float]) + mo_coeff_ref: NDArray[float] = np.array(self.cc.mo_coeff, copy=True, dtype=types[float]) mo_coeff_ref = self.mo_to_correlated(mo_coeff_ref) u_tot = None @@ -154,7 +158,8 @@ def kernel(self) -> float: # Run CC calculation: e_prev = self.cc.e_tot with lib.temporary_env(self.cc, log=NullLogger()): - self.cc.__init__( + self.cc.__class__.__init__( + self.cc, self.mf, log=self.cc.log, ansatz=self.cc.ansatz, @@ -247,7 +252,7 @@ def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> pass @abstractmethod - def mo_to_correlated(self, mo_coeff: T) -> T: + def mo_to_correlated(self, mo_coeff: Any) -> Any: """Transform the MO coefficients into the correlated basis. Args: @@ -259,7 +264,7 @@ def mo_to_correlated(self, mo_coeff: T) -> T: pass @abstractmethod - def mo_update_correlated(self, mo_coeff: T, mo_coeff_corr: T) -> T: + def mo_update_correlated(self, mo_coeff: Any, mo_coeff_corr: Any) -> Any: """Update the correlated slice of a set of MO coefficients. Args: @@ -272,7 +277,7 @@ def mo_update_correlated(self, mo_coeff: T, mo_coeff_corr: T) -> T: pass @abstractmethod - def update_coefficients(self, u_tot: AmplitudeType, mo_coeff_new: T, mo_coeff_ref: T) -> T: + def update_coefficients(self, u_tot: AmplitudeType, mo_coeff_new: Any, mo_coeff_ref: Any) -> Any: """Update the MO coefficients. Args: diff --git a/ebcc/opt/gbrueckner.py b/ebcc/opt/gbrueckner.py index 761c4292..ca5759dc 100644 --- a/ebcc/opt/gbrueckner.py +++ b/ebcc/opt/gbrueckner.py @@ -8,13 +8,13 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types +from ebcc.core.precision import types, astype from ebcc.opt.base import BaseBruecknerEBCC if TYPE_CHECKING: from typing import Optional - from ebcc.cc.gebcc import AmplitudeType + from ebcc.cc.gebcc import AmplitudeType, GEBCC from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -28,6 +28,9 @@ class BruecknerGEBCC(BaseBruecknerEBCC): options: Options for the EOM calculation. """ + # Attributes + cc: GEBCC + def get_rotation_matrix( self, u_tot: Optional[AmplitudeType] = None, @@ -51,7 +54,7 @@ def get_rotation_matrix( if u_tot is None: u_tot = np.eye(self.cc.space.ncorr, dtype=types[float]) - t1_block = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) + t1_block: NDArray[float] = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) t1_block[: self.cc.space.ncocc, self.cc.space.ncocc :] = -t1 t1_block[self.cc.space.ncocc :, : self.cc.space.ncocc] = t1.T @@ -120,7 +123,8 @@ def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> """ if not amplitudes: amplitudes = self.cc.amplitudes - return np.linalg.norm(amplitudes["t1"]) + weight: float = types[float](np.linalg.norm(amplitudes["t1"])) + return weight def mo_to_correlated(self, mo_coeff: NDArray[float]) -> NDArray[float]: """Transform the MO coefficients into the correlated basis. diff --git a/ebcc/opt/rbrueckner.py b/ebcc/opt/rbrueckner.py index f36f2864..c083a75a 100644 --- a/ebcc/opt/rbrueckner.py +++ b/ebcc/opt/rbrueckner.py @@ -8,13 +8,13 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types +from ebcc.core.precision import types, astype from ebcc.opt.base import BaseBruecknerEBCC if TYPE_CHECKING: from typing import Optional - from ebcc.cc.rebcc import AmplitudeType + from ebcc.cc.rebcc import AmplitudeType, REBCC from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -28,6 +28,9 @@ class BruecknerREBCC(BaseBruecknerEBCC): options: Options for the EOM calculation. """ + # Attributes + cc: REBCC + def get_rotation_matrix( self, u_tot: Optional[AmplitudeType] = None, @@ -51,7 +54,7 @@ def get_rotation_matrix( if u_tot is None: u_tot = np.eye(self.cc.space.ncorr, dtype=types[float]) - t1_block = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) + t1_block: NDArray[float] = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) t1_block[: self.cc.space.ncocc, self.cc.space.ncocc :] = -t1 t1_block[self.cc.space.ncocc :, : self.cc.space.ncocc] = t1.T @@ -120,7 +123,8 @@ def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> """ if not amplitudes: amplitudes = self.cc.amplitudes - return np.linalg.norm(amplitudes["t1"]) + weight: float = types[float](np.linalg.norm(amplitudes["t1"])) + return weight def mo_to_correlated(self, mo_coeff: NDArray[float]) -> NDArray[float]: """Transform the MO coefficients into the correlated basis. diff --git a/ebcc/opt/ubrueckner.py b/ebcc/opt/ubrueckner.py index 5f1fa272..49d047a1 100644 --- a/ebcc/opt/ubrueckner.py +++ b/ebcc/opt/ubrueckner.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.uebcc import AmplitudeType + from ebcc.cc.uebcc import AmplitudeType, UEBCC from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -28,6 +28,9 @@ class BruecknerUEBCC(BaseBruecknerEBCC): options: Options for the EOM calculation. """ + # Attributes + cc: UEBCC + def get_rotation_matrix( self, u_tot: Optional[AmplitudeType] = None, @@ -54,7 +57,7 @@ def get_rotation_matrix( bb=np.eye(self.cc.space[1].ncorr, dtype=types[float]), ) - t1_block = util.Namespace( + t1_block: Namespace[NDArray[float]] = util.Namespace( aa=np.zeros((self.cc.space[0].ncorr, self.cc.space[0].ncorr), dtype=types[float]), bb=np.zeros((self.cc.space[1].ncorr, self.cc.space[1].ncorr), dtype=types[float]), ) @@ -143,11 +146,12 @@ def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> """ if not amplitudes: amplitudes = self.cc.amplitudes - norm_a = np.linalg.norm(amplitudes["t1"].aa) - norm_b = np.linalg.norm(amplitudes["t1"].bb) - return np.linalg.norm([norm_a, norm_b]) + weight_a: float = types[float](np.linalg.norm(amplitudes["t1"].aa)) + weight_b: float = types[float](np.linalg.norm(amplitudes["t1"].bb)) + weight: float = types[float](np.linalg.norm([weight_a, weight_b])) + return weight - def mo_to_correlated(self, mo_coeff: tuple[NDArray[float]]) -> tuple[NDArray[float]]: + def mo_to_correlated(self, mo_coeff: tuple[NDArray[float], NDArray[float]]) -> tuple[NDArray[float], NDArray[float]]: """Transform the MO coefficients into the correlated basis. Args: @@ -162,8 +166,8 @@ def mo_to_correlated(self, mo_coeff: tuple[NDArray[float]]) -> tuple[NDArray[flo ) def mo_update_correlated( - self, mo_coeff: tuple[NDArray[float]], mo_coeff_corr: tuple[NDArray[float]] - ) -> tuple[NDArray[float]]: + self, mo_coeff: tuple[NDArray[float], NDArray[float]], mo_coeff_corr: tuple[NDArray[float], NDArray[float]] + ) -> tuple[NDArray[float], NDArray[float]]: """Update the correlated slice of a set of MO coefficients. Args: @@ -180,9 +184,9 @@ def mo_update_correlated( def update_coefficients( self, u_tot: AmplitudeType, - mo_coeff: tuple[NDArray[float]], - mo_coeff_ref: tuple[NDArray[float]], - ) -> tuple[NDArray[float]]: + mo_coeff: tuple[NDArray[float], NDArray[float]], + mo_coeff_ref: tuple[NDArray[float], NDArray[float]], + ) -> tuple[NDArray[float], NDArray[float]]: """Update the MO coefficients. Args: diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index 50a15624..e9d1ec50 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -54,12 +54,12 @@ def __setattr__(self, key: str, val: T) -> None: """Set an attribute.""" return self.__setitem__(key, val) - def __getitem__(self, key: str) -> T: + def __getitem__(self, key: str) -> Any: """Get an item.""" - value: T = self.__dict__["_members"][key] + value: Any = self.__dict__["_members"][key] return value - def __getattr__(self, key: str) -> T: + def __getattr__(self, key: str) -> Any: """Get an attribute.""" if key in self.__dict__: return self.__dict__[key] @@ -104,7 +104,7 @@ def keys(self) -> KeysView[str]: def values(self) -> ValuesView[T]: """Get values of the namespace as a dictionary.""" - return self._members.values + return self._members.values() def items(self) -> ItemsView[str, T]: """Get items of the namespace as a dictionary.""" diff --git a/ebcc/util/permutations.py b/ebcc/util/permutations.py index d4b77135..03b98649 100644 --- a/ebcc/util/permutations.py +++ b/ebcc/util/permutations.py @@ -322,13 +322,13 @@ def compress_axes( array = array.reshape([sizes[char] ** subscript.count(char) for char in sorted(set(subscript))]) # For each axis type, get the necessary lower-triangular indices: - indices = [ + indices_ndim = [ tril_indices_ndim(sizes[char], subscript.count(char), include_diagonal=include_diagonal) for char in sorted(set(subscript)) ] indices = [ np.ravel_multi_index(ind, (sizes[char],) * subscript.count(char)) - for ind, char in zip(indices, sorted(set(subscript))) + for ind, char in zip(indices_ndim, sorted(set(subscript))) ] # Apply the indices: From 868b56984087022a7162a70339a528d0024a40c7 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 22:35:57 +0100 Subject: [PATCH 24/37] Fully linted --- ebcc/cc/base.py | 22 ++++++++++++++++------ ebcc/cc/gebcc.py | 5 ++--- ebcc/cc/rebcc.py | 7 +++---- ebcc/cc/uebcc.py | 13 +++++-------- ebcc/eom/geom.py | 4 ++-- ebcc/eom/reom.py | 4 ++-- ebcc/eom/ueom.py | 20 ++++++++++++++------ ebcc/ham/base.py | 2 +- ebcc/opt/base.py | 6 ++++-- ebcc/opt/gbrueckner.py | 8 +++++--- ebcc/opt/rbrueckner.py | 8 +++++--- ebcc/opt/ubrueckner.py | 10 +++++++--- ebcc/util/misc.py | 1 - ebcc/util/permutations.py | 4 ++-- 14 files changed, 68 insertions(+), 46 deletions(-) diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index a9dff88d..eeebd136 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -16,12 +16,12 @@ from ebcc.core.precision import astype, types if TYPE_CHECKING: - from typing import Any, Callable, Literal, Optional, TypeVar, Union, Generic + from typing import Any, Callable, Literal, Optional, Union from pyscf.scf.hf import SCF from ebcc.core.logging import Logger - from ebcc.ham.base import BaseERIs, BaseFock, BaseElectronBoson + from ebcc.ham.base import BaseElectronBoson, BaseERIs, BaseFock from ebcc.numpy.typing import NDArray from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace @@ -128,8 +128,12 @@ def __init__( # Parameters: self.log = default_log if log is None else log self.mf = self._convert_mf(mf) - self._mo_coeff: Optional[NDArray[float]] = np.asarray(mo_coeff).astype(types[float]) if mo_coeff is not None else None - self._mo_occ: Optional[NDArray[float]] = np.asarray(mo_occ).astype(types[float]) if mo_occ is not None else None + self._mo_coeff: Optional[NDArray[float]] = ( + np.asarray(mo_coeff).astype(types[float]) if mo_coeff is not None else None + ) + self._mo_occ: Optional[NDArray[float]] = ( + np.asarray(mo_occ).astype(types[float]) if mo_occ is not None else None + ) # Ansatz: if isinstance(ansatz, Ansatz): @@ -533,7 +537,9 @@ def init_lams( """ pass - def _get_amps(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + def _get_amps( + self, amplitudes: Optional[Namespace[AmplitudeType]] = None + ) -> Namespace[AmplitudeType]: """Get the cluster amplitudes, initialising if required.""" if not amplitudes: amplitudes = self.amplitudes @@ -541,7 +547,11 @@ def _get_amps(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Na amplitudes = self.init_amps() return amplitudes - def _get_lams(self, lambdas: Optional[Namespace[AmplitudeType]] = None, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> Namespace[AmplitudeType]: + def _get_lams( + self, + lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[AmplitudeType]] = None, + ) -> Namespace[AmplitudeType]: """Get the cluster lambda amplitudes, initialising if required.""" if not lambdas: lambdas = self.lambdas diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index e273b5b1..79d614c6 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -11,10 +11,10 @@ from ebcc.cc.base import BaseEBCC from ebcc.core.precision import types from ebcc.eom import EA_GEOM, EE_GEOM, IP_GEOM +from ebcc.ham.elbos import GElectronBoson from ebcc.ham.eris import GERIs from ebcc.ham.fock import GFock from ebcc.ham.space import Space -from ebcc.ham.elbos import GElectronBoson from ebcc.opt.gbrueckner import BruecknerGEBCC if TYPE_CHECKING: @@ -23,7 +23,6 @@ from pyscf.scf.ghf import GHF from pyscf.scf.hf import SCF - from ebcc.cc.base import BaseOptions from ebcc.cc.rebcc import REBCC from ebcc.cc.uebcc import UEBCC from ebcc.numpy.typing import NDArray @@ -672,7 +671,7 @@ def energy_sum(self, *args: str, signs_dict: Optional[dict[str, str]] = None) -> Returns: Sum of energies. """ - subscript, = args + (subscript,) = args n = 0 def next_char() -> str: diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 482834c5..0de6f6aa 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -10,18 +10,17 @@ from ebcc.core.precision import types from ebcc.eom import EA_REOM, EE_REOM, IP_REOM from ebcc.ham.cderis import RCDERIs +from ebcc.ham.elbos import RElectronBoson from ebcc.ham.eris import RERIs from ebcc.ham.fock import RFock from ebcc.ham.space import Space -from ebcc.ham.elbos import RElectronBoson from ebcc.opt.rbrueckner import BruecknerREBCC if TYPE_CHECKING: - from typing import Any, Optional, Union, TypeAlias + from typing import Any, Optional, TypeAlias, Union from pyscf.scf.hf import RHF, SCF - from ebcc.cc.base import BaseOptions from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -448,7 +447,7 @@ def energy_sum(self, *args: str, signs_dict: Optional[dict[str, str]] = None) -> Returns: Sum of energies. """ - subscript, = args + (subscript,) = args n = 0 def next_char() -> str: diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index cb23d5e1..c451dd74 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -10,19 +10,18 @@ from ebcc.core.precision import types from ebcc.eom import EA_UEOM, EE_UEOM, IP_UEOM from ebcc.ham.cderis import UCDERIs +from ebcc.ham.elbos import UElectronBoson from ebcc.ham.eris import UERIs from ebcc.ham.fock import UFock from ebcc.ham.space import Space -from ebcc.ham.elbos import UElectronBoson from ebcc.opt.ubrueckner import BruecknerUEBCC if TYPE_CHECKING: - from typing import Any, Optional, Union, TypeAlias + from typing import Any, Optional, TypeAlias, Union from pyscf.scf.hf import SCF from pyscf.scf.uhf import UHF - from ebcc.cc.base import BaseOptions from ebcc.cc.rebcc import REBCC from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -577,9 +576,7 @@ def make_eb_coup_rdm( return dm_eb - def energy_sum( - self, *args: str, signs_dict: Optional[dict[str, str]] = None - ) -> NDArray[float]: + def energy_sum(self, *args: str, signs_dict: Optional[dict[str, str]] = None) -> NDArray[float]: """Get a direct sum of energies. Args: @@ -676,7 +673,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): shape = (self.nbos,) * n - size = self.nbos ** n + size = self.nbos**n amplitudes[name] = vector[i0 : i0 + size].reshape(shape) i0 += size @@ -757,7 +754,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): shape = (self.nbos,) * n - size = self.nbos ** n + size = self.nbos**n lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) i0 += size diff --git a/ebcc/eom/geom.py b/ebcc/eom/geom.py index 0fe417c1..3cd56b90 100644 --- a/ebcc/eom/geom.py +++ b/ebcc/eom/geom.py @@ -7,14 +7,14 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types, astype +from ebcc.core.precision import astype, types from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM if TYPE_CHECKING: from typing import Optional - from ebcc.ham.space import Space from ebcc.cc.gebcc import GEBCC, AmplitudeType, ERIsInputType + from ebcc.ham.space import Space from ebcc.numpy.typing import NDArray from ebcc.util import Namespace diff --git a/ebcc/eom/reom.py b/ebcc/eom/reom.py index b10ecff9..b22c8f32 100644 --- a/ebcc/eom/reom.py +++ b/ebcc/eom/reom.py @@ -7,14 +7,14 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types, astype +from ebcc.core.precision import astype, types from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM if TYPE_CHECKING: from typing import Optional - from ebcc.ham.space import Space from ebcc.cc.rebcc import REBCC, AmplitudeType, ERIsInputType + from ebcc.ham.space import Space from ebcc.numpy.typing import NDArray from ebcc.util import Namespace diff --git a/ebcc/eom/ueom.py b/ebcc/eom/ueom.py index 74ed633b..973411f0 100644 --- a/ebcc/eom/ueom.py +++ b/ebcc/eom/ueom.py @@ -6,14 +6,14 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types, astype +from ebcc.core.precision import astype, types from ebcc.eom.base import BaseEA_EOM, BaseEE_EOM, BaseEOM, BaseIP_EOM if TYPE_CHECKING: from typing import Optional - from ebcc.ham.space import Space from ebcc.cc.uebcc import UEBCC, AmplitudeType, ERIsInputType + from ebcc.ham.space import Space from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -117,7 +117,9 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: bras_tmp.a.append(self.amplitudes_to_vector(*amps_a)) bras_tmp.b.append(self.amplitudes_to_vector(*amps_b)) - bras: Namespace[NDArray[float]] = util.Namespace(a=np.array(bras_tmp.a), b=np.array(bras_tmp.b)) + bras: Namespace[NDArray[float]] = util.Namespace( + a=np.array(bras_tmp.a), b=np.array(bras_tmp.b) + ) return bras @@ -173,7 +175,9 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: kets_tmp.a.append(self.amplitudes_to_vector(*amps_a)) kets_tmp.b.append(self.amplitudes_to_vector(*amps_b)) - kets: Namespace[NDArray[float]] = util.Namespace(a=np.array(kets_tmp.a), b=np.array(kets_tmp.b)) + kets: Namespace[NDArray[float]] = util.Namespace( + a=np.array(kets_tmp.a), b=np.array(kets_tmp.b) + ) return kets @@ -316,7 +320,9 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: bras_tmp.a.append(self.amplitudes_to_vector(*amps_a)) bras_tmp.b.append(self.amplitudes_to_vector(*amps_b)) - bras: Namespace[NDArray[float]] = util.Namespace(a=np.array(bras_tmp.a), b=np.array(bras_tmp.b)) + bras: Namespace[NDArray[float]] = util.Namespace( + a=np.array(bras_tmp.a), b=np.array(bras_tmp.b) + ) return bras @@ -372,7 +378,9 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: kets_tmp.a.append(self.amplitudes_to_vector(*amps_a)) kets_tmp.b.append(self.amplitudes_to_vector(*amps_b)) - kets: Namespace[NDArray[float]] = util.Namespace(a=np.array(kets_tmp.a), b=np.array(kets_tmp.b)) + kets: Namespace[NDArray[float]] = util.Namespace( + a=np.array(kets_tmp.a), b=np.array(kets_tmp.b) + ) return kets diff --git a/ebcc/ham/base.py b/ebcc/ham/base.py index 36c1fbbb..b0124e9e 100644 --- a/ebcc/ham/base.py +++ b/ebcc/ham/base.py @@ -8,7 +8,7 @@ from ebcc.util import Namespace if TYPE_CHECKING: - from typing import TypeVar, Optional + from typing import Optional, TypeVar from ebcc.cc.base import BaseEBCC from ebcc.numpy.typing import NDArray diff --git a/ebcc/opt/base.py b/ebcc/opt/base.py index e1d6cb06..1a125fa7 100644 --- a/ebcc/opt/base.py +++ b/ebcc/opt/base.py @@ -17,8 +17,8 @@ if TYPE_CHECKING: from typing import Any, Optional, TypeVar - from ebcc.numpy.typing import NDArray from ebcc.cc.base import AmplitudeType, BaseEBCC + from ebcc.numpy.typing import NDArray from ebcc.util import Namespace T = TypeVar("T") @@ -277,7 +277,9 @@ def mo_update_correlated(self, mo_coeff: Any, mo_coeff_corr: Any) -> Any: pass @abstractmethod - def update_coefficients(self, u_tot: AmplitudeType, mo_coeff_new: Any, mo_coeff_ref: Any) -> Any: + def update_coefficients( + self, u_tot: AmplitudeType, mo_coeff_new: Any, mo_coeff_ref: Any + ) -> Any: """Update the MO coefficients. Args: diff --git a/ebcc/opt/gbrueckner.py b/ebcc/opt/gbrueckner.py index ca5759dc..3278179a 100644 --- a/ebcc/opt/gbrueckner.py +++ b/ebcc/opt/gbrueckner.py @@ -8,13 +8,13 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types, astype +from ebcc.core.precision import types from ebcc.opt.base import BaseBruecknerEBCC if TYPE_CHECKING: from typing import Optional - from ebcc.cc.gebcc import AmplitudeType, GEBCC + from ebcc.cc.gebcc import GEBCC, AmplitudeType from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -54,7 +54,9 @@ def get_rotation_matrix( if u_tot is None: u_tot = np.eye(self.cc.space.ncorr, dtype=types[float]) - t1_block: NDArray[float] = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) + t1_block: NDArray[float] = np.zeros( + (self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float] + ) t1_block[: self.cc.space.ncocc, self.cc.space.ncocc :] = -t1 t1_block[self.cc.space.ncocc :, : self.cc.space.ncocc] = t1.T diff --git a/ebcc/opt/rbrueckner.py b/ebcc/opt/rbrueckner.py index c083a75a..c9a6f6b8 100644 --- a/ebcc/opt/rbrueckner.py +++ b/ebcc/opt/rbrueckner.py @@ -8,13 +8,13 @@ from ebcc import numpy as np from ebcc import util -from ebcc.core.precision import types, astype +from ebcc.core.precision import types from ebcc.opt.base import BaseBruecknerEBCC if TYPE_CHECKING: from typing import Optional - from ebcc.cc.rebcc import AmplitudeType, REBCC + from ebcc.cc.rebcc import REBCC, AmplitudeType from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -54,7 +54,9 @@ def get_rotation_matrix( if u_tot is None: u_tot = np.eye(self.cc.space.ncorr, dtype=types[float]) - t1_block: NDArray[float] = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float]) + t1_block: NDArray[float] = np.zeros( + (self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float] + ) t1_block[: self.cc.space.ncocc, self.cc.space.ncocc :] = -t1 t1_block[self.cc.space.ncocc :, : self.cc.space.ncocc] = t1.T diff --git a/ebcc/opt/ubrueckner.py b/ebcc/opt/ubrueckner.py index 49d047a1..bc098b0f 100644 --- a/ebcc/opt/ubrueckner.py +++ b/ebcc/opt/ubrueckner.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.uebcc import AmplitudeType, UEBCC + from ebcc.cc.uebcc import UEBCC, AmplitudeType from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -151,7 +151,9 @@ def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> weight: float = types[float](np.linalg.norm([weight_a, weight_b])) return weight - def mo_to_correlated(self, mo_coeff: tuple[NDArray[float], NDArray[float]]) -> tuple[NDArray[float], NDArray[float]]: + def mo_to_correlated( + self, mo_coeff: tuple[NDArray[float], NDArray[float]] + ) -> tuple[NDArray[float], NDArray[float]]: """Transform the MO coefficients into the correlated basis. Args: @@ -166,7 +168,9 @@ def mo_to_correlated(self, mo_coeff: tuple[NDArray[float], NDArray[float]]) -> t ) def mo_update_correlated( - self, mo_coeff: tuple[NDArray[float], NDArray[float]], mo_coeff_corr: tuple[NDArray[float], NDArray[float]] + self, + mo_coeff: tuple[NDArray[float], NDArray[float]], + mo_coeff_corr: tuple[NDArray[float], NDArray[float]], ) -> tuple[NDArray[float], NDArray[float]]: """Update the correlated slice of a set of MO coefficients. diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index e9d1ec50..f7d8ec5e 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -4,7 +4,6 @@ import time from collections.abc import MutableMapping -from dataclasses import dataclass from typing import TYPE_CHECKING, Generic, TypeVar if TYPE_CHECKING: diff --git a/ebcc/util/permutations.py b/ebcc/util/permutations.py index 03b98649..99dc9f4f 100644 --- a/ebcc/util/permutations.py +++ b/ebcc/util/permutations.py @@ -9,7 +9,7 @@ import numpy as np if TYPE_CHECKING: - from typing import Any, Generator, Hashable, Iterable, Optional, TypeVar, Union + from typing import Any, Generator, Hashable, Iterable, Optional, TypeVar from ebcc.numpy.typing import NDArray @@ -228,7 +228,7 @@ def is_mixed_spin(spin: Iterable[Hashable]) -> bool: def combine_subscripts( *subscripts: str, - sizes: dict[tuple[str, ...], int] = {}, + sizes: Optional[dict[tuple[str, ...], int]] = None, ) -> tuple[str, dict[str, int]]: """Combine subscripts into new unique subscripts for functions such as `compress_axes`. From 48bab941dce1257458b2ee836709d5cdb4047e7d Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 23:03:14 +0100 Subject: [PATCH 25/37] Fix electron-boson shift --- ebcc/cc/gebcc.py | 2 +- ebcc/cc/rebcc.py | 2 +- ebcc/cc/uebcc.py | 2 +- ebcc/ham/fock.py | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 79d614c6..e38963cb 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -1016,7 +1016,7 @@ def get_fock(self) -> GFock: Returns: Fock matrix. """ - return self.Fock(self, array=self.bare_fock) + 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. diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 0de6f6aa..ba715d61 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -772,7 +772,7 @@ def get_fock(self) -> RFock: Returns: Fock matrix. """ - return self.Fock(self, array=self.bare_fock) + 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. diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index c451dd74..1e077d98 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -1051,7 +1051,7 @@ def get_fock(self) -> UFock: Returns: Fock matrix. """ - return self.Fock(self, array=(self.bare_fock.aa, self.bare_fock.bb)) + 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. diff --git a/ebcc/ham/fock.py b/ebcc/ham/fock.py index cf0f22c9..fdc0045c 100644 --- a/ebcc/ham/fock.py +++ b/ebcc/ham/fock.py @@ -48,7 +48,7 @@ def __getitem__(self, key: str) -> NDArray[float]: if self.shift: xi = self.xi - g = self.g.__getattr__(f"b{key}") + g = self.g.__getattr__(f"b{key}").copy() g += self.g.__getattr__(f"b{key[::-1]}").transpose(0, 2, 1) self._members[key] -= np.einsum("I,Ipq->pq", xi, g) @@ -133,7 +133,7 @@ def __getitem__(self, key: str) -> NDArray[float]: if self.shift: xi = self.xi - g = self.g.__getattr__(f"b{key}") + g = self.g.__getattr__(f"b{key}").copy() g += self.g.__getattr__(f"b{key[::-1]}").transpose(0, 2, 1) self._members[key] -= np.einsum("I,Ipq->pq", xi, g) From a43173eabbcdbbb8d1ee2fd06d77a1d99742754c Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 23:04:38 +0100 Subject: [PATCH 26/37] Forget to add the files all this time --- ebcc/core/__init__.py | 1 + ebcc/core/ansatz.py | 317 +++++++++++++++++++++++++++++++++++++++++ ebcc/core/damping.py | 47 ++++++ ebcc/core/dump.py | 260 +++++++++++++++++++++++++++++++++ ebcc/core/logging.py | 128 +++++++++++++++++ ebcc/core/precision.py | 76 ++++++++++ ebcc/ham/elbos.py | 99 +++++++++++++ 7 files changed, 928 insertions(+) create mode 100644 ebcc/core/__init__.py create mode 100644 ebcc/core/ansatz.py create mode 100644 ebcc/core/damping.py create mode 100644 ebcc/core/dump.py create mode 100644 ebcc/core/logging.py create mode 100644 ebcc/core/precision.py create mode 100644 ebcc/ham/elbos.py diff --git a/ebcc/core/__init__.py b/ebcc/core/__init__.py new file mode 100644 index 00000000..e20152a0 --- /dev/null +++ b/ebcc/core/__init__.py @@ -0,0 +1 @@ +"""Core functionality.""" diff --git a/ebcc/core/ansatz.py b/ebcc/core/ansatz.py new file mode 100644 index 00000000..08a55265 --- /dev/null +++ b/ebcc/core/ansatz.py @@ -0,0 +1,317 @@ +"""Ansatz definition.""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING + +from ebcc import METHOD_TYPES, util + +if TYPE_CHECKING: + from types import ModuleType + from typing import Optional + +named_ansatzes = { + "MP2": ("MP2", "", 0, 0), + "MP3": ("MP3", "", 0, 0), + "CCD": ("CCD", "", 0, 0), + "CCSD": ("CCSD", "", 0, 0), + "CCSDT": ("CCSDT", "", 0, 0), + "CCSDTQ": ("CCSDTQ", "", 0, 0), + "CCSD(T)": ("CCSD(T)", "", 0, 0), + "CC2": ("CC2", "", 0, 0), + "CC3": ("CC3", "", 0, 0), + "QCISD": ("QCISD", "", 0, 0), + "DCD": ("DCD", "", 0, 0), + "DCSD": ("DCSD", "", 0, 0), + "CCSDt'": ("CCSDt'", "", 0, 0), + "CCSD-S-1-1": ("CCSD", "S", 1, 1), + "CCSD-SD-1-1": ("CCSD", "SD", 1, 1), + "CCSD-SD-1-2": ("CCSD", "SD", 1, 2), +} + + +def name_to_identifier(name: str) -> str: + """Convert an ansatz name to an identifier. + + The identifier is used as for the filename of the module containing the generated equations, + where the name may contain illegal characters. + + Args: + name: Name of the ansatz. + + Returns: + Identifier for the ansatz. + + Examples: + >>> name_to_identifier("CCSD(T)") + 'CCSDxTx' + >>> name_to_identifier("CCSD-SD-1-2") + 'CCSD_SD_1_2' + """ + iden = name.replace("(", "x").replace(")", "x") + iden = iden.replace("[", "y").replace("]", "y") + iden = iden.replace("-", "_") + iden = iden.replace("'", "p") + return iden + + +def identifity_to_name(iden: str) -> str: + """Convert an ansatz identifier to a name. + + Inverse operation of `name_to_identifier`. + + Args: + iden: Identifier for the ansatz. + + Returns: + Name of the ansatz. + + Examples: + >>> identifier_to_name("CCSDxTx") + 'CCSD(T)' + >>> identifier_to_name("CCSD_SD_1_2") + 'CCSD-SD-1-2' + """ + name = iden.replace("-", "_") + while "x" in name: + name = name.replace("x", "(", 1).replace("x", ")", 1) + while "y" in name: + name = name.replace("y", "(", 1).replace("y", ")", 1) + name = name.replace("p", "'") + return name + + +class Ansatz: + """Ansatz class. + + Attributes: + fermion_ansatz: Fermionic ansatz. + boson_ansatz: Rank of bosonic excitations. + fermion_coupling_rank: Rank of fermionic term in coupling. + boson_coupling_rank: Rank of bosonic term in coupling. + density_fitting: Use density fitting. + module_name: Name of the module containing the generated equations. + """ + + def __init__( + self, + fermion_ansatz: str = "CCSD", + boson_ansatz: str = "", + fermion_coupling_rank: int = 0, + boson_coupling_rank: int = 0, + density_fitting: bool = False, + module_name: Optional[str] = None, + ) -> None: + """Initialise the ansatz. + + Args: + fermion_ansatz: Fermionic ansatz. + boson_ansatz: Rank of bosonic excitations. + fermion_coupling_rank: Rank of fermionic term in coupling. + boson_coupling_rank: Rank of bosonic term in coupling. + density_fitting: Use density fitting. + module_name: Name of the module containing the generated equations. + """ + self.fermion_ansatz = fermion_ansatz + self.boson_ansatz = boson_ansatz + self.fermion_coupling_rank = fermion_coupling_rank + self.boson_coupling_rank = boson_coupling_rank + self.density_fitting = density_fitting + self.module_name = module_name + + def _get_eqns(self, prefix: str) -> ModuleType: + """Get the module containing the generated equations.""" + if self.module_name is None: + name = prefix + name_to_identifier(self.name) + else: + name = self.module_name + return importlib.import_module("ebcc.codegen.%s" % name) + + @classmethod + def from_string(cls, string: str, density_fitting: bool = False) -> Ansatz: + """Build an `Ansatz` from a string for the default ansatzes. + + Args: + string: Input string. + density_fitting: Use density fitting. + + Returns: + Ansatz object. + """ + if string not in named_ansatzes: + raise util.ModelNotImplemented(string) + return cls(*named_ansatzes[string], density_fitting=density_fitting) + + def __repr__(self) -> str: + """Get a string with the name of the method.""" + name = "" + if self.density_fitting: + name += "DF" + name += self.fermion_ansatz + if self.boson_ansatz: + name += "-%s" % self.boson_ansatz + if self.fermion_coupling_rank or self.boson_coupling_rank: + name += "-%d" % self.fermion_coupling_rank + name += "-%d" % self.boson_coupling_rank + return name + + @property + def name(self) -> str: + """Get the name of the ansatz.""" + return repr(self) + + @property + def has_perturbative_correction(self) -> bool: + """Get a boolean indicating if the ansatz includes a perturbative correction e.g. CCSD(T). + + Returns: + perturbative: Boolean indicating if the ansatz is perturbatively corrected. + """ + return any( + "(" in ansatz and ")" in ansatz for ansatz in (self.fermion_ansatz, self.boson_ansatz) + ) + + @property + def is_one_shot(self) -> bool: + """Get a boolean indicating whether the ansatz is a one-shot energy calculation e.g. MP2. + + Returns: + one_shot: Boolean indicating if the ansatz is a one-shot energy calculation. + """ + return all( + ansatz.startswith("MP") or ansatz == "" + for ansatz in (self.fermion_ansatz, self.boson_ansatz) + ) + + def fermionic_cluster_ranks(self, spin_type: str = "G") -> list[tuple[str, str, int]]: + """Get a list of cluster operator ranks for the fermionic space. + + Args: + spin_type: Spin type of the cluster operator. + + Returns: + List of cluster operator ranks, each element is a tuple containing the name, the slices + and the rank. + """ + ranks: list[tuple[str, str, int]] = [] + if not self.fermion_ansatz: + return ranks + + notations = { + "S": [("t1", "ov", 1)], + "D": [("t2", "oovv", 2)], + "T": [("t3", "ooovvv", 3)], + "t": [("t3", "ooOvvV", 3)], + "t'": [("t3", "OOOVVV", 3)], + } + if spin_type == "R": + notations["Q"] = [("t4a", "oooovvvv", 4), ("t4b", "oooovvvv", 4)] + else: + notations["Q"] = [("t4", "oooovvvv", 4)] + notations["2"] = notations["S"] + notations["D"] + notations["3"] = notations["2"] + notations["T"] + notations["4"] = notations["3"] + notations["Q"] + + # Remove any perturbative corrections + op = self.fermion_ansatz + while "(" in op: + start = op.index("(") + end = op.index(")") + op = op[:start] + if (end + 1) < len(op): + op += op[end + 1 :] + + # Check in order of longest to shortest string in case one + # method name starts with a substring equal to the name of + # another method + for method_type in sorted(METHOD_TYPES, key=len)[::-1]: + if op.startswith(method_type): + op = op.replace(method_type, "", 1) + break + + # If it's MP we only ever need to initialise second-order + # amplitudes + if method_type == "MP": + op = "D" + + # Determine the ranks + for key in sorted(notations.keys(), key=len)[::-1]: + if key in op: + ranks += notations[key] + op = op.replace(key, "") + + # Check there are no duplicates + if len(ranks) != len(set(ranks)): + raise util.ModelNotImplemented("Duplicate ranks in %s" % self.fermion_ansatz) + + # Sort the ranks by the cluster operator dimension + ranks = sorted(ranks, key=lambda x: x[2]) + + return ranks + + def bosonic_cluster_ranks(self, spin_type: str = "G") -> list[tuple[str, str, int]]: + """Get a list of cluster operator ranks for the bosonic space. + + Args: + spin_type: Spin type of the cluster operator. + + Returns: + List of cluster operator ranks, each element is a tuple containing the name, the slices + and the rank. + """ + ranks: list[tuple[str, str, int]] = [] + if not self.boson_ansatz: + return ranks + + notations = { + "S": [("s1", "b", 1)], + "D": [("s2", "bb", 2)], + "T": [("s3", "bbb", 3)], + } + notations["2"] = notations["S"] + notations["D"] + notations["3"] = notations["2"] + notations["T"] + + # Remove any perturbative corrections + op = self.boson_ansatz + while "(" in op: + start = op.index("(") + end = op.index(")") + op = op[:start] + if (end + 1) < len(op): + op += op[end + 1 :] + + # Determine the ranks + for key in sorted(notations.keys(), key=len)[::-1]: + if key in op: + ranks += notations[key] + op = op.replace(key, "") + + # Check there are no duplicates + if len(ranks) != len(set(ranks)): + raise util.ModelNotImplemented("Duplicate ranks in %s" % self.boson_ansatz) + + # Sort the ranks by the cluster operator dimension + ranks = sorted(ranks, key=lambda x: x[2]) + + return ranks + + def coupling_cluster_ranks(self, spin_type: str = "G") -> list[tuple[str, str, int, int]]: + """Get a list of cluster operator ranks for the coupling between fermions and bosons. + + Args: + spin_type: Spin type of the cluster operator. + + Returns: + List of cluster operator ranks, each element is a tuple containing the name, the slices + and the rank. + """ + ranks = [] + + for fermion_rank in range(1, self.fermion_coupling_rank + 1): + for boson_rank in range(1, self.boson_coupling_rank + 1): + name = f"u{fermion_rank}{boson_rank}" + key = "b" * boson_rank + "o" * fermion_rank + "v" * fermion_rank + ranks.append((name, key, fermion_rank, boson_rank)) + + return ranks diff --git a/ebcc/core/damping.py b/ebcc/core/damping.py new file mode 100644 index 00000000..8dff8480 --- /dev/null +++ b/ebcc/core/damping.py @@ -0,0 +1,47 @@ +"""Damping and DIIS control.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyscf.lib import diis + +if TYPE_CHECKING: + from ebcc.numpy.typing import NDArray + + +class DIIS(diis.DIIS): + """Direct inversion in the iterative subspace. + + Attributes: + space: The number of vectors to store in the DIIS space. + min_space: The minimum number of vectors to store in the DIIS space. + damping: The damping factor to apply to the extrapolated vector. + """ + + def __init__(self, space: int = 6, min_space: int = 1, damping: float = 0.0) -> None: + """Initialize the DIIS object. + + Args: + space: The number of vectors to store in the DIIS space. + min_space: The minimum number of vectors to store in the DIIS space. + damping: The damping factor to apply to the extrapolated vector. + """ + super().__init__(incore=True) + self.verbose = 0 + self.space = space + self.min_space = min_space + self.damping = damping + + def update(self, x: NDArray[float], xerr: NDArray[float] = None) -> NDArray[float]: + """Extrapolate a vector.""" + x = super().update(x, xerr=xerr) + + # Apply damping + if self.damping: + nd = self.get_num_vec() + if nd > 1: + xprev = self.get_vec(self.get_num_vec() - 1) + x = (1.0 - self.damping) * x + self.damping * xprev + + return x diff --git a/ebcc/core/dump.py b/ebcc/core/dump.py new file mode 100644 index 00000000..a6497ef8 --- /dev/null +++ b/ebcc/core/dump.py @@ -0,0 +1,260 @@ +"""File dumping and reading functionality.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyscf import scf +from pyscf.lib.chkfile import dump, dump_mol, load, load_mol + +from ebcc import util +from ebcc.core.ansatz import Ansatz +from ebcc.ham.space import Space + +if TYPE_CHECKING: + from typing import Any, Optional, Union + + from ebcc.cc.base import BaseEBCC + from ebcc.space.logging import Logger + + +class Dump: + """File handler for reading and writing EBCC calculations. + + Args: + name: The name of the file. + """ + + def __init__(self, name: str) -> None: + """Initialise the file handler. + + Args: + name: The name of the file. + """ + self.name = name + + def write(self, ebcc: BaseEBCC) -> None: + """Write the EBCC object to the file. + + Args: + ebcc: The EBCC object to write. + """ + # Write the options + dic = {} + for key, val in ebcc.options.__dict__.items(): + if val is not None: + dic[key] = val + dump(self.name, "options", dic) + + # Write the mean-field data + dic = { + "e_tot": ebcc.mf.e_tot, + "mo_energy": ebcc.mf.mo_energy, + "mo_coeff": ebcc.mf.mo_coeff, + "mo_occ": ebcc.mf.mo_occ, + } + dump_mol(ebcc.mf.mol, self.name) + dump(self.name, "mean-field", dic) + + # Write the MOs used + dic = { + "mo_coeff": ebcc.mo_coeff, + "mo_occ": ebcc.mo_occ, + } + dump(self.name, "mo", dic) + + # Write the ansatz + dic = { + "fermion_ansatz": ebcc.ansatz.fermion_ansatz, + "boson_ansatz": ebcc.ansatz.boson_ansatz, + "fermion_coupling_rank": ebcc.ansatz.fermion_coupling_rank, + "boson_coupling_rank": ebcc.ansatz.boson_coupling_rank, + } + if ebcc.ansatz.module_name is not None: + dic["module_name"] = ebcc.ansatz.module_name + dump(self.name, "ansatz", dic) + + # Write the space + if ebcc.spin_type == "U": + dic = { + "occupied": (ebcc.space[0]._occupied, ebcc.space[1]._occupied), + "frozen": (ebcc.space[0]._frozen, ebcc.space[1]._frozen), + "active": (ebcc.space[0]._active, ebcc.space[1]._active), + } + else: + dic = { + "occupied": ebcc.space._occupied, + "frozen": ebcc.space._frozen, + "active": ebcc.space._active, + } + dump(self.name, "space", dic) + + # Write the bosonic parameters + dic = {} + if ebcc.omega is not None: + dic["omega"] = ebcc.omega + if ebcc.bare_g is not None: + dic["bare_g"] = ebcc.bare_g + if ebcc.bare_G is not None: + dic["bare_G"] = ebcc.bare_G + dump(self.name, "bosons", dic) + + # Write the Fock matrix + # TODO write the Fock matrix class instead + + # Write miscellaneous data + kwargs: dict[str, Any] = { + "spin_type": ebcc.spin_type, + } + if ebcc.e_corr is not None: + kwargs["e_corr"] = ebcc.e_corr + if ebcc.converged is not None: + kwargs["converged"] = ebcc.converged + if ebcc.converged_lambda is not None: + kwargs["converged_lambda"] = ebcc.converged_lambda + dump(self.name, "misc", kwargs) + + # Write the amplitudes + if ebcc.spin_type == "U": + if ebcc.amplitudes is not None: + dump( + self.name, + "amplitudes", + { + key: ({**val} if isinstance(val, (util.Namespace, dict)) else val) + for key, val in ebcc.amplitudes.items() + }, + ) + if ebcc.lambdas is not None: + dump( + self.name, + "lambdas", + { + key: ({**val} if isinstance(val, (util.Namespace, dict)) else val) + for key, val in ebcc.lambdas.items() + }, + ) + else: + if ebcc.amplitudes is not None: + dump(self.name, "amplitudes", {**ebcc.amplitudes}) + if ebcc.lambdas is not None: + dump(self.name, "lambdas", {**ebcc.lambdas}) + + def read(self, cls: type[BaseEBCC], log: Optional[Logger] = None) -> BaseEBCC: + """Load the file to an EBCC object. + + Args: + cls: EBCC class to load the file to. + log: Logger to assign to the EBCC object. + + Returns: + The EBCC object loaded from the file. + """ + # Load the options + dic = load(self.name, "options") + options = cls.Options() + for key, val in dic.items(): + setattr(options, key, val) + + # Load the miscellaneous data + misc = load(self.name, "misc") + spin_type = misc.pop("spin_type").decode("ascii") + + # Load the mean-field data + mf_cls = {"G": scf.GHF, "U": scf.UHF, "R": scf.RHF}[spin_type] + mol = load_mol(self.name) + dic = load(self.name, "mean-field") + mf = mf_cls(mol) + mf.__dict__.update(dic) + + # Load the MOs used + dic = load(self.name, "mo") + mo_coeff = dic.get("mo_coeff", None) + mo_occ = dic.get("mo_occ", None) + + # Load the ansatz + dic = load(self.name, "ansatz") + module_name = dic.get("module_name", None) + if isinstance(module_name, str): + module_name = module_name.encode("ascii") + ansatz = Ansatz( + dic.get("fermion_ansatz", b"CCSD").decode("ascii"), + dic.get("boson_ansatz", b"").decode("ascii"), + dic.get("fermion_coupling_rank", 0), + dic.get("boson_coupling_rank", 0), + module_name, + ) + + # Load the space + dic = load(self.name, "space") + space: Union[Space, tuple[Space, Space]] + if spin_type == "U": + space = ( + Space( + dic.get("occupied", None)[0], + dic.get("frozen", None)[0], + dic.get("active", None)[0], + ), + Space( + dic.get("occupied", None)[1], + dic.get("frozen", None)[1], + dic.get("active", None)[1], + ), + ) + else: + space = Space( + dic.get("occupied", None), + dic.get("frozen", None), + dic.get("active", None), + ) + + # Load the bosonic parameters + dic = load(self.name, "bosons") + omega = dic.get("omega", None) + bare_g = dic.get("bare_g", None) + bare_G = dic.get("bare_G", None) + + # Load the Fock matrix + # TODO load the Fock matrix class instead + + # Load the amplitudes + amplitudes = load(self.name, "amplitudes") + lambdas = load(self.name, "lambdas") + if spin_type == "U": + if amplitudes is not None: + amplitudes = { + key: (util.Namespace(**val) if isinstance(val, dict) else val) + for key, val in amplitudes.items() + } + amplitudes = util.Namespace(**amplitudes) + if lambdas is not None: + lambdas = { + key: (util.Namespace(**val) if isinstance(val, dict) else val) + for key, val in lambdas.items() + } + lambdas = util.Namespace(**lambdas) + else: + if amplitudes is not None: + amplitudes = util.Namespace(**amplitudes) + if lambdas is not None: + lambdas = util.Namespace(**lambdas) + + # Initialise the EBCC object + cc = cls( + mf, + log=log, + ansatz=ansatz, + space=space, + omega=omega, + g=bare_g, + G=bare_G, + mo_coeff=mo_coeff, + mo_occ=mo_occ, + # fock=fock, + options=options, + ) + cc.__dict__.update(misc) + cc.amplitudes = amplitudes + cc.lambdas = lambdas + + return cc diff --git a/ebcc/core/logging.py b/ebcc/core/logging.py new file mode 100644 index 00000000..1627d60d --- /dev/null +++ b/ebcc/core/logging.py @@ -0,0 +1,128 @@ +"""Logging.""" + +from __future__ import annotations + +import logging +import os +import subprocess +import sys +from typing import TYPE_CHECKING, cast + +from ebcc import __version__ +from ebcc.util import Namespace + +if TYPE_CHECKING: + from typing import Any + +HEADER = """ _ + | | + ___ | |__ ___ ___ + / _ \| '_ \ / __| / __| + | __/| |_) || (__ | (__ + \___||_.__/ \___| \___| +%s""" # noqa: W605 + + +class Logger(logging.Logger): + """Logger with a custom output level.""" + + def __init__(self, name: str, level: int = logging.INFO) -> None: + """Initialise the logger.""" + super().__init__(name, level) + + def output(self, msg: str, *args: Any, **kwargs: Any) -> None: + """Output a message at the `"OUTPUT"` level.""" + if self.isEnabledFor(25): + self._log(25, msg, args, **kwargs) + + +logging.setLoggerClass(Logger) +logging.addLevelName(25, "OUTPUT") + + +default_log = Logger("ebcc") +default_log.setLevel(logging.INFO) +default_log.addHandler(logging.StreamHandler(sys.stderr)) + + +class NullLogger(Logger): + """A logger that does nothing.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Initialise the logger.""" + super().__init__("null") + + def _log(self, *args, **kwargs): # type: ignore + pass + + +def init_logging(log: Logger) -> None: + """Initialise the logging with a header.""" + if globals().get("_EBCC_LOG_INITIALISED", False): + return + + # Print header + header_size = max([len(line) for line in HEADER.split("\n")]) + space = " " * (header_size - len(__version__)) + log.info(f"{ANSI.B}{HEADER}{ANSI.R}" % f"{space}{ANSI.B}{__version__}{ANSI.R}") + + # Print versions of dependencies and ebcc + def get_git_hash(directory: str) -> str: + git_directory = os.path.join(directory, ".git") + cmd = ["git", "--git-dir=%s" % git_directory, "rev-parse", "--short", "HEAD"] + try: + git_hash = subprocess.check_output( + cmd, universal_newlines=True, stderr=subprocess.STDOUT + ).rstrip() + except subprocess.CalledProcessError: + git_hash = "N/A" + return git_hash + + import numpy + import pyscf + + log.info("numpy:") + log.info(" > Version: %s" % numpy.__version__) + log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(numpy.__file__), ".."))) + + log.info("pyscf:") + log.info(" > Version: %s" % pyscf.__version__) + log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(pyscf.__file__), ".."))) + + log.info("ebcc:") + log.info(" > Version: %s" % __version__) + log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(__file__), ".."))) + + # Environment variables + log.info("OMP_NUM_THREADS = %s" % os.environ.get("OMP_NUM_THREADS", "")) + + log.debug("") + + globals()["_EBCC_LOG_INITIALISED"] = True + + +def _check_output(*args: Any, **kwargs: Any) -> bytes: + """Call a command. + + If the return code is non-zero, an empty `bytes` object is returned. + """ + try: + return cast(bytes, subprocess.check_output(*args, **kwargs)) + except subprocess.CalledProcessError: + return bytes() + + +ANSI = Namespace( + B="\x1b[1m", + H="\x1b[3m", + R="\x1b[m\x0f", + U="\x1b[4m", + b="\x1b[34m", + c="\x1b[36m", + g="\x1b[32m", + k="\x1b[30m", + m="\x1b[35m", + r="\x1b[31m", + w="\x1b[37m", + y="\x1b[33m", +) diff --git a/ebcc/core/precision.py b/ebcc/core/precision.py new file mode 100644 index 00000000..a616e1b7 --- /dev/null +++ b/ebcc/core/precision.py @@ -0,0 +1,76 @@ +"""Floating point precision control.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING + +from ebcc import numpy as np + +if TYPE_CHECKING: + from collections.abc import Iterator + from typing import Type, TypeVar + + T = TypeVar("T", float, complex) + + +types: dict[type, type] +if TYPE_CHECKING: + types = { + float: float, + complex: complex, + } +else: + types = { + float: np.float64, + complex: np.complex128, + } + + +def astype(value: T, dtype: Type[T]) -> T: + """Cast a value to the current floating point type. + + Args: + value: The value to cast. + dtype: The type to cast to. + + Returns: + The value cast to the current floating point type. + """ + out: T = types[dtype](value) + return out + + +def set_precision(**kwargs: type) -> None: + """Set the floating point type. + + Args: + float: The floating point type to use. + complex: The complex type to use. + """ + types[float] = kwargs.get("float", types[float]) + types[complex] = kwargs.get("complex", types[complex]) + + +@contextmanager +def precision(**kwargs: type) -> Iterator[None]: + """Context manager for setting the floating point precision. + + Args: + float: The floating point type to use. + complex: The complex type to use. + """ + old = { + "float": types[float], + "complex": types[complex], + } + set_precision(**kwargs) + yield + set_precision(**old) + + +@contextmanager +def single_precision() -> Iterator[None]: + """Context manager for setting the floating point precision to single precision.""" + with precision(float=np.float32, complex=np.complex64): + yield diff --git a/ebcc/ham/elbos.py b/ebcc/ham/elbos.py new file mode 100644 index 00000000..5653293c --- /dev/null +++ b/ebcc/ham/elbos.py @@ -0,0 +1,99 @@ +"""Electron-boson coupling matrix containers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ebcc import numpy as np +from ebcc.ham.base import BaseElectronBoson + +if TYPE_CHECKING: + from ebcc.numpy.typing import NDArray + + +class RElectronBoson(BaseElectronBoson): + """Restricted electron-boson coupling matrices. + + Attributes: + cc: Coupled cluster object. + space: Space object. + array: Electron-boson coupling matrix in the MO basis. + """ + + _members: dict[str, NDArray[float]] + + def __getitem__(self, key: str) -> NDArray[float]: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + Electron-boson coupling matrix for the given spaces. + """ + if key not in self._members: + assert key[0] == "b" + i = self.space[0].mask(key[1]) + j = self.space[1].mask(key[2]) + self._members[key] = self.array[:, i][:, :, j].copy() + return self._members[key] + + +class UElectronBoson(BaseElectronBoson): + """Unrestricted electron-boson coupling matrices. + + Attributes: + cc: Coupled cluster object. + space: Space object. + array: Electron-boson coupling matrix in the MO basis. + """ + + _members: dict[str, RElectronBoson] + + def __getitem__(self, key: str) -> RElectronBoson: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + Electron-boson coupling matrix for the given spin. + """ + 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 np.asarray(self.array).ndim == 4 else self.array, + space=(self.space[0][i], self.space[1][i]), + ) + return self._members[key] + + +class GElectronBoson(BaseElectronBoson): + """Generalised electron-boson coupling matrices. + + Attributes: + cc: Coupled cluster object. + space: Space object. + array: Electron-boson coupling matrix in the MO basis. + """ + + _members: dict[str, NDArray[float]] + + def __getitem__(self, key: str) -> NDArray[float]: + """Just-in-time getter. + + Args: + key: Key to get. + + Returns: + Electron-boson coupling matrix for the given spaces. + """ + if key not in self._members: + assert key[0] == "b" + i = self.space[0].mask(key[1]) + j = self.space[1].mask(key[2]) + self._members[key] = self.array[:, i][:, :, j].copy() + return self._members[key] From a682468674732640d05abf09ef08c451d4bbcf6b Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 23:05:43 +0100 Subject: [PATCH 27/37] Add mypy to linting ci --- .github/workflows/ci.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index b6f22ea0..010379e2 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -38,6 +38,7 @@ jobs: python -m black ebcc/ --diff --check --verbose python -m isort ebcc/ --diff --check-only --verbose python -m flake8 ebcc/ --verbose + python -m mypy ebcc/ --verbose - name: Run unit tests run: | python -m pip install pytest pytest-cov From 244c44c0eece46bbe2ec9f210c5c27bf47a88f75 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 23:14:52 +0100 Subject: [PATCH 28/37] Go away --- ebcc/gebcc.py | 482 ----------- ebcc/rebcc.py | 2236 ------------------------------------------------- 2 files changed, 2718 deletions(-) delete mode 100644 ebcc/gebcc.py delete mode 100644 ebcc/rebcc.py diff --git a/ebcc/gebcc.py b/ebcc/gebcc.py deleted file mode 100644 index 22feac62..00000000 --- a/ebcc/gebcc.py +++ /dev/null @@ -1,482 +0,0 @@ -"""General electron-boson coupled cluster.""" - -from pyscf import lib, scf - -from ebcc import geom -from ebcc import numpy as np -from ebcc import uebcc, util -from ebcc.brueckner import BruecknerGEBCC -from ebcc.eris import GERIs -from ebcc.fock import GFock -from ebcc.precision import types -from ebcc.rebcc import REBCC -from ebcc.space import Space - - -@util.has_docstring -class GEBCC(REBCC, metaclass=util.InheritDocstrings): - __doc__ = __doc__.replace("Restricted", "Generalised", 1) - - ERIs = GERIs - Fock = GFock - Brueckner = BruecknerGEBCC - - @staticmethod - def _convert_mf(mf): - if isinstance(mf, scf.ghf.GHF): - return mf - # NOTE: First convert to UHF - otherwise conversions from - # RHF->GHF and UHF->GHF may have inconsistent ordering - hf = mf.to_uhf().to_ghf() - if hasattr(mf, "xc"): - hf.e_tot = hf.energy_tot() - return hf - - @classmethod - def from_uebcc(cls, ucc): - """Initialise a `GEBCC` object from an `UEBCC` object. - - Parameters - ---------- - ucc : UEBCC - The UEBCC object to initialise from. - - Returns - ------- - gcc : GEBCC - The GEBCC object. - """ - - orbspin = scf.addons.get_ghf_orbspin(ucc.mf.mo_energy, ucc.mf.mo_occ, False) - nocc = ucc.space[0].nocc + ucc.space[1].nocc - nvir = ucc.space[0].nvir + ucc.space[1].nvir - nbos = ucc.nbos - sa = np.where(orbspin == 0)[0] - sb = np.where(orbspin == 1)[0] - - occupied = np.zeros((nocc + nvir,), dtype=bool) - occupied[sa] = ucc.space[0]._occupied.copy() - occupied[sb] = ucc.space[1]._occupied.copy() - frozen = np.zeros((nocc + nvir,), dtype=bool) - frozen[sa] = ucc.space[0]._frozen.copy() - frozen[sb] = ucc.space[1]._frozen.copy() - active = np.zeros((nocc + nvir,), dtype=bool) - active[sa] = ucc.space[0]._active.copy() - active[sb] = ucc.space[1]._active.copy() - space = Space(occupied, frozen, active) - - slices = util.Namespace( - a=util.Namespace(**{k: np.where(orbspin[space.mask(k)] == 0)[0] for k in "oOivVa"}), - b=util.Namespace(**{k: np.where(orbspin[space.mask(k)] == 1)[0] for k in "oOivVa"}), - ) - - if ucc.bare_g is not None: - if np.asarray(ucc.bare_g).ndim == 3: - bare_g_a = bare_g_b = ucc.bare_g - else: - bare_g_a, bare_g_b = ucc.bare_g - g = np.zeros((ucc.nbos, ucc.nmo * 2, ucc.nmo * 2)) - g[np.ix_(range(ucc.nbos), sa, sa)] = bare_g_a.copy() - g[np.ix_(range(ucc.nbos), sb, sb)] = bare_g_b.copy() - else: - g = None - - gcc = cls( - ucc.mf, - log=ucc.log, - ansatz=ucc.ansatz, - space=space, - omega=ucc.omega, - g=g, - G=ucc.bare_G, - options=ucc.options, - ) - - gcc.e_corr = ucc.e_corr - gcc.converged = ucc.converged - gcc.converged_lambda = ucc.converged_lambda - - has_amps = ucc.amplitudes is not None - has_lams = ucc.lambdas is not None - - if has_amps: - amplitudes = util.Namespace() - - for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): - shape = tuple(space.size(k) for k in key) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - for comb in util.generate_spin_combinations(n, unique=True): - done = set() - for lperm, lsign in util.permutations_with_signs(tuple(range(n))): - for uperm, usign in util.permutations_with_signs(tuple(range(n))): - combn = util.permute_string(comb[:n], lperm) - combn += util.permute_string(comb[n:], uperm) - if combn in done: - continue - mask = np.ix_(*[slices[s][k] for s, k in zip(combn, key)]) - transpose = tuple(lperm) + tuple(p + n for p in uperm) - amp = ( - getattr(ucc.amplitudes[name], comb).transpose(transpose) - * lsign - * usign - ) - for perm, sign in util.permutations_with_signs(tuple(range(n))): - transpose = tuple(perm) + tuple(range(n, 2 * n)) - if util.permute_string(comb[:n], perm) == comb[:n]: - amplitudes[name][mask] += amp.transpose(transpose).copy() * sign - done.add(combn) - - for name, key, n in ucc.ansatz.bosonic_cluster_ranks(spin_type=ucc.spin_type): - amplitudes[name] = ucc.amplitudes[name].copy() - - for name, key, nf, nb in ucc.ansatz.coupling_cluster_ranks(spin_type=ucc.spin_type): - shape = (nbos,) * nb + tuple(space.size(k) for k in key[nb:]) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - for comb in util.generate_spin_combinations(nf): - done = set() - for lperm, lsign in util.permutations_with_signs(tuple(range(nf))): - for uperm, usign in util.permutations_with_signs(tuple(range(nf))): - combn = util.permute_string(comb[:nf], lperm) - combn += util.permute_string(comb[nf:], uperm) - if combn in done: - continue - mask = np.ix_( - *([range(nbos)] * nb), - *[slices[s][k] for s, k in zip(combn, key[nb:])], - ) - transpose = ( - tuple(range(nb)) - + tuple(p + nb for p in lperm) - + tuple(p + nb + nf for p in uperm) - ) - amp = ( - getattr(ucc.amplitudes[name], comb).transpose(transpose) - * lsign - * usign - ) - for perm, sign in util.permutations_with_signs(tuple(range(nf))): - transpose = ( - tuple(range(nb)) - + tuple(p + nb for p in perm) - + tuple(range(nb + nf, nb + 2 * nf)) - ) - if util.permute_string(comb[:nf], perm) == comb[:nf]: - amplitudes[name][mask] += amp.transpose(transpose).copy() * sign - done.add(combn) - - gcc.amplitudes = amplitudes - - if has_lams: - lambdas = gcc.init_lams() # Easier this way - but have to build ERIs... - - for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): - lname = name.replace("t", "l") - shape = tuple(space.size(k) for k in key[n:] + key[:n]) - lambdas[lname] = np.zeros(shape, dtype=types[float]) - for comb in util.generate_spin_combinations(n, unique=True): - done = set() - for lperm, lsign in util.permutations_with_signs(tuple(range(n))): - for uperm, usign in util.permutations_with_signs(tuple(range(n))): - combn = util.permute_string(comb[:n], lperm) - combn += util.permute_string(comb[n:], uperm) - if combn in done: - continue - mask = np.ix_(*[slices[s][k] for s, k in zip(combn, key[n:] + key[:n])]) - transpose = tuple(lperm) + tuple(p + n for p in uperm) - amp = ( - getattr(ucc.lambdas[lname], comb).transpose(transpose) - * lsign - * usign - ) - for perm, sign in util.permutations_with_signs(tuple(range(n))): - transpose = tuple(perm) + tuple(range(n, 2 * n)) - if util.permute_string(comb[:n], perm) == comb[:n]: - lambdas[lname][mask] += amp.transpose(transpose).copy() * sign - done.add(combn) - - for name, key, n in ucc.ansatz.bosonic_cluster_ranks(spin_type=ucc.spin_type): - lname = "l" + name - lambdas[lname] = ucc.lambdas[lname].copy() - - for name, key, nf, nb in ucc.ansatz.coupling_cluster_ranks(spin_type=ucc.spin_type): - lname = "l" + name - shape = (nbos,) * nb + tuple( - space.size(k) for k in key[nb + nf :] + key[nb : nb + nf] - ) - lambdas[lname] = np.zeros(shape, dtype=types[float]) - for comb in util.generate_spin_combinations(nf, unique=True): - done = set() - for lperm, lsign in util.permutations_with_signs(tuple(range(nf))): - for uperm, usign in util.permutations_with_signs(tuple(range(nf))): - combn = util.permute_string(comb[:nf], lperm) - combn += util.permute_string(comb[nf:], uperm) - if combn in done: - continue - mask = np.ix_( - *([range(nbos)] * nb), - *[ - slices[s][k] - for s, k in zip(combn, key[nb + nf :] + key[nb : nb + nf]) - ], - ) - transpose = ( - tuple(range(nb)) - + tuple(p + nb for p in lperm) - + tuple(p + nb + nf for p in uperm) - ) - amp = ( - getattr(ucc.lambdas[lname], comb).transpose(transpose) - * lsign - * usign - ) - for perm, sign in util.permutations_with_signs(tuple(range(nf))): - transpose = ( - tuple(range(nb)) - + tuple(p + nb for p in perm) - + tuple(range(nb + nf, nb + 2 * nf)) - ) - if util.permute_string(comb[:nf], perm) == comb[:nf]: - lambdas[lname][mask] += amp.transpose(transpose).copy() * sign - done.add(combn) - - gcc.lambdas = lambdas - - return gcc - - @classmethod - def from_rebcc(cls, rcc): - """ - Initialise a `GEBCC` object from an `REBCC` object. - - Parameters - ---------- - rcc : REBCC - The REBCC object to initialise from. - - Returns - ------- - gcc : GEBCC - The GEBCC object. - """ - - ucc = uebcc.UEBCC.from_rebcc(rcc) - gcc = cls.from_uebcc(ucc) - - return gcc - - @util.has_docstring - def init_amps(self, eris=None): - eris = self.get_eris(eris) - amplitudes = util.Namespace() - - # Build T amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - if n == 1: - amplitudes[name] = getattr(self.fock, key) / self.energy_sum(key) - elif n == 2: - amplitudes[name] = getattr(eris, key) / self.energy_sum(key) - else: - shape = tuple(self.space.size(k) for k in key) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - if self.boson_ansatz: - # Only true for real-valued couplings: - h = self.g - H = self.G - - # Build S amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - if n == 1: - amplitudes[name] = -H / self.omega - else: - shape = (self.nbos,) * n - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - # Build U amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - if n == 1: - amplitudes[name] = h[key] / self.energy_sum(key) - else: - shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - return amplitudes - - @util.has_docstring - def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): - func, kwargs = self._load_function( - "make_rdm2_f", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm = func(**kwargs) - - if hermitise: - dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) - dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(1, 0, 3, 2)) - - return dm - - @util.has_docstring - def excitations_to_vector_ip(self, *excitations): - vectors = [] - m = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[:-1] - vectors.append(util.compress_axes(key, excitations[m]).ravel()) - m += 1 - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return np.concatenate(vectors) - - @util.has_docstring - def excitations_to_vector_ee(self, *excitations): - vectors = [] - m = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - vectors.append(util.compress_axes(key, excitations[m]).ravel()) - m += 1 - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return np.concatenate(vectors) - - @util.has_docstring - def vector_to_excitations_ip(self, vector): - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[:-1] - size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) - shape = tuple(self.space.size(k) for k in key) - vn_tril = vector[i0 : i0 + size] - vn = util.decompress_axes(key, vn_tril, shape=shape) - excitations.append(vn) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - @util.has_docstring - def vector_to_excitations_ea(self, vector): - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[n:] + key[: n - 1] - size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) - shape = tuple(self.space.size(k) for k in key) - vn_tril = vector[i0 : i0 + size] - vn = util.decompress_axes(key, vn_tril, shape=shape) - excitations.append(vn) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - @util.has_docstring - def vector_to_excitations_ee(self, vector): - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)}) - shape = tuple(self.space.size(k) for k in key) - vn_tril = vector[i0 : i0 + size] - vn = util.decompress_axes(key, vn_tril, shape=shape) - excitations.append(vn) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - @util.has_docstring - def get_mean_field_G(self): - val = lib.einsum("Ipp->I", self.g.boo) - val -= self.xi * self.omega - - if self.bare_G is not None: - val += self.bare_G - - return val - - def get_eris(self, eris=None): - """ - Get blocks of the ERIs. - - Parameters - ---------- - eris : np.ndarray or ERIs, optional. - Electronic repulsion integrals, either in the form of a - dense array or an ERIs object. Default value is `None`. - - Returns - ------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.ERIs()`. - """ - if (eris is None) or isinstance(eris, np.ndarray): - return self.ERIs(self, array=eris) - else: - return eris - - @util.has_docstring - def ip_eom(self, options=None, **kwargs): - return geom.IP_GEOM(self, options=options, **kwargs) - - @util.has_docstring - def ea_eom(self, options=None, **kwargs): - return geom.EA_GEOM(self, options=options, **kwargs) - - @util.has_docstring - def ee_eom(self, options=None, **kwargs): - return geom.EE_GEOM(self, options=options, **kwargs) - - @property - @util.has_docstring - def xi(self): - if self.options.shift: - xi = lib.einsum("Iii->I", self.g.boo) - xi /= self.omega - if self.bare_G is not None: - xi += self.bare_G / self.omega - else: - xi = np.zeros_like(self.omega) - return xi - - @property - @util.has_docstring - def spin_type(self): - return "G" diff --git a/ebcc/rebcc.py b/ebcc/rebcc.py deleted file mode 100644 index 424c35fb..00000000 --- a/ebcc/rebcc.py +++ /dev/null @@ -1,2236 +0,0 @@ -"""Restricted electron-boson coupled cluster.""" - -import dataclasses - -from pyscf import lib - -from ebcc import default_log, init_logging -from ebcc import numpy as np -from ebcc import reom, util -from ebcc.ansatz import Ansatz -from ebcc.brueckner import BruecknerREBCC -from ebcc.cderis import RCDERIs -from ebcc.damping import DIIS -from ebcc.dump import Dump -from ebcc.eris import RERIs -from ebcc.fock import RFock -from ebcc.logging import ANSI -from ebcc.precision import types -from ebcc.space import Space - - -class EBCC: - """Base class for EBCC.""" - - pass - - -@dataclasses.dataclass -class Options: - """ - Options for EBCC calculations. - - Attributes - ---------- - shift : bool, optional - If `True`, 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. Default value is `True`. - e_tol : float, optional - Threshold for convergence in the correlation energy. Default value - is `1e-8`. - t_tol : float, optional - Threshold for convergence in the amplitude norm. Default value is - `1e-8`. - max_iter : int, optional - Maximum number of iterations. Default value is `200`. - diis_space : int, optional - Number of amplitudes to use in DIIS extrapolation. Default value is - `12`. - damping : float, optional - Damping factor for DIIS extrapolation. Default value is `0.0`. - """ - - shift: bool = True - e_tol: float = 1e-8 - t_tol: float = 1e-8 - max_iter: int = 200 - diis_space: int = 12 - damping: float = 0.0 - - -class REBCC(EBCC): - r""" - Restricted electron-boson coupled cluster class. - - Parameters - ---------- - mf : pyscf.scf.hf.SCF - PySCF mean-field object. - log : logging.Logger, optional - Log to print output to. Default value is the global logger which - outputs to `sys.stderr`. - ansatz : str or Ansatz, optional - Overall ansatz, as a string representation or an `Ansatz` object. - If `None`, construct from the individual ansatz parameters, - otherwise override them. Default value is `None`. - space : Space, optional - Space object defining the size of frozen, correlated and active - fermionic spaces. If `None`, all fermionic degrees of freedom are - assumed correlated. Default value is `None`. - omega : numpy.ndarray (nbos,), optional - Bosonic frequencies. Default value is `None`. - g : numpy.ndarray (nbos, nmo, nmo), optional - Electron-boson coupling matrix corresponding to the bosonic - annihilation operator i.e. - - .. math:: g_{xpq} p^\dagger q b - - The creation part is assume to be the fermionic transpose of this - tensor to retain hermiticity in the overall Hamiltonian. Default - value is `None`. - G : numpy.ndarray (nbos,), optional - Boson non-conserving term of the Hamiltonian i.e. - - .. math:: G_x (b^\dagger + b) - - Default value is `None`. - mo_coeff : numpy.ndarray, optional - Molecular orbital coefficients. Default value is `mf.mo_coeff`. - mo_occ : numpy.ndarray, optional - Molecular orbital occupancies. Default value is `mf.mo_occ`. - fock : util.Namespace, optional - Fock input. Default value is calculated using `get_fock()`. - options : dataclasses.dataclass, optional - Object containing the options. Default value is `Options()`. - **kwargs : dict - Additional keyword arguments used to update `options`. - - Attributes - ---------- - mf : pyscf.scf.hf.SCF - PySCF mean-field object. - log : logging.Logger - Log to print output to. - options : dataclasses.dataclass - Object containing the options. - e_corr : float - Correlation energy. - amplitudes : Namespace - Cluster amplitudes. - lambdas : Namespace - Cluster lambda amplitudes. - converged : bool - Whether the coupled cluster equations converged. - converged_lambda : bool - Whether the lambda coupled cluster equations converged. - omega : numpy.ndarray (nbos,) - Bosonic frequencies. - g : util.Namespace - Namespace containing blocks of the electron-boson coupling matrix. - Each attribute should be a length-3 string of `b`, `o` or `v` - signifying whether the corresponding axis is bosonic, occupied, or - virtual. - G : numpy.ndarray (nbos,) - Mean-field boson non-conserving term of the Hamiltonian. - bare_G : numpy.ndarray (nbos,) - Boson non-conserving term of the Hamiltonian. - fock : util.Namespace - Namespace containing blocks of the Fock matrix. Each attribute - should be a length-2 string of `o` or `v` signifying whether the - corresponding axis is occupied or virtual. - bare_fock : numpy.ndarray (nmo, nmo) - The mean-field Fock matrix in the MO basis. - xi : numpy.ndarray (nbos,) - Shift in bosonic operators to diagonalise the phononic Hamiltonian. - const : float - Shift in the energy from moving to polaritonic basis. - name : str - Name of the method. - - Methods - ------- - init_amps(eris=None) - Initialise the amplitudes. - init_lams(amplitudes=None) - Initialise the lambda amplitudes. - kernel(eris=None) - Run the coupled cluster calculation. - solve_lambda(amplitudes=None, eris=None) - Solve the lambda coupled cluster equations. - energy(eris=None, amplitudes=None) - Compute the correlation energy. - update_amps(eris=None, amplitudes=None) - Update the amplitudes. - update_lams(eris=None, amplitudes=None, lambdas=None) - Update the lambda amplitudes. - make_sing_b_dm(eris=None, amplitudes=None, lambdas=None) - Build the single boson density matrix. - make_rdm1_b(eris=None, amplitudes=None, lambdas=None, unshifted=None, - hermitise=None) - Build the bosonic one-particle reduced density matrix. - make_rdm1_f(eris=None, amplitudes=None, lambdas=None, hermitise=None) - Build the fermionic one-particle reduced density matrix. - make_rdm2_f(eris=None, amplitudes=None, lambdas=None, hermitise=None) - Build the fermionic two-particle reduced density matrix. - make_eb_coup_rdm(eris=None, amplitudes=None, lambdas=None, - unshifted=True, hermitise=True) - Build the electron-boson coupling reduced density matrices. - hbar_matvec_ip(r1, r2, eris=None, amplitudes=None) - Compute the product between a state vector and the EOM Hamiltonian - for the IP. - hbar_matvec_ea(r1, r2, eris=None, amplitudes=None) - Compute the product between a state vector and the EOM Hamiltonian - for the EA. - hbar_matvec_ee(r1, r2, eris=None, amplitudes=None) - Compute the product between a state vector and the EOM Hamiltonian - for the EE. - make_ip_mom_bras(eris=None, amplitudes=None, lambdas=None) - Get the bra IP vectors to construct EOM moments. - make_ea_mom_bras(eris=None, amplitudes=None, lambdas=None) - Get the bra EA vectors to construct EOM moments. - make_ee_mom_bras(eris=None, amplitudes=None, lambdas=None) - Get the bra EE vectors to construct EOM moments. - make_ip_mom_kets(eris=None, amplitudes=None, lambdas=None) - Get the ket IP vectors to construct EOM moments. - make_ea_mom_kets(eris=None, amplitudes=None, lambdas=None) - Get the ket EA vectors to construct EOM moments. - make_ee_mom_kets(eris=None, amplitudes=None, lambdas=None) - Get the ket EE vectors to construct EOM moments. - amplitudes_to_vector(amplitudes) - Construct a vector containing all of the amplitudes used in the - given ansatz. - vector_to_amplitudes(vector) - Construct all of the amplitudes used in the given ansatz from a - vector. - lambdas_to_vector(lambdas) - Construct a vector containing all of the lambda amplitudes used in - the given ansatz. - vector_to_lambdas(vector) - Construct all of the lambdas used in the given ansatz from a - vector. - excitations_to_vector_ip(*excitations) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the IP. - excitations_to_vector_ea(*excitations) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EA. - excitations_to_vector_ee(*excitations) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EE. - vector_to_excitations_ip(vector) - Construct all of the excitation amplitudes used in the given - ansatz from a vector for the IP. - vector_to_excitations_ea(vector) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EA. - vector_to_excitations_ee(vector) - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EE. - get_mean_field_G() - Get the mean-field boson non-conserving term of the Hamiltonian. - get_g(g) - Get the blocks of the electron-boson coupling matrix corresponding - to the bosonic annihilation operator. - get_fock() - Get the blocks of the Fock matrix, shifted due to bosons where the - ansatz requires. - get_eris() - Get blocks of the ERIs. - """ - - Options = Options - ERIs = RERIs - Fock = RFock - CDERIs = RCDERIs - Brueckner = BruecknerREBCC - - def __init__( - self, - mf, - log=None, - ansatz="CCSD", - space=None, - omega=None, - g=None, - G=None, - mo_coeff=None, - mo_occ=None, - fock=None, - options=None, - **kwargs, - ): - # Options: - if options is None: - options = self.Options() - self.options = options - for key, val in kwargs.items(): - setattr(self.options, key, val) - - # Parameters: - self.log = default_log if log is None else log - self.mf = self._convert_mf(mf) - self._mo_coeff = np.asarray(mo_coeff).astype(types[float]) if mo_coeff is not None else None - self._mo_occ = np.asarray(mo_occ).astype(types[float]) if mo_occ is not None else None - - # Ansatz: - if isinstance(ansatz, Ansatz): - self.ansatz = ansatz - else: - self.ansatz = Ansatz.from_string( - ansatz, density_fitting=getattr(self.mf, "with_df", None) is not None - ) - self._eqns = self.ansatz._get_eqns(self.spin_type) - - # Space: - if space is not None: - self.space = space - else: - self.space = self.init_space() - - # Boson parameters: - if bool(self.fermion_coupling_rank) != bool(self.boson_coupling_rank): - raise ValueError( - "Fermionic and bosonic coupling ranks must both be zero, or both non-zero." - ) - self.omega = omega.astype(types[float]) if omega is not None else None - self.bare_g = g.astype(types[float]) if g is not None else None - self.bare_G = G.astype(types[float]) if G is not None else None - if self.boson_ansatz != "": - self.g = self.get_g(g) - self.G = self.get_mean_field_G() - if self.options.shift: - self.log.info(" > Energy shift due to polaritonic basis: %.10f", self.const) - else: - assert self.nbos == 0 - self.options.shift = False - self.g = None - self.G = None - - # Fock matrix: - if fock is None: - self.fock = self.get_fock() - else: - self.fock = fock - - # Attributes: - self.e_corr = None - self.amplitudes = None - self.converged = False - self.lambdas = None - self.converged_lambda = False - - # Logging: - init_logging(self.log) - self.log.info(f"\n{ANSI.B}{ANSI.U}{self.name}{ANSI.R}") - self.log.debug(f"{ANSI.B}{'*' * len(self.name)}{ANSI.R}") - self.log.debug("") - self.log.info(f"{ANSI.B}Options{ANSI.R}:") - self.log.info(f" > e_tol: {ANSI.y}{self.options.e_tol}{ANSI.R}") - self.log.info(f" > t_tol: {ANSI.y}{self.options.t_tol}{ANSI.R}") - self.log.info(f" > max_iter: {ANSI.y}{self.options.max_iter}{ANSI.R}") - self.log.info(f" > diis_space: {ANSI.y}{self.options.diis_space}{ANSI.R}") - self.log.info(f" > damping: {ANSI.y}{self.options.damping}{ANSI.R}") - self.log.debug("") - self.log.info(f"{ANSI.B}Ansatz{ANSI.R}: {ANSI.m}{self.ansatz}{ANSI.R}") - self.log.debug("") - self.log.info(f"{ANSI.B}Space{ANSI.R}: {ANSI.m}{self.space}{ANSI.R}") - self.log.debug("") - - def kernel(self, eris=None): - """ - Run the coupled cluster calculation. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - - Returns - ------- - e_cc : float - Correlation energy. - """ - - # Start a timer: - timer = util.Timer() - - # Get the ERIs: - eris = self.get_eris(eris) - - # Get the amplitude guesses: - if self.amplitudes is None: - amplitudes = self.init_amps(eris=eris) - else: - amplitudes = self.amplitudes - - # Get the initial energy: - e_cc = self.energy(amplitudes=amplitudes, eris=eris) - - self.log.output("Solving for excitation amplitudes.") - self.log.debug("") - self.log.info( - f"{ANSI.B}{'Iter':>4s} {'Energy (corr.)':>16s} {'Energy (tot.)':>18s} " - f"{'Δ(Energy)':>13s} {'Δ(Ampl.)':>13s}{ANSI.R}" - ) - self.log.info(f"{0:4d} {e_cc:16.10f} {e_cc + self.e_hf:18.10f}") - - if not self.ansatz.is_one_shot: - # Set up DIIS: - diis = DIIS() - diis.space = self.options.diis_space - diis.damping = self.options.damping - - converged = False - for niter in range(1, self.options.max_iter + 1): - # Update the amplitudes, extrapolate with DIIS and - # calculate change: - amplitudes_prev = amplitudes - amplitudes = self.update_amps(amplitudes=amplitudes, eris=eris) - vector = self.amplitudes_to_vector(amplitudes) - vector = diis.update(vector) - amplitudes = self.vector_to_amplitudes(vector) - dt = np.linalg.norm(vector - self.amplitudes_to_vector(amplitudes_prev), ord=np.inf) - - # Update the energy and calculate change: - e_prev = e_cc - e_cc = self.energy(amplitudes=amplitudes, eris=eris) - de = abs(e_prev - e_cc) - - # Log the iteration: - converged_e = de < self.options.e_tol - converged_t = dt < self.options.t_tol - self.log.info( - f"{niter:4d} {e_cc:16.10f} {e_cc + self.e_hf:18.10f}" - f" {[ANSI.r, ANSI.g][converged_e]}{de:13.3e}{ANSI.R}" - f" {[ANSI.r, ANSI.g][converged_t]}{dt:13.3e}{ANSI.R}" - ) - - # Check for convergence: - converged = converged_e and converged_t - if converged: - self.log.debug("") - self.log.output(f"{ANSI.g}Converged.{ANSI.R}") - break - else: - self.log.debug("") - self.log.warning(f"{ANSI.r}Failed to converge.{ANSI.R}") - - # Include perturbative correction if required: - if self.ansatz.has_perturbative_correction: - self.log.debug("") - self.log.info("Computing perturbative energy correction.") - e_pert = self.energy_perturbative(amplitudes=amplitudes, eris=eris) - e_cc += e_pert - self.log.info(f"E(pert) = {e_pert:.10f}") - - else: - converged = True - - # Update attributes: - self.e_corr = e_cc - self.amplitudes = amplitudes - self.converged = converged - - self.log.debug("") - self.log.output(f"E(corr) = {self.e_corr:.10f}") - self.log.output(f"E(tot) = {self.e_tot:.10f}") - self.log.debug("") - self.log.debug("Time elapsed: %s", timer.format_time(timer())) - self.log.debug("") - - return e_cc - - def solve_lambda(self, amplitudes=None, eris=None): - """ - Solve the lambda coupled cluster equations. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - """ - - # Start a timer: - timer = util.Timer() - - # Get the ERIs: - eris = self.get_eris(eris) - - # Get the amplitudes: - if amplitudes is None: - amplitudes = self.amplitudes - if amplitudes is None: - amplitudes = self.init_amps(eris=eris) - - # If needed, precompute the perturbative part of the lambda - # amplitudes: - if self.ansatz.has_perturbative_correction: - lambdas_pert = self.update_lams(eris=eris, amplitudes=amplitudes, perturbative=True) - else: - lambdas_pert = None - - # Get the lambda amplitude guesses: - if self.lambdas is None: - lambdas = self.init_lams(amplitudes=amplitudes) - else: - lambdas = self.lambdas - - # Set up DIIS: - diis = DIIS() - diis.space = self.options.diis_space - diis.damping = self.options.damping - - self.log.output("Solving for de-excitation (lambda) amplitudes.") - self.log.debug("") - self.log.info(f"{ANSI.B}{'Iter':>4s} {'Δ(Ampl.)':>13s}{ANSI.R}") - - converged = False - for niter in range(1, self.options.max_iter + 1): - # Update the lambda amplitudes, extrapolate with DIIS and - # calculate change: - lambdas_prev = lambdas - lambdas = self.update_lams( - amplitudes=amplitudes, - lambdas=lambdas, - lambdas_pert=lambdas_pert, - eris=eris, - ) - vector = self.lambdas_to_vector(lambdas) - vector = diis.update(vector) - lambdas = self.vector_to_lambdas(vector) - dl = np.linalg.norm(vector - self.lambdas_to_vector(lambdas_prev), ord=np.inf) - - # Log the iteration: - converged = dl < self.options.t_tol - self.log.info(f"{niter:4d} {[ANSI.r, ANSI.g][converged]}{dl:13.3e}{ANSI.R}") - - # Check for convergence: - if converged: - self.log.debug("") - self.log.output(f"{ANSI.g}Converged.{ANSI.R}") - break - else: - self.log.debug("") - self.log.warning(f"{ANSI.r}Failed to converge.{ANSI.R}") - - self.log.debug("") - self.log.debug("Time elapsed: %s", timer.format_time(timer())) - self.log.debug("") - self.log.debug("") - - # Update attributes: - self.lambdas = lambdas - self.converged_lambda = converged - - def brueckner(self, *args, **kwargs): - """ - Run a Brueckner orbital coupled cluster calculation. - - Returns - ------- - e_cc : float - Correlation energy. - """ - - bcc = self.Brueckner(self, *args, **kwargs) - - return bcc.kernel() - - def write(self, file): - """ - Write the EBCC data to a file. - - Parameters - ---------- - file : str - Path of file to write to. - """ - - writer = Dump(file) - writer.write(self) - - @classmethod - def read(cls, file, log=None): - """ - Read the data from a file. - - Parameters - ---------- - file : str - Path of file to read from. - log : Logger, optional - Logger to assign to the EBCC object. - - Returns - ------- - ebcc : EBCC - The EBCC object loaded from the file. - """ - - reader = Dump(file) - cc = reader.read(cls=cls, log=log) - - return cc - - @staticmethod - def _convert_mf(mf): - """ - Convert the input PySCF mean-field object to the one required for - the current class. - """ - hf = mf.to_rhf() - if hasattr(mf, "xc"): - hf.e_tot = hf.energy_tot() - return hf - - def _load_function(self, name, eris=False, amplitudes=False, lambdas=False, **kwargs): - """ - Load a function from the generated code, and return a dict of - arguments. - """ - - if not (eris is False): - eris = self.get_eris(eris) - else: - eris = None - - dicts = [] - - if not (amplitudes is False): - if amplitudes is None: - amplitudes = self.amplitudes - if amplitudes is None: - amplitudes = self.init_amps(eris=eris) - dicts.append(amplitudes) - - if not (lambdas is False): - if lambdas is None: - lambdas = self.lambdas - if lambdas is None: - self.log.warning("Using Λ = T* for %s", name) - lambdas = self.init_lams(amplitudes=amplitudes) - dicts.append(lambdas) - - if kwargs: - dicts.append(kwargs) - - func = getattr(self._eqns, name, None) - - if func is None: - raise util.ModelNotImplemented("%s for rank = %s" % (name, self.name)) - - kwargs = self._pack_codegen_kwargs(*dicts, eris=eris) - - return func, kwargs - - def _pack_codegen_kwargs(self, *extra_kwargs, eris=None): - """ - Pack all the possible keyword arguments for generated code - into a dictionary. - """ - # TODO change all APIs to take the space object instead of - # nocc, nvir, nbos, etc. - - eris = self.get_eris(eris) - - omega = np.diag(self.omega) if self.omega is not None else None - - kwargs = dict( - f=self.fock, - v=eris, - g=self.g, - G=self.G, - w=omega, - space=self.space, - nocc=self.space.ncocc, # FIXME rename? - nvir=self.space.ncvir, # FIXME rename? - nbos=self.nbos, - ) - if isinstance(eris, self.CDERIs): - kwargs["naux"] = self.mf.with_df.get_naoaux() - for kw in extra_kwargs: - if kw is not None: - kwargs.update(kw) - - return kwargs - - def init_space(self): - """ - Initialise the default `Space` object. - - Returns - ------- - space : Space - Space object in which all fermionic degrees of freedom are - considered inactive. - """ - - space = Space( - self.mo_occ > 0, - np.zeros_like(self.mo_occ, dtype=bool), - np.zeros_like(self.mo_occ, dtype=bool), - ) - - return space - - def init_amps(self, eris=None): - """ - Initialise the amplitudes. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - - Returns - ------- - amplitudes : Namespace - Cluster amplitudes. - """ - - eris = self.get_eris(eris) - amplitudes = util.Namespace() - - # Build T amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - if n == 1: - amplitudes[name] = self.fock[key] / self.energy_sum(key) - elif n == 2: - key_t = key[0] + key[2] + key[1] + key[3] - amplitudes[name] = eris[key_t].swapaxes(1, 2) / self.energy_sum(key) - else: - shape = tuple(self.space.size(k) for k in key) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - if self.boson_ansatz: - # Only true for real-valued couplings: - h = self.g - H = self.G - - # Build S amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - if n == 1: - amplitudes[name] = -H / self.omega - else: - shape = (self.nbos,) * n - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - # Build U amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - if nb == 1: - amplitudes[name] = h[key] / self.energy_sum(key) - else: - shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) - amplitudes[name] = np.zeros(shape, dtype=types[float]) - - return amplitudes - - def init_lams(self, amplitudes=None): - """ - Initialise the lambda amplitudes. - - Parameters - ---------- - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - lambdas : Namespace - Cluster lambda amplitudes. - """ - - if amplitudes is None: - amplitudes = self.amplitudes - lambdas = util.Namespace() - - # Build L amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - lname = name.replace("t", "l") - perm = list(range(n, 2 * n)) + list(range(n)) - lambdas[lname] = amplitudes[name].transpose(perm) - - # Build LS amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - lname = "l" + name - lambdas[lname] = amplitudes[name] - - # Build LU amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - lname = "l" + name - perm = list(range(nb)) + [nb + 1, nb] - lambdas[lname] = amplitudes[name].transpose(perm) - - return lambdas - - def energy(self, eris=None, amplitudes=None): - """ - Compute the correlation energy. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - e_cc : float - Correlation energy. - """ - - func, kwargs = self._load_function( - "energy", - eris=eris, - amplitudes=amplitudes, - ) - - return types[float](func(**kwargs).real) - - def energy_perturbative(self, eris=None, amplitudes=None, lambdas=None): - """ - Compute the perturbative part to the correlation energy. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - e_pert : float - Perturbative correction to the correlation energy. - """ - - func, kwargs = self._load_function( - "energy_perturbative", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return types[float](func(**kwargs).real) - - def update_amps(self, eris=None, amplitudes=None): - """ - Update the amplitudes. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - amplitudes : Namespace - Updated cluster amplitudes. - """ - - func, kwargs = self._load_function( - "update_amps", - eris=eris, - amplitudes=amplitudes, - ) - res = func(**kwargs) - res = {key.rstrip("new"): val for key, val in res.items()} - - # Divide T amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - res[name] /= self.energy_sum(key) - res[name] += amplitudes[name] - - # Divide S amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - res[name] /= self.energy_sum(key) - res[name] += amplitudes[name] - - # Divide U amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - res[name] /= self.energy_sum(key) - res[name] += amplitudes[name] - - return res - - def update_lams( - self, - eris=None, - amplitudes=None, - lambdas=None, - lambdas_pert=None, - perturbative=False, - ): - """ - Update the lambda amplitudes. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - perturbative : bool, optional - Whether to compute the perturbative part of the lambda - amplitudes. Default value is `False`. - - Returns - ------- - lambdas : Namespace - Updated cluster lambda amplitudes. - """ - # TODO active - - if lambdas_pert is not None: - lambdas.update(lambdas_pert) - - func, kwargs = self._load_function( - "update_lams%s" % ("_perturbative" if perturbative else ""), - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - res = func(**kwargs) - res = {key.rstrip("new"): val for key, val in res.items()} - - # Divide T amplitudes: - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - lname = name.replace("t", "l") - res[lname] /= self.energy_sum(key[n:] + key[:n]) - if not perturbative: - res[lname] += lambdas[lname] - - # Divide S amplitudes: - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - lname = "l" + name - res[lname] /= self.energy_sum(key[n:] + key[:n]) - if not perturbative: - res[lname] += lambdas[lname] - - # Divide U amplitudes: - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - if nf != 1: - raise util.ModelNotImplemented - lname = "l" + name - res[lname] /= self.energy_sum(key[:nb] + key[nb + nf :] + key[nb : nb + nf]) - if not perturbative: - res[lname] += lambdas[lname] - - if perturbative: - res = {key + "pert": val for key, val in res.items()} - - return res - - def make_sing_b_dm(self, eris=None, amplitudes=None, lambdas=None): - r""" - Build the single boson density matrix: - - ..math :: \langle b^+ \rangle - - and - - ..math :: \langle b \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - dm_b : numpy.ndarray (nbos,) - Single boson density matrix. - """ - - func, kwargs = self._load_function( - "make_sing_b_dm", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_rdm1_b(self, eris=None, amplitudes=None, lambdas=None, unshifted=True, hermitise=True): - r""" - Build the bosonic one-particle reduced density matrix: - - ..math :: \langle b^+ b \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - unshifted : bool, optional - If `self.shift` is `True`, then `unshifted=True` applies the - reverse transformation such that the bosonic operators are - defined with respect to the unshifted bosons. Default value is - `True`. Has no effect if `self.shift` is `False`. - hermitise : bool, optional - Force Hermiticity in the output. Default value is `True`. - - Returns - ------- - rdm1_b : numpy.ndarray (nbos, nbos) - Bosonic one-particle reduced density matrix. - """ - - func, kwargs = self._load_function( - "make_rdm1_b", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm = func(**kwargs) - - if hermitise: - dm = 0.5 * (dm + dm.T) - - if unshifted and self.options.shift: - dm_cre, dm_ann = self.make_sing_b_dm() - xi = self.xi - dm[np.diag_indices_from(dm)] -= xi * (dm_cre + dm_ann) - xi**2 - - return dm - - def make_rdm1_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): - r""" - Build the fermionic one-particle reduced density matrix: - - ..math :: \langle i^+ j \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - hermitise : bool, optional - Force Hermiticity in the output. Default value is `True`. - - Returns - ------- - rdm1_f : numpy.ndarray (nmo, nmo) - Fermionic one-particle reduced density matrix. - """ - - func, kwargs = self._load_function( - "make_rdm1_f", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm = func(**kwargs) - - if hermitise: - dm = 0.5 * (dm + dm.T) - - return dm - - def make_rdm2_f(self, eris=None, amplitudes=None, lambdas=None, hermitise=True): - r""" - Build the fermionic two-particle reduced density matrix: - - ..math :: \Gamma_{ijkl} = \langle i^+ j^+ l k \rangle - - which is stored in Chemist's notation. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - hermitise : bool, optional - Force Hermiticity in the output. Default value is `True`. - - Returns - ------- - rdm2_f : numpy.ndarray (nmo, nmo, nmo, nmo) - Fermionic two-particle reduced density matrix. - """ - - func, kwargs = self._load_function( - "make_rdm2_f", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm = func(**kwargs) - - if hermitise: - dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) - dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(1, 0, 3, 2)) - - return dm - - def make_eb_coup_rdm( - self, - eris=None, - amplitudes=None, - lambdas=None, - unshifted=True, - hermitise=True, - ): - r""" - Build the electron-boson coupling reduced density matrices: - - ..math :: \langle b^+ i^+ j \rangle - - and - - ..math :: \langle b i^+ j \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - unshifted : bool, optional - If `self.shift` is `True`, then `unshifted=True` applies the - reverse transformation such that the bosonic operators are - defined with respect to the unshifted bosons. Default value is - `True`. Has no effect if `self.shift` is `False`. - hermitise : bool, optional - Force Hermiticity in the output. Default value is `True`. - - Returns - ------- - dm_eb : numpy.ndarray (2, nbos, nmo, nmo) - Electron-boson coupling reduce density matrices. First - index corresponds to creation and second to annihilation - of the bosonic index. - """ - - func, kwargs = self._load_function( - "make_eb_coup_rdm", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - dm_eb = func(**kwargs) - - if hermitise: - dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) - dm_eb[1] = dm_eb[0].transpose(0, 2, 1).copy() - - if unshifted and self.options.shift: - rdm1_f = self.make_rdm1_f(hermitise=hermitise) - shift = util.einsum("x,ij->xij", self.xi, rdm1_f) - dm_eb -= shift[None] - - return dm_eb - - def hbar_matvec_ip(self, r1, r2, eris=None, amplitudes=None): - """ - Compute the product between a state vector and the EOM Hamiltonian - for the IP. - - Parameters - ---------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector. Keys are - strings of the name of each vector, and values are arrays - whose dimension depends on the particular sector. - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector resulting from - the matrix-vector product with the input vectors. Keys are - strings of the name of each vector, and values are arrays whose - dimension depends on the particular sector. - """ - # TODO generalise vectors input - - func, kwargs = self._load_function( - "hbar_matvec_ip", - eris=eris, - amplitudes=amplitudes, - r1=r1, - r2=r2, - ) - - return func(**kwargs) - - def hbar_matvec_ea(self, r1, r2, eris=None, amplitudes=None): - """ - Compute the product between a state vector and the EOM Hamiltonian - for the EA. - - Parameters - ---------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector. Keys are - strings of the name of each vector, and values are arrays - whose dimension depends on the particular sector. - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector resulting from - the matrix-vector product with the input vectors. Keys are - strings of the name of each vector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "hbar_matvec_ea", - eris=eris, - amplitudes=amplitudes, - r1=r1, - r2=r2, - ) - - return func(**kwargs) - - def hbar_matvec_ee(self, r1, r2, eris=None, amplitudes=None): - """ - Compute the product between a state vector and the EOM Hamiltonian - for the EE. - - Parameters - ---------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector. Keys are - strings of the name of each vector, and values are arrays - whose dimension depends on the particular sector. - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - vectors : dict of (str, numpy.ndarray) - Dictionary containing the vectors in each sector resulting from - the matrix-vector product with the input vectors. Keys are - strings of the name of each vector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "hbar_matvec_ee", - eris=eris, - amplitudes=amplitudes, - r1=r1, - r2=r2, - ) - - return func(**kwargs) - - def make_ip_mom_bras(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the bra IP vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - bras : dict of (str, numpy.ndarray) - Dictionary containing the bra vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ip_mom_bras", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ea_mom_bras(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the bra EA vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - bras : dict of (str, numpy.ndarray) - Dictionary containing the bra vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ea_mom_bras", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ee_mom_bras(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the bra EE vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - bras : dict of (str, numpy.ndarray) - Dictionary containing the bra vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ee_mom_bras", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ip_mom_kets(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the ket IP vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - kets : dict of (str, numpy.ndarray) - Dictionary containing the ket vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ip_mom_kets", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ea_mom_kets(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the ket IP vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - kets : dict of (str, numpy.ndarray) - Dictionary containing the ket vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ea_mom_kets", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ee_mom_kets(self, eris=None, amplitudes=None, lambdas=None): - """ - Get the ket EE vectors to construct EOM moments. - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - kets : dict of (str, numpy.ndarray) - Dictionary containing the ket vectors in each sector. Keys are - strings of the name of each sector, and values are arrays whose - dimension depends on the particular sector. - """ - - func, kwargs = self._load_function( - "make_ee_mom_kets", - eris=eris, - amplitudes=amplitudes, - lambdas=lambdas, - ) - - return func(**kwargs) - - def make_ip_1mom(self, eris=None, amplitudes=None, lambdas=None): - r""" - Build the first fermionic hole single-particle moment. - - .. math:: T_{pq} = \langle c_p^+ (H - E_0) c_q \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - mom : numpy.ndarray (nmo, nmo) - Array of the first moment. - """ - - raise util.ModelNotImplemented # TODO - - def make_ea_1mom(self, eris=None, amplitudes=None, lambdas=None): - r""" - Build the first fermionic particle single-particle moment. - - .. math:: T_{pq} = \langle c_p (H - E_0) c_q^+ \rangle - - Parameters - ---------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.get_eris()`. - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - mom : numpy.ndarray (nmo, nmo) - Array of the first moment. - """ - - raise util.ModelNotImplemented # TODO - - def ip_eom(self, options=None, **kwargs): - """Get the IP EOM object.""" - return reom.IP_REOM(self, options=options, **kwargs) - - def ea_eom(self, options=None, **kwargs): - """Get the EA EOM object.""" - return reom.EA_REOM(self, options=options, **kwargs) - - def ee_eom(self, options=None, **kwargs): - """Get the EE EOM object.""" - return reom.EE_REOM(self, options=options, **kwargs) - - def amplitudes_to_vector(self, amplitudes): - """ - Construct a vector containing all of the amplitudes used in the - given ansatz. - - Parameters - ---------- - amplitudes : Namespace, optional - Cluster amplitudes. Default value is generated using - `self.init_amps()`. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the amplitudes flattened and - concatenated. Size depends on the ansatz. - """ - - vectors = [] - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - vectors.append(amplitudes[name].ravel()) - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - vectors.append(amplitudes[name].ravel()) - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - vectors.append(amplitudes[name].ravel()) - - return np.concatenate(vectors) - - def vector_to_amplitudes(self, vector): - """ - Construct all of the amplitudes used in the given ansatz from a - vector. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the amplitudes flattened - and concatenated. Size depends on the ansatz. - - Returns - ------- - amplitudes : Namespace - Cluster amplitudes. - """ - - amplitudes = util.Namespace() - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - amplitudes[name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - shape = (self.nbos,) * n - size = np.prod(shape) - amplitudes[name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) - size = np.prod(shape) - amplitudes[name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - return amplitudes - - def lambdas_to_vector(self, lambdas): - """ - Construct a vector containing all of the lambda amplitudes used in - the given ansatz. - - Parameters - ---------- - lambdas : Namespace, optional - Cluster lambda amplitudes. Default value is generated using - `self.init_lams()`. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the lambdas flattened and - concatenated. Size depends on the ansatz. - """ - - vectors = [] - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - vectors.append(lambdas[name.replace("t", "l")].ravel()) - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - vectors.append(lambdas["l" + name].ravel()) - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - vectors.append(lambdas["l" + name].ravel()) - - return np.concatenate(vectors) - - def vector_to_lambdas(self, vector): - """ - Construct all of the lambdas used in the given ansatz from a - vector. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the lambdas flattened and - concatenated. Size depends on the ansatz. - - Returns - ------- - lambdas : Namespace - Cluster lambda amplitudes. - """ - - lambdas = util.Namespace() - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - lname = name.replace("t", "l") - key = key[n:] + key[:n] - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - lambdas[lname] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - shape = (self.nbos,) * n - size = np.prod(shape) - lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - key = key[:nb] + key[nb + nf :] + key[nb : nb + nf] - shape = (self.nbos,) * nb + tuple(self.space.size(k) for k in key[nb:]) - size = np.prod(shape) - lambdas["l" + name] = vector[i0 : i0 + size].reshape(shape) - i0 += size - - return lambdas - - def excitations_to_vector_ip(self, *excitations): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the IP. - - Parameters - ---------- - *excitations : iterable of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - """ - - vectors = [] - m = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - vectors.append(excitations[m].ravel()) - m += 1 - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return np.concatenate(vectors) - - def excitations_to_vector_ea(self, *excitations): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EA. - - Parameters - ---------- - *excitations : iterable of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - """ - return self.excitations_to_vector_ip(*excitations) - - def excitations_to_vector_ee(self, *excitations): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EE. - - Parameters - ---------- - *excitations : iterable of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - - Returns - ------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - """ - return self.excitations_to_vector_ip(*excitations) - - def vector_to_excitations_ip(self, vector): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the IP. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - - Returns - ------- - excitations : tuple of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - """ - - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[:-1] - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - excitations.append(vector[i0 : i0 + size].reshape(shape)) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def vector_to_excitations_ea(self, vector): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EA. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - - Returns - ------- - excitations : tuple of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - """ - - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - key = key[n:] + key[: n - 1] - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - excitations.append(vector[i0 : i0 + size].reshape(shape)) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def vector_to_excitations_ee(self, vector): - """ - Construct a vector containing all of the excitation amplitudes - used in the given ansatz for the EE. - - Parameters - ---------- - vector : numpy.ndarray - Single vector consisting of all the excitations flattened and - concatenated. Size depends on the ansatz. - - Returns - ------- - excitations : tuple of numpy.ndarray - Dictionary containing the excitations. Keys are strings of the - name of each excitations, and values are arrays whose dimension - depends on the particular excitation amplitude. - """ - - excitations = [] - i0 = 0 - - for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - shape = tuple(self.space.size(k) for k in key) - size = np.prod(shape) - excitations.append(vector[i0 : i0 + size].reshape(shape)) - i0 += size - - for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type): - raise util.ModelNotImplemented - - return tuple(excitations) - - def get_mean_field_G(self): - """ - Get the mean-field boson non-conserving term of the Hamiltonian. - - Returns - ------- - G_mf : numpy.ndarray (nbos,) - Mean-field boson non-conserving term of the Hamiltonian. - """ - - # FIXME should this also sum in frozen orbitals? - val = lib.einsum("Ipp->I", self.g.boo) * 2.0 - val -= self.xi * self.omega - - if self.bare_G is not None: - val += self.bare_G - - return val - - def get_g(self, g): - """ - Get blocks of the electron-boson coupling matrix corresponding to - the bosonic annihilation operator. - - Parameters - ---------- - g : numpy.ndarray (nbos, nmo, nmo) - Array of the electron-boson coupling matrix. - - Returns - ------- - g : Namespace - Namespace containing blocks of the electron-boson coupling - matrix. Each attribute should be a length-3 string of - `b`, `o` or `v` signifying whether the corresponding axis - is bosonic, occupied or virtual. - """ - - slices = { - "x": self.space.correlated, - "o": self.space.correlated_occupied, - "v": self.space.correlated_virtual, - "O": self.space.active_occupied, - "V": self.space.active_virtual, - "i": self.space.inactive_occupied, - "a": self.space.inactive_virtual, - } - - class Blocks(util.Namespace): - def __getitem__(selffer, key): - assert key[0] == "b" - i = slices[key[1]] - j = slices[key[2]] - return g[:, i][:, :, j].copy() - - __getattr__ = __getitem__ - - return Blocks() - - def get_fock(self): - """ - Get blocks of the Fock matrix, shifted due to bosons where the - ansatz requires. - - Returns - ------- - fock : Namespace - Namespace containing blocks of the Fock matrix. Each attribute - should be a length-2 string of `o` or `v` signifying whether - the corresponding axis is occupied or virtual. - """ - return self.Fock(self, array=self.bare_fock) - - def get_eris(self, eris=None): - """Get blocks of the ERIs. - - Parameters - ---------- - eris : np.ndarray or ERIs, optional. - Electronic repulsion integrals, either in the form of a dense - array or an `ERIs` object. Default value is `None`. - - Returns - ------- - eris : ERIs, optional - Electronic repulsion integrals. Default value is generated - using `self.ERIs()`. - """ - if (eris is None) or isinstance(eris, np.ndarray): - if (isinstance(eris, np.ndarray) and eris.ndim == 3) or getattr( - self.mf, "with_df", None - ): - return self.CDERIs(self, array=eris) - else: - return self.ERIs(self, array=eris) - else: - return eris - - @property - def bare_fock(self): - """ - Get the mean-field Fock matrix in the MO basis, including frozen - parts. - - Returns - ------- - bare_fock : numpy.ndarray (nmo, nmo) - The mean-field Fock matrix in the MO basis. - """ - - fock_ao = self.mf.get_fock().astype(types[float]) - mo_coeff = self.mo_coeff - - fock = util.einsum("pq,pi,qj->ij", fock_ao, mo_coeff, mo_coeff) - - return fock - - @property - def xi(self): - """ - Get the shift in bosonic operators to diagonalise the photonic - Hamiltonian. - - Returns - ------- - xi : numpy.ndarray (nbos,) - Shift in bosonic operators to diagonalise the phononic - Hamiltonian. - """ - - if self.options.shift: - xi = lib.einsum("Iii->I", self.g.boo) * 2.0 - xi /= self.omega - if self.bare_G is not None: - xi += self.bare_G / self.omega - else: - xi = np.zeros_like(self.omega, dtype=types[float]) - - return xi - - @property - def const(self): - """ - Get the shift in the energy from moving to polaritonic basis. - - Returns - ------- - const : float - Shift in the energy from moving to polaritonic basis. - """ - if self.options.shift: - return lib.einsum("I,I->", self.omega, self.xi**2) - else: - return 0.0 - - @property - def fermion_ansatz(self): - """Get a string representation of the fermion ansatz.""" - return self.ansatz.fermion_ansatz - - @property - def boson_ansatz(self): - """Get a string representation of the boson ansatz.""" - return self.ansatz.boson_ansatz - - @property - def fermion_coupling_rank(self): - """Get an integer representation of the fermion coupling rank.""" - return self.ansatz.fermion_coupling_rank - - @property - def boson_coupling_rank(self): - """Get an integer representation of the boson coupling rank.""" - return self.ansatz.boson_coupling_rank - - @property - def name(self): - """Get a string representation of the method name.""" - return self.spin_type + self.ansatz.name - - @property - def spin_type(self): - """Get a string represent of the spin channel.""" - return "R" - - @property - def mo_coeff(self): - """ - Get the molecular orbital coefficients. - - Returns - ------- - mo_coeff : numpy.ndarray - Molecular orbital coefficients. - """ - if self._mo_coeff is None: - return np.asarray(self.mf.mo_coeff).astype(types[float]) - return self._mo_coeff - - @property - def mo_occ(self): - """ - Get the molecular orbital occupancies. - - Returns - ------- - mo_occ : numpy.ndarray - Molecular orbital occupancies. - """ - if self._mo_occ is None: - return np.asarray(self.mf.mo_occ).astype(types[float]) - return self._mo_occ - - @property - def nmo(self): - """Get the number of MOs.""" - return self.space.nmo - - @property - def nocc(self): - """Get the number of occupied MOs.""" - return self.space.nocc - - @property - def nvir(self): - """Get the number of virtual MOs.""" - return self.space.nvir - - @property - def nbos(self): - """ - Get the number of bosonic degrees of freedom. - - Returns - ------- - nbos : int - Number of bosonic degrees of freedom. - """ - if self.omega is None: - return 0 - return self.omega.shape[0] - - def energy_sum(self, subscript, signs_dict=None): - """ - Get a direct sum of energies. - - Parameters - ---------- - subscript : str - The direct sum subscript, where each character indicates the - sector for each energy. For the default slice characters, see - `Space`. Occupied degrees of freedom are assumed to be - positive, virtual and bosonic negative (the signs can be - changed via the `signs_dict` keyword argument). - signs_dict : dict, optional - Dictionary defining custom signs for each sector. If `None`, - initialised such that `["o", "O", "i"]` are positive, and - `["v", "V", "a", "b"]` negative. Default value is `None`. - - Returns - ------- - energy_sum : numpy.ndarray - Array of energy sums. - """ - - n = 0 - - def next_char(): - nonlocal n - if n < 26: - char = chr(ord("a") + n) - else: - char = chr(ord("A") + n) - n += 1 - return char - - if signs_dict is None: - signs_dict = {} - for k, s in zip("vVaoOib", "---+++-"): - if k not in signs_dict: - signs_dict[k] = s - - energies = [] - for key in subscript: - if key == "b": - energies.append(self.omega) - else: - energies.append(np.diag(self.fock[key + key])) - - subscript = "".join([signs_dict[k] + next_char() for k in subscript]) - energy_sum = lib.direct_sum(subscript, *energies) - - return energy_sum - - @property - def e_hf(self): - """ - Return the mean-field energy. - - Returns - ------- - e_hf : float - Mean-field energy. - """ - return types[float](self.mf.e_tot) - - @property - def e_tot(self): - """ - Return the total energy (mean-field plus correlation). - - Returns - ------- - e_tot : float - Total energy. - """ - return types[float](self.e_hf) + self.e_corr - - @property - def t1(self): - """Get the T1 amplitudes.""" - return self.amplitudes["t1"] - - @property - def t2(self): - """Get the T2 amplitudes.""" - return self.amplitudes["t2"] - - @property - def t3(self): - """Get the T3 amplitudes.""" - return self.amplitudes["t3"] - - @property - def l1(self): - """Get the L1 amplitudes.""" - return self.lambdas["l1"] - - @property - def l2(self): - """Get the L2 amplitudes.""" - return self.lambdas["l2"] - - @property - def l3(self): - """Get the L3 amplitudes.""" - return self.lambdas["l3"] From 1728ed8c36cc3e38ec96d9bb6f77527284730346 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Mon, 5 Aug 2024 23:20:02 +0100 Subject: [PATCH 29/37] Fix MutableMapping for 3.8, it's not quite eol --- ebcc/util/misc.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index f7d8ec5e..22cac218 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -2,10 +2,15 @@ from __future__ import annotations +import sys import time -from collections.abc import MutableMapping from typing import TYPE_CHECKING, Generic, TypeVar +if sys.version_info >= (3, 9): + from collections.abc import MutableMapping +else: + from typing import MutableMapping + if TYPE_CHECKING: from typing import Any, ItemsView, Iterator, KeysView, ValuesView From a275bb151e84d6b7ba244f13ede99b0a1e1384e3 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Tue, 6 Aug 2024 07:45:12 +0100 Subject: [PATCH 30/37] Just drop 3.8 tbh --- .github/workflows/ci.yaml | 4 ++-- ebcc/util/misc.py | 6 +----- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 010379e2..ebbf5575 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -15,10 +15,10 @@ jobs: fail-fast: false matrix: include: - - {python-version: "3.8", os: ubuntu-latest, documentation: True} - - {python-version: "3.9", os: ubuntu-latest, documentation: False} + - {python-version: "3.9", os: ubuntu-latest, documentation: True} - {python-version: "3.10", os: ubuntu-latest, documentation: False} - {python-version: "3.11", os: ubuntu-latest, documentation: False} + - {python-version: "3.12", os: ubuntu-latest, documentation: False} steps: - uses: actions/checkout@v2 diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index 22cac218..aee0a5ee 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -5,11 +5,7 @@ import sys import time from typing import TYPE_CHECKING, Generic, TypeVar - -if sys.version_info >= (3, 9): - from collections.abc import MutableMapping -else: - from typing import MutableMapping +from collections.abc import MutableMapping if TYPE_CHECKING: from typing import Any, ItemsView, Iterator, KeysView, ValuesView From bf870cfdc3bb4df7effaa5193a730ca887bd7a5b Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Tue, 6 Aug 2024 07:48:30 +0100 Subject: [PATCH 31/37] linting --- ebcc/util/misc.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ebcc/util/misc.py b/ebcc/util/misc.py index aee0a5ee..f7d8ec5e 100644 --- a/ebcc/util/misc.py +++ b/ebcc/util/misc.py @@ -2,10 +2,9 @@ from __future__ import annotations -import sys import time -from typing import TYPE_CHECKING, Generic, TypeVar from collections.abc import MutableMapping +from typing import TYPE_CHECKING, Generic, TypeVar if TYPE_CHECKING: from typing import Any, ItemsView, Iterator, KeysView, ValuesView From c85c325e6190ec0f089398625e90bf24bdcce579 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Tue, 6 Aug 2024 07:52:06 +0100 Subject: [PATCH 32/37] Stop using outdated scipy --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 4d971a00..52fe80fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ classifiers = [ ] dependencies = [ "numpy>=1.19.0", - "scipy<=1.10.0", + "scipy>=1.11.0", "pyscf>=2.0.0", "h5py>=3.0.0", ] From a750b6ff74262e1240656b56d3cdfb4b39f52534 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Tue, 6 Aug 2024 07:59:23 +0100 Subject: [PATCH 33/37] Fix amplitudes checks --- ebcc/cc/gebcc.py | 2 +- ebcc/cc/rebcc.py | 2 +- ebcc/cc/uebcc.py | 2 +- ebcc/ham/space.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 0a7ec63c..6ddd2ecb 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -422,7 +422,7 @@ def init_lams( Returns: Initial cluster lambda amplitudes. """ - if amplitudes is None: + if not amplitudes: amplitudes = self.amplitudes lambdas: Namespace[AmplitudeType] = util.Namespace() diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index db5ed629..7f42da73 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -198,7 +198,7 @@ def init_lams( Returns: Initial cluster lambda amplitudes. """ - if amplitudes is None: + if not amplitudes: amplitudes = self.amplitudes lambdas: Namespace[AmplitudeType] = util.Namespace() diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index e6fe9cb1..74a2723f 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -305,7 +305,7 @@ def init_lams( Returns: Initial cluster lambda amplitudes. """ - if amplitudes is None: + if not amplitudes: amplitudes = self.amplitudes lambdas: Namespace[AmplitudeType] = util.Namespace() diff --git a/ebcc/ham/space.py b/ebcc/ham/space.py index 314b594c..d0d62f64 100644 --- a/ebcc/ham/space.py +++ b/ebcc/ham/space.py @@ -358,7 +358,7 @@ def construct_fno_space( """ # Get the MP2 1RDM solver = MP2(mf) - if amplitudes is None: + if not amplitudes: solver.kernel() dm1 = solver.make_rdm1() else: From 96e134af33c3b3a867e85b1b1f1b87c411739331 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Tue, 6 Aug 2024 08:06:43 +0100 Subject: [PATCH 34/37] AmplitudeType -> SpinArrayType --- ebcc/cc/base.py | 163 +++++++++++++++++++++-------------------- ebcc/cc/gebcc.py | 78 ++++++++++---------- ebcc/cc/rebcc.py | 64 ++++++++-------- ebcc/cc/uebcc.py | 88 +++++++++++----------- ebcc/eom/base.py | 28 +++---- ebcc/eom/geom.py | 26 +++---- ebcc/eom/reom.py | 32 ++++---- ebcc/eom/ueom.py | 64 ++++++++-------- ebcc/ham/space.py | 4 +- ebcc/opt/base.py | 16 ++-- ebcc/opt/gbrueckner.py | 16 ++-- ebcc/opt/rbrueckner.py | 16 ++-- ebcc/opt/ubrueckner.py | 16 ++-- 13 files changed, 308 insertions(+), 303 deletions(-) diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py index eeebd136..c3e8ea41 100644 --- a/ebcc/cc/base.py +++ b/ebcc/cc/base.py @@ -26,8 +26,13 @@ from ebcc.opt.base import BaseBruecknerEBCC from ebcc.util import Namespace + """Defines the type for the `eris` argument in functions.""" ERIsInputType = Any - AmplitudeType = Any + + """Defines the type for arrays, including spin labels.""" + SpinArrayType = Any + + """Defines the type for the spaces, including spin labels.""" SpaceType = Any @@ -79,8 +84,8 @@ class BaseEBCC(ABC): # Attributes space: SpaceType - amplitudes: Namespace[AmplitudeType] - lambdas: Namespace[AmplitudeType] + amplitudes: Namespace[SpinArrayType] + lambdas: Namespace[SpinArrayType] g: Optional[BaseElectronBoson] G: Optional[NDArray[float]] @@ -309,7 +314,7 @@ def kernel(self, eris: Optional[ERIsInputType] = None) -> float: def solve_lambda( self, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, eris: Optional[ERIsInputType] = None, ) -> None: """Solve for the lambda amplitudes. @@ -468,8 +473,8 @@ def _load_function( self, name: str, eris: Optional[Union[ERIsInputType, Literal[False]]] = False, - amplitudes: Optional[Union[Namespace[AmplitudeType], Literal[False]]] = False, - lambdas: Optional[Union[Namespace[AmplitudeType], Literal[False]]] = False, + amplitudes: Optional[Union[Namespace[SpinArrayType], Literal[False]]] = False, + lambdas: Optional[Union[Namespace[SpinArrayType], Literal[False]]] = False, **kwargs: Any, ) -> tuple[Callable[..., Any], dict[str, Any]]: """Load a function from the generated code, and return the arguments.""" @@ -512,7 +517,7 @@ def _pack_codegen_kwargs( pass @abstractmethod - def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[SpinArrayType]: """Initialise the cluster amplitudes. Args: @@ -525,8 +530,8 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude @abstractmethod def init_lams( - self, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Initialise the cluster lambda amplitudes. Args: @@ -538,8 +543,8 @@ def init_lams( pass def _get_amps( - self, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Get the cluster amplitudes, initialising if required.""" if not amplitudes: amplitudes = self.amplitudes @@ -549,9 +554,9 @@ def _get_amps( def _get_lams( self, - lambdas: Optional[Namespace[AmplitudeType]] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> Namespace[AmplitudeType]: + lambdas: Optional[Namespace[SpinArrayType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + ) -> Namespace[SpinArrayType]: """Get the cluster lambda amplitudes, initialising if required.""" if not lambdas: lambdas = self.lambdas @@ -562,7 +567,7 @@ def _get_lams( def energy( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, ) -> float: """Calculate the correlation energy. @@ -584,8 +589,8 @@ def energy( def energy_perturbative( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, ) -> float: """Calculate the perturbative correction to the correlation energy. @@ -610,8 +615,8 @@ def energy_perturbative( def update_amps( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> Namespace[AmplitudeType]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + ) -> Namespace[SpinArrayType]: """Update the cluster amplitudes. Args: @@ -627,11 +632,11 @@ def update_amps( def update_lams( self, eris: ERIsInputType = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - lambdas_pert: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + lambdas_pert: Optional[Namespace[SpinArrayType]] = None, perturbative: bool = False, - ) -> Namespace[AmplitudeType]: + ) -> Namespace[SpinArrayType]: """Update the cluster lambda amplitudes. Args: @@ -649,8 +654,8 @@ def update_lams( def make_sing_b_dm( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, ) -> NDArray[float]: r"""Make the single boson density matrix :math:`\langle b \rangle`. @@ -674,8 +679,8 @@ def make_sing_b_dm( def make_rdm1_b( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, unshifted: bool = True, hermitise: bool = True, ) -> NDArray[float]: @@ -714,8 +719,8 @@ def make_rdm1_b( def make_rdm1_f( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, ) -> Any: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. @@ -735,8 +740,8 @@ def make_rdm1_f( def make_rdm2_f( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, ) -> Any: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. @@ -756,8 +761,8 @@ def make_rdm2_f( def make_eb_coup_rdm( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, unshifted: bool = True, hermitise: bool = True, ) -> Any: @@ -786,10 +791,10 @@ def make_eb_coup_rdm( def hbar_matvec_ip( self, - *excitations: AmplitudeType, + *excitations: SpinArrayType, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, AmplitudeType]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, SpinArrayType]: """Compute the product between a state vector and the IP-EOM Hamiltonian. Args: @@ -810,15 +815,15 @@ def hbar_matvec_ip( r1=r1, r2=r2, ) - res: tuple[AmplitudeType, AmplitudeType] = func(**kwargs) + res: tuple[SpinArrayType, SpinArrayType] = func(**kwargs) return res def hbar_matvec_ea( self, - *excitations: AmplitudeType, + *excitations: SpinArrayType, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, AmplitudeType]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, SpinArrayType]: """Compute the product between a state vector and the EA-EOM Hamiltonian. Args: @@ -839,15 +844,15 @@ def hbar_matvec_ea( r1=r1, r2=r2, ) - res: tuple[AmplitudeType, AmplitudeType] = func(**kwargs) + res: tuple[SpinArrayType, SpinArrayType] = func(**kwargs) return res def hbar_matvec_ee( self, - *excitations: AmplitudeType, + *excitations: SpinArrayType, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, AmplitudeType]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, SpinArrayType]: """Compute the product between a state vector and the EE-EOM Hamiltonian. Args: @@ -868,15 +873,15 @@ def hbar_matvec_ee( r1=r1, r2=r2, ) - res: tuple[AmplitudeType, AmplitudeType] = func(**kwargs) + res: tuple[SpinArrayType, SpinArrayType] = func(**kwargs) return res def make_ip_mom_bras( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, ...]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, ...]: """Get the bra vectors to construct IP-EOM moments. Args: @@ -893,15 +898,15 @@ def make_ip_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - res: tuple[AmplitudeType, ...] = func(**kwargs) + res: tuple[SpinArrayType, ...] = func(**kwargs) return res def make_ea_mom_bras( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, ...]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, ...]: """Get the bra vectors to construct EA-EOM moments. Args: @@ -918,15 +923,15 @@ def make_ea_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - res: tuple[AmplitudeType, ...] = func(**kwargs) + res: tuple[SpinArrayType, ...] = func(**kwargs) return res def make_ee_mom_bras( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, ...]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, ...]: """Get the bra vectors to construct EE-EOM moments. Args: @@ -943,15 +948,15 @@ def make_ee_mom_bras( amplitudes=amplitudes, lambdas=lambdas, ) - res: tuple[AmplitudeType, ...] = func(**kwargs) + res: tuple[SpinArrayType, ...] = func(**kwargs) return res def make_ip_mom_kets( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, ...]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, ...]: """Get the ket vectors to construct IP-EOM moments. Args: @@ -968,15 +973,15 @@ def make_ip_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - res: tuple[AmplitudeType, ...] = func(**kwargs) + res: tuple[SpinArrayType, ...] = func(**kwargs) return res def make_ea_mom_kets( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, ...]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, ...]: """Get the ket vectors to construct EA-EOM moments. Args: @@ -993,15 +998,15 @@ def make_ea_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - res: tuple[AmplitudeType, ...] = func(**kwargs) + res: tuple[SpinArrayType, ...] = func(**kwargs) return res def make_ee_mom_kets( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - ) -> tuple[AmplitudeType, ...]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + ) -> tuple[SpinArrayType, ...]: """Get the ket vectors to construct EE-EOM moments. Args: @@ -1018,7 +1023,7 @@ def make_ee_mom_kets( amplitudes=amplitudes, lambdas=lambdas, ) - res: tuple[AmplitudeType, ...] = func(**kwargs) + res: tuple[SpinArrayType, ...] = func(**kwargs) return res @abstractmethod @@ -1036,7 +1041,7 @@ def energy_sum(self, *args: str, signs_dict: Optional[dict[str, str]] = None) -> pass @abstractmethod - def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[float]: + def amplitudes_to_vector(self, amplitudes: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -1048,7 +1053,7 @@ def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[ pass @abstractmethod - def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[SpinArrayType]: """Construct a namespace of amplitudes from a vector. Args: @@ -1060,7 +1065,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp pass @abstractmethod - def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float]: + def lambdas_to_vector(self, lambdas: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the lambda amplitudes used in the given ansatz. Args: @@ -1072,7 +1077,7 @@ def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float] pass @abstractmethod - def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[SpinArrayType]: """Construct a namespace of lambda amplitudes from a vector. Args: @@ -1084,7 +1089,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: pass @abstractmethod - def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the IP-EOM excitations. Args: @@ -1096,7 +1101,7 @@ def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> ND pass @abstractmethod - def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ea(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the EA-EOM excitations. Args: @@ -1108,7 +1113,7 @@ def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> ND pass @abstractmethod - def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the EE-EOM excitations. Args: @@ -1122,7 +1127,7 @@ def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> ND @abstractmethod def vector_to_excitations_ip( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of IP-EOM excitations from a vector. Args: @@ -1136,7 +1141,7 @@ def vector_to_excitations_ip( @abstractmethod def vector_to_excitations_ea( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of EA-EOM excitations from a vector. Args: @@ -1150,7 +1155,7 @@ def vector_to_excitations_ea( @abstractmethod def vector_to_excitations_ee( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of EE-EOM excitations from a vector. Args: diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py index 6ddd2ecb..3d96788b 100644 --- a/ebcc/cc/gebcc.py +++ b/ebcc/cc/gebcc.py @@ -29,7 +29,7 @@ from ebcc.util import Namespace ERIsInputType: TypeAlias = Union[GERIs, NDArray[float]] - AmplitudeType: TypeAlias = NDArray[float] + SpinArrayType: TypeAlias = NDArray[float] class GEBCC(BaseEBCC): @@ -168,7 +168,7 @@ def from_uebcc(cls, ucc: UEBCC) -> GEBCC: has_lams = bool(ucc.lambdas) if has_amps: - amplitudes: Namespace[AmplitudeType] = util.Namespace() + amplitudes: Namespace[SpinArrayType] = util.Namespace() for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): shape = tuple(space.size(k) for k in key) @@ -362,7 +362,7 @@ def _pack_codegen_kwargs( return kwargs - def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[SpinArrayType]: """Initialise the cluster amplitudes. Args: @@ -372,7 +372,7 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude Initial cluster amplitudes. """ eris = self.get_eris(eris) - amplitudes: Namespace[AmplitudeType] = util.Namespace() + amplitudes: Namespace[SpinArrayType] = util.Namespace() # Build T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -412,8 +412,8 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude return amplitudes def init_lams( - self, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Initialise the cluster lambda amplitudes. Args: @@ -424,7 +424,7 @@ def init_lams( """ if not amplitudes: amplitudes = self.amplitudes - lambdas: Namespace[AmplitudeType] = util.Namespace() + lambdas: Namespace[SpinArrayType] = util.Namespace() # Build L amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -450,8 +450,8 @@ def init_lams( def update_amps( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> Namespace[AmplitudeType]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + ) -> Namespace[SpinArrayType]: """Update the cluster amplitudes. Args: @@ -467,7 +467,7 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res: Namespace[AmplitudeType] = func(**kwargs) + res: Namespace[SpinArrayType] = func(**kwargs) res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: @@ -492,11 +492,11 @@ def update_amps( def update_lams( self, eris: ERIsInputType = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - lambdas_pert: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + lambdas_pert: Optional[Namespace[SpinArrayType]] = None, perturbative: bool = False, - ) -> Namespace[AmplitudeType]: + ) -> Namespace[SpinArrayType]: """Update the cluster lambda amplitudes. Args: @@ -521,7 +521,7 @@ def update_lams( amplitudes=amplitudes, lambdas=lambdas, ) - res: Namespace[AmplitudeType] = func(**kwargs) + res: Namespace[SpinArrayType] = func(**kwargs) res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: @@ -555,10 +555,10 @@ def update_lams( def make_rdm1_f( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -576,7 +576,7 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm: AmplitudeType = func(**kwargs) + dm: SpinArrayType = func(**kwargs) if hermitise: dm = 0.5 * (dm + dm.T) @@ -586,10 +586,10 @@ def make_rdm1_f( def make_rdm2_f( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -607,7 +607,7 @@ def make_rdm2_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm: AmplitudeType = func(**kwargs) + dm: SpinArrayType = func(**kwargs) if hermitise: dm = 0.5 * (dm.transpose(0, 1, 2, 3) + dm.transpose(2, 3, 0, 1)) @@ -618,11 +618,11 @@ def make_rdm2_f( def make_eb_coup_rdm( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, unshifted: bool = True, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: r"""Make the electron-boson coupling reduced density matrix. .. math:: @@ -650,7 +650,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - dm_eb: AmplitudeType = func(**kwargs) + dm_eb: SpinArrayType = func(**kwargs) if hermitise: dm_eb[0] = 0.5 * (dm_eb[0] + dm_eb[1].transpose(0, 2, 1)) @@ -704,7 +704,7 @@ def next_char() -> str: return energy_sum - def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[float]: + def amplitudes_to_vector(self, amplitudes: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -726,7 +726,7 @@ def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[ return np.concatenate(vectors) - def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[SpinArrayType]: """Construct a namespace of amplitudes from a vector. Args: @@ -735,7 +735,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp Returns: Cluster amplitudes. """ - amplitudes: Namespace[AmplitudeType] = util.Namespace() + amplitudes: Namespace[SpinArrayType] = util.Namespace() i0 = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -758,7 +758,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp return amplitudes - def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float]: + def lambdas_to_vector(self, lambdas: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the lambda amplitudes used in the given ansatz. Args: @@ -780,7 +780,7 @@ def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float] return np.concatenate(vectors) - def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[SpinArrayType]: """Construct a namespace of lambda amplitudes from a vector. Args: @@ -789,7 +789,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: Returns: Cluster lambda amplitudes. """ - lambdas: Namespace[AmplitudeType] = util.Namespace() + lambdas: Namespace[SpinArrayType] = util.Namespace() i0 = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -815,7 +815,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: return lambdas - def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the IP-EOM excitations. Args: @@ -840,7 +840,7 @@ def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> ND return np.concatenate(vectors) - def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ea(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the EA-EOM excitations. Args: @@ -851,7 +851,7 @@ def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> ND """ return self.excitations_to_vector_ip(*excitations) - def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the EE-EOM excitations. Args: @@ -877,7 +877,7 @@ def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> ND def vector_to_excitations_ip( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of IP-EOM excitations from a vector. Args: @@ -908,7 +908,7 @@ def vector_to_excitations_ip( def vector_to_excitations_ea( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of EA-EOM excitations from a vector. Args: @@ -939,7 +939,7 @@ def vector_to_excitations_ea( def vector_to_excitations_ee( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of EE-EOM excitations from a vector. Args: diff --git a/ebcc/cc/rebcc.py b/ebcc/cc/rebcc.py index 7f42da73..08623a1e 100644 --- a/ebcc/cc/rebcc.py +++ b/ebcc/cc/rebcc.py @@ -25,7 +25,7 @@ from ebcc.util import Namespace ERIsInputType: TypeAlias = Union[RERIs, RCDERIs, NDArray[float]] - AmplitudeType: TypeAlias = NDArray[float] + SpinArrayType: TypeAlias = NDArray[float] class REBCC(BaseEBCC): @@ -137,7 +137,7 @@ def _pack_codegen_kwargs( return kwargs - def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[SpinArrayType]: """Initialise the cluster amplitudes. Args: @@ -147,7 +147,7 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude Initial cluster amplitudes. """ eris = self.get_eris(eris) - amplitudes: Namespace[AmplitudeType] = util.Namespace() + amplitudes: Namespace[SpinArrayType] = util.Namespace() # Build T amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -188,8 +188,8 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude return amplitudes def init_lams( - self, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Initialise the cluster lambda amplitudes. Args: @@ -200,7 +200,7 @@ def init_lams( """ if not amplitudes: amplitudes = self.amplitudes - lambdas: Namespace[AmplitudeType] = util.Namespace() + lambdas: Namespace[SpinArrayType] = util.Namespace() # Build L amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -226,8 +226,8 @@ def init_lams( def update_amps( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> Namespace[AmplitudeType]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + ) -> Namespace[SpinArrayType]: """Update the cluster amplitudes. Args: @@ -243,7 +243,7 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res: Namespace[AmplitudeType] = func(**kwargs) + res: Namespace[SpinArrayType] = func(**kwargs) res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: @@ -268,11 +268,11 @@ def update_amps( def update_lams( self, eris: ERIsInputType = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - lambdas_pert: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + lambdas_pert: Optional[Namespace[SpinArrayType]] = None, perturbative: bool = False, - ) -> Namespace[AmplitudeType]: + ) -> Namespace[SpinArrayType]: """Update the cluster lambda amplitudes. Args: @@ -297,7 +297,7 @@ def update_lams( amplitudes=amplitudes, lambdas=lambdas, ) - res: Namespace[AmplitudeType] = func(**kwargs) + res: Namespace[SpinArrayType] = func(**kwargs) res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: @@ -331,8 +331,8 @@ def update_lams( def make_rdm1_f( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, ) -> NDArray[float]: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. @@ -362,8 +362,8 @@ def make_rdm1_f( def make_rdm2_f( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, ) -> NDArray[float]: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. @@ -394,8 +394,8 @@ def make_rdm2_f( def make_eb_coup_rdm( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, unshifted: bool = True, hermitise: bool = True, ) -> NDArray[float]: @@ -480,7 +480,7 @@ def next_char() -> str: return energy_sum - def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[float]: + def amplitudes_to_vector(self, amplitudes: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -502,7 +502,7 @@ def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[ return np.concatenate(vectors) - def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[SpinArrayType]: """Construct a namespace of amplitudes from a vector. Args: @@ -511,7 +511,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp Returns: Cluster amplitudes. """ - amplitudes: Namespace[AmplitudeType] = util.Namespace() + amplitudes: Namespace[SpinArrayType] = util.Namespace() i0 = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -534,7 +534,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp return amplitudes - def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float]: + def lambdas_to_vector(self, lambdas: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the lambda amplitudes used in the given ansatz. Args: @@ -556,7 +556,7 @@ def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float] return np.concatenate(vectors) - def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[SpinArrayType]: """Construct a namespace of lambda amplitudes from a vector. Args: @@ -565,7 +565,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: Returns: Cluster lambda amplitudes. """ - lambdas: Namespace[AmplitudeType] = util.Namespace() + lambdas: Namespace[SpinArrayType] = util.Namespace() i0 = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -591,7 +591,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: return lambdas - def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the IP-EOM excitations. Args: @@ -615,7 +615,7 @@ def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> ND return np.concatenate(vectors) - def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ea(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the EA-EOM excitations. Args: @@ -626,7 +626,7 @@ def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> ND """ return self.excitations_to_vector_ip(*excitations) - def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the EE-EOM excitations. Args: @@ -639,7 +639,7 @@ def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> ND def vector_to_excitations_ip( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of IP-EOM excitations from a vector. Args: @@ -668,7 +668,7 @@ def vector_to_excitations_ip( def vector_to_excitations_ea( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of EA-EOM excitations from a vector. Args: @@ -697,7 +697,7 @@ def vector_to_excitations_ea( def vector_to_excitations_ee( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of EE-EOM excitations from a vector. Args: diff --git a/ebcc/cc/uebcc.py b/ebcc/cc/uebcc.py index 74a2723f..66d6b1a1 100644 --- a/ebcc/cc/uebcc.py +++ b/ebcc/cc/uebcc.py @@ -27,7 +27,7 @@ from ebcc.util import Namespace ERIsInputType: TypeAlias = Union[UERIs, UCDERIs, tuple[NDArray[float], ...]] - AmplitudeType: TypeAlias = Union[NDArray[float], Namespace[NDArray[float]]] # S_{n} has no spin + SpinArrayType: TypeAlias = Union[NDArray[float], Namespace[NDArray[float]]] # S_{n} has no spin class UEBCC(BaseEBCC): @@ -131,7 +131,7 @@ def from_rebcc(cls, rcc: REBCC) -> UEBCC: has_lams = bool(rcc.lambdas) if has_amps: - amplitudes: Namespace[AmplitudeType] = util.Namespace() + amplitudes: Namespace[SpinArrayType] = util.Namespace() for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): amplitudes[name] = util.Namespace() @@ -153,7 +153,7 @@ def from_rebcc(cls, rcc: REBCC) -> UEBCC: ucc.amplitudes = amplitudes if has_lams: - lambdas: AmplitudeType = util.Namespace() + lambdas: SpinArrayType = util.Namespace() for name, key, n in ucc.ansatz.fermionic_cluster_ranks(spin_type=ucc.spin_type): lname = name.replace("t", "l") @@ -224,7 +224,7 @@ def _pack_codegen_kwargs( return kwargs - def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[SpinArrayType]: """Initialise the cluster amplitudes. Args: @@ -234,11 +234,11 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude Initial cluster amplitudes. """ eris = self.get_eris(eris) - amplitudes: Namespace[AmplitudeType] = util.Namespace() + amplitudes: Namespace[SpinArrayType] = util.Namespace() # Build T amplitudes for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - tn: AmplitudeType = util.Namespace() + tn: SpinArrayType = util.Namespace() for comb in util.generate_spin_combinations(n, unique=True): if n == 1: tn[comb] = self.fock[comb][key] / self.energy_sum(key, comb) @@ -295,8 +295,8 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[Amplitude return amplitudes def init_lams( - self, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Initialise the cluster lambda amplitudes. Args: @@ -307,7 +307,7 @@ def init_lams( """ if not amplitudes: amplitudes = self.amplitudes - lambdas: Namespace[AmplitudeType] = util.Namespace() + lambdas: Namespace[SpinArrayType] = util.Namespace() # Build L amplitudes: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): @@ -337,8 +337,8 @@ def init_lams( def update_amps( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - ) -> Namespace[AmplitudeType]: + amplitudes: Optional[Namespace[SpinArrayType]] = None, + ) -> Namespace[SpinArrayType]: """Update the cluster amplitudes. Args: @@ -354,7 +354,7 @@ def update_amps( eris=eris, amplitudes=amplitudes, ) - res: Namespace[AmplitudeType] = func(**kwargs) + res: Namespace[SpinArrayType] = func(**kwargs) res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: @@ -390,11 +390,11 @@ def update_amps( def update_lams( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, - lambdas_pert: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, + lambdas_pert: Optional[Namespace[SpinArrayType]] = None, perturbative: bool = False, - ) -> Namespace[AmplitudeType]: + ) -> Namespace[SpinArrayType]: """Update the cluster lambda amplitudes. Args: @@ -417,7 +417,7 @@ def update_lams( lambdas=lambdas, lambdas_pert=lambdas_pert, ) - res: Namespace[AmplitudeType] = func(**kwargs) + res: Namespace[SpinArrayType] = func(**kwargs) res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()}) # Divide T amplitudes: @@ -457,10 +457,10 @@ def update_lams( def make_rdm1_f( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: r"""Make the one-particle fermionic reduced density matrix :math:`\langle i^+ j \rangle`. Args: @@ -478,7 +478,7 @@ def make_rdm1_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm: AmplitudeType = func(**kwargs) + dm: SpinArrayType = func(**kwargs) if hermitise: dm.aa = 0.5 * (dm.aa + dm.aa.T) @@ -489,10 +489,10 @@ def make_rdm1_f( def make_rdm2_f( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: r"""Make the two-particle fermionic reduced density matrix :math:`\langle i^+j^+lk \rangle`. Args: @@ -510,7 +510,7 @@ def make_rdm2_f( amplitudes=amplitudes, lambdas=lambdas, ) - dm: AmplitudeType = func(**kwargs) + dm: SpinArrayType = func(**kwargs) if hermitise: @@ -531,11 +531,11 @@ def transpose2(dm: NDArray[float]) -> NDArray[float]: def make_eb_coup_rdm( self, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, - lambdas: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, + lambdas: Optional[Namespace[SpinArrayType]] = None, unshifted: bool = True, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: r"""Make the electron-boson coupling reduced density matrix. .. math:: @@ -563,7 +563,7 @@ def make_eb_coup_rdm( amplitudes=amplitudes, lambdas=lambdas, ) - dm_eb: AmplitudeType = func(**kwargs) + dm_eb: SpinArrayType = func(**kwargs) if hermitise: dm_eb.aa[0] = 0.5 * (dm_eb.aa[0] + dm_eb.aa[1].transpose(0, 2, 1)) @@ -621,7 +621,7 @@ def next_char() -> str: return energy_sum - def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[float]: + def amplitudes_to_vector(self, amplitudes: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -649,7 +649,7 @@ def amplitudes_to_vector(self, amplitudes: Namespace[AmplitudeType]) -> NDArray[ return np.concatenate(vectors) - def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[SpinArrayType]: """Construct a namespace of amplitudes from a vector. Args: @@ -658,7 +658,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp Returns: Cluster amplitudes. """ - amplitudes: Namespace[AmplitudeType] = util.Namespace() + amplitudes: Namespace[SpinArrayType] = util.Namespace() i0 = 0 sizes: dict[tuple[str, ...], int] = { (o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab") @@ -698,7 +698,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[AmplitudeTyp return amplitudes - def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float]: + def lambdas_to_vector(self, lambdas: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the lambda amplitudes used in the given ansatz. Args: @@ -728,7 +728,7 @@ def lambdas_to_vector(self, lambdas: Namespace[AmplitudeType]) -> NDArray[float] return np.concatenate(vectors) - def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: + def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[SpinArrayType]: """Construct a namespace of lambda amplitudes from a vector. Args: @@ -737,7 +737,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: Returns: Cluster lambda amplitudes. """ - lambdas: AmplitudeType = util.Namespace() + lambdas: SpinArrayType = util.Namespace() i0 = 0 sizes: dict[tuple[str, ...], int] = { (o, s): self.space[i].size(o) for o in "ovOVia" for i, s in enumerate("ab") @@ -781,7 +781,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[AmplitudeType]: return lambdas - def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the IP-EOM excitations. Args: @@ -808,7 +808,7 @@ def excitations_to_vector_ip(self, *excitations: Namespace[AmplitudeType]) -> ND return np.concatenate(vectors) - def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ea(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the EA-EOM excitations. Args: @@ -836,7 +836,7 @@ def excitations_to_vector_ea(self, *excitations: Namespace[AmplitudeType]) -> ND return np.concatenate(vectors) - def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> NDArray[float]: + def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]: """Construct a vector containing all of the EE-EOM excitations. Args: @@ -865,7 +865,7 @@ def excitations_to_vector_ee(self, *excitations: Namespace[AmplitudeType]) -> ND def vector_to_excitations_ip( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of IP-EOM excitations from a vector. Args: @@ -882,7 +882,7 @@ def vector_to_excitations_ip( for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): key = key[:-1] - amp: AmplitudeType = util.Namespace() + amp: SpinArrayType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True, unique=True): subscript, csizes = util.combine_subscripts(key, spin, sizes=sizes) size = util.get_compressed_size(subscript, **csizes) @@ -909,7 +909,7 @@ def vector_to_excitations_ip( def vector_to_excitations_ea( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of EA-EOM excitations from a vector. Args: @@ -926,7 +926,7 @@ def vector_to_excitations_ea( for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): key = key[n:] + key[: n - 1] - amp: AmplitudeType = util.Namespace() + amp: SpinArrayType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True, unique=True): subscript, csizes = util.combine_subscripts(key, spin, sizes=sizes) size = util.get_compressed_size(subscript, **csizes) @@ -953,7 +953,7 @@ def vector_to_excitations_ea( def vector_to_excitations_ee( self, vector: NDArray[float] - ) -> tuple[Namespace[AmplitudeType], ...]: + ) -> tuple[Namespace[SpinArrayType], ...]: """Construct a namespace of EE-EOM excitations from a vector. Args: @@ -969,7 +969,7 @@ def vector_to_excitations_ee( } for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp: AmplitudeType = util.Namespace() + amp: SpinArrayType = util.Namespace() for spin in util.generate_spin_combinations(n): subscript, csizes = util.combine_subscripts(key, spin, sizes=sizes) size = util.get_compressed_size(subscript, **csizes) diff --git a/ebcc/eom/base.py b/ebcc/eom/base.py index 4adcb1e6..263f2133 100644 --- a/ebcc/eom/base.py +++ b/ebcc/eom/base.py @@ -16,7 +16,7 @@ if TYPE_CHECKING: from typing import Any, Callable, Optional, Union - from ebcc.cc.base import AmplitudeType, BaseEBCC, ERIsInputType, SpaceType + from ebcc.cc.base import BaseEBCC, ERIsInputType, SpaceType, SpinArrayType from ebcc.core.ansatz import Ansatz from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -121,7 +121,7 @@ def name(self) -> str: return f"{self.excitation_type.upper()}-EOM-{self.spin_type}{self.ansatz.name}" @abstractmethod - def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: + def amplitudes_to_vector(self, *amplitudes: SpinArrayType) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -133,7 +133,7 @@ def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: pass @abstractmethod - def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudeType, ...]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[SpinArrayType, ...]: """Construct amplitudes from a vector. Args: @@ -172,7 +172,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: pass @abstractmethod - def bras(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + def bras(self, eris: Optional[ERIsInputType] = None) -> Namespace[SpinArrayType]: """Get the bra vectors. Args: @@ -184,7 +184,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType] pass @abstractmethod - def kets(self, eris: Optional[ERIsInputType] = None) -> Namespace[AmplitudeType]: + def kets(self, eris: Optional[ERIsInputType] = None) -> Namespace[SpinArrayType]: """Get the ket vectors. Args: @@ -246,7 +246,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: pass @abstractmethod - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" pass @@ -358,9 +358,9 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: """Construct the moments of the EOM Hamiltonian. Args: @@ -398,7 +398,7 @@ def excitation_type(self) -> str: """Get the type of excitation.""" return "ip" - def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: + def amplitudes_to_vector(self, *amplitudes: SpinArrayType) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -409,7 +409,7 @@ def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: """ return self.ebcc.excitations_to_vector_ip(*amplitudes) - def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudeType, ...]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[SpinArrayType, ...]: """Construct amplitudes from a vector. Args: @@ -445,7 +445,7 @@ def excitation_type(self) -> str: """Get the type of excitation.""" return "ea" - def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: + def amplitudes_to_vector(self, *amplitudes: SpinArrayType) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -456,7 +456,7 @@ def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: """ return self.ebcc.excitations_to_vector_ea(*amplitudes) - def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudeType, ...]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[SpinArrayType, ...]: """Construct amplitudes from a vector. Args: @@ -492,7 +492,7 @@ def excitation_type(self) -> str: """Get the type of excitation.""" return "ee" - def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: + def amplitudes_to_vector(self, *amplitudes: SpinArrayType) -> NDArray[float]: """Construct a vector containing all of the amplitudes used in the given ansatz. Args: @@ -503,7 +503,7 @@ def amplitudes_to_vector(self, *amplitudes: AmplitudeType) -> NDArray[float]: """ return self.ebcc.excitations_to_vector_ee(*amplitudes) - def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[AmplitudeType, ...]: + def vector_to_amplitudes(self, vector: NDArray[float]) -> tuple[SpinArrayType, ...]: """Construct amplitudes from a vector. Args: diff --git a/ebcc/eom/geom.py b/ebcc/eom/geom.py index 3cd56b90..9faede45 100644 --- a/ebcc/eom/geom.py +++ b/ebcc/eom/geom.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.gebcc import GEBCC, AmplitudeType, ERIsInputType + from ebcc.cc.gebcc import GEBCC, ERIsInputType, SpinArrayType from ebcc.ham.space import Space from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -39,7 +39,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(np.abs(diag)) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.ravel(), r1.ravel()) return astype(weight, float) @@ -64,7 +64,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -79,7 +79,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return bras - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -98,7 +98,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian. @@ -150,7 +150,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(np.abs(diag)) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.ravel(), r1.ravel()) return astype(weight, float) @@ -175,7 +175,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -190,7 +190,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return bras - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -209,7 +209,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian. @@ -261,7 +261,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(diag) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.ravel(), r1.ravel()) return astype(weight, float) @@ -285,7 +285,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -303,7 +303,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return bras - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -328,7 +328,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, diagonal_only: bool = True, ) -> NDArray[float]: diff --git a/ebcc/eom/reom.py b/ebcc/eom/reom.py index b22c8f32..8c179f05 100644 --- a/ebcc/eom/reom.py +++ b/ebcc/eom/reom.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.rebcc import REBCC, AmplitudeType, ERIsInputType + from ebcc.cc.rebcc import REBCC, ERIsInputType, SpinArrayType from ebcc.ham.space import Space from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -39,7 +39,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(np.abs(diag)) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.ravel(), r1.ravel()) return astype(weight, float) @@ -64,7 +64,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -79,7 +79,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return bras - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -98,9 +98,9 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: """Construct the moments of the EOM Hamiltonian. Args: @@ -150,7 +150,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(np.abs(diag)) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.ravel(), r1.ravel()) return astype(weight, float) @@ -175,7 +175,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -190,7 +190,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return bras - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -209,9 +209,9 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: """Construct the moments of the EOM Hamiltonian. Args: @@ -261,7 +261,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(diag) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.ravel(), r1.ravel()) return astype(weight, float) @@ -285,7 +285,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -303,7 +303,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: ) return bras - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -328,10 +328,10 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, diagonal_only: bool = True, - ) -> AmplitudeType: + ) -> SpinArrayType: """Construct the moments of the EOM Hamiltonian. Args: diff --git a/ebcc/eom/ueom.py b/ebcc/eom/ueom.py index 973411f0..550cd66a 100644 --- a/ebcc/eom/ueom.py +++ b/ebcc/eom/ueom.py @@ -12,7 +12,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.uebcc import UEBCC, AmplitudeType, ERIsInputType + from ebcc.cc.uebcc import UEBCC, ERIsInputType, SpinArrayType from ebcc.ham.space import Space from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -38,7 +38,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(np.abs(diag)) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.a.ravel(), r1.a.ravel()) + np.dot(r1.b.ravel(), r1.b.ravel()) return astype(weight, float) @@ -56,7 +56,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): key = key[:-1] - spin_part: AmplitudeType = util.Namespace() + spin_part: SpinArrayType = util.Namespace() for comb in util.generate_spin_combinations(n, excited=True): spin_part[comb] = self.ebcc.energy_sum(key, comb) parts.append(spin_part) @@ -66,7 +66,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -79,13 +79,13 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: bras_tmp: Namespace[list[NDArray[float]]] = util.Namespace(a=[], b=[]) for i in range(self.nmo): - amps_a: list[AmplitudeType] = [] - amps_b: list[AmplitudeType] = [] + amps_a: list[SpinArrayType] = [] + amps_b: list[SpinArrayType] = [] m = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_a: AmplitudeType = util.Namespace() - amp_b: AmplitudeType = util.Namespace() + amp_a: SpinArrayType = util.Namespace() + amp_b: SpinArrayType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True): shape = tuple(self.space["ab".index(s)].ncocc for s in spin[:n]) + tuple( self.space["ab".index(s)].ncvir for s in spin[n:] @@ -123,7 +123,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: return bras - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -137,13 +137,13 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: for i in range(self.nmo): j = (Ellipsis, i) - amps_a: list[AmplitudeType] = [] - amps_b: list[AmplitudeType] = [] + amps_a: list[SpinArrayType] = [] + amps_b: list[SpinArrayType] = [] m = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_a: AmplitudeType = util.Namespace() - amp_b: AmplitudeType = util.Namespace() + amp_a: SpinArrayType = util.Namespace() + amp_b: SpinArrayType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True): shape = tuple(self.space["ab".index(s)].ncocc for s in spin[:n]) + tuple( self.space["ab".index(s)].ncvir for s in spin[n:] @@ -185,7 +185,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian. @@ -241,7 +241,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(np.abs(diag)) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.a.ravel(), r1.a.ravel()) + np.dot(r1.b.ravel(), r1.b.ravel()) return astype(weight, float) @@ -259,7 +259,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): key = key[n:] + key[: n - 1] - spin_part: AmplitudeType = util.Namespace() + spin_part: SpinArrayType = util.Namespace() for comb in util.generate_spin_combinations(n, excited=True): spin_part[comb] = -self.ebcc.energy_sum(key, comb) parts.append(spin_part) @@ -269,7 +269,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -282,13 +282,13 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: bras_tmp: Namespace[list[NDArray[float]]] = util.Namespace(a=[], b=[]) for i in range(self.nmo): - amps_a: list[AmplitudeType] = [] - amps_b: list[AmplitudeType] = [] + amps_a: list[SpinArrayType] = [] + amps_b: list[SpinArrayType] = [] m = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_a: AmplitudeType = util.Namespace() - amp_b: AmplitudeType = util.Namespace() + amp_a: SpinArrayType = util.Namespace() + amp_b: SpinArrayType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True): shape = tuple(self.space["ab".index(s)].ncvir for s in spin[:n]) + tuple( self.space["ab".index(s)].ncocc for s in spin[n:] @@ -326,7 +326,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: return bras - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -340,13 +340,13 @@ def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: for i in range(self.nmo): j = (Ellipsis, i) - amps_a: list[AmplitudeType] = [] - amps_b: list[AmplitudeType] = [] + amps_a: list[SpinArrayType] = [] + amps_b: list[SpinArrayType] = [] m = 0 for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - amp_a: AmplitudeType = util.Namespace() - amp_b: AmplitudeType = util.Namespace() + amp_a: SpinArrayType = util.Namespace() + amp_b: SpinArrayType = util.Namespace() for spin in util.generate_spin_combinations(n, excited=True): shape = tuple(self.space["ab".index(s)].ncvir for s in spin[:n]) + tuple( self.space["ab".index(s)].ncocc for s in spin[n:] @@ -388,7 +388,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, ) -> NDArray[float]: """Construct the moments of the EOM Hamiltonian. @@ -444,7 +444,7 @@ def _argsort_guesses(self, diag: NDArray[float]) -> NDArray[int]: arg = np.argsort(diag) return arg - def _quasiparticle_weight(self, r1: AmplitudeType) -> float: + def _quasiparticle_weight(self, r1: SpinArrayType) -> float: """Get the quasiparticle weight.""" weight: float = np.dot(r1.aa.ravel(), r1.aa.ravel()) + np.dot(r1.bb.ravel(), r1.bb.ravel()) return astype(weight, float) @@ -461,7 +461,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: parts = [] for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type): - spin_part: AmplitudeType = util.Namespace() + spin_part: SpinArrayType = util.Namespace() for comb in util.generate_spin_combinations(n): spin_part[comb] = self.ebcc.energy_sum(key, comb) parts.append(spin_part) @@ -471,7 +471,7 @@ def diag(self, eris: Optional[ERIsInputType] = None) -> NDArray[float]: return self.amplitudes_to_vector(*parts) - def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def bras(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the bra vectors. Args: @@ -482,7 +482,7 @@ def bras(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: """ raise util.ModelNotImplemented("EE moments for UEBCC not working.") - def kets(self, eris: Optional[ERIsInputType] = None) -> AmplitudeType: + def kets(self, eris: Optional[ERIsInputType] = None) -> SpinArrayType: """Get the ket vectors. Args: @@ -497,7 +497,7 @@ def moments( self, nmom: int, eris: Optional[ERIsInputType] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, hermitise: bool = True, diagonal_only: bool = True, ) -> NDArray[float]: diff --git a/ebcc/ham/space.py b/ebcc/ham/space.py index d0d62f64..ce730dcc 100644 --- a/ebcc/ham/space.py +++ b/ebcc/ham/space.py @@ -15,7 +15,7 @@ from pyscf.scf.hf import SCF - from ebcc.cc.base import AmplitudeType + from ebcc.cc.base import SpinArrayType from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -340,7 +340,7 @@ def construct_fno_space( mf: SCF, occ_tol: Optional[float] = 1e-5, occ_frac: Optional[float] = None, - amplitudes: Optional[Namespace[AmplitudeType]] = None, + amplitudes: Optional[Namespace[SpinArrayType]] = None, ) -> Union[RConstructSpaceReturnType, UConstructSpaceReturnType]: """Construct a frozen natural orbital space. diff --git a/ebcc/opt/base.py b/ebcc/opt/base.py index 1a125fa7..5aad5157 100644 --- a/ebcc/opt/base.py +++ b/ebcc/opt/base.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: from typing import Any, Optional, TypeVar - from ebcc.cc.base import AmplitudeType, BaseEBCC + from ebcc.cc.base import BaseEBCC, SpinArrayType from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -206,10 +206,10 @@ def kernel(self) -> float: @abstractmethod def get_rotation_matrix( self, - u_tot: Optional[AmplitudeType] = None, + u_tot: Optional[SpinArrayType] = None, diis: Optional[DIIS] = None, - t1: Optional[AmplitudeType] = None, - ) -> tuple[AmplitudeType, AmplitudeType]: + t1: Optional[SpinArrayType] = None, + ) -> tuple[SpinArrayType, SpinArrayType]: """Update the rotation matrix. Also returns the total rotation matrix. @@ -226,8 +226,8 @@ def get_rotation_matrix( @abstractmethod def transform_amplitudes( - self, u: AmplitudeType, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, u: SpinArrayType, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Transform the amplitudes into the Brueckner orbital basis. Args: @@ -240,7 +240,7 @@ def transform_amplitudes( pass @abstractmethod - def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> float: + def get_t1_norm(self, amplitudes: Optional[Namespace[SpinArrayType]] = None) -> float: """Get the norm of the T1 amplitude. Args: @@ -278,7 +278,7 @@ def mo_update_correlated(self, mo_coeff: Any, mo_coeff_corr: Any) -> Any: @abstractmethod def update_coefficients( - self, u_tot: AmplitudeType, mo_coeff_new: Any, mo_coeff_ref: Any + self, u_tot: SpinArrayType, mo_coeff_new: Any, mo_coeff_ref: Any ) -> Any: """Update the MO coefficients. diff --git a/ebcc/opt/gbrueckner.py b/ebcc/opt/gbrueckner.py index 3278179a..92f74513 100644 --- a/ebcc/opt/gbrueckner.py +++ b/ebcc/opt/gbrueckner.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.gebcc import GEBCC, AmplitudeType + from ebcc.cc.gebcc import GEBCC, SpinArrayType from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -33,10 +33,10 @@ class BruecknerGEBCC(BaseBruecknerEBCC): def get_rotation_matrix( self, - u_tot: Optional[AmplitudeType] = None, + u_tot: Optional[SpinArrayType] = None, diis: Optional[DIIS] = None, - t1: Optional[AmplitudeType] = None, - ) -> tuple[AmplitudeType, AmplitudeType]: + t1: Optional[SpinArrayType] = None, + ) -> tuple[SpinArrayType, SpinArrayType]: """Update the rotation matrix. Also returns the total rotation matrix. @@ -76,8 +76,8 @@ def get_rotation_matrix( return u, u_tot def transform_amplitudes( - self, u: AmplitudeType, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, u: SpinArrayType, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Transform the amplitudes into the Brueckner orbital basis. Args: @@ -114,7 +114,7 @@ def transform_amplitudes( return self.cc.amplitudes - def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> float: + def get_t1_norm(self, amplitudes: Optional[Namespace[SpinArrayType]] = None) -> float: """Get the norm of the T1 amplitude. Args: @@ -155,7 +155,7 @@ def mo_update_correlated( return mo_coeff def update_coefficients( - self, u_tot: AmplitudeType, mo_coeff: NDArray[float], mo_coeff_ref: NDArray[float] + self, u_tot: SpinArrayType, mo_coeff: NDArray[float], mo_coeff_ref: NDArray[float] ) -> NDArray[float]: """Update the MO coefficients. diff --git a/ebcc/opt/rbrueckner.py b/ebcc/opt/rbrueckner.py index c9a6f6b8..4f485555 100644 --- a/ebcc/opt/rbrueckner.py +++ b/ebcc/opt/rbrueckner.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.rebcc import REBCC, AmplitudeType + from ebcc.cc.rebcc import REBCC, SpinArrayType from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -33,10 +33,10 @@ class BruecknerREBCC(BaseBruecknerEBCC): def get_rotation_matrix( self, - u_tot: Optional[AmplitudeType] = None, + u_tot: Optional[SpinArrayType] = None, diis: Optional[DIIS] = None, - t1: Optional[AmplitudeType] = None, - ) -> tuple[AmplitudeType, AmplitudeType]: + t1: Optional[SpinArrayType] = None, + ) -> tuple[SpinArrayType, SpinArrayType]: """Update the rotation matrix. Also returns the total rotation matrix. @@ -76,8 +76,8 @@ def get_rotation_matrix( return u, u_tot def transform_amplitudes( - self, u: AmplitudeType, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, u: SpinArrayType, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Transform the amplitudes into the Brueckner orbital basis. Args: @@ -114,7 +114,7 @@ def transform_amplitudes( return self.cc.amplitudes - def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> float: + def get_t1_norm(self, amplitudes: Optional[Namespace[SpinArrayType]] = None) -> float: """Get the norm of the T1 amplitude. Args: @@ -155,7 +155,7 @@ def mo_update_correlated( return mo_coeff def update_coefficients( - self, u_tot: AmplitudeType, mo_coeff: NDArray[float], mo_coeff_ref: NDArray[float] + self, u_tot: SpinArrayType, mo_coeff: NDArray[float], mo_coeff_ref: NDArray[float] ) -> NDArray[float]: """Update the MO coefficients. diff --git a/ebcc/opt/ubrueckner.py b/ebcc/opt/ubrueckner.py index bc098b0f..f2b0d8c8 100644 --- a/ebcc/opt/ubrueckner.py +++ b/ebcc/opt/ubrueckner.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from typing import Optional - from ebcc.cc.uebcc import UEBCC, AmplitudeType + from ebcc.cc.uebcc import UEBCC, SpinArrayType from ebcc.core.damping import DIIS from ebcc.numpy.typing import NDArray from ebcc.util import Namespace @@ -33,10 +33,10 @@ class BruecknerUEBCC(BaseBruecknerEBCC): def get_rotation_matrix( self, - u_tot: Optional[AmplitudeType] = None, + u_tot: Optional[SpinArrayType] = None, diis: Optional[DIIS] = None, - t1: Optional[AmplitudeType] = None, - ) -> tuple[AmplitudeType, AmplitudeType]: + t1: Optional[SpinArrayType] = None, + ) -> tuple[SpinArrayType, SpinArrayType]: """Update the rotation matrix. Also returns the total rotation matrix. @@ -96,8 +96,8 @@ def get_rotation_matrix( return u, u_tot def transform_amplitudes( - self, u: AmplitudeType, amplitudes: Optional[Namespace[AmplitudeType]] = None - ) -> Namespace[AmplitudeType]: + self, u: SpinArrayType, amplitudes: Optional[Namespace[SpinArrayType]] = None + ) -> Namespace[SpinArrayType]: """Transform the amplitudes into the Brueckner orbital basis. Args: @@ -135,7 +135,7 @@ def transform_amplitudes( return self.cc.amplitudes - def get_t1_norm(self, amplitudes: Optional[Namespace[AmplitudeType]] = None) -> float: + def get_t1_norm(self, amplitudes: Optional[Namespace[SpinArrayType]] = None) -> float: """Get the norm of the T1 amplitude. Args: @@ -187,7 +187,7 @@ def mo_update_correlated( def update_coefficients( self, - u_tot: AmplitudeType, + u_tot: SpinArrayType, mo_coeff: tuple[NDArray[float], NDArray[float]], mo_coeff_ref: tuple[NDArray[float], NDArray[float]], ) -> tuple[NDArray[float], NDArray[float]]: From cb3afe1ae8641beec988b4de008bfdff12d7213f Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Tue, 6 Aug 2024 08:08:56 +0100 Subject: [PATCH 35/37] Update python version requirement --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 52fe80fb..a46ef0b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ keywords = [ "ccsd", ] readme = "README.md" -requires-python = ">=3.7" +requires-python = ">=3.9" classifiers = [ "Development Status :: 2 - Pre-Alpha", "Intended Audience :: Science/Research", From 8719310f5ff6d179727ddb2e44b7048e1b9f031d Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Tue, 6 Aug 2024 08:13:21 +0100 Subject: [PATCH 36/37] Remove some unnecessary type: ignore --- ebcc/__init__.py | 6 +++--- ebcc/util/einsumfunc.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/ebcc/__init__.py b/ebcc/__init__.py index 3fe846c5..a8576cb6 100644 --- a/ebcc/__init__.py +++ b/ebcc/__init__.py @@ -58,7 +58,7 @@ if TYPE_CHECKING: from typing import Any, Callable - from pyscf.scf.hf import SCF # type: ignore + from pyscf.scf.hf import SCF from ebcc.cc.base import BaseEBCC @@ -68,7 +68,7 @@ def EBCC(mf: SCF, *args: Any, **kwargs: Any) -> BaseEBCC: """Construct an EBCC object for the given mean-field object.""" - from pyscf import scf # type: ignore + from pyscf import scf if isinstance(mf, scf.uhf.UHF): return UEBCC(mf, *args, **kwargs) @@ -83,7 +83,7 @@ def EBCC(mf: SCF, *args: Any, **kwargs: Any) -> BaseEBCC: def _factory(ansatz: str) -> Callable[[SCF, Any, Any], BaseEBCC]: """Constructor for some specific ansatz.""" - from pyscf import scf # type: ignore + from pyscf import scf def constructor(mf: SCF, *args: Any, **kwargs: Any) -> BaseEBCC: """Construct an EBCC object for the given mean-field object.""" diff --git a/ebcc/util/einsumfunc.py b/ebcc/util/einsumfunc.py index 4b5ed4ff..902cb863 100644 --- a/ebcc/util/einsumfunc.py +++ b/ebcc/util/einsumfunc.py @@ -5,15 +5,15 @@ import ctypes from typing import TYPE_CHECKING -from pyscf.lib import direct_sum, dot # type: ignore # noqa: F401 -from pyscf.lib import einsum as pyscf_einsum # type: ignore # noqa: F401 +from pyscf.lib import direct_sum, dot # noqa: F401 +from pyscf.lib import einsum as pyscf_einsum # noqa: F401 from ebcc import numpy as np if TYPE_CHECKING: from typing import Any, TypeVar, Union - from ebcc.numpy.typing import NDArray # type: ignore + from ebcc.numpy.typing import NDArray T = TypeVar("T") @@ -298,7 +298,7 @@ def einsum(*operands: Any, **kwargs: Any) -> Union[T, NDArray[T]]: optimize = kwargs.pop("optimize", True) args = list(args) path_kwargs = dict(optimize=optimize, einsum_call=True) - contractions = np.einsum_path(subscript, *args, **path_kwargs)[1] # type: ignore + contractions = np.einsum_path(subscript, *args, **path_kwargs)[1] for contraction in contractions: inds, idx_rm, einsum_str, remain = list(contraction[:4]) contraction_args = [args.pop(x) for x in inds] From 91c716c7ed67a54454a060aec468d44f2e50bed5 Mon Sep 17 00:00:00 2001 From: Oliver Backhouse Date: Tue, 6 Aug 2024 08:41:28 +0100 Subject: [PATCH 37/37] Fix outcore mean-field --- ebcc/ham/cderis.py | 7 ++++++- ebcc/ham/eris.py | 27 ++++++++++----------------- 2 files changed, 16 insertions(+), 18 deletions(-) diff --git a/ebcc/ham/cderis.py b/ebcc/ham/cderis.py index 8dfb0381..c18dae78 100644 --- a/ebcc/ham/cderis.py +++ b/ebcc/ham/cderis.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING -from pyscf import ao2mo +from pyscf import ao2mo, lib from ebcc import numpy as np from ebcc import util @@ -50,6 +50,11 @@ def __getitem__(self, key: str, e2: Optional[bool] = False) -> NDArray[float]: raise KeyError("Key must be of length 3 or 4.") 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 key_e2 not in self._members: s = 0 if not e2 else 2 coeffs = [ diff --git a/ebcc/ham/eris.py b/ebcc/ham/eris.py index bd7571ad..3f6c45a4 100644 --- a/ebcc/ham/eris.py +++ b/ebcc/ham/eris.py @@ -43,11 +43,7 @@ def __getitem__(self, key: str) -> NDArray[float]: self.mo_coeff[i][:, self.space[i].mask(k)].astype(np.float64) for i, k in enumerate(key) ] - block = ao2mo.incore.general( - self.cc.mf._eri, - coeffs, - compact=False, - ) + block = ao2mo.kernel(self.cc.mf.mol, coeffs, compact=False, max_memory=1e6) block = block.reshape([c.shape[-1] for c in coeffs]) self._members[key] = block.astype(types[float]) return self._members[key] @@ -92,14 +88,11 @@ def __getitem__(self, key: str) -> RERIs: array = array.transpose(2, 3, 0, 1) elif isinstance(self.cc.mf._eri, tuple): # Support spin-dependent integrals in the mean-field - array = ao2mo.incore.general( - self.cc.mf._eri[ij], - [ - self.mo_coeff[x][y].astype(np.float64) - for y, x in enumerate(sorted((i, i, j, j))) - ], - compact=False, - ) + coeffs = [ + self.mo_coeff[x][y].astype(np.float64) + for y, x in enumerate(sorted((i, i, j, j))) + ] + array = ao2mo.kernel(self.cc.mf.mol, coeffs, compact=False, max_memory=1e6) if key == "bbaa": array = array.transpose(2, 3, 0, 1) array = array.astype(types[float]) @@ -136,10 +129,10 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: if self.array is None: mo_a = [mo[: self.cc.mf.mol.nao].astype(np.float64) for mo in self.mo_coeff] mo_b = [mo[self.cc.mf.mol.nao :].astype(np.float64) for mo in self.mo_coeff] - array = ao2mo.kernel(self.cc.mf._eri, mo_a) - array += ao2mo.kernel(self.cc.mf._eri, mo_b) - array += ao2mo.kernel(self.cc.mf._eri, mo_a[:2] + mo_b[2:]) - array += ao2mo.kernel(self.cc.mf._eri, mo_b[:2] + mo_a[2:]) + array = ao2mo.kernel(self.cc.mf.mol, mo_a) + array += ao2mo.kernel(self.cc.mf.mol, mo_b) + array += ao2mo.kernel(self.cc.mf.mol, mo_a[:2] + mo_b[2:]) + array += ao2mo.kernel(self.cc.mf.mol, mo_b[:2] + mo_a[2:]) array = ao2mo.addons.restore(1, array, self.cc.nmo).reshape((self.cc.nmo,) * 4) array = array.astype(types[float]) array = array.transpose(0, 2, 1, 3) - array.transpose(0, 2, 3, 1)