From 69024e546dfa6f0b495df0893885389d7856bf70 Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Tue, 30 Jun 2026 10:04:25 +0200 Subject: [PATCH 1/9] add Basis shallow_copy --- src/rydstate/basis/basis_base.py | 7 +++++++ src/rydstate/basis/basis_mqdt.py | 8 ++++++++ 2 files changed, 15 insertions(+) diff --git a/src/rydstate/basis/basis_base.py b/src/rydstate/basis/basis_base.py index 87e9ec9d..69763fbc 100644 --- a/src/rydstate/basis/basis_base.py +++ b/src/rydstate/basis/basis_base.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy from abc import ABC from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload @@ -28,6 +29,12 @@ def __init__(self, species: str) -> None: def __len__(self) -> int: return len(self.states) + 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 + @overload def filter_states( self, qn: str, value: tuple[float, float], *, delta: float = 1e-10, keep_unknown: bool = False diff --git a/src/rydstate/basis/basis_mqdt.py b/src/rydstate/basis/basis_mqdt.py index d2f335e3..87afbe96 100644 --- a/src/rydstate/basis/basis_mqdt.py +++ b/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, From f2e40f908cc37a0e3818795c1dc06f0e072e6b55 Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Mon, 20 Jul 2026 19:37:39 +0200 Subject: [PATCH 2/9] angular ket add parity to repr if unknown --- src/rydstate/angular/angular_ket.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/rydstate/angular/angular_ket.py b/src/rydstate/angular/angular_ket.py index 78b23888..683648b3 100644 --- a/src/rydstate/angular/angular_ket.py +++ b/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})" From 212bcf7e3b8f5a0c5efecbbbba7cd8e69c5cd44f Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Mon, 20 Jul 2026 19:44:16 +0200 Subject: [PATCH 3/9] make Unknown and NotSet comparable to numbers --- src/rydstate/angular/utils.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/src/rydstate/angular/utils.py b/src/rydstate/angular/utils.py index 7aeb473e..a04d0c63 100644 --- a/src/rydstate/angular/utils.py +++ b/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): From 7db9499c04b02b25f67f61375cbb247bb1029fb4 Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Mon, 20 Jul 2026 19:38:28 +0200 Subject: [PATCH 4/9] Basis add filter by parity or label substring --- docs/examples/mqdt/mqdt_exp_qn.ipynb | 6 +-- src/rydstate/basis/basis_base.py | 43 +++++++--------------- src/rydstate/rydberg_state/rydberg_base.py | 4 ++ 3 files changed, 21 insertions(+), 32 deletions(-) diff --git a/docs/examples/mqdt/mqdt_exp_qn.ipynb b/docs/examples/mqdt/mqdt_exp_qn.ipynb index 33754252..dec5d421 100644 --- a/docs/examples/mqdt/mqdt_exp_qn.ipynb +++ b/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", diff --git a/src/rydstate/basis/basis_base.py b/src/rydstate/basis/basis_base.py index 69763fbc..489040a3 100644 --- a/src/rydstate/basis/basis_base.py +++ b/src/rydstate/basis/basis_base.py @@ -8,12 +8,12 @@ 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) @@ -35,39 +35,24 @@ def shallow_copy(self) -> Self: new_basis.states = self.states.copy() return new_basis - @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 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: diff --git a/src/rydstate/rydberg_state/rydberg_base.py b/src/rydstate/rydberg_state/rydberg_base.py index 57ebf3de..bc535387 100644 --- a/src/rydstate/rydberg_state/rydberg_base.py +++ b/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: From 8644940b6505502c2b22c3d9c0d92ce9e8fa84af Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Mon, 20 Jul 2026 19:47:34 +0200 Subject: [PATCH 5/9] MQDT fix allow general AngularKet as outer channel --- src/rydstate/species/fmodel.py | 11 +++++++---- src/rydstate/species/mqdt.py | 7 +++++-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/rydstate/species/fmodel.py b/src/rydstate/species/fmodel.py index e0884b58..61eb8e36 100644 --- a/src/rydstate/species/fmodel.py +++ b/src/rydstate/species/fmodel.py @@ -7,13 +7,14 @@ 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 from rydstate.angular.utils import AllKnown from rydstate.species.mqdt import MQDT from rydstate.species.utils import RydbergRitzParameters @@ -40,7 +41,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]] @@ -344,9 +345,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] diff --git a/src/rydstate/species/mqdt.py b/src/rydstate/species/mqdt.py index 74f1780f..b4314a42 100644 --- a/src/rydstate/species/mqdt.py +++ b/src/rydstate/species/mqdt.py @@ -4,13 +4,14 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, ClassVar, overload +from rydstate.angular.utils import is_not_set from rydstate.metaclass_cache import CachedABCMeta from rydstate.species.fmodel import FModelSQDT from rydstate.species.utils import get_all_subclasses from rydstate.units import ureg if TYPE_CHECKING: - from rydstate.angular.angular_ket import AngularKetFJ + from rydstate.angular.angular_ket import AngularKetBase from rydstate.angular.core_ket import CoreKet from rydstate.species.fmodel import FModel from rydstate.units import PintFloat @@ -90,8 +91,10 @@ def reference_ionization_energy_au(self) -> float: return self.get_ionization_threshold(self.reference_core_ket, unit="a.u.") return min(self.get_ionization_threshold(core_ket, unit="a.u.") for core_ket in self.ionization_threshold_dict) - def get_mqdt_models(self, outer_channel: AngularKetFJ[Any]) -> list[FModel]: + def get_mqdt_models(self, outer_channel: AngularKetBase[Any]) -> list[FModel]: """Return a list of MQDT models for the outer_channel.""" + if not is_not_set(outer_channel.m): + raise ValueError("The m quantum number of the outer_channel must be NotSet.") models = [ model for model in self.models From 9dc57e6a1a7b2e1708586da18307ed732fcf7851 Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Mon, 20 Jul 2026 16:24:01 +0200 Subject: [PATCH 6/9] simplify fmodel and basis mqdt --- src/rydstate/basis/basis_mqdt.py | 36 +++++++++++++------------------- src/rydstate/species/fmodel.py | 24 --------------------- 2 files changed, 15 insertions(+), 45 deletions(-) diff --git a/src/rydstate/basis/basis_mqdt.py b/src/rydstate/basis/basis_mqdt.py index 87afbe96..fa0da9e4 100644 --- a/src/rydstate/basis/basis_mqdt.py +++ b/src/rydstate/basis/basis_mqdt.py @@ -123,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). @@ -155,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 @@ -190,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( @@ -200,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 diff --git a/src/rydstate/species/fmodel.py b/src/rydstate/species/fmodel.py index 61eb8e36..14e8cda8 100644 --- a/src/rydstate/species/fmodel.py +++ b/src/rydstate/species/fmodel.py @@ -292,30 +292,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. From 18cac874d2070a96625c1a3445e588553931a6f2 Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Mon, 20 Jul 2026 19:46:13 +0200 Subject: [PATCH 7/9] mqdt and fmodel add get_core_kets --- src/rydstate/species/fmodel.py | 6 ++++++ src/rydstate/species/mqdt.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/src/rydstate/species/fmodel.py b/src/rydstate/species/fmodel.py index 14e8cda8..561146a6 100644 --- a/src/rydstate/species/fmodel.py +++ b/src/rydstate/species/fmodel.py @@ -15,6 +15,7 @@ from types import ModuleType from rydstate.angular.angular_ket import AngularKetBase + 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 @@ -72,6 +73,11 @@ def nu_max(self) -> float: """Maximum nu for which the model is valid.""" return self.nu_range[1] + 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]: ... diff --git a/src/rydstate/species/mqdt.py b/src/rydstate/species/mqdt.py index b4314a42..c043d2ea 100644 --- a/src/rydstate/species/mqdt.py +++ b/src/rydstate/species/mqdt.py @@ -104,6 +104,13 @@ def get_mqdt_models(self, outer_channel: AngularKetBase[Any]) -> list[FModel]: models = [FModelSQDT(self.species, outer_channel, mqdt=self)] return models + def get_core_kets(self) -> list[CoreKet]: + """Return a list of relevant core kets of all FModels.""" + core_kets = set() + for model in self.models: + core_kets.update(model.get_core_kets()) + return sorted(core_kets, key=lambda ket: (ket.l_c, ket.j_c, ket.f_c, str(ket.label))) + def get_mqdt(species: str, tag: str | None = None) -> MQDT: """Get an instance of the subclass of MQDT for the given species and tag.""" From fdd2d4b199e1141f0d2665f65c24638112143801 Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Tue, 21 Jul 2026 10:17:09 +0200 Subject: [PATCH 8/9] FModel add fj_channels property --- src/rydstate/species/fmodel.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/rydstate/species/fmodel.py b/src/rydstate/species/fmodel.py index 561146a6..2587836b 100644 --- a/src/rydstate/species/fmodel.py +++ b/src/rydstate/species/fmodel.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from types import ModuleType - from rydstate.angular.angular_ket import AngularKetBase + 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 @@ -73,6 +73,11 @@ 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} From b34bfa02cc2ff849e912dad824268ad84c99b0d4 Mon Sep 17 00:00:00 2001 From: johannes-moegerle Date: Tue, 21 Jul 2026 12:00:23 +0200 Subject: [PATCH 9/9] [tests] add more tests --- tests/test_angular_utils.py | 70 +++++++++++++++++++++++++++++++++++++ tests/test_basis.py | 23 ++++++++++++ tests/test_basis_mqdt.py | 12 +++++++ tests/test_mqdt_models.py | 24 +++++++++++++ 4 files changed, 129 insertions(+) create mode 100644 tests/test_angular_utils.py diff --git a/tests/test_angular_utils.py b/tests/test_angular_utils.py new file mode 100644 index 00000000..2a2cc6eb --- /dev/null +++ b/tests/test_angular_utils.py @@ -0,0 +1,70 @@ +import math +from typing import Any + +import pytest +from rydstate.angular.utils import NotSet, Unknown + +REAL_NUMBERS = [-math.inf, -3, 0, 0.5, 2, 1e100, math.inf] + + +@pytest.mark.parametrize("singleton", [Unknown, NotSet]) +def test_singleton_sorts_above_every_real_number(singleton: type[Any]) -> None: + """Unknown and NotSet compare greater than every real number, including inf.""" + for number in REAL_NUMBERS: + assert singleton > number + assert singleton >= number + assert not singleton < number + assert not singleton <= number + # reflected operators (number on the left-hand side) + assert number < singleton + assert number <= singleton + assert not number > singleton + assert not number >= singleton + assert singleton != number # type: ignore [comparison-overlap] + assert number != singleton # type: ignore [comparison-overlap] + + +@pytest.mark.parametrize("singleton", [Unknown, NotSet]) +def test_singleton_comparison_with_itself(singleton: type[Any]) -> None: + other = singleton + assert singleton == other + assert singleton <= other + assert singleton >= other + assert not singleton < other + assert not singleton > other + + +def test_singletons_are_not_comparable_to_each_other() -> None: + assert Unknown != NotSet # type: ignore [comparison-overlap] + with pytest.raises(TypeError): + _ = Unknown < NotSet + + +@pytest.mark.parametrize("singleton", [Unknown, NotSet]) +def test_singleton_arithmetic_returns_singleton(singleton: type[Any]) -> None: + """Binary and unary arithmetic with real numbers propagates the singleton.""" + assert singleton + 1 is singleton + assert 1 + singleton is singleton + assert singleton - 0.5 is singleton + assert 0.5 - singleton is singleton + assert singleton * 2 is singleton + assert 2 * singleton is singleton + assert +singleton is singleton + assert -singleton is singleton + assert singleton + singleton is singleton + + +def test_singleton_arithmetic_with_non_number_raises() -> None: + with pytest.raises(TypeError): + _ = Unknown + "a" + with pytest.raises(TypeError): + _ = Unknown + NotSet + + +def test_sorting_with_unknown() -> None: + """Unknown values sort after all real numbers, e.g. in filter and sort keys.""" + assert sorted([2.0, Unknown, 1.0]) == [1.0, 2.0, Unknown] # type: ignore [type-var] + assert max(3, Unknown) is Unknown + assert min(3, Unknown) == 3 + # a range filter like in BasisBase.filter_states always excludes Unknown + assert not 0 <= Unknown <= math.inf diff --git a/tests/test_basis.py b/tests/test_basis.py index a79babec..a61e9e48 100644 --- a/tests/test_basis.py +++ b/tests/test_basis.py @@ -75,3 +75,26 @@ def test_alkaline_basis(species: str) -> None: me_matrix = basis.calc_reduced_matrix_elements(basis, "electric_dipole", unit="e a0") assert np.shape(me_matrix) == (len(basis.states), len(basis.states)) assert np.count_nonzero(me_matrix) > 0 + + +def test_shallow_copy() -> None: + """A shallow copy has an independent states list but shares the state objects.""" + basis = BasisSQDT("H", n=(1, 5), coupling_scheme="LS") + copied = basis.shallow_copy() + assert copied.states is not basis.states + assert all(a is b for a, b in zip(copied.states, basis.states, strict=True)) + + n_states = len(basis) + copied.filter_states("l_r", 0) + assert len(copied) < n_states + assert len(basis) == n_states + + +def test_filter_states_parity() -> None: + """filter_states("parity", ...) splits the basis into even and odd states.""" + basis = BasisSQDT("H", n=(1, 5), coupling_scheme="LS") + even = basis.shallow_copy().filter_states("parity", 1) + odd = basis.shallow_copy().filter_states("parity", -1) + assert len(even) + len(odd) == len(basis) + assert all(state.angular.l_r % 2 == 0 for state in even.states) + assert all(state.angular.l_r % 2 == 1 for state in odd.states) diff --git a/tests/test_basis_mqdt.py b/tests/test_basis_mqdt.py index 2c2059af..f2ae6780 100644 --- a/tests/test_basis_mqdt.py +++ b/tests/test_basis_mqdt.py @@ -75,3 +75,15 @@ def test_mqdt_basis_sort_and_filter() -> None: basis.filter_states("nu", (27.0, 28.0)) assert len(basis) < n_before assert all(27.0 - 1e-10 <= s.nu <= 28.0 + 1e-10 for s in basis.states) + + +def test_mqdt_basis_filter_by_label() -> None: + """filter_states_label keeps only states with a matching channel label.""" + basis = BasisMQDT("Yb174", nu=(25, 26), f_tot=(0, 1)) + n_states = len(basis) + assert n_states > 0 + + labeled = basis.shallow_copy().filter_states_label("4f13 5d") + assert 0 < len(labeled) < n_states + for state in labeled.states: + assert any(ket.angular.label is not None and "4f13 5d" in ket.angular.label for ket in state.rydberg_kets) diff --git a/tests/test_mqdt_models.py b/tests/test_mqdt_models.py index 3441c5b1..398f281d 100644 --- a/tests/test_mqdt_models.py +++ b/tests/test_mqdt_models.py @@ -198,3 +198,27 @@ def test_all_models_found_by_get_mqdt_models(mqdt: MQDT) -> None: if model.species == mqdt.species and model.full_name not in found_model_names ] assert not missing, f"{mqdt!r}: {len(missing)} models not reachable via get_mqdt_models: {missing}" + + +def test_fj_channels(model: FModel) -> None: + """fj_channels decomposes every outer channel into FJ kets with the model's f_tot.""" + fj_channels = model.fj_channels + assert len(fj_channels) >= len(model.outer_channels) + assert all(isinstance(ket, AngularKetFJ) for ket in fj_channels) + assert all(ket.f_tot == model.f_tot for ket in fj_channels) + + +def test_fmodel_get_core_kets(model: FModel) -> None: + """get_core_kets returns the sorted unique core kets of the outer channels.""" + core_kets = model.get_core_kets() + assert len(core_kets) == len(set(core_kets)) + assert set(core_kets) == {ch.get_core_ket() for ch in model.outer_channels} + sort_keys = [(ket.l_c, ket.j_c, ket.f_c, str(ket.label)) for ket in core_kets] + assert sort_keys == sorted(sort_keys) + + +def test_mqdt_get_core_kets(mqdt: MQDT) -> None: + """MQDT.get_core_kets returns the union of core kets over all its models.""" + core_kets = mqdt.get_core_kets() + assert len(core_kets) == len(set(core_kets)) + assert set(core_kets) == {ck for model in mqdt.models for ck in model.get_core_kets()}