From 82b106057a08854857bee72453b31008b63eb64f Mon Sep 17 00:00:00 2001
From: johannes-moegerle <johannes.moegerle@itp3.uni-stuttgart.de>
Date: Fri, 28 Nov 2025 05:54:50 +0100
Subject: [PATCH 1/4] rename RadialState -> RadialKet

---
 docs/examples.rst                             |  2 +-
 .../radial/hydrogen_wavefunction.ipynb        |  6 +--
 src/rydstate/radial/__init__.py               |  4 +-
 .../radial/{radial_state.py => radial_ket.py} | 18 ++++-----
 src/rydstate/radial/wavefunction.py           | 32 ++++++++--------
 src/rydstate/rydberg_state.py                 | 38 +++++++++----------
 tests/test_radial_matrix_element.py           |  4 +-
 7 files changed, 50 insertions(+), 54 deletions(-)
 rename src/rydstate/radial/{radial_state.py => radial_ket.py} (94%)

diff --git a/docs/examples.rst b/docs/examples.rst
index 5ba8c86..aff1b31 100644
--- a/docs/examples.rst
+++ b/docs/examples.rst
@@ -5,7 +5,7 @@ Examples
 RadialState
 -----------
 
-Some examples demonstrating the usage of the RadialState class, which uses the Numerov method for solving the radial Schrödinger equation.
+Some examples demonstrating the usage of the RadialKet class, which uses the Numerov method for solving the radial Schrödinger equation.
 
 .. nbgallery::
    examples/radial/hydrogen_wavefunction
diff --git a/docs/examples/radial/hydrogen_wavefunction.ipynb b/docs/examples/radial/hydrogen_wavefunction.ipynb
index 64c99cd..4f41872 100644
--- a/docs/examples/radial/hydrogen_wavefunction.ipynb
+++ b/docs/examples/radial/hydrogen_wavefunction.ipynb
@@ -9,7 +9,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -18,7 +18,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.radial import RadialState"
+    "from rydstate.radial import RadialKet"
    ]
   },
   {
@@ -27,7 +27,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "state = RadialState(\"H_textbook\", nu=10, l_r=5)\n",
+    "state = RadialKet(\"H_textbook\", nu=10, l_r=5)\n",
     "state.set_n_for_sanity_check(10)\n",
     "state.create_model()\n",
     "state.create_grid(dz=1e-2)\n",
diff --git a/src/rydstate/radial/__init__.py b/src/rydstate/radial/__init__.py
index 4e436b5..3b21898 100644
--- a/src/rydstate/radial/__init__.py
+++ b/src/rydstate/radial/__init__.py
@@ -1,15 +1,15 @@
 from rydstate.radial.grid import Grid
 from rydstate.radial.model import Model, PotentialType
 from rydstate.radial.numerov import run_numerov_integration
+from rydstate.radial.radial_ket import RadialKet
 from rydstate.radial.radial_matrix_element import calc_radial_matrix_element_from_w_z
-from rydstate.radial.radial_state import RadialState
 from rydstate.radial.wavefunction import Wavefunction, WavefunctionNumerov, WavefunctionWhittaker
 
 __all__ = [
     "Grid",
     "Model",
     "PotentialType",
-    "RadialState",
+    "RadialKet",
     "Wavefunction",
     "WavefunctionNumerov",
     "WavefunctionWhittaker",
diff --git a/src/rydstate/radial/radial_state.py b/src/rydstate/radial/radial_ket.py
similarity index 94%
rename from src/rydstate/radial/radial_state.py
rename to src/rydstate/radial/radial_ket.py
index ccac5f3..782a137 100644
--- a/src/rydstate/radial/radial_state.py
+++ b/src/rydstate/radial/radial_ket.py
@@ -21,7 +21,7 @@
 logger = logging.getLogger(__name__)
 
 
-class RadialState:
+class RadialKet:
     r"""Class representing a radial Rydberg state."""
 
     def __init__(
@@ -30,7 +30,7 @@ def __init__(
         nu: float,
         l_r: int,
     ) -> None:
-        r"""Initialize the radial state.
+        r"""Initialize the radial ket.
 
         Args:
             species: Atomic species.
@@ -197,11 +197,11 @@ def create_wavefunction(
         self._wavefunction.apply_sign_convention(sign_convention)
         self._grid = self._wavefunction.grid
 
-    def calc_overlap(self, other: RadialState, *, integration_method: INTEGRATION_METHODS = "sum") -> float:
-        r"""Calculate the overlap <self|other> of two radial states.
+    def calc_overlap(self, other: RadialKet, *, integration_method: INTEGRATION_METHODS = "sum") -> float:
+        r"""Calculate the overlap <self|other> of two radial kets.
 
         Args:
-            other: Other radial state
+            other: Other radial ket
             integration_method: Integration method to use
 
         Returns:
@@ -212,17 +212,17 @@ def calc_overlap(self, other: RadialState, *, integration_method: INTEGRATION_ME
 
     @overload
     def calc_matrix_element(
-        self, other: RadialState, k_radial: int, *, integration_method: INTEGRATION_METHODS = "sum"
+        self, other: RadialKet, k_radial: int, *, integration_method: INTEGRATION_METHODS = "sum"
     ) -> PintFloat: ...
 
     @overload
     def calc_matrix_element(
-        self, other: RadialState, k_radial: int, unit: str, *, integration_method: INTEGRATION_METHODS = "sum"
+        self, other: RadialKet, k_radial: int, unit: str, *, integration_method: INTEGRATION_METHODS = "sum"
     ) -> float: ...
 
     def calc_matrix_element(
         self,
-        other: RadialState,
+        other: RadialKet,
         k_radial: int,
         unit: str | None = None,
         *,
@@ -241,7 +241,7 @@ def calc_matrix_element(
         and w(z) = z^{-1/2} \tilde{u}(z^2) = (r/_a_0)^{1/4} \sqrt{a_0} r R(r).
 
         Args:
-            other: Other radial state
+            other: Other radial ket
             k_radial: Power of r in the matrix element
                 (default=0, this corresponds to the overlap integral \int dr r^2 R_1(r) R_2(r))
             unit: Unit of the returned matrix element, default None returns a Pint quantity.
diff --git a/src/rydstate/radial/wavefunction.py b/src/rydstate/radial/wavefunction.py
index 2ec3b7b..d7b8109 100644
--- a/src/rydstate/radial/wavefunction.py
+++ b/src/rydstate/radial/wavefunction.py
@@ -14,7 +14,7 @@
 
 if TYPE_CHECKING:
     from rydstate.radial import Grid, Model
-    from rydstate.radial.radial_state import RadialState
+    from rydstate.radial.radial_ket import RadialKet
     from rydstate.units import NDArray
 
 logger = logging.getLogger(__name__)
@@ -27,17 +27,17 @@ class Wavefunction(ABC):
 
     def __init__(
         self,
-        radial_state: RadialState,
+        radial_ket: RadialKet,
         grid: Grid,
     ) -> None:
         """Create a Wavefunction object.
 
         Args:
-            radial_state: The RadialState object.
+            radial_ket: The RadialKet object.
             grid: The grid object.
 
         """
-        self.radial_state = radial_state
+        self.radial_ket = radial_ket
         self.grid = grid
 
         self._w_list: NDArray | None = None
@@ -92,8 +92,8 @@ def apply_sign_convention(self, sign_convention: WavefunctionSignConvention) ->
                 break
 
         if sign_convention == "n_l_1":
-            assert self.radial_state.n is not None, "n must be given to apply the n_l_1 sign convention."
-            if current_outer_sign != (-1) ** (self.radial_state.n - self.radial_state.l_r - 1):
+            assert self.radial_ket.n is not None, "n must be given to apply the n_l_1 sign convention."
+            if current_outer_sign != (-1) ** (self.radial_ket.n - self.radial_ket.l_r - 1):
                 self._w_list = -self._w_list
         elif sign_convention == "positive_at_outer_bound":
             if current_outer_sign != 1:
@@ -115,19 +115,19 @@ class WavefunctionNumerov(Wavefunction):
 
     def __init__(
         self,
-        radial_state: RadialState,
+        radial_ket: RadialKet,
         grid: Grid,
         model: Model,
     ) -> None:
         """Create a Wavefunction object.
 
         Args:
-            radial_state: The RadialState object.
+            radial_ket: The RadialKet object.
             grid: The grid object.
             model: The model object.
 
         """
-        super().__init__(radial_state, grid)
+        super().__init__(radial_ket, grid)
         self.model = model
 
     def integrate(self, run_backward: bool = True, w0: float = 1e-10, *, _use_njit: bool = True) -> None:
@@ -172,8 +172,8 @@ def integrate(self, run_backward: bool = True, w0: float = 1e-10, *, _use_njit:
         # and not like in the rest of this class, i.e. y = w(z) and x = z
         grid = self.grid
 
-        species = self.radial_state.species
-        energy_au = calc_energy_from_nu(species.reduced_mass_au, self.radial_state.nu)
+        species = self.radial_ket.species
+        energy_au = calc_energy_from_nu(species.reduced_mass_au, self.radial_ket.nu)
         v_eff = self.model.calc_total_effective_potential(grid.x_list)
         glist = 8 * species.reduced_mass_au * grid.z_list * grid.z_list * (energy_au - v_eff)
 
@@ -201,7 +201,7 @@ def integrate(self, run_backward: bool = True, w0: float = 1e-10, *, _use_njit:
             y0, y1 = 0, w0
             x_start, x_stop, dx = grid.z_min, grid.z_max, grid.dz
             g_list_directed = glist
-            n = self.radial_state.n if self.radial_state.n is not None else self.radial_state.nu
+            n = self.radial_ket.n if self.radial_ket.n is not None else self.radial_ket.nu
             x_min = math.sqrt(n * (n + 15))
 
         if _use_njit:
@@ -238,7 +238,7 @@ def sanity_check(self, z_stop: float, run_backward: bool) -> bool:  # noqa: C901
         warning_msgs: list[str] = []
 
         grid = self.grid
-        state = self.radial_state
+        state = self.radial_ket
 
         # Check and Correct if divergence of the wavefunction
         w_list_abs = np.abs(self.w_list)
@@ -283,7 +283,7 @@ def sanity_check(self, z_stop: float, run_backward: bool) -> bool:  # noqa: C901
         elif n <= 16:
             tol = 2e-3
 
-        species = self.radial_state.species
+        species = self.radial_ket.species
         if species.number_valence_electrons == 2:
             # For divalent atoms the inner boundary is less well behaved ...
             tol = 2e-2
@@ -338,8 +338,8 @@ def sanity_check(self, z_stop: float, run_backward: bool) -> bool:  # noqa: C901
 class WavefunctionWhittaker(Wavefunction):
     def integrate(self) -> None:
         logger.warning("Using Whittaker to get the wavefunction is not recommended! Use this only for comparison.")
-        l = self.radial_state.l_r
-        nu = self.radial_state.nu
+        l = self.radial_ket.l_r
+        nu = self.radial_ket.nu
 
         whitw_vectorized = np.vectorize(whitw, otypes=[float])
         whitw_list = whitw_vectorized(nu, l + 0.5, 2 * self.grid.x_list / nu)
diff --git a/src/rydstate/rydberg_state.py b/src/rydstate/rydberg_state.py
index c3eab6e..7fa6a90 100644
--- a/src/rydstate/rydberg_state.py
+++ b/src/rydstate/rydberg_state.py
@@ -10,7 +10,7 @@
 
 from rydstate.angular import AngularKetJJ, AngularKetLS
 from rydstate.angular.utils import try_trivial_spin_addition
-from rydstate.radial import RadialState
+from rydstate.radial import RadialKet
 from rydstate.species.species_object import SpeciesObject
 from rydstate.species.utils import calc_energy_from_nu
 from rydstate.units import BaseQuantities, MatrixElementOperatorRanks, ureg
@@ -33,11 +33,13 @@ def __str__(self) -> str:
 
     @property
     @abstractmethod
-    def radial(self) -> RadialState: ...
+    def radial(self) -> RadialKet:
+        """The radial part of the Rydberg electron."""
 
     @property
     @abstractmethod
-    def angular(self) -> AngularKetBase: ...
+    def angular(self) -> AngularKetBase:
+        """The angular/spin part of the Rydberg electron."""
 
     @abstractmethod
     def get_nu(self) -> float:
@@ -227,15 +229,13 @@ def __init__(
 
     @cached_property
     def angular(self) -> AngularKetLS:
-        """The angular/spin state of the Rydberg electron."""
         return AngularKetLS(l_r=self.l, j_tot=self.j, m=self.m, f_tot=self.f, species=self.species)
 
     @cached_property
-    def radial(self) -> RadialState:
-        """The radial state of the Rydberg electron."""
-        radial_state = RadialState(self.species, nu=self.get_nu(), l_r=self.l)
-        radial_state.set_n_for_sanity_check(self.n)
-        return radial_state
+    def radial(self) -> RadialKet:
+        radial_ket = RadialKet(self.species, nu=self.get_nu(), l_r=self.l)
+        radial_ket.set_n_for_sanity_check(self.n)
+        return radial_ket
 
     def __repr__(self) -> str:
         species, n, l, j, f, m = self.species, self.n, self.l, self.j, self.f, self.m
@@ -290,17 +290,15 @@ def __init__(
 
     @cached_property
     def angular(self) -> AngularKetLS:
-        """The angular/spin state of the Rydberg electron."""
         return AngularKetLS(
             l_r=self.l, s_tot=self.s_tot, j_tot=self.j_tot, f_tot=self.f_tot, m=self.m, species=self.species
         )
 
     @cached_property
-    def radial(self) -> RadialState:
-        """The radial state of the Rydberg electron."""
-        radial_state = RadialState(self.species, nu=self.get_nu(), l_r=self.l)
-        radial_state.set_n_for_sanity_check(self.n)
-        return radial_state
+    def radial(self) -> RadialKet:
+        radial_ket = RadialKet(self.species, nu=self.get_nu(), l_r=self.l)
+        radial_ket.set_n_for_sanity_check(self.n)
+        return radial_ket
 
     def __repr__(self) -> str:
         species, n, l, s_tot, j_tot, f_tot, m = self.species, self.n, self.l, self.s_tot, self.j_tot, self.f_tot, self.m
@@ -358,17 +356,15 @@ def __init__(
 
     @cached_property
     def angular(self) -> AngularKetJJ:
-        """The angular/spin state of the Rydberg electron."""
         return AngularKetJJ(
             l_r=self.l, j_r=self.j_r, j_tot=self.j_tot, f_tot=self.f_tot, m=self.m, species=self.species
         )
 
     @cached_property
-    def radial(self) -> RadialState:
-        """The radial state of the Rydberg electron."""
-        radial_state = RadialState(self.species, nu=self.get_nu(), l_r=self.l)
-        radial_state.set_n_for_sanity_check(self.n)
-        return radial_state
+    def radial(self) -> RadialKet:
+        radial_ket = RadialKet(self.species, nu=self.get_nu(), l_r=self.l)
+        radial_ket.set_n_for_sanity_check(self.n)
+        return radial_ket
 
     def __repr__(self) -> str:
         species, n, l, j_r, j_tot, f_tot, m = self.species, self.n, self.l, self.j_r, self.j_tot, self.f_tot, self.m
diff --git a/tests/test_radial_matrix_element.py b/tests/test_radial_matrix_element.py
index 709f0aa..552e5fd 100644
--- a/tests/test_radial_matrix_element.py
+++ b/tests/test_radial_matrix_element.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from rydstate.radial import RadialState
+from rydstate.radial import RadialKet
 from rydstate.rydberg_state import RydbergStateAlkali
 from rydstate.species import SpeciesObject
 
@@ -61,7 +61,7 @@ def test_circular_expectation_value(species_name: str, n: int, l: int, j_tot: fl
     species = SpeciesObject.from_name(species_name)
     nu = species.calc_nu(n, l, j_tot)
 
-    state = RadialState(species, nu=nu, l_r=l)
+    state = RadialKet(species, nu=nu, l_r=l)
     state.set_n_for_sanity_check(n)
     state.create_wavefunction()
 

From 5988f294f94f13d5687bba46a5c15036328de9b8 Mon Sep 17 00:00:00 2001
From: johannes-moegerle <johannes.moegerle@itp3.uni-stuttgart.de>
Date: Fri, 19 Dec 2025 14:40:37 +0100
Subject: [PATCH 2/4] rename RydbergState -> RydbergStateSQDT

---
 README.md                                     |  2 +-
 docs/examples/benchmark/benchmark_njit.ipynb  | 10 +++---
 .../compare_dipole_matrix_element.ipynb       |  9 +++---
 .../compare_model_potentials.ipynb            | 20 ++++++------
 .../compare_radial_matrix_element.ipynb       |  9 +++---
 .../comparisons/compare_wavefunctions.ipynb   | 10 +++---
 .../comparisons/compare_whittaker.ipynb       | 22 ++++++-------
 .../comparisons/compare_z_min_cutoff.ipynb    |  6 ++--
 docs/examples/dipole_matrix_elements.ipynb    |  8 ++---
 .../radial/rubidium_wavefunction.ipynb        |  6 ++--
 src/rydstate/__init__.py                      | 12 ++++---
 src/rydstate/rydberg/__init__.py              | 13 ++++++++
 .../rydberg_sqdt.py}                          | 32 +++++++++++++++----
 src/rydstate/species/species_object.py        |  2 +-
 tests/test_all_elements.py                    | 14 ++++----
 tests/test_hydrogen.py                        |  4 +--
 tests/test_matrix_elements.py                 |  4 +--
 tests/test_radial_matrix_element.py           |  6 ++--
 18 files changed, 111 insertions(+), 78 deletions(-)
 create mode 100644 src/rydstate/rydberg/__init__.py
 rename src/rydstate/{rydberg_state.py => rydberg/rydberg_sqdt.py} (92%)

diff --git a/README.md b/README.md
index 7f396d1..15a6aee 100644
--- a/README.md
+++ b/README.md
@@ -83,7 +83,7 @@ This package relies on quantum defects provided by the community. Consider citin
 ## Using custom quantum defects
 To use custom quantum defects (or quantum defects for a new species), you can simply create a subclass of `rydstate.species.species_object.SpeciesObject` (e.g. `class CustomRubidium(SpeciesObject):`) with a custom species name (e.g. `name = "Custom_Rb"`).
 Then, similarly to `rydstate.species.rubidium.py` you can define the quantum defects (and model potential parameters, ...) for your species.
-Finally, you can use the custom species by simply calling `rydstate.RydbergStateAlkali("Custom_Rb", n=50, l=0, j=1/2, m=1/2)` (the code will look for all subclasses of `SpeciesObject` until it finds one with the species name "Custom_Rb").
+Finally, you can use the custom species by simply calling `rydstate.RydbergStateSQDTAlkali("Custom_Rb", n=50, l=0, j=1/2, m=1/2)` (the code will look for all subclasses of `SpeciesObject` until it finds one with the species name "Custom_Rb").
 
 
 ## License
diff --git a/docs/examples/benchmark/benchmark_njit.ipynb b/docs/examples/benchmark/benchmark_njit.ipynb
index 77bc3bc..c042006 100644
--- a/docs/examples/benchmark/benchmark_njit.ipynb
+++ b/docs/examples/benchmark/benchmark_njit.ipynb
@@ -17,7 +17,7 @@
     "\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali\n",
+    "from rydstate import RydbergStateSQDTAlkali\n",
     "\n",
     "test_cases: list[tuple[str, int, int, bool]] = [\n",
     "    # species, n, l, use_njit\n",
@@ -28,7 +28,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -41,21 +41,21 @@
     "    \"\"\"\n",
     "    # run the integration once to compile the numba function\n",
     "    species, n, l, use_njit = test_cases[0]\n",
-    "    state = RydbergStateAlkali(species, n, l, j=l + 0.5)\n",
+    "    state = RydbergStateSQDTAlkali(species, n, l, j=l + 0.5)\n",
     "    state.radial.create_wavefunction(_use_njit=True)\n",
     "\n",
     "    results = []\n",
     "    for species, n, l, use_njit in test_cases:\n",
     "        # Setup the test function\n",
     "        stmt = (\n",
-    "            \"state = RydbergStateAlkali(species, n, l, j=l+0.5)\\n\"\n",
+    "            \"state = RydbergStateSQDTAlkali(species, n, l, j=l+0.5)\\n\"\n",
     "            \"state.radial.create_grid(dz=1e-3)\\n\"\n",
     "            \"state.radial.create_wavefunction(_use_njit=use_njit)\"\n",
     "        )\n",
     "\n",
     "        # Time the integration multiple times and take average/std\n",
     "        globals_dict = {\n",
-    "            \"RydbergStateAlkali\": RydbergStateAlkali,\n",
+    "            \"RydbergStateSQDTAlkali\": RydbergStateSQDTAlkali,\n",
     "            \"species\": species,\n",
     "            \"n\": n,\n",
     "            \"l\": l,\n",
diff --git a/docs/examples/comparisons/compare_dipole_matrix_element.ipynb b/docs/examples/comparisons/compare_dipole_matrix_element.ipynb
index bfbfce5..2691e3b 100644
--- a/docs/examples/comparisons/compare_dipole_matrix_element.ipynb
+++ b/docs/examples/comparisons/compare_dipole_matrix_element.ipynb
@@ -9,7 +9,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -18,8 +18,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali\n",
-    "from rydstate.units import ureg"
+    "from rydstate import RydbergStateSQDTAlkali, ureg"
    ]
   },
   {
@@ -92,8 +91,8 @@
     "for qn1, qn2 in zip(qn1_list, qn2_list):\n",
     "    print(f\"n={qn1[0]}\", end=\"\\r\")\n",
     "    q = round(qn2[-1] - qn1[-1])\n",
-    "    state_i = RydbergStateAlkali(\"Rb\", n=qn1[0], l=qn1[1], j=qn1[2], m=qn1[3])\n",
-    "    state_f = RydbergStateAlkali(\"Rb\", n=qn2[0], l=qn2[1], j=qn2[2], m=qn2[3])\n",
+    "    state_i = RydbergStateSQDTAlkali(\"Rb\", n=qn1[0], l=qn1[1], j=qn1[2], m=qn1[3])\n",
+    "    state_f = RydbergStateSQDTAlkali(\"Rb\", n=qn2[0], l=qn2[1], j=qn2[2], m=qn2[3])\n",
     "    dipole_me = state_i.calc_matrix_element(state_f, \"electric_dipole\", q, unit=\"a.u.\")\n",
     "    matrixelements.append(dipole_me)\n",
     "\n",
diff --git a/docs/examples/comparisons/compare_model_potentials.ipynb b/docs/examples/comparisons/compare_model_potentials.ipynb
index 5c0a6d4..3b9a7df 100644
--- a/docs/examples/comparisons/compare_model_potentials.ipynb
+++ b/docs/examples/comparisons/compare_model_potentials.ipynb
@@ -9,13 +9,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali, RydbergStateAlkalineLS"
+    "from rydstate import RydbergStateSQDTAlkali, RydbergStateSQDTAlkalineLS"
    ]
   },
   {
@@ -41,15 +41,15 @@
     }
    ],
    "source": [
-    "state = RydbergStateAlkali(\"Rb\", n=40, l=0, j=0.5)\n",
+    "state = RydbergStateSQDTAlkali(\"Rb\", n=40, l=0, j=0.5)\n",
     "\n",
-    "states: dict[str, RydbergStateAlkali] = {}\n",
+    "states: dict[str, RydbergStateSQDTAlkali] = {}\n",
     "\n",
     "\n",
-    "states[\"model_potential_marinescu_1993\"] = RydbergStateAlkali(state.species, n=state.n, l=state.l, j=state.j)\n",
+    "states[\"model_potential_marinescu_1993\"] = RydbergStateSQDTAlkali(state.species, n=state.n, l=state.l, j=state.j)\n",
     "states[\"model_potential_marinescu_1993\"].radial.create_model(potential_type=\"model_potential_marinescu_1993\")\n",
     "\n",
-    "states[\"model_potential_fei_2009\"] = RydbergStateAlkali(state.species, n=state.n, l=state.l, j=state.j)\n",
+    "states[\"model_potential_fei_2009\"] = RydbergStateSQDTAlkali(state.species, n=state.n, l=state.l, j=state.j)\n",
     "states[\"model_potential_fei_2009\"].radial.create_model(potential_type=\"model_potential_fei_2009\")\n",
     "\n",
     "for label, state in states.items():\n",
@@ -122,17 +122,17 @@
     }
    ],
    "source": [
-    "state = RydbergStateAlkalineLS(\"Sr88\", n=8, l=0, j_tot=0, s_tot=0)\n",
+    "state = RydbergStateSQDTAlkalineLS(\"Sr88\", n=8, l=0, j_tot=0, s_tot=0)\n",
     "\n",
-    "states: dict[str, RydbergStateAlkalineLS] = {}\n",
+    "states: dict[str, RydbergStateSQDTAlkalineLS] = {}\n",
     "\n",
     "\n",
-    "states[\"model_potential_marinescu_1993\"] = RydbergStateAlkalineLS(\n",
+    "states[\"model_potential_marinescu_1993\"] = RydbergStateSQDTAlkalineLS(\n",
     "    state.species, n=state.n, l=state.l, j_tot=state.j_tot, s_tot=state.s_tot\n",
     ")\n",
     "states[\"model_potential_marinescu_1993\"].radial.create_model(potential_type=\"model_potential_marinescu_1993\")\n",
     "\n",
-    "states[\"model_potential_fei_2009\"] = RydbergStateAlkalineLS(\n",
+    "states[\"model_potential_fei_2009\"] = RydbergStateSQDTAlkalineLS(\n",
     "    state.species, n=state.n, l=state.l, j_tot=state.j_tot, s_tot=state.s_tot\n",
     ")\n",
     "states[\"model_potential_fei_2009\"].radial.create_model(potential_type=\"model_potential_fei_2009\")\n",
diff --git a/docs/examples/comparisons/compare_radial_matrix_element.ipynb b/docs/examples/comparisons/compare_radial_matrix_element.ipynb
index b8f960c..095ba37 100644
--- a/docs/examples/comparisons/compare_radial_matrix_element.ipynb
+++ b/docs/examples/comparisons/compare_radial_matrix_element.ipynb
@@ -9,7 +9,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -18,8 +18,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali\n",
-    "from rydstate.units import ureg"
+    "from rydstate import RydbergStateSQDTAlkali, ureg"
    ]
   },
   {
@@ -64,8 +63,8 @@
     "    results[key] = []\n",
     "    for qn1, qn2 in zip(qn1_list, qn2_list):\n",
     "        print(f\"n={qn1[0]}\", end=\"\\r\")\n",
-    "        state_i = RydbergStateAlkali(species, qn1[0], qn1[1], j=qn1[2])\n",
-    "        state_f = RydbergStateAlkali(species, qn2[0], qn2[1], j=qn2[2])\n",
+    "        state_i = RydbergStateSQDTAlkali(species, qn1[0], qn1[1], j=qn1[2])\n",
+    "        state_f = RydbergStateSQDTAlkali(species, qn2[0], qn2[1], j=qn2[2])\n",
     "        radial_me = state_i.radial.calc_matrix_element(state_f.radial, 1, unit=\"a.u.\")\n",
     "        results[key].append(radial_me)\n",
     "\n",
diff --git a/docs/examples/comparisons/compare_wavefunctions.ipynb b/docs/examples/comparisons/compare_wavefunctions.ipynb
index af490fa..699739c 100644
--- a/docs/examples/comparisons/compare_wavefunctions.ipynb
+++ b/docs/examples/comparisons/compare_wavefunctions.ipynb
@@ -9,14 +9,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali"
+    "from rydstate import RydbergStateSQDTAlkali"
    ]
   },
   {
@@ -44,7 +44,7 @@
    "source": [
     "results[\"rydstate\"] = []\n",
     "for qn in qns:\n",
-    "    state = RydbergStateAlkali(\"Rb\", n=qn[0], l=qn[1], j=qn[2])\n",
+    "    state = RydbergStateSQDTAlkali(\"Rb\", n=qn[0], l=qn[1], j=qn[2])\n",
     "    state.radial.create_grid()\n",
     "    state.radial.create_wavefunction()\n",
     "    results[\"rydstate\"].append(\n",
@@ -125,7 +125,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -142,8 +142,8 @@
     "# the small difference of the wavefunctions explains the difference in the radial matrix element of circular states\n",
     "# (see also the compare_radial_matrix_element and compare_dipole_matrix_element notebooks)\n",
     "\n",
+    "from rydstate import ureg\n",
     "from rydstate.radial.radial_matrix_element import calc_radial_matrix_element_from_w_z\n",
-    "from rydstate.units import ureg\n",
     "\n",
     "to_mum = ureg.Quantity(1, \"bohr_radius\").to(\"micrometer\").magnitude\n",
     "\n",
diff --git a/docs/examples/comparisons/compare_whittaker.ipynb b/docs/examples/comparisons/compare_whittaker.ipynb
index ddd36c7..2bab6ab 100644
--- a/docs/examples/comparisons/compare_whittaker.ipynb
+++ b/docs/examples/comparisons/compare_whittaker.ipynb
@@ -21,7 +21,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate import RydbergStateAlkali\n",
+    "from rydstate import RydbergStateSQDTAlkali\n",
     "\n",
     "logging.basicConfig(level=logging.INFO, format=\"%(levelname)s %(filename)s: %(message)s\")"
    ]
@@ -41,7 +41,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "states: dict[str, RydbergStateAlkali] = {}"
+    "states: dict[str, RydbergStateSQDTAlkali] = {}"
    ]
   },
   {
@@ -62,7 +62,7 @@
     }
    ],
    "source": [
-    "state = RydbergStateAlkali(\"Rb\", n=21, l=0, j=0.5)\n",
+    "state = RydbergStateSQDTAlkali(\"Rb\", n=21, l=0, j=0.5)\n",
     "\n",
     "state.radial.create_model(potential_type=\"model_potential_marinescu_1993\")\n",
     "state.radial.create_wavefunction(\"numerov\")\n",
@@ -70,7 +70,7 @@
     "\n",
     "# Using Numerov without model potentials will lead to some warnings,\n",
     "# since the resulting wavefunction does not pass all heuristic checks\n",
-    "state_without_mp = RydbergStateAlkali(state.species, state.n, state.l, state.j)\n",
+    "state_without_mp = RydbergStateSQDTAlkali(state.species, state.n, state.l, state.j)\n",
     "state_without_mp.radial.create_model(potential_type=\"coulomb\")\n",
     "state_without_mp.radial.create_wavefunction(\"numerov\")\n",
     "states[\"Numerov without Model Potentials\"] = state_without_mp"
@@ -91,7 +91,7 @@
     }
    ],
    "source": [
-    "state_whittaker = RydbergStateAlkali(state.species, state.n, state.l, state.j)\n",
+    "state_whittaker = RydbergStateSQDTAlkali(state.species, state.n, state.l, state.j)\n",
     "state_whittaker.radial.create_grid(x_min=state.radial.grid.x_min, x_max=state.radial.grid.x_max)\n",
     "state_whittaker.radial.create_wavefunction(\"whittaker\")\n",
     "states[\"Whittaker\"] = state_whittaker"
@@ -188,15 +188,15 @@
     }
    ],
    "source": [
-    "state1 = RydbergStateAlkali(\"Rb\", n=10, l=0, j=0.5)\n",
-    "state2 = RydbergStateAlkali(\"Rb\", n=9, l=1, j=1.5)\n",
+    "state1 = RydbergStateSQDTAlkali(\"Rb\", n=10, l=0, j=0.5)\n",
+    "state2 = RydbergStateSQDTAlkali(\"Rb\", n=9, l=1, j=1.5)\n",
     "\n",
     "dipole_me = state1.radial.calc_matrix_element(state2.radial, 1)\n",
     "print(f\"Numerov with model potentials: {dipole_me}\", flush=True)\n",
     "\n",
-    "_state1 = RydbergStateAlkali(state1.species, state1.n, state1.l, state1.j)\n",
+    "_state1 = RydbergStateSQDTAlkali(state1.species, state1.n, state1.l, state1.j)\n",
     "_state1.radial.create_model(potential_type=\"coulomb\")\n",
-    "_state2 = RydbergStateAlkali(state2.species, state2.n, state2.l, state2.j)\n",
+    "_state2 = RydbergStateSQDTAlkali(state2.species, state2.n, state2.l, state2.j)\n",
     "_state2.radial.create_model(potential_type=\"coulomb\")\n",
     "\n",
     "dipole_me = _state1.radial.calc_matrix_element(_state2.radial, 1)\n",
@@ -206,10 +206,10 @@
     "# to avoid integrating over the diverging peak at the origin (see plots above)\n",
     "xmin1, xmax1 = _state1.radial.grid.x_min, _state1.radial.grid.x_max\n",
     "xmin2, xmax2 = _state2.radial.grid.x_min, _state2.radial.grid.x_max\n",
-    "_state1 = RydbergStateAlkali(state1.species, state1.n, state1.l, state1.j)\n",
+    "_state1 = RydbergStateSQDTAlkali(state1.species, state1.n, state1.l, state1.j)\n",
     "_state1.radial.create_grid(x_min=xmin1, x_max=xmax1)\n",
     "_state1.radial.create_wavefunction(\"whittaker\")\n",
-    "_state2 = RydbergStateAlkali(state2.species, state2.n, state2.l, state2.j)\n",
+    "_state2 = RydbergStateSQDTAlkali(state2.species, state2.n, state2.l, state2.j)\n",
     "_state2.radial.create_grid(x_min=xmin2, x_max=xmax2)\n",
     "_state2.radial.create_wavefunction(\"whittaker\")\n",
     "\n",
diff --git a/docs/examples/comparisons/compare_z_min_cutoff.ipynb b/docs/examples/comparisons/compare_z_min_cutoff.ipynb
index 1fff6bf..cdcaa1c 100644
--- a/docs/examples/comparisons/compare_z_min_cutoff.ipynb
+++ b/docs/examples/comparisons/compare_z_min_cutoff.ipynb
@@ -9,14 +9,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali"
+    "from rydstate import RydbergStateSQDTAlkali"
    ]
   },
   {
@@ -57,7 +57,7 @@
     "z_i_dict = {\"hydrogen\": [], \"classical\": [], \"rydstate cutoff\": []}\n",
     "for qn in qn_list:\n",
     "    print(f\"n={qn[0]}\", end=\"\\r\")\n",
-    "    state = RydbergStateAlkali(\"Rb\", n=qn[0], l=qn[1], j=qn[2])\n",
+    "    state = RydbergStateSQDTAlkali(\"Rb\", n=qn[0], l=qn[1], j=qn[2])\n",
     "\n",
     "    hydrogen_z_i = state.radial.model.calc_hydrogen_turning_point_z(state.n, state.l)\n",
     "    z_i_dict[\"hydrogen\"].append(hydrogen_z_i)\n",
diff --git a/docs/examples/dipole_matrix_elements.ipynb b/docs/examples/dipole_matrix_elements.ipynb
index c4bc544..9ce7391 100644
--- a/docs/examples/dipole_matrix_elements.ipynb
+++ b/docs/examples/dipole_matrix_elements.ipynb
@@ -9,13 +9,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali"
+    "from rydstate import RydbergStateSQDTAlkali"
    ]
   },
   {
@@ -34,8 +34,8 @@
     }
    ],
    "source": [
-    "state_i = RydbergStateAlkali(\"Rb\", 60, 2, j=3 / 2, m=1 / 2)\n",
-    "state_f = RydbergStateAlkali(\"Rb\", 60, 3, j=5 / 2, m=1 / 2)\n",
+    "state_i = RydbergStateSQDTAlkali(\"Rb\", 60, 2, j=3 / 2, m=1 / 2)\n",
+    "state_f = RydbergStateSQDTAlkali(\"Rb\", 60, 3, j=5 / 2, m=1 / 2)\n",
     "\n",
     "kappa = 1\n",
     "radial = state_i.radial.calc_matrix_element(state_f.radial, k_radial=1)\n",
diff --git a/docs/examples/radial/rubidium_wavefunction.ipynb b/docs/examples/radial/rubidium_wavefunction.ipynb
index de64f17..b8cfad1 100644
--- a/docs/examples/radial/rubidium_wavefunction.ipynb
+++ b/docs/examples/radial/rubidium_wavefunction.ipynb
@@ -18,7 +18,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali\n",
+    "from rydstate import RydbergStateSQDTAlkali\n",
     "\n",
     "logging.basicConfig(level=logging.INFO, format=\"%(levelname)s %(filename)s: %(message)s\")\n",
     "logging.getLogger(\"rydstate\").setLevel(logging.DEBUG)"
@@ -30,7 +30,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "state = RydbergStateAlkali(\"Rb\", n=130, l=129, j=129.5)\n",
+    "state = RydbergStateSQDTAlkali(\"Rb\", n=130, l=129, j=129.5)\n",
     "state.radial.create_wavefunction()\n",
     "\n",
     "turning_points = {\n",
@@ -45,7 +45,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "hydrogen = RydbergStateAlkali(\"H_textbook\", n=state.n, l=state.l, j=state.j)\n",
+    "hydrogen = RydbergStateSQDTAlkali(\"H_textbook\", n=state.n, l=state.l, j=state.j)\n",
     "hydrogen.radial.create_model()\n",
     "hydrogen.radial.create_wavefunction()"
    ]
diff --git a/src/rydstate/__init__.py b/src/rydstate/__init__.py
index 00b275b..bca5aec 100644
--- a/src/rydstate/__init__.py
+++ b/src/rydstate/__init__.py
@@ -1,11 +1,15 @@
 from rydstate import angular, radial, species
-from rydstate.rydberg_state import RydbergStateAlkali, RydbergStateAlkalineJJ, RydbergStateAlkalineLS
+from rydstate.rydberg import (
+    RydbergStateSQDTAlkali,
+    RydbergStateSQDTAlkalineJJ,
+    RydbergStateSQDTAlkalineLS,
+)
 from rydstate.units import ureg
 
 __all__ = [
-    "RydbergStateAlkali",
-    "RydbergStateAlkalineJJ",
-    "RydbergStateAlkalineLS",
+    "RydbergStateSQDTAlkali",
+    "RydbergStateSQDTAlkalineJJ",
+    "RydbergStateSQDTAlkalineLS",
     "angular",
     "radial",
     "species",
diff --git a/src/rydstate/rydberg/__init__.py b/src/rydstate/rydberg/__init__.py
new file mode 100644
index 0000000..96c1a24
--- /dev/null
+++ b/src/rydstate/rydberg/__init__.py
@@ -0,0 +1,13 @@
+from rydstate.rydberg.rydberg_sqdt import (
+    RydbergStateSQDTAlkali,
+    RydbergStateSQDTAlkalineJJ,
+    RydbergStateSQDTAlkalineLS,
+    RydbergStateSQDTBase,
+)
+
+__all__ = [
+    "RydbergStateSQDTAlkali",
+    "RydbergStateSQDTAlkalineJJ",
+    "RydbergStateSQDTAlkalineLS",
+    "RydbergStateSQDTBase",
+]
diff --git a/src/rydstate/rydberg_state.py b/src/rydstate/rydberg/rydberg_sqdt.py
similarity index 92%
rename from src/rydstate/rydberg_state.py
rename to src/rydstate/rydberg/rydberg_sqdt.py
index 7fa6a90..6ad94d6 100644
--- a/src/rydstate/rydberg_state.py
+++ b/src/rydstate/rydberg/rydberg_sqdt.py
@@ -11,7 +11,7 @@
 from rydstate.angular import AngularKetJJ, AngularKetLS
 from rydstate.angular.utils import try_trivial_spin_addition
 from rydstate.radial import RadialKet
-from rydstate.species.species_object import SpeciesObject
+from rydstate.species import SpeciesObject
 from rydstate.species.utils import calc_energy_from_nu
 from rydstate.units import BaseQuantities, MatrixElementOperatorRanks, ureg
 
@@ -25,7 +25,7 @@
 logger = logging.getLogger(__name__)
 
 
-class RydbergStateBase(ABC):
+class RydbergStateSQDTBase(ABC):
     species: SpeciesObject
 
     def __str__(self) -> str:
@@ -70,7 +70,7 @@ def get_energy(self, unit: str | None = None) -> PintFloat | float:
             return energy
         return energy.to(unit, "spectroscopy").magnitude
 
-    def calc_reduced_overlap(self, other: RydbergStateBase) -> float:
+    def calc_reduced_overlap(self, other: RydbergStateSQDTBase) -> float:
         """Calculate the reduced overlap <self|other> (ignoring the magnetic quantum number m)."""
         radial_overlap = self.radial.calc_overlap(other.radial)
         angular_overlap = self.angular.calc_reduced_overlap(other.angular)
@@ -187,7 +187,7 @@ def calc_matrix_element(
         return prefactor * reduced_matrix_element
 
 
-class RydbergStateAlkali(RydbergStateBase):
+class RydbergStateSQDTAlkali(RydbergStateSQDTBase):
     """Create an Alkali Rydberg state, including the radial and angular states."""
 
     def __init__(
@@ -198,6 +198,7 @@ def __init__(
         j: float | None = None,
         f: float | None = None,
         m: float | None = None,
+        nu: float | None = None,
     ) -> None:
         r"""Initialize the Rydberg state.
 
@@ -210,6 +211,8 @@ def __init__(
               Optional, only needed if the species supports hyperfine structure (i.e. species.i_c is not None or 0).
             m: Total magnetic quantum number.
               Optional, only needed for concrete angular matrix elements.
+            nu: Effective principal quantum number of the rydberg electron.
+              Optional, if not given it will be calculated from n, l, j.
 
         """
         if isinstance(species, str):
@@ -221,6 +224,7 @@ def __init__(
         self.j = try_trivial_spin_addition(l, 0.5, j, "j")
         self.f = try_trivial_spin_addition(self.j, i_c, f, "f")
         self.m = m
+        self._nu = nu
 
         if species.number_valence_electrons != 1:
             raise ValueError(f"The species {species.name} is not an alkali atom.")
@@ -242,10 +246,12 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j=}, {f=}, {m=})"
 
     def get_nu(self) -> float:
+        if self._nu is not None:
+            return self._nu
         return self.species.calc_nu(self.n, self.l, self.j, s_tot=1 / 2)
 
 
-class RydbergStateAlkalineLS(RydbergStateBase):
+class RydbergStateSQDTAlkalineLS(RydbergStateSQDTBase):
     """Create an Alkaline Rydberg state, including the radial and angular states."""
 
     def __init__(
@@ -257,6 +263,7 @@ def __init__(
         j_tot: int | None = None,
         f_tot: float | None = None,
         m: float | None = None,
+        nu: float | None = None,
     ) -> None:
         r"""Initialize the Rydberg state.
 
@@ -270,6 +277,8 @@ def __init__(
               Optional, only needed if the species supports hyperfine structure (i.e. species.i_c is not None or 0).
             m: Total magnetic quantum number.
               Optional, only needed for concrete angular matrix elements.
+            nu: Effective principal quantum number of the rydberg electron.
+              Optional, if not given it will be calculated from n, l, j_tot, s_tot.
 
         """
         if isinstance(species, str):
@@ -282,6 +291,7 @@ def __init__(
         self.j_tot = try_trivial_spin_addition(l, s_tot, j_tot, "j_tot")
         self.f_tot = try_trivial_spin_addition(self.j_tot, i_c, f_tot, "f_tot")
         self.m = m
+        self._nu = nu
 
         if species.number_valence_electrons != 2:
             raise ValueError(f"The species {species.name} is not an alkaline atom.")
@@ -305,10 +315,12 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {s_tot=}, {j_tot=}, {f_tot=}, {m=})"
 
     def get_nu(self) -> float:
+        if self._nu is not None:
+            return self._nu
         return self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=self.s_tot)
 
 
-class RydbergStateAlkalineJJ(RydbergStateBase):
+class RydbergStateSQDTAlkalineJJ(RydbergStateSQDTBase):
     """Create an Alkaline Rydberg state, including the radial and angular states."""
 
     def __init__(
@@ -320,6 +332,7 @@ def __init__(
         j_tot: int | None = None,
         f_tot: float | None = None,
         m: float | None = None,
+        nu: float | None = None,
     ) -> None:
         r"""Initialize the Rydberg state.
 
@@ -334,6 +347,8 @@ def __init__(
               (i.e. species.i_c is not None and species.i_c != 0).
             m: Total magnetic quantum number.
               Optional, only needed for concrete angular matrix elements.
+            nu: Effective principal quantum number of the rydberg electron.
+              Optional, if not given it will be calculated from n, l, j_tot.
 
         """
         if isinstance(species, str):
@@ -347,6 +362,7 @@ def __init__(
         self.j_tot = try_trivial_spin_addition(self.j_r, s_c, j_tot, "j_tot")
         self.f_tot = try_trivial_spin_addition(self.j_tot, i_c, f_tot, "f_tot")
         self.m = m
+        self._nu = nu
 
         if species.number_valence_electrons != 2:
             raise ValueError(f"The species {species.name} is not an alkaline atom.")
@@ -371,11 +387,13 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j_r=}, {j_tot=}, {f_tot=}, {m=})"
 
     def get_nu(self) -> float:
+        if self._nu is not None:
+            return self._nu
         nu_singlet = self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=0)
         nu_triplet = self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=1)
         if abs(nu_singlet - nu_triplet) > 1e-10:
             raise ValueError(
-                "RydbergStateAlkalineJJ is intended for high-l states only, "
+                "RydbergStateSQDTAlkalineJJ is intended for high-l states only, "
                 "where the quantum defects are the same for singlet and triplet states."
             )
         return nu_singlet
diff --git a/src/rydstate/species/species_object.py b/src/rydstate/species/species_object.py
index 5ac1e01..71b9928 100644
--- a/src/rydstate/species/species_object.py
+++ b/src/rydstate/species/species_object.py
@@ -178,7 +178,7 @@ def from_name(cls, name: str) -> SpeciesObject:
         This approach allows for easy extension of the library with new species.
         A user can even subclass SpeciesObject in his code (without modifying the rydstate library),
         e.g. `class CustomRubidium(SpeciesObject): name = "Custom_Rb" ...`
-        and then use the new species by calling RydbergStateAlkali("Custom_Rb", ...)
+        and then use the new species by calling RydbergStateSQDTAlkali("Custom_Rb", ...)
 
         Args:
             name: The species name (e.g. "Rb").
diff --git a/tests/test_all_elements.py b/tests/test_all_elements.py
index 8fb56b2..9a6df4d 100644
--- a/tests/test_all_elements.py
+++ b/tests/test_all_elements.py
@@ -1,11 +1,11 @@
 from typing import TYPE_CHECKING
 
 import pytest
-from rydstate.rydberg_state import RydbergStateAlkali, RydbergStateAlkalineLS
+from rydstate import RydbergStateSQDTAlkali, RydbergStateSQDTAlkalineLS
 from rydstate.species import SpeciesObject
 
 if TYPE_CHECKING:
-    from rydstate.rydberg_state import RydbergStateBase
+    from rydstate.rydberg.rydberg_sqdt import RydbergStateSQDTBase
 
 
 @pytest.mark.parametrize("species_name", SpeciesObject.get_available_species())
@@ -13,16 +13,16 @@ def test_magnetic(species_name: str) -> None:
     """Test magnetic units."""
     species = SpeciesObject.from_name(species_name)
 
-    state: RydbergStateBase
+    state: RydbergStateSQDTBase
     if species.number_valence_electrons == 1:
         if species.i_c is None:
-            state = RydbergStateAlkali(species, n=50, l=0)
+            state = RydbergStateSQDTAlkali(species, n=50, l=0)
         else:
-            state = RydbergStateAlkali(species, n=50, l=0, f=species.i_c + 0.5)
+            state = RydbergStateSQDTAlkali(species, n=50, l=0, f=species.i_c + 0.5)
         state.radial.create_wavefunction()
         with pytest.raises(ValueError, match="j must be set"):
-            RydbergStateAlkali(species, n=50, l=1)
+            RydbergStateSQDTAlkali(species, n=50, l=1)
     elif species.number_valence_electrons == 2 and species._quantum_defects is not None:  # noqa: SLF001
         for s_tot in [0, 1]:
-            state = RydbergStateAlkalineLS(species, n=50, l=1, s_tot=s_tot, j_tot=1 + s_tot)
+            state = RydbergStateSQDTAlkalineLS(species, n=50, l=1, s_tot=s_tot, j_tot=1 + s_tot)
             state.radial.create_wavefunction()
diff --git a/tests/test_hydrogen.py b/tests/test_hydrogen.py
index a28ee13..9ea6863 100644
--- a/tests/test_hydrogen.py
+++ b/tests/test_hydrogen.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from rydstate.rydberg_state import RydbergStateAlkali
+from rydstate import RydbergStateSQDTAlkali
 from sympy.abc import r as sympy_r
 from sympy.physics import hydrogen as sympy_hydrogen
 from sympy.utilities.lambdify import lambdify
@@ -28,7 +28,7 @@
 def test_hydrogen_wavefunctions(species: str, n: int, l: int, run_backward: bool) -> None:
     """Test that numerov integration matches sympy's analytical hydrogen wavefunctions."""
     # Setup atom
-    state = RydbergStateAlkali(species, n=n, l=l, j=l + 0.5)
+    state = RydbergStateSQDTAlkali(species, n=n, l=l, j=l + 0.5)
 
     # Run the numerov integration
     state.radial.create_wavefunction("numerov", run_backward=run_backward, sign_convention="n_l_1")
diff --git a/tests/test_matrix_elements.py b/tests/test_matrix_elements.py
index ed3c37d..2f98c57 100644
--- a/tests/test_matrix_elements.py
+++ b/tests/test_matrix_elements.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from rydstate.rydberg_state import RydbergStateAlkali
+from rydstate import RydbergStateSQDTAlkali
 from rydstate.units import BaseUnits, ureg
 
 
@@ -10,7 +10,7 @@ def test_magnetic(l: int) -> None:
     g_s = 2.002319304363
     g_l = 1
 
-    state = RydbergStateAlkali("Rb", n=max(l + 1, 10), l=l, j=l + 0.5, m=l + 0.5)
+    state = RydbergStateSQDTAlkali("Rb", n=max(l + 1, 10), l=l, j=l + 0.5, m=l + 0.5)
 
     # Check that for m = j_tot = l + s_tot the magnetic matrix element is - mu_B * (g_l * l + g_s * s_tot)
     mu = state.calc_matrix_element(state, "magnetic_dipole", q=0)
diff --git a/tests/test_radial_matrix_element.py b/tests/test_radial_matrix_element.py
index 552e5fd..74d9b70 100644
--- a/tests/test_radial_matrix_element.py
+++ b/tests/test_radial_matrix_element.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
+from rydstate import RydbergStateSQDTAlkali
 from rydstate.radial import RadialKet
-from rydstate.rydberg_state import RydbergStateAlkali
 from rydstate.species import SpeciesObject
 
 
@@ -28,8 +28,8 @@ def test_circular_matrix_element(species: str, n: int, dn: int, dl: int) -> None
 
     matrix_element = {}
     for _species in [species, "H_textbook"]:
-        state_i = RydbergStateAlkali(_species, n=n, l=l1, j=l1 + 0.5)
-        state_f = RydbergStateAlkali(_species, n=n + dn, l=l2, j=l2 + 0.5)
+        state_i = RydbergStateSQDTAlkali(_species, n=n, l=l1, j=l1 + 0.5)
+        state_f = RydbergStateSQDTAlkali(_species, n=n + dn, l=l2, j=l2 + 0.5)
         matrix_element[_species] = state_i.radial.calc_matrix_element(state_f.radial, 1, unit="bohr")
 
     assert np.isclose(matrix_element[species], matrix_element["H_textbook"], rtol=1e-4)

From 370ccb0e2df6bea768900226d42bb7aa9a33e3dd Mon Sep 17 00:00:00 2001
From: johannes-moegerle <johannes.moegerle@itp3.uni-stuttgart.de>
Date: Fri, 19 Dec 2025 14:42:16 +0100
Subject: [PATCH 3/4] make RydbergStateSQDT a proper clss and simplify

---
 src/rydstate/__init__.py             |   2 +
 src/rydstate/angular/angular_ket.py  |  50 ++++++
 src/rydstate/rydberg/__init__.py     |   4 +-
 src/rydstate/rydberg/rydberg_sqdt.py | 233 +++++++++++++++------------
 tests/test_all_elements.py           |  14 +-
 5 files changed, 186 insertions(+), 117 deletions(-)

diff --git a/src/rydstate/__init__.py b/src/rydstate/__init__.py
index bca5aec..4088e9f 100644
--- a/src/rydstate/__init__.py
+++ b/src/rydstate/__init__.py
@@ -1,5 +1,6 @@
 from rydstate import angular, radial, species
 from rydstate.rydberg import (
+    RydbergStateSQDT,
     RydbergStateSQDTAlkali,
     RydbergStateSQDTAlkalineJJ,
     RydbergStateSQDTAlkalineLS,
@@ -7,6 +8,7 @@
 from rydstate.units import ureg
 
 __all__ = [
+    "RydbergStateSQDT",
     "RydbergStateSQDTAlkali",
     "RydbergStateSQDTAlkalineJJ",
     "RydbergStateSQDTAlkalineLS",
diff --git a/src/rydstate/angular/angular_ket.py b/src/rydstate/angular/angular_ket.py
index bb01fd3..62ceeba 100644
--- a/src/rydstate/angular/angular_ket.py
+++ b/src/rydstate/angular/angular_ket.py
@@ -708,3 +708,53 @@ 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.
+
+    """
+    if all(qn is None for qn in [j_c, f_c, j_r]):
+        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
+        )
+    if all(qn is None for qn in [s_tot, l_tot, f_c]):
+        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
+        )
+    if all(qn is None for qn in [s_tot, l_tot, j_tot]):
+        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.")
diff --git a/src/rydstate/rydberg/__init__.py b/src/rydstate/rydberg/__init__.py
index 96c1a24..d9bdec4 100644
--- a/src/rydstate/rydberg/__init__.py
+++ b/src/rydstate/rydberg/__init__.py
@@ -1,13 +1,13 @@
 from rydstate.rydberg.rydberg_sqdt import (
+    RydbergStateSQDT,
     RydbergStateSQDTAlkali,
     RydbergStateSQDTAlkalineJJ,
     RydbergStateSQDTAlkalineLS,
-    RydbergStateSQDTBase,
 )
 
 __all__ = [
+    "RydbergStateSQDT",
     "RydbergStateSQDTAlkali",
     "RydbergStateSQDTAlkalineJJ",
     "RydbergStateSQDTAlkalineLS",
-    "RydbergStateSQDTBase",
 ]
diff --git a/src/rydstate/rydberg/rydberg_sqdt.py b/src/rydstate/rydberg/rydberg_sqdt.py
index 6ad94d6..da50022 100644
--- a/src/rydstate/rydberg/rydberg_sqdt.py
+++ b/src/rydstate/rydberg/rydberg_sqdt.py
@@ -2,14 +2,12 @@
 
 import logging
 import math
-from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import TYPE_CHECKING, overload
 
 import numpy as np
 
-from rydstate.angular import AngularKetJJ, AngularKetLS
-from rydstate.angular.utils import try_trivial_spin_addition
+from rydstate.angular.angular_ket import quantum_numbers_to_angular_ket
 from rydstate.radial import RadialKet
 from rydstate.species import SpeciesObject
 from rydstate.species.utils import calc_energy_from_nu
@@ -18,32 +16,119 @@
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    from rydstate.angular.angular_ket import AngularKetBase
+    from rydstate.angular.angular_ket import AngularKetBase, AngularKetJJ, AngularKetLS
     from rydstate.units import MatrixElementOperator, PintFloat
 
 
 logger = logging.getLogger(__name__)
 
 
-class RydbergStateSQDTBase(ABC):
+class RydbergStateSQDT:
     species: SpeciesObject
 
+    def __init__(
+        self,
+        species: str | SpeciesObject,
+        n: int | None = None,
+        nu: float | None = None,
+        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,
+    ) -> None:
+        r"""Initialize the Rydberg state.
+
+        Args:
+            species: Atomic species.
+            n: Principal quantum number of the rydberg electron.
+            nu: Effective principal quantum number of the rydberg electron.
+              Optional, if not given it will be calculated from n, l, j_tot, s_tot.
+            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.
+
+        """
+        if isinstance(species, str):
+            species = SpeciesObject.from_name(species)
+        self.species = species
+
+        self._qns = dict(  # noqa: C408
+            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,
+            s_tot=s_tot,
+            l_tot=l_tot,
+            j_tot=j_tot,
+            f_tot=f_tot,
+            m=m,
+        )
+
+        self.n = n
+        self._nu = nu
+        if nu is None and n is None:
+            raise ValueError("Either n or nu must be given to initialize the Rydberg state.")
+
+    def __repr__(self) -> str:
+        species, n, nu = self.species.name, self.n, self.nu
+        n_str = f", {n=}" if n is not None else ""
+        return f"{self.__class__.__name__}({species=}{n_str}, {nu=}, {self.angular})"
+
     def __str__(self) -> str:
         return self.__repr__()
 
-    @property
-    @abstractmethod
+    @cached_property
     def radial(self) -> RadialKet:
         """The radial part of the Rydberg electron."""
+        radial_ket = RadialKet(self.species, nu=self.nu, l_r=self.angular.l_r)
+        if self.n is not None:
+            radial_ket.set_n_for_sanity_check(self.n)
+            s_tot_list = [self.angular.get_qn("s_tot")] if "s_tot" in self.angular.quantum_number_names else [0, 1]
+            for s_tot in s_tot_list:
+                if not self.species.is_allowed_shell(self.n, self.angular.l_r, s_tot=s_tot):
+                    raise ValueError(
+                        f"The shell (n={self.n}, l_r={self.angular.l_r}, s_tot={s_tot}) "
+                        f"is not allowed for the species {self.species}."
+                    )
+        return radial_ket
 
-    @property
-    @abstractmethod
+    @cached_property
     def angular(self) -> AngularKetBase:
         """The angular/spin part of the Rydberg electron."""
+        return quantum_numbers_to_angular_ket(species=self.species, **self._qns)  # type: ignore [arg-type]
 
-    @abstractmethod
-    def get_nu(self) -> float:
-        """Get the effective principal quantum number nu (for alkali atoms also known as n*) for the Rydberg state."""
+    @cached_property
+    def nu(self) -> float:
+        """The effective principal quantum number nu (for alkali atoms also known as n*) for the Rydberg state."""
+        if self._nu is not None:
+            return self._nu
+        assert self.n is not None
+        if any(qn not in self.angular.quantum_number_names for qn in ["j_tot", "s_tot"]):
+            raise ValueError("j_tot and s_tot must be defined to calculate nu from n.")
+        return self.species.calc_nu(
+            self.n, self.angular.l_r, self.angular.get_qn("j_tot"), s_tot=self.angular.get_qn("s_tot")
+        )
 
     @overload
     def get_energy(self, unit: None = None) -> PintFloat: ...
@@ -61,8 +146,7 @@ def get_energy(self, unit: str | None = None) -> PintFloat | float:
 
         where `\mu = R_M/R_\infty` is the reduced mass and `\nu` the effective principal quantum number.
         """
-        nu = self.get_nu()
-        energy_au = calc_energy_from_nu(self.species.reduced_mass_au, nu)
+        energy_au = calc_energy_from_nu(self.species.reduced_mass_au, self.nu)
         if unit == "a.u.":
             return energy_au
         energy: PintFloat = energy_au * BaseQuantities["energy"]
@@ -70,7 +154,7 @@ def get_energy(self, unit: str | None = None) -> PintFloat | float:
             return energy
         return energy.to(unit, "spectroscopy").magnitude
 
-    def calc_reduced_overlap(self, other: RydbergStateSQDTBase) -> float:
+    def calc_reduced_overlap(self, other: RydbergStateSQDT) -> float:
         """Calculate the reduced overlap <self|other> (ignoring the magnetic quantum number m)."""
         radial_overlap = self.radial.calc_overlap(other.radial)
         angular_overlap = self.angular.calc_reduced_overlap(other.angular)
@@ -187,9 +271,11 @@ def calc_matrix_element(
         return prefactor * reduced_matrix_element
 
 
-class RydbergStateSQDTAlkali(RydbergStateSQDTBase):
+class RydbergStateSQDTAlkali(RydbergStateSQDT):
     """Create an Alkali Rydberg state, including the radial and angular states."""
 
+    angular: AngularKetLS
+
     def __init__(
         self,
         species: str | SpeciesObject,
@@ -215,45 +301,23 @@ def __init__(
               Optional, if not given it will be calculated from n, l, j.
 
         """
-        if isinstance(species, str):
-            species = SpeciesObject.from_name(species)
-        self.species = species
-        i_c = species.i_c if species.i_c is not None else 0
-        self.n = n
+        super().__init__(species=species, n=n, nu=nu, l_r=l, j_tot=j, f_tot=f, m=m)
+
         self.l = l
-        self.j = try_trivial_spin_addition(l, 0.5, j, "j")
-        self.f = try_trivial_spin_addition(self.j, i_c, f, "f")
+        self.j = self.angular.j_tot
+        self.f = self.angular.f_tot
         self.m = m
-        self._nu = nu
-
-        if species.number_valence_electrons != 1:
-            raise ValueError(f"The species {species.name} is not an alkali atom.")
-        if not species.is_allowed_shell(n, l):
-            raise ValueError(f"The shell ({n=}, {l=}) is not allowed for the species {self.species}.")
-
-    @cached_property
-    def angular(self) -> AngularKetLS:
-        return AngularKetLS(l_r=self.l, j_tot=self.j, m=self.m, f_tot=self.f, species=self.species)
-
-    @cached_property
-    def radial(self) -> RadialKet:
-        radial_ket = RadialKet(self.species, nu=self.get_nu(), l_r=self.l)
-        radial_ket.set_n_for_sanity_check(self.n)
-        return radial_ket
 
     def __repr__(self) -> str:
         species, n, l, j, f, m = self.species, self.n, self.l, self.j, self.f, self.m
         return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j=}, {f=}, {m=})"
 
-    def get_nu(self) -> float:
-        if self._nu is not None:
-            return self._nu
-        return self.species.calc_nu(self.n, self.l, self.j, s_tot=1 / 2)
 
-
-class RydbergStateSQDTAlkalineLS(RydbergStateSQDTBase):
+class RydbergStateSQDTAlkalineLS(RydbergStateSQDT):
     """Create an Alkaline Rydberg state, including the radial and angular states."""
 
+    angular: AngularKetLS
+
     def __init__(
         self,
         species: str | SpeciesObject,
@@ -281,48 +345,24 @@ def __init__(
               Optional, if not given it will be calculated from n, l, j_tot, s_tot.
 
         """
-        if isinstance(species, str):
-            species = SpeciesObject.from_name(species)
-        self.species = species
-        i_c = species.i_c if species.i_c is not None else 0
-        self.n = n
+        super().__init__(species=species, n=n, nu=nu, l_r=l, s_tot=s_tot, j_tot=j_tot, f_tot=f_tot, m=m)
+
         self.l = l
-        self.s_tot = s_tot
-        self.j_tot = try_trivial_spin_addition(l, s_tot, j_tot, "j_tot")
-        self.f_tot = try_trivial_spin_addition(self.j_tot, i_c, f_tot, "f_tot")
+        self.s_tot = self.angular.s_tot
+        self.j_tot = self.angular.j_tot
+        self.f_tot = self.angular.f_tot
         self.m = m
-        self._nu = nu
-
-        if species.number_valence_electrons != 2:
-            raise ValueError(f"The species {species.name} is not an alkaline atom.")
-        if not species.is_allowed_shell(n, l, s_tot=s_tot):
-            raise ValueError(f"The shell ({n=}, {l=}) is not allowed for the species {self.species}.")
-
-    @cached_property
-    def angular(self) -> AngularKetLS:
-        return AngularKetLS(
-            l_r=self.l, s_tot=self.s_tot, j_tot=self.j_tot, f_tot=self.f_tot, m=self.m, species=self.species
-        )
-
-    @cached_property
-    def radial(self) -> RadialKet:
-        radial_ket = RadialKet(self.species, nu=self.get_nu(), l_r=self.l)
-        radial_ket.set_n_for_sanity_check(self.n)
-        return radial_ket
 
     def __repr__(self) -> str:
         species, n, l, s_tot, j_tot, f_tot, m = self.species, self.n, self.l, self.s_tot, self.j_tot, self.f_tot, self.m
         return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {s_tot=}, {j_tot=}, {f_tot=}, {m=})"
 
-    def get_nu(self) -> float:
-        if self._nu is not None:
-            return self._nu
-        return self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=self.s_tot)
-
 
-class RydbergStateSQDTAlkalineJJ(RydbergStateSQDTBase):
+class RydbergStateSQDTAlkalineJJ(RydbergStateSQDT):
     """Create an Alkaline Rydberg state, including the radial and angular states."""
 
+    angular: AngularKetJJ
+
     def __init__(
         self,
         species: str | SpeciesObject,
@@ -351,44 +391,23 @@ def __init__(
               Optional, if not given it will be calculated from n, l, j_tot.
 
         """
-        if isinstance(species, str):
-            species = SpeciesObject.from_name(species)
-        self.species = species
-        s_r, s_c = 1 / 2, 1 / 2
-        i_c = species.i_c if species.i_c is not None else 0
-        self.n = n
-        self.l = l
-        self.j_r = try_trivial_spin_addition(l, s_r, j_r, "j_r")
-        self.j_tot = try_trivial_spin_addition(self.j_r, s_c, j_tot, "j_tot")
-        self.f_tot = try_trivial_spin_addition(self.j_tot, i_c, f_tot, "f_tot")
-        self.m = m
-        self._nu = nu
+        super().__init__(species=species, n=n, nu=nu, l_r=l, j_r=j_r, j_tot=j_tot, f_tot=f_tot, m=m)
 
-        if species.number_valence_electrons != 2:
-            raise ValueError(f"The species {species.name} is not an alkaline atom.")
-        for s_tot in [0, 1]:
-            if not species.is_allowed_shell(n, l, s_tot=s_tot):
-                raise ValueError(f"The shell ({n=}, {l=}) is not allowed for the species {self.species}.")
-
-    @cached_property
-    def angular(self) -> AngularKetJJ:
-        return AngularKetJJ(
-            l_r=self.l, j_r=self.j_r, j_tot=self.j_tot, f_tot=self.f_tot, m=self.m, species=self.species
-        )
-
-    @cached_property
-    def radial(self) -> RadialKet:
-        radial_ket = RadialKet(self.species, nu=self.get_nu(), l_r=self.l)
-        radial_ket.set_n_for_sanity_check(self.n)
-        return radial_ket
+        self.l = self.angular.l_r
+        self.j_r = self.angular.j_r
+        self.j_tot = self.angular.j_tot
+        self.f_tot = self.angular.f_tot
+        self.m = self.angular.m
 
     def __repr__(self) -> str:
         species, n, l, j_r, j_tot, f_tot, m = self.species, self.n, self.l, self.j_r, self.j_tot, self.f_tot, self.m
         return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j_r=}, {j_tot=}, {f_tot=}, {m=})"
 
-    def get_nu(self) -> float:
+    @cached_property
+    def nu(self) -> float:
         if self._nu is not None:
             return self._nu
+        assert self.n is not None
         nu_singlet = self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=0)
         nu_triplet = self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=1)
         if abs(nu_singlet - nu_triplet) > 1e-10:
diff --git a/tests/test_all_elements.py b/tests/test_all_elements.py
index 9a6df4d..6321b05 100644
--- a/tests/test_all_elements.py
+++ b/tests/test_all_elements.py
@@ -5,24 +5,22 @@
 from rydstate.species import SpeciesObject
 
 if TYPE_CHECKING:
-    from rydstate.rydberg.rydberg_sqdt import RydbergStateSQDTBase
+    from rydstate import RydbergStateSQDT
 
 
 @pytest.mark.parametrize("species_name", SpeciesObject.get_available_species())
 def test_magnetic(species_name: str) -> None:
     """Test magnetic units."""
     species = SpeciesObject.from_name(species_name)
+    i_c = species.i_c if species.i_c is not None else 0
 
-    state: RydbergStateSQDTBase
+    state: RydbergStateSQDT
     if species.number_valence_electrons == 1:
-        if species.i_c is None:
-            state = RydbergStateSQDTAlkali(species, n=50, l=0)
-        else:
-            state = RydbergStateSQDTAlkali(species, n=50, l=0, f=species.i_c + 0.5)
+        state = RydbergStateSQDTAlkali(species, n=50, l=0, f=i_c + 0.5)
         state.radial.create_wavefunction()
-        with pytest.raises(ValueError, match="j must be set"):
+        with pytest.raises(ValueError, match="j_tot must be set"):
             RydbergStateSQDTAlkali(species, n=50, l=1)
     elif species.number_valence_electrons == 2 and species._quantum_defects is not None:  # noqa: SLF001
         for s_tot in [0, 1]:
-            state = RydbergStateSQDTAlkalineLS(species, n=50, l=1, s_tot=s_tot, j_tot=1 + s_tot)
+            state = RydbergStateSQDTAlkalineLS(species, n=50, l=1, s_tot=s_tot, j_tot=1 + s_tot, f_tot=s_tot + 1 + i_c)
             state.radial.create_wavefunction()

From a442c685780a779afca4018134c5d2c780d1c619 Mon Sep 17 00:00:00 2001
From: johannes-moegerle <johannes.moegerle@itp3.uni-stuttgart.de>
Date: Fri, 19 Dec 2025 16:13:25 +0100
Subject: [PATCH 4/4] add rydberg base

---
 src/rydstate/rydberg/rydberg_base.py | 26 +++++++++++++++++++++++
 src/rydstate/rydberg/rydberg_sqdt.py | 31 ++++++++++++++++++----------
 2 files changed, 46 insertions(+), 11 deletions(-)
 create mode 100644 src/rydstate/rydberg/rydberg_base.py

diff --git a/src/rydstate/rydberg/rydberg_base.py b/src/rydstate/rydberg/rydberg_base.py
new file mode 100644
index 0000000..e13e00c
--- /dev/null
+++ b/src/rydstate/rydberg/rydberg_base.py
@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from rydstate.angular import AngularState
+    from rydstate.angular.angular_ket import AngularKetBase
+    from rydstate.units import MatrixElementOperator, PintFloat
+
+logger = logging.getLogger(__name__)
+
+
+class RydbergStateBase(ABC):
+    @property
+    @abstractmethod
+    def angular(self) -> AngularState[Any] | AngularKetBase: ...
+
+    @abstractmethod
+    def calc_reduced_overlap(self, other: RydbergStateBase) -> float: ...
+
+    @abstractmethod
+    def calc_reduced_matrix_element(
+        self, other: RydbergStateBase, operator: MatrixElementOperator, unit: str | None = None
+    ) -> PintFloat | float: ...
diff --git a/src/rydstate/rydberg/rydberg_sqdt.py b/src/rydstate/rydberg/rydberg_sqdt.py
index da50022..0ca1a9e 100644
--- a/src/rydstate/rydberg/rydberg_sqdt.py
+++ b/src/rydstate/rydberg/rydberg_sqdt.py
@@ -9,13 +9,12 @@
 
 from rydstate.angular.angular_ket import quantum_numbers_to_angular_ket
 from rydstate.radial import RadialKet
+from rydstate.rydberg.rydberg_base import RydbergStateBase
 from rydstate.species import SpeciesObject
 from rydstate.species.utils import calc_energy_from_nu
 from rydstate.units import BaseQuantities, MatrixElementOperatorRanks, ureg
 
 if TYPE_CHECKING:
-    from typing_extensions import Self
-
     from rydstate.angular.angular_ket import AngularKetBase, AngularKetJJ, AngularKetLS
     from rydstate.units import MatrixElementOperator, PintFloat
 
@@ -23,7 +22,7 @@
 logger = logging.getLogger(__name__)
 
 
-class RydbergStateSQDT:
+class RydbergStateSQDT(RydbergStateBase):
     species: SpeciesObject
 
     def __init__(
@@ -154,22 +153,27 @@ def get_energy(self, unit: str | None = None) -> PintFloat | float:
             return energy
         return energy.to(unit, "spectroscopy").magnitude
 
-    def calc_reduced_overlap(self, other: RydbergStateSQDT) -> float:
+    def calc_reduced_overlap(self, other: RydbergStateBase) -> float:
         """Calculate the reduced overlap <self|other> (ignoring the magnetic quantum number m)."""
+        if not isinstance(other, RydbergStateSQDT):
+            raise NotImplementedError("Reduced overlap only implemented between RydbergStateSQDT states.")
+
         radial_overlap = self.radial.calc_overlap(other.radial)
         angular_overlap = self.angular.calc_reduced_overlap(other.angular)
         return radial_overlap * angular_overlap
 
-    @overload
+    @overload  # type: ignore [override]
     def calc_reduced_matrix_element(
-        self, other: Self, operator: MatrixElementOperator, unit: None = None
+        self, other: RydbergStateBase, operator: MatrixElementOperator, unit: None = None
     ) -> PintFloat: ...
 
     @overload
-    def calc_reduced_matrix_element(self, other: Self, operator: MatrixElementOperator, unit: str) -> float: ...
+    def calc_reduced_matrix_element(
+        self, other: RydbergStateBase, operator: MatrixElementOperator, unit: str
+    ) -> float: ...
 
     def calc_reduced_matrix_element(
-        self, other: Self, operator: MatrixElementOperator, unit: str | None = None
+        self, other: RydbergStateBase, operator: MatrixElementOperator, unit: str | None = None
     ) -> PintFloat | float:
         r"""Calculate the reduced matrix element.
 
@@ -192,6 +196,9 @@ def calc_reduced_matrix_element(
             The reduced matrix element for the given operator.
 
         """
+        if not isinstance(other, RydbergStateSQDT):
+            raise NotImplementedError("Reduced matrix element only implemented between RydbergStateSQDT states.")
+
         if operator not in MatrixElementOperatorRanks:
             raise ValueError(
                 f"Operator {operator} not supported, must be one of {list(MatrixElementOperatorRanks.keys())}."
@@ -234,13 +241,15 @@ def calc_reduced_matrix_element(
         return matrix_element.to(unit).magnitude
 
     @overload
-    def calc_matrix_element(self, other: Self, operator: MatrixElementOperator, q: int) -> PintFloat: ...
+    def calc_matrix_element(self, other: RydbergStateSQDT, operator: MatrixElementOperator, q: int) -> PintFloat: ...
 
     @overload
-    def calc_matrix_element(self, other: Self, operator: MatrixElementOperator, q: int, unit: str) -> float: ...
+    def calc_matrix_element(
+        self, other: RydbergStateSQDT, operator: MatrixElementOperator, q: int, unit: str
+    ) -> float: ...
 
     def calc_matrix_element(
-        self, other: Self, operator: MatrixElementOperator, q: int, unit: str | None = None
+        self, other: RydbergStateSQDT, operator: MatrixElementOperator, q: int, unit: str | None = None
     ) -> PintFloat | float:
         r"""Calculate the matrix element.