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

1390 lines
57 KiB
Python

import copy
import json
import time
from datetime import datetime
from math import ceil
from typing import List, Tuple, Optional, Callable
import secrets
import os
import cv2
import numpy as np
import redis
from scipy import ndimage
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,
ZoomModeEnum)
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
print(f"{self.__cfg.state}")
if self.__cfg.state == BeamlineStateEnum.BeamLocation:
self.__cfg.zoom_mode = ZoomModeEnum.BeamLocation
else:
self.__cfg.zoom_mode = ZoomModeEnum.User
zoom_settings = self.__cfg.zoom_settings.get_camera_settings(val)
print(f"Setting zoom {val} {zoom_settings}")
self.samcam_settings = zoom_settings
#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}")
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)
self.__cfg.state_busy = False
#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 rebuild_array_from_scan_results(self,
scan_results: List,
value_field: str,
array_shape: Optional[tuple] = None,
nx_field: str = 'nx',
ny_field: str = 'ny',
default_value: float = 0.0,
threshold: Optional[float] = None,
condition_func: Optional[Callable] = None,
apply_filter_before: bool = True
) -> np.ndarray:
# Extract coordinates and values
positions = []
values = []
for result in scan_results:
nx = getattr(result, nx_field)
ny = getattr(result, ny_field)
value = getattr(result, value_field)
# Skip if coordinates are None
if nx is None or ny is None:
continue
positions.append((int(nx), int(ny))) # Note: (row, col) = (ny, nx)
if not value:
value = 0.0
values.append(float(value))
if not positions:
raise ValueError("No valid positions found in scan results")
# Determine array shape
if array_shape is None:
max_row = max(pos[0] for pos in positions)
max_col = max(pos[1] for pos in positions)
array_shape = (max_row + 1, max_col + 1)
# Initialize array with default values
result_array = np.full(array_shape, default_value, dtype=float)
# Apply pre-filtering if requested
if apply_filter_before:
filtered_data = []
for pos, val in zip(positions, values):
keep_value = True
# Apply threshold filter
if threshold is not None and val < threshold:
keep_value = False
# Apply custom condition
if condition_func is not None and not condition_func(val):
keep_value = False
if keep_value:
filtered_data.append((pos, val))
else:
filtered_data.append((pos, 0.0))
# Fill array with filtered values
for pos, val in filtered_data:
if 0 <= pos[0] < array_shape[0] and 0 <= pos[1] < array_shape[1]:
result_array[pos[0], pos[1]] = val
else:
# Fill array first, then apply filters
for pos, val in zip(positions, values):
if 0 <= pos[0] < array_shape[0] and 0 <= pos[1] < array_shape[1]:
result_array[pos[0], pos[1]] = val
# Apply post-filtering
if threshold is not None:
result_array[result_array < threshold] = 0.0
if condition_func is not None:
mask = np.vectorize(condition_func)(result_array)
result_array[~mask] = 0.0
return result_array
def create_quality_filtered_array(self,
scan_results: List,
value_field: str,
min_spots: Optional[int] = None,
min_efficiency: Optional[float] = 1.0,
min_background: Optional[float] = None,
exclude_ice: Optional[bool] = True,
**kwargs
) -> np.ndarray:
"""
Create array with comprehensive quality filtering
"""
def quality_condition(result, min_bkg, min_spots, min_efficiency):
if exclude_ice and (result.spots_ice / max(result.spots_low_res, 1.0)) == 1.0:
#print(f"all ice for {result.number}")
return False
# if (result.spots_ice / result.spots_low_res) > 5.0:
# return False
if exclude_ice and result.spots_ice > result.spots * 0.8: # More than 50% ice
#print(f"more than 80% ice for {result.number}")
return False
if result.index:
#print(f"index is True for {result.number}")
return True
if result.spots < min_spots:
# print(f"{result.spots} is less than {min_spots} for {result.number}")
return False
if result.spots_low_res < min_background:
return False
if result.efficiency < min_efficiency:
return False
return True
# Filter results first
filtered_results = []
if min_spots is None:
min_spots = min((result.spots for result in scan_results if result.spots is not None), default=1)
if min_background is None:
min_background = min((result.bkg for result in scan_results if result.bkg is not None), default=1)
if min_efficiency is None:
min_efficiency = 1.0
for result in scan_results:
if result.nx is not None and result.ny is not None:
if quality_condition(result, min_bkg=min_background, min_spots=min_spots,
min_efficiency=min_efficiency):
filtered_results.append(result)
else:
# Create a copy with zero value for filtered positions
import copy
zero_result = copy.copy(result)
setattr(zero_result, value_field, 0)
filtered_results.append(zero_result)
return self.rebuild_array_from_scan_results(filtered_results, value_field, **kwargs)
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 = 50
grid.file_prefix = f"{old_prefix}_{grid.omega_deg}deg"
grid.grid_size_mm = Coordinate(x=geom.beam_size_mm.x, y=geom.beam_size_mm.y * 0.5)
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)
images = result.images
output_data = {
'timestamp': time.ctime(),
'scan_results': [result.model_dump() for result in images],
'total_results': len([result for result in images])
}
if self.sample is not None and self.sample.db_id is not None:
with open(f'{self.sample.db_id}_scan_results.json', 'w') as f:
json.dump(output_data, f, indent=2)
# if any(img.index for img in images):
# print("COM by indexed spots")
# result_array= self.create_quality_filtered_array(images, 'spots_indexed', min_spots=None,
# min_efficiency=1.0, min_background=None)
# else:
print("COM by low res spots")
result_array = self.create_quality_filtered_array(images, 'spots_low_res', min_spots=None,
min_efficiency=1.0, min_background=None)
if r.n_x == 1:
print('vertical scan')
max_image = max(images, key=lambda img: img.spots_low_res)
new_y = max_image.ny
com = (0, new_y)
else:
print('horizontal scan')
com = ndimage.center_of_mass(result_array)
print(f"Center of mass: {com}")
if np.isnan(com[1]) or np.isnan(com[0]):
print("Center of mass is nan")
com = None
grid_mm_x=None
grid_mm_y=None
else:
grid_mm_x = com[0] * r.grid_size_mm.x
grid_mm_y = com[1] * r.grid_size_mm.y
if r.n_x == 1:
grid_mm_x += (0.5 * r.grid_size_mm.x)
for image in images:
if image.nx == round(com[0]) and image.ny == round(com[1]):
print(f"com found for image: {image.number}")
if com is not None and grid_mm_x is not None and grid_mm_y is not None:
new_delta_mm = self.sample_geometry.smargon_nudge(Coordinate(x=grid_mm_x, y=grid_mm_y))
print(f"new delta mm: {new_delta_mm}, new grid x: {grid_mm_x}, new grid y: {grid_mm_y}")
else:
print(f"using old method as COM is none or nan")
new_delta_mm = self.identify_crystal_raster(result, r)
if new_delta_mm is not None:
print(f'{time.ctime()}, moving SMARGON to target new delta mm {r.smargon.sh_mm + new_delta_mm} mm')
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
self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter
zoom_settings = self.__cfg.zoom_settings.z
for zoom_value in zoom_settings:
exposure = zoom_settings[zoom_value].exposure
gain = zoom_settings[zoom_value].gain
self.__devs.samcam_settings = SampleCameraSettings(exposure=exposure, gain=gain)
self.__devs.zoom_sync(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(zoom_value, exposure, gain, image)
bgr_array = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imwrite(f"bkg{zoom_value:.0f}_{exposure*1000:.0f}_{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 val < self.__devs.dtz_low or val > self.__devs.dtz_high:
self.__cfg.state_busy = False
raise RuntimeError(f"dtz={val} outside limits {self.__devs.dtz_low} to {self.__devs.dtz_high}")
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, 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.exposure, s.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"curr_image.jpg", gray_with_feature)
cv2.imwrite(f"diff.jpg", diff_image)
cv2.imwrite(f"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
self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter
zoom_settings = self.__cfg.zoom_settings.z
for zoom_value in zoom_settings:
exposure = zoom_settings[zoom_value].exposure
gain = zoom_settings[zoom_value].gain
max_attempt = 2
attempt = 0
found = 0
angles = (0, -45, -90)
while attempt < max_attempt:
print('new loop center settings')
self.__devs.samcam_settings = SampleCameraSettings(exposure=exposure, gain=gain)
self.__devs.zoom_sync(zoom_value)
print(f'zoom={zoom_value},gain={gain}, exp={exposure}')
if sample_id is not None:
self.save_screenshot_db(sample_id, f"pre_alc")
print(sample_id)
found_flag = False
found_angle = None
targets_found_this_attempt = 0
exp = int(exposure * 1000)
gn = int(gain)
for loop, angle in enumerate(angles):
print(f"Moving to new omega: {angle}")
if sample_id is not None and i == 0:
self.save_screenshot_db(sample_id, f"pre_alc_{sample_id}_{angle}_{zoom_value:.0f}_{exp}_{gn}")
self.__devs.aerotech.move(angle, wait=True)
time.sleep(1)
filename = None
if sample_id is not None:
filename = f"{sample_id}_{angle}_{zoom_value:.0f}_{exp}_{gn}"
target = self.__ml_loop_centre_box(sample_id, filename)
if target is not None:
found_flag = True
found_angle = angle
found += 1
targets_found_this_attempt += 1
self.__devs.smargon.target = target
self.__devs.smargon.wait(60)
if sample_id is not None:
print(sample_id)
self.__aare.sample_centered(self.__cfg.current_sample)
time.sleep(0.1)
self.save_screenshot_db(sample_id, f"{sample_id}_{angle}_{zoom_value:.0f}_{exp}_{gn}")
if targets_found_this_attempt == 0:
raise LoopCenteringFailed
else:
if found >= 3 or targets_found_this_attempt >= 3:
print(f"sucessfully found {found} or {targets_found_this_attempt} targets in {attempt} attempts")
return True
if found_flag and found_angle is not None:
print(f"found a target at angle {found_angle} in attempt {attempt}")
angles = (found_angle, found_angle + 45, found_angle + 90)
attempt += 1
print(f"attempt {attempt} of {max_attempt}")
if attempt >= max_attempt:
raise LoopCenteringFailed
#i += 1
print("alc success")
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):
self.__aare.alc_failed(self.sample)
raise LoopCenteringFailed
self.__cfg.state_busy = False
except Exception:
self.__cfg.zoom_mode = ZoomModeEnum.User
self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom)
self.__cfg.state_busy = False
raise
self.__cfg.zoom_mode = ZoomModeEnum.User
self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom)
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")
low_resolution_flag=False #False for high resolution
if low_resolution_flag:
det_dist = 300.0
else:
det_dist = 110.0
if self.__auto_center(RasterGridRequest(
exp_time_s=0.02,
file_prefix=sample_prefix + f"_{hex_string}",
smargon= SmargonCoordinate(),
n_x=1,
n_y=1,
dtz=det_dist,
grid_size_mm=Coordinate(x=geom.beam_size_mm.x * 0.5, y=geom.beam_size_mm.y * 0.5),
omega_deg=0
)):
self.__cfg.zoom_mode = ZoomModeEnum.User
self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom)
self.__rotation( RotationScanRequest(start_omega_deg=0,
dtz=det_dist,
file_prefix=sample_prefix + f"_{hex_string}",
exp_time_s=0.04,
incr_omega_deg=0.2,
steps=900,))
else:
self.__aare.axc_failed(sample)
self.zoom = 1
self.__cfg.zoom_mode = ZoomModeEnum.User
self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom)
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 not self.__cfg.state_busy:
raise Exception("Beamline should be busy")
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,
poni_rot1_rad=-0.001396263,
poni_rot2_rad=-0.003839724,
)
@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.full_flux,
sample_camera=self.__devs.samcam_settings,
name=self.__bl,
transmission=self.__devs.transmission.get(),
zoom=self.__devs.zoom,
commissioning_mode=self.__cfg.commissioning_mode,
dtz_min=self.__devs.dtz_low,
dtz_max=self.__devs.dtz_high,
)
@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