329 lines
12 KiB
Python
329 lines
12 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
|
|
import numpy as np
|
|
import redis
|
|
import redis_lock
|
|
from aaredaqlib.coordinate import Coordinate
|
|
from aaredaqlib.models import (
|
|
BeamlineSettingsModel,
|
|
BeamCenterCoeffModel,
|
|
LoopCenteringZoomModel,
|
|
def_loop_centering_zoom,
|
|
SampleShortInfo,
|
|
SessionStatus,
|
|
BeamlineStateEnum, SessionsStateEnum,
|
|
)
|
|
|
|
from aaredaqlib.beamline import MXBeamline
|
|
|
|
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"{self.__bl}-redis.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_status(self, session: int) -> SessionStatus:
|
|
return SessionStatus(session=self.session_state(session),
|
|
current_pgroup=self.pgroup)
|
|
|
|
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_mark(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)
|