"""Aurora plot dialog for detector saturation inspection."""
from __future__ import annotations
import os
from typing import Optional
import numpy as np
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.colors import LogNorm
from matplotlib.figure import Figure
from PyQt6.QtCore import QObject, QThread, pyqtSignal
from PyQt6.QtWidgets import (
QCheckBox,
QComboBox,
QDialog,
QHBoxLayout,
QLabel,
QMessageBox,
QPushButton,
QVBoxLayout,
)
from coseda.peakfinding.aurora import (
AURORA_OUTPUT_SUBFOLDER,
AURORA_WINDOW_SIZE,
compute_aurora_histogram_for_files,
)
WINDOW_SIZE_OPTIONS = (1, 3, 5, 7, 9)
[docs]
class AuroraPlotWorker(QObject):
"""Worker object for computing aurora histograms outside the UI thread."""
finished = pyqtSignal(object)
failed = pyqtSignal(str)
progress = pyqtSignal(str, float, int)
def __init__(self, file_specs: list[dict], force_recompute: bool, window_size: int):
super().__init__()
self.file_specs = file_specs
self.force_recompute = force_recompute
self.window_size = window_size
[docs]
def run(self):
try:
result = compute_aurora_histogram_for_files(
self.file_specs,
window_size=self.window_size,
force_recompute=self.force_recompute,
progress_callback=self.progress.emit,
)
except Exception as exc:
self.failed.emit(str(exc))
return
self.finished.emit(result)
[docs]
class AuroraPlotWindow(QDialog):
"""Small non-modal window that renders peak max intensity versus detector radius."""
def __init__(
self,
parent=None,
hdf5_path: Optional[str] = None,
ini_path: Optional[str] = None,
workspace_files: Optional[list[dict]] = None,
):
super().__init__(parent)
self.setWindowTitle("Aurora Plot")
self.current_file = {"hdf5_path": hdf5_path, "ini_path": ini_path}
self.workspace_files = list(workspace_files or [])
self._worker_thread = None
self._worker = None
self.resize(900, 620)
self._init_ui()
self.start_plot(force_recompute=False)
@property
def hdf5_path(self):
return self.current_file.get("hdf5_path")
def _init_ui(self):
layout = QVBoxLayout(self)
controls = QHBoxLayout()
self.all_files_checkbox = QCheckBox("All workspace files")
self.all_files_checkbox.setEnabled(len(self.workspace_files) > 1)
self.all_files_checkbox.stateChanged.connect(lambda _state: self.start_plot(force_recompute=False))
controls.addWidget(self.all_files_checkbox)
controls.addWidget(QLabel("Window"))
self.window_size_combo = QComboBox()
for size in WINDOW_SIZE_OPTIONS:
self.window_size_combo.addItem(f"{size}x{size}", size)
default_index = WINDOW_SIZE_OPTIONS.index(AURORA_WINDOW_SIZE)
self.window_size_combo.setCurrentIndex(default_index)
self.window_size_combo.currentIndexChanged.connect(lambda _index: self.start_plot(force_recompute=False))
controls.addWidget(self.window_size_combo)
self.recompute_button = QPushButton("Recompute")
self.recompute_button.clicked.connect(lambda: self.start_plot(force_recompute=True))
controls.addWidget(self.recompute_button)
self.status_label = QLabel("")
controls.addWidget(self.status_label, stretch=1)
close_button = QPushButton("Close")
close_button.clicked.connect(self.close)
controls.addWidget(close_button)
layout.addLayout(controls)
self.figure = Figure(figsize=(8, 5), tight_layout=True)
self.canvas = FigureCanvas(self.figure)
layout.addWidget(self.canvas)
def _selected_file_specs(self):
if self.all_files_checkbox.isChecked() and self.workspace_files:
return self.workspace_files
return [self.current_file]
def _window_size(self) -> int:
return int(self.window_size_combo.currentData())
def _plot_png_path(self, file_count: int) -> str:
window_size = self._window_size()
if file_count == 1:
h5_path = self._selected_file_specs()[0]["hdf5_path"]
folder = os.path.join(os.path.dirname(os.path.abspath(h5_path)), AURORA_OUTPUT_SUBFOLDER)
stem = os.path.splitext(os.path.basename(h5_path))[0]
return os.path.join(folder, f"{stem}_aurora_peakmax_{window_size}x{window_size}.png")
h5_path = self.current_file.get("hdf5_path") or self._selected_file_specs()[0]["hdf5_path"]
folder = os.path.join(os.path.dirname(os.path.abspath(h5_path)), AURORA_OUTPUT_SUBFOLDER)
return os.path.join(folder, f"workspace_aurora_peakmax_{window_size}x{window_size}.png")
[docs]
def start_plot(self, force_recompute: bool):
file_specs = self._selected_file_specs()
if not any(spec.get("hdf5_path") and os.path.exists(spec.get("hdf5_path")) for spec in file_specs):
QMessageBox.warning(self, "No HDF5 file", "Please open an HDF5 file first.")
return
if self._worker_thread is not None:
return
self.recompute_button.setEnabled(False)
self.all_files_checkbox.setEnabled(False)
self.window_size_combo.setEnabled(False)
self.status_label.setText("Preparing...")
self.figure.clear()
ax = self.figure.add_subplot(111)
ax.text(0.5, 0.5, "Preparing aurora plot...", ha="center", va="center", transform=ax.transAxes)
ax.set_axis_off()
self.canvas.draw_idle()
thread = QThread(self)
worker = AuroraPlotWorker(file_specs, force_recompute, self._window_size())
worker.moveToThread(thread)
thread.started.connect(worker.run)
worker.progress.connect(self._on_progress)
worker.finished.connect(self._on_result)
worker.failed.connect(self._on_failed)
worker.finished.connect(thread.quit)
worker.failed.connect(thread.quit)
worker.finished.connect(worker.deleteLater)
worker.failed.connect(worker.deleteLater)
thread.finished.connect(thread.deleteLater)
thread.finished.connect(self._clear_worker_refs)
self._worker_thread = thread
self._worker = worker
thread.start()
def _set_controls_enabled(self):
self.recompute_button.setEnabled(True)
self.all_files_checkbox.setEnabled(len(self.workspace_files) > 1)
self.window_size_combo.setEnabled(True)
def _on_progress(self, phase: str, current: float, total: int):
if total > 0:
pct = int(round(float(current) / total * 100))
self.status_label.setText(f"{phase}: {pct}%")
else:
self.status_label.setText(phase)
def _on_result(self, result: dict):
hist = np.asarray(result["hist"], dtype=np.float64)
masked = np.ma.masked_less_equal(hist.T, 0)
self.figure.clear()
ax = self.figure.add_subplot(111)
mesh = ax.pcolormesh(
result["x_edges"],
result["y_edges"],
masked,
cmap="viridis",
norm=LogNorm(vmin=1, vmax=max(float(np.max(hist)), 1.0)),
shading="auto",
)
ax.set_xlabel(result["x_label"])
ax.set_ylabel(result["y_label"])
bit_depth = result.get("bit_depth_label", "unknown")
if result.get("file_count", 1) == 1:
ax.set_title(f"{os.path.basename(self._selected_file_specs()[0]['hdf5_path'])} | image {bit_depth}")
else:
ax.set_title(f"Workspace ({result['file_count']} files) | image {bit_depth}")
cbar = self.figure.colorbar(mesh, ax=ax)
cbar.set_label("Peak count")
ax.set_xlim(float(result["x_edges"][0]), float(result["x_edges"][-1]))
ax.set_ylim(float(result["y_edges"][0]), float(result["y_edges"][-1]))
self.canvas.draw_idle()
total_peaks = int(result.get("total_peaks", 0))
file_count = int(result.get("file_count", 1))
cache_hits = int(result.get("cache_hits", 0))
png_path = self._plot_png_path(file_count)
try:
os.makedirs(os.path.dirname(png_path), exist_ok=True)
self.figure.savefig(png_path, dpi=200)
save_message = f", saved {os.path.basename(png_path)}"
except Exception as exc:
save_message = ", PNG save failed"
QMessageBox.warning(self, "Aurora Plot", f"Could not save plot PNG:\n{exc}")
self.status_label.setText(
f"{total_peaks:,} peaks, {file_count} file(s), image {bit_depth}, "
f"{cache_hits} cache hit(s), pixel radius{save_message}"
)
self._set_controls_enabled()
def _on_failed(self, message: str):
self.figure.clear()
ax = self.figure.add_subplot(111)
ax.text(0.5, 0.5, "Aurora plot unavailable", ha="center", va="center", transform=ax.transAxes)
ax.set_axis_off()
self.canvas.draw_idle()
self.status_label.setText("Failed")
self._set_controls_enabled()
QMessageBox.warning(self, "Aurora Plot", message)
def _clear_worker_refs(self):
self._worker_thread = None
self._worker = None
[docs]
def closeEvent(self, event):
if self._worker_thread is not None:
QMessageBox.information(
self,
"Aurora Plot",
"The aurora plot is still being calculated. Please wait for it to finish.",
)
event.ignore()
return
super().closeEvent(event)