# SPDX-FileCopyrightText: 2024 PairInteraction Developers
# SPDX-License-Identifier: LGPL-3.0-or-later
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Literal, overload
import numpy as np
from scipy.sparse import csr_matrix
from pairinteraction import _backend
from pairinteraction.basis.basis_base import BasisBase
from pairinteraction.database import Database
from pairinteraction.enums import get_cpp_operator_type, parity_to_int
from pairinteraction.ket import KetAtom, KetAtomReal
from pairinteraction.state import StateAtom, StateAtomReal
from pairinteraction.units import QuantityArray, QuantityScalar, QuantitySparse
if TYPE_CHECKING:
from collections.abc import Sequence
from typing_extensions import Self
from pairinteraction.enums import OperatorType, Parity
from pairinteraction.units import NDArray, PintArray, PintFloat, PintSparse
[docs]
class BasisAtom(BasisBase[KetAtom, StateAtom]):
"""Basis for a single atom.
Add all KetAtom objects that match the given quantum numbers to the basis.
The initial coefficients matrix is a unit matrix, i.e. the first basis state is the first ket, etc.
The BasisAtom coefficients matrix will always be square,
i.e. the number of kets is equal to the number of states.
Examples:
>>> import pairinteraction as pi
>>> ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
>>> energy_min, energy_max = ket.get_energy(unit="GHz") - 100, ket.get_energy(unit="GHz") + 100
>>> basis = pi.BasisAtom("Rb", n=(57, 63), l=(0, 3), energy=(energy_min, energy_max), energy_unit="GHz")
>>> print(basis)
BasisAtom('Rb', n=(57, 63), l=(0, 3), energy=(1008911.9216, 1009111.9216), energy_unit='GHz')
"""
_cpp: _backend.BasisAtomComplex
_cpp_creator = _backend.BasisAtomCreatorComplex
_ket_class = KetAtom
_state_class = StateAtom
_args: dict[str, Any] | None = None
[docs]
def __init__( # noqa: C901, PLR0912
self,
species: str,
n: tuple[int, int] | None = None,
nu: tuple[float, float] | None = None,
nui: tuple[float, float] | None = None,
l: tuple[float, float] | None = None,
s: tuple[float, float] | None = None,
j: tuple[float, float] | None = None,
l_ryd: tuple[float, float] | None = None,
j_ryd: tuple[float, float] | None = None,
f: tuple[float, float] | None = None,
m: tuple[float, float] | None = None,
energy: tuple[float, float] | tuple[PintFloat, PintFloat] | None = None,
energy_unit: str | None = None,
parity: Parity | None = None,
additional_kets: Sequence[KetAtom] | None = None,
*,
database: Database | None = None,
mode: Literal["exact", "fuzzy"] | float = "fuzzy",
) -> None:
"""Create a basis for a single atom.
Args:
species: The species of the atom.
n: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
nu: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
nui: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
l: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
s: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
j: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
l_ryd: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
j_ryd: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
f: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
m: tuple of (min, max) values for this quantum number. Default None, i.e. add all available states.
energy: tuple of (min, max) value for the energy. Default None, i.e. add all available states.
energy_unit: In which unit the energy values are given, e.g. "GHz".
Default None, i.e. energy is provided as pint object.
parity: The parity of the states to consider. Default None, i.e. add all available states.
additional_kets: List of additional kets to add to the basis. Default None.
database: Which database to use. Default None, i.e. use the global database instance.
mode: Specifies how restrictions on expectation-value quantum numbers are applied.
``"fuzzy"`` is equal to 2 and includes states whose expectation-value overlaps the requested range
within two standard deviations.
``"exact"`` is equal to 0 and includes only states whose expectation-value itself lie in the range.
A non-negative number sets the factor applied to the standard deviation explicitly.
Default ``"fuzzy"``.
"""
self._args = {"species": species}
creator = self._cpp_creator()
creator.set_species(species)
if n is not None and not all(isinstance(x, int) or x.is_integer() for x in n):
raise ValueError("Quantum numbers n must be integers.")
quantum_numbers = {
"n": n,
"nu": nu,
"nui": nui,
"s": s,
"l": l,
"j": j,
"l_ryd": l_ryd,
"j_ryd": j_ryd,
"f": f,
"m": m,
}
for name, value in quantum_numbers.items():
if value is not None:
self._args[name] = value
creator.restrict_quantum_number(name, *value)
if parity is not None:
self._args["parity"] = parity
parity_int = parity_to_int(parity)
creator.restrict_quantum_number("parity", parity_int, parity_int)
if energy is not None:
self._args.update({"energy": energy, "energy_unit": energy_unit})
min_energy_au = QuantityScalar.convert_user_to_au(energy[0], energy_unit, "energy")
max_energy_au = QuantityScalar.convert_user_to_au(energy[1], energy_unit, "energy")
creator.restrict_energy(min_energy_au, max_energy_au)
if database is None:
if Database.get_global_database() is None:
Database.initialize_global_database()
database = Database.get_global_database()
if additional_kets is not None:
self._args["additional_kets"] = additional_kets
for ket in additional_kets:
creator.add_ket(ket._cpp)
if mode == "fuzzy":
quantum_number_standard_deviation_factor = 2.0
elif mode == "exact":
self._args["mode"] = mode
quantum_number_standard_deviation_factor = 0.0
else:
self._args["mode"] = mode
msg = "mode must be 'exact', 'fuzzy', or a non-negative number."
try:
quantum_number_standard_deviation_factor = float(mode)
except (TypeError, ValueError) as err:
raise ValueError(msg) from err
if quantum_number_standard_deviation_factor < 0:
raise ValueError(msg)
creator.set_quantum_number_standard_deviation_factor(quantum_number_standard_deviation_factor)
self._cpp = creator.create(database._cpp)
self._post_init()
[docs]
@classmethod
def from_kets(
cls: type[Self],
kets: KetAtom | Sequence[KetAtom],
delta_n: int | None = None,
delta_nu: float | None = None,
delta_nui: float | None = None,
delta_l: float | None = None,
delta_s: float | None = None,
delta_j: float | None = None,
delta_l_ryd: float | None = None,
delta_j_ryd: float | None = None,
delta_f: int | None = None,
delta_m: int | None = None,
delta_energy: float | PintFloat | None = None,
delta_energy_unit: str | None = None,
parity: Parity | None = None,
database: Database | None = None,
additional_kets: Sequence[KetAtom] | None = None,
*,
mode: Literal["exact", "fuzzy"] | float = "fuzzy",
) -> Self:
"""Create a BasisAtom from one or more kets and quantum number deltas.
Currently a single big basis including all kets for the quantum numbers
from min_value - delta to max_value + delta is returned.
In the future this might change to return a basis including all states around the given kets +/- delta,
but not necessarily all states between the given kets.
For each quantum number, pass the corresponding ``delta_*`` argument to include
all states within ``[min_value - delta, max_value + delta]``, where
``min_value`` / ``max_value`` are the extremes across all provided kets.
If no ``delta_*`` is given for a quantum number, that quantum number is left
unrestricted.
Args:
kets: The ket(s) around which the basis should be centered.
delta_n: Half-width of the n window (integer steps).
Default None means no restriction on n.
delta_nu: Half-width of the nu window.
Default None means no restriction on nu.
delta_nui: Half-width of the nui window.
Default None means no restriction on nui.
delta_l: Half-width of the l window.
Default None means no restriction on l.
delta_s: Half-width of the s window.
Default None means no restriction on s.
delta_j: Half-width of the j window.
Default None means no restriction on j.
delta_l_ryd: Half-width of the l_ryd window.
Default None means no restriction on l_ryd.
delta_j_ryd: Half-width of the j_ryd window.
Default None means no restriction on j_ryd.
delta_f: Half-width of the f window (integer steps).
Default None means no restriction on f.
delta_m: Half-width of the m window (integer steps).
Default None means no restriction on m.
delta_energy: Half-width of the energy window around the energies of the
provided kets. Default None means no energy restriction.
delta_energy_unit: Unit for ``delta_energy`` (e.g. ``"GHz"``).
Default None means pint quantities are used.
parity: Restrict to states with this parity.
Default None means no parity restriction.
database: Database instance to use.
Default None uses the global database.
additional_kets: Extra kets to force-include in the basis.
Default None.
mode: Passed to :class:`BasisAtom`. Default ``"fuzzy"``.
Returns:
A new :class:`BasisAtom` centered around the provided kets.
Examples:
>>> import pairinteraction as pi
>>> ket1 = pi.KetAtom("Rb", n=60, l=0, m=0.5)
>>> ket2 = pi.KetAtom("Rb", n=59, l=0, m=0.5)
>>> basis = pi.BasisAtom.from_kets([ket1, ket2], delta_n=2, delta_l=1)
>>> basis.species
'Rb'
>>> all(57 <= k.n <= 62 for k in basis.kets)
True
"""
if isinstance(kets, KetAtom):
kets = [kets]
kets = list(kets)
if len(kets) == 0:
raise ValueError("kets must not be empty.")
if len({ket.species for ket in kets}) > 1:
raise ValueError(f"All kets must have the same species, but got: {sorted({ket.species for ket in kets})}.")
def get_range(name: str, delta: float | None) -> tuple[float, float] | None:
if delta is None:
return None
if name == "energy":
values = [ket.get_energy(unit=delta_energy_unit) for ket in kets]
else:
values = [getattr(ket, name) for ket in kets]
return (min(values) - delta, max(values) + delta)
return cls(
species=kets[0].species,
n=get_range("n", delta_n), # type: ignore [arg-type]
nu=get_range("nu", delta_nu),
nui=get_range("nui", delta_nui),
l=get_range("l", delta_l),
s=get_range("s", delta_s),
j=get_range("j", delta_j),
l_ryd=get_range("l_ryd", delta_l_ryd),
j_ryd=get_range("j_ryd", delta_j_ryd),
f=get_range("f", delta_f),
m=get_range("m", delta_m),
energy=get_range("energy", delta_energy), # type: ignore [arg-type]
energy_unit=delta_energy_unit,
parity=parity,
database=database,
additional_kets=additional_kets,
mode=mode,
)
def __repr__(self) -> str:
if self._args is None:
return super().__repr__()
args_str: list[str] = []
for k, v in self._args.items():
if k == "species":
args_str.append(repr(v))
elif k == "energy":
args_str.append(f"{k}=({v[0]:.4f}, {v[1]:.4f})")
else:
args_str.append(f"{k}={v!r}")
return f"{type(self).__name__}({', '.join(args_str)})"
@property
def database(self) -> Database:
"""The database used for this object."""
return self.get_ket(0).database
@property
def species(self) -> str:
"""The atomic species."""
return self.get_ket(0).species
@overload
def get_amplitudes(self, other: KetAtom | StateAtom) -> NDArray: ...
@overload
def get_amplitudes(self, other: BasisAtom) -> csr_matrix: ...
[docs]
def get_amplitudes(self, other: KetAtom | StateAtom | BasisAtom) -> NDArray | csr_matrix:
return self.get_matrix_elements(other, "identity", 0, unit="")
@overload
def get_overlaps(self, other: KetAtom | StateAtom) -> NDArray: ...
@overload
def get_overlaps(self, other: BasisAtom) -> csr_matrix: ...
[docs]
def get_overlaps(self, other: KetAtom | StateAtom | BasisAtom) -> NDArray | csr_matrix:
amplitudes = self.get_amplitudes(other)
if isinstance(amplitudes, csr_matrix):
return amplitudes.multiply(amplitudes.conj()).real # type: ignore [no-any-return]
return np.abs(amplitudes) ** 2
@overload
def get_matrix_elements(
self, other: KetAtom | StateAtom, operator: OperatorType, q: int, unit: None = None
) -> PintArray: ...
@overload
def get_matrix_elements(self, other: KetAtom | StateAtom, operator: OperatorType, q: int, unit: str) -> NDArray: ...
@overload
def get_matrix_elements(
self, other: BasisAtom, operator: OperatorType, q: int, unit: None = None
) -> PintSparse: ...
@overload
def get_matrix_elements(self, other: BasisAtom, operator: OperatorType, q: int, unit: str) -> csr_matrix: ...
[docs]
def get_matrix_elements(
self, other: KetAtom | StateAtom | BasisAtom, operator: OperatorType, q: int, unit: str | None = None
) -> NDArray | PintArray | csr_matrix | PintSparse:
cpp_op = get_cpp_operator_type(operator)
matrix_elements_au: NDArray
if isinstance(other, KetAtom):
other = other.to_state()
if isinstance(other, StateAtom):
matrix_elements_au = self._cpp.get_matrix_elements(other._cpp, cpp_op, q).toarray().ravel()
matrix_elements_au = np.real_if_close(matrix_elements_au)
return QuantityArray.convert_au_to_user(matrix_elements_au, operator, unit)
if isinstance(other, BasisAtom):
matrix_elements_sparse_au = self._cpp.get_matrix_elements(other._cpp, cpp_op, q)
matrix_elements_sparse_au.data = np.real_if_close(matrix_elements_sparse_au.data)
return QuantitySparse.convert_au_to_user(matrix_elements_sparse_au, operator, unit)
raise TypeError(f"Unknown type: {type(other)=}")
class BasisAtomReal(BasisAtom):
_cpp: _backend.BasisAtomReal # type: ignore [assignment]
_cpp_creator = _backend.BasisAtomCreatorReal # type: ignore [assignment]
_ket_class = KetAtomReal
_state_class = StateAtomReal
def get_cpp_basis_atom_from_kets(kets: Sequence[KetAtom], *, real: bool) -> _backend.BasisAtomComplex:
"""Create a cpp BasisAtom object containing only the given kets.
Like for the _cpp attributes, the return type is annotated as the complex variant,
although the real variant is returned if real=True.
"""
if len(kets) == 0:
raise ValueError("Cannot create a basis with zero kets.")
creator = _backend.BasisAtomCreatorReal() if real else _backend.BasisAtomCreatorComplex()
for ket in kets:
creator.add_ket(ket._cpp)
return creator.create(kets[0].database._cpp) # type: ignore [return-value]