Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3           1 : from __future__ import annotations
       4             : 
       5           1 : import logging
       6           1 : from dataclasses import dataclass
       7           1 : from typing import TYPE_CHECKING
       8             : 
       9           1 : import pairinteraction as pi
      10           1 : from pairinteraction_gui.calculate.calculate_base import as_string, dict_to_repl
      11             : 
      12             : if TYPE_CHECKING:
      13             :     import numpy as np
      14             :     from numpy.typing import NDArray
      15             : 
      16             :     from pairinteraction_gui.config.ket_config import QuantumNumbers
      17             :     from pairinteraction_gui.page import LifetimesPage
      18             : 
      19           1 : logger = logging.getLogger(__name__)
      20             : 
      21             : 
      22           1 : @dataclass
      23           1 : class ParametersLifetimes:
      24             :     """Parameters for the lifetimes calculation."""
      25             : 
      26           1 :     species: str
      27           1 :     quantum_numbers: QuantumNumbers
      28           1 :     temperature: float
      29             : 
      30           1 :     @classmethod
      31           1 :     def from_page(cls, page: LifetimesPage) -> ParametersLifetimes:
      32             :         """Create ParametersLifetimes from the GUI page state."""
      33           1 :         return cls(
      34             :             species=page.ket_config.get_species(),
      35             :             quantum_numbers=page.ket_config.get_quantum_numbers(),
      36             :             temperature=page.ket_config.get_temperature(),
      37             :         )
      38             : 
      39           1 :     def to_replacement_dict(self) -> dict[str, str]:
      40             :         """Return a dictionary with the parameters for template replacement."""
      41           1 :         return {
      42             :             "$SPECIES_0": as_string(self.species),
      43             :             "$QUANTUM_NUMBERS_0": dict_to_repl(self.quantum_numbers),
      44             :             "$TEMPERATURE": str(self.temperature),
      45             :         }
      46             : 
      47             : 
      48           1 : @dataclass
      49           1 : class KetData:
      50             :     """Picklable representation of a KetAtom for inter-process transfer."""
      51             : 
      52           1 :     n: int
      53           1 :     label: str
      54             : 
      55           1 :     def __str__(self) -> str:
      56           0 :         return self.label
      57             : 
      58             : 
      59           1 : @dataclass
      60           1 : class ResultsLifetimes:
      61             :     """Results for the lifetimes calculation."""
      62             : 
      63           1 :     kets_sp: list[KetData]
      64           1 :     transition_rates_sp: NDArray[np.float64]
      65           1 :     kets_bbr: list[KetData]
      66           1 :     transition_rates_bbr: NDArray[np.float64]
      67           1 :     lifetime: float
      68             : 
      69             : 
      70           1 : def calculate_lifetimes(parameters: ParametersLifetimes) -> ResultsLifetimes:
      71             :     """Calculate the transition rates for a given ket."""
      72           1 :     return _calculate_lifetimes(parameters)
      73             : 
      74             : 
      75           1 : def _calculate_lifetimes(parameters: ParametersLifetimes) -> ResultsLifetimes:
      76             :     """Make the unwrapped function available for testing."""
      77           1 :     ket = pi.KetAtom(parameters.species, **parameters.quantum_numbers)
      78           1 :     kets_sp, rates_sp = ket.get_spontaneous_transition_rates(unit="1/ms")
      79           1 :     kets_bbr, rates_bbr = ket.get_black_body_transition_rates(parameters.temperature, "K", unit="1/ms")
      80           1 :     lifetime = ket.get_lifetime(parameters.temperature, temperature_unit="K", unit="mus")
      81           1 :     return ResultsLifetimes(
      82             :         kets_sp=[KetData(n=int(k.n), label=str(k)) for k in kets_sp],
      83             :         transition_rates_sp=rates_sp,
      84             :         kets_bbr=[KetData(n=int(k.n), label=str(k)) for k in kets_bbr],
      85             :         transition_rates_bbr=rates_bbr,
      86             :         lifetime=lifetime,
      87             :     )