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