diff --git a/src/rydstate/basis/basis_base.py b/src/rydstate/basis/basis_base.py index ed259cf2..ec39277c 100644 --- a/src/rydstate/basis/basis_base.py +++ b/src/rydstate/basis/basis_base.py @@ -28,12 +28,6 @@ def __init__(self, species: str) -> None: def __len__(self) -> int: return len(self.states) - def copy(self) -> Self: - new_basis = self.__class__.__new__(self.__class__) - new_basis.__dict__ = self.__dict__.copy() - new_basis.states = list(self.states) - return new_basis - @overload def filter_states( self, qn: str, value: tuple[float, float], *, delta: float = 1e-10, keep_unknown: bool = False diff --git a/src/rydstate/cli.py b/src/rydstate/cli.py index b9772f96..cf82bfbc 100644 --- a/src/rydstate/cli.py +++ b/src/rydstate/cli.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) -def main() -> None: +def main() -> None: # noqa: C901, PLR0912, PLR0915 """Entry point for the generate_database script.""" parser = argparse.ArgumentParser( description="Generate a database, containing energies and matrix elements, for a given species.", @@ -25,7 +25,6 @@ def main() -> None: default=None, type=int, help="The minimal principal quantum number n for the states to be included in the database. " - "This is used for species, where the low lying states do not converge nicely, so we exclude those states. " "Default 1 will start with the ground state configuration of the specific species (e.g. n=5 for Rb).", ) parser.add_argument( @@ -35,16 +34,29 @@ def main() -> None: help="The maximum principal quantum number n for the states to be included in the database.", ) parser.add_argument( - "--max-delta-n", + "--nu-min", + default=None, + type=int, + help="The minimal effective principal quantum number nu for the states to be included in the database. " + "Default 0 will include all low lying states.", + ) + parser.add_argument( + "--nu-max", + default=None, + type=int, + help="The maximum effective principal quantum number nu for the states to be included in the database.", + ) + parser.add_argument( + "--max-delta-nu", default=float("inf"), type=float, - help="The maximum difference in principal quantum number n for matrix elements to be calculated.", + help="The maximum difference in effective principal quantum number nu for matrix elements to be calculated.", ) parser.add_argument( - "--all-n-up-to", + "--all-nu-up-to", default=float("inf"), type=float, - help="Calculate all matrix elements where at least one state has principal quantum number n " + help="Calculate all matrix elements where at least one state has effective principal quantum number nu " "smaller than or equal to this value.", ) parser.add_argument( @@ -81,11 +93,14 @@ def main() -> None: if ( args.n_min is not None or args.n_max is not None - or args.max_delta_n != float("inf") - or args.all_n_up_to != float("inf") + or args.nu_min is not None + or args.nu_max is not None + or args.max_delta_nu != float("inf") + or args.all_nu_up_to != float("inf") ): parser.error( - "--n-min, --n-max, --max-delta-n, and --all-n-up-to are only valid when generating a species database." + "--n-min, --n-max, --nu-min, --nu-max, --max-delta-nu, and --all-nu-up-to are only valid " + "when generating a species database." ) elif args.f_max is not None: parser.error("--f-max is only valid when generating the misc database.") @@ -108,10 +123,25 @@ def main() -> None: parser.error("--f-max is required when generating the misc database.") create_tables_for_misc(f_max=args.f_max, kappa_max=3) else: - n_min = args.n_min if args.n_min is not None else 1 - if args.n_max is None: - parser.error("--n-max is required when generating a species database.") - create_tables_for_one_species(args.species, n_min, args.n_max, args.max_delta_n, args.all_n_up_to) + if args.n_max is None and args.nu_max is None: + parser.error("At least one of --n-max or --nu-max must be provided.") + + if args.n_min is None and args.n_max is None: + n = None + else: + n_min = args.n_min if args.n_min is not None else 1 + n_max = args.n_max if args.n_max is not None else int(args.nu_max) + 10 + n = (n_min, n_max) + + if args.nu_min is None and args.nu_max is None: + nu = None + else: + nu_min = args.nu_min if args.nu_min is not None else 0 + nu_max = args.nu_max if args.nu_max is not None else args.n_max + nu = (nu_min, nu_max) + create_tables_for_one_species( + args.species, n=n, nu=nu, max_delta_nu=args.max_delta_nu, all_nu_up_to=args.all_nu_up_to + ) logger.info("Time taken: %.2f seconds", time.perf_counter() - time_start) diff --git a/src/rydstate/generate_database/generate_database.py b/src/rydstate/generate_database/generate_database.py index 166c8b80..ea1b872d 100644 --- a/src/rydstate/generate_database/generate_database.py +++ b/src/rydstate/generate_database/generate_database.py @@ -4,19 +4,16 @@ import sqlite3 from importlib.resources import files from pathlib import Path +from typing import Any import numpy as np import pandas as pd from rydstate import __version__ from rydstate.angular.wigner_symbols import calc_wigner_3j +from rydstate.basis.basis_mqdt import BasisMQDT from rydstate.basis.basis_sqdt import BasisSQDT -from rydstate.generate_database.generate_matrix_elements_table import ( - _calc_radial_matrix_element_cached, - calc_reduced_angular_matrix_element_cached, - generate_matrix_elements_tables, - get_radial_state_cached, -) +from rydstate.generate_database.generate_matrix_elements_table import generate_matrix_elements_tables from rydstate.generate_database.generate_misc_table import generate_wigner_table from rydstate.generate_database.generate_states_table import generate_states_table @@ -27,25 +24,39 @@ def create_tables_for_one_species( - species: str, - n_min: int, - n_max: int, - max_delta_n: float = np.inf, - all_n_up_to: float = np.inf, + species_specifier: str, + n: tuple[int, int] | None = None, + nu: tuple[float, float] | None = None, + max_delta_nu: float = np.inf, + all_nu_up_to: float = np.inf, ) -> None: """Create the database tables for a given species in the current directory.""" - logger.info("Start creating database for %s", species) - logger.info("n-min=%d, n-max=%d", n_min, n_max) - logger.info("max_delta_n=%s, all_n_up_to=%s", max_delta_n, all_n_up_to) + logger.info("Start creating database for %s", species_specifier) + logger.info("n-range=%s", n) + logger.info("nu-range=%s", nu) + logger.info("max_delta_nu=%s, all_nu_up_to=%s", max_delta_nu, all_nu_up_to) logger.info("rydstate.__version__=%s", __version__) # create the database and populate the states and matrix elements tables db_file = Path("database.db") with sqlite3.connect(db_file) as conn: conn.executescript(DATABASE_SQL_FILE.read_text(encoding="utf-8")) - basis = BasisSQDT(species, n=(n_min, n_max), coupling_scheme="LS") + basis: BasisSQDT[Any] | BasisMQDT + species = species_specifier.removesuffix("_mqdt").removesuffix("_sqdt") + if species_specifier.endswith("_mqdt"): + if nu is None: + raise ValueError("nu must be provided for MQDT basis") + basis = BasisMQDT(species, nu=nu) + if n is not None: + basis.filter_states("n", n) + else: + if n is None: + raise ValueError("n must be provided for SQDT basis") + basis = BasisSQDT(species, n=n, coupling_scheme="LS") + if nu is not None: + basis.filter_states("nu", nu) generate_states_table(basis, conn) - generate_matrix_elements_tables(basis, conn, max_delta_n, all_n_up_to) + generate_matrix_elements_tables(basis, conn, max_delta_nu, all_nu_up_to, free_memory=True) logger.info("Size of %s: %.6f megabytes", db_file, db_file.stat().st_size * 1e-6) # convert the tables to parquet files @@ -67,12 +78,6 @@ def create_tables_for_one_species( with Path("log").open("a") as buf: table.info(buf=buf) - logger.info( - "calc_reduced_angular_matrix_element_cached: %s", calc_reduced_angular_matrix_element_cached.cache_info() - ) - logger.info("_calc_radial_matrix_element_cached: %s", _calc_radial_matrix_element_cached.cache_info()) - logger.info("get_radial_state_cached: %s", get_radial_state_cached.cache_info()) - def create_tables_for_misc(f_max: float, kappa_max: int = 3) -> None: """Create misc databases, i.e. the wigner table in the current directory.""" diff --git a/src/rydstate/generate_database/generate_matrix_elements_table.py b/src/rydstate/generate_database/generate_matrix_elements_table.py index 9b021e17..04a89a4d 100644 --- a/src/rydstate/generate_database/generate_matrix_elements_table.py +++ b/src/rydstate/generate_database/generate_matrix_elements_table.py @@ -1,21 +1,19 @@ from __future__ import annotations import logging -import math -from functools import lru_cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any -from rydstate.angular import AngularKetLS -from rydstate.radial.radial_ket import RadialKet -from rydstate.species.potential import get_potential_class +import numpy as np + +from rydstate.angular.angular_ket import AngularKetBase +from rydstate.angular.utils import is_unknown from rydstate.units import MatrixElementOperatorRanks if TYPE_CHECKING: import sqlite3 - from rydstate.angular.utils import AllKnown, AngularOperatorType - from rydstate.basis.basis_sqdt import BasisSQDT - from rydstate.rydberg_state import RydbergStateSQDT + from rydstate.basis import BasisMQDT, BasisSQDT + from rydstate.rydberg_state.rydberg_base import RydbergStateBase from rydstate.units import MatrixElementOperator logger = logging.getLogger(__name__) @@ -35,32 +33,61 @@ } +def get_l_r_difference(state1: RydbergStateBase, state2: RydbergStateBase) -> int: + """Calculate the minimal difference in l_r between two states.""" + angular1, angular2 = state1.angular, state2.angular + if isinstance(angular1, AngularKetBase) and isinstance(angular2, AngularKetBase): + return abs(angular1.l_r - angular2.l_r) # type: ignore [no-any-return] + + if isinstance(angular1, AngularKetBase): + angular1 = angular1.to_state() + if isinstance(angular2, AngularKetBase): + angular2 = angular2.to_state() + + if any(ket1 == ket2 for ket1 in angular1.kets for ket2 in angular2.kets): + # this also checks states with l_r = unknown + return 0 + + all_diffs = [ + ket1.l_r - ket2.l_r + for ket1 in angular1.kets + for ket2 in angular2.kets + if not is_unknown(ket1.l_r) and not is_unknown(ket2.l_r) + ] + if len(all_diffs) == 0: + raise RuntimeError("Could not calculate l_r difference, this should not happen.") + return int(np.min(np.abs(all_diffs))) + + def generate_matrix_elements_tables( # noqa: C901 - basis: BasisSQDT[AngularKetLS[AllKnown]], + basis: BasisMQDT | BasisSQDT[Any], conn: sqlite3.Connection | None = None, - max_delta_n: float = float("inf"), - all_n_up_to: float = float("inf"), + max_delta_nu: float = float("inf"), + all_nu_up_to: float = float("inf"), + *, + free_memory: bool = False, ) -> dict[str, list[tuple[int, int, float]]]: """Populate matrix element tables for all relevant pairs of states.""" - if basis.coupling_scheme != "LS": - raise ValueError("Only LS coupling scheme is supported for now.") - k_angular_max = max(MatrixElementOperatorRanks[op][1] for op in MATRIX_ELEMENTS_OF_INTEREST.values()) basis.sort_states("nu") # sort by nu == sort by energy list_of_id_state = list(enumerate(basis.states)) - list_of_id_state = sorted(list_of_id_state, key=lambda x: (x[1].angular.l_r, x[1].n, x[0])) + list_of_id_state = sorted(list_of_id_state, key=lambda x: (x[1].angular.calc_exp_qn("l_r"), x[1].nu, x[0])) matrix_elements: dict[str, list[tuple[int, int, float]]] = {tkey: [] for tkey in MATRIX_ELEMENTS_OF_INTEREST} for i, (id1, state1) in enumerate(list_of_id_state): for id2, state2 in list_of_id_state[i:]: - if abs(state1.angular.l_r - state2.angular.l_r) > k_angular_max: + if get_l_r_difference(state1, state2) > k_angular_max: # If the difference in l is larger than k_angular_max, no matrix elements have to be calculated continue - if all(n > all_n_up_to for n in [state1.n, state2.n]) and abs(state1.n - state2.n) > max_delta_n: - # If delta_n is larger than max_delta_n, we dont calculate the matrix elements anymore, + if ( + all(nu > all_nu_up_to for nu in [state1.nu, state2.nu]) + and abs(state1.nu - state2.nu) > max_delta_nu + 0.5 + ): + # If delta_nu is larger than max_delta_nu (+0.5 to not lose states compared to previous max_delta_n) + # we dont calculate the matrix elements anymore, # since these are so small, that they are usually not relevant for further calculations - # However, we keep all dipole interactions with small n (we choose all_n_up_to as a cutoff) + # However, we keep all dipole interactions with small n (we choose all_nu_up_to as a cutoff) # since these are relevant for the spontaneous decay rates continue @@ -76,6 +103,9 @@ def generate_matrix_elements_tables( # noqa: C901 for tkey, me in me_one_pair.items(): matrix_elements[tkey].append((id_tuple[1], id_tuple[0], me)) + if free_memory: + state1.free_memory() + for key, mes in matrix_elements.items(): matrix_elements[key] = sorted(mes) assert len(mes) == 0 or len(COLUMNS) == len(mes[0]) @@ -91,96 +121,11 @@ def generate_matrix_elements_tables( # noqa: C901 def calc_matrix_elements_one_pair( - state1: RydbergStateSQDT[AngularKetLS[AllKnown]], - state2: RydbergStateSQDT[AngularKetLS[AllKnown]], - matrix_elements_of_interest: dict[str, MatrixElementOperator], + state1: RydbergStateBase, state2: RydbergStateBase, matrix_elements_of_interest: dict[str, MatrixElementOperator] ) -> dict[str, float]: matrix_elements: dict[str, float] = {} for tkey, operator in matrix_elements_of_interest.items(): - k_radial, k_angular = MatrixElementOperatorRanks[operator] - - if operator == "magnetic_dipole": - # 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 - ) - g_l = 1 - value_l_tot = calc_reduced_angular_matrix_element_cached( - state1.angular.quantum_numbers, 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 - ) - prefactor = math.sqrt(4 * math.pi / (2 * k_angular + 1)) # e in atomic units is 1 - else: - raise NotImplementedError(f"Operator {operator} not implemented.") - - if angular_matrix_element == 0: - continue - - radial_matrix_element_au = calc_radial_matrix_element_cached( - state1.species, - *(state1.n, state1.nu, state1.angular.l_r), - *(state2.n, state2.nu, state2.angular.l_r), - k_radial, - ) - if radial_matrix_element_au == 0: - continue - - matrix_elements[tkey] = prefactor * radial_matrix_element_au * angular_matrix_element - + me = state2.calc_reduced_matrix_element(state1, operator, unit="a.u.") + if me != 0: + matrix_elements[tkey] = me return matrix_elements - - -@lru_cache(maxsize=100_000) -def calc_reduced_angular_matrix_element_cached( - qns1: tuple[float, ...], - qns2: tuple[float, ...], - operator: AngularOperatorType, - k_angular: int, -) -> float: - ket1: AngularKetLS[AllKnown] = AngularKetLS(*qns1) # type: ignore[call-overload] - ket2: AngularKetLS[AllKnown] = AngularKetLS(*qns2) # type: ignore[call-overload] - # ket2 is the final state and ket1 the initial state - # ket2.calc_reduced_matrix_element(ket1, T, k) gives the reduced matrix element - return ket2.calc_reduced_matrix_element(ket1, operator, k_angular) - - -def calc_radial_matrix_element_cached( - species: str, n1: int, nu1: float, l1: int, n2: int, nu2: float, l2: int, k_radial: int -) -> float: - if k_radial == 0 and nu1 == nu2: - return 1 if l1 == l2 else 0 - - if (nu1, l1) > (nu2, l2): # for better use of the cache and since the radial matrix element is symmetric - return _calc_radial_matrix_element_cached(species, n2, nu2, l2, n1, nu1, l1, k_radial) - - return _calc_radial_matrix_element_cached(species, n1, nu1, l1, n2, nu2, l2, k_radial) - - -# Cache size should be at least on the order of 4 * (all_n_up_to + 2 * max_delta_n) -# however, for the first n until n=all_n_up_to we need an even larger cache size -@lru_cache(maxsize=50_000) -def _calc_radial_matrix_element_cached( - species: str, n1: int, nu1: float, l1: int, n2: int, nu2: float, l2: int, k_radial: int -) -> float: - state1 = get_radial_state_cached(species, n1, nu1, l1) - state2 = get_radial_state_cached(species, n2, nu2, l2) - # state2 is the final state and state1 the initial state - return state2.calc_matrix_element(state1, k_radial, unit="a.u.") - - -# Cache size should be one the order of N_MAX * 4 * 2 -# (since for each initial state we loop over all l' = l, l+1, l+2 and l+3 final states (and all j final)) -@lru_cache(maxsize=2_000) -def get_radial_state_cached(species: str, n: int, nu: float, l: int) -> RadialKet: - """Get the cached rydberg state (where the wavefunction was already calculated).""" - potential = get_potential_class(species)(l) - state = RadialKet(nu, potential, n_expected=n, sign_convention="n_l_1") - state.integrate_wavefunction() - return state diff --git a/src/rydstate/generate_database/generate_states_table.py b/src/rydstate/generate_database/generate_states_table.py index dfafdc24..c895b0ae 100644 --- a/src/rydstate/generate_database/generate_states_table.py +++ b/src/rydstate/generate_database/generate_states_table.py @@ -1,16 +1,17 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any -from rydstate.angular.angular_ket import AngularKetLS +import numpy as np + +from rydstate.rydberg_state.rydberg_sqdt import RydbergStateSQDT if TYPE_CHECKING: import sqlite3 - from rydstate.angular.utils import AllKnown - from rydstate.basis import BasisSQDT - from rydstate.rydberg_state.rydberg_sqdt import RydbergStateSQDT + from rydstate.basis import BasisMQDT, BasisSQDT + from rydstate.rydberg_state.rydberg_base import RydbergStateBase logger = logging.getLogger(__name__) @@ -41,12 +42,10 @@ def generate_states_table( - basis: BasisSQDT[AngularKetLS[AllKnown]], + basis: BasisMQDT | BasisSQDT[Any], conn: sqlite3.Connection | None = None, ) -> list[tuple[float | int | str | bool, ...]]: """Populate the states table for a given species and n-range using BasisSQDT.""" - if basis.coupling_scheme != "LS": - raise ValueError("Only LS coupling scheme is supported for now.") basis.sort_states("nu") # sort by nu == sort by energy states_data: list[tuple[float | int | str | bool, ...]] = [] @@ -64,39 +63,34 @@ def generate_states_table( return states_data -def get_state_data(ids: int, state: RydbergStateSQDT[AngularKetLS[AllKnown]]) -> tuple[float | int | str | bool, ...]: +def get_state_data(ids: int, state: RydbergStateBase) -> tuple[float | int | str | bool, ...]: """Get the data for a given state as a tuple.""" - angular_ket = state.angular - if not isinstance(angular_ket, AngularKetLS): - raise TypeError("Only AngularKetLS is supported for now") - - angular_state = angular_ket.to_state() - - parity = -1 if angular_ket.l_tot % 2 == 1 else 1 + angular = state.angular + underspecified_channel_contribution = sum(abs(coeff) ** 2 for coeff, ket in state if ket.angular.contains_unknown) - is_j_total_momentum = state.element_properties.i_c == 0 - is_calculated_with_mqdt = False + n = state.n if isinstance(state, RydbergStateSQDT) else 0 - return ( + data = ( ids, # id state.get_energy("a.u."), # energy - parity, # parity = (-1)^l_tot - state.n, # n: quantum number - state.nu, # nu = NStar for sqdt - angular_ket.f_tot, # f: quantum number - state.nu, # exp_nui = nu for sqdt - angular_ket.l_tot, # exp_l = l - angular_ket.j_tot, # exp_j = j - angular_ket.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 - 0, # std_j = 0 - 0, # 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 - is_j_total_momentum, # is_j_total_momentum = True for no hyperfine splitting - is_calculated_with_mqdt, # is_calculated_with_mqdt = False for sqdt - 0, # underspecified_channel_contribution = 0 for sqdt + angular.parity, # parity = (-1)^l_tot + n, # n: quantum number + state.nu, # nu + angular.f_tot, # f_tot + state.calc_exp_qn("nui"), # exp_nui + angular.calc_exp_qn("l_tot"), # exp_l + angular.calc_exp_qn("j_tot"), # exp_j + angular.calc_exp_qn("s_tot"), # exp_s + angular.calc_exp_qn("l_r"), # exp_l_ryd + angular.calc_exp_qn("j_r"), # exp_j_ryd = j for sqdt only one valence electron + state.calc_std_qn("nui"), # std_nui = 0 + angular.calc_std_qn("l_tot"), # std_l + angular.calc_std_qn("j_tot"), # std_j + angular.calc_std_qn("s_tot"), # std_s + angular.calc_std_qn("l_r"), # std_l_ryd + angular.calc_std_qn("j_r"), # std_j_ryd + bool(angular.i_c == 0), # is_j_total_momentum + bool(len(state.rydberg_kets) > 1), # is_calculated_with_mqdt + underspecified_channel_contribution, # underspecified_channel_contribution = 0 for sqdt ) + return tuple(x.item() if isinstance(x, np.generic) else x for x in data) diff --git a/src/rydstate/rydberg_state/rydberg_base.py b/src/rydstate/rydberg_state/rydberg_base.py index 9b7991f6..f0610a7d 100644 --- a/src/rydstate/rydberg_state/rydberg_base.py +++ b/src/rydstate/rydberg_state/rydberg_base.py @@ -50,6 +50,16 @@ def __init__(self) -> None: def __iter__(self) -> Iterator[tuple[float, RydbergKet]]: return zip(self.coefficients, self.rydberg_kets, strict=True).__iter__() + def free_memory(self) -> None: + """Release the cached radial and angular data to reduce memory usage. + + This drops the references to the (potentially large) radial wavefunctions of the rydberg kets. + After calling this, matrix elements and overlaps can no longer be calculated for this state. + """ + for rydberg_ket in self.rydberg_kets: + rydberg_ket.__dict__.pop("radial", None) + rydberg_ket.__dict__.pop("angular", None) + @property def norm(self) -> float: """Return the norm of the state (should be 1).""" diff --git a/src/rydstate/rydberg_state/rydberg_ket.py b/src/rydstate/rydberg_state/rydberg_ket.py index 3ae25196..371ab0d5 100644 --- a/src/rydstate/rydberg_state/rydberg_ket.py +++ b/src/rydstate/rydberg_state/rydberg_ket.py @@ -6,6 +6,7 @@ import numpy as np +from rydstate.angular.utils import is_unknown from rydstate.units import MatrixElementOperatorRanks, ureg if TYPE_CHECKING: @@ -39,6 +40,11 @@ def __str__(self) -> str: def calc_reduced_overlap(self, other: RydbergKet) -> float: """Calculate the reduced overlap (ignoring the magnetic quantum number m).""" + if is_unknown(self.radial.l_r) and is_unknown(other.radial.l_r): + return 1 if abs(self.radial.nu - other.radial.nu) < 1e-10 else 0 + if is_unknown(self.radial.l_r) or is_unknown(other.radial.l_r): + return 0 + angular_overlap = self.angular.calc_reduced_overlap(other.angular) if angular_overlap == 0: return 0 @@ -82,6 +88,9 @@ def calc_reduced_matrix_element( f"Operator {operator} not supported, must be one of {list(MatrixElementOperatorRanks.keys())}." ) + if is_unknown(self.radial.l_r) or is_unknown(other.radial.l_r): + return 0 + k_radial, k_angular = MatrixElementOperatorRanks[operator] if operator == "magnetic_dipole": diff --git a/src/rydstate/rydberg_state/rydberg_sqdt.py b/src/rydstate/rydberg_state/rydberg_sqdt.py index dee39a5e..e4bd8fe3 100644 --- a/src/rydstate/rydberg_state/rydberg_sqdt.py +++ b/src/rydstate/rydberg_state/rydberg_sqdt.py @@ -218,7 +218,7 @@ def _energy_au(self) -> float: # type: ignore [override] @cached_property def radial(self) -> RadialKet: """The radial part of the Rydberg electron.""" - return RadialKet(self.nu, self.potential, n_expected=self.n) + return RadialKet(self.nu, self.potential, n_expected=self.n, sign_convention="n_l_1") @cached_property def coefficients(self) -> NDArray: # type: ignore [override] @@ -228,6 +228,13 @@ def coefficients(self) -> NDArray: # type: ignore [override] def rydberg_kets(self) -> list[RydbergKet]: # type: ignore [override] return [RydbergKet(self.angular, self.radial)] + def free_memory(self) -> None: + super().free_memory() + # For SQDT the radial ket is held by these cached properties of the state itself, + # so they have to be dropped as well to actually release the radial wavefunction. + self.__dict__.pop("rydberg_kets", None) + self.__dict__.pop("radial", None) + @overload def get_radial_energy(self, unit: None = None) -> PintFloat: ... diff --git a/tests/test_basis.py b/tests/test_basis.py index edd2b435..a79babec 100644 --- a/tests/test_basis.py +++ b/tests/test_basis.py @@ -32,15 +32,6 @@ def test_alkali_basis(species: str) -> None: assert np.count_nonzero(me_matrix) > 0 -def test_basis_copy() -> None: - basis = BasisSQDT("Sr88", n=(30, 30), coupling_scheme="LS") - basis_copy = basis.copy() - assert basis_copy.coupling_scheme == "LS" - assert basis_copy.species is basis.species - assert basis_copy.states == basis.states - assert basis_copy.states is not basis.states - - def test_sqdt_basis_m_range() -> None: basis = BasisSQDT("H", n=(5, 5), f_tot=(0.5, 0.5), l_r=(0, 0), m=(-0.5, 0.5), coupling_scheme="LS") invalid_basis = BasisSQDT("H", n=(5, 5), f_tot=(0.5, 0.5), l_r=(0, 0), m=(0, 0), coupling_scheme="LS") diff --git a/tests/test_free_memory.py b/tests/test_free_memory.py new file mode 100644 index 00000000..11f2b7aa --- /dev/null +++ b/tests/test_free_memory.py @@ -0,0 +1,70 @@ +# ruff: noqa: SLF001 + +from __future__ import annotations + +import gc +import weakref + +from rydstate import RydbergStateSQDTAlkali +from rydstate.basis import BasisMQDT +from rydstate.radial import RadialKet + + +def _cache_key(radial: RadialKet) -> tuple[tuple[str, object], ...]: + """Return the key under which the given radial ket is stored in RadialKet._instances.""" + return next(key for key, value in RadialKet._instances.items() if value is radial) + + +def test_free_memory_releases_radial_kets_sqdt() -> None: + RadialKet.clear_cached_instances() + + state = RydbergStateSQDTAlkali("Rb", n=50, l=0, j=0.5) + + # access radial and angular kets to ensure they are created (and the radial ket is cached) + radial_refs: list[weakref.ref[RadialKet]] = [] + cache_keys = [] + for ket in state.rydberg_kets: + radial, _angular = ket.radial, ket.angular + radial_refs.append(weakref.ref(radial)) + cache_keys.append(_cache_key(radial)) + del radial, _angular, ket + + assert all(ref() is not None for ref in radial_refs) + assert all(key in RadialKet._instances for key in cache_keys) + + state.free_memory() + gc.collect() + + # the radial kets are no longer referenced and have been evicted from the cache + assert all(ref() is None for ref in radial_refs) + assert all(key not in RadialKet._instances for key in cache_keys) + + +def test_free_memory_releases_radial_kets_mqdt() -> None: + RadialKet.clear_cached_instances() + + basis = BasisMQDT("Yb174", nu=(50, 52), l_r=(0, 2), m=(0, 0)) + # pick a genuine multi-channel state so several radial kets are involved + state = max(basis.states, key=lambda s: len(s.rydberg_kets)) + assert len(state.rydberg_kets) > 1 + + # access radial and angular kets to ensure they are created (and the radial kets are cached) + radial_refs: list[weakref.ref[RadialKet]] = [] + cache_keys = [] + for ket in state.rydberg_kets: + radial, _angular = ket.radial, ket.angular + radial_refs.append(weakref.ref(radial)) + cache_keys.append(_cache_key(radial)) + del radial, _angular, ket + + assert all(ref() is not None for ref in radial_refs) + assert all(key in RadialKet._instances for key in cache_keys) + + # drop the basis so sibling states do not keep the (potentially shared) radial kets alive + del basis + state.free_memory() + gc.collect() + + # the radial kets are no longer referenced and have been evicted from the cache + assert all(ref() is None for ref in radial_refs) + assert all(key not in RadialKet._instances for key in cache_keys) diff --git a/tests/test_generate_database.py b/tests/test_generate_database.py index ed0ae9c6..865c4f75 100644 --- a/tests/test_generate_database.py +++ b/tests/test_generate_database.py @@ -1,12 +1,12 @@ from __future__ import annotations import sqlite3 -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np import pytest from rydstate import RydbergStateSQDTAlkali -from rydstate.basis.basis_sqdt import BasisSQDT +from rydstate.basis import BasisMQDT, BasisSQDT from rydstate.generate_database.generate_database import DATABASE_SQL_FILE from rydstate.generate_database.generate_matrix_elements_table import generate_matrix_elements_tables from rydstate.generate_database.generate_misc_table import generate_wigner_table @@ -16,7 +16,11 @@ from collections.abc import Generator -TEST_SPECIES = ["H", "Li", "Na", "K", "Rb", "Cs", "Sr88", "Yb174"] +TEST_SPECIES_SPECIFIER = [ + *["H", "Li", "Na", "K", "Rb", "Cs"], + *["Sr88_sqdt", "Yb174_sqdt"], + *["Sr87_mqdt", "Sr88_mqdt", "Yb171_mqdt", "Yb173_mqdt", "Yb174_mqdt"], +] @pytest.fixture @@ -45,10 +49,14 @@ def test_get_state_data_for_sqdt_alkali_state() -> None: assert row[12:] == (0, 0, 0, 0, 0, 0, True, False, 0) -@pytest.mark.parametrize("species", TEST_SPECIES) -def test_generate_states_table(species: str, conn: sqlite3.Connection) -> None: - basis = BasisSQDT(species, n=(20, 21), coupling_scheme="LS") - basis.filter_states("l_r", (0, 2)) +@pytest.mark.parametrize("species_specifier", TEST_SPECIES_SPECIFIER) +def test_generate_states_table(species_specifier: str, conn: sqlite3.Connection) -> None: + species = species_specifier.removesuffix("_mqdt").removesuffix("_sqdt") + basis: BasisMQDT | BasisSQDT[Any] + if species_specifier.endswith("_mqdt"): + basis = BasisMQDT(species, nu=(50, 52), l_r=(0, 2)) + else: + basis = BasisSQDT(species, n=(50, 52), l_r=(0, 2), coupling_scheme="LS") basis.sort_states("nu") rows = generate_states_table(basis, conn=conn) @@ -56,27 +64,28 @@ def test_generate_states_table(species: str, conn: sqlite3.Connection) -> None: assert len(basis.states) > 2 assert len(rows) == len(basis.states) - data = np.array(conn.execute("SELECT n, exp_l_ryd, exp_s FROM states").fetchall()) - assert np.allclose(data[:, 0], basis.calc_exp_qn("n")) + data = np.array(conn.execute("SELECT nu, exp_l_ryd, exp_s FROM states").fetchall()) + assert np.allclose(data[:, 0], basis.calc_exp_qn("nu")) assert np.allclose(data[:, 1], basis.calc_exp_qn("l_r")) assert np.allclose(data[:, 2], basis.calc_exp_qn("s_tot")) -@pytest.mark.parametrize("species", TEST_SPECIES) -def test_generate_matrix_elements_table(species: str, conn: sqlite3.Connection) -> None: - basis = BasisSQDT(species, n=(20, 21), coupling_scheme="LS") - basis.filter_states("l_r", (0, 2)) +@pytest.mark.parametrize("species_specifier", TEST_SPECIES_SPECIFIER) +def test_generate_matrix_elements_table(species_specifier: str, conn: sqlite3.Connection) -> None: + species = species_specifier.removesuffix("_mqdt").removesuffix("_sqdt") + basis: BasisMQDT | BasisSQDT[Any] + if species_specifier.endswith("_mqdt"): + basis = BasisMQDT(species, nu=(50, 52), l_r=(0, 2)) + else: + basis = BasisSQDT(species, n=(50, 52), l_r=(0, 2), coupling_scheme="LS") basis.sort_states("nu") - rows_by_table = generate_matrix_elements_tables(basis, conn=conn) + rows_by_table = generate_matrix_elements_tables(basis, conn=conn, free_memory=False) for rows in rows_by_table.values(): assert len(rows) > 2 states = basis.states for row in rows_by_table["matrix_elements_d"]: - state1, state2 = states[row[0]], states[row[1]] - sign1 = (-1) ** (state1.radial.n_expected - state1.radial.l_r - 1) # type: ignore [operator] - sign2 = (-1) ** (state2.radial.n_expected - state2.radial.l_r - 1) # type: ignore [operator] - reference = state2.calc_reduced_matrix_element(state1, "electric_dipole", unit="a.u.") * sign1 * sign2 + reference = states[row[1]].calc_reduced_matrix_element(states[row[0]], "electric_dipole", unit="a.u.") assert np.isclose(row[2], reference)