Various improvements

docs/examples/mqdt/mqdt_exp_qn.ipynb

@@ -68,7 +68,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "10bb4d52",
    "metadata": {},
    "outputs": [
@@ -82,8 +82,8 @@
     }
    ],
    "source": [
-    "basis.filter_states(\"j_tot\", (0, float(\"inf\")), keep_unknown=False)\n",
-    "basis.filter_states(\"f_c\", (0, float(\"inf\")), keep_unknown=False)\n",
+    "basis.filter_states(\"j_tot\", (0, float(\"inf\")))\n",
+    "basis.filter_states(\"f_c\", (0, float(\"inf\")))\n",
     "basis.filter_states(\"f_tot\", f_tot)\n",
     "basis.filter_states(\"l_r\", l_r, delta=0.5)\n",
     "basis.filter_states(\"l_c\", 0, delta=0.3)\n",

src/rydstate/angular/angular_ket.py

@@ -241,6 +241,8 @@ def __repr__(self) -> str:
         args = ", ".join(f"{qn}={val}" for qn, val in zip(self.quantum_number_names, self.quantum_numbers, strict=True))
         if not is_not_set(self.m):
             args += f", m={self.m}"
+        if is_unknown(self.l_r) or is_unknown(self.l_c):
+            args += f", parity={self.parity}"
         if self.label is not None:
             args += f", label={self.label}"
         return f"{self.__class__.__name__}({args})"

src/rydstate/angular/utils.py

@@ -2,6 +2,7 @@
 
 import contextlib
 import functools
+import numbers
 import typing as t
 from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypeVar, get_args
 
@@ -54,10 +55,33 @@ def lru_cache(maxsize: int) -> Callable[[Callable[P, R]], Callable[P, R]]:
 ]
 
 
+@functools.total_ordering
 class _Meta(type(t.Protocol)):  # type: ignore [misc]
     def __repr__(cls) -> str:
         return str(cls.__name__)
 
+    def __lt__(cls, other: object) -> bool:
+        if other is cls:
+            return False
+        if isinstance(other, numbers.Real):
+            return False  # objects derived from this metaclass always sort above every real number
+        return NotImplemented
+
+    def _binary_operation(cls, other: object) -> Any:  # noqa: ANN401
+        if other is cls or isinstance(other, numbers.Real):
+            return cls
+        return NotImplemented
+
+    __add__ = __radd__ = _binary_operation
+    __sub__ = __rsub__ = _binary_operation
+    __mul__ = __rmul__ = _binary_operation
+    __truediv__ = __rtruediv__ = _binary_operation
+
+    def _unary_operation(cls) -> Any:  # noqa: ANN401
+        return cls
+
+    __pos__ = __neg__ = _unary_operation
+
 
 @t.runtime_checkable
 class NotSet(t.Protocol, metaclass=_Meta):

src/rydstate/basis/basis_base.py

@@ -1,18 +1,19 @@
 from __future__ import annotations
 
+import copy
 from abc import ABC
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload
 
 import numpy as np
 from typing_extensions import Self
 
-from rydstate.angular.utils import is_angular_momentum_quantum_number, is_unknown
 from rydstate.rydberg_state.rydberg_base import RydbergState
 from rydstate.species import get_element_properties
 from rydstate.units import ureg
 
 if TYPE_CHECKING:
+    from rydstate.angular.utils import Unknown
     from rydstate.units import MatrixElementOperator, NDArray, PintArray, PintFloat
 
 _RydbergState = TypeVar("_RydbergState", bound=RydbergState)
@@ -28,39 +29,30 @@ def __init__(self, species: str) -> None:
     def __len__(self) -> int:
         return len(self.states)
 
-    @overload
-    def filter_states(
-        self, qn: str, value: tuple[float, float], *, delta: float = 1e-10, keep_unknown: bool = False
-    ) -> Self: ...
-
-    @overload
-    def filter_states(self, qn: str, value: float, *, delta: float = 1e-10, keep_unknown: bool = False) -> Self: ...
+    def shallow_copy(self) -> Self:
+        """Return a shallow copy of the basis (with its own independent list of states)."""
+        new_basis = copy.copy(self)
+        new_basis.states = self.states.copy()
+        return new_basis
 
-    def filter_states(
-        self, qn: str, value: float | tuple[float, float], *, delta: float = 1e-10, keep_unknown: bool = False
-    ) -> Self:
+    def filter_states(self, qn: str, value: float | Unknown | tuple[float, float], *, delta: float = 1e-10) -> Self:
         if isinstance(value, Sequence) and len(value) == 2:
             qn_min = value[0] - delta
             qn_max = value[1] + delta
         else:
-            qn_min = value - delta
-            qn_max = value + delta
-
-        if is_angular_momentum_quantum_number(qn):
-            new_states: list[_RydbergState] = []
-            for state in self.states:
-                qn_value = state.calc_exp_qn(qn)
-                if is_unknown(qn_value):
-                    if keep_unknown:
-                        new_states.append(state)
-                elif qn_min <= qn_value <= qn_max:
-                    new_states.append(state)
-            self.states = new_states
-        elif qn in ["n", "nu"]:
-            self.states = [state for state in self.states if qn_min <= getattr(state, qn) <= qn_max]
-        else:
-            raise ValueError(f"Unknown quantum number {qn}")
+            qn_min = value - delta  # type: ignore [operator]
+            qn_max = value + delta  # type: ignore [operator]
 
+        self.states = [state for state in self.states if qn_min <= state.calc_exp_qn(qn) <= qn_max]
+        return self
+
+    def filter_states_label(self, substring: str) -> Self:
+        """Filter the basis states by a substring in their label."""
+        self.states = [
+            state
+            for state in self.states
+            if any(substring in ket.angular.label for ket in state.rydberg_kets if ket.angular.label is not None)
+        ]
         return self
 
     def sort_states(self, *qns: str) -> Self:

src/rydstate/basis/basis_mqdt.py

@@ -15,6 +15,8 @@
 from rydstate.species import MQDT, FModelSQDT, Potential, get_mqdt, get_potential_class
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
     from rydstate.species import FModel
 
 
@@ -75,6 +77,12 @@ def __init__(
         self._init_models(max_l_r, f_tot, l_r, include_sqdt_fallback_models=include_sqdt_fallback_models)
         self._init_states(nu, m)
 
+    def shallow_copy(self) -> Self:
+        """Return a shallow copy of the basis (with its own independent list of states)."""
+        new_basis = super().shallow_copy()
+        new_basis.models = self.models.copy()
+        return new_basis
+
     def _init_models(
         self,
         max_l_r: int,
@@ -115,29 +123,27 @@ def _init_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)
-            if len(_states) == 0:
+            states = get_mqdt_states_from_fmodel(model, nu_range, m_range, self.potential_class)
+            if len(states) == 0:
                 logger.debug("  no states found for model %s", model.name)
             else:
                 logger.debug(
                     "  model %s: nu_min=%s, nu_max=%s, total states=%d",
                     model.name,
-                    _states[0].nu,
-                    _states[-1].nu,
-                    len(_states),
+                    states[0].nu,
+                    states[-1].nu,
+                    len(states),
                 )
-            self.states.extend(_states)
+            self.states.extend(states)
 
         self.states.sort(key=lambda state: state.nu)
 
 
-def get_mqdt_states_from_fmodel(  # noqa: C901, PLR0912
+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],
-    *,
-    overwrite_model_limits: bool = False,
 ) -> list[RydbergStateMQDT]:
     """Calculate MQDT states from an FModel by finding zeros of det(M-matrix).
 
@@ -147,24 +153,19 @@ def get_mqdt_states_from_fmodel(  # noqa: C901, PLR0912
         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.
-        overwrite_model_limits: If True, use nu_min/nu_max directly without clamping to
-            the model's validity range.
 
     Returns:
         List of :class:`RydbergStateMQDT` objects, one per root of det(M).
 
     """
-    nu_min, nu_max = nu_range
-    if overwrite_model_limits:
-        if nu_min is None or nu_max is None:
-            raise ValueError("nu_min and nu_max must be given if overwrite_model_limits is True.")
-    else:
-        nu_min = model.nu_min if nu_min is None else max(nu_min, model.nu_min)
-        nu_max = model.nu_max if nu_max is None else min(nu_max, model.nu_max)
+    nu_min = max(nu_range[0], model.nu_min)
+    nu_max = min(nu_range[1], model.nu_max)
     if np.isinf(nu_max):
         raise ValueError("nu_max must be finite to calculate MQDT states.")
+    if nu_min > nu_max:
+        return []
 
-    nu_list = find_roots(model.calc_det_scaled_m_matrix, nu_min, nu_max)
+    nu_list = find_roots(lambda nu: np.linalg.det(model.calc_scaled_m_matrix(nu)), nu_min, nu_max)
     if len(nu_list) == 0:
         logger.warning(
             "No MQDT states found in the range nu_min=%s, nu_max=%s for model %s", nu_min, nu_max, model.name
@@ -182,7 +183,8 @@ def get_mqdt_states_from_fmodel(  # noqa: C901, PLR0912
 
     states: list[RydbergStateMQDT] = []
     for nu in nu_list:
-        det_mmat = model.calc_det_m_matrix(nu)
+        mmat = 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(
@@ -192,7 +194,7 @@ def get_mqdt_states_from_fmodel(  # noqa: C901, PLR0912
             )
 
         nuis = model.calc_channel_nuis(nu)
-        coefficients = calc_nullvector(model.calc_m_matrix(nu))
+        coefficients = calc_nullvector(mmat)
         if coefficients is None:
             logger.warning("Failed to calculate nullvector for nu=%s, skipping this state.", nu)
             continue

src/rydstate/rydberg_state/rydberg_base.py

@@ -283,6 +283,8 @@ def calc_matrix_element(
     def calc_exp_qn(self, qn: str) -> float:
         if qn == "nu":
             return self.nu
+        if qn == "parity":
+            return self.parity
 
         if is_angular_momentum_quantum_number(qn):
             if qn not in self.rydberg_kets[0].angular.quantum_number_names:
@@ -296,6 +298,8 @@ def calc_exp_qn(self, qn: str) -> float:
     def calc_std_qn(self, qn: str) -> float:
         if qn == "nu":
             return 0
+        if qn == "parity":
+            return 0
 
         if is_angular_momentum_quantum_number(qn):
             if qn not in self.rydberg_kets[0].angular.quantum_number_names:

src/rydstate/species/fmodel.py

@@ -7,13 +7,15 @@
 
 import numpy as np
 
+from rydstate.angular.utils import is_not_set
 from rydstate.species.element_properties import get_element_properties
 from rydstate.species.utils import calc_energy_from_nu, calc_modified_ritz_formula_in_nu, calc_nu_from_energy
 
 if TYPE_CHECKING:
     from types import ModuleType
 
-    from rydstate.angular.angular_ket import AngularKetBase, AngularKetFJ, AngularKetJJ, AngularKetLS
+    from rydstate.angular.angular_ket import AngularKetBase, AngularKetFJ
+    from rydstate.angular.core_ket import CoreKet
     from rydstate.angular.utils import AllKnown
     from rydstate.species.mqdt import MQDT
     from rydstate.species.utils import RydbergRitzParameters
@@ -40,7 +42,7 @@ class FModel:
     inner_channels: ClassVar[list[AngularKetBase[Any]]]
     """List of inner channels in the MQDT model."""
 
-    outer_channels: ClassVar[list[AngularKetFJ[Any] | AngularKetJJ[Any] | AngularKetLS[Any]]]
+    outer_channels: ClassVar[list[AngularKetBase[Any]]]
     """List of outer channels in the MQDT model."""
 
     eigen_quantum_defects: ClassVar[list[RydbergRitzParameters]]
@@ -71,6 +73,16 @@ def nu_max(self) -> float:
         """Maximum nu for which the model is valid."""
         return self.nu_range[1]
 
+    @cached_property
+    def fj_channels(self) -> list[AngularKetFJ[Any]]:
+        """Return a list of FJ channels in the model."""
+        return [ket_fj for angular_ket in self.outer_channels for ket_fj in angular_ket.to_state("FJ").kets]
+
+    def get_core_kets(self) -> list[CoreKet]:
+        """Return a list of relevant core kets of the model."""
+        core_kets = {channel.get_core_ket() for channel in self.outer_channels}
+        return sorted(core_kets, key=lambda ket: (ket.l_c, ket.j_c, ket.f_c, str(ket.label)))
+
     @overload
     def get_ionization_thresholds(self, unit: None = None) -> list[PintFloat]: ...
 
@@ -291,30 +303,6 @@ def calc_scaled_m_matrix(self, nu: float) -> NDArray:
         nuis = self.calc_channel_nuis(nu)
         return np.array(np.diag(np.sin(np.pi * nuis)) + np.diag(np.cos(np.pi * nuis)) @ kmat)
 
-    def calc_det_m_matrix(self, nu: float) -> float:
-        """Calculate the determinant of the M-matrix at a given nu value.
-
-        Args:
-            nu: Effective principal quantum number with reference to the lowest ionization threshold.
-
-        Returns:
-            Determinant of the M-matrix at the given nu value.
-
-        """
-        return float(np.linalg.det(self.calc_m_matrix(nu)))
-
-    def calc_det_scaled_m_matrix(self, nu: float) -> float:
-        """Calculate the determinant of the scaled M-matrix at a given nu value.
-
-        Args:
-            nu: Effective principal quantum number with reference to the lowest ionization threshold.
-
-        Returns:
-            Determinant of the scaled M-matrix at the given nu value.
-
-        """
-        return float(np.linalg.det(self.calc_scaled_m_matrix(nu)))
-
 
 def get_fmodels(module: ModuleType, species: str) -> list[type[FModel]]:
     """Return all FModel subclasses defined in ``module`` that match the given species.
@@ -344,9 +332,11 @@ def get_fmodels(module: ModuleType, species: str) -> list[type[FModel]]:
 
 
 class FModelSQDT(FModel):
-    def __init__(self, species: str, channel: AngularKetFJ[AllKnown], mqdt: MQDT) -> None:
+    def __init__(self, species: str, channel: AngularKetBase[AllKnown], mqdt: MQDT) -> None:
+        if not is_not_set(channel.m):
+            raise ValueError("The m quantum number of the channel must be NotSet.")
         self.species = species  # type: ignore [misc]
-        self.name = f"SQDT l_r={channel.l_r}, j_r={channel.j_r}, f_tot={channel.f_tot}, nu > {channel.l_r + 1}"  # type: ignore [misc]
+        self.name = f"SQDT {channel}, nu >= {channel.l_r + 1}"  # type: ignore [misc]
         self.f_tot = channel.f_tot  # type: ignore [misc]
         self.nu_range = (channel.l_r + 1, math.inf)  # type: ignore [misc]
         self.inner_channels = [channel]  # type: ignore [misc]