Source code for cosedaUI.unitcell_window

"""Dialog for plotting unit cell histograms from CrystFEL stream files."""

from __future__ import annotations

import os
from typing import Optional

from PyQt6.QtCore import Qt, QTimer
from PyQt6.QtWidgets import (
    QCheckBox,
    QComboBox,
    QDialog,
    QHBoxLayout,
    QLabel,
    QMessageBox,
    QPushButton,
    QSpinBox,
    QVBoxLayout,
)
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure

try:
    from CFstreamparser import parse_stream_file, unit_cell_histograms
except ImportError:  # pragma: no cover - handled at runtime
    parse_stream_file = None
    unit_cell_histograms = None


[docs] class UnitCellHistogramWindow(QDialog): """Plot unit cell parameter histograms for the current CrystFEL stream.""" def __init__( self, parent=None, stream_path: Optional[str] = None, indexing_manager=None, ): super().__init__(parent) self.indexing_manager = indexing_manager # Prefer explicit stream_path, otherwise fall back to manager self.stream_path = stream_path or getattr(indexing_manager, "stream_path", None) self._auto_refresh_enabled = False self._refresh_interval_sec = 60 self._countdown_remaining = 0 self._status_base = "" self._refresh_timer = QTimer(self) self._refresh_timer.timeout.connect(self._on_plot) self._countdown_timer = QTimer(self) self._countdown_timer.timeout.connect(self._tick_countdown) self._countdown_timer.setInterval(1000) # 1s ticks self.setWindowTitle("Unit Cell Histograms") self.resize(960, 640) if parse_stream_file is None or unit_cell_histograms is None: QMessageBox.warning( self, "CFstreamparser missing", "CFstreamparser is required to parse stream files and plot unit cell histograms.\n" "Install it and try again.", ) self.close() return self._build_ui() self._auto_plot_if_available() # --- UI setup ----------------------------------------------------- def _build_ui(self): layout = QVBoxLayout(self) # Options row opts_layout = QHBoxLayout() bins_label = QLabel("Bins:") self.bins_spin = QSpinBox() self.bins_spin.setRange(5, 500) self.bins_spin.setValue(60) self.per_image_cb = QCheckBox("One solution per image") self.per_image_cb.setChecked(True) self.multi_solution_combo = QComboBox() self.multi_solution_combo.addItem("Best (most reflections)", "best") self.multi_solution_combo.addItem("Highest resolution", "highest_res") self.multi_solution_combo.addItem("Lowest resolution", "lowest_res") opts_layout.addWidget(bins_label) opts_layout.addWidget(self.bins_spin) opts_layout.addSpacing(12) opts_layout.addWidget(self.per_image_cb) opts_layout.addSpacing(12) opts_layout.addWidget(QLabel("Multi-solution choice:")) opts_layout.addWidget(self.multi_solution_combo) opts_layout.addStretch() layout.addLayout(opts_layout) # Plot + status controls controls = QHBoxLayout() self.plot_btn = QPushButton("Plot histograms") self.plot_btn.clicked.connect(self._on_plot) controls.addWidget(self.plot_btn) self.auto_refresh_cb = QCheckBox("Auto refresh") self.auto_refresh_cb.stateChanged.connect( lambda state: self.start_auto_refresh() if state == Qt.CheckState.Checked else self.stop_auto_refresh() ) controls.addWidget(self.auto_refresh_cb) controls.addStretch() self.status_label = QLabel("") self.status_label.setStyleSheet("color: #555;") controls.addWidget(self.status_label) layout.addLayout(controls) # Matplotlib canvas self.figure = Figure(figsize=(9, 5)) self.canvas = FigureCanvas(self.figure) layout.addWidget(self.canvas) # --- Actions ------------------------------------------------------ def _auto_plot_if_available(self): """Auto-render when we already have a valid stream path.""" if self.stream_path and os.path.exists(self.stream_path): self._on_plot() else: # Keep controls enabled so auto-refresh can poll for the file self.plot_btn.setEnabled(bool(self.stream_path)) if self.auto_refresh_cb is not None: self.auto_refresh_cb.setEnabled(True) self._set_status("Waiting for stream file...") def _load_parsed_stream(self, path): """Return a parsed stream, reusing the indexing manager when possible.""" # Reuse the main window's parser to avoid reparsing if ( self.indexing_manager and os.path.exists(path) and os.path.abspath(getattr(self.indexing_manager, "stream_path", "") or "") == os.path.abspath(path) ): if not self.indexing_manager.ensure_loaded(): err = getattr(self.indexing_manager, "get_last_error", lambda: None)() detail = f": {err}" if err else "" QMessageBox.critical( self, "Stream parse failed", f"Failed to parse {path}{detail}", ) return None try: return self.indexing_manager.get_parsed_stream() except Exception: # Fall through and try direct parse pass try: return parse_stream_file(path) except Exception as exc: QMessageBox.critical( self, "Stream parse failed", f"Failed to parse {path}:\n{exc}", ) return None
[docs] def set_stream_path(self, path: Optional[str], auto_plot: bool = False): """Update the linked stream path (called when selection changes externally).""" self.stream_path = path has_file = bool(path and os.path.exists(path)) self.plot_btn.setEnabled(has_file) self.auto_refresh_cb.setEnabled(True) if not path: self._set_status("No stream file for selected run.") self._clear_canvas() return if not has_file: self._set_status("Waiting for stream file...") self._clear_canvas() # Keep auto-refresh running to pick it up when file appears if auto_plot and self._auto_refresh_enabled: self._reset_countdown() return if auto_plot or self._auto_refresh_enabled: self._on_plot()
def _on_plot(self): if self._auto_refresh_enabled: self._countdown_remaining = self._refresh_interval_sec path = (self.stream_path or "").strip() if not path: self._set_status("No stream file for selected run.") self._clear_canvas() return if not os.path.exists(path): self._set_status("Waiting for stream file...") self._clear_canvas() return parsed = self._load_parsed_stream(path) if parsed is None: return bins = int(self.bins_spin.value()) per_image = self.per_image_cb.isChecked() multi_solution = str(self.multi_solution_combo.currentData()) length_scale = 10.0 # fixed: nm -> Å hists = unit_cell_histograms( parsed, bins=bins, multi_solution=multi_solution, per_image=per_image, length_scale=length_scale, ) self._draw_histograms(hists) self._set_status(f"Plotted from {os.path.basename(path)}") self._reset_countdown() # --- Plot helpers ------------------------------------------------- def _clear_canvas(self): self.figure.clear() self.canvas.draw_idle() def _draw_histograms(self, hists): params = ["a", "b", "c", "alpha", "beta", "gamma"] units = {"a": "Å", "b": "Å", "c": "Å", "alpha": "deg", "beta": "deg", "gamma": "deg"} # Clear previous axes to avoid stacking plots on re-render self.figure.clear() axes = self.figure.subplots(2, 3, squeeze=False).ravel() for ax, param in zip(axes, params): ax.clear() counts, edges = hists.get(param, ([], [])) label = f"{param} ({units[param]})" if counts and len(edges) >= 2: try: width = edges[1] - edges[0] if not width or width <= 0: width = 0.8 except Exception: width = 0.8 ax.bar( edges[:-1], counts, width=width, align="edge", color="tab:blue", edgecolor="tab:blue", alpha=0.55, ) ax.set_ylabel("Count") else: ax.text(0.5, 0.5, "No data", ha="center", va="center", color="#666", transform=ax.transAxes) ax.set_yticks([]) ax.set_xlabel(label) # Ensure angle axes span at least 20 degrees for readability if param in ("alpha", "beta", "gamma") and len(edges) >= 2: xmin, xmax = edges[0], edges[-1] if xmax - xmin < 20: mid = 0.5 * (xmin + xmax) ax.set_xlim(mid - 10, mid + 10) ax.grid(True, linestyle=":", alpha=0.4) self.figure.tight_layout() self.canvas.draw_idle() # --- Auto refresh helpers ----------------------------------------
[docs] def start_auto_refresh(self, interval_sec: int = 60): """Begin auto-refreshing plots every interval_sec with countdown.""" self._auto_refresh_enabled = True self._refresh_interval_sec = max(1, int(interval_sec)) # Restart timers cleanly self._refresh_timer.stop() self._countdown_timer.stop() self._refresh_timer.setInterval(self._refresh_interval_sec * 1000) self._countdown_timer.setInterval(1000) self._refresh_timer.start() self._countdown_timer.start() # Keep checkbox in sync without re-entering if hasattr(self, "auto_refresh_cb"): self.auto_refresh_cb.blockSignals(True) self.auto_refresh_cb.setChecked(True) self.auto_refresh_cb.blockSignals(False) self._reset_countdown() # Trigger an immediate attempt even if file not present yet if self.stream_path: self._on_plot() else: self._set_status("Waiting for stream file...")
[docs] def stop_auto_refresh(self): self._auto_refresh_enabled = False self._refresh_timer.stop() self._countdown_timer.stop() self._countdown_remaining = None if hasattr(self, "auto_refresh_cb"): self.auto_refresh_cb.blockSignals(True) self.auto_refresh_cb.setChecked(False) self.auto_refresh_cb.blockSignals(False) self._render_status()
def _reset_countdown(self): if not self._auto_refresh_enabled: return self._countdown_remaining = self._refresh_interval_sec self._render_status() def _tick_countdown(self): if not self._auto_refresh_enabled: return if self._countdown_remaining > 0: self._countdown_remaining -= 1 self._render_status() def _set_status(self, text: str): self._status_base = text self._render_status() def _render_status(self): msg = self._status_base or "" if self._auto_refresh_enabled: remaining = self._countdown_remaining if self._countdown_remaining is not None else self._refresh_interval_sec msg = (msg + " | " if msg else "") + f"Refresh in {int(remaining):02d}s" self.status_label.setText(msg)
[docs] def closeEvent(self, event): try: self.stop_auto_refresh() except Exception: pass super().closeEvent(event)