feat: forward v2 scan results decision to DB and move calculation to analysis stage

src/aare/daq/aaredb.py

@@ -10,9 +10,9 @@ import requests
 from aarecommon.config.logger import setup_logger
 from aarecommon.config.logger_events import log_timing
 from aarecommon.math.coordinate import Coordinate
-from aarecommon.math.find_xtal import compute_crystal_score_array
 from aarecommon.math.sample_geometry import SampleGeometryModel
 from aarecommon.models.beamline import MXBeamline
+from aarecommon.models.gridscan_decision import GridScanDecision
 from aarecommon.models.models import (
     DAQStatusModel,
     PuckLoadedInfo,
@@ -303,20 +303,27 @@ class AareWrapper:
         sample: SampleShortInfo | None,
         raster_result: ScanResult,
         raster_request: RasterGridRequest,
+        analysis_image_scores: list[float | None] | None,
         geom: SampleGeometryModel,
-        com: CenterOfMassModel | None,
+        com: CenterOfMassModel,
         beam_mark_pxl: tuple[float, float],
+        decision: GridScanDecision | None,
     ):
 
         if sample is None:
+            logger.error("No sample for DB gridscan ingestion request!")
             return
 
-        payload_model = self.format_gridscan_payload(
-            sample, raster_result, raster_request, geom, com, beam_mark_pxl
-        )
-        if payload_model is None:
-            return
-        payload = payload_model.model_dump()
+        payload = self.format_gridscan_payload(
+            sample,
+            raster_result,
+            raster_request,
+            geom,
+            com,
+            analysis_image_scores,
+            beam_mark_pxl,
+            decision,
+        ).model_dump_json()
 
         url = f"{self._host}/protected_router/gridscan_runner/ingest"
         headers = {
@@ -327,14 +334,13 @@ class AareWrapper:
             url,
             auth=(os.getenv("AAREDB_USERNAME"), os.getenv("AAREDB_PASSWORD")),
             headers=headers,
-            data=json.dumps(payload),
+            data=payload,
             timeout=30,
             verify=self._ssl_ca_cert,
             cert=(self._cert_file, self._key_file),
         )
         response.raise_for_status()
-
-        logger.info(f"Response status code: {response.status_code}")
+        logger.info(f"AareDB gridscan ingestion response status code: {response.status_code}")
 
     def format_gridscan_payload(
         self,
@@ -342,64 +348,40 @@ class AareWrapper:
         raster_result: ScanResult,
         raster_request: RasterGridRequest,
         geom: SampleGeometryModel,
-        com: CenterOfMassModel | None,
+        com: CenterOfMassModel,
+        analysis_image_scores: list[float | None] | None,
         beam_mark_pxl: tuple[float, float],
-    ) -> RasterPayloadModel | None:
+        decision: GridScanDecision | None,
+    ) -> RasterPayloadModel:
+        cell_size_pxl = Coordinate(
+            x=raster_request.grid_size_mm.x / geom.pixel_in_mm,
+            y=raster_request.grid_size_mm.y / geom.pixel_in_mm,
+        )
 
-        try:
-            cell_size_pxl = Coordinate(
-                x=raster_request.grid_size_mm.x / geom.pixel_in_mm,
-                y=raster_request.grid_size_mm.y / geom.pixel_in_mm,
-            )
+        if raster_request.smargon_top_left is None:
+            beam_loc = geom.smargon.sh_mm
+        else:
+            beam_loc = raster_request.smargon_top_left.sh_mm
 
-            if raster_request.smargon_top_left is None:
-                beam_loc = geom.smargon.sh_mm
-            else:
-                beam_loc = raster_request.smargon_top_left.sh_mm
+        start_pxl = geom.smargon_to_picture(beam_loc)
 
-            start_pxl = geom.smargon_to_picture(beam_loc)
+        com_grid_pxl = com.get_com_pxl(raster_request, geom)
+        center_pxl = Coordinate(x=start_pxl.x + com_grid_pxl.x, y=start_pxl.y + com_grid_pxl.y)
 
-            if com:
-                com_grid_pxl = com.get_com_pxl(raster_request, geom)
-                x = start_pxl.x + com_grid_pxl.x
-                y = start_pxl.y + com_grid_pxl.y
-                center_pxl = Coordinate(x=x, y=y)
-            else:
-                center_pxl = None
-
-            try:
-                score_arr = compute_crystal_score_array(raster_result.images)
-                score = [
-                    float(score_arr[img.nx, img.ny])
-                    if img.nx is not None and img.ny is not None
-                    else None
-                    for img in raster_result.images
-                ]
-            except Exception as e:
-                logger.warning(
-                    f"raster score computation failed, sending null score: {e}", exc_info=True
-                )
-                score = None
-
-            payload = RasterPayloadModel(
-                request=raster_request,
-                result=raster_result,
-                sample_id=sample.db_id,
-                attach_image=True,
-                centre_of_mass=com,
-                raster_score=score,  # TODO: Check whether AareDB needs to change for taking this input
-                start_pxl=start_pxl,
-                center_pxl=center_pxl,
-                cell_size_pxl=cell_size_pxl,
-                beam_mark_pxl=beam_mark_pxl,
-                beam_size_mm=geom.beam_size_mm,
-            )
-
-            return payload
-
-        except Exception as e:
-            logger.error(e)
-            raise
+        return RasterPayloadModel(
+            request=raster_request,
+            result=raster_result,
+            sample_id=sample.db_id,
+            attach_image=True,
+            centre_of_mass=com,
+            raster_score=analysis_image_scores,
+            start_pxl=start_pxl,
+            center_pxl=center_pxl,
+            cell_size_pxl=cell_size_pxl,
+            beam_mark_pxl=beam_mark_pxl,
+            beam_size_mm=geom.beam_size_mm,
+            decision=decision,
+        )
 
     @log_timing(logger, "AareDB call")
     def ingest_scan(

src/aare/daq/daq.py

@@ -40,6 +40,7 @@ from aarecommon.errors.exception_handler import (
 )
 from aarecommon.math.coordinate import AerotechCoordinate, Coordinate, SmargonCoordinate
 from aarecommon.math.diffraction_geometry import DiffractionGeometry
+from aarecommon.math.find_xtal import CenterOfMassModel
 from aarecommon.math.sample_geometry import SampleGeometryModel
 from aarecommon.models.automation import (
     AutomationProgress,
@@ -63,10 +64,11 @@ from aarecommon.models.models import (
     SimpleScanParameters,
     ZoomModeEnum,
 )
-from aarecommon.models.raster_grid import CompletedRasterGrid, RasterGridRequest
+from aarecommon.models.raster_grid import CompletedRasterGrid, GridScanDecision, RasterGridRequest
 from aarecommon.models.rotation_scan import CompletedRotationScan, RotationScanRequest
 from aarecommon.models.tell import TellPhaseEnum, TellStateModel
 from aareDB import SampleEventType
+from jfjoch_client import ScanResult
 from redis.lock import Lock as RedisLock
 
 from aare.beamline_dispatch.protocols import BeamlineDispatch
@@ -214,15 +216,26 @@ class _DAQScanIngestor:
         )
 
     def ingest_gridscan(
-        self, *, sample, raster_result, raster_request, geom, com, beam_mark_pxl
+        self,
+        *,
+        sample: SampleShortInfo | None,
+        raster_result: ScanResult,
+        raster_request: RasterGridRequest,
+        analysis_image_scores: list[float | None] | None,
+        geom: SampleGeometryModel,
+        com: CenterOfMassModel,
+        beam_mark_pxl: tuple[float, float],
+        decision: GridScanDecision | None,
     ) -> None:
         self._daq._aare.ingest_gridscan(
             sample=sample,
             raster_result=raster_result,
             raster_request=raster_request,
+            analysis_image_scores=analysis_image_scores,
             geom=geom,
             com=com,
             beam_mark_pxl=beam_mark_pxl,
+            decision=decision,
         )
 
 

src/aare/daq/operations/common/services.py

@@ -1,6 +1,11 @@
 from dataclasses import dataclass
 from typing import Protocol
 
+from aarecommon.math.find_xtal import CenterOfMassModel
+from aarecommon.models.models import SampleGeometryModel, SampleShortInfo
+from aarecommon.models.raster_grid import GridScanDecision, RasterGridRequest
+from jfjoch_client import ScanResult
+
 from aare.daq.config import BeamlineStateEnum
 from aare.daq.operations.screenshot.service import ScreenshotService
 
@@ -29,7 +34,16 @@ class ScanIngestor(Protocol):
     def ingest_scan(self, *, sample, result, geom, beam_mark_pxl) -> None: ...
 
     def ingest_gridscan(
-        self, *, sample, raster_result, raster_request, geom, com, beam_mark_pxl
+        self,
+        *,
+        sample: SampleShortInfo | None,
+        raster_result: ScanResult,
+        raster_request: RasterGridRequest,
+        analysis_image_scores: list[float | None] | None,
+        geom: SampleGeometryModel,
+        com: CenterOfMassModel,
+        beam_mark_pxl: tuple[float, float],
+        decision: GridScanDecision | None,
     ) -> None: ...
 
 
@@ -87,15 +101,26 @@ class ScanIngestionService:
         )
 
     def ingest_gridscan(
-        self, *, sample, raster_result, raster_request, geom, com, beam_mark_pxl
+        self,
+        *,
+        sample: SampleShortInfo | None,
+        raster_result: ScanResult,
+        raster_request: RasterGridRequest,
+        analysis_image_scores: list[float | None] | None,
+        geom: SampleGeometryModel,
+        com: CenterOfMassModel,
+        beam_mark_pxl: tuple[float, float],
+        decision: GridScanDecision | None,
     ) -> None:
         self.ingestor.ingest_gridscan(
             sample=sample,
             raster_result=raster_result,
             raster_request=raster_request,
+            analysis_image_scores=analysis_image_scores,
             geom=geom,
             com=com,
             beam_mark_pxl=beam_mark_pxl,
+            decision=decision,
         )
 
 

src/aare/daq/operations/raster/service.py

@@ -1,9 +1,7 @@
 import copy
-import re
 import time
 from math import ceil, floor
 
-from aarecommon.config.beamline import cfg_get
 from aarecommon.config.config_models import GridscanAnalysisMode
 from aarecommon.config.logger_events import (
     geom_log_context,
@@ -45,13 +43,13 @@ from aare.daq.operations.common.ml_bounding_box import (
     get_ml_bounding_box,
 )
 from aare.daq.operations.raster.models import RasterContext
-from aare.devices.area_detector import AutoEnum
 
 
 @dataclass
 class _AnalysisResults:
     center_smargon_coord: SmargonCoordinate
     com: CenterOfMassModel
+    scores: list[float | None] | None
     v2_results: GridScanResult | None = None
 
 
@@ -59,7 +57,7 @@ class RasterService:
     def __init__(self, *, context: RasterContext, logger):
         self.ctx = context
         self.logger = logger
-        self.ingestor = self.ctx.services.ingestion or self.ctx.deps.aare
+        self.ingestor = self.ctx.deps.aare
         self.send_sample_event = (
             self.ctx.services.events.send
             if self.ctx.services.events is not None
@@ -78,6 +76,13 @@ class RasterService:
     ):
         if mode == GridscanAnalysisMode.FindXtal:
             com = raster_highest_score(scan_result.images)
+            _scores = compute_crystal_score_array(scan_result.images)
+            scores = [
+                float(_scores[img.nx, img.ny])
+                if img.nx is not None and img.ny is not None
+                else None
+                for img in scan_result.images
+            ]
             smargon_coord = self._smargon_offset_from_scan_result(scan_result, request, com)
             v2_results = None
         else:
@@ -104,7 +109,8 @@ class RasterService:
             else:
                 com = center_of_grid(request.n_x, request.n_y)
                 smargon_coord = self._smargon_offset_from_scan_result(scan_result, request, com)
-        return _AnalysisResults(smargon_coord, com, v2_results)
+            scores = None
+        return _AnalysisResults(smargon_coord, com, scores, v2_results)
 
     def _smargon_offset_from_scan_result(
         self, scan_result: ScanResult, request: RasterGridRequest, com: CenterOfMassModel
@@ -120,8 +126,8 @@ class RasterService:
                     "centre_offset_x_mm": target_coor_offset.x,
                     "centre_offset_y_mm": target_coor_offset.y,
                     "centre_offset_z_mm": target_coor_offset.z,
-                    "grid_half_width_x_mm": x,
-                    "grid_half_height_y_mm": y,
+                    "grid_half_width_x_mm": request.grid_size_mm.x / 2,
+                    "grid_half_height_y_mm": request.grid_size_mm.y / 2,
                     "top_left_x_mm": getattr(request.smargon_top_left.sh_mm, "x", None),
                     "top_left_y_mm": getattr(request.smargon_top_left.sh_mm, "y", None),
                     "top_left_z_mm": getattr(request.smargon_top_left.sh_mm, "z", None),
@@ -529,18 +535,7 @@ class RasterService:
             self.ctx.deps.devs.smargon_wait(timeout=180)
 
             self.logger.info(
-                "Moved Smargon to raster centre",
-                extra=merge_log_context(
-                    sample_log_context(self.ctx.sample),
-                    raster_request_log_context(request),
-                    {
-                        "centre_sh_x_mm": target_smargon.sh_mm.x,
-                        "centre_sh_y_mm": target_smargon.sh_mm.y,
-                        "centre_sh_z_mm": target_smargon.sh_mm.z,
-                        "centre_phi_deg": target_smargon.phi_deg,
-                        "centre_chi_deg": target_smargon.chi_deg,
-                    },
-                ),
+                f"Moved Smargon to raster centre: {analysis_results.center_smargon_coord}"
             )
 
             sample_id = (
@@ -590,7 +585,9 @@ class RasterService:
             self.ctx.deps.cfg.last_best_res = get_best_res(result_list=scan_result.images)
             self.ctx.deps.cfg.last_best_b_factor = get_best_b_factor(result_list=scan_result.images)
             return CompletedRasterGridElem(
-                request=copy.deepcopy(request), result=scan_result, centre_of_mass=com
+                request=copy.deepcopy(request),
+                result=scan_result,
+                centre_of_mass=analysis_results.com,
             )
 
         except Exception: