diff --git a/README.md b/README.md index 309c807..de535d1 100644 --- a/README.md +++ b/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 diff --git a/docs/api_reference.rst b/docs/api_reference.rst index 6a24941..3b91c9a 100644 --- a/docs/api_reference.rst +++ b/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 diff --git a/docs/examples/comparisons/compare_nist_energy_levels_data.ipynb b/docs/examples/comparisons/compare_nist_energy_levels_data.ipynb index b8af67c..1359f45 100644 --- a/docs/examples/comparisons/compare_nist_energy_levels_data.ipynb +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index 53eb18e..0bf8001 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/rydstate/angular/angular_ket.py b/src/rydstate/angular/angular_ket.py index 45a1c62..ccc0b71 100644 --- a/src/rydstate/angular/angular_ket.py +++ b/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: diff --git a/src/rydstate/basis/basis_sqdt.py b/src/rydstate/basis/basis_sqdt.py index fac839f..8892db1 100644 --- a/src/rydstate/basis/basis_sqdt.py +++ b/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") diff --git a/src/rydstate/rydberg/rydberg_sqdt.py b/src/rydstate/rydberg/rydberg_sqdt.py index 5fe65ec..13ce75a 100644 --- a/src/rydstate/rydberg/rydberg_sqdt.py +++ b/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,7 +25,7 @@ class RydbergStateSQDT(RydbergStateBase): - species: SpeciesObject + species: SpeciesObjectSQDT """The atomic species of the Rydberg state.""" angular: AngularKetBase @@ -33,7 +33,7 @@ class RydbergStateSQDT(RydbergStateBase): 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,13 +145,14 @@ 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 @@ -159,14 +160,9 @@ 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, diff --git a/src/rydstate/species/__init__.py b/src/rydstate/species/__init__.py index 6959029..6a96fb4 100644 --- a/src/rydstate/species/__init__.py +++ b/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", diff --git a/src/rydstate/species/species_object.py b/src/rydstate/species/species_object.py index 59d99ab..82568be 100644 --- a/src/rydstate/species/species_object.py +++ b/src/rydstate/species/species_object.py @@ -2,19 +2,14 @@ import inspect import logging -import re from abc import ABC -from fractions import Fraction from functools import cache, cached_property from typing import TYPE_CHECKING, ClassVar, overload -import numpy as np - -from rydstate.species.utils import calc_nu_from_energy, convert_electron_configuration from rydstate.units import rydberg_constant, ureg if TYPE_CHECKING: - from pathlib import Path + from typing_extensions import Self from rydstate.radial.model import PotentialType from rydstate.units import PintFloat @@ -38,21 +33,6 @@ class SpeciesObject(ABC): """Nuclear spin, (default None to ignore hyperfine structure, will be treated like i_c = 0).""" number_valence_electrons: ClassVar[int] """Number of valence electrons (i.e. 1 for alkali atoms and 2 for alkaline earth atoms).""" - ground_state_shell: ClassVar[tuple[int, int]] - """Shell (n, l) describing the electronic ground state configuration.""" - _additional_allowed_shells: ClassVar[list[tuple[int, int]]] = [] - """Additional allowed shells (n, l), which (n, l) is smaller than the ground state shell.""" - - _core_electron_configuration: ClassVar[str] - """Electron configuration of the core electrons, e.g. 4p6 for Rb or 5s for Sr.""" - _ionization_energy: tuple[float, float | None, str] - """Ionization energy with uncertainty and unit: (value, uncertainty, unit).""" - - # Parameters for the extended Rydberg Ritz formula, see calc_nu - _quantum_defects: ClassVar[dict[tuple[int, float, float], tuple[float, float, float, float, float]] | None] = None - """Dictionary containing the quantum defects for each (l, j_tot, s_tot) combination, i.e. - _quantum_defects[(l,j_tot,s_tot)] = (d0, d2, d4, d6, d8) - """ _corrected_rydberg_constant: tuple[float, float | None, str] r"""Corrected Rydberg constant stored as (value, uncertainty, unit)""" @@ -82,102 +62,18 @@ class SpeciesObject(ABC): defined in: Y. Fei et al., Chin. Phys. B 18, 4349 (2009), https://iopscience.iop.org/article/10.1088/1674-1056/18/10/025 """ - _nist_energy_levels_file: Path | None = None - """Path to the NIST energy levels file for this species. - The file should be directly downloaded from https://physics.nist.gov/PhysRefData/ASD/levels_form.html - in the 'Tab-delimited' format and in units of Hartree. - """ - - def __init__(self) -> None: - """Initialize an species instance. - - Use this init method to set up additional properties and data for the species, - like loading NIST energy levels from a file. - - """ - self._nist_energy_levels: dict[tuple[int, int, float, float], float] = {} - if self._nist_energy_levels_file is not None: - self._setup_nist_energy_levels(self._nist_energy_levels_file) - - def _setup_nist_energy_levels(self, file: Path) -> None: # noqa: C901, PLR0912 - """Set up NIST energy levels from a file. - - This method should be called in the constructor to load the NIST energy levels - from the specified file. It reads the file and prepares the data for further use. - - Args: - file: Path to the NIST energy levels file. - n_max: Maximum principal quantum number for which to load the NIST energy levels. - For large quantum numbers, the NIST data is not accurate enough - (it does not even show fine structure splitting), - so we limit the maximum principal quantum number to 15 by default. - - """ - if not file.exists(): - raise ValueError(f"NIST energy data file {file} does not exist.") - - header = file.read_text().splitlines()[0] - if "Level (Hartree)" not in header: - raise ValueError( - f"NIST energy data file {file} not given in Hartree, please download the data in units of Hartree." - ) - - data = np.loadtxt(file, skiprows=1, dtype=str, quotechar='"', delimiter="\t") - # data[i] := (Configuration, Term, J, Prefix, Energy, Suffix, Uncertainty, Reference) - core_config_parts = convert_electron_configuration(self._core_electron_configuration) - - for row in data: - if re.match(r"^([A-Z])", row[0]): - # Skip rows, where the first column starts with an element symbol - continue - - try: - config_parts = convert_electron_configuration(row[0]) - except ValueError: - # Skip rows with invalid electron configuration format - # (they usually correspond to core configurations, that are not the ground state configuration) - # e.g. strontium "4d.(2D<3/2>).4f" - continue - if sum(part[2] for part in config_parts) != sum(part[2] for part in core_config_parts) + 1: - # Skip configurations, where the number of electrons does not match the core configuration + 1 - continue - - for part in core_config_parts: - if part in config_parts: - config_parts.remove(part) - elif (part[0], part[1], part[2] + 1) in config_parts: - config_parts.remove((part[0], part[1], part[2] + 1)) - config_parts.append((part[0], part[1], 1)) - else: - break - if sum(part[2] for part in config_parts) != 1: - # Skip configurations, where the inner electrons are not in the ground state configuration - continue - n, l = config_parts[0][:2] - - multiplicity = int(row[1][0]) - s_tot = (multiplicity - 1) / 2 - - j_tot_list = [float(Fraction(j_str)) for j_str in row[2].split(",")] - for j_tot in j_tot_list: - energy = float(row[4]) - self._nist_energy_levels[(n, l, j_tot, s_tot)] = energy - - if len(self._nist_energy_levels) == 0: - raise ValueError(f"No NIST energy levels found for species {self.name} in file {file}.") - def __repr__(self) -> str: return f"{self.__class__.__name__}()" @classmethod @cache - def from_name(cls, name: str) -> SpeciesObject: + def from_name(cls: type[Self], name: str) -> Self: """Create an instance of the species class from the species name. 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: @@ -196,7 +92,7 @@ def from_name(cls, name: str) -> SpeciesObject: ) @classmethod - def _get_concrete_subclasses(cls) -> list[type[SpeciesObject]]: + def _get_concrete_subclasses(cls: type[Self]) -> list[type[Self]]: subclasses = [] for subclass in cls.__subclasses__(): if not inspect.isabstract(subclass) and hasattr(subclass, "name"): @@ -216,64 +112,6 @@ def get_available_species(cls) -> list[str]: """ return sorted([subclass.name for subclass in cls._get_concrete_subclasses()]) - def is_allowed_shell(self, n: int, l: int, s_tot: float | None = None) -> bool: - """Check if the quantum numbers describe an allowed shell. - - I.e. whether the shell is above the ground state shell. - - Args: - n: Principal quantum number - l: Orbital angular momentum quantum number - s_tot: Total spin quantum number - - Returns: - True if the quantum numbers specify a shell equal to or above the ground state shell, False otherwise. - - """ - if s_tot is None: - if self.number_valence_electrons > 1: - raise ValueError("s_tot must be specified for species with more than one valence electron.") - s_tot = self.number_valence_electrons / 2 - if (self.number_valence_electrons / 2) % 1 != s_tot % 1 or s_tot > self.number_valence_electrons / 2: - raise ValueError(f"Invalid spin {s_tot=} for {self.name}.") - - if (n, l) == self.ground_state_shell: - return s_tot != 1 # For alkaline earth atoms, the triplet state of the ground state shell is not allowed - if n < 1 or l < 0 or l >= n: - raise ValueError(f"Invalid shell: (n={n}, l={l}). Must be n >= 1 and 0 <= l <= n-1.") - if (n, l) >= self.ground_state_shell: - return True - return (n, l) in self._additional_allowed_shells - - @overload - def get_ionization_energy(self, unit: None = None) -> PintFloat: ... - - @overload - def get_ionization_energy(self, unit: str) -> float: ... - - def get_ionization_energy(self, unit: str | None = "hartree") -> PintFloat | float: - """Return the ionization energy in the desired unit. - - Args: - unit: Desired unit for the ionization energy. Default is atomic units "hartree". - - Returns: - Ionization energy in the desired unit. - - """ - ionization_energy: PintFloat = ureg.Quantity(self._ionization_energy[0], self._ionization_energy[2]) - ionization_energy = ionization_energy.to("hartree", "spectroscopy") - if unit is None: - return ionization_energy - if unit == "a.u.": - return ionization_energy.magnitude - return ionization_energy.to(unit, "spectroscopy").magnitude - - @cached_property - def ionization_energy_au(self) -> float: - """Ionization energy in atomic units (Hartree).""" - return self.get_ionization_energy("hartree") - @overload def get_corrected_rydberg_constant(self, unit: None = None) -> PintFloat: ... @@ -325,68 +163,3 @@ def reduced_mass_au(self) -> float: """ return self.get_corrected_rydberg_constant("hartree") / rydberg_constant.to("hartree").m - - def calc_nu( - self, - n: int, - l: int, - j_tot: float, - s_tot: float | None = None, - *, - use_nist_data: bool = True, - nist_n_max: int = 15, - ) -> float: - r"""Calculate the effective principal quantum number nu of a Rydberg state with the given n, l, j_tot and s_tot. - - I.e. either look up the energy for low lying states in the nist data (if use_nist_data is True), - and calculate nu from the energy via (see also `calc_nu_from_energy`): - - .. math:: - \nu = \sqrt{\frac{1}{2} \frac{\mu/m_e}{-E/E_H}} - - Or calculate nu via the quantum defect theory, - where nu is defined as series expansion :math:`\nu = n^* = n - \delta_{lj}(n)` - with the quantum defect - - .. math:: - \delta_{lj}(n) = d0_{lj} + d2_{lj} / [n - d0_{lj}(n)]^2 + d4_{lj} / [n - \delta_{lj}(n)]^4 + ... - - References: - - On a New Law of Series Spectra, Ritz; DOI: 10.1086/141591, https://ui.adsabs.harvard.edu/abs/1908ApJ....28..237R/abstract - - Rydberg atoms, Gallagher; DOI: 10.1088/0034-4885/51/2/001, (Eq. 16.19) - - Args: - n: The principal quantum number of the Rydberg state. - l: The orbital angular momentum quantum number of the Rydberg state. - j_tot: The total angular momentum quantum number of the Rydberg state. - s_tot: The total spin quantum number of the Rydberg state. - use_nist_data: Whether to use NIST energy data. - Default is True. - nist_n_max: Maximum principal quantum number for which to use the NIST energy data. - Default is 15. - - """ - if s_tot is None: - if self.number_valence_electrons != 1: - raise ValueError("s_tot must be specified for species with more than one valence electron.") - s_tot = 0.5 - if (s_tot % 1) != ((self.number_valence_electrons / 2) % 1): - raise ValueError(f"Invalid spin {s_tot=} for {self.name}.") - if j_tot % 1 != (l + s_tot) % 1: - raise ValueError(f"Invalid quantum numbers: ({l=}, {j_tot=}, {s_tot=})") - - if n <= nist_n_max and use_nist_data: # try to use NIST data - if (n, l, j_tot, s_tot) in self._nist_energy_levels: - energy_au = self._nist_energy_levels[(n, l, j_tot, s_tot)] - energy_au -= self.ionization_energy_au # use the cached ionization energy for better performance - return calc_nu_from_energy(self.reduced_mass_au, energy_au) - logger.debug( - "NIST energy levels for (n=%d, l=%d, j_tot=%s, s_tot=%s) not found, using quantum defect theory.", - *(n, l, j_tot, s_tot), - ) - - if self._quantum_defects is None: - raise ValueError(f"No quantum defect data available for species {self.name}.") - d0, d2, d4, d6, d8 = self._quantum_defects.get((l, j_tot, s_tot), (0, 0, 0, 0, 0)) - delta_nlj = d0 + d2 / (n - d0) ** 2 + d4 / (n - d0) ** 4 + d6 / (n - d0) ** 6 + d8 / (n - d0) ** 8 - return n - delta_nlj diff --git a/src/rydstate/species/sqdt/__init__.py b/src/rydstate/species/sqdt/__init__.py new file mode 100644 index 0000000..a890f5f --- /dev/null +++ b/src/rydstate/species/sqdt/__init__.py @@ -0,0 +1,25 @@ +from rydstate.species.sqdt.cesium import Cesium +from rydstate.species.sqdt.hydrogen import Hydrogen, HydrogenTextBook +from rydstate.species.sqdt.lithium import Lithium +from rydstate.species.sqdt.potassium import Potassium +from rydstate.species.sqdt.rubidium import Rubidium +from rydstate.species.sqdt.sodium import Sodium +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT +from rydstate.species.sqdt.strontium import Strontium87, Strontium88 +from rydstate.species.sqdt.ytterbium import Ytterbium171, Ytterbium173, Ytterbium174 + +__all__ = [ + "Cesium", + "Hydrogen", + "HydrogenTextBook", + "Lithium", + "Potassium", + "Rubidium", + "Sodium", + "SpeciesObjectSQDT", + "Strontium87", + "Strontium88", + "Ytterbium171", + "Ytterbium173", + "Ytterbium174", +] diff --git a/src/rydstate/species/cesium.py b/src/rydstate/species/sqdt/cesium.py similarity index 94% rename from src/rydstate/species/cesium.py rename to src/rydstate/species/sqdt/cesium.py index a41c0cd..4eb85a2 100644 --- a/src/rydstate/species/cesium.py +++ b/src/rydstate/species/sqdt/cesium.py @@ -1,10 +1,10 @@ from pathlib import Path from typing import ClassVar -from rydstate.species.species_object import SpeciesObject +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT -class Cesium(SpeciesObject): +class Cesium(SpeciesObjectSQDT): name = "Cs" Z = 55 number_valence_electrons = 1 diff --git a/src/rydstate/species/hydrogen.py b/src/rydstate/species/sqdt/hydrogen.py similarity index 84% rename from src/rydstate/species/hydrogen.py rename to src/rydstate/species/sqdt/hydrogen.py index 063494a..d14e522 100644 --- a/src/rydstate/species/hydrogen.py +++ b/src/rydstate/species/sqdt/hydrogen.py @@ -1,10 +1,10 @@ from typing import ClassVar -from rydstate.species.species_object import SpeciesObject +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT from rydstate.units import rydberg_constant -class Hydrogen(SpeciesObject): +class Hydrogen(SpeciesObjectSQDT): name = "H" Z = 1 number_valence_electrons = 1 @@ -20,7 +20,7 @@ class Hydrogen(SpeciesObject): _corrected_rydberg_constant = (109677.58340280356, None, "1/cm") -class HydrogenTextBook(SpeciesObject): +class HydrogenTextBook(SpeciesObjectSQDT): """Hydrogen from QM textbook with infinite nucleus mass and no spin orbit coupling.""" name = "H_textbook" diff --git a/src/rydstate/species/lithium.py b/src/rydstate/species/sqdt/lithium.py similarity index 94% rename from src/rydstate/species/lithium.py rename to src/rydstate/species/sqdt/lithium.py index cdb7697..c617092 100644 --- a/src/rydstate/species/lithium.py +++ b/src/rydstate/species/sqdt/lithium.py @@ -1,10 +1,10 @@ from pathlib import Path from typing import ClassVar -from rydstate.species.species_object import SpeciesObject +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT -class Lithium(SpeciesObject): +class Lithium(SpeciesObjectSQDT): name = "Li" Z = 3 number_valence_electrons = 1 diff --git a/src/rydstate/species/nist_energy_levels/cesium.txt b/src/rydstate/species/sqdt/nist_energy_levels/cesium.txt similarity index 100% rename from src/rydstate/species/nist_energy_levels/cesium.txt rename to src/rydstate/species/sqdt/nist_energy_levels/cesium.txt diff --git a/src/rydstate/species/nist_energy_levels/lithium.txt b/src/rydstate/species/sqdt/nist_energy_levels/lithium.txt similarity index 100% rename from src/rydstate/species/nist_energy_levels/lithium.txt rename to src/rydstate/species/sqdt/nist_energy_levels/lithium.txt diff --git a/src/rydstate/species/nist_energy_levels/potassium.txt b/src/rydstate/species/sqdt/nist_energy_levels/potassium.txt similarity index 100% rename from src/rydstate/species/nist_energy_levels/potassium.txt rename to src/rydstate/species/sqdt/nist_energy_levels/potassium.txt diff --git a/src/rydstate/species/nist_energy_levels/rubidium.txt b/src/rydstate/species/sqdt/nist_energy_levels/rubidium.txt similarity index 100% rename from src/rydstate/species/nist_energy_levels/rubidium.txt rename to src/rydstate/species/sqdt/nist_energy_levels/rubidium.txt diff --git a/src/rydstate/species/nist_energy_levels/sodium.txt b/src/rydstate/species/sqdt/nist_energy_levels/sodium.txt similarity index 100% rename from src/rydstate/species/nist_energy_levels/sodium.txt rename to src/rydstate/species/sqdt/nist_energy_levels/sodium.txt diff --git a/src/rydstate/species/nist_energy_levels/strontium.txt b/src/rydstate/species/sqdt/nist_energy_levels/strontium.txt similarity index 100% rename from src/rydstate/species/nist_energy_levels/strontium.txt rename to src/rydstate/species/sqdt/nist_energy_levels/strontium.txt diff --git a/src/rydstate/species/nist_energy_levels/ytterbium.txt b/src/rydstate/species/sqdt/nist_energy_levels/ytterbium.txt similarity index 100% rename from src/rydstate/species/nist_energy_levels/ytterbium.txt rename to src/rydstate/species/sqdt/nist_energy_levels/ytterbium.txt diff --git a/src/rydstate/species/potassium.py b/src/rydstate/species/sqdt/potassium.py similarity index 94% rename from src/rydstate/species/potassium.py rename to src/rydstate/species/sqdt/potassium.py index 589d2f4..0683476 100644 --- a/src/rydstate/species/potassium.py +++ b/src/rydstate/species/sqdt/potassium.py @@ -1,10 +1,10 @@ from pathlib import Path from typing import ClassVar -from rydstate.species.species_object import SpeciesObject +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT -class Potassium(SpeciesObject): +class Potassium(SpeciesObjectSQDT): name = "K" Z = 19 number_valence_electrons = 1 diff --git a/src/rydstate/species/rubidium.py b/src/rydstate/species/sqdt/rubidium.py similarity index 95% rename from src/rydstate/species/rubidium.py rename to src/rydstate/species/sqdt/rubidium.py index 1842c59..035afa1 100644 --- a/src/rydstate/species/rubidium.py +++ b/src/rydstate/species/sqdt/rubidium.py @@ -1,10 +1,10 @@ from pathlib import Path from typing import ClassVar -from rydstate.species.species_object import SpeciesObject +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT -class _RubidiumAbstract(SpeciesObject): +class _RubidiumAbstract(SpeciesObjectSQDT): Z = 37 number_valence_electrons = 1 ground_state_shell = (5, 0) diff --git a/src/rydstate/species/sodium.py b/src/rydstate/species/sqdt/sodium.py similarity index 95% rename from src/rydstate/species/sodium.py rename to src/rydstate/species/sqdt/sodium.py index 90fbe7b..77ed9de 100644 --- a/src/rydstate/species/sodium.py +++ b/src/rydstate/species/sqdt/sodium.py @@ -1,10 +1,10 @@ from pathlib import Path from typing import ClassVar -from rydstate.species.species_object import SpeciesObject +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT -class Sodium(SpeciesObject): +class Sodium(SpeciesObjectSQDT): name = "Na" Z = 11 number_valence_electrons = 1 diff --git a/src/rydstate/species/sqdt/species_object_sqdt.py b/src/rydstate/species/sqdt/species_object_sqdt.py new file mode 100644 index 0000000..d41d04c --- /dev/null +++ b/src/rydstate/species/sqdt/species_object_sqdt.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +import logging +import re +from fractions import Fraction +from functools import cached_property +from typing import TYPE_CHECKING, ClassVar, overload + +import numpy as np + +from rydstate.angular.angular_ket import AngularKetLS +from rydstate.species.species_object import SpeciesObject +from rydstate.species.utils import calc_nu_from_energy, convert_electron_configuration +from rydstate.units import ureg + +if TYPE_CHECKING: + from pathlib import Path + + from rydstate.angular.angular_ket import AngularKetBase + from rydstate.units import PintFloat + +logger = logging.getLogger(__name__) + + +class SpeciesObjectSQDT(SpeciesObject): + """Abstract base class for all species objects. + + For the electronic ground state configurations and sorted shells, + see e.g. https://www.webelements.com/atoms.html + + """ + + ground_state_shell: ClassVar[tuple[int, int]] + """Shell (n, l) describing the electronic ground state configuration.""" + _additional_allowed_shells: ClassVar[list[tuple[int, int]]] = [] + """Additional allowed shells (n, l), which (n, l) is smaller than the ground state shell.""" + + _core_electron_configuration: ClassVar[str] + """Electron configuration of the core electrons, e.g. 4p6 for Rb or 5s for Sr.""" + _ionization_energy: tuple[float, float | None, str] + """Ionization energy with uncertainty and unit: (value, uncertainty, unit).""" + + # Parameters for the extended Rydberg Ritz formula, see calc_nu + _quantum_defects: ClassVar[dict[tuple[int, float, float], tuple[float, float, float, float, float]] | None] = None + """Dictionary containing the quantum defects for each (l, j_tot, s_tot) combination, i.e. + _quantum_defects[(l,j_tot,s_tot)] = (d0, d2, d4, d6, d8) + """ + + _nist_energy_levels_file: Path | None = None + """Path to the NIST energy levels file for this species. + The file should be directly downloaded from https://physics.nist.gov/PhysRefData/ASD/levels_form.html + in the 'Tab-delimited' format and in units of Hartree. + """ + + def __init__(self) -> None: + """Initialize an species instance. + + Use this init method to set up additional properties and data for the species, + like loading NIST energy levels from a file. + + """ + self._nist_energy_levels: dict[tuple[int, int, float, float], float] = {} + if self._nist_energy_levels_file is not None: + self._setup_nist_energy_levels(self._nist_energy_levels_file) + + def _setup_nist_energy_levels(self, file: Path) -> None: # noqa: C901, PLR0912 + """Set up NIST energy levels from a file. + + This method should be called in the constructor to load the NIST energy levels + from the specified file. It reads the file and prepares the data for further use. + + Args: + file: Path to the NIST energy levels file. + n_max: Maximum principal quantum number for which to load the NIST energy levels. + For large quantum numbers, the NIST data is not accurate enough + (it does not even show fine structure splitting), + so we limit the maximum principal quantum number to 15 by default. + + """ + if not file.exists(): + raise ValueError(f"NIST energy data file {file} does not exist.") + + header = file.read_text().splitlines()[0] + if "Level (Hartree)" not in header: + raise ValueError( + f"NIST energy data file {file} not given in Hartree, please download the data in units of Hartree." + ) + + data = np.loadtxt(file, skiprows=1, dtype=str, quotechar='"', delimiter="\t") + # data[i] := (Configuration, Term, J, Prefix, Energy, Suffix, Uncertainty, Reference) + core_config_parts = convert_electron_configuration(self._core_electron_configuration) + + for row in data: + if re.match(r"^([A-Z])", row[0]): + # Skip rows, where the first column starts with an element symbol + continue + + try: + config_parts = convert_electron_configuration(row[0]) + except ValueError: + # Skip rows with invalid electron configuration format + # (they usually correspond to core configurations, that are not the ground state configuration) + # e.g. strontium "4d.(2D<3/2>).4f" + continue + if sum(part[2] for part in config_parts) != sum(part[2] for part in core_config_parts) + 1: + # Skip configurations, where the number of electrons does not match the core configuration + 1 + continue + + for part in core_config_parts: + if part in config_parts: + config_parts.remove(part) + elif (part[0], part[1], part[2] + 1) in config_parts: + config_parts.remove((part[0], part[1], part[2] + 1)) + config_parts.append((part[0], part[1], 1)) + else: + break + if sum(part[2] for part in config_parts) != 1: + # Skip configurations, where the inner electrons are not in the ground state configuration + continue + n, l = config_parts[0][:2] + + multiplicity = int(row[1][0]) + s_tot = (multiplicity - 1) / 2 + + j_tot_list = [float(Fraction(j_str)) for j_str in row[2].split(",")] + for j_tot in j_tot_list: + energy = float(row[4]) + self._nist_energy_levels[(n, l, j_tot, s_tot)] = energy + + if len(self._nist_energy_levels) == 0: + raise ValueError(f"No NIST energy levels found for species {self.name} in file {file}.") + + def __repr__(self) -> str: + return f"{self.__class__.__name__}()" + + def is_allowed_shell(self, n: int, l: int, s_tot: float | None = None) -> bool: + """Check if the quantum numbers describe an allowed shell. + + I.e. whether the shell is above the ground state shell. + + Args: + n: Principal quantum number + l: Orbital angular momentum quantum number + s_tot: Total spin quantum number + + Returns: + True if the quantum numbers specify a shell equal to or above the ground state shell, False otherwise. + + """ + if s_tot is None: + if self.number_valence_electrons > 1: + raise ValueError("s_tot must be specified for species with more than one valence electron.") + s_tot = self.number_valence_electrons / 2 + if (self.number_valence_electrons / 2) % 1 != s_tot % 1 or s_tot > self.number_valence_electrons / 2: + raise ValueError(f"Invalid spin {s_tot=} for {self.name}.") + + if (n, l) == self.ground_state_shell: + return s_tot != 1 # For alkaline earth atoms, the triplet state of the ground state shell is not allowed + if n < 1 or l < 0 or l >= n: + raise ValueError(f"Invalid shell: (n={n}, l={l}). Must be n >= 1 and 0 <= l <= n-1.") + if (n, l) >= self.ground_state_shell: + return True + return (n, l) in self._additional_allowed_shells + + @overload + def get_ionization_energy(self, unit: None = None) -> PintFloat: ... + + @overload + def get_ionization_energy(self, unit: str) -> float: ... + + def get_ionization_energy(self, unit: str | None = "hartree") -> PintFloat | float: + """Return the ionization energy in the desired unit. + + Args: + unit: Desired unit for the ionization energy. Default is atomic units "hartree". + + Returns: + Ionization energy in the desired unit. + + """ + ionization_energy: PintFloat = ureg.Quantity(self._ionization_energy[0], self._ionization_energy[2]) + ionization_energy = ionization_energy.to("hartree", "spectroscopy") + if unit is None: + return ionization_energy + if unit == "a.u.": + return ionization_energy.magnitude + return ionization_energy.to(unit, "spectroscopy").magnitude + + @cached_property + def ionization_energy_au(self) -> float: + """Ionization energy in atomic units (Hartree).""" + return self.get_ionization_energy("hartree") + + def calc_nu( + self, + n: int, + angular_ket: AngularKetBase, + *, + use_nist_data: bool = True, + nist_n_max: int = 15, + ) -> float: + r"""Calculate the effective principal quantum number nu of a Rydberg state with the given n, l, j_tot and s_tot. + + I.e. either look up the energy for low lying states in the nist data (if use_nist_data is True), + and calculate nu from the energy via (see also `calc_nu_from_energy`): + + .. math:: + \nu = \sqrt{\frac{1}{2} \frac{\mu/m_e}{-E/E_H}} + + Or calculate nu via the quantum defect theory, + where nu is defined as series expansion :math:`\nu = n^* = n - \delta_{lj}(n)` + with the quantum defect + + .. math:: + \delta_{lj}(n) = d0_{lj} + d2_{lj} / [n - d0_{lj}(n)]^2 + d4_{lj} / [n - \delta_{lj}(n)]^4 + ... + + References: + - On a New Law of Series Spectra, Ritz; DOI: 10.1086/141591, https://ui.adsabs.harvard.edu/abs/1908ApJ....28..237R/abstract + - Rydberg atoms, Gallagher; DOI: 10.1088/0034-4885/51/2/001, (Eq. 16.19) + + Args: + n: The principal quantum number of the Rydberg state. + angular_ket: The angular ket specifying l, j_tot, and s_tot of the Rydberg state. + use_nist_data: Whether to use NIST energy data. + Default is True. + nist_n_max: Maximum principal quantum number for which to use the NIST energy data. + Default is 15. + + """ + if not isinstance(angular_ket, AngularKetLS): + raise NotImplementedError("calc_nu is only implemented for AngularKetLS.") + + l, j_tot, s_tot = angular_ket.l_r, angular_ket.j_tot, angular_ket.s_r + + if n <= nist_n_max and use_nist_data: # try to use NIST data + if (n, l, j_tot, s_tot) in self._nist_energy_levels: + energy_au = self._nist_energy_levels[(n, l, j_tot, s_tot)] + energy_au -= self.ionization_energy_au # use the cached ionization energy for better performance + return calc_nu_from_energy(self.reduced_mass_au, energy_au) + logger.debug( + "NIST energy levels for (n=%d, l=%d, j_tot=%s, s_tot=%s) not found, using quantum defect theory.", + *(n, l, j_tot, s_tot), + ) + + if self._quantum_defects is None: + raise ValueError(f"No quantum defect data available for species {self.name}.") + quantum_defects = list(self._quantum_defects.get((l, j_tot, s_tot), [])) + if len(quantum_defects) > 5: + raise ValueError(f"Quantum defects for {self.name} must be a list with up to 5 elements.") + + d0, d2, d4, d6, d8 = quantum_defects + [0] * (5 - len(quantum_defects)) + delta_nlj = d0 + d2 / (n - d0) ** 2 + d4 / (n - d0) ** 4 + d6 / (n - d0) ** 6 + d8 / (n - d0) ** 8 + return n - delta_nlj diff --git a/src/rydstate/species/strontium.py b/src/rydstate/species/sqdt/strontium.py similarity index 96% rename from src/rydstate/species/strontium.py rename to src/rydstate/species/sqdt/strontium.py index c20a43f..4aa1c46 100644 --- a/src/rydstate/species/strontium.py +++ b/src/rydstate/species/sqdt/strontium.py @@ -3,11 +3,11 @@ from pathlib import Path from typing import ClassVar -from rydstate.species.species_object import SpeciesObject +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT from rydstate.units import electron_mass, rydberg_constant -class _StrontiumAbstract(SpeciesObject): +class _StrontiumAbstract(SpeciesObjectSQDT): Z = 38 number_valence_electrons = 2 ground_state_shell = (5, 0) diff --git a/src/rydstate/species/ytterbium.py b/src/rydstate/species/sqdt/ytterbium.py similarity index 97% rename from src/rydstate/species/ytterbium.py rename to src/rydstate/species/sqdt/ytterbium.py index 6c0a37a..3097b09 100644 --- a/src/rydstate/species/ytterbium.py +++ b/src/rydstate/species/sqdt/ytterbium.py @@ -3,11 +3,11 @@ from pathlib import Path from typing import ClassVar -from rydstate.species.species_object import SpeciesObject +from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT from rydstate.units import electron_mass, rydberg_constant -class _YtterbiumAbstract(SpeciesObject): +class _YtterbiumAbstract(SpeciesObjectSQDT): Z = 70 number_valence_electrons = 2 ground_state_shell = (6, 0) diff --git a/tests/test_all_elements.py b/tests/test_all_elements.py index 762b28f..ca4d7a0 100644 --- a/tests/test_all_elements.py +++ b/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 diff --git a/tests/test_radial_matrix_element.py b/tests/test_radial_matrix_element.py index 74d9b70..24b183d 100644 --- a/tests/test_radial_matrix_element.py +++ b/tests/test_radial_matrix_element.py @@ -1,8 +1,9 @@ import numpy as np import pytest 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( @@ -58,8 +59,9 @@ def test_circular_expectation_value(species_name: str, n: int, l: int, j_tot: fl _{nl} = 1/2 (3 n^2 - l(l+1)) _{nl} = n^2/2 (5 n^2 - 3 l(l+1) + 1) """ - species = SpeciesObject.from_name(species_name) - nu = species.calc_nu(n, l, j_tot) + 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) state = RadialKet(species, nu=nu, l_r=l) state.set_n_for_sanity_check(n)