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