Skip to content

Commit

Permalink
Merge pull request #79 from BoothGroup/numpy_typing
Browse files Browse the repository at this point in the history 
Fixes to numpy static typing
  • Loading branch information
@obackhouse
obackhouse committed Sep 14, 2024
2 parents 9fe3fc0 + 2d2a692 commit 27b71e4
Show file tree
Hide file tree
Showing 68 changed files with 1,236 additions and 844 deletions.
74 changes: 41 additions & 33 deletions ebcc/cc/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,17 @@
if TYPE_CHECKING:
from typing import Any, Callable, Literal, Optional, Union

from numpy import float64
from numpy.typing import NDArray
from pyscf.scf.hf import SCF

from ebcc.core.logging import Logger
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

T = float64

"""Defines the type for the `eris` argument in functions."""
ERIsInputType = Any

Expand Down Expand Up @@ -74,8 +77,12 @@ class BaseEBCC(ABC):
space: SpaceType
amplitudes: Namespace[SpinArrayType]
lambdas: Namespace[SpinArrayType]
fock: BaseFock
g: Optional[BaseElectronBoson]
G: Optional[NDArray[float]]
G: Optional[NDArray[T]]
omega: Optional[NDArray[T]]
bare_g: Optional[NDArray[T]]
bare_G: Optional[NDArray[T]]

def __init__(
self,
Expand All @@ -84,11 +91,11 @@ def __init__(
ansatz: Optional[Union[Ansatz, str]] = "CCSD",
options: Optional[BaseOptions] = None,
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,
omega: Optional[NDArray[T]] = None,
g: Optional[NDArray[T]] = None,
G: Optional[NDArray[T]] = None,
mo_coeff: Optional[NDArray[T]] = None,
mo_occ: Optional[NDArray[T]] = None,
fock: Optional[BaseFock] = None,
**kwargs: Any,
) -> None:
Expand Down Expand Up @@ -121,11 +128,11 @@ 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_coeff: Optional[NDArray[T]] = (
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_occ: Optional[NDArray[T]] = (
mo_occ.astype(types[float]) if mo_occ is not None else None
)

# Ansatz:
Expand Down Expand Up @@ -249,7 +256,7 @@ def kernel(self, eris: Optional[ERIsInputType] = None) -> float:
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)
dt = np.max(np.abs(vector - self.amplitudes_to_vector(amplitudes_prev)))

# Update the energy and calculate change:
e_prev = e_cc
Expand Down Expand Up @@ -353,7 +360,7 @@ def solve_lambda(
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)
dl = np.max(np.abs(vector - self.lambdas_to_vector(lambdas_prev)))

# Log the iteration:
converged = bool(dl < self.options.t_tol)
Expand Down Expand Up @@ -576,7 +583,7 @@ def energy(
eris=eris,
amplitudes=amplitudes,
)
res: float = func(**kwargs).real
res: float = func(**kwargs).real.item()
return astype(res, float)

def energy_perturbative(
Expand All @@ -601,7 +608,7 @@ def energy_perturbative(
amplitudes=amplitudes,
lambdas=lambdas,
)
res: float = func(**kwargs).real
res: float = func(**kwargs).real.item()
return astype(res, float)

@abstractmethod
Expand Down Expand Up @@ -649,7 +656,7 @@ def make_sing_b_dm(
eris: Optional[ERIsInputType] = None,
amplitudes: Optional[Namespace[SpinArrayType]] = None,
lambdas: Optional[Namespace[SpinArrayType]] = None,
) -> NDArray[float]:
) -> NDArray[T]:
r"""Make the single boson density matrix :math:`\langle b \rangle`.
Args:
Expand All @@ -666,7 +673,7 @@ def make_sing_b_dm(
amplitudes=amplitudes,
lambdas=lambdas,
)
res: NDArray[float] = func(**kwargs)
res: NDArray[T] = func(**kwargs)
return res

def make_rdm1_b(
Expand All @@ -676,7 +683,7 @@ def make_rdm1_b(
lambdas: Optional[Namespace[SpinArrayType]] = None,
unshifted: bool = True,
hermitise: bool = True,
) -> NDArray[float]:
) -> NDArray[T]:
r"""Make the one-particle boson reduced density matrix :math:`\langle b^+ c \rangle`.
Args:
Expand All @@ -696,15 +703,15 @@ def make_rdm1_b(
amplitudes=amplitudes,
lambdas=lambdas,
)
dm: NDArray[float] = func(**kwargs)
dm: NDArray[T] = 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
dm_b = util.einsum("ni->i", self.make_sing_b_dm())
dm -= util.einsum("ij,i->ij", np.eye(dm.shape[0]), xi * dm_b - xi**2.0)

return dm

Expand Down Expand Up @@ -783,7 +790,7 @@ def make_eb_coup_rdm(
pass

@abstractmethod
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[T]:
"""Get a direct sum of energies.
Args:
Expand All @@ -797,7 +804,7 @@ def energy_sum(self, *args: str, signs_dict: Optional[dict[str, str]] = None) ->
pass

@abstractmethod
def amplitudes_to_vector(self, amplitudes: Namespace[SpinArrayType]) -> NDArray[float]:
def amplitudes_to_vector(self, amplitudes: Namespace[SpinArrayType]) -> NDArray[T]:
"""Construct a vector containing all of the amplitudes used in the given ansatz.
Args:
Expand All @@ -809,7 +816,7 @@ def amplitudes_to_vector(self, amplitudes: Namespace[SpinArrayType]) -> NDArray[
pass

@abstractmethod
def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
def vector_to_amplitudes(self, vector: NDArray[T]) -> Namespace[SpinArrayType]:
"""Construct a namespace of amplitudes from a vector.
Args:
Expand All @@ -821,7 +828,7 @@ def vector_to_amplitudes(self, vector: NDArray[float]) -> Namespace[SpinArrayTyp
pass

@abstractmethod
def lambdas_to_vector(self, lambdas: Namespace[SpinArrayType]) -> NDArray[float]:
def lambdas_to_vector(self, lambdas: Namespace[SpinArrayType]) -> NDArray[T]:
"""Construct a vector containing all of the lambda amplitudes used in the given ansatz.
Args:
Expand All @@ -833,7 +840,7 @@ def lambdas_to_vector(self, lambdas: Namespace[SpinArrayType]) -> NDArray[float]
pass

@abstractmethod
def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
def vector_to_lambdas(self, vector: NDArray[T]) -> Namespace[SpinArrayType]:
"""Construct a namespace of lambda amplitudes from a vector.
Args:
Expand Down Expand Up @@ -865,7 +872,7 @@ def boson_coupling_rank(self) -> int:
return self.ansatz.boson_coupling_rank

@abstractmethod
def init_space(self) -> Any:
def init_space(self) -> SpaceType:
"""Initialise the fermionic space.
Returns:
Expand Down Expand Up @@ -927,7 +934,7 @@ def bare_fock(self) -> Any:

@property
@abstractmethod
def xi(self) -> NDArray[float]:
def xi(self) -> NDArray[T]:
"""Get the shift in the bosonic operators to diagonalise the photon Hamiltonian.
Returns:
Expand All @@ -943,29 +950,30 @@ 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)
assert self.omega is not None
return util.einsum("I,I->", self.omega, self.xi**2.0).item()
return 0.0

@property
def mo_coeff(self) -> NDArray[float]:
def mo_coeff(self) -> NDArray[T]:
"""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.mf.mo_coeff.astype(types[float]) # type: ignore
return self._mo_coeff

@property
def mo_occ(self) -> NDArray[float]:
def mo_occ(self) -> NDArray[T]:
"""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.mf.mo_occ.astype(types[float]) # type: ignore
return self._mo_occ

@property
Expand Down
Loading

0 comments on commit 27b71e4

Please sign in to comment.