improve angular ket

Author: johannes-moegerle

src/ryd_numerov/angular/angular_ket.py

@@ -37,11 +37,14 @@ class AngularKetBase(ABC):
     """Base class for a angular ket (i.e. a simple canonical spin ketstate)."""
 
     # We use __slots__ to prevent dynamic attributes and make the objects immutable after initialization
-    __slots__ = ("i_c", "s_c", "l_c", "s_r", "l_r", "f_tot", "m", "_initialized")
+    __slots__ = ("i_c", "s_c", "l_c", "s_r", "l_r", "f_tot", "m", "quantum_numbers", "_initialized")
 
-    spin_quantum_number_names: ClassVar[set[AngularMomentumQuantumNumbers]]
+    quantum_number_names: ClassVar[tuple[AngularMomentumQuantumNumbers, ...]]
     """Names of all well defined spin quantum numbers (without the magnetic quantum number m) in this class."""
 
+    quantum_numbers: tuple[float, ...]
+    """The quantum numbers corresponding to the quantum_number_names (without the magnetic quantum number m)."""
+
     coupled_quantum_numbers: ClassVar[
         dict[AngularMomentumQuantumNumbers, tuple[AngularMomentumQuantumNumbers, AngularMomentumQuantumNumbers]]
     ]
@@ -109,6 +112,8 @@ def __init__(
         self.m = None if m is None else float(m)
 
     def _post_init(self) -> None:
+        self.quantum_numbers = tuple(getattr(self, qn) for qn in self.quantum_number_names)
+
         self._initialized = True
 
         self.sanity_check()
@@ -138,7 +143,7 @@ def __setattr__(self, key: str, value: object) -> None:
         super().__setattr__(key, value)
 
     def __repr__(self) -> str:
-        args = ", ".join(f"{k}={v}" for k, v in self.spin_quantum_numbers_dict.items())
+        args = ", ".join(f"{qn}={val}" for qn, val in zip(self.quantum_number_names, self.quantum_numbers))
         if self.m is not None:
             args += f", m={self.m}"
         return f"{self.__class__.__name__}({args})"
@@ -153,24 +158,20 @@ def __eq__(self, other: object) -> bool:
             return False
         if self.m != other.m:
             return False
-        return all(self.get_qn(q) == other.get_qn(q) for q in self.spin_quantum_numbers_dict)
+        return self.quantum_numbers == other.quantum_numbers
 
     def __hash__(self) -> int:
         return hash(
             (
-                tuple((k, v) for k, v in self.spin_quantum_numbers_dict.items()),
+                self.quantum_number_names,
+                self.quantum_numbers,
                 self.m,
             )
         )
 
-    @property
-    def spin_quantum_numbers_dict(self) -> dict[AngularMomentumQuantumNumbers, float | int]:
-        """Return the spin quantum numbers (i.e. without the magnetic quantum number) as dictionary."""
-        return {q: getattr(self, q) for q in self.spin_quantum_number_names}
-
     def get_qn(self, qn: AngularMomentumQuantumNumbers) -> float:
         """Get the value of a quantum number by name."""
-        if qn not in self.spin_quantum_number_names:
+        if qn not in self.quantum_number_names:
             raise ValueError(f"Quantum number {qn} not found in {self!r}.")
         return getattr(self, qn)  # type: ignore [no-any-return]
 
@@ -205,7 +206,7 @@ def calc_reduced_overlap(self, other: AngularKetBase) -> float:
         If the kets are of different types, the overlap is calculated using the corresponding
         Clebsch-Gordan coefficients (/ Wigner-j symbols).
         """
-        for q in self.spin_quantum_number_names & other.spin_quantum_number_names:
+        for q in set(self.quantum_number_names) & set(other.quantum_number_names):
             if self.get_qn(q) != other.get_qn(q):
                 return 0
 
@@ -258,14 +259,14 @@ def calc_reduced_matrix_element(  # noqa: C901
 
         if type(self) is not type(other):
             return self.to_state().calc_reduced_matrix_element(other.to_state(), operator, kappa)
-        if operator in get_args(AngularMomentumQuantumNumbers) and operator not in self.spin_quantum_number_names:
+        if operator in get_args(AngularMomentumQuantumNumbers) and operator not in self.quantum_number_names:
             return self.to_state().calc_reduced_matrix_element(other.to_state(), operator, kappa)
 
         qn_name: AngularMomentumQuantumNumbers
         if operator == "SPHERICAL":
             qn_name = "l_r"
             complete_reduced_matrix_element = calc_reduced_spherical_matrix_element(self.l_r, other.l_r, kappa)
-        elif operator in self.spin_quantum_number_names:
+        elif operator in self.quantum_number_names:
             if not kappa == 1:
                 raise ValueError("Only kappa=1 is supported for spin operators.")
             qn_name = operator  # type: ignore [assignment]
@@ -334,7 +335,7 @@ def _kronecker_delta_non_involved_spins(self, other: AngularKetBase, qn: Angular
         This means return 0 if any of the quantum numbers,
         that are not qn or a coupled quantum number resulting from qn differ between self and other.
         """
-        if qn not in self.spin_quantum_number_names:
+        if qn not in self.quantum_number_names:
             raise ValueError(f"Quantum number {qn} is not a valid angular momentum quantum number for {self!r}.")
 
         resulting_qns = {qn}
@@ -350,7 +351,7 @@ def _kronecker_delta_non_involved_spins(self, other: AngularKetBase, qn: Angular
                     f"_kronecker_delta_non_involved_spins: {last_qn} not found in coupled_quantum_numbers."
                 )
 
-        non_involved_qns = self.spin_quantum_number_names - resulting_qns
+        non_involved_qns = set(self.quantum_number_names) - resulting_qns
         for _qn in non_involved_qns:
             if self.get_qn(_qn) != other.get_qn(_qn):
                 return 0
@@ -396,7 +397,7 @@ class AngularKetLS(AngularKetBase):
     """Spin ket in LS coupling."""
 
     __slots__ = ("s_tot", "l_tot", "j_tot")
-    spin_quantum_number_names: ClassVar = {"i_c", "s_c", "l_c", "s_r", "l_r", "s_tot", "l_tot", "j_tot", "f_tot"}
+    quantum_number_names: ClassVar = ("i_c", "s_c", "l_c", "s_r", "l_r", "s_tot", "l_tot", "j_tot", "f_tot")
     coupled_quantum_numbers: ClassVar = {
         "s_tot": ("s_c", "s_r"),
         "l_tot": ("l_c", "l_r"),
@@ -518,7 +519,7 @@ class AngularKetJJ(AngularKetBase):
     """Spin ket in JJ coupling."""
 
     __slots__ = ("j_c", "j_r", "j_tot")
-    spin_quantum_number_names: ClassVar = {"i_c", "s_c", "l_c", "s_r", "l_r", "j_c", "j_r", "j_tot", "f_tot"}
+    quantum_number_names: ClassVar = ("i_c", "s_c", "l_c", "s_r", "l_r", "j_c", "j_r", "j_tot", "f_tot")
     coupled_quantum_numbers: ClassVar = {
         "j_c": ("s_c", "l_c"),
         "j_r": ("s_r", "l_r"),
@@ -652,7 +653,7 @@ class AngularKetFJ(AngularKetBase):
     """Spin ket in FJ coupling."""
 
     __slots__ = ("j_c", "f_c", "j_r")
-    spin_quantum_number_names: ClassVar = {"i_c", "s_c", "l_c", "s_r", "l_r", "j_c", "f_c", "j_r", "f_tot"}
+    quantum_number_names: ClassVar = ("i_c", "s_c", "l_c", "s_r", "l_r", "j_c", "f_c", "j_r", "f_tot")
     coupled_quantum_numbers: ClassVar = {
         "j_c": ("s_c", "l_c"),
         "f_c": ("i_c", "j_c"),

src/ryd_numerov/angular/angular_state.py

@@ -97,9 +97,9 @@ def calc_exp_qn(self, q: AngularMomentumQuantumNumbers) -> float:
             q: The quantum number to calculate the expectation value for.
 
         """
-        if q not in self.kets[0].spin_quantum_number_names:
+        if q not in self.kets[0].quantum_number_names:
             for ket_class in [AngularKetLS, AngularKetJJ, AngularKetFJ]:
-                if q in ket_class.spin_quantum_number_names:
+                if q in ket_class.quantum_number_names:
                     return self._to_coupling_scheme(ket_class.coupling_scheme).calc_exp_qn(q)
 
         qs = np.array([ket.get_qn(q) for ket in self.kets])
@@ -115,9 +115,9 @@ def calc_std_qn(self, q: AngularMomentumQuantumNumbers) -> float:
             q: The quantum number to calculate the standard deviation for.
 
         """
-        if q not in self.kets[0].spin_quantum_number_names:
+        if q not in self.kets[0].quantum_number_names:
             for ket_class in [AngularKetLS, AngularKetJJ, AngularKetFJ]:
-                if q in ket_class.spin_quantum_number_names:
+                if q in ket_class.quantum_number_names:
                     return self._to_coupling_scheme(ket_class.coupling_scheme).calc_std_qn(q)
 
         qs = np.array([ket.get_qn(q) for ket in self.kets])
@@ -156,12 +156,9 @@ def calc_reduced_matrix_element(
         """
         if isinstance(other, AngularKetBase):
             other = other.to_state()
-        if (
-            operator in get_args(AngularMomentumQuantumNumbers)
-            and operator not in self.kets[0].spin_quantum_number_names
-        ):
+        if operator in get_args(AngularMomentumQuantumNumbers) and operator not in self.kets[0].quantum_number_names:
             for ket_class in [AngularKetLS, AngularKetJJ, AngularKetFJ]:
-                if operator in ket_class.spin_quantum_number_names:
+                if operator in ket_class.quantum_number_names:
                     return self._to_coupling_scheme(ket_class.coupling_scheme).calc_reduced_matrix_element(
                         other, operator, kappa
                     )

tests/test_angular_matrix_elements.py

@@ -79,13 +79,14 @@ def test_reduced_identity(ket: AngularKetBase) -> None:
     state_jj = ket.to_jj()
     state_fj = ket.to_fj()
 
-    for op in state_ls.kets[0].spin_quantum_number_names:
+    op: AngularMomentumQuantumNumbers
+    for op in state_ls.kets[0].quantum_number_names:
         assert np.isclose(reduced_identity, state_ls.calc_reduced_matrix_element(state_ls, "identity_" + op, kappa=0))  # type: ignore [arg-type]
 
-    for op in state_jj.kets[0].spin_quantum_number_names:
+    for op in state_jj.kets[0].quantum_number_names:
         assert np.isclose(reduced_identity, state_jj.calc_reduced_matrix_element(state_jj, "identity_" + op, kappa=0))  # type: ignore [arg-type]
 
-    for op in state_fj.kets[0].spin_quantum_number_names:
+    for op in state_fj.kets[0].quantum_number_names:
         assert np.isclose(reduced_identity, state_fj.calc_reduced_matrix_element(state_fj, "identity_" + op, kappa=0))  # type: ignore [arg-type]