RydbergStateSQDT add from_angular_ket

Author: johannes-moegerle

src/rydstate/rydberg/rydberg_base.py

@@ -13,9 +13,7 @@
 
 
 class RydbergStateBase(ABC):
-    @property
-    @abstractmethod
-    def angular(self) -> AngularState[Any] | AngularKetBase: ...
+    angular: AngularState[Any] | AngularKetBase
 
     @abstractmethod
     def calc_reduced_overlap(self, other: RydbergStateBase) -> float: ...

src/rydstate/rydberg/rydberg_sqdt.py

@@ -24,6 +24,10 @@
 
 class RydbergStateSQDT(RydbergStateBase):
     species: SpeciesObject
+    """The atomic species of the Rydberg state."""
+
+    angular: AngularKetBase
+    """The angular/spin part of the Rydberg electron."""
 
     def __init__(
         self,
@@ -69,7 +73,8 @@ def __init__(
             species = SpeciesObject.from_name(species)
         self.species = species
 
-        self._qns = dict(  # noqa: C408
+        self.angular = quantum_numbers_to_angular_ket(
+            species=self.species,
             s_c=s_c,
             l_c=l_c,
             j_c=j_c,
@@ -89,6 +94,30 @@ def __init__(
         if nu is None and n is None:
             raise ValueError("Either n or nu must be given to initialize the Rydberg state.")
 
+    @classmethod
+    def from_angular_ket(
+        cls,
+        species: str | SpeciesObject,
+        angular_ket: AngularKetBase,
+        n: int | None = None,
+        nu: float | None = None,
+    ) -> RydbergStateSQDT:
+        """Initialize the Rydberg state from an angular ket."""
+        obj = cls.__new__(cls)
+
+        if isinstance(species, str):
+            species = SpeciesObject.from_name(species)
+        obj.species = species
+
+        obj.n = n
+        obj._nu = nu  # noqa: SLF001
+        if nu is None and n is None:
+            raise ValueError("Either n or nu must be given to initialize the Rydberg state.")
+
+        obj.angular = angular_ket
+
+        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 ""
@@ -100,6 +129,11 @@ def __str__(self) -> str:
     @cached_property
     def radial(self) -> RadialKet:
         """The radial part of the Rydberg electron."""
+        if "l_r" not in self.angular.quantum_number_names:
+            raise ValueError(
+                f"l_r must be defined in the angular ket to access the radial ket, but angular={self.angular}."
+            )
+
         radial_ket = RadialKet(self.species, nu=self.nu, l_r=self.angular.l_r)
         if self.n is not None:
             radial_ket.set_n_for_sanity_check(self.n)
@@ -112,11 +146,6 @@ def radial(self) -> RadialKet:
                     )
         return radial_ket
 
-    @cached_property
-    def angular(self) -> AngularKetBase:
-        """The angular/spin part of the Rydberg electron."""
-        return quantum_numbers_to_angular_ket(species=self.species, **self._qns)  # type: ignore [arg-type]
-
     @cached_property
     def nu(self) -> float:
         """The effective principal quantum number nu (for alkali atoms also known as n*) for the Rydberg state."""