WIP start adding exp for operator_ from squared_qn

Author: johannes-moegerle

src/rydstate/rydberg_state/rydberg_base.py

@@ -288,6 +288,13 @@ def calc_exp_qn(self, qn: str) -> float:
         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])
@@ -311,6 +318,9 @@ def calc_std_qn(self, qn: str) -> float:
                 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])