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, Union
5 :
6 1 : from PySide6.QtWidgets import (
7 : QLabel,
8 : )
9 :
10 1 : from pairinteraction import (
11 : complex as pi_complex,
12 : real as pi_real,
13 : )
14 1 : from pairinteraction_gui.config.base_config import BaseConfig
15 1 : from pairinteraction_gui.qobjects import NamedStackedWidget, QnItemDouble, QnItemInt, WidgetV
16 1 : from pairinteraction_gui.qobjects.item import RangeItem
17 1 : from pairinteraction_gui.theme import label_error_theme, label_theme
18 1 : from pairinteraction_gui.utils import DatabaseMissingError, NoStateFoundError, get_species_type
19 1 : from pairinteraction_gui.worker import MultiThreadWorker
20 :
21 : if TYPE_CHECKING:
22 : from pairinteraction_gui.page import OneAtomPage, TwoAtomsPage
23 : from pairinteraction_gui.qobjects.item import _QnItem
24 :
25 :
26 1 : class BasisConfig(BaseConfig):
27 : """Section for configuring the basis."""
28 :
29 1 : title = "Basis"
30 1 : page: Union["OneAtomPage", "TwoAtomsPage"]
31 :
32 1 : def setupWidget(self) -> None:
33 1 : self.stacked_basis: list[NamedStackedWidget[RestrictionsBase]] = []
34 1 : self.basis_label: list[QLabel] = []
35 :
36 1 : def setupOneBasisAtom(self) -> None:
37 : """Set up the UI components for a single basis atom."""
38 1 : atom = len(self.stacked_basis)
39 :
40 : # Create stacked widget for different species configurations
41 1 : stacked_basis = NamedStackedWidget[RestrictionsBase]()
42 1 : stacked_basis.addNamedWidget(RestrictionsSQDT(), "sqdt")
43 1 : stacked_basis.addNamedWidget(RestrictionsMQDT(), "mqdt")
44 :
45 1 : for _, widget in stacked_basis.items():
46 1 : for _, item in widget.items.items():
47 1 : item.connectAll(lambda atom=atom: self.update_basis_label(atom)) # type: ignore [misc]
48 1 : self.layout().addWidget(stacked_basis)
49 :
50 : # Add a label to display the current basis
51 1 : basis_label = QLabel()
52 1 : basis_label.setStyleSheet(label_theme)
53 1 : basis_label.setWordWrap(True)
54 1 : self.layout().addWidget(basis_label)
55 :
56 : # Store the widgets for later access
57 1 : self.stacked_basis.append(stacked_basis)
58 1 : self.basis_label.append(basis_label)
59 1 : self.update_basis_label(atom)
60 :
61 1 : def update_basis_label(self, atom: int) -> None:
62 1 : worker = MultiThreadWorker(self.get_basis, atom)
63 :
64 1 : def update_result(basis: Union["pi_real.BasisAtom", "pi_complex.BasisAtom"]) -> None:
65 1 : self.basis_label[atom].setText(str(basis) + f"\n ⇒ Basis consists of {basis.number_of_kets} kets")
66 1 : self.basis_label[atom].setStyleSheet(label_theme)
67 :
68 1 : worker.signals.result.connect(update_result)
69 :
70 1 : def update_error(err: Exception) -> None:
71 1 : if isinstance(err, NoStateFoundError):
72 1 : self.basis_label[atom].setText("Ket of interest wrong quantum numbers, first fix those.")
73 0 : elif isinstance(err, DatabaseMissingError):
74 0 : self.basis_label[atom].setText(
75 : "Database required but not downloaded. Please select a different state of interest."
76 : )
77 : else:
78 0 : self.basis_label[atom].setText(str(err))
79 1 : self.basis_label[atom].setStyleSheet(label_error_theme)
80 :
81 1 : worker.signals.error.connect(update_error)
82 :
83 1 : worker.start()
84 :
85 1 : def get_qn_restrictions(self, atom: int) -> dict[str, tuple[float, float]]:
86 : """Return the quantum number restrictions to construct a BasisAtom."""
87 0 : ket = self.page.ket_config.get_ket_atom(atom)
88 0 : basis_widget = self.stacked_basis[atom].currentWidget()
89 0 : delta_qns: dict[str, float] = {
90 : key: item.value() for key, item in basis_widget.items.items() if item.isChecked()
91 : }
92 :
93 0 : qns: dict[str, tuple[float, float]] = {}
94 0 : for key, value in delta_qns.items():
95 0 : qn: float = getattr(ket, key)
96 0 : qns[key] = (qn - value, qn + value)
97 :
98 0 : return qns
99 :
100 1 : def get_basis(
101 : self, atom: int, dtype: Literal["real", "complex"] = "real"
102 : ) -> Union["pi_real.BasisAtom", "pi_complex.BasisAtom"]:
103 : """Return the basis of interest."""
104 0 : ket = self.page.ket_config.get_ket_atom(atom)
105 0 : qn_restrictions = self.get_qn_restrictions(atom)
106 0 : if dtype == "real":
107 0 : return pi_real.BasisAtom(ket.species, **qn_restrictions) # type: ignore [arg-type]
108 0 : return pi_complex.BasisAtom(ket.species, **qn_restrictions) # type: ignore [arg-type]
109 :
110 1 : def get_quantum_number_deltas(self, atom: int = 0) -> dict[str, float]:
111 : """Return the quantum number deltas for the basis of interest."""
112 1 : stacked_basis = self.stacked_basis[atom].currentWidget()
113 1 : return {key: item.value() for key, item in stacked_basis.items.items() if item.isChecked()}
114 :
115 1 : def on_species_changed(self, atom: int, species: str) -> None:
116 : """Handle species selection change."""
117 0 : species_type = get_species_type(species)
118 0 : if "mqdt" in species_type:
119 0 : self.stacked_basis[atom].setCurrentNamedWidget("mqdt")
120 : else:
121 0 : self.stacked_basis[atom].setCurrentNamedWidget("sqdt")
122 0 : self.update_basis_label(atom)
123 :
124 :
125 1 : class BasisConfigOneAtom(BasisConfig):
126 1 : def setupWidget(self) -> None:
127 1 : super().setupWidget()
128 1 : self.setupOneBasisAtom()
129 :
130 :
131 1 : class BasisConfigTwoAtoms(BasisConfig):
132 1 : def setupWidget(self) -> None:
133 1 : super().setupWidget()
134 :
135 1 : self.layout().addWidget(QLabel("<b>Atom 1</b>"))
136 1 : self.setupOneBasisAtom()
137 1 : self.layout().addSpacing(15)
138 :
139 1 : self.layout().addWidget(QLabel("<b>Atom 2</b>"))
140 1 : self.setupOneBasisAtom()
141 1 : self.layout().addSpacing(15)
142 :
143 1 : self.layout().addWidget(QLabel("<b>Pair Basis Restrictions</b>"))
144 1 : self.pair_delta_energy = QnItemDouble(
145 : self,
146 : "ΔEnergy",
147 : vdefault=5,
148 : vmin=0,
149 : unit="GHz",
150 : tooltip="Restriction for the pair energy difference to the state of interest",
151 : )
152 1 : self.layout().addWidget(self.pair_delta_energy)
153 1 : self.pair_m_range = RangeItem(
154 : self,
155 : "Total m",
156 : tooltip_label="pair total angular momentum m",
157 : checked=False,
158 : )
159 1 : self.layout().addWidget(self.pair_m_range)
160 :
161 1 : self.basis_pair_label = QLabel()
162 1 : self.basis_pair_label.setStyleSheet(label_theme)
163 1 : self.basis_pair_label.setWordWrap(True)
164 1 : self.layout().addWidget(self.basis_pair_label)
165 :
166 1 : def update_basis_pair_label(self, basis_pair_label: str) -> None:
167 : """Update the quantum state label with current values."""
168 0 : self.basis_pair_label.setText(basis_pair_label)
169 0 : self.basis_pair_label.setStyleSheet(label_theme)
170 :
171 1 : def clear_basis_pair_label(self) -> None:
172 : """Clear the basis pair label."""
173 0 : self.basis_pair_label.setText("")
174 :
175 :
176 1 : class RestrictionsBase(WidgetV):
177 : """Base class for quantum number configuration."""
178 :
179 1 : margin = (10, 0, 10, 0)
180 1 : spacing = 5
181 :
182 1 : items: dict[str, "_QnItem[Any]"]
183 :
184 1 : def postSetupWidget(self) -> None:
185 1 : for _key, item in self.items.items():
186 1 : self.layout().addWidget(item)
187 :
188 :
189 1 : class RestrictionsSQDT(RestrictionsBase):
190 : """Configuration for alkali atoms using SQDT."""
191 :
192 1 : def setupWidget(self) -> None:
193 1 : self.items = {}
194 1 : self.items["n"] = QnItemInt(self, "Δn", vdefault=3, tooltip="Restriction for the Principal quantum number n")
195 1 : self.items["l"] = QnItemInt(self, "Δl", vdefault=2, tooltip="Restriction for the Orbital angular momentum l")
196 1 : self.items["j"] = QnItemInt(self, "Δj", tooltip="Restriction for the Total angular momentum j", checked=False)
197 1 : self.items["m"] = QnItemInt(self, "Δm", tooltip="Restriction for the Magnetic quantum number m", checked=False)
198 :
199 :
200 1 : class RestrictionsMQDT(RestrictionsBase):
201 : """Configuration for alkali atoms using SQDT."""
202 :
203 1 : def setupWidget(self) -> None:
204 1 : self.items = {}
205 1 : self.items["nu"] = QnItemDouble(
206 : self, "Δnu", vdefault=4, tooltip="Restriction for the Effective principal quantum number nu"
207 : )
208 1 : self.items["s"] = QnItemDouble(self, "Δs", vdefault=0.5, tooltip="Restriction for the Spin s")
209 1 : self.items["j"] = QnItemDouble(self, "Δj", vdefault=3, tooltip="Restriction for the Total angular momentum j")
210 :
211 1 : self.items["f"] = QnItemInt(
212 : self, "Δf", vdefault=5, tooltip="Restriction for the Total angular momentum f", checked=False
213 : )
214 1 : self.items["m"] = QnItemInt(
215 : self, "Δm", vdefault=5, tooltip="Restriction for the Magnetic quantum number m", checked=False
216 : )
|
