IMOD processing bypass and file copying

Dockerfiles/tomo_align

@@ -1,4 +1,4 @@
-# This Dockerfile is used for GPU AreTomo2 processing
+# This Dockerfile is used for GPU AreTomo3 or IMOD processing
 FROM rockylinux:8 AS conda-build
 
 # Set up conda environment
@@ -34,15 +34,16 @@ ENV PATH=/install/venv/bin:${PATH}
 ENV LD_LIBRARY_PATH=/install/venv/lib:${LD_LIBRARY_PATH}
 
 # Install IMOD
+RUN yum groupinstall "Development Tools" -y
+RUN yum install mesa-libGL -y
 RUN mkdir imod_install && \
     curl https://bio3d.colorado.edu/imod/AMD64-RHEL5/imod_5.1.9_RHEL8-64_CUDA12.0.sh > imod_5.1.9_RHEL8-64_CUDA12.0.sh && \
     chmod +x imod_5.1.9_RHEL8-64_CUDA12.0.sh && \
-    ln -s /install/venv/bin/python /usr/bin/python3 && \
     ./imod_5.1.9_RHEL8-64_CUDA12.0.sh -dir imod_install -skip -y
 ENV PATH=/imod_install/IMOD/bin:${PATH}
 ENV IMOD_DIR=/imod_install/IMOD
 
 # Install AreTomo
-COPY --chown="${userid}":"${groupid}" packages/AreTomo2 /AreTomo2
-ENV PATH=/AreTomo2:${PATH}
-RUN chmod +x /AreTomo2
+COPY --chown="${userid}":"${groupid}" packages/AreTomo3 /AreTomo3
+ENV PATH=/AreTomo3:${PATH}
+RUN chmod +x /AreTomo3

recipes/ispyb/sxt-aretomo.json

@@ -22,6 +22,7 @@
       "success": 7
     },
     "parameters": {
+      "copy_output": true,
       "dark_tol": 0,
       "manual_tilt_offset": "{manual_tilt_offset}",
       "out_bin": 1,

recipes/ispyb/sxt-imod-beads-wbp.json

@@ -22,6 +22,7 @@
       "success": 7
     },
     "parameters": {
+      "copy_output": true,
       "manual_tilt_offset": "{manual_tilt_offset}",
       "out_bin": 1,
       "patch": 0,
@@ -30,7 +31,8 @@
       "stack_file": "{stack_file}",
       "tilt_axis": 0,
       "txrm_file": "{txrm_file}",
-      "wbp": 1
+      "wbp": 1,
+      "xrm_reference": "{xrm_reference}"
     },
     "queue": "tomo_align_imod",
     "service": "ImodTomoAlign"

recipes/ispyb/sxt-imod-patch-wbp.json

@@ -22,6 +22,7 @@
       "success": 7
     },
     "parameters": {
+      "copy_output": true,
       "manual_tilt_offset": "{manual_tilt_offset}",
       "out_bin": 1,
       "patch": 1,
@@ -30,7 +31,8 @@
       "stack_file": "{stack_file}",
       "tilt_axis": 0,
       "txrm_file": "{txrm_file}",
-      "wbp": 1
+      "wbp": 1,
+      "xrm_reference": "{xrm_reference}"
     },
     "queue": "tomo_align_imod",
     "service": "ImodTomoAlign"

src/cryoemservices/services/denoise.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+import shutil
 import subprocess
 from pathlib import Path
 from typing import List, Optional
@@ -48,6 +49,7 @@ class DenoiseParameters(BaseModel):
     patch_padding: Optional[int] = None  # 48
     device: Optional[int] = None  # -2
     cleanup_output: bool = True
+    copy_output: bool = False
     visits_for_slurm: Optional[list] = ["bi", "cm", "nr", "nt"]
     relion_options: RelionServiceOptions
 
@@ -305,6 +307,7 @@ def denoise(self, rw, header: dict, message: dict):
             "tomogram": str(denoised_full_path),
             "output_dir": str(segmentation_dir),
             "pixel_size": str(denoise_params.relion_options.pixel_size_downscaled),
+            "copy_output": denoise_params.copy_output,
             "relion_options": dict(denoise_params.relion_options),
         }
         cryolo_parameters = {
@@ -325,6 +328,18 @@ def denoise(self, rw, header: dict, message: dict):
         }
         rw.send_to("ispyb_connector", ispyb_parameters)
 
+        # Optionally copy output file
+        if denoise_params.copy_output:
+            # Take file name for Relion-type projects, or folder name for SXT-style
+            tomo_name = (
+                denoised_full_path.name
+                if re.match(".*/job[0-9]+/.*", str(denoised_full_path))
+                else f"{denoised_full_path.parent.parent}_denoised.mrc"
+            )
+            shutil.copy(
+                denoised_full_path, denoised_full_path.parent.parent.parent / tomo_name
+            )
+
         self.log.info(f"Done denoising for {denoise_params.volume}")
         rw.transport.ack(header)
         return

src/cryoemservices/services/membrain_seg.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import re
+import shutil
 import subprocess
 from pathlib import Path
 from typing import Optional
@@ -36,6 +38,7 @@ class MembrainSegParameters(BaseModel):
     segmentation_threshold: float = 0.0
     cleanup_output: bool = True
     submit_to_slurm: bool = False
+    copy_output: bool = False
     relion_options: RelionServiceOptions
 
 
@@ -266,6 +269,16 @@ def membrain_seg(self, rw, header: dict, message: dict):
             },
         )
 
+        # Optionally copy output file
+        if membrain_seg_params.copy_output:
+            # Take file name for Relion-type projects, or folder name for SXT-style
+            tomo_name = (
+                segmented_path.name
+                if re.match(".*/job[0-9]+/.*", str(segmented_path))
+                else f"{segmented_path.parent.parent}_segmented.mrc"
+            )
+            shutil.copy(segmented_path, segmented_path.parent.parent.parent / tomo_name)
+
         self.log.info(f"Done segmentation for {membrain_seg_params.tomogram}")
         rw.transport.ack(header)
         return

src/cryoemservices/services/tomo_align_aretomo.py

@@ -4,6 +4,7 @@
 import json
 import os.path
 import re
+import shutil
 import subprocess
 import time
 from pathlib import Path
@@ -135,6 +136,7 @@ class AreTomoParameters(BaseModel):
     interpolation_correction: int | None = None
     dark_tol: float | None = None
     manual_tilt_offset: float | None = None
+    copy_output: bool = False
     visits_for_slurm: list | None = ["bi", "cm", "nr", "nt"]
     relion_options: RelionServiceOptions
 
@@ -835,6 +837,7 @@ def _tilt(file_list_for_tilts):
                 "volume": str(aretomo_output_path),
                 "output_dir": str(denoise_dir),
                 "relion_options": dict(tomo_params.relion_options),
+                "copy_output": tomo_params.copy_output,
             },
         )
 
@@ -852,6 +855,22 @@ def _tilt(file_list_for_tilts):
         ):
             tmp_file.unlink()
 
+        # Optionally copy output file
+        if tomo_params.copy_output:
+            # Take file name for Relion-type projects, or folder name for SXT-style
+            stack_name = (
+                Path(tomo_params.stack_file).name
+                if re.match(".*/job[0-9]+/.*", tomo_params.stack_file)
+                else f"{Path(tomo_params.stack_file).parent.parent}_stack.mrc"
+            )
+            tomo_name = (
+                aretomo_output_path.name
+                if re.match(".*/job[0-9]+/.*", str(aretomo_output_path))
+                else f"{aretomo_output_path.parent.parent}_volume.mrc"
+            )
+            shutil.copy(Path(tomo_params.stack_file), project_dir.parent / stack_name)
+            shutil.copy(aretomo_output_path, project_dir.parent / tomo_name)
+
         # Update success processing status
         rw.send_to("success", {})
         self.log.info(f"Done tomogram alignment for {tomo_params.stack_file}")

src/cryoemservices/services/tomo_align_imod.py

@@ -1,3 +1,4 @@
+import shutil
 import subprocess
 from pathlib import Path
 from typing import Optional
@@ -32,6 +33,7 @@ class ImodTomoParameters(BaseModel):
     local_alignment: int = 0
     flip_vol: int = 1
     manual_tilt_offset: Optional[float] = None
+    copy_output: bool = False
     cpus: int = 4
 
 
@@ -150,6 +152,7 @@ def tomo_align(self, rw, header: dict, message: dict):
                 str(adoc_file),
                 "-cpus",
                 str(tomo_params.cpus),
+                "-bypass",
             ],
             capture_output=True,
         )
@@ -272,6 +275,7 @@ def tomo_align(self, rw, header: dict, message: dict):
             {
                 "volume": str(imod_output_path),
                 "output_dir": str(imod_output_path.parent.parent / "Denoise"),
+                "copy_output": tomo_params.copy_output,
                 "relion_options": {},
             },
         )
@@ -290,6 +294,14 @@ def tomo_align(self, rw, header: dict, message: dict):
         ):
             tmp_file.unlink()
 
+        # Optionally copy output file
+        if tomo_params.copy_output:
+            shutil.copy(
+                imod_output_path,
+                imod_output_path.parent.parent.parent
+                / f"{imod_output_path.parent.parent}_volume.mrc",
+            )
+
         # Update success processing status
         rw.send_to("success", {})
         self.log.info(f"Done tomogram alignment for {tomo_params.stack_file}")

tests/services/test_denoise.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import subprocess
 from unittest import mock
 
 import pytest
@@ -188,6 +189,7 @@ def test_denoise_local_topaz_service(
             "tomogram": f"{tmp_path}/Denoise/job007/denoised/test_stack_aretomo.denoised.mrc",
             "output_dir": f"{tmp_path}/Segmentation/job008/tomograms",
             "pixel_size": "1.0",
+            "copy_output": False,
             "relion_options": output_relion_options,
         },
     )
@@ -216,9 +218,18 @@ def test_denoise_local_subprocess_service(
     This should call the mock subprocess then send messages on to
     the membrain-seg and images services.
     """
-    mock_subprocess().returncode = 0
-    mock_subprocess().stdout = "stdout".encode("ascii")
-    mock_subprocess().stderr = "stderr".encode("ascii")
+
+    def touch_denoise_output(*args, **kwargs):
+        (tmp_path / "Denoise/job007/denoised").mkdir(parents=True, exist_ok=True)
+        (tmp_path / "Denoise/job007/denoised/test_stack_aretomo.denoised.mrc").touch()
+        return subprocess.CompletedProcess(
+            "",
+            returncode=0,
+            stdout="stdout".encode("ascii"),
+            stderr="stderr".encode("ascii"),
+        )
+
+    mock_subprocess.side_effect = touch_denoise_output
 
     header = {
         "message-id": mock.sentinel,
@@ -250,6 +261,7 @@ def test_denoise_local_subprocess_service(
         "patch_size": 96,
         "patch_padding": 48,
         "device": "-2",
+        "copy_output": True,
         "relion_options": {"pixel_size_downscaled": 1},
     }
     output_relion_options = dict(RelionServiceOptions())
@@ -310,7 +322,10 @@ def test_denoise_local_subprocess_service(
         "-d",
         "-2",
     ]
-    mock_subprocess.assert_any_call(denoise_command, capture_output=True)
+    mock_subprocess.assert_called_once_with(denoise_command, capture_output=True)
+
+    # Check output copy
+    assert (tmp_path / "Denoise/test_stack_aretomo.denoised.mrc").is_file()
 
     # Check the images service request
     assert offline_transport.send.call_count == 6
@@ -356,6 +371,7 @@ def test_denoise_local_subprocess_service(
             "tomogram": f"{tmp_path}/Denoise/job007/denoised/test_stack_aretomo.denoised.mrc",
             "output_dir": f"{tmp_path}/Segmentation/job008/tomograms",
             "pixel_size": "1.0",
+            "copy_output": True,
             "relion_options": output_relion_options,
         },
     )
@@ -568,6 +584,7 @@ def write_denoised_files(**kwargs):
             "tomogram": f"{tmp_path}/cm12345-6/Denoise/job007/denoised/test_stack_aretomo.denoised.mrc",
             "output_dir": f"{tmp_path}/cm12345-6/Segmentation/job008/tomograms",
             "pixel_size": "1.0",
+            "copy_output": False,
             "relion_options": output_relion_options,
         },
     )
@@ -654,6 +671,7 @@ def test_denoise_local_topaz_service_rerun(
             "tomogram": f"{tmp_path}/Denoise/job007/denoised/test_stack_aretomo.denoised.mrc",
             "output_dir": f"{tmp_path}/Segmentation/job008/tomograms",
             "pixel_size": "1.0",
+            "copy_output": False,
             "relion_options": output_relion_options,
         },
     )

tests/services/test_membrain_seg.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import subprocess
 from unittest import mock
 
 import pytest
@@ -159,9 +160,20 @@ def test_membrain_seg_service_local_subprocess(
     Send a test message to membrain-seg for the local version
     This should call the mock subprocess then send messages to the images service.
     """
-    mock_subprocess().returncode = 0
-    mock_subprocess().stdout = "stdout".encode("ascii")
-    mock_subprocess().stderr = "stderr".encode("ascii")
+
+    def touch_membrain_output(*args, **kwargs):
+        (
+            tmp_path
+            / "Segmentation/job008/tomograms/test_stack_aretomo.denoised_segmented.mrc"
+        ).touch()
+        return subprocess.CompletedProcess(
+            "",
+            returncode=0,
+            stdout="stdout".encode("ascii"),
+            stderr="stderr".encode("ascii"),
+        )
+
+    mock_subprocess.side_effect = touch_membrain_output
 
     header = {
         "message-id": mock.sentinel,
@@ -179,6 +191,7 @@ def test_membrain_seg_service_local_subprocess(
         "window_size": 100,
         "connected_component_threshold": 2,
         "segmentation_threshold": 4,
+        "copy_output": True,
         "relion_options": {},
     }
     output_relion_options = dict(RelionServiceOptions())
@@ -213,8 +226,12 @@ def test_membrain_seg_service_local_subprocess(
         "--store-probabilities",
         "--store-connected-components",
     ]
-    assert mock_subprocess.call_count == 4
-    mock_subprocess.assert_any_call(membrain_command, capture_output=True)
+    mock_subprocess.assert_called_once_with(membrain_command, capture_output=True)
+
+    # Check output copy
+    assert (
+        tmp_path / "Segmentation/test_stack_aretomo.denoised_segmented.mrc"
+    ).is_file()
 
     # Check the images service request
     assert offline_transport.send.call_count == 5