commit 177d9fab191ba6a932888047ba550502a45a14a0 Author: Filip Leonarski Date: Wed May 14 10:06:30 2025 +0200 Initial commit diff --git a/src/.gitignore b/src/.gitignore new file mode 100644 index 00000000..62c89355 --- /dev/null +++ b/src/.gitignore @@ -0,0 +1 @@ +.idea/ \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/aaredaq/__init__.py b/src/aaredaq/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/aaredaq/aaredb.py b/src/aaredaq/aaredb.py new file mode 100644 index 00000000..3f80fa9f --- /dev/null +++ b/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) \ No newline at end of file diff --git a/src/aaredaq/auth.py b/src/aaredaq/auth.py new file mode 100644 index 00000000..a31db894 --- /dev/null +++ b/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) diff --git a/src/aaredaq/beamline.py b/src/aaredaq/beamline.py new file mode 100644 index 00000000..44a7d77e --- /dev/null +++ b/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 diff --git a/src/aaredaq/best.pt b/src/aaredaq/best.pt new file mode 100644 index 00000000..8ad56d3d Binary files /dev/null and b/src/aaredaq/best.pt differ diff --git a/src/aaredaq/config.py b/src/aaredaq/config.py new file mode 100644 index 00000000..a0daff3e --- /dev/null +++ b/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) diff --git a/src/aaredaq/daq.py b/src/aaredaq/daq.py new file mode 100644 index 00000000..5b6abd05 --- /dev/null +++ b/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 + ) + diff --git a/src/aaredaq/devices.py b/src/aaredaq/devices.py new file mode 100644 index 00000000..ea2a69fb --- /dev/null +++ b/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() diff --git a/src/aaredaq/mlbox.py b/src/aaredaq/mlbox.py new file mode 100644 index 00000000..a850af50 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/aaredaq/models.py b/src/aaredaq/models.py new file mode 100644 index 00000000..b143742a --- /dev/null +++ b/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 diff --git a/src/aaredaq/server.py b/src/aaredaq/server.py new file mode 100644 index 00000000..45f6361a --- /dev/null +++ b/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() diff --git a/src/aaredaq/workflows.py b/src/aaredaq/workflows.py new file mode 100644 index 00000000..04fb200e --- /dev/null +++ b/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()) diff --git a/src/mxlibs3/__init__.py b/src/mxlibs3/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/mxlibs3/aerotech.py b/src/mxlibs3/aerotech.py new file mode 100644 index 00000000..8e34fbe5 --- /dev/null +++ b/src/mxlibs3/aerotech.py @@ -0,0 +1,1191 @@ +""" +``Aerotech`` --- Aerotech control software +****************************************** + +This module provides an object to control the Aerotech Abr rotational stage. + +The following constants are defined. + +Measurement phases: + + Aerotech.ABR_DONE + Aerotech.ABR_READY + Aerotech.ABR_BUSY + +Axis mode + + Aerotech.DIRECT_MODE + Aerotech.MEASURING_MODE + +Shutter + + Aerotech.SHUTTER_OPEN + Aerotech.SHUTTER_CLOSE + +Methods in the Abr class +======================== + +Aerotech.is_homed() +Aerotech.do_homing(wait=True) +Aerotech.get_ready(ostart=None, orange=None, etime=None, wait=True) +Aerotech.is_done() +Aerotech.is_ready() +Aerotech.is_busy() +Aerotech.stop() +Aerotech.start_exposure() +Aerotech.wait_status(status) +Aerotech.set_mode(, value) +Aerotech.get_mode() +Aerotech.set_direct_mode() +Aerotech.set_measuring_mode() +Aerotech.move(angle, wait=False, speed=None) +Aerotech.set_shutter(state) + +Attribute Access in Abr class +============================= + +The Abr class overwrites the getattr and setattr methods to provide a pythonic +way of controlling the Abr. + +The following properties are implemented: + +velo + + Sets the velocity of rotation: ``-ES-DF1:ROTX-SETV`` + +omega + + Move the omega angle without any wait: ``-ES-DF1:ROTX-VAL`` + +exposure_time + + Sets the PV for the measurement's exposure time: ``-ES-OSC:ETIME`` + +start_angle + + Sets the PV for the measurement's starting angle: ``-ES-OSC:START-POS`` + +oscillation_angle + + Sets the PV for the measurement's oscillation angle: ``-ES-OSC:RANGE`` + +shutter + + Controls the shutter: ``-ES-PH1:SET`` and ``-ES-PH1:GET`` + +mode + + Controls the axis mode: ``-ES-DF1:AXES-MODE`` + +measurement_state + + Controls the PV for the measurement state: ``-ES-OSC:DONE`` + + +Examples +======== + + import Aerotech + abr = Aerotech.Abr() + + # move omega to 270.0 degrees + abr.omega = 270.0 + + # move omega to 180 degrees and wait for movement to finish + abr.move(180, wait=True) + + # move omega to 3000 degrees at 360 degrees/s and wait for movement to finish + abr.move(3000, speed=360, wait=True) + + # stop any movement + abr.stop() # this function only returns after the STATUS is back to OK + +""" + +import math +import time + +from epics import PV, poll +from epics.ca import pend_io + +from aaredaq.beamline import MXBeamline + +ABR_DONE = 0 +ABR_READY = 1 +ABR_BUSY = 2 + +GRID_SCAN_BUSY = 0 +GRID_SCAN_DONE = 1 + +DIRECT_MODE = 0 +MEASURING_MODE = 1 + +SHUTTER_CLOSE = 0 +SHUTTER_OPEN = 1 + +FULL_PERIOD = 0 +HALF_PERIOD = 1 + + +def pv_wait(pv, val, timeout=10.0): + timeout = timeout + time.time() + timeisup = False + while val != pv.get() and not timeisup: + poll(0.01) + timeisup = time.time() > timeout + if timeisup: + raise RuntimeWarning( + "timeout waiting for %s to reach %s" % (pv.pvname, str(val)) + ) + + +CMD_NONE = 0 +CMD_RASTER_SCAN_SIMPLE = 1 +CMD_MEASURE_STANDARD = 2 +CMD_VERTICAL_LINE_SCAN = 3 +CMD_SCREENING = 4 +CMD_SUPER_FAST_OMEGA = 5 +CMD_STILL_WEDGE = 6 +CMD_STILLS = 7 +CMD_REPEAT_SINGLE_OSCILLATION = 8 +CMD_SINGLE_OSCILLATION = 9 +CMD_OLD_FASHIONED = 10 +CMD_RASTER_SCAN = 11 +CMD_JET_ROTATION = 12 +CMD_X_HELICAL = 13 +CMD_X_RUNSEQ = 14 +CMD_JUNGFRAU = 15 +CMD_MSOX = 16 +CMD_SLIT_SCAN = 17 +CMD_RASTER_SCAN_STILL = 18 +CMD_SCAN_SASTT = 19 +CMD_SCAN_SASTT_V2 = 20 +CMD_SCAN_SASTT_V3 = 21 + +AXIS_OMEGA = 1 +AXIS_GMX = 2 +AXIS_GMY = 3 +AXIS_GMZ = 4 +AXIS_STY = 5 +AXIS_STZ = 6 + + +class Abr: + def __init__(self, bl: MXBeamline): + BEAMLINE = "X06DA" # FIXME + self.__pv_stop = PV(f"{BEAMLINE}-ES-AERO:TSK-STOP") + self.__pv_lock = PV(f"{BEAMLINE}-ES-DF1:LOCK") + self.__pv_direct_mode = PV(f"{BEAMLINE}-ES-DF1:MODE-DIRECT") + self.set_direct_mode = self.__set_direct_mode_v2 + self.lock_device = self.__lock_device_v2 + self.lock = self.__lock_device_v2 + + self.unlock = self.set_direct_mode + + self.__pv_shutter_set = PV(f"{BEAMLINE}-ES-PH1:SET") + self.__pv_shutter_get = PV(f"{BEAMLINE}-ES-PH1:GET") + self.__pv_axes_mode = PV(f"{BEAMLINE}-ES-DF1:AXES-MODE") + self.__pv_omega_set = PV(f"{BEAMLINE}-ES-DF1:OMEGA-VAL") + self.__pv_omega_get = PV(f"{BEAMLINE}-ES-DF1:OMEGA-RBV") + self.__pv_omega_done = PV(f"{BEAMLINE}-ES-DF1:OMEGA-DONE") + self.__pv_gmx_setp = PV(f"{BEAMLINE}-ES-DF1:GMX-SETP") + self.__pv_gmx_getp = PV(f"{BEAMLINE}-ES-DF1:GMX-GETP") + self.__pv_gmx_val = PV(f"{BEAMLINE}-ES-DF1:GMX-VAL") + self.__pv_gmx_rbv = PV(f"{BEAMLINE}-ES-DF1:GMX-RBV") + self.__pv_gmx_velo = PV(f"{BEAMLINE}-ES-DF1:GMX-SETV") + self.__pv_gmx_done = PV(f"{BEAMLINE}-ES-DF1:GMX-DONE") + self.__pv_gmy_setp = PV(f"{BEAMLINE}-ES-DF1:GMY-SETP") + self.__pv_gmy_getp = PV(f"{BEAMLINE}-ES-DF1:GMY-GETP") + self.__pv_gmy_set = PV(f"{BEAMLINE}-ES-DF1:GMY-VAL") + self.__pv_gmy_get = PV(f"{BEAMLINE}-ES-DF1:GMY-RBV") + self.__pv_gmy_done = PV(f"{BEAMLINE}-ES-DF1:GMY-DONE") + self.__pv_gmz_setp = PV(f"{BEAMLINE}-ES-DF1:GMZ-SETP") + self.__pv_gmz_getp = PV(f"{BEAMLINE}-ES-DF1:GMZ-GETP") + self.__pv_gmz_set = PV(f"{BEAMLINE}-ES-DF1:GMZ-VAL") + self.__pv_gmz_get = PV(f"{BEAMLINE}-ES-DF1:GMZ-RBV") + self.__pv_gmz_done = PV(f"{BEAMLINE}-ES-DF1:GMZ-DONE") + self.__pv_omega_velo = PV(f"{BEAMLINE}-ES-DF1:OMEGA-SETV") + self.__pv_ostart = PV(f"{BEAMLINE}-ES-OSC:START-POS") + self.__pv_orange = PV(f"{BEAMLINE}-ES-OSC:RANGE") + self.__pv_etime = PV(f"{BEAMLINE}-ES-OSC:ETIME") + self.__pv_get_ready = PV(f"{BEAMLINE}-ES-OSC:READY.PROC") + self.__pv_start = PV(f"{BEAMLINE}-ES-OSC:START") + self.__pv_phase = PV(f"{BEAMLINE}-ES-OSC:DONE") + self.__pv_stat = PV(f"{BEAMLINE}-ES-AERO:STAT") + self.__pv_ishomed = PV(f"{BEAMLINE}-ES-DF1:OMEGA-AS00") + self.__pv_dohome = PV(f"{BEAMLINE}-ES-DF1:OMEGA-HOME") + self.__pv_omega_status = PV(f"{BEAMLINE}-ES-DF1:OMEGA-STAT") + self._raster_x_start = PV(f"{BEAMLINE}-ES-GRD:GMX-START") + self._raster_x_end = PV(f"{BEAMLINE}-ES-GRD:GMX-END") + self._raster_omega = PV(f"{BEAMLINE}-ES-OSC:#M-START") + self._raster_osc = PV(f"{BEAMLINE}-ES-GRD:ANGLE") + self._raster_celltime = PV(f"{BEAMLINE}-ES-GRD:CELL-TIME") + self._raster_y_start = PV(f"{BEAMLINE}-ES-OSC:#STY-START") + self._raster_y_step = PV(f"{BEAMLINE}-ES-OSC:#STY-END") + self._raster_z_start = PV(f"{BEAMLINE}-ES-OSC:#STZ-START") + self._raster_z_step = PV(f"{BEAMLINE}-ES-OSC:#STZ-END") + self._raster_columns = PV(f"{BEAMLINE}-ES-GRD:COLUMNS") + self._raster_rows = PV(f"{BEAMLINE}-ES-GRD:ROWS") + self._raster_delay = PV(f"{BEAMLINE}-ES-GRD:RAST-DLY") + self._raster_osc_mode = PV(f"{BEAMLINE}-ES-GRD:SET-MODE") + self._raster_velo = PV(f"{BEAMLINE}-ES-GRD:#X-VEL") + self._raster_get_ready = PV(f"{BEAMLINE}-ES-GRD:READY.PROC") + self._raster_grid_start = PV(f"{BEAMLINE}-ES-GRD:START.PROC") + self._raster_grid_next = PV(f"{BEAMLINE}-ES-GRD:NEXT-ROW") + self._raster_scan_done = PV(f"{BEAMLINE}-ES-GRD:SCAN-DONE") + self._raster_row_done = PV(f"{BEAMLINE}-ES-GRD:ROW-DONE") + self._raster_grid_done = PV(f"{BEAMLINE}-ES-GRD:DONE") + self._var_1 = PV(f"{BEAMLINE}-ES-PSO:VAR-1") + self._var_2 = PV(f"{BEAMLINE}-ES-PSO:VAR-2") + self._var_3 = PV(f"{BEAMLINE}-ES-PSO:VAR-3") + self._var_4 = PV(f"{BEAMLINE}-ES-PSO:VAR-4") + self._var_5 = PV(f"{BEAMLINE}-ES-PSO:VAR-5") + self._var_6 = PV(f"{BEAMLINE}-ES-PSO:VAR-6") + self._var_7 = PV(f"{BEAMLINE}-ES-PSO:VAR-7") + self._var_8 = PV(f"{BEAMLINE}-ES-PSO:VAR-8") + self._var_9 = PV(f"{BEAMLINE}-ES-PSO:VAR-9") + self._var_10 = PV(f"{BEAMLINE}-ES-PSO:VAR-10") + self._zcmd = PV(f"{BEAMLINE}-ES-PSO:CMD") + self.__start_command = PV(f"{BEAMLINE}-ES-PSO:START-TEST.PROC") + self._cmd_2 = PV(f"{BEAMLINE}-ES-PSO:STOP-TEST.PROC") + self._gmy_done = PV(f"{BEAMLINE}-ES-DF1:GMY-DONE") + self._gmx_offset = PV(f"{BEAMLINE}-ES-DF1:GMX-OFF") + + pend_io(10) + + def raster_setup(self, positions, columns, angle, etime): + """configures parameters for a raster scan + + positions : a list of xyz indicating the positions of each cell + columns: an integer with how many columns + angle : the oscillation angle for each row + etime : the exposure time per cell + """ + self.set_direct_mode() + if columns == 1: + raise RuntimeWarning( + "Fast scans are not available with vertical line scans." + ) + + rows = len(positions) / columns + x1, y1, z1 = tuple(positions[0]) # first cell on first row + x2, y2, z2 = tuple(positions[1]) # second cell on first row + x4, y4, z4 = tuple(positions[columns - 1]) # last cell on first row + + if rows > 1: + x3, y3, z3 = tuple(positions[columns]) # first cell on second row + ystep = y3 - y1 + zstep = z3 - z1 + gmy_step = math.sqrt(pow(y3 - y1, 2) + pow(z3 - z1, 2)) + else: + ystep = 0.010 + zstep = 0.010 + gmy_step = 0.010 + + half_cell = abs(x2 - x1) / 2.0 + start_x = x1 - half_cell + end_x = x4 + half_cell + + self._raster_starting_x = start_x + self._raster_starting_y = y1 + self._raster_starting_z = z1 + self._raster_step_y = ystep + self._raster_step_z = zstep + self._raster_current_row = 0 + self._raster_num_rows = rows + + self._raster_x_start.put(start_x) + self._raster_x_end.put(end_x) + self._raster_omega.put(self.position) + self._raster_osc.put(angle) + self._raster_celltime.put(etime) + self._raster_y_start.put(y1) + self._raster_z_start.put(z1) + self._raster_y_step.put(ystep) + self._raster_z_step.put(zstep) + self._raster_columns.put(columns) + self._raster_rows.put(rows) + + def raster_next_row_yz(self): + r = self._raster_current_row + self._raster_current_row = r + 1 + + y, z = ( + r * self._raster_step_y + self._raster_starting_y, + r * self._raster_step_z + self._raster_starting_z, + ) + + if r < self._raster_num_rows: + return (y, z) + else: + return (None, None) + + def raster_get_y_step(self): + return self._raster_y_step.get() + + def raster_get_z_step(self): + return self._raster_z_step.get() + + def raster_get_ready(self): + self._raster_get_ready.put(1) + for n in range(10): + try: + pv_wait(self._raster_grid_done, ABR_READY, timeout=2) + except RuntimeWarning as e: + print(str(e), end=" ") + print(" --- trying again.") + self._raster_get_ready.put(1, wait=True) + + def raster_scan_row(self): + """start rastering a new row, either first or following rows""" + if ABR_READY == self._raster_grid_done.get(): + self._raster_grid_start.put(1) + poll(0.05) + else: + self._raster_grid_next.put(1) + + def is_scan_done(self): + return 1 == self._raster_scan_done.get() + + def wait_scan_done(self, timeout=60.0): + poll(0.05) + pv_wait(self._raster_scan_done, 1, timeout=timeout) + + def raster_is_row_scanning(self): + return 1 == self._raster_scan_done.get() + + def raster_wait_for_row(self): + poll(0.05) + pv_wait(self._raster_row_done, 1, timeout=60) + + def command(self, cmd): + self._zcmd.put(cmd) + poll(0.1) + + def start_command(self): + poll(0.1) + self.__start_command.put(1) + poll(0.1) + + def measure_standard(self, start, wedge, etime, ready_rate=500.0): + """measure a standard wedge from `start` to `start` + `wedge` during `etime` seconds""" + self.command(CMD_MEASURE_STANDARD) + self._var_1.put(start) + self._var_2.put(wedge) + self._var_3.put(etime) + self._var_4.put(ready_rate) + self.start_command() + + def measure_scan_slits(self, delta_x, delta_y, velocity): + """measure a standard wedge from `start` to `start` + `wedge` during `etime` seconds""" + self.command(CMD_SLIT_SCAN) + self._var_1.put(delta_x) + self._var_2.put(delta_y) + self._var_3.put(velocity) + self.start_command() + + def measure_raster_simple(self, etime, cell_width, cell_height, ncols, nrows): + """measures the specified grid assuming the goniometer is currently at the CENTER of TOP-LEFT CELL + + etime: a float in seconds + cell_width: a float in mm + cell_height: a float in mm + ncols: an int + nrows: an int + + """ + # self.reset() + self.command(CMD_RASTER_SCAN_SIMPLE) + self._raster_scan_done.put(0) # make SCAN DONE go into busy state + self._var_1.put(etime) + self._var_2.put(cell_width) + self._var_3.put(cell_height) + self._var_4.put(int(ncols)) + self._var_5.put(int(nrows)) + self._var_6.put(0) + self._var_7.put(0) + self._var_8.put(0) + self._raster_num_rows = nrows + self.start_command() + + def measure_sastt( + self, + etime, + cell_width, + cell_height, + ncols, + nrows, + even_tweak=0.0, + odd_tweak=0.0, + version=1, + ): + """measures the specified grid assuming the goniometer is currently at the CENTER of the required grid + + etime: a float in seconds + cell_width: a float in mm + cell_height: a float in mm + ncols: an int + nrows: an int + + odd_tweak, even_tweak: float, a distance to be added to the start of + the open shutter command. used on in the scan mode + version=3 + + - these values can be either positive/negative + - they should be small, i.e. just to tweak a trriggering event + by minor amounts + version: an int (1, 2, 3) + 1 = original snake scan, single PSO window + 2 = scan always from LEFT---RIGHT + 3 = snake scan alternating PSO window for even/odd rows + + """ + # self.reset() + if version == 1: + self.command(CMD_SCAN_SASTT) + elif version == 2: + self.command(CMD_SCAN_SASTT_V2) + elif version == 3: + self.command(CMD_SCAN_SASTT_V3) + else: + raise RuntimeError("non existing SAS-TT scan mode") + + self._raster_scan_done.put(0) # make SCAN DONE go into busy state + self._var_1.put(etime) + self._var_2.put(cell_width) + self._var_3.put(cell_height) + self._var_4.put(int(ncols)) + self._var_5.put(int(nrows)) + self._var_6.put(even_tweak) + self._var_7.put(odd_tweak) + self._var_8.put(0) + self._raster_num_rows = nrows + self.start_command() + + def measure_vertical_line(self, cellHeight, numCells, exposureTime): + """measure a vertical line using the GMY motor + + The line is nummCells * cellHeight long and the exposureTime is per cell. + + `cellHeight` in mm + `numCells` an integer + `exposureTime` in seconds / per cellHeight + """ + self.command(CMD_VERTICAL_LINE_SCAN) + self._var_1.put(cellHeight) + self._var_2.put(numCells) + self._var_3.put(exposureTime) + self.start_command() + + def measure_screening(self, start, oscangle, etime, degrees, frames, delta=0.5): + self.command(CMD_SCREENING) + self._var_1.put(start) + self._var_2.put(oscangle) + self._var_3.put(etime) + self._var_4.put(degrees) + self._var_5.put(frames) + self._var_6.put(delta) + self.start_command() + + def measure_fast_omega(self, start, wedge, etime, rate=200.0): + self.command(CMD_SUPER_FAST_OMEGA) + self._var_1.put(start) + self._var_2.put(wedge) + self._var_3.put(etime) + self._var_4.put(rate) + self.start_command() + + def measure_still_wedge( + self, start, wedge, etime, oscangle, sleep_after_shutter_close=0.010 + ): + self.command(CMD_STILL_WEDGE) + self._var_1.put(start) + self._var_2.put(wedge) + self._var_3.put(etime) + self._var_4.put(oscangle) + self._var_5.put(sleep_after_shutter_close) + self.start_command() + + def measure_stills(self, number_of_frames, exposure_time, idle=0.0): + """collect still images, i.e. no omega rotation + + number_of_frames => an integer + exposure_time => a float + idle => sleep in between each exposure + + IMPORTANT: use idle=0.0 to prevent shutter open/close + + """ + self.command(CMD_STILLS) + self._var_1.put(number_of_frames) + self._var_2.put(exposure_time) + self._var_3.put(idle) + self.start_command() + + def measure_repeat_singles(self, frames, etime, oscangle, sleep=0.0, delta=0.5): + self.command(CMD_REPEAT_SINGLE_OSCILLATION) + self._var_1.put(frames) + self._var_2.put(etime) + self._var_3.put(oscangle) + self._var_4.put(sleep) + self._var_5.put(delta) + self.start_command() + + def measure_single( + self, start_angle, oscillation_angle, exposure_time, delta=0.0, settle=0.0 + ): + self.command(CMD_SINGLE_OSCILLATION) + self._var_1.put(start_angle) + self._var_2.put(oscillation_angle) + self._var_3.put(exposure_time) + self._var_4.put(delta) + self._var_5.put(settle) + self.start_command() + + def measure_raster(self, positions, columns, angle, etime): + """raster via generic experiment interface""" + self.reset() + if len(positions) == 1: + raise RuntimeWarning("Raster scan with one cell makes no sense") + self.command(CMD_RASTER_SCAN) + + rows = len(positions) / columns + x1, y1, z1 = tuple(positions[0]) # first cell on first row + x2, y2, z2 = tuple(positions[1]) # second cell on first row + x4, y4, z4 = tuple(positions[columns - 1]) # last cell on first row + + if rows > 1: + x3, y3, z3 = tuple(positions[columns]) # first cell on second row + ystep = y3 - y1 + zstep = z3 - z1 + gmy_step = math.sqrt(pow(y3 - y1, 2) + pow(z3 - z1, 2)) + else: + ystep = 0.010 + zstep = 0.010 + gmy_step = 0.010 + + half_cell = abs(x2 - x1) / 2.0 + start_x = x1 - half_cell + end_x = x4 + half_cell + + self._raster_starting_x = start_x + self._raster_starting_y = y1 + self._raster_starting_z = z1 + self._raster_step_y = ystep + self._raster_step_z = zstep + self._raster_current_row = 0 + self._raster_num_rows = rows + + self._var_1.put(self.position) + self._var_2.put(start_x) + self._var_3.put(end_x) + self._var_4.put(columns) + self._var_5.put(rows) + self._var_6.put(angle) + self._var_7.put(etime) + self._var_8.put(HALF_PERIOD) + self._var_9.put(self._gmx_offset.get()) + self._var_10.put(gmy_step) + + self.start_command() + + def measure_raster_still(self, positions, columns, etime, delay): + """raster via generic experiment interface + + positions: list of coordinates for each cell + columns: how many columns in grid + etime: exposure time / cell_width + delay: delay configured in delay generator + to trigger detector after aerotech; + must be greater than actual shutter delay + """ + self.reset() + if len(positions) == 1: + raise RuntimeWarning("Raster scan with one cell makes no sense") + self.command(CMD_RASTER_SCAN_STILL) + + rows = len(positions) / columns + x1, y1, z1 = tuple(positions[0]) # first cell on first row + x2, y2, z2 = tuple(positions[1]) # second cell on first row + x4, y4, z4 = tuple(positions[columns - 1]) # last cell on first row + + if rows > 1: + x3, y3, z3 = tuple(positions[columns]) # first cell on second row + ystep = y3 - y1 + zstep = z3 - z1 + gmy_step = math.sqrt(pow(y3 - y1, 2) + pow(z3 - z1, 2)) + else: + ystep = 0.010 + zstep = 0.010 + gmy_step = 0.010 + + half_cell = abs(x2 - x1) / 2.0 + start_x = x1 - half_cell + end_x = x4 + half_cell + + self._raster_starting_x = start_x + self._raster_starting_y = y1 + self._raster_starting_z = z1 + self._raster_step_y = ystep + self._raster_step_z = zstep + self._raster_current_row = 0 + self._raster_num_rows = rows + + self._var_1.put(self.position) + self._var_2.put(start_x) + self._var_3.put(end_x) + self._var_4.put(columns) + self._var_5.put(rows) + self._var_6.put(delay) + self._var_7.put(etime) + self._var_8.put(HALF_PERIOD) + self._var_9.put(self._gmx_offset.get()) + self._var_10.put(gmy_step) + + self.start_command() + + def measure_jungfrau(self, columns, rows, width, height, etime, sleep=0.100): + """for Jungfrau, grid from current position""" + self.reset() + self.command(CMD_JUNGFRAU) + self._var_1.put(columns) + self._var_2.put(rows) + self._var_3.put(height) + self._var_4.put(width) + self._var_5.put(etime) + self._var_6.put(sleep) + self._var_7.put(0.0) + self._var_8.put(0.0) + self._var_9.put(0.0) + self._var_10.put(0.0) + + self.start_command() + if sleep > 0: + timeout = 5.0 + (rows * sleep) + (columns * rows * etime) + else: + timeout = None + + if timeout: + pv_wait(self._raster_scan_done, GRID_SCAN_DONE, timeout) + self.set_direct_mode() + + def measure_msox(self, start, wedge, etime, num_datasets=1, rest_time=0.0): + """for msox + start = start omega + wedge = total range to measure + etime = total time for wedge + num_datasets = how many times to repeat + rest_time = how many seconds to rest in between datasets + """ + self.reset() + self.command(CMD_MSOX) + self._var_1.put(start) + self._var_2.put(wedge) + self._var_3.put(etime) + self._var_4.put(num_datasets) + self._var_5.put(rest_time) + self._var_6.put(0) + self._var_7.put(0) + self._var_8.put(0) + self._var_9.put(0) + self._var_10.put(0) + + self.start_command() + + def next_row(self): + """start rastering a new row""" + self._raster_grid_next.put(1) + + def configure_stills(self, frames, etime, sleep=0.0): + self._var_1.put(frames) + self._var_2.put(etime) + self._var_3.put(sleep) + + def reset(self): + self.command(CMD_NONE) + self.set_direct_mode() + + def raster_2d(self): + pass + + def is_ioc_ok(self): + return 0 == self.__pv_stat.get() + + def is_ready_and_willing(self): + return self.is_ioc_ok() and self.is_homed() + + def is_omega_ok(self): + return 0 == self.__pv_omega_status.get() + + def is_homed(self): + """Are the motors homed?""" + return 1 == self.__pv_ishomed.get() + + def do_homing(self, wait=True): + """Execute the homing procedure. + + Executes the homing procedure and waits (default) until it is completed. + + PARAMETERS + `wait` true / false if the routine is to wait for the homing to finish. + """ + self.__pv_dohome.put(1, wait=True) + poll(1.0) + if not wait: + return + while not self.is_homed(): + poll(0.2) + + def __setattr__(self, name, value): + """Overwrite __setattr__ for pythonic access. + + This is in principle unnecessary but it results in clear code. + + Carefull when modifying this! + """ + if name == "velo": + self.__pv_omega_velo.put(value, wait=True) + elif name == "omega": + self.__pv_omega_set.put(value, wait=True) + elif name == "position": + self.__pv_omega_set.put(value, wait=True) + elif name == "exposure_time": + self.__pv_etime.put(value, wait=True) + elif name == "start_angle": + self.__pv_ostart.put(value, wait=True) + elif name == "oscillation_angle": + self.__pv_orange.put(value, wait=True) + elif name == "measurement_state": + self.__pv_phase.put(value, wait=True) + elif name == "shutter": + self.set_shutter(value) + elif name == "mode": + self.set_mode(value) + else: + self.__dict__[name] = value + + def __getattr__(self, name): + """Overwrite __getattr__ for pythonic access. + + This is in principle unnecessary but it results in clear code. + + Carefull when modifying this! + """ + if name == "velo": + value = self.__pv_omega_velo.get() + elif name == "omega": + value = self.__pv_omega_get.get() + elif name == "position": + value = self.__pv_omega_get.get() + elif name == "exposure_time": + value = self.__pv_etime.get() + elif name == "start_angle": + value = self.__pv_ostart.get() + elif name == "oscillation_angle": + value = self.__pv_orange.get() + elif name == "measurement_state": + value = self.__pv_phase.get() + elif name == "shutter": + value = self.__pv_shutter_get.get() + elif name == "mode": + value = self.get_mode() + else: + raise AttributeError("unknown attribute") + return value + + def get_ready(self, ostart=None, orange=None, etime=None, wait=True): + self.wait_for_movements() + if self.measurement_state == ABR_BUSY: + raise RuntimeError("ABR is not DONE!!!!") + + if self.measurement_state == ABR_READY: + self.measurement_state = ABR_DONE + + if ostart is not None: + self.__pv_ostart.put(ostart, wait=True) + if orange is not None: + self.__pv_orange.put(orange, wait=True) + if etime is not None: + self.__pv_etime.put(etime, wait=True) + + self.__pv_get_ready.put(1, wait=True) + + poll(0.1) + + if wait: + for n in range(10): + try: + self.wait_status(ABR_READY, timeout=5) + except RuntimeWarning as e: + print(str(e), end=" ") + print(" --- trying ready again.") + self.__pv_get_ready.put(1, wait=True) + + def wait_for_movements(self, timeout=60.0): + timeout = timeout + time.time() + timeisup = False + + try: + moving = self.is_moving() + except: + moving = True + + while not timeisup and moving: + poll(0.1) + try: + moving = self.is_moving() + except: + print( + "Aerotech() Failed to retrieve moving state for axes omega, gmx, gmy, gmz" + ) + moving = True + timeisup = timeout < time.time() + + if timeisup: + raise RuntimeWarning("timeout waiting for all axis to stop moving") + + def is_moving(self): + return not ( + self.__pv_omega_done.get() + and self.__pv_gmx_done.get() + and self.__pv_gmy_done.get() + and self.__pv_gmz_done.get() + ) + + def is_done(self): + """Is the measurement DONE?""" + return ABR_DONE == self.__pv_phase.get() + + def is_ready(self): + """Is the measurement READY to start?""" + return ABR_READY == self.__pv_phase.get() + + def is_busy(self): + """Is the measurement currently BUSY (executing)?""" + return ABR_BUSY == self.__pv_phase.get() + + def stop(self, wait_after_reload=1.0): + """Stops current motions; reloads aerobasic programs; goes DIRECT + + This will stop movements in both DIRECT and MEASURING modes. During the + stop the `status` temporarely goes to ERROR but reverts to OK after a + couple of seconds. + + """ + self.__pv_stop.put(1, wait=True) + poll(wait_after_reload) + self.reset() + poll(0.1) + + reload_programs = stop + + def start_exposure(self): + """Starts the previously configured exposure.""" + self.wait_for_movements() + self.__pv_start.put(1) + for n in range(10): + try: + self.wait_status(ABR_BUSY, timeout=1) + except RuntimeWarning as e: + print(str(e), end=" ") + print(" --- trying start again.") + self.__pv_start.put(1, wait=True) + + def wait_status(self, status, timeout=60.0): + """Wait for the Aertotech IOC to reach the desired `status`. + + PARAMETERS + `status` can be any the three ABR_DONE, ABR_READY or ABR_BUSY. + + NO RETURN + """ + st = self.__pv_phase.get() + timeout = timeout + time.time() + timeisup = False + while not timeisup and status != st: + poll(0.05) + st = self.__pv_phase.get() + timeisup = timeout < time.time() + + if timeisup: + raise RuntimeWarning( + "timeout waiting for abr status to change to %d: current status => %d" + % (status, st) + ) + + def get_mode(self, as_string=False): + return self.__pv_axes_mode.get(as_string=as_string) + + def is_measuring_mode(self): + return MEASURING_MODE == self.get_mode() + + def __set_direct_mode_v2(self): + """Sets the Abr to direct mode.""" + self.__pv_direct_mode.put(37, wait=True) + poll(0.1) + + def __lock_device_v1(self): + self.__pv_axes_mode.put(MEASURING_MODE, wait=True) + poll(0.1) + + def __lock_device_v2(self): + self.__pv_lock.put(1, wait=True) + poll(0.1) + + def lock_device(self): + self.__pv_lock.put(1, wait=True) + poll(0.1) + + def is_locked(self): + lpv = self.__pv_lock + if lpv is not None: + return 0 != self.get_mode() + else: + return 1 == self.get_mode() + + def velo(self, axis: int, speed: float | None = None) -> float: + """set's the speed on the given axis to speed""" + pv_velo: PV + if AXIS_GMX == axis: + velo_pv = self.__pv_gmx_velo + min, max = 0.01, 100.0 + elif AXIS_GMY == axis: + velo_pv = self.__pv_gmy_velo + min, max = 0.01, 10.0 + elif AXIS_GMZ == axis: + velo_pv = self.__pv_gmz_velo + min, max = 0.01, 10.0 + elif AXIS_OMEGA == axis: + velo_pv = self.__pv_omega_velo + min, max = 0.01, 720.0 + else: + raise Exception("unknown axis index") + + if speed is not None and velo_pv is not None: + if not (min <= speed <= max): + raise Exception( + f"requested speed {speed} outside range {min} < speed < {max}" + ) + velo_pv.put(speed) + poll(0.05) + return velo_pv.get() # type: ignore + + def select_axis(self, axis): + """returns PVs and limits for the given axis + + axis: Aerotech.AXIS_{OMEGA=1, GMX=2, GMY=3, GMZ=4} + + return (val_pv, rbv_pv, velo_pv, done_pv, velomin, velomax) + """ + + if AXIS_GMX == axis: + val_pv = self.__pv_gmx_val + rbv_pv = self.__pv_gmx_rbv + velo_pv = self.__pv_gmx_velo + done_pv = self.__pv_gmx_done + min, max = 0.01, 100.0 + elif AXIS_GMY == axis: + val_pv = self.__pv_gmy_val + rbv_pv = self.__pv_gmy_rbv + velo_pv = self.__pv_gmy_velo + done_pv = self.__pv_gmy_done + min, max = 0.01, 10.0 + elif AXIS_GMZ == axis: + val_pv = self.__pv_gmz_val + rbv_pv = self.__pv_gmz_rbv + velo_pv = self.__pv_gmz_velo + done_pv = self.__pv_gmz_done + min, max = 0.01, 10.0 + elif AXIS_OMEGA == axis: + val_pv = self.__pv_omega_set + rbv_pv = self.__pv_omega_get + velo_pv = self.__pv_omega_velo + done_pv = self.__pv_omega_done + min, max = 0.01, 720.0 + else: + raise Exception("unknown axis index") + + return (val_pv, rbv_pv, velo_pv, done_pv, min, max) + + def incr(self, axis, val, wait=False, velo=None): + (val_pv, rbv_pv, velo_pv, done_pv, min, max) = self.select_axis(axis) + cur = rbv_pv.get() # type: ignore + cv = velo_pv.get() # type: ignore + + if velo: + self.velo(axis, velo) + + val_pv.put(cur + val) + if wait: + poll(0.2) + pv_wait(done_pv, 1) + + if velo: + self.velo(axis, cv) + + def incr_x(self, val, wait=False, velo=None): + """increments GMX by val""" + self.incr(AXIS_GMX, val, wait=wait, velo=velo) + + def incr_y(self, val, wait=False, velo=None): + """increments GMY by val""" + self.incr(AXIS_GMY, val, wait=wait, velo=velo) + + def incr_z(self, val, wait=False, velo=None): + """increments GMZ by val""" + self.incr(AXIS_GMZ, val, wait=wait, velo=velo) + + def incr_omega(self, val, wait=False, velo=None): + """increments OMEGA by val""" + self.incr(AXIS_OMEGA, val, wait=wait, velo=velo) + + def move_x(self, val, wait=False, velo=None): + if velo is not None: + oldv = self.__pv_gmx_velo.get() + self.__pv_gmx_velo.put(velo) + velo = oldv + self.move_axis(AXIS_GMX, val, wait) + if velo is not None: + self.__pv_gmx_velo.put(velo) + + def move_y(self, val, wait=False): + self.move_axis(AXIS_GMY, val, wait) + + def move_z(self, val, wait=False): + self.move_axis(AXIS_GMZ, val, wait) + + def move_axis(self, axis, pos, wait=False): + axis_motor = None + if DIRECT_MODE != self.mode: + self.set_direct_mode() + if axis == AXIS_OMEGA: + self.move(pos) + else: + if axis == AXIS_GMX: + axis_motor = self.__pv_gmx_val.put + elif axis == AXIS_GMY: + axis_motor = self.__pv_gmy_set.put + elif axis == AXIS_GMZ: + axis_motor = self.__pv_gmz_set.put + if axis_motor: + axis_motor(pos) + + if wait: + self.wait_for_movements() + + def rbv(self): + x = self.__pv_gmx_rbv.get() + y = self.__pv_gmy_get.get() + z = self.__pv_gmz_get.get() + omega = self.__pv_omega_get.get() + return {"gmx": x, "gmy": y, "gmz": z, "omega": omega} + + def move(self, angle, *, wait=False, speed=None, direct=False): + """Move omega to `angle` at `speed` deg/s and `wait` (or not). + + If a `speed` is specified, it will be used for the movement but the speed + setting will be reverted to its old value afterwards. + + If `wait` is true, the routine waits until the omega position is close + to within 0.01 degrees to the target or that a timeout occured. The + timeout is the time required for the movement plus 5 seconds. + + In case of a timeout a RuntimeWarning is raised. + + PARAMETERS: + `angle` [degrees] a float specifying where to move to. + No limits. + `speed` [degrees/s] a float specifying speed to move at. + Limits: 0 - 360 + `wait` [boolean] wait or not for movement to finish. + """ + angle = float(angle) + if speed is not None: + speed = float(speed) + if DIRECT_MODE != self.mode: + if direct: + self.set_direct_mode() + else: + raise RuntimeWarning("abr is not in direct mode!") + oget = self.__pv_omega_get + oset = self.__pv_omega_set + vel = self.__pv_omega_velo + odone = self.__pv_omega_done + + cur = oget.get() + ovelo = vel.get() + if speed is not None: + vel.put(speed, wait=True) + poll(0.01) + velo = vel.get() + oset.put(angle, wait=True) + if not wait: + vel.put(ovelo, wait=True) + return + poll(0.01) + timeout = 30.0 + time.time() + abs(angle - cur) / velo # type: ignore + waited_time = timeout - 10.0 + while waited_time < timeout and abs(oget.get() - angle) > 0.01: # type: ignore + poll(0.05) + if not self.is_ioc_ok(): + raise RuntimeError("Aerotech IOC has an error status.") + if not self.is_omega_ok(): + raise RuntimeError("Aerotech OMEGA axis has an error status.") + waited_time = time.time() + if waited_time > timeout: + raise RuntimeError("Timeout on OMEGA movement.") + vel.put(ovelo, wait=True) + + def set_shutter(self, state): + if self.__pv_axes_mode.get(): + print("ABR is not in direct mode; cannot manipulate shutter") + return False + state = str(state).lower() + if state not in ["1", "0", "closed", "open"]: + print("unknown shutter state requested") + return None + elif state in ["1", "open"]: + state = 1 + elif state == ["0", "closed"]: + state = 0 + self.__pv_shutter_set.put(state, wait=True) + return state == self.__pv_shutter_get.get() + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="aerotech client") + parser.add_argument( + "--stop", help="stop/reset/restart aerotech", action="store_true" + ) + parser.add_argument( + "--sanitize-speeds", + help="reset speeds on axis to sane defaults", + action="store_true", + ) + parser.add_argument( + "--measure-standard", help="measure a standard dataset", action="store" + ) + parser.add_argument("--omega", help="move omega to...", action="store") + parser.add_argument("--gmx", help="move gmx to...", action="store") + parser.add_argument("--gmy", help="move gmy to...", action="store") + parser.add_argument("--gmz", help="move gmz to...", action="store") + args = parser.parse_args() + + abr = Abr() + if args.stop: + abr.stop() + + elif args.sanitize_speeds: + abr.velo(AXIS_GMX, 50.0) + abr.velo(AXIS_GMY, 10.0) + abr.velo(AXIS_GMZ, 10.0) + abr.velo(AXIS_OMEGA, 180.0) + + elif args.gmx: + abr.move_x(args.gmx) + + elif args.measure_standard: + start_angle, total_range, total_time = args.measure_standard.split(",") + abr.measure_standard(start_angle, total_range, total_time) diff --git a/src/mxlibs3/area_detector.py b/src/mxlibs3/area_detector.py new file mode 100644 index 00000000..1c2de069 --- /dev/null +++ b/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) diff --git a/src/mxlibs3/enum_pv.py b/src/mxlibs3/enum_pv.py new file mode 100755 index 00000000..ce500b47 --- /dev/null +++ b/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}]") diff --git a/src/mxlibs3/mx_lib.py b/src/mxlibs3/mx_lib.py new file mode 100644 index 00000000..9a0e30cc --- /dev/null +++ b/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}") diff --git a/src/mxlibs3/my_motor.py b/src/mxlibs3/my_motor.py new file mode 100644 index 00000000..3727b3d9 --- /dev/null +++ b/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 diff --git a/src/mxlibs3/non_standard.py b/src/mxlibs3/non_standard.py new file mode 100644 index 00000000..02dbba57 --- /dev/null +++ b/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") diff --git a/src/mxlibs3/pshell_client.py b/src/mxlibs3/pshell_client.py new file mode 100644 index 00000000..aaa223bd --- /dev/null +++ b/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=""): + """Returns help or auto-completion strings. + + Args: + input(optional, str): - ":" for control commands + - "" 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)) diff --git a/src/mxlibs3/smargon.py b/src/mxlibs3/smargon.py new file mode 100644 index 00000000..53b786e2 --- /dev/null +++ b/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) diff --git a/src/mxlibs3/tell_client.py b/src/mxlibs3/tell_client.py new file mode 100755 index 00000000..97e6c842 --- /dev/null +++ b/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 diff --git a/src/mxlibs3/workflow_tools.py b/src/mxlibs3/workflow_tools.py new file mode 100755 index 00000000..494412df --- /dev/null +++ b/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) diff --git a/src/pyproject.toml b/src/pyproject.toml new file mode 100644 index 00000000..3d26412a --- /dev/null +++ b/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"