More angular matrix elements

src/rydstate/angular/angular_ket.py

@@ -7,9 +7,10 @@
 
 from rydstate.angular.angular_matrix_element import (
     calc_prefactor_of_operator_in_coupled_scheme,
-    calc_reduced_identity_matrix_element,
+    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_") and kappa != 0:
-            raise ValueError("Only kappa=0 is supported for identity 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):
@@ -675,7 +676,17 @@ def _calc_reduced_matrix_element(  # noqa: C901
         elif is_angular_momentum_quantum_number(operator):
             complete_reduced_matrix_element = calc_reduced_spin_matrix_element(qn_self, qn_other)
         elif operator.startswith("identity_"):
-            complete_reduced_matrix_element = calc_reduced_identity_matrix_element(qn_self, qn_other)
+            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,
+            # so <state| raw_value_x |state> yields the raw value of x (and raw_value_x_2 yields x^2).
+            exponent = 2 if operator.endswith("_2") else 1
+            complete_reduced_matrix_element = calc_reduced_raw_quantum_number_matrix_element(
+                qn_self, qn_other, exponent
+            )
         else:
             raise NotImplementedError(f"calc_reduced_matrix_element is not implemented for operator {operator}.")
 

src/rydstate/angular/angular_matrix_element.py

@@ -67,9 +67,36 @@ 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_identity_matrix_element(s_final: float, s_initial: float) -> float:
-    r"""Calculate the reduced identity matrix element (s_final || \id || s_initial).
+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).
+
+    For exponent = 0, this is the reduced matrix element of the identity operator (s_final || \id || s_initial).
 
     We follow the convention from Edmonds 1985 "Angular Momentum in Quantum Mechanics"
     (using equation (5.4.1) and (3.7.9)).
@@ -82,14 +109,17 @@ def calc_reduced_identity_matrix_element(s_final: float, s_initial: float) -> fl
     Args:
         s_final: The spin quantum number of the final state.
         s_initial: The spin quantum number of the initial state.
+        exponent: The exponent of the quantum number (0 for the identity operator).
 
     Returns:
-        The reduced matrix element :math:`(s_final || \id || s_initial)`.
+        The reduced matrix element :math:`(s_final || s_final^exponent || s_initial)`.
 
     """
     if s_final != s_initial:
         return 0
-    return math.sqrt(2 * s_final + 1)
+    if exponent == 0:
+        return math.sqrt(2 * s_final + 1)
+    return math.sqrt(2 * s_final + 1) * s_final**exponent
 
 
 @lru_cache(maxsize=100_000)

src/rydstate/angular/utils.py

@@ -45,12 +45,57 @@ 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",
+    "raw_value_l_c",
+    "raw_value_s_r",
+    "raw_value_l_r",
+    "raw_value_s_tot",
+    "raw_value_l_tot",
+    "raw_value_j_c",
+    "raw_value_j_r",
+    "raw_value_j_tot",
+    "raw_value_f_c",
+    "raw_value_f_tot",
+]
+RawValueOperators2 = Literal[
+    "raw_value_i_c_2",
+    "raw_value_s_c_2",
+    "raw_value_l_c_2",
+    "raw_value_s_r_2",
+    "raw_value_l_r_2",
+    "raw_value_s_tot_2",
+    "raw_value_l_tot_2",
+    "raw_value_j_c_2",
+    "raw_value_j_r_2",
+    "raw_value_j_tot_2",
+    "raw_value_f_c_2",
+    "raw_value_f_tot_2",
+]
 
 AngularOperatorType = Literal[
     "spherical",
     "spherical_core",
     AngularMomentumQuantumNumbers,
     IdentityOperators,
+    SquaredSpinOperators,
+    RawValueOperators,
+    RawValueOperators2,
 ]
 
 
@@ -208,6 +253,11 @@ 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")
 
     if not is_angular_momentum_quantum_number(qn):
         raise ValueError(f"Invalid operator {operator}.")

src/rydstate/rydberg_state/rydberg_base.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import math
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, overload
 
@@ -284,6 +285,16 @@ def calc_exp_qn(self, qn: str) -> float:
         if qn == "nu":
             return self.nu
 
+        if qn.startswith("raw_value_"):
+            return self.calc_matrix_element(self, qn, q=0, unit="a.u.")  # type: ignore [call-overload,no-any-return]
+
+        if qn.startswith("operator_"):
+            qn_name = qn[len("operator_") :]
+            exp_q2 = self.calc_matrix_element(self, "squared_" + qn_name, q=0, unit="a.u.")  # type: ignore [call-overload]
+            # exp_q2 returns the reduced matrix element of the operator \hat{S}^2
+            # (which for the concrete (non reduced) matrix element gives -> S(S+1))
+            return -0.5 + math.sqrt(1 / 4 + exp_q2)
+
         if is_angular_momentum_quantum_number(qn):
             if qn not in self.rydberg_kets[0].angular.quantum_number_names:
                 coupling_scheme = get_coupling_scheme_for_quantum_number(qn, [self.coupling_scheme])
@@ -297,6 +308,19 @@ def calc_std_qn(self, qn: str) -> float:
         if qn == "nu":
             return 0
 
+        if qn.startswith("raw_value_"):
+            exp_q = self.calc_matrix_element(self, qn, q=0, unit="a.u.")  # type: ignore [call-overload]
+            exp_q2 = self.calc_matrix_element(self, qn + "_2", q=0, unit="a.u.")  # type: ignore [call-overload]
+            if abs(exp_q2 - exp_q**2) < 1e-10:
+                return 0
+            if exp_q2 - exp_q**2 < 0:
+                logger.warning("Got negative variance for quantum number %s: %.3e. Returning 0.", qn, exp_q2 - exp_q**2)
+                return 0
+            return math.sqrt(exp_q2 - exp_q**2)
+
+        if qn.startswith("operator_"):
+            raise NotImplementedError("Standard deviation for operator quantum numbers is not implemented.")
+
         if is_angular_momentum_quantum_number(qn):
             if qn not in self.rydberg_kets[0].angular.quantum_number_names:
                 coupling_scheme = get_coupling_scheme_for_quantum_number(qn, [self.coupling_scheme])

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_"):
+            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

@@ -124,6 +124,67 @@ def test_scalar_matrix_element_independent_of_m(ket: AngularKetBase[AllKnown]) -
             assert np.isclose(val_m, val_not_set), f"{operator=}, {m=}, {val_m=}, {val_not_set=}"
 
 
+@pytest.mark.parametrize("ket", TEST_KETS)
+def test_reduced_raw_value(ket: AngularKetBase[AllKnown]) -> None:
+    # In its native coupling scheme every quantum number is definite, so raw_value_x = x^exponent * identity and its
+    # reduced matrix element is x^exponent * sqrt(2 * f_tot + 1). This pins down the Wigner-Eckart normalization: the
+    # full matrix element <ket| raw_value_x |ket> then evaluates to the raw value of x (and raw_value_x_2 to x^2).
+    reduced_identity = np.sqrt(2 * ket.f_tot + 1)
+
+    op: AngularMomentumQuantumNumbers
+    for op in ket.quantum_number_names:
+        raw = ket.get_qn(op)
+        assert np.isclose(raw * reduced_identity, ket.calc_reduced_matrix_element(ket, "raw_value_" + op, kappa=0))  # type: ignore [arg-type]
+        assert np.isclose(
+            raw**2 * reduced_identity,
+            ket.calc_reduced_matrix_element(ket, "raw_value_" + op + "_2", kappa=0),  # type: ignore [arg-type]
+        )
+        for m in np.arange(-ket.f_tot, ket.f_tot + 1):
+            ket_m = ket.replace_m(m)
+            assert np.isclose(raw, ket_m.calc_matrix_element(ket_m, "raw_value_" + op, kappa=0, q=0))  # type: ignore [arg-type]
+            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