"""Workspace-wide unit-cell histogram and Gaussian fit dialog."""
from __future__ import annotations
import os
from typing import Optional
import numpy as np
from PyQt6.QtCore import Qt
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
from matplotlib.widgets import SpanSelector
from coseda.unitcell_fit import (
ANGLE_PARAMS,
PARAMS,
discover_workspace_unit_cell_streams,
fit_gaussian_to_values,
update_workspace_refined_cell_parameter,
)
try:
from CFstreamparser import parse_stream_file, unit_cell_parameters
except ImportError: # pragma: no cover - handled at runtime
parse_stream_file = None
unit_cell_parameters = None
[docs]
class WorkspaceUnitCellFitWindow(QDialog):
"""Aggregate unit-cell values from all workspace files and fit selected peaks."""
def __init__(
self,
parent=None,
workspace_ini_paths: Optional[list[str]] = None,
workspace_path: Optional[str] = None,
):
super().__init__(parent)
self.workspace_ini_paths = list(workspace_ini_paths or [])
self.workspace_path = workspace_path
self.sources = []
self.params = {param: [] for param in PARAMS}
self.fit_results = {}
self.selected_param = None
self._selectors = []
self._axes_by_param = {}
self.setWindowTitle("Workspace Unit Cell Fit")
self.resize(1040, 720)
if parse_stream_file is None or unit_cell_parameters is None:
QMessageBox.warning(
self,
"CFstreamparser missing",
"CFstreamparser is required to parse stream files and fit unit-cell histograms.\n"
"Install it and try again.",
)
self.close()
return
self._build_ui()
self._load_and_plot()
[docs]
def set_workspace_ini_paths(self, paths: list[str]):
self.workspace_ini_paths = list(paths or [])
self.fit_results.clear()
self.selected_param = None
self._load_and_plot()
[docs]
def set_workspace_context(self, paths: list[str], workspace_path: Optional[str]):
self.workspace_path = workspace_path
self.set_workspace_ini_paths(paths)
def _build_ui(self):
layout = QVBoxLayout(self)
options = QHBoxLayout()
options.addWidget(QLabel("Bins:"))
self.bins_spin = QSpinBox()
self.bins_spin.setRange(5, 500)
self.bins_spin.setValue(60)
self.bins_spin.valueChanged.connect(self._draw_histograms)
options.addWidget(self.bins_spin)
self.per_image_cb = QCheckBox("One solution per image")
self.per_image_cb.setChecked(True)
self.per_image_cb.stateChanged.connect(lambda _state: self._load_and_plot())
options.addWidget(self.per_image_cb)
options.addWidget(QLabel("Multi-solution choice:"))
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")
self.multi_solution_combo.currentIndexChanged.connect(lambda _idx: self._load_and_plot())
options.addWidget(self.multi_solution_combo)
self.refresh_btn = QPushButton("Refresh")
self.refresh_btn.clicked.connect(self._load_and_plot)
options.addWidget(self.refresh_btn)
self.update_cell_btn = QPushButton("Update workspace refined cell")
self.update_cell_btn.setEnabled(False)
self.update_cell_btn.clicked.connect(self._update_cell_files)
options.addWidget(self.update_cell_btn)
options.addStretch()
layout.addLayout(options)
self.fit_label = QLabel("Drag across a peak in any histogram to fit a Gaussian.")
self.fit_label.setTextInteractionFlags(Qt.TextInteractionFlag.TextSelectableByMouse)
layout.addWidget(self.fit_label)
self.figure = Figure(figsize=(10, 6))
self.canvas = FigureCanvas(self.figure)
layout.addWidget(self.canvas)
self.status_label = QLabel("")
self.status_label.setStyleSheet("color: #555;")
self.status_label.setTextInteractionFlags(Qt.TextInteractionFlag.TextSelectableByMouse)
layout.addWidget(self.status_label)
def _set_busy(self, busy: bool):
for widget in (
self.bins_spin,
self.per_image_cb,
self.multi_solution_combo,
self.refresh_btn,
self.update_cell_btn,
):
widget.setEnabled(not busy)
if not busy:
self.update_cell_btn.setEnabled(bool(self.selected_param and self.fit_results.get(self.selected_param)))
def _load_and_plot(self):
self._set_busy(True)
self.status_label.setText("Loading workspace streams...")
self.figure.clear()
self.canvas.draw_idle()
try:
self.sources = discover_workspace_unit_cell_streams(self.workspace_ini_paths)
self.params = {param: [] for param in PARAMS}
if not self.sources:
self.fit_label.setText("No indexingintegration or ICI streams were found for the workspace files.")
self.status_label.setText(f"Workspace files: {len(self.workspace_ini_paths)} | Streams: 0")
self.fit_results.clear()
self.selected_param = None
self._draw_histograms()
return
multi_solution = str(self.multi_solution_combo.currentData())
per_image = self.per_image_cb.isChecked()
failed_streams = []
for source in self.sources:
try:
parsed = parse_stream_file(source.stream_path)
stream_params = unit_cell_parameters(
parsed,
multi_solution=multi_solution,
per_image=per_image,
length_scale=10.0,
)
for param in PARAMS:
self.params[param].extend(stream_params.get(param, []))
except Exception as exc:
failed_streams.append(f"{os.path.basename(source.stream_path)}: {exc}")
total_values = sum(len(self.params[param]) for param in PARAMS)
status = (
f"Workspace files: {len(self.workspace_ini_paths)} | "
f"Streams: {len(self.sources)} | Values: {total_values}"
)
if failed_streams:
status += f" | Skipped streams: {len(failed_streams)}"
self.status_label.setText(status)
self._draw_histograms()
except Exception as exc:
QMessageBox.critical(self, "Unit Cell Fit", f"Failed to load workspace unit-cell values:\n{exc}")
finally:
self._set_busy(False)
def _draw_histograms(self):
if not hasattr(self, "figure"):
return
self.figure.clear()
self._selectors = []
self._axes_by_param = {}
axes = self.figure.subplots(2, 3, squeeze=False).ravel()
bins = int(self.bins_spin.value())
units = {param: ("deg" if param in ANGLE_PARAMS else "Å") for param in PARAMS}
for ax, param in zip(axes, PARAMS):
self._axes_by_param[param] = ax
values = np.asarray(self.params.get(param, []), dtype=float)
values = values[np.isfinite(values)]
ax.set_xlabel(f"{param} ({units[param]})")
ax.set_ylabel("Count")
ax.grid(True, linestyle=":", alpha=0.4)
if values.size:
counts, edges = np.histogram(values, bins=bins)
width = edges[1] - edges[0] if len(edges) > 1 else 0.8
ax.bar(
edges[:-1],
counts,
width=width,
align="edge",
color="tab:blue",
edgecolor="tab:blue",
alpha=0.55,
)
fit = self.fit_results.get(param)
if fit:
ax.set_title(f"mean {fit['mean']:.5g} | std {fit['sigma']:.4g}")
else:
ax.set_title(f"N={values.size}")
else:
ax.text(0.5, 0.5, "No data", ha="center", va="center", color="#666", transform=ax.transAxes)
ax.set_yticks([])
fit = self.fit_results.get(param)
if fit:
xmin, xmax = fit["range"]
ax.axvspan(xmin, xmax, color="tab:orange", alpha=0.14)
ax.plot(fit["curve_x"], fit["curve_y"], color="tab:red", linewidth=2)
ax.axvline(fit["mean"], color="tab:red", linestyle="--", linewidth=1)
selector = SpanSelector(
ax,
lambda xmin, xmax, selected=param: self._on_range_selected(selected, xmin, xmax),
"horizontal",
useblit=True,
props={"facecolor": "tab:orange", "alpha": 0.25},
interactive=True,
)
self._selectors.append(selector)
if param in ANGLE_PARAMS and values.size:
xmin, xmax = ax.get_xlim()
if xmax - xmin < 20:
mid = 0.5 * (xmin + xmax)
ax.set_xlim(mid - 10, mid + 10)
self.figure.tight_layout()
self.canvas.draw_idle()
def _on_range_selected(self, param: str, xmin: float, xmax: float):
values = self.params.get(param, [])
try:
fit = fit_gaussian_to_values(values, (xmin, xmax), bins=int(self.bins_spin.value()))
except ValueError as exc:
self.fit_label.setText(str(exc))
self.update_cell_btn.setEnabled(False)
return
self.fit_results[param] = fit
self.selected_param = param
self.fit_label.setText(
f"{param} fit: mean {fit['mean']:.5g}, "
f"std {fit['sigma']:.4g}, FWHM {fit['fwhm']:.4g}, "
f"N={fit['selected_count']} [{fit['range'][0]:.5g}, {fit['range'][1]:.5g}]"
)
self.update_cell_btn.setEnabled(True)
self._draw_histograms()
def _update_cell_files(self):
if not self.selected_param:
return
fit = self.fit_results.get(self.selected_param)
if not fit:
return
if not self.workspace_path:
QMessageBox.warning(self, "Update workspace refined cell", "No workspace file is available.")
return
param = self.selected_param
value = fit["mean"]
reply = QMessageBox.question(
self,
"Update workspace refined cell",
f"Update workspace {param} to {value:.6g}?\n\n"
"This stores the refined value in the workspace. Existing indexing run cell files are not changed.",
QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
QMessageBox.StandardButton.No,
)
if reply != QMessageBox.StandardButton.Yes:
return
try:
update_workspace_refined_cell_parameter(self.workspace_path, param, value)
except Exception as exc:
QMessageBox.warning(self, "Update workspace refined cell", f"Failed to update workspace cell:\n{exc}")
return
QMessageBox.information(self, "Update workspace refined cell", "Workspace refined cell updated.")
self.status_label.setText(self.status_label.text() + " | Workspace refined cell updated")