Search for a txrm reference is no internal one is present (#826)

If a txrm does not have an internal reference, instead of dismissing it search for a reference file in xrm format. Takes the most recent xrm which is not a mosaic image.

Still skips processing for txrms with less than 20 frames as these will be reference stacks.

Also changes the tilt series tags to remove angle naming, and skips transfer of "New folder"

Pairs with DiamondLightSource/cryoem-services#271

Author: stephen-riggs

src/murfey/client/contexts/sxt.py

@@ -1,4 +1,5 @@
 import logging
+from datetime import datetime
 from pathlib import Path
 from typing import Any
 
@@ -13,11 +14,50 @@
 )
 from murfey.client.instance_environment import MurfeyInstanceEnvironment
 from murfey.util.client import capture_post
+from murfey.util.models import File
 from murfey.util.tomo import midpoint
 
 logger = logging.getLogger("murfey.client.contexts.sxt")
 
 
+def _find_reference(txrm_file: Path) -> Path | None:
+    """Find a suitable reference to apply to the given txrm file"""
+    # Look for xrm files in the txrm folder, reverse sorted by time
+    candidates = []
+    for gf in txrm_file.parent.glob("*.xrm"):
+        candidates.append(
+            File(
+                name=gf.name,
+                description="",
+                size=gf.stat().st_size / 1e6,
+                timestamp=datetime.fromtimestamp(gf.stat().st_mtime),
+                full_path=str(gf),
+            )
+        )
+    candidates.sort(key=lambda x: x.timestamp, reverse=True)
+    for ref_option in candidates:
+        mosaic_size = 1
+        with OleFileIO(ref_option.full_path) as xrm_ole:
+            # Find images which are not mosaics (txrm spec typos this as mosiac)
+            if xrm_ole.exists("ImageInfo/MosiacRows") and xrm_ole.exists(
+                "ImageInfo/MosiacColumns"
+            ):
+                mosaic_size = int(
+                    np.frombuffer(
+                        xrm_ole.openstream("ImageInfo/MosiacRows").getvalue(), np.int32
+                    )[0]
+                    * np.frombuffer(
+                        xrm_ole.openstream("ImageInfo/MosiacColumns").getvalue(),
+                        np.int32,
+                    )[0]
+                )
+        if mosaic_size == 0:
+            logger.info(f"Found reference {ref_option}")
+            return Path(ref_option.full_path)
+    logger.warning(f"No reference found for {txrm_file}")
+    return None
+
+
 class SXTContext(Context):
     def __init__(
         self,
@@ -128,6 +168,7 @@ def post_transfer(
                 return False
 
             # Read the tilt angles and pixel size from the txrm
+            angles: list = []
             metadata: dict[str, Any] = {
                 "source": str(self._basepath),
                 "tilt_series_tag": transferred_file.stem,
@@ -206,16 +247,30 @@ def post_transfer(
                         )
                         metadata["energy"] = int(round(axis_values[energy_index]))
 
-            if not metadata.get("has_reference", False):
+            if (
+                not metadata.get("has_reference", False)
+                and metadata.get("tilt_series_length", len(angles)) < 20
+            ):
+                # References are collected with only 10 frames
                 logger.debug(f"Reference image {transferred_file} not processed")
                 return True
+            elif not metadata.get("has_reference", False):
+                reference_file = _find_reference(transferred_file)
+            else:
+                reference_file = None
 
+            if "@" in transferred_file.stem:
+                tilt_series_tag = "_".join(
+                    transferred_file.stem.split("@")[0].split("_")[:-1]
+                )
+            else:
+                tilt_series_tag = transferred_file.stem
             visit_index = transferred_file.parent.parts.index(environment.visit)
             destination_search_dir = "/".join(
                 transferred_file.parts[: visit_index + 2]
             ).replace("//", "/")
             self.register_sxt_data_collection(
-                tilt_series=transferred_file.stem,
+                tilt_series=tilt_series_tag,
                 data_collection_parameters=metadata,
                 file_extension=transferred_file.suffix,
                 image_directory=str(
@@ -238,6 +293,15 @@ def post_transfer(
                 transferred_file,
                 Path(self._machine_config.get("rsync_basepath", "")),
             )
+            if reference_file:
+                reference_file_transferred_to = _file_transferred_to(
+                    environment,
+                    source,
+                    reference_file,
+                    Path(self._machine_config.get("rsync_basepath", "")),
+                )
+            else:
+                reference_file_transferred_to = None
             capture_post(
                 base_url=str(environment.url.geturl()),
                 router_name="workflow.sxt_router",
@@ -247,7 +311,7 @@ def post_transfer(
                 visit_name=environment.visit,
                 session_id=environment.murfey_session,
                 data={
-                    "tag": transferred_file.stem,
+                    "tag": tilt_series_tag,
                     "source": destination_search_dir,
                     "pixel_size": round(
                         metadata.get("pixel_size", 100), 2
@@ -257,6 +321,9 @@ def post_transfer(
                         "tilt_series_length", len(angles)
                     ),
                     "txrm": str(file_transferred_to),
+                    "xrm_reference": str(reference_file_transferred_to)
+                    if reference_file_transferred_to
+                    else None,
                 },
             )
         return True

src/murfey/client/watchdir_multigrid.py

@@ -93,7 +93,9 @@ def _handle_fractions(self, directory: Path):
     def _process(self):
         while not self._stopping:
             for d in self._basepath.glob("*"):
-                if d.name in self._machine_config["create_directories"]:
+                if d.name.startswith("New folder"):
+                    self._seen_dirs.append(d)
+                elif d.name in self._machine_config["create_directories"]:
                     if d.is_dir() and d not in self._seen_dirs:
                         self.notify(
                             d,

src/murfey/workflows/sxt/process_sxt_tilt_series.py

@@ -28,6 +28,7 @@ class SXTTiltSeriesInfo(BaseModel):
     tilt_series_length: int
     pixel_size: float
     tilt_offset: int
+    xrm_reference: str | None
 
 
 def process_sxt_tilt_series_workflow(
@@ -98,6 +99,7 @@ def process_sxt_tilt_series_workflow(
         "recipes": ["sxt-aretomo"],
         "parameters": {
             "txrm_file": tilt_series_info.txrm,
+            "xrm_reference": tilt_series_info.xrm_reference or "",
             "dcid": collected_ids[1].id,
             "appid": collected_ids[3].id,
             "stack_file": str(stack_file),

tests/client/contexts/test_sxt.py

@@ -54,7 +54,7 @@ def test_sxt_context_txrm(mock_ole_file, mock_post, tmp_path):
         np.array([2048], dtype=np.int32).tobytes(),  # Image Height
         np.array([1.5], dtype=np.float32).tobytes(),  # Exposure time
         np.array([1000], dtype=np.float32).tobytes(),  # Mag
-        np.array([5], dtype=np.int32).tobytes(),  # Image count
+        np.array([200], dtype=np.int32).tobytes(),  # Image count
         np.array([0, 519, 2, 3], dtype=np.float32).tobytes(),  # Motor Pos (energy)
     ]
 
@@ -122,7 +122,7 @@ def test_sxt_context_txrm(mock_ole_file, mock_post, tmp_path):
             "voltage": 0,
             "axis_start": -55,
             "axis_end": 65,
-            "tilt_series_length": 5,
+            "tilt_series_length": 200,
         },
         headers={"Authorization": "Bearer "},
     )
@@ -143,8 +143,120 @@ def test_sxt_context_txrm(mock_ole_file, mock_post, tmp_path):
             "source": f"{tmp_path}/cm12345-6/grid1",
             "pixel_size": 100.1,
             "tilt_offset": 5,
-            "tilt_series_length": 5,
+            "tilt_series_length": 200,
             "txrm": str(tmp_path / "destination/cm12345-6/grid1/example.txrm"),
+            "xrm_reference": None,
+        },
+        headers={"Authorization": "Bearer "},
+    )
+
+
+@patch("requests.post")
+@patch("murfey.client.contexts.sxt.OleFileIO")
+def test_sxt_context_txrm_external_ref(mock_ole_file, mock_post, tmp_path):
+    mock_post().status_code = 200
+    exists_return = [False]  # False for reference, then True
+    exists_return.extend([True for i in range(20)])
+    mock_ole_file().__enter__().exists.side_effect = exists_return
+    # Motor position names
+    mock_ole_file().__enter__().openstream().read.return_value = (
+        "\x00Val1\x00\x00Energy\x00".encode()
+    )
+    # Metadata encoded arrays
+    mock_ole_file().__enter__().openstream().getvalue.side_effect = [
+        np.array([-55, -25, 5, 35, 65], dtype=np.float32).tobytes(),  # Angles
+        np.array([0.01001], dtype=np.float32).tobytes(),  # Pixel size
+        np.array([1024], dtype=np.int32).tobytes(),  # Image Width
+        np.array([2048], dtype=np.int32).tobytes(),  # Image Height
+        np.array([1.5], dtype=np.float32).tobytes(),  # Exposure time
+        np.array([1000], dtype=np.float32).tobytes(),  # Mag
+        np.array([200], dtype=np.int32).tobytes(),  # Image count
+        np.array([0, 519, 2, 3], dtype=np.float32).tobytes(),  # Motor Pos (energy)
+        np.array([0], dtype=np.int32).tobytes(),  # Mosaic size
+        np.array([0], dtype=np.int32).tobytes(),  # Mosaic size
+    ]
+
+    # xrm file as reference
+    (tmp_path / "cm12345-6/grid1").mkdir(parents=True)
+    (tmp_path / "cm12345-6/grid1/ref.xrm").touch()
+
+    env = MurfeyInstanceEnvironment(
+        url=urlparse("http://localhost:8000"),
+        client_id=0,
+        sources=[tmp_path / "cm12345-6/grid1"],
+        default_destinations={
+            f"{tmp_path}/cm12345-6/grid1": f"{tmp_path}/destination/cm12345-6/grid1"
+        },
+        instrument_name="",
+        visit="cm12345-6",
+        murfey_session=1,
+    )
+    context = SXTContext("zeiss", tmp_path / "cm12345-6/grid1", {}, "")
+    context.post_transfer(
+        tmp_path / "cm12345-6/grid1/example_-60to60@0.5.txrm",
+        required_position_files=[],
+        required_strings=["fractions"],
+        environment=env,
+    )
+
+    mock_ole_file.assert_any_call(
+        str(tmp_path / "cm12345-6/grid1/example_-60to60@0.5.txrm")
+    )
+    mock_ole_file.assert_any_call(str(tmp_path / "cm12345-6/grid1/ref.xrm"))
+
+    assert mock_post.call_count == 5
+    mock_post.assert_any_call(
+        "http://localhost:8000/workflow/visits/cm12345-6/sessions/1/register_data_collection_group",
+        json={
+            "experiment_type_id": 47,
+            "tag": f"{tmp_path}/cm12345-6/grid1",
+        },
+        headers={"Authorization": "Bearer "},
+    )
+    mock_post.assert_any_call(
+        "http://localhost:8000/workflow/visits/cm12345-6/sessions/1/start_data_collection",
+        json={
+            "experiment_type": "sxt",
+            "file_extension": ".txrm",
+            "acquisition_software": "zeiss",
+            "image_directory": f"{tmp_path}/destination/cm12345-6/grid1",
+            "data_collection_tag": "example",
+            "source": f"{tmp_path}/cm12345-6/grid1",
+            "tag": "example",
+            "pixel_size_on_image": str(100.1 * 1e-10),
+            "image_size_x": 1024,
+            "image_size_y": 2048,
+            "magnification": 1000,
+            "energy": 519,
+            "voltage": 0,
+            "axis_start": -55,
+            "axis_end": 65,
+            "tilt_series_length": 200,
+        },
+        headers={"Authorization": "Bearer "},
+    )
+    mock_post.assert_any_call(
+        "http://localhost:8000/workflow/visits/cm12345-6/sessions/1/register_processing_job",
+        json={
+            "tag": "example",
+            "source": f"{tmp_path}/cm12345-6/grid1",
+            "recipe": "sxt-aretomo",
+            "experiment_type": "sxt",
+        },
+        headers={"Authorization": "Bearer "},
+    )
+    mock_post.assert_any_call(
+        "http://localhost:8000/workflow/sxt/visits/cm12345-6/sessions/1/sxt_tilt_series",
+        json={
+            "tag": "example",
+            "source": f"{tmp_path}/cm12345-6/grid1",
+            "pixel_size": 100.1,
+            "tilt_offset": 5,
+            "tilt_series_length": 200,
+            "txrm": str(
+                tmp_path / "destination/cm12345-6/grid1/example_-60to60@0.5.txrm"
+            ),
+            "xrm_reference": str(tmp_path / "destination/cm12345-6/grid1/ref.xrm"),
         },
         headers={"Authorization": "Bearer "},
     )

tests/workflows/sxt/test_process_sxt_tilt_series.py

@@ -65,6 +65,7 @@ def test_process_new_sxt_tilt_series(
         tag="tomogram_tag",
         source="/path/to/tomogram_source",
         txrm=f"{tmp_path}/cm12345-6/raw/tomogram_tag.txrm",
+        xrm_reference=f"{tmp_path}/cm12345-6/raw/ref.xrm",
         tilt_series_length=5,
         pixel_size=100,
         tilt_offset=1,
@@ -84,6 +85,7 @@ def test_process_new_sxt_tilt_series(
         {
             "parameters": {
                 "txrm_file": f"{tmp_path}/cm12345-6/raw/tomogram_tag.txrm",
+                "xrm_reference": f"{tmp_path}/cm12345-6/raw/ref.xrm",
                 "dcid": dc_id,
                 "appid": app_id,
                 "stack_file": f"{tmp_path}/cm12345-6/processed/raw/Tomograms/tomogram_tag_stack.mrc",