import logging
from functools import cached_property
from typing import TYPE_CHECKING, Literal, Optional, Union, overload
from ryd_numerov.rydberg import _CommonRydbergState
from ryd_numerov.units import BaseQuantities
if TYPE_CHECKING:
from typing_extensions import Self
from ryd_numerov.units import PintFloat
logger = logging.getLogger(__name__)
[docs]
class RydbergStateMQDT(_CommonRydbergState):
def __init__(
self,
species: str,
*,
nu: Optional[float] = None,
n: Optional[int] = None,
l: Optional[int] = None,
energy_au: Optional[float] = None,
) -> None:
r"""Initialize the Rydberg state.
Args:
species: Atomic species.
nu: Effective principal quantum number of the rydberg electron,
which is used to calculate the energy of the state.
n: Principal quantum number of the rydberg electron.
l: Orbital angular momentum quantum number of the rydberg electron.
energy_au: The energy of the Rydberg state in atomic units ("hartree").
Either `nu` or `energy_au` must be provided.
"""
self.species = species
self.n = n
assert l is not None, "l must be set"
self.l = l
self._energy_au: float
if nu is not None and energy_au is not None:
raise ValueError("Only one of nu or energy_au can be given.")
if nu is None and energy_au is None:
raise ValueError("Either nu or energy_au must be provided.")
if nu is not None:
self._energy_au = -0.5 * self.element.reduced_mass_factor / nu**2
elif energy_au is not None:
self._energy_au = energy_au
self.sanity_check()
def __repr__(self) -> str:
species, nu, n, l = self.species, self.nu, self.n, self.l
return f"{self.__class__.__name__}({species}, {nu=}, {n=}, {l=})"
def __str__(self) -> str:
return self.get_label("ket")
[docs]
def copy(self) -> "Self":
"""Create a copy of the Rydberg state."""
return self.__class__(self.species, nu=self.nu, n=self.n, l=self.l)
[docs]
def get_label(self, fmt: Literal["raw", "ket", "bra"]) -> str:
"""Label representing the ket.
Args:
fmt: The format of the label, i.e. whether to return the raw label, or the label in ket or bra notation.
Returns:
The label of the ket in the given format.
"""
l_dict = {0: "S", 1: "P", 2: "D", 3: "F", 4: "G", 5: "H"}
raw = f"{self.species}:{self.nu:.3f},{l_dict.get(self.l, self.l)}"
if self.n is not None:
raw += f"(n={self.n})"
if fmt == "raw":
return raw
if fmt == "ket":
return f"|{raw}⟩"
if fmt == "bra":
return f"⟨{raw}|"
raise ValueError(f"Unknown fmt {fmt}")
@cached_property
def nu(self) -> float:
"""Return the effective quantum number nu = n*."""
return self.get_n_star()
[docs]
def sanity_check(self) -> None:
"""Check that the quantum numbers are valid."""
msgs: list[str] = []
nu, n, l = self.nu, self.n, self.l
if not nu > 0:
msgs.append(f"nu must be larger than 0, but is {nu=}")
if n is not None:
if not isinstance(n, int):
msgs.append(f"n must be an integer, but {n=}")
if not n >= 1:
msgs.append(f"n must be larger than 0, but is {n=}")
if not nu <= n:
msgs.append(f"nu must be smaller (or equal) than n, but {nu=}, {n=}")
if not isinstance(l, int):
msgs.append(f"l must be an integer, but {l=}")
if n is not None and not 0 <= l <= n - 1:
msgs.append(f"l must be between 0 and n - 1, but {l=}, {n=}")
# TODO check is_allowed_shell like in RydbergStateSQDT (we dont know s_tot here ...)
for msg in msgs:
logger.error(msg)
if msgs:
raise ValueError(f"Invalid Rydberg state {self!r}")
@overload
def get_energy(self, unit: None = None) -> "PintFloat": ...
@overload
def get_energy(self, unit: str) -> float: ...
[docs]
def get_energy(self, unit: Optional[str] = None) -> Union["PintFloat", float]:
if unit == "a.u.":
return self._energy_au
energy: PintFloat = self._energy_au * BaseQuantities["ENERGY"]
if unit is None:
return energy
return energy.to(unit, "spectroscopy").magnitude