Add recipe and tomogram processing for B24 tomogram alignment

pyproject.toml

@@ -47,6 +47,7 @@ dependencies = [
     "starfile",
     "stomp-py>8.1.1",
     "tifffile", # CLEM workflow
+    "txrm2tiff",
     "workflows>=3",
 ]
 [project.optional-dependencies]

recipes/b24-tomo-align.json

@@ -0,0 +1,100 @@
+{
+  "1": {
+    "output": 2,
+    "parameters": {
+      "dcid": "{dcid}",
+      "ispyb_command": "insert_tomogram",
+      "program_id": "{appid}",
+      "store_result": "ispyb_tomogram_id"
+    },
+    "queue": "ispyb_connector",
+    "service": "EMISPyB"
+  },
+  "2": {
+    "output": {
+      "denoise": 5,
+      "failure": 3,
+      "images": 6,
+      "ispyb_connector": 4,
+      "movie": 6,
+      "projxy": 6,
+      "projxz": 6,
+      "success": 7
+    },
+    "parameters": {
+      "dark_tol": 0,
+      "manual_tilt_offset": 0,
+      "out_bin": 1,
+      "pixel_size": 100,
+      "relion_options": {},
+      "stack_file": "{stack_file}",
+      "tilt_axis": 0,
+      "txrm_file": "{txrm_file}",
+      "wbp": 1
+    },
+    "queue": "tomo_align",
+    "service": "TomoAlign"
+  },
+  "3": {
+    "parameters": {
+      "ispyb_command": "update_processing_status",
+      "status_message": "processing failure",
+      "program_id": "{appid}",
+      "status": "failure"
+    },
+    "queue": "ispyb_connector",
+    "service": "EMISPyB"
+  },
+  "4": {
+    "parameters": {
+      "dcid": "{dcid}",
+      "ispyb_command": "multipart_message",
+      "program_id": "{appid}",
+      "tomogram_id": "$ispyb_tomogram_id"
+    },
+    "queue": "ispyb_connector",
+    "service": "EMISPyB"
+  },
+  "5": {
+    "output": {
+      "images": 6,
+      "ispyb_connector": 9,
+      "movie": 6,
+      "segmentation": 8
+    },
+    "queue": "denoise",
+    "service": "Denoise"
+  },
+  "6": {
+    "queue": "images",
+    "service": "Images"
+  },
+  "7": {
+    "parameters": {
+      "ispyb_command": "update_processing_status",
+      "status_message": "processing successful",
+      "program_id": "{appid}",
+      "status": "success"
+    },
+    "queue": "ispyb_connector",
+    "service": "EMISPyB"
+  },
+  "8": {
+    "output": {
+      "images": 6,
+      "ispyb_connector": 9,
+      "movie": 6
+    },
+    "queue": "segmentation",
+    "service": "MembrainSeg"
+  },
+  "9": {
+    "parameters": {
+      "ispyb_command": "insert_processed_tomogram",
+      "tomogram_id": "$ispyb_tomogram_id"
+    },
+    "queue": "ispyb_connector",
+    "service": "EMISPyB"
+  },
+  "start": [[1, []]]
+}

src/cryoemservices/services/tomo_align.py

@@ -48,6 +48,7 @@ class TomoParameters(BaseModel):
     aretomo_version: Literal[2, 3] = 3
     stack_file: str = Field(..., min_length=1)
     pixel_size: float
+    txrm_file: Optional[str] = None
     path_pattern: Optional[str] = None
     input_file_list: Optional[str] = None
     vol_z: Optional[int] = None
@@ -86,12 +87,15 @@ class TomoParameters(BaseModel):
     def check_only_one_is_provided(cls, values):
         input_file_list = values.get("input_file_list")
         path_pattern = values.get("path_pattern")
-        if not input_file_list and not path_pattern:
-            raise ValueError("input_file_list or path_pattern must be provided")
-        if input_file_list and path_pattern:
+        txrm_file = values.get("txrm_file")
+        if not input_file_list and not path_pattern and not txrm_file:
             raise ValueError(
-                "Message must only include one of 'path_pattern' and 'input_file_list'."
-                " Both are set or one has been set by the recipe."
+                "input_file_list or path_pattern or txrm_file must be provided"
+            )
+        if input_file_list and path_pattern and txrm_file:
+            raise ValueError(
+                "Message must only include one of "
+                "'path_pattern', 'input_file_list' or 'txrm_file'."
             )
         return values
 
@@ -127,6 +131,7 @@ class TomoAlign(CommonService):
 
     # Values to extract for ISPyB
     input_file_list_of_lists: List[Any]
+    tilt_angles: dict
     refined_tilts: List[float]
     x_shift: List[float]
     y_shift: List[float]
@@ -139,6 +144,7 @@ class TomoAlign(CommonService):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.tilt_angles = {}
         self.refined_tilts = []
         self.rot_centre_z_list = []
 
@@ -254,15 +260,13 @@ def _tilt(file_list_for_tilts):
             # If no volume provided, set it to the maximum we allow
             tomo_params.vol_z = tomo_params.max_vol
 
-        # Update the relion options
-        tomo_params.relion_options = update_relion_options(
-            tomo_params.relion_options, dict(tomo_params)
-        )
-        tomo_params.relion_options.pixel_size_downscaled = (
-            tomo_params.pixel_size * tomo_params.out_bin
-        )
+        # Get the names of the output files expected
+        alignment_output_dir = Path(tomo_params.stack_file).parent
+        Path(tomo_params.stack_file).parent.mkdir(parents=True, exist_ok=True)
+        stack_name = str(Path(tomo_params.stack_file).stem)
 
         # Convert a path pattern into a file list
+        self.tilt_angles = {}
         if tomo_params.path_pattern:
             directory = Path(tomo_params.path_pattern).parent
 
@@ -276,10 +280,34 @@ def _tilt(file_list_for_tilts):
         elif tomo_params.input_file_list:
             file_list = ast.literal_eval(tomo_params.input_file_list)
             self.input_file_list_of_lists = file_list
+        elif tomo_params.txrm_file and Path(tomo_params.txrm_file).is_file():
+            tomo_params.pixel_size = self.convert_txrm_to_stack(
+                tomo_params.txrm_file, tomo_params.stack_file
+            )
+            self.input_file_list_of_lists = []
+            # Check we now have the expected stack file
+            if not Path(tomo_params.stack_file).is_file():
+                self.log.warning("Stack file generation failed")
+                rw.transport.nack(header)
+                return
+        else:
+            self.log.warning(f"Invalid input or {tomo_params.txrm_file} is not a file")
+            rw.transport.nack(header)
+            return
 
-        self.log.info(f"Input list {self.input_file_list_of_lists}")
+        self.log.info(
+            f"Input list {self.input_file_list_of_lists or tomo_params.txrm_file}"
+        )
         self.input_file_list_of_lists.sort(key=_tilt)
 
+        # Update the relion options
+        tomo_params.relion_options = update_relion_options(
+            tomo_params.relion_options, dict(tomo_params)
+        )
+        tomo_params.relion_options.pixel_size_downscaled = (
+            tomo_params.pixel_size * tomo_params.out_bin
+        )
+
         # Find all the tilt angles and remove duplicates
         tilt_dict: dict = {}
         for tilt in self.input_file_list_of_lists:
@@ -333,23 +361,26 @@ def _tilt(file_list_for_tilts):
             self.input_file_list_of_lists.remove(self.input_file_list_of_lists[index])
 
         # Find the input image dimensions
-        with mrcfile.open(self.input_file_list_of_lists[0][0]) as mrc:
-            mrc_header = mrc.header
-        # x and y get flipped on tomogram creation
+        if self.input_file_list_of_lists:
+            # Read the first image if a list is given
+            with mrcfile.open(self.input_file_list_of_lists[0][0]) as mrc:
+                mrc_header = mrc.header
+
+        else:
+            # Read shape of the stack if using txrm
+            with mrcfile.open(tomo_params.stack_file) as mrc:
+                mrc_header = mrc.header
         tomo_params.relion_options.tomo_size_x = int(mrc_header["nx"])
         tomo_params.relion_options.tomo_size_y = int(mrc_header["ny"])
+
+        # Scale size of output
         scaled_x_size = tomo_params.relion_options.tomo_size_x / int(
             tomo_params.out_bin
         )
         scaled_y_size = tomo_params.relion_options.tomo_size_y / int(
             tomo_params.out_bin
         )
 
-        # Get the names of the output files expected
-        alignment_output_dir = Path(tomo_params.stack_file).parent
-        Path(tomo_params.stack_file).parent.mkdir(parents=True, exist_ok=True)
-        stack_name = str(Path(tomo_params.stack_file).stem)
-
         project_dir_search = re.search(".+/job[0-9]+/", tomo_params.stack_file)
         job_num_search = re.search("/job[0-9]+", tomo_params.stack_file)
         if project_dir_search and job_num_search:
@@ -360,38 +391,35 @@ def _tilt(file_list_for_tilts):
             rw.transport.nack(header)
             return
 
-        # Stack the tilts with newstack
-        newstack_path = alignment_output_dir / f"{stack_name}_newstack.txt"
-        newstack_result = self.newstack(tomo_params, newstack_path)
-        if newstack_result.returncode:
-            self.log.error(
-                f"Newstack failed with exitcode {newstack_result.returncode}:\n"
-                + newstack_result.stderr.decode("utf8", "replace")
-            )
-            rw.transport.nack(header)
-            return
+        if self.input_file_list_of_lists:
+            # Stack the tilts with newstack
+            newstack_path = alignment_output_dir / f"{stack_name}_newstack.txt"
+            newstack_result = self.newstack(tomo_params, newstack_path)
+            if newstack_result.returncode:
+                self.log.error(
+                    f"Newstack failed with exitcode {newstack_result.returncode}:\n"
+                    + newstack_result.stderr.decode("utf8", "replace")
+                )
+                rw.transport.nack(header)
+                return
 
         # Set up the angle file needed for dose weighting
-        angle_file = (
-            Path(tomo_params.stack_file).parent
-            / f"{Path(tomo_params.stack_file).stem}_TLT.txt"
-        )
-        tilt_angles = {}
+        angle_file = Path(tomo_params.stack_file).parent / f"{stack_name}_TLT.txt"
         for i in range(len(self.input_file_list_of_lists)):
             tilt_index = _get_tilt_number_v5_12(
                 Path(self.input_file_list_of_lists[i][0])
             )
             tilt_index -= len(np.where(removed_tilt_numbers < tilt_index)[0])
-            tilt_angles[tilt_index] = float(self.input_file_list_of_lists[i][1])
+            self.tilt_angles[tilt_index] = float(self.input_file_list_of_lists[i][1])
         with open(angle_file, "w") as angfile:
-            for tilt_id in tilt_angles.keys():
+            for tilt_id in self.tilt_angles.keys():
                 if tomo_params.aretomo_version == 3 and tomo_params.manual_tilt_offset:
                     # AreTomo3 performs better with pre-shifted angles
                     angfile.write(
-                        f"{tilt_angles[tilt_id] + tomo_params.manual_tilt_offset:.2f}  {int(tilt_id)}\n"
+                        f"{self.tilt_angles[tilt_id] + tomo_params.manual_tilt_offset:.2f}  {int(tilt_id)}\n"
                     )
                 else:
-                    angfile.write(f"{tilt_angles[tilt_id]:.2f}  {int(tilt_id)}\n")
+                    angfile.write(f"{self.tilt_angles[tilt_id]:.2f}  {int(tilt_id)}\n")
 
         # Do alignment with AreTomo
         aretomo_output_path = alignment_output_dir / f"{stack_name}_Vol.mrc"
@@ -436,7 +464,9 @@ def _tilt(file_list_for_tilts):
             node_creator_parameters = {
                 "experiment_type": "tomography",
                 "job_type": self.job_type,
-                "input_file": self.input_file_list_of_lists[0][0],
+                "input_file": self.input_file_list_of_lists[0][0]
+                if self.input_file_list_of_lists
+                else tomo_params.txrm_file,
                 "output_file": str(aretomo_output_path),
                 "relion_options": dict(tomo_params.relion_options),
                 "command": " ".join(aretomo_command),
@@ -577,18 +607,19 @@ def _tilt(file_list_for_tilts):
         im_diff = 0
         # TiltImageAlignment (one per movie)
         node_creator_params_list = []
-        (project_dir / f"ExcludeTiltImages/job{job_number - 2:03}").mkdir(
-            parents=True, exist_ok=True
-        )
-        if not (
-            project_dir / f"ExcludeTiltImages/job{job_number - 2:03}/tilts"
-        ).is_symlink():
-            (
+        if self.input_file_list_of_lists:
+            (project_dir / f"ExcludeTiltImages/job{job_number - 2:03}").mkdir(
+                parents=True, exist_ok=True
+            )
+            if not (
                 project_dir / f"ExcludeTiltImages/job{job_number - 2:03}/tilts"
-            ).symlink_to(project_dir / "MotionCorr/job002/Movies")
-        (project_dir / f"AlignTiltSeries/job{job_number - 1:03}").mkdir(
-            parents=True, exist_ok=True
-        )
+            ).is_symlink():
+                (
+                    project_dir / f"ExcludeTiltImages/job{job_number - 2:03}/tilts"
+                ).symlink_to(project_dir / "MotionCorr/job002/Movies")
+            (project_dir / f"AlignTiltSeries/job{job_number - 1:03}").mkdir(
+                parents=True, exist_ok=True
+            )
         if self.rot:
             tomo_params.relion_options.tilt_axis_angle = self.rot
         for im, movie in enumerate(self.input_file_list_of_lists):
@@ -741,7 +772,6 @@ def _tilt(file_list_for_tilts):
                     self.thickness_pixels * tomo_params.pixel_size
                 )
             rw.send_to("images", images_call_params)
-            print(images_call_params)
 
         # Forward results to denoise service
         self.log.info(f"Sending to denoise service {aretomo_output_path}")
@@ -775,6 +805,42 @@ def _tilt(file_list_for_tilts):
         self.log.info(f"Done tomogram alignment for {tomo_params.stack_file}")
         rw.transport.ack(header)
 
+    def convert_txrm_to_stack(self, txrm_file: str, stack_file: str) -> float:
+        from txrm2tiff.inspector import Inspector
+        from txrm2tiff.main import convert_and_save
+        from txrm2tiff.txrm import open_txrm
+        from txrm2tiff.txrm_functions.general import read_stream
+        from txrm2tiff.xradia_properties.enums import XrmDataTypes
+
+        # Read the tilt angles and pixel size from the txrm
+        with open_txrm(
+            txrm_file, load_images=False, load_reference=False, strict=False
+        ) as txrm:
+            inspector = Inspector(txrm)
+            angles = read_stream(
+                inspector.txrm.ole,
+                "ImageInfo/Angles",
+                XrmDataTypes.XRM_FLOAT,
+                strict=True,
+            )
+            pixel_size_microns = read_stream(
+                inspector.txrm.ole,
+                "ImageInfo/PixelSize",
+                XrmDataTypes.XRM_FLOAT,
+                strict=True,
+            )
+        self.tilt_angles = dict(enumerate(angles))
+
+        # Convert the txrm to a tiff stack, then convert to mrc for aretomo
+        tifftomo = Path(stack_file).with_suffix(".tiff")
+        convert_and_save(txrm_file, str(tifftomo), custom_reference=None)
+        # Let this run, and check later if the output file exists
+        subprocess.run(["tif2mrc", str(tifftomo), stack_file])
+        tifftomo.unlink(missing_ok=True)
+        if pixel_size_microns:
+            return pixel_size_microns[0] * 1000
+        return 100
+
     def newstack(self, tomo_parameters: TomoParameters, newstack_path: Path):
         """
         Construct file containing a list of files