Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : import logging
       5           1 : from typing import TYPE_CHECKING
       6             : 
       7           1 : import mplcursors
       8           1 : import numpy as np
       9             : 
      10           1 : from pairinteraction_gui.config import KetConfigLifetimes
      11           1 : from pairinteraction_gui.page.base_page import SimulationPage
      12           1 : from pairinteraction_gui.plotwidget.plotwidget import PlotWidget
      13           1 : from pairinteraction_gui.qobjects import show_status_tip
      14             : 
      15             : if TYPE_CHECKING:
      16             :     import pairinteraction.real as pi
      17             : 
      18           1 : logger = logging.getLogger(__name__)
      19             : 
      20             : 
      21           1 : class LifetimesPage(SimulationPage):
      22             :     """Page for calculating lifetimes."""
      23             : 
      24           1 :     title = "Lifetimes"
      25           1 :     tooltip = "Calculate the lifetimes and transition rates for a specified ket."
      26             : 
      27           1 :     ket_config: KetConfigLifetimes
      28             : 
      29           1 :     def setupWidget(self) -> None:
      30           1 :         self.plotwidget = PlotWidget(self)
      31           1 :         self.layout().addWidget(self.plotwidget)
      32           1 :         super().setupWidget()
      33             : 
      34           1 :         show_status_tip(self, "Ready", timeout=1)
      35             : 
      36             :         # all attributes of instance BaseConfig will be added to the toolbox in postSetupWidget
      37           1 :         self.ket_config = KetConfigLifetimes(self)
      38             : 
      39           1 :     def calculate(self) -> None:
      40             :         # since this calculation is rather fast, we can just call it directly and dont have to use a different process
      41           0 :         ket = self.ket_config.get_ket_atom(0)
      42           0 :         temperature = self.ket_config.get_temperature()
      43           0 :         self.kets_sp, self.transition_rates_sp = ket.get_spontaneous_transition_rates(unit="1/ms")
      44           0 :         self.kets_bbr, self.transition_rates_bbr = ket.get_black_body_transition_rates(temperature, "K", unit="1/ms")
      45             : 
      46           1 :     def update_plot(self) -> None:
      47           0 :         ax = self.plotwidget.canvas.ax
      48           0 :         ax.clear()
      49             : 
      50           0 :         n_list = np.arange(0, np.max([s.n for s in self.kets_bbr]) + 1)
      51           0 :         sorted_rates: dict[str, dict[int, list[tuple[pi.KetAtom, float]]]] = {}
      52           0 :         for key, kets, rates in [
      53             :             ("BBR", self.kets_bbr, self.transition_rates_bbr),
      54             :             ("SP", self.kets_sp, self.transition_rates_sp),
      55             :         ]:
      56           0 :             sorted_rates[key] = {n: [] for n in n_list}
      57           0 :             for i, s in enumerate(kets):
      58           0 :                 sorted_rates[key][s.n].append((s, rates[i]))
      59           0 :         self.sorted_rates = sorted_rates
      60             : 
      61           0 :         rates_summed = {key: [sum(rates for _, rates in sorted_rates[key][n]) for n in n_list] for key in sorted_rates}
      62           0 :         bar_sp = ax.bar(n_list, rates_summed["SP"], label="Spontaneous Decay", color="blue", alpha=0.8)
      63           0 :         bar_bbr = ax.bar(n_list, rates_summed["BBR"], label="Black Body Radiation", color="red", alpha=0.8)
      64           0 :         self.artists = (bar_sp, bar_bbr)
      65           0 :         ax.legend()
      66             : 
      67           0 :         ax.set_xlabel("Principal Quantum Number $n$")
      68           0 :         ax.set_ylabel(r"Transition Rates (1 / ms)")
      69             : 
      70           0 :         self.add_cursor()
      71             : 
      72           0 :         self.plotwidget.canvas.draw()
      73             : 
      74           1 :     def add_cursor(self) -> None:
      75             :         """Add interactive cursor to the plot."""
      76             :         # Remove any existing cursors to avoid duplicates
      77           0 :         if hasattr(self, "mpl_cursor"):
      78           0 :             if hasattr(self.mpl_cursor, "remove"):  # type: ignore
      79           0 :                 self.mpl_cursor.remove()  # type: ignore
      80           0 :             del self.mpl_cursor  # type: ignore
      81             : 
      82           0 :         self.mpl_cursor = mplcursors.cursor(
      83             :             self.artists,
      84             :             hover=mplcursors.HoverMode.Transient,
      85             :             annotation_kwargs={
      86             :                 "bbox": {"boxstyle": "round,pad=0.5", "fc": "white", "alpha": 0.9, "ec": "gray"},
      87             :                 "arrowprops": {"arrowstyle": "->", "connectionstyle": "arc3", "color": "gray"},
      88             :             },
      89             :         )
      90             : 
      91           0 :         @self.mpl_cursor.connect("add")
      92           0 :         def on_add(sel: mplcursors.Selection) -> None:
      93           0 :             label = sel.artist.get_label()
      94           0 :             x, y, width, height = sel.artist[sel.index].get_bbox().bounds
      95             : 
      96           0 :             n = round(x + width / 2)
      97           0 :             key = "BBR" if "Black Body" in label else "SP"
      98           0 :             state_text = "\n".join(f"{s}: {r:.5f}/ms" for (s, r) in self.sorted_rates[key][n])
      99           0 :             text = f"{label} to n={n}:\n{state_text}"
     100             : 
     101           0 :             sel.annotation.set(text=text, position=(0, 20), anncoords="offset points")
     102           0 :             sel.annotation.xy = (x + width / 2, y + height / 2)