diff --git a/src/rydstate/__init__.py b/src/rydstate/__init__.py index 052da1b4..e5cceb4b 100644 --- a/src/rydstate/__init__.py +++ b/src/rydstate/__init__.py @@ -1,10 +1,11 @@ from rydstate import angular, basis, radial, rydberg_state, species -from rydstate.basis import BasisMQDT, BasisSQDT +from rydstate.basis import BasisMQDT, BasisOQDT, BasisSQDT from rydstate.rydberg_state import RydbergStateMQDT, RydbergStateSQDT, RydbergStateSQDTAlkali from rydstate.units import ureg __all__ = [ "BasisMQDT", + "BasisOQDT", "BasisSQDT", "RydbergStateMQDT", "RydbergStateSQDT", diff --git a/src/rydstate/basis/__init__.py b/src/rydstate/basis/__init__.py index 57f41d8f..9f08f0c4 100644 --- a/src/rydstate/basis/__init__.py +++ b/src/rydstate/basis/__init__.py @@ -1,5 +1,6 @@ from rydstate.basis.basis_base import BasisBase from rydstate.basis.basis_mqdt import BasisMQDT +from rydstate.basis.basis_oqdt import BasisOQDT from rydstate.basis.basis_sqdt import BasisSQDT -__all__ = ["BasisBase", "BasisMQDT", "BasisSQDT"] +__all__ = ["BasisBase", "BasisMQDT", "BasisOQDT", "BasisSQDT"] diff --git a/src/rydstate/basis/basis_mqdt.py b/src/rydstate/basis/basis_mqdt.py index fa0da9e4..bcfc70a5 100644 --- a/src/rydstate/basis/basis_mqdt.py +++ b/src/rydstate/basis/basis_mqdt.py @@ -18,6 +18,7 @@ from typing_extensions import Self from rydstate.species import FModel + from rydstate.units import NDArray logger = logging.getLogger(__name__) @@ -38,6 +39,7 @@ def __init__( # potential and mqdt parameters mqdt: MQDT | str | None = None, potential_class: type[Potential] | str | None = None, + scale_off_diagonal: float | None = None, ) -> None: """Initialize the MQDT basis. @@ -61,6 +63,7 @@ def __init__( potential_class: The potential class to use for the radial ket. Either a a potential class or a string representing the tag of the potential class to use. + scale_off_diagonal: If provided, scale the off-diagonal elements of the M-matrix by this factor. """ super().__init__(species) @@ -75,7 +78,7 @@ def __init__( # and for high l_r the quantum defects are 0, so n = nu max_l_r = int(nu[1]) self._init_models(max_l_r, f_tot, l_r, include_sqdt_fallback_models=include_sqdt_fallback_models) - self._init_states(nu, m) + self._init_states(nu, m, scale_off_diagonal=scale_off_diagonal) def shallow_copy(self) -> Self: """Return a shallow copy of the basis (with its own independent list of states).""" @@ -118,12 +121,16 @@ def _init_states( self, nu_range: tuple[float, float], m_range: tuple[float, float] | None | NotSet, + *, + scale_off_diagonal: float | None = None, ) -> None: logger.debug("Calculating MQDT states...") self.states = [] for model in self.models: logger.debug(" calculating states for model %s with nu_range=%s", model.name, nu_range) - states = get_mqdt_states_from_fmodel(model, nu_range, m_range, self.potential_class) + states = get_mqdt_states_from_fmodel( + model, nu_range, m_range, self.potential_class, scale_off_diagonal=scale_off_diagonal + ) if len(states) == 0: logger.debug(" no states found for model %s", model.name) else: @@ -144,6 +151,8 @@ def get_mqdt_states_from_fmodel( # noqa: C901 nu_range: tuple[float, float], m_range: tuple[float, float] | None | NotSet, potential_class: type[Potential], + *, + scale_off_diagonal: float | None = None, ) -> list[RydbergStateMQDT]: """Calculate MQDT states from an FModel by finding zeros of det(M-matrix). @@ -153,6 +162,7 @@ def get_mqdt_states_from_fmodel( # noqa: C901 m_range: Tuple of (m_min, m_max) for the magnetic quantum number range. NotSet will only include states with m=NotSet. potential_class: The potential class to use for the radial ket. + scale_off_diagonal: If provided, scale the off-diagonal elements of the M-matrix by this factor. Returns: List of :class:`RydbergStateMQDT` objects, one per root of det(M). @@ -165,7 +175,15 @@ def get_mqdt_states_from_fmodel( # noqa: C901 if nu_min > nu_max: return [] - nu_list = find_roots(lambda nu: np.linalg.det(model.calc_scaled_m_matrix(nu)), nu_min, nu_max) + def calc_scaled_off_diagonal(mmat: NDArray) -> NDArray: + if scale_off_diagonal is None: + return mmat + mmat_diag = np.diag(np.diag(mmat)) + return scale_off_diagonal * (mmat - mmat_diag) + mmat_diag + + nu_list = find_roots( + lambda nu: np.linalg.det(calc_scaled_off_diagonal(model.calc_scaled_m_matrix(nu))), nu_min, nu_max + ) if len(nu_list) == 0: logger.warning( "No MQDT states found in the range nu_min=%s, nu_max=%s for model %s", nu_min, nu_max, model.name @@ -183,7 +201,7 @@ def get_mqdt_states_from_fmodel( # noqa: C901 states: list[RydbergStateMQDT] = [] for nu in nu_list: - mmat = model.calc_m_matrix(nu) + mmat = calc_scaled_off_diagonal(model.calc_m_matrix(nu)) det_mmat = np.linalg.det(mmat) if abs(det_mmat) > 1e-6: # this can happen, because we use the scaled M-matrix to find roots ... @@ -227,8 +245,10 @@ def get_mqdt_states_from_fmodel( # noqa: C901 for m in get_m_range(model.f_tot, m_range): rydberg_kets = [ RydbergKet(model.species, angular_ket.replace_m(m), radial_ket) - for angular_ket, radial_ket in zip(angular_kets_fj, radial_kets_fj, strict=True) + for i, (angular_ket, radial_ket) in enumerate(zip(angular_kets_fj, radial_kets_fj, strict=True)) + if abs(coefficients_all[i]) > 1e-10 ] + coefficients_all = [coeff for coeff in coefficients_all if abs(coeff) > 1e-10] states.append( RydbergStateMQDT( model.species, diff --git a/src/rydstate/basis/basis_oqdt.py b/src/rydstate/basis/basis_oqdt.py new file mode 100644 index 00000000..4743a027 --- /dev/null +++ b/src/rydstate/basis/basis_oqdt.py @@ -0,0 +1,279 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import numpy as np + +from rydstate.angular import NotSet +from rydstate.angular.angular_ket import AngularKetFJ, AngularKetJJ, AngularKetLS +from rydstate.angular.utils import get_possible_quantum_number_values, is_unknown +from rydstate.basis.basis_base import BasisBase +from rydstate.basis.utils import get_m_range, is_allowed_qn +from rydstate.linalg import find_roots +from rydstate.radial.radial_ket import RadialDummy, RadialKet +from rydstate.rydberg_state.rydberg_ket import RydbergKet +from rydstate.rydberg_state.rydberg_mqdt import RydbergStateMQDT +from rydstate.species.mqdt import MQDT, get_mqdt +from rydstate.species.potential import Potential, get_potential_class + +if TYPE_CHECKING: + from rydstate.angular.angular_ket import AngularKetBase + from rydstate.angular.core_ket import CoreKet + from rydstate.species.fmodel import FModel + + +logger = logging.getLogger(__name__) + + +class BasisOQDT(BasisBase[RydbergStateMQDT]): + """Basis for Rydberg states in the FJ coupling scheme, using Outer Channel QDT.""" + + states: list[RydbergStateMQDT] + + def __init__( + self, + species: str, + nu: tuple[float, float], + *, + f_tot: tuple[float, float] | None = None, + l_r: tuple[int, int] | None = None, + m: tuple[float, float] | None | NotSet = NotSet, + # potential and qdt parameters + mqdt: MQDT | str | None = None, + potential_class: type[Potential] | str | None = None, + ) -> None: + super().__init__(species) + self.mqdt = mqdt if isinstance(mqdt, MQDT) else get_mqdt(species, tag=mqdt) + + if isinstance(potential_class, type) and issubclass(potential_class, Potential): + self.potential_class = potential_class + else: + self.potential_class = get_potential_class(species, tag=potential_class) + + self._init_states(nu, f_tot, l_r, m) + + def _init_states( + self, + nu_range: tuple[float, float], + f_tot_range: tuple[float, float] | None, + l_r_range: tuple[int, int] | None, + m_range: tuple[float, float] | None | NotSet, + ) -> None: + self.states = [] + + for core_ket in self.mqdt.get_core_kets(): + logger.info("Generating states for core ket: %s", core_ket) + if is_unknown(core_ket.l_c): + self._add_states_unknown(core_ket, nu_range, f_tot_range, m_range) + continue + + for l_r in range(int(nu_range[1]) + 10): + if not is_allowed_qn(l_r_range, l_r): + continue + + if not core_ket.contains_unknown: + self._add_states_fj(core_ket, nu_range, l_r, f_tot_range, m_range) + elif not is_unknown(core_ket.j_c): + self._add_states_jj(core_ket, nu_range, l_r, f_tot_range, m_range) + elif not is_unknown(core_ket.l_c): + self._add_states_ls(core_ket, nu_range, l_r, f_tot_range, m_range) + else: + raise RuntimeError("This should never happen.") + self.sort_states("nu") + + def _add_states_fj( + self, + core_ket: CoreKet, + nu_range: tuple[float, float], + l_r: int, + f_tot_range: tuple[float, float] | None, + m_range: tuple[float, float] | None | NotSet = NotSet, + ) -> None: + s_r = 0.5 + + for j_r in get_possible_quantum_number_values(l_r, s_r, None): + for f_tot in get_possible_quantum_number_values(core_ket.f_c, j_r, None): + if is_unknown(f_tot): + raise ValueError("Cannot determine f_tot for BasisOQDT (FJ).") + if not is_allowed_qn(f_tot_range, f_tot): + continue + + angular_ket = AngularKetFJ( + l_c=core_ket.l_c, + j_c=core_ket.j_c, + l_r=l_r, + j_r=j_r, + f_c=core_ket.f_c, + f_tot=f_tot, + species=self.species, + ) + + self._create_states(nu_range, angular_ket, m_range) + + def _add_states_jj( + self, + core_ket: CoreKet, + nu_range: tuple[float, float], + l_r: int, + f_tot_range: tuple[float, float] | None, + m_range: tuple[float, float] | None | NotSet = NotSet, + ) -> None: + assert is_unknown(core_ket.f_c) + assert not is_unknown(core_ket.j_c) + s_r = 0.5 + + for j_r in get_possible_quantum_number_values(l_r, s_r, None): + for j_tot in get_possible_quantum_number_values(core_ket.j_c, j_r, None): + for f_tot in get_possible_quantum_number_values(j_tot, core_ket.i_c, None): + if is_unknown(f_tot): + raise ValueError("Cannot determine f_tot for BasisOQDT (JJ).") + if not is_allowed_qn(f_tot_range, f_tot): + continue + + angular_ket = AngularKetJJ( + l_c=core_ket.l_c, + j_c=core_ket.j_c, + l_r=l_r, + j_r=j_r, + j_tot=j_tot, + f_tot=f_tot, + species=self.species, + ) + + self._create_states(nu_range, angular_ket, m_range) + + def _add_states_ls( + self, + core_ket: CoreKet, + nu_range: tuple[float, float], + l_r: int, + f_tot_range: tuple[float, float] | None, + m_range: tuple[float, float] | None | NotSet = NotSet, + ) -> None: + assert is_unknown(core_ket.j_c) + s_r = 0.5 + + for s_tot in get_possible_quantum_number_values(core_ket.s_c, s_r, None): + for l_tot in get_possible_quantum_number_values(core_ket.l_c, l_r, None): + for j_tot in get_possible_quantum_number_values(s_tot, l_tot, None): + for f_tot in get_possible_quantum_number_values(j_tot, core_ket.i_c, None): + if is_unknown(f_tot): + raise ValueError("Cannot determine f_tot for BasisOQDT (LS).") + if not is_allowed_qn(f_tot_range, f_tot): + continue + + angular_ket = AngularKetLS( + l_c=core_ket.l_c, + l_r=l_r, + s_tot=s_tot, + l_tot=l_tot, # type: ignore [arg-type] + j_tot=j_tot, + f_tot=f_tot, + species=self.species, + ) + + self._create_states(nu_range, angular_ket, m_range) + + def _add_states_unknown( + self, + core_ket: CoreKet, + nu_range: tuple[float, float], + f_tot_range: tuple[float, float] | None, + m_range: tuple[float, float] | None | NotSet = NotSet, + ) -> None: + assert is_unknown(core_ket.l_c) + + fmodels = [model for model in self.mqdt.models if core_ket in model.get_core_kets()] + angular_kets = { + channel for model in fmodels for channel in model.outer_channels if channel.get_core_ket() == core_ket + } + + for angular_ket in angular_kets: + if not is_allowed_qn(f_tot_range, angular_ket.f_tot): + continue + self._create_states(nu_range, angular_ket, m_range) + + def _create_states( + self, + nu_range: tuple[float, float], + angular_ket: AngularKetBase[Any], + m_range: tuple[float, float] | None | NotSet, + ) -> None: + for model in self.mqdt.get_mqdt_models(angular_ket): + if angular_ket not in model.outer_channels: + logger.warning( + "Angular ket %s has overlap but was not found in model %s outer channels.", angular_ket, model.name + ) + continue + states = get_oqdt_states_from_fmodel(angular_ket, model, nu_range, m_range, self.potential_class) + self.states.extend(states) + + +def get_oqdt_states_from_fmodel( + angular_ket: AngularKetBase[Any], + model: FModel, + nu_range: tuple[float, float], + m_range: tuple[float, float] | None | NotSet, + potential_class: type[Potential], +) -> list[RydbergStateMQDT]: + """Calculate MQDT states from an FModel by finding zeros of the OQDT condition (see model._calc_oqdt_condition). + + Args: + angular_ket: The angular ket to compute states for. + model: The MQDT model to compute states for. + nu_range: Tuple of (nu_min, nu_max) for the search range. + m_range: Tuple of (m_min, m_max) for the magnetic quantum number range. + NotSet will only include states with m=NotSet. + potential_class: The potential class to use for the radial ket. + + Returns: + List of :class:`RydbergStateMQDT` objects, one per root of the OQDT condition. + + """ + nu_min = max(nu_range[0], model.nu_min) + nu_max = min(nu_range[1], model.nu_max) + if np.isinf(nu_max): + raise ValueError("nu_max must be finite to calculate MQDT states.") + if nu_min > nu_max: + return [] + + ind = next((i for i, ket in enumerate(model.outer_channels) if ket == angular_ket), None) + if ind is None: + raise ValueError(f"Angular ket {angular_ket} not found in model {model.name} outer channels.") + nu_list = find_roots(lambda nu: model._calc_oqdt_condition(nu, ind), nu_min, nu_max) # noqa: SLF001 + + if len(nu_list) == 0: + logger.warning( + "No MQDT states found in the range nu_min=%s, nu_max=%s for model %s", nu_min, nu_max, model.name + ) + return [] + + states: list[RydbergStateMQDT] = [] + for nu in nu_list: + nuis = model.calc_channel_nuis(nu) + nui = nuis[ind] + + radial: RadialKet | RadialDummy + if not is_unknown(angular_ket.l_r): + potential = potential_class(angular_ket.l_r) + radial = RadialKet(float(nui), potential, sign_convention="positive_at_outer_bound") + else: + radial = RadialDummy(1.0, nui) + + energy_au = model.calc_energy_au(nu) + for m in get_m_range(model.f_tot, m_range): + rydberg_ket = RydbergKet(model.species, angular_ket.replace_m(m), radial) + states.append( + RydbergStateMQDT( + model.species, + [1], + [rydberg_ket], + nu=nu, + energy_au=energy_au, + mqdt=model.mqdt, + potential_class=potential_class, + ) + ) + + return states diff --git a/src/rydstate/radial/radial_base.py b/src/rydstate/radial/radial_base.py index 5cd1dddc..021e21c4 100644 --- a/src/rydstate/radial/radial_base.py +++ b/src/rydstate/radial/radial_base.py @@ -226,7 +226,7 @@ def _calc_matrix_element_au( if self._is_dummy or other._is_dummy: if self._is_dummy is not other._is_dummy: return 0.0 - if k_radial == 0 and abs(self.nu - other.nu) < 1e-10: # type: ignore [attr-defined] + if k_radial == 0 and abs(self.nu - other.nu) < 1: # type: ignore [attr-defined] return self._coeff.conjugate() * other._coeff # type: ignore [attr-defined,no-any-return] # if not k_radial == 0 or nu are not the same we cant compute the matrix element and simply return 0 return 0.0 diff --git a/src/rydstate/species/fmodel.py b/src/rydstate/species/fmodel.py index 2587836b..8d000701 100644 --- a/src/rydstate/species/fmodel.py +++ b/src/rydstate/species/fmodel.py @@ -1,6 +1,7 @@ from __future__ import annotations import inspect +import logging import math from functools import cached_property from typing import TYPE_CHECKING, Any, ClassVar, overload @@ -21,6 +22,8 @@ from rydstate.species.utils import RydbergRitzParameters from rydstate.units import NDArray, PintFloat +logger = logging.getLogger(__name__) + class FModel: """Class to store the parameters of a MQDT model for a given species.""" @@ -303,6 +306,26 @@ def calc_scaled_m_matrix(self, nu: float) -> NDArray: nuis = self.calc_channel_nuis(nu) return np.array(np.diag(np.sin(np.pi * nuis)) + np.diag(np.cos(np.pi * nuis)) @ kmat) + def _calc_oqdt_condition(self, nu: float, ind: int) -> float: + """Calculate the condition for OQDT. + + This is equivalent to the determinant of the scaled M-matrix for MQDT, + where we assume the K-matrix to be diagonal in the outer channel frame + (i.e. in the outer channel frame we set the off-diagonal elements of the K-matrix to zero). + This simplifies the determinant condition to uncoupled equations for each outer channel. + + Args: + nu: Effective principal quantum number with reference to the lowest ionization threshold. + ind: Index of the channel for which to calculate the OQDT condition. + + Returns: + OQDT condition for the specified channel at the given nu value. + + """ + k = self.calc_k_matrix(nu)[ind, ind] + nui = self.calc_channel_nuis(nu)[ind] + return float(np.sin(np.pi * nui) + np.cos(np.pi * nui) * k) + def get_fmodels(module: ModuleType, species: str) -> list[type[FModel]]: """Return all FModel subclasses defined in ``module`` that match the given species. diff --git a/src/rydstate/species/strontium/sqdt_data.py b/src/rydstate/species/strontium/sqdt_data.py index e80a5a3a..2759da84 100644 --- a/src/rydstate/species/strontium/sqdt_data.py +++ b/src/rydstate/species/strontium/sqdt_data.py @@ -44,3 +44,8 @@ class SQDTStrontium88(SQDT): (3, 3.0, 1): (0.119, -2.0, 100, 0.0, 0.0), # [3] (3, 4.0, 1): (0.120, -2.4, 120, 0.0, 0.0), # [3] } + + +class SQDTStrontium87(SQDT): + species = "Sr87" + is_default = True diff --git a/src/rydstate/species/ytterbium/sqdt_data.py b/src/rydstate/species/ytterbium/sqdt_data.py index 3d8e43e1..9ca20466 100644 --- a/src/rydstate/species/ytterbium/sqdt_data.py +++ b/src/rydstate/species/ytterbium/sqdt_data.py @@ -47,3 +47,13 @@ class SQDTYtterbium174(SQDT): # (4, 4.0, "+"): (0.0262659964, 0.0254568575, 0.0, 0.0, 0.0), # [3] S8 # (4, 4.0, "-"): (-0.148808463, -0.134219071, 0.0, 0.0, 0.0), # [3] S8 } + + +class SQDTYtterbium171(SQDT): + species = "Yb171" + is_default = True + + +class SQDTYtterbium173(SQDT): + species = "Yb173" + is_default = True