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