diff --git a/common/src/aaredaqlib/find_xtal.py b/common/src/aaredaqlib/find_xtal.py new file mode 100644 index 00000000..a40ecec2 --- /dev/null +++ b/common/src/aaredaqlib/find_xtal.py @@ -0,0 +1,288 @@ +from typing import List, Optional, Callable + +import numpy as np +from scipy import ndimage + +from aaredaqlib.models import CrystalSize +from aaredaqlib.raster_grid import RasterGridRequest + + +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})") + return grid_mm_x, grid_mm_y + + else: + return None, 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 get_xtal_size(crystal_size, result_array, r:RasterGridRequest): + # Optional: get bounding box of the largest object + labeled_array, num_objects = ndimage.label(result_array) + areas = ndimage.sum(np.ones_like(result_array, dtype=np.int32), labeled_array, + index=range(1, num_objects + 1)) + largest_idx = int(np.argmax(areas)) + 1 # +1 because labels start at 1 + largest_area = int(areas[largest_idx - 1]) + print(f"Largest object label: {largest_idx}, area (px): {largest_area}") + + object_mask = labeled_array == largest_idx + rows = np.any(object_mask, axis=1) + cols = np.any(object_mask, axis=0) + row_min, row_max = np.where(rows)[0][[0, -1]] + col_min, col_max = np.where(cols)[0][[0, -1]] + print(f"Largest bbox: width={col_max - col_min}, height={row_max - row_min}") + print( + f"Largest bbox: width={(col_max - col_min) * r.grid_size_mm.x}, y={(row_max - row_min) * r.grid_size_mm.y}") + try: + if r.n_x == 1: + + + + crystal_size = crystal_size + + crystal_size = CrystalSize(x=crystal_size.x, y=crystal_size.y, + z=(col_max - col_min) * r.grid_size_mm.y * 1000) + + else: + crystal_size = CrystalSize(x=(row_max - row_min) * r.grid_size_mm.x * 1000, + y=(col_max - col_min) * r.grid_size_mm.y * 1000, + z=crystal_size.z) + except ValueError as e: + print(f"error calculating xtal size: {e}") + crystal_size = CrystalSize(x=0,y=0,z=0) + + return crystal_size + + +def get_best_b_factor(result_list: List): + if not result_list: + return None + best_b_factor = min((img for img in result_list if img.b is not None), + key=lambda img: img.b, + default=None) + print(f"Best b: {best_b_factor.b}") + return best_b_factor + +def get_best_res(result_list: List): + if not result_list: + return None + best_res = min((img for img in result_list if img.res is not None), + key=lambda img: img.res, + default=None) + print(f"Best res: {best_res}") + return best_res + +def com_nan_check(com): + if np.isnan(com[0]) or np.isnan(com[1]): + return False + else: + return True + +def get_result_list_from_com(images, com): + if not com_nan_check(com): + return None + cx, cy = com[::-1] + start_x, end_x = round(cx - 1), round(cx + 1) + start_y, end_y = round(cy - 1), round(cy + 1) + print(f"range x {start_x} {end_x}, y {start_y} {end_y}") + result_list = [img for img in images + if start_x <= img.nx <= end_x and start_y <= img.ny <= end_y] + return result_list + +def get_com_image_number(com, images): + for image in images: + if image.nx == round(com[0]) and image.ny == round(com[1]): + print(f"com found for image: {image.number}") + return + +def get_grid_mm_from_com(com, r:RasterGridRequest): + if not com_nan_check(com): + 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) + return grid_mm_x, grid_mm_y + + +def raster_centre_of_mass(self, result_array, r:RasterGridRequest): + print('horizontal scan') + com = ndimage.center_of_mass(result_array) + print(f"Center of mass: {com}") + grid_mm_x, grid_mm_y = self.get_grid_mm_from_com(com, r) + return grid_mm_x, grid_mm_y, com + diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py index 1cc6649a..d02295f5 100644 --- a/daq/src/aaredaq/daq.py +++ b/daq/src/aaredaq/daq.py @@ -22,6 +22,8 @@ from aaredaq.mlbox import MlBox from aaredaqlib.beamline import MXBeamline from aaredaqlib.coordinate import Coordinate, SmargonCoordinate from aaredaqlib.diffraction_geometry import DiffractionGeometry +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, @@ -363,202 +365,22 @@ class AareDAQ: 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 @@ -584,132 +406,6 @@ class AareDAQ: else: return None - def raster_centre_of_mass(self, images, r:RasterGridRequest, result): - # 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) - - print('horizontal scan') - com = ndimage.center_of_mass(result_array) - - print(f"Center of mass: {com}") - - try: - labeled_array, num_objects = ndimage.label(result_array) - areas = ndimage.sum(np.ones_like(result_array, dtype=np.int32), labeled_array, - index=range(1, num_objects + 1)) - largest_idx = int(np.argmax(areas)) + 1 # +1 because labels start at 1 - largest_area = int(areas[largest_idx - 1]) - print(f"Largest object label: {largest_idx}, area (px): {largest_area}") - - # Optional: get bounding box of largest object - object_mask = labeled_array == largest_idx - rows = np.any(object_mask, axis=1) - cols = np.any(object_mask, axis=0) - row_min, row_max = np.where(rows)[0][[0, -1]] - col_min, col_max = np.where(cols)[0][[0, -1]] - print(f"Largest bbox: width={col_max - col_min}, height={row_max - row_min}") - print( - f"Largest bbox: width={(col_max - col_min) * r.grid_size_mm.x}, y={(row_max - row_min) * r.grid_size_mm.y}") - - if r.n_x == 1: - - if self.crystal_size is None: - self.crystal_size = CrystalSize(x=0,y=0,z=0) - - crystal_size = self.crystal_size - - crystal_size = CrystalSize(x=crystal_size.x, y=crystal_size.y, - z=(col_max - col_min) * r.grid_size_mm.y * 1000) - - else: - crystal_size = CrystalSize(x=(row_max - row_min) * r.grid_size_mm.x * 1000, - y=(col_max - col_min) * r.grid_size_mm.y * 1000, - z=0) - except Exception as e: - crystal_size = CrystalSize(x=0, y=0, z=0) - - - self.crystal_size = crystal_size - - if r.n_x == 1 and (np.isnan(com[1]) or np.isnan(com[0])): - print('vertical scan') - try: - max_image = max(images, key=lambda img: img.spots_low_res) - com = (0, max_image.ny) - except: - print("no spots") - - 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 - best_res = None - best_b_factor = 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) - - cx, cy = com[::-1] # com=(y, x) -> (x, y), rounded once - start_x, end_x = round(cx - 1), round(cx + 1) - start_y, end_y = round(cy - 1), round(cy + 1) - print(f"range x {start_x} {end_x}, y {start_y} {end_y}") - # Collect images in the 3x3 neighborhood around the center - res_list = [img for img in images - if start_x <= img.nx <= end_x and start_y <= img.ny <= end_y] - print(res_list) - # Best by res, skipping None - best_res = min((img for img in res_list if img.res is not None), - key=lambda img: img.res, - default=None) - print(f"Best res: {best_res}") - best_b_factor = min((img for img in res_list if img.b is not None), - key=lambda img: img.b, - default=None) - print(f"Best b: {best_b_factor}") - - for image in images: - if image.nx == round(com[0]) and image.ny == round(com[1]): - print(f"com found for image: {image.number}") - try: - if best_res is not None and best_res.res is not None: - # store the numeric resolution on the config/session so it appears in status - print(f"best res: {best_res.res}") - self.last_best_res = float(best_res.res) - else: - print(f'res is None') - self.last_best_res = None - except Exception as e: - print(f'error with last_best_res {e}') - self.last_best_res = None - try: - if best_b_factor is not None and best_b_factor.b is not None: - print(f"best res: {best_res.b}") - self.last_best_b_factor = float(best_b_factor.b) - else: - print(f'last_best_b_factor is None') - self.last_best_b_factor = None - except Exception as e: - print(f'error with last_best_b_factor {e}') - self.last_best_b_factor = None - - 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) - - return new_delta_mm, best_res def __raster(self, r: RasterGridRequest) -> CompletedRasterGridElem: max_time = r.exp_time_s * r.n_y * r.n_x + 60 @@ -768,7 +464,27 @@ class AareDAQ: with open(filename, 'w') as f: json.dump(output_data, f, indent=2) print('before centre_of_mass') - new_delta_mm, best_res = self.raster_centre_of_mass(images, r, result) + 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) + 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) + + result_list = get_result_list_from_com(images, com) + + self.last_best_b_factor = get_best_b_factor(result_list) + self.last_best_res = get_best_res(result_list) + + 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 r.n_x == 1: + new_delta_mm = self.sample_geometry.smargon_nudge(Coordinate(x=0, y=grid_mm_y)) + else: + new_delta_mm = self.sample_geometry.smargon_nudge(Coordinate(x=grid_mm_x, y=grid_mm_y)) + 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')