Species sqdt

README.md

@@ -81,9 +81,9 @@ 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.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").
+To use custom quantum defects (or quantum defects for a new species), you can simply create a subclass of `rydstate.species.SpeciesObjectSQDT` (e.g. `class CustomRubidium(SpeciesObjectSQDT):`) with a custom species name (e.g. `name = "Custom_Rb"`).
+Then, similarly to `rydstate.species.sqdt.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.RydbergStateSQDTAlkali("Custom_Rb", n=50, l=0, j=1/2, m=1/2)` (the code will look for all subclasses of `SpeciesObjectSQDT` until it finds one with the species name "Custom_Rb").
 
 
 ## License

docs/api_reference.rst

@@ -61,7 +61,7 @@ All the available classes, methods and functions are documented below:
 .. autosummary::
     :toctree: _autosummary/
 
-       species.SpeciesObject
+       species.SpeciesObjectSQDT
        species.HydrogenTextBook
        species.Hydrogen
        species.Lithium

docs/examples/comparisons/compare_nist_energy_levels_data.ipynb

@@ -9,7 +9,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -19,12 +19,13 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
+    "from rydstate.angular import AngularKetLS\n",
     "from rydstate.species import Strontium88"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -47,8 +48,9 @@
     "                continue\n",
     "\n",
     "            labels.append(f\"{n}{l_int2str.get(l, ',' + str(l))}_{j_tot:.1f}\")\n",
-    "            nus_with_nist.append(species.calc_nu(n, l, j_tot, s_tot, use_nist_data=True, nist_n_max=60))\n",
-    "            nus_without_nist.append(species.calc_nu(n, l, j_tot, s_tot, use_nist_data=False))\n",
+    "            angular_ket = AngularKetLS(l_r=l, j_tot=j_tot, s_tot=s_tot, species=species)\n",
+    "            nus_with_nist.append(species.calc_nu(n, angular_ket, use_nist_data=True, nist_n_max=60))\n",
+    "            nus_without_nist.append(species.calc_nu(n, angular_ket, use_nist_data=False))\n",
     "            n_list.append(n)\n",
     "\n",
     "nus_with_nist = np.array(nus_with_nist)\n",

pyproject.toml

@@ -85,7 +85,7 @@ version = {attr = "rydstate.__version__"}
 where = ["src"]
 
 [tool.setuptools.package-data]
-rydstate = ["species/nist_energy_levels/*.txt"]
+rydstate = ["species/sqdt/nist_energy_levels/*.txt"]
 
 [tool.check-wheel-contents]
 toplevel = ["rydstate"]

src/rydstate/angular/angular_ket.py

@@ -90,9 +90,9 @@ def __init__(
         """
         if species is not None:
             if isinstance(species, str):
-                from rydstate.species import SpeciesObject  # noqa: PLC0415
+                from rydstate.species import SpeciesObjectSQDT  # noqa: PLC0415
 
-                species = SpeciesObject.from_name(species)
+                species = SpeciesObjectSQDT.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
             if i_c is not None and i_c != species_i_c:

src/rydstate/basis/basis_sqdt.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -12,16 +11,16 @@
     RydbergStateSQDTAlkalineJJ,
     RydbergStateSQDTAlkalineLS,
 )
-
-if TYPE_CHECKING:
-    from rydstate.species.species_object import SpeciesObject
+from rydstate.species import SpeciesObjectSQDT
 
 logger = logging.getLogger(__name__)
 
 
 class BasisSQDTAlkali(BasisBase[RydbergStateSQDTAlkali]):
-    def __init__(self, species: str | SpeciesObject, n_min: int = 1, n_max: int | None = None) -> None:
-        super().__init__(species)
+    def __init__(self, species: str | SpeciesObjectSQDT, n_min: int = 1, n_max: int | None = None) -> None:
+        if isinstance(species, str):
+            species = SpeciesObjectSQDT.from_name(species)
+        self.species = species
 
         if n_max is None:
             raise ValueError("n_max must be given")
@@ -41,8 +40,10 @@ def __init__(self, species: str | SpeciesObject, n_min: int = 1, n_max: int | No
 
 
 class BasisSQDTAlkalineLS(BasisBase[RydbergStateSQDTAlkalineLS]):
-    def __init__(self, species: str | SpeciesObject, n_min: int = 1, n_max: int | None = None) -> None:
-        super().__init__(species)
+    def __init__(self, species: str | SpeciesObjectSQDT, n_min: int = 1, n_max: int | None = None) -> None:
+        if isinstance(species, str):
+            species = SpeciesObjectSQDT.from_name(species)
+        self.species = species
 
         if n_max is None:
             raise ValueError("n_max must be given")
@@ -64,8 +65,10 @@ def __init__(self, species: str | SpeciesObject, n_min: int = 1, n_max: int | No
 
 
 class BasisSQDTAlkalineJJ(BasisBase[RydbergStateSQDTAlkalineJJ]):
-    def __init__(self, species: str | SpeciesObject, n_min: int = 0, n_max: int | None = None) -> None:
-        super().__init__(species)
+    def __init__(self, species: str | SpeciesObjectSQDT, n_min: int = 0, n_max: int | None = None) -> None:
+        if isinstance(species, str):
+            species = SpeciesObjectSQDT.from_name(species)
+        self.species = species
 
         if n_max is None:
             raise ValueError("n_max must be given")
@@ -94,8 +97,10 @@ def __init__(self, species: str | SpeciesObject, n_min: int = 0, n_max: int | No
 
 
 class BasisSQDTAlkalineFJ(BasisBase[RydbergStateSQDTAlkalineFJ]):
-    def __init__(self, species: str | SpeciesObject, n_min: int = 0, n_max: int | None = None) -> None:
-        super().__init__(species)
+    def __init__(self, species: str | SpeciesObjectSQDT, n_min: int = 0, n_max: int | None = None) -> None:
+        if isinstance(species, str):
+            species = SpeciesObjectSQDT.from_name(species)
+        self.species = species
 
         if n_max is None:
             raise ValueError("n_max must be given")

src/rydstate/rydberg/rydberg_sqdt.py

@@ -10,7 +10,7 @@
 from rydstate.angular.utils 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 import SpeciesObjectSQDT
 from rydstate.species.utils import calc_energy_from_nu
 from rydstate.units import BaseQuantities, MatrixElementOperatorRanks, ureg
 
@@ -25,15 +25,15 @@
 
 
 class RydbergStateSQDT(RydbergStateBase):
-    species: SpeciesObject
+    species: SpeciesObjectSQDT
     """The atomic species of the Rydberg state."""
 
     angular: AngularKetBase
     """The angular/spin part of the Rydberg electron."""
 
     def __init__(
         self,
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         n: int | None = None,
         nu: float | None = None,
         s_c: float | None = None,
@@ -72,7 +72,7 @@ def __init__(
 
         """
         if isinstance(species, str):
-            species = SpeciesObject.from_name(species)
+            species = SpeciesObjectSQDT.from_name(species)
         self.species = species
 
         self.angular = quantum_numbers_to_angular_ket(
@@ -104,7 +104,7 @@ def _set_qn_as_attributes(self) -> None:
     @classmethod
     def from_angular_ket(
         cls: type[Self],
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         angular_ket: AngularKetBase,
         n: int | None = None,
         nu: float | None = None,
@@ -113,7 +113,7 @@ def from_angular_ket(
         obj = cls.__new__(cls)
 
         if isinstance(species, str):
-            species = SpeciesObject.from_name(species)
+            species = SpeciesObjectSQDT.from_name(species)
         obj.species = species
 
         obj.n = n
@@ -145,28 +145,24 @@ def radial(self) -> RadialKet:
         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}."
-                    )
+            if isinstance(self.species, SpeciesObjectSQDT):
+                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
 
     @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 isinstance(self.species, SpeciesObject), "nu must be given if not sqdt"
+        assert isinstance(self.species, SpeciesObjectSQDT), "nu must be given if not sqdt"
         assert self.n is not None, "either nu or n must be given"
-
-        if "j_tot" not in self.angular.quantum_number_names or "s_tot" not in self.angular.quantum_number_names:
-            raise RuntimeError("j_tot and s_tot must be defined in the angular ket 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")
-        )
+        return self.species.calc_nu(self.n, self.angular)
 
     @property
     def nu_ref(self) -> float:
@@ -330,7 +326,7 @@ class RydbergStateSQDTAlkali(RydbergStateSQDT):
 
     def __init__(
         self,
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         n: int,
         l: int,
         j: float | None = None,
@@ -377,7 +373,7 @@ class RydbergStateSQDTAlkalineLS(RydbergStateSQDT):
 
     def __init__(
         self,
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         n: int,
         l: int,
         s_tot: int,
@@ -426,7 +422,7 @@ class RydbergStateSQDTAlkalineJJ(RydbergStateSQDT):
 
     def __init__(
         self,
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         n: int,
         l: int,
         j_r: float,
@@ -475,7 +471,7 @@ class RydbergStateSQDTAlkalineFJ(RydbergStateSQDT):
 
     def __init__(
         self,
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         n: int,
         l: int,
         j_r: float,

src/rydstate/species/__init__.py

@@ -1,12 +1,19 @@
-from rydstate.species.cesium import Cesium
-from rydstate.species.hydrogen import Hydrogen, HydrogenTextBook
-from rydstate.species.lithium import Lithium
-from rydstate.species.potassium import Potassium
-from rydstate.species.rubidium import Rubidium
-from rydstate.species.sodium import Sodium
 from rydstate.species.species_object import SpeciesObject
-from rydstate.species.strontium import Strontium87, Strontium88
-from rydstate.species.ytterbium import Ytterbium171, Ytterbium173, Ytterbium174
+from rydstate.species.sqdt import (
+    Cesium,
+    Hydrogen,
+    HydrogenTextBook,
+    Lithium,
+    Potassium,
+    Rubidium,
+    Sodium,
+    SpeciesObjectSQDT,
+    Strontium87,
+    Strontium88,
+    Ytterbium171,
+    Ytterbium173,
+    Ytterbium174,
+)
 
 __all__ = [
     "Cesium",
@@ -17,6 +24,7 @@
     "Rubidium",
     "Sodium",
     "SpeciesObject",
+    "SpeciesObjectSQDT",
     "Strontium87",
     "Strontium88",
     "Ytterbium171",