Add BasisOQDT

docs/examples/mqdt/utils.py

@@ -0,0 +1,44 @@
+import numpy as np
+
+
+def build_segments(y_lists: list[list[float]]) -> list[list[tuple[int, int]]]:
+    """Connect points of neighbouring x-columns into continuous line segments.
+
+    Between two adjacent columns each point is matched to at most one point in
+    the next column, greedily choosing the pairs with the smallest y difference
+    first. The resulting one-to-one links form simple paths (segments); a new
+    segment starts wherever a point has no match on its left.
+    """
+    cols = [np.asarray(ys, dtype=float) for ys in y_lists]
+    edges = {}  # (i, j) -> (i + 1, k): links a point to its match in the next column
+    for i in range(len(cols) - 1):
+        left, right = cols[i], cols[i + 1]
+        if left.size == 0 or right.size == 0:
+            continue
+        # |Δy| for every left/right pair at once, then visit pairs closest first
+        dist = np.abs(left[:, None] - right[None, :])
+        r = right.size
+        used_left = np.zeros(left.size, dtype=bool)
+        used_right = np.zeros(r, dtype=bool)
+        remaining = min(left.size, r)
+        for idx in np.argsort(dist, axis=None):
+            j, k = divmod(int(idx), r)
+            if used_left[j] or used_right[k]:
+                continue
+            used_left[j] = used_right[k] = True
+            edges[(i, j)] = (i + 1, k)
+            remaining -= 1
+            if remaining == 0:  # every point in the smaller column is matched
+                break
+
+    targets = set(edges.values())
+    segments = []
+    for i, ys in enumerate(cols):
+        for j in range(ys.size):
+            if (i, j) in targets:
+                continue  # not a segment start, it continues an earlier one
+            seg = [(i, j)]
+            while seg[-1] in edges:
+                seg.append(edges[seg[-1]])
+            segments.append(seg)
+    return segments

src/rydstate/__init__.py

@@ -1,10 +1,11 @@
 from rydstate import angular, basis, radial, rydberg_state, species
-from rydstate.basis import BasisMQDT, BasisSQDT
+from rydstate.basis import BasisMQDT, BasisOQDT, BasisSQDT
 from rydstate.rydberg_state import RydbergStateMQDT, RydbergStateSQDT, RydbergStateSQDTAlkali
 from rydstate.units import ureg
 
 __all__ = [
     "BasisMQDT",
+    "BasisOQDT",
     "BasisSQDT",
     "RydbergStateMQDT",
     "RydbergStateSQDT",

src/rydstate/basis/__init__.py

@@ -1,5 +1,6 @@
 from rydstate.basis.basis_base import BasisBase
 from rydstate.basis.basis_mqdt import BasisMQDT
+from rydstate.basis.basis_oqdt import BasisOQDT
 from rydstate.basis.basis_sqdt import BasisSQDT
 
-__all__ = ["BasisBase", "BasisMQDT", "BasisSQDT"]
+__all__ = ["BasisBase", "BasisMQDT", "BasisOQDT", "BasisSQDT"]

src/rydstate/basis/basis_mqdt.py

@@ -18,6 +18,7 @@
     from typing_extensions import Self
 
     from rydstate.species import FModel
+    from rydstate.units import NDArray
 
 
 logger = logging.getLogger(__name__)
@@ -37,6 +38,7 @@ def __init__(
         # potential and mqdt parameters
         mqdt: MQDT | str | None = None,
         potential_class: type[Potential] | str | None = None,
+        scale_off_diagonal: float | None = None,
     ) -> None:
         """Initialize the MQDT basis.
 
@@ -58,6 +60,7 @@ def __init__(
             potential_class: The potential class to use for the radial ket.
                 Either a a potential class
                 or a string representing the tag of the potential class to use.
+            scale_off_diagonal: If provided, scale the off-diagonal elements of the M-matrix by this factor.
 
         """
         super().__init__(species)
@@ -72,7 +75,7 @@ def __init__(
         # and for high l_r the quantum defects are 0, so n = nu
         max_l_r = int(nu[1])
         self._init_models(max_l_r, f_tot, l_r)
-        self._init_states(nu, m)
+        self._init_states(nu, m, scale_off_diagonal=scale_off_diagonal)
 
     def shallow_copy(self) -> Self:
         """Return a shallow copy of the basis (with its own independent list of states)."""
@@ -110,12 +113,16 @@ def _init_states(
         self,
         nu_range: tuple[float, float],
         m_range: tuple[float, float] | None | NotSet,
+        *,
+        scale_off_diagonal: float | None = None,
     ) -> None:
         logger.debug("Calculating MQDT states...")
         self.states = []
         for model in self.models:
             logger.debug("  calculating states for model %s with nu_range=%s", model.name, nu_range)
-            states = get_mqdt_states_from_fmodel(model, nu_range, m_range, self.potential_class)
+            states = get_mqdt_states_from_fmodel(
+                model, nu_range, m_range, self.potential_class, scale_off_diagonal=scale_off_diagonal
+            )
             if len(states) == 0:
                 logger.debug("  no states found for model %s", model.name)
             else:
@@ -131,11 +138,13 @@ def _init_states(
         self.states.sort(key=lambda state: state.nu)
 
 
-def get_mqdt_states_from_fmodel(
+def get_mqdt_states_from_fmodel(  # noqa: C901
     model: FModel,
     nu_range: tuple[float, float],
     m_range: tuple[float, float] | None | NotSet,
     potential_class: type[Potential],
+    *,
+    scale_off_diagonal: float | None = None,
 ) -> list[RydbergStateMQDT]:
     """Calculate MQDT states from an FModel by finding zeros of det(M-matrix).
 
@@ -145,6 +154,7 @@ def get_mqdt_states_from_fmodel(
         m_range: Tuple of (m_min, m_max) for the magnetic quantum number range.
             NotSet will only include states with m=NotSet.
         potential_class: The potential class to use for the radial ket.
+        scale_off_diagonal: If provided, scale the off-diagonal elements of the M-matrix by this factor.
 
     Returns:
         List of :class:`RydbergStateMQDT` objects, one per root of det(M).
@@ -155,7 +165,15 @@ def get_mqdt_states_from_fmodel(
     if np.isinf(nu_max):
         raise ValueError("nu_max must be finite to calculate MQDT states.")
 
-    nu_list = find_roots(lambda nu: np.linalg.det(model.calc_scaled_m_matrix(nu)), nu_min, nu_max)
+    def calc_scaled_off_diagonal(mmat: NDArray) -> NDArray:
+        if scale_off_diagonal is None:
+            return mmat
+        mmat_diag = np.diag(np.diag(mmat))
+        return scale_off_diagonal * (mmat - mmat_diag) + mmat_diag
+
+    nu_list = find_roots(
+        lambda nu: np.linalg.det(calc_scaled_off_diagonal(model.calc_scaled_m_matrix(nu))), nu_min, nu_max
+    )
     if len(nu_list) == 0:
         if nu_max - nu_min > 1.0:
             logger.warning(
@@ -174,8 +192,18 @@ def get_mqdt_states_from_fmodel(
 
     states: list[RydbergStateMQDT] = []
     for nu in nu_list:
+        mmat = calc_scaled_off_diagonal(model.calc_m_matrix(nu))
+        det_mmat = np.linalg.det(mmat)
+        if abs(det_mmat) > 1e-6:
+            # this can happen, because we use the scaled M-matrix to find roots ...
+            logger.warning(
+                "%s: Found a root of det(M) that is not actually a root (nu=%s, det(M)=%s). "
+                "Keeping this state, but you should treat it with caution.",
+                *(model.full_name, nu, det_mmat),
+            )
+
         nuis = model.calc_channel_nuis(nu)
-        coefficients = calc_nullvector(model.calc_scaled_m_matrix(nu))
+        coefficients = calc_nullvector(calc_scaled_off_diagonal(model.calc_scaled_m_matrix(nu)))
         coefficients = np.array(
             [coeff * (nui ** (3 / 2)) / np.cos(np.pi * nui) for coeff, nui in zip(coefficients, nuis, strict=True)]
         )
@@ -205,8 +233,10 @@ def get_mqdt_states_from_fmodel(
         for m in get_m_range(model.f_tot, m_range):
             rydberg_kets = [
                 RydbergKet(model.species, angular_ket.replace_m(m), radial_ket)
-                for angular_ket, radial_ket in zip(angular_kets_fj, radial_kets_fj, strict=True)
+                for i, (angular_ket, radial_ket) in enumerate(zip(angular_kets_fj, radial_kets_fj, strict=True))
+                if abs(coefficients_all[i]) > 1e-10
             ]
+            coefficients_all = [coeff for coeff in coefficients_all if abs(coeff) > 1e-10]
             states.append(
                 RydbergStateMQDT(
                     model.species,