start adding AngularKetDummy

Author: johannes-moegerle

src/rydstate/angular/__init__.py

@@ -1,8 +1,9 @@
 from rydstate.angular import utils
-from rydstate.angular.angular_ket import AngularKetFJ, AngularKetJJ, AngularKetLS
+from rydstate.angular.angular_ket import AngularKetDummy, AngularKetFJ, AngularKetJJ, AngularKetLS
 from rydstate.angular.angular_state import AngularState
 
 __all__ = [
+    "AngularKetDummy",
     "AngularKetFJ",
     "AngularKetJJ",
     "AngularKetLS",

src/rydstate/angular/angular_ket.py

@@ -24,14 +24,14 @@
 
 if TYPE_CHECKING:
     import juliacall
-    from typing_extensions import Self
+    from typing_extensions import Never, Self
 
     from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers, AngularOperatorType
     from rydstate.angular.angular_state import AngularState
 
 logger = logging.getLogger(__name__)
 
-CouplingScheme = Literal["LS", "JJ", "FJ"]
+CouplingScheme = Literal["LS", "JJ", "FJ", "Dummy"]
 
 
 class InvalidQuantumNumbersError(ValueError):
@@ -224,6 +224,9 @@ def to_state(self, coupling_scheme: Literal["JJ"]) -> AngularState[AngularKetJJ]
     @overload
     def to_state(self, coupling_scheme: Literal["FJ"]) -> AngularState[AngularKetFJ]: ...
 
+    @overload
+    def to_state(self, coupling_scheme: Literal["Dummy"]) -> Never: ...
+
     @overload
     def to_state(self: Self) -> AngularState[Self]: ...
 
@@ -372,6 +375,10 @@ def calc_reduced_overlap(self, other: AngularKetBase) -> float:
 
         kets = [self, other]
 
+        # Dummy overlaps
+        if any(isinstance(s, AngularKetDummy) for s in kets):
+            return int(self == other)
+
         # JJ - FJ overlaps
         if any(isinstance(s, AngularKetJJ) for s in kets) and any(isinstance(s, AngularKetFJ) for s in kets):
             jj = next(s for s in kets if isinstance(s, AngularKetJJ))
@@ -414,6 +421,10 @@ def calc_reduced_matrix_element(  # noqa: C901
         if not is_angular_operator_type(operator):
             raise NotImplementedError(f"calc_reduced_matrix_element is not implemented for operator {operator}.")
 
+        # Dummy matrix elements
+        if any(isinstance(s, AngularKetDummy) for s in [self, other]):
+            return 0
+
         if type(self) is not type(other):
             return self.to_state().calc_reduced_matrix_element(other.to_state(), operator, kappa)
         if is_angular_momentum_quantum_number(operator) and operator not in self.quantum_number_names:
@@ -739,6 +750,62 @@ def sanity_check(self, msgs: list[str] | None = None) -> None:
         super().sanity_check(msgs)
 
 
+class AngularKetDummy(AngularKetBase):
+    """Dummy spin ket for unknown quantum numbers."""
+
+    __slots__ = ("name",)
+    quantum_number_names: ClassVar = ("f_tot",)
+    coupled_quantum_numbers: ClassVar = {}
+    coupling_scheme = "Dummy"
+
+    name: str
+    """Name of the dummy ket."""
+
+    def __init__(
+        self,
+        name: str,
+        f_tot: float,
+        m: float | None = None,
+    ) -> None:
+        """Initialize the Spin ket."""
+        self.name = name
+
+        self.f_tot = f_tot
+        self.m = m
+
+        super()._post_init()
+
+    def sanity_check(self, msgs: list[str] | None = None) -> None:
+        """Check that the quantum numbers are valid."""
+        msgs = msgs if msgs is not None else []
+
+        if self.m is not None 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:
+            msg = "\n  ".join(msgs)
+            raise InvalidQuantumNumbersError(self, msg)
+
+    def __repr__(self) -> str:
+        args = f"{self.name}, f_tot={self.f_tot}"
+        if self.m is not None:
+            args += f", m={self.m}"
+        return f"{self.__class__.__name__}({args})"
+
+    def __str__(self) -> str:
+        return self.__repr__().replace("AngularKet", "")
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, AngularKetBase):
+            raise NotImplementedError(f"Cannot compare {self!r} with {other!r}.")
+        if not isinstance(other, AngularKetDummy):
+            return False
+        return self.name == other.name and self.f_tot == other.f_tot and self.m == other.m
+
+    def __hash__(self) -> int:
+        return hash((self.name, self.f_tot, self.m))
+
+
 def julia_qn_to_dict(qn: juliacall.AnyValue) -> dict[str, float]:
     """Convert MQDT Julia quantum numbers to dict object."""
     if "fjQuantumNumbers" in str(qn):

src/rydstate/angular/angular_state.py

@@ -8,6 +8,7 @@
 
 from rydstate.angular.angular_ket import (
     AngularKetBase,
+    AngularKetDummy,
     AngularKetFJ,
     AngularKetJJ,
     AngularKetLS,
@@ -17,7 +18,7 @@
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
 
-    from typing_extensions import Self
+    from typing_extensions import Never, Self
 
     from rydstate.angular.angular_ket import CouplingScheme
     from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers, AngularOperatorType
@@ -32,22 +33,28 @@ class AngularState(Generic[_AngularKet]):
     def __init__(
         self, coefficients: Sequence[float], kets: Sequence[_AngularKet], *, warn_if_not_normalized: bool = True
     ) -> None:
-        self.coefficients = np.array(coefficients)
-        self.kets = kets
+        """Initialize an angular state as a linear combination of angular kets.
+
+        All kets must be of the same type (coupling scheme), and no duplicate kets are allowed.
+        Dummy kets (AngularKetDummy) are ignored in the state representation,
+        however adding them is recommended for normalization purposes.
+        """
+        self._coefficients = np.array(coefficients)
+        self._kets = kets
         self._warn_if_not_normalized = warn_if_not_normalized
 
         if len(coefficients) != len(kets):
             raise ValueError("Length of coefficients and kets must be the same.")
         if len(kets) == 0:
             raise ValueError("At least one ket must be provided.")
-        if not all(type(ket) is type(kets[0]) for ket in kets):
+        if not all(type(ket) is type(self.kets[0]) for ket in self.kets):
             raise ValueError("All kets must be of the same type.")
-        if len(set(kets)) != len(kets):
-            raise ValueError("AngularState initialized with duplicate kets.")
+        if len(set(self.kets)) != len(self.kets):
+            raise ValueError("AngularState initialized with duplicate kets: %s", self.kets)
         if abs(self.norm - 1) > 1e-10 and warn_if_not_normalized:
             logger.warning("AngularState initialized with non-normalized coefficients: %s, %s", coefficients, kets)
         if self.norm > 1:
-            self.coefficients /= self.norm
+            self._coefficients /= self.norm
 
     def __iter__(self) -> Iterator[tuple[float, _AngularKet]]:
         return zip(self.coefficients, self.kets).__iter__()
@@ -60,6 +67,16 @@ def __str__(self) -> str:
         terms = [f"{coeff}*{ket!s}" for coeff, ket in self]
         return f"{', '.join(terms)}"
 
+    @property
+    def kets(self) -> list[_AngularKet]:
+        return [ket for ket in self._kets if not isinstance(ket, AngularKetDummy)]
+
+    @property
+    def coefficients(self) -> np.ndarray:
+        return np.array(
+            [coeff for coeff, ket in zip(self._coefficients, self._kets) if not isinstance(ket, AngularKetDummy)]
+        )
+
     @property
     def coupling_scheme(self) -> CouplingScheme:
         """Return the coupling scheme of the state."""
@@ -68,7 +85,7 @@ def coupling_scheme(self) -> CouplingScheme:
     @property
     def norm(self) -> float:
         """Return the norm of the state (should be 1)."""
-        return np.linalg.norm(self.coefficients)  # type: ignore [return-value]
+        return np.linalg.norm(self._coefficients)  # type: ignore [return-value]
 
     @overload
     def to(self, coupling_scheme: Literal["LS"]) -> AngularState[AngularKetLS]: ...
@@ -79,6 +96,9 @@ def to(self, coupling_scheme: Literal["JJ"]) -> AngularState[AngularKetJJ]: ...
     @overload
     def to(self, coupling_scheme: Literal["FJ"]) -> AngularState[AngularKetFJ]: ...
 
+    @overload
+    def to(self, coupling_scheme: Literal["Dummy"]) -> Never: ...
+
     def to(self, coupling_scheme: CouplingScheme) -> AngularState[Any]:
         """Convert to specified coupling scheme.