GUI: Raster grid manager - work in progress

2025-06-11 12:02:12 +02:00
parent 6c920c8fd4
commit 228d96bf7b
2 changed files with 179 additions and 5 deletions
@@ -1,5 +1,5 @@
 import math
-from typing import Annotated
+from typing import Annotated, List

 from pydantic import Field, AfterValidator, BaseModel

@@ -11,16 +11,15 @@ def normalize_angle(angle_deg: float) -> float:
    normalized_angle = math.fmod(angle_deg + 180, 360) - 180.0
    return normalized_angle

-
 class RasterGridRequest(BaseModel):
    dtz: float | None = None
    transmission: float | None = None
    file_prefix: str | None = None
    exp_time_s: float

-    # Number of grid elements
-    n_x: Annotated[int, Field(ge=1)]
-    n_y: Annotated[int, Field(ge=1)]
+    # Number of grid elements - 0 is allowed for empty grid
+    n_x: Annotated[int, Field(ge=0)]
+    n_y: Annotated[int, Field(ge=0)]

    # Size of grid elements
    grid_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)]
@@ -56,3 +55,12 @@ class RasterGridRequest(BaseModel):
        ):
            return True
        return False
+
+
+class CompletedRasterGrid(BaseModel):
+    request: RasterGridRequest
+    bkg_estimate: List[float]
+    spot_count: List[int]
+    indexed: List[int]
+    b_factor: List[float]
+    mosaicity: List[float]
@@ -0,0 +1,166 @@
+import math
+
+from PySide6.QtCore import QObject, Signal, Slot, QPointF
+from PySide6.QtGui import QPainter, QPen, QColor, QBrush
+from PySide6.QtCore import Qt, QRect
+from typing import List, Optional
+
+from aaredaqlib.coordinate import Coordinate, SmargonCoordinate
+from aaredaqlib.models import DAQStatusModel
+from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid
+
+
+def float_to_viridis_brush(value: float, alpha: int = 127) -> QBrush:
+    """
+    Convert float value to a QBrush with Viridis colormap
+    Args:
+        value: Float between 0 and 1
+        alpha: Transparency (0-255), default 127 for 50% transparency
+    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
+    ]
+
+    # Clamp value between 0 and 1
+    value = max(0.0, min(1.0, value))
+
+    # Find the two colors to interpolate between
+    index = value * (len(viridis_colors) - 1)
+    i = int(index)
+    t = index - i
+
+    if i == len(viridis_colors) - 1:
+        color = viridis_colors[-1]
+    else:
+        # Linear interpolation between two adjacent colors
+        c1 = viridis_colors[i]
+        c2 = viridis_colors[i + 1]
+        color = tuple(int(c1[j] + (c2[j] - c1[j]) * t) for j in range(3))
+
+    return QBrush(QColor(color[0], color[1], color[2], alpha))
+
+def normalize_angle(angle_deg: float) -> float:
+    normalized_angle = math.fmod(angle_deg + 180, 360) - 180.0
+    return normalized_angle
+
+class RasterGridManager(QObject):
+    grid_updated = Signal()
+    grid_scan = Signal(RasterGridRequest)
+
+    def __init__(self, geom):
+        super().__init__()
+        self.__geom = geom
+        self.__sample = None
+        self.__active_grid : RasterGridRequest = RasterGridRequest(
+            n_x= 0,
+            n_y= 0,
+            smargon = self.__geom.smargon,
+            grid_size_mm=Coordinate(x=0.1, y=0.1),
+            omega_deg=self.__geom.omega_deg,
+            exp_time_s=1.0
+        )
+        self.__completed_grids : List[CompletedRasterGrid] = []
+
+    def is_grid_visible(self, grid: RasterGridRequest):
+        if (
+            abs(normalize_angle(grid.omega_deg - self.__geom.omega_deg)) < 1.0
+            and abs(grid.smargon.phi_deg - self.__geom.smargon.phi_deg) < 1.0
+            and abs(grid.smargon.chi_deg - self.__geom.smargon.chi_deg) < 1.0
+            and grid.n_x > 0
+            and grid.n_y > 0
+        ):
+            return True
+        return False
+
+    @Slot(DAQStatusModel)
+    def update_daq_status(self, s: DAQStatusModel):
+        self.__geom = s.geom
+
+    @Slot(QPointF, QPointF)
+    def update_raster_grid(self, start: QPointF, end: QPointF):
+
+        delta_x = (end.x() - start.x()) * self.__geom.pixel_in_mm
+        delta_y = (end.y() - start.y()) * self.__geom.pixel_in_mm
+
+        self.__active_grid.n_x = round(abs(delta_x / self.__active_grid.grid_mm.x))
+        self.__active_grid.n_y = round(abs(delta_y / self.__active_grid.grid_mm.y))
+
+        c = self.__geom.picture_to_sample(Coordinate(x=start.x(), y=start.y()))
+
+        if delta_x < 0:
+            c.x -= self.__active_grid.n_x * self.__active_grid.grid_mm.x
+        if delta_y < 0:
+            c.y -= self.__active_grid.n_y * self.__active_grid.grid_mm.y
+
+        self.__active_grid.smargon = SmargonCoordinate(
+            sh_mm=self.__geom.beamline_to_smargon(c),
+            phi_deg=self.__geom.smargon.phi_deg,
+            chi_deg=self.__geom.smargon.chi_deg,
+        )
+
+        self.__active_grid.omega_deg = self.__geom.omega_deg
+
+    @Slot(float)
+    def update_exposure_time(self, exp_time_s: float):
+        self.__active_grid.exp_time_s = exp_time_s
+
+    @Slot(float, float)
+    def update_grid_size(self, grid_size_mm_x: float, grid_size_mm_y: float):
+        if (grid_size_mm_x <= 0) or (grid_size_mm_y <= 0):
+            raise ValueError("Grid size must be positive")
+
+        self.__active_grid.n_x = round(self.__active_grid.n_x * self.__active_grid.grid_size_mm.x / grid_size_mm_x)
+        self.__active_grid.n_y = round(self.__active_grid.n_y * self.__active_grid.grid_size_mm.y / grid_size_mm_y)
+
+        self.__active_grid.grid_size_mm = Coordinate(x=grid_size_mm_x, y=grid_size_mm_y)
+
+    @Slot()
+    def run_grid_scan(self):
+        self.grid_scan.emit(self.__active_grid)
+        self.__active_grid.n_x = 0
+        self.__active_grid.n_y = 0
+
+    def draw_grid(self, painter: QPainter):
+        self._draw_grid(painter, self.__active_grid)
+        for i in self.__completed_grids:
+            self._draw_grid(painter, i.request, i.spot_count)
+
+    def _draw_grid(self, painter: QPainter, grid: RasterGridRequest, values: List[float] | List[int] | None = None):
+        if not self.is_grid_visible(grid):
+            return
+
+        painter.setPen(QPen(QColor(114, 159, 207), 2, Qt.PenStyle.SolidLine))
+
+        min_value = min((x for x in values if not math.isnan(x)), default=0)
+        max_value = max((x for x in values if not math.isnan(x)), default=1)
+
+        g = grid.grid_size_pxl(self.__geom)
+        c0 = grid.start_pxl(self.__geom)
+        for x in range(self.__active_grid.n_x):
+            for y in range(self.__active_grid.n_y):
+                pxl = x + y * self.__active_grid.n_x
+                if len(values) <= pxl or math.isnan(values[pxl]):
+                    painter.setBrush(Qt.BrushStyle.NoBrush)
+                else:
+                    painter.setBrush(QBrush(float_to_viridis_brush((values[pxl] - min_value) / (max_value - min_value), alpha=127)))
+
+                c = self.__geom.sample_to_picture(
+                    c0 + Coordinate(x=x * grid.grid_size_mm.x,
+                                    y=y * grid.grid_size_mm.y)
+                )
+                painter.drawRect(QRect(int(c.x), int(c.y), int(g.x), int(g.y)))
+
+        painter.setBrush(Qt.BrushStyle.NoBrush)