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