diff --git a/src/rydstate/angular/angular_ket.py b/src/rydstate/angular/angular_ket.py index 446d3f96..8bff75ec 100644 --- a/src/rydstate/angular/angular_ket.py +++ b/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 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}.") diff --git a/src/rydstate/angular/angular_matrix_element.py b/src/rydstate/angular/angular_matrix_element.py index 56356505..327af9dd 100644 --- a/src/rydstate/angular/angular_matrix_element.py +++ b/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) diff --git a/src/rydstate/angular/utils.py b/src/rydstate/angular/utils.py index 7aeb473e..cdf1a980 100644 --- a/src/rydstate/angular/utils.py +++ b/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}.") diff --git a/src/rydstate/rydberg_state/rydberg_base.py b/src/rydstate/rydberg_state/rydberg_base.py index 57ebf3de..cd7d9ab7 100644 --- a/src/rydstate/rydberg_state/rydberg_base.py +++ b/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]) diff --git a/src/rydstate/rydberg_state/rydberg_ket.py b/src/rydstate/rydberg_state/rydberg_ket.py index 0db0dee7..501b4c3b 100644 --- a/src/rydstate/rydberg_state/rydberg_ket.py +++ b/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 diff --git a/tests/test_angular_matrix_elements.py b/tests/test_angular_matrix_elements.py index cfb3b24b..5ef952d7 100644 --- a/tests/test_angular_matrix_elements.py +++ b/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 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(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 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) + + # must also equal the sum over all final states and components q of the squared + # matrix elements of the (rank-1) spin operator: = sum_{f,q} ||^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