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