Lifetimes

docs/examples.rst

@@ -26,6 +26,7 @@ Some examples demonstrating the usage of the RydbergState class.
 
 .. nbgallery::
    examples/dipole_matrix_elements
+   examples/lifetimes
 
 
 Comparisons

src/rydstate/angular/__init__.py

@@ -1,11 +1,13 @@
 from rydstate.angular import utils
 from rydstate.angular.angular_ket import AngularKetFJ, AngularKetJJ, AngularKetLS
 from rydstate.angular.angular_state import AngularState
+from rydstate.angular.utils import NotSet
 
 __all__ = [
     "AngularKetFJ",
     "AngularKetJJ",
     "AngularKetLS",
     "AngularState",
+    "NotSet",
     "utils",
 ]

src/rydstate/angular/angular_ket.py

@@ -14,8 +14,10 @@
 )
 from rydstate.angular.utils import (
     InvalidQuantumNumbersError,
+    NotSet,
     check_spin_addition_rule,
     get_possible_quantum_number_values,
+    is_not_set,
     minus_one_pow,
     try_trivial_spin_addition,
 )
@@ -67,9 +69,9 @@ class AngularKetBase(ABC):
 
     f_tot: float
     """Total atom angular quantum number (including nuclear, core electron and rydberg electron contributions)."""
-    m: float | None
+    m: float | NotSet
     """Magnetic quantum number, which is the projection of `f_tot` onto the quantization axis.
-    If None, only reduced matrix elements can be calculated.
+    If NotSet, only reduced matrix elements can be calculated.
     """
 
     def __init__(
@@ -80,7 +82,7 @@ def __init__(
         s_r: float = 0.5,
         l_r: int | None = None,
         f_tot: float | None = None,  # noqa: ARG002
-        m: float | None = None,
+        m: float | NotSet = NotSet,
         species: str | SpeciesObject | None = None,
     ) -> None:
         """Initialize the Spin ket.
@@ -94,10 +96,9 @@ def __init__(
 
                 species = SpeciesObjectSQDT.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
-            if i_c is not None and i_c != species_i_c:
+            if i_c is not None and i_c != species.i_c_number:
                 raise ValueError(f"Nuclear spin i_c={i_c} does not match the species {species} with i_c={species.i_c}.")
-            i_c = species_i_c
+            i_c = species.i_c_number
             s_c = 0.5 * (species.number_valence_electrons - 1)
         if i_c is None:
             raise ValueError("Nuclear spin i_c must be set or a species must be given.")
@@ -114,7 +115,7 @@ def __init__(
         self.l_r = int(l_r)
 
         # f_tot will be set in the subclasses
-        self.m = None if m is None else float(m)
+        self.m = NotSet if is_not_set(m) else float(m)
 
     def _post_init(self) -> None:
         self.quantum_numbers = tuple(getattr(self, qn) for qn in self.quantum_number_names)
@@ -132,7 +133,7 @@ def sanity_check(self, msgs: list[str] | None = None) -> None:
         if self.s_r != 0.5:
             msgs.append(f"Rydberg electron spin s_r must be 1/2, but {self.s_r=}")
 
-        if self.m is not None and not -self.f_tot <= self.m <= self.f_tot:
+        if not is_not_set(self.m) and not -self.f_tot <= self.m <= self.f_tot:
             msgs.append(f"m must be between -f_tot and f_tot, but {self.f_tot=}, {self.m=}")
 
         if msgs:
@@ -149,7 +150,7 @@ def __setattr__(self, key: str, value: object) -> None:
 
     def __repr__(self) -> str:
         args = ", ".join(f"{qn}={val}" for qn, val in zip(self.quantum_number_names, self.quantum_numbers, strict=True))
-        if self.m is not None:
+        if not is_not_set(self.m):
             args += f", m={self.m}"
         return f"{self.__class__.__name__}({args})"
 
@@ -158,7 +159,7 @@ def __str__(self) -> str:
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, AngularKetBase):
-            raise NotImplementedError(f"Cannot compare {self!r} with {other!r}.")
+            return NotImplemented
         if type(self) is not type(other):
             return False
         if self.m != other.m:
@@ -466,15 +467,16 @@ def calc_matrix_element(self, other: AngularKetBase, operator: AngularOperatorTy
             The dimensionless angular matrix element.
 
         """
-        if self.m is None or other.m is None:
-            raise ValueError("m must be set to calculate the matrix element.")
+        if is_not_set(self.m) or is_not_set(other.m):
+            raise RuntimeError("m must be set to calculate the matrix element.")
 
         prefactor = self._calc_wigner_eckart_prefactor(other, kappa, q)
         reduced_matrix_element = self.calc_reduced_matrix_element(other, operator, kappa)
         return prefactor * reduced_matrix_element
 
     def _calc_wigner_eckart_prefactor(self, other: AngularKetBase, kappa: int, q: int) -> float:
-        assert self.m is not None and other.m is not None, "m must be set to calculate the Wigner-Eckart prefactor."  # noqa: PT018
+        if is_not_set(self.m) or is_not_set(other.m):
+            raise RuntimeError("m must be set to calculate the Wigner-Eckart prefactor.")
         return minus_one_pow(self.f_tot - self.m) * calc_wigner_3j(self.f_tot, kappa, other.f_tot, -self.m, q, other.m)
 
     def _kronecker_delta_non_involved_spins(self, other: AngularKetBase, qn: AngularMomentumQuantumNumbers) -> int:
@@ -572,7 +574,7 @@ def __init__(
         l_tot: int | None = None,
         j_tot: float | None = None,
         f_tot: float | None = None,
-        m: float | None = None,
+        m: float | NotSet = NotSet,
         species: str | SpeciesObject | None = None,
     ) -> None:
         """Initialize the Spin ket."""
@@ -635,7 +637,7 @@ def __init__(
         j_r: float | None = None,
         j_tot: float | None = None,
         f_tot: float | None = None,
-        m: float | None = None,
+        m: float | NotSet = NotSet,
         species: str | SpeciesObject | None = None,
     ) -> None:
         """Initialize the Spin ket."""
@@ -698,7 +700,7 @@ def __init__(
         f_c: float | None = None,
         j_r: float | None = None,
         f_tot: float | None = None,
-        m: float | None = None,
+        m: float | NotSet = NotSet,
         species: str | SpeciesObject | None = None,
     ) -> None:
         """Initialize the Spin ket."""

src/rydstate/angular/angular_state.py

@@ -13,6 +13,7 @@
     AngularKetLS,
 )
 from rydstate.angular.angular_matrix_element import is_angular_momentum_quantum_number
+from rydstate.angular.utils import is_not_set
 
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
@@ -209,8 +210,8 @@ def calc_matrix_element(
                     "Different m values are not supported yet for AngularState.calc_matrix_element."
                 )
 
-        if self.kets[0].m is None or other.kets[0].m is None:
-            raise ValueError("m must be set for all kets to calculate the matrix element.")
+        if is_not_set(self.kets[0].m) or is_not_set(other.kets[0].m):
+            raise RuntimeError("m must be set for all kets to calculate the matrix element.")
 
         prefactor = self.kets[0]._calc_wigner_eckart_prefactor(other.kets[0], kappa, q)  # noqa: SLF001
         reduced_matrix_element = self.calc_reduced_matrix_element(other, operator, kappa)

src/rydstate/angular/utils.py

@@ -1,17 +1,39 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Literal
+import typing as t
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 
 if TYPE_CHECKING:
+    from typing_extensions import TypeIs
+
     from rydstate.angular.angular_ket import AngularKetBase
     from rydstate.species.species_object import SpeciesObject
 
 CouplingScheme = Literal["LS", "JJ", "FJ"]
 
 
+@t.runtime_checkable
+class NotSet(t.Protocol):
+    """Singleton for a not set value and type at the same time.
+
+    See Also:
+    https://stackoverflow.com/questions/77571796/how-to-create-singleton-object-which-could-be-used-both-as-type-and-value-simi
+
+    """
+
+    @staticmethod
+    def __not_set() -> None: ...
+
+
+def is_not_set(obj: Any) -> TypeIs[NotSet]:  # noqa: ANN401
+    """Check if the obj is the NotSet singleton."""
+    return obj is NotSet
+
+
 class InvalidQuantumNumbersError(ValueError):
     def __init__(self, ket: AngularKetBase, msg: str = "") -> None:
         _msg = f"Invalid quantum numbers for {ket!r}"
@@ -61,6 +83,7 @@ def get_possible_quantum_number_values(s_1: float, s_2: float, s_tot: float | No
     return [float(s) for s in np.arange(abs(s_1 - s_2), s_1 + s_2 + 1, 1)]
 
 
+@lru_cache(maxsize=1_000)
 def quantum_numbers_to_angular_ket(
     species: str | SpeciesObject,
     s_c: float | None = None,
@@ -74,7 +97,7 @@ def quantum_numbers_to_angular_ket(
     l_tot: int | None = None,
     j_tot: float | None = None,
     f_tot: float | None = None,
-    m: float | None = None,
+    m: float | NotSet = NotSet,
 ) -> AngularKetBase:
     r"""Return an AngularKet object in the corresponding coupling scheme from the given quantum numbers.