various small improvements

Author: johannes-moegerle

src/rydstate/angular/angular_ket.py

@@ -20,14 +20,16 @@
     try_trivial_spin_addition,
 )
 from rydstate.angular.wigner_symbols import calc_wigner_3j, clebsch_gordan_6j, clebsch_gordan_9j
-from rydstate.species import SpeciesObject
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from typing_extensions import Self
 
     from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers, AngularOperatorType
     from rydstate.angular.angular_state import AngularState
     from rydstate.angular.utils import CouplingScheme
+    from rydstate.species import SpeciesObject
 
 logger = logging.getLogger(__name__)
 
@@ -83,12 +85,13 @@ def __init__(
     ) -> None:
         """Initialize the Spin ket.
 
-        species:
-        Atomic species, e.g. 'Rb87'.
-        Not used for calculation, only for convenience to infer the core electron spin and nuclear spin quantum numbers.
+        Atomic species, e.g. 'Rb87', will not be used for calculation,
+        only for convenience to infer the core electron spin and nuclear spin quantum numbers.
         """
         if species is not None:
             if isinstance(species, str):
+                from rydstate.species import SpeciesObject  # noqa: PLC0415
+
                 species = SpeciesObject.from_name(species)
             # use i_c = 0 for species without defined nuclear spin (-> ignore hyperfine)
             species_i_c = species.i_c if species.i_c is not None else 0
@@ -228,10 +231,8 @@ def to_state(self, coupling_scheme: CouplingScheme | None = None) -> AngularStat
             The angular state in the specified coupling scheme.
 
         """
-        from rydstate.angular.angular_state import AngularState  # noqa: PLC0415
-
         if coupling_scheme is None or coupling_scheme == self.coupling_scheme:
-            return AngularState([1], [self])
+            return self._create_angular_state([1], [self])
         if coupling_scheme == "LS":
             return self._to_state_ls()
         if coupling_scheme == "JJ":
@@ -271,9 +272,7 @@ def _to_state_ls(self) -> AngularState[AngularKetLS]:
                         kets.append(ls_ket)
                         coefficients.append(coeff)
 
-        from rydstate.angular.angular_state import AngularState  # noqa: PLC0415
-
-        return AngularState(coefficients, kets)
+        return self._create_angular_state(coefficients, kets)
 
     def _to_state_jj(self) -> AngularState[AngularKetJJ]:
         """Convert a single ket to state in JJ coupling."""
@@ -306,9 +305,7 @@ def _to_state_jj(self) -> AngularState[AngularKetJJ]:
                         kets.append(jj_ket)
                         coefficients.append(coeff)
 
-        from rydstate.angular.angular_state import AngularState  # noqa: PLC0415
-
-        return AngularState(coefficients, kets)
+        return self._create_angular_state(coefficients, kets)
 
     def _to_state_fj(self) -> AngularState[AngularKetFJ]:
         """Convert a single ket to state in FJ coupling."""
@@ -341,6 +338,10 @@ def _to_state_fj(self) -> AngularState[AngularKetFJ]:
                         kets.append(fj_ket)
                         coefficients.append(coeff)
 
+        return self._create_angular_state(coefficients, kets)
+
+    def _create_angular_state(self, coefficients: Sequence[float], kets: Sequence[AngularKetBase]) -> AngularState[Any]:
+        """Create an AngularState from coefficients and kets."""
         from rydstate.angular.angular_state import AngularState  # noqa: PLC0415
 
         return AngularState(coefficients, kets)

src/rydstate/angular/angular_state.py

@@ -21,6 +21,7 @@
 
     from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers, AngularOperatorType
     from rydstate.angular.utils import CouplingScheme
+    from rydstate.units import NDArray
 
 logger = logging.getLogger(__name__)
 
@@ -30,7 +31,11 @@
 
 class AngularState(Generic[_AngularKet]):
     def __init__(
-        self, coefficients: Sequence[float], kets: Sequence[_AngularKet], *, warn_if_not_normalized: bool = True
+        self,
+        coefficients: Sequence[float] | NDArray,
+        kets: Sequence[_AngularKet],
+        *,
+        warn_if_not_normalized: bool = True,
     ) -> None:
         self.coefficients = np.array(coefficients)
         self.kets = kets

src/rydstate/basis/basis_base.py

@@ -50,7 +50,7 @@ def filter_states(self, qn: str, value: float | tuple[float, float], *, delta: f
 
         if is_angular_momentum_quantum_number(qn):
             self.states = [state for state in self.states if qn_min <= state.angular.calc_exp_qn(qn) <= qn_max]
-        elif qn in ["n", "nu", "nu_energy"]:
+        elif qn in ["n", "nu", "nu_ref"]:
             self.states = [state for state in self.states if qn_min <= getattr(state, qn) <= qn_max]
         else:
             raise ValueError(f"Unknown quantum number {qn}")
@@ -68,18 +68,18 @@ def sort_states(self, *qns: str) -> Self:
         self.states = [self.states[i] for i in sorted_indices]
         return self
 
-    def calc_exp_qn(self, qn: str) -> list[float]:
+    def calc_exp_qn(self, qn: str) -> NDArray:
         if is_angular_momentum_quantum_number(qn):
-            return [state.angular.calc_exp_qn(qn) for state in self.states]
-        if qn in ["n", "nu", "nu_energy"]:
-            return [getattr(state, qn) for state in self.states]
+            return np.array([state.angular.calc_exp_qn(qn) for state in self.states])
+        if qn in ["n", "nu", "nu_ref"]:
+            return np.array([getattr(state, qn) for state in self.states])
         raise ValueError(f"Unknown quantum number {qn}")
 
-    def calc_std_qn(self, qn: str) -> list[float]:
+    def calc_std_qn(self, qn: str) -> NDArray:
         if is_angular_momentum_quantum_number(qn):
-            return [state.angular.calc_std_qn(qn) for state in self.states]
-        if qn in ["n", "nu", "nu_energy"]:
-            return [0 for state in self.states]
+            return np.array([state.angular.calc_std_qn(qn) for state in self.states])
+        if qn in ["n", "nu", "nu_ref"]:
+            return np.zeros(len(self.states))
         raise ValueError(f"Unknown quantum number {qn}")
 
     def calc_reduced_overlap(self, other: RydbergStateBase) -> NDArray:

src/rydstate/basis/basis_sqdt.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -12,11 +13,14 @@
     RydbergStateSQDTAlkalineLS,
 )
 
+if TYPE_CHECKING:
+    from rydstate.species.species_object import SpeciesObject
+
 logger = logging.getLogger(__name__)
 
 
 class BasisSQDTAlkali(BasisBase[RydbergStateSQDTAlkali]):
-    def __init__(self, species: str, n_min: int = 1, n_max: int | None = None) -> None:
+    def __init__(self, species: str | SpeciesObject, n_min: int = 1, n_max: int | None = None) -> None:
         super().__init__(species)
 
         if n_max is None:
@@ -37,7 +41,7 @@ def __init__(self, species: str, n_min: int = 1, n_max: int | None = None) -> No
 
 
 class BasisSQDTAlkalineLS(BasisBase[RydbergStateSQDTAlkalineLS]):
-    def __init__(self, species: str, n_min: int = 1, n_max: int | None = None) -> None:
+    def __init__(self, species: str | SpeciesObject, n_min: int = 1, n_max: int | None = None) -> None:
         super().__init__(species)
 
         if n_max is None:
@@ -60,7 +64,7 @@ def __init__(self, species: str, n_min: int = 1, n_max: int | None = None) -> No
 
 
 class BasisSQDTAlkalineJJ(BasisBase[RydbergStateSQDTAlkalineJJ]):
-    def __init__(self, species: str, n_min: int = 0, n_max: int | None = None) -> None:
+    def __init__(self, species: str | SpeciesObject, n_min: int = 0, n_max: int | None = None) -> None:
         super().__init__(species)
 
         if n_max is None:
@@ -90,7 +94,7 @@ def __init__(self, species: str, n_min: int = 0, n_max: int | None = None) -> No
 
 
 class BasisSQDTAlkalineFJ(BasisBase[RydbergStateSQDTAlkalineFJ]):
-    def __init__(self, species: str, n_min: int = 0, n_max: int | None = None) -> None:
+    def __init__(self, species: str | SpeciesObject, n_min: int = 0, n_max: int | None = None) -> None:
         super().__init__(species)
 
         if n_max is None:

src/rydstate/rydberg/rydberg_base.py

@@ -22,3 +22,8 @@ def calc_reduced_overlap(self, other: RydbergStateBase) -> float: ...
     def calc_reduced_matrix_element(
         self, other: RydbergStateBase, operator: MatrixElementOperator, unit: str | None = None
     ) -> PintFloat | float: ...
+
+    @property
+    @abstractmethod
+    def nu_ref(self) -> float:
+        """The reference effective principal quantum number nu_ref."""

src/rydstate/rydberg/rydberg_sqdt.py

@@ -15,6 +15,8 @@
 from rydstate.units import BaseQuantities, MatrixElementOperatorRanks, ureg
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
     from rydstate.angular.angular_ket import AngularKetBase, AngularKetFJ, AngularKetJJ, AngularKetLS
     from rydstate.units import MatrixElementOperator, PintFloat
 
@@ -94,14 +96,19 @@ def __init__(
         if nu is None and n is None:
             raise ValueError("Either n or nu must be given to initialize the Rydberg state.")
 
+        self._set_qn_as_attributes()
+
+    def _set_qn_as_attributes(self) -> None:
+        pass
+
     @classmethod
     def from_angular_ket(
-        cls,
+        cls: type[Self],
         species: str | SpeciesObject,
         angular_ket: AngularKetBase,
         n: int | None = None,
         nu: float | None = None,
-    ) -> RydbergStateSQDT:
+    ) -> Self:
         """Initialize the Rydberg state from an angular ket."""
         obj = cls.__new__(cls)
 
@@ -115,13 +122,14 @@ def from_angular_ket(
             raise ValueError("Either n or nu must be given to initialize the Rydberg state.")
 
         obj.angular = angular_ket
+        obj._set_qn_as_attributes()  # noqa: SLF001
 
         return obj
 
     def __repr__(self) -> str:
         species, n, nu = self.species.name, self.n, self.nu
         n_str = f", {n=}" if n is not None else ""
-        return f"{self.__class__.__name__}({species=}{n_str}, {nu=}, {self.angular})"
+        return f"{self.__class__.__name__}({species}{n_str}, {nu=}, {self.angular})"
 
     def __str__(self) -> str:
         return self.__repr__()
@@ -151,13 +159,19 @@ def nu(self) -> float:
         """The effective principal quantum number nu (for alkali atoms also known as n*) for the Rydberg state."""
         if self._nu is not None:
             return self._nu
-        assert self.n is not None
-        if any(qn not in self.angular.quantum_number_names for qn in ["j_tot", "s_tot"]):
-            raise ValueError("j_tot and s_tot must be defined to calculate nu from n.")
+        assert isinstance(self.species, SpeciesObject), "nu must be given if not sqdt"
+        assert self.n is not None, "either nu or n must be given"
+
+        if "j_tot" not in self.angular.quantum_number_names or "s_tot" not in self.angular.quantum_number_names:
+            raise RuntimeError("j_tot and s_tot must be defined in the angular ket to calculate nu from n.")
         return self.species.calc_nu(
             self.n, self.angular.l_r, self.angular.get_qn("j_tot"), s_tot=self.angular.get_qn("s_tot")
         )
 
+    @property
+    def nu_ref(self) -> float:
+        return self.nu
+
     @overload
     def get_energy(self, unit: None = None) -> PintFloat: ...
 
@@ -342,15 +356,18 @@ def __init__(
         """
         super().__init__(species=species, n=n, nu=nu, l_r=l, j_tot=j, f_tot=f, m=m)
 
-        self.l = l
+    def _set_qn_as_attributes(self) -> None:
+        self.l = self.angular.l_r
         self.j = self.angular.j_tot
         self.f = self.angular.f_tot
-        self.m = m
+        self.m = self.angular.m
 
     def __repr__(self) -> str:
-        species, n, l, j, f, m = self.species, self.n, self.l, self.j, self.f, self.m
+        species, n, nu = self.species.name, self.n, self.nu
+        l, j, f, m = self.l, self.j, self.f, self.m
+        n_str = f", {n=}" if n is not None else ""
         f_string = f", {f=}" if self.species.i_c not in (None, 0) else ""
-        return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j=}{f_string}, {m=})"
+        return f"{self.__class__.__name__}({species}{n_str}, {nu=}, {l=}, {j=}{f_string}, {m=})"
 
 
 class RydbergStateSQDTAlkalineLS(RydbergStateSQDT):
@@ -388,15 +405,18 @@ def __init__(
         """
         super().__init__(species=species, n=n, nu=nu, l_r=l, s_tot=s_tot, j_tot=j_tot, f_tot=f_tot, m=m)
 
-        self.l = l
+    def _set_qn_as_attributes(self) -> None:
+        self.l = self.angular.l_r
         self.s_tot = self.angular.s_tot
         self.j_tot = self.angular.j_tot
         self.f_tot = self.angular.f_tot
-        self.m = m
+        self.m = self.angular.m
 
     def __repr__(self) -> str:
-        species, n, l, s_tot, j_tot, f_tot, m = self.species, self.n, self.l, self.s_tot, self.j_tot, self.f_tot, self.m
-        return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {s_tot=}, {j_tot=}, {f_tot=}, {m=})"
+        species, n, nu = self.species.name, self.n, self.nu
+        l, s_tot, j_tot, f_tot, m = self.l, self.s_tot, self.j_tot, self.f_tot, self.m
+        n_str = f", {n=}" if n is not None else ""
+        return f"{self.__class__.__name__}({species}{n_str}, {nu=}, {l=}, {s_tot=}, {j_tot=}, {f_tot=}, {m=})"
 
 
 class RydbergStateSQDTAlkalineJJ(RydbergStateSQDT):
@@ -434,29 +454,18 @@ def __init__(
         """
         super().__init__(species=species, n=n, nu=nu, l_r=l, j_r=j_r, j_tot=j_tot, f_tot=f_tot, m=m)
 
+    def _set_qn_as_attributes(self) -> None:
         self.l = self.angular.l_r
         self.j_r = self.angular.j_r
         self.j_tot = self.angular.j_tot
         self.f_tot = self.angular.f_tot
         self.m = self.angular.m
 
     def __repr__(self) -> str:
-        species, n, l, j_r, j_tot, f_tot, m = self.species, self.n, self.l, self.j_r, self.j_tot, self.f_tot, self.m
-        return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j_r=}, {j_tot=}, {f_tot=}, {m=})"
-
-    @cached_property
-    def nu(self) -> float:
-        if self._nu is not None:
-            return self._nu
-        assert self.n is not None
-        nus = [self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=s_tot) for s_tot in [0, 1]]
-
-        if any(abs(nu - nus[0]) > 1e-10 for nu in nus[1:]):
-            raise ValueError(
-                "RydbergStateSQDTAlkalineJJ is intended for high-l states only, "
-                "where the quantum defects are the same for singlet and triplet states."
-            )
-        return nus[0]
+        species, n, nu = self.species.name, self.n, self.nu
+        l, j_r, j_tot, f_tot, m = self.l, self.j_r, self.j_tot, self.f_tot, self.m
+        n_str = f", {n=}" if n is not None else ""
+        return f"{self.__class__.__name__}({species}{n_str}, {nu=}, {l=}, {j_r=}, {j_tot=}, {f_tot=}, {m=})"
 
 
 class RydbergStateSQDTAlkalineFJ(RydbergStateSQDT):
@@ -494,30 +503,17 @@ def __init__(
         """
         super().__init__(species=species, n=n, nu=nu, l_r=l, j_r=j_r, f_c=f_c, f_tot=f_tot, m=m)
 
+    def _set_qn_as_attributes(self) -> None:
         self.l = self.angular.l_r
         self.j_r = self.angular.j_r
         self.f_c = self.angular.f_c
         self.f_tot = self.angular.f_tot
         self.m = self.angular.m
 
     def __repr__(self) -> str:
-        species, n, l, j_r, f_c, f_tot, m = self.species, self.n, self.l, self.j_r, self.f_c, self.f_tot, self.m
-        return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j_r=}, {f_c=}, {f_tot=}, {m=})"
-
-    @cached_property
-    def nu(self) -> float:
-        if self._nu is not None:
-            return self._nu
-        assert self.n is not None
-        nus = [
-            self.species.calc_nu(self.n, self.l, float(j_tot), s_tot=s_tot)
-            for s_tot in [0, 1]
-            for j_tot in np.arange(abs(self.j_r - 1 / 2), self.j_r + 1 / 2 + 1)
-        ]
-
-        if any(abs(nu - nus[0]) > 1e-10 for nu in nus[1:]):
-            raise ValueError(
-                "RydbergStateSQDTAlkalineFJ is intended for high-l states only, "
-                "where the quantum defects are the same for singlet and triplet states."
-            )
-        return nus[0]
+        species, n, nu = self.species.name, self.n, self.nu
+        l, j_r, f_c, f_tot, m = self.l, self.j_r, self.f_c, self.f_tot, self.m
+        l_c, j_c = self.angular.l_c, self.angular.j_c
+        core_string = f", {l_c=}, {j_c=}" if l_c != 0 else ""
+        n_str = f", {n=}" if n is not None else ""
+        return f"{self.__class__.__name__}({species}{n_str}, {nu=}{core_string}, {l=}, {j_r=}, {f_c=}, {f_tot=}, {m=})"