diff --git a/RydState b/RydState index 8a96fe0..03fa30d 160000 --- a/RydState +++ b/RydState @@ -1 +1 @@ -Subproject commit 8a96fe05c309be228a12aaf1b93f0a894f799f08 +Subproject commit 03fa30da5e5802f0d77064a7c09c4736656771a3 diff --git a/src/generate_database_sqdt/main.py b/src/generate_database_sqdt/main.py index fb85a6f..f37b638 100644 --- a/src/generate_database_sqdt/main.py +++ b/src/generate_database_sqdt/main.py @@ -203,15 +203,15 @@ def populate_states_table( state.get_nu(), # nu = NStar for sqdt angular_ket.f_tot, # f: quantum number state.get_nu(), # exp_nui = nu for sqdt - angular_ket.l_tot, # exp_l = l + angular_state.calc_exp_qn("l_tot"), # exp_l = l angular_ket.j_tot, # exp_j = j - angular_ket.s_tot, # exp_s = s + angular_state.calc_exp_qn("s_tot"), # exp_s = s angular_ket.l_r, # exp_l_ryd = l for sqdt angular_state.calc_exp_qn("j_r"), # exp_j_ryd = j for sqdt only one valence electron 0, # std_nui = 0 - 0, # std_l = 0 + angular_state.calc_std_qn("l_tot"), # std_l = 0 0, # std_j = 0 - 0, # std_s = 0 + angular_state.calc_std_qn("s_tot"), # std_s = 0 0, # std_l_ryd = 0 angular_state.calc_std_qn("j_r"), # std_j_ryd = 0 for sqdt and only one valence electron "True", # is_j_total_momentum = True for no hyperfine splitting diff --git a/src/generate_database_sqdt/utils.py b/src/generate_database_sqdt/utils.py index 46587f9..cf1de38 100644 --- a/src/generate_database_sqdt/utils.py +++ b/src/generate_database_sqdt/utils.py @@ -5,20 +5,24 @@ from typing import TYPE_CHECKING import numpy as np -from rydstate import RydbergStateAlkali, RydbergStateAlkalineLS -from rydstate.angular import AngularKetLS +from rydstate import RydbergStateAlkali, RydbergStateAlkalineJJ, RydbergStateAlkalineLS +from rydstate.angular import AngularKetJJ, AngularKetLS from rydstate.radial import RadialState from rydstate.species import SpeciesObject from rydstate.units import MatrixElementOperatorRanks if TYPE_CHECKING: + from rydstate.angular.angular_ket import CouplingScheme from rydstate.angular.angular_matrix_element import AngularOperatorType from rydstate.units import MatrixElementOperator +MIN_L_TO_USE_JJ_COUPLING = 6 + + def get_sorted_list_of_states( species_name: str, n_min: int, n_max: int -) -> list[RydbergStateAlkali] | list[RydbergStateAlkalineLS]: +) -> list[RydbergStateAlkali | RydbergStateAlkalineLS | RydbergStateAlkalineJJ]: """Create a list of quantum numbers sorted by their state energies.""" species = SpeciesObject.from_name(species_name) if species.number_valence_electrons == 1: @@ -45,18 +49,35 @@ def _get_sorted_list_of_states_alkali(species: SpeciesObject, n_min: int, n_max: return sorted(list_of_states, key=lambda x: x.get_energy("a.u.")) -def _get_sorted_list_of_states_alkaline(species: SpeciesObject, n_min: int, n_max: int) -> list[RydbergStateAlkalineLS]: +def _get_sorted_list_of_states_alkaline( # noqa: C901, PLR0912 + species: SpeciesObject, n_min: int, n_max: int +) -> list[RydbergStateAlkalineLS | RydbergStateAlkalineJJ]: i_c = species.i_c if species.i_c is not None else 0 - list_of_states: list[RydbergStateAlkalineLS] = [] - for s in [0, 1]: - for n in range(n_min, n_max + 1): - for l in range(n): - if not species.is_allowed_shell(n, l, s): + s_r, s_c = 1 / 2, 1 / 2 + list_of_states: list[RydbergStateAlkalineLS | RydbergStateAlkalineJJ] = [] + for n in range(n_min, n_max + 1): + for l in range(n): + if l < MIN_L_TO_USE_JJ_COUPLING: # low-l states use LS coupling + for s_tot in [0, 1]: + if not species.is_allowed_shell(n, l, s_tot): + continue + for j_tot in range(abs(l - s_tot), l + s_tot + 1): + for f in np.arange(abs(j_tot - i_c), j_tot + i_c + 1): + state = RydbergStateAlkalineLS(species, n, l, s_tot=s_tot, j_tot=j_tot, f_tot=float(f)) + list_of_states.append(state) + + else: # high-l states use jj coupling + if not all(species.is_allowed_shell(n, l, s_tot) for s_tot in [0, 1]): + if any(species.is_allowed_shell(n, l, s_tot) for s_tot in [0, 1]): + raise RuntimeError("JJ coupling singlet and triplet differ in is_allowed_shell check.") continue - for j in range(abs(l - s), l + s + 1): - for f in np.arange(abs(j - i_c), j + i_c + 1): - state = RydbergStateAlkalineLS(species, n, l, s_tot=s, j_tot=j, f_tot=float(f)) - list_of_states.append(state) + for j_r in [l - s_r, l + s_r]: + if not (j_r + s_c).is_integer(): + raise RuntimeError("Non-integer j_r + s_c encountered.") + for j_tot in range(int(abs(j_r - s_c)), int(j_r + s_c + 1)): + for f in np.arange(abs(j_tot - i_c), j_tot + i_c + 1): + state = RydbergStateAlkalineJJ(species, n, l, j_r=j_r, j_tot=j_tot, f_tot=float(f)) + list_of_states.append(state) return sorted(list_of_states, key=lambda x: x.get_energy("a.u.")) @@ -74,18 +95,33 @@ def calc_matrix_element_one_pair( # Magnetic dipole operator: mu = - mu_B (g_l + g_s ) g_s = 2.0023192 value_s_tot = calc_reduced_angular_matrix_element_cached( - state1.angular.quantum_numbers, state2.angular.quantum_numbers, "s_tot", k_angular + state1.angular.coupling_scheme, + state1.angular.quantum_numbers, + state2.angular.coupling_scheme, + state2.angular.quantum_numbers, + "s_tot", + k_angular, ) g_l = 1 value_l_tot = calc_reduced_angular_matrix_element_cached( - state1.angular.quantum_numbers, state2.angular.quantum_numbers, "l_tot", k_angular + state1.angular.coupling_scheme, + state1.angular.quantum_numbers, + state2.angular.coupling_scheme, + state2.angular.quantum_numbers, + "l_tot", + k_angular, ) angular_matrix_element = g_s * value_s_tot + g_l * value_l_tot prefactor = -0.5 # - mu_B in atomic units elif operator in ["electric_dipole", "electric_quadrupole", "electric_octupole", "electric_quadrupole_zero"]: angular_matrix_element = calc_reduced_angular_matrix_element_cached( - state1.angular.quantum_numbers, state2.angular.quantum_numbers, "spherical", k_angular + state1.angular.coupling_scheme, + state1.angular.quantum_numbers, + state2.angular.coupling_scheme, + state2.angular.quantum_numbers, + "spherical", + k_angular, ) prefactor = math.sqrt(4 * math.pi / (2 * k_angular + 1)) # e in atomic units is 1 else: @@ -110,13 +146,16 @@ def calc_matrix_element_one_pair( @lru_cache(maxsize=100_000) def calc_reduced_angular_matrix_element_cached( + coupling_scheme1: CouplingScheme, qns1: tuple[float, ...], + coupling_scheme2: CouplingScheme, qns2: tuple[float, ...], operator: AngularOperatorType, k_angular: int, ) -> float: - ket1 = AngularKetLS(*qns1) # type: ignore[arg-type] - ket2 = AngularKetLS(*qns2) # type: ignore[arg-type] + ket_classes = {"LS": AngularKetLS, "JJ": AngularKetJJ} + ket1 = ket_classes[coupling_scheme1](*qns1) # type: ignore[arg-type] + ket2 = ket_classes[coupling_scheme2](*qns2) # type: ignore[arg-type] return ket2.calc_reduced_matrix_element(ket1, operator, k_angular)