Fix slot number determination logic for the FIB workflow

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:

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

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):
     """

src/murfey/workflows/fib/register_atlas.py

@@ -1,18 +1,20 @@
 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
 from pydantic import BaseModel, computed_field, model_validator
 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,15 +364,26 @@ 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
             )
             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}",

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

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="",
     )