From 6d42b1de5a418a9453cd12e0fb97eb7e1fae2458 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Tue, 23 Jun 2026 09:50:37 +0200 Subject: [PATCH] raster_grid_manager.py: fixed lag with large grid, now generates a single bitmap --- .../gui/scan_logic/raster_grid_manager.py | 183 ++++++++++++++++-- 1 file changed, 166 insertions(+), 17 deletions(-) diff --git a/src/aare/gui/scan_logic/raster_grid_manager.py b/src/aare/gui/scan_logic/raster_grid_manager.py index 61584627..627d7646 100644 --- a/src/aare/gui/scan_logic/raster_grid_manager.py +++ b/src/aare/gui/scan_logic/raster_grid_manager.py @@ -1,8 +1,9 @@ import math from enum import Enum +import numpy as np from PySide6.QtCore import QObject, Signal, Slot, QPointF, QRectF, QLineF -from PySide6.QtGui import QPainter, QPen, QColor, QBrush +from PySide6.QtGui import QPainter, QPen, QColor, QBrush, QImage from PySide6.QtCore import Qt, QRect from typing import List, Tuple @@ -29,6 +30,37 @@ class RasterGridMetric(Enum): SPOTS_INDEXED = 8 SPOTS_ICE_LOW_RES = 9 +# Viridis colormap colors (from matplotlib), dark purple -> yellow +_VIRIDIS_COLORS = [ + (68, 1, 84), # Dark purple + (72, 35, 116), + (64, 67, 135), + (52, 94, 141), + (41, 120, 142), + (32, 144, 140), + (34, 167, 132), + (68, 190, 112), + (121, 209, 81), + (189, 223, 38), + (253, 231, 37) # Yellow +] + +# Float (N, 3) lookup table for vectorised colour mapping. +_VIRIDIS_LUT = np.array(_VIRIDIS_COLORS, dtype=np.float64) + + +def viridis_rgb_array(norm: np.ndarray) -> np.ndarray: + """Map an array of normalized values in [0, 1] to (..., 3) uint8 RGB using + the Viridis colormap, vectorised with the same linear interpolation as + float_to_viridis_brush.""" + pos = np.clip(norm, 0.0, 1.0) * (len(_VIRIDIS_LUT) - 1) + lo = np.clip(np.floor(pos).astype(np.intp), 0, len(_VIRIDIS_LUT) - 1) + hi = np.clip(lo + 1, 0, len(_VIRIDIS_LUT) - 1) + t = (pos - lo)[..., None] + c = _VIRIDIS_LUT[lo] * (1.0 - t) + _VIRIDIS_LUT[hi] * t + return np.clip(np.rint(c), 0, 255).astype(np.uint8) + + def float_to_viridis_brush(value: float, alpha: int = 127) -> QBrush: """ Convert float value to a QBrush with Viridis colormap @@ -38,20 +70,7 @@ def float_to_viridis_brush(value: float, alpha: int = 127) -> QBrush: Returns: QBrush with corresponding Viridis color """ - # Viridis colormap colors (from matplotlib) - viridis_colors = [ - (68, 1, 84), # Dark purple - (72, 35, 116), - (64, 67, 135), - (52, 94, 141), - (41, 120, 142), - (32, 144, 140), - (34, 167, 132), - (68, 190, 112), - (121, 209, 81), - (189, 223, 38), - (253, 231, 37) # Yellow - ] + viridis_colors = _VIRIDIS_COLORS # Clamp value between 0 and 1 value = max(0.0, min(1.0, value)) @@ -121,6 +140,12 @@ class RasterGridManager(QObject): ) self.__completed_grids : List[CompletedRasterGridElem] = [] + # Cache of pre-rendered heatmap bitmaps, keyed by (id(grid_elem), metric). + # Each entry is (QImage, backing ndarray); the ndarray must be kept alive + # because QImage shares its buffer without copying. Rebuilt only when the + # data or metric changes, not on every repaint (sample move / zoom). + self.__heatmap_cache: dict[tuple[int, "RasterGridMetric"], tuple[QImage, np.ndarray]] = {} + @property def active_grid(self) -> RasterGridRequest: return self.__active_grid @@ -300,7 +325,7 @@ class RasterGridManager(QObject): self.__loaded_image_index != grid.result.images[cell].number): self.__loaded_image_prefix = grid.result.file_prefix self.__loaded_image_index = grid.result.images[cell].number - print(f"Load {grid.result.file_prefix} {grid.result.images[cell].number}") + logger.debug(f"Load {grid.result.file_prefix} {grid.result.images[cell].number}") #self.image_selected.emit(grid.result.file_prefix, grid.result.images[cell].number) self.image_selected.emit(self.__detector_url, grid.result.images[cell].number) @@ -448,7 +473,127 @@ class RasterGridManager(QObject): v = [self.spot_ice_ratio(obj) for obj in i.result.images] case RasterGridMetric.SPOTS_INDEXED: v = [obj.spots_indexed for obj in i.result.images] - self._draw_grid(painter, i.request, v, alpha, visible_rect) + self._draw_completed_heatmap( + painter, i, v, alpha, visible_rect, (id(i), self.__metric) + ) + + def _invalidate_heatmap_cache(self) -> None: + """Drop all cached heatmap bitmaps. Call whenever the completed-grid data + set or the active metric changes; cheap to rebuild lazily on next paint.""" + self.__heatmap_cache.clear() + + def _heatmap_image( + self, + cache_key: tuple, + grid: RasterGridRequest, + values: List[float] | List[int], + ) -> QImage | None: + """Return a cached n_x*n_y heatmap bitmap for this grid, building it once + on a cache miss. One pixel per cell; colours baked at full opacity with + the alpha channel encoding validity (overlay transparency is applied at + blit time via painter opacity, so it never invalidates the cache).""" + cached = self.__heatmap_cache.get(cache_key) + if cached is not None: + return cached[0] + + built = self._build_heatmap_image(grid.n_x, grid.n_y, values) + if built is None: + return None + + self.__heatmap_cache[cache_key] = built + return built[0] + + def _build_heatmap_image( + self, + n_x: int, + n_y: int, + values: List[float] | List[int], + ) -> tuple[QImage, np.ndarray] | None: + if n_x <= 0 or n_y <= 0: + return None + + count = n_x * n_y + + # Cell values are laid out row-major: idx = x + y * n_x (matches the + # per-cell loop and load_image()/get_grid_coord() hit-testing). + flat = np.array( + [np.nan if v is None else v for v in values[:count]], + dtype=np.float64, + ) + if flat.size < count: + flat = np.concatenate([flat, np.full(count - flat.size, np.nan)]) + + valid = np.isfinite(flat) & (flat >= 0) + if valid.any(): + vmin = float(flat[valid].min()) + vmax = float(flat[valid].max()) + else: + vmin, vmax = 0.0, 1.0 + diff = (vmax - vmin) or 1.0 + + norm = np.clip((flat - vmin) / diff, 0.0, 1.0) + norm[~valid] = 0.0 + + rgba = np.zeros((count, 4), dtype=np.uint8) + rgba[:, :3] = viridis_rgb_array(norm) + rgba[:, 3] = np.where(valid, 255, 0).astype(np.uint8) + + # (n_y rows, n_x cols, RGBA); contiguous so QImage can share the buffer. + rgba = np.ascontiguousarray(rgba.reshape(n_y, n_x, 4)) + image = QImage( + rgba.data, n_x, n_y, 4 * n_x, QImage.Format.Format_RGBA8888 + ) + return image, rgba + + def _draw_completed_heatmap( + self, + painter: QPainter, + elem: CompletedRasterGridElem, + values: List[float] | List[int], + alpha: int, + visible_rect: QRectF | None, + cache_key: tuple, + ) -> None: + """Draw a completed grid by blitting its cached heatmap bitmap, scaled to + the current grid bounds. The bounds follow the sample geometry (x/y/z) + and the view transform handles zoom, so a move/zoom only changes the + destination rect - no per-cell repaint.""" + grid = elem.request + if not self._is_grid_visible(grid): + return + if alpha < 0 or alpha > 255: + return + + bounds = self._grid_bounds_rect(grid) + if bounds is None: + return + + if ( + visible_rect is not None + and not visible_rect.isEmpty() + and not bounds.intersects(visible_rect) + ): + return + + image = self._heatmap_image(cache_key, grid, values) + if image is None or image.isNull(): + # Fall back to per-cell rendering if the bitmap could not be built. + self._draw_grid(painter, grid, values, alpha, visible_rect) + return + + painter.save() + painter.setRenderHint(QPainter.RenderHint.Antialiasing, False) + # Nearest-neighbour scaling keeps cells crisp and blocky like the + # original fillRect rendering. + painter.setRenderHint(QPainter.RenderHint.SmoothPixmapTransform, False) + painter.setOpacity(alpha / 255.0) + painter.drawImage(bounds, image) + painter.setOpacity(1.0) + + painter.setPen(QPen(QColor(114, 159, 207, min(255, alpha + 40)), 1, Qt.PenStyle.SolidLine)) + painter.setBrush(Qt.BrushStyle.NoBrush) + painter.drawRect(bounds) + painter.restore() #TODO make sure draw_grid is visualising the grid correctly, correct orientation, correct x/y labelling!!!! def _draw_grid( @@ -697,10 +842,12 @@ class RasterGridManager(QObject): @Slot(RasterGridMetric) def metric(self, r: RasterGridMetric): self.__metric = r + self._invalidate_heatmap_cache() @Slot() def clear_completed_grids(self): self.__completed_grids.clear() + self._invalidate_heatmap_cache() self.completed_grid_updated.emit() @Slot(CompletedRasterGrid) @@ -711,6 +858,7 @@ class RasterGridManager(QObject): last_raster = raster self.__active_grid.n_x = 0 self.__active_grid.n_y = 0 + self._invalidate_heatmap_cache() if last_raster is not None and last_raster.result.file_prefix is not None: com = last_raster.centre_of_mass logger.info(f"COM: {com}") @@ -728,6 +876,7 @@ class RasterGridManager(QObject): def completed_grid_scan_delete(self, row: int): if 0 <= row < len(self.__completed_grids): del self.__completed_grids[row] + self._invalidate_heatmap_cache() self.completed_grid_updated.emit() @Slot(int)