1459 lines
59 KiB
Python
1459 lines
59 KiB
Python
import copy
|
|
import json
|
|
import math
|
|
import time
|
|
from datetime import datetime
|
|
from math import ceil
|
|
from typing import List, Tuple, Optional, Callable, Dict
|
|
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.logger_config import setup_logger
|
|
from aaredaqlib.find_xtal import raster_centre_of_mass, create_quality_filtered_array, identify_crystal_raster, \
|
|
get_result_list_from_com, get_best_b_factor, get_best_res, get_xtal_size
|
|
from aaredaqlib.models import (
|
|
SampleShortInfo,
|
|
PuckLoadedInfo,
|
|
SampleShortInfoList,
|
|
DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, AutofocusSettings, BoundingBoxModel,
|
|
ZoomModeEnum, CrystalSize, SimpleScanParameters)
|
|
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
|
|
|
|
logger = setup_logger(__name__, "/tmp/mxlogs")
|
|
|
|
class TransformationInvalidException(Exception):
|
|
def __init__(self, message="Transformation is not implemented"):
|
|
super().__init__(message)
|
|
self.message = message
|
|
logger.error(f"{message}", extra={"exception:" : Exception})
|
|
|
|
def __str__(self):
|
|
return self.message
|
|
|
|
|
|
class LoopCenteringFailed(Exception):
|
|
#logger.error(f"Loop centering failed", extra={"exception:" : 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
|
|
self.crystal_size = CrystalSize(x=0,y=0,z=0)
|
|
self.last_best_b_factor = None
|
|
self.last_best_res = None
|
|
self.auto_params = SimpleScanParameters()
|
|
|
|
@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:
|
|
logger.error(f"Cannot explicitly move to busy state", extra={"target":target, "state":self.__cfg.state})
|
|
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 smart_params(self) -> SimpleScanParameters:
|
|
return self.auto_params
|
|
|
|
@smart_params.setter
|
|
def smart_params(self, params: SimpleScanParameters):
|
|
self.auto_params = params
|
|
|
|
@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
|
|
logger.error("Omega has to be between -720 and 720 degrees")
|
|
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 = 1
|
|
# 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
|
|
|
|
self.crystal_size = CrystalSize(x=0, y=0, z=0)
|
|
self.last_best_b_factor = None
|
|
self.last_best_res = None
|
|
|
|
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 __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
|
|
logger.info(f"raster status: {status}")
|
|
logger.info(f"raster grid request: {r}")
|
|
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': [image.model_dump() for image in images],
|
|
'total_results': len([image for image in images])
|
|
}
|
|
if self.sample is not None and self.sample.db_id is not None:
|
|
if r.n_x == 1:
|
|
filename=f'{self.sample.db_id}_scan_results_vertical.json'
|
|
else:
|
|
filename=f'{self.sample.db_id}_scan_results_horizontal.json'
|
|
with open(filename, 'w') as f:
|
|
json.dump(output_data, f, indent=2)
|
|
print('before centre_of_mass')
|
|
print("COM by low res spots")
|
|
|
|
|
|
result_array = create_quality_filtered_array(images, 'spots_low_res', min_spots=None,
|
|
min_efficiency=1.0, min_background=None, min_low_res_spots=10.0)
|
|
self.crystal_size = get_xtal_size(self.crystal_size, result_array, r)
|
|
grid_mm_x, grid_mm_y, com = raster_centre_of_mass(result_array, r)
|
|
|
|
if grid_mm_x is None or grid_mm_y is None:
|
|
print(f"using old method as COM is none or nan")
|
|
grid_mm_x, grid_mm_y = identify_crystal_raster(result, r)
|
|
|
|
if grid_mm_x or grid_mm_y:
|
|
|
|
if r.n_x == 1:
|
|
new_delta_mm = self.sample_geometry.smargon_nudge(Coordinate(x=0, y=grid_mm_y))
|
|
|
|
else:
|
|
result_list = get_result_list_from_com(images, com)
|
|
print(result_list)
|
|
self.last_best_b_factor = get_best_b_factor(result_list)
|
|
self.last_best_res = get_best_res(result_list)
|
|
print("b_factor: ", self.last_best_b_factor, " best_res: ", self.last_best_res)
|
|
new_delta_mm = self.sample_geometry.smargon_nudge(Coordinate(x=grid_mm_x, y=grid_mm_y))
|
|
|
|
else:
|
|
new_delta_mm = None
|
|
|
|
print('after centre_of_mass')
|
|
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, pref_class=(3,0))
|
|
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
|
|
|
|
cls, 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
|
|
|
|
if cls == 0: # loop_all
|
|
centre_y = y1 + (y2 - y1)/2
|
|
centre_x = x1
|
|
elif cls == 1: # pin
|
|
centre_y = y1 + (y2 - y1)/2
|
|
centre_x = x1
|
|
elif cls == 2 or cls == 3: #crystal or loop_face
|
|
centre_y = y1 + (y2 - y1)/2
|
|
centre_x = x1 + (x2 - x1)/2
|
|
else:
|
|
print(f"unknown box class {cls}")
|
|
return None
|
|
|
|
coord = geom.picture_to_smargon(Coordinate(x=centre_x, y=centre_y))
|
|
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 box_height_from_tuple(self, box: Tuple[float, float, float, float]) -> float:
|
|
x1, y1, x2, y2 = box
|
|
return abs(y2 - y1)
|
|
|
|
def box_area_from_tuple(self, box: Tuple[float, float, float, float]) -> float:
|
|
x1, y1, x2, y2 = box
|
|
return abs(y2 - y1) * abs(x2-x1)
|
|
|
|
def prepare_height_samples(self, boxes_by_angle: Dict[float, Tuple[float, float, float, float]], area = False) -> List[
|
|
Tuple[float, float]]:
|
|
# angles in degrees -> (theta_rad, height)
|
|
samples = []
|
|
for deg, box in boxes_by_angle.items():
|
|
if not area:
|
|
h = self.box_height_from_tuple(box)
|
|
else:
|
|
h = self.box_area_from_tuple(box)
|
|
samples.append((math.radians(deg), h))
|
|
return samples
|
|
|
|
def fit_area_vs_angle(self, areas_by_angle_deg: List[Tuple[float, float]]) -> Tuple[Callable[[float], float], float, dict]:
|
|
"""
|
|
Fit area(θ) = A + B*cos(θ - φ) using a linear fit on cos/sin terms.
|
|
Input:
|
|
areas_by_angle_deg: { angle_deg: area }
|
|
Returns:
|
|
(area_fn, best_angle_deg, params)
|
|
area_fn(theta_deg) -> predicted area
|
|
best_angle_deg: angle (deg) maximizing fitted curve (in [0, 360))
|
|
params: {"A": A, "B": B, "phi_rad": phi}
|
|
"""
|
|
if not areas_by_angle_deg:
|
|
return (lambda _: 0.0, 0.0, {"A": 0.0, "B": 0.0, "phi_rad": 0.0})
|
|
|
|
samples = [(math.radians(deg), float(area)) for deg, area in areas_by_angle_deg]
|
|
|
|
if len(samples) < 3:
|
|
# Fallback: constant model at mean, choose best measured angle
|
|
mean_area = sum(a for _, a in samples) / len(samples)
|
|
best_measured = max(areas_by_angle_deg, key=lambda kv: kv[1])[0] % 360
|
|
return (lambda _: mean_area, best_measured, {"A": mean_area, "B": 0.0, "phi_rad": 0.0})
|
|
|
|
# Accumulate sums for normal equations
|
|
n = len(samples)
|
|
sum1 = n
|
|
sum_cos = sum(math.cos(t) for t, _ in samples)
|
|
sum_sin = sum(math.sin(t) for t, _ in samples)
|
|
sum_y = sum(y for _, y in samples)
|
|
sum_cos2 = sum((math.cos(t)) ** 2 for t, _ in samples)
|
|
sum_sin2 = sum((math.sin(t)) ** 2 for t, _ in samples)
|
|
sum_cossin = sum(math.cos(t) * math.sin(t) for t, _ in samples)
|
|
sum_ycos = sum(y * math.cos(t) for t, y in samples)
|
|
sum_ysin = sum(y * math.sin(t) for t, y in samples)
|
|
|
|
# Solve for [A, C, S] in:
|
|
# [ n sum_cos sum_sin ] [A] = [ sum_y ]
|
|
# [ sum_cos sum_cos2 sum_cossin ] [C] [ sum_ycos]
|
|
# [ sum_sin sum_cossin sum_sin2 ] [S] [ sum_ysin]
|
|
def det3(m):
|
|
return (m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
|
|
- m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
|
|
+ m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]))
|
|
|
|
M = [
|
|
[sum1, sum_cos, sum_sin],
|
|
[sum_cos, sum_cos2, sum_cossin],
|
|
[sum_sin, sum_cossin, sum_sin2],
|
|
]
|
|
b = [sum_y, sum_ycos, sum_ysin]
|
|
|
|
def replace_col(M, col, vec):
|
|
R = [row[:] for row in M]
|
|
for i in range(3):
|
|
R[i][col] = vec[i]
|
|
return R
|
|
|
|
D = det3(M) or 1e-12
|
|
A = det3(replace_col(M, 0, b)) / D
|
|
C = det3(replace_col(M, 1, b)) / D
|
|
S = det3(replace_col(M, 2, b)) / D
|
|
|
|
B = math.hypot(C, S)
|
|
phi = math.atan2(S, C) # C = B cosφ, S = B sinφ
|
|
|
|
def area_fn(theta_deg: float) -> float:
|
|
return A + B * math.cos(math.radians(theta_deg) - phi)
|
|
|
|
# Best angle occurs at θ = φ (convert to degrees, normalize)
|
|
best_angle_deg = (math.degrees(phi)) % 360
|
|
|
|
return area_fn, best_angle_deg, {"A": A, "B": B, "phi_rad": phi}
|
|
|
|
def fit_cosine_height(self, samples: List[Tuple[float, float]]) -> Tuple[float, float, float]:
|
|
"""
|
|
Fit h(θ) = A + C*cosθ + S*sinθ via linear least squares, then convert to A + B*cos(θ - φ).
|
|
Returns (A, B, phi) where phi in radians.
|
|
"""
|
|
if len(samples) < 3:
|
|
# fallback: constant model
|
|
A = sum(h for _, h in samples) / max(1, len(samples))
|
|
return (A, 0.0, 0.0)
|
|
|
|
# Build normal equations for [A, C, S]
|
|
sum1 = len(samples)
|
|
sum_cos = sum(math.cos(t) for t, _ in samples)
|
|
sum_sin = sum(math.sin(t) for t, _ in samples)
|
|
sum_h = sum(h for _, h in samples)
|
|
sum_cos2 = sum(math.cos(t) ** 2 for t, _ in samples)
|
|
sum_sin2 = sum(math.sin(t) ** 2 for t, _ in samples)
|
|
sum_cossin = sum(math.cos(t) * math.sin(t) for t, _ in samples)
|
|
sum_hcos = sum(h * math.cos(t) for t, h in samples)
|
|
sum_hsin = sum(h * math.sin(t) for t, h in samples)
|
|
|
|
# Solve 3x3 linear system:
|
|
# [ sum1 sum_cos sum_sin ] [A] = [ sum_h ]
|
|
# [ sum_cos sum_cos2 sum_cossin ] [C] [ sum_hcos ]
|
|
# [ sum_sin sum_cossin sum_sin2 ] [S] [ sum_hsin ]
|
|
# Use Cramer's rule or a tiny solver since numpy may not be allowed externally.
|
|
def det3(m):
|
|
return (m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
|
|
- m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
|
|
+ m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]))
|
|
|
|
M = [
|
|
[sum1, sum_cos, sum_sin],
|
|
[sum_cos, sum_cos2, sum_cossin],
|
|
[sum_sin, sum_cossin, sum_sin2],
|
|
]
|
|
bA = [sum_h, sum_hcos, sum_hsin]
|
|
|
|
# Matrices with columns replaced
|
|
def replace_col(M, col_idx, vec):
|
|
R = [row[:] for row in M]
|
|
for i in range(3):
|
|
R[i][col_idx] = vec[i]
|
|
return R
|
|
|
|
D = det3(M) or 1e-12
|
|
A = det3(replace_col(M, 0, bA)) / D
|
|
C = det3(replace_col(M, 1, bA)) / D
|
|
S = det3(replace_col(M, 2, bA)) / D
|
|
|
|
# Convert A + C cosθ + S sinθ to A + B cos(θ - φ)
|
|
B = math.hypot(C, S)
|
|
phi = math.atan2(S, C) # since C = B cosφ, S = B sinφ
|
|
|
|
return (A, B, phi)
|
|
|
|
def __face_detection_sequence(self, zoom_value: int = 280):
|
|
self.__set_state(BeamlineStateEnum.SampleAlignment)
|
|
self.__devs.lamp_light = 2.5
|
|
self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter
|
|
zoom_settings = self.__cfg.zoom_settings.z
|
|
zoom_value = 280
|
|
print('new loop')
|
|
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)
|
|
boxes = {}
|
|
angles = (90, 75, 60, 45, 30, 15, 0, -15, -30, -45, -60, -75, -90)
|
|
for angle in angles:
|
|
self.__devs.aerotech.move(angle, wait=True)
|
|
curr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR)
|
|
box = self.__mlbox.predict(curr_image, "", pref_class = (3,0))
|
|
if box:
|
|
cls, x1, y1, x2, y2 = box
|
|
if cls == 0:
|
|
#logger.info(f"loop all for angle {angle}")
|
|
boxes[angle] = (x1, y1, x2, y2)
|
|
|
|
if not boxes:
|
|
print("no boxes found")
|
|
return
|
|
|
|
samples = self.prepare_height_samples(boxes, area=True)
|
|
#logger.info(f"height vs angle samples: {samples}")
|
|
#area_fn, best_angle_deg, params = self.fit_area_vs_angle(samples)
|
|
A, B, phi = self.fit_cosine_height(samples)
|
|
#logger.info(f"cosine fit: A={A}, B={B}, phi={phi}")
|
|
#A, B, phi = params["A"], params["B"], params["phi_rad"]
|
|
|
|
def height_deg(theta_deg: float) -> float:
|
|
return A + B * math.cos(math.radians(theta_deg) - phi)
|
|
|
|
search_grid = range(-90, 90, 1)
|
|
best_fit_angle = max(search_grid, key=lambda d: height_deg(d))
|
|
best_fit_height = height_deg(best_fit_angle)
|
|
|
|
flat_face_angle = max(boxes.keys(), key=lambda a: self.box_height_from_tuple(boxes[a]))
|
|
flat_face_box = boxes[flat_face_angle]
|
|
#logger.info(f'Best fitted angle: {best_fit_angle}, fitted height: {best_fit_height:.3f}')
|
|
#logger.info(f'Flat face angle: {flat_face_angle}, box: {flat_face_box}')
|
|
|
|
# Move to fitted best angle and return both measured and fitted info and predictor
|
|
self.__devs.aerotech.move(best_fit_angle, wait=True)
|
|
return
|
|
|
|
#return flat_face_angle, flat_face_box
|
|
|
|
|
|
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")
|
|
break
|
|
if found_flag is not None 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")
|
|
self.__face_detection_sequence()
|
|
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
|
|
|
|
@property
|
|
def reference_tools(self) -> SampleShortInfoList:
|
|
return self.__cfg.reference_tools
|
|
|
|
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
|
|
)):
|
|
params = self.smart_params
|
|
print(params)
|
|
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:
|
|
print("auto center failed")
|
|
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 as e:
|
|
print(f"Error in measure: {e}")
|
|
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,
|
|
last_best_res = self.last_best_res,
|
|
last_best_b_factor = self.last_best_b_factor,
|
|
crystal_size = self.crystal_size
|
|
)
|
|
|
|
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
|