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 : import contextlib
6 1 : import logging
7 1 : import math
8 1 : from typing import TYPE_CHECKING, Any
9 :
10 1 : import numpy as np
11 1 : from matplotlib.cm import ScalarMappable
12 1 : from matplotlib.collections import LineCollection
13 1 : from matplotlib.colors import Normalize
14 1 : from PySide6.QtGui import QPalette
15 1 : from PySide6.QtWidgets import QHBoxLayout
16 1 : from scipy.optimize import curve_fit
17 :
18 1 : from pairinteraction.state.state_atom import StateAtom
19 1 : from pairinteraction.visualization.colormaps import alphamagma
20 1 : from pairinteraction_gui.plotwidget.canvas import MatplotlibCanvas
21 1 : from pairinteraction_gui.plotwidget.navigation_toolbar import CustomNavigationToolbar
22 1 : from pairinteraction_gui.qobjects import WidgetV
23 1 : from pairinteraction_gui.qobjects.events import show_status_tip
24 1 : from pairinteraction_gui.theme import theme_manager
25 :
26 : if TYPE_CHECKING:
27 : from collections.abc import Callable, Sequence
28 : from typing import Concatenate
29 :
30 : import matplotlib as mpl
31 : from numpy.typing import NDArray
32 :
33 : from pairinteraction.state import StateBase
34 : from pairinteraction_gui.calculate.calculate_base import Parameters, Results
35 : from pairinteraction_gui.calculate.calculate_lifetimes import KetData, ParametersLifetimes, ResultsLifetimes
36 : from pairinteraction_gui.page import SimulationPage
37 :
38 1 : logger = logging.getLogger(__name__)
39 :
40 :
41 1 : class PlotWidget(WidgetV):
42 : """Widget for displaying plots with controls."""
43 :
44 1 : margin = (0, 0, 0, 0)
45 1 : spacing = 15
46 1 : _annotations: dict[Any, mpl.text.Annotation]
47 :
48 1 : def __init__(self, parent: SimulationPage) -> None:
49 : """Initialize the base section."""
50 1 : self.page = parent
51 1 : super().__init__(parent)
52 :
53 1 : self._annotations = {}
54 1 : self._click_cids: list[int] = []
55 :
56 1 : def setupWidget(self) -> None:
57 1 : self.canvas = MatplotlibCanvas(self)
58 1 : self.navigation_toolbar = CustomNavigationToolbar(self.canvas, self)
59 1 : self.navigation_toolbar.setObjectName("PlotNavigationToolBar")
60 :
61 1 : top_layout = QHBoxLayout()
62 1 : top_layout.addStretch(1)
63 1 : top_layout.addWidget(self.navigation_toolbar)
64 1 : self.layout().addLayout(top_layout)
65 :
66 1 : self.layout().addWidget(self.canvas, stretch=1)
67 :
68 1 : def clear(self) -> None:
69 1 : self.canvas.ax.clear()
70 1 : self.canvas.draw_idle()
71 1 : self.clear_annotations()
72 1 : self.disconnect_click()
73 :
74 1 : def clear_annotations(self) -> None:
75 1 : for ann in self._annotations.values():
76 0 : with contextlib.suppress(NotImplementedError):
77 0 : ann.remove() # artist may already be gone if ax.clear() was called
78 1 : self._annotations.clear()
79 1 : self.canvas.draw_idle()
80 :
81 1 : def disconnect_click(self) -> None:
82 1 : for cid in self._click_cids:
83 0 : self.canvas.mpl_disconnect(cid)
84 1 : self._click_cids = []
85 :
86 1 : def connect_click(self, on_click: Callable[[mpl.backend_bases.MouseEvent], None]) -> None:
87 : """Connect a click handler that fires on mouse release, but only for genuine clicks."""
88 : # Maximum distance (in pixels) the mouse may travel between press and release
89 : # for the event to still count as a click rather than a drag (e.g. a zoom rectangle or a pan).
90 1 : click_move_threshold = 5
91 :
92 1 : press_positions: dict[Any, tuple[float, float]] = {}
93 :
94 1 : def on_press(event: mpl.backend_bases.MouseEvent) -> None:
95 0 : if event.x is None or event.y is None:
96 0 : return
97 0 : press_positions[event.button] = (event.x, event.y)
98 :
99 1 : def on_release(event: mpl.backend_bases.MouseEvent) -> None:
100 0 : start = press_positions.pop(event.button, None)
101 0 : if start is None or event.x is None or event.y is None:
102 0 : return
103 0 : if math.hypot(event.x - start[0], event.y - start[1]) > click_move_threshold:
104 0 : return # this was a drag (e.g. a zoom rectangle), not a click
105 0 : on_click(event)
106 :
107 1 : def on_figure_leave(_event: mpl.backend_bases.LocationEvent) -> None:
108 : # Forget pending presses whose release we will not see, a stale entry would otherwise
109 : # be used as the starting point of some later, unrelated release.
110 0 : press_positions.clear()
111 :
112 1 : self.disconnect_click()
113 1 : self._click_cids = [
114 : self.canvas.mpl_connect("button_press_event", on_press),
115 : self.canvas.mpl_connect("button_release_event", on_release),
116 : self.canvas.mpl_connect("figure_leave_event", on_figure_leave),
117 : ]
118 :
119 :
120 1 : class PlotEnergies(PlotWidget):
121 : """Plotwidget for plotting energy levels."""
122 :
123 1 : parameters: Parameters[Any] | None = None
124 1 : results: Results | None = None
125 1 : _annotations: dict[int, mpl.text.Annotation]
126 :
127 1 : def __init__(self, parent: SimulationPage) -> None:
128 1 : super().__init__(parent)
129 :
130 1 : self.fit_idx = 0
131 1 : self.fit_type = ""
132 1 : self.fit_data_highlight: mpl.collections.PathCollection | None = None
133 1 : self.fit_curve: Sequence[mpl.lines.Line2D] | None = None
134 :
135 1 : def setupWidget(self) -> None:
136 1 : super().setupWidget()
137 :
138 1 : window_color = theme_manager.get_palette().color(QPalette.ColorRole.Window).name()
139 :
140 1 : self.canvas.fig.set_facecolor(window_color)
141 1 : self.canvas.fig.set_layout_engine(
142 : "constrained",
143 : w_pad=0.2,
144 : h_pad=0.2,
145 : wspace=0.0,
146 : hspace=0.0,
147 : )
148 1 : mappable = ScalarMappable(cmap=alphamagma, norm=Normalize(vmin=0, vmax=1))
149 1 : cbar = self.canvas.fig.colorbar(mappable, ax=self.canvas.ax, label="Overlap with state of interest", aspect=60)
150 1 : cbar.ax.set_zorder(0)
151 1 : self.canvas.ax.set_zorder(1)
152 :
153 1 : def plot(self, parameters: Parameters[Any], results: Results) -> None:
154 1 : self.clear()
155 :
156 1 : show_status_tip(self, "Plotting energy curves...")
157 1 : ax = self.canvas.ax
158 1 : ax.set_xmargin(0)
159 :
160 : # store data to allow fitting later on
161 1 : self.parameters = parameters
162 1 : self.results = results
163 :
164 1 : x_values = parameters.get_x_values()
165 1 : energies = results.energies
166 1 : x_repeated = np.repeat(x_values, [len(es) for es in energies])
167 :
168 1 : if len({len(es) for es in energies}) <= 1: # check if homogeneous shape
169 1 : segments = [np.column_stack([x_values, es]) for es in np.transpose(energies)]
170 1 : ax.add_collection(LineCollection(segments, colors="0.75", linewidths=0.25, zorder=-10))
171 1 : ax.autoscale_view() # add_collection does not rescale the axes on its own
172 : else: # inhomogeneous shape
173 0 : ax.plot(x_repeated, np.hstack(energies), c="0.75", ls="None", marker=".", zorder=-10)
174 :
175 1 : show_status_tip(self, "Plotting overlaps...")
176 :
177 : # Flatten the arrays for scatter plot
178 : # (dont use numpy.flatten, etc. to also handle inhomogeneous shapes)
179 1 : energies_flattened = np.hstack(energies)
180 1 : overlaps_flattened = np.hstack(results.ket_overlaps)
181 :
182 1 : min_overlap = 1e-4
183 1 : inds: NDArray[Any] = np.argwhere(overlaps_flattened > min_overlap).ravel()
184 1 : inds = inds[np.argsort(overlaps_flattened[inds])]
185 :
186 1 : if len(inds) > 0:
187 1 : ax.scatter(
188 : x_repeated[inds],
189 : energies_flattened[inds],
190 : c=overlaps_flattened[inds],
191 : s=15,
192 : vmin=0,
193 : vmax=1,
194 : cmap=alphamagma,
195 : )
196 :
197 1 : ax.set_xlabel(parameters.get_x_label())
198 1 : ax.set_ylabel("Energy (GHz)")
199 :
200 1 : def setup_annotations(self, parameters: Parameters[Any], results: Results) -> None:
201 : """Connect click-based state annotation to the energy plot."""
202 1 : energies = results.energies
203 1 : overlaps = results.ket_overlaps
204 1 : x_values = parameters.get_x_values()
205 :
206 1 : self._point_index_map: list[tuple[int, int]] = []
207 1 : all_x: list[float] = []
208 1 : all_y: list[float] = []
209 1 : all_overlaps: list[float] = []
210 1 : for idx in range(len(energies)):
211 1 : x = x_values[idx]
212 1 : for idstate, (energy, overlap) in enumerate(zip(energies[idx], overlaps[idx], strict=True)):
213 1 : all_x.append(x)
214 1 : all_y.append(float(energy))
215 1 : all_overlaps.append(float(overlap))
216 1 : self._point_index_map.append((idx, idstate))
217 1 : pts_data = np.column_stack([all_x, all_y]) if all_x else np.empty((0, 2))
218 1 : pts_overlaps = np.array(all_overlaps)
219 :
220 1 : def on_click(event: mpl.backend_bases.MouseEvent) -> None:
221 0 : if event.inaxes is not self.canvas.ax or event.button not in [1, 3] or len(pts_data) == 0:
222 0 : return
223 0 : if event.button == 3: # right click clears annotations
224 0 : self.clear_annotations()
225 0 : return
226 :
227 0 : pts_pos = self.canvas.ax.transData.transform(pts_data)
228 0 : click_pos = np.array([event.x, event.y])
229 0 : dists = np.hypot(pts_pos[:, 0] - click_pos[0], pts_pos[:, 1] - click_pos[1])
230 0 : candidates = np.flatnonzero(dists <= 10) # threshold in pixels
231 0 : if len(candidates) == 0:
232 0 : self.clear_annotations()
233 0 : return
234 0 : selected = int(candidates[np.argmax(pts_overlaps[candidates])])
235 0 : if selected in self._annotations:
236 0 : self._annotations[selected].remove()
237 0 : del self._annotations[selected]
238 0 : self.canvas.draw_idle()
239 0 : return
240 0 : idstep, idstate = self._point_index_map[selected]
241 0 : state: StateBase[Any] = results.systems[idstep].get_eigenbasis().get_state(idstate)
242 0 : label = state.get_label().replace(" + ", "\n + ").replace(" - ", "\n - ")
243 0 : xlim = self.canvas.ax.get_xlim()
244 0 : ylim = self.canvas.ax.get_ylim()
245 0 : x_frac = (pts_data[selected, 0] - xlim[0]) / (xlim[1] - xlim[0])
246 0 : y_frac = (pts_data[selected, 1] - ylim[0]) / (ylim[1] - ylim[0])
247 0 : x_offset = -100 if isinstance(state, StateAtom) else -250
248 0 : x_offset = x_offset if x_frac > 0.5 else 0
249 0 : y_offset = 15 + 10 * label.count("\n")
250 0 : y_offset = -y_offset if y_frac > 0.5 else y_offset
251 0 : ann = self.canvas.ax.annotate(
252 : label,
253 : xy=(pts_data[selected, 0], pts_data[selected, 1]),
254 : xytext=(x_offset, y_offset),
255 : textcoords="offset points",
256 : va="center",
257 : bbox={"boxstyle": "round,pad=0.5", "fc": "white", "alpha": 0.9, "ec": "gray"},
258 : arrowprops={"arrowstyle": "->", "connectionstyle": "arc3", "color": "gray"},
259 : clip_on=False,
260 : )
261 0 : ann.set_in_layout(False)
262 0 : self._annotations[selected] = ann
263 0 : self.canvas.draw_idle()
264 :
265 1 : self.connect_click(on_click)
266 1 : self.navigation_toolbar._home_callbacks = [self.clear_annotations]
267 :
268 1 : def fit(self, fit_type: str = "c6") -> None: # noqa: PLR0912, PLR0915, C901
269 : """Fits a potential curve and displays the fit values.
270 :
271 : Args:
272 : fit_type: Type of fit to perform. Options are:
273 : c6: E = E0 + C6 * r^6
274 : c3: E = E0 + C3 * r^3
275 : c3+c6: E = E0 + C3 * r^3 + C6 * r^6
276 :
277 : Iterative calls will iterate through the potential curves
278 :
279 : """
280 1 : if self.parameters is None or self.results is None:
281 0 : logger.warning("No data to fit.")
282 0 : return
283 :
284 1 : energies = self.results.energies
285 1 : x_values = np.array(self.parameters.get_x_values())
286 1 : overlaps_list = self.results.ket_overlaps
287 :
288 : fit_func: Callable[Concatenate[NDArray[Any], ...], NDArray[Any]]
289 1 : if fit_type == "c6":
290 1 : fit_func = fit_c6
291 1 : fitlabel = "E0 = {0:.3f} GHz\nC6 = {1:.3f} GHz*µm^6"
292 1 : elif fit_type == "c3":
293 1 : fit_func = fit_c3
294 1 : fitlabel = "E0 = {0:.3f} GHz\nC3 = {1:.3f} GHz*µm^3"
295 1 : elif fit_type == "c3+c6":
296 1 : fit_func = fit_c3_c6
297 1 : fitlabel = "E0 = {0:.3f} GHz\nC3 = {1:.3f} GHz*µm^3\nC6 = {2:.3f} GHz*µm^6"
298 : else:
299 0 : raise ValueError(f"Unknown fit type: {fit_type}")
300 :
301 : # increase the selected potential curve by one if we use the same fit type
302 1 : if self.fit_type == fit_type:
303 1 : self.fit_idx = (self.fit_idx + 1) % len(energies)
304 : else:
305 1 : self.fit_idx = 1
306 1 : self.fit_type = fit_type
307 :
308 : # We want to follow the potential curves. The ordering of energies is just by value, so we
309 : # need to follow the curve somehow. We go right to left, start at the nth largest value, keep our
310 : # index as long as the difference in overlap is less than a factor 2 or less than 5% total difference.
311 : # Otherwise, we search until we find an overlap that is less than a factor 2 different.
312 : # This is of course a simple heuristic, a more sophisticated approach would do some global optimization
313 : # of the curves. This approach is simple, fast and robust, but curves may e.g. merge.
314 : # This does not at all take into account the line shapes of the curves. There is also no trade-off
315 : # between overlap being close and not doing jumps.
316 1 : idxs = [np.argpartition(overlaps_list[0], -self.fit_idx)[-self.fit_idx]]
317 1 : last_overlap = overlaps_list[0][idxs[-1]]
318 1 : for overlaps in overlaps_list[1:]:
319 1 : idx = idxs[-1]
320 1 : overlap = overlaps[idx]
321 1 : if 0.5 * last_overlap < overlap < 2 * last_overlap or abs(overlap - last_overlap) < 0.05:
322 : # we keep the current index
323 1 : idxs.append(idx)
324 1 : last_overlap = overlap
325 : else:
326 : # we search until we find an overlap that is less than a factor 2 different
327 1 : possible_options = np.argwhere(
328 : np.logical_and(overlaps > 0.5 * last_overlap, overlaps < 2 * last_overlap)
329 : ).flatten()
330 1 : if len(possible_options) == 0:
331 : # there is no state in that range - our best bet is to keep the current index
332 1 : idxs.append(idx)
333 1 : last_overlap = overlap
334 : else:
335 : # we select the closest possible option
336 1 : best_option = np.argmin(np.abs(possible_options - idx))
337 1 : idxs.append(possible_options[best_option])
338 1 : last_overlap = overlaps[idxs[-1]]
339 :
340 : # this could be a call to np.take_along_axis if the sizes match, but the handling of inhomogeneous shapes
341 : # in the plot() function makes me worry they won't, so I go for a slower python for loop...
342 1 : energies_fit = np.array([energy[idx] for energy, idx in zip(energies, idxs, strict=True)])
343 :
344 : # stop highlighting the previous fit
345 1 : if self.fit_data_highlight is not None:
346 1 : self.fit_data_highlight.remove()
347 1 : if self.fit_curve is not None:
348 1 : for curve in self.fit_curve:
349 1 : curve.remove()
350 :
351 1 : self.fit_data_highlight = self.canvas.ax.scatter(x_values, energies_fit, c="green", s=5)
352 :
353 1 : try:
354 1 : fit_params = curve_fit(fit_func, x_values, energies_fit)[0]
355 0 : except (RuntimeError, TypeError):
356 0 : logger.warning("Curve fit failed.")
357 : else:
358 1 : self.fit_curve = self.canvas.ax.plot(
359 : x_values,
360 : fit_func(x_values, *fit_params),
361 : c="green",
362 : linestyle="dashed",
363 : lw=2,
364 : label=fitlabel.format(*fit_params),
365 : )
366 1 : self.canvas.ax.legend()
367 :
368 1 : self.canvas.draw_idle()
369 :
370 1 : def clear(self) -> None:
371 1 : super().clear()
372 1 : self.reset_fit()
373 :
374 1 : def reset_fit(self) -> None:
375 : """Clear fit output and reset fit index."""
376 : # restart at first potential curve
377 1 : self.fit_idx = 0
378 : # and also remove any previous highlighting/fit display
379 1 : self.fit_data_highlight = None
380 1 : self.fit_curve = None
381 :
382 :
383 1 : class PlotLifetimes(PlotWidget):
384 : """Plotwidget for plotting lifetime/transition rate bar charts."""
385 :
386 1 : _annotations: dict[tuple[str, int], mpl.text.Annotation]
387 :
388 1 : def setupWidget(self) -> None:
389 1 : super().setupWidget()
390 :
391 1 : window_color = theme_manager.get_palette().color(QPalette.ColorRole.Window).name()
392 1 : self.canvas.fig.set_facecolor(window_color)
393 1 : self.canvas.fig.set_layout_engine(
394 : "constrained",
395 : w_pad=0.2,
396 : h_pad=0.2,
397 : wspace=0.0,
398 : hspace=0.0,
399 : )
400 :
401 1 : def plot(self, parameters: ParametersLifetimes, results: ResultsLifetimes) -> None:
402 1 : self.clear()
403 1 : ax = self.canvas.ax
404 :
405 1 : show_status_tip(self, "Preparing transition rates...")
406 1 : labels = ["Spontaneous Decay", "Black Body Radiation"]
407 1 : n_list = np.arange(0, np.max([s.n for s in results.kets_bbr + results.kets_sp] + [0]) + 1)
408 1 : sorted_rates: dict[str, dict[int, list[tuple[KetData, float]]]] = {}
409 1 : for key, kets, rates in [
410 : (labels[0], results.kets_sp, results.transition_rates_sp),
411 : (labels[1], results.kets_bbr, results.transition_rates_bbr),
412 : ]:
413 1 : sorted_rates[key] = {n: [] for n in n_list}
414 1 : for i, s in enumerate(kets):
415 1 : sorted_rates[key][s.n].append((s, rates[i]))
416 1 : self.sorted_rates = sorted_rates
417 1 : rates_summed = {key: [sum(r for _, r in sorted_rates[key][n]) for n in n_list] for key in sorted_rates}
418 :
419 1 : show_status_tip(self, "Plotting transition rates...")
420 1 : self.artists: list[mpl.container.BarContainer] = []
421 1 : for label, color in zip(labels, ["blue", "red"], strict=True):
422 1 : bar = ax.bar(n_list, rates_summed[label], label=label, color=color, alpha=0.8)
423 1 : self.artists.append(bar)
424 1 : ax.legend()
425 :
426 1 : ax.set_xlabel("Principal Quantum Number $n$")
427 1 : ax.set_ylabel(r"Transition Rates (1 / ms)")
428 :
429 1 : def setup_annotations(self, parameters: ParametersLifetimes, results: ResultsLifetimes) -> None: # noqa: C901
430 : """Add click-based annotations to the plot."""
431 0 : show_status_tip(self, "Adding transition rate annotations...")
432 :
433 0 : self._bar_data: list[tuple[str, int, mpl.patches.Rectangle]] = []
434 0 : for container in reversed(self.artists):
435 0 : label = container.get_label()
436 0 : if label is None:
437 0 : continue
438 0 : for rect in container.patches:
439 0 : n = round(rect.get_x() + rect.get_width() / 2)
440 0 : self._bar_data.append((label, n, rect))
441 :
442 0 : def on_click(event: mpl.backend_bases.MouseEvent) -> None:
443 0 : if event.inaxes is not self.canvas.ax or event.button not in [1, 3]:
444 0 : return
445 0 : if event.button == 3: # right click clears annotations
446 0 : self.clear_annotations()
447 0 : return
448 :
449 0 : if event.xdata is None or event.ydata is None:
450 0 : return
451 0 : click_coords = np.array([event.xdata, event.ydata])
452 :
453 0 : for _label, _n, rect in self._bar_data:
454 0 : x, y = rect.get_x(), rect.get_y()
455 0 : if x <= click_coords[0] <= x + rect.get_width() and y <= click_coords[1] <= y + rect.get_height():
456 0 : label, n = _label, _n
457 0 : break
458 : else: # no break -> no bar found
459 0 : self.clear_annotations()
460 0 : return
461 :
462 : # if we click the same bar again, remove the annotation
463 0 : if (label, n) in self._annotations:
464 0 : self._annotations[(label, n)].remove()
465 0 : del self._annotations[(label, n)]
466 0 : self.canvas.draw_idle()
467 0 : return
468 :
469 0 : state_text = "\n".join(f" - {s}: {r:.5f}/ms" for (s, r) in self.sorted_rates[label][n])
470 0 : text = f"{label} to n={n}:\n{state_text}"
471 0 : xlim = self.canvas.ax.get_xlim()
472 0 : ylim = self.canvas.ax.get_ylim()
473 0 : bar_cx = rect.get_x() + rect.get_width() / 2
474 0 : bar_top = rect.get_y() + rect.get_height()
475 0 : x_frac = (bar_cx - xlim[0]) / (xlim[1] - xlim[0])
476 0 : y_frac = (bar_top - ylim[0]) / (ylim[1] - ylim[0])
477 0 : x_offset = -200 if x_frac > 0.5 else 25
478 0 : y_offset = -50 if y_frac > 0.5 else 50
479 0 : ann = self.canvas.ax.annotate(
480 : text,
481 : xy=(bar_cx, bar_top),
482 : xytext=(x_offset, y_offset),
483 : textcoords="offset points",
484 : bbox={"boxstyle": "round,pad=0.5", "fc": "white", "alpha": 0.9, "ec": "gray"},
485 : arrowprops={"arrowstyle": "->", "connectionstyle": "arc3", "color": "gray"},
486 : clip_on=False,
487 : )
488 0 : ann.set_in_layout(False)
489 0 : self._annotations[(label, n)] = ann
490 0 : self.canvas.draw_idle()
491 :
492 0 : self.connect_click(on_click)
493 0 : self.navigation_toolbar._home_callbacks = [self.clear_annotations]
494 :
495 :
496 1 : def fit_c3(x: NDArray[Any], /, e0: float, c3: float) -> NDArray[Any]:
497 1 : return e0 + c3 / x**3
498 :
499 :
500 1 : def fit_c6(x: NDArray[Any], /, e0: float, c6: float) -> NDArray[Any]:
501 1 : return e0 + c6 / x**6
502 :
503 :
504 1 : def fit_c3_c6(x: NDArray[Any], /, e0: float, c3: float, c6: float) -> NDArray[Any]:
505 1 : return e0 + c3 / x**3 + c6 / x**6
|
