Initial commit

src/.gitignore

@@ -0,0 +1 @@
+.idea/

src/aaredaq/aaredb.py

@@ -0,0 +1,132 @@
+import io
+from typing import List, Optional
+
+import numpy as np
+import aareDBclient
+import cv2
+import requests
+from aareDBclient import SetTellPosition, SampleEventCreate, SetTellPositionRequest
+
+from aaredaq.beamline import MXBeamline
+from aaredaq.models import (
+    SampleShortInfo,
+    PuckLoadedInfo,
+    PuckInfo,
+    DewarAddress,
+    SampleShortInfoList,
+)
+
+
+class AareWrapper:
+    def __init__(
+        self,
+        bl: MXBeamline,
+        host: str = "https://mx-aare-test.psi.ch:1492",
+    ):
+        configuration = aareDBclient.Configuration(host=host)
+        configuration.verify_ssl = False  # Disable SSL verification
+
+        self.client = aareDBclient.ApiClient(configuration)
+        self.__host = host
+        self.__puck_api = aareDBclient.PucksApi(self.client)
+        self.__sample_api = aareDBclient.SamplesApi(self.client)
+        self.__bl = bl
+
+    def get_pucks_beamline(self) -> List[PuckInfo]:
+        puck_list = self.__puck_api.get_pucks_by_slot_pucks_slot_slot_identifier_get(
+            slot_identifier=self.__bl.value.upper()
+        )
+        ret = []
+        for i in puck_list:
+            if i.tell_position is not None and len(i.tell_position) == 2:
+                d = DewarAddress(segment=i.tell_position[0], pos=i.tell_position[1])
+            else:
+                d = None
+            ret.append(
+                PuckInfo(
+                    db_id=i.id,
+                    puck_name=i.puck_name,
+                    dewar_name=i.dewar_name,
+                    user=i.pgroup,
+                    location=d,
+                )
+            )
+        return ret
+
+    def set_pucks_beamline(self, input_list: List[PuckLoadedInfo]):
+        o = []
+
+        for i in input_list:
+            t = SetTellPosition(
+                puck_name=i.puck_name,
+                segment=i.location.segment,
+                puck_in_segment=i.location.pos,
+            )
+            o.append(t)
+        payload = SetTellPositionRequest(pucks = o, tell=self.__bl.value.upper())
+        ret = self.__puck_api.set_tell_positions_pucks_set_tell_positions_put(
+            set_tell_position_request=payload,
+        )
+        print(ret)
+
+    def get_sample_info(self) -> SampleShortInfoList:
+        sample_list = (
+            self.__puck_api.get_pucks_with_tell_position_pucks_with_tell_position_get(tell=self.__bl.value.upper())
+        )
+
+        ret = []
+        for p in sample_list:
+            for s in p.samples:
+                dewar_address = None
+                if p.tell_position is not None and len(p.tell_position) == 2:
+                    dewar_address = DewarAddress(
+                        segment=p.tell_position[0], pos=p.tell_position[1]
+                    )
+
+                ret.append(
+                    SampleShortInfo(
+                        db_id=s.id,
+                        puck_name=p.puck_name,
+                        dewar_name="Unknown" if p.dewar_name is None else p.dewar_name,
+                        sample_name=s.sample_name,
+                        user=p.pgroup,
+                        pin=s.position,
+                        location=dewar_address,
+                        priority=1.0 if s.priority is None else s.priority,
+                        comment=s.comments,
+                        mount_count=s.mount_count or 0
+                    )
+                )
+        return SampleShortInfoList(s=ret)
+
+    def sample_mounted(self, s: Optional[SampleShortInfo]):
+        if s is not None:
+            try:
+                self.__sample_api.create_sample_event_samples_samples_sample_id_events_post(
+                    sample_id=s.db_id,
+                    sample_event_create=SampleEventCreate(event_type="Mounted"),
+                )
+            except Exception as e:
+                print(e)
+
+    def sample_unmounted(self, s: Optional[SampleShortInfo]):
+        if s is not None:
+            try:
+                self.__sample_api.create_sample_event_samples_samples_sample_id_events_post(
+                    sample_id=s.db_id,
+                    sample_event_create=SampleEventCreate(event_type="Unmounted"),
+                )
+            except Exception as e:
+                print(e)
+
+    def upload_image(self, sample_id: int, filename: str, bgr_image: np.ndarray):
+        _, buffer = cv2.imencode('.jpg', bgr_image)
+        jpeg_bytes = io.BytesIO(buffer)
+        url = f"{self.__host}/samples/{sample_id}/upload-images"
+        headers = {
+            "accept": "application/json"
+        }
+        response = requests.post(url,
+                                 files={'uploaded_file': (filename + ".jpg", jpeg_bytes, "image/jpeg")},
+                                 verify=False, headers=headers)
+        print(response.status_code)

src/aaredaq/auth.py

@@ -0,0 +1,99 @@
+import grp
+import os
+import pwd
+from datetime import datetime, timedelta, UTC
+from typing import List
+
+import jwt
+from fastapi import HTTPException, status
+from fastapi.security import OAuth2PasswordRequestForm
+from pydantic import BaseModel
+
+from aaredaq.config import BeamlineConfig
+
+SECRET_KEY = "secret_key_not_for_production"
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 24 * 60
+SESSION_EXPIRE_SECONDS = 60 * 10
+
+STAFF_GROUP = "unx-MXgroup"
+
+
+class TokenData(BaseModel):
+    sub: str # Username
+    pgroups: List[str]
+    session: int
+    staff: bool = False
+
+
+def create_access_token(token: TokenData):
+    to_encode = token.model_dump()
+    to_encode.update({"exp": datetime.now(UTC) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)})
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+
+def authenticate_user(cfg: BeamlineConfig, form_data: OAuth2PasswordRequestForm) -> str:
+    user_info = pwd.getpwnam(form_data.username)
+
+    groups = os.getgrouplist(user_info.pw_name, user_info.pw_gid)
+    supplementary_groups = [grp.getgrgid(g).gr_name for g in groups]
+
+    pgroups = [group for group in supplementary_groups if group.startswith('p')]
+    staff = "unx-MXgroup" in supplementary_groups
+    pgroups.append("p20000")
+    token = TokenData(sub=form_data.username,
+                      pgroups=pgroups,
+                      staff=staff,
+                      session=cfg.generate_session())
+
+    return create_access_token(token)
+
+
+def parse_token(token: str) -> TokenData:
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        token = TokenData(**payload)
+        return token
+    except jwt.PyJWTError as e:
+        print(f"JWT error {e}")
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid token",
+        )
+
+
+def check_jwt_ro(cfg: BeamlineConfig, data: TokenData) -> None:
+    active_pgroup = cfg.pgroup
+    if not data.staff and (active_pgroup is None or active_pgroup not in data.pgroups):
+        print("Not member of a currently active p-group.")
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Not member of a currently active p-group.",
+        )
+
+
+def check_jwt_rw(cfg: BeamlineConfig, data: TokenData) -> None:
+    check_jwt_ro(cfg, data)
+
+    try:
+        cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS)
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Another session is active.",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+
+def check_jwt_staff(data: TokenData) -> None:
+    if not data.staff:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Not member of the MX staff.",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+
+def force_current_sesion(cfg: BeamlineConfig, data: TokenData) -> None:
+    cfg.force_set_active_session(data.session, SESSION_EXPIRE_SECONDS)

src/aaredaq/beamline.py

@@ -0,0 +1,14 @@
+import os
+from enum import Enum
+
+
+class MXBeamline(Enum):
+    X06SA = "X06SA"
+    X10SA = "X10SA"
+    X06DA = "X06DA"
+    SIMULATED = "SIMULATED"
+
+
+def mx_beamline() -> MXBeamline:
+    name = os.getenv("BEAMLINE_XNAME", "SIMULATED").strip().upper()
+    return MXBeamline[name] if name in MXBeamline.__members__ else MXBeamline.SIMULATED

src/aaredaq/config.py

@@ -0,0 +1,323 @@
+import base64
+import io
+import json
+
+import numpy as np
+import redis
+import redis_lock
+from aaredaqlib.coordinate import Coordinate
+
+from aaredaq.beamline import MXBeamline
+from aaredaq.models import (
+    BeamlineSettingsModel,
+    BeamCenterCoeffModel,
+    LoopCenteringZoomModel,
+    def_loop_centering_zoom,
+    SampleShortInfo,
+    BeamlineStateEnum, SessionsStateEnum,
+)
+
+ABR_POS_ALIGN_DEF = Coordinate(x=-18, y=-0.266, z=0)
+ABR_POS_MOUNT = Coordinate(x=-18, y=0, z=0)
+ABR_OMEGA_MOUNT = 0.0
+
+# Serialize the NumPy array to a Base64 string
+def numpy_to_base64(array: np.ndarray) -> str:
+    buffer = io.BytesIO()  # Create an in-memory buffer
+    # Save the array as a binary file to the buffer
+    np.save(buffer, array, allow_pickle=False)
+    # Reset buffer's position to the beginning
+    buffer.seek(0)
+    # Base64 encode and convert bytes to string
+    return base64.b64encode(buffer.read()).decode()
+
+
+# Retrieve and decode the array
+def base64_to_numpy(encoded_str: str | None) -> np.ndarray | None:
+    if encoded_str is None:
+        return None
+    decoded = base64.b64decode(encoded_str)  # Base64 decode
+    buffer = io.BytesIO(decoded)  # Convert bytes to a buffer
+    return np.load(buffer)  # Load buffer as a NumPy array
+
+
+class BeamlineBusyException(Exception):
+    pass
+
+
+class BeamlineConfig:
+    """
+    Manages the configuration and state of a beamline system by interacting with a Redis
+    backend. Provides methods for updating status, settings, and related state details.
+
+    This class encapsulates the logic needed to manage beamline states and corresponding
+    operations such as setting busy flags, updating or retrieving settings, and handling
+    different beamline-related data such as background images. By connecting to a Redis
+    database, it ensures that the beamline's state and configurations are stored and
+    retrieved efficiently. Each beamline instance is identified by a unique name, and all
+    operations are performed with atomic safety using Redis locking mechanisms.
+
+    Attributes:
+        __bl (str): The beamline's unique identifier or name.
+        __client (redis.Redis): Redis client instance used for interacting with the datastore.
+    """
+
+    def __init__(self, bl: MXBeamline):
+        self.__bl = bl.value.lower()
+        if bl is MXBeamline.SIMULATED:
+            host = "localhost"
+        else:
+            host = f"mx-aare-test.psi.ch"
+        self.__client = redis.Redis(host=host, port=6379, db=0, decode_responses=True)
+
+    # Session and authentication management
+
+    def generate_session(self) -> int:
+        return int(self.__client.incr(f"{self.__bl}:session"))
+
+    @property
+    def active_session(self) -> int | None:
+        tmp = self.__client.get(f"{self.__bl}:active_session")
+        if tmp is None:
+            return None
+        return int(tmp)
+
+    def session_state(self, session: int) -> SessionsStateEnum:
+        curr_session = self.active_session
+        if curr_session is None:
+            return SessionsStateEnum.Vacant
+        elif curr_session == session:
+            return SessionsStateEnum.OwnedByYou
+        else:
+            return SessionsStateEnum.OwnedByElse
+
+    def try_set_active_session(self, session: int, expiry_sec: int) -> None:
+        with redis_lock.Lock(
+            self.__client, f"{self.__bl}:active_session_lock", expire=10
+        ):
+            tmp = self.__client.get(f"{self.__bl}:active_session")
+            if tmp is None:
+                self.__client.set(f"{self.__bl}:active_session", session)
+            elif int(tmp) != session:
+                raise Exception(
+                    "There is already active session with different id. Try again later."
+                )
+            self.__client.expire(f"{self.__bl}:active_session", expiry_sec)
+
+    def end_active_session(self, session: int) -> None:
+        with redis_lock.Lock(
+            self.__client, f"{self.__bl}:active_session_lock", expire=10
+        ):
+            if int(self.__client.get(f"{self.__bl}:active_session")) == session:
+                self.__client.delete(f"{self.__bl}:active_session")
+
+    def force_set_active_session(self, session: int, expiry_sec: int) -> None:
+        # Ensure that there is no active try-set for active session
+        with redis_lock.Lock(
+            self.__client, f"{self.__bl}:active_session_lock", expire=10
+        ):
+            self.__client.set(f"{self.__bl}:active_session", session)
+            self.__client.expire(f"{self.__bl}:active_session", expiry_sec)
+
+    @property
+    def tell_mount_count(self) -> int:
+        return int(self.__client.incr(f"{self.__bl}:tell_mount_count"))
+
+    @property
+    def pgroup(self) -> str | None:
+        tmp = self.__client.get(f"{self.__bl}:pgroup")
+        return tmp
+
+    @pgroup.setter
+    def pgroup(self, pgroup: str | None) -> None:
+        if pgroup is None:
+            self.__client.delete(f"{self.__bl}:pgroup")
+        else:
+            self.__client.set(f"{self.__bl}:pgroup", pgroup)
+
+    def pixel_to_mm(self, zoom: float) -> float:
+        cfg = self.settings
+        return 1.0 / (
+            cfg.camera_translation_factor_b
+            * np.exp(cfg.camera_translation_factor_a * zoom)
+        )
+
+    @property
+    def settings(self) -> BeamlineSettingsModel:
+        tmp = self.__client.get(f"{self.__bl}:settings")
+        if tmp is None:
+            raise Exception("Beamline key not found")
+
+        data_dict = json.loads(tmp)
+        model = BeamlineSettingsModel(**data_dict)
+        return model
+
+    @settings.setter
+    def settings(self, data: BeamlineSettingsModel):
+        self.__client.set(f"{self.__bl}:settings", data.model_dump_json())
+
+    def get_alc_bkg(self, zoom: float) -> np.ndarray | None:
+        return base64_to_numpy(self.__client.get(f"{self.__bl}:bkg{zoom:.1f}"))
+
+    def put_alc_bkg(self, zoom: float, data: np.ndarray):
+        self.__client.set(f"{self.__bl}:bkg{zoom:.1f}", numpy_to_base64(data))
+
+    @property
+    def state(self) -> BeamlineStateEnum:
+        raw_value = self.__client.get(f"{self.__bl}:state")
+
+        if raw_value is None:
+            return BeamlineStateEnum.Maintenance
+        try:
+            int_value = int(raw_value)  # Ensure it's an integer
+            return BeamlineStateEnum(int_value)  # Convert to BeamlineStateEnum
+        except (ValueError, KeyError):
+            raise ValueError(
+                f"Invalid 'mx_state' value: {raw_value}. Expected integer corresponding to a BeamlineStateEnum."
+            )
+
+    @state.setter
+    def state(self, state: BeamlineStateEnum):
+        self.__client.set(f"{self.__bl}:state", state.value)
+
+    @property
+    def state_busy(self) -> bool:
+        return self.__client.get(f"{self.__bl}:busy") is not None
+
+    @state_busy.setter
+    def state_busy(self, i: bool):
+        if i:
+            self.__client.set(f"{self.__bl}:busy", "1")
+        else:
+            self.__client.delete(f"{self.__bl}:busy")
+
+    @property
+    def current_sample(self) -> SampleShortInfo | None:
+        tmp = self.__client.get(f"{self.__bl}:current_sample")
+        if tmp is None:
+            return None
+
+        data_dict = json.loads(tmp)
+        return SampleShortInfo(**data_dict)
+
+    @current_sample.setter
+    def current_sample(self, sample: SampleShortInfo | None):
+        if sample is None:
+            self.__client.delete(f"{self.__bl}:current_sample")
+        else:
+            self.__client.set(f"{self.__bl}:current_sample", sample.model_dump_json())
+
+    @property
+    def beam_center_coeff(self) -> BeamCenterCoeffModel:
+        tmp = self.__client.get(f"{self.__bl}:beam_center_camera")
+        if tmp is None:
+            return BeamCenterCoeffModel()
+
+        data_dict = json.loads(tmp)
+        return BeamCenterCoeffModel(**data_dict)
+
+    @beam_center_coeff.setter
+    def beam_center_coeff(self, data: BeamCenterCoeffModel):
+        self.__client.set(f"{self.__bl}:beam_center_camera", data.model_dump_json())
+
+    @property
+    def alc_zoom_settings(self) -> LoopCenteringZoomModel:
+        tmp = self.__client.get(f"{self.__bl}:alc_zoom_settings")
+        if tmp is None:
+            return def_loop_centering_zoom()
+
+        data_dict = json.loads(tmp)
+        return LoopCenteringZoomModel(**data_dict)
+
+    @alc_zoom_settings.setter
+    def alc_zoom_settings(self, data: LoopCenteringZoomModel):
+        self.__client.set(f"{self.__bl}:alc_zoom_settings", data.model_dump_json())
+
+    @property
+    def abr_meas_pos(self) -> Coordinate:
+        tmp = self.__client.get(f"{self.__bl}:abr_meas_pos")
+        if tmp is None:
+            return ABR_POS_ALIGN_DEF
+
+        data_dict = json.loads(tmp)
+        return Coordinate(**data_dict)
+
+    @abr_meas_pos.setter
+    def abr_meas_pos(self, data: Coordinate):
+        self.__client.set(f"{self.__bl}:abr_meas_pos", data.model_dump_json())
+
+
+    @property
+    def dtz(self) -> float | None:
+        tmp = self.__client.get(f"{self.__bl}:dtz")
+        if tmp is None:
+            return None
+        return float(tmp)
+
+    @dtz.setter
+    def dtz(self, dtz: float):
+        self.__client.set(f"{self.__bl}:dtz", dtz)
+
+    def listen_changes(self) -> redis.client.PubSub:
+        self.__client.config_set("notify-keyspace-events", "KEA")
+        pubsub = self.__client.pubsub()
+        pubsub.psubscribe(f"__keyspace@0__:{self.__bl}:*")
+        return pubsub
+
+    # Atomic check if beamline is busy and if not set state to busy
+    def try_set_busy(self, timeout: int | None = None):
+        with redis_lock.Lock(self.__client, f"{self.__bl}:move_state_lock", expire=10):
+            if self.state_busy:
+                raise BeamlineBusyException("Beamline is busy")
+            self.state_busy = True
+            if timeout is not None:
+                self.__client.expire(f"{self.__bl}:busy", timeout)
+
+    def set_busy(self, target: BeamlineStateEnum, timeout: int | None = None):
+        self.try_set_busy(timeout=timeout)
+        curr_state = self.state
+        if curr_state != target:
+            self.state_busy = False
+            raise Exception("Beamline is not in a proper state")
+
+    def start_moving(
+        self, target: BeamlineStateEnum, timeout: int | None = None
+    ) -> BeamlineStateEnum:
+        self.try_set_busy(timeout=timeout)
+        curr_state = self.state
+        if target == curr_state:
+            self.state_busy = False
+        else:
+            self.state = BeamlineStateEnum.Moving
+        return curr_state
+
+    def clear_mark_beam(self):
+        self.__client.delete(f"{self.__bl}:beam_mark")
+
+    def mark_beam(self, x_pxl: float, y_pxl: float, zoom: float):
+        self.__client.hset(
+            f"{self.__bl}:beam_mark",
+            mapping={f"{zoom}": json.dumps({"x": x_pxl, "y": y_pxl})},
+        )
+        vals = self.__client.hgetall(f"{self.__bl}:beam_mark")
+
+        if len(vals) >= 3:
+            zooms = []
+            x_pxls = []
+            y_pxls = []
+            for k, v in vals.items():
+                zooms.append(float(k))
+                x_pxls.append(float(json.loads(v)["x"]))
+                y_pxls.append(float(json.loads(v)["y"]))
+
+            model = BeamCenterCoeffModel(
+                coeff_x=tuple(np.polyfit(zooms, x_pxls, deg=2)),
+                coeff_y=tuple(np.polyfit(zooms, y_pxls, deg=2)),
+            )
+            self.beam_center_coeff = model
+
+    def get_beam_center(self, zoom: float) -> tuple[float, float]:
+        model = self.beam_center_coeff
+        x_pxl = np.polyval(model.coeff_x, zoom)
+        y_pxl = np.polyval(model.coeff_y, zoom)
+        return float(x_pxl), float(y_pxl)

src/aaredaq/daq.py

@@ -0,0 +1,720 @@
+import time
+from math import ceil
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+import redis
+from aaredaqlib.coordinate import Coordinate, SmargonCoordinate
+from aaredaqlib.helical_geometry import HelicalScanRequest
+from aaredaqlib.raster_grid import RasterGridRequest
+from aaredaqlib.rotation_scan import RotationScanRequest
+from aaredaqlib.sample_geometry import SampleGeometryModel
+from aaredaqlib.screening_scan import ScreeningScanRequest
+
+from aaredaq import workflows
+from aaredaq.beamline import MXBeamline
+from aaredaq.config import BeamlineConfig, ABR_POS_MOUNT
+from aaredaq.config import BeamlineStateEnum
+from aaredaq.devices import BeamlineDevices
+from aaredaq.mlbox import MlBox
+from aaredaq.models import (
+    SampleShortInfo,
+    PuckLoadedInfo,
+    SampleShortInfoList,
+    DAQStatusModel, )
+
+
+class TransformationInvalidException(Exception):
+    def __init__(self, message="Transformation is not implemented"):
+        super().__init__(message)
+        self.message = message
+
+    def __str__(self):
+        return self.message
+
+
+class LoopCenteringFailed(Exception):
+    pass
+
+
+class AareDAQ:
+    def __init__(self, cfg: BeamlineConfig, bl: MXBeamline):
+        self.last_time = 0.0
+        self.__cfg = cfg
+        self.__devs = BeamlineDevices(bl)
+        self.__mlbox = MlBox()
+
+    @property
+    def state(self) -> BeamlineStateEnum:
+        return self.__cfg.state
+
+    @property
+    def busy(self) -> bool:
+        return self.__cfg.state_busy
+
+    @state.setter
+    def state(self, target: BeamlineStateEnum):
+        if target == BeamlineStateEnum.Moving:
+            raise Exception("Cannot explicitly move to busy state")
+
+        start = time.perf_counter()
+        self.__cfg.try_set_busy(timeout=300)
+        self.__set_state(target)
+        self.__cfg.state_busy = False
+
+        end = time.perf_counter()
+
+        self.last_time = end - start
+
+    @property
+    def omega(self) -> float:
+        return self.__devs.aerotech.omega
+
+    @omega.setter
+    def omega(self, val: float):
+        self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
+        if -720 < val < 720:
+            self.__devs.aerotech.move(val, wait=True, speed=180.0, direct=True)
+            self.__cfg.state_busy = False
+        else:
+            self.__cfg.state_busy = False
+            raise ValueError("Omega has to be between -720 and 720 degrees (for now)")
+
+    @property
+    def zoom(self) -> float:
+        return self.__devs.zoom
+
+    @zoom.setter
+    def zoom(self, val: float):
+        self.__devs.zoom = val
+
+    @property
+    def light(self) -> float:
+        val = self.__devs.lamp_light.value
+        if val <= 1.0:
+            return 0
+        elif val >= 2.5:
+            return 100.0
+        else:
+            return (val - 1.0) / 1.5 * 100.0
+
+    @light.setter
+    def light(self, f: float):
+        conv = (f / 100.0 * 1.5) + 1.0
+        print(f"Light {f} -> {conv}")
+        self.__devs.lamp_light.move(conv, wait=True)
+
+    @property
+    def sample(self) -> SampleShortInfo | None:
+        if self.__devs.tell.get_mounted_sample() is None:
+            return None
+        return self.__cfg.current_sample
+
+    def tweak_abr_meas_pos(self, c: Coordinate):
+        self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
+        try:
+            new_meas_pos = self.__cfg.abr_meas_pos + c
+            self.__cfg.abr_meas_pos = new_meas_pos
+            self.__devs.abr_pos = new_meas_pos
+            self.__cfg.state_busy = False
+        except:
+            self.__cfg.state_busy = False
+            raise
+
+    def save_abr_meas_pos(self):
+        self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
+        try:
+            self.__cfg.abr_meas_pos = self.__devs.abr_pos
+            print(self.__cfg.abr_meas_pos.y)
+            self.__cfg.state_busy = False
+        except:
+            self.__cfg.state_busy = False
+            raise
+
+    def __mount(self, target: SampleShortInfo | None):
+        self.__set_state(BeamlineStateEnum.RobotSampleExchange)
+        self.__devs.smargon.move_home(wait=True)
+        self.__devs.abr_pos = ABR_POS_MOUNT
+        time.sleep(0.5)
+        print(f"Smargon x={self.__devs.smargon.readback.sh_mm.x}, y={self.__devs.smargon.readback.sh_mm.y}, z={self.__devs.smargon.readback.sh_mm.z}")
+        print(f"ABR x={self.__devs.abr_pos.x}, y={self.__devs.abr_pos.y}, z={self.__devs.abr_pos.z}  omega={self.__devs.aerotech.omega}")
+        if self.__devs.bsz.value < 24.0:
+            raise Exception("Beamstop Z below 24.0 mm - potentially unsafe with mounting")
+        if self.__devs.magnet_position_sensor.value != 0:
+            raise Exception("Magnet position sensor is not in position")
+        self.__devs.tell.check_enable_motion()
+        self.__devs.tell.set_in_mount_position(True)
+        curr_sample = self.__cfg.current_sample
+
+        self.__devs.tell.wait_ready()
+
+        if target is None:
+            self.__devs.tell.unmount(wait=True, timeout=360)
+
+            if curr_sample is not None:
+                self.__devs.aare.sample_unmounted(curr_sample)
+            self.__cfg.current_sample = None
+        else:
+            if target.location is None:
+                raise Exception("Sample not loaded into dewar")
+            if curr_sample is not None:
+                self.__devs.aare.sample_unmounted(curr_sample)
+            self.__devs.tell.mount(
+                address=target.tell_address(),
+                force=True,
+                auto_unmount=True,
+                read_dm=False,
+                wait=True,
+                timeout=360,
+            )
+
+            self.__devs.aare.sample_mounted(target)
+            self.__cfg.current_sample = target
+
+    @sample.setter
+    def sample(self, target: SampleShortInfo | None):
+        # This will set internally state to SampleExchange, but will return to SampleAlignment
+        # before exiting
+        start = time.perf_counter()
+
+        self.__cfg.try_set_busy(timeout=360)
+
+        try:
+            self.__mount(target)
+            self.__cfg.state_busy = False
+        except Exception as e:
+            self.__cfg.state_busy = False
+            # If unmount succeeded, but mount failed
+            if self.__devs.tell.get_mounted_sample() is None:
+                self.__cfg.current_sample = None
+            raise e
+
+        self.__set_state(BeamlineStateEnum.SampleAlignment)
+
+        end = time.perf_counter()
+        self.last_time = end - start
+
+    @property
+    def camera_image(self) -> np.ndarray | None:
+        image = self.__devs.sample_cam.get_image(gray=False)
+        return image
+
+    @property
+    def camera_image_gray(self) -> np.ndarray | None:
+        image = self.__devs.sample_cam.get_image(gray=True)
+        return image
+
+    def list_loaded_pucks(self) -> List[PuckLoadedInfo]:
+        return self.__devs.tell.get_detected_pucks()
+
+    def __raster(self, r: RasterGridRequest):
+        max_time = r.exp_time_s * r.n_y * r.n_x + 60
+        self.__set_state(BeamlineStateEnum.DataCollection)
+        self.__devs.aerotech.move(r.omega_deg, wait=True, speed=360.0)
+        save_smargon_position = self.__devs.smargon.readback
+        self.__devs.smargon.target = r.smargon
+        self.__devs.aerotech.measure_raster_simple(
+            r.exp_time_s, r.grid_size_mm.x, r.grid_size_mm.y, r.n_x, r.n_y
+        )
+        self.__devs.aerotech.wait_scan_done(max_time)
+        self.__devs.aerotech.reset()
+        self.__devs.smargon.target = save_smargon_position
+
+    def measure_raster(self, r: RasterGridRequest):
+        max_time = r.exp_time_s * r.n_y * r.n_x + 60
+        self.__cfg.try_set_busy(timeout=ceil(max_time + 360))
+        try:
+            self.__raster(r)
+            self.__set_state(BeamlineStateEnum.SampleAlignment)
+            self.__cfg.state_busy = False
+        except Exception as e:
+            self.__cfg.state_busy = False
+            raise e
+
+    def measure_helical(self, h: HelicalScanRequest) -> float:
+        start = time.perf_counter()
+
+        self.__cfg.try_set_busy(timeout=ceil(h.time_sec + 360))
+        self.__set_state(BeamlineStateEnum.DataCollection)
+
+        try:
+            self.__devs.smargon.target = h.start
+            self.__devs.aerotech.measure_standard(
+                h.omega_start, h.total_omega, h.time_sec
+            )
+
+            smargon_time_step = h.time_sec / float(h.steps)
+            pos_step = (h.end.sh_mm - h.start.sh_mm) * (1.0 / float(h.steps))
+
+            for i in range(h.steps):
+                self.__devs.smargon.target = SmargonCoordinate(
+                    sh_mm=h.start.sh_mm + pos_step * i
+                )
+                time.sleep(smargon_time_step)
+
+            self.__devs.aerotech.wait_scan_done(h.time_sec + 60)
+            self.__devs.aerotech.reset()
+            self.__cfg.state_busy = False
+        except Exception as e:
+            self.__cfg.state_busy = False
+            raise e
+        end = time.perf_counter()
+        return end - start
+
+    def __rotation(self, request: RotationScanRequest):
+        self.__set_state(BeamlineStateEnum.DataCollection)
+        self.__devs.aerotech.measure_standard(
+            request.start_omega_deg, request.total_omega_deg, request.total_time_sec
+        )
+        self.__devs.aerotech.wait_scan_done(request.total_time_sec + 60)
+        self.__devs.aerotech.reset()
+
+    def measure_rotation(self, request: RotationScanRequest):
+        self.__cfg.try_set_busy(timeout=ceil(request.total_time_sec + 360))
+
+        try:
+            self.__rotation(request)
+            self.__set_state(BeamlineStateEnum.SampleAlignment)
+            self.__cfg.state_busy = False
+        except Exception as e:
+            self.__cfg.state_busy = False
+            raise e
+
+    def __screening(self, request: ScreeningScanRequest):
+        print(request)
+        self.__devs.aerotech.measure_screening(request.omega_start_deg,
+                                               request.oscrange_deg,
+                                               request.exp_time_s,
+                                               request.omega_step_deg,
+                                               request.steps)
+        self.__devs.aerotech.wait_scan_done(request.exp_time_s * request.steps + 60)
+        self.__devs.aerotech.reset()
+
+    def measure_screening(self, request: ScreeningScanRequest) -> float:
+        start = time.perf_counter()
+
+        self.__cfg.try_set_busy(timeout=ceil(360))
+        self.__set_state(BeamlineStateEnum.DataCollection)
+
+        try:
+            self.__screening(request)
+            self.__set_state(BeamlineStateEnum.SampleAlignment)
+            self.__cfg.state_busy = False
+        except Exception as e:
+            self.__cfg.state_busy = False
+            raise e
+        end = time.perf_counter()
+        return end - start
+
+    def get_background(self):
+        #if self.sample is not None:
+        #    raise Exception("Background cannot be measured when sample is mounted")
+
+        self.__cfg.try_set_busy(timeout=360)
+        self.__set_state(BeamlineStateEnum.SampleAlignment)
+        try:
+            self.__devs.lamp_light.move("on", wait=True)
+
+            for s in self.__cfg.alc_zoom_settings.z:
+                self.__devs.sample_cam.setup(expo=s.sam_cam_exp, gain=s.sam_cam_gain)
+                self.__devs.zoom = s.zoom_value
+                time.sleep(5.0) # Wait for settings to stabilize
+
+                image = self.__devs.sample_cam.get_image(gray=False)
+                self.__cfg.put_alc_bkg(s.zoom_value, image)
+                bgr_array = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+                cv2.imwrite(f"bkg{s.zoom_value:.0f}.jpg", bgr_array)
+            self.__cfg.state_busy = False
+            self.__devs.sample_cam.set_auto()
+        except Exception as e:
+            self.__cfg.state_busy = False
+            raise e
+
+
+    @property
+    def dtz(self) -> float:
+        tmp = self.__cfg.dtz
+        if tmp is None:
+            return 150.0
+        else:
+            return tmp
+
+    @dtz.setter
+    def dtz(self, val: float):
+        self.__cfg.try_set_busy(timeout=360)
+        state = self.__cfg.state
+        if state == BeamlineStateEnum.DataCollection:
+            self.__cfg.state_busy = False
+            raise RuntimeError("Cannot set dtz during data collection")
+        elif state == BeamlineStateEnum.SampleAlignment:
+            self.__devs.dtz.move(val)
+        self.__cfg.dtz = val
+        self.__cfg.state_busy = False
+
+    @property
+    def smargon(self) -> SmargonCoordinate:
+        return self.__devs.smargon.readback
+
+    @smargon.setter
+    def smargon(self, sc: SmargonCoordinate):
+        self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
+        try:
+            self.__devs.smargon.target = sc
+            self.__devs.smargon.wait()
+            self.__cfg.state_busy = False
+        except Exception as e:
+            self.__cfg.state_busy = False
+            raise e
+        pass
+
+    def mark_beam(self, x_pxl: float, y_pxl: float):
+        self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
+        try:
+            self.__cfg.mark_beam(x_pxl, y_pxl, self.__devs.zoom)
+            self.__cfg.state_busy = False
+        except Exception as e:
+            self.__cfg.state_busy = False
+            raise e
+
+    def clear_mark_beam(self):
+        self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
+        self.__cfg.clear_mark_beam()
+        self.__cfg.state_busy = False
+
+    @property
+    def sample_geometry(self) -> SampleGeometryModel:
+        zoom = self.__devs.zoom
+        abr_pos = self.__cfg.abr_meas_pos
+        return SampleGeometryModel(
+            beam_location_pxl=self.__cfg.beam_center_coeff.apply(zoom),
+            pixel_in_mm=self.__cfg.pixel_to_mm(zoom),
+            omega_deg=self.__devs.aerotech.omega,
+            smargon=self.__devs.smargon.readback,
+            beam_size_mm=Coordinate(x=0.040, y=0.040),
+            aerotech=Coordinate(
+                x=self.__devs.gmx.value - abr_pos.x,
+                y=self.__devs.gmy.value - abr_pos.y,
+                z=self.__devs.gmz.value - abr_pos.z,
+            ),
+            aerotech_meas=self.__devs.abr_pos
+        )
+
+    def get_beam_center(self):
+        return self.__cfg.get_beam_center(self.__devs.zoom)
+
+    def listen_changes(self) -> redis.client.PubSub:
+        return self.__cfg.listen_changes()
+
+    def __ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None, move: bool = False) -> None | Tuple[float, float, float, float]:
+        time.sleep(0.2) # Just to be sure image is stable
+        #curr_image = self.camera_image
+        #box = self.__mlbox.predict(curr_image)
+        bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR)
+        box = self.__mlbox.predict(bgr_image)
+        if box is None:
+            if filename is not None:
+                cv2.imwrite(f"{filename}_no_detection.jpg", bgr_image)
+            return None
+        x1, y1, x2, y2 = box
+        if filename is not None:
+            cv2.rectangle(bgr_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
+            cv2.imwrite(f"{filename}.jpg", bgr_image)
+            if sample_id is not None:
+                self.__devs.aare.upload_image(sample_id, filename, bgr_image)
+        if move:
+            geom = self.sample_geometry
+            coord = geom.picture_to_smargon(Coordinate(x=x1, y=y1))
+            self.__devs.smargon.target = SmargonCoordinate(sh_mm=coord)
+            self.__devs.smargon.wait(60)
+        return box
+
+    def ml_bounding_box(self,  sample_id: int | None = None, filename: str | None = None, move: bool = False) -> None | Tuple[float, float, float, float]:
+        try:
+            self.__cfg.try_set_busy(timeout=360)
+            box = self.__ml_bounding_box(sample_id, filename, move)
+            if not box:
+                raise LoopCenteringFailed
+            self.__cfg.state_busy = False
+            return box
+        except Exception:
+            self.__cfg.state_busy = False
+            raise
+
+    def __loop_center(self, filename: str | None = "") -> SmargonCoordinate:
+        curr_image = self.camera_image
+
+        gray_without_feature = cv2.cvtColor(
+            self.__cfg.get_alc_bkg(self.zoom), cv2.COLOR_RGB2GRAY
+        )
+        gray_with_feature = cv2.cvtColor(curr_image, cv2.COLOR_RGB2GRAY)
+
+        # Subtract the "without feature" image from the "with feature" image
+        diff_image = cv2.absdiff(gray_with_feature, gray_without_feature)
+
+        # Threshold the difference to isolate the feature
+        _, thresh = cv2.threshold(diff_image, 70, 255, cv2.THRESH_BINARY)
+        # thresh = cv2.adaptiveThreshold(diff_image, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
+
+        if filename is not None:
+            cv2.imwrite(f"{filename}_diff.jpg", diff_image)
+            cv2.imwrite(f"{filename}_thresh.jpg", thresh)
+
+        # Find contours of the detected feature
+        contours, _ = cv2.findContours(
+            thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        )
+
+        leftmost_point = None
+
+        # Loop through each contour
+        for contour in contours:
+            # Loop through each point in the current contour
+            for point in contour:
+                x, y = point[0]  # Point is a nested array [ [x, y] ]
+
+                # Check if this is the leftmost point
+                if leftmost_point is None or x < leftmost_point[0]:
+                    leftmost_point = (x, y)
+
+        if leftmost_point is None:
+            raise LoopCenteringFailed()
+
+        geom = self.sample_geometry
+
+        coord = geom.picture_to_smargon(
+            Coordinate(x=leftmost_point[0], y=leftmost_point[1])
+        )
+        return SmargonCoordinate(sh_mm=coord)
+
+    def __loop_center_sequence(self, sample_id: int | None = None) -> bool:
+        self.__set_state(BeamlineStateEnum.SampleAlignment)
+        self.__devs.lamp_light.move("on", wait=True)
+        try:
+            i = 0
+            for s in self.__cfg.alc_zoom_settings.z:
+                self.__devs.sample_cam.setup(expo=s.sam_cam_exp, gain=s.sam_cam_gain)
+                self.__devs.zoom_sync(s.zoom_value)
+                if i % 2 == 0:
+                    angles = (0, 90)
+                else:
+                    angles = (90, 0)
+                for angle in angles:
+                    self.__devs.aerotech.move(angle, wait=True)
+                    filename = None
+                    if sample_id is not None:
+                        filename = f"{angle}_{s.zoom_value:.0f}"
+                    self.__devs.smargon.target = self.__loop_center(filename)
+                    self.__devs.smargon.wait(60)
+                    if sample_id is not None:
+                        time.sleep(0.1)
+                        self.save_screenshot_db(sample_id, f"{angle}_{s.zoom_value:.0f}")
+                i += 1
+            return True
+        except LoopCenteringFailed:
+            return False
+
+    def auto_loop_center(self, sample_id: int | None = None) -> float:
+        start = time.perf_counter()
+        try:
+            self.__cfg.try_set_busy(timeout=360)
+            if not self.__loop_center_sequence(sample_id):
+                raise LoopCenteringFailed
+            self.__cfg.state_busy = False
+        except Exception:
+            self.__cfg.state_busy = False
+            raise
+
+        end = time.perf_counter()
+        return end - start
+
+    def save_screenshot(self, filename: str):
+        time.sleep(0.2) # Wait 200 ms to ensure camera image is stable
+        bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(f"{filename}.jpg", bgr_image)
+
+    def save_screenshot_db(self, sample_id: int, filename: str):
+        time.sleep(0.2)  # Wait 200 ms to ensure camera image is stable
+        bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR)
+        self.__devs.aare.upload_image(sample_id, filename, bgr_image)
+
+    @property
+    def sample_spreadsheet(self) -> SampleShortInfoList:
+        return self.__devs.aare.get_sample_info()
+
+    def sample_spreadsheet_user(self, pgroup: str) -> SampleShortInfoList:
+        sample = self.sample_spreadsheet
+        sample.s = list(filter(lambda x: x.user == pgroup, sample.s))
+        return sample
+
+    def measure(self, sample: SampleShortInfo) -> float:
+        start = time.perf_counter()
+        if sample.location is not None:
+            sample_name = "{}{}S{:02d}_{}".format(
+                sample.location.segment,
+                sample.location.pos,
+                sample.pin,
+                sample.sample_name
+            )
+        else:
+            sample_name = sample.sample_name
+        try:
+            self.__cfg.try_set_busy(timeout=360)
+
+            self.__mount(sample)
+            self.__set_state(BeamlineStateEnum.SampleAlignment)
+            self.save_screenshot_db(sample.db_id, "mount.jpeg")
+            self.__loop_center_sequence(sample.db_id)
+            self.save_screenshot_db(sample.db_id, "after_lc.jpeg")
+
+            self.__set_state(BeamlineStateEnum.DataCollection)
+
+            self.__devs.aerotech.move(0, wait=True)
+
+            self.__ml_bounding_box(sample.db_id, f"bb_raster_0", move=True)
+
+            self.__raster(
+                RasterGridRequest(
+                    n_x=25,
+                    n_y=15,
+                    grid_size_mm=Coordinate(x=0.02, y=0.02),
+                    smargon=SmargonCoordinate(),
+                    omega_deg=0,
+                    exp_time_s=0.01,
+                )
+            )
+
+            self.__devs.aerotech.move(90, wait=True)
+
+            self.__ml_bounding_box(sample.db_id, f"bb_raster_90", move=True)
+
+            self.__raster(
+                RasterGridRequest(
+                    n_x=25,
+                    n_y=10,
+                    grid_size_mm=Coordinate(x=0.02, y=0.02),
+                    smargon=SmargonCoordinate(),
+                    omega_deg=90,
+                    exp_time_s=0.01,
+                )
+            )
+
+            self.__rotation(
+                RotationScanRequest(total_omega_deg=360,
+                                    start_omega_deg=0,
+                                    total_time_sec=18))
+
+            self.save_screenshot_db(sample.db_id, "after_dc.jpeg")
+
+            self.__set_state(BeamlineStateEnum.RobotSampleExchange)
+            self.__cfg.state_busy = False
+        except Exception:
+            self.__cfg.state_busy = False
+            raise
+
+        end = time.perf_counter()
+        return end - start
+
+    def __set_state(self, target: BeamlineStateEnum):
+        # __set_state assumes that beamline is already in busy state
+        # it will apply a proper transformation and change state afterward
+        # specifically:
+        #    1. If target state is maintenance, just go there
+        #    2. If target state is same as current, nothing will happen
+        #    3. If target state cannot be reached, exception is raised and current state is kept
+        #    4. If exception is raised during transformation, state is set to maintenance
+        #    5. If transformation goes OK, target state is set
+        #
+        # Busy state will be cleared only, if exception is raised.
+        #
+        curr_state = self.__cfg.state
+
+        if target == BeamlineStateEnum.Maintenance:
+            self.__cfg.state = BeamlineStateEnum.Maintenance
+        elif target != curr_state:
+            self.__cfg.state = BeamlineStateEnum.Moving
+            try:
+                match curr_state:
+                    case BeamlineStateEnum.Maintenance:
+                        if target == BeamlineStateEnum.SampleExchange:
+                            workflows.m2se(self.__devs, self.__cfg)
+                        else:
+                            raise TransformationInvalidException()
+                    case BeamlineStateEnum.SampleExchange:
+                        if target == BeamlineStateEnum.SampleAlignment:
+                            workflows.se2sa(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.BeamLocation:
+                            workflows.se2sa(self.__devs, self.__cfg)
+                            workflows.sa2bl(self.__devs, self.__cfg)
+                        else:
+                            raise TransformationInvalidException()
+                    case BeamlineStateEnum.DewarTransfer:
+                        if target == BeamlineStateEnum.SampleAlignment:
+                            workflows.dh2sa(self.__devs, self.__cfg)
+                        else:
+                            raise TransformationInvalidException()
+                    case BeamlineStateEnum.DataCollection:
+                        if target == BeamlineStateEnum.SampleExchange:
+                            workflows.dc2se(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.SampleAlignment:
+                            workflows.dc2sa(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.RobotSampleExchange:
+                            workflows.dc2rse(self.__devs, self.__cfg)
+                        else:
+                            raise TransformationInvalidException()
+                    case BeamlineStateEnum.BeamLocation:
+                        if target == BeamlineStateEnum.SampleAlignment:
+                            workflows.bl2sa(self.__devs, self.__cfg)
+                        if target == BeamlineStateEnum.SampleExchange:
+                            workflows.bl2sa(self.__devs, self.__cfg)
+                            workflows.sa2se(self.__devs, self.__cfg)
+                        else:
+                            raise TransformationInvalidException()
+                    case BeamlineStateEnum.SampleAlignment:
+                        if target == BeamlineStateEnum.DewarTransfer:
+                            workflows.sa2dh(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.SampleExchange:
+                            workflows.sa2se(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.RobotSampleExchange:
+                            workflows.sa2rse(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.DataCollection:
+                            workflows.sa2dc(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.XrayFluorescence:
+                            workflows.sa2xrf(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.BeamLocation:
+                            workflows.sa2bl(self.__devs, self.__cfg)
+                        else:
+                            raise TransformationInvalidException()
+                    case BeamlineStateEnum.XrayFluorescence:
+                        if target == BeamlineStateEnum.SampleAlignment:
+                            workflows.xrf2sa(self.__devs, self.__cfg)
+                        else:
+                            raise TransformationInvalidException()
+                    case BeamlineStateEnum.RobotSampleExchange:
+                        if target == BeamlineStateEnum.SampleAlignment:
+                            workflows.se2sa(self.__devs, self.__cfg)
+                        else:
+                            raise TransformationInvalidException()
+                self.__cfg.state = target
+            except TransformationInvalidException as e:
+                self.__cfg.state = curr_state
+                self.__cfg.state_busy = False
+                raise e
+            except Exception as e:
+                self.__cfg.state = BeamlineStateEnum.Maintenance
+                self.__cfg.state_busy = False
+                raise e
+
+
+    @property
+    def status(self) -> DAQStatusModel:
+        # session_status should be set by FastAPI server
+        return DAQStatusModel(
+            geom=self.sample_geometry,
+            state=self.state,
+            busy=self.busy,
+            light=self.light,
+            current_pgroup=self.__cfg.pgroup,
+            sample=self.sample
+        )
+

src/aaredaq/devices.py

@@ -0,0 +1,219 @@
+from aaredaqlib.coordinate import Coordinate
+from epics import PV
+
+from aaredaq.aaredb import AareWrapper
+from aaredaq.beamline import MXBeamline
+from mxlibs3 import aerotech, smargon
+from mxlibs3.area_detector import epicsAD
+from mxlibs3.enum_pv import EnumPv
+from mxlibs3.my_motor import MyMotor
+from mxlibs3.non_standard import NonStandard
+from mxlibs3.tell_client import TellClient
+from mxlibs3.workflow_tools import wait_position
+
+
+#from bec_lib.client import BECClient
+#from bec_lib.service_config import ServiceConfig
+
+class BeamlineDevices:
+    def __init__(self, beamline: MXBeamline):
+        #service_config = ServiceConfig(redis={"host": "x06da-bec-001.psi.ch", "port": 6379})
+        #client = BECClient(service_config, name="Filips-Custom-Client")
+        #client.start()
+
+        #self.__bec_dev = client.device_manager.devices
+        #self.__bec_scans = client.scans
+
+        self.aare = AareWrapper(beamline)
+
+        BEAMLINE = beamline.value.upper()
+
+        self.tell = TellClient(beamline)
+        self.smargon = smargon.Smargon(beamline)
+
+        self.magnet_position_sensor = PV(f"{BEAMLINE}-ES-DF1:CBOX-CMP1")
+        self.magnet_position_sensor_readout = PV(f"{BEAMLINE}-ES-DF1:CBOX-USER1")
+
+        self.dtz = MyMotor(f"{BEAMLINE}-ES-DET:TRZ1")
+
+        # self.dty = MyMotor(f"{BEAMLINE}-MOCK-DET:TRY1")
+
+        self.bsz = MyMotor(f"{BEAMLINE}-ES-BS:TRZ1")
+
+        self.bsx = MyMotor(f"{BEAMLINE}-ES-BS:TRX1")
+        self.bsy = MyMotor(f"{BEAMLINE}-ES-BS:TRY1")
+
+        self.aerotech = aerotech.Abr(beamline)
+
+        self.aerotech.reload_programs()
+
+        self.gmx = NonStandard(
+            name="GMX",
+            setpv=f"{BEAMLINE}-ES-DF1:GMX-VAL",
+            getpv=f"{BEAMLINE}-ES-DF1:GMX-RBV",
+            speed=(f"{BEAMLINE}-ES-DF1:GMX-SETV", f"{BEAMLINE}-ES-DF1:GMX-SETV"),
+            move_done_when=(f"{BEAMLINE}-ES-DF1:GMX-DONE", 1),
+        )
+
+        self.gmy = NonStandard(
+            name="GMY",
+            setpv=f"{BEAMLINE}-ES-DF1:GMY-VAL",
+            getpv=f"{BEAMLINE}-ES-DF1:GMY-RBV",
+            speed=(f"{BEAMLINE}-ES-DF1:GMY-SETV", f"{BEAMLINE}-ES-DF1:GMY-SETV"),
+            move_done_when=(f"{BEAMLINE}-ES-DF1:GMY-DONE", 1),
+        )
+
+        self.gmz = NonStandard(
+            name="GMZ",
+            setpv=f"{BEAMLINE}-ES-DF1:GMZ-VAL",
+            getpv=f"{BEAMLINE}-ES-DF1:GMZ-RBV",
+            speed=(f"{BEAMLINE}-ES-DF1:GMZ-SETV", f"{BEAMLINE}-ES-DF1:GMZ-SETV"),
+            move_done_when=(f"{BEAMLINE}-ES-DF1:GMZ-DONE", 1),
+        )
+
+        self.lamp_light = NonStandard(
+            name="FrontLight",
+            setpv=f"{BEAMLINE}-ES-FL:SET-BRGHT",
+            getpv=f"{BEAMLINE}-ES-FL:SET-BRGHT",
+            tolerance=0.01,
+            predefs={
+                "on": 2.5,
+                "half": 1.9,
+                "off": 1.0,
+            },
+        )
+
+        self.cryo = NonStandard(
+            name="CryoJet",
+            setpv=f"{BEAMLINE}-ES-CJ:TRX1",
+            getpv=f"{BEAMLINE}-ES-CJ:TRX1.RBV",
+            move_done_when=(f"{BEAMLINE}-ES-CJ:TRX1.DMOV", 1),
+            speed=(f"{BEAMLINE}-ES-CJ:TRX1.VELO", f"{BEAMLINE}-ES-CJ:TRX1.VELO"),
+        )
+
+        self.__zoom = NonStandard(
+            name="Zoom",
+            setpv=f"{BEAMLINE}-ES-SAMCAM:ZOOM.VAL",
+            getpv=f"{BEAMLINE}-ES-SAMCAM:ZOOM.RBV",
+            move_done_when=(f"{BEAMLINE}-ES-SAMCAM:ZOOM.DMOV", 1),
+            stoppv=(f"{BEAMLINE}-ES-SAMCAM:ZOOM.STOP", 1),
+        )
+
+        # shutter = EnumPv(
+        #    {
+        #        "name": "ExposureBoxShutter",
+        #        "pv": f"{BEAMLINE}-MOCK-PH1:SET",
+        #        "readback": f"{BEAMLINE}-MOCK-PH1:SET",
+        #        "states": {"closed": "closed", "opened": "opened"},
+        #    }
+        # )
+
+        # xrf_stage = EnumPv(
+        #    {
+        #        "name": "KetekStage",
+        #        "pv": f"{BEAMLINE}-MOCK-FD:SET-POS",
+        #        "readback": f"{BEAMLINE}-MOCK-FD:GET-POS",
+        #        "states": {"park": "park", "measure": "measure"},
+        #    }
+        # )
+
+        self.__reflector = EnumPv(
+            {
+                "name": "BackReflector",
+                "pv": f"{BEAMLINE}-ES-BL:SET-POS",
+                "readback": f"{BEAMLINE}-ES-BL:GET-POS",
+                "states": {
+                    "park": "park",
+                    "measure": "measure",
+                },
+            }
+        )
+
+        # detector_cover = EnumPv(
+        #    {
+        #        "name": "DetectorCover",
+        #        "pv": f"{BEAMLINE}-MOCK-DETCOV:SET",
+        #        "readback": f"{BEAMLINE}-MOCK-DETCOV:GET",
+        #        "timeout": 10.0,
+        #        "move_done_when": (
+        #            f"{BEAMLINE}-MOCK-DETCOV:GET",
+        #            ("Closed", "Opened", "Error"),
+        #        ),
+        #        "states": {
+        #            "opened": "opened",
+        #            "closed": "closed",
+        #        },
+        #    }
+        # )
+
+        self.__beamstop_stage = EnumPv(
+            {
+                "name": "Beamstop_Stage",
+                "pv": f"{BEAMLINE}-ES-BS:SET-POS",
+                "readback": f"{BEAMLINE}-ES-BS:GET-POS",
+                "states": {
+                    "park": "park",
+                    "measure": "measure",
+                },
+            }
+        )
+
+        self.sample_cam = epicsAD(f"{BEAMLINE}-SAMCAM:")
+
+    @property
+    def zoom(self) -> float:
+        return self.__zoom.value
+
+    @zoom.setter
+    def zoom(self, value: float):
+        self.__zoom.move(value, wait=False)
+        # status = self.__bec_dev.samzoom.move(value)
+        # status.wait()
+
+    def zoom_sync(self, value: float):
+        self.__zoom.move(value, wait=True)
+
+    @property
+    def reflector_up(self) -> bool:
+        return bool(self.__reflector.position_is("measure"))
+
+    @reflector_up.setter
+    def reflector_up(self, value: bool):
+        if value:
+            self.__reflector.move("measure", wait=True)
+        else:
+            self.__reflector.move("park", wait=True)
+
+    @property
+    def beamstop_stage_up(self) -> bool:
+        return bool(self.__beamstop_stage.position_is("measure"))
+
+    @beamstop_stage_up.setter
+    def beamstop_stage_up(self, value: bool):
+        if value:
+            self.__beamstop_stage.move("measure", wait=True)
+        else:
+            self.__beamstop_stage.move("park", wait=True)
+
+    @property
+    def abr_pos(self) -> Coordinate:
+        return Coordinate(x=self.gmx.readback, y=self.gmy.readback, z=self.gmz.readback)
+
+    @abr_pos.setter
+    def abr_pos(self, pos: Coordinate):
+        print(f"Going to x {pos.x} y {pos.y} z {pos.z}")
+        self.gmx.move(pos.x, wait=True)
+        self.gmy.move(pos.y, wait=True)
+        self.gmz.move(pos.z, wait=True)
+
+        wait_position(self.gmx, pos.x, 0.005)
+        wait_position(self.gmy, pos.y, 0.005)
+        wait_position(self.gmz, pos.z, 0.005)
+
+    def enable_motors(self):
+        self.aerotech.unlock()
+        self.dtz.refresh()  # Why ??
+        self.bsz.refresh()  # Why ??
+
+    def disable_motors(self):
+        self.aerotech.lock()

src/aaredaq/mlbox.py

@@ -0,0 +1,22 @@
+from typing import Tuple
+
+from ultralytics import YOLO
+
+class MlBox:
+    def __init__(self, model_path: str = "/home/leonarski_f/mxlibs3/src/aaredaq/best.pt"):
+        self.__model = YOLO(model_path)
+
+    @staticmethod
+    def __get_pred(results) -> None | Tuple[float, float, float, float]:
+        for r in results:
+            for box in r.boxes.data.tolist():
+                x1, y1, x2, y2, conf, cls = box
+                if int(cls) == 0:
+                    return x1, y1, x2, y2
+        return None
+
+    def predict(self, image) -> None | Tuple[float, float, float, float]:
+        results = self.__model.predict(source=image, conf=0.5)
+        pred = self.__get_pred(results)
+
+        return pred

src/aaredaq/models.py

@@ -0,0 +1,157 @@
+from enum import Enum
+from typing import Annotated, Literal, Tuple, List, Optional
+
+from pydantic import BaseModel, Field
+
+from aaredaqlib.coordinate import Coordinate
+from aaredaqlib.sample_geometry import SampleGeometryModel
+
+
+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
+
+
+# From database to TELL after loading
+class SampleShortInfo(BaseModel):
+    db_id: int
+    puck_name: str
+    dewar_name: str
+    sample_name: str
+    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
+    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}"
+
+
+class SampleShortInfoList(BaseModel):
+    s: List[SampleShortInfo]
+
+
+class BeamCenterCoeffModel(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 LoopCenteringZoomModelElem(BaseModel):
+    zoom_value: float
+    sam_cam_gain: float
+    sam_cam_exp: float
+
+
+class LoopCenteringZoomModel(BaseModel):
+    z: List[LoopCenteringZoomModelElem]
+
+
+# zoom gain exp (X06DA)
+#    1  50 0.05
+#  200  50 0.05
+#  500  70 0.05
+#  700 150 0.05
+#  800 100 0.10
+# 1000 200 0.10
+
+
+def def_loop_centering_zoom() -> LoopCenteringZoomModel:
+    return LoopCenteringZoomModel(
+        z=[
+            LoopCenteringZoomModelElem(
+                zoom_value=200, sam_cam_gain=50.0, sam_cam_exp=0.05
+            )
+        ]
+    )
+
+
+class BeamlineStateEnum(Enum):
+    Maintenance = 1
+    SampleExchange = 2
+    SampleAlignment = 3
+    DataCollection = 4
+    DewarTransfer = 5
+    XrayFluorescence = 6
+    BeamLocation = 7
+    Moving = 8
+    RobotSampleExchange = 9
+    # Busy state is used whenever transition between two states happen or mounting/data collection procedure happens
+
+
+class SessionsStateEnum(Enum):
+    OwnedByYou = 1
+    OwnedByElse = 2
+    Vacant = 3
+
+class DAQStatusModel(BaseModel):
+    geom: SampleGeometryModel
+    state: BeamlineStateEnum
+    busy: bool
+    light: Annotated[float, Field(ge=0.0, le=100.0)]
+    session: SessionsStateEnum = SessionsStateEnum.Vacant
+    sample: SampleShortInfo | None = None
+    current_pgroup: str | None
+    staff: bool = False
+
+
+class BeamlineSettingsModel(BaseModel):
+    # Float properties
+    cryojet_park_position: float = 12.0
+    cryojet_measurement_position: float = 5.0
+    dtz_collection: float = 130
+    dtz_park: float = 150.0
+    dtz_wash_sample_distance: float = 150.0
+    dtz_bsz_safety_margin: float = 50.0
+    bsz: float = 25.0
+
+    camera_max_magnification: float = 1.0
+    camera_min_magnification: float = 500.0
+    camera_translation_factor_a: float = 0.00253
+    camera_translation_factor_b: float = 255.0
+
+    cryojet_in_use: bool = True

src/aaredaq/server.py

@@ -0,0 +1,349 @@
+import io
+from typing import Tuple
+
+import urllib3
+from fastapi import HTTPException
+import logging
+
+import cv2
+import uvicorn
+from fastapi import FastAPI, Depends
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
+from starlette.responses import StreamingResponse
+from urllib3.exceptions import InsecureRequestWarning
+
+from aaredaq import auth
+from aaredaq.auth import TokenData
+from aaredaq.beamline import mx_beamline
+from aaredaq.config import BeamlineConfig
+from aaredaq.daq import AareDAQ
+from aaredaq.models import SampleShortInfo, DAQStatusModel, BeamlineStateEnum, BeamlineSettingsModel, \
+    SampleShortInfoList
+from aaredaqlib.coordinate import SmargonCoordinate, Coordinate
+from aaredaqlib.helical_geometry import HelicalScanRequest
+from aaredaqlib.raster_grid import RasterGridRequest
+from aaredaqlib.rotation_scan import RotationScanRequest
+from aaredaqlib.sample_geometry import SampleGeometryModel
+from aaredaqlib.screening_scan import ScreeningScanRequest
+
+app = FastAPI()
+
+# OAuth2 setup
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+bl = mx_beamline()
+cfg = BeamlineConfig(bl)
+daq = AareDAQ(cfg, bl)
+
+
+@app.post("/token")
+async def login(form_data: OAuth2PasswordRequestForm = Depends()):
+    data = auth.authenticate_user(cfg, form_data)
+    return {"access_token": data, "token_type": "bearer"}
+
+
+@app.get("/status")
+async def status(token: str = Depends(oauth2_scheme)) -> DAQStatusModel:
+    data = auth.parse_token(token)
+    auth.check_jwt_ro(cfg, data)
+    ret = daq.status
+    ret.session = cfg.session_state(data.session)
+    ret.staff = data.staff
+    return ret
+
+
+@app.get("/beamline/geometry")
+async def sample_geometry(
+        token: str = Depends(oauth2_scheme),
+) -> SampleGeometryModel:
+    auth.check_jwt_ro(cfg, auth.parse_token(token))
+    return daq.sample_geometry
+
+
+@app.put("/beamline/omega")
+async def omega(val: float, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.omega = val
+    return "OK"
+
+
+@app.put("/beamline/light")
+async def light(val: float, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.light = val
+    return "OK"
+
+
+@app.put("/beamline/zoom")
+async def zoom(val: float, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.zoom = val
+    return "OK"
+
+
+@app.put("/beamline/smargon")
+async def smargon(val: SmargonCoordinate, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.smargon = val
+    return "OK"
+
+@app.post("/beamline/tweak_abr_meas_pos")
+async def tweak_abr_meas_pos(val: Coordinate, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_staff(auth.parse_token(token))
+    daq.tweak_abr_meas_pos(val)
+    return "OK"
+
+@app.post("/beamline/save_abr_meas_pos")
+async def save_abr_meas_pos(token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_staff(auth.parse_token(token))
+    daq.save_abr_meas_pos()
+    return "OK"
+
+@app.post("/beam_center/mark")
+async def mark_beam_center(x: float, y: float, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_staff(auth.parse_token(token))
+    daq.mark_beam(x, y)
+    return "OK"
+
+@app.post("/beam_center/clear")
+async def clear_beam_center(token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_staff(auth.parse_token(token))
+    daq.clear_mark_beam()
+    return "OK"
+
+async def get_image(token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_ro(cfg, auth.parse_token(token))
+
+    _, encoded_image = cv2.imencode(
+        ".jpg", daq.camera_image
+    )  # Encodes the image in JPEG format
+    image_bytes = io.BytesIO(
+        encoded_image.tobytes()
+    )  # Convert OpenCV byte format to a file-like object
+
+    return StreamingResponse(image_bytes, media_type="image/jpeg")
+
+
+# TELL procedures
+@app.get("/tell/sample")
+async def sample(token: str = Depends(oauth2_scheme)) -> SampleShortInfo:
+    token_data = auth.parse_token(token)
+    auth.check_jwt_ro(cfg, auth.parse_token(token))
+    s = daq.sample
+    if token_data.staff or token_data.group == s.user:
+        return daq.sample
+    else:
+        return SampleShortInfo(
+            sample_name="Other user sample",
+            puck_name="",
+            dewar_name="",
+            db_id = -1,
+            pin = sample.pin,
+            location=sample.location
+        )
+
+@app.post("/tell/mount")
+async def mount(dbid: int, token: str = Depends(oauth2_scheme)):
+    token_data = auth.parse_token(token)
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    #st = daq.sample_spreadsheet
+    st = SampleShortInfoList(s=[])
+
+    index = -1
+    for i in range(len(st.s)):
+        if st.s[i].db_id == dbid:
+            index = i
+
+    if index == -1:
+        raise RuntimeError("Sample not found")
+    if token_data.staff or token_data.group == st.s[index].user:
+        daq.sample = st.s[index]
+        return "OK"
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Sample belongs to a different user.",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+
+@app.post("/tell/unmount")
+async def unmount(token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.sample = None
+    return "OK"
+
+
+@app.get("/tell/spreadsheet")
+async def spreadsheet(token: str = Depends(oauth2_scheme)):
+    data = auth.parse_token(token)
+    if data.staff:
+        return daq.sample_spreadsheet
+    else:
+        return daq.sample_spreadsheet_user(data.group)
+
+
+# State transitions
+@app.post("/state/dewar_exchange")
+async def dewar_exchange(token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.state = BeamlineStateEnum.DewarTransfer
+    return "OK"
+
+
+@app.post("/state/sample_exchange")
+async def dewar_exchange(token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.state = BeamlineStateEnum.SampleExchange
+    return "OK"
+
+
+@app.post("/state/sample_alignment")
+async def sample_alignment(token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.state = BeamlineStateEnum.SampleAlignment
+
+
+# Scans
+@app.post("/scan/raster")
+async def raster(val: RasterGridRequest, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.measure_raster(val)
+    return "OK"
+
+
+@app.post("/scan/rotation")
+async def rotation(val: RotationScanRequest, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.measure_rotation(val)
+    return "OK"
+
+
+@app.post("/scan/screening")
+async def screening(val: ScreeningScanRequest, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.measure_screening(val)
+    return "OK"
+
+
+@app.post("/scan/helical")
+async def helical(val: HelicalScanRequest, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.measure_helical(val)
+    return "OK"
+
+
+@app.post("/scan/auto")
+async def auto(s: SampleShortInfo, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    runtime = daq.measure(s)
+    return f"{runtime:0.3f}"
+
+# ALC routines
+@app.post("/alc/background")
+async def alc_background(token: str = Depends(oauth2_scheme)) -> str:
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.get_background()
+    return "OK"
+
+
+@app.post("/alc/center_loop")
+async def alc_center_loop(token: str = Depends(oauth2_scheme)) -> str:
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    daq.auto_loop_center()
+    return "OK"
+
+@app.post("/alc/ml_bounding_box")
+async def alc_ml_bounding_box(token: str = Depends(oauth2_scheme), move: bool = False, filename: str = "") -> Tuple[float, float, float, float] | None:
+    auth.check_jwt_rw(cfg, auth.parse_token(token))
+    return daq.ml_bounding_box(move=move)
+
+# Access management
+@app.get("/access/pgroup")
+async def pgroup(token: str = Depends(oauth2_scheme)) -> str:
+    auth.check_jwt_ro(cfg, auth.parse_token(token))
+    return cfg.pgroup
+
+
+@app.put("/access/pgroup")
+async def set_pgroup(val: str, token: str = Depends(oauth2_scheme)) -> str:
+    auth.check_jwt_staff(auth.parse_token(token))
+    cfg.pgroup = val
+    return "OK"
+
+
+@app.delete("/access/pgroup")
+async def del_pgroup(token: str = Depends(oauth2_scheme)) -> str:
+    auth.check_jwt_staff(auth.parse_token(token))
+    cfg.pgroup = None
+    return "OK"
+
+
+@app.post("/access/end_session")
+async def end_session(token: str = Depends(oauth2_scheme)) -> str:
+    # End active session will only delete session, if it is equal to token value
+    # so no need to check R/W permissions
+    token_data = auth.parse_token(token)
+    cfg.end_active_session(token_data.session)
+    return "OK"
+
+
+@app.post("/access/force_current_session")
+async def force_current_session(token: str = Depends(oauth2_scheme)) -> str:
+    data = auth.parse_token(token)
+    # Counterintuitive, this operation requires only R/O permission
+    # as this is actually acquiring R/W permissions
+    auth.check_jwt_ro(cfg, data)
+    auth.force_current_sesion(cfg, data)
+    return "OK"
+
+@app.get("/beamline/settings")
+async def get_settings(token: str = Depends(oauth2_scheme)) -> BeamlineSettingsModel:
+    auth.check_jwt_staff(auth.parse_token(token))
+    return cfg.settings
+
+@app.put("/beamline/settings")
+async def put_settings(s: BeamlineSettingsModel, token: str = Depends(oauth2_scheme)):
+    auth.check_jwt_staff(auth.parse_token(token))
+    cfg.settings = s
+
+LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "default": {
+            "()": "uvicorn.logging.DefaultFormatter",
+            "fmt": "%(levelprefix)s %(message)s",
+            "use_colors": True,
+        },
+    },
+    "handlers": {
+        "default": {
+            "formatter": "default",
+            "class": "logging.StreamHandler",
+            "stream": "ext://sys.stdout",
+        },
+    },
+    "loggers": {
+        "uvicorn": {
+            "handlers": ["default"],
+            "level": "WARNING",
+        },
+        "uvicorn.access": {
+            "handlers": ["default"],
+            "level": "WARNING",
+        },
+    },
+}
+
+
+def main():
+    # Remove in production!
+    urllib3.disable_warnings(InsecureRequestWarning)
+
+    # Run the application using uvicorn
+    uvicorn.run("aaredaq.server:app", host="0.0.0.0", port=5210, workers=4, log_config=LOGGING_CONFIG)
+
+
+if __name__ == "__main__":
+    main()

src/aaredaq/workflows.py

@@ -0,0 +1,291 @@
+import time
+
+from epics import poll
+
+from aaredaq.config import ABR_POS_MOUNT, ABR_OMEGA_MOUNT
+from aaredaq.devices import BeamlineDevices
+from aaredaq.config import BeamlineConfig
+from aaredaqlib.coordinate import Coordinate
+from mxlibs3.mx_lib import pv_wait
+
+
+def move_bsz(devs: BeamlineDevices, target: float):
+    if abs(target - devs.bsz.position) > 0.1:
+        beamstop_stage_measure = devs.beamstop_stage_up
+        reflector_measure = devs.reflector_up
+        devs.reflector_up = False
+        devs.beamstop_stage_up = True
+
+        devs.bsz.move(target, wait=True)
+
+        if reflector_measure:
+            devs.reflector_up = True
+
+        if not beamstop_stage_measure:
+            devs.beamstop_stage_up = False
+
+
+def wait_for_dc_devices(devs: BeamlineDevices):
+    timeout = time.time() + 60
+    timeisup = False
+    # devs.detector_cover.wait()
+    while not timeisup and (
+            not devs.dtz.done_moving
+    ):
+        poll(0.1)
+        timeisup = timeout < time.time()
+    if timeisup:
+        if not devs.dtz.done_moving:
+            print("DTZ still moving.")
+        raise RuntimeError("Timeout moving devs. Dropping to maintenance.")
+
+
+def wait_for_se_devices(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    _cryo = hub.cryojet_in_use
+    _cryopark = hub.cryojet_park_position
+
+    if _cryo and not devs.cryo.position_is(_cryopark):
+        raise RuntimeError("Cryojet did not reach far position")
+
+
+def common_2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
+    devs.smargon.move_home(wait=True)
+    devs.aerotech.move(ABR_OMEGA_MOUNT, wait=True, direct=True)
+    devs.abr_pos = ABR_POS_MOUNT
+    devs.aerotech.lock()
+
+
+def m2se(devs: BeamlineDevices, cfg: BeamlineConfig):
+    print("executing MAINTENANCE -> sample exchange")
+    hub = cfg.settings
+
+    _dtzpark = hub.dtz_park
+    _cryopark = hub.cryojet_park_position
+    _cryo = hub.cryojet_in_use
+
+    # devs.detector_cover.move("closed")
+    if not devs.aerotech.is_ready_and_willing():
+        devs.aerotech.stop()
+        time.sleep(5.0)
+        raise RuntimeError("Aerotech device is not ready")
+
+    if devs.dtz.position < _dtzpark:
+        devs.dtz.move(_dtzpark)
+
+    devs.gmx.set_speed(100.0)
+    devs.abr_pos = ABR_POS_MOUNT
+    devs.aerotech.move(ABR_OMEGA_MOUNT)
+
+    if _cryo:
+        devs.cryo.move(_cryopark)
+
+    devs.reflector_up = False
+
+    devs.bsx.move(0.0)
+    devs.bsy.move(0.0)
+    move_bsz(devs, hub.bsz)
+
+    devs.beamstop_stage_up = False
+
+
+def sa2se(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    _cryopark = hub.cryojet_park_position
+    _dtzpark = hub.dtz_park
+
+    if devs.cryo.position < _cryopark:
+        devs.cryo.move(_cryopark)
+
+    devs.reflector_up = False
+    devs.beamstop_stage_up = False
+    devs.smargon.move_home(wait=True)
+    if devs.dtz.position < _dtzpark:
+        devs.dtz.move(_dtzpark)
+
+
+def sa2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
+    common_2rse(devs, cfg)
+
+
+def dc2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
+    if not devs.aerotech.is_ready_and_willing():
+        devs.aerotech.stop()
+        time.sleep(5.0)
+        raise RuntimeError("Aerotech device is not ready")
+    common_2rse(devs, cfg)
+
+
+def se2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    _cryo = hub.cryojet_in_use
+    cryo_meas = hub.cryojet_measurement_position
+    _dtz_safety = hub.dtz_bsz_safety_margin
+
+    devs.aerotech.unlock()
+    devs.aerotech.set_direct_mode()
+    devs.reflector_up = True
+    devs.beamstop_stage_up = True
+
+    devs.abr_pos = cfg.abr_meas_pos
+
+    if _cryo:
+        devs.cryo.move(cryo_meas)
+
+    devs.dtz.move(cfg.dtz)
+
+
+def sa2dc(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+
+    _cryo = hub.cryojet_in_use
+    cryo_meas = hub.cryojet_measurement_position
+
+    # devs.reflector = False
+
+    if _cryo:
+        devs.cryo.move(cryo_meas)
+
+    #devs.reflector.wait("park")
+    devs.beamstop_stage_up = True
+    devs.dtz.move(cfg.dtz, wait=True)
+    pv_wait(devs.dtz, None, tolerance=1.0)
+    wait_for_dc_devices(devs)
+
+
+def dc2se(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    _cryo = hub.cryojet_in_use
+    _dtzpark = hub.dtz_park
+
+    devs.cryo.move(hub.cryojet_park_position)
+    devs.aerotech.move(ABR_OMEGA_MOUNT, speed=180, wait=True, direct=True)
+
+    # devs.detector_cover.move("closed")
+
+    if not devs.aerotech.is_ready_and_willing():
+        devs.aerotech.stop()
+        time.sleep(5.0)
+        raise RuntimeError("Aerotech device is not ready")
+
+    if devs.dtz.position < _dtzpark:
+        devs.dtz.move(_dtzpark)
+
+
+def dc2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+
+    devs.reflector_up = True
+
+    devs.gmz.move(cfg.abr_meas_pos.z, wait=True)
+
+
+def sa2xrf(devs: BeamlineDevices, cfg: BeamlineConfig):
+    """sample alignment to XrfCollection"""
+    if devs.dtz.position < 150.0:
+        devs.dtz.move(150.0)
+    devs.reflector_up = False
+    # devs.xrf_stage.move("measure")
+    devs.beamstop_stage_up = True
+    move_bsz(devs, 70.0)
+
+
+def xrf2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
+    """XrfCollection to sample alignment"""
+    if devs.dtz.position < 150.0:
+        devs.dtz.move(150.0)
+    devs.reflector_up = False
+    # devs.xrf_stage.move("measure")
+    devs.beamstop_stage_up = True
+    move_bsz(devs, 70.0)
+
+
+def sa2ws(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    _dtzwash = hub.dtz_wash_sample_distance
+    devs.cryo.move(hub.cryojet_park_position)
+    if devs.dtz.position < _dtzwash:
+        devs.dtz.move(_dtzwash)
+    devs.reflector_up = False
+    devs.beamstop_stage_up = False
+
+
+def ws2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    cryo_meas = hub.cryojet_measurement_position
+    cryo_user = hub.cryojet_measurement_user
+    _cryo = hub.cryojet_in_use
+
+    if cryo_user >= cryo_meas:
+        cryo_meas = cryo_user
+    devs.dtz.move(cfg.dtz)
+    if _cryo:
+        devs.cryo.move(cryo_meas)
+    devs.reflector_up = True
+    devs.beamstop_stage_up = True
+
+
+def sa2ba(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+
+    if hub.cryojet_in_use:
+        devs.cryo.move(hub.cryojet_park_position)
+    devs.gmx.move(-38.0)
+    devs.lamp_light.move("on")
+    devs.beamstop_stage_up = True
+    move_bsz(devs, 0.0)
+
+
+def ba2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    devs.lamp_light.move("half")
+    devs.beamstop_stage_up = True
+    move_bsz(devs, hub.bsz)
+    devs.reflector_up = True
+    devs.abr_pos = cfg.abr_meas_pos
+
+
+def sa2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
+    devs.cryo.move(20)
+    devs.abr_pos = cfg.abr_meas_pos + Coordinate(x=1)
+    devs.beamstop_stage_up = True
+    devs.reflector_up = True
+
+
+def bl2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    move_bsz(devs, hub.bsz)
+    devs.reflector_up = True
+    devs.abr_pos = cfg.abr_meas_pos
+
+
+def bl2ba(devs: BeamlineDevices, cfg: BeamlineConfig):
+    """beam location to beamstop alignment"""
+    devs.gmx.move(-38)
+    devs.reflector_up = True
+    # devs.shutter.move("closed")
+    time.sleep(0.5)
+    move_bsz(devs, 0)
+
+
+def ba2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
+    hub = cfg.settings
+    move_bsz(devs, hub.bsz)
+    devs.abr_pos = cfg.abr_meas_pos + Coordinate(x=1)
+    devs.reflector_up = True
+
+
+def sa2dh(devs: BeamlineDevices, cfg: BeamlineConfig):
+    if devs.tell.get_mounted_sample() is not None:
+        try:
+            # Best effort try to unmount
+            devs.tell.unmount(wait=True)
+        except Exception as e:
+            print(f"Error for unmounting: {e}")
+    devs.tell.dry(wait_cold=-1, wait=False)
+    # devs.tell.move_park(wait=True)
+    devs.tell.set_samples_info(devs.aare.get_pucks_beamline())
+
+
+def dh2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
+    devs.aare.set_pucks_beamline(devs.tell.get_detected_pucks())

src/mxlibs3/area_detector.py

@@ -0,0 +1,121 @@
+import epics
+import numpy as np
+
+
+class epicsAD(object):
+    def __init__(self, prefix, cam="cam1:", image="image1:"):
+        self.img = None
+        self.monitored = False
+
+        self.gain = epics.PV(prefix + cam + "Gain")
+        self.acquire = epics.PV(prefix + cam + "Acquire")
+        self.color = epics.PV(prefix + cam + "ColorMode")
+        self.gain_mode = epics.PV(prefix + cam + "GainAuto")
+        # 0 - manual, GainMode, 2 - auto GainAuto
+        self.expo_mode = epics.PV(prefix + cam + "ExposureAuto")
+        # 0 - manual, ExposureMode, 2 - auto ExposureAutoself.gain = epics.PV(prefix + cam + "Gain")
+        self.expo = epics.PV(prefix + cam + "AcquireTime")
+
+        self.ndim = epics.PV(prefix + image + "NDimensions_RBV")
+        self.dim0 = epics.PV(prefix + image + "ArraySize0_RBV")
+        self.dim1 = epics.PV(prefix + image + "ArraySize1_RBV")
+        self.dim2 = epics.PV(prefix + image + "ArraySize2_RBV")
+        self.uid = epics.PV(prefix + image + "UniqueId_RBV")
+        self.data = epics.PV(prefix + image + "ArrayData")
+
+        try:
+            epics.ca.pend_io()
+        except:
+            pass
+
+    def monitorCB(self, epicsArgs, userArgs):
+        self.img = epicsArgs["pv_value"]
+        self.img.dtype = np.uint8
+        self.img.shape = self.shape
+        self.monitored = True
+
+    def setMonitor(self):
+        ndim = self.ndim.get()
+        dim0 = self.dim0.get()
+        dim1 = self.dim1.get()
+        dim2 = self.dim2.get()
+        if ndim == 3:
+            size = dim0 * dim1 * dim2
+            self.shape = (dim2, dim1, dim0)
+        else:
+            size = dim0 * dim1
+            self.shape = (dim1, dim0)
+        self.data.add_masked_array_event(None, size, None, self.monitorCB)
+
+    def clearMonitor(self):
+        self.data.clear_event()
+        self.monitored = False
+
+    def getShape(self):
+        ndim = self.ndim.get()
+        dim0 = self.dim0.get()
+        dim1 = self.dim1.get()
+        dim2 = self.dim2.get()
+
+        if ndim == 3:
+            shape = (dim2, dim1, dim0)
+        else:
+            shape = (dim1, dim0)
+        return shape
+
+    def get_image_center(self):
+        shape = self.getShape()
+        return (shape[0] / 2, shape[1] / 2)
+
+    def get_image(self, gray=True) -> np.ndarray:
+        if self.monitored:
+            return self.img
+
+        ndim = self.ndim.get()
+        dim0 = self.dim0.get()
+        dim1 = self.dim1.get()
+        dim2 = self.dim2.get()
+        if ndim == 3:
+            size = dim0 * dim1 * dim2
+        else:
+            size = dim0 * dim1
+        data = self.data.get(count=size)
+        data.dtype = np.uint8
+
+        if ndim == 3:
+            data.shape = (dim2, dim1, dim0)
+            # convert to grayscale
+            if gray:
+                data = (
+                    data[:, :, 0] * 299.0 / 1000.0
+                    + data[:, :, 1] * 587.0 / 1000.0
+                    + data[:, :, 2] * 114.0 / 1000.0
+                )
+        else:
+            data.shape = (dim1, dim0)
+
+        return data
+
+    def setup(self, gain: float = 15, expo: float = 0.025):
+        # self.color.put(0) # black and white
+        self.acquire.put(0, wait=True)
+        self.gain_mode.put(0)  # manual gain control
+        self.gain.put(gain)
+        self.expo_mode.put(0)  # manual exposure control
+        self.expo.put(expo)  # 20 hz acquisition
+        self.acquire.put(1)
+
+    def set_auto(self):
+        self.acquire.put(0, wait=True)
+        self.gain_mode.put(2)  # automatic gain control
+        self.expo_mode.put(2)  # automatic exposure control
+        self.acquire.put(1)
+
+    def restore(self, gain: float = 5, expo: float = 0.025):
+        # Put this values, so the camera faster resores proper values in Auto mode
+        self.acquire.put(0, wait=True)
+        self.gain.put(gain)
+        self.expo.put(expo)
+        self.gain_mode.put(2)  # automatic gain control
+        self.expo_mode.put(2)  # automatic exposure control
+        self.acquire.put(1)

src/mxlibs3/enum_pv.py

@@ -0,0 +1,126 @@
+from typing import cast
+
+from epics import PV
+
+from .mx_lib import ValueWaitTimeout, pv_wait
+
+
+class EnumPv(object):
+    def __init__(self, args):
+        mandatory_args = {"pv", "readback", "states"}
+        supplied_args = set(args.keys())
+        if not mandatory_args.issubset(supplied_args):
+            raise RuntimeError("Missing arguments: %s" % (", ".join(mandatory_args.difference(supplied_args))))
+
+        if type(args["states"]) is not dict:
+            raise RuntimeError("states must be a dictionary mapping put/readback for waiting purposes")
+
+        self.val = PV(args["pv"])
+        self.rbv = PV(args["readback"])
+
+        if "name" in args:
+            self.device_name = args["name"]
+        else:
+            self.device_name = self.val.pvname
+
+        if not hasattr(self.val, "enum_strs"):
+            raise RuntimeError("PV %s does not have enum_strs" % self.val.pvname)
+        if not hasattr(self.rbv, "enum_strs"):
+            raise RuntimeError("PV %s does not have enum_strs" % self.rbv.pvname)
+
+        self.__target = None
+
+        p = args.get("move_done_when")
+        if p is not None:
+            self.__wait_pv = PV(p[0])
+            self.__wait_target = p[1]
+        else:
+            self.__wait_pv = None
+            self.__wait_target = None
+
+        self.states = {}
+        for k, v in list(args["states"].items()):
+            self.states[k.upper()] = v.upper()
+
+        self.val_enums = [x.upper() for x in self.val.enum_strs]  # type: ignore
+        self.rbv_enums = [x.upper() for x in self.rbv.enum_strs]  # type: ignore
+
+        if "timeout" not in args:
+            self.__timeout = 60  # seconds
+        else:
+            self.__timeout = args["timeout"]
+
+    def __str__(self):
+        return f"<{self.device_name} at {self.position}>"
+
+    def __repr__(self):
+        s = (
+            f"<{self.device_name} at {self.position} an {self.__class__.__name__} "
+            f"instance at {hex(id(self))} positions = {self.positions}>"
+        )
+        return s
+
+    def has_position(self, position):
+        return position.upper() in self.states
+
+    @property
+    def positions(self):
+        return self.rbv_enums
+
+    def __position(self) -> str:
+        return self.rbv.enum_strs[self.rbv.get()]  # type: ignore
+
+    position = property(__position)
+
+    def position_is(self, position) -> bool:
+        if type(position) in (int, float):
+            return int(position) == self.get()
+        else:
+            return str(position).upper() == str(self.__position()).upper()
+
+    def get(self, req_type=None) -> str | int:
+        if req_type:
+            res = self.position
+        else:
+            res = cast(int, self.rbv.get())
+        return res
+
+    def put(self, value, wait=False) -> None:
+        self.val.put(value)
+        self.__target = self.rbv_enums.index(self.states[self.val_enums[value]])
+        if wait:
+            self.wait()
+
+    def move(self, position, wait=False) -> None:
+        if type(position) is not int:
+            if not self.has_position(position):
+                raise ValueError(f"Invalid position {position} for {self.device_name}")
+            target = self.val_enums.index(position.upper())
+        else:
+            target = position
+
+        self.put(target, wait=wait)
+
+    def equal_dbr_string(self, val) -> bool:
+        value = self.rbv.get(as_string=True)
+        return cast(str, value).lower() == val.lower()
+
+    def equal_dbr_int(self, val) -> bool:
+        return val == self.rbv.get()
+
+    def wait(self, target=None):
+        if target is None:
+            target = self.__target
+
+        if self.__wait_pv:
+            pv = self.__wait_pv
+        else:
+            pv = self.rbv
+
+        if target is None:
+            target = self.__wait_target
+
+        try:
+            pv_wait(pv, target, timeout=self.__timeout, verbose=True)
+        except ValueWaitTimeout:
+            raise ValueWaitTimeout(f"Timeout waiting for device to reach [{target}], currently at [{self.position}]")

src/mxlibs3/mx_lib.py

@@ -0,0 +1,251 @@
+import datetime
+import re
+import time
+from typing import Callable, Union
+
+from epics import PV, Motor, poll
+
+
+def timestamp():
+    """
+    Returns a fixed width string (15 characters) with a timestamp
+    in the format (24Hour:Minute:Second.Microsecond).
+
+    Example: 13:09:43.009508
+    """
+    x = datetime.datetime(1, 1, 1).now()
+    return "%4d-%02d-%02d %02d:%02d:%02d,%03d" % (
+        x.year,
+        x.month,
+        x.day,
+        x.hour,
+        x.minute,
+        x.second,
+        x.microsecond / 1000,
+    )
+
+
+def itoa(x, base=10):
+    is_negative = x < 0
+    if is_negative:
+        x = -x
+    digits = []
+    while x > 0:
+        x, last_digit = divmod(x, base)
+        digits.append("0123456789abcdefghijklmnopqrstuvwxyz"[last_digit])
+    if is_negative:
+        digits.append("-")
+    digits.reverse()
+    return "".join(digits)
+
+
+def wait_for_movement_to_finish(*motors):
+    """
+    Wait for all {motors} passed in argument to finish movement.
+
+    Requirements: epics.Motor -- each {motor} is an instance of this
+                  class.
+
+    The overall timeout is the maximum of (.RBV - .VAL)/.VELO for each
+    concerned {motor}.
+
+    Returns: nothing
+    """
+    # timeout = 1.5 * max([abs(motor.get_position(readback=True) -
+    #                         motor.get_position()) / motor.slew_speed
+    #                     for motor in motors])
+    #
+    poll(0.3)
+    longest = 0.0
+    for motor in motors:
+        time_to_target = motor.readback / motor.slew_speed
+        longest = time_to_target if time_to_target > longest else longest
+
+    timeout = time.time() + 1.5 * longest
+    done = False
+    while not done and time.time() < timeout:
+        done = all([m.done_moving for m in motors])
+
+    if time.time() > timeout:
+        print("TIMEOUT waiting for motors to be done moving; current motor positions:")
+        for m in motors:
+            print(f"...{m._prefix[:-1]} ({m.units})  VAL={m.drive:.1f}; RBV={m.readback:.1f}")
+        raise TimeoutError("timeout waiting for movement to finish.")
+
+
+class ValueWaitTimeout(Exception):
+    pass
+
+
+def pv_wait(pv, value, *, timeout=30.0, polling=0.2, tolerance=None, verbose=False):
+    if not (isinstance(pv, PV) or isinstance(pv, Motor)):
+        raise ValueError("wait what!? what pv?!")
+
+    if isinstance(pv, Motor):
+
+        def checker(m, target=None, tolerance=None):
+            if tolerance is None:
+                tolerance = m.get("RDBD")
+            if target is None:
+                target = m.drive
+
+            diff = abs(target - m.readback)
+            return m.done_moving and diff < tolerance
+
+        wait_motor_position(pv, checker, target=value, tolerance=tolerance, timeout=timeout, polling=polling)
+
+    elif is_epics_type(pv, "double") or is_epics_type(pv, "time_double"):
+        wait_float_condition(pv, value, tolerance=tolerance, timeout=timeout, polling=polling)
+
+    elif is_epics_type(pv, "enum") or is_epics_type(pv, "time_enum"):
+        wait_enum_condition(pv, value, timeout=timeout, polling=polling)
+
+    elif is_epics_type(pv, "string") or is_epics_type(pv, "time_string"):
+        wait_string_condition(pv, value, timeout=timeout, polling=polling)
+    else:
+        raise ValueError(f"pv_wait() don't know how to handle {pv.pvname} of type {pv.type}")
+
+
+def is_epics_type(pv: PV, pv_type: str) -> bool:
+    if isinstance(pv_type, type):
+        pv_type = pv_type.__name__
+
+    return pv_type == pv.type
+
+
+def wait_string_condition(pv, target: Union[str, re.Pattern], *, timeout=60.0, polling=0.1):
+    if not (isinstance(pv, PV) or "string" not in pv.type):
+        raise AttributeError("argument 'pv' must be an epics.PV of type string")
+
+    if not isinstance(target, re.Pattern):
+        target = re.compile(f"^{target}$", re.IGNORECASE)
+
+    tout = time.time() + timeout
+
+    while time.time() < tout:
+        if target.match(pv.char_value):
+            return
+        poll(polling)
+
+    if time.time() > tout:
+        raise TimeoutError(f"timeout waiting for string {pv.pvname} == {target}; actual value == {pv.char_value}")
+
+
+def wait_float_condition(pv: PV, value, *, timeout: float = 60.0, **kwargs):
+    """wait until an epics.PV of type double reaches value
+    pv: epics.PV
+        the PV enum on which you want to wait on
+
+    value: float
+        the target value
+
+    tolerance: float or None
+        the tolerance to accept when comparing values, if None (default) we
+        try to figure an appropriate value
+
+
+    timeout: double
+        a timeout in seconds
+
+    return: nothing
+
+    raises: TimeoutError if a timeout occurs
+    """
+    if not (isinstance(pv, PV) or "double" not in pv.type):
+        raise AttributeError("argument 'pv' must be an epics.PV of type double")
+
+    try:
+        value = float(value)
+    except ValueError:
+        raise AttributeError("argument 'value' must be a number")
+
+    tolerance = kwargs.get("tolerance", pow(10, -(pv.precision - 1)))  # type: ignore
+    polling = kwargs.get("polling", 0.1)
+
+    tout = time.time() + timeout
+    while time.time() < tout:
+        if abs(pv.value - value) < tolerance:
+            return
+        poll(polling)
+
+    if time.time() > tout:
+        raise TimeoutError(f"timeout waiting for enum {pv.pvname} == {value}")
+
+
+def wait_motor_position(motor: Motor, tester: Callable, *, timeout: float = 50.0, polling: float = 0.1, **kwargs):
+    """wait until an epics.Motor reaches a condition
+    motor: epics.Motor
+        the motor on which you want to wait on
+
+    tester: Callable
+        a Callable which is passed motor and **kwargs
+           => Callable(motor, **kwargs)
+
+    timeout: double
+        a timeout in seconds
+
+    **kwargs: extra arguments to tester
+
+    return: nothing
+
+    raises: TimeoutError if a timeout occurs
+    """
+    if not isinstance(tester, Callable):
+        raise RuntimeError("argument 'tester' must be a function")
+
+    try:
+        move_time = abs(motor.drive - motor.readback) / motor.speed
+    except Exception:
+        move_time = 1.0  # in case we're dealing with weird motor record
+
+    tout = move_time + time.time() + timeout
+
+    while time.time() < tout:
+        if tester(motor, **kwargs):
+            return
+        poll(polling)
+
+    if time.time() > tout:
+        raise TimeoutError(f"timeout waiting for a condition on {motor} {motor.drive} != {motor.readback}")
+
+
+def wait_enum_condition(pv: PV, value: Union[str, int, re.Pattern], *, timeout: float = 60.0, polling=0.1, **kwargs):
+    """wait until an epics.PV enum reaches value
+    pv: epics.PV
+        the PV enum on which you want to wait on
+
+    value: int or str
+        if int check against index of current value
+
+    timeout: double
+        a timeout in seconds
+
+    return: nothing
+
+    raises: TimeoutError if a timeout occurs
+    """
+    if not (isinstance(pv, PV) or "enum" != pv.type[-4:].lower()):
+        raise AttributeError("argument 'pv' must be an epics.PV of type enum")
+
+    if not (isinstance(value, str) or isinstance(value, int) or isinstance(value, re.Pattern)):
+        raise AttributeError("argument 'value' must be either an int, str, or re.Pattern")
+
+    if type(value) is int:
+        tester = lambda pv: value == pv.get()  # noqa: E731
+    elif type(value) is str:
+        tester = lambda pv: str(value) == pv.get(as_string=True).lower()  # noqa: E731
+        value = str(value).lower()  # it's already a str :-/
+    elif isinstance(value, re.Pattern):
+        tester = lambda pv: value.match(pv.get(as_string=True))  # noqa: E731
+    else:
+        raise AttributeError("argument 'value' must be either an int, str, or re.Pattern")
+
+    tout = time.time() + timeout
+
+    while time.time() < tout:
+        if tester(pv):
+            return
+        poll(polling)
+
+    if time.time() > tout:
+        raise TimeoutError(f"timeout waiting for enum {pv.pvname} == {value}")

src/mxlibs3/my_motor.py

@@ -0,0 +1,53 @@
+import time
+
+from epics import Motor
+
+
+class MyMotor(Motor):
+    def __init__(self, name):
+        super().__init__(name.upper())
+
+    @property
+    def speed(self):
+        return self.slew_speed
+
+    @speed.setter
+    def speed(self, speed):
+        self.slew_speed = speed
+
+    @property
+    def pvname(self):
+        return self._prefix[:-1]
+
+    def __readback(self):
+        return self.readback
+
+    position = property(__readback)
+
+    def __value(self):
+        return self.drive
+
+    value = property(__value)
+
+
+class BogusMotor:
+    def __init__(self, name, **kwargs):
+        self.name = name
+        self.readback = kwargs.get("readback", 0.0)
+        self.target = kwargs.get("target", 0.0)
+        self.velo = kwargs.get("velo", 10.0)
+
+    @property
+    def position(self):
+        return self.readback
+
+    def move(self, target, wait=False):
+        if wait:
+            time.sleep(abs(target - self.readback) / self.velo)
+        self.readback = target
+
+    def refresh(self):
+        pass
+
+    def get_position(self, *args, **kwargs):
+        return self.readback

src/mxlibs3/non_standard.py

@@ -0,0 +1,281 @@
+import time
+from typing import Callable, cast
+
+from epics import PV, poll
+from epics.ca import pend_io
+
+
+def is_number(x):
+    return isinstance(x, (int, float))
+
+
+class NonStandard(object):
+    def __init__(self, **kwargs):
+        mandatory = ["name", "setpv", "getpv"]
+        if set(mandatory) != set(kwargs.keys()) & set(mandatory):
+            raise RuntimeError(
+                "missing mandatory argument(s): {}".format(str(list(set(mandatory) - set(kwargs.keys()))))
+            )
+        self.name = kwargs.get("name")
+        self.__target_pos: float | int | str = 0
+        self.device_name = self.name
+        self.__setpv = PV(kwargs["setpv"])
+        self.__getpv = PV(kwargs["getpv"])
+        self.__timeout_margin = kwargs.get("timeout", 2.0)  # add this safety margin to timeout calc from speed
+        self.__timeout = kwargs.get("timeout", 60.0)
+        self.__format = kwargs.get("format")
+
+        p = kwargs.get("speed")
+        if p is not None:
+            if is_number(p):
+                self.__speed_set, self.__speed_get = p, p
+            elif type(p) in [tuple, list] and len(p) == 2:
+                self.__speed_set = PV(p[0])
+                self.__speed_get = PV(p[1])
+            else:
+                raise Exception(f"{self.name} => don't know how to parse speed parameter")
+        else:
+            self.__speed_set, self.__speed_get = None, None
+        #
+        # move_done_when a tuple, first item the pv to read
+        # second item, the int type indicating that it is DONE
+        #
+        p = kwargs.get("move_done_when")
+        if p is not None:
+            self.__move_flag = PV(p[0])
+            self.__move_flag_done = p[1]
+        else:
+            self.__move_flag = None
+
+        #
+        # stoppv a tuple, first item the pv to read
+        # second item, the int type to put to stop the motor
+        #
+        p = kwargs.get("stoppv")
+        if p is not None:
+            self.__stoppv = PV(p[0])
+            self.__stoppv_done = p[1]
+        else:
+            self.__stoppv = None
+
+        #
+        # still_moving_when a tuple, first item the pv to read
+        # second item, the int type indicating that it is still moving.
+        #
+        p = kwargs.get("still_moving_when")
+        if p is not None:
+            self.__still_moving_flag = PV(p[0])
+            self.__still_moving_flag_value = p[1]
+        else:
+            self.__still_moving_flag = None
+
+        p = kwargs.get("predefs")
+        if p:
+            self.__predefs = p
+            self.positions = list(p.keys())
+        else:
+            self.__predefs = {}
+            self.positions = []
+
+        pend_io(5.0)
+
+        typ = self.__ret_type = self.__getpv.type
+
+        match str(typ):
+            case "double" | "time_double":
+                self.still_moving = self.__moving_float
+            case "int" | "time_int":
+                self.still_moving = self.__moving_int
+            case "string" | "time_string":
+                self.still_moving = self.__moving_str
+            case _:
+                self.still_moving = self.__moving_unknown
+
+        tolerance = kwargs.get("tolerance", None)
+        if tolerance is not None and tolerance > 0:
+            self.tolerance = tolerance
+        elif tolerance is not None and tolerance < 0:
+            self.still_moving = self.__moving_unknown
+
+        if self.__move_flag:
+            self.still_moving = self.is_move_not_done
+        elif self.__still_moving_flag:
+            self.still_moving = self.still_moving_when
+
+    def __str__(self):
+        if self.positions:
+            return "<{} at {}({}) Pre-defined positions {}>".format(
+                self.name, self.__position(predef=True), self.position, self.positions
+            )
+        else:
+            return "<{} at {}>".format(self.name, self.__position())
+
+    def __repr__(self):
+        if self.positions:
+            return "<{} instance at {}: {} at {}({}) Pre-defined positions {}>".format(
+                self.__class__.__name__,
+                hex(id(self)),
+                self.name,
+                self.__position(predef=True),
+                self.position,
+                self.positions,
+            )
+        else:
+            return "<{} instance at {}: {} at {}>".format(
+                self.__class__.__name__, hex(id(self)), self.name, self.__position()
+            )
+
+    def still_moving_when(self):
+        return self.__still_moving_flag_value == self.__still_moving_flag.get()  # type: ignore
+
+    def is_moving(self):
+        return not self.is_move_done()
+
+    def is_move_done(self):
+        return self.__move_flag_done == self.__move_flag.get()  # type: ignore
+
+    def is_move_not_done(self):
+        return not self.is_move_done()
+
+    def __moving_float(self):
+        if type(self.__target_pos) is not float:
+            raise RuntimeError("target position is not an float")
+        current_value = cast(float, self.__getpv.get())
+        return self.tolerance < abs(cast(float, self.__target_pos) - current_value)
+
+    def __moving_int(self):
+        if type(self.__target_pos) is not int:
+            raise RuntimeError("target position is not an integer")
+        return self.__target_pos != self.__getpv.get()
+
+    def __moving_str(self):
+        if type(self.__target_pos) is not str:
+            raise RuntimeError("target position is not a string")
+        current_value = cast(str, self.__getpv.get())
+        return self.__target_pos.upper() != current_value.upper()
+
+    def __moving_unknown(self):
+        return False
+
+    def __position(self, predef: bool = False):
+        """returns either a predefined position or a readback
+
+        :rtype: Union[str, float, int]
+        """
+        cur = self.readback
+        if not predef:
+            return cur
+        # loop over predefined positions, if one found and
+        # readback matches it's value return "position"
+        for pos_label, val_or_cbk in self.__predefs.items():
+            if self.__ret_type is bytes:
+                tst = val_or_cbk == cur
+            elif type(val_or_cbk) is tuple:
+                # a tuple (Callable, arg-to-Callable)
+                val = val_or_cbk[0](*val_or_cbk[1])
+                tst = 0.1 > abs(val - cur)  # FIXME sooo bad
+            else:
+                tst = 0.1 > abs(val_or_cbk - cur)
+            if tst:
+                return pos_label
+        return "unknown"
+
+    def __readback(self):
+        return self.__getpv.get()
+
+    readback = property(__readback)
+    position = property(__position)
+
+    def has_position(self, position: str) -> bool:
+        return position.lower() in [p.lower() for p in self.positions]
+
+    def position_is(self, position, tolerance=0.1):
+        pos = self.__predefs.get(position, position)
+        cur = self.position
+
+        print(f"position_is: {pos} == {cur}")
+
+        if type(pos) is tuple:
+            pos = pos[0](*pos[1])
+
+        print(f"position_is: {pos} == {cur}")
+
+        if self.__ret_type is bytes:
+            print("checking bytes")
+            return pos == cur
+        else:
+            print(f"checking numbers: abs({pos} - {cur}) < {tolerance}")
+            return abs(pos - cur) < tolerance
+
+    def __value(self):
+        val = self.__getpv.get()
+        if self.__format is not None:
+            val = self.__format % (val,)
+        return val
+
+    value = property(__value)
+
+    def __raw_value(self):
+        return self.__getpv.get()
+
+    raw_value = property(__raw_value)
+
+    def __char_value(self):
+        return str(self.value)
+
+    char_value = property(__char_value)
+
+    def set_speed(self, speed):
+        if isinstance(self.__speed_set, PV):
+            self.__speed_set.put(speed)
+        elif is_number(self.__speed_set):
+            self.__speed_set = speed
+            self.__speed_get = speed
+
+    def get_speed(self):
+        if is_number(self.__speed_get):
+            return self.__speed_get
+        elif isinstance(self.__speed_get, PV):
+            return self.__speed_get.value
+        else:
+            return None
+
+    def get_timeout(self):
+        if not self.__speed_get:
+            timeout = self.__timeout
+        else:
+            speed = self.get_speed()
+            curp = self.position
+            timeout = self.__timeout_margin + (abs(curp - self.__target_pos) / speed)
+        return timeout
+
+    def wait(self):
+        pos = self.__target_pos
+        timeout = time.time() + self.get_timeout()
+
+        while time.time() < timeout and self.still_moving():
+            poll(0.1)
+        if time.time() > timeout:
+            msg = "timeout occurred when moving %s to %s" % (self.name, pos)
+            print(msg)
+            raise RuntimeWarning(msg)
+
+    def move(self, pos, relative=False, wait=False):
+        pos = self.__predefs.get(pos, pos)
+        if type(pos) is tuple:
+            pos = pos[0](*pos[1])
+
+        if relative:
+            pos += self.__getpv.get()
+
+        self.__target_pos = pos
+        self.__setpv.put(pos)
+        if wait:
+            time.sleep(0.1)
+            self.wait()
+
+    def stop(self):
+        if self.__stoppv:
+            self.__stoppv.put(self.__stoppv_done)
+        else:
+            raise RuntimeError("stop pv not configured for this motor")

src/mxlibs3/pshell_client.py

@@ -0,0 +1,292 @@
+import json
+from urllib.parse import quote
+
+import requests
+
+
+class PShellClient:
+    def __init__(self, url):
+        self.url = url
+        self.sse_event_loop_thread = None
+        self.subscribed_events = None
+        self.event_callback = None
+
+    def _get_response(self, response, is_json=True):
+        if response.status_code != 200:
+            raise Exception(response.text)
+        return json.loads(response.text) if is_json else response.text
+
+    def _get_binary_response(self, response):
+        if response.status_code != 200:
+            raise Exception(response.text)
+        return response.raw.read()
+
+    def get_version(self):
+        """Return application version.
+
+        Args:
+
+        Returns:
+            String with application version.
+
+        """
+        return self._get_response(requests.get(url=self.url + "/version"), False)
+
+    def get_config(self):
+        """Return application configuration.
+
+        Args:
+
+        Returns:
+            Dictionary.
+        """
+        return self._get_response(requests.get(url=self.url + "/config"))
+
+    def get_state(self):
+        """Return application state.
+
+        Args:
+
+        Returns:
+            String: Invalid, Initializing,Ready, Paused, Busy, Disabled, Closing, Fault, Offline
+        """
+        return self._get_response(requests.get(url=self.url + "/state"))
+
+    def get_logs(self):
+        """Return application logs.
+
+        Args:
+
+        Returns:
+            List of logs.
+            Format of each log: [date, time, origin, level, description]
+
+        """
+        return self._get_response(requests.get(url=self.url + "/logs"))
+
+    def get_history(self, index):
+        """Access console command history.
+
+        Args:
+            index(int): Index of history entry (0 is the most recent)
+
+        Returns:
+            History entry
+
+        """
+        return self._get_response(requests.get(url=self.url + "/history/" + str(index)), False)
+
+    def get_script(self, path):
+        """Return script.
+
+        Args:
+            path(str): Script path (absolute or relative to script folder)
+
+        Returns:
+            String with file contents.
+
+        """
+        return self._get_response(requests.get(url=self.url + "/script/" + str(path)), False)
+
+    def get_devices(self):
+        """Return global devices.
+
+        Args:
+
+        Returns:
+            List of devices.
+            Format of each device record: [name, type, state, value, age]
+
+        """
+        return self._get_response(requests.get(url=self.url + "/devices"))
+
+    def abort(self, command_id=None):
+        """Abort execution of command
+
+        Args:
+            command_id(optional, int): id of the command to be aborted.
+                                       if None (default), aborts the foreground execution.
+
+        Returns:
+
+        """
+        if command_id is None:
+            requests.get(url=self.url + "/abort")
+        else:
+            return requests.get(url=self.url + "/abort/" + str(command_id))
+
+    def reinit(self):
+        """Reinitialize the software.
+
+        Args:
+
+        Returns:
+
+        """
+        requests.get(url=self.url + "/reinit")
+
+    def stop(self):
+        """Stop all devices implementing the 'Stoppable' interface.
+
+        Args:
+
+        Returns:
+
+        """
+        requests.get(url=self.url + "/stop")
+
+    def update(self):
+        """Update all global devices.
+
+        Args:
+
+        Returns:
+
+        """
+        requests.get(url=self.url + "/update")
+
+    def eval(self, statement):
+        """Evaluates a statement in the interpreter.
+           If the statement finishes by '&', it is executed in background.
+           Otherwise statement is executed in foreground (exclusive).
+
+        Args:
+            statement(str): input statement
+
+        Returns:
+            String containing the console return.
+            If an exception is produces in the interpretor, it is re-thrown here.
+        """
+        statement = quote(statement)
+        return self._get_response(requests.get(url=self.url + "/eval/" + statement), False)
+
+    def run(self, script, pars=None, background=False):
+        """Executes script in the interpreter.
+
+        Args:
+            script(str): name of the script (absolute or relative to the script base folder). Extension may be omitted.
+            pars(optional, list or dict): if a list is given, it sets sys.argv for the script.
+                                          If a dict is given, it sets global variable for the script.
+            background(optional, bool): if True script is executed in background.
+
+        Returns:
+            Return value of the script.
+            If an exception is produces in the interpretor, it is re-thrown here.
+        """
+        return self._get_response(
+            requests.put(
+                url=self.url + "/run", json={"script": script, "pars": pars, "background": background, "async": False}
+            )
+        )
+
+    def start_eval(self, statement):
+        """Starts evaluation of a statement in the interpreter.
+           If the statement finishes by '&', it is executed in background.
+           Otherwise statement is executed in foreground (exclusive).
+
+        Args:
+            statement(str): input statement
+
+        Returns:
+            Command id (int), which is used to retrieve command execution status/result (get_result).
+        """
+        statement = quote(statement)
+        return int(self._get_response(requests.get(url=self.url + "/evalAsync/" + statement), False))
+
+    def start_run(self, script, pars=None, background=False):
+        """Starts execution of a script in the interpreter.
+
+        Args:
+            script(str): name of the script (absolute or relative to the script base folder). Extension may be omitted.
+            pars(optional, list or dict): if a list is given, it sets sys.argv for the script.
+                                          If a dict is given, it sets global variable for the script.
+            background(optional, bool): if True script is executed in background.
+
+        Returns:
+            Command id (int), which is used to retrieve command execution status/result (get_result).
+        """
+        return int(
+            self._get_response(
+                requests.put(
+                    url=self.url + "/run",
+                    json={"script": script, "pars": pars, "background": background, "async": True},
+                )
+            )
+        )
+
+    def get_result(self, command_id=-1):
+        """Gets status/result of a command executed asynchronously (start_eval and start_run).
+
+        Args:
+            command_id(optional, int): command id. If equals to -1 (default) return status/result of the foreground task.
+
+        Returns:
+            Dictionary with the fields: 'id' (int): command id
+                                        'status' (str): unlaunched, invalid, removed, running, aborted, failed or completed.
+                                        'exception' (str): if status equals 'failed', holds exception string.
+                                        'return' (obj): if status equals 'completed', holds return value of script (start_run)
+                                                        or console return (start_eval)
+        """
+        return self._get_response(requests.get(url=self.url + "/result/" + str(command_id)))
+
+    def help(self, input="<builtins>"):
+        """Returns help or auto-completion strings.
+
+        Args:
+            input(optional, str): - ":" for control commands
+                                  - "<builtins>" for builtin functions
+                                  - "devices" for device names
+                                  - builtin function name for function help
+                                  - else contains entry for auto-completion
+
+        Returns:
+            List
+
+        """
+        return self._get_response(requests.get(url=self.url + "/autocompletion/" + input))
+
+    def get_contents(self, path=None):
+        """Returns contents of data path.
+
+        Args:
+            path(optional, str): Path to data relative to data home path.
+                                 - Folder
+                                 - File
+                                 - File (data root) |  internal path
+                                 - internal path (on currently open data root)
+
+        Returns:
+            List of contents
+
+        """
+        return self._get_response(
+            requests.get(url=self.url + "/contents" + ("" if path is None else ("/" + path))), False
+        )
+
+    def get_data(self, path, type="txt"):
+        """Returns data on a given path.
+
+        Args:
+            path(str): Path to data relative to data home path.
+                                 - File (data root) |  internal path
+                                 - internal path (on currently open data root)
+            type(optional, str): txt, "json", "bin", "bs"
+
+        Returns:
+            Data accordind to selected format/.
+
+        """
+        if type == "json":
+            return self._get_response(requests.get(url=self.url + "/data-json/" + path), True)
+        elif type == "bin":
+            return self._get_binary_response(requests.get(url=self.url + "/data-bin/" + path, stream=True))
+
+        return self._get_response(requests.get(url=self.url + "/data" + ("" if path is None else ("/" + path))), False)
+
+    def print_logs(self):
+        for log_line in self.get_logs():
+            print("%s %s %-20s %-8s %s" % tuple(log_line))
+
+    def print_devices(self):
+        for log_line in self.get_devices():
+            print("%-16s %-32s %-10s %-32s %s" % tuple(log_line))

src/mxlibs3/smargon.py

@@ -0,0 +1,132 @@
+from enum import Enum
+from time import sleep, time
+
+import requests
+
+from aaredaq.beamline import MXBeamline
+from aaredaqlib.coordinate import SmargonCoordinate, Coordinate
+
+
+class SmargonMode(Enum):
+    UNINITIALIZED = 0
+    INITIALIZING = 1
+    READY = 2
+    ERROR = 99
+
+
+class Smargon(object):
+    SMARGON_HOME = SmargonCoordinate(
+        sh_mm=Coordinate(x=0, y=0, z=18), phi_deg=0, chi_deg=0
+    )
+
+    def __init__(self, bl: MXBeamline):
+        if bl == MXBeamline.X06DA:
+            self.__simulated = False
+            self.__base = "http://x06da-smargopolo.psi.ch:3000"
+        elif bl == MXBeamline.SIMULATED:
+            self.__simulated = True
+            self.__pos = self.SMARGON_HOME
+        else:
+            raise Exception("unknown beamline")
+
+    def gonget(self, thing: str) -> dict:
+        """issue a GET for some API component on the smargopolo server"""
+        cmd = f"{self.__base}/{thing}"
+        r = requests.get(cmd)
+        if not r.ok:
+            raise Exception(
+                f"error getting {thing}; server returned {r.status_code} => {r.reason}"
+            )
+        return r.json()
+
+    def gonput(self, thing: str):
+        cmd = f"{self.__base}/{thing}"
+        r = requests.put(cmd)
+        if not r.ok:
+            raise Exception(
+                f"error putting {thing}; server returned {r.status_code} => {r.reason}"
+            )
+
+    def move_home(self, wait=False) -> None:
+        self.target = self.SMARGON_HOME
+        if wait:
+            self.wait()
+
+    @property
+    def mode(self) -> SmargonMode:
+        return SmargonMode.INITIALIZING
+
+    @mode.setter
+    def mode(self, mode: SmargonMode):
+        if self.__simulated:
+            return
+        self.gonput(f"mode?mode={mode}")
+
+    def initialize(self):
+        self.mode = SmargonMode.UNINITIALIZED
+        sleep(0.1)
+        self.mode = SmargonMode.INITIALIZING
+
+    def enable_correction(self):
+        if self.__simulated:
+            return
+
+        self.gonput("corr_type?corr_type=1")
+
+    def disable_correction(self):
+        if self.__simulated:
+            return
+
+        self.gonput("corr_type?corr_type=0")
+
+    @property
+    def readback(self) -> SmargonCoordinate:
+        if self.__simulated:
+            return self.__pos
+
+        scs = self.gonget("readbackSCS")
+        return SmargonCoordinate(
+            sh_mm=Coordinate(x=scs["SHX"], y=scs["SHY"], z=scs["SHZ"]),
+            phi_deg=scs["PHI"],
+            chi_deg=scs["CHI"],
+        )
+
+    @property
+    def target(self) -> SmargonCoordinate:
+        if self.__simulated:
+            return self.__pos
+
+        scs = self.gonget("targetSCS")
+        return SmargonCoordinate(
+            sh_mm=Coordinate(x=scs["SHX"], y=scs["SHY"], z=scs["SHZ"]),
+            phi_deg=scs["PHI"],
+            chi_deg=scs["CHI"],
+        )
+
+    @target.setter
+    def target(self, coord: SmargonCoordinate):
+        if self.__simulated:
+            self.__pos = coord
+            return
+
+        target_string = ""
+        if coord.sh_mm is not None:
+            target_string += "&SHX={:.5f}&SHY={:.5f}&SHZ={:.5f}".format(
+                coord.sh_mm.x, coord.sh_mm.y, coord.sh_mm.z
+            )
+        if coord.chi_deg is not None:
+            target_string += "&CHI={:.5f}".format(coord.chi_deg)
+        if coord.phi_deg is not None:
+            target_string += "&PHI={:.5f}".format(coord.phi_deg)
+        if target_string:
+            self.gonput(f"targetSCS?{target_string}")
+
+    def wait(self, timeout=60.0, tol=0.01, poll_time=0.01):
+        target = self.target
+        timeout = timeout + time()
+        while time() < timeout:
+            if target.eq(self.readback, tol):
+                break
+            if time() > timeout:
+                raise TimeoutError("Timed out waiting for Smargon to reach target")
+            sleep(poll_time)

src/mxlibs3/tell_client.py

@@ -0,0 +1,593 @@
+import json
+import logging
+import random
+import re
+import time
+from typing import List
+from urllib.parse import urlparse
+
+import requests
+
+from aaredaq.beamline import MXBeamline  # noqa: F401
+from aaredaq.models import (
+    PuckLoadedInfo,
+    SampleShortInfo,
+    DewarAddress,
+    SampleDewarAddress,
+    PuckInfo,
+)
+from mxlibs3.pshell_client import PShellClient
+
+logger = logging.getLogger(__name__)
+
+
+class TellMountFailedException(Exception):
+    pass
+
+
+class TellCommandWhileBusyException(Exception):
+    pass
+
+
+VALID_DEWAR_POSITIONS = [f"{p}{n}" for n in "12345" for p in "ABCDEF"]
+
+POSITION_PARK = "pPark"
+POSITION_COLD = "pCold"
+POSITION_AUX = "pAux"
+POSITION_DEWAR = "pDewar"
+POSITION_HOME = "pHome"
+POSITION_HEATER = "pHeatB"
+
+
+class TellClient:
+    def __init__(self, bl: MXBeamline):
+        if bl == MXBeamline.SIMULATED:
+            self.__simulation = True
+        elif bl == MXBeamline.X06DA:
+            self.__simulation = False
+            self.__url = "http://x06da-tell.psi.ch:22222"
+
+        if self.__simulation:
+            print("TELL P-Shell in SIMULATION mode")
+        else:
+            print(f"Connecting TELL p-shell service at {self.__url} ...", end="")
+            hostname = urlparse(self.__url).hostname
+            try:
+                requests.get(f"{self.__url}/history/0", timeout=1.0)
+            except ConnectionError:
+                print(f"...connection to {hostname} failed")
+                raise
+            except requests.ReadTimeout:
+                print(f"...PShell service {hostname} is down")
+                raise
+            self.pshell = PShellClient(self.__url)
+        self._simulated_samples_info = {}
+        self._simulated_detected_pucks = []
+        self._simulated_mounted_sample = ""
+        self._simulated_current = 30.0
+        self._simulated_suppress = True
+        self._simulated_state = "Ready"
+        self._simulated_offset = 0.0
+        self._aborted = False
+        self.state = self.get_state()
+        self.debug = False
+        self._last_cmd_id = -1
+
+    def smc_get_current(self) -> float:
+        if self.__simulation:
+            return 30.0
+        return json.loads(self.pshell.eval("smart_magnet.get_current()&"))
+
+    def smc_set_current(self, current: float):
+        if self.__simulation:
+            return
+        self.pshell.eval(f"smart_magnet.set_current({current:.1f})&")
+
+    def smc_get_suppress(self) -> bool:
+        if self.__simulation:
+            return self._simulated_suppress
+        return json.loads(self.pshell.eval("smart_magnet.get_supress()&"))
+
+    def smc_set_suppress(self, state: bool):
+        if self.__simulation:
+            self._simulated_suppress = state
+            return
+        self.pshell.eval(f"smart_magnet.set_supress({state})&")
+
+    def check_smc(self):
+        self.pshell.eval("smart_magnet.set_supress(False)&")
+        self.pshell.eval("smart_magnet.set_resting_current()&")
+        self.pshell.eval("smart_magnet.check_mounted(idle_time=1.0, timeout=1.0")
+
+    def magnet_blower(self, state):
+        """turn on/off blower to the magnet for de-icying purposes
+
+        state = false (blower is off)
+        state = true (blower is on)
+        """
+        if X06DA:
+            self.pshell.eval(f"set_pin_cleaner({state})&")
+        else:
+            cmd = "true" if state else "false"
+            self.pshell.eval(f'robot.evaluate("doFOut1={cmd}")&')
+
+    def get_state(self):
+        if self.__simulation:
+            return self._simulated_state
+        self.state = self.pshell.get_state()
+        return self.state
+
+    def get_result(self, command_id=-1):
+        if self.__simulation:
+            return {
+                "id": self._last_cmd_id,
+                "status": "completed",
+                "exception": "",
+                "return": (True, "PINCODE6546"),
+            }
+        return self.pshell.get_result(command_id)
+
+    def wait_ready(self):
+        if self.__simulation:
+            time.sleep(3.0)
+            self._simulated_state = "Ready"
+            return
+        count = 0
+        # Monitors event but polls every second just n case an event is missed
+        while True:
+            if self.state != "Busy":
+                break
+            time.sleep(0.01)
+            count = count + 1
+            if count >= 100:
+                count = 0
+                self.get_state()
+        if self.state != "Ready":
+            raise Exception("Invalid state: " + str(self.state))
+
+    def set_in_mount_position(self, value):
+        if self.__simulation:
+            return
+        self.pshell.eval("in_mount_position = " + str(value) + "&")
+
+    def is_in_mount_position(self):
+        if self.__simulation:
+            return True
+        return self.pshell.eval("in_mount_position&").lower() == "true"
+
+    def set_simulated_mounted_sample(self, info):
+        self._simulated_mounted_sample = info
+
+    def set_simulated_samples_info(self, samples):
+        self._simulated_samples_info = samples
+
+    def set_simulated_detected_pucks(self, pucks):
+        self._simulated_detected_pucks = pucks
+
+    def get_samples_info(self) -> List[SampleShortInfo]:
+        if self.__simulation:
+            j = self._simulated_samples_info  # FIXME
+        else:
+            j = json.loads(self.pshell.eval("get_samples_info()&"))
+
+        output: List[SampleShortInfo] = []
+        for i in j:
+            if len(i["puckAddress"]) == 2:
+                dewar_location = DewarAddress(
+                    segment=i["puckAddress"][0], pos=i["puckAddress"][1]
+                )
+            else:
+                dewar_location = None
+
+            output.append(
+                SampleShortInfo(
+                    puck_name=i["puckBarcode"],
+                    dewar_name=i["dewarName"],
+                    sample_name=i["sampleName"],
+                    pin=i["samplePosition"],
+                    user=i["userName"],
+                    location=dewar_location,
+                )
+            )
+        return output
+
+    def set_samples_info(self, info: List[PuckInfo]):
+        if self.__simulation:
+            return
+
+        j = []
+        for x in info:
+            if x.location is None:
+                puck_address = ""
+            else:
+                puck_address = "{:1s}{:1d}".format(x.location.segment, x.location.pos)
+
+            j.append(
+                {
+                    "userName": x.user,
+                    "dewarName": x.dewar_name,
+                    "puckName": x.puck_name,
+                    "puckType": "Unipuck",
+                    "puckAddress": puck_address,
+                    "puckBarcode": x.puck_name,
+                    "sampleBarcode": "",
+                    "sampleMountCount": 0,
+                    "sampleName": "",
+                    "samplePosition": 1,
+                    "sampleStatus": "",
+                }
+            )
+
+        self.pshell.run("data/set_samples_info", pars=[json.dumps(j)], background=True)
+        # self.pshell.eval("set_samples_info(" + json.dumps(info) + ")&")
+
+    def start_cmd(self, cmd, *argv):
+        if self.__simulation:
+            return 666  # FIXME should be different?
+        cmd = cmd + "("
+        for a in argv:
+            cmd = cmd + (("'" + a + "'") if type(a) is str else str(a)) + ", "
+        cmd = cmd + ")"
+        ret = self.pshell.start_eval(cmd)
+        self.get_state()
+        return ret
+
+    def wait_cmd(self, cmd):
+        if self.__simulation:
+            return {True, "BARCODE_BIGUS"}
+        self.wait_ready()
+        result = self.get_result(cmd)
+        # print (result)
+        if result["exception"] is not None:
+            raise Exception(result["exception"])
+        return result["return"]
+
+    def is_cmd_completed(self, cmd):
+        if self.__simulation:
+            return True
+        return self.get_result(cmd)["status"] != "running"
+
+    def wait_mount_complete(self, timeout: float = 360):
+        if self.__simulation:
+            time.sleep(1.0)
+            return
+        timeisup = timeout + time.time()
+        while time.time() < timeisup:
+            if not self.is_busy():
+                break
+            time.sleep(0.01)
+
+    def check_command_ok(self, timeout: float = 360.0, msg: str = ""):
+        self.wait_mount_complete(timeout)
+        result = self.get_result(self._last_cmd_id)
+        if "completed" != result["status"]:
+            raise TellMountFailedException(f"{msg} {result}")
+
+    def estimate_mounting_time(self, segment) -> int:
+        try:
+            current_mounted = self.get_mounted_sample()
+            gripper_in_cold = self.is_in_cold()
+
+            if current_mounted is None:
+                unmount_needs_drying = 0  # might not have anything
+                unmount_needs_cooling = 0
+            else:
+                segment_in_cold = current_mounted.puck.segment in "ABCDEF"
+                unmount_needs_drying = int(gripper_in_cold and not segment_in_cold)
+                unmount_needs_cooling = int(not gripper_in_cold and segment_in_cold)
+
+            mount_needs_cooling = int(segment in "ABCDEF" and not gripper_in_cold)
+            mount_needs_drying = int(segment not in "ABCDEF" and gripper_in_cold)
+
+            needs_cooling = mount_needs_cooling + unmount_needs_cooling
+            needs_drying = mount_needs_drying + unmount_needs_drying
+            return needs_cooling * 30 + needs_drying * 120
+        except:
+            return 0
+
+    def mount(
+        self,
+        address: SampleDewarAddress,
+        force: bool = False,  # kept for future
+        read_dm: bool = False,  # read data matrix
+        auto_unmount: bool = False,  # single command, if False it will raise exception
+        wait: bool = False,  # blocking operation
+        timeout: float = 60.0,
+    ):
+        print(address)
+        SampleDewarAddress.model_validate(address)
+
+        segment = address.puck.segment
+        puck = address.puck.pos
+        sample = address.pin
+
+        if self.__simulation:
+            self._last_cmd_id = random.randint(1000, 9999)
+            print(
+                f"simulated mount({segment}, {puck}, {sample}) -> {self._last_cmd_id}"
+            )
+            self._simulated_mounted_sample = f"{segment}{puck}{sample}"
+            if random.random() < 0.1:
+                print("simulated failed mount")
+                raise TellMountFailedException(f"mount failed for {segment}{puck}")
+
+            return self._last_cmd_id
+
+        if self.is_busy():
+            raise TellCommandWhileBusyException("mount received while robot is busy")
+
+        self._last_cmd_id = self.start_cmd(
+            "mount", segment, puck, sample, force, read_dm, auto_unmount
+        )
+
+        wait_timeout = timeout + self.estimate_mounting_time(segment)
+
+        if wait:
+            self.check_command_ok(
+                timeout=wait_timeout, msg=f"Mount failed for {segment}{puck}-{sample}: "
+            )
+
+        return self._last_cmd_id
+
+    def unmount(self, force=False, wait=False, timeout=360.0):
+        # Force has a meaning, will unmount even if smart magnet is not detecting sample
+        if self.__simulation:
+            print("simulated unmount")
+            return
+
+        if self.is_busy():
+            raise TellCommandWhileBusyException("mount received while robot is busy")
+
+        self._last_cmd_id = self.start_cmd("unmount", None, None, None, force)
+
+        if wait:
+            self.check_command_ok(timeout=timeout, msg="Unmount failed: ")
+
+        return self._last_cmd_id
+
+    def scan_pin(self, segment, puck, sample, force=False):
+        return self.start_cmd("scan_pin", segment, puck, sample, force)
+
+    def scan_puck(self, segment, puck, force=False):
+        return self.start_cmd("scan_puck", segment, puck, force)
+
+    def dry(self, heat_time=None, speed=None, wait_cold=None, wait=False):
+        if self.__simulation:
+            time.sleep(5.0)
+        self._last_cmd_id = self.start_cmd("dry", heat_time, speed, wait_cold)
+        if wait:
+            self.check_command_ok(timeout=360.0, msg=f"Dry failed")
+
+    def move_park(self, wait=False):
+        if self.__simulation:
+            return
+        self._last_cmd_id = self.start_cmd("move_park")
+
+        if wait:
+            self.check_command_ok(timeout=360.0, msg=f"Move to park failed")
+
+    def move_cold(self, reset_timestamp=False):
+        if self.__simulation:
+            return
+        return self.start_cmd("move_cold", reset_timestamp)
+
+    def trash(self):
+        if self.__simulation:
+            return
+        return self.start_cmd("trash_sample")
+
+    def abort_cmd(self):
+        if self.__simulation:
+            print("simulated abort")
+            return
+        self.pshell.abort()
+        self.pshell.eval("robot.stop_task()&")
+
+    def set_gonio_mount_position(self, homing=False):
+        if self.__simulation:
+            return
+        if homing:
+            self.pshell.eval("home_fast_table()")
+        self.pshell.eval("set_mount_position()")
+
+    def set_setting(self, key: str, value: str):
+        self.pshell.eval(f"set_setting('{key}', '{value}')&")
+
+    def get_setting(self, key: str) -> str:
+        return self.pshell.eval(f"get_setting('{key}')&")
+
+    def enable_room_temperature(self):
+        self.set_setting("room_temperature_enabled", "true")
+
+    def disable_room_temperature(self):
+        self.set_setting("room_temperature_enabled", "false")
+
+    def get_mounted_sample(self) -> SampleDewarAddress | None:
+        if self.__simulation:
+            ret = self._simulated_mounted_sample
+        else:
+            ret = self.pshell.eval("get_setting('mounted_sample_position')&").strip()
+
+        if len(ret) == 0:
+            return None
+
+        match = re.match(r"([A-Z])(\d)(\d{1,2})", ret)
+
+        if match:
+            segment, puck, sample = match.groups()
+            dewar_location = DewarAddress(segment=segment, pos=int(puck))
+            return SampleDewarAddress(puck=dewar_location, pin=int(sample))
+        else:
+            print("Failed to decode string!")
+
+    def get_system_check(self):
+        if self.__simulation:
+            if random.random() < 0.1:
+                raise RuntimeError("get_system_check_failed")
+            return "Ok"
+        return self.pshell.eval("system_check_msg()&")
+
+    def get_robot_state(self):
+        if self.__simulation:
+            return "Ready"
+        return self.pshell.eval("robot.state&")
+
+    def get_robot_status(self):
+        if self.__simulation:
+            return {
+                "powered": True,
+                "settled": True,
+                "speed": 100,
+                "empty": True,
+                "mode": "remote",
+                "task": None,
+                "pos": "pCold",
+                "open": True,
+                "status": "move",
+            }
+
+        status = self.pshell.eval("robot.take()&")
+        return eval(status)  # FIXME ALL functions must return a valid JSON object
+
+    def get_speed(self) -> float:
+        if self.__simulation:
+            if random.random() < 0.1:
+                return random.choice([1, 5, 25, 50, 75, 90])
+            return 100.0
+        speed = self.get_robot_status()["speed"]
+        return float(speed)
+
+    def get_detected_pucks(self) -> List[PuckLoadedInfo]:
+        if self.__simulation:
+            j = self._simulated_detected_pucks
+        else:
+            j = json.loads(self.pshell.eval("get_pucks_info()&"))
+
+        output = []
+
+        for i in j:
+            if i["puckState"] == "Present":
+                puck_address = i["puckAddress"]
+                if len(puck_address) == 2:
+                    output.append(
+                        PuckLoadedInfo(
+                            puck_name=i["puckBarcode"],
+                            location=DewarAddress(
+                                segment=puck_address[0], pos=int(puck_address[1])
+                            ),
+                        ),
+                    )
+        return output
+
+    def set_pin_offset(self, value):
+        if self.__simulation:
+            print(f"simulated set_pin_offset {value}")
+            self._simulated_offset = value
+            return
+        self.pshell.eval("set_pin_offset(" + str(value) + ")&")
+
+    def get_pin_offset(self):
+        if self.__simulation:
+            print(f"simulated get_pin_offset -> {self._simulated_offset}")
+            return self._simulated_offset
+        try:
+            offset = float(self.pshell.eval("get_pin_offset()&"))
+        except Exception:
+            offset = 0.0
+        return offset
+
+    def get_current(self):
+        if self.__simulation:
+            return self._simulated_current
+        current = self.pshell.eval("smart_magnet.get_current_rb()&")
+        return float(current)
+
+    def set_current(self, current):
+        if self.__simulation:
+            self._simulated_current = current
+            return
+        self.pshell.eval("smart_magnet.set_current({:.1f})&".format(current))
+        current = self.pshell.eval("smart_magnet.get_current_rb()&")
+        return float(current)
+
+    def print_info(self):
+        print("State: " + str(self.get_state()))
+        print("Mounted sample: " + str(self.get_mounted_sample()))
+        print("System check: " + str(self.get_system_check()))
+        print("Robot state: " + str(self.get_robot_state()))
+        print("Robot status: ")
+        status = self.get_robot_status()
+        status = status
+        for k, v in status.items():
+            print(f"{k:>10s}: {v}")
+        print("Pin offset: " + str(self.get_pin_offset()))
+        print("Mount position: " + str(self.is_in_mount_position()))
+        print("")
+
+    def is_powered(self):
+        if self.__simulation:
+            return True
+        return self.get_robot_status()["powered"]
+
+    def check_enable_motion(self):
+        if self.__simulation:
+            if random.random() < 0.1:
+                raise RuntimeError("check_enable_motion failed")
+        if not self.is_powered():
+            self.pshell.eval("enable_motion()&")
+
+    def is_in_park(self):
+        if self.__simulation:
+            return True
+        return self.is_position(POSITION_PARK)
+
+    def is_in_home(self):
+        if self.__simulation:
+            return False
+        return self.is_position(POSITION_HOME)
+
+    def is_in_cold(self):
+        if self.__simulation:
+            return False
+        return self.is_position(POSITION_COLD)
+
+    def is_position(self, position: str) -> bool:
+        return position == self.get_robot_status()["pos"]
+
+    def get_task(self):
+        """
+        robot_status = {
+            'powered': False,
+            'settled': True,
+            'speed': 10,
+            'empty': True,
+            'mode': 'remote',
+            'task': None,
+            'pos': 'pPark',
+            'open': True,
+            'status': 'hold'
+        }
+        :return:
+        """
+        status = self.get_robot_status()
+        return status["task"]
+
+    def is_ready(self):
+        return "ready" == self.get_state().lower()
+
+    def is_busy(self):
+        return "busy" == self.get_state().lower()
+
+
+def is_true(value):
+    """check if argument is semantically true"""
+    value = str(value).lower()
+    return value != "0" or value in ("true", "yes", "on", "enabled")
+
+
+def is_false(value):
+    return not is_true(value)
+
+
+def is_valid_dewar_position(position):
+    return position in VALID_DEWAR_POSITIONS

src/mxlibs3/workflow_tools.py

@@ -0,0 +1,41 @@
+import time
+
+from epics import poll
+
+
+def wait_position(motor, target, tolerance=None, timeout=60.0):
+    position = motor.readback
+    if isinstance(position, float) and tolerance is None:
+        tst = "%f %s"
+        ltst = lambda x, y, z: eval(tst % (x, y))  # noqa: E731
+    elif isinstance(position, float) and tolerance is not None:
+        tst = "abs(%f - %f) < %f"
+        ltst = lambda x, y, z: eval(tst % (x, y, z))  # noqa: E731
+    elif isinstance(position, bytes):
+        tst = "'%s' == '%s'"
+        ltst = lambda x, y, z: eval(tst % (x.lower(), y.lower()))  # noqa: E731
+    else:
+        raise RuntimeError("BUG: could not understand arguments: ")
+
+    timeout = time.time() + timeout
+    timeisup = False
+    n = 20
+    condition = ltst(motor.readback, target, tolerance)
+    while not timeisup and not condition:
+        poll(0.1)
+        n = n + 1
+        if n > 20:
+            n = 0
+            print("waiting_position test: %s (%s, %s, %s)" % (tst, str(motor.readback), str(target), str(tolerance)))
+        timeisup = timeout < time.time()
+        condition = ltst(motor.readback, target, tolerance)
+
+    if timeisup:
+        msg = "Timeout when waiting for %s to reach %s with tolerance %s. Device was at: %s" % (
+            motor,
+            str(target),
+            str(tolerance),
+            str(motor.readback),
+        )
+        print(msg)
+        raise RuntimeError(msg)

src/pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "aaredaq"
+version = "0.0.1"
+description = "AareDAQ data acquisition server"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "requests~=2.32.3",
+    "pyepics~=3.5.7",
+    "redis~=6.1.0",
+    "python-redis-lock~=4.0.0",
+    "PyJWT==2.10.1",
+    "fastapi==0.115.12",
+    "uvicorn==0.34.2",
+    "ultralytics==8.3.133",
+    "aareDB==0.1.1a4",
+    "aaredaqlib==0.0.2",
+    "jfjoch_client==1.0.0rc37"
+]
+
+[lint]
+ignore = ["F401", "F541", "W503", "W504"]
+
+[build-system]
+requires = ["setuptools>=75.6.0"]
+build-backend = "setuptools.build_meta"