diff --git a/src/murfey/client/contexts/fib.py b/src/murfey/client/contexts/fib.py index c24e11142..b5db2c13e 100644 --- a/src/murfey/client/contexts/fib.py +++ b/src/murfey/client/contexts/fib.py @@ -5,12 +5,12 @@ import xml.etree.ElementTree as ET from dataclasses import dataclass, field from pathlib import Path -from typing import Callable, Type, TypeVar +from typing import Callable, Type, TypeVar, cast from murfey.client.context import Context from murfey.client.instance_environment import MurfeyInstanceEnvironment from murfey.util.client import capture_post -from murfey.util.fib import number_from_name +from murfey.util.fib import get_slot_number, number_from_name from murfey.util.models import ( LamellaSiteInfo, MillingStepInfo, @@ -489,12 +489,24 @@ def _determine_output_dir( # Determine the slot number slot_number: int | None = None for stage_name in reversed(STAGE_POSITION_NAMES.keys()): - if (stage_info := getattr(site_info.stage_info, stage_name, None)) is None: - continue - if stage_info.slot_number is None: + stage_values: StagePositionValues | None = getattr( + site_info.stage_info, stage_name, None + ) + if stage_values is None: continue else: - slot_number = stage_info.slot_number + rotation_offset = cast( + float, + self._machine_config.get("calibrations", {}).get( + "rotation_offset", 0 + ), + ) + slot_number = get_slot_number( + x=stage_values.x, + y=stage_values.y, + rotation=stage_values.rotation, + rotation_offset=rotation_offset, + ) break # Early exit if no slot number if slot_number is None: diff --git a/src/murfey/util/fib.py b/src/murfey/util/fib.py index fc0b0767f..48f4be53c 100644 --- a/src/murfey/util/fib.py +++ b/src/murfey/util/fib.py @@ -2,6 +2,7 @@ General functinos specific to the FIB workflow """ +import math from pathlib import Path @@ -27,3 +28,27 @@ def number_from_name(name: str) -> int: return int(stem[stem.rfind("(") + 1 : -1]) # Names without '()' or '#' should return 1 return 1 + + +def get_slot_number( + x: float | None = None, + y: float | None = None, + # Angles are in degrees + rotation: float | None = None, + rotation_offset: float = -75, +): + """ + In the Aquilos, the stage position values corresponding to slots 1 and 2 are + taken at a fixed stage rotation; at different stage rotation values, the x- + and y- ranges corresponding to slots 1 and 2 will change. This function will + rotate the provided stage values into the correct reference frame and return + the slot number. + """ + if x is not None and y is not None and rotation is not None: + # Rotate the xy-coordinates to reference frame + theta = math.radians(rotation - rotation_offset) + sin = math.sin(theta) + cos = math.cos(theta) + x_rot = (x * cos) - (y * sin) + return 1 if x_rot < 0 else 2 + return None diff --git a/src/murfey/util/models.py b/src/murfey/util/models.py index cdf6898c3..339b089ad 100644 --- a/src/murfey/util/models.py +++ b/src/murfey/util/models.py @@ -5,7 +5,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional -from pydantic import BaseModel, computed_field, field_validator +from pydantic import BaseModel, field_validator """ ======================================================================================= @@ -109,12 +109,6 @@ class StagePositionValues(BaseModel): rotation: float | None = None tilt_alpha: float | None = None - @computed_field - def slot_number(self) -> int | None: - if self.x is None: - return None - return 1 if self.x < 0 else 2 - class StagePositionInfo(BaseModel): """ diff --git a/src/murfey/workflows/fib/register_atlas.py b/src/murfey/workflows/fib/register_atlas.py index 03df5dc19..37fec21d6 100644 --- a/src/murfey/workflows/fib/register_atlas.py +++ b/src/murfey/workflows/fib/register_atlas.py @@ -1,10 +1,11 @@ import logging +import math import traceback import xml.etree.ElementTree as ET from functools import cached_property from importlib.metadata import entry_points from pathlib import Path -from typing import Any +from typing import Any, cast import numpy as np import PIL.Image @@ -12,7 +13,8 @@ from sqlmodel import Session, select import murfey.util.db as MurfeyDB -from murfey.util.fib import number_from_name +from murfey.util.config import get_machine_config +from murfey.util.fib import get_slot_number, number_from_name logger = logging.getLogger("murfey.workflows.fib.register_atlas") @@ -39,6 +41,7 @@ class FIBAtlasMetadata(BaseModel): pos_y: float pos_z: float rotation: float # Radians + slot_number: int tilt_alpha: float # Radians tilt_beta: float # Radians # Image dimensions @@ -66,17 +69,6 @@ def pixel_size(self) -> float: """ return 0.5 * (self.pixel_size_x + self.pixel_size_y) - # mypy doesn't support decorators on @property - @computed_field # type: ignore - @cached_property - def slot_number(self) -> int: - """ - Decide on a slot number for the site being inspected. From observation, - the x-position is entirely negative for one slot and entirely positive - for the other. - """ - return 1 if self.pos_x < 0 else 2 - # mypy doesn't support decorators on @property @computed_field # type: ignore @cached_property @@ -100,7 +92,7 @@ def site_name(self) -> str: return f"{self.project_name}--slot_{self.slot_number}" -def _parse_metadata(file: Path, visit_name: str): +def _parse_metadata(file: Path, visit_name: str, rotation_offset: float): """ Parses through the atlas image's tags to extract the relevant metadata """ @@ -127,31 +119,39 @@ def _parse_metadata(file: Path, visit_name: str): raise ValueError(f"Could not find required metadata in file {file}") # Extract key values from metadata + extracted: dict[str, Any] = { + key: node.text if (node := xml_metadata.find(node_path)) is not None else None + for key, node_path in ( + ("voltage", ".//Optics/AccelerationVoltage"), + ("shift_x", ".//Optics/BeamShift/X"), + ("shift_y", ".//Optics/BeamShift/Y"), + ("len_x", ".//Optics/ScanFieldOfView/X"), + ("len_y", ".//Optics/ScanFieldOfView/Y"), + ("pos_x", ".//StageSettings/StagePosition/X"), + ("pos_y", ".//StageSettings/StagePosition/Y"), + ("pos_z", ".//StageSettings/StagePosition/Z"), + # Angles are in radians + ("rotation", ".//StageSettings/StagePosition/Rotation"), + ("tilt_alpha", ".//StageSettings/StagePosition/Tilt/Alpha"), + ("tilt_beta", ".//StageSettings/StagePosition/Tilt/Beta"), + ("pixels_x", ".//BinaryResult/ImageSize/X"), + ("pixels_y", ".//BinaryResult/ImageSize/Y"), + ("pixel_size_x", ".//BinaryResult/PixelSize/X"), + ("pixel_size_y", ".//BinaryResult/PixelSize/Y"), + ) + } + # Calculate the slot number + extracted["slot_number"] = get_slot_number( + x=float(extracted["pos_x"]), + y=float(extracted["pos_y"]), + rotation=math.degrees(float(extracted["rotation"])), # Convert to degrees + rotation_offset=rotation_offset, + ) + # Return the parsed Pydantic model return FIBAtlasMetadata( visit_name=visit_name, file=file, - **{ - key: node.text - if (node := xml_metadata.find(node_path)) is not None - else None - for key, node_path in ( - ("voltage", ".//Optics/AccelerationVoltage"), - ("shift_x", ".//Optics/BeamShift/X"), - ("shift_y", ".//Optics/BeamShift/Y"), - ("len_x", ".//Optics/ScanFieldOfView/X"), - ("len_y", ".//Optics/ScanFieldOfView/Y"), - ("pos_x", ".//StageSettings/StagePosition/X"), - ("pos_y", ".//StageSettings/StagePosition/Y"), - ("pos_z", ".//StageSettings/StagePosition/Z"), - ("rotation", ".//StageSettings/StagePosition/Rotation"), - ("tilt_alpha", ".//StageSettings/StagePosition/Tilt/Alpha"), - ("tilt_beta", ".//StageSettings/StagePosition/Tilt/Beta"), - ("pixels_x", ".//BinaryResult/ImageSize/X"), - ("pixels_y", ".//BinaryResult/ImageSize/Y"), - ("pixel_size_x", ".//BinaryResult/PixelSize/X"), - ("pixel_size_y", ".//BinaryResult/PixelSize/Y"), - ) - }, + **extracted, ) @@ -364,6 +364,7 @@ def run( ) ).one() visit_name = murfey_session.visit + instrument_name = murfey_session.instrument_name except Exception: logger.error( "Exception encountered while querying Murfey database", exc_info=True @@ -371,8 +372,18 @@ def run( return {"success": False, "requeue": False} try: + # Load the machine config + machine_config = get_machine_config(instrument_name)[instrument_name] + rotation_offset: float = cast( + float, machine_config.calibrations.get("rotation_offset", 0) + ) + # Extract metadata from Electron Snapshot image - metadata = _parse_metadata(fib_info.atlas_file, visit_name) + metadata = _parse_metadata( + fib_info.atlas_file, + visit_name=visit_name, + rotation_offset=rotation_offset, + ) except Exception: logger.error( f"Error extracting metadata from file {fib_info.atlas_file}", diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py index e3d794f69..84fb4d6ae 100644 --- a/src/murfey/workflows/fib/register_milling_progress.py +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -3,12 +3,14 @@ import json import logging from importlib.metadata import entry_points -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from sqlmodel import Session as SQLModelSession, select import murfey.util.db as MurfeyDB from murfey.server import _transport_object +from murfey.util.config import get_machine_config +from murfey.util.fib import get_slot_number from murfey.util.models import ( GridSquareParameters, LamellaSiteInfo, @@ -344,10 +346,6 @@ def run(message: dict[str, Any], murfey_db: SQLModelSession): "Could not construct lookup tags; no stage position information found" ) return {"success": False, "requeue": False} - if latest_stage_position.slot_number is None: - logger.error("Could not construct lookup tags; 'slot_number' is missing") - return {"success": False, "requeue": False} - slot_number = latest_stage_position.slot_number # Milling step information if site_info.steps is None: @@ -364,6 +362,25 @@ def run(message: dict[str, Any], murfey_db: SQLModelSession): ).one() visit_name = murfey_session.visit instrument_name = murfey_session.instrument_name + + # Load the machine config + machine_config = get_machine_config(instrument_name)[instrument_name] + rotation_offset = cast( + float, machine_config.calibrations.get("rotation_offset", 0) + ) + + # Calculate the slot number + slot_number = get_slot_number( + x=latest_stage_position.x, + y=latest_stage_position.y, + rotation=latest_stage_position.rotation, + rotation_offset=rotation_offset, + ) + if slot_number is None: + logger.error( + "Could not construct lookup tags; 'slot_number' is missing" + ) + return {"success": False, "requeue": False} except Exception: logger.error( "Exception encountered while querying Murfey database", exc_info=True diff --git a/tests/client/contexts/test_fib.py b/tests/client/contexts/test_fib.py index 6c469457a..857133d83 100644 --- a/tests/client/contexts/test_fib.py +++ b/tests/client/contexts/test_fib.py @@ -60,6 +60,11 @@ def visit_dir(tmp_path: Path): return tmp_path / visit_name +@pytest.fixture +def mock_machine_config(): + return {"calibrations": {"rotation_offset": -75}} + + def _create_stage_position_node(stage_values: dict[str, str]): stage_position_node = ET.Element("StagePosition") for key, value in stage_values.items(): @@ -481,6 +486,7 @@ def test_handle_autotem_metadata( test_params: tuple[bool, bool, bool, bool, bool, bool, bool, bool], tmp_path: Path, visit_dir: Path, + mock_machine_config: dict, ): # Unpack test params ( @@ -528,7 +534,7 @@ def test_handle_autotem_metadata( context = FIBContext( acquisition_software="autotem", basepath=basepath, - machine_config={}, + machine_config=mock_machine_config, token="", ) if has_drift_correction_images: @@ -640,6 +646,7 @@ def test_make_drift_correction_gif( test_params: tuple[bool, bool, bool, bool, bool, bool, bool], tmp_path: Path, visit_dir: Path, + mock_machine_config: dict, fib_autotem_dc_images: list[Path], ): # Unpack test params @@ -688,7 +695,7 @@ def test_make_drift_correction_gif( context = FIBContext( acquisition_software="autotem", basepath=basepath, - machine_config={}, + machine_config=mock_machine_config, token="", ) @@ -704,7 +711,11 @@ def test_make_drift_correction_gif( if has_stage_position: stage_dict: dict[str, dict] = {"preparation_site": {}} if has_stage_values: - stage_dict["preparation_site"] = {"x": 0.003} + stage_dict["preparation_site"] = { + "x": 0.003, + "y": 0.003, + "rotation": -75, + } metadata_dict["stage_info"] = stage_dict if has_site_info: context._site_info[lamella_num] = LamellaSiteInfo(**metadata_dict) @@ -861,6 +872,7 @@ def test_fib_autotem_context( mocker: MockerFixture, visit_dir: Path, test_params: tuple[bool, str], + mock_machine_config: dict, ): # Unpack test params is_manual, trigger = test_params @@ -902,7 +914,7 @@ def test_fib_autotem_context( context = FIBContext( acquisition_software="autotem", basepath=basepath, - machine_config={}, + machine_config=mock_machine_config, token="", ) @@ -936,6 +948,7 @@ def test_fib_maps_context( mocker: MockerFixture, tmp_path: Path, visit_dir: Path, + mock_machine_config: dict, fib_maps_images: list[Path], ): # Mock the environment @@ -961,7 +974,7 @@ def test_fib_maps_context( context = FIBContext( acquisition_software="maps", basepath=basepath, - machine_config={}, + machine_config=mock_machine_config, token="", ) diff --git a/tests/workflows/fib/test_register_atlas.py b/tests/workflows/fib/test_register_atlas.py index 455a62ad4..f94821e06 100644 --- a/tests/workflows/fib/test_register_atlas.py +++ b/tests/workflows/fib/test_register_atlas.py @@ -12,6 +12,7 @@ from sqlmodel import Session as SQLModelSession, select as sm_select import murfey.util.db as MurfeyDB +from murfey.util.fib import get_slot_number from murfey.workflows.fib.register_atlas import FIBAtlasMetadata, _parse_metadata, run from tests.conftest import ExampleVisit @@ -152,8 +153,10 @@ def create_electron_snapshot_metadata( 0.0003, # Y 0.01, # Z -1.309, # Rotation + -75, # Rotation offset 0.8, # Alpha tilt 0, # Beta tilt + 2, # Expected slot number 3072, # Image size X 2048, # Y 1e-6, # Pixel size X @@ -171,9 +174,11 @@ def create_electron_snapshot_metadata( -0.003, # Stage X 0.0003, # Y 0.01, # Z - 1.309, # Rotation + 1.833, # Rotation + -75, # Rotation offset 0, # Alpha tilt 0, # Beta tilt + 2, # Expected slot number 3072, # Image size X 2048, # Y 1e-6, # Pixel size X @@ -198,6 +203,8 @@ def test_parse_metadata( float, float, float, + float, + int, int, int, float, @@ -219,8 +226,10 @@ def test_parse_metadata( pos_y, pos_z, rotation, + rotation_offset, tilt_alpha, tilt_beta, + expected_slot_number, pixels_x, pixels_y, pixel_size_x, @@ -234,7 +243,6 @@ def test_parse_metadata( / image_name / f"{image_name}.tiff" ) - slot_number = 1 if pos_x < 0 else 2 # Mock the results of opening an image file xml_string = create_electron_snapshot_metadata( @@ -264,7 +272,7 @@ def test_parse_metadata( ) # Run the function and check that output is correct - parsed = _parse_metadata(file, visit_name) + parsed = _parse_metadata(file, visit_name, rotation_offset) assert parsed.visit_name == visit_name assert parsed.file == file @@ -283,8 +291,8 @@ def test_parse_metadata( assert parsed.pixels_y == pixels_y assert parsed.pixel_size_x == pixel_size_x assert parsed.pixel_size_y == pixel_size_y - assert parsed.slot_number == slot_number - assert parsed.site_name == f"{project_name}--slot_{slot_number}" + assert parsed.slot_number == expected_slot_number + assert parsed.site_name == f"{project_name}--slot_{expected_slot_number}" assert parsed.pixel_size == 0.5 * (pixel_size_x + pixel_size_y) @@ -299,6 +307,7 @@ def test_run_with_db( ispyb_db_session: SQLAlchemySession, mock_ispyb_credentials, ): + rotation_offset = -75 test_files = ( visit_dir / "maps/LayersData/Layer/Electron Snapshot/Electron Snapshot.tiff", visit_dir @@ -319,6 +328,19 @@ def test_run_with_db( murfey_db_session.add(session_entry) murfey_db_session.commit() + # Mock the MachineConfig + mock_machine_config = MagicMock( + calibrations={ + "rotation_offset": rotation_offset, + } + ) + mocker.patch( + "murfey.workflows.fib.register_atlas.get_machine_config", + return_value={ + instrument_name: mock_machine_config, + }, + ) + # Mock the ISPyB connection where the TransportManager class is located mock_security_config = MagicMock() mock_security_config.ispyb_credentials = mock_ispyb_credentials @@ -352,25 +374,35 @@ def test_run_with_db( # Mock the metadata returned from the image file import murfey.workflows.fib.register_atlas + for test_file in test_files: + extracted = { + "voltage": 2000, + "shift_x": 0, + "shift_y": 0, + "len_x": 0.003072, + "len_y": 0.002048, + "pos_x": 0.003, + "pos_y": 0.0003, + "pos_z": 0.01, + "rotation": -1.309, + "tilt_alpha": 0.8, + "tilt_beta": 0, + "pixels_x": 3072, + "pixels_y": 2048, + "pixel_size_x": 1e-6, + "pixel_size_y": 1e-6, + } + extracted["slot_number"] = get_slot_number( + x=extracted["pos_x"], + y=extracted["pos_y"], + rotation=extracted["rotation"], + rotation_offset=rotation_offset, + ) mock_metadata = [ FIBAtlasMetadata( visit_name=visit_name, file=test_file, - voltage=2000, - shift_x=0, - shift_y=0, - len_x=0.003072, - len_y=0.002048, - pos_x=0.003, - pos_y=0.0003, - pos_z=0.01, - rotation=-1.309, - tilt_alpha=0.8, - tilt_beta=0, - pixels_x=3072, - pixels_y=2048, - pixel_size_x=1e-6, - pixel_size_y=1e-6, + **extracted, ) for test_file in test_files ] diff --git a/tests/workflows/fib/test_register_milling_progress.py b/tests/workflows/fib/test_register_milling_progress.py index 606f5ba79..6cbe80c96 100644 --- a/tests/workflows/fib/test_register_milling_progress.py +++ b/tests/workflows/fib/test_register_milling_progress.py @@ -299,7 +299,6 @@ "z": 0.0323644854106331, "rotation": 285.003247202109, "tilt_alpha": 25.9996646026832, - "slot_number": 1, }, "chunk_site": { "x": -0.0030037500000000003, @@ -307,7 +306,6 @@ "z": 0.032350405092592606, "rotation": 285.003247202109, "tilt_alpha": -0.000134158926728586, - "slot_number": 1, }, "thinning_site": { "x": -0.0030037500000000003, @@ -315,7 +313,6 @@ "z": 0.032350405092592606, "rotation": 285.003247202109, "tilt_alpha": -0.000134158926728586, - "slot_number": 1, }, "chunk_coincidence_params": { "x": -0.0030048260286678298, @@ -323,7 +320,6 @@ "z": 0.0323400707790533, "rotation": 285.003247202109, "tilt_alpha": -0.000134158926728586, - "slot_number": 1, }, "thinning_params": { "x": -0.0030037500000000003, @@ -331,7 +327,6 @@ "z": 0.032350405092592606, "rotation": 285.003247202109, "tilt_alpha": -0.000134158926728586, - "slot_number": 1, }, } site_info = { @@ -369,9 +364,21 @@ def test_run_with_db( murfey_db_session.add(session_entry) murfey_db_session.commit() + # Mock the MachineConfig + mock_machine_config = MagicMock( + calibrations={ + "rotation_offset": -75, + } + ) + mocker.patch( + "murfey.workflows.fib.register_milling_progress.get_machine_config", + return_value={ + instrument_name: mock_machine_config, + }, + ) + # Mock the ISPyB connection where the TransportManager class is located - mock_security_config = MagicMock() - mock_security_config.ispyb_credentials = mock_ispyb_credentials + mock_security_config = MagicMock(ispyb_credentials=mock_ispyb_credentials) mocker.patch( "murfey.server.ispyb.get_security_config", return_value=mock_security_config,