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.QtWidgets import QCheckBox, QLabel
8 :
9 1 : import pairinteraction as pi_complex
10 1 : import pairinteraction.real as pi_real
11 1 : from pairinteraction_gui.config.base_config import BaseConfig
12 1 : from pairinteraction_gui.qobjects import NamedStackedWidget, QnItemDouble, QnItemInt, WidgetV
13 1 : from pairinteraction_gui.qobjects.item import RangeItem
14 1 : from pairinteraction_gui.utils import DatabaseMissingError, NoStateFoundError, get_species_type
15 1 : from pairinteraction_gui.worker import MultiThreadWorker
16 :
17 : if TYPE_CHECKING:
18 : from PySide6.QtGui import QShowEvent
19 : from PySide6.QtWidgets import QWidget
20 :
21 : from pairinteraction_gui.config.ket_config import QuantumNumbers
22 : from pairinteraction_gui.page import OneAtomPage, TwoAtomsPage
23 : from pairinteraction_gui.qobjects.item import _QnItem
24 :
25 :
26 1 : class QuantumNumberRestrictions(TypedDict, total=False):
27 1 : n: tuple[int, int]
28 1 : nu: tuple[float, float]
29 1 : l: tuple[float, float]
30 1 : s: tuple[float, float]
31 1 : j: tuple[float, float]
32 1 : l_ryd: tuple[float, float]
33 1 : f: tuple[float, float]
34 1 : m: tuple[float, float]
35 :
36 :
37 1 : class BasisConfig(BaseConfig):
38 : """Section for configuring the basis."""
39 :
40 1 : title = "Basis"
41 1 : page: OneAtomPage | TwoAtomsPage
42 :
43 1 : def setupWidget(self) -> None:
44 1 : self.stacked_basis_list: list[NamedStackedWidget[RestrictionsBase]] = []
45 1 : self.basis_label_list: list[QLabel] = []
46 :
47 1 : def setupOneBasisAtom(self) -> None:
48 : """Set up the UI components for a single basis atom."""
49 1 : atom = len(self.stacked_basis_list)
50 :
51 1 : stacked_basis = NamedStackedWidget[RestrictionsBase]()
52 1 : self.layout().addWidget(stacked_basis)
53 :
54 : # Add a label to display the current basis
55 1 : basis_label = QLabel()
56 1 : basis_label.setWordWrap(True)
57 1 : self.layout().addWidget(basis_label)
58 :
59 : # Store the widgets for later access
60 1 : self.stacked_basis_list.append(stacked_basis)
61 1 : self.basis_label_list.append(basis_label)
62 1 : self._set_theme_role(basis_label, "info")
63 1 : self.update_basis_label(atom)
64 :
65 1 : def update_basis_label(self, atom: int) -> None:
66 1 : worker = MultiThreadWorker(self.get_basis, atom)
67 :
68 1 : def update_result(basis: pi_real.BasisAtom | pi_complex.BasisAtom) -> None:
69 1 : self.basis_label_list[atom].setText(str(basis) + f"\n ⇒ Basis consists of {basis.number_of_kets} kets")
70 1 : self._set_theme_role(self.basis_label_list[atom], "info")
71 :
72 1 : worker.signals.result.connect(update_result)
73 :
74 1 : def update_error(err: Exception) -> None:
75 1 : if isinstance(err, NoStateFoundError):
76 1 : self.basis_label_list[atom].setText("Ket of interest wrong quantum numbers, first fix those.")
77 0 : elif isinstance(err, DatabaseMissingError):
78 0 : self.basis_label_list[atom].setText(
79 : "Database required but not downloaded. Please select a different state of interest."
80 : )
81 : else:
82 0 : self.basis_label_list[atom].setText(str(err))
83 1 : self._set_theme_role(self.basis_label_list[atom], "error")
84 :
85 1 : worker.signals.error.connect(update_error)
86 :
87 1 : worker.start()
88 :
89 1 : def get_quantum_number_restrictions(self, atom: int) -> QuantumNumberRestrictions:
90 : """Return the quantum number restrictions to construct a BasisAtom."""
91 1 : ket = self.page.ket_config.get_ket_atom(atom)
92 1 : qns = self.page.ket_config.get_quantum_numbers(atom)
93 1 : delta_qns = self.get_quantum_number_deltas(atom)
94 :
95 1 : qn_restrictions: dict[str, tuple[float, float]] = {}
96 1 : for key, delta in delta_qns.items():
97 1 : if key in qns:
98 1 : qn = qns[key] # type: ignore [literal-required]
99 0 : elif hasattr(ket, key):
100 0 : qn = getattr(ket, key)
101 : else:
102 0 : raise ValueError(f"Quantum number {key} not found in quantum_numbers or KetAtom.")
103 1 : qn_restrictions[key] = (qn - delta, qn + delta)
104 :
105 1 : return qn_restrictions # type: ignore [return-value]
106 :
107 1 : def get_basis(
108 : self, atom: int, dtype: Literal["real", "complex"] = "real"
109 : ) -> pi_real.BasisAtom | pi_complex.BasisAtom:
110 : """Return the basis of interest."""
111 1 : ket = self.page.ket_config.get_ket_atom(atom)
112 1 : qn_restrictions = self.get_quantum_number_restrictions(atom)
113 1 : if dtype == "real":
114 1 : return pi_real.BasisAtom(ket.species, **qn_restrictions)
115 0 : return pi_complex.BasisAtom(ket.species, **qn_restrictions)
116 :
117 1 : def get_quantum_number_deltas(self, atom: int = 0) -> QuantumNumbers:
118 : """Return the quantum number deltas for the basis of interest."""
119 1 : basis_widget = self.stacked_basis_list[atom].currentWidget()
120 1 : return {key: item.value() for key, item in basis_widget.items.items() if item.isChecked()} # type: ignore [return-value]
121 :
122 1 : def on_species_changed(self, atom: int, species: str) -> None:
123 : """Handle species selection change."""
124 1 : if species not in self.stacked_basis_list[atom]._widgets:
125 1 : restrictions_widget = RestrictionsBase.from_species(atom, species, parent=self)
126 1 : self.stacked_basis_list[atom].addNamedWidget(restrictions_widget, species)
127 1 : for _, item in restrictions_widget.items.items():
128 1 : item.connectAll(lambda atom=atom: self.update_basis_label(atom)) # type: ignore [misc]
129 :
130 1 : self.stacked_basis_list[atom].setCurrentNamedWidget(species)
131 1 : self.update_basis_label(atom)
132 :
133 1 : def showEvent(self, event: QShowEvent) -> None:
134 0 : super().showEvent(event)
135 0 : for i in range(len(self.stacked_basis_list)):
136 0 : self.update_basis_label(i)
137 :
138 1 : @staticmethod
139 1 : def _set_theme_role(label: QLabel, role: str) -> None:
140 1 : label.setProperty("themeRole", role)
141 1 : style = label.style()
142 1 : style.unpolish(label)
143 1 : style.polish(label)
144 1 : label.update()
145 :
146 :
147 1 : class BasisConfigOneAtom(BasisConfig):
148 1 : def setupWidget(self) -> None:
149 1 : super().setupWidget()
150 1 : self.setupOneBasisAtom()
151 :
152 :
153 1 : class BasisConfigTwoAtoms(BasisConfig):
154 1 : def setupWidget(self) -> None:
155 1 : super().setupWidget()
156 :
157 1 : self.layout().addWidget(QLabel("<b>Atom 1</b>"))
158 1 : self.setupOneBasisAtom()
159 1 : self.layout().addSpacing(15)
160 :
161 1 : self.layout().addWidget(QLabel("<b>Atom 2</b>"))
162 1 : self.setupOneBasisAtom()
163 1 : self.layout().addSpacing(15)
164 :
165 1 : self.layout().addWidget(QLabel("<b>Pair Basis Restrictions</b>"))
166 1 : self.pair_delta_energy = QnItemDouble(
167 : self,
168 : "ΔEnergy",
169 : vdefault=5,
170 : vmin=0,
171 : unit="GHz",
172 : tooltip="Restriction for the pair energy difference to the state of interest",
173 : )
174 1 : self.layout().addWidget(self.pair_delta_energy)
175 1 : self.pair_m_range = RangeItem(
176 : self,
177 : "Total m",
178 : tooltip_label="pair total angular momentum m",
179 : checked=False,
180 : )
181 1 : self.layout().addWidget(self.pair_m_range)
182 :
183 1 : self.basis_pair_label = QLabel()
184 1 : self._set_theme_role(self.basis_pair_label, "info")
185 1 : self.basis_pair_label.setWordWrap(True)
186 1 : self.layout().addWidget(self.basis_pair_label)
187 :
188 1 : def update_basis_pair_label(self, basis_pair_label: str) -> None:
189 : """Update the quantum state label with current values."""
190 0 : self.basis_pair_label.setText(basis_pair_label)
191 0 : self._set_theme_role(self.basis_pair_label, "info")
192 :
193 1 : def clear_basis_pair_label(self) -> None:
194 : """Clear the basis pair label."""
195 0 : self.basis_pair_label.setText("")
196 :
197 :
198 1 : class RestrictionsBase(WidgetV):
199 : """Base class for quantum number configuration."""
200 :
201 1 : margin = (10, 0, 10, 0)
202 1 : spacing = 5
203 :
204 1 : atom: int
205 1 : items: dict[str, _QnItem[Any]]
206 :
207 1 : def postSetupWidget(self) -> None:
208 1 : for _key, item in self.items.items():
209 1 : self.layout().addWidget(item)
210 :
211 1 : for item in self.items.values():
212 1 : if isinstance(item.checkbox, QCheckBox):
213 1 : item.checkbox.setObjectName(f"atom{self.atom}_{item.checkbox.objectName()}")
214 1 : item.spinbox.setObjectName(f"atom{self.atom}_{item.spinbox.objectName()}")
215 :
216 1 : @classmethod
217 1 : def from_species(cls, atom: int, species: str, parent: QWidget | None = None) -> RestrictionsBase:
218 : """Create a quantum number restriction configuration from the species name."""
219 1 : species_type = get_species_type(species)
220 1 : if species_type == "sqdt_duplet":
221 1 : return RestrictionsSQDT(atom, parent, s_type="halfint", s=0.5)
222 0 : if species_type == "sqdt_singlet":
223 0 : return RestrictionsSQDT(atom, parent, s_type="int", s=0)
224 0 : if species_type == "sqdt_triplet":
225 0 : return RestrictionsSQDT(atom, parent, s_type="int", s=1)
226 0 : if species_type == "mqdt_halfint":
227 0 : return RestrictionsMQDT(atom, parent, f_type="halfint", i=0.5)
228 0 : if species_type == "mqdt_int":
229 0 : return RestrictionsMQDT(atom, parent, f_type="int", i=0)
230 :
231 0 : raise ValueError(f"Unknown species type: {species_type}")
232 :
233 :
234 1 : class RestrictionsSQDT(RestrictionsBase):
235 : """Configuration atoms using SQDT."""
236 :
237 1 : def __init__(
238 : self,
239 : atom: int,
240 : parent: QWidget | None = None,
241 : *,
242 : s_type: Literal["int", "halfint"],
243 : s: float,
244 : ) -> None:
245 1 : assert s_type in ("int", "halfint"), "s_type must be int or halfint"
246 1 : self.atom = atom
247 1 : self.s_type = s_type
248 1 : self.s = s
249 1 : self.items = {}
250 1 : super().__init__(parent)
251 :
252 1 : def setupWidget(self) -> None:
253 1 : self.items["n"] = QnItemInt(self, "Δn", vdefault=3, tooltip="Restriction for the principal quantum number n")
254 1 : self.items["l"] = QnItemInt(self, "Δl", vdefault=2, tooltip="Restriction for the orbital angular momentum l")
255 1 : if self.s != 0:
256 1 : self.items["j"] = QnItemInt(
257 : self, "Δj", tooltip="Restriction for the total angular momentum j", checked=False
258 : )
259 1 : self.items["m"] = QnItemInt(self, "Δm", tooltip="Restriction for the magnetic quantum number m", checked=False)
260 :
261 :
262 1 : class RestrictionsMQDT(RestrictionsBase):
263 : """Configuration for alkali atoms using SQDT."""
264 :
265 1 : def __init__(
266 : self,
267 : atom: int,
268 : parent: QWidget | None = None,
269 : *,
270 : f_type: Literal["int", "halfint"],
271 : i: float,
272 : ) -> None:
273 0 : assert f_type in ("int", "halfint"), "f_type must be int or halfint"
274 0 : self.atom = atom
275 0 : self.f_type = f_type
276 0 : self.i = i
277 0 : self.items = {}
278 0 : super().__init__(parent)
279 :
280 1 : def setupWidget(self) -> None:
281 0 : self.items["nu"] = QnItemDouble(
282 : self, "Δnu", vdefault=4, tooltip="Restriction for the effective principal quantum number nu"
283 : )
284 0 : self.items["s"] = QnItemDouble(self, "Δs", vdefault=0.5, tooltip="Restriction for the spin s")
285 :
286 0 : if self.i == 0:
287 0 : key = "j"
288 0 : description = "Restriction for the total angular momentum j (j=f for I=0)"
289 : else:
290 0 : key = "f"
291 0 : description = "Restriction for the total angular momentum f"
292 0 : self.items[key] = QnItemInt(self, "Δ" + key, vdefault=3, tooltip=description)
293 :
294 0 : self.items["m"] = QnItemInt(
295 : self, "Δm", vdefault=5, tooltip="Restriction for the magnetic quantum number m", checked=False
296 : )
297 :
298 0 : self.items["l_ryd"] = QnItemDouble(
299 : self,
300 : "Δl_ryd",
301 : vdefault=3,
302 : tooltip="Restriction for the orbital angular momentum l_ryd of the Rydberg electron",
303 : checked=False,
304 : )
|
