From f6f68f17e52ea3fd7fcefbe4f350c22b641971c2 Mon Sep 17 00:00:00 2001 From: perl_d Date: Thu, 3 Sep 2026 16:38:52 +0200 Subject: [PATCH] feat: forward v2 scan results decision to DB and move calculation to analysis stage --- src/aare/daq/aaredb.py | 110 +++++++++------------ src/aare/daq/daq.py | 17 +++- src/aare/daq/operations/common/services.py | 29 +++++- src/aare/daq/operations/raster/service.py | 14 ++- 4 files changed, 100 insertions(+), 70 deletions(-) diff --git a/src/aare/daq/aaredb.py b/src/aare/daq/aaredb.py index 70f14953..ea84ea69 100644 --- a/src/aare/daq/aaredb.py +++ b/src/aare/daq/aaredb.py @@ -13,6 +13,7 @@ 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, @@ -35,6 +36,7 @@ from aareDB import ( ) from aareDB import Detector as DetectorParameters from jfjoch_client.models import ScanResult +from numpy.typing import NDArray from pydantic import StrictInt logger = setup_logger("aareDAQ") @@ -303,20 +305,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 +336,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 +350,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( diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index a9adf94e..c9b836ea 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -41,6 +41,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, @@ -64,10 +65,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 aare.beamline_dispatch.protocols import BeamlineDispatch from aare.daq import workflows @@ -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, ) diff --git a/src/aare/daq/operations/common/services.py b/src/aare/daq/operations/common/services.py index 0b9252ef..20baf241 100644 --- a/src/aare/daq/operations/common/services.py +++ b/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, ) diff --git a/src/aare/daq/operations/raster/service.py b/src/aare/daq/operations/raster/service.py index 9c0b03ad..d157aaae 100644 --- a/src/aare/daq/operations/raster/service.py +++ b/src/aare/daq/operations/raster/service.py @@ -45,13 +45,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 +59,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 +78,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 +111,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