add squared angular operator

Author: johannes-moegerle

src/rydstate/angular/angular_ket.py

@@ -10,6 +10,7 @@
     calc_reduced_raw_quantum_number_matrix_element,
     calc_reduced_spherical_matrix_element,
     calc_reduced_spin_matrix_element,
+    calc_reduced_spin_squared_matrix_element,
 )
 from rydstate.angular.core_ket import CoreKet
 from rydstate.angular.utils import (
@@ -646,7 +647,7 @@ def calc_reduced_matrix_element(
             cache[cache_key] = self._calc_reduced_matrix_element(other, operator, kappa)
         return cache[cache_key]
 
-    def _calc_reduced_matrix_element(  # noqa: C901
+    def _calc_reduced_matrix_element(  # noqa: C901, PLR0912
         self: Self, other: AngularKetBase[Any], operator: AngularOperatorType, kappa: int
     ) -> float:
         if not is_angular_operator_type(operator):
@@ -663,8 +664,8 @@ def _calc_reduced_matrix_element(  # noqa: C901
 
         if is_angular_momentum_quantum_number(operator) and kappa != 1:
             raise ValueError("Only kappa=1 is supported for spin operators.")
-        if operator.startswith(("identity_", "raw_value_")) and kappa != 0:
-            raise ValueError("Only kappa=0 is supported for identity/raw_value operators.")
+        if operator.startswith(("identity_", "raw_value_", "squared_")) and kappa != 0:
+            raise ValueError("Only kappa=0 is supported for identity/raw_value/squared operators.")
 
         qn_self, qn_other = self.get_qn(qn_name), other.get_qn(qn_name)
         if is_unknown(qn_self) or is_unknown(qn_other):
@@ -676,6 +677,8 @@ def _calc_reduced_matrix_element(  # noqa: C901
             complete_reduced_matrix_element = calc_reduced_spin_matrix_element(qn_self, qn_other)
         elif operator.startswith("identity_"):
             complete_reduced_matrix_element = calc_reduced_raw_quantum_number_matrix_element(qn_self, qn_other, 0)
+        elif operator.startswith("squared_"):
+            complete_reduced_matrix_element = calc_reduced_spin_squared_matrix_element(qn_self, qn_other)
         elif operator.startswith("raw_value_"):
             # raw_value_x is the diagonal scalar operator x^exponent; reusing the same reduced matrix element as
             # identity (exponent=0) ensures the coupled-scheme prefactor and Wigner-Eckart normalization are applied,

src/rydstate/angular/angular_matrix_element.py

@@ -67,6 +67,31 @@ def calc_reduced_spin_matrix_element(s_final: float, s_initial: float) -> float:
     return math.sqrt((2 * s_final + 1) * (s_final + 1) * s_final)
 
 
+def calc_reduced_spin_squared_matrix_element(s_final: float, s_initial: float) -> float:
+    r"""Calculate the reduced squared spin matrix element (s_final || \hat{s}^2 || s_initial).
+
+    The squared spin operator \hat{s}^2 = \hat{s} \cdot \hat{s} is a scalar (rank kappa=0) operator.
+    Within a subspace of fixed spin quantum number s it is proportional to the identity,
+    \hat{s}^2 = s(s+1) \id, so with the convention from Edmonds 1985 "Angular Momentum in Quantum Mechanics"
+    (using equation (5.4.1) and (3.7.9)) its reduced matrix element is given by:
+
+    .. math::
+        (s_final || \hat{s}^2 || s_initial)
+            = s_initial (s_initial + 1) \sqrt{2 * s_initial + 1} * \delta_{s_final, s_initial}
+
+    Args:
+        s_final: The spin quantum number of the final state.
+        s_initial: The spin quantum number of the initial state.
+
+    Returns:
+        The reduced matrix element :math:`(s_final || \hat{s}^2 || s_initial)`.
+
+    """
+    if s_final != s_initial:
+        return 0
+    return s_initial * (s_initial + 1) * math.sqrt(2 * s_initial + 1)
+
+
 @lru_cache(maxsize=1_000)
 def calc_reduced_raw_quantum_number_matrix_element(s_final: float, s_initial: float, exponent: int) -> float:
     r"""Calculate the reduced matrix element (s_final || s_final^exponent || s_initial).

src/rydstate/angular/utils.py

@@ -45,6 +45,20 @@ def lru_cache(maxsize: int) -> Callable[[Callable[P, R]], Callable[P, R]]:
     "identity_f_c",
     "identity_f_tot",
 ]
+SquaredSpinOperators = Literal[
+    "squared_i_c",
+    "squared_s_c",
+    "squared_l_c",
+    "squared_s_r",
+    "squared_l_r",
+    "squared_s_tot",
+    "squared_l_tot",
+    "squared_j_c",
+    "squared_j_r",
+    "squared_j_tot",
+    "squared_f_c",
+    "squared_f_tot",
+]
 RawValueOperators = Literal[
     "raw_value_i_c",
     "raw_value_s_c",
@@ -79,6 +93,7 @@ def lru_cache(maxsize: int) -> Callable[[Callable[P, R]], Callable[P, R]]:
     "spherical_core",
     AngularMomentumQuantumNumbers,
     IdentityOperators,
+    SquaredSpinOperators,
     RawValueOperators,
     RawValueOperators2,
 ]
@@ -238,6 +253,8 @@ def get_qn_name_from_operator(operator: AngularOperatorType) -> AngularMomentumQ
         qn = "l_c"
     elif operator.startswith("identity_"):
         qn = operator.removeprefix("identity_")
+    elif operator.startswith("squared_"):
+        qn = operator.removeprefix("squared_")
     elif operator.startswith("raw_value_"):
         qn = operator.removeprefix("raw_value_")
         qn = qn.removesuffix("_2")

src/rydstate/rydberg_state/rydberg_ket.py

@@ -254,7 +254,7 @@ def _get_ks(self, operator: MatrixElementOperator) -> tuple[int, int]:
             return MatrixElementOperatorRanks[operator]
         if is_angular_operator_type(operator):
             k_radial = 0
-            if operator.startswith(("identity_", "raw_value_")):
+            if operator.startswith(("identity_", "raw_value_", "squared_")):
                 return k_radial, 0
             if is_angular_momentum_quantum_number(operator):
                 return k_radial, 1

tests/test_angular_matrix_elements.py

@@ -145,6 +145,46 @@ def test_reduced_raw_value(ket: AngularKetBase[AllKnown]) -> None:
             assert np.isclose(raw**2, ket_m.calc_matrix_element(ket_m, "raw_value_" + op + "_2", kappa=0, q=0))  # type: ignore [arg-type]
 
 
+@pytest.mark.parametrize("ket", TEST_KETS)
+def test_reduced_spin_squared(ket: AngularKetBase[AllKnown]) -> None:
+    op: AngularMomentumQuantumNumbers
+    coupling_schemes: list[CouplingScheme] = ["LS", "JJ", "FJ"]
+    for scheme in coupling_schemes:
+        state = ket.to_state(scheme)
+        for op in state.kets[0].quantum_number_names:
+            # the squared spin operator is diagonal in the scheme's own basis with eigenvalue qn * (qn + 1),
+            # so for a superposition the reduced matrix element is the weighted average over the components
+            exp_squared = sum(coeff**2 * s_ket.get_qn(op) * (s_ket.get_qn(op) + 1) for coeff, s_ket in state)
+            reduced_squared = exp_squared * np.sqrt(2 * ket.f_tot + 1)
+            assert np.isclose(reduced_squared, state.calc_reduced_matrix_element(state, "squared_" + op, kappa=0))  # type: ignore [arg-type]
+
+
+def test_spin_squared_expectation_value() -> None:
+    """The expectation value of s^2 must be s(s+1) for good quantum numbers and match the sum rule otherwise."""
+    ket = AngularKetFJ(l_r=0, j_r=0.5, f_c=1, m=0.5, f_tot=0.5, species="Yb171")
+
+    # s_r and s_c are good quantum numbers in every ket, so <s^2> = s(s+1) = 3/4
+    assert np.isclose(ket.calc_matrix_element(ket, "squared_s_r", 0, q=0), 0.75)
+    assert np.isclose(ket.calc_matrix_element(ket, "squared_s_c", 0, q=0), 0.75)
+
+    # s_tot is not a good quantum number of an FJ ket, so <s_tot^2> is a weighted average
+    # over the LS decomposition: sum_i |c_i|^2 * s_tot_i * (s_tot_i + 1)
+    expected = sum(coeff**2 * ls_ket.s_tot * (ls_ket.s_tot + 1) for coeff, ls_ket in ket.to_state("LS"))
+    assert np.isclose(ket.calc_matrix_element(ket, "squared_s_tot", 0, q=0), expected)
+
+    # <s^2> must also equal the sum over all final states and components q of the squared
+    # matrix elements of the (rank-1) spin operator: <s^2> = sum_{f,q} |<f|s_q|ket>|^2
+    finals = [
+        AngularKetFJ(l_r=0, j_r=0.5, f_c=f_c, m=m / 2, f_tot=f_tot, species="Yb171")
+        for f_c in (0, 1)
+        for f_tot in {abs(f_c - 0.5), f_c + 0.5}
+        for m in range(-int(2 * f_tot), int(2 * f_tot) + 1, 2)
+    ]
+    for op in ("s_r", "s_c"):
+        sum_rule = sum(f.calc_matrix_element(ket, op, 1, q=q) ** 2 for f in finals for q in (-1, 0, 1))
+        assert np.isclose(sum_rule, ket.calc_matrix_element(ket, "squared_" + op, 0, q=0))  # type: ignore [arg-type]
+
+
 @pytest.mark.parametrize(("ket1", "ket2"), TEST_KET_PAIRS)
 def test_matrix_elements_in_different_coupling_schemes(
     ket1: AngularKetBase[AllKnown], ket2: AngularKetBase[AllKnown], coupling_scheme: CouplingScheme