RydbergStateMQDT add principal quantum number n

Author: johannes-moegerle

src/rydstate/generate_database/generate_states_table.py

@@ -5,8 +5,6 @@
 
 import numpy as np
 
-from rydstate.rydberg_state.rydberg_sqdt import RydbergStateSQDT
-
 if TYPE_CHECKING:
     from rydstate.basis import BasisMQDT, BasisSQDT
     from rydstate.rydberg_state.rydberg_base import RydbergStateBase
@@ -63,13 +61,11 @@ def get_state_data(ids: int, state: RydbergStateBase) -> tuple[float | int | str
     angular = state.angular
     underspecified_channel_contribution = sum(abs(coeff) ** 2 for coeff, ket in state if ket.angular.contains_unknown)
 
-    n = state.n if isinstance(state, RydbergStateSQDT) else 0
-
     data = (
         ids,  # id
         state.get_energy("a.u."),  # energy
         angular.parity,  # parity = (-1)^l_tot
-        n,  # n: quantum number
+        state.n,  # n
         state.nu,  # nu
         angular.f_tot,  # f_tot
         state.calc_exp_qn("nui"),  # exp_nui

src/rydstate/rydberg_state/rydberg_base.py

@@ -45,6 +45,13 @@ class RydbergStateBase(ABC):
     _energy_au: float
     """The energy of the Rydberg state in atomic units (Hartree)."""
 
+    n: int
+    """The principal quantum number n of the Rydberg state.
+
+    For MQDT states, we define the corresponding principal quantum number n via the number of nodes
+    in the radial wavefunction of the most dominant channel.
+    """
+
     def __init__(self) -> None:
         if abs(self.norm - 1) > 1e-10:
             raise ValueError(
@@ -209,10 +216,7 @@ def calc_exp_qn(self, qn: str) -> float:
         if qn == "nu":
             return self.nu
         if qn == "n":
-            n = getattr(self, "n", None)
-            if n is None:
-                raise ValueError(f"{self} has no quantum number n")
-            return n  # type: ignore [no-any-return]
+            return self.n
         if qn == "nui":
             return float(sum([abs(coeff) ** 2 * ket.radial.nu / self.norm**2 for coeff, ket in self]))
         raise ValueError(f"Unknown quantum number {qn}")

src/rydstate/rydberg_state/rydberg_mqdt.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import logging
+from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 from rydstate.angular import AngularKetFJ, AngularState
+from rydstate.angular.utils import is_unknown
 from rydstate.rydberg_state.rydberg_base import RydbergStateBase
 
 if TYPE_CHECKING:
@@ -67,3 +69,23 @@ def __str__(self) -> str:
     def mqdt(self) -> MQDT:
         """Return the MQDT object used to calculate this state."""
         return self.model.mqdt
+
+    @cached_property
+    def n(self) -> int:  # type: ignore [override]
+        """Return the corresponding principal quantum number n of the state.
+
+        We define the corresponding principal quantum number n for MQDT states via the nodes of
+        the main contributing rydberg ket (nodes = n - l_r - 1).
+        For FModelSQDT states, the quantum defect is zero, so the channel dependent effective quantum number nui
+        is already an integer and we simply round it to the nearest integer.
+        """
+        defects = self.model.eigen_quantum_defects
+        if (
+            len(defects) == 1 and np.isscalar(defects[0]) and abs(defects[0]) < 1e-10  # type: ignore [arg-type]
+        ):
+            return round(self.nui[0])
+
+        main_ket = max(
+            [(coeff, ket) for coeff, ket in self if not is_unknown(ket.angular.l_r)], key=lambda x: abs(x[0])
+        )[1]
+        return main_ket.radial.nodes + main_ket.angular.l_r + 1

tests/test_rydberg_state_mqdt.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+import itertools
 from typing import TYPE_CHECKING
 
 import numpy as np
 import pytest
 from rydstate import BasisMQDT, RydbergStateSQDT
+from rydstate.angular.utils import is_unknown
+from rydstate.species import FModelSQDT
 
 if TYPE_CHECKING:
     from rydstate import RydbergStateMQDT
@@ -134,3 +137,58 @@ def test_matrix_element_between_mqdt_and_sqdt_state(basis: BasisMQDT) -> None:
     me = s_mqdt.calc_reduced_matrix_element(p_sqdt, "electric_dipole", unit="e a0")
     assert np.isfinite(me)
     assert abs(me) > 0.0
+
+
+def test_n_of_sqdt_fallback_model_uses_channel_nui() -> None:
+    """For SQDT fallback models n is given by the channel nui, not by nu."""
+    basis = BasisMQDT("Yb171", nu=(78.0, 81.0), l_r=(5, 5), m=(0.5, 0.5))
+    assert len(basis.states) > 0
+
+    states_with_shifted_nu = 0
+    for state in basis.states:
+        assert isinstance(state.model, FModelSQDT)
+        assert len(state.rydberg_kets) == 1
+
+        nui = state.nui[0]
+        assert abs(nui - round(nui)) < 1e-9
+        assert state.n == round(nui)
+
+        if state.n != round(state.nu):
+            states_with_shifted_nu += 1
+
+    assert states_with_shifted_nu > 0
+
+
+@pytest.mark.parametrize(
+    ("species", "m", "l_r", "model_name"),
+    [
+        ("Yb174", (0, 0), (0, 0), "S J=0, nu > 2"),
+        ("Yb174", (0, 0), (2, 2), "D J=2, nu > 5"),
+        ("Yb171", (0.5, 0.5), (1, 1), "P F=1/2, nu > 5.7"),
+        ("Yb171", (0.5, 0.5), (1, 1), "P F=3/2, nu > 10"),
+    ],
+)
+def test_n_increments_by_one_along_a_rydberg_series(
+    species: str, m: tuple[float, float], l_r: tuple[int, int], model_name: str
+) -> None:
+    """Within one Rydberg series of an MQDT model, n increases by exactly one from state to state.
+
+    A single MQDT model describes several interleaved Rydberg series (one per channel), so the states
+    of a model have to be grouped by their dominant channel before comparing consecutive n.
+    """
+    basis = BasisMQDT(species, nu=(40.0, 46.0), l_r=l_r, m=m)
+    states = [state for state in basis.states if state.model.name == model_name]
+    assert len(states) > 0
+
+    def dominant_channel(state: RydbergStateMQDT) -> int:
+        """Return the index of the channel with the largest coefficient (ignoring unknown channels)."""
+        known = [(abs(coeff), i) for i, (coeff, ket) in enumerate(state) if not is_unknown(ket.angular.l_r)]
+        return max(known)[1]
+
+    states.sort(key=lambda state: (dominant_channel(state), state.nu))
+    for channel, series in itertools.groupby(states, key=dominant_channel):
+        n_list = [state.n for state in series]
+        assert len(n_list) > 1, f"series of channel {channel} is too short to test"
+        assert all(n_next - n == 1 for n, n_next in itertools.pairwise(n_list)), (
+            f"n does not increase by one along the series of channel {channel}: {n_list}"
+        )