raster_grid_manager.py: fixed lag with large grid, now generates a single bitmap

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)