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):

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
+  ]
+}

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

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"
 )