diff --git a/src/aare/beamline_dispatch/default/beamline_dispatch.py b/src/aare/beamline_dispatch/default/beamline_dispatch.py index 03175a62..7444b34a 100644 --- a/src/aare/beamline_dispatch/default/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/default/beamline_dispatch.py @@ -2,10 +2,12 @@ import json import os from abc import ABC from importlib.resources import files +from typing import Protocol +import numpy as np from aarecommon.config.beamline import MXBeamline -from aarecommon.math.beam_center import BeamCenterFromDetectorStage from aarecommon.models.beam_centre import BeamCentre +from numpy.typing import ArrayLike, NDArray from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch, Geometry @@ -19,6 +21,12 @@ class DefaultAuthDispatch(AuthDispatch): return key +class BeamCenterModel(Protocol): + def predict( + self, det_z_mm: ArrayLike, det_y_mm: ArrayLike + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ... + + class DefaultGeometry(Geometry): def __init__(self, beamline: MXBeamline) -> None: super().__init__() @@ -28,7 +36,7 @@ class DefaultGeometry(Geometry): self._beamline = beamline @property - def beam_centre_model(self) -> BeamCenterFromDetectorStage: + def beam_centre_model(self) -> BeamCenterModel: return self._model.model @@ -36,9 +44,11 @@ class DefaultDispatch(BeamlineDispatch, ABC): """Default implementation for anything which can vary between beamlines and/or simulation. Should be safe and fail rather than assuming anything.""" - def __init__(self, beamline: MXBeamline) -> None: + def __init__( + self, beamline: MXBeamline, geo_class: type[DefaultGeometry] = DefaultGeometry + ) -> None: self._auth = DefaultAuthDispatch() - self._geo = DefaultGeometry(beamline=beamline) + self._geo = geo_class(beamline=beamline) @property def auth(self): diff --git a/src/aare/beamline_dispatch/x10sa/beam_centers.json b/src/aare/beamline_dispatch/x10sa/beam_centers.json new file mode 100644 index 00000000..a67755c1 --- /dev/null +++ b/src/aare/beamline_dispatch/x10sa/beam_centers.json @@ -0,0 +1,10 @@ +{ + "det_z": [190, 200, 250, 300, 350, 400, 450, 500, 550], + "center_x": [ + 2081.3, 2081.37, 2081.29, 2081.81, 2081.51, 2080.22, 2079.81, 2080.24, + 2078.89 + ], + "center_y": [ + 2242.0, 2242.99, 2242.72, 2242.05, 2241.46, 2240.96, 2239.9, 2239.8, 2238.23 + ] +} diff --git a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py index 60e40813..a5121512 100644 --- a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py @@ -1,8 +1,16 @@ +import json +from importlib.resources import files from typing import Any +import numpy as np from aarecommon.config.beamline import MXBeamline +from numpy.typing import ArrayLike, NDArray -from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch +from aare.beamline_dispatch.default.beamline_dispatch import ( + BeamCenterModel, + DefaultDispatch, + DefaultGeometry, +) from aare.beamline_dispatch.protocols import BecMacros @@ -61,9 +69,38 @@ class X10SaBecMacros(BecMacros): def auto_exposure(): ... +class _1dModel(BeamCenterModel): + def __init__(self, x_fit: NDArray[np.float64], y_fit: NDArray[np.float64]) -> None: + self._x_fit = x_fit + self._y_fit = y_fit + + def predict( + self, det_z_mm: ArrayLike, det_y_mm: ArrayLike + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + return ( + self._x_fit[1] + self._x_fit[0] * det_z_mm, + self._y_fit[1] + self._y_fit[0] * det_z_mm, + ) + + +class X10saGeometry(DefaultGeometry): + def __init__(self, beamline: MXBeamline) -> None: + super().__init__(beamline) + with open(str(files("aare.beamline_dispatch.x10sa") / "beam_centers.json")) as f: + data = json.loads(f.read()) + x_fit = np.polyfit(data["det_z"], data["center_x"], 1) + y_fit = np.polyfit(data["det_z"], data["center_y"], 1) + self._model = _1dModel(x_fit, y_fit) + self._beamline = beamline + + @property + def beam_centre_model(self) -> BeamCenterModel: + return self._model + + class X10saDispatch(DefaultDispatch): def __init__(self, beamline: MXBeamline) -> None: - super().__init__(beamline=beamline) + super().__init__(beamline=beamline, geo_class=X10saGeometry) self._bec_macros = X10SaBecMacros() @property diff --git a/tests/unit/utils/test_beamline_dispatch.py b/tests/unit/utils/test_beamline_dispatch.py index 0aa48866..5da1bbd5 100644 --- a/tests/unit/utils/test_beamline_dispatch.py +++ b/tests/unit/utils/test_beamline_dispatch.py @@ -1,6 +1,7 @@ import importlib.util from unittest.mock import patch +import numpy as np import pytest from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch @@ -37,6 +38,17 @@ def test_pxii_dispatch_can_be_instantiated(): assert isinstance(dispatch, X10saDispatch) +@pytest.mark.skipif( + importlib.util.find_spec("pxii_bec") is None, reason="run only for pxii flavour" +) +def test_pxii_dispatch_beam_center(): + with patch.dict("os.environ", {"BEAMLINE": "X10SA"}): + dispatch = get_beamline_dispatch() + beam_center = dispatch.geo.beam_centre_model.predict(300, 62) + assert np.isclose(beam_center[0], 2081.8, atol=1) + assert np.isclose(beam_center[1], 2242.05, atol=1) + + @pytest.mark.skipif( importlib.util.find_spec("pxiii_bec") is None, reason="run only for pxiii flavour" )