diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py index 724b8891..bba73950 100644 --- a/daq/src/aaredaq/daq.py +++ b/daq/src/aaredaq/daq.py @@ -1,14 +1,16 @@ import copy +import json import time from datetime import datetime from math import ceil -from typing import List, Tuple +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 @@ -390,6 +392,149 @@ class AareDAQ: 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: @@ -416,9 +561,9 @@ class AareDAQ: self.__devs.aerotech.move(grid.omega_deg, wait=True) grid.n_x = 1 - grid.n_y = 100 + 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 * 0.8, y=geom.beam_size_mm.y * 0.2) + 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) @@ -470,8 +615,58 @@ class AareDAQ: #if self.sample is not None and self.sample.db_id is not None: # self.__aare.ingest_gridscan(self.sample, result, r) - new_delta_mm = self.identify_crystal_raster(result, r) + 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) @@ -797,9 +992,9 @@ class AareDAQ: print(f'Thresh value: {thresh_value}') if filename is not None: #cv2.imwrite(f"{filename}_curr_image_colour.jpg", curr_image) - cv2.imwrite(f"{filename}_curr_image.jpg", gray_with_feature) - cv2.imwrite(f"{filename}_diff.jpg", diff_image) - cv2.imwrite(f"{filename}_thresh.jpg", thresh) + cv2.imwrite(f"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 @@ -840,47 +1035,72 @@ class AareDAQ: try: i = 0 for s in self.__cfg.alc_zoom_settings.z: - print('new loop center settings') - self.__devs.samcam_settings = SampleCameraSettings(exposure=s.sam_cam_exp, gain=s.sam_cam_gain) - self.__devs.zoom_sync(s.zoom_value) - print(f'zoom={s.zoom_value},gain={s.sam_cam_gain}, exp={s.sam_cam_exp}') - if sample_id is not None: - self.save_screenshot_db(sample_id, f"pre_alc") - print(sample_id) - if i % 2 == 0: - angles = (0, -45, -90, -135) - else: - angles = (-135, -90, -45, 0) - for angle in angles: - print(f"Moving to new omega: {angle}") - if sample_id is not None and i == 0: - exp = int(s.sam_cam_exp * 1000) - gain = int(s.sam_cam_gain) - self.save_screenshot_db(sample_id, f"pre_alc_{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}") - self.__devs.aerotech.move(angle, wait=True) - time.sleep(1) - filename = None - if sample_id is not None: - if not os.path.exists(f"{sample_id}_{s.zoom_value:.0f}"): - os.mkdir(f"{sample_id}_{s.zoom_value:.0f}") - if not os.path.exists(f"{sample_id}_{s.zoom_value:.0f}_negative"): - os.mkdir(f"{sample_id}_{s.zoom_value:.0f}_negative") - if angle > 0: - filename = f"{sample_id}_{s.zoom_value:.0f}/{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}" - else: - filename = f"{sample_id}_{s.zoom_value:.0f}_negative/{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}" + max_attempt = 2 + attempt = 0 + found = 0 + angles = (0, -45, -90) - target = self.__ml_loop_centre_box(sample_id, filename) - if target is None: - continue - self.__devs.smargon.target = target #self.__loop_center(s, filename, iteration=i) - self.__devs.smargon.wait(60) - print(sample_id) - if sample_id is not None: - self.__aare.sample_centered(self.__cfg.current_sample) - time.sleep(0.1) - self.save_screenshot_db(sample_id, f"{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}") - i += 1 + while attempt < max_attempt: + print('new loop center settings') + self.__devs.samcam_settings = SampleCameraSettings(exposure=s.sam_cam_exp, gain=s.sam_cam_gain) + self.__devs.zoom_sync(s.zoom_value) + print(f'zoom={s.zoom_value},gain={s.sam_cam_gain}, exp={s.sam_cam_exp}') + + if sample_id is not None: + self.save_screenshot_db(sample_id, f"pre_alc") + print(sample_id) + + found_flag = False + found_angle = None + targets_found_this_attempt = 0 + exp = int(s.sam_cam_exp * 1000) + gain = int(s.sam_cam_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}_{s.zoom_value:.0f}_{exp}_{gain}") + + self.__devs.aerotech.move(angle, wait=True) + time.sleep(1) + filename = None + if sample_id is not None: + filename = f"{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}" + + 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}_{s.zoom_value:.0f}_{exp}_{gain}") + + + 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}")