"""Interactive stage map viewer for navigating frames and intensities."""
import configparser
import os
from coseda.logging_utils import log_print
from PyQt6.QtWidgets import (
QDialog,
QVBoxLayout,
QHBoxLayout,
QLabel,
QComboBox,
QPushButton,
QCheckBox,
QLineEdit,
QGroupBox,
QSizePolicy,
QMessageBox,
QSlider,
QDoubleSpinBox,
QWidget
)
from PyQt6.QtCore import pyqtSignal, Qt, QSize
from PyQt6.QtGui import QPixmap, QPainter, QColor, QFont
from PIL import Image
from PIL.ImageQt import ImageQt
import h5py
import numpy as np
from scipy.spatial import KDTree
import matplotlib.pyplot as plt
from .image_tools import HoverLabel
VIRTUAL_DARK_FIELD_OPTION = "Virtual dark field"
ATLAS_LEFT_PANEL_WIDTH = 300
ATLAS_BACKLASH_SECTION = "Parameters"
ATLAS_BACKLASH_KEY = "atlas_backlash_offset_um"
CONTRAST_LABELS = {
'frame_mean_intensities': "Mean intensity",
'frame_total_intensities': "Total intensity",
VIRTUAL_DARK_FIELD_OPTION: "Virtual dark field",
}
[docs]
class MapWindow(QDialog):
# Signal to update the main window when a point on the plot is clicked
point_clicked_signal = pyqtSignal(int)
def __init__(self, parent=None, hdf5_path=None, ini_path=None):
"""Initialize the map dialog and draw the initial map image if data is present."""
super().__init__(parent)
self.setWindowTitle("Atlas")
self.hdf5_path = hdf5_path
self.ini_path = ini_path
self.showcryst = None
self.current_frame_index = None
self.zdim = None
self._radial_bins = 0
self._radial_center = 0
self._radial_width = 1
self._backlash_offset_um = self._read_ini_backlash_offset()
self._streak_directions = {}
self.streak_ids = None
self._generate_streaks()
self.init_ui()
def _generate_streaks(self):
"""Detect scan streaks and write streak_id / streak_frame /
streak_endpoints into the HDF5 file when the map window is opened.
Runs on open so the map always has up-to-date streak assignments. Any
failure is logged and swallowed so it never blocks the window."""
if not self.hdf5_path:
return
try:
from coseda.map.streaks import assign_streaks
assign_streaks(self.hdf5_path, write=True)
except Exception as exc:
log_print(f"Streak assignment failed: {exc}", level="warning")
[docs]
def set_hdf5_path(self, hdf5_path, ini_path=None):
"""Switch the map to a different HDF5 file (e.g. the main window changed file)."""
if not hdf5_path or hdf5_path == self.hdf5_path:
return
self.hdf5_path = hdf5_path
self.ini_path = ini_path
self._backlash_offset_um = self._read_ini_backlash_offset()
if hasattr(self, "backlash_offset_spin"):
self.backlash_offset_spin.blockSignals(True)
self.backlash_offset_spin.setValue(self._backlash_offset_um)
self.backlash_offset_spin.blockSignals(False)
self._base_pixmap = None
self.raster = None
self.hover_info.setText("Position: -\nValue: -")
self.setWindowTitle(f"Atlas — {os.path.basename(hdf5_path)}")
self._generate_streaks() # write streak assignments for the new file
self._reload_controls() # repopulate tilt / dataset combos for the new file
if self.zdim is not None:
self.plot_data()
else:
self._base_pixmap = None
self.image_label.setText("Intensity datasets are not available.")
self.map_caption.setText("")
def _reload_controls(self):
"""Repopulate the alpha-tilt and dataset selectors for the current file."""
self.alphatilt_combo.blockSignals(True)
self.alphatilt_combo.clear()
self.alphatilt_combo.addItems(self.get_unique_alphatilts(self.hdf5_path))
self.alphatilt_combo.blockSignals(False)
self._set_default_contrast_selection()
self.print_option_combo.blockSignals(True)
self.print_option_combo.clear()
self._populate_contrast_combo()
self.print_option_combo.blockSignals(False)
self._configure_vdf_controls()
def _populate_contrast_combo(self):
"""Fill the contrast combo with readable labels and dataset keys as item data."""
for dataset in self.available_intensity_datasets:
self.print_option_combo.addItem(CONTRAST_LABELS.get(dataset, dataset), dataset)
if self.zdim is not None:
index = self.print_option_combo.findData(self.zdim)
if index >= 0:
self.print_option_combo.setCurrentIndex(index)
def _set_default_contrast_selection(self):
"""Select mean intensity by default, falling back to other available contrasts."""
self.available_intensity_datasets = self.get_available_intensity_datasets(self.hdf5_path)
if 'frame_mean_intensities' in self.available_intensity_datasets:
self.zdim = 'frame_mean_intensities'
elif 'frame_total_intensities' in self.available_intensity_datasets:
self.zdim = 'frame_total_intensities'
elif VIRTUAL_DARK_FIELD_OPTION in self.available_intensity_datasets:
self.zdim = VIRTUAL_DARK_FIELD_OPTION
else:
self.zdim = None
def _read_ini_backlash_offset(self):
"""Read the atlas backlash offset from the selected INI file."""
if not self.ini_path or not os.path.exists(self.ini_path):
return 0.0
config = configparser.ConfigParser()
try:
config.read(self.ini_path)
return config.getfloat(ATLAS_BACKLASH_SECTION, ATLAS_BACKLASH_KEY, fallback=0.0)
except Exception as exc:
log_print(f"Failed to read atlas backlash offset from INI: {exc}", level="warning")
return 0.0
def _write_ini_backlash_offset(self):
"""Persist the atlas backlash offset to the selected INI file."""
if not self.ini_path or not os.path.exists(self.ini_path):
return
config = configparser.ConfigParser()
try:
config.read(self.ini_path)
if ATLAS_BACKLASH_SECTION not in config:
config.add_section(ATLAS_BACKLASH_SECTION)
config.set(ATLAS_BACKLASH_SECTION, ATLAS_BACKLASH_KEY, f"{self._backlash_offset_um:.6g}")
with open(self.ini_path, "w") as handle:
config.write(handle)
except Exception as exc:
log_print(f"Failed to write atlas backlash offset to INI: {exc}", level="warning")
def _read_radial_bin_count(self):
"""Return the number of radial bins available for virtual dark field."""
if not self.hdf5_path:
return 0
try:
with h5py.File(self.hdf5_path, 'r') as file:
data_group = file.get('entry/data')
if data_group is None or 'frame_radial_intensities' not in data_group:
return 0
radial = data_group['frame_radial_intensities']
if radial.ndim != 2:
return 0
return int(radial.shape[1])
except Exception:
return 0
def _configure_vdf_controls(self):
"""Show virtual dark-field controls only when radial profiles are selected."""
if not hasattr(self, 'vdf_group'):
return
bins = self._read_radial_bin_count()
self._radial_bins = bins
has_radial = bins > 0
use_vdf = self.zdim == VIRTUAL_DARK_FIELD_OPTION and has_radial
self.vdf_group.setVisible(use_vdf)
self.vdf_center_slider.blockSignals(True)
self.vdf_width_slider.blockSignals(True)
self.vdf_center_slider.setEnabled(use_vdf)
self.vdf_width_slider.setEnabled(use_vdf)
if hasattr(self, 'vdf_overlay_checkbox'):
self.vdf_overlay_checkbox.setEnabled(use_vdf)
if has_radial:
default_center = max(0, bins // 2)
default_width = max(1, bins // 10)
if self._radial_center <= 0 or self._radial_center >= bins:
self._radial_center = default_center
if self._radial_width <= 0 or self._radial_width > bins:
self._radial_width = default_width
self.vdf_center_slider.setRange(0, bins - 1)
self.vdf_width_slider.setRange(1, bins)
self.vdf_center_slider.setValue(self._radial_center)
self.vdf_width_slider.setValue(self._radial_width)
else:
self._radial_center = 0
self._radial_width = 1
self.vdf_center_slider.setRange(0, 0)
self.vdf_width_slider.setRange(1, 1)
self.vdf_center_slider.setValue(0)
self.vdf_width_slider.setValue(1)
self.vdf_center_slider.blockSignals(False)
self.vdf_width_slider.blockSignals(False)
self._update_vdf_labels()
self._sync_vdf_overlay()
def _update_vdf_labels(self):
"""Refresh virtual detector slider labels."""
if not hasattr(self, 'vdf_center_label'):
return
if self._radial_bins <= 0:
self.vdf_center_label.setText("Center radius: - px")
self.vdf_width_label.setText("Annulus width: - px")
return
self.vdf_center_label.setText(f"Center radius: {self._radial_center} px")
self.vdf_width_label.setText(f"Annulus width: {self._radial_width} px")
def _vdf_bin_bounds(self):
"""Return half-open radial-bin bounds for the virtual detector annulus."""
half_width = self._radial_width / 2.0
lo = int(np.floor(self._radial_center - half_width))
hi = int(np.ceil(self._radial_center + half_width))
lo = max(0, lo)
hi = min(self._radial_bins, max(lo + 1, hi))
return lo, hi
def _contrast_label(self):
"""Return a readable label for the current contrast."""
return CONTRAST_LABELS.get(self.zdim, str(self.zdim))
def _atlas_contrast_lines(self):
"""Return compact lines describing the active contrast."""
if self.zdim != VIRTUAL_DARK_FIELD_OPTION:
return [f"Contrast: {self._contrast_label()}"]
lo, hi = self._vdf_bin_bounds()
return [
f"Contrast: {self._contrast_label()}",
f"Detector: center {self._radial_center} px, width {self._radial_width} px",
f"Radial bins: {lo}-{hi - 1} px",
]
def _atlas_backlash_lines(self):
"""Return caption lines for the persistent atlas backlash correction."""
if not self._backlash_active():
return []
return [f"Backlash: left streaks {self._backlash_offset_um:+.2f} µm"]
@staticmethod
def _format_atlas_value(value):
"""Format atlas values without long labels changing the side-panel width."""
if not np.isfinite(value):
return "-"
magnitude = abs(value)
if magnitude >= 10000 or (0 < magnitude < 0.01):
return f"{value:.3g}"
return f"{value:.1f}"
@staticmethod
def _read_endpoint_directions(data_group):
"""Read streak direction metadata when available."""
if 'streak_endpoints' not in data_group:
return {}
endpoints = data_group['streak_endpoints'][:]
if endpoints.dtype.names is None:
return {}
if 'streak_id' not in endpoints.dtype.names or 'direction' not in endpoints.dtype.names:
return {}
return {
int(row['streak_id']): int(row['direction'])
for row in endpoints
if int(row['streak_id']) >= 0 and int(row['direction']) != 0
}
def _build_streak_directions(self, x, valid):
"""Build per-streak scan directions, falling back to x start/stop if needed."""
directions = dict(getattr(self, "_streak_directions", {}))
if self.streak_id is None:
return directions
valid = np.asarray(valid, dtype=bool) & (self.streak_id >= 0)
for sid in np.unique(self.streak_id[valid]):
sid = int(sid)
if sid in directions:
continue
idx = np.flatnonzero(valid & (self.streak_id == sid))
if idx.size < 2:
continue
dx = float(x[idx[-1]] - x[idx[0]])
if dx != 0:
directions[sid] = 1 if dx > 0 else -1
self._streak_directions = directions
return directions
def _left_streak_mask(self, x, valid):
"""Return dense-frame mask for left-going streaks."""
if self.streak_id is None:
return np.zeros(len(x), dtype=bool)
directions = self._build_streak_directions(x, valid)
left_ids = {sid for sid, direction in directions.items() if direction < 0}
if not left_ids:
return np.zeros(len(x), dtype=bool)
return np.isin(self.streak_id, list(left_ids))
def _backlash_offset_m(self):
"""Return the active left-streak x offset in metres."""
return self._backlash_offset_um * 1e-6 if self._backlash_active() else 0.0
def _backlash_active(self):
"""Return True when a nonzero atlas backlash correction should be applied."""
return abs(self._backlash_offset_um) > 1e-12
def _corrected_stagepos_x(self):
"""Return stage x with the left-streak backlash offset applied."""
offset_m = self._backlash_offset_m()
if offset_m == 0.0:
return self.stagepos_x
corrected = self.stagepos_x.astype(float, copy=True)
corrected[self._left_streak_mask(self.stagepos_x, self._raster_valid)] += offset_m
return corrected
def _x_lookup_from_display(self, x_display, raster_row):
"""Convert displayed x back to real stage x for nearest-frame lookup."""
offset_m = self._backlash_offset_m()
if offset_m == 0.0 or self.streak_ids is None:
return x_display
if not (0 <= raster_row < len(self.streak_ids)):
return x_display
sid = int(self.streak_ids[raster_row])
if self._streak_directions.get(sid, 0) < 0:
return x_display - offset_m
return x_display
def _read_contrast_values(self, data_group):
"""Read or derive the selected per-frame contrast values."""
if self.zdim != VIRTUAL_DARK_FIELD_OPTION:
return data_group[self.zdim][:].astype(float)
radial = data_group['frame_radial_intensities'][:].astype(float)
lo, hi = self._vdf_bin_bounds()
return np.nansum(radial[:, lo:hi], axis=1)
def _sync_vdf_overlay(self):
"""Update or clear the main-window virtual detector overlay."""
parent = self.parent()
overlay_checkbox = getattr(self, 'vdf_overlay_checkbox', None)
active = (
self.zdim == VIRTUAL_DARK_FIELD_OPTION
and self._radial_bins > 0
and overlay_checkbox is not None
and overlay_checkbox.isChecked()
)
if active and hasattr(parent, "set_vdf_overlay"):
parent.set_vdf_overlay(self._radial_center, self._radial_width, enabled=True)
elif hasattr(parent, "clear_vdf_overlay"):
parent.clear_vdf_overlay()
[docs]
def init_ui(self):
"""Build the left-column control/info group boxes and the map image."""
# Create a horizontal layout for the two-column structure
main_layout = QHBoxLayout(self)
# Left column: stacked labelled group boxes (like the main window).
left_col = QVBoxLayout()
# --- Controls ---
control_group = QGroupBox("Control Elements")
control_layout = QVBoxLayout()
selector_width = 180
label_width = 92
alphatilt_row = QHBoxLayout()
alphatilt_label = QLabel("Alpha tilt:")
alphatilt_label.setFixedWidth(label_width)
alphatilt_row.addWidget(alphatilt_label)
self.alphatilt_combo = QComboBox()
self.alphatilt_combo.addItems(self.get_unique_alphatilts(self.hdf5_path))
self.alphatilt_combo.setMaximumWidth(selector_width)
alphatilt_row.addWidget(self.alphatilt_combo)
alphatilt_row.addStretch(1)
control_layout.addLayout(alphatilt_row)
contrast_row = QHBoxLayout()
contrast_label = QLabel("Contrast:")
contrast_label.setFixedWidth(label_width)
contrast_row.addWidget(contrast_label)
self.print_option_combo = QComboBox()
self._set_default_contrast_selection()
self._populate_contrast_combo()
self.print_option_combo.currentIndexChanged.connect(self.handle_print_option_change)
self.print_option_combo.setMaximumWidth(selector_width)
contrast_row.addWidget(self.print_option_combo)
contrast_row.addStretch(1)
control_layout.addLayout(contrast_row)
self.vdf_group = QGroupBox("Virtual dark field")
vdf_layout = QVBoxLayout()
self.vdf_center_label = QLabel("Center: - px")
self.vdf_center_slider = QSlider(Qt.Orientation.Horizontal)
self.vdf_center_slider.setMinimum(0)
self.vdf_center_slider.setMaximum(0)
self.vdf_center_slider.valueChanged.connect(self.on_vdf_slider_changed)
self.vdf_width_label = QLabel("Width: - px")
self.vdf_width_slider = QSlider(Qt.Orientation.Horizontal)
self.vdf_width_slider.setMinimum(1)
self.vdf_width_slider.setMaximum(1)
self.vdf_width_slider.valueChanged.connect(self.on_vdf_slider_changed)
self.vdf_overlay_checkbox = QCheckBox("Show detector overlay")
self.vdf_overlay_checkbox.setChecked(True)
self.vdf_overlay_checkbox.stateChanged.connect(self.on_vdf_overlay_toggled)
vdf_layout.addWidget(self.vdf_center_label)
vdf_layout.addWidget(self.vdf_center_slider)
vdf_layout.addWidget(self.vdf_width_label)
vdf_layout.addWidget(self.vdf_width_slider)
vdf_layout.addWidget(self.vdf_overlay_checkbox)
self.vdf_group.setLayout(vdf_layout)
control_layout.addWidget(self.vdf_group)
self._configure_vdf_controls()
self.backlash_group = QGroupBox("Backlash correction")
backlash_layout = QVBoxLayout()
backlash_row = QHBoxLayout()
backlash_label = QLabel("Left streak offset:")
backlash_label.setFixedWidth(120)
backlash_row.addWidget(backlash_label)
self.backlash_offset_spin = QDoubleSpinBox()
self.backlash_offset_spin.setRange(-1000.0, 1000.0)
self.backlash_offset_spin.setDecimals(2)
self.backlash_offset_spin.setSingleStep(0.10)
self.backlash_offset_spin.setSuffix(" µm")
self.backlash_offset_spin.setValue(self._backlash_offset_um)
self.backlash_offset_spin.setMaximumWidth(120)
self.backlash_offset_spin.valueChanged.connect(self.on_backlash_offset_changed)
backlash_row.addWidget(self.backlash_offset_spin)
backlash_row.addStretch(1)
backlash_layout.addLayout(backlash_row)
self.backlash_group.setLayout(backlash_layout)
control_layout.addWidget(self.backlash_group)
self.crystallinity_checkbox = QCheckBox("Highlight Diffracting Regions")
self.crystallinity_checkbox.stateChanged.connect(self.on_crystallinity_checkbox_toggle)
control_layout.addWidget(self.crystallinity_checkbox)
self.crystallinity_threshold_entry = QLineEdit()
self.crystallinity_threshold_entry.setPlaceholderText("Crystallinity Threshold")
self.crystallinity_threshold_entry.setDisabled(True)
self.crystallinity_threshold_entry.textChanged.connect(self.on_crystallinity_threshold_changed)
control_layout.addWidget(self.crystallinity_threshold_entry)
control_group.setLayout(control_layout)
left_col.addWidget(control_group)
# --- Map info ---
map_info_group = QGroupBox("Atlas info")
map_info_layout = QVBoxLayout()
self.map_caption = QLabel("")
self.map_caption.setWordWrap(True)
self.map_caption.setTextInteractionFlags(Qt.TextInteractionFlag.TextSelectableByMouse)
map_info_layout.addWidget(self.map_caption)
map_info_group.setLayout(map_info_layout)
left_col.addWidget(map_info_group)
# --- Cursor ---
cursor_group = QGroupBox("Cursor")
cursor_layout = QVBoxLayout()
self.hover_info = QLabel("Position: -\nValue: -")
self.hover_info.setTextInteractionFlags(Qt.TextInteractionFlag.TextSelectableByMouse)
cursor_layout.addWidget(self.hover_info)
cursor_group.setLayout(cursor_layout)
left_col.addWidget(cursor_group)
# Close button, then a stretch so the column stays top-aligned.
close_button = QPushButton("Close")
close_button.clicked.connect(self.close)
left_col.addWidget(close_button)
left_col.addStretch(1)
left_panel = QWidget()
left_panel.setLayout(left_col)
left_panel.setFixedWidth(ATLAS_LEFT_PANEL_WIDTH)
main_layout.addWidget(left_panel)
# Right column: the map as a proper image (like the main-window frame).
map_layout = QVBoxLayout()
self.image_label = HoverLabel()
self.image_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
self.image_label.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
self.image_label.setMinimumSize(320, 320)
self.image_label.setStyleSheet("QLabel { background: transparent; }")
self.image_label.hovered.connect(self.on_hover)
self.image_label.mousePressEvent = self.on_image_click
map_layout.addWidget(self.image_label)
main_layout.addLayout(map_layout, stretch=1)
self.setLayout(main_layout)
# Rendering state.
self._base_pixmap = None # full-res colormapped QPixmap (no scalebar)
self._scaled_size = (0, 0)
self.raster = None
self.extent = None
self.row_y = None
# Plot the data when opening the window if zdim is set
if self.zdim is not None:
self.plot_data()
else:
QMessageBox.warning(
self,
"Dataset Missing",
"Intensity datasets are not available."
)
[docs]
def handle_print_option_change(self):
"""Update z-dimension based on selection and replot."""
self.zdim = self.print_option_combo.currentData()
self._configure_vdf_controls()
self.plot_data()
[docs]
def on_vdf_slider_changed(self):
"""Update virtual dark-field annulus settings and refresh the map."""
self._radial_center = int(self.vdf_center_slider.value())
self._radial_width = max(1, int(self.vdf_width_slider.value()))
self._update_vdf_labels()
self._sync_vdf_overlay()
if self.zdim == VIRTUAL_DARK_FIELD_OPTION:
self.plot_data()
[docs]
def on_vdf_overlay_toggled(self):
"""Show/hide the main-window virtual detector overlay."""
self._sync_vdf_overlay()
[docs]
def on_backlash_offset_changed(self, value):
"""Update the persistent left-streak backlash offset and refresh the map."""
self._backlash_offset_um = float(value)
self._write_ini_backlash_offset()
self.plot_data()
[docs]
def on_crystallinity_checkbox_toggle(self):
"""Enable/disable crystallinity thresholding and refresh the plot."""
if self.crystallinity_checkbox.isChecked():
self.crystallinity_threshold_entry.setDisabled(False)
try:
self.showcryst = float(self.crystallinity_threshold_entry.text() or 0)
except ValueError:
self.showcryst = None
else:
self.crystallinity_threshold_entry.setDisabled(True)
self.showcryst = None
self.plot_data()
[docs]
def on_crystallinity_threshold_changed(self, text):
"""Parse user threshold and refresh plot when highlighting diffracting regions."""
if self.crystallinity_checkbox.isChecked():
try:
self.showcryst = float(text)
except ValueError:
self.showcryst = None # Invalid input
self.plot_data()
# REMOVED: Handlers for Highlight Stripped Frames Checkbox and Threshold Entry
# def on_highlight_stripped_checkbox_toggle(self, state):
# if self.highlight_stripped_checkbox.isChecked():
# self.stripped_threshold_entry.setDisabled(False)
# self.highlight_stripped = True
# try:
# self.stripped_threshold = float(self.stripped_threshold_entry.text() or 0)
# except ValueError:
# self.stripped_threshold = None
# else:
# self.stripped_threshold_entry.setDisabled(True)
# self.highlight_stripped = False
# self.stripped_threshold = None
# self.plot_data()
# def on_stripped_threshold_changed(self, text):
# if self.highlight_stripped_checkbox.isChecked():
# try:
# self.stripped_threshold = float(text)
# except ValueError:
# self.stripped_threshold = None # Invalid input
# self.plot_data()
[docs]
def plot_data(self):
"""Load the scan, rasterize it, and render it as a map image."""
if self.zdim is None:
QMessageBox.warning(
self,
"Dataset Missing",
"No Z dimension selected for plotting."
)
return
with h5py.File(self.hdf5_path, 'r') as workingfile:
data_group = workingfile['entry/data']
contrast_dataset = (
'frame_radial_intensities'
if self.zdim == VIRTUAL_DARK_FIELD_OPTION else self.zdim
)
# Required datasets
stagepos_x_dataset = (
'stagepos_x_refined'
if 'stagepos_x_refined' in data_group else 'stagepos_x'
)
stagepos_y_dataset = (
'stagepos_y_refined'
if 'stagepos_y_refined' in data_group else 'stagepos_y'
)
required_datasets = [
stagepos_x_dataset,
stagepos_y_dataset,
contrast_dataset,
'index'
]
missing_datasets = [ds for ds in required_datasets if ds not in data_group]
if missing_datasets:
QMessageBox.warning(
self,
"Dataset Missing",
f"The following datasets are missing in the HDF5 file: {missing_datasets}"
)
return
# Read exception datasets (full length)
self.stagepos_x = data_group[stagepos_x_dataset][:]
self.stagepos_y = data_group[stagepos_y_dataset][:]
self.frame_indices = np.arange(len(self.stagepos_x))
# Read index dataset
self.index_dataset = data_group['index'][:]
self.available_frames_mask = self.index_dataset != -1 # True for available frames
# Streak assignment (written on open by _generate_streaks); enables the raster view.
self.streak_id = data_group['streak_id'][:] if 'streak_id' in data_group else None
self._streak_directions = self._read_endpoint_directions(data_group)
# Per-frame stage validity: every saved frame keeps its stage position
# (dense, untouched by stripping) even if its image was later dumped.
valid_stagepos = (
np.isfinite(self.stagepos_x) & np.isfinite(self.stagepos_y)
& (self.stagepos_x != 999999.0) & (self.stagepos_y != 999999.0)
)
valid_pos = np.flatnonzero(valid_stagepos)
self._raster_valid = valid_stagepos
# Read or derive the display contrast and place it on the dense frame axis.
zdimension_reduced = self._read_contrast_values(data_group)
self.zdimension_full = np.full(len(self.stagepos_x), np.nan)
if len(zdimension_reduced) == len(self.stagepos_x):
self.zdimension_full = zdimension_reduced.astype(float) # already dense
elif len(zdimension_reduced) == len(valid_pos):
# Compact intensities are ordered by the valid-stagepos frames in dense
# order -- robust to stripping (which shrinks index, not stagepos), so
# dumped frames still get their intensity and the atlas stays complete.
self.zdimension_full[valid_pos] = zdimension_reduced
else:
# Fallback: map via the current index (unstripped/aligned case).
self.zdimension_full[self.available_frames_mask] = \
zdimension_reduced[self.index_dataset[self.available_frames_mask]]
# Build KD-tree for available frames
self.available_positions = np.column_stack((self.stagepos_x[self.available_frames_mask], self.stagepos_y[self.available_frames_mask]))
self.available_indices = self.frame_indices[self.available_frames_mask]
self.kd_tree = KDTree(self.available_positions)
# Read nPeaks if needed for crystallinity filter
if self.showcryst is not None and 'nPeaks' in data_group:
nPeaks = data_group['nPeaks'][:]
else:
nPeaks = None
stagepos_x_for_map = self._corrected_stagepos_x()
# Diffraction overlay: stage positions of available frames with nPeaks >= threshold.
if self.showcryst is not None and nPeaks is not None:
npeaks_full = np.full(len(self.stagepos_x), np.nan)
if len(nPeaks) == len(self.stagepos_x): # dense nPeaks
npeaks_full = nPeaks.astype(float)
else: # compact -> dense via index
npeaks_full[self.available_frames_mask] = \
nPeaks[self.index_dataset[self.available_frames_mask]]
diff_mask = self.available_frames_mask & (npeaks_full >= self.showcryst)
self._diffraction_xy = np.column_stack(
[stagepos_x_for_map[diff_mask], self.stagepos_y[diff_mask]])
else:
self._diffraction_xy = None
if self.streak_id is None:
self._base_pixmap = None
self.image_label.setText("Streak assignment unavailable.\nReopen after processing.")
self.map_caption.setText("")
return
try:
from coseda.map.streaks import rasterize_arrays
out = rasterize_arrays(
self.streak_id, stagepos_x_for_map, self.stagepos_y,
self.zdimension_full, valid=self._raster_valid,
)
except Exception as exc:
log_print(f"Rasterization failed: {exc}", level="warning")
self._base_pixmap = None
self.image_label.setText(f"Rasterization failed:\n{exc}")
self.map_caption.setText("")
return
self.raster = out["raster"]
self.extent = out["extent"]
self.row_y = out["row_y"]
self.streak_ids = out["streak_ids"]
out["left_streak_offset_m"] = self._backlash_offset_m()
self._x_coords = np.linspace(self.extent[0], self.extent[1], self.raster.shape[1])
# Persist direct dataset atlases into the file. VDF contrast is derived
# interactively from sliders, so avoid writing a new atlas on every drag.
if self.zdim != VIRTUAL_DARK_FIELD_OPTION:
try:
from coseda.map.streaks import write_atlas
write_atlas(self.hdf5_path, self.zdim, out)
except Exception as exc:
log_print(f"Atlas save failed: {exc}", level="warning")
self._render_image()
def _current_cmap(self):
"""Reuse the main window's selected colormap, defaulting to inferno_r."""
combo = getattr(self.parent(), "cmap_combo", None)
name = combo.currentText() if combo is not None else None
if name:
try:
plt.get_cmap(name)
return name
except Exception:
pass
return "inferno_r"
def _render_image(self):
"""Colormap the raster into an RGBA image; the scalebar is drawn on scaling."""
if self.raster is None:
return
raster = self.raster
finite = raster[np.isfinite(raster)]
vmin = float(np.percentile(finite, 1.0)) if finite.size else 0.0
vmax = float(np.percentile(finite, 99.0)) if finite.size else 1.0
if not np.isfinite(vmax) or vmax <= vmin:
vmax = vmin + 1.0
norm = np.clip((raster - vmin) / (vmax - vmin), 0.0, 1.0)
cmap = plt.get_cmap(self._current_cmap())
rgba = (cmap(np.nan_to_num(norm)) * 255).astype(np.uint8) # (rows, cols, 4)
rgba[~np.isfinite(raster)] = (0, 0, 0, 0) # gaps transparent
rgba = np.flipud(rgba) # image top = max y
img = Image.fromarray(rgba, "RGBA")
self._base_pixmap = QPixmap.fromImage(ImageQt(img))
x0, x1, y0, y1 = self.extent
caption_lines = [
f"Size: {raster.shape[0]} streaks × {raster.shape[1]} px",
f"X span: {x0 * 1e6:.1f} to {x1 * 1e6:.1f} µm",
f"Y span: {y0 * 1e6:.1f} to {y1 * 1e6:.1f} µm",
*self._atlas_contrast_lines(),
*self._atlas_backlash_lines(),
f"Range: {self._format_atlas_value(vmin)} to {self._format_atlas_value(vmax)}",
]
self.map_caption.setText("\n".join(caption_lines))
self._update_pixmap()
def _update_pixmap(self):
"""Scale the base image to the PHYSICAL aspect ratio and bake the scalebar.
The raster's pixel-count aspect (n_cols/n_rows) varies with the frames-per-
streak, so scaling by it distorts the map. Instead we fit a box whose aspect
is the physical x-span : y-span (like imshow aspect='equal'), so the same
scan area always renders with the same shape regardless of sampling density.
"""
if self._base_pixmap is None or self.extent is None:
return
lw, lh = self.image_label.width(), self.image_label.height()
if lw < 4 or lh < 4:
return
x0, x1, y0, y1 = self.extent
x_span, y_span = abs(x1 - x0), abs(y1 - y0)
aspect = (x_span / y_span) if y_span > 0 else (
self._base_pixmap.width() / max(self._base_pixmap.height(), 1))
if lw / lh > aspect: # label wider than the map -> height-limited
disp_h, disp_w = lh, int(round(lh * aspect))
else: # width-limited
disp_w, disp_h = lw, int(round(lw / aspect))
scaled = self._base_pixmap.scaled(
max(1, disp_w), max(1, disp_h),
Qt.AspectRatioMode.IgnoreAspectRatio, Qt.TransformationMode.FastTransformation,
)
self._draw_diffraction(scaled)
self._draw_scalebar(scaled)
self._scaled_size = (scaled.width(), scaled.height())
self.image_label.setPixmap(scaled)
def _draw_diffraction(self, pixmap):
"""Overlay diffracting frames (nPeaks >= threshold) as crystalmap-style red squares."""
xy = getattr(self, "_diffraction_xy", None)
if xy is None or len(xy) == 0 or self.extent is None or self.raster is None:
return
x0, x1, y0, y1 = self.extent
x_span, y_span = (x1 - x0), (y1 - y0)
if x_span == 0 or y_span == 0:
return
W, H = pixmap.width(), pixmap.height()
hs = max(2.0, W / self.raster.shape[1] * 1.1) # ~ one raster pixel wide
color = QColor(255, 40, 40, 217) # rgba(255,40,40,0.85)
painter = QPainter(pixmap)
try:
for x, y in xy:
px = (x - x0) / x_span * W
py = (y1 - y) / y_span * H # image top = max y
painter.fillRect(int(px - hs / 2), int(py - hs / 2),
int(hs), int(hs), color)
finally:
painter.end()
def _draw_scalebar(self, pixmap):
"""Draw a 1-2-5 scalebar (nm/µm/mm) bottom-left, crisp at display resolution."""
if self.extent is None:
return
x_span_m = abs(self.extent[1] - self.extent[0]) # stage units are metres
if not (x_span_m > 0):
return
W, H = pixmap.width(), pixmap.height()
target_nm = 0.15 * x_span_m * 1e9
steps = [10, 20, 50, 100, 200, 500, 1e3, 2e3, 5e3, 1e4, 2e4, 5e4, 1e5, 2e5, 5e5, 1e6]
nice = steps[0]
for s in steps:
if s <= target_nm:
nice = s
bar_px = max(1, int((nice * 1e-9) / x_span_m * W))
label = (f"{nice:g} nm" if nice < 1e3 else
f"{nice / 1e3:g} µm" if nice < 1e6 else f"{nice / 1e6:g} mm")
painter = QPainter(pixmap)
try:
pad = 6
thick = max(3, int(H * 0.008))
font = QFont()
font.setPixelSize(max(11, int(W * 0.02)))
painter.setFont(font)
fm = painter.fontMetrics()
tw = fm.horizontalAdvance(label)
th = fm.height()
box_w = max(bar_px, tw) + 2 * pad
box_h = thick + th + 3 * pad
margin_x = max(12, int(W * 0.05)) # further inward horizontally
margin_y = max(6, int(H * 0.02))
box_left = margin_x - pad
box_top = H - margin_y - box_h
painter.fillRect(box_left, box_top, box_w, box_h, QColor(0, 0, 0, 140))
bar_top = box_top + pad
painter.fillRect(margin_x, bar_top, bar_px, thick, QColor(255, 255, 255))
painter.setPen(QColor(255, 255, 255))
painter.drawText(margin_x, bar_top + thick + pad + fm.ascent(), label)
finally:
painter.end()
[docs]
def resizeEvent(self, event):
super().resizeEvent(event)
self._update_pixmap()
def _widget_to_cell(self, px, py):
"""Map a label pixel to (x_m, y_m, raster_row, col, value), or None if outside."""
if self.raster is None or self._base_pixmap is None:
return None
sw, sh = self._scaled_size
if sw < 1 or sh < 1:
return None
off_x = (self.image_label.width() - sw) / 2.0
off_y = (self.image_label.height() - sh) / 2.0
ix, iy = px - off_x, py - off_y
if not (0 <= ix < sw and 0 <= iy < sh):
return None
n_rows, n_cols = self.raster.shape
col = min(int(ix / sw * n_cols), n_cols - 1)
raster_row = n_rows - 1 - min(int(iy / sh * n_rows), n_rows - 1) # image top = max y
return (float(self._x_coords[col]), float(self.row_y[raster_row]),
raster_row, col, float(self.raster[raster_row, col]))
[docs]
def on_hover(self, px, py):
"""Update the cursor info box with stage position + value under the cursor."""
cell = self._widget_to_cell(px, py)
if cell is None:
self.hover_info.setText("Position: -\nValue: -")
return
x_m, y_m, _r, _c, val = cell
val_txt = "-" if not np.isfinite(val) else f"{val:.4g}"
self.hover_info.setText(
f"x = {x_m * 1e6:.2f} µm\ny = {y_m * 1e6:.2f} µm\nvalue = {val_txt}"
)
[docs]
def on_image_click(self, event):
"""Jump the main window to the nearest available frame under the click."""
if self.raster is None:
return
pos = event.position()
cell = self._widget_to_cell(pos.x(), pos.y())
if cell is None:
return
x_lookup = self._x_lookup_from_display(cell[0], cell[2])
closest = self.find_closest_available_frame(x_lookup, cell[1])
if closest is not None:
self.point_clicked_signal.emit(int(self.index_dataset[closest]))
[docs]
def find_closest_available_frame(self, x_clicked, y_clicked):
"""Return the nearest available frame index to the clicked stage coordinates."""
# Query the KD-tree for the nearest available frame
distance, idx = self.kd_tree.query([x_clicked, y_clicked])
closest_frame_index = self.available_indices[idx]
# Optionally, set a maximum allowable distance
MAX_DISTANCE = np.inf # Or set a finite value if desired
if distance <= MAX_DISTANCE:
return closest_frame_index
else:
return None
[docs]
def get_unique_alphatilts(self, hdf5_path):
"""Sorted unique alpha tilts, rounded to 1 decimal, without the 999999 placeholder."""
with h5py.File(hdf5_path, 'r') as file:
data_group = file['entry/data']
if 'alphatilt' not in data_group:
return []
alphatilt = data_group['alphatilt'][:]
values = np.unique(alphatilt)
values = values[np.isfinite(values) & (values != 999999)] # drop placeholder / NaN
rounded = sorted({(round(float(v), 1) or 0.0) for v in values}) # 1 decimal, dedup, no -0.0
return [f"{v:.1f}" for v in rounded]
[docs]
def get_available_intensity_datasets(self, hdf5_path):
"""Return available atlas contrast sources."""
datasets = []
with h5py.File(hdf5_path, 'r') as file:
data_group = file['entry/data']
if 'frame_mean_intensities' in data_group:
datasets.append('frame_mean_intensities')
if 'frame_total_intensities' in data_group:
datasets.append('frame_total_intensities')
if 'frame_radial_intensities' in data_group:
datasets.append(VIRTUAL_DARK_FIELD_OPTION)
return datasets
# Cleanup on close
[docs]
def closeEvent(self, event):
"""Release the rendered image on dialog close."""
self._base_pixmap = None
parent = self.parent()
if hasattr(parent, "clear_vdf_overlay"):
parent.clear_vdf_overlay()
event.accept()