diff --git a/src/rydstate/__init__.py b/src/rydstate/__init__.py index 4db6a45..7edfb28 100644 --- a/src/rydstate/__init__.py +++ b/src/rydstate/__init__.py @@ -1,8 +1,9 @@ from rydstate import angular, radial, species -from rydstate.basis import BasisSQDTAlkali, BasisSQDTAlkalineLS +from rydstate.basis import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS from rydstate.rydberg import ( RydbergStateSQDT, RydbergStateSQDTAlkali, + RydbergStateSQDTAlkalineFJ, RydbergStateSQDTAlkalineJJ, RydbergStateSQDTAlkalineLS, ) @@ -10,9 +11,12 @@ __all__ = [ "BasisSQDTAlkali", + "BasisSQDTAlkalineFJ", + "BasisSQDTAlkalineJJ", "BasisSQDTAlkalineLS", "RydbergStateSQDT", "RydbergStateSQDTAlkali", + "RydbergStateSQDTAlkalineFJ", "RydbergStateSQDTAlkalineJJ", "RydbergStateSQDTAlkalineLS", "angular", diff --git a/src/rydstate/basis/__init__.py b/src/rydstate/basis/__init__.py index e2c87e3..08200ec 100644 --- a/src/rydstate/basis/__init__.py +++ b/src/rydstate/basis/__init__.py @@ -1,3 +1,3 @@ -from rydstate.basis.basis_sqdt import BasisSQDTAlkali, BasisSQDTAlkalineLS +from rydstate.basis.basis_sqdt import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS -__all__ = ["BasisSQDTAlkali", "BasisSQDTAlkalineLS"] +__all__ = ["BasisSQDTAlkali", "BasisSQDTAlkalineFJ", "BasisSQDTAlkalineJJ", "BasisSQDTAlkalineLS"] diff --git a/src/rydstate/basis/basis_sqdt.py b/src/rydstate/basis/basis_sqdt.py index d936b58..ced0b2d 100644 --- a/src/rydstate/basis/basis_sqdt.py +++ b/src/rydstate/basis/basis_sqdt.py @@ -1,9 +1,18 @@ from __future__ import annotations +import logging + import numpy as np from rydstate.basis.basis_base import BasisBase -from rydstate.rydberg import RydbergStateSQDTAlkali, RydbergStateSQDTAlkalineLS +from rydstate.rydberg import ( + RydbergStateSQDTAlkali, + RydbergStateSQDTAlkalineFJ, + RydbergStateSQDTAlkalineJJ, + RydbergStateSQDTAlkalineLS, +) + +logger = logging.getLogger(__name__) class BasisSQDTAlkali(BasisBase[RydbergStateSQDTAlkali]): @@ -48,3 +57,63 @@ def __init__(self, species: str, n_min: int = 1, n_max: int | None = None) -> No species, n=n, l=l, s_tot=s_tot, j_tot=j_tot, f_tot=float(f_tot) ) self.states.append(state) + + +class BasisSQDTAlkalineJJ(BasisBase[RydbergStateSQDTAlkalineJJ]): + def __init__(self, species: str, n_min: int = 0, n_max: int | None = None) -> None: + super().__init__(species) + + if n_max is None: + raise ValueError("n_max must be given") + + i_c = self.species.i_c if self.species.i_c is not None else 0 + j_c = 0.5 + s_r = 0.5 + self.states = [] + for n in range(n_min, n_max + 1): + for l_r in range(n): + if self.species.is_allowed_shell(n, l_r, 0) != self.species.is_allowed_shell(n, l_r, 1): + logger.warning( + "For l=%d, n=%d one of the singlet/triplet states is not allowed. " + "In JJ coupling the state does not exist, thus skipping this shell", + *(l_r, n), + ) + if not all(self.species.is_allowed_shell(n, l_r, s_tot) for s_tot in [0, 1]): + continue + for j_r in np.arange(abs(l_r - s_r), l_r + s_r + 1): + for j_tot in range(int(abs(j_r - j_c)), int(j_r + j_c + 1)): + for f_tot in np.arange(abs(j_tot - i_c), j_tot + i_c + 1): + state = RydbergStateSQDTAlkalineJJ( + species, n=n, l=l_r, j_r=float(j_r), j_tot=j_tot, f_tot=float(f_tot) + ) + self.states.append(state) + + +class BasisSQDTAlkalineFJ(BasisBase[RydbergStateSQDTAlkalineFJ]): + def __init__(self, species: str, n_min: int = 0, n_max: int | None = None) -> None: + super().__init__(species) + + if n_max is None: + raise ValueError("n_max must be given") + + i_c = self.species.i_c if self.species.i_c is not None else 0 + j_c = 0.5 + s_r = 0.5 + self.states = [] + for n in range(n_min, n_max + 1): + for l_r in range(n): + if self.species.is_allowed_shell(n, l_r, 0) != self.species.is_allowed_shell(n, l_r, 1): + logger.warning( + "For l=%d, n=%d one of the singlet/triplet states is not allowed. " + "In FJ coupling the state does not exist, thus skipping this shell", + *(l_r, n), + ) + if not all(self.species.is_allowed_shell(n, l_r, s_tot) for s_tot in [0, 1]): + continue + for j_r in np.arange(abs(l_r - s_r), l_r + s_r + 1): + for f_c in np.arange(abs(j_c - i_c), j_c + i_c + 1): + for f_tot in np.arange(abs(f_c - j_r), f_c + j_r + 1): + state = RydbergStateSQDTAlkalineFJ( + species, n=n, l=l_r, j_r=float(j_r), f_c=float(f_c), f_tot=float(f_tot) + ) + self.states.append(state) diff --git a/src/rydstate/rydberg/__init__.py b/src/rydstate/rydberg/__init__.py index d9bdec4..d3dcc50 100644 --- a/src/rydstate/rydberg/__init__.py +++ b/src/rydstate/rydberg/__init__.py @@ -1,6 +1,7 @@ from rydstate.rydberg.rydberg_sqdt import ( RydbergStateSQDT, RydbergStateSQDTAlkali, + RydbergStateSQDTAlkalineFJ, RydbergStateSQDTAlkalineJJ, RydbergStateSQDTAlkalineLS, ) @@ -8,6 +9,7 @@ __all__ = [ "RydbergStateSQDT", "RydbergStateSQDTAlkali", + "RydbergStateSQDTAlkalineFJ", "RydbergStateSQDTAlkalineJJ", "RydbergStateSQDTAlkalineLS", ] diff --git a/src/rydstate/rydberg/rydberg_sqdt.py b/src/rydstate/rydberg/rydberg_sqdt.py index 538db0e..5ee3e83 100644 --- a/src/rydstate/rydberg/rydberg_sqdt.py +++ b/src/rydstate/rydberg/rydberg_sqdt.py @@ -15,7 +15,7 @@ from rydstate.units import BaseQuantities, MatrixElementOperatorRanks, ureg if TYPE_CHECKING: - from rydstate.angular.angular_ket import AngularKetBase, AngularKetJJ, AngularKetLS + from rydstate.angular.angular_ket import AngularKetBase, AngularKetFJ, AngularKetJJ, AngularKetLS from rydstate.units import MatrixElementOperator, PintFloat @@ -418,11 +418,74 @@ 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: + nus = [self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=s_tot) for s_tot in [0, 1]] + + if any(abs(nu - nus[0]) > 1e-10 for nu in nus[1:]): raise ValueError( "RydbergStateSQDTAlkalineJJ is intended for high-l states only, " "where the quantum defects are the same for singlet and triplet states." ) - return nu_singlet + return nus[0] + + +class RydbergStateSQDTAlkalineFJ(RydbergStateSQDT): + """Create an Alkaline Rydberg state, including the radial and angular states.""" + + angular: AngularKetFJ + + def __init__( + self, + species: str | SpeciesObject, + n: int, + l: int, + j_r: float, + f_c: float | None = None, + f_tot: float | None = None, + m: float | None = None, + nu: float | None = None, + ) -> None: + r"""Initialize the Rydberg state. + + Args: + species: Atomic species. + n: Principal quantum number of the rydberg electron. + l: Orbital angular momentum quantum number of the rydberg electron. + j_r: Total angular momentum quantum number of the Rydberg electron. + f_c: Total angular momentum quantum number of the core (core electron + nucleus). + f_tot: Total angular momentum quantum number of the atom (rydberg electron + core) + 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. + + """ + super().__init__(species=species, n=n, nu=nu, l_r=l, j_r=j_r, f_c=f_c, f_tot=f_tot, m=m) + + self.l = self.angular.l_r + self.j_r = self.angular.j_r + self.f_c = self.angular.f_c + self.f_tot = self.angular.f_tot + self.m = self.angular.m + + def __repr__(self) -> str: + species, n, l, j_r, f_c, f_tot, m = self.species, self.n, self.l, self.j_r, self.f_c, self.f_tot, self.m + return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j_r=}, {f_c=}, {f_tot=}, {m=})" + + @cached_property + def nu(self) -> float: + if self._nu is not None: + return self._nu + assert self.n is not None + nus = [ + self.species.calc_nu(self.n, self.l, float(j_tot), s_tot=s_tot) + for s_tot in [0, 1] + for j_tot in np.arange(abs(self.j_r - 1 / 2), self.j_r + 1 / 2 + 1) + ] + + if any(abs(nu - nus[0]) > 1e-10 for nu in nus[1:]): + raise ValueError( + "RydbergStateSQDTAlkalineFJ is intended for high-l states only, " + "where the quantum defects are the same for singlet and triplet states." + ) + return nus[0] diff --git a/tests/test_basis.py b/tests/test_basis.py index d8ae608..bbef57a 100644 --- a/tests/test_basis.py +++ b/tests/test_basis.py @@ -1,7 +1,11 @@ +from typing import TYPE_CHECKING, Any + import numpy as np import pytest -from rydstate import BasisSQDTAlkali -from rydstate.basis.basis_sqdt import BasisSQDTAlkalineLS +from rydstate import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS + +if TYPE_CHECKING: + from rydstate.basis.basis_base import BasisBase @pytest.mark.parametrize("species_name", ["Rb", "Na", "H"]) @@ -62,3 +66,13 @@ def test_alkaline_basis(species_name: str) -> None: me_matrix = basis.calc_reduced_matrix_elements(basis, "electric_dipole", unit="e a0") assert np.shape(me_matrix) == (len(basis.states), len(basis.states)) assert np.count_nonzero(me_matrix) > 0 + + basis = BasisSQDTAlkalineLS(species_name, n_min=30, n_max=35) + basis.filter_states("l_r", (6, 10)) + for basis_class in [BasisSQDTAlkalineJJ, BasisSQDTAlkalineFJ]: + basis2: BasisBase[Any] = basis_class(species_name, n_min=30, n_max=35) # type: ignore [assignment] + basis2.filter_states("l_r", (6, 10)) + assert len(basis2.states) == len(basis.states) + trafo = basis.calc_reduced_overlaps(basis2) + trafo_inv = basis2.calc_reduced_overlaps(basis) + assert np.allclose(trafo @ trafo_inv, np.eye(len(basis.states)), atol=1e-3)