use SpeciesObjectSQDT where appropriate

Author: johannes-moegerle

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

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,21 +145,22 @@ 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"
         return self.species.calc_nu(self.n, self.angular)
 
@@ -325,7 +326,7 @@ class RydbergStateSQDTAlkali(RydbergStateSQDT):
 
     def __init__(
         self,
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         n: int,
         l: int,
         j: float | None = None,
@@ -372,7 +373,7 @@ class RydbergStateSQDTAlkalineLS(RydbergStateSQDT):
 
     def __init__(
         self,
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         n: int,
         l: int,
         s_tot: int,
@@ -421,7 +422,7 @@ class RydbergStateSQDTAlkalineJJ(RydbergStateSQDT):
 
     def __init__(
         self,
-        species: str | SpeciesObject,
+        species: str | SpeciesObjectSQDT,
         n: int,
         l: int,
         j_r: float,
@@ -470,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/species_object.py

@@ -71,8 +71,8 @@ def from_name(cls: type[Self], name: str) -> Self:
 
         This method searches through all subclasses of SpeciesObject until it finds one with a matching species name.
         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" ...`
+        A user can even subclass SpeciesObjectSQDT in his code (without modifying the rydstate library),
+        e.g. `class CustomRubidium(SpeciesObjectSQDT): name = "Custom_Rb" ...`
         and then use the new species by calling RydbergStateSQDTAlkali("Custom_Rb", ...)
 
         Args:

tests/test_all_elements.py

@@ -2,16 +2,15 @@
 
 import pytest
 from rydstate import RydbergStateSQDTAlkali, RydbergStateSQDTAlkalineLS
-from rydstate.species import SpeciesObject
+from rydstate.species import SpeciesObjectSQDT
 
 if TYPE_CHECKING:
     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)
+@pytest.mark.parametrize("species_name", SpeciesObjectSQDT.get_available_species())
+def test_sqdt_species(species_name: str) -> None:
+    species = SpeciesObjectSQDT.from_name(species_name)
     i_c = species.i_c if species.i_c is not None else 0
 
     state: RydbergStateSQDT

tests/test_radial_matrix_element.py

@@ -3,7 +3,7 @@
 from rydstate import RydbergStateSQDTAlkali
 from rydstate.angular import AngularKetLS
 from rydstate.radial import RadialKet
-from rydstate.species import SpeciesObject
+from rydstate.species import SpeciesObjectSQDT
 
 
 @pytest.mark.parametrize(
@@ -59,7 +59,7 @@ def test_circular_expectation_value(species_name: str, n: int, l: int, j_tot: fl
         <r>_{nl} = 1/2 (3 n^2 - l(l+1))
         <r^2>_{nl} = n^2/2 (5 n^2 - 3 l(l+1) + 1)
     """
-    species = SpeciesObject.from_name(species_name)
+    species = SpeciesObjectSQDT.from_name(species_name)
     angular_ket = AngularKetLS(l_r=l, j_tot=j_tot, species=species)
     nu = species.calc_nu(n, angular_ket)