mx/AareDAQ
0
1
Fork 0
Code Issues 52 Pull Requests 1 Actions Packages Projects Releases 65 Wiki Activity

logging: attempted to standardise logging, in particular timing of functions with a wrapper, focus on raster and timing at the moment.

Browse Source
This commit is contained in:
appleb_m
2026-04-24 14:53:09 +02:00
parent c73075fb7f
commit b8ee008477
5 changed files with 476 additions and 320 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aare/common/find_xtal.py
+4 -1
View File
@@ -4,11 +4,13 @@ import numpy as np
from scipy import ndimage
from aare.common.models import CrystalSize
from aare.common.logger_events import log_timing
from aare.common.raster_grid import RasterGridRequest, CenterOfMassModel
from aare.common.logger_config import setup_logger
logger = setup_logger('aareDAQ')
@log_timing(logger, "Identify crystal raster")
def identify_crystal_raster(result, r: RasterGridRequest) -> CenterOfMassModel | None:
images = result.images
if images and any(getattr(img, "spots_low_res", 0) for img in images):
@@ -51,6 +53,7 @@ def identify_crystal_raster(result, r: RasterGridRequest) -> CenterOfMassModel |
else:
return None
@log_timing(logger, "Rebuild array from scan results")
def rebuild_array_from_scan_results(scan_results: List,
value_field: str,
array_shape: Optional[tuple] = None,
@@ -74,7 +77,7 @@ def rebuild_array_from_scan_results(scan_results: List,
if nx is None or ny is None:
continue
positions.append((int(nx), int(ny))) # Note: (row, col) = (ny, nx)
positions.append((int(ny), int(nx))) # Corrected: (row, col) = (ny, nx)
if not value:
value = 0.0
values.append(float(value))
src/aare/common/logger_events.py
+158
View File
@@ -0,0 +1,158 @@
import functools
import time
from typing import Any, Callable
import logging
from aare.common.raster_grid import RasterGridRequest
from aare.common.rotation_scan import RotationScanRequest
def log_timing(
logger: logging.Logger,
message_prefix: str = "",
level: int = logging.DEBUG,
extra: dict[str, Any] | None = None,
) -> Callable:
"""Decorator to time a function call and log the duration."""
def decorator(func: Callable) -> Callable:
@functools.wraps(func)
def wrapper(*args: Any, **kwargs: Any) -> Any:
start = time.perf_counter()
prefix = f"{message_prefix}: " if message_prefix else ""
logger.log(level, f"{prefix}Starting {func.__name__}")
try:
result = func(*args, **kwargs)
duration = time.perf_counter() - start
logger.log(
level,
f"{prefix}Finished {func.__name__} in {duration:.4f} seconds",
extra=merge_log_context(extra, duration_s=duration)
)
return result
except Exception as e:
duration = time.perf_counter() - start
logger.log(
level,
f"{prefix}{func.__name__} FAILED after {duration:.4f} seconds with error: {e}",
extra=merge_log_context(extra, duration_s=duration)
)
raise
return wrapper
return decorator
def merge_log_context(*parts: dict[str, Any] | None, **extra: Any) -> dict[str, Any]:
payload: dict[str, Any] = {}
for part in parts:
if part:
payload.update(part)
payload.update(extra)
return payload
def sample_log_context(sample) -> dict[str, Any]:
return {
"sample_id": getattr(sample, "db_id", None),
"sample_name": getattr(sample, "sample_name", None),
}
def raster_request_log_context(request: RasterGridRequest | None) -> dict[str, Any]:
if request is None:
return {}
return {
"file_prefix": request.file_prefix,
"omega_deg": request.omega_deg,
"dtz": request.dtz,
"transmission": request.transmission,
"n_x": request.n_x,
"n_y": request.n_y,
"grid_size_x_mm": getattr(request.grid_size_mm, "x", None),
"grid_size_y_mm": getattr(request.grid_size_mm, "y", None),
}
def rotation_request_log_context(
request: RotationScanRequest | None,
*,
total_time_s: float | None = None,
) -> dict[str, Any]:
if request is None:
return {}
payload = {
"file_prefix": request.file_prefix,
"start_omega_deg": request.start_omega_deg,
"dtz": request.dtz,
"exp_time_s": request.exp_time_s,
"incr_omega_deg": request.incr_omega_deg,
"steps": request.steps,
"transmission": request.transmission,
"screening": getattr(request, "screening", None),
}
if total_time_s is not None:
payload["total_time_s"] = total_time_s
return payload
def geom_log_context(geom, *, prefix: str = "") -> dict[str, Any]:
smargon = getattr(geom, "smargon", None)
sh_mm = getattr(smargon, "sh_mm", None)
return {
f"{prefix}omega_deg": getattr(geom, "omega_deg", None),
f"{prefix}beam_x_pxl": getattr(getattr(geom, "beam_location_pxl", None), "x", None),
f"{prefix}beam_y_pxl": getattr(getattr(geom, "beam_location_pxl", None), "y", None),
f"{prefix}pixel_in_mm": getattr(geom, "pixel_in_mm", None),
f"{prefix}sh_x_mm": getattr(sh_mm, "x", None),
f"{prefix}sh_y_mm": getattr(sh_mm, "y", None),
f"{prefix}sh_z_mm": getattr(sh_mm, "z", None),
}
def ml_bundle_meta_log_context(
*,
target_point: tuple[float, float] | None = None,
focus: float | None = None,
) -> dict[str, Any]:
return {
"target_point": target_point,
"focus": focus,
}
def log_ml_bundle_meta(
logger: logging.Logger,
context: str,
*,
target_point: tuple[float, float] | None = None,
focus: float | None = None,
) -> None:
if target_point is None and focus is None:
return
logger.debug(
"ML bundle metadata",
extra=merge_log_context(
{"context": context},
ml_bundle_meta_log_context(target_point=target_point, focus=focus),
),
)
def log_duration(
logger: logging.Logger,
message: str,
duration_s: float,
*,
level: int = logging.INFO,
extra: dict[str, Any] | None = None,
) -> None:
logger.log(
level,
message,
extra=merge_log_context(extra, duration_s=duration_s),
)
src/aare/daq/aaredb.py
+10 -26
View File
@@ -25,6 +25,7 @@ from aareDB import (
from aare.common.coordinate import Coordinate
from aare.common.logger_config import setup_logger
from aare.common.logger_events import log_timing
from aare.common.models import (
SampleShortInfo,
PuckLoadedInfo,
@@ -40,23 +41,6 @@ from jfjoch_client.models import ScanResult
logger = setup_logger("aareDAQ")
def time_db_call(func):
"""Decorator to time AareDB calls and log the duration."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
start = time.perf_counter()
logger.debug(f"Starting AareDB call: {func.__name__}")
try:
result = func(*args, **kwargs)
duration = time.perf_counter() - start
logger.debug(f"Finished AareDB call: {func.__name__} in {duration:.4f} seconds")
return result
except Exception as e:
duration = time.perf_counter() - start
logger.debug(f"AareDB call: {func.__name__} FAILED after {duration:.4f} seconds with error: {e}")
raise
return wrapper
class AareWrapper:
def __init__(
@@ -91,7 +75,7 @@ class AareWrapper:
self.__cert_file = configuration.cert_file
self.__key_file = configuration.key_file
@time_db_call
@log_timing(logger, "AareDB call")
def set_pucks_beamline(self, input_list: List[PuckLoadedInfo]):
o = []
@@ -108,7 +92,7 @@ class AareWrapper:
)
logger.debug(ret)
@time_db_call
@log_timing(logger, "AareDB call")
def create_manual_sample(self, s: SampleShortInfo):
from aareDB.models import ManualSampleCreate
@@ -123,7 +107,7 @@ class AareWrapper:
except Exception as e:
logger.error(f"Error inserting sample: {e}")
@time_db_call
@log_timing(logger, "AareDB call")
def send_sample_event(
self,
s: Optional[SampleShortInfo],
@@ -149,7 +133,7 @@ class AareWrapper:
except Exception as e:
logger.error(f"Error sending sample event {event_type!s} to db: {e}")
@time_db_call
@log_timing(logger, "AareDB call")
def upload_image(self, sample_id: int, filename: str, bgr_image: np.ndarray, message: Optional[str] = None):
_, buffer = cv2.imencode('.jpg', bgr_image)
jpeg_bytes = io.BytesIO(buffer)
@@ -170,7 +154,7 @@ class AareWrapper:
response = requests.post(url, **request_kwargs)
logger.debug(f"Response status code: {response.status_code}")
@time_db_call
@log_timing(logger, "AareDB call")
def upload_jpg(self, sample_id: int, filename: str, jpg_image):
logger.debug(f"jppg_image of type: {type(jpg_image)}")
url = f"{self.__host}/protected_router/sample_runner/{sample_id}/upload-images"
@@ -185,7 +169,7 @@ class AareWrapper:
headers=headers)
logger.debug(f"Response status code: {response.status_code}")
@time_db_call
@log_timing(logger, "AareDB call")
def create_rotation_run(self, s: Optional[SampleShortInfo], r:RotationScanRequest, d:DAQStatusModel):
if s is None:
return
@@ -261,7 +245,7 @@ class AareWrapper:
except Exception as e:
logger.error(e)
@time_db_call
@log_timing(logger, "AareDB call")
def create_gridscan_run(self, s: Optional[SampleShortInfo], r:RasterGridRequest, d:DAQStatusModel):
if s is None:
return
@@ -326,7 +310,7 @@ class AareWrapper:
except Exception as e:
logger.debug(e)
@time_db_call
@log_timing(logger, "AareDB call")
def ingest_gridscan(self, sample: Optional[SampleShortInfo], raster_result: ScanResult,
raster_request: RasterGridRequest, geom: SampleGeometryModel,
com: Optional[CenterOfMassModel], beam_mark_pxl:tuple[float,float]):
@@ -397,7 +381,7 @@ class AareWrapper:
logger.error(e)
raise e
@time_db_call
@log_timing(logger, "AareDB call")
def ingest_scan(self, sample: Optional[SampleShortInfo], result: ScanResult,
geom: SampleGeometryModel, beam_mark_pxl:tuple[float,float]):
src/aare/daq/daq.py
+302 -293
View File
@@ -26,6 +26,16 @@ from aare.common.beamline import MXBeamline
from aare.common.coordinate import Coordinate, SmargonCoordinate, AerotechCoordinate
from aare.common.diffraction_geometry import DiffractionGeometry
from aare.common.logger_config import setup_logger
from aare.common.logger_events import (
geom_log_context,
log_duration,
log_timing,
log_ml_bundle_meta,
merge_log_context,
raster_request_log_context,
rotation_request_log_context,
sample_log_context,
)
from aare.common.models import (
SampleShortInfo,
PuckLoadedInfo,
@@ -205,19 +215,6 @@ class AareDAQ:
progress.success = success
self._emit_automation_progress(progress)
def _log_ml_bundle_meta(
self,
context: str,
*,
target_point: tuple[float, float] | None = None,
focus: float | None = None,
) -> None:
if target_point is None and focus is None:
return
logger.debug(
f"ML bundle metadata for {context}: target_point={target_point}, focus={focus}"
)
#--------------------------------------------
# Operation Handlers
#--------------------------------------------
@@ -381,19 +378,11 @@ class AareDAQ:
try:
logger.info(
"Starting raster sequence",
extra={
"auto_center": auto_center,
"file_prefix": grid_request.file_prefix,
"omega_deg": grid_request.omega_deg,
"dtz": grid_request.dtz,
"transmission": grid_request.transmission,
"n_x": grid_request.n_x,
"n_y": grid_request.n_y,
"grid_size_x_mm": getattr(grid_request.grid_size_mm, "x", None),
"grid_size_y_mm": getattr(grid_request.grid_size_mm, "y", None),
"sample_id": getattr(self.sample, "db_id", None),
"sample_name": getattr(self.sample, "sample_name", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(grid_request),
{"auto_center": auto_center},
),
)
if auto_center:
@@ -404,12 +393,10 @@ class AareDAQ:
if result is None:
logger.error(
"Raster sequence returned no result after auto-centering",
extra={
"file_prefix": grid_request.file_prefix,
"omega_deg": grid_request.omega_deg,
"sample_id": getattr(self.sample, "db_id", None),
"sample_name": getattr(self.sample, "sample_name", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(grid_request),
),
)
else:
self.__setup_datacollection(request=grid_request)
@@ -421,12 +408,14 @@ class AareDAQ:
if result is not None:
logger.info(
"Raster sequence completed",
extra={
"auto_center": auto_center,
"file_prefix": grid_request.file_prefix,
"result_count": len(result.r) if hasattr(result, "r") and result.r is not None else None,
"sample_id": getattr(self.sample, "db_id", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(grid_request),
{
"auto_center": auto_center,
"result_count": len(result.r) if hasattr(result, "r") and result.r is not None else None,
},
),
)
return result
@@ -434,12 +423,10 @@ class AareDAQ:
except JFJochCommunicationError as e:
logger.exception(
"Raster sequence failed due to JFJoch communication error",
extra={
"file_prefix": grid_request.file_prefix,
"omega_deg": grid_request.omega_deg,
"sample_id": getattr(self.sample, "db_id", None),
"sample_name": getattr(self.sample, "sample_name", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(grid_request),
),
)
self._handle_operation_error(
operation=DAQOperation.RASTER,
@@ -452,17 +439,11 @@ class AareDAQ:
except Exception as e:
logger.exception(
"Raster sequence failed",
extra={
"auto_center": auto_center,
"file_prefix": grid_request.file_prefix,
"omega_deg": grid_request.omega_deg,
"dtz": grid_request.dtz,
"transmission": grid_request.transmission,
"n_x": grid_request.n_x,
"n_y": grid_request.n_y,
"sample_id": getattr(self.sample, "db_id", None),
"sample_name": getattr(self.sample, "sample_name", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(grid_request),
{"auto_center": auto_center},
),
)
self._handle_operation_error(
operation=DAQOperation.RASTER,
@@ -1017,18 +998,17 @@ class AareDAQ:
logger.info(
"Starting auto-center raster workflow",
extra={
"sample_id": getattr(sample, "db_id", None),
"sample_name": getattr(sample, "sample_name", None),
"omega_deg": geom.omega_deg,
"current_sh_x_mm": geom.smargon.sh_mm.x,
"current_sh_y_mm": geom.smargon.sh_mm.y,
"current_sh_z_mm": geom.smargon.sh_mm.z,
"beam_x_pxl": geom.beam_location_pxl.x,
"beam_y_pxl": geom.beam_location_pxl.y,
"pixel_in_mm": geom.pixel_in_mm,
"file_prefix": old_prefix,
},
extra=merge_log_context(
sample_log_context(sample),
geom_log_context(geom, prefix="current_"),
{
"omega_deg": geom.omega_deg,
"beam_x_pxl": geom.beam_location_pxl.x,
"beam_y_pxl": geom.beam_location_pxl.y,
"pixel_in_mm": geom.pixel_in_mm,
"file_prefix": old_prefix,
},
),
)
r = self.__ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg:.2f}deg")
@@ -1036,11 +1016,13 @@ class AareDAQ:
if r is None:
logger.warning(
"No ML bounding box found at primary angle during auto-center raster",
extra={
"sample_id": getattr(sample, "db_id", None),
"omega_deg": geom.omega_deg,
"file_prefix": old_prefix,
},
extra=merge_log_context(
sample_log_context(sample),
{
"omega_deg": geom.omega_deg,
"file_prefix": old_prefix,
},
),
)
self.__devs.aerotech_omega = geom.omega_deg + 90.0
time.sleep(0.2)
@@ -1049,17 +1031,19 @@ class AareDAQ:
if r is not None:
logger.info(
"ML bounding box found for auto-center raster",
extra={
"sample_id": getattr(sample, "db_id", None),
"ml_omega_deg": r.omega_deg,
"ml_n_x": r.n_x,
"ml_n_y": r.n_y,
"ml_grid_size_x_mm": getattr(r.grid_size_mm, "x", None),
"ml_grid_size_y_mm": getattr(r.grid_size_mm, "y", None),
"ml_top_left_x_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "x", None),
"ml_top_left_y_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "y", None),
"ml_top_left_z_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "z", None),
},
extra=merge_log_context(
sample_log_context(sample),
{
"ml_omega_deg": r.omega_deg,
"ml_n_x": r.n_x,
"ml_n_y": r.n_y,
"ml_grid_size_x_mm": getattr(r.grid_size_mm, "x", None),
"ml_grid_size_y_mm": getattr(r.grid_size_mm, "y", None),
"ml_top_left_x_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "x", None),
"ml_top_left_y_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "y", None),
"ml_top_left_z_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "z", None),
},
),
)
self.__set_state(BeamlineStateEnum.DataCollection)
geom = self.sample_geometry
@@ -1073,15 +1057,15 @@ class AareDAQ:
logger.info(
"Running first auto-center raster",
extra={
"file_prefix": grid.file_prefix,
"omega_deg": grid.omega_deg,
"n_x": grid.n_x,
"n_y": grid.n_y,
"top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None),
"top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None),
"top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(grid),
{
"top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None),
"top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None),
"top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None),
},
),
)
res1 = self.__raster(grid)
@@ -1107,40 +1091,37 @@ class AareDAQ:
logger.info(
"Prepared second auto-center raster",
extra={
"file_prefix": grid.file_prefix,
"omega_deg": grid.omega_deg,
"top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None),
"top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None),
"top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None),
"n_x": grid.n_x,
"n_y": grid.n_y,
"grid_size_x_mm": getattr(grid.grid_size_mm, "x", None),
"grid_size_y_mm": getattr(grid.grid_size_mm, "y", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(grid),
{
"top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None),
"top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None),
"top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None),
},
),
)
logger.info(
"Running second auto-center raster",
extra={
"file_prefix": grid.file_prefix,
"omega_deg": grid.omega_deg,
"n_x": grid.n_x,
"n_y": grid.n_y,
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(grid),
),
)
res2 = self.__raster(grid)
return CompletedRasterGrid(r=[res1, res2])
else:
logger.error(
"Auto-center raster aborted because no ML bounding box was found at either angle",
extra={
"sample_id": getattr(sample, "db_id", None),
"sample_name": getattr(sample, "sample_name", None),
"primary_omega_deg": geom.omega_deg,
"secondary_omega_deg": geom.omega_deg + 90.0,
"file_prefix": old_prefix,
},
extra=merge_log_context(
sample_log_context(sample),
{
"primary_omega_deg": geom.omega_deg,
"secondary_omega_deg": geom.omega_deg + 90.0,
"file_prefix": old_prefix,
},
),
)
return None
@@ -1284,19 +1265,11 @@ class AareDAQ:
status = self.status
logger.info(
"Starting raster acquisition",
extra={
"file_prefix": request.file_prefix,
"omega_deg": request.omega_deg,
"dtz": request.dtz,
"transmission": request.transmission,
"n_x": request.n_x,
"n_y": request.n_y,
"grid_size_x_mm": getattr(request.grid_size_mm, "x", None),
"grid_size_y_mm": getattr(request.grid_size_mm, "y", None),
"sample_id": getattr(self.sample, "db_id", None),
"sample_name": getattr(self.sample, "sample_name", None),
"state": getattr(status, "state", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{"state": getattr(status, "state", None)},
),
)
total_time = request.exp_time_s * request.n_x * request.n_y + request.n_y * 0.3
@@ -1340,18 +1313,20 @@ class AareDAQ:
logger.info(
"Calculated raster centre offset",
extra={
"file_prefix": request.file_prefix,
"omega_deg": request.omega_deg,
"centre_offset_x_mm": grid_centre_offset.x,
"centre_offset_y_mm": grid_centre_offset.y,
"centre_offset_z_mm": grid_centre_offset.z,
"grid_half_width_x_mm": x,
"grid_half_height_y_mm": y,
"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),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{
"centre_offset_x_mm": grid_centre_offset.x,
"centre_offset_y_mm": grid_centre_offset.y,
"centre_offset_z_mm": grid_centre_offset.z,
"grid_half_width_x_mm": x,
"grid_half_height_y_mm": y,
"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),
},
),
)
logger.info(f"moving Smargon to grid centre offset {grid_centre_offset}")
@@ -1366,15 +1341,17 @@ class AareDAQ:
logger.info(
"Moved Smargon to raster centre",
extra={
"file_prefix": request.file_prefix,
"omega_deg": request.omega_deg,
"centre_sh_x_mm": grid_centre_smargon.sh_mm.x,
"centre_sh_y_mm": grid_centre_smargon.sh_mm.y,
"centre_sh_z_mm": grid_centre_smargon.sh_mm.z,
"centre_phi_deg": grid_centre_smargon.phi_deg,
"centre_chi_deg": grid_centre_smargon.chi_deg,
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{
"centre_sh_x_mm": grid_centre_smargon.sh_mm.x,
"centre_sh_y_mm": grid_centre_smargon.sh_mm.y,
"centre_sh_z_mm": grid_centre_smargon.sh_mm.z,
"centre_phi_deg": grid_centre_smargon.phi_deg,
"centre_chi_deg": grid_centre_smargon.chi_deg,
},
),
)
if self.__cfg.simulated_detector:
@@ -1389,15 +1366,11 @@ class AareDAQ:
if scan_result is None:
logger.error(
"JFJoch returned no ScanResult for raster",
extra={
"file_prefix": request.file_prefix,
"omega_deg": request.omega_deg,
"n_x": request.n_x,
"n_y": request.n_y,
"exp_time_s": request.exp_time_s,
"sample_id": getattr(self.sample, "db_id", None),
"sample_name": getattr(self.sample, "sample_name", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{"exp_time_s": request.exp_time_s},
),
)
sample_id = self.sample.db_id if self.sample and self.sample.db_id is not None else None
@@ -1435,12 +1408,14 @@ class AareDAQ:
logger.info(
"Raster finished",
extra={
"file_prefix": request.file_prefix,
"omega_deg": request.omega_deg,
"scan_result_is_none": scan_result is None,
"sample_id": sample_id,
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{
"scan_result_is_none": scan_result is None,
"sample_id": sample_id,
},
),
)
return CompletedRasterGridElem(
@@ -1452,32 +1427,35 @@ class AareDAQ:
except Exception as e:
logger.exception(
"Failed during raster",
extra={
"file_prefix": request.file_prefix,
"omega_deg": request.omega_deg,
"n_x": request.n_x,
"n_y": request.n_y,
"exp_time_s": request.exp_time_s,
"sample_id": getattr(self.sample, "db_id", None),
"sample_name": getattr(self.sample, "sample_name", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
{"exp_time_s": request.exp_time_s},
),
)
raise
def measure_raster(self, r: RasterGridRequest, auto_center: bool) -> CompletedRasterGrid:
def measure_raster(self, request: RasterGridRequest, auto_center: bool) -> CompletedRasterGrid:
"""
Execute a raster scan.
Args:
r: RasterGridRequest parameters for the scan.
request: RasterGridRequest parameters for the scan.
auto_center: Boolean flag indicating if this is part of an automated sequence.
Returns:
CompletedRasterGrid result.
"""
logger.info(
"Received raster scan request",
extra=merge_log_context(
sample_log_context(self.sample),
raster_request_log_context(request),
),
)
self.__cfg.try_set_busy(timeout=ceil(360))
try:
result = self._execute_raster_sequence(r, auto_center=auto_center)
result = self._execute_raster_sequence(request, auto_center=auto_center)
if result is None:
raise RasterScanException("Raster scan failed")
@@ -1567,19 +1545,10 @@ class AareDAQ:
total_time = request.exp_time_s * request.steps
logger.info(
"Received rotation scan request",
extra={
"file_prefix": request.file_prefix,
"start_omega_deg": request.start_omega_deg,
"dtz": request.dtz,
"exp_time_s": request.exp_time_s,
"incr_omega_deg": request.incr_omega_deg,
"steps": request.steps,
"transmission": request.transmission,
"screening": getattr(request, "screening", None),
"sample_id": getattr(self.sample, "db_id", None),
"sample_name": getattr(self.sample, "sample_name", None),
"total_time_s": total_time,
},
extra=merge_log_context(
sample_log_context(self.sample),
rotation_request_log_context(request, total_time_s=total_time),
),
)
self.__cfg.try_set_busy(timeout=ceil(total_time + 360))
@@ -1696,7 +1665,8 @@ class AareDAQ:
)
m = prediction_result.box
bundle_image = prediction_result.image
self._log_ml_bundle_meta(
log_ml_bundle_meta(
logger,
f"ml_bounding_box:{filename or 'unnamed'}",
target_point=prediction_result.target_point,
focus=prediction_result.focus,
@@ -1724,19 +1694,21 @@ class AareDAQ:
geom = self.sample_geometry
logger.info(
"ML bounding box selected",
extra={
"sample_id": sample_id,
"ml_image_name": filename,
"omega_deg": geom.omega_deg,
"beam_x_pxl": geom.beam_location_pxl.x,
"beam_y_pxl": geom.beam_location_pxl.y,
"pixel_in_mm": geom.pixel_in_mm,
"box_x1": x1,
"box_y1": y1,
"box_x2": x2,
"box_y2": y2,
"target_point": prediction_result.target_point,
},
extra=merge_log_context(
sample_log_context(self.sample),
{
"sample_id": sample_id,
"ml_image_name": filename,
},
geom_log_context(geom),
{
"box_x1": x1,
"box_y1": y1,
"box_x2": x2,
"box_y2": y2,
"target_point": prediction_result.target_point,
},
),
)
start_coord = geom.picture_to_smargon(Coordinate(x=x1, y=y1))
@@ -1746,19 +1718,22 @@ class AareDAQ:
logger.info(
"Converted ML bounding box to raster request",
extra={
"sample_id": sample_id,
"ml_image_name": filename,
"start_sh_x_mm": start_coord.x,
"start_sh_y_mm": start_coord.y,
"start_sh_z_mm": start_coord.z,
"grid_size_x_mm": grid_size.x,
"grid_size_y_mm": grid_size.y,
"n_x": n_x,
"n_y": n_y,
"smargon_phi_deg": geom.smargon.phi_deg,
"smargon_chi_deg": geom.smargon.chi_deg,
},
extra=merge_log_context(
sample_log_context(self.sample),
{
"sample_id": sample_id,
"ml_image_name": filename,
"start_sh_x_mm": start_coord.x,
"start_sh_y_mm": start_coord.y,
"start_sh_z_mm": start_coord.z,
"grid_size_x_mm": grid_size.x,
"grid_size_y_mm": grid_size.y,
"n_x": n_x,
"n_y": n_y,
"smargon_phi_deg": geom.smargon.phi_deg,
"smargon_chi_deg": geom.smargon.chi_deg,
},
),
)
return RasterGridRequest(
@@ -1789,7 +1764,8 @@ class AareDAQ:
)
boxes = prediction_result.predictions
bundle_image = prediction_result.image
self._log_ml_bundle_meta(
log_ml_bundle_meta(
logger,
f"ml_loop_centre_box:{filename or 'unnamed'}",
target_point=prediction_result.target_point,
focus=prediction_result.focus,
@@ -1927,6 +1903,7 @@ class AareDAQ:
return
#TODO operator function similar to mount, loop_center, raster and rotation
@log_timing(logger, "Face detection sequence")
def __face_detection_sequence(self, steps: int = 14, step_size: int = 15, face_min_ratio: float = 0.3) -> dict:
self.__set_state(BeamlineStateEnum.SampleAlignment)
self.__devs.lamp_light = 2.5
@@ -1948,21 +1925,25 @@ class AareDAQ:
logger.debug(f'moving to angle: {angle}')
rotate_time = time.perf_counter()
self.__devs.aerotech_omega = angle
logger.info(f"time to rotate 15 degrees: {time.perf_counter() - rotate_time}")
log_duration(
logger,
"Completed Aerotech move during face detection",
time.perf_counter() - rotate_time,
extra={"angle_deg": angle},
)
box_time = time.perf_counter()
prediction_result: MLBoxPredictionResult = self.__mlbox.predict(
preferred_class=(3, 0),
return_image=True,
return_bundle_meta=True
)
m = prediction_result.box
self._log_ml_bundle_meta(
log_ml_bundle_meta(
logger,
f"face_detection_angle_{angle}",
target_point=prediction_result.target_point,
focus=prediction_result.focus,
)
logger.info(f"time to predict: {time.perf_counter() - box_time}")
if not m or not m.box:
logger.info(f"no box found for angle {angle}")
@@ -2116,12 +2097,12 @@ class AareDAQ:
return_image=False,
return_bundle_meta=True,
)
self._log_ml_bundle_meta(
log_ml_bundle_meta(
logger,
f"auto_center_line_scan_{omega_deg:.2f}deg",
target_point=prediction_result.target_point,
focus=prediction_result.focus,
)
target_point = prediction_result.target_point
prediction_box = prediction_result.box
@@ -2138,39 +2119,48 @@ class AareDAQ:
)
logger.info(
"Using ML target y for second auto-center raster",
extra={
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"beam_y_pxl": beam_y_pxl,
"target_y_pxl": target_y_pxl,
"y_delta_mm": y_delta_mm,
"threshold_mm": y_retarget_threshold_mm,
"target_sh_x_mm": line_scan_centre.x,
"target_sh_y_mm": line_scan_centre.y,
"target_sh_z_mm": line_scan_centre.z,
},
extra=merge_log_context(
sample_log_context(self.sample),
{
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"beam_y_pxl": beam_y_pxl,
"target_y_pxl": target_y_pxl,
"y_delta_mm": y_delta_mm,
"threshold_mm": y_retarget_threshold_mm,
"target_sh_x_mm": line_scan_centre.x,
"target_sh_y_mm": line_scan_centre.y,
"target_sh_z_mm": line_scan_centre.z,
},
),
)
else:
logger.info(
"Keeping beam-centred y line scan because ML target y shift is small",
extra={
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"beam_y_pxl": beam_y_pxl,
"target_y_pxl": target_y_pxl,
"y_delta_mm": y_delta_mm,
"threshold_mm": y_retarget_threshold_mm,
},
extra=merge_log_context(
sample_log_context(self.sample),
{
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"beam_y_pxl": beam_y_pxl,
"target_y_pxl": target_y_pxl,
"y_delta_mm": y_delta_mm,
"threshold_mm": y_retarget_threshold_mm,
},
),
)
else:
logger.info(
"No ML target point for second auto-center raster; using beam-centred line scan",
extra={
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"beam_x_pxl": beam_x_pxl,
"beam_y_pxl": beam_y_pxl,
},
extra=merge_log_context(
sample_log_context(self.sample),
{
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"beam_x_pxl": beam_x_pxl,
"beam_y_pxl": beam_y_pxl,
},
),
)
if prediction_box is not None and prediction_box.box is not None and grid_size_mm.y > 0:
@@ -2179,26 +2169,32 @@ class AareDAQ:
n_y = max(1, int(ceil(padded_height_mm / grid_size_mm.y)))
logger.info(
"Computed second auto-center raster y size from ML box height",
extra={
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"box_height_pxl": box_height_pxl,
"pixel_in_mm": geom.pixel_in_mm,
"grid_size_y_mm": grid_size_mm.y,
"padding_fraction_each_side": y_padding_fraction_each_side,
"padded_height_mm": padded_height_mm,
"computed_n_y": n_y,
},
extra=merge_log_context(
sample_log_context(self.sample),
{
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"box_height_pxl": box_height_pxl,
"pixel_in_mm": geom.pixel_in_mm,
"grid_size_y_mm": grid_size_mm.y,
"padding_fraction_each_side": y_padding_fraction_each_side,
"padded_height_mm": padded_height_mm,
"computed_n_y": n_y,
},
),
)
else:
logger.info(
"Using default y size for second auto-center raster",
extra={
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"default_n_y": default_n_y,
"has_prediction_box": prediction_box is not None and prediction_box.box is not None,
},
extra=merge_log_context(
sample_log_context(self.sample),
{
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"default_n_y": default_n_y,
"has_prediction_box": prediction_box is not None and prediction_box.box is not None,
},
),
)
offset = Coordinate(
@@ -2213,21 +2209,25 @@ class AareDAQ:
logger.info(
"Prepared second auto-center raster top-left from helper",
extra={
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"n_y": n_y,
"offset_x_mm": offset.x,
"offset_y_mm": offset.y,
"offset_z_mm": offset.z,
"top_left_x_mm": getattr(top_left.sh_mm, "x", None),
"top_left_y_mm": getattr(top_left.sh_mm, "y", None),
"top_left_z_mm": getattr(top_left.sh_mm, "z", None),
},
extra=merge_log_context(
sample_log_context(self.sample),
{
"file_prefix": file_prefix,
"omega_deg": omega_deg,
"n_y": n_y,
"offset_x_mm": offset.x,
"offset_y_mm": offset.y,
"offset_z_mm": offset.z,
"top_left_x_mm": getattr(top_left.sh_mm, "x", None),
"top_left_y_mm": getattr(top_left.sh_mm, "y", None),
"top_left_z_mm": getattr(top_left.sh_mm, "z", None),
},
),
)
return top_left, n_y
@log_timing(logger, "Loop center sequence")
def __loop_center_sequence(
self,
sample_id: Optional[int] = None,
@@ -2237,7 +2237,6 @@ class AareDAQ:
found_classes_count: dict[int, int] = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0}
try:
start = time.perf_counter()
for zoom_iter, zoom_value in enumerate([200]):
max_attempt = 2
attempt = 0
@@ -2258,7 +2257,12 @@ class AareDAQ:
logger.debug(f"Moving to new omega: {angle}")
time_to_move_aerotech = time.perf_counter()
self.__devs.aerotech_omega = angle
logger.info(f"time to move AEROTECH: {time.perf_counter() - time_to_move_aerotech}")
log_duration(
logger,
"Completed Aerotech move during loop centering",
time.perf_counter() - time_to_move_aerotech,
extra={"angle_deg": angle, "zoom": zoom_value},
)
filename = f"{sample_id}_{angle}_{zoom_value:.0f}" if sample_id is not None else None
@@ -2270,7 +2274,12 @@ class AareDAQ:
return_image=True,
return_bundle_meta=True
)
logger.info(f"time to PREDICT: {time.perf_counter() - time_to_get_pred}")
log_duration(
logger,
"Completed ML prediction during loop centering",
time.perf_counter() - time_to_get_pred,
extra={"angle_deg": angle, "zoom": zoom_value},
)
boxes = prediction_result.predictions
pred_target_point = prediction_result.target_point
@@ -2281,7 +2290,8 @@ class AareDAQ:
))
bundle_image = prediction_result.image
self._log_ml_bundle_meta(
log_ml_bundle_meta(
logger,
f"loop_center_angle_{angle}_zoom_{zoom_value:.0f}",
target_point=prediction_result.target_point,
focus=prediction_result.focus,
@@ -2319,7 +2329,12 @@ class AareDAQ:
time_to_move_smargon = time.perf_counter()
self.__devs.smargon_pos = target
self.__devs.smargon_wait(60)
logger.info(f"time to move smargon: {time.perf_counter() - time_to_move_smargon}")
log_duration(
logger,
"Completed Smargon move during loop centering",
time.perf_counter() - time_to_move_smargon,
extra={"angle_deg": angle, "zoom": zoom_value},
)
if sample_id is not None:
if trace_all_alc_moves:
self._append_smargon_trace(
@@ -2349,7 +2364,6 @@ class AareDAQ:
f"sucessfully found {targets_found_this_attempt} targets in attempt {attempt + 1} ")
break
logger.error(f"{attempt} exceeds max attempts {max_attempt}")
logger.debug(f"time to loop center: {time.perf_counter() - start}")
raise LoopCenteringFailed
#i += 1
@@ -2358,7 +2372,6 @@ class AareDAQ:
if sample_id is not None:
logger.info(f"sample {sample_id} centered")
self._append_smargon_trace(sample_id=sample_id, event="alc_success")
logger.debug(f"time to loop center: {time.perf_counter() - start}")
return LoopCenteringResult(success=True)
except Exception as e:
@@ -2384,6 +2397,7 @@ class AareDAQ:
error=e
)
@log_timing(logger, "Auto loop center")
def auto_loop_center(self, sample:Optional[SampleShortInfo]=None) -> float:
"""
Automatically center the loop using ML-based detection.
@@ -2597,6 +2611,7 @@ class AareDAQ:
end = time.perf_counter()
return end - start
@log_timing(logger, "Measure Sequence")
def measure(self, sample: SampleShortInfo) -> float:
"""
Main measurement sequence for a given sample.
@@ -2669,16 +2684,10 @@ class AareDAQ:
if raster_result is None:
logger.error(
"Raster result was None during automation",
extra={
"sample_id": getattr(sample, "db_id", None),
"sample_name": getattr(sample, "sample_name", None),
"file_prefix": raster_grid.file_prefix,
"omega_deg": raster_grid.omega_deg,
"dtz": raster_grid.dtz,
"transmission": raster_grid.transmission,
"n_x": raster_grid.n_x,
"n_y": raster_grid.n_y,
},
extra=merge_log_context(
sample_log_context(sample),
raster_request_log_context(raster_grid),
),
)
self._mark_progress_failed(progress, WorkflowStateKind.RASTER, "Raster failed")
return self._end_operation(start, DAQOperation.RASTER, error=True)
src/aare/daq/mlbox.py
+2
View File
@@ -12,6 +12,7 @@ from aare.common.aarelc_infer import AareLCInferWrapper
from aare.common.beamline import MXBeamline
from aare.common.models import MLBoxModel, MLOutputModel, MLBoxType, BoundingBoxModel
from aare.common.logger_config import setup_logger
from aare.common.logger_events import log_timing
logger=setup_logger("aareDAQ")
@@ -162,6 +163,7 @@ class MlBox:
bundle_meta = self._extract_bundle_meta(prediction)
return MLBundle(predictions=predictions, image=image, bundle_meta=bundle_meta)
@log_timing(logger, "Collect best bundle")
def _collect_best_bundle(self, attempts: int = RETRY_COUNT) -> MLBundle:
best_predictions: MLOutputModel | None = None
best_image: np.ndarray | None = None