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              :     )