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 : from typing import TYPE_CHECKING, Any, Literal, TypedDict
6 :
7 1 : from PySide6.QtCore import Signal
8 1 : from PySide6.QtWidgets import QComboBox, QLabel
9 :
10 1 : import pairinteraction as pi
11 1 : from pairinteraction_gui.app import Application
12 1 : from pairinteraction_gui.config.base_config import BaseConfig
13 1 : from pairinteraction_gui.qobjects import (
14 : NamedStackedWidget,
15 : QnItemDouble,
16 : QnItemHalfInt,
17 : QnItemInt,
18 : WidgetForm,
19 : WidgetV,
20 : )
21 1 : from pairinteraction_gui.theme import label_error_theme, label_theme
22 1 : from pairinteraction_gui.utils import (
23 : AVAILABLE_SPECIES,
24 : DatabaseMissingError,
25 : NoStateFoundError,
26 : get_custom_error,
27 : get_species_type,
28 : )
29 1 : from pairinteraction_gui.worker import MultiThreadWorker
30 :
31 : if TYPE_CHECKING:
32 : from PySide6.QtWidgets import QWidget
33 :
34 : from pairinteraction_gui.page.lifetimes_page import LifetimesPage
35 : from pairinteraction_gui.qobjects.item import _QnItem
36 :
37 1 : CORE_SPIN_DICT = {"Sr87": 9 / 2, "Sr88": 0, "Yb171": 1 / 2, "Yb173": 5 / 2, "Yb174": 0}
38 :
39 :
40 1 : class QuantumNumbers(TypedDict, total=False):
41 1 : n: int
42 1 : nu: float
43 1 : l: float
44 1 : s: float
45 1 : j: float
46 1 : l_ryd: float
47 1 : f: float
48 1 : m: float
49 :
50 :
51 1 : class KetConfig(BaseConfig):
52 : """Section for configuring the ket of interest."""
53 :
54 1 : title = "State of Interest"
55 :
56 1 : signal_species_changed = Signal(int, str)
57 :
58 1 : def setupWidget(self) -> None:
59 1 : self.species_combo_list: list[QComboBox] = []
60 1 : self.stacked_qn_list: list[NamedStackedWidget[QnBase]] = []
61 1 : self.ket_label_list: list[QLabel] = []
62 :
63 1 : @property
64 1 : def n_atoms(self) -> int:
65 : """Return the number of atoms configured."""
66 1 : return len(self.species_combo_list)
67 :
68 1 : def setupOneKetAtom(self) -> None:
69 : """Set up the UI components for a single ket atom."""
70 1 : atom = len(self.species_combo_list)
71 :
72 : # Species selection
73 1 : species_widget = WidgetForm()
74 1 : species_combo = QComboBox()
75 1 : species_combo.addItems(AVAILABLE_SPECIES) # TODO get available species from pairinteraction
76 1 : species_combo.setToolTip("Select the atomic species")
77 1 : species_widget.layout().addRow("Species", species_combo)
78 1 : species_combo.currentTextChanged.connect(
79 : lambda species, atom=atom: self.signal_species_changed.emit(atom, species)
80 : )
81 1 : self.signal_species_changed.connect(self.on_species_changed)
82 1 : self.layout().addWidget(species_widget)
83 :
84 1 : stacked_qn = NamedStackedWidget[QnBase]()
85 1 : self.layout().addWidget(stacked_qn)
86 :
87 : # Add a label to display the current ket
88 1 : ket_label = QLabel()
89 1 : ket_label.setStyleSheet(label_theme)
90 1 : ket_label.setWordWrap(True)
91 1 : self.layout().addWidget(ket_label)
92 :
93 : # Store the widgets for later access
94 1 : self.species_combo_list.append(species_combo)
95 1 : self.stacked_qn_list.append(stacked_qn)
96 1 : self.ket_label_list.append(ket_label)
97 1 : self.on_qnitem_changed(atom)
98 :
99 1 : def get_species(self, atom: int = 0) -> str:
100 : """Return the selected species of the ... atom."""
101 1 : return self.species_combo_list[atom].currentText()
102 :
103 1 : def get_quantum_numbers(self, atom: int = 0) -> QuantumNumbers:
104 : """Return the quantum numbers of the ... atom."""
105 1 : qn_widget = self.stacked_qn_list[atom].currentWidget()
106 1 : return {key: item.value() for key, item in qn_widget.items.items() if item.isChecked()} # type: ignore [return-value]
107 :
108 1 : def get_ket_atom(self, atom: int, *, ask_download: bool = False) -> pi.KetAtom:
109 : """Return the ket of interest of the ... atom."""
110 1 : species = self.get_species(atom)
111 1 : qns = self.get_quantum_numbers(atom)
112 :
113 1 : try:
114 1 : return pi.KetAtom(species, **qns)
115 1 : except Exception as err:
116 1 : err = get_custom_error(err)
117 1 : if ask_download and isinstance(err, DatabaseMissingError):
118 0 : Application.signals.ask_download_database.emit(err.species)
119 1 : raise err
120 :
121 1 : def on_species_changed(self, atom: int, species: str) -> None:
122 : """Handle species selection change."""
123 1 : if species not in self.stacked_qn_list[atom]._widgets:
124 1 : qn_widget = QnBase.from_species(species, parent=self)
125 1 : self.stacked_qn_list[atom].addNamedWidget(qn_widget, species)
126 1 : for item in qn_widget.items.values():
127 1 : item.connectAll(lambda atom=atom: self.on_qnitem_changed(atom)) # type: ignore [misc]
128 :
129 1 : self.stacked_qn_list[atom].setCurrentNamedWidget(species)
130 1 : self.on_qnitem_changed(atom)
131 :
132 1 : def on_qnitem_changed(self, atom: int) -> None:
133 : """Update the ket label with current values."""
134 1 : try:
135 1 : ket = self.get_ket_atom(atom, ask_download=True)
136 1 : self.ket_label_list[atom].setText(str(ket))
137 1 : self.ket_label_list[atom].setStyleSheet(label_theme)
138 1 : except Exception as err:
139 1 : if isinstance(err, NoStateFoundError):
140 1 : self.ket_label_list[atom].setText("No ket found. Please select different quantum numbers.")
141 1 : elif isinstance(err, DatabaseMissingError):
142 0 : self.ket_label_list[atom].setText(
143 : "Database required but not downloaded. Please select a different species."
144 : )
145 : else:
146 1 : self.ket_label_list[atom].setText(str(err))
147 1 : self.ket_label_list[atom].setStyleSheet(label_error_theme)
148 :
149 :
150 1 : class KetConfigOneAtom(KetConfig):
151 1 : def setupWidget(self) -> None:
152 1 : super().setupWidget()
153 1 : self.setupOneKetAtom()
154 :
155 :
156 1 : class KetConfigLifetimes(KetConfig):
157 1 : worker_label: MultiThreadWorker | None = None
158 1 : worker_plot: MultiThreadWorker | None = None
159 :
160 1 : page: LifetimesPage
161 :
162 1 : def setupWidget(self) -> None:
163 1 : super().setupWidget()
164 1 : self.setupOneKetAtom()
165 1 : self.layout().addSpacing(15)
166 :
167 1 : self.item_temperature = QnItemDouble(
168 : self,
169 : "Temperature",
170 : vdefault=300,
171 : unit="K",
172 : tooltip="Temperature in Kelvin (0K considers only spontaneous decay)",
173 : )
174 1 : self.item_temperature.connectAll(self.update_lifetime_label)
175 1 : self.layout().addWidget(self.item_temperature)
176 1 : self.layout().addSpacing(15)
177 :
178 : # Add a label to display the lifetime
179 1 : self.lifetime_label = QLabel()
180 1 : self.lifetime_label.setStyleSheet(label_theme)
181 1 : self.lifetime_label.setWordWrap(True)
182 1 : self.layout().addWidget(self.lifetime_label)
183 :
184 1 : self.update_lifetime_label()
185 :
186 1 : def get_temperature(self) -> float:
187 0 : return self.item_temperature.value(default=0)
188 :
189 1 : def update_lifetime_label(self) -> None:
190 1 : if self.worker_label and self.worker_label.isRunning():
191 0 : self.worker_label.quit()
192 0 : self.worker_label.wait()
193 :
194 1 : def get_lifetime() -> float:
195 0 : ket = self.get_ket_atom(0, ask_download=True)
196 0 : temperature = self.get_temperature()
197 0 : return ket.get_lifetime(temperature, temperature_unit="K", unit="mus")
198 :
199 1 : def update_result(lifetime: float) -> None:
200 0 : self.lifetime_label.setText(f"Lifetime: {lifetime:.3f} μs")
201 0 : self.lifetime_label.setStyleSheet(label_theme)
202 :
203 1 : def update_error(err: Exception) -> None:
204 1 : self.lifetime_label.setText("Ket not found.")
205 1 : self.lifetime_label.setStyleSheet(label_error_theme)
206 1 : self.page.plotwidget.clear()
207 :
208 1 : self.worker_label = MultiThreadWorker(get_lifetime)
209 1 : self.worker_label.signals.result.connect(update_result)
210 1 : self.worker_label.signals.error.connect(update_error)
211 1 : self.worker_label.start()
212 :
213 1 : if self.worker_plot and self.worker_plot.isRunning():
214 0 : self.worker_plot.quit()
215 0 : self.worker_plot.wait()
216 :
217 1 : self.worker_plot = MultiThreadWorker(self.page.calculate)
218 1 : self.worker_plot.signals.result.connect(self.page.update_plot)
219 1 : self.worker_plot.start()
220 :
221 1 : def on_qnitem_changed(self, atom: int) -> None:
222 1 : super().on_qnitem_changed(atom)
223 1 : if hasattr(self, "item_temperature"): # not yet initialized the first time this method is called
224 1 : self.update_lifetime_label()
225 :
226 :
227 1 : class KetConfigTwoAtoms(KetConfig):
228 1 : def setupWidget(self) -> None:
229 1 : super().setupWidget()
230 :
231 1 : self.layout().addWidget(QLabel("<b>Atom 1</b>"))
232 1 : self.setupOneKetAtom()
233 1 : self.layout().addSpacing(15)
234 :
235 1 : self.layout().addWidget(QLabel("<b>Atom 2</b>"))
236 1 : self.setupOneKetAtom()
237 :
238 :
239 1 : class QnBase(WidgetV):
240 : """Base class for quantum number configuration."""
241 :
242 1 : margin = (10, 0, 10, 0)
243 1 : spacing = 5
244 :
245 1 : items: dict[str, _QnItem[Any]]
246 :
247 1 : def postSetupWidget(self) -> None:
248 1 : for item in self.items.values():
249 1 : self.layout().addWidget(item)
250 :
251 1 : @classmethod
252 1 : def from_species(cls, species: str, parent: QWidget | None = None) -> QnBase:
253 : """Create a quantum number configuration from the species name."""
254 1 : species_type = get_species_type(species)
255 1 : if species_type == "sqdt_duplet":
256 1 : return QnSQDT(parent, s_type="halfint", s=0.5)
257 0 : if species_type == "sqdt_singlet":
258 0 : return QnSQDT(parent, s_type="int", s=0)
259 0 : if species_type == "sqdt_triplet":
260 0 : return QnSQDT(parent, s_type="int", s=1)
261 :
262 0 : element = species.split("_")[0]
263 0 : if species_type == "mqdt_halfint":
264 0 : i = CORE_SPIN_DICT.get(element, 0.5)
265 0 : return QnMQDT(parent, f_type="halfint", i=i)
266 0 : if species_type == "mqdt_int":
267 0 : i = CORE_SPIN_DICT.get(element, 0)
268 0 : return QnMQDT(parent, f_type="int", i=i)
269 :
270 0 : raise ValueError(f"Unknown species type: {species_type}")
271 :
272 :
273 1 : class QnSQDT(QnBase):
274 : """Configuration for atoms using SQDT."""
275 :
276 1 : def __init__(
277 : self,
278 : parent: QWidget | None = None,
279 : *,
280 : s_type: Literal["int", "halfint"],
281 : s: float,
282 : ) -> None:
283 1 : assert s_type in ("int", "halfint"), "s_type must be int or halfint"
284 1 : self.s_type = s_type
285 1 : self.s = s
286 1 : self.items = {}
287 1 : super().__init__(parent)
288 :
289 1 : def setupWidget(self) -> None:
290 1 : self.items["n"] = QnItemInt(self, "n", vmin=1, vdefault=80, tooltip="Principal quantum number n")
291 1 : self.items["l"] = QnItemInt(self, "l", vmin=0, tooltip="Orbital angular momentum l")
292 :
293 1 : s = self.s
294 1 : if self.s_type == "int":
295 0 : if s != 0:
296 0 : self.items["j"] = QnItemInt(self, "j", vmin=int(s), vdefault=int(s), tooltip="Total angular momentum j")
297 0 : self.items["m"] = QnItemInt(self, "m", vmin=-999, vmax=999, vdefault=0, tooltip="Magnetic quantum number m")
298 : else:
299 1 : self.items["j"] = QnItemHalfInt(self, "j", vmin=s, vdefault=s, tooltip="Total angular momentum j")
300 1 : self.items["m"] = QnItemHalfInt(
301 : self, "m", vmin=-999.5, vmax=999.5, vdefault=0.5, tooltip="Magnetic quantum number m"
302 : )
303 :
304 :
305 1 : class QnMQDT(QnBase):
306 : """Configuration for alkaline-earth atoms using MQDT."""
307 :
308 1 : def __init__(
309 : self,
310 : parent: QWidget | None = None,
311 : *,
312 : f_type: Literal["int", "halfint"],
313 : i: float,
314 : ) -> None:
315 0 : assert f_type in ("int", "halfint"), "f_type must be int or halfint"
316 0 : self.f_type = f_type
317 0 : self.i = i
318 0 : self.items = {}
319 0 : super().__init__(parent)
320 :
321 1 : def setupWidget(self) -> None:
322 0 : self.items["nu"] = QnItemDouble(
323 : self, "nu", vmin=1, vdefault=80, vstep=1, tooltip="Effective principal quantum number nu"
324 : )
325 0 : self.items["s"] = QnItemDouble(self, "s", vmin=0, vmax=1, vstep=0.1, tooltip="Spin s")
326 :
327 0 : if self.i == 0:
328 0 : key = "j"
329 0 : description = "Total angular momentum j (j=f for I=0)"
330 : else:
331 0 : key = "f"
332 0 : description = "Total angular momentum f"
333 0 : if self.f_type == "int":
334 0 : self.items[key] = QnItemInt(self, key, vmin=0, vdefault=int(self.i), tooltip=description)
335 0 : self.items["m"] = QnItemInt(self, "m", vmin=-999, vmax=999, vdefault=0, tooltip="Magnetic quantum number m")
336 : else:
337 0 : self.items[key] = QnItemHalfInt(self, key, vmin=0.5, vdefault=self.i, tooltip=description)
338 0 : self.items["m"] = QnItemHalfInt(
339 : self, "m", vmin=-999.5, vmax=999.5, vdefault=0.5, tooltip="Magnetic quantum number m"
340 : )
341 :
342 0 : self.items["l_ryd"] = QnItemDouble(
343 : self,
344 : "l_ryd",
345 : vmin=0,
346 : vstep=1,
347 : tooltip="Orbital angular momentum l_ryd of the Rydberg electron",
348 : checked=False,
349 : )
|
