Files
AareDAQ/daq/src/aaredaq/daq.py
T

1130 lines
46 KiB
Python

import copy
import time
from datetime import datetime
from math import ceil
from typing import List, Tuple
import secrets
import os
import cv2
import numpy as np
import redis
from aaredaq import workflows
from aaredaq.aaredb import AareWrapper
from aaredaq.autofocus import calculate_focus_measure
from aaredaq.config import BeamlineConfig, ABR_POS_MOUNT
from aaredaq.config import BeamlineStateEnum
from aaredaq.devices import BeamlineDevices
from aaredaq.mlbox import MlBox
from aaredaqlib.beamline import MXBeamline
from aaredaqlib.coordinate import Coordinate, SmargonCoordinate
from aaredaqlib.diffraction_geometry import DiffractionGeometry
from aaredaqlib.models import (
SampleShortInfo,
PuckLoadedInfo,
SampleShortInfoList,
DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, AutofocusSettings, BoundingBoxModel,
LoopCenteringZoomModelElem)
from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid, CompletedRasterGridElem
from aaredaqlib.rotation_scan import RotationScanRequest, CompletedRotationScan
from aaredaqlib.sample_geometry import SampleGeometryModel
from mxlibs3.jfjoch import JFJochWrapper
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()
self.__jfjoch = JFJochWrapper(bl)
self.__bl = bl.value.upper()
self.__aare = AareWrapper(bl)
self.__saved_box = None
@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)
self.__saved_box = None
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.__saved_box = None
self.__devs.zoom = val
#if self.__cfg.state == BeamlineStateEnum.BeamLocation:
# self.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(val)
self.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(val)
#elf.__cfg.state_busy = False
@property
def light(self) -> float:
val = self.__devs.lamp_light
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 = conv
@property
def sample(self) -> SampleShortInfo | None:
if self.__cfg.current_sample is not None and self.__cfg.current_sample.location is None:
return self.__cfg.current_sample
if self.__devs.tell.get_mounted_sample() is None:
return None
return self.__cfg.current_sample
@property
def samcam_settings(self) -> SampleCameraSettings:
return self.__devs.samcam_settings
@samcam_settings.setter
def samcam_settings(self, s: SampleCameraSettings):
self.__devs.samcam_settings = s
def autofocus(self, f: AutofocusSettings):
pos = self.__cfg.abr_meas_pos
z_min = pos.z - f.z_range_um / 2.0
z_max = pos.z + f.z_range_um / 2.0
measures = []
for i in range(f.z_steps):
pos.z = z_min + i * (z_max - z_min) / f.z_steps
self.__devs.abr_pos = pos
self.__devs.sample_cam.get_single_image()
image = self.__devs.sample_cam.get_image(gray=True)
val = calculate_focus_measure(image, f.center_x_pxl, f.center_y_pxl, f.radius_pxl)
measures.append((pos.z, val))
best_pos, _ = max(measures, key=lambda x: x[1])
pos.z = best_pos
self.__devs.abr_pos = pos
self.__cfg.abr_meas_pos = pos
self.__devs.sample_cam.collect_auto()
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.__saved_box = None
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
self.__cfg.state_busy = False
except:
self.__cfg.state_busy = False
raise
def goto_abr_meas_pos(self):
self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
try:
self.__devs.abr_pos = self.__cfg.abr_meas_pos
self.__cfg.state_busy = False
except:
self.__cfg.state_busy = False
raise
def create_sample(self, target: SampleShortInfo):
self.__cfg.try_set_busy(timeout=360)
try:
curr_sample = self.__cfg.current_sample
if curr_sample is not None and curr_sample.location is not None:
raise Exception("Sample from TELL is loaded")
self.__aare.create_manual_sample(target)
self.__cfg.current_sample = target
self.__cfg.state_busy = False
except:
self.__cfg.state_busy = False
raise
def __mount(self, target: SampleShortInfo | None):
self.__set_state(BeamlineStateEnum.RobotSampleExchange)
self.__saved_box = None
print(f"{time.ctime()} moving smargon to home")
self.__devs.smargon.move_home(wait=True)
print(f"{time.ctime()} moving aerotech to mount position")
self.__devs.abr_pos = ABR_POS_MOUNT
time.sleep(0.5)
print(f"{time.ctime()} checking beamstop")
if self.__devs.bsz.value < 24.0:
raise Exception("Beamstop Z below 24.0 mm - potentially unsafe with mounting")
print(f"{time.ctime()} checking magnet postion sensor positon: {self.__devs.magnet_position_sensor_readout.value}")
print(f"{time.ctime()} checking smargon position: {self.__devs.smargon.readback}")
print(f"{time.ctime()} checking abr position: {self.__devs.abr_pos}")
if self.__devs.magnet_position_sensor.value != 0:
time.sleep(5)
print("!!!!!!!!!!!!!!!!!!!!!MAGNET CONTROLLER BROKE AGAIN!!!!!!!!!!!!!!!!!!")
print(f"{time.ctime()} checking magnet position sensor positon: {self.__devs.magnet_position_sensor_readout.value}")
print(f"{time.ctime()} checking smargon position: {self.__devs.smargon.readback}")
print(f"{time.ctime()} checking abr position: {self.__devs.abr_pos}")
if self.__devs.magnet_position_sensor.value != 0:
print(time.ctime())
start = time.time()
end = start + 360
while self.__devs.magnet_position_sensor.value != 0:
time.sleep(1)
print(f"checking magnet position sensor positon: {self.__devs.magnet_position_sensor_readout.value}")
print(f"{time.ctime()} checking smargon position: {self.__devs.smargon.readback}")
print(f"{time.ctime()} checking abr position: {self.__devs.abr_pos}")
if time.time() > end:
raise Exception(f"Magnet position sensor is not in position: {self.__devs.magnet_position_sensor_readout.value}")
#Reenable TELL after doors locked
print(f"{time.ctime()} enable tell motion after doors locked")
self.__devs.tell.check_enable_motion()
#self.__devs.tell.wait_mount_complete()
if self.__devs.tell.is_in_park():
print(f"{time.ctime()} moving tell to cold")
self.__devs.tell.move_cold(wait=True)
time.sleep(30)
# self.__devs.tell.check_enable_motion()
print(f"{time.ctime()} moving tell to mount position")
self.__devs.tell.set_in_mount_position(True)
curr_sample = self.__cfg.current_sample
print(f"{time.ctime()} waiting for tell to be ready")
self.__devs.tell.wait_ready()
if curr_sample is not None and curr_sample.location is None:
print(f"{time.ctime()} unmounting current sample")
# TODO: Smarter way to know manual sample has be removed from goniometer
self.__aare.sample_unmounted(curr_sample)
self.__cfg.current_sample = None
elif target is None:
print(f"{time.ctime()} moving tell to unmount with no target")
self.__devs.tell.unmount(wait=True, timeout=360)
self.__aare.sample_unmounted(curr_sample)
self.__cfg.current_sample = None
else:
# reset zoom
self.zoom = 280
# mount sample
print(f"{time.ctime()} moving tell to unmount")
self.__aare.sample_unmounted(curr_sample)
print(f"{time.ctime()} moving tell to mount")
self.__devs.tell.mount(
address=target.tell_address(),
force=True,
auto_unmount=True,
read_dm=False,
wait=True,
timeout=360,
)
print(f"{time.ctime()} post tell mount, pre db input")
self.__aare.sample_mounted(target)
self.__cfg.current_sample = target
if target is not None and target.db_id is not None:
self.save_screenshot_db(target.db_id, f"sample_mounted")
# TODO: Need to know if dry needs to happen
@sample.setter
def sample(self, target: SampleShortInfo | None):
# This will set internally state to SampleExchange, but will return to SampleAlignment
# before exiting
self.__cfg.try_set_busy(timeout=360)
try:
curr_sample = self.__cfg.current_sample
curr_sample_is_manual = False
if curr_sample is not None and curr_sample.location is None:
self.__cfg.current_sample = None
curr_sample_is_manual = True
if target is not None or not curr_sample_is_manual:
self.__mount(target)
if target is not None and target.db_id is not None:
#self.__aare.sample_mounted(self.__cfg.current_sample)
self.save_screenshot_db(target.db_id , f"post_mount_{target.db_id }_{self.__devs.zoom}")
self.__cfg.state_busy = False
except Exception as e:
self.__cfg.state_busy = False
# If unmount succeeded, but mount failed
mounted_sample = self.__devs.tell.get_mounted_sample()
if target is not None:
print(f"TELL: Error mounting, TELL believes current sample is {mounted_sample}, while expected {target.location}")
if mounted_sample is None:
print("TELL: After mounting there is no sample on gonio according to TELL")
self.__cfg.current_sample = None
if target is not None:
raise e
elif (target.location is not None
and mounted_sample.puck.pos == target.location.pos
and mounted_sample.puck.segment == target.location.segment
and mounted_sample.pin == target.pin):
print("TELL: Actually mounting was correct according to TELL, so setting current sample to target")
self.__cfg.current_sample = target
else:
raise e
self.__set_state(BeamlineStateEnum.SampleAlignment)
if target is not None and target.location is not None:
self.__loop_center_sequence(target.db_id)
@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 identify_crystal_raster(self, result, r: RasterGridRequest):
images = result.images
if images and any(img.spots > 0 for img in images):
indexed_images = [img for img in images if img.index and img.spots_low_res > 4 and img.bkg > 4.5]
if indexed_images:
#indexed_images = [img for img in indexed_images if img.spots_indexed > 10]
images = indexed_images
filtered_images = [img for img in images
if img.spots_ice is not None and img.spots_low_res > 4 and (
img.spots_ice / img.spots_low_res) < 5.0
and (img.spots_ice / img.spots_low_res) != 1]
if indexed_images:
print(f"Find image by maximum number of spots indexed")
max_image = max(images, key=lambda img: img.spots_indexed)
max_spots = max_image.spots_indexed
max_images = [img for img in images if img.spots_indexed == max_spots]
max_image = max_images[len(max_images) // 2]
print(f"Image with maximum spots_low_res: {max_image}")
print(f"Maximum spots_indexed value: {max_image.spots_indexed}")
print(f"Maximum spots_low_res value: {max_image.spots_low_res}")
else:
print(f"Find image by maximum number of low resolution spots")
max_image = max(images, key=lambda img: img.spots_low_res)
print(f"Image with maximum spots_low_res: {max_image}")
print(f"Maximum spots_low_res value: {max_image.spots_low_res}")
grid_mm_x = max_image.nx * r.grid_size_mm.x
grid_mm_y = max_image.ny * r.grid_size_mm.y
print(f"Grid coordinates in mm: ({grid_mm_x}, {grid_mm_y})")
delta_mm = self.sample_geometry.smargon_nudge(Coordinate(x=grid_mm_x, y=grid_mm_y))
return delta_mm
else:
return None
def __auto_center(self, grid: RasterGridRequest) -> CompletedRasterGrid | None:
sample = self.sample
if sample is None:
raise Exception("Sample must be mounted to auto center")
old_prefix = grid.file_prefix
geom = self.sample_geometry
r = self.__ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg:.2f}deg")
if r is None:
self.__devs.aerotech.move(geom.omega_deg + 90.0, wait=True)
time.sleep(0.2)
r = self.__ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg + 90.0:.2f}deg")
if r is not None:
geom = self.sample_geometry
grid.smargon = r.smargon
grid.grid_size_mm = r.grid_size_mm
grid.n_x = r.n_x
grid.n_y = r.n_y
grid.file_prefix = f"{old_prefix}_{grid.omega_deg}deg"
grid.omega_deg = geom.omega_deg
res1 = self.__raster(grid)
grid.omega_deg += 90
self.__devs.aerotech.move(grid.omega_deg, wait=True)
grid.n_x = 1
grid.n_y = 100
grid.file_prefix = f"{old_prefix}_{grid.omega_deg}deg"
grid.grid_size_mm = Coordinate(x=geom.beam_size_mm.x * 0.8, y=geom.beam_size_mm.y * 0.2)
offset = Coordinate(x=0, y=-grid.n_y * grid.grid_size_mm.y / 2.0)
geom = self.sample_geometry
grid.smargon.sh_mm = geom.smargon.sh_mm + geom.smargon_nudge(offset)
res2 = self.__raster(grid)
return CompletedRasterGrid(r=[res1, res2])
else:
return None
def __raster(self, r: RasterGridRequest) -> CompletedRasterGridElem:
max_time = r.exp_time_s * r.n_y * r.n_x + 60
if r.dtz is not None:
self.__cfg.dtz = r.dtz
self.__set_state(BeamlineStateEnum.DataCollection)
save_smargon_position = self.__devs.smargon.readback
if r.n_x > 1:
delta_mm = self.sample_geometry.smargon_nudge(Coordinate(x=0, y=r.grid_size_mm.y / 2))
else:
delta_mm = self.sample_geometry.smargon_nudge(Coordinate(x=r.grid_size_mm.x / 2, y=0))
self.__devs.smargon.target = SmargonCoordinate(sh_mm=r.smargon.sh_mm + delta_mm,
phi_deg=r.smargon.phi_deg,
chi_deg=r.smargon.chi_deg)
if r.transmission is not None:
self.__devs.transmission.set(r.transmission, wait=False)
self.__devs.aerotech.move(r.omega_deg, wait=True, speed=360.0)
self.__devs.transmission.wait()
self.__devs.smargon.wait()
status = self.status
self.__jfjoch.measure_raster(r, status)
self.__aare.create_gridscan_run(self.sample, r, status)
if self.sample is not None and self.sample.db_id is not None:
self.save_screenshot_db(self.sample.db_id, "before_raster")
if r.n_x == 1:
self.__devs.aerotech.measure_vertical_line(r.exp_time_s, r.grid_size_mm.y, r.n_y)
else:
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.abr_pos = self.__cfg.abr_meas_pos
result = self.__jfjoch.wait_till_done(60)
#if self.sample is not None and self.sample.db_id is not None:
# self.__aare.ingest_gridscan(self.sample, result, r)
new_delta_mm = self.identify_crystal_raster(result, r)
if new_delta_mm is not None:
self.__devs.smargon.target = SmargonCoordinate(sh_mm=r.smargon.sh_mm + new_delta_mm,
phi_deg=r.smargon.phi_deg,
chi_deg=r.smargon.chi_deg)
else:
print("Auto finding optimal image failed due to no images found. Using previous position.")
self.__devs.smargon.target = save_smargon_position
return CompletedRasterGridElem(request=copy.deepcopy(r), result=result)
def measure_raster(self, r: RasterGridRequest, auto: bool) -> CompletedRasterGrid:
self.__cfg.try_set_busy(timeout=ceil(360))
try:
if auto:
result = self.__auto_center(r)
else:
raster_result = self.__raster(r)
result = CompletedRasterGrid(r=[raster_result])
self.__set_state(BeamlineStateEnum.SampleAlignment)
self.__cfg.state_busy = False
return result
except Exception as e:
self.__set_state(BeamlineStateEnum.SampleAlignment)
self.__cfg.state_busy = False
raise e
def __rotation(self, request: RotationScanRequest) -> CompletedRotationScan:
if request.dtz is not None:
self.__cfg.dtz = request.dtz
self.save_screenshot_db(self.sample.db_id, f"pre_rotation{self.sample.db_id}_{self.__devs.zoom}")
self.__set_state(BeamlineStateEnum.DataCollection)
if request.transmission is not None:
self.__devs.transmission.set(request.transmission, wait=False)
if request.start is not None:
self.__devs.smargon.target = request.start
self.__devs.transmission.wait()
if request.start is not None:
self.__devs.smargon.wait()
status = self.status
self.__jfjoch.measure_rotation(request, status)
self.__aare.create_rotation_run(self.sample, request, status)
if self.sample is not None and self.sample.db_id is not None:
self.save_screenshot_db(self.sample.db_id, "before_dc")
if request.screening:
self.__devs.aerotech.measure_screening(request.start_omega_deg,
request.wedge_omega_deg,
request.exp_time_s,
request.incr_omega_deg,
request.steps)
self.__devs.aerotech.wait_scan_done(request.exp_time_s * request.steps + 60)
else:
total_omega = request.incr_omega_deg * request.steps
total_time = request.exp_time_s * request.steps
self.__devs.aerotech.measure_standard(
request.start_omega_deg, total_omega, total_time
)
if request.start is not None and request.end is not None:
smargon_time_step = request.time_sec / float(request.steps)
pos_step = (request.end.sh_mm - request.start.sh_mm) * (1.0 / float(request.steps))
for i in range(request.steps):
self.__devs.smargon.target = SmargonCoordinate(
sh_mm=request.start.sh_mm + pos_step * i
)
time.sleep(smargon_time_step)
self.__devs.aerotech.wait_scan_done(total_time + 60)
self.__devs.aerotech.reset()
result = self.__jfjoch.wait_till_done(60)
self.__aare.sample_collected(self.sample)
if self.sample is not None and self.sample.db_id is not None:
self.save_screenshot_db(self.sample.db_id, "after_dc")
return CompletedRotationScan(request=copy.deepcopy(request), result=result)
def measure_rotation(self, request: RotationScanRequest) -> CompletedRotationScan:
total_time = request.exp_time_s * request.steps
self.__cfg.try_set_busy(timeout=ceil(total_time + 360))
try:
result = self.__rotation(request)
self.__set_state(BeamlineStateEnum.SampleAlignment)
self.__cfg.state_busy = False
return result
except Exception as e:
self.__set_state(BeamlineStateEnum.SampleAlignment)
self.__cfg.state_busy = False
raise e
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 = 2.5
for s in self.__cfg.alc_zoom_settings.z:
self.__devs.samcam_settings = SampleCameraSettings(exposure=s.sam_cam_exp, gain=s.sam_cam_gain)
self.__devs.zoom_sync(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, s.sam_cam_exp, s.sam_cam_gain, image)
bgr_array = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imwrite(f"bkg{s.zoom_value:.0f}_{s.sam_cam_exp*1000:.0f}_{s.sam_cam_gain:.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, wait=False)
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.__saved_box = None
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.BeamLocation)
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.BeamLocation)
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
try:
aerotech_coord = 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,
)
sample_geom = SampleGeometryModel(
beam_location_pxl=self.__cfg.beam_mark_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=self.__cfg.beam_size_mm,
aerotech=aerotech_coord,
aerotech_meas=self.__devs.abr_pos
)
return sample_geom
except Exception as e:
try:
test_x = self.__devs.gmx.value - abr_pos.x
print(f"Error creating sample geometry model: {e}")
print(f"Error not ABR, check smargon then restart server")
raise Exception(f"Error getting sample geometry {e}")
except:
print(f"Error getting sample geometry: {e}")
print(f"!!!!! ABR MAY HAVE DISCONNECTED !!!!!")
print(f"!!!!! CHECK ABR EPICS PANEL FOR ERRORS!!!!!")
print("!!!!! IF NOT RESTART GUI !!!!!")
return SampleGeometryModel(
beam_location_pxl=self.__cfg.beam_mark_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=self.__cfg.beam_size_mm,
aerotech=Coordinate(x=0.0, y=0.0, z=0.0),
aerotech_meas=self.__devs.abr_pos
)
@property
def beam_center(self) -> Tuple[float, float]:
return self.__cfg.beam_center
@beam_center.setter
def beam_center(self, val: Tuple[float, float]):
self.__cfg.beam_center = val
@property
def beam_size_mm(self) -> Coordinate:
return self.__cfg.beam_size_mm
@beam_size_mm.setter
def beam_size_mm(self, val: Coordinate):
self.__cfg.beam_size_mm = val
def get_beam_mark(self):
return self.__cfg.get_beam_mark(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) -> RasterGridRequest | None:
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)
self.__aare.upload_image(sample_id, f"{filename}_no_detection", 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)
self.__aare.upload_image(sample_id, filename, bgr_image)
geom = self.sample_geometry
start_coord = geom.picture_to_smargon(Coordinate(x=x1, y=y1))
grid_size = Coordinate(x=geom.beam_size_mm.x * 0.8, y=geom.beam_size_mm.y * 0.8)
n_x = abs(ceil((x2 - x1) * geom.pixel_in_mm / grid_size.x))
n_y = abs(ceil((y2 - y1) * geom.pixel_in_mm / grid_size.y))
return RasterGridRequest(
exp_time_s=0.02,
transmission=1.0,
smargon= SmargonCoordinate(chi_deg = geom.smargon.chi_deg,
phi_deg= geom.smargon.phi_deg,
sh_mm=start_coord),
n_x=n_x,
n_y=n_y,
grid_size_mm= grid_size,
omega_deg=geom.omega_deg
)
def __ml_loop_centre_box(self, sample_id: int | None = None, filename: str | None = None) -> SmargonCoordinate | None:
time.sleep(0.2) # Just to be sure image is stable
bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR)
box = self.__mlbox.predict(bgr_image, filename)
if box is None:
if filename is not None:
#cv2.imwrite(f"{filename}_no_detection.jpg", bgr_image)
self.__aare.upload_image(sample_id, f"{filename}_no_detection", 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)
self.__aare.upload_image(sample_id, filename, bgr_image)
geom = self.sample_geometry
centre_coord = y1 + (y2 - y1)/2
coord = geom.picture_to_smargon(Coordinate(x=x1, y=centre_coord))
return SmargonCoordinate(sh_mm=coord)
def ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None) -> RasterGridRequest | None:
try:
self.__cfg.try_set_busy(timeout=360)
r = self.__ml_bounding_box(sample_id, filename)
self.__cfg.state_busy = False
return r
except Exception:
self.__cfg.state_busy = False
raise
def __loop_center(self, s: LoopCenteringZoomModelElem, filename: str | None = "", iteration: int = 0) -> SmargonCoordinate:
# TODO: tidy up!
curr_image = self.camera_image
gray_without_feature = cv2.cvtColor(
self.__cfg.get_alc_bkg(self.zoom, s.sam_cam_exp, s.sam_cam_gain), 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)
thresh_value, thresh = cv2.threshold(diff_image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
print(f'Thresh value: {thresh_value}')
if filename is not None:
#cv2.imwrite(f"{filename}_curr_image_colour.jpg", curr_image)
cv2.imwrite(f"{filename}_curr_image.jpg", gray_with_feature)
cv2.imwrite(f"{filename}_diff.jpg", diff_image)
cv2.imwrite(f"{filename}_thresh.jpg", thresh)
#cv2.imwrite(f"{filename}_adaptive.tiff", adapt_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
if iteration == 0:
coord = geom.picture_to_smargon(
Coordinate(x=leftmost_point[0], y=leftmost_point[1])
)
else:
#TODO only move x after first iteration
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 = 2.5
try:
i = 0
for s in self.__cfg.alc_zoom_settings.z:
print('new loop center settings')
self.__devs.samcam_settings = SampleCameraSettings(exposure=s.sam_cam_exp, gain=s.sam_cam_gain)
self.__devs.zoom_sync(s.zoom_value)
print(f'zoom={s.zoom_value},gain={s.sam_cam_gain}, exp={s.sam_cam_exp}')
if sample_id is not None:
self.save_screenshot_db(sample_id, f"pre_alc")
print(sample_id)
if i % 2 == 0:
angles = (0, -45, -90, -135)
else:
angles = (-135, -90, -45, 0)
for angle in angles:
print(f"Moving to new omega: {angle}")
if sample_id is not None and i == 0:
exp = int(s.sam_cam_exp * 1000)
gain = int(s.sam_cam_gain)
self.save_screenshot_db(sample_id, f"pre_alc_{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}")
self.__devs.aerotech.move(angle, wait=True)
time.sleep(1)
filename = None
if sample_id is not None:
if not os.path.exists(f"{sample_id}_{s.zoom_value:.0f}"):
os.mkdir(f"{sample_id}_{s.zoom_value:.0f}")
if not os.path.exists(f"{sample_id}_{s.zoom_value:.0f}_negative"):
os.mkdir(f"{sample_id}_{s.zoom_value:.0f}_negative")
if angle > 0:
filename = f"{sample_id}_{s.zoom_value:.0f}/{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}"
else:
filename = f"{sample_id}_{s.zoom_value:.0f}_negative/{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}"
target = self.__ml_loop_centre_box(sample_id, filename)
if target is None:
continue
self.__devs.smargon.target = target #self.__loop_center(s, filename, iteration=i)
self.__devs.smargon.wait(60)
print(sample_id)
if sample_id is not None:
self.__aare.sample_centered(self.__cfg.current_sample)
time.sleep(0.1)
self.save_screenshot_db(sample_id, f"{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}")
i += 1
return True
except Exception as e:
print(f"Error in loop centering: {e}")
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.__aare.upload_image(sample_id, filename, bgr_image)
@property
def sample_spreadsheet(self) -> SampleShortInfoList:
return self.__cfg.spreadsheet
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()
formatted_date = datetime.now().strftime('%Y%m%d')
sample_prefix = "{}/{}/{:02d}/{}".format(
formatted_date,
sample.puck_name,
sample.pin,
sample.sample_name
)
try:
self.__cfg.try_set_busy(timeout=360)
geom = self.sample_geometry
print(f"{time.ctime()} starting mount {sample.db_id}")
self.__mount(sample)
if not self.__loop_center_sequence(sample.db_id):
self.__aare.alc_failed(sample)
print("alc failed")
self.__cfg.state_busy = False
end = time.perf_counter()
return end - start
#raise LoopCenteringFailed
self.zoom = 500
time.sleep(0.5)
hex_string = secrets.token_hex(3) # 3 bytes = 6 hex characters
print(hex_string.upper())
# self.__ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg:.2f}deg")
if self.__auto_center(RasterGridRequest(
exp_time_s=0.01,
file_prefix=sample_prefix + f"_{hex_string}",
smargon= SmargonCoordinate(),
n_x=1,
n_y=1,
dtz=110.0,
grid_size_mm=Coordinate(x=geom.beam_size_mm.x * 0.8, y=geom.beam_size_mm.y * 0.8),
omega_deg=0
)):
self.__rotation( RotationScanRequest(start_omega_deg=0,
dtz=110.0,
file_prefix=sample_prefix + f"_{hex_string}",
exp_time_s=0.003,
incr_omega_deg=0.4,
steps=900,))
else:
self.__aare.axc_failed(sample)
self.zoom = 1
self.__cfg.state_busy = False
except Exception:
self.__aare.sample_failed(sample)
self.__cfg.state_busy = False
end = time.perf_counter()
return end - start
#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.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)
elif 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.rse2sa(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 shutter(self) -> bool:
return self.__devs.shutter
@shutter.setter
def shutter(self, v: bool):
self.__devs.shutter = v
@property
def diffraction_geometry(self) -> DiffractionGeometry:
det_cfg = self.__jfjoch.detector()
return DiffractionGeometry(
energy_keV=self.__devs.energy_kev,
dtz_mm=self.__devs.dtz.value,
detector_size_pxl=(det_cfg.width, det_cfg.height),
pixel_size_mm=det_cfg.pixel_size_mm,
beam_center_pxl=self.__cfg.beam_center,
detector_description=det_cfg.description,
detector_serial_number=det_cfg.serial_number
)
@property
def beamline_status(self) -> BeamlineStatus:
return BeamlineStatus(
ring_current_mA=self.__devs.ring_current,
light=self.light,
cryojet_K=self.__devs.cryojet_temp,
shutter_open=self.__devs.shutter,
flux_ph_s=self.__devs.flux,
sample_camera=self.__devs.samcam_settings,
name=self.__bl,
transmission=self.__devs.transmission.get(),
zoom=self.__devs.zoom
)
@property
def status(self) -> DAQStatusModel:
# session should be set by FastAPI server
return DAQStatusModel(
state=self.state,
busy=self.busy,
geom=self.sample_geometry,
bl=self.beamline_status,
sample=self.sample,
session=SessionStatus(),
diffraction=self.diffraction_geometry,
box=self.__saved_box
)
def cancel(self):
if self.__cfg.state == BeamlineStateEnum.DataCollection:
self.__devs.aerotech.stop()
self.__jfjoch.cancel()
def anneal(self, time_s: float):
self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
try:
self.__devs.anneal(time_s)
finally:
self.__cfg.state_busy = False