Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : from __future__ import annotations
       5             : 
       6           1 : from typing import TYPE_CHECKING
       7             : 
       8           1 : import numpy as np
       9           1 : import pytest
      10           1 : from pairinteraction import ureg
      11           1 : from scipy.optimize import curve_fit
      12             : 
      13             : if TYPE_CHECKING:
      14             :     from pairinteraction.units import NDArray
      15             : 
      16             :     from .utils import PairinteractionModule
      17             : 
      18           1 : from .utils import no_log_propagation, skip_value_check_if_shrunk_database
      19             : 
      20             : # The shrunk database that comes with the repository only contains a few low-lying states,
      21             : # so that the decay of a Rydberg state is only captured partially and the calculated lifetimes are far off.
      22             : # The tests are run nevertheless so that all code is executed, but the checks of the resulting values are skipped
      23             : 
      24             : 
      25           1 : def test_lifetime(pi_module: PairinteractionModule) -> None:
      26             :     """Test calculating the lifetime of a state."""
      27           1 :     ket = pi_module.KetAtom("Yb174_mqdt", n=64, l=0, j=0, m=0)
      28             : 
      29           1 :     lifetime1 = ket.get_lifetime(temperature=300, temperature_unit="K", unit="us")
      30           1 :     lifetime2 = ket.get_lifetime(temperature=300 * ureg.K, unit="us")
      31           1 :     lifetime3 = ket.get_lifetime(temperature=300 * ureg.K)
      32           1 :     lifetime4 = ket.get_lifetime(unit="us")
      33             : 
      34           1 :     assert lifetime1 == lifetime2 == lifetime3.to(ureg.us).magnitude < lifetime4
      35             : 
      36           1 :     skip_value_check_if_shrunk_database()
      37           0 :     assert pytest.approx(lifetime1, rel=0.05) == 45.7  # NOSONAR
      38           0 :     assert pytest.approx(lifetime4, rel=0.05) == 59.3  # NOSONAR
      39             : 
      40             : 
      41           1 : def test_lifetime_scaling(pi_module: PairinteractionModule) -> None:
      42             :     """Test the scaling of the lifetime with the principal quantum number."""
      43             : 
      44           1 :     def fit_function(x: NDArray, a: float, b: float) -> NDArray:
      45           1 :         return a * x + b
      46             : 
      47           1 :     n_list = list(range(58, 68, 1))
      48             : 
      49             :     # S states
      50           1 :     kets = [pi_module.KetAtom("Rb", n=n, l=0, j=0.5, m=0.5) for n in n_list]
      51           1 :     nu = [ket.nu for ket in kets]
      52           1 :     with no_log_propagation("cpp"):  # surpress low n state warnings from lifetime calculation
      53           1 :         lifetimes = [ket.get_lifetime(unit="us") for ket in kets]
      54           1 :     popt, _ = curve_fit(fit_function, np.log(nu), np.log(lifetimes))
      55             : 
      56           1 :     skip_value_check_if_shrunk_database()
      57           0 :     assert np.isclose(popt[0], 3, atol=0.02)
      58             : 
      59             :     # Circular states
      60             :     # In contrast to the states above, they cannot be obtained from the shrunk database at all
      61             :     # because the shrunk database only contains states with a small quantum number l_ryd.
      62           0 :     kets = [pi_module.KetAtom("Rb", n=n, l=n - 1, j=n - 0.5, m=n - 0.5) for n in n_list]
      63           0 :     nu = [ket.nu for ket in kets]
      64           0 :     with no_log_propagation("cpp"):  # surpress low n state warnings from lifetime calculation
      65           0 :         lifetimes = [ket.get_lifetime(unit="us") for ket in kets]
      66           0 :     popt, _ = curve_fit(fit_function, np.log(nu), np.log(lifetimes))
      67           0 :     assert np.isclose(popt[0], 5, atol=0.02)
      68             : 
      69             : 
      70           1 : def test_transition_rates(pi_module: PairinteractionModule) -> None:
      71             :     """Test calculating transition rates to other states."""
      72           1 :     ket = pi_module.KetAtom("Yb174_mqdt", n=60, l=0, j=0, m=0)
      73             : 
      74           1 :     with no_log_propagation("cpp"):  # surpress low n state warnings from lifetime calculation
      75           1 :         lifetime = ket.get_lifetime(temperature=300, temperature_unit="K", unit="us")
      76           1 :         kets_sp, rates_sp = ket.get_spontaneous_transition_rates(unit="MHz")
      77           1 :         kets_bbr, rates_bbr = ket.get_black_body_transition_rates(temperature=300, temperature_unit="K", unit="MHz")
      78             : 
      79           1 :     assert len(rates_sp) == len(kets_sp)
      80           1 :     assert len(rates_bbr) == len(kets_bbr)
      81           1 :     assert np.isclose(1 / (sum(rates_sp) + sum(rates_bbr)), lifetime)