Rename elements to species

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 element), you can simply create a subclass of `ryd_numerov.elements.base_element.BaseElement` (e.g. `class CustomRubidium(BaseElement):`) with a custom species name (e.g. `species = "Custom_Rb"`).
-Then, similarly to `ryd_numerov.elements.rubidium.py` you can define the quantum defects (and model potential parameters, ...) for your element.
-Finally, you can use the custom element by simply calling `ryd_numerov.RydbergStateAlkali("Custom_Rb", n=50, l=0, j=1/2, m=1/2)` (the code will look for all subclasses of `BaseElement` 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 `ryd_numerov.species.species_object.SpeciesObject` (e.g. `class CustomRubidium(SpeciesObject):`) with a custom species name (e.g. `name = "Custom_Rb"`).
+Then, similarly to `ryd_numerov.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 `ryd_numerov.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").
 
 
 ## License

pyproject.toml

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

src/ryd_numerov/__init__.py

@@ -1,4 +1,4 @@
-from ryd_numerov import angular, elements, radial
+from ryd_numerov import angular, radial, species
 from ryd_numerov.rydberg_state import RydbergStateAlkali, RydbergStateAlkaliHyperfine, RydbergStateAlkalineLS
 from ryd_numerov.units import ureg
 
@@ -7,8 +7,8 @@
     "RydbergStateAlkaliHyperfine",
     "RydbergStateAlkalineLS",
     "angular",
-    "elements",
     "radial",
+    "species",
     "ureg",
 ]
 

src/ryd_numerov/angular/angular_ket.py

@@ -21,7 +21,7 @@
     minus_one_pow,
     try_trivial_spin_addition,
 )
-from ryd_numerov.elements import BaseElement
+from ryd_numerov.species import SpeciesObject
 
 if TYPE_CHECKING:
     from ryd_numerov.angular.angular_state import AngularState
@@ -86,7 +86,7 @@ def __init__(
         l_r: int | None = None,
         f_tot: float | None = None,  # noqa: ARG002
         m: float | None = None,
-        species: str | None = None,
+        species: str | SpeciesObject | None = None,
     ) -> None:
         """Initialize the Spin ket.
 
@@ -95,11 +95,12 @@ def __init__(
         Not used for calculation, only for convenience to infer the core electron spin and nuclear spin quantum numbers.
         """
         if species is not None:
-            element = BaseElement.from_species(species)
-            if i_c is not None and i_c != element.i_c:
-                raise ValueError(f"Nuclear spin i_c={i_c} does not match the element {species} with i_c={element.i_c}.")
-            i_c = element.i_c
-            s_c = 0.5 * (element.number_valence_electrons - 1)
+            if isinstance(species, str):
+                species = SpeciesObject.from_name(species)
+            if i_c is not None and i_c != species.i_c:
+                raise ValueError(f"Nuclear spin i_c={i_c} does not match the species {species} with i_c={species.i_c}.")
+            i_c = species.i_c
+            s_c = 0.5 * (species.number_valence_electrons - 1)
         if i_c is None:
             raise ValueError("Nuclear spin i_c must be set or a species must be given.")
         self.i_c = float(i_c)
@@ -548,7 +549,7 @@ def __init__(
         j_tot: float | None = None,
         f_tot: float | None = None,
         m: float | None = None,
-        species: str | None = None,
+        species: str | SpeciesObject | None = None,
     ) -> None:
         """Initialize the Spin ket."""
         super().__init__(i_c, s_c, l_c, s_r, l_r, f_tot, m, species)
@@ -611,7 +612,7 @@ def __init__(
         j_tot: float | None = None,
         f_tot: float | None = None,
         m: float | None = None,
-        species: str | None = None,
+        species: str | SpeciesObject | None = None,
     ) -> None:
         """Initialize the Spin ket."""
         super().__init__(i_c, s_c, l_c, s_r, l_r, f_tot, m, species)
@@ -674,7 +675,7 @@ def __init__(
         j_r: float | None = None,
         f_tot: float | None = None,
         m: float | None = None,
-        species: str | None = None,
+        species: str | SpeciesObject | None = None,
     ) -> None:
         """Initialize the Spin ket."""
         super().__init__(i_c, s_c, l_c, s_r, l_r, f_tot, m, species)

src/ryd_numerov/radial/model.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from ryd_numerov.elements import BaseElement
+from ryd_numerov.species import SpeciesObject
 
 if TYPE_CHECKING:
     from ryd_numerov.units import NDArray
@@ -24,7 +24,7 @@ class Model:
 
     def __init__(
         self,
-        species: str,
+        species: str | SpeciesObject,
         l: int,
         potential_type: PotentialType | None = None,
     ) -> None:
@@ -36,11 +36,13 @@ def __init__(
             potential_type: Which potential to use for the model.
 
         """
-        self.element = BaseElement.from_species(species)
+        if isinstance(species, str):
+            species = SpeciesObject.from_name(species)
+        self.species = species
         self.l = l
 
         if potential_type is None:
-            potential_type = self.element.potential_type_default
+            potential_type = self.species.potential_type_default
             if potential_type is None:
                 potential_type = "coulomb"
         if potential_type not in get_args(PotentialType):
@@ -91,23 +93,23 @@ def calc_model_potential_marinescu_1993(self, x: XType) -> XType:
             V_{mp,marinescu}: The four parameter potential V_{mp,marinescu}(x) in atomic units.
 
         """
-        parameter_dict = self.element.model_potential_parameter_marinescu_1993
+        parameter_dict = self.species.model_potential_parameter_marinescu_1993
         if len(parameter_dict) == 0:
-            raise ValueError("No parametric model potential parameters defined for this element.")
+            raise ValueError(f"No parametric model potential parameters defined for the species {self.species}.")
         # default to parameters for the maximum l
         a1, a2, a3, a4 = parameter_dict.get(self.l, parameter_dict[max(parameter_dict.keys())])
         exp_a1 = np.exp(-a1 * x)
         exp_a2 = np.exp(-a2 * x)
-        z_nl: XType = 1 + (self.element.Z - 1) * exp_a1 - x * (a3 + a4 * x) * exp_a2
+        z_nl: XType = 1 + (self.species.Z - 1) * exp_a1 - x * (a3 + a4 * x) * exp_a2
         v_c = -z_nl / x
 
-        alpha_c = self.element.alpha_c_marinescu_1993
+        alpha_c = self.species.alpha_c_marinescu_1993
         if alpha_c == 0:
             v_p = 0
         else:
-            r_c_dict = self.element.r_c_dict_marinescu_1993
+            r_c_dict = self.species.r_c_dict_marinescu_1993
             if len(r_c_dict) == 0:
-                raise ValueError("No parametric model potential parameters defined for this element.")
+                raise ValueError(f"No parametric model potential parameters defined for the species {self.species}.")
             # default to x_c for the maximum l
             x_c = r_c_dict.get(self.l, r_c_dict[max(r_c_dict.keys())])
             x2: XType = x * x
@@ -137,10 +139,10 @@ def calc_model_potential_fei_2009(self, x: XType) -> XType:
             V_{mp,fei}: The four parameter potential V_{mp,fei}(x) in atomic units.
 
         """
-        delta, alpha, beta, gamma = self.element.model_potential_parameter_fei_2009
+        delta, alpha, beta, gamma = self.species.model_potential_parameter_fei_2009
         with np.errstate(over="ignore"):
             denom: XType = 1 - alpha + alpha * np.exp(beta * x**delta + gamma * x ** (2.0 * delta))
-            return -1 / x - (self.element.Z - 1) / (x * denom)
+            return -1 / x - (self.species.Z - 1) / (x * denom)
 
     def calc_effective_potential_centrifugal(self, x: XType) -> XType:
         r"""Calculate the effective centrifugal potential V_l(x) in atomic units.
@@ -160,7 +162,7 @@ def calc_effective_potential_centrifugal(self, x: XType) -> XType:
 
         """
         x2 = x * x
-        return (1 / self.element.reduced_mass_factor) * self.l * (self.l + 1) / (2 * x2)
+        return (1 / self.species.reduced_mass_au) * self.l * (self.l + 1) / (2 * x2)
 
     def calc_effective_potential_sqrt(self, x: XType) -> XType:
         r"""Calculate the effective potential V_sqrt(x) from the sqrt transformation in atomic units.
@@ -181,7 +183,7 @@ def calc_effective_potential_sqrt(self, x: XType) -> XType:
 
         """
         x2 = x * x
-        return (1 / self.element.reduced_mass_factor) * (3 / 32) / x2
+        return (1 / self.species.reduced_mass_au) * (3 / 32) / x2
 
     def calc_total_effective_potential(self, x: XType) -> XType:
         r"""Calculate the total effective potential V_eff(x) in atomic units.

src/ryd_numerov/radial/radial_state.py

@@ -6,11 +6,12 @@
 
 import numpy as np
 
-from ryd_numerov.elements import BaseElement
 from ryd_numerov.radial.grid import Grid
 from ryd_numerov.radial.model import Model
 from ryd_numerov.radial.radial_matrix_element import calc_radial_matrix_element_from_w_z
 from ryd_numerov.radial.wavefunction import WavefunctionNumerov, WavefunctionWhittaker
+from ryd_numerov.species import SpeciesObject
+from ryd_numerov.species.utils import calc_energy_from_nu
 from ryd_numerov.units import ureg
 
 if TYPE_CHECKING:
@@ -27,7 +28,7 @@ class RadialState:
 
     def __init__(
         self,
-        species: str,
+        species: str | SpeciesObject,
         nu: float,
         l_r: int,
     ) -> None:
@@ -40,6 +41,8 @@ def __init__(
             l_r: Orbital angular momentum quantum number of the rydberg electron.
 
         """
+        if isinstance(species, str):
+            species = SpeciesObject.from_name(species)
         self.species = species
 
         self.n: int | None = None
@@ -75,7 +78,7 @@ def set_n_for_sanity_check(self, n: int) -> None:
     def __repr__(self) -> str:
         species, nu, l_r, n = self.species, self.nu, self.l_r, self.n
         n_str = "" if n is None else f", ({n=})"
-        return f"{self.__class__.__name__}({species}, {nu=}, {l_r=}{n_str})"
+        return f"{self.__class__.__name__}({species.name}, {nu=}, {l_r=}{n_str})"
 
     def __str__(self) -> str:
         return self.__repr__()
@@ -131,8 +134,7 @@ def create_grid(
             if self.l_r <= 10:
                 z_min = 0.0
             else:
-                element = BaseElement.from_species(self.species)
-                energy_au = element.calc_energy_from_nu(self.nu)
+                energy_au = calc_energy_from_nu(self.species.reduced_mass_au, self.nu)
                 z_min = self.model.calc_turning_point_z(energy_au)
                 z_min = math.sqrt(0.5) * z_min - 3  # see also compare_z_min_cutoff.ipynb
         else:

src/ryd_numerov/radial/wavefunction.py

@@ -9,8 +9,8 @@
 from mpmath import whitw
 from scipy.special import gamma
 
-from ryd_numerov.elements.base_element import BaseElement
 from ryd_numerov.radial.numerov import _run_numerov_integration_python, run_numerov_integration
+from ryd_numerov.species.utils import calc_energy_from_nu
 
 if TYPE_CHECKING:
     from ryd_numerov.radial import Grid, Model
@@ -172,10 +172,10 @@ 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
 
-        element = BaseElement.from_species(self.radial_state.species)
-        energy_au = element.calc_energy_from_nu(self.radial_state.nu)
+        species = self.radial_state.species
+        energy_au = calc_energy_from_nu(species.reduced_mass_au, self.radial_state.nu)
         v_eff = self.model.calc_total_effective_potential(grid.x_list)
-        glist = 8 * element.reduced_mass_factor * grid.z_list * grid.z_list * (energy_au - v_eff)
+        glist = 8 * species.reduced_mass_au * grid.z_list * grid.z_list * (energy_au - v_eff)
 
         if run_backward:
             # During the Numerov integration we define the wavefunction such that it should always stop