Various improvements

.github/workflows/testing.yml

@@ -10,7 +10,7 @@ jobs:
   pytest-with-nbmake:
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
     runs-on: ubuntu-latest
     timeout-minutes: 30
     steps:

.pre-commit-config.yaml

@@ -9,20 +9,20 @@ repos:
       - id: check-merge-conflict
 
   - repo: https://github.com/Lucas-C/pre-commit-hooks
-    rev: v1.5.5
+    rev: v1.5.6
     hooks:
       - id: remove-tabs
         exclude: 'Makefile|nist_energy_levels/.*\.txt$'
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.0
+    rev: v0.15.5
     hooks:
       - id: ruff
         args: ["--fix", "--exit-non-zero-on-fix"]
       - id: ruff-format
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.18.2
+    rev: v1.19.1
     hooks:
       - id: mypy
         additional_dependencies: ["numpy >= 2.0", "pint >= 0.25.1"]

README.md

@@ -16,7 +16,7 @@
 
 The *RydState* software calculates properties of Rydberg states.
 We especially focus on the calculation of the radial wavefunction of Rydberg states via the Numerov method.
-The software can be installed via pip (requires Python >= 3.9):
+The software can be installed via pip (requires Python >= 3.10):
 
 ```bash
 pip install rydstate

pyproject.toml

@@ -24,7 +24,6 @@ classifiers = [
   "Programming Language :: Python",
   "Programming Language :: Python :: 3",
   "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.9",
   "Programming Language :: Python :: 3.10",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
@@ -33,7 +32,7 @@ classifiers = [
   "Topic :: Scientific/Engineering :: Physics",
   "Typing :: Typed",
 ]
-requires-python = ">= 3.9"
+requires-python = ">= 3.10"
 dependencies = [
     "numpy >= 2.0",
     "numba >= 0.60",
@@ -47,7 +46,6 @@ dependencies = [
 tests = [
   "pytest >= 8.0",
   "nbmake >= 1.3",
-  "rydstate[comparison]",
 ]
 docs = [
   "sphinx >= 7",
@@ -72,7 +70,7 @@ mypy = [
 
 [dependency-groups]
 dev = [
-  "rydstate[docs,tests,comparison,jupyter,mypy]",
+  "rydstate[docs,tests,jupyter,mypy]",
   "check-wheel-contents >= 0.6",
 ]
 
@@ -108,7 +106,7 @@ addopts = [
 
 [tool.ruff]
 line-length = 120
-target-version = "py39"
+target-version = "py310"
 extend-include = ["*.ipynb"]
 
 [tool.ruff.lint]

src/rydstate/angular/angular_ket.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import contextlib
 import logging
 from abc import ABC
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, overload
@@ -14,32 +13,26 @@
     is_angular_operator_type,
 )
 from rydstate.angular.utils import (
+    InvalidQuantumNumbersError,
     check_spin_addition_rule,
     get_possible_quantum_number_values,
     minus_one_pow,
     try_trivial_spin_addition,
 )
 from rydstate.angular.wigner_symbols import calc_wigner_3j, clebsch_gordan_6j, clebsch_gordan_9j
-from rydstate.species import SpeciesObject
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from typing_extensions import Self
 
     from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers, AngularOperatorType
     from rydstate.angular.angular_state import AngularState
+    from rydstate.angular.utils import CouplingScheme
+    from rydstate.species import SpeciesObject
 
 logger = logging.getLogger(__name__)
 
-CouplingScheme = Literal["LS", "JJ", "FJ"]
-
-
-class InvalidQuantumNumbersError(ValueError):
-    def __init__(self, ket: AngularKetBase, msg: str = "") -> None:
-        _msg = f"Invalid quantum numbers for {ket!r}"
-        if len(msg) > 0:
-            _msg += f"\n  {msg}"
-        super().__init__(_msg)
-
 
 class AngularKetBase(ABC):
     """Base class for a angular ket (i.e. a simple canonical spin ketstate)."""
@@ -92,12 +85,13 @@ def __init__(
     ) -> None:
         """Initialize the Spin ket.
 
-        species:
-        Atomic species, e.g. 'Rb87'.
-        Not used for calculation, only for convenience to infer the core electron spin and nuclear spin quantum numbers.
+        Atomic species, e.g. 'Rb87', will not be used for calculation,
+        only for convenience to infer the core electron spin and nuclear spin quantum numbers.
         """
         if species is not None:
             if isinstance(species, str):
+                from rydstate.species import SpeciesObject  # noqa: PLC0415
+
                 species = SpeciesObject.from_name(species)
             # use i_c = 0 for species without defined nuclear spin (-> ignore hyperfine)
             species_i_c = species.i_c if species.i_c is not None else 0
@@ -154,7 +148,7 @@ def __setattr__(self, key: str, value: object) -> None:
         super().__setattr__(key, value)
 
     def __repr__(self) -> str:
-        args = ", ".join(f"{qn}={val}" for qn, val in zip(self.quantum_number_names, self.quantum_numbers))
+        args = ", ".join(f"{qn}={val}" for qn, val in zip(self.quantum_number_names, self.quantum_numbers, strict=True))
         if self.m is not None:
             args += f", m={self.m}"
         return f"{self.__class__.__name__}({args})"
@@ -237,10 +231,8 @@ def to_state(self, coupling_scheme: CouplingScheme | None = None) -> AngularStat
             The angular state in the specified coupling scheme.
 
         """
-        from rydstate.angular.angular_state import AngularState  # noqa: PLC0415
-
         if coupling_scheme is None or coupling_scheme == self.coupling_scheme:
-            return AngularState([1], [self])
+            return self._create_angular_state([1], [self])
         if coupling_scheme == "LS":
             return self._to_state_ls()
         if coupling_scheme == "JJ":
@@ -280,9 +272,7 @@ def _to_state_ls(self) -> AngularState[AngularKetLS]:
                         kets.append(ls_ket)
                         coefficients.append(coeff)
 
-        from rydstate.angular.angular_state import AngularState  # noqa: PLC0415
-
-        return AngularState(coefficients, kets)
+        return self._create_angular_state(coefficients, kets)
 
     def _to_state_jj(self) -> AngularState[AngularKetJJ]:
         """Convert a single ket to state in JJ coupling."""
@@ -315,9 +305,7 @@ def _to_state_jj(self) -> AngularState[AngularKetJJ]:
                         kets.append(jj_ket)
                         coefficients.append(coeff)
 
-        from rydstate.angular.angular_state import AngularState  # noqa: PLC0415
-
-        return AngularState(coefficients, kets)
+        return self._create_angular_state(coefficients, kets)
 
     def _to_state_fj(self) -> AngularState[AngularKetFJ]:
         """Convert a single ket to state in FJ coupling."""
@@ -350,6 +338,10 @@ def _to_state_fj(self) -> AngularState[AngularKetFJ]:
                         kets.append(fj_ket)
                         coefficients.append(coeff)
 
+        return self._create_angular_state(coefficients, kets)
+
+    def _create_angular_state(self, coefficients: Sequence[float], kets: Sequence[AngularKetBase]) -> AngularState[Any]:
+        """Create an AngularState from coefficients and kets."""
         from rydstate.angular.angular_state import AngularState  # noqa: PLC0415
 
         return AngularState(coefficients, kets)
@@ -736,55 +728,3 @@ def sanity_check(self, msgs: list[str] | None = None) -> None:
             msgs.append(f"{self.f_c=}, {self.j_r=}, {self.f_tot=} don't satisfy spin addition rule.")
 
         super().sanity_check(msgs)
-
-
-def quantum_numbers_to_angular_ket(
-    species: str | SpeciesObject,
-    s_c: float | None = None,
-    l_c: int = 0,
-    j_c: float | None = None,
-    f_c: float | None = None,
-    s_r: float = 0.5,
-    l_r: int | None = None,
-    j_r: float | None = None,
-    s_tot: float | None = None,
-    l_tot: int | None = None,
-    j_tot: float | None = None,
-    f_tot: float | None = None,
-    m: float | None = None,
-) -> AngularKetBase:
-    r"""Return an AngularKet object in the corresponding coupling scheme from the given quantum numbers.
-
-    Args:
-        species: Atomic species.
-        s_c: Spin quantum number of the core electron (0 for Alkali, 0.5 for divalent atoms).
-        l_c: Orbital angular momentum quantum number of the core electron.
-        j_c: Total angular momentum quantum number of the core electron.
-        f_c: Total angular momentum quantum number of the core (core electron + nucleus).
-        s_r: Spin quantum number of the rydberg electron (always 0.5).
-        l_r: Orbital angular momentum quantum number of the rydberg electron.
-        j_r: Total angular momentum quantum number of the rydberg electron.
-        s_tot: Total spin quantum number of all electrons.
-        l_tot: Total orbital angular momentum quantum number of all electrons.
-        j_tot: Total angular momentum quantum number of all electrons.
-        f_tot: Total angular momentum quantum number of the atom (rydberg electron + core).
-        m: Total magnetic quantum number.
-          Optional, only needed for concrete angular matrix elements.
-
-    """
-    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
-        return AngularKetLS(
-            s_c=s_c, l_c=l_c, s_r=s_r, l_r=l_r, s_tot=s_tot, l_tot=l_tot, j_tot=j_tot, f_tot=f_tot, m=m, species=species
-        )
-
-    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
-        return AngularKetJJ(
-            s_c=s_c, l_c=l_c, j_c=j_c, s_r=s_r, l_r=l_r, j_r=j_r, j_tot=j_tot, f_tot=f_tot, m=m, species=species
-        )
-
-    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
-        return AngularKetFJ(
-            s_c=s_c, l_c=l_c, j_c=j_c, f_c=f_c, s_r=s_r, l_r=l_r, j_r=j_r, f_tot=f_tot, m=m, species=species
-        )
-
-    raise ValueError("Invalid combination of angular quantum numbers provided.")

src/rydstate/angular/angular_matrix_element.py

@@ -2,15 +2,16 @@
 
 import math
 from functools import lru_cache
-from typing import TYPE_CHECKING, Callable, Literal, TypeVar, get_args
+from typing import TYPE_CHECKING, Literal, TypeGuard, TypeVar, get_args
 
 import numpy as np
-from typing_extensions import TypeGuard
 
 from rydstate.angular.utils import minus_one_pow
 from rydstate.angular.wigner_symbols import calc_wigner_3j, calc_wigner_6j
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from typing_extensions import ParamSpec
 
     P = ParamSpec("P")

src/rydstate/angular/angular_state.py

@@ -19,8 +19,9 @@
 
     from typing_extensions import Self
 
-    from rydstate.angular.angular_ket import CouplingScheme
     from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers, AngularOperatorType
+    from rydstate.angular.utils import CouplingScheme
+    from rydstate.units import NDArray
 
 logger = logging.getLogger(__name__)
 
@@ -30,7 +31,11 @@
 
 class AngularState(Generic[_AngularKet]):
     def __init__(
-        self, coefficients: Sequence[float], kets: Sequence[_AngularKet], *, warn_if_not_normalized: bool = True
+        self,
+        coefficients: Sequence[float] | NDArray,
+        kets: Sequence[_AngularKet],
+        *,
+        warn_if_not_normalized: bool = True,
     ) -> None:
         self.coefficients = np.array(coefficients)
         self.kets = kets
@@ -49,7 +54,7 @@ def __init__(
             self.coefficients /= self.norm
 
     def __iter__(self) -> Iterator[tuple[float, _AngularKet]]:
-        return zip(self.coefficients, self.kets).__iter__()
+        return zip(self.coefficients, self.kets, strict=True).__iter__()
 
     def __repr__(self) -> str:
         terms = [f"{coeff}*{ket!r}" for coeff, ket in self]

src/rydstate/angular/utils.py

@@ -1,7 +1,24 @@
 from __future__ import annotations
 
+import contextlib
+from typing import TYPE_CHECKING, Literal
+
 import numpy as np
 
+if TYPE_CHECKING:
+    from rydstate.angular.angular_ket import AngularKetBase
+    from rydstate.species.species_object import SpeciesObject
+
+CouplingScheme = Literal["LS", "JJ", "FJ"]
+
+
+class InvalidQuantumNumbersError(ValueError):
+    def __init__(self, ket: AngularKetBase, msg: str = "") -> None:
+        _msg = f"Invalid quantum numbers for {ket!r}"
+        if len(msg) > 0:
+            _msg += f"\n  {msg}"
+        super().__init__(_msg)
+
 
 def minus_one_pow(n: float) -> int:
     """Calculate (-1)^n for an integer n and raise an error if n is not an integer."""
@@ -42,3 +59,57 @@ def get_possible_quantum_number_values(s_1: float, s_2: float, s_tot: float | No
     if s_tot is not None:
         return [s_tot]
     return [float(s) for s in np.arange(abs(s_1 - s_2), s_1 + s_2 + 1, 1)]
+
+
+def quantum_numbers_to_angular_ket(
+    species: str | SpeciesObject,
+    s_c: float | None = None,
+    l_c: int = 0,
+    j_c: float | None = None,
+    f_c: float | None = None,
+    s_r: float = 0.5,
+    l_r: int | None = None,
+    j_r: float | None = None,
+    s_tot: float | None = None,
+    l_tot: int | None = None,
+    j_tot: float | None = None,
+    f_tot: float | None = None,
+    m: float | None = None,
+) -> AngularKetBase:
+    r"""Return an AngularKet object in the corresponding coupling scheme from the given quantum numbers.
+
+    Args:
+        species: Atomic species.
+        s_c: Spin quantum number of the core electron (0 for Alkali, 0.5 for divalent atoms).
+        l_c: Orbital angular momentum quantum number of the core electron.
+        j_c: Total angular momentum quantum number of the core electron.
+        f_c: Total angular momentum quantum number of the core (core electron + nucleus).
+        s_r: Spin quantum number of the rydberg electron (always 0.5).
+        l_r: Orbital angular momentum quantum number of the rydberg electron.
+        j_r: Total angular momentum quantum number of the rydberg electron.
+        s_tot: Total spin quantum number of all electrons.
+        l_tot: Total orbital angular momentum quantum number of all electrons.
+        j_tot: Total angular momentum quantum number of all electrons.
+        f_tot: Total angular momentum quantum number of the atom (rydberg electron + core).
+        m: Total magnetic quantum number.
+          Optional, only needed for concrete angular matrix elements.
+
+    """
+    from rydstate.angular.angular_ket import AngularKetFJ, AngularKetJJ, AngularKetLS  # noqa: PLC0415
+
+    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
+        return AngularKetLS(
+            s_c=s_c, l_c=l_c, s_r=s_r, l_r=l_r, s_tot=s_tot, l_tot=l_tot, j_tot=j_tot, f_tot=f_tot, m=m, species=species
+        )
+
+    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
+        return AngularKetJJ(
+            s_c=s_c, l_c=l_c, j_c=j_c, s_r=s_r, l_r=l_r, j_r=j_r, j_tot=j_tot, f_tot=f_tot, m=m, species=species
+        )
+
+    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
+        return AngularKetFJ(
+            s_c=s_c, l_c=l_c, j_c=j_c, f_c=f_c, s_r=s_r, l_r=l_r, j_r=j_r, f_tot=f_tot, m=m, species=species
+        )
+
+    raise ValueError("Invalid combination of angular quantum numbers provided.")