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