diff --git a/src/murfey/client/contexts/sxt.py b/src/murfey/client/contexts/sxt.py index b65b94499..2e58993c0 100644 --- a/src/murfey/client/contexts/sxt.py +++ b/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 diff --git a/src/murfey/client/watchdir_multigrid.py b/src/murfey/client/watchdir_multigrid.py index 611480e80..c155c5b55 100644 --- a/src/murfey/client/watchdir_multigrid.py +++ b/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, diff --git a/src/murfey/workflows/sxt/process_sxt_tilt_series.py b/src/murfey/workflows/sxt/process_sxt_tilt_series.py index 67009cbd0..a3e64b399 100644 --- a/src/murfey/workflows/sxt/process_sxt_tilt_series.py +++ b/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), diff --git a/tests/client/contexts/test_sxt.py b/tests/client/contexts/test_sxt.py index e24ff4b52..d1705ad29 100644 --- a/tests/client/contexts/test_sxt.py +++ b/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 "}, ) diff --git a/tests/workflows/sxt/test_process_sxt_tilt_series.py b/tests/workflows/sxt/test_process_sxt_tilt_series.py index f45c4be0d..3a34ea73c 100644 --- a/tests/workflows/sxt/test_process_sxt_tilt_series.py +++ b/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",