commit 3c27cbe75f51d8d21015943af334139c2a386362 Author: David Perl Date: Thu Jul 2 09:56:37 2026 +0200 refactor: extract common code from aaredaq diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e75b720 --- /dev/null +++ b/.gitignore @@ -0,0 +1,185 @@ +**/*_venv +**/.idea +*.log +**/__pycache__ +**/.DS_Store +**/out +**/.vscode +**/.pytest_cache +**/*.egg* + +# file writer data +**.h5 + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Output from end2end testing +tests/reference_failures/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/**/_build/ +docs/**/autodoc/ +docs/**/_autosummary/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +**.prof + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + + +# Development tools +tombi.toml +.ruff_cache/ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..baad9a4 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["setuptools>=75.6.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "aarecommon" +version = "0.1" +description = "Common components (models, calculations) for Aare packages" +readme = "README.md" +requires-python = ">=3.11" +dependencies = [ + "numpy", + "opencv-python", + "pillow" +] + +[project.optional-dependencies] +test = [ + "pytest", +] + +[[tool.uv.index]] +name = "psi" +url = "https://gitea.psi.ch/api/packages/mx/pypi/simple" + +[tool.pytest.ini_options] +testpaths = ["tests"] +norecursedirs = ["scripts", ".venv", "logs", "docs"] +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" diff --git a/src/aarecommon/__init__.py b/src/aarecommon/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/aarecommon/config/__init__.py b/src/aarecommon/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/aarecommon/config/beamline.py b/src/aarecommon/config/beamline.py new file mode 100644 index 0000000..3f83ed8 --- /dev/null +++ b/src/aarecommon/config/beamline.py @@ -0,0 +1,120 @@ +import os +from pathlib import Path +import yaml +from typing import Any, Dict +from aarecommon.models.beamline import MXBeamline + + +class BeamlineYAMLConfig: + def __init__(self): + self.beamline: MXBeamline = mx_beamline() + self.config = self._load_config() + + def _load_config(self) -> Dict[str, Any]: + config_dir = Path(__file__).parent / "config" + yaml_file = config_dir / f"{self.beamline.value.lower()}.yaml" + + if not yaml_file.exists(): + raise FileNotFoundError(f"Config file not found: {yaml_file}") + + with open(yaml_file, "r") as f: + return yaml.safe_load(f) + + def get(self, key: str, default: Any = None) -> Any: + return self.config.get(key, default) + + +def mx_beamline() -> MXBeamline: + name = os.getenv("BEAMLINE") + if name is None: + return MXBeamline.SIMULATED + name = name.strip().upper() + return MXBeamline[name] if name in MXBeamline.__members__ else MXBeamline.SIMULATED + + +def get_jfjoch_url(bl: MXBeamline) -> str: + """Centralized URL resolution for JFJoch services.""" + match bl: + case MXBeamline.X10SA: + return cfg_get("daq.hardware.jfjoch_url", "http://sls-gpu-002:8080") + case MXBeamline.X06DA: + return cfg_get("daq.hardware.jfjoch_url", "http://sls-gpu-001:8080") + case MXBeamline.SIMULATED: + return cfg_get("daq.hardware.jfjoch_url", "http://localhost:8080") + case MXBeamline.X06SA: + raise NotImplementedError("X06SA beamline not supported yet") + case _: + raise ValueError(f"unknown beamline {bl}") + + +def get_beamline_config() -> BeamlineYAMLConfig: + beamline_config = BeamlineYAMLConfig() + return beamline_config + + +def cfg_get(path: str, default: Any = None) -> Any: + """ + Hierarchical getter for YAML entries using dotted paths, e.g.: + cfg_get("endpoints.smargon_base") + cfg_get("epics.pv_prefix") + """ + beamline_config = get_beamline_config() + node = beamline_config.config + for part in path.split("."): + if not isinstance(node, dict) or part not in node: + return default + node = node[part] + return node + + +def jfjoch_url() -> str | None: + return cfg_get("shared.jfjoch.jfjoch_url") + + +def smargon_url() -> str | None: + return cfg_get("daq.hardware.smargon_url") + + +def aerotech_url() -> str | None: + return cfg_get("daq.hardware.aerotech_url") + + +def tell_url() -> str | None: + return cfg_get("daq.hardware.tell_url") + + +def bec_host() -> str | None: + return cfg_get("daq.hardware.bec_url") + + +def redis_host() -> str | None: + return cfg_get("daq.hardware.redis_url") + + +def gui_sample_camera_zmq() -> str | None: + return cfg_get("gui.cameras.sample_camera_zmq_url") + + +def gui_prediction_zmq() -> str | None: + return cfg_get("gui.cameras.prediction_zmq_url") + + +def gui_beamline_camera_addr() -> str | None: + return cfg_get("gui.cameras.beamline_camera_url") + + +def gui_gonio_camera_addr() -> str | None: + return cfg_get("gui.cameras.gonio_camera_url") + + +def gui_gonio_camera_id() -> int | None: + return cfg_get("gui.cameras.gonio_camera_id") + + +def daq_base_url() -> str | None: + return cfg_get("gui.daq.daq_url") + + +# Global instance for easy import +if __name__ == "__main__": + beamline_config = BeamlineYAMLConfig() diff --git a/src/aarecommon/config/logger.py b/src/aarecommon/config/logger.py new file mode 100644 index 0000000..f39de0b --- /dev/null +++ b/src/aarecommon/config/logger.py @@ -0,0 +1,180 @@ +import logging +import logging.config +import os +from pathlib import Path + +import yaml + + +class IgnoreSuccessfulStatusAccessFilter(logging.Filter): + # Should be disabled if debugging status calls. + # In particular if this is causing server/GUI lag + _SUPPRESSED_SUCCESS_MESSAGES = ( + '"GET /status HTTP/1.1" 200', + '"GET /sample/reference_tools HTTP/1.1" 200', + '"GET /sample/spreadsheet HTTP/1.1" 200', + '"GET /local_contact/device_state HTTP/1.1" 200', + '"POST /scan/smart_params HTTP/1.1" 200', + ) + + _SUPPRESSED_EXPECTED_403_MESSAGES = ( + '"GET /sse/face_detection HTTP/1.1" 403', + '"GET /sse/automation_progress HTTP/1.1" 403', + ) + + def filter(self, record: logging.LogRecord) -> bool: + message = record.getMessage() + return not any( + entry in message + for entry in ( + *self._SUPPRESSED_SUCCESS_MESSAGES, + *self._SUPPRESSED_EXPECTED_403_MESSAGES, + ) + ) + + +def get_uvicorn_logging_config() -> dict: + return { + "version": 1, + "disable_existing_loggers": False, + "filters": { + "ignore_successful_status_access": { + "()": "aarecommon.logger_config.IgnoreSuccessfulStatusAccessFilter", + }, + }, + "formatters": { + "default": { + "()": "uvicorn.logging.DefaultFormatter", + "fmt": "%(levelprefix)s %(message)s", + "use_colors": True, + }, + }, + "handlers": { + "default": { + "formatter": "default", + "class": "logging.StreamHandler", + "stream": "ext://sys.stdout", + }, + "access": { + "formatter": "default", + "class": "logging.StreamHandler", + "stream": "ext://sys.stdout", + "filters": ["ignore_successful_status_access"], + }, + }, + "loggers": { + "uvicorn": { + "handlers": ["default"], + "level": "DEBUG", + "propagate": False, + }, + "uvicorn.access": { + "handlers": ["access"], + "level": "INFO", + "propagate": False, + }, + }, + } + + +def get_config_path(env: str = "dev") -> Path: + module_dir = Path(__file__).parent + config_filename = f"logging_{env}.yaml" + config_path = module_dir / "logging_configs" / config_filename + + if not config_path.exists(): + raise FileNotFoundError(f"Logging config not found: {config_path}") + + return config_path + + +# TODO fix logging. +def setup_logger( + name="aareDAQ", + base_dir: str | None = f"~/tmp/mxlogs/", + config_path: str | None = None, +): + # switch to production mode using: $ APP_ENV=prod python main.py + env = os.getenv("APP_ENV", "dev") # default: dev + + if config_path is None: + config_file = get_config_path(env) + else: + config_file = Path(config_path) + + with open(config_file, "r") as f: + config = yaml.safe_load(f.read()) + + # Force base directory to be under the user's home dir + effective_base = base_dir + f"{name}" or f"/tmp/logs/mxlogs/{name}" + effective_base = os.path.abspath( + os.path.expanduser(os.path.expandvars(effective_base)) + ) + + # Ensure directories for file handlers exist + handlers = config.get("handlers", {}) + + for h in handlers.values(): + filename = h.get("filename") + if not filename: + continue + abs_filename = os.path.join(effective_base, os.path.basename(filename)) + + h["filename"] = abs_filename + log_dir = os.path.dirname(abs_filename) + if log_dir and not os.path.exists(log_dir): + os.makedirs(log_dir, exist_ok=True) + + logging.config.dictConfig(config) + logging.getLogger("redis_lock").setLevel(logging.WARNING) + logging.getLogger("urllib3").setLevel(logging.WARNING) + logging.getLogger("matplotlib.font_manager").setLevel(logging.WARNING) + logging.getLogger("aaredaq").setLevel(logging.DEBUG) # DAQ + logging.getLogger("aaregui").setLevel(logging.INFO) + return logging.getLogger(name) + + +def find_existing_formatter( + default_fmt="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + default_datefmt=None, +): + root = logging.getLogger() + for h in root.handlers: + if h.formatter: + return h.formatter + + for logger_name, logger in logging.Logger.manager.loggerDict.items(): + if isinstance(logger, logging.Logger): + for h in logger.handlers: + if h.formatter: + return h.formatter + + return logging.Formatter(default_fmt, datefmt=default_datefmt) + + +def attach_to_logger(logger_name: str = "", handler: logging.Handler = None): + logging.getLogger(logger_name).addHandler(handler) + + +if __name__ == "__main__": + os.environ["APP_ENV"] = "dev" # or "prod" + # os.environ["LOG_DIR"] = "/tmp/mxlogs" # optional; matches your setup code + + log = setup_logger( + "aareDAQ", base_dir="/tmp/test_logs" + ) # use the configured logger name + + log.debug("debug smoke test") + log.error("error smoke test") + log.info("info smoke test") + log.warning("warning smoke test") + log.critical("critical smoke test") + + for h in log.handlers: + if hasattr(h, "flush"): + h.flush() + + for h in log.handlers: + log.debug("debug test") + log.error("error test") + print(type(h), getattr(h, "baseFilename", None), h.level) diff --git a/src/aarecommon/config/logger_events.py b/src/aarecommon/config/logger_events.py new file mode 100644 index 0000000..4171b0f --- /dev/null +++ b/src/aarecommon/config/logger_events.py @@ -0,0 +1,165 @@ +import functools +import logging +import time +from typing import Any, Callable + +from aarecommon.models.raster_grid import RasterGridRequest +from aarecommon.models.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), + ) diff --git a/src/aarecommon/config/logging_configs/logging_dev.yaml b/src/aarecommon/config/logging_configs/logging_dev.yaml new file mode 100644 index 0000000..571696b --- /dev/null +++ b/src/aarecommon/config/logging_configs/logging_dev.yaml @@ -0,0 +1,45 @@ +version: 1 +disable_existing_loggers: False + +formatters: + detailed: + format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + +handlers: + console: + class: logging.StreamHandler + level: DEBUG + formatter: detailed + stream: ext://sys.stdout + + app_file: + class: logging.handlers.RotatingFileHandler + level: DEBUG + formatter: detailed + filename: app.log + maxBytes: 1440000 #not to overflow a diskette + backupCount: 10 + encoding: utf8 + + error_file: + class: logging.handlers.RotatingFileHandler + level: ERROR + formatter: detailed + filename: errors.log + maxBytes: 1440000 + backupCount: 5 + encoding: utf8 + +loggers: + aareDAQ: + level: DEBUG + handlers: [console, app_file, error_file] + propagate: yes + aareGUI: + level: DEBUG + handlers: [ console, app_file, error_file ] + propagate: yes + +root: + level: DEBUG + handlers: [console] \ No newline at end of file diff --git a/src/aarecommon/config/logging_configs/logging_prod.yaml b/src/aarecommon/config/logging_configs/logging_prod.yaml new file mode 100644 index 0000000..d9e42fd --- /dev/null +++ b/src/aarecommon/config/logging_configs/logging_prod.yaml @@ -0,0 +1,39 @@ +version: 1 +disable_existing_loggers: False + +formatters: + simple: + format: "%(asctime)s - %(levelname)s - %(message)s" + +handlers: + app_file: + class: logging.handlers.RotatingFileHandler + level: INFO + formatter: simple + filename: logs/app.log + maxBytes: 1000000 + backupCount: 10 + encoding: utf8 + + error_file: + class: logging.handlers.RotatingFileHandler + level: ERROR + formatter: simple + filename: logs/errors.log + maxBytes: 500000 + backupCount: 5 + encoding: utf8 + +loggers: + aareDAQ: + level: DEBUG + handlers: [console, app_file, error_file] + propagate: no + aareGUI: + level: DEBUG + handlers: [ console, app_file, error_file ] + propagate: n + +root: + level: WARNING + handlers: [] diff --git a/src/aarecommon/config/logging_configs/simulated.yaml b/src/aarecommon/config/logging_configs/simulated.yaml new file mode 100644 index 0000000..ed36a09 --- /dev/null +++ b/src/aarecommon/config/logging_configs/simulated.yaml @@ -0,0 +1,7 @@ +beamline: "SIMULATED" +smargon_base: "simulated" +aerotech_base: "simulated" +zmq_camera_url: "simulated" +jfjoch_url: "simulated" +tell_url: "simulated" +redis_host: "localhost" \ No newline at end of file diff --git a/src/aarecommon/config/logging_configs/x06da.yaml b/src/aarecommon/config/logging_configs/x06da.yaml new file mode 100644 index 0000000..43aaa2e --- /dev/null +++ b/src/aarecommon/config/logging_configs/x06da.yaml @@ -0,0 +1,64 @@ +beamline_id: "X06DA" + +gui: + cameras: + sample_camera_zmq_url: "tcp://sls-gpu-003:9093" + prediction_zmq_url: "tcp://sls-gpu-003:9093" + beamline_camera_url: "x06da-axis-1.psi.ch" + secondary_beamline_camera_url: "axis-accc8e9a2995.psi.ch" + gonio_camera_url: "axis-server-es.psi.ch" + gonio_camera_id: "1" + + daq: + daq_url: "https://mx-x06da-queue-01.psi.ch" + cert_path: "/sls/x06da/misc/.cert/6d.crt" + +shared: + jfjoch: + jfjoch_url: "http://sls-gpu-001:8080" + +daq: + hardware: + smargon_url: "http://x06da-smargopolo.psi.ch:3000" + smargon_frontend_url: "http://x06da-smargopolo.psi.ch:8080" + aerotech_url: "http://mx-x06da-queue-01.psi.ch:5234" # Adjust if needed + tell_url: "http://x06da-tell.psi.ch:22222" + bec_url: "x06da-bec-001.psi.ch" + redis_url: "x06da-redis.psi.ch" + aarelc_url: "http://sls-gpu-003:9094" + dtz_safe_position: null + lens_magnification: 10 #or 5 currently + default_detector_distance_minimum: 86 + default_detector_distance_maximum: 900 + zoom_min: 1 # camera zoom travel limits (used to clamp auto-center zoom-to-fit) + zoom_max: 1000 + sam_cam: + #TODO wire this + white_balance_ratio = 1.233 + + db: + aaredb_url: "https://mx-aaredb-dmz-01.psi.ch/dispatcher" + + data_collection_settings: + default_raster_scan_settings: + exp_time_s: 0.02 + transmission: 1.0 + dtz: 150.0 + + default_rotation_settings: + exp_time_s: 0.02 + transmission: 1.0 + dtz: 150.0 + start_omega_deg: 0.0 + increment_omega_deg: 0.2 + steps: 1800 + + detector_limit_modifier: 2.0 + maximum_flux: 4e11 + + auto_raster: + grid_padding_fraction_x: 0.15 # pad the 1st grid scan by this fraction of its size per side in x (min 1 cell) + grid_padding_fraction_y: 0.15 # ... and the TOP in y; shifts smargon_top_left outward (cells before cell 0) + grid_padding_fraction_y_bottom: 0.30 # pad the BOTTOM of the grid (far end of n_y) more; defaults to grid_padding_fraction_y + include_crystal: false # extend the grid to cover crystals outside the loop box + line_scan_y_padding_fraction: 0.15 # pad the 2nd-stage vertical line scan height by this per side (10-20%) diff --git a/src/aarecommon/config/logging_configs/x06sa.yaml b/src/aarecommon/config/logging_configs/x06sa.yaml new file mode 100644 index 0000000..72d52b1 --- /dev/null +++ b/src/aarecommon/config/logging_configs/x06sa.yaml @@ -0,0 +1,41 @@ +beamline_id: "X06SA" +gui: + cameras: + sample_camera_zmq_url: "tcp://x06sa-pserv-01:9089" + prediction_zmq_url: "" + beamline_camera_url: "" + gonio_camera_url: "" + gonio_camera_id: "" + + daq: + daq_url: "https://mx-x06sa-queue-01.psi.ch" + cert_path: "/sls/x06sa/misc/.cert/6s.crt" + +shared: + jfjoch: + jfjoch_url: "" + +daq: + hardware: + smargon_url: "http://x06sa-smargopolo.psi.ch:3000" + aerotech_url: "http://x06sa-queue-01.psi.ch:5234" # Adjust if needed + tell_url: "" + bec_url: "x06sa-bec-001.psi.ch" + redis_url: "" + + db: + aaredb_url: "https://mx-aaredb-dmz-01.psi.ch/dispatcher" + + data_collection_settings: + default_raster_scan_settings: + exp_time_s: 0.02 + transmission: 1.0 + dtz: 150.0 + + default_rotation_settings: + exp_time_s: 0.02 + transmission: 1.0 + dtz: 150.0 + start_omega_deg: 0.0 + increment_omega_deg: 0.2 + steps: 1800 diff --git a/src/aarecommon/config/logging_configs/x10sa.yaml b/src/aarecommon/config/logging_configs/x10sa.yaml new file mode 100644 index 0000000..60e5c98 --- /dev/null +++ b/src/aarecommon/config/logging_configs/x10sa.yaml @@ -0,0 +1,51 @@ +beamline_id: "X10SA" + +gui: + cameras: + sample_camera_zmq_url: "tcp://sls-gpu-003:9089" #"tcp://x10sa-spark-01:9091" # + prediction_zmq_url: "tcp://sls-gpu-003:9089" #"tcp://x10sa-spark-01:9091" # + beamline_camera_url: "axis-accc8eb02488.psi.ch" + gonio_camera_url: "axis-accc8ea5e463.psi.ch" + gonio_camera_id: "1" + + daq: + daq_url: "https://mx-x10sa-queue-01.psi.ch" + cert_path: "/sls/x10sa/misc/.cert/10s.crt" + +shared: + jfjoch: + jfjoch_url: "http://sls-gpu-002:8080" + +daq: + hardware: + smargon_url: "http://x10sa-smargopolo.psi.ch:3000" + smargon_frontend_url: "http://x10sa-smargopolo.psi.ch:8080" + aerotech_url: "http://mx-x10sa-queue-01.psi.ch:5234" # Adjust if needed + tell_url: "http://PC17488:22222" + bec_url: "x10sa-bec-001.psi.ch" + redis_url: "x10sa-redis.psi.ch" + aarelc_url: "http://sls-gpu-003:9090" + dtz_safe_position: 300 + lens_magnification: 10 #or 5 currently + default_detector_distance_minimum: 150 + default_detector_distance_maximum: 900 + + db: + aaredb_url: "https://mx-aaredb-dmz-01.psi.ch/dispatcher" + + data_collection_settings: + default_raster_scan_settings: + exp_time_s: 0.01 + transmission: 1.0 + dtz: 200 + + default_rotation_settings: + exp_time_s: 0.01 + transmission: 1.0 + dtz: 200 + start_omega_deg: 0.0 + increment_omega_deg: 0.2 + steps: 1800 + + detector_distance_limit_modifier: 2.0 + maximum_flux: 1e12 diff --git a/src/aarecommon/errors/__init__.py b/src/aarecommon/errors/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/aarecommon/errors/codes.py b/src/aarecommon/errors/codes.py new file mode 100644 index 0000000..195045a --- /dev/null +++ b/src/aarecommon/errors/codes.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import re +from enum import StrEnum + + +class AuthErrorCode(StrEnum): + """ + Stable, machine-readable error codes used across the API. + + Rules: + - never rename an existing value (treat as public API) + - only add new values + """ + + # Generic / defaults + AUTHENTICATION_ERROR = "AUTHENTICATION_ERROR" + AUTHENTICATION_FAILED = "AUTHENTICATION_FAILED" + FORBIDDEN = "FORBIDDEN" + HTTP_ERROR = "HTTP_ERROR" + INTERNAL_SERVER_ERROR = "INTERNAL_SERVER_ERROR" + + # Auth/JWT + INVALID_TOKEN = "INVALID_TOKEN" + SESSION_ALREADY_ACTIVE = "SESSION_ALREADY_ACTIVE" + + # Authorization + NOT_STAFF = "NOT_STAFF" + NOT_IN_ACTIVE_PGROUP = "NOT_IN_ACTIVE_PGROUP" + NOT_BATON_HOLDER = "NOT_BATON_HOLDER" + + +class AareErrorCode(StrEnum): + """New error-code enum -- one value per concrete exception class in the + AareException hierarchy. + + Naming: SCREAMING_SNAKE_CASE that mirrors the class name (e.g. + ``TellCommunicationError`` → ``TELL_COMMUNICATION_ERROR``). + + Rules: + - never rename an existing value (treat as public API) + - only add new values + - one value per concrete exception class; family base classes do NOT get + values (they're abstract for isinstance matching only) + """ + + # Tell family + TELL_COMMUNICATION_ERROR = "TELL_COMMUNICATION_ERROR" + TELL_CONNECTION_EXCEPTION = "TELL_CONNECTION_EXCEPTION" + CRITICAL_TELL_EXCEPTION = "CRITICAL_TELL_EXCEPTION" + WARNING_TELL_EXCEPTION = "WARNING_TELL_EXCEPTION" + TELL_COMMAND_WHILE_BUSY_EXCEPTION = "TELL_COMMAND_WHILE_BUSY_EXCEPTION" + MOUNTING_FAILED = "MOUNTING_FAILED" + UNMOUNTING_FAILED = "UNMOUNTING_FAILED" + + # Other device-comm families + SMARGON_COMMUNICATION_ERROR = "SMARGON_COMMUNICATION_ERROR" + AEROTECH_COMMUNICATION_ERROR = "AEROTECH_COMMUNICATION_ERROR" + JF_JOCH_COMMUNICATION_ERROR = "JF_JOCH_COMMUNICATION_ERROR" + BEC_COMMUNICATION_ERROR = "BEC_COMMUNICATION_ERROR" + AARE_DB_COMMUNICATION_ERROR = "AARE_DB_COMMUNICATION_ERROR" + + # Beamline-state family + STATE_TRANSITION_FAILED = "STATE_TRANSITION_FAILED" + MAINTENANCE_STATE_EXCEPTION = "MAINTENANCE_STATE_EXCEPTION" + BEAMLINE_BUSY_EXCEPTION = "BEAMLINE_BUSY_EXCEPTION" + BEAMLINE_BUSY_TIMEOUT_EXCEPTION = "BEAMLINE_BUSY_TIMEOUT_EXCEPTION" + + # Data-collection family + DATA_COLLECTION_EXCEPTION = "DATA_COLLECTION_EXCEPTION" + RASTER_SCAN_EXCEPTION = "RASTER_SCAN_EXCEPTION" + + # Other automation + LOOP_CENTERING_FAILED = "LOOP_CENTERING_FAILED" + AXC_FAILED = "AXC_FAILED" + AUTO_RASTER_SAMPLE_SKIPPED = "AUTO_RASTER_SAMPLE_SKIPPED" + TRANSFORMATION_INVALID_EXCEPTION = "TRANSFORMATION_INVALID_EXCEPTION" + MAGNET_POSITION_SENSOR_ERORR = ( + "MAGNET_POSITION_SENSOR_ERORR" # NOTE: class name "Erorr" has a typo; preserved + ) + SMART_MAGNET_FAULT_EXCEPTION = "SMART_MAGNET_FAULT_EXCEPTION" + DOOR_SAFETY_ERROR = "DOOR_SAFETY_ERROR" + + # User errors + MANUAL_MOUNT_EXCEPTION = "MANUAL_MOUNT_EXCEPTION" + SAMPLE_EXCEPTION = "SAMPLE_EXCEPTION" + + # Auth errors -- usually surfaced via finer-grained AuthErrorCode, but + # these provide a class-level fallback when no specific auth code is set. + AUTHENTICATION_EXCEPTION = "AUTHENTICATION_EXCEPTION" + USER_RIGHTS_EXCEPTION = "USER_RIGHTS_EXCEPTION" + + # Fallback for unmapped exceptions caught by the bare-Exception handler + INTERNAL_ERROR = "INTERNAL_ERROR" + + +_CAMEL_BOUNDARY_1 = re.compile(r"(.)([A-Z][a-z]+)") +_CAMEL_BOUNDARY_2 = re.compile(r"([a-z0-9])([A-Z])") + + +def code_for_exception_class(cls_name: str) -> str: + """Convert a CamelCase exception class name to SCREAMING_SNAKE_CASE. + + Examples: + TellCommunicationError → TELL_COMMUNICATION_ERROR + AareDBCommunicationError → AARE_DB_COMMUNICATION_ERROR + AXCFailed → AXC_FAILED + """ + s1 = _CAMEL_BOUNDARY_1.sub(r"\1_\2", cls_name) + s2 = _CAMEL_BOUNDARY_2.sub(r"\1_\2", s1) + return s2.upper() + + +_ERROR_CODE_HELP: dict[str, str] = { + # Auth/JWT + AuthErrorCode.AUTHENTICATION_ERROR: ( + "Generic authentication problem. Usually means the request lacked valid credentials " + "(expired/invalid token, missing Authorization header, etc.)." + ), + AuthErrorCode.AUTHENTICATION_FAILED: ( + "Authentication failed during login/token creation. Typically incorrect credentials " + "or an inability to validate the user." + ), + AuthErrorCode.INVALID_TOKEN: ( + "The provided token could not be decoded/validated (bad signature, expired, malformed). " + "Re-authenticate to obtain a new token." + ), + AuthErrorCode.SESSION_ALREADY_ACTIVE: ( + "A different session currently owns control. Use “force current session” (if allowed) " + "or wait for the active session to expire/end." + ), + AuthErrorCode.NOT_BATON_HOLDER: ( + "A different session currently owns the baton. Request the baton or wait for active session to expire/end" + ), + # Authorization + AuthErrorCode.FORBIDDEN: ( + "Generic permissions failure. The user is authenticated but not allowed to perform this action." + ), + AuthErrorCode.NOT_STAFF: ( + "This action requires staff privileges. Log in with a staff account or ask staff to perform it." + ), + AuthErrorCode.NOT_IN_ACTIVE_PGROUP: ( + "You are not a member of the currently active p-group. Change p-group or use an account " + "that belongs to the active group." + ), + # Generic + AuthErrorCode.HTTP_ERROR: ( + "Generic HTTP error wrapper. The server returned an HTTPException that wasn’t mapped to a more specific code." + ), + AuthErrorCode.INTERNAL_SERVER_ERROR: ( + "Unhandled server error. Check server logs for a stack trace and context." + ), +} + + +def error_code_help(code: str) -> str | None: + """ + Return a human help message for a code string, if known. + Accepts either enum value strings or raw strings. + """ + if not code: + return None + return _ERROR_CODE_HELP.get(str(code)) + + +def export_error_code_help() -> dict[str, str]: + """ + Export help text as {"CODE": "help text", ...} + """ + return {str(k): str(v) for k, v in _ERROR_CODE_HELP.items()} + + +def export_error_codes_grouped() -> dict[str, dict[str, str]]: + """ + Export codes grouped by enum class name: + + { + "AuthErrorCode": {"INVALID_TOKEN": "INVALID_TOKEN", ...}, + "AareErrorCode": {"TELL_COMMUNICATION_ERROR": "TELL_COMMUNICATION_ERROR", ...} + } + """ + enums: tuple[type[StrEnum], ...] = (AuthErrorCode, AareErrorCode) + return {e.__name__: {c.name: str(c.value) for c in e} for e in enums} + + +def export_error_codes() -> dict[str, str]: + """ + Backwards-compatible, flat export used by older clients/tests/docs: + + {"INVALID_TOKEN": "INVALID_TOKEN", ...} + + NOTE: This intentionally exports only AuthErrorCode to avoid breaking + existing consumers that assume a flat map and/or specific keys. + """ + return {c.name: str(c.value) for c in AuthErrorCode} diff --git a/src/aarecommon/errors/exception_handler.py b/src/aarecommon/errors/exception_handler.py new file mode 100644 index 0000000..96a8ff4 --- /dev/null +++ b/src/aarecommon/errors/exception_handler.py @@ -0,0 +1,599 @@ +from __future__ import annotations + +import time +from typing import ClassVar + +from aarecommon.errors.codes import AuthErrorCode +from aarecommon.config.logger import setup_logger + +logger = setup_logger("aareDAQ") + + +class AareException(Exception): + """Root of the aare exception hierarchy. + + Class-level ``critical`` is the default; pass ``critical=...`` to the + constructor of an "optionally critical" subclass to override on a single + raise site. + """ + + critical: ClassVar[bool] = False + + def __init__(self, *args, critical: bool | None = None, **kwargs): + super().__init__(*args, **kwargs) + if critical is not None: + self.critical = critical + + +class AutomationError(AareException): + """Errors that happen during a DAQ operation; route through automation flows.""" + + +class AareUserError(AareException): + """Bad input / mode misuse; routes through the input-correction flow.""" + + +class AareAuthError(AareException): + """Authentication / authorization failures; routes through re-auth flow.""" + + +# --------------------------------------------------------------------------- +# Device-family base classes +# +# These exist purely so watchers can match a whole device family via +# ``isinstance(e, TellException)`` etc. Per the plan they are intentionally +# empty -- do not add behavior. +# --------------------------------------------------------------------------- + + +class TellException(AutomationError): + pass + + +class SmargonException(AutomationError): + pass + + +class AerotechException(AutomationError): + pass + + +class JFJochException(AutomationError): + pass + + +class BECException(AutomationError): + pass + + +class BeamlineStateException(AutomationError): + pass + + +# --------------------------------------------------------------------------- +# Concrete automation exceptions +# --------------------------------------------------------------------------- + + +class TransformationInvalidException(AutomationError): + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Transformation is not implemented", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class StateTransitionFailed(BeamlineStateException): + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Beamline state transition failed", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class MaintenanceStateException(BeamlineStateException): + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Beamline is in Maintenance state", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class DataCollectionException(AutomationError): + """Group parent for data-collection-time exceptions; also raisable directly.""" + + def __init__( + self, message: str = "Data collection failed", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class RasterScanException(DataCollectionException): + def __init__( + self, message: str = "Data collection failed", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class AutoRasterSampleSkipped(AutomationError): + """Raised when automation should skip the current sample because auto-raster is too large.""" + + +class LoopCenteringFailed(AutomationError): + def __init__( + self, + message: str = "Loop Centering did not detect a sample", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class UnmountingFailed(TellException): + """Sample failed to unmount. Optionally critical via instance flag.""" + + def __init__( + self, + message: str = "A sample was not unmounted", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class MountingFailed(TellException): + """Sample failed to mount. Optionally critical via instance flag.""" + + def __init__( + self, message: str = "A sample was not mounted", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class ManualMountException(AareUserError): + """Custom exception for manual mounting""" + + def __init__( + self, message: str = "Manual mounting failed", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class SmartMagnetFaultException(AutomationError): + """Custom exception for smart magnet fault""" + + critical: ClassVar[bool] = True + + def __init__( + self, message: str = "Smart magnet fault", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class DoorSafetyError(AutomationError): + """Hutch personnel-safety system does not permit robot motion. + + Raised before a mount/unmount when ``…-EH1-PSYS:PROHIBITED-STATE`` is not in + the prohibited state (door safety could not be activated, so the robot will + not move) or when ``…-EH1-PSYS:ALARM-STATE`` reports an active alarm. Always + critical so automation halts and the GUI shows a pop-up. + """ + + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Door safety could not be activated", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class TellCommandWhileBusyException(TellException): + """Custom exception for trying to move Tell when it is busy""" + + def __init__(self, message: str = "Tell is busy", *, critical: bool | None = None): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class TellConnectionException(TellException): + """Custom exception for connection problems""" + + critical: ClassVar[bool] = True + + def __init__( + self, message: str = "Lost connection to Tell", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class WarningTellException(TellException): + def __init__( + self, message: str = "Warning error in TELL", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class CriticalTellException(TellException): + critical: ClassVar[bool] = True + + def __init__( + self, message: str = "Critical error in TELL", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class AXCFailed(AutomationError): + def __init__( + self, + message: str = "Auto X-ray centering failed", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class BeamlineBusyException(BeamlineStateException): + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Beamline is in busy state", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class BeamlineBusyTimeoutException(BeamlineStateException): + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Beamline busy state expired during operation", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class SampleException(AareUserError): + """Requested sample could not be located. + + NOTE: not listed in the redesign hierarchy (§1.2); parented under + AareUserError pending confirmation -- "sample not found" reads as bad + input rather than a runtime failure.""" + + def __init__( + self, message: str = "Sample not found", *, critical: bool | None = None + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +# --------------------------------------------------------------------------- +# Auth exceptions +# --------------------------------------------------------------------------- + + +class AuthenticationException(AareAuthError): + def __init__( + self, + message: str = "Authentication failed.", + *, + status_code: int = 401, + headers: dict[str, str] | None = None, + code: AuthErrorCode = AuthErrorCode.AUTHENTICATION_FAILED, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + self.status_code = status_code + self.headers = headers + self.code = code + + def __str__(self) -> str: + return self.message + + +class UserRightsException(AareAuthError): + _last_log_ts_by_message: dict[str, float] = {} + _throttle_window_s = 30.0 + + def __init__( + self, + message: str = "User does not have rights to perform this action.", + *, + status_code: int = 403, + headers: dict[str, str] | None = None, + code: AuthErrorCode = AuthErrorCode.FORBIDDEN, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + self.status_code = status_code + self.headers = headers + self.code = code + + now = time.monotonic() + last_ts = self._last_log_ts_by_message.get(message, 0.0) + if (now - last_ts) >= self._throttle_window_s: + self._last_log_ts_by_message[message] = now + logger.warning(message, extra={"exception:": Exception}) + + def __str__(self) -> str: + return self.message + + +# --------------------------------------------------------------------------- +# Device communication exceptions +# --------------------------------------------------------------------------- + + +class SmargonCommunicationError(SmargonException): + """ + Raised when Smargon HTTP communication fails (connection refused, timeout, bad HTTP status, etc). + Keep the original exception in `__cause__` by using `raise ... from e`. + """ + + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Smargon communication error", + *, + endpoint: str | None = None, + base_url: str | None = None, + operation: str | None = None, + status_code: int | None = None, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + self.endpoint = endpoint + self.base_url = base_url + self.operation = operation + self.status_code = status_code + + def __str__(self) -> str: + return self.message + + +class TellCommunicationError(TellException): + """ + Raised when TELL HTTP/PShell communication fails (timeouts, connection refused, etc). + Intended to be caught centrally by FastAPI exception handlers. + """ + + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "TELL communication error", + *, + endpoint: str | None = None, + base_url: str | None = None, + operation: str | None = None, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + self.endpoint = endpoint + self.base_url = base_url + self.operation = operation + + def __str__(self) -> str: + return self.message + + +class JFJochCommunicationError(JFJochException): + """ + Raised when JFJoch HTTP/API communication fails. + Intended for scan-time fallbacks and GUI-visible alerts. + """ + + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "JFJoch communication error", + *, + operation: str | None = None, + endpoint: str | None = None, + base_url: str | None = None, + status_code: int | None = None, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + self.operation = operation + self.endpoint = endpoint + self.base_url = base_url + self.status_code = status_code + + def __str__(self) -> str: + return self.message + + +class AareDBCommunicationError(AutomationError): + """Raised when AareDB HTTPS communication fails. + + Default not critical -- DB hiccups don't always halt automation. Raise with + ``critical=True`` at call sites where a DB read failure should escalate.""" + + def __init__( + self, + message: str = "AareDB communication error", + *, + operation: str | None = None, + endpoint: str | None = None, + base_url: str | None = None, + status_code: int | None = None, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + self.operation = operation + self.endpoint = endpoint + self.base_url = base_url + self.status_code = status_code + + def __str__(self) -> str: + return self.message + + +class AerotechCommunicationError(AerotechException): + """ + Raised when Aerotech HTTP/API communication fails (connection refused, timeout, bad HTTP status, etc). + Keep the original exception in `__cause__` by using `raise ... from e`. + """ + + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Aerotech communication error", + *, + endpoint: str | None = None, + base_url: str | None = None, + operation: str | None = None, + status_code: int | None = None, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + self.endpoint = endpoint + self.base_url = base_url + self.operation = operation + self.status_code = status_code + + def __str__(self) -> str: + return self.message + + +class MagnetPositionSensorErorr(AutomationError): + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "Magnet position sensor error", + *, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + + def __str__(self) -> str: + return self.message + + +class BECCommunicationError(BECException): + critical: ClassVar[bool] = True + + def __init__( + self, + message: str = "BEC communication error", + *, + exception: Exception | None = None, + operation: str | None = None, + endpoint: str | None = None, + base_url: str | None = None, + critical: bool | None = None, + ): + super().__init__(message, critical=critical) + self.message = message + self.operation = operation + self.exception = exception + self.endpoint = endpoint + self.base_url = base_url + + def __str__(self) -> str: + return self.message diff --git a/src/aarecommon/lc_infer_wrapper.py b/src/aarecommon/lc_infer_wrapper.py new file mode 100644 index 0000000..a1bec5d --- /dev/null +++ b/src/aarecommon/lc_infer_wrapper.py @@ -0,0 +1,144 @@ +import io + +import cv2 +import numpy as np +from aarecommon.models.beamline import MXBeamline +from aarecommon.config.beamline import cfg_get, mx_beamline +from aarelcinfer_client import AuthenticatedClient +from aarelcinfer_client.api import beam, config, predictions +from aarelcinfer_client.models import LatestPredictionModel, RuntimeConfigPatchModel +from PIL import Image + + +class AareLCInferWrapper: + def __init__( + self, + bl: MXBeamline, + secret: str = "1s3ng@rd", + ): + if bl == MXBeamline.X10SA or bl == MXBeamline.X06DA: + host = cfg_get("daq.hardware.aarelc_url") + if host is None: + raise Exception("AareLCInferWrapper: AareLC URL not configured") + elif bl == MXBeamline.X06SA: + raise NotImplementedError(f"AareLCInferWrapper not implemented for {bl}") + elif bl == MXBeamline.SIMULATED: + raise NotImplementedError(f"AareLCInferWrapper not implemented for {bl}") + else: + raise Exception(f"Unknown beamline {bl}") + + self.client = AuthenticatedClient(base_url=host, api_key=secret) + self.client.headers["X-API-Key"] = secret + self._host = host + self._api_config = config + self._api_predictions = predictions + self._api_beam = beam + + def get_config(self) -> config.ConfigSnapshotResponse: + return self._api_config.get_config(self.client) + + def update_config( + self, patch: RuntimeConfigPatchModel + ) -> config.ConfigUpdateResponse: + return self._api_config.update_config(self.client, patch) + + def get_latest_prediction(self) -> LatestPredictionModel: + return self._api_predictions.get_latest_prediction(self.client) + + def get_latest_frame_png(self): + return self._api_predictions.get_latest_frame_png(self.client) + + def get_latest_prediction_bundle(self): + return self._api_predictions.get_latest_prediction_bundle(self.client) + + def send_samcam_details(self, beam_mark, beam_dimensions): + return self._api_beam.set_beam_mark(beam_mark, beam_dimensions) + + +if __name__ == "__main__": + wrapper = AareLCInferWrapper(bl=mx_beamline()) + + try: + print("Fetching config...") + config = wrapper.get_config() + print(f"Config: {config}") + print("Updating config...") + patch = RuntimeConfigPatchModel( + conf=0.4 + # infer_scale: float | None = Field(default=None, gt=0.0, le=1.0) + # skip: int | None = Field(default=None, ge=0) + # max_fps: float | None = Field(default=None, ge=0.0) + # device: str | None = None + # publish_enabled: bool | None = None + # compute_target_point: bool | None = None + # focus_enabled: bool | None = None + # focus_epics_enabled: bool | None = None + # focus_pv: str | None = None + # focus_every: int | None = Field(default=None, ge=1) + # focus_scale: float | None = Field(default=None, gt=0.0, le=1.0) + # focus_pv_min_period_ms: float | None = Field(default=None, ge=0.0) + # tracker: Literal["none", "bytetrack", "botsort"] | None = None + # pt: str | None = None + # engine: str | None = None + # task: Literal["auto", "detect", "segment"] | None = None + ) + response = wrapper.update_config(patch) + print("Config updated!") + print(f"Model reloaded: {getattr(response, 'model_reloaded', False)}") + except Exception as e: + print(f"Error updating config: {e}") + + try: + print("Fetching bundle...") + bundle = wrapper.get_latest_prediction_bundle() + + jpeg_bytes = bundle.image_jpeg + + # The client seems to already return the validated model as 'metadata' + prediction = bundle.metadata + + # Safety check: if it's still a dict for some reason, validate it; otherwise use as is + if isinstance(prediction, dict): + prediction = LatestPredictionModel.model_validate(prediction) + + print( + f"Received bundle. Image: {len(jpeg_bytes)} bytes. Detections: {len(prediction.boxes)}" + ) + + # Decode Image + img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB") + img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) + + # Draw Overlay using the model attributes + print(prediction.target_point.x, prediction.target_point.y) + print("boxes: ", prediction.boxes) + for det in prediction.boxes: + x1, y1 = int(det.x1), int(det.y1) + x2, y2 = int(det.x2), int(det.y2) + + cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText( + img, + f"{det.label} {det.conf:.2f}", + (x1, max(20, y1 - 8)), + cv2.FONT_HERSHEY_SIMPLEX, + 0.6, + (0, 255, 0), + 2, + ) + + if det.poly: + # Convert to numpy and shift coordinates from relative to absolute + pts = np.array(det.poly, dtype=np.int32) + pts[:, 0] += x1 # Shift X + pts[:, 1] += y1 # Shift Y + + pts = pts.reshape((-1, 1, 2)) + cv2.polylines(img, [pts], isClosed=True, color=(0, 0, 255), thickness=2) + + # cv2.imshow("Latest Prediction Bundle", img) + # cv2.waitKey(0) + # cv2.destroyAllWindows() + + except Exception as e: + print(f"Error processing prediction bundle: {e}") diff --git a/src/aarecommon/math/__init__.py b/src/aarecommon/math/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/aarecommon/math/autofocus.py b/src/aarecommon/math/autofocus.py new file mode 100644 index 0000000..fd821c8 --- /dev/null +++ b/src/aarecommon/math/autofocus.py @@ -0,0 +1,66 @@ +import cv2 +import numpy as np + + +def focus_measure_edges( + gray: np.ndarray, mask: np.ndarray | None = None, verbose: bool = False +) -> float: + # mild denoise (optional but usually stabilizes the curve) + gray = cv2.GaussianBlur(gray, (3, 3), 0) + + gx = cv2.Scharr(gray, cv2.CV_64F, 1, 0) + gy = cv2.Scharr(gray, cv2.CV_64F, 0, 1) + g2 = gx * gx + gy * gy + + roi = g2[mask] if mask is not None else g2.reshape(-1) + if roi.size == 0: + return 0.0 + + # threshold relative to median -> knocks out noise floor + t = float(np.median(roi) * 3.0) + strong = roi[roi > t] + + if strong.size == 0: + return 0.0 + if verbose: + mask_sum = mask.sum() if mask is not None else roi.size + print( + f"mask pixels: {mask_sum}, " + f"focus={strong.mean():.2f}" + f"strong_size={strong.size}" + ) + + return float(strong.mean()) # higher = sharper + + +def focus_measure_blob_size(gray: np.ndarray, mask: np.ndarray | None = None) -> float: + """ + Measures sharpness for a single bright blob. + Higher = sharper (smaller blob). + """ + g = gray.astype(np.float64) + + if mask is not None: + g = np.where(mask, g, 0.0) + + # Background subtraction is crucial for blob metrics + # Use a large-ish blur as background estimate (tune ksize to your scale) + bg = cv2.GaussianBlur(g, (0, 0), sigmaX=10.0, sigmaY=10.0) + s = g - bg + s[s < 0] = 0.0 + + total = float(s.sum()) + if total <= 0: + return 0.0 + + h, w = s.shape + y, x = np.mgrid[0:h, 0:w] + + cx = float((s * x).sum() / total) + cy = float((s * y).sum() / total) + + # intensity-weighted second central moment (variance) + var = float((s * ((x - cx) ** 2 + (y - cy) ** 2)).sum() / total) + + # smaller var => sharper, so invert + return float(1.0 / (var + 1e-9)) diff --git a/src/aarecommon/math/coordinate.py b/src/aarecommon/math/coordinate.py new file mode 100644 index 0000000..6f215b9 --- /dev/null +++ b/src/aarecommon/math/coordinate.py @@ -0,0 +1,122 @@ +from typing import Optional + +import numpy as np +from pydantic import BaseModel + + +class Coordinate(BaseModel): + x: float = 0.0 + y: float = 0.0 + z: float = 0.0 + + # Overload + + def __add__(self, other: "Coordinate") -> "Coordinate": + if isinstance(other, Coordinate): + return Coordinate( + x=self.x + other.x, y=self.y + other.y, z=self.z + other.z + ) + return NotImplemented + + # Overload - + def __sub__(self, other: "Coordinate") -> "Coordinate": + if isinstance(other, Coordinate): + return Coordinate( + x=self.x - other.x, y=self.y - other.y, z=self.z - other.z + ) + return NotImplemented + + # Overload * + def __mul__(self, other): + if isinstance(other, (int, float)): # Scalar multiplication + return Coordinate(x=self.x * other, y=self.y * other, z=self.z * other) + elif isinstance(other, Coordinate): # Dot product + return self.x * other.x + self.y * other.y + self.z * other.z + return NotImplemented + + # Overload / + def __truediv__(self, scalar: float) -> "Coordinate": + if isinstance(scalar, (int, float)) and scalar != 0: # Avoid division by zero + return Coordinate(x=self.x / scalar, y=self.y / scalar, z=self.z / scalar) + elif scalar == 0: + raise ValueError("Cannot divide by zero") + return NotImplemented + + def normalize(self) -> "Coordinate": + magnitude = np.sqrt(self.x**2 + self.y**2 + self.z**2) + if magnitude == 0: + raise ValueError("Cannot normalize a zero-magnitude vector.") + return self / magnitude + + def rotate(self, angle_deg: float, axis: str) -> "Coordinate": + """ + Rotate the coordinate around a specified axis by a given angle in degrees. + + :param angle_deg: The angle of rotation in degrees. + :param axis: The axis to rotate around ('x', 'y', or 'z'). + :return: A new Coordinate after rotation. + """ + angle_rad = np.radians(angle_deg) # Convert angle to radians + c = np.cos(angle_rad) + s = np.sin(angle_rad) + + if axis == "x": + # Rotate around x-axis (affects y, z) + y_new = self.y * c - self.z * s + z_new = self.y * s + self.z * c + return Coordinate(x=self.x, y=y_new, z=z_new) + elif axis == "y": + # Rotate around y-axis (affects x, z) + x_new = self.x * c + self.z * s + z_new = -self.x * s + self.z * c + return Coordinate(x=x_new, y=self.y, z=z_new) + elif axis == "z": + # Rotate around z-axis (affects x, y) + x_new = self.x * c - self.y * s + y_new = self.x * s + self.y * c + return Coordinate(x=x_new, y=y_new, z=self.z) + else: + raise ValueError("Invalid axis. Choose 'x', 'y', or 'z'.") + + +class SmargonCoordinate(BaseModel): + sh_mm: Optional[Coordinate] = None # SH coordinate in Smargon + phi_deg: Optional[float] = None + chi_deg: Optional[float] = None + + def eq(self, other: "SmargonCoordinate", tol: float) -> bool: + return ( + abs(self.sh_mm.x - other.sh_mm.x) < tol + and abs(self.sh_mm.y - other.sh_mm.y) < tol + and abs(self.sh_mm.z - other.sh_mm.z) < tol + and abs(self.phi_deg - other.phi_deg) < tol + and abs(self.chi_deg - other.chi_deg) < tol + ) + + def __eq__(self, other: object) -> bool: + if isinstance(other, SmargonCoordinate): + return self.eq(other, tol=0.1) + return NotImplemented + + +def positive_coords(value: Coordinate) -> Coordinate: + if value.x <= 0 or value.y <= 0: + raise ValueError("Coordinates must be positive") + return value + + +class AerotechCoordinate(BaseModel): + at_mm: Optional[Coordinate] = None + omega_deg: Optional[float] = None + + def eq(self, other: "AerotechCoordinate", tol: float) -> bool: + return ( + abs(self.at_mm.x - other.at_mm.x) < tol + and abs(self.at_mm.y - other.at_mm.y) < tol + and abs(self.at_mm.z - other.at_mm.z) < tol + and abs(self.omega_deg - other.omega_deg) < tol + ) + + def __eq__(self, other: object) -> bool: + if isinstance(other, AerotechCoordinate): + return self.eq(other, tol=0.01) + return NotImplemented diff --git a/src/aarecommon/math/diffraction_geometry.py b/src/aarecommon/math/diffraction_geometry.py new file mode 100644 index 0000000..146f8ad --- /dev/null +++ b/src/aarecommon/math/diffraction_geometry.py @@ -0,0 +1,87 @@ +import math +from typing import Annotated, Tuple + +from pydantic import BaseModel, Field + + +class DiffractionGeometry(BaseModel): + energy_keV: Annotated[float, Field(gt=1.0, lt=100.0)] + dtz_mm: Annotated[float, Field(gt=10.0, lt=5000.0)] + pixel_size_mm: Annotated[float, Field(ge=0.05, le=0.5)] + beam_center_pxl: Tuple[float, float] + detector_size_pxl: Tuple[int, int] + detector_description: str + detector_serial_number: str + poni_rot1_rad: float + poni_rot2_rad: float + + @property + def detector_max_radius_pxl(self): + x0 = self.detector_size_pxl[0] - self.beam_center_pxl[0] + x1 = self.beam_center_pxl[0] + y0 = self.detector_size_pxl[1] - self.beam_center_pxl[1] + y1 = self.beam_center_pxl[1] + return max(abs(x0), abs(x1), abs(y0), abs(y1)) + + @property + def detector_radius_mm(self): + return self.detector_max_radius_pxl * self.pixel_size_mm + + @property + def wavelength_angstrom(self): + return 12.398 / self.energy_keV + + @property + def max_resolution_angstrom(self): + return self.resolution_angstrom(self.dtz_mm) + + def resolution_angstrom(self, exp_dtz_mm: float) -> float: + if exp_dtz_mm <= 0: + raise ValueError(f"Detector distance must be positive {exp_dtz_mm}") + + theta = math.atan(self.detector_radius_mm / exp_dtz_mm) * 0.5 + return self.wavelength_angstrom / (2 * math.sin(theta)) + + def calc_dtz_mm(self, exp_resolution_angstrom: float) -> float: + if exp_resolution_angstrom <= 0: + raise ValueError("Resolution must be positive") + x = self.wavelength_angstrom / (2 * exp_resolution_angstrom) + if x >= 1.0 or x <= -1.0: + return 0.0 + theta = math.asin(x) + return self.detector_radius_mm / math.tan(2 * theta) + + +if __name__ == "__main__": + from aare.daq.config import BeamlineConfig + from aare.devices.jfjoch import JFJochWrapper + from aarecommon.config.beamline import mx_beamline + + bl = mx_beamline() + client = JFJochWrapper(bl) + cfg = BeamlineConfig(bl) + print(cfg) + det_cfg = client.detector() + print(det_cfg) + print(f"pixel_size: {det_cfg.pixel_size_mm:.4f} mm") + print(f"Detector height: {det_cfg.height:.2f} pixel") + print(f"Detector width: {det_cfg.width:.2f} pixel") + print(f"beam centre = ({cfg.beam_center[0]:.2f}, {cfg.beam_center[1]:.2f})") + geom = DiffractionGeometry( + energy_keV=12.4, + dtz_mm=200.0, + pixel_size_mm=det_cfg.pixel_size_mm, + beam_center_pxl=(det_cfg.width / 2, det_cfg.height / 2), + detector_size_pxl=(det_cfg.width, det_cfg.height), + detector_description=det_cfg.description, + detector_serial_number=det_cfg.serial_number, + poni_rot1_rad=-0.001396263, + poni_rot2_rad=-0.003839724, + ) + print(f"geom.max_resolution_angstrom: {geom.max_resolution_angstrom:.2f} Angstrom") + print(f"geom.calc_dtz_mm(3.9): {geom.calc_dtz_mm(3.9):.2f} mm") + + print( + f"geom.resolution_angstrom(1244.42): {geom.resolution_angstrom(1244.42):.2f} Angstrom" + ) + print(f"geom.detector_radius_mm: {geom.detector_radius_mm:.2f} mm") diff --git a/src/aarecommon/math/find_xtal.py b/src/aarecommon/math/find_xtal.py new file mode 100644 index 0000000..7a04626 --- /dev/null +++ b/src/aarecommon/math/find_xtal.py @@ -0,0 +1,702 @@ +from typing import Callable, List, Optional + +import numpy as np +from aarecommon.config.logger import setup_logger +from aarecommon.models.models import CrystalSize +from aarecommon.models.raster_grid import CenterOfMassModel, RasterGridRequest +from scipy import ndimage + +logger = setup_logger("aareDAQ") + + +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): + logger.debug(f"Find image by maximum number of low resolution spots") + max_image = max(images, key=lambda img: img.spots_low_res) + logger.debug(f"Image with maximum spots_low_res: {max_image}") + logger.debug(f"Maximum spots_low_res value: {max_image.spots_low_res}") + logger.debug( + f"Maximum image found at grid coordiantes {max_image.nx}, {max_image.ny}" + ) + logger.debug( + f"Maximum image found at umL {max_image.nx * r.grid_size_mm.x}, {max_image.ny * r.grid_size_mm.y}" + ) + com = CenterOfMassModel( + n_x=max_image.nx, n_y=max_image.ny, max_image=max_image.number + ) + com_mm = com.get_com_mm(r) + grid_mm_x = com_mm.x + grid_mm_y = com_mm.y + + logger.debug(f"Grid coordinates in mm: x={grid_mm_x}, y={grid_mm_y}") + return com + else: + return None + + +def rebuild_array_from_scan_results( + scan_results: List, + value_field: str, + array_shape: Optional[tuple] = None, + nx_field: str = "nx", + ny_field: str = "ny", + default_value: float = 0.0, + threshold: Optional[float] = None, + condition_func: Optional[Callable] = None, + apply_filter_before: bool = True, +) -> np.ndarray: + positions = [] + values = [] + + for result in scan_results: + nx = getattr(result, nx_field) + ny = getattr(result, ny_field) + value = getattr(result, value_field) + + # Skip if coordinates are None + if nx is None or ny is None: + continue + + positions.append((int(nx), int(ny))) # Note: (row, col) = (ny, nx) + if not value: + value = 0.0 + values.append(float(value)) + + if not positions: + logger.error("No valid positions found in scan results") + raise ValueError("No valid positions found in scan results") + + # Determine array shape + if array_shape is None: + max_row = max(pos[0] for pos in positions) + max_col = max(pos[1] for pos in positions) + array_shape = (max_row + 1, max_col + 1) + + # Initialize array with default values + result_array = np.full(array_shape, default_value, dtype=float) + + # Apply pre-filtering if requested + if apply_filter_before: + filtered_data = [] + for pos, val in zip(positions, values): + keep_value = True + + # Apply threshold filter + if threshold is not None and val < threshold: + keep_value = False + + # Apply custom condition + if condition_func is not None and not condition_func(val): + keep_value = False + + if keep_value: + filtered_data.append((pos, val)) + else: + filtered_data.append((pos, 0.0)) + + # Fill array with filtered values + for pos, val in filtered_data: + if 0 <= pos[0] < array_shape[0] and 0 <= pos[1] < array_shape[1]: + result_array[pos[0], pos[1]] = val + else: + # Fill array first, then apply filters + for pos, val in zip(positions, values): + if 0 <= pos[0] < array_shape[0] and 0 <= pos[1] < array_shape[1]: + result_array[pos[0], pos[1]] = val + + # Apply post-filtering + if threshold is not None: + result_array[result_array < threshold] = 0.0 + + if condition_func is not None: + mask = np.vectorize(condition_func)(result_array) + result_array[~mask] = 0.0 + + return result_array + + +def create_quality_filtered_array( + scan_results: List, + value_field: str, + min_spots: Optional[int] = None, + min_efficiency: Optional[float] = 1.0, + min_background: Optional[float] = None, + exclude_ice: Optional[bool] = True, + min_low_res_spots: Optional[float] = 10.0, + **kwargs, +) -> np.ndarray: + """ + Create array with comprehensive quality filtering + """ + + def quality_condition( + result, + max_spots_low_res, + min_background, + min_spots, + min_efficiency, + min_low_res_spots, + max_filter: float = 0.3, + ): + if exclude_ice and (result.spots_ice / max(result.spots_low_res, 1.0)) == 1.0: + # print(f"all ice for {result.number}") + return False + if min_low_res_spots and result.spots_low_res < min_low_res_spots: + return False + if result.spots_low_res < (max_spots_low_res * max_filter): + return False + if exclude_ice and result.spots_ice > result.spots * 0.8: # More than 50% ice + # print(f"more than 80% ice for {result.number}") + return False + if result.index: + # print(f"index is True for {result.number}") + return True + if min_spots and result.spots < min_spots: + # print(f"{result.spots} is less than {min_spots} for {result.number}") + return False + if result.spots_low_res < min_background: + return False + + if result.efficiency < min_efficiency: + return False + + return True + + # Filter results first + filtered_results = [] + + if min_spots is None: + min_spots = min( + (result.spots for result in scan_results if result.spots is not None), + default=1, + ) + if min_low_res_spots is None: + min_low_res_spots = min( + ( + result.spots_low_res + for result in scan_results + if result.spots_low_res is not None + ), + default=1, + ) + if min_background is None: + min_background = min( + (result.bkg for result in scan_results if result.bkg is not None), default=1 + ) + if min_efficiency is None: + min_efficiency = 1.0 + max_spots_low_res = max( + (result.spots_low_res for result in scan_results if result.spots is not None), + default=1, + ) + + for result in scan_results: + if result.nx is not None and result.ny is not None: + if quality_condition( + result, + max_spots_low_res=max_spots_low_res, + min_background=min_background, + min_spots=min_spots, + min_efficiency=min_efficiency, + min_low_res_spots=min_low_res_spots, + ): + filtered_results.append(result) + else: + # Create a copy with zero value for filtered positions + import copy + + zero_result = copy.copy(result) + setattr(zero_result, value_field, 0) + filtered_results.append(zero_result) + + return rebuild_array_from_scan_results(filtered_results, value_field, **kwargs) + + +def get_xtal_size(crystal_size, result_array, r: RasterGridRequest): + # Optional: get bounding box of the largest object + try: + labeled_array, num_objects = ndimage.label(result_array) + areas = ndimage.sum( + np.ones_like(result_array, dtype=np.int32), + labeled_array, + index=range(1, num_objects + 1), + ) + largest_idx = int(np.argmax(areas)) + 1 # +1 because labels start at 1 + largest_area = int(areas[largest_idx - 1]) + logger.info(f"Largest object label: {largest_idx}, area (px): {largest_area}") + + object_mask = labeled_array == largest_idx + rows = np.any(object_mask, axis=1) + cols = np.any(object_mask, axis=0) + row_min, row_max = np.where(rows)[0][[0, -1]] + col_min, col_max = np.where(cols)[0][[0, -1]] + logger.info( + f"Largest bbox: width={col_max - col_min}, height={row_max - row_min}" + ) + logger.info( + f"Largest bbox: width={(col_max - col_min) * r.grid_size_mm.x}, y={(row_max - row_min) * r.grid_size_mm.y}" + ) + if r.n_x == 1: + logger.debug("calculating z") + crystal_size = CrystalSize( + x=crystal_size.x, + y=crystal_size.y, + z=(col_max - col_min) * r.grid_size_mm.y * 1000, + ) + logger.info( + f"Crystal Size: x={crystal_size.x}, y={crystal_size.y}, z={crystal_size.z}" + ) + + else: + logger.debug("calculating x and y") + crystal_size = CrystalSize( + x=(row_max - row_min) * r.grid_size_mm.x * 1000, + y=(col_max - col_min) * r.grid_size_mm.y * 1000, + z=crystal_size.z, + ) + logger.info( + f"Crystal Size: x={crystal_size.x}, y={crystal_size.y}, z={crystal_size.z}" + ) + except ValueError as e: + logger.error(f"error calculating xtal size: {e}") + crystal_size = CrystalSize(x=0, y=0, z=0) + + return crystal_size + + +def get_best_b_factor(result_list: List): + if not result_list: + return None + best_b_factor = min( + (img for img in result_list if img.b is not None), + key=lambda img: img.b, + default=None, + ) + if best_b_factor is None: + return None + logger.info(f"Best b: {best_b_factor.b}") + return best_b_factor.b + + +def get_best_res(result_list: List): + if not result_list: + return None + best_res = min( + (img for img in result_list if img.res is not None), + key=lambda img: img.res, + default=None, + ) + if best_res is None: + return None + logger.info(f"Best res: {best_res.res}") + return best_res.res + + +def com_nan_check(com): + if np.isnan(com.n_x) or np.isnan(com.n_y): + return False + else: + return True + + +def get_result_list_from_com(images, com: CenterOfMassModel): + if not com_nan_check(com): + return None + cx, cy = com.n_x, com.n_y + start_x, end_x = round(cx - 1), round(cx + 1) + start_y, end_y = round(cy - 1), round(cy + 1) + logger.info(f"range x {start_x} {end_x}, y {start_y} {end_y}") + result_list = [ + img + for img in images + if start_x <= img.nx <= end_x and start_y <= img.ny <= end_y + ] + return result_list + + +def get_com_image_number(com, images): + for image in images: + if image.nx == round(com[0]) and image.ny == round(com[1]): + logger.info(f"com found for image: {image.number}") + return image.number + return None + + +def _to_com_model(coords, images, label: str) -> CenterOfMassModel | None: + """Validate (n_x, n_y) grid coords and wrap them in a CenterOfMassModel. + + Shared by raster_centre_of_mass and raster_highest_score, which differ only + in how they pick the target cell. + """ + if coords and not np.isnan(coords[0]) and not np.isnan(coords[1]): + logger.info(f"{label}: {coords}") + max_image = get_com_image_number(coords, images) + return CenterOfMassModel(n_x=coords[0], n_y=coords[1], max_image=max_image) + elif coords and np.isnan(coords[0]) and np.isnan(coords[1]): + logger.warning(f"{label} is nan: {coords[0]}, {coords[1]}") + return None + else: + logger.warning(f"No valid {label} found") + return None + + +def raster_centre_of_mass(result_array, images) -> CenterOfMassModel | None: + # grid_mm_x and grid_mm_y are relative to the top left corner of raster grid + com = ndimage.center_of_mass(result_array) + return _to_com_model(com, images, "Center of mass") + + +def raster_highest_score(images) -> CenterOfMassModel | None: + """Target the grid cell with the highest crystal score (compute_crystal_score_array).""" + score_array = compute_crystal_score_array(images) + return _to_com_model(_max_cell(score_array), images, "Highest score") + + +def has_sufficient_low_res_spots( + result_array: np.ndarray, + min_spots_low_res: float, +) -> bool: + if result_array is None or result_array.size == 0: + logger.warning("Empty result_array passed to has_sufficient_low_res_spots") + return False + + max_val = float(np.nanmax(result_array)) + logger.info(f"Max spots_low_res in raster result_array: {max_val}") + + if max_val < min_spots_low_res: + logger.info( + "Raster rejected: max spots_low_res " + f"{max_val} < required {min_spots_low_res}" + ) + return False + + return True + + +def compute_crystal_score_array( + scan_results: List, + w_bkg: float = 0.25, + w_low_res: float = 0.55, + w_indexed: float = 0.20, +) -> np.ndarray: + """Combine bkg (25%), spots_low_res (55%), and spots_indexed (20%) into a 0–100 score. + + Each field is min-max normalised to [0, 100] within the grid before weighting, + so the final score is the probability (0–100) that a pixel belongs to a crystal. + """ + arr_bkg = rebuild_array_from_scan_results(scan_results, "bkg") + arr_low = rebuild_array_from_scan_results(scan_results, "spots_low_res") + arr_idx = rebuild_array_from_scan_results(scan_results, "spots_indexed") + + def _norm(a: np.ndarray) -> np.ndarray: + mn, mx = float(a.min()), float(a.max()) + if mx == mn: + return np.zeros_like(a, dtype=float) + return (a - mn) / (mx - mn) * 100.0 + + return ( + w_bkg * _norm(arr_bkg) + w_low_res * _norm(arr_low) + w_indexed * _norm(arr_idx) + ) + + +def _max_cell(arr: np.ndarray) -> tuple[int, int]: + """Return (nx, ny) of the cell with the highest value.""" + idx = np.unravel_index(np.argmax(arr), arr.shape) + return int(idx[0]), int(idx[1]) + + +def _draw_panel( + ax, + arr: np.ndarray, + mask: np.ndarray, + label: str, + threshold: float, + grid_size_mm: Optional[tuple[float, float]] = None, + cbar_label: str = "spots_low_res", +) -> None: + """Shared helper: heatmap + crystal contour + max-cell square on one Axes. + + If grid_size_mm=(step_x_mm, step_y_mm) is provided, the crystal size in µm + is computed via get_xtal_size and shown in the panel title. + """ + import matplotlib.patches as mpatches + import matplotlib.pyplot as plt + from aarecommon.coordinate import Coordinate + + n_nx, n_ny = arr.shape + max_nx, max_ny = _max_cell(arr) + n_cells = int(mask.sum()) + + im = ax.imshow( + arr.T, + origin="lower", + cmap="viridis", + aspect="equal", + extent=[-0.5, n_nx - 0.5, -0.5, n_ny - 0.5], + ) + plt.colorbar(im, ax=ax, label=cbar_label, fraction=0.046, pad=0.04) + ax.contour( + np.arange(n_nx), + np.arange(n_ny), + mask.T.astype(float), + levels=[0.5], + colors="cyan", + linewidths=1.8, + ) + rect = mpatches.Rectangle( + (max_nx - 0.5, max_ny - 0.5), + 1, + 1, + linewidth=2, + edgecolor="red", + facecolor="none", + ) + ax.add_patch(rect) + + size_line = "" + if grid_size_mm is not None and n_cells > 0: + r = RasterGridRequest( + exp_time_s=0.0, + n_x=n_nx, + n_y=n_ny, + grid_size_mm=Coordinate(x=grid_size_mm[0], y=grid_size_mm[1]), + smargon_top_left=None, + ) + xtal_size = get_xtal_size(CrystalSize(x=0, y=0, z=0), mask.astype(float), r) + size_line = f"\n{xtal_size.x:.0f}×{xtal_size.y:.0f} µm" + + ax.set_title( + f"{label}\nthresh≈{threshold:.1f} | {n_cells} cells{size_line}", + fontsize=8, + ) + ax.set_xlabel("nx", fontsize=7) + ax.set_ylabel("ny", fontsize=7) + ax.tick_params(labelsize=6) + + +# ── Clustering / thresholding methods ───────────────────────────────────────── + + +def crystal_mask_corner_background(arr: np.ndarray) -> tuple[np.ndarray, float]: + """Threshold = max(far-corner value, floor=10). + + Simple, parameter-free. Returns too many cells when corners are zero + because any nonzero value passes. + """ + n_nx, n_ny = arr.shape + corners = [arr[0, 0], arr[n_nx - 1, 0], arr[0, n_ny - 1], arr[n_nx - 1, n_ny - 1]] + threshold = max(float(np.max(corners)), 10.0) + return arr > threshold, threshold + + +def crystal_mask_otsu_nonzero(arr: np.ndarray) -> tuple[np.ndarray, float]: + """Otsu threshold computed only on the nonzero values. + + Finds the natural gap in the signal distribution. + Ignores the large mass of background zeros so the threshold is + placed within the diffraction-signal population. + """ + from skimage.filters import threshold_otsu + + nonzero = arr[arr > 0] + if nonzero.size == 0: + return np.zeros_like(arr, dtype=bool), 0.0 + threshold = float(threshold_otsu(nonzero)) + return arr > threshold, threshold + + +def crystal_mask_mean_sigma( + arr: np.ndarray, n_sigma: float = 0.5 +) -> tuple[np.ndarray, float]: + """Threshold = mean + n_sigma * std of nonzero values. + + n_sigma=0.5 keeps cells within ~1/2 std above average signal. + Raise n_sigma to tighten the region around the strongest-diffracting core. + """ + nonzero = arr[arr > 0] + if nonzero.size == 0: + return np.zeros_like(arr, dtype=bool), 0.0 + threshold = float(nonzero.mean() + n_sigma * nonzero.std()) + return arr > threshold, threshold + + +def crystal_mask_signal_percentile( + arr: np.ndarray, percentile: float = 60.0 +) -> tuple[np.ndarray, float]: + """Threshold = given percentile of nonzero values. + + percentile=60 keeps the top 40 % of signal cells; raise to sharpen. + Unlike mean+sigma this is robust to heavy-tailed distributions. + """ + nonzero = arr[arr > 0] + if nonzero.size == 0: + return np.zeros_like(arr, dtype=bool), 0.0 + threshold = float(np.percentile(nonzero, percentile)) + return arr > threshold, threshold + + +def crystal_mask_dbscan( + arr: np.ndarray, min_signal: float = 10.0, eps: float = 1.5, min_samples: int = 3 +) -> tuple[np.ndarray, float]: + """DBSCAN spatial clustering on cells with signal > min_signal. + + Groups adjacent diffracting cells into clusters; the largest cluster + (by total signal weight) is labelled as the crystal. + eps=1.5 connects cells that are one grid step apart (including diagonal). + """ + from sklearn.cluster import DBSCAN + + ys, xs = np.where(arr > min_signal) + if len(ys) == 0: + return np.zeros_like(arr, dtype=bool), min_signal + + coords = np.column_stack([ys, xs]).astype(float) + labels = DBSCAN(eps=eps, min_samples=min_samples).fit_predict(coords) + + best_label, best_weight = -1, -1.0 + for lbl in set(labels): + if lbl == -1: + continue + weight = float(arr[ys[labels == lbl], xs[labels == lbl]].sum()) + if weight > best_weight: + best_label, best_weight = lbl, weight + + mask = np.zeros_like(arr, dtype=bool) + if best_label != -1: + sel = labels == best_label + mask[ys[sel], xs[sel]] = True + return mask, min_signal + + +def crystal_mask_kmeans( + arr: np.ndarray, n_clusters: int = 3 +) -> tuple[np.ndarray, float]: + """K-means on signal values (1-D feature). + + Splits cells into n_clusters groups; the cluster with the highest + centroid is the crystal. n_clusters=3 separates background / fringe / crystal. + """ + from sklearn.cluster import KMeans + + flat = arr.flatten().reshape(-1, 1) + km = KMeans(n_clusters=n_clusters, random_state=0, n_init="auto").fit(flat) + crystal_label = int(np.argmax(km.cluster_centers_.flatten())) + threshold = float(sorted(km.cluster_centers_.flatten())[-2]) + mask = km.labels_.reshape(arr.shape) == crystal_label + return mask, threshold + + +CRYSTAL_METHOD_MAP: dict = { + "corner": ("Corner background", crystal_mask_corner_background), + "otsu": ("Otsu (nonzero)", crystal_mask_otsu_nonzero), + "mean_sigma": ( + "Mean + 0.5σ", + lambda arr: crystal_mask_mean_sigma(arr, n_sigma=0.5), + ), + "percentile": ( + "Top-40% signal", + lambda arr: crystal_mask_signal_percentile(arr, percentile=60.0), + ), + "dbscan": ("DBSCAN (spatial)", crystal_mask_dbscan), + "kmeans": ("K-means (k=3)", crystal_mask_kmeans), +} + + +def compare_crystal_methods( + results: List, + arr: Optional[np.ndarray] = None, + grid_size_mm: Optional[tuple[float, float]] = None, +) -> None: + """Plot each crystal-detection method side-by-side for visual comparison. + + Cyan contour = detected crystal region. + Red square = cell with maximum spots_low_res. + grid_size_mm = (step_x_mm, step_y_mm) — when provided, crystal size in µm + is calculated via get_xtal_size and shown in each panel title. + """ + import matplotlib.pyplot as plt + + if arr is None: + arr = rebuild_array_from_scan_results(results, "spots_low_res") + + methods = [ + ("Corner background", *crystal_mask_corner_background(arr)), + ("Otsu (nonzero)", *crystal_mask_otsu_nonzero(arr)), + ("Mean + 0.5σ", *crystal_mask_mean_sigma(arr, n_sigma=0.5)), + ("Top-40% signal", *crystal_mask_signal_percentile(arr, percentile=60.0)), + ("DBSCAN (spatial)", *crystal_mask_dbscan(arr)), + ("K-means (k=3)", *crystal_mask_kmeans(arr)), + ] + + fig, axes = plt.subplots(2, 3, figsize=(14, 9)) + for ax, (label, mask, threshold) in zip(axes.flat, methods): + _draw_panel(ax, arr, mask, label, threshold, grid_size_mm=grid_size_mm) + + fig.suptitle( + "Crystal region detection — method comparison\n" + "cyan = crystal outline | red square = max spots_low_res cell", + fontsize=10, + ) + plt.tight_layout() + plt.show() + + +def compare_crystal_methods_scored( + results: List, + score_arr: Optional[np.ndarray] = None, + method: Optional[str] = None, + grid_size_mm: Optional[tuple[float, float]] = None, + w_bkg: float = 0.20, + w_low_res: float = 0.60, + w_indexed: float = 0.20, +) -> None: + """Plot each crystal-detection method applied to the composite 0–100 crystal score. + + The score combines bkg (w_bkg), spots_low_res (w_low_res), and spots_indexed + (w_indexed), each min-max normalised. All six methods are shown side-by-side. + + Args: + results: scan result list (used to build score if score_arr is None) + score_arr: pre-computed score array; computed from results when None + method: if given (one of 'corner','otsu','mean_sigma','percentile', + 'dbscan','kmeans'), that panel is highlighted with a yellow border + grid_size_mm: (step_x_mm, step_y_mm) for crystal-size annotation + w_bkg, w_low_res, w_indexed: composite score weights (must sum to 1.0) + """ + import matplotlib.pyplot as plt + + if score_arr is None: + score_arr = compute_crystal_score_array( + results, w_bkg=w_bkg, w_low_res=w_low_res, w_indexed=w_indexed + ) + + method_entries = [ + (name, label, fn) for name, (label, fn) in CRYSTAL_METHOD_MAP.items() + ] + + fig, axes = plt.subplots(2, 3, figsize=(15, 9)) + for ax, (name, label, fn) in zip(axes.flat, method_entries): + mask, threshold = fn(score_arr) + _draw_panel( + ax, + score_arr, + mask, + label, + threshold, + grid_size_mm=grid_size_mm, + cbar_label="crystal score (0–100)", + ) + if method and name == method: + for spine in ax.spines.values(): + spine.set_edgecolor("yellow") + spine.set_linewidth(3) + + weight_note = f"bkg×{w_bkg:.0%} + spots_low_res×{w_low_res:.0%} + spots_indexed×{w_indexed:.0%}" + fig.suptitle( + f"Crystal region detection on composite score [{weight_note}]\n" + "cyan = crystal outline | red square = max-score cell | score range 0–100", + fontsize=10, + ) + plt.tight_layout() + plt.show() diff --git a/src/aarecommon/math/raster_grid.py b/src/aarecommon/math/raster_grid.py new file mode 100644 index 0000000..92b21ea --- /dev/null +++ b/src/aarecommon/math/raster_grid.py @@ -0,0 +1,49 @@ +def grid_to_image_id(grid_x: int, grid_y: int, number_of_cols: int) -> int: + """ + Get the associated image id for a given grid position (grid_x, grid_y) + using serpentine row ordering. + + grid_x: 1-based column index + grid_y: 1-based row index + number_of_cols: total number of columns in the grid + """ + if number_of_cols < 1: + raise ValueError("number_of_cols must be >= 1") + if grid_x < 1: + raise ValueError("grid_x must be >= 1") + if grid_y < 1: + raise ValueError("grid_y must be >= 1") + if grid_x > number_of_cols: + raise ValueError("grid_x cannot be greater than number_of_cols") + + col = grid_x - 1 + row = grid_y - 1 + + if row % 2 == 0: # left -> right + return row * number_of_cols + col + else: # right -> left + return row * number_of_cols + (number_of_cols - 1 - col) + + +def image_id_to_grid(image_id: int, number_of_cols: int) -> tuple[int, int]: + """ + Get the associated grid position (grid_x, grid_y) for a given image id + using serpentine row ordering. + """ + if number_of_cols < 1: + raise ValueError("number_of_cols must be >= 1") + if image_id < 0: + raise ValueError("image_id must be >= 0") + + row = image_id // number_of_cols + pos_in_row = image_id % number_of_cols + + if row % 2 == 0: # left -> right + col = pos_in_row + else: # right -> left + col = number_of_cols - 1 - pos_in_row + + grid_x = col + 1 + grid_y = row + 1 + + return grid_x, grid_y diff --git a/src/aarecommon/math/sample_geometry.py b/src/aarecommon/math/sample_geometry.py new file mode 100644 index 0000000..5185b01 --- /dev/null +++ b/src/aarecommon/math/sample_geometry.py @@ -0,0 +1,97 @@ +from typing import Annotated + +import numpy as np +from aarecommon.math.coordinate import Coordinate, SmargonCoordinate, positive_coords +from pydantic import AfterValidator, BaseModel, Field + + +class SampleGeometryModel(BaseModel): + # Beam location in the camera coordinates + beam_location_pxl: Coordinate + + # 1 pixel in mm + pixel_in_mm: Annotated[float, Field(gt=0.0)] + + aerotech: Coordinate + + aerotech_meas: Coordinate + # Current smargon coordinates/target + smargon: SmargonCoordinate + + omega_deg: float + + # Beam size in mm + beam_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)] + + @property + def beam_size_pxl(self) -> Coordinate: + return self.beam_size_mm / self.pixel_in_mm + + def smargon_z(self, pic_y: float) -> Coordinate: + y_mm = (pic_y - self.beam_location_pxl.y) * self.pixel_in_mm + return self.smargon.sh_mm + self.smargon_nudge(Coordinate(z=y_mm)) + + # This goes from image coordinates to relative beam coordinates in image plane + def picture_to_sample(self, pic_pxl: Coordinate) -> Coordinate: + # Relative to beam location + beam_mm = (pic_pxl - self.beam_location_pxl) * self.pixel_in_mm + self.aerotech + return beam_mm + + # This goes from relative beam coordinates in image plane to image + def sample_to_picture(self, beam_mm: Coordinate) -> Coordinate: + beam_pxl = (beam_mm - self.aerotech) / self.pixel_in_mm + return beam_pxl + self.beam_location_pxl + + def translate_smargon(self, coord: Coordinate) -> SmargonCoordinate: + return SmargonCoordinate( + sh_mm=self.smargon.sh_mm + self.smargon_nudge(coord), + phi_deg=self.smargon.phi_deg, + chi_deg=self.smargon.chi_deg, + ) + + def smargon_nudge(self, coord: Coordinate) -> Coordinate: + phi = np.radians(np.around(self.smargon.phi_deg, decimals=1)) + chi = np.radians(np.around(self.smargon.chi_deg, decimals=1)) + omega = np.radians(np.around(self.omega_deg, decimals=1)) + + offset = Coordinate() + co = np.cos(omega) + so = np.sin(omega) + cp = np.cos(phi) + sp = np.sin(phi) + cc = np.cos(chi) + sc = np.sin(chi) + + offset.x = ( + -coord.z * co * sp + - coord.y * so * sp + + coord.x * cp * sc + - coord.y * co * cc * cp + + coord.z * so * cc * cp + ) + offset.y = ( + -coord.z * co * cp + - coord.y * so * cp + - coord.x * sc * sp + + coord.y * co * cc * sp + - coord.z * so * cc * sp + ) + offset.z = -coord.x * cc - coord.y * co * sc + coord.z * so * sc + return offset + + def beamline_to_smargon(self, coord: Coordinate) -> Coordinate: + return self.smargon.sh_mm + self.smargon_nudge(coord) + + def smargon_to_beamline(self, coord: Coordinate) -> Coordinate: + rel = coord - self.smargon.sh_mm + + vec_x = self.smargon_nudge(Coordinate(x=1)).normalize() + vec_y = self.smargon_nudge(Coordinate(y=1)).normalize() + + return Coordinate(x=rel * vec_x, y=rel * vec_y) + + def picture_to_smargon(self, coord: Coordinate) -> Coordinate: + return self.beamline_to_smargon(self.picture_to_sample(coord)) + + def smargon_to_picture(self, coord: Coordinate) -> Coordinate: + return self.sample_to_picture(self.smargon_to_beamline(coord)) diff --git a/src/aarecommon/math/simulate_raster.py b/src/aarecommon/math/simulate_raster.py new file mode 100644 index 0000000..b1f510f --- /dev/null +++ b/src/aarecommon/math/simulate_raster.py @@ -0,0 +1,226 @@ +""" +Synthetic raster scan data for no-beam / offline testing. + +Entry point +----------- +generate_no_beam_scan_result(request, seed=None) -> ScanResult + +The returned ScanResult is structurally identical to a real one, so all +find_xtal.py functions consume it without modification. + +Pre-defined seeds +----------------- +Each seed fixes cluster geometry so results are reproducible. The six +scenarios below cover the main cases needed for unit-testing crystal-finding +algorithms. + +Seed Scenario +---- -------- +1 Single crystal, centred – baseline positive detection +2 Single crystal, off-centre – tests COM accuracy near an edge +3 Two well-separated crystals – multi-crystal detection +4 Two overlapping crystals – segmentation challenge +5 Two clusters at ~90 ° – twinned / differently oriented crystal +6 Pure background only – true negative (no crystal) +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List + +import numpy as np +from aarecommon.models.raster_grid import RasterGridRequest +from jfjoch_client import ScanResult +from jfjoch_client.models.scan_result_images_inner import ScanResultImagesInner + +# --------------------------------------------------------------------------- +# Cluster geometry descriptor +# --------------------------------------------------------------------------- + + +@dataclass +class _ClusterParams: + """Fractional coordinates and shape of one elliptical crystal cluster. + + cx, cy – cluster centre as a fraction of (n_x, n_y) [0..1] + ax, ay – semi-axes as a fraction of (n_x, n_y) + theta – rotation of the ellipse in radians + peak – peak spots_low_res at cluster centre (integer) + """ + + cx: float + cy: float + ax: float + ay: float + theta: float + peak: int + + +@dataclass +class _SeedConfig: + clusters: List[_ClusterParams] = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Pre-defined seed catalogue +# --------------------------------------------------------------------------- + +SEED_CATALOGUE: dict[int, _SeedConfig] = { + # --- 1: single crystal, centred, compact, strong signal ---------------- + 1: _SeedConfig( + clusters=[ + _ClusterParams(cx=0.50, cy=0.50, ax=0.15, ay=0.15, theta=0.0, peak=120), + ] + ), + # --- 2: single crystal, off-centre, slightly elongated ----------------- + 2: _SeedConfig( + clusters=[ + _ClusterParams(cx=0.25, cy=0.70, ax=0.12, ay=0.18, theta=0.35, peak=90), + ] + ), + # --- 3: two well-separated crystals ------------------------------------ + 3: _SeedConfig( + clusters=[ + _ClusterParams(cx=0.20, cy=0.20, ax=0.12, ay=0.12, theta=0.0, peak=100), + _ClusterParams(cx=0.75, cy=0.72, ax=0.15, ay=0.10, theta=0.2, peak=80), + ] + ), + # --- 4: two overlapping crystals (centres ~1 sigma apart) -------------- + 4: _SeedConfig( + clusters=[ + _ClusterParams(cx=0.40, cy=0.45, ax=0.18, ay=0.15, theta=0.0, peak=110), + _ClusterParams(cx=0.58, cy=0.55, ax=0.16, ay=0.18, theta=0.5, peak=95), + ] + ), + # --- 5: two clusters with ~90° different orientations (twinned) -------- + 5: _SeedConfig( + clusters=[ + _ClusterParams(cx=0.30, cy=0.40, ax=0.25, ay=0.08, theta=0.0, peak=105), + _ClusterParams(cx=0.68, cy=0.62, ax=0.08, ay=0.25, theta=0.0, peak=100), + ] + ), + # --- 6: pure background, no crystal – true negative -------------------- + 6: _SeedConfig(clusters=[]), +} + + +# --------------------------------------------------------------------------- +# Core generation +# --------------------------------------------------------------------------- + + +def _gaussian_cluster( + ix: np.ndarray, + iy: np.ndarray, + p: _ClusterParams, + n_x: int, + n_y: int, +) -> np.ndarray: + """Return a 2-D Gaussian signal array for one cluster. + + ix, iy are integer coordinate grids with shape (n_x, n_y). + """ + cx = p.cx * (n_x - 1) + cy = p.cy * (n_y - 1) + + dx = (ix - cx) / max(p.ax * n_x, 0.5) + dy = (iy - cy) / max(p.ay * n_y, 0.5) + + cos_t = np.cos(p.theta) + sin_t = np.sin(p.theta) + dx_r = dx * cos_t + dy * sin_t + dy_r = -dx * sin_t + dy * cos_t + + return p.peak * np.exp(-2.0 * (dx_r**2 + dy_r**2)) + + +def generate_no_beam_scan_result( + request: RasterGridRequest, + seed: int | None = None, +) -> ScanResult: + """Build a synthetic ScanResult for no-beam / offline operation. + + Parameters + ---------- + request: + The grid request whose n_x, n_y, and file_prefix are used. + seed: + Integer 1-6 selects a pre-defined scenario from SEED_CATALOGUE. + Any other value (or None) draws cluster parameters randomly using + the seed as an RNG seed (None → fully random). + + Returns + ------- + ScanResult + Populated with one ScanResultImagesInner per grid cell, with + realistic background noise and optional crystal cluster(s). + """ + n_x = max(request.n_x, 1) + n_y = max(request.n_y, 1) + + rng = np.random.default_rng(seed) + + # Coordinate grids + ix, iy = np.meshgrid(np.arange(n_x), np.arange(n_y), indexing="ij") + + # Background: Poisson noise, mean ≈ 3 counts + bkg = rng.poisson(lam=3.0, size=(n_x, n_y)).astype(float) + + # Crystal signal layer (spots_low_res) + signal = np.zeros((n_x, n_y), dtype=float) + + if seed in SEED_CATALOGUE: + cfg = SEED_CATALOGUE[seed] + clusters = cfg.clusters + else: + # Random fallback: 1 or 2 clusters + n_clusters = rng.integers(1, 3) + clusters = [ + _ClusterParams( + cx=rng.uniform(0.15, 0.85), + cy=rng.uniform(0.15, 0.85), + ax=rng.uniform(0.12, 0.30), + ay=rng.uniform(0.12, 0.30), + theta=rng.uniform(0, np.pi), + peak=int(rng.integers(50, 150)), + ) + for _ in range(n_clusters) + ] + + for p in clusters: + signal += _gaussian_cluster(ix, iy, p, n_x, n_y) + + # Add Poisson noise on top of the crystal signal + spots_low_res = rng.poisson(lam=np.maximum(signal, 0)).astype(int) + + # Assemble image list – one entry per grid cell in serpentine order: + # row 0 left→right, row 1 right→left, row 2 left→right, … + images: list[ScanResultImagesInner] = [] + img_number = 0 + for xi in range(n_x): + yi_range = range(n_y) + for yi in yi_range: + images.append( + ScanResultImagesInner( + number=img_number, + nx=xi, + ny=yi, + efficiency=1.0, + bkg=float(bkg[xi, yi]), + spots=int(spots_low_res[xi, yi]), + spots_low_res=int(spots_low_res[xi, yi]), + spots_indexed=0, + spots_ice=0, + index=0, + b=0.0, + res=None, + pixel_sum=None, + max=None, + sat=None, + err=None, + ) + ) + img_number += 1 + + return ScanResult(file_prefix=request.file_prefix, images=images) diff --git a/src/aarecommon/models/__init__.py b/src/aarecommon/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/aarecommon/models/aerotech.py b/src/aarecommon/models/aerotech.py new file mode 100644 index 0000000..a290e08 --- /dev/null +++ b/src/aarecommon/models/aerotech.py @@ -0,0 +1,154 @@ +from enum import Enum +from typing import Optional + +from aarecommon.models.rotation_scan import RotationScanRequest +from pydantic import BaseModel, ConfigDict, Field + + +class TaskEnum(Enum): + TASK_0 = 0 + TASK_1 = 1 + TASK_2 = 2 + TASK_3 = 3 + TASK_4 = 4 + TASK_5 = 5 + TASK_6 = 6 + TASK_7 = 7 + TASK_8 = 8 + TASK_9 = 9 + + +class AxisEnum(Enum): + X = "x" + Y = "y" + Z = "z" + OMEGA = "u" + + +class AerotechRunEnum(Enum): + STOP = 0 + START = 1 + RUN = 2 + LOAD = 3 + PAUSE = 4 + RESET = 5 + + +class VariableTypeEnum(Enum): + INT = 0 + REAL = 1 + STRING = 2 + + +class AerotechAxisStatus(BaseModel): + enabled: bool + fault: int + homed: bool + is_fault: bool + moving: bool + position: float + status: int + velocity: float + + +class AerotechStatus(BaseModel): + state: str + x: Optional[AerotechAxisStatus] = None + y: Optional[AerotechAxisStatus] = None + z: Optional[AerotechAxisStatus] = None + u: Optional[AerotechAxisStatus] = None + + model_config = ConfigDict(extra="allow") + + def __str__(self) -> str: + return self.to_pretty_string() + + def to_pretty_string(self) -> str: + def axis_line(name: str, axis: AerotechAxisStatus | None) -> str: + if axis is None: + return f"{name.upper()}: unavailable" + return ( + f"{name.upper():>2} | pos={axis.position:>12.6f} | " + f"homed={axis.homed!s:<5} | moving={axis.moving!s:<5} | " + f"enabled={axis.enabled!s:<5} | fault={axis.fault} | " + f"faulted={axis.is_fault!s:<5} | vel={axis.velocity:>10.6f}" + ) + + return "\n".join( + [ + f"STATE: {self.state}", + axis_line("x", self.x), + axis_line("y", self.y), + axis_line("z", self.z), + axis_line("u", self.u), + ] + ) + + def to_compact_string(self) -> str: + axes = [] + for name in ("x", "y", "z", "u"): + axis = getattr(self, name) + if axis is not None: + axes.append( + f"{name}={axis.position:.4f} " + f"({'H' if axis.homed else 'NH'}, {'M' if axis.moving else '-'})" + ) + return f"state={self.state} | " + " | ".join(axes) + + def to_colored_string(self) -> str: + # ANSI colors for terminal use + RESET = "\033[0m" + BOLD = "\033[1m" + CYAN = "\033[36m" + GREEN = "\033[32m" + YELLOW = "\033[33m" + RED = "\033[31m" + + def color_bool(value: bool) -> str: + return f"{GREEN}True{RESET}" if value else f"{RED}False{RESET}" + + def axis_line(name: str, axis: AerotechAxisStatus | None) -> str: + if axis is None: + return f"{YELLOW}{name.upper()}: unavailable{RESET}" + return ( + f"{BOLD}{name.upper()}{RESET} | " + f"pos={CYAN}{axis.position:>12.6f}{RESET} | " + f"homed={color_bool(axis.homed)} | " + f"moving={color_bool(axis.moving)} | " + f"enabled={color_bool(axis.enabled)} | " + f"fault={axis.fault} | " + f"faulted={color_bool(axis.is_fault)} | " + f"vel={axis.velocity:>10.6f}" + ) + + return "\n".join( + [ + f"{BOLD}STATE:{RESET} {CYAN}{self.state}{RESET}", + axis_line("x", self.x), + axis_line("y", self.y), + axis_line("z", self.z), + axis_line("u", self.u), + ] + ) + + +class AerotechTarget(BaseModel): + x: Optional[float] = None + y: Optional[float] = None + z: Optional[float] = None + u: Optional[float] = None + + def to_payload(self) -> dict: + return self.model_dump(exclude_none=True) + + +class AerotechRotationScanRequest(RotationScanRequest): + rotation_deg: float + time_sec: float + start_pos_deg: float + async_move: bool = Field(default=False, alias="async") + + model_config = ConfigDict(populate_by_name=True) + + def to_payload(self) -> dict: + return self.model_dump(by_alias=True, exclude_none=True) diff --git a/src/aarecommon/models/auth.py b/src/aarecommon/models/auth.py new file mode 100644 index 0000000..2960c5a --- /dev/null +++ b/src/aarecommon/models/auth.py @@ -0,0 +1,66 @@ +from enum import Enum +from pydantic import BaseModel +from datetime import datetime + + +class BatonRequestStatus(Enum): + PENDING = "pending" + ACCEPTED = "accepted" + REFUSED = "refused" + TIMEOUT = "timeout" + CANCELLED = "cancelled" + + +class BatonHolderInfo(BaseModel): + """Information about the current baton holder.""" + + username: str + session: int + is_staff: bool + pgroup: str | None = None + + +class BatonRequest(BaseModel): + """A request from one user to take the baton from another.""" + + request_id: str + requester_username: str + requester_session: int + requester_is_staff: bool + holder_username: str | None = None + holder_session: int | None = None + created_at: float # Unix timestamp + timeout_seconds: int = 30 + status: BatonRequestStatus = BatonRequestStatus.PENDING + + +class BatonTransferQueue(BaseModel): + """Queued baton transfer waiting for beamline to be available.""" + + target_session: int + target_username: str + target_is_staff: bool + target_pgroup: str | None = None + queued_at: float # Unix timestamp + reason: str = "beamline_busy" + + +class BatonStatus(BaseModel): + """Full baton status for GUI display.""" + + holder: BatonHolderInfo | None = None + pending_request: BatonRequest | None = None + queued_transfer: BatonTransferQueue | None = None + you_are_holder: bool = False + you_have_pending_request: bool = False + incoming_request: bool = False + allow_non_staff_request: bool = False + + +def get_user(): + import os, getpass + + try: + return os.getlogin() + except OSError: + return os.environ.get("USER") or getpass.getuser() diff --git a/src/aarecommon/models/automation.py b/src/aarecommon/models/automation.py new file mode 100644 index 0000000..05362d2 --- /dev/null +++ b/src/aarecommon/models/automation.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime, timezone +from enum import Enum +from typing import Any, Literal + + +class WorkflowStateKind(str, Enum): + MOUNT = "mount" + LOOP_CENTRE = "loop_centre" + RASTER = "raster" + DATA_COLLECTION = "data_collection" + FINAL = "Paused/Finished" + + +class StepStatus(str, Enum): + PENDING = "pending" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + SKIPPED = "skipped" + PAUSED = "paused" + + +@dataclass +class StepState: + step: WorkflowStateKind + status: StepStatus = StepStatus.PENDING + message: str = "" + started_at: float | None = None + completed_at: float | None = None + error_code: str | None = None + + +@dataclass +class AutomationProgress: + current_step: str | None + steps: list[StepState] = field(default_factory=list) + events: list[LogEvent] = field(default_factory=list) + finished: bool = False + success: bool | None = None + samples_in_queue: int = 0 + avg_time_per_sample: float = 0.0 + current_sample_name: str = "" + + def append_event( + self, + *, + level: Literal["INFO", "WARNING", "ERROR"], + code: str, + message: str, + exception_class: str | None = None, + sample_id: int | None = None, + context: dict[str, Any] | None = None, + ) -> None: + self.events.append( + LogEvent( + ts=datetime.now(timezone.utc), + level=level, + code=code, + exception_class=exception_class, + message=message, + sample_id=sample_id, + context=dict(context or {}), + ) + ) + + +@dataclass +class LogEvent: + ts: datetime + level: Literal["INFO", "WARNING", "ERROR"] + code: str + exception_class: str | None + message: str + sample_id: int | None + context: dict[str, Any] = field(default_factory=dict) diff --git a/src/aarecommon/models/beamline.py b/src/aarecommon/models/beamline.py new file mode 100644 index 0000000..26cfa25 --- /dev/null +++ b/src/aarecommon/models/beamline.py @@ -0,0 +1,8 @@ +from enum import StrEnum + + +class MXBeamline(StrEnum): + X06SA = "X06SA" + X10SA = "X10SA" + X06DA = "X06DA" + SIMULATED = "SIMULATED" diff --git a/src/aarecommon/models/models.py b/src/aarecommon/models/models.py new file mode 100644 index 0000000..e527610 --- /dev/null +++ b/src/aarecommon/models/models.py @@ -0,0 +1,749 @@ +import re +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Annotated, List, Literal, Optional, Tuple + +from aarecommon.models.beamline import MXBeamline +from aarecommon.math.coordinate import Coordinate, positive_coords +from aarecommon.math.diffraction_geometry import DiffractionGeometry +from aarecommon.math.sample_geometry import SampleGeometryModel +from aarecommon.models.tell import TellStateModel +from jfjoch_client.models.scan_result import ScanResult +from pydantic import ( + AfterValidator, + AliasChoices, + BaseModel, + ConfigDict, + Field, + field_validator, +) + + +class StagePositionEnum(Enum): + MEASURE = 0 + PARK = 1 + DOWN = 2 + UNKNOWN = 3 + + +class TokenData(BaseModel): + sub: str # Username + pgroups: List[str] + session: int + staff: bool = False + + +class DewarAddress(BaseModel): + segment: Literal["A", "B", "C", "D", "E", "F", "X", "R"] + pos: Annotated[int, Field(ge=1, le=5)] + + +class SampleDewarAddress(BaseModel): + puck: DewarAddress + pin: Annotated[int, Field(ge=1, le=16)] + + +# From the database for puck loading +class PuckInfo(BaseModel): + db_id: int + puck_name: str + dewar_name: str + user: str = "" + location: Optional[DewarAddress] = None + + +# From TELL to database after loading +class PuckLoadedInfo(BaseModel): + puck_name: str + location: DewarAddress + + +class DataCollectionParameters(BaseModel): + model_config = ConfigDict(from_attributes=True, populate_by_name=True) + + directory: Optional[str] = None + oscillation: Optional[float] = None # Only accept positive float + exposure: Optional[float] = None # Only accept positive floats between 0 and 1 + totalangle: Optional[int] = Field( # was totalrange + default=None, validation_alias=AliasChoices("totalangle", "totalrange") + ) # Only accept positive integers between 0 and 360 + transmission: Optional[int] = ( + None # Only accept positive integers between 0 and 100 + ) + targetresolution: Optional[float] = None # Only accept positive float + beamsize: Optional[str] = None + aperture: Optional[int] = None # Optional string field + datacollectiontype: Optional[str] = ( + None # Only accept "standard", other types might be added later + ) + processingpipeline: Optional[str] = ( + "" # Only accept "gopy", "autoproc", "xia2dials" + ) + spacegroupnumber: Optional[int] = ( + None # Only accept positive integers between 1 and 230 + ) + unitcell: Optional[str] = Field( # was cellparameters + default=None, validation_alias=AliasChoices("unitcell", "cellparameters") + ) # Must be a set of six positive floats or integers + rescutkey: Optional[str] = None # Only accept "is" or "cchalf" + rescutvalue: Optional[float] = ( + None # Must be a positive float if rescutkey is provided + ) + processingresolution: Optional[float] = Field( # was userresolution + default=None, + validation_alias=AliasChoices("processingresolution", "userresolution"), + ) + pdbid: Optional[str] = ( + "" # Accepts either the format of the protein data bank code or {provided} + ) + autoprocfull: Optional[bool] = None + procfull: Optional[bool] = None + adpenabled: Optional[bool] = None + noano: Optional[bool] = None + ffcscampaign: Optional[bool] = None + trustedhigh: Optional[float] = None # Should be a float between 0 and 2.0 + autoprocextraparams: Optional[str] = None # Optional string field + chiphiangles: Optional[float] = None # Optional float field between 0 and 30 + dose: Optional[float] = None # Optional float field + cloud: bool = True + pdbmodel: Optional[str] = None + + @field_validator("directory", mode="after") + @classmethod + def directory_characters(cls, v): + + # Default directory value if empty + if not v: # Handles None or empty cases + default_value = "{date}/{prefix}" + return default_value + + # Strip trailing slashes and store original value for comparison + v = str(v).strip("/") # Ensure it's a string and no trailing slashes + original_value = v + + # Replace spaces with underscores + v = v.replace(" ", "_") + + # Validate directory pattern with macros and allowed characters + valid_macros = [ + # Current macros + "{puck}", + "{position}", + "{prefix}", + "{date}", + "{run}", + "{beamline}", + # Legacy macros — accepted for back-compat, no longer documented + "{sgPuck}", + "{sgPosition}", + "{sgPrefix}", + "{sgPriority}", + "{protein}", + "{method}", + ] + valid_macro_pattern = re.compile( + "|".join(re.escape(macro) for macro in valid_macros) + ) + + # Check if the value contains valid macros + allowed_chars_pattern = "[a-z0-9_.+-/]" + v_without_macros = valid_macro_pattern.sub("macro", v) + + allowed_path_pattern = re.compile( + f"^(({allowed_chars_pattern}+|macro)*/*)*$", re.IGNORECASE + ) + if not allowed_path_pattern.match(v_without_macros): + raise ValueError( + f"'{v}' is not valid. Value must be a valid path or macro." + ) + return v + + @field_validator("unitcell", mode="before") + @classmethod + def unitcell_format(cls, v): + if v: + tokens = v.replace(",", " ").split() + try: + values = [float(i) for i in tokens] + except ValueError: + raise ValueError( + f"'{v}' is not valid. " + "Value must be six positive floats or integers (space or comma separated)." + ) + if len(values) != 6 or any(val <= 0 for val in values): + raise ValueError( + f"'{v}' is not valid. " + "Value must be six positive floats or integers (space or comma separated)." + ) + return v + + @field_validator("cloud", mode="before") + @classmethod + def coerce_cloud_default(cls, v): + if v in ("", None): + return True + if isinstance(v, bool): + return v + v_str = str(v).strip().lower() + if v_str in {"true", "yes", "1"}: + return True + if v_str in {"false", "no", "0"}: + return False + raise ValueError("cloud must be blank for default, or True/False") + + +# From database to TELL after loading +class SampleShortInfo(BaseModel): + db_id: int + puck_name: str + dewar_name: str + sample_name: str + run_number: int + aaredb_params: Optional[DataCollectionParameters] = None + user: str = "" + pin: Annotated[int, Field(ge=1, le=16)] + location: DewarAddress | None = None + priority: Optional[float] = 1.0 + comment: Optional[str] = None + mount_count: int = 0 + rotation_count: int = 0 + raster_count: int = 0 + screening_count: int = 0 + + def tell_address(self) -> SampleDewarAddress: + return SampleDewarAddress(puck=self.location, pin=self.pin) + + def loc_str(self) -> str: + if self.location is None: + return "-" + else: + return f"{self.location.segment}{self.location.pos}-{self.pin}" + + def loc_str_sort(self) -> str: + if self.location is None: + return "" + else: + return f"{self.location.segment}{self.location.pos}-{self.pin:02d}" + + @classmethod + def from_dict(cls, data: dict): + return cls(**data) + + +class SampleShortInfoList(BaseModel): + s: List[SampleShortInfo] + + +class BeamMarkCoeffModel(BaseModel): + """ + Model to calculate beam center for a given zoom level based on a quadratic approximation + For both x and y it contains tuple of coefficients a, b, c (a*zoom^2 + b*zoom + c) + Default is beam center in 1000,1000 at any zoom level + """ + + coeff_x: Tuple[float, float, float] = (0, 0, 1000.0) + coeff_y: Tuple[float, float, float] = (0, 0, 1000.0) + + def apply(self, zoom: float): + return Coordinate( + x=self.coeff_x[0] * zoom**2 + self.coeff_x[1] * zoom + self.coeff_x[2], + y=self.coeff_y[0] * zoom**2 + self.coeff_y[1] * zoom + self.coeff_y[2], + ) + + +class FluorescenceSpectrumParameterModel(BaseModel): + erase: bool = True + acq_time_s: float + transmission: Annotated[float, Field(ge=0, le=1)] | None = None + + +class FluorescenceSpectrumOutputModel(BaseModel): + bkg: list[float] | None = None + spectrum: list[float] + energy_eV: list[float] + average_dead_time: Annotated[float, Field(ge=0, le=1)] + + +class MLBoxType(Enum): + LOOP_ALL = 0 + PIN = 1 + CRYSTAL = 2 + LOOP_FACE = 3 + ICE = 4 + NEEDLE = 5 + + +class BoundingBoxModel(BaseModel): + top_x: float + top_y: float + bottom_x: float + bottom_y: float + + +class MLBoxModel(BaseModel): + cls: MLBoxType + box: BoundingBoxModel + conf: float + + @staticmethod + def from_tuple( + cls: MLBoxType, box_tuple: tuple[float, float, float, float], conf: float + ) -> "MLBoxModel": + x1, y1, x2, y2 = box_tuple + return MLBoxModel( + cls=cls, + box=BoundingBoxModel( + top_x=float(x1), top_y=float(y1), bottom_x=float(x2), bottom_y=float(y2) + ), + conf=float(conf), + ) + + +class MLOutputModel(BaseModel): + boxes: dict[str, MLBoxModel] = {} + + @staticmethod + def get_class_str(cls: MLBoxType) -> str: + if cls == MLBoxType.LOOP_ALL: + return "Loop_all" + if cls == MLBoxType.PIN: + return "Pin" + if cls == MLBoxType.CRYSTAL: + return "Crystal" + if cls == MLBoxType.LOOP_FACE: + return "Loop_face" + if cls == MLBoxType.ICE: + return "Ice" + if cls == MLBoxType.NEEDLE: + return "Needle" + return "Unknown" + + def _next_unique_key(self, base: str) -> str: + if base not in self.boxes: + return base + i = 2 + while f"{base}_{i}" in self.boxes: + i += 1 + return f"{base}_{i}" + + def add_box( + self, + cls: MLBoxType, + box_tuple: tuple[float, float, float, float], + conf: float | None = None, + ) -> str: + key_base = self.get_class_str(cls) + key = self._next_unique_key(key_base) + self.boxes[key] = MLBoxModel.from_tuple(cls, box_tuple, conf) + return key + + def get_box_model(self, key: str) -> MLBoxModel | None: + return self.boxes.get(key) + + def get_box_tuple(self, key: str) -> tuple[float, float, float, float] | None: + m = self.get_box_model(key) + if not m or not m.box: + return None + return (m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y) + + def get_box_tuple_with_conf( + self, key: str + ) -> tuple[float, float, float, float, float] | None: + m = self.get_box_model(key) + if not m or not m.box: + return None + conf = float(m.conf) if m.conf is not None else 0.0 + return (m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y, conf) + + def get_keys_for_class(self, cls: MLBoxType) -> list[str]: + base = self.get_class_str(cls) + return [ + k + for k, v in self.boxes.items() + if v.cls == cls and (k == base or k.startswith(base + "_")) + ] + + def get_models_for_class(self, cls: MLBoxType) -> list[MLBoxModel]: + keys = self.get_keys_for_class(cls) + return [self.boxes[k] for k in keys] + + def get_tuples_for_class( + self, cls: MLBoxType + ) -> list[tuple[float, float, float, float]]: + out: list[tuple[float, float, float, float]] = [] + for m in self.get_models_for_class(cls): + if m.box: + out.append((m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y)) + return out + + def get_tuples_with_conf_for_class( + self, cls: MLBoxType + ) -> list[tuple[float, float, float, float, float]]: + out: list[tuple[float, float, float, float, float]] = [] + for m in self.get_models_for_class(cls): + if m.box: + out.append( + ( + m.box.top_x, + m.box.top_y, + m.box.bottom_x, + m.box.bottom_y, + float(m.conf), + ) + ) + return out + + def get_best_for_class(self, cls: MLBoxType) -> MLBoxModel | None: + models = self.get_models_for_class(cls) + if not models: + return None + return max(models, key=lambda m: m.conf if m.conf is not None else 0.0) + + +class BeamlineStateEnum(Enum): + Maintenance = 1 + SampleExchange = 2 + SampleAlignment = 3 + DataCollection = 4 + DewarTransfer = 5 + XrayFluorescence = 6 + BeamLocation = 7 + Moving = 8 + RobotSampleExchange = 9 + XtalSnapshot = 10 + BeamstopAlignment = 11 + FluxMeasurement = 12 + + def display_name(self) -> str: + return { + BeamlineStateEnum.Maintenance: "Maintenance", + BeamlineStateEnum.SampleExchange: "Sample exchange", + BeamlineStateEnum.SampleAlignment: "Sample alignment", + BeamlineStateEnum.DataCollection: "Data collection", + BeamlineStateEnum.DewarTransfer: "Dewar transfer", + BeamlineStateEnum.XrayFluorescence: "X-ray fluorescence", + BeamlineStateEnum.BeamLocation: "Beam location", + BeamlineStateEnum.Moving: "Moving", + BeamlineStateEnum.RobotSampleExchange: "Robot sample exchange", + BeamlineStateEnum.XtalSnapshot: "Xtal snapshot", + BeamlineStateEnum.BeamstopAlignment: "Beamstop alignment", + BeamlineStateEnum.FluxMeasurement: "Flux measurement", + }.get(self, "-") + + +class DAQOperation(str, Enum): + AUTOMATION = "automation" + MOUNT = "mount" + UNMOUNT = "unmount" + LOOP_CENTERING = "loop_centering" + FACE_CENTERING = "face_centering" + RASTER = "raster" + ROTATION = "rotation" + MEASURE = "measure" + + +class SessionsStateEnum(Enum): + Vacant = 0 + OwnedByYou = 1 + OwnedByElse = 2 + PendingYouToElse = 3 + PendingElseToYou = 4 + + +class SampleCameraSettings(BaseModel): + gain: float + exposure: float + + +class ZoomModeEnum(Enum): + User = 1 + LoopCenter = 2 + BeamLocation = 3 + + +class ZoomModel(BaseModel): + z: dict[float, SampleCameraSettings] + + def get_camera_settings( + self, zoom_value: float + ) -> SampleCameraSettings: + if not self.z: + raise ValueError("No zoom data available") + + if zoom_value in self.z: + elem = self.z[zoom_value] + return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure) + + sorted_zooms = sorted(self.z.keys()) + + if zoom_value <= sorted_zooms[0]: + elem = self.z[sorted_zooms[0]] + return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure) + + if zoom_value >= sorted_zooms[-1]: + elem = self.z[sorted_zooms[-1]] + return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure) + print("interpolating zoom") + for i in range(len(sorted_zooms) - 1): + if sorted_zooms[i] <= zoom_value <= sorted_zooms[i + 1]: + lower_zoom = sorted_zooms[i] + upper_zoom = sorted_zooms[i + 1] + + lower_elem = self.z[lower_zoom] + upper_elem = self.z[upper_zoom] + + t = (zoom_value - lower_zoom) / (upper_zoom - lower_zoom) + + interpolated_gain = lower_elem.gain + t * ( + upper_elem.gain - lower_elem.gain + ) + interpolated_exp = lower_elem.exposure + t * ( + upper_elem.exposure - lower_elem.exposure + ) + + return SampleCameraSettings( + gain=interpolated_gain, exposure=interpolated_exp + ) + + closest = min(self.z.keys(), key=lambda x: abs(x - zoom_value)) + elem = self.z[closest] + return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure) + + +def def_zoom(beamline) -> ZoomModel: + print(f"user zoom for {beamline}") + if beamline == MXBeamline.X06DA: + return ZoomModel( + z={ + 1: SampleCameraSettings(gain=0, exposure=0.05), + 280: SampleCameraSettings(gain=0, exposure=0.05), + 500: SampleCameraSettings(gain=0, exposure=0.1), + 700: SampleCameraSettings(gain=0, exposure=0.15), + 800: SampleCameraSettings(gain=0, exposure=0.2), + 1000: SampleCameraSettings(gain=0, exposure=0.25), + } + ) + elif beamline == MXBeamline.X10SA: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) + elif beamline == MXBeamline.X06SA: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) + elif beamline == MXBeamline.SIMULATED: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) + else: + raise ValueError(f"Invalid beamline: {beamline}") + + +def def_bl_zoom(beamline) -> ZoomModel: + if beamline == MXBeamline.X06DA: + # Sensible starting presets so the beam is visible at every zoom; these + # are tuned live and persisted to Redis from the GUI (config.zoom_settings). + return ZoomModel( + z={ + 1: SampleCameraSettings(gain=0, exposure=0.05), + 280: SampleCameraSettings(gain=0, exposure=0.05), + 500: SampleCameraSettings(gain=0, exposure=0.1), + 700: SampleCameraSettings(gain=0, exposure=0.15), + 800: SampleCameraSettings(gain=0, exposure=0.2), + 1000: SampleCameraSettings(gain=0, exposure=0.25), + } + ) + elif beamline == MXBeamline.X10SA: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) + elif beamline == MXBeamline.X06SA: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) + elif beamline == MXBeamline.SIMULATED: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)}) + else: + raise ValueError(f"Invalid beamline: {beamline}") + + +def def_loop_centering_zoom(beamline) -> ZoomModel: + if beamline == MXBeamline.X06DA: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)}) + elif beamline == MXBeamline.X10SA: + return ZoomModel( + z={ + 1: SampleCameraSettings(gain=0, exposure=0.002), + 280: SampleCameraSettings(gain=0, exposure=0.002), + } + ) + elif beamline == MXBeamline.X06SA: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)}) + elif beamline == MXBeamline.SIMULATED: + return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)}) + else: + raise ValueError(f"Invalid beamline: {beamline}") + + +def zoom_manager( + mode: ZoomModeEnum = ZoomModeEnum.User, beamline: MXBeamline = None +) -> ZoomModel | None: + if beamline is None or beamline == MXBeamline.SIMULATED: + print("SIMULATED") + return def_zoom(beamline) + if mode == ZoomModeEnum.BeamLocation: + return def_bl_zoom(beamline) + elif mode == ZoomModeEnum.LoopCenter: + return def_loop_centering_zoom(beamline) + elif mode == ZoomModeEnum.User: + return def_zoom(beamline) + else: + raise ValueError(f"Invalid zoom mode: {mode}") + + +class AutofocusSettings(BaseModel): + center_x_pxl: float | None # Use beam center + center_y_pxl: float | None # Use beam center + radius_pxl: float + z_range_um: float + z_steps: int + + +class BeamlineStatus(BaseModel): + name: str + ring_current_mA: float + front_light: Annotated[float, Field(ge=0.0, le=100.0)] + back_light: Annotated[float, Field(ge=0.0, le=100.0)] + cryojet_K: float + shutter_open: bool + exp_shutter_open: bool | None + flux_ph_s: float + sample_camera: SampleCameraSettings + transmission: Annotated[float, Field(ge=0.0, le=1.0)] | None + zoom: float + commissioning_mode: bool + dtz_min: float + dtz_max: float + # Hutch personnel-safety system state. ``pss_prohibited`` is True when the + # hutch is interlocked so the robot may move (PROHIBITED-STATE); the GUI + # blocks a mount when it is False. ``pss_alarm`` is True when ALARM-STATE + # != 0 (warning). Defaults keep older payloads/constructors valid and avoid + # the GUI false-blocking when an old server omits the field. + pss_prohibited: bool = True + pss_alarm: bool = False + + +class SessionStatus(BaseModel): + session: SessionsStateEnum = SessionsStateEnum.Vacant + current_pgroup: str | None = None + staff: bool = False + + +class OpenGuiSessionInfo(BaseModel): + session: int + username: str + staff: bool = False + last_seen_ts: float + last_interaction_ts: float | None = None + close_requested: bool = False + close_requested_by: str | None = None + close_requested_at: float | None = None + close_grace_seconds: int | None = None + holds_baton: bool = False + + +class CrystalSize(BaseModel): + x: float = 0.0 + y: float = 0.0 + z: float = 0.0 + + +class DAQStatusModel(BaseModel): + geom: SampleGeometryModel + diffraction: DiffractionGeometry + bl: BeamlineStatus + state: BeamlineStateEnum + busy: bool + sample: SampleShortInfo | None = None + session: SessionStatus + open_guis: list[OpenGuiSessionInfo] = [] + box: BoundingBoxModel | None = None + last_best_res: float | None = None + last_best_b_factor: float | None = None + crystal_size: CrystalSize = CrystalSize(x=0, y=0, z=0) + + tell_connected: bool = True + tell_error: str | None = None + tell_state: TellStateModel | None = None + + smargon_connected: bool = True + smargon_error: str | None = None + + aerotech_connected: bool = True + aerotech_error: str | None = None + + +class BeamlineSettingsModel(BaseModel): + dtz_max: float | None = 1600.0 + dtz_min: float | None = 120.0 + dtz_collection: float | None = 130.0 + dtz_park: float | None = 150.0 + dtz_wash_sample_distance: float | None = 150.0 + dtz_bsz_safety_margin: float | None = 50.0 + bsz: float | None = 25.0 + + camera_max_magnification: float | None = 1.0 + camera_min_magnification: float | None = 500.0 + camera_translation_factor_a: float | None = 0.00253 + camera_translation_factor_b: float | None = 512.0 + + +class CryojetSettingsModel(BaseModel): + cryojet_park_position: float | None = 12.0 + cryojet_measurement_position: float | None = 5.0 + cryojet_in_use: bool | None = True + + +class SimpleStrategyInputModel(BaseModel): + last_best_res: float | None = None + last_best_b_factor: float | None = None + crystal_size: CrystalSize = CrystalSize(x=0, y=0, z=0) + angular_range: int = 360 + start_angle: float = 0 + incr_omega_deg: float = 0.2 + d_vis: float = 1.5 + current_temp_k: float = 100 + filename: Optional[str] = None + + +class SimpleScanParameters(BaseModel): + filename: Optional[str] = None + dtz: float = 150 + exp_time_s: float = 0.02 + start_omega_deg: float = 0 + incr_omega_deg: float = 0.2 + steps: int = 1800 + transmission: Annotated[float, Field(ge=0.0, le=1.0)] = 1.0 + last_best_res: float | None = None + last_best_b_factor: float | None = None + crystal_size: CrystalSize = CrystalSize(x=0, y=0, z=0) + flux_ph_s: float | None = None + calculated_dose_Mgy: float | None = None + calculated_dose_rate_MGy_s: float | None = None + xtal_size_dose_rate_MGy_s: float | None = None + target_dose_MGy: float | None = None + beam_size_x_um: float | None = None + beam_size_y_um: float | None = None + dose_rate_MGy_s: float | None = None + d_vis: float | None = None + d_tar: float | None = None + + +class ScanResultPayloadModel(BaseModel): + result: ScanResult + sample_id: int + attach_image: bool = True + beam_mark_pxl: tuple[float, float] + beam_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)] + + +class RecoveryActionRequest(BaseModel): + confirmation_code: str + + +@dataclass +class LoopCenteringResult: + success: bool + comment: str | None = None + error: Exception | None = None diff --git a/src/aarecommon/models/raster_grid.py b/src/aarecommon/models/raster_grid.py new file mode 100644 index 0000000..19ceadb --- /dev/null +++ b/src/aarecommon/models/raster_grid.py @@ -0,0 +1,95 @@ +from typing import Annotated, List + +import numpy as np +from aarecommon.math.coordinate import Coordinate, SmargonCoordinate, positive_coords +from aarecommon.math.sample_geometry import SampleGeometryModel +from jfjoch_client.models.scan_result import ScanResult +from pydantic import AfterValidator, BaseModel, Field + + +class RasterGridRequest(BaseModel): + dtz: float | None = None + transmission: float | None = None + file_prefix: str | None = None + exp_time_s: float + + # 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)] + + # Coordinates of top left corner in Smargon (SH...) coordinates. + # If None, measure from the current position. + smargon_top_left: SmargonCoordinate | None + + # omega angle + omega_deg: float = 0.0 + + visible: bool = True + + def get_image_number(self) -> int: + return self.n_x * self.n_y + + def grid_size_pxl(self, geom: SampleGeometryModel) -> Coordinate: + return self.grid_size_mm / geom.pixel_in_mm + + def start_pxl(self, geom: SampleGeometryModel) -> Coordinate: + if self.smargon_top_left is None: + return geom.smargon_to_beamline(geom.smargon.sh_mm) + return geom.smargon_to_beamline(self.smargon_top_left.sh_mm) + + +class CenterOfMassModel(BaseModel): + n_x: float + n_y: float + max_image: int | None = None + + @classmethod + def model_validate_maybe(cls, n_x: float, n_y: float) -> "CenterOfMassModel | None": + # Return None if either coordinate is NaN + if np.isnan(n_x) or np.isnan(n_y): + return None + return cls(n_x=float(n_x), n_y=float(n_y)) + + def get_com_mm(self, r: RasterGridRequest) -> Coordinate: + # defines COM in mm at beam position, relative to the top left corner of grid + return Coordinate( + x=(self.n_x + 0.5) * r.grid_size_mm.x, y=(self.n_y + 0.5) * r.grid_size_mm.y + ) + + def get_com_pxl( + self, r: RasterGridRequest, geom: SampleGeometryModel + ) -> Coordinate: + # defines COM in pixels at beam position, relative to the top left corner of grid + com_mm = self.get_com_mm(r) + return Coordinate(x=com_mm.x / geom.pixel_in_mm, y=com_mm.y / geom.pixel_in_mm) + + def com_nan_check(self): + # True if both values are finite (not NaN) + return not (np.isnan(self.n_x) or np.isnan(self.n_y)) + + +class CompletedRasterGridElem(BaseModel): + request: RasterGridRequest + result: ScanResult + centre_of_mass: CenterOfMassModel | None + + +class CompletedRasterGrid(BaseModel): + r: List[CompletedRasterGridElem] + + +class RasterPayloadModel(BaseModel): + request: RasterGridRequest + result: ScanResult + sample_id: int + attach_image: bool = True + centre_of_mass: CenterOfMassModel | None = None + raster_score: List[float | None] | None = None + center_pxl: Coordinate | None + start_pxl: Coordinate + cell_size_pxl: Annotated[Coordinate, AfterValidator(positive_coords)] + beam_mark_pxl: tuple[float, float] + beam_size_mm: Annotated[Coordinate, AfterValidator(positive_coords)] diff --git a/src/aarecommon/models/rotation_scan.py b/src/aarecommon/models/rotation_scan.py new file mode 100644 index 0000000..c71e218 --- /dev/null +++ b/src/aarecommon/models/rotation_scan.py @@ -0,0 +1,25 @@ +from aarecommon.math.coordinate import SmargonCoordinate +from jfjoch_client.models.scan_result import ScanResult +from pydantic import BaseModel + + +class RotationScanRequest(BaseModel): + dtz: float | None = None + transmission: float | None = None + file_prefix: str | None = None + exp_time_s: float + + start: SmargonCoordinate | None = None # None = start where you are + end: SmargonCoordinate | None = None + + start_omega_deg: float = 0 + incr_omega_deg: float + steps: int + + screening: bool = False + wedge_omega_deg: float | None = None + + +class CompletedRotationScan(BaseModel): + request: RotationScanRequest + result: ScanResult diff --git a/src/aarecommon/models/tell.py b/src/aarecommon/models/tell.py new file mode 100644 index 0000000..80e27e3 --- /dev/null +++ b/src/aarecommon/models/tell.py @@ -0,0 +1,62 @@ +from enum import Enum + +from pydantic import BaseModel + + +class TellActivityEnum(str, Enum): + IDLE = "idle" + MOUNTING = "mounting" + UNMOUNTING = "unmounting" + DRYING = "drying" + COOLING = "cooling" + ERROR = "error" + + def display_name(self) -> str: + return { + TellActivityEnum.IDLE: "Idle", + TellActivityEnum.MOUNTING: "Mounting", + TellActivityEnum.UNMOUNTING: "Unmounting", + TellActivityEnum.DRYING: "Drying", + TellActivityEnum.COOLING: "Cooling", + TellActivityEnum.ERROR: "Error", + }.get(self, str(self.value).capitalize()) + + +class TellPhaseEnum(str, Enum): + IDLE = "idle" + PREPARING = "preparing" + AUTO_UNMOUNT = "auto_unmount" + RETURNING_OLD_SAMPLE = "returning_old_sample" + OLD_SAMPLE_RETURNED = "old_sample_returned" + PICKING_NEW_SAMPLE = "picking_new_sample" + PLACING_NEW_SAMPLE = "placing_new_sample" + FINALIZING = "finalizing" + COMPLETE = "complete" + FAILED = "failed" + + def display_name(self) -> str: + return { + TellPhaseEnum.IDLE: "Idle", + TellPhaseEnum.PREPARING: "Preparing", + TellPhaseEnum.AUTO_UNMOUNT: "Auto-unmount", + TellPhaseEnum.RETURNING_OLD_SAMPLE: "Returning old sample", + TellPhaseEnum.OLD_SAMPLE_RETURNED: "Old sample returned", + TellPhaseEnum.PICKING_NEW_SAMPLE: "Picking new sample", + TellPhaseEnum.PLACING_NEW_SAMPLE: "Placing sample on gonio", + TellPhaseEnum.FINALIZING: "Finalizing", + TellPhaseEnum.COMPLETE: "Complete", + TellPhaseEnum.FAILED: "Failed", + }.get(self, str(self.value).replace("_", " ").capitalize()) + + +class TellStateModel(BaseModel): + activity: TellActivityEnum = TellActivityEnum.IDLE + message: str | None = None + last_event_class: str | None = None + last_event_value: str | None = None + last_update_ts: str | None = None + mount_success: bool | None = None + mount_error: str | None = None + sample_position: str | None = None + operation: str | None = None + phase: TellPhaseEnum | None = None diff --git a/src/aarecommon/recurrence_watcher.py b/src/aarecommon/recurrence_watcher.py new file mode 100644 index 0000000..43750d6 --- /dev/null +++ b/src/aarecommon/recurrence_watcher.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, TypeAlias + +from aarecommon.errors import exception_handler +from aarecommon.errors.exception_handler import ( + AareException, + LoopCenteringFailed, + SmargonException, + TellException, + TransformationInvalidException, +) + +ExceptionType: TypeAlias = type[AareException] + + +@dataclass(frozen=True) +class WatcherTrip: + name: str + threshold: int + streak: int + exception_class: str + + def message(self) -> str: + return ( + f"{self.streak} consecutive {self.name} errors across samples " + f"({self.exception_class}) — automation halted" + ) + + +class RecurrenceWatcher: + def __init__(self, *, name: str, observes: ExceptionType, threshold: int): + self.name = name + self._observes = observes + self.threshold = max(1, int(threshold)) + self._streak = 0 + + @property + def streak(self) -> int: + return self._streak + + def reset(self) -> None: + self._streak = 0 + + def observe(self, exception_class: type | None) -> bool: + if exception_class is None: + self._streak = 0 + return False + if not isinstance(exception_class, type) or not issubclass( + exception_class, self._observes + ): + return False + self._streak += 1 + return self._streak >= self.threshold + + def maybe_trip(self, exception_class: type | None) -> WatcherTrip | None: + if not self.observe(exception_class): + return None + class_name = ( + exception_class.__name__ if isinstance(exception_class, type) else "Unknown" + ) + return WatcherTrip( + name=self.name, + threshold=self.threshold, + streak=self._streak, + exception_class=class_name, + ) + + +DEFAULT_WATCHERS: tuple[tuple[str, ExceptionType, int], ...] = ( + ("tell", TellException, 5), + ("alc", LoopCenteringFailed, 3), + ("transformation", TransformationInvalidException, 3), + ("smargon", SmargonException, 3), +) + + +def create_default_watchers( + overrides: dict[str, int] | None = None, +) -> list[RecurrenceWatcher]: + thresholds = dict(overrides or {}) + return [ + RecurrenceWatcher( + name=name, + observes=observes, + threshold=thresholds.get(name, threshold), + ) + for name, observes, threshold in DEFAULT_WATCHERS + ] + + +def redis_key_to_env_var(key: str) -> str: + """Adapter from canonical Redis key shape to an env-var-safe name. + + ``aare:watchers:{bl}:{watcher}:threshold`` → + ``AARE_WATCHERS_{BL}_{WATCHER}_THRESHOLD``. + + Used by the GUI today since it has no Redis client; the loader's + ``get_value`` argument keeps the Redis key shape authoritative so the + DAQ side can later swap in ``cfg.get`` without churn (plan §7).""" + return key.replace(":", "_").upper() + + +def load_watcher_threshold_overrides( + *, + beamline: str, + watcher_names: list[str], + get_value: Callable[[str], str | bytes | None], +) -> dict[str, int]: + overrides: dict[str, int] = {} + for watcher_name in watcher_names: + key = f"aare:watchers:{beamline}:{watcher_name}:threshold" + raw_value = get_value(key) + if raw_value is None: + continue + if isinstance(raw_value, bytes): + raw_value = raw_value.decode("utf-8", errors="ignore") + try: + parsed = int(str(raw_value).strip()) + except (TypeError, ValueError): + continue + if parsed > 0: + overrides[watcher_name] = parsed + return overrides + + +_EXCEPTION_CLASS_BY_NAME: dict[str, type] = { + name: klass + for name, klass in vars(exception_handler).items() + if isinstance(klass, type) and issubclass(klass, AareException) +} + + +def resolve_exception_class(exception_class_name: str | None) -> type | None: + if not exception_class_name: + return None + return _EXCEPTION_CLASS_BY_NAME.get(exception_class_name) diff --git a/tests/test_aare_exception.py b/tests/test_aare_exception.py new file mode 100644 index 0000000..82ce9be --- /dev/null +++ b/tests/test_aare_exception.py @@ -0,0 +1,258 @@ +"""Tests for the AareException hierarchy introduced in Phase 1 of the +exception-handling redesign. + +These cover: +- class-level ``critical`` default +- instance-level override via ``critical=...`` kwarg +- correct parent/family relationships so watchers can match by family +""" + +from __future__ import annotations + +import pytest +from aarecommon.errors.exception_handler import ( + AareAuthError, + AareDBCommunicationError, + AareException, + AareUserError, + AerotechCommunicationError, + AerotechException, + AuthenticationException, + AutomationError, + AutoRasterSampleSkipped, + AXCFailed, + BeamlineBusyException, + BeamlineBusyTimeoutException, + BeamlineStateException, + BECCommunicationError, + BECException, + CriticalTellException, + DataCollectionException, + JFJochCommunicationError, + JFJochException, + LoopCenteringFailed, + MagnetPositionSensorErorr, + MaintenanceStateException, + ManualMountException, + MountingFailed, + RasterScanException, + SampleException, + SmargonCommunicationError, + SmargonException, + SmartMagnetFaultException, + StateTransitionFailed, + TellCommandWhileBusyException, + TellCommunicationError, + TellConnectionException, + TellException, + TransformationInvalidException, + UnmountingFailed, + UserRightsException, + WarningTellException, +) + +# --------------------------------------------------------------------------- +# Class-level criticality defaults +# --------------------------------------------------------------------------- + + +def test_aare_exception_class_critical_default_false(): + assert AareException.critical is False + + +def test_phase3_classes_have_critical_default_true(): + for cls in ( + TellCommunicationError, + TellConnectionException, + CriticalTellException, + SmargonCommunicationError, + AerotechCommunicationError, + JFJochCommunicationError, + BECCommunicationError, + StateTransitionFailed, + MaintenanceStateException, + BeamlineBusyException, + BeamlineBusyTimeoutException, + MagnetPositionSensorErorr, + SmartMagnetFaultException, + TransformationInvalidException, + ): + assert cls.critical is True, ( + f"{cls.__name__} should default to critical=True in Phase 3" + ) + + +def test_other_core_classes_keep_critical_default_false(): + for cls in ( + AutomationError, + AareUserError, + AareAuthError, + TellException, + SmargonException, + AerotechException, + JFJochException, + BECException, + BeamlineStateException, + MountingFailed, + UnmountingFailed, + LoopCenteringFailed, + TellCommandWhileBusyException, + WarningTellException, + AareDBCommunicationError, + ): + assert cls.critical is False, f"{cls.__name__} should remain critical=False" + + +# --------------------------------------------------------------------------- +# Instance-level critical override +# --------------------------------------------------------------------------- + + +def test_instance_critical_override_true(): + exc = MountingFailed("mount fail", critical=True) + assert exc.critical is True + # class default unchanged + assert MountingFailed.critical is False + + +def test_instance_critical_override_false(): + exc = MountingFailed("mount fail", critical=False) + assert exc.critical is False + + +def test_instance_critical_none_keeps_class_default(): + # Omitting critical uses class default + exc = MountingFailed("mount fail") + assert exc.critical is False # class default + + +def test_aaredb_communication_error_instance_critical_override(): + exc = AareDBCommunicationError("db down", critical=True, operation="GET") + assert exc.critical is True + assert exc.operation == "GET" + + +# --------------------------------------------------------------------------- +# Routing-root membership +# --------------------------------------------------------------------------- + + +def test_automation_errors_are_automation_error(): + for cls in ( + LoopCenteringFailed, + MountingFailed, + UnmountingFailed, + AXCFailed, + TellCommunicationError, + SmargonCommunicationError, + AerotechCommunicationError, + JFJochCommunicationError, + BECCommunicationError, + AareDBCommunicationError, + BeamlineBusyException, + RasterScanException, + MagnetPositionSensorErorr, + SmartMagnetFaultException, + TransformationInvalidException, + AutoRasterSampleSkipped, + ): + exc = cls() if cls is not AutoRasterSampleSkipped else cls("skipped") + assert isinstance(exc, AutomationError), ( + f"{cls.__name__} should be AutomationError" + ) + assert isinstance(exc, AareException), f"{cls.__name__} should be AareException" + + +def test_auth_errors_are_aare_auth_error(): + for cls in (AuthenticationException, UserRightsException): + exc = cls() + assert isinstance(exc, AareAuthError) + assert isinstance(exc, AareException) + + +def test_user_errors_are_aare_user_error(): + for cls in (ManualMountException, SampleException): + exc = cls() + assert isinstance(exc, AareUserError) + assert isinstance(exc, AareException) + + +# --------------------------------------------------------------------------- +# Family-base membership (drives watcher matching) +# --------------------------------------------------------------------------- + + +def test_tell_family_membership(): + for cls in ( + TellCommunicationError, + TellConnectionException, + CriticalTellException, + TellCommandWhileBusyException, + WarningTellException, + MountingFailed, + UnmountingFailed, + ): + exc = cls() + assert isinstance(exc, TellException), ( + f"{cls.__name__} should be in Tell family" + ) + assert isinstance(exc, AutomationError) + + +def test_smargon_family_membership(): + exc = SmargonCommunicationError() + assert isinstance(exc, SmargonException) + assert isinstance(exc, AutomationError) + + +def test_aerotech_family_membership(): + exc = AerotechCommunicationError() + assert isinstance(exc, AerotechException) + assert isinstance(exc, AutomationError) + + +def test_jfjoch_family_membership(): + exc = JFJochCommunicationError() + assert isinstance(exc, JFJochException) + assert isinstance(exc, AutomationError) + + +def test_bec_family_membership(): + exc = BECCommunicationError() + assert isinstance(exc, BECException) + assert isinstance(exc, AutomationError) + + +def test_beamline_state_family_membership(): + for cls in ( + StateTransitionFailed, + MaintenanceStateException, + BeamlineBusyException, + BeamlineBusyTimeoutException, + ): + exc = cls() + assert isinstance(exc, BeamlineStateException) + assert isinstance(exc, AutomationError) + + +def test_data_collection_family_membership(): + exc = RasterScanException() + assert isinstance(exc, DataCollectionException) + assert isinstance(exc, AutomationError) + + +# --------------------------------------------------------------------------- +# Backward compat: existing __str__/.message contracts preserved +# --------------------------------------------------------------------------- + + +def test_message_attribute_preserved(): + exc = LoopCenteringFailed("custom message") + assert exc.message == "custom message" + assert str(exc) == "custom message" + + +def test_default_messages_preserved(): + assert str(MountingFailed()) == "A sample was not mounted" + assert str(UnmountingFailed()) == "A sample was not unmounted" + assert str(LoopCenteringFailed()) == "Loop Centering did not detect a sample" diff --git a/tests/test_aerotech_models.py b/tests/test_aerotech_models.py new file mode 100644 index 0000000..e682e6b --- /dev/null +++ b/tests/test_aerotech_models.py @@ -0,0 +1,84 @@ +from aarecommon.models.aerotech import ( + AerotechAxisStatus, + AerotechRotationScanRequest, + AerotechRunEnum, + AerotechStatus, + AerotechTarget, + AxisEnum, + TaskEnum, + VariableTypeEnum, +) + + +def test_enums(): + assert TaskEnum.TASK_0.value == 0 + assert AxisEnum.X.value == "x" + assert AerotechRunEnum.START.value == 1 + assert VariableTypeEnum.REAL.value == 1 + + +def test_aerotech_axis_status(): + status = AerotechAxisStatus( + enabled=True, + fault=0, + homed=True, + is_fault=False, + moving=False, + position=10.0, + status=1, + velocity=0.0, + ) + assert status.position == 10.0 + assert status.enabled is True + + +def test_aerotech_status_strings(): + axis_status = AerotechAxisStatus( + enabled=True, + fault=0, + homed=True, + is_fault=False, + moving=False, + position=1.234567, + status=1, + velocity=0.1, + ) + status = AerotechStatus(state="READY", x=axis_status) + + pretty = status.to_pretty_string() + assert "STATE: READY" in pretty + assert " X | pos= 1.234567" in pretty + assert "Y: unavailable" in pretty + + compact = status.to_compact_string() + assert "state=READY" in compact + assert "x=1.2346 (H, -)" in compact + + colored = status.to_colored_string() + assert "STATE:" in colored + assert "READY" in colored + assert "pos=" in colored + + assert str(status) == pretty + + +def test_aerotech_target(): + target = AerotechTarget(x=10.0, y=20.0) + payload = target.to_payload() + assert payload == {"x": 10.0, "y": 20.0} + assert "z" not in payload + + +def test_aerotech_rotation_scan_request(): + request = AerotechRotationScanRequest( + rotation_deg=360.0, + time_sec=10.0, + start_pos_deg=0.0, + async_move=True, + exp_time_s=0.1, + incr_omega_deg=1.0, + steps=360, + ) + payload = request.to_payload() + assert payload["rotation_deg"] == 360.0 + assert payload["async"] is True diff --git a/tests/test_autofocus_tools.py b/tests/test_autofocus_tools.py new file mode 100644 index 0000000..97d64f6 --- /dev/null +++ b/tests/test_autofocus_tools.py @@ -0,0 +1,80 @@ +import numpy as np +from aarecommon.math.autofocus import focus_measure_blob_size, focus_measure_edges + + +def test_focus_measure_edges_all_zeros(): + gray = np.zeros((100, 100), dtype=np.uint8) + assert focus_measure_edges(gray) == 0.0 + + +def test_focus_measure_edges_sharp_vs_blurry(): + # Sharp image (large blocks to survive GaussianBlur) + sharp = np.zeros((100, 100), dtype=np.uint8) + sharp[:, 0:50] = 255 + + # Blurry image (flat) + blurry = np.full((100, 100), 128, dtype=np.uint8) + + fm_sharp = focus_measure_edges(sharp, verbose=True) + fm_blurry = focus_measure_edges(blurry) + + assert fm_sharp > fm_blurry + assert fm_blurry == 0.0 + + +def test_focus_measure_edges_with_mask(): + gray = np.zeros((100, 100), dtype=np.uint8) + gray[40:60, 40:60] = 255 + + mask = np.zeros((100, 100), dtype=bool) + mask[40:60, 40:60] = True + + fm_with_mask = focus_measure_edges(gray, mask=mask) + assert fm_with_mask > 0 + + empty_mask = np.zeros((100, 100), dtype=bool) + assert focus_measure_edges(gray, mask=empty_mask) == 0.0 + + +def test_focus_measure_edges_verbose(capsys): + gray = np.zeros((100, 100), dtype=np.uint8) + gray[40:60, 40:60] = 255 + mask = np.ones((100, 100), dtype=bool) + + focus_measure_edges(gray, mask=mask, verbose=True) + captured = capsys.readouterr() + assert "focus=" in captured.out + + +def test_focus_measure_blob_size_all_zeros(): + gray = np.zeros((100, 100), dtype=np.uint8) + assert focus_measure_blob_size(gray) == 0.0 + + +def test_focus_measure_blob_size_sharp_vs_blurry(): + # Small sharp blob + sharp = np.zeros((100, 100), dtype=np.uint8) + sharp[50, 50] = 255 + + # Larger blurry blob + blurry = np.zeros((100, 100), dtype=np.uint8) + blurry[45:55, 45:55] = 255 + + fm_sharp = focus_measure_blob_size(sharp) + fm_blurry = focus_measure_blob_size(blurry) + + assert fm_sharp > fm_blurry + + +def test_focus_measure_blob_size_with_mask(): + gray = np.zeros((100, 100), dtype=np.uint8) + gray[50, 50] = 255 + + mask = np.zeros((100, 100), dtype=bool) + mask[50, 50] = True + + fm_with_mask = focus_measure_blob_size(gray, mask=mask) + assert fm_with_mask > 0 + + empty_mask = np.zeros((100, 100), dtype=bool) + assert focus_measure_blob_size(gray, mask=empty_mask) == 0.0 diff --git a/tests/test_beamline.py b/tests/test_beamline.py new file mode 100644 index 0000000..acc8ea5 --- /dev/null +++ b/tests/test_beamline.py @@ -0,0 +1,26 @@ +import os +from unittest.mock import patch + +from aarecommon.models.beamline import MXBeamline +from aarecommon.config.beamline import mx_beamline + + +def test_mx_beamline_default(): + with patch.dict(os.environ, {}, clear=True): + # If BEAMLINE is not set, it should return SIMULATED + assert mx_beamline() == MXBeamline.SIMULATED + + +def test_mx_beamline_x10sa(): + with patch.dict(os.environ, {"BEAMLINE": "x10sa"}): + assert mx_beamline() == MXBeamline.X10SA + + +def test_mx_beamline_x06sa(): + with patch.dict(os.environ, {"BEAMLINE": "X06SA "}): + assert mx_beamline() == MXBeamline.X06SA + + +def test_mx_beamline_invalid(): + with patch.dict(os.environ, {"BEAMLINE": "INVALID"}): + assert mx_beamline() == MXBeamline.SIMULATED diff --git a/tests/test_coordinate.py b/tests/test_coordinate.py new file mode 100644 index 0000000..c43f073 --- /dev/null +++ b/tests/test_coordinate.py @@ -0,0 +1,153 @@ +import numpy as np +import pytest +from aarecommon.math.coordinate import ( + AerotechCoordinate, + Coordinate, + SmargonCoordinate, + positive_coords, +) + + +def test_coordinate_addition(): + c1 = Coordinate(x=1, y=2, z=3) + c2 = Coordinate(x=4, y=5, z=6) + c3 = c1 + c2 + assert c3.x == 5 + assert c3.y == 7 + assert c3.z == 9 + + +def test_coordinate_subtraction(): + c1 = Coordinate(x=5, y=7, z=9) + c2 = Coordinate(x=1, y=2, z=3) + c3 = c1 - c2 + assert c3.x == 4 + assert c3.y == 5 + assert c3.z == 6 + + +def test_coordinate_multiplication_scalar(): + c1 = Coordinate(x=1, y=2, z=3) + c2 = c1 * 2.0 + assert c2.x == 2.0 + assert c2.y == 4.0 + assert c2.z == 6.0 + + +def test_coordinate_dot_product(): + c1 = Coordinate(x=1, y=2, z=3) + c2 = Coordinate(x=4, y=5, z=6) + dot = c1 * c2 + assert dot == (1 * 4 + 2 * 5 + 3 * 6) + + +def test_coordinate_division(): + c1 = Coordinate(x=2, y=4, z=6) + c2 = c1 / 2.0 + assert c2.x == 1.0 + assert c2.y == 2.0 + assert c2.z == 3.0 + + +def test_coordinate_division_by_zero(): + c1 = Coordinate(x=2, y=4, z=6) + with pytest.raises(ValueError, match="Cannot divide by zero"): + _ = c1 / 0.0 + + +def test_coordinate_normalize(): + c1 = Coordinate(x=3, y=0, z=4) + c2 = c1.normalize() + assert c2.x == 0.6 + assert c2.y == 0.0 + assert c2.z == 0.8 + + +def test_coordinate_normalize_zero(): + c1 = Coordinate(x=0, y=0, z=0) + with pytest.raises(ValueError, match="Cannot normalize a zero-magnitude vector."): + c1.normalize() + + +def test_coordinate_rotate_x(): + c1 = Coordinate(x=1, y=1, z=0) + # Rotate 90 degrees around X. (1, 1, 0) -> (1, 0, 1) + c2 = c1.rotate(90, "x") + assert np.isclose(c2.x, 1) + assert np.isclose(c2.y, 0) + assert np.isclose(c2.z, 1) + + +def test_coordinate_rotate_y(): + c1 = Coordinate(x=1, y=0, z=1) + # Rotate 90 degrees around Y. (1, 0, 1) -> (1, 0, -1) + c2 = c1.rotate(90, "y") + assert np.isclose(c2.x, 1) + assert np.isclose(c2.y, 0) + assert np.isclose(c2.z, -1) + + +def test_coordinate_rotate_z(): + c1 = Coordinate(x=1, y=0, z=0) + # Rotate 90 degrees around Z. (1, 0, 0) -> (0, 1, 0) + c2 = c1.rotate(90, "z") + assert np.isclose(c2.x, 0) + assert np.isclose(c2.y, 1) + assert np.isclose(c2.z, 0) + + +def test_coordinate_rotate_invalid_axis(): + c1 = Coordinate(x=1, y=1, z=1) + with pytest.raises(ValueError, match="Invalid axis"): + c1.rotate(90, "w") + + +def test_smargon_coordinate_equality(): + s1 = SmargonCoordinate(sh_mm=Coordinate(x=1, y=1, z=1), phi_deg=10, chi_deg=20) + s2 = SmargonCoordinate( + sh_mm=Coordinate(x=1.05, y=0.95, z=1.01), phi_deg=10.05, chi_deg=19.95 + ) + assert s1 == s2 + + s3 = SmargonCoordinate(sh_mm=Coordinate(x=2, y=1, z=1), phi_deg=10, chi_deg=20) + assert s1 != s3 + + +def test_aerotech_coordinate_equality(): + a1 = AerotechCoordinate(at_mm=Coordinate(x=1, y=1, z=1), omega_deg=10) + a2 = AerotechCoordinate( + at_mm=Coordinate(x=1.005, y=0.995, z=1.001), omega_deg=10.005 + ) + assert a1 == a2 + + a3 = AerotechCoordinate(at_mm=Coordinate(x=1.1, y=1, z=1), omega_deg=10) + assert a1 != a3 + + +def test_positive_coords(): + c1 = Coordinate(x=1, y=1, z=1) + assert positive_coords(c1) == c1 + + with pytest.raises(ValueError, match="Coordinates must be positive"): + positive_coords(Coordinate(x=-1, y=1, z=1)) + + with pytest.raises(ValueError, match="Coordinates must be positive"): + positive_coords(Coordinate(x=1, y=-1, z=1)) + + +def test_coordinate_unsupported_ops(): + c1 = Coordinate(x=1, y=1, z=1) + assert c1.__add__(1) == NotImplemented + assert c1.__sub__(1) == NotImplemented + assert c1.__mul__("string") == NotImplemented + assert c1.__truediv__("string") == NotImplemented + + +def test_smargon_coordinate_eq_not_implemented(): + s1 = SmargonCoordinate(sh_mm=Coordinate(x=1, y=1, z=1), phi_deg=10, chi_deg=20) + assert s1.__eq__(1) == NotImplemented + + +def test_aerotech_coordinate_eq_not_implemented(): + a1 = AerotechCoordinate(at_mm=Coordinate(x=1, y=1, z=1), omega_deg=10) + assert a1.__eq__(1) == NotImplemented diff --git a/tests/test_data_collection_parameters.py b/tests/test_data_collection_parameters.py new file mode 100644 index 0000000..15362a4 --- /dev/null +++ b/tests/test_data_collection_parameters.py @@ -0,0 +1,74 @@ +import pytest +from aarecommon.models.models import DataCollectionParameters +from pydantic import ValidationError + + +def test_directory_defaults_when_missing(): + params = DataCollectionParameters() + assert params.directory is None + + +def test_directory_blank_defaults_to_macro_path(): + params = DataCollectionParameters(directory="") + assert params.directory == "{date}/{prefix}" + + +def test_directory_spaces_are_replaced(): + params = DataCollectionParameters(directory="my folder/run 1") + assert params.directory == "my_folder/run_1" + + +def test_directory_rejects_invalid_characters(): + with pytest.raises(ValidationError): + DataCollectionParameters(directory="bad|path") + + +def test_exposure_accepts_values_above_1_after_refactor(): + params = DataCollectionParameters(exposure=1.5) + assert params.exposure == 1.5 + + +def test_cloud_blank_defaults_to_true(): + params = DataCollectionParameters(cloud="") + assert params.cloud is True + + +def test_directory_accepts_valid_macros(): + params = DataCollectionParameters(directory="{date}/{prefix}/run") + assert params.directory == "{date}/{prefix}/run" + + +def test_aperture_accepts_float_string(): + params = DataCollectionParameters(aperture="2.0") + assert params.aperture == 2 + + +def test_processingpipeline_accepts_unknown_value_after_refactor(): + params = DataCollectionParameters(processingpipeline="xia2") + assert params.processingpipeline == "xia2" + + +def test_datacollectionparameters_accepts_legacy_aliases(): + params = DataCollectionParameters( + totalrange=180, + cellparameters="10 20 30 90 90 120", + userresolution=1.4, + ) + assert params.totalangle == 180 + assert params.unitcell == "10 20 30 90 90 120" + assert params.processingresolution == 1.4 + + +def test_datacollectionparameters_accepts_new_fields(): + params = DataCollectionParameters( + totalangle=90, + unitcell="11,22,33,90,90,120", + processingresolution=1.2, + pdbmodel="model.pdb", + cloud=False, + ) + assert params.totalangle == 90 + assert params.unitcell == "11,22,33,90,90,120" + assert params.processingresolution == 1.2 + assert params.pdbmodel == "model.pdb" + assert params.cloud is False diff --git a/tests/test_diffraction_geometry.py b/tests/test_diffraction_geometry.py new file mode 100644 index 0000000..2cf3c37 --- /dev/null +++ b/tests/test_diffraction_geometry.py @@ -0,0 +1,68 @@ +import pytest +from aarecommon.math.diffraction_geometry import DiffractionGeometry + + +@pytest.fixture +def sample_dg(): + return DiffractionGeometry( + energy_keV=12.4, + dtz_mm=100.0, + pixel_size_mm=0.172, + beam_center_pxl=(1000.0, 1000.0), + detector_size_pxl=(2000, 2000), + detector_description="Eiger 16M", + detector_serial_number="E-123", + poni_rot1_rad=0.0, + poni_rot2_rad=0.0, + ) + + +def test_detector_max_radius_pxl(sample_dg): + # center (1000, 1000), size (2000, 2000) + # x0 = 2000-1000 = 1000 + # x1 = 1000 + # y0 = 2000-1000 = 1000 + # y1 = 1000 + # max = 1000 + assert sample_dg.detector_max_radius_pxl == 1000.0 + + +def test_detector_radius_mm(sample_dg): + # 1000 * 0.172 = 172.0 + assert sample_dg.detector_radius_mm == 172.0 + + +def test_wavelength_angstrom(sample_dg): + # 12.398 / 12.4 = 0.9998387... + assert sample_dg.wavelength_angstrom == pytest.approx(0.9998387) + + +def test_resolution_angstrom(sample_dg): + # dtz = 100 + # radius = 172 + # theta = atan(172/100) * 0.5 = atan(1.72) * 0.5 = 1.044 * 0.5 = 0.522 rad + # res = 0.9998 / (2 * sin(0.522)) = 0.9998 / (2 * 0.498) = 1.003 + res = sample_dg.resolution_angstrom(100.0) + assert res > 0 + assert res == pytest.approx(1.002469, abs=1e-5) + + with pytest.raises(ValueError): + sample_dg.resolution_angstrom(0) + + +def test_max_resolution_angstrom(sample_dg): + assert sample_dg.max_resolution_angstrom == sample_dg.resolution_angstrom( + sample_dg.dtz_mm + ) + + +def test_calc_dtz_mm(sample_dg): + res = sample_dg.resolution_angstrom(100.0) + dtz = sample_dg.calc_dtz_mm(res) + assert dtz == pytest.approx(100.0) + + with pytest.raises(ValueError): + sample_dg.calc_dtz_mm(-1) + + # test x >= 1.0 case: wavelength / (2*res) >= 1.0 -> res <= wavelength / 2 + assert sample_dg.calc_dtz_mm(0.0001) == 0.0 diff --git a/tests/test_error_codes.py b/tests/test_error_codes.py new file mode 100644 index 0000000..80e41ca --- /dev/null +++ b/tests/test_error_codes.py @@ -0,0 +1,124 @@ +from aarecommon.errors.codes import ( + AareErrorCode, + AuthErrorCode, + code_for_exception_class, + error_code_help, + export_error_codes, + export_error_codes_grouped, +) + + +def test_error_code_help_returns_string(): + help_text = error_code_help(AuthErrorCode.AUTHENTICATION_FAILED) + assert isinstance(help_text, str) + assert len(help_text) > 0 + + +def test_error_code_help_unknown_code(): + help_text = error_code_help("UNKNOWN_CODE") + assert help_text is None + + +def test_export_error_codes_contains_known_codes(): + exported = export_error_codes() + assert AuthErrorCode.AUTHENTICATION_FAILED.name in exported + + +def test_code_for_exception_class_basic(): + assert ( + code_for_exception_class("TellCommunicationError") == "TELL_COMMUNICATION_ERROR" + ) + assert code_for_exception_class("MountingFailed") == "MOUNTING_FAILED" + assert code_for_exception_class("LoopCenteringFailed") == "LOOP_CENTERING_FAILED" + + +def test_code_for_exception_class_acronyms(): + assert code_for_exception_class("AXCFailed") == "AXC_FAILED" + assert ( + code_for_exception_class("AareDBCommunicationError") + == "AARE_DB_COMMUNICATION_ERROR" + ) + assert ( + code_for_exception_class("JFJochCommunicationError") + == "JF_JOCH_COMMUNICATION_ERROR" + ) + + +def test_code_for_each_concrete_exception_is_in_aare_error_code_enum(): + """Every code we'd produce from the rebuilt hierarchy must exist in the + AareErrorCode enum -- otherwise clients have no symbol to branch on.""" + from aarecommon.errors.exception_handler import ( + AareDBCommunicationError, + AerotechCommunicationError, + AuthenticationException, + AutoRasterSampleSkipped, + AXCFailed, + BeamlineBusyException, + BeamlineBusyTimeoutException, + BECCommunicationError, + CriticalTellException, + DataCollectionException, + JFJochCommunicationError, + LoopCenteringFailed, + MagnetPositionSensorErorr, + MaintenanceStateException, + ManualMountException, + MountingFailed, + RasterScanException, + SampleException, + SmargonCommunicationError, + SmartMagnetFaultException, + StateTransitionFailed, + TellCommandWhileBusyException, + TellCommunicationError, + TellConnectionException, + TransformationInvalidException, + UnmountingFailed, + UserRightsException, + WarningTellException, + ) + + valid = {c.value for c in AareErrorCode} + classes = [ + TellCommunicationError, + TellConnectionException, + CriticalTellException, + WarningTellException, + TellCommandWhileBusyException, + MountingFailed, + UnmountingFailed, + SmargonCommunicationError, + AerotechCommunicationError, + JFJochCommunicationError, + BECCommunicationError, + AareDBCommunicationError, + StateTransitionFailed, + MaintenanceStateException, + BeamlineBusyException, + BeamlineBusyTimeoutException, + DataCollectionException, + RasterScanException, + LoopCenteringFailed, + AXCFailed, + AutoRasterSampleSkipped, + TransformationInvalidException, + MagnetPositionSensorErorr, + SmartMagnetFaultException, + ManualMountException, + SampleException, + AuthenticationException, + UserRightsException, + ] + missing = [] + for cls in classes: + code = code_for_exception_class(cls.__name__) + if code not in valid: + missing.append((cls.__name__, code)) + assert not missing, f"Codes missing from AareErrorCode: {missing}" + + +def test_export_error_codes_grouped_includes_aare_error_code(): + exported = export_error_codes_grouped() + assert "AareErrorCode" in exported + assert "AuthErrorCode" in exported + assert "TELL_COMMUNICATION_ERROR" in exported["AareErrorCode"] diff --git a/tests/test_exception_handler.py b/tests/test_exception_handler.py new file mode 100644 index 0000000..bbbac75 --- /dev/null +++ b/tests/test_exception_handler.py @@ -0,0 +1,196 @@ +import pytest +from aarecommon.errors.exception_handler import ( + AareDBCommunicationError, + AerotechCommunicationError, + AuthenticationException, + AuthErrorCode, + AXCFailed, + BeamlineBusyException, + CriticalTellException, + DataCollectionException, + JFJochCommunicationError, + LoopCenteringFailed, + MagnetPositionSensorErorr, + ManualMountException, + MountingFailed, + RasterScanException, + SampleException, + SmargonCommunicationError, + SmartMagnetFaultException, + TellCommandWhileBusyException, + TellCommunicationError, + TellConnectionException, + TransformationInvalidException, + UnmountingFailed, + UserRightsException, + WarningTellException, +) + + +def test_data_collection_exception_message(): + exc = DataCollectionException("Custom error") + assert str(exc) == "Custom error" + + +def test_data_collection_exception_default_message(): + exc = DataCollectionException() + assert str(exc) == "Data collection failed" + + +def test_authentication_exception_properties(): + exc = AuthenticationException( + "Failed", status_code=403, code=AuthErrorCode.FORBIDDEN + ) + assert exc.status_code == 403 + assert exc.code == AuthErrorCode.FORBIDDEN + assert "Failed" in str(exc) + + +def test_tell_communication_error_str(): + exc = TellCommunicationError("Timeout", endpoint="/state", operation="GET") + assert str(exc) == "Timeout" + assert exc.endpoint == "/state" + assert exc.operation == "GET" + + +def test_transformation_invalid_exception(): + exc = TransformationInvalidException() + assert "Transformation is not implemented" in str(exc) + + +def test_raster_scan_exception(): + exc = RasterScanException("Raster failed") + assert "Raster failed" in str(exc) + + +def test_loop_centering_failed(): + exc = LoopCenteringFailed() + assert "Loop Centering did not detect a sample" in str(exc) + + +def test_unmounting_failed(): + exc = UnmountingFailed() + assert "A sample was not unmounted" in str(exc) + + +def test_mounting_failed(): + exc = MountingFailed() + assert "A sample was not mounted" in str(exc) + + +def test_manual_mount_exception(): + exc = ManualMountException() + assert "Manual mounting failed" in str(exc) + + +def test_smart_magnet_fault_exception(): + exc = SmartMagnetFaultException() + assert "Smart magnet fault" in str(exc) + + +def test_tell_command_while_busy_exception(): + exc = TellCommandWhileBusyException() + assert "Tell is busy" in str(exc) + + +def test_tell_connection_exception(): + exc = TellConnectionException() + assert "Lost connection to Tell" in str(exc) + + +def test_warning_tell_exception(): + exc = WarningTellException() + assert "Warning error in TELL" in str(exc) + + +def test_critical_tell_exception(): + exc = CriticalTellException() + assert "Critical error in TELL" in str(exc) + + +def test_axc_failed(): + exc = AXCFailed() + assert "Auto X-ray centering failed" in str(exc) + + +def test_beamline_busy_exception(): + exc = BeamlineBusyException() + assert "Beamline is in busy state" in str(exc) + + +def test_sample_exception(): + exc = SampleException() + assert "Sample not found" in str(exc) + + +def test_user_rights_exception(): + exc = UserRightsException(code=AuthErrorCode.NOT_STAFF) + assert exc.status_code == 403 + assert exc.code == AuthErrorCode.NOT_STAFF + assert "User does not have rights" in str(exc) + + +def test_smargon_communication_error(): + exc = SmargonCommunicationError("Conn error", endpoint="/move", status_code=500) + assert str(exc) == "Conn error" + assert exc.endpoint == "/move" + assert exc.status_code == 500 + + +def test_jfjoch_communication_error(): + exc = JFJochCommunicationError("JFJoch error", operation="POST") + assert str(exc) == "JFJoch error" + assert exc.operation == "POST" + + +def test_aaredb_communication_error(): + exc = AareDBCommunicationError("DB error") + assert "DB error" in str(exc) + + +def test_aerotech_communication_error(): + exc = AerotechCommunicationError("Aerotech error") + assert "Aerotech error" in str(exc) + + +def test_magnet_position_sensor_error(): + exc = MagnetPositionSensorErorr() + assert "Magnet position sensor error" in str(exc) + + +@pytest.mark.parametrize( + "factory", + [ + lambda: TransformationInvalidException("t"), + lambda: LoopCenteringFailed("lc"), + lambda: UnmountingFailed("un"), + lambda: MountingFailed("mt"), + lambda: ManualMountException("mm"), + lambda: SmartMagnetFaultException("sm"), + lambda: TellCommandWhileBusyException("tb"), + lambda: TellConnectionException("tc"), + lambda: WarningTellException("wt"), + lambda: CriticalTellException("ct"), + lambda: AXCFailed("ax"), + lambda: BeamlineBusyException("bb"), + lambda: SampleException("se"), + lambda: AuthenticationException("ae"), + lambda: TellCommunicationError("tce", endpoint="/x", operation="GET"), + lambda: SmargonCommunicationError("sce", endpoint="/m", status_code=500), + lambda: JFJochCommunicationError("jce", operation="POST"), + lambda: AareDBCommunicationError("dbce"), + lambda: AerotechCommunicationError("ace"), + lambda: MagnetPositionSensorErorr(), + ], +) +def test_exception_construction_emits_no_logs(factory, caplog): + """Constructing an Aare exception must not emit log records. + + Logging is the responsibility of the catch site (server handler, + best-effort wrapper, or callsite that decides to swallow). Emitting + on construction double-logs and pollutes the WARNING-vs-ERROR split + for best-effort flows.""" + caplog.clear() + with caplog.at_level("DEBUG"): + factory() + assert caplog.records == [] diff --git a/tests/test_find_xtal.py b/tests/test_find_xtal.py new file mode 100644 index 0000000..fab48d3 --- /dev/null +++ b/tests/test_find_xtal.py @@ -0,0 +1,178 @@ +from unittest.mock import MagicMock + +import numpy as np +import pytest +from aarecommon.math.find_xtal import ( + compute_crystal_score_array, + create_quality_filtered_array, + get_best_b_factor, + get_best_res, + get_xtal_size, + has_sufficient_low_res_spots, + identify_crystal_raster, + raster_centre_of_mass, + raster_highest_score, + rebuild_array_from_scan_results, +) +from aarecommon.models.models import Coordinate, CrystalSize +from aarecommon.models.raster_grid import CenterOfMassModel, RasterGridRequest + + +@pytest.fixture +def mock_raster_results(): + results = [] + for i in range(9): + res = MagicMock() + res.nx = i % 3 + res.ny = i // 3 + res.number = i + res.spots_low_res = float(i) + res.spots = i + 10 + res.spots_ice = 1 + res.index = False + res.efficiency = 1.0 + res.bkg = 5.0 + res.b = 20.0 + i + res.res = 2.0 - i * 0.1 + results.append(res) + return results + + +@pytest.fixture +def raster_request(): + return RasterGridRequest( + exp_time_s=0.1, + n_x=3, + n_y=3, + grid_size_mm=Coordinate(x=0.02, y=0.02), + smargon_top_left=None, + ) + + +def test_identify_crystal_raster(mock_raster_results, raster_request): + mock_result = MagicMock() + mock_result.images = mock_raster_results + + com = identify_crystal_raster(mock_result, raster_request) + assert isinstance(com, CenterOfMassModel) + # The max spots_low_res is 8.0 at (2, 2) + assert com.n_x == 2 + assert com.n_y == 2 + assert com.max_image == 8 + + +def test_rebuild_array_from_scan_results(mock_raster_results): + arr = rebuild_array_from_scan_results( + mock_raster_results, "spots_low_res", array_shape=(3, 3) + ) + assert arr.shape == (3, 3) + assert arr[0, 0] == 0.0 + assert arr[2, 2] == 8.0 + + +def test_create_quality_filtered_array(mock_raster_results): + # Testing create_quality_filtered_array with a more lenient filter + arr = create_quality_filtered_array( + mock_raster_results, + "spots", + min_low_res_spots=0.0, + min_background=0.0, + array_shape=(3, 3), + ) + # i=8 has nx=2, ny=2 -> row=2, col=2 + assert arr[2, 2] == 18.0 + + +def test_get_xtal_size(raster_request): + # 3x3 array where only center is 1 + arr = np.zeros((3, 3)) + arr[1, 1] = 1 + + size = CrystalSize(x=0, y=0, z=0) + new_size = get_xtal_size(size, arr, raster_request) + # row_max-row_min = 0, so size 0? Wait, it should probably be at least 1 grid unit if present. + # The code does (row_max - row_min) * grid_size_mm.x * 1000 + # If row_min=1, row_max=1, then size is 0. + # This might be a bug in find_xtal.py if it doesn't account for the pixel itself. + assert new_size.x == 0 + + +def test_get_best_b_factor(mock_raster_results): + assert get_best_b_factor(mock_raster_results) == 20.0 + assert get_best_b_factor([]) is None + + +def test_get_best_res(mock_raster_results): + # min res. i=8 -> res = 2.0 - 0.8 = 1.2 + assert pytest.approx(get_best_res(mock_raster_results)) == 1.2 + assert get_best_res([]) is None + + +def test_raster_centre_of_mass(mock_raster_results): + arr = np.zeros((3, 3)) + arr[1, 1] = 10.0 + com = raster_centre_of_mass(arr, mock_raster_results) + assert com.n_x == 1.0 + assert com.n_y == 1.0 + + +def test_has_sufficient_low_res_spots(): + arr = np.array([[1, 2], [3, 4]]) + assert has_sufficient_low_res_spots(arr, 3.0) is True + assert has_sufficient_low_res_spots(arr, 5.0) is False + assert has_sufficient_low_res_spots(None, 1.0) is False + + +@pytest.fixture +def mock_score_results(): + # 2x2 grid; cell (1, 1) is the strongest crystal across bkg, spots_low_res + # and spots_indexed (the three inputs to compute_crystal_score_array). + data = [ + # nx, ny, bkg, spots_low_res, spots_indexed + (0, 0, 1.0, 1.0, 0.0), + (1, 0, 2.0, 2.0, 0.0), + (0, 1, 2.0, 2.0, 0.0), + (1, 1, 10.0, 20.0, 5.0), + ] + results = [] + for i, (nx, ny, bkg, low, idx) in enumerate(data): + res = MagicMock() + res.nx, res.ny, res.number = nx, ny, i + res.bkg, res.spots_low_res, res.spots_indexed = bkg, low, idx + results.append(res) + return results + + +def test_compute_crystal_score_array(mock_score_results): + score = compute_crystal_score_array(mock_score_results) + assert score.shape == (2, 2) + # weights sum to 1.0 and each input is min-max normalised to [0, 100] + assert score.min() >= 0.0 and score.max() <= 100.0 + # (1, 1) is max in all three inputs -> 100; weak corner (0, 0) is min in all -> 0 + assert score[1, 1] == pytest.approx(100.0) + assert score[0, 0] == pytest.approx(0.0) + assert np.unravel_index(np.argmax(score), score.shape) == (1, 1) + + +def test_compute_crystal_score_array_weights(): + # Cell A is the sole max in spots_low_res (weight 0.55); cell B is the sole + # max in spots_indexed (weight 0.20). A must outscore B. + def _cell(nx, n, low, idx): + r = MagicMock() + r.nx, r.ny, r.number = nx, 0, n + r.bkg, r.spots_low_res, r.spots_indexed = 0.0, low, idx + return r + + score = compute_crystal_score_array( + [_cell(0, 0, 10.0, 0.0), _cell(1, 1, 0.0, 10.0)] + ) + assert score[0, 0] > score[1, 0] + assert score[0, 0] == pytest.approx(55.0) + assert score[1, 0] == pytest.approx(20.0) + + +def test_raster_highest_score(mock_score_results): + com = raster_highest_score(mock_score_results) + assert com.n_x == 1.0 + assert com.n_y == 1.0 + assert com.max_image == 3 diff --git a/tests/test_logger_events.py b/tests/test_logger_events.py new file mode 100644 index 0000000..c575b3f --- /dev/null +++ b/tests/test_logger_events.py @@ -0,0 +1,60 @@ +import logging +import time +from unittest.mock import MagicMock + +import pytest +from aarecommon.config.logger_events import log_timing, merge_log_context + + +def test_log_timing_success(): + logger = MagicMock(spec=logging.Logger) + + @log_timing(logger, message_prefix="Test", level=logging.INFO) + def sample_func(x): + time.sleep(0.01) + return x * 2 + + result = sample_func(21) + + assert result == 42 + assert logger.log.call_count == 2 + # First call: Starting + logger.log.assert_any_call(logging.INFO, "Test: Starting sample_func") + # Second call: Finished (check that it contains the message and duration_s in extra) + args, kwargs = logger.log.call_args_list[1] + assert args[0] == logging.INFO + assert "Finished sample_func in" in args[1] + assert "duration_s" in kwargs["extra"] + assert kwargs["extra"]["duration_s"] >= 0.01 + + +def test_log_timing_failure(): + logger = MagicMock(spec=logging.Logger) + + @log_timing(logger, level=logging.ERROR) + def failing_func(): + time.sleep(0.01) + raise ValueError("Something went wrong") + + with pytest.raises(ValueError, match="Something went wrong"): + failing_func() + + assert logger.log.call_count == 2 + # First call: Starting + logger.log.assert_any_call(logging.ERROR, "Starting failing_func") + # Second call: FAILED + args, kwargs = logger.log.call_args_list[1] + assert args[0] == logging.ERROR + assert "failing_func FAILED after" in args[1] + assert "Something went wrong" in args[1] + assert "duration_s" in kwargs["extra"] + + +def test_merge_log_context(): + ctx1 = {"a": 1, "b": 2} + ctx2 = {"b": 3, "c": 4} + merged = merge_log_context(ctx1, ctx2, d=5) + assert merged == {"a": 1, "b": 3, "c": 4, "d": 5} + + merged_none = merge_log_context(ctx1, None) + assert merged_none == {"a": 1, "b": 2} diff --git a/tests/test_mlbox_model.py b/tests/test_mlbox_model.py new file mode 100644 index 0000000..5d72070 --- /dev/null +++ b/tests/test_mlbox_model.py @@ -0,0 +1,26 @@ +from aarecommon.models.models import MLBoxType, MLOutputModel + + +def test_add_box_generates_unique_keys(): + model = MLOutputModel() + key1 = model.add_box(MLBoxType.CRYSTAL, (1, 2, 3, 4), 0.8) + key2 = model.add_box(MLBoxType.CRYSTAL, (5, 6, 7, 8), 0.9) + + assert key1 == "Crystal" + assert key2 == "Crystal_2" + + +def test_get_best_for_class_returns_highest_confidence(): + model = MLOutputModel() + model.add_box(MLBoxType.PIN, (1, 1, 2, 2), 0.3) + model.add_box(MLBoxType.PIN, (3, 3, 4, 4), 0.7) + + best = model.get_best_for_class(MLBoxType.PIN) + + assert best is not None + assert best.conf == 0.7 + + +def test_get_best_for_class_returns_none_when_missing(): + model = MLOutputModel() + assert model.get_best_for_class(MLBoxType.CRYSTAL) is None diff --git a/tests/test_models_extra.py b/tests/test_models_extra.py new file mode 100644 index 0000000..2d683ac --- /dev/null +++ b/tests/test_models_extra.py @@ -0,0 +1,109 @@ +import pytest +from aarecommon.models.models import ( + BeamlineStateEnum, + BeamMarkCoeffModel, + DataCollectionParameters, + DewarAddress, + MLBoxType, + MLOutputModel, + SampleShortInfo, +) + + +def test_sample_short_info_methods(): + info = SampleShortInfo( + db_id=1, + puck_name="puck1", + dewar_name="dew1", + sample_name="sample1", + run_number=1, + aaredb_params=DataCollectionParameters( + totalangle=180, + processingresolution=1.5, + cloud=True, + ), + user="group1", + pin=3, + location=DewarAddress(segment="A", pos=2), + ) + + addr = info.tell_address() + assert addr.puck.segment == "A" + assert addr.puck.pos == 2 + assert addr.pin == 3 + + assert info.loc_str() == "A2-3" + assert info.loc_str_sort() == "A2-03" + assert info.aaredb_params is not None + assert info.aaredb_params.totalangle == 180 + assert info.aaredb_params.processingresolution == 1.5 + + info_no_loc = info.model_copy(update={"location": None}) + assert info_no_loc.loc_str() == "-" + assert info_no_loc.loc_str_sort() == "" + + data = { + "db_id": 2, + "puck_name": "puck2", + "dewar_name": "dew2", + "sample_name": "sample2", + "run_number": 2, + "aaredb_params": { + "totalrange": 120, + "userresolution": 1.8, + "cloud": "", + }, + "user": "group2", + "pin": 4, + "location": {"segment": "B", "pos": 5}, + } + info2 = SampleShortInfo.from_dict(data) + assert info2.db_id == 2 + assert info2.location.segment == "B" + assert info2.aaredb_params is not None + assert info2.aaredb_params.totalangle == 120 + assert info2.aaredb_params.processingresolution == 1.8 + assert info2.aaredb_params.cloud is True + + +def test_beam_mark_coeff_model_apply(): + model = BeamMarkCoeffModel(coeff_x=(1.0, 2.0, 5.0), coeff_y=(3.0, 4.0, 6.0)) + res = model.apply(10.0) + assert res.x == 125.0 + assert res.y == 346.0 + + +def test_ml_output_model_extra_methods(): + model = MLOutputModel() + key = model.add_box(MLBoxType.CRYSTAL, (1, 2, 3, 4), 0.8) + + box_model = model.get_box_model(key) + assert box_model.conf == 0.8 + + assert model.get_box_tuple(key) == (1, 2, 3, 4) + assert model.get_box_tuple("NonExistent") is None + + assert model.get_box_tuple_with_conf(key) == (1, 2, 3, 4, 0.8) + assert model.get_box_tuple_with_conf("NonExistent") is None + + tuples = model.get_tuples_for_class(MLBoxType.CRYSTAL) + assert len(tuples) == 1 + assert tuples[0] == (1, 2, 3, 4) + + tuples_conf = model.get_tuples_with_conf_for_class(MLBoxType.CRYSTAL) + assert len(tuples_conf) == 1 + assert tuples_conf[0] == (1, 2, 3, 4, 0.8) + + assert MLOutputModel.get_class_str(MLBoxType.LOOP_ALL) == "Loop_all" + assert MLOutputModel.get_class_str(MLBoxType.PIN) == "Pin" + assert MLOutputModel.get_class_str(MLBoxType.CRYSTAL) == "Crystal" + assert MLOutputModel.get_class_str(MLBoxType.LOOP_FACE) == "Loop_face" + assert MLOutputModel.get_class_str(MLBoxType.ICE) == "Ice" + assert MLOutputModel.get_class_str(MLBoxType.NEEDLE) == "Needle" + assert MLOutputModel.get_class_str(100) == "Unknown" + + +def test_beamline_state_enum_display_name(): + assert BeamlineStateEnum.SampleExchange.display_name() == "Sample exchange" + assert BeamlineStateEnum.Moving.display_name() == "Moving" + assert BeamlineStateEnum.display_name(None) == "-" diff --git a/tests/test_raster_grid_common.py b/tests/test_raster_grid_common.py new file mode 100644 index 0000000..38d7e5a --- /dev/null +++ b/tests/test_raster_grid_common.py @@ -0,0 +1,99 @@ +import pytest +from aarecommon.math.raster_grid import grid_to_image_id, image_id_to_grid + + +def test_grid_to_image_id_single_cell(): + assert grid_to_image_id(grid_x=1, grid_y=1, number_of_cols=1) == 0 + + +def test_grid_to_image_id_first_row_left_to_right(): + assert grid_to_image_id(grid_x=1, grid_y=1, number_of_cols=5) == 0 + assert grid_to_image_id(grid_x=2, grid_y=1, number_of_cols=5) == 1 + assert grid_to_image_id(grid_x=3, grid_y=1, number_of_cols=5) == 2 + assert grid_to_image_id(grid_x=4, grid_y=1, number_of_cols=5) == 3 + assert grid_to_image_id(grid_x=5, grid_y=1, number_of_cols=5) == 4 + + +def test_grid_to_image_id_second_row_right_to_left(): + assert grid_to_image_id(grid_x=5, grid_y=2, number_of_cols=5) == 5 + assert grid_to_image_id(grid_x=4, grid_y=2, number_of_cols=5) == 6 + assert grid_to_image_id(grid_x=3, grid_y=2, number_of_cols=5) == 7 + assert grid_to_image_id(grid_x=2, grid_y=2, number_of_cols=5) == 8 + assert grid_to_image_id(grid_x=1, grid_y=2, number_of_cols=5) == 9 + + +def test_image_id_to_grid_first_row_left_to_right(): + assert image_id_to_grid(0, 5) == (1, 1) + assert image_id_to_grid(1, 5) == (2, 1) + assert image_id_to_grid(2, 5) == (3, 1) + assert image_id_to_grid(3, 5) == (4, 1) + assert image_id_to_grid(4, 5) == (5, 1) + + +def test_image_id_to_grid_second_row_right_to_left(): + assert image_id_to_grid(5, 5) == (5, 2) + assert image_id_to_grid(6, 5) == (4, 2) + assert image_id_to_grid(7, 5) == (3, 2) + assert image_id_to_grid(8, 5) == (2, 2) + assert image_id_to_grid(9, 5) == (1, 2) + + +def test_round_trip_conversion_for_multiple_grid_shapes(): + for number_of_cols in (1, 2, 3, 5, 8): + for grid_y in range(1, 7): + for grid_x in range(1, number_of_cols + 1): + image_id = grid_to_image_id( + grid_x=grid_x, + grid_y=grid_y, + number_of_cols=number_of_cols, + ) + assert image_id_to_grid(image_id, number_of_cols) == (grid_x, grid_y) + + +def test_known_3x3_serpentine_mapping(): + expected = { + (1, 1): 0, + (2, 1): 1, + (3, 1): 2, + (3, 2): 3, + (2, 2): 4, + (1, 2): 5, + (1, 3): 6, + (2, 3): 7, + (3, 3): 8, + } + for grid_pos, image_id in expected.items(): + assert grid_to_image_id(grid_pos[0], grid_pos[1], 3) == image_id + assert image_id_to_grid(image_id, 3) == grid_pos + + +def test_grid_to_image_id_rejects_zero_columns(): + with pytest.raises(ValueError, match="number_of_cols must be >= 1"): + grid_to_image_id(grid_x=1, grid_y=1, number_of_cols=0) + + +def test_grid_to_image_id_rejects_zero_grid_x(): + with pytest.raises(ValueError, match="grid_x must be >= 1"): + grid_to_image_id(grid_x=0, grid_y=1, number_of_cols=5) + + +def test_grid_to_image_id_rejects_zero_grid_y(): + with pytest.raises(ValueError, match="grid_y must be >= 1"): + grid_to_image_id(grid_x=1, grid_y=0, number_of_cols=5) + + +def test_grid_to_image_id_rejects_grid_x_larger_than_number_of_cols(): + with pytest.raises( + ValueError, match="grid_x cannot be greater than number_of_cols" + ): + grid_to_image_id(grid_x=6, grid_y=1, number_of_cols=5) + + +def test_image_id_to_grid_rejects_zero_columns(): + with pytest.raises(ValueError, match="number_of_cols must be >= 1"): + image_id_to_grid(image_id=0, number_of_cols=0) + + +def test_image_id_to_grid_rejects_negative_image_id(): + with pytest.raises(ValueError, match="image_id must be >= 0"): + image_id_to_grid(image_id=-1, number_of_cols=5) diff --git a/tests/test_recurrence_watcher.py b/tests/test_recurrence_watcher.py new file mode 100644 index 0000000..9aa360c --- /dev/null +++ b/tests/test_recurrence_watcher.py @@ -0,0 +1,84 @@ +from aarecommon.errors.exception_handler import ( + LoopCenteringFailed, + MountingFailed, + SampleException, + TellException, +) +from aarecommon.recurrence_watcher import ( + RecurrenceWatcher, + create_default_watchers, + load_watcher_threshold_overrides, + redis_key_to_env_var, + resolve_exception_class, +) + + +def test_recurrence_watcher_trips_at_threshold_and_resets_on_none(): + watcher = RecurrenceWatcher(name="tell", observes=TellException, threshold=2) + + assert watcher.observe(MountingFailed) is False + assert watcher.observe(MountingFailed) is True + assert watcher.streak == 2 + + assert watcher.observe(None) is False + assert watcher.streak == 0 + + +def test_recurrence_watcher_ignores_non_matching_exception(): + watcher = RecurrenceWatcher(name="tell", observes=TellException, threshold=2) + assert watcher.observe(SampleException) is False + assert watcher.streak == 0 + + +def test_load_watcher_threshold_overrides_reads_valid_values_only(): + values = { + "aare:watchers:mx:tell:threshold": "7", + "aare:watchers:mx:alc:threshold": b"3", + "aare:watchers:mx:smargon:threshold": "not-an-int", + } + + overrides = load_watcher_threshold_overrides( + beamline="mx", + watcher_names=["tell", "alc", "smargon"], + get_value=values.get, + ) + + assert overrides == {"tell": 7, "alc": 3} + + +def test_create_default_watchers_applies_overrides(): + watchers = { + watcher.name: watcher for watcher in create_default_watchers({"alc": 4}) + } + assert watchers["alc"].threshold == 4 + assert watchers["tell"].threshold > 0 + + +def test_resolve_exception_class_maps_known_classes(): + assert resolve_exception_class("MountingFailed") is MountingFailed + assert resolve_exception_class("LoopCenteringFailed") is LoopCenteringFailed + assert resolve_exception_class("UnknownClass") is None + + +def test_redis_key_to_env_var_converts_colons_and_uppercases(): + assert redis_key_to_env_var("aare:watchers:default:tell:threshold") == ( + "AARE_WATCHERS_DEFAULT_TELL_THRESHOLD" + ) + assert redis_key_to_env_var("aare:watchers:mx:alc:threshold") == ( + "AARE_WATCHERS_MX_ALC_THRESHOLD" + ) + + +def test_load_watcher_threshold_overrides_through_env_var_adapter(monkeypatch): + monkeypatch.setenv("AARE_WATCHERS_DEFAULT_TELL_THRESHOLD", "9") + monkeypatch.setenv("AARE_WATCHERS_DEFAULT_ALC_THRESHOLD", "bad") + + import os + + overrides = load_watcher_threshold_overrides( + beamline="default", + watcher_names=["tell", "alc", "smargon"], + get_value=lambda k: os.getenv(redis_key_to_env_var(k)), + ) + + assert overrides == {"tell": 9} diff --git a/tests/test_zoom_model_camera_settings.py b/tests/test_zoom_model_camera_settings.py new file mode 100644 index 0000000..f5c794a --- /dev/null +++ b/tests/test_zoom_model_camera_settings.py @@ -0,0 +1,47 @@ +import pytest +from aarecommon.models.models import SampleCameraSettings, ZoomModel + + +def test_get_camera_settings_exact_match(): + model = ZoomModel( + z={ + 100: SampleCameraSettings(gain=1.0, exposure=0.1), + 200: SampleCameraSettings(gain=2.0, exposure=0.2), + } + ) + + settings = model.get_camera_settings(100) + assert settings.gain == 1.0 + assert settings.exposure == 0.1 + + +def test_get_camera_settings_interpolates(): + model = ZoomModel( + z={ + 100: SampleCameraSettings(gain=0.0, exposure=0.1), + 200: SampleCameraSettings(gain=2.0, exposure=0.3), + } + ) + + settings = model.get_camera_settings(150) + assert settings.gain == 1.0 + assert settings.exposure == 0.2 + + +def test_get_camera_settings_raises_when_empty(): + model = ZoomModel(z={}) + with pytest.raises(ValueError): + model.get_camera_settings(100) + + +def test_get_camera_settings_below_minimum_returns_first(): + model = ZoomModel( + z={ + 100: SampleCameraSettings(gain=1.0, exposure=0.1), + 200: SampleCameraSettings(gain=2.0, exposure=0.2), + } + ) + + settings = model.get_camera_settings(50) + assert settings.gain == 1.0 + assert settings.exposure == 0.1