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 : j_ryd: tuple[float, float]
34 1 : l_core: tuple[float, float]
35 1 : j_core: tuple[float, float]
36 1 : f_core: tuple[float, float]
37 1 : f: tuple[float, float]
38 1 : m: tuple[float, float]
39 :
40 :
41 1 : class BasisConfig(BaseConfig):
42 : """Section for configuring the basis."""
43 :
44 1 : title = "Basis"
45 1 : page: OneAtomPage | TwoAtomsPage
46 :
47 1 : def setupWidget(self) -> None:
48 1 : self.stacked_basis_list: list[NamedStackedWidget[RestrictionsBase]] = []
49 1 : self.basis_label_list: list[QLabel] = []
50 :
51 1 : def setupOneBasisAtom(self) -> None:
52 : """Set up the UI components for a single basis atom."""
53 1 : stacked_basis = NamedStackedWidget[RestrictionsBase]()
54 1 : self.layout().addWidget(stacked_basis)
55 :
56 : # Add a label to display the current basis
57 1 : basis_label = QLabel()
58 1 : basis_label.setWordWrap(True)
59 1 : self.layout().addWidget(basis_label)
60 :
61 : # Store the widgets for later access
62 1 : self.stacked_basis_list.append(stacked_basis)
63 1 : self.basis_label_list.append(basis_label)
64 1 : self._set_theme_role(basis_label, "info")
65 :
66 1 : def update_basis_label(self, atom: int) -> None:
67 1 : worker = MultiThreadWorker(self.get_basis, atom)
68 :
69 1 : def update_result(basis: pi_real.BasisAtom | pi_complex.BasisAtom) -> None:
70 1 : self.basis_label_list[atom].setText(str(basis) + f"\n ⇒ Basis consists of {basis.number_of_kets} kets")
71 1 : self._set_theme_role(self.basis_label_list[atom], "info")
72 :
73 1 : worker.signals.result.connect(update_result)
74 :
75 1 : def update_error(err: Exception) -> None:
76 0 : if isinstance(err, NoStateFoundError):
77 0 : self.basis_label_list[atom].setText("Ket of interest wrong quantum numbers, first fix those.")
78 0 : elif isinstance(err, DatabaseMissingError):
79 0 : self.basis_label_list[atom].setText(
80 : "Database required but not downloaded. Please select a different state of interest."
81 : )
82 : else:
83 0 : self.basis_label_list[atom].setText(str(err))
84 0 : self._set_theme_role(self.basis_label_list[atom], "error")
85 :
86 1 : worker.signals.error.connect(update_error)
87 :
88 1 : worker.start()
89 :
90 1 : def get_quantum_number_restrictions(self, atom: int) -> QuantumNumberRestrictions:
91 : """Return the quantum number restrictions to construct a BasisAtom."""
92 1 : ket = self.page.ket_config.get_ket_atom(atom)
93 1 : qns = self.page.ket_config.get_quantum_numbers(atom)
94 1 : delta_qns = self.get_quantum_number_deltas(atom)
95 :
96 1 : qn_restrictions: dict[str, tuple[float, float]] = {}
97 1 : for key, delta in delta_qns.items():
98 1 : if key in qns:
99 1 : qn = qns[key] # type: ignore [literal-required]
100 0 : elif hasattr(ket, key):
101 0 : qn = getattr(ket, key)
102 : else:
103 0 : raise ValueError(f"Quantum number {key} not found in quantum_numbers or KetAtom.")
104 1 : qn_restrictions[key] = (qn - delta, qn + delta)
105 :
106 1 : return qn_restrictions # type: ignore [return-value]
107 :
108 1 : def get_basis(
109 : self, atom: int, dtype: Literal["real", "complex"] = "real"
110 : ) -> pi_real.BasisAtom | pi_complex.BasisAtom:
111 : """Return the basis of interest."""
112 1 : ket = self.page.ket_config.get_ket_atom(atom)
113 1 : qn_restrictions = self.get_quantum_number_restrictions(atom)
114 1 : if dtype == "real":
115 1 : return pi_real.BasisAtom(ket.species, **qn_restrictions)
116 0 : return pi_complex.BasisAtom(ket.species, **qn_restrictions)
117 :
118 1 : def get_quantum_number_deltas(self, atom: int = 0) -> QuantumNumbers:
119 : """Return the quantum number deltas for the basis of interest."""
120 1 : basis_widget = self.stacked_basis_list[atom].currentWidget()
121 1 : return {key: item.value() for key, item in basis_widget.items.items() if item.isChecked()} # type: ignore [return-value]
122 :
123 1 : def on_species_changed(self, atom: int, species: str) -> None:
124 : """Handle species selection change."""
125 1 : if species not in self.stacked_basis_list[atom]._widgets:
126 1 : restrictions_widget = RestrictionsBase.from_species(atom, species, parent=self)
127 1 : self.stacked_basis_list[atom].addNamedWidget(restrictions_widget, species)
128 1 : for _, item in restrictions_widget.items.items():
129 1 : item.connectAll(lambda atom=atom: self.update_basis_label(atom)) # type: ignore [misc]
130 :
131 1 : self.stacked_basis_list[atom].setCurrentNamedWidget(species)
132 1 : self.update_basis_label(atom)
133 :
134 1 : def showEvent(self, event: QShowEvent) -> None:
135 0 : super().showEvent(event)
136 0 : for i in range(len(self.stacked_basis_list)):
137 0 : self.update_basis_label(i)
138 :
139 1 : @staticmethod
140 1 : def _set_theme_role(label: QLabel, role: str) -> None:
141 1 : label.setProperty("themeRole", role)
142 1 : style = label.style()
143 1 : style.unpolish(label)
144 1 : style.polish(label)
145 1 : label.update()
146 :
147 :
148 1 : class BasisConfigOneAtom(BasisConfig):
149 1 : def setupWidget(self) -> None:
150 1 : super().setupWidget()
151 1 : self.setupOneBasisAtom()
152 :
153 :
154 1 : class BasisConfigTwoAtoms(BasisConfig):
155 1 : def setupWidget(self) -> None:
156 1 : super().setupWidget()
157 :
158 1 : self.layout().addWidget(QLabel("<b>Atom 1</b>"))
159 1 : self.setupOneBasisAtom()
160 1 : self.layout().addSpacing(15)
161 :
162 1 : self.layout().addWidget(QLabel("<b>Atom 2</b>"))
163 1 : self.setupOneBasisAtom()
164 1 : self.layout().addSpacing(15)
165 :
166 1 : self.layout().addWidget(QLabel("<b>Pair Basis Restrictions</b>"))
167 1 : self.pair_delta_energy = QnItemDouble(
168 : self,
169 : "ΔEnergy",
170 : vdefault=5,
171 : vmin=0,
172 : unit="GHz",
173 : tooltip="Restriction for the pair energy difference to the state of interest",
174 : )
175 1 : self.layout().addWidget(self.pair_delta_energy)
176 1 : self.pair_m_range = RangeItem(
177 : self,
178 : "Total m",
179 : tooltip_label="pair total angular momentum m",
180 : checked=False,
181 : )
182 1 : self.layout().addWidget(self.pair_m_range)
183 :
184 1 : self.basis_pair_label = QLabel()
185 1 : self._set_theme_role(self.basis_pair_label, "info")
186 1 : self.basis_pair_label.setWordWrap(True)
187 1 : self.layout().addWidget(self.basis_pair_label)
188 :
189 1 : def update_basis_pair_label(self, basis_pair_label: str) -> None:
190 : """Update the quantum state label with current values."""
191 0 : self.basis_pair_label.setText(basis_pair_label)
192 0 : self._set_theme_role(self.basis_pair_label, "info")
193 :
194 1 : def clear_basis_pair_label(self) -> None:
195 : """Clear the basis pair label."""
196 0 : self.basis_pair_label.setText("")
197 :
198 :
199 1 : class RestrictionsBase(WidgetV):
200 : """Base class for quantum number configuration."""
201 :
202 1 : margin = (10, 0, 10, 0)
203 1 : spacing = 5
204 :
205 1 : atom: int
206 1 : items: dict[str, _QnItem[Any]]
207 :
208 1 : def postSetupWidget(self) -> None:
209 1 : for _key, item in self.items.items():
210 1 : self.layout().addWidget(item)
211 :
212 1 : for item in self.items.values():
213 1 : if isinstance(item.checkbox, QCheckBox):
214 1 : item.checkbox.setObjectName(f"atom{self.atom}_{item.checkbox.objectName()}")
215 1 : item.spinbox.setObjectName(f"atom{self.atom}_{item.spinbox.objectName()}")
216 :
217 1 : @classmethod
218 1 : def from_species(cls, atom: int, species: str, parent: QWidget | None = None) -> RestrictionsBase:
219 : """Create a quantum number restriction configuration from the species name."""
220 1 : species_type = get_species_type(species)
221 1 : if species_type in ("sqdt_monovalent", "sqdt_divalent"):
222 1 : return RestrictionsSQDT(atom, parent, species_type=species_type)
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 : species_type: Literal["sqdt_monovalent", "sqdt_divalent"],
240 : ) -> None:
241 1 : self.atom = atom
242 1 : self.species_type = species_type
243 1 : self.items = {}
244 1 : super().__init__(parent)
245 :
246 1 : def setupWidget(self) -> None:
247 1 : self.items["n"] = QnItemInt(self, "Δn", vdefault=3, tooltip="Restriction for the principal quantum number n")
248 1 : self.items["l"] = QnItemInt(self, "Δl", vdefault=2, tooltip="Restriction for the orbital angular momentum l")
249 1 : if self.species_type == "sqdt_divalent":
250 0 : self.items["s"] = QnItemInt(
251 : self,
252 : "Δs",
253 : tooltip="Restriction for the spin s (Δs=0 restricts the basis to the singlet or triplet sector)",
254 : )
255 1 : self.items["j"] = QnItemInt(self, "Δj", tooltip="Restriction for the total angular momentum j", checked=False)
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 : )
|
