Extract particles from tomograms

pyproject.toml

@@ -96,6 +96,7 @@ torch = [
     EMISPyB = "cryoemservices.services.ispyb_connector:EMISPyB"
     Extract = "cryoemservices.services.extract:Extract"
     ExtractClass = "cryoemservices.services.extract_class:ExtractClass"
+    ExtractSubTomo = "cryoemservices.services.extract_2d_subtomo:ExtractSubTomoFor2D"
     IceBreaker = "cryoemservices.services.icebreaker:IceBreaker"
     Images = "cryoemservices.services.images:Images"
     MembrainSeg = "cryoemservices.services.membrain_seg:MembrainSeg"

src/cryoemservices/pipeliner_plugins/combine_star_files.py

@@ -44,7 +44,7 @@ def combine_star_files(files_to_process: List[Path], output_dir: Path):
     ):
         for line_counter in range(50):
             line = full_starfile.readline()
-            if line.startswith("opticsGroup"):
+            if "opticsGroup" in line:
                 reference_optics = line.split()
             if not line:
                 break
@@ -71,7 +71,7 @@ def combine_star_files(files_to_process: List[Path], output_dir: Path):
                 optics_line = added_starfile.readline()
                 if not optics_line:
                     raise IndexError(f"Cannot find optics group in {split_file}")
-                if optics_line.startswith("opticsGroup"):
+                if "opticsGroup" in optics_line:
                     new_optics = optics_line.split()
                     break
 
@@ -104,8 +104,14 @@ def combine_star_files(files_to_process: List[Path], output_dir: Path):
                 particle_line = added_starfile.readline()
                 if not particle_line:
                     break
+                if "opticsGroup" in particle_line:
+                    # Skip the optics group
+                    continue
+                if particle_line.startswith(("_", "loop_", "data_", "#")):
+                    # Skip all block and loop header lines
+                    continue
                 particle_split_line = particle_line.split()
-                if len(particle_split_line) > 0 and particle_split_line[0][0].isdigit():
+                if len(particle_split_line) > 0:
                     file_particles_count += 1
                     total_particles += 1
                     particles_file.write(particle_line)

src/cryoemservices/services/cryolo.py

@@ -23,7 +23,7 @@ class CryoloParameters(BaseModel):
     input_path: str = Field(..., min_length=1)
     output_path: str = Field(..., min_length=1)
     experiment_type: str
-    pixel_size: Optional[float] = None
+    pixel_size: float
     cryolo_box_size: int = 160
     cryolo_model_weights: str = "gmodel_phosnet_202005_N63_c17.h5"
     cryolo_threshold: float = 0.3
@@ -40,6 +40,7 @@ class CryoloParameters(BaseModel):
     mc_uuid: Optional[int] = None
     app_id: Optional[int] = None
     picker_uuid: Optional[int] = None
+    raw_tomogram: Optional[str] = None
     relion_options: RelionServiceOptions
     ctf_values: dict = {}
 
@@ -53,13 +54,12 @@ def is_spa_or_tomo(cls, experiment):
     @model_validator(mode="after")
     def check_spa_has_uuids_and_pixel_size(self):
         if self.experiment_type == "spa" and (
-            self.mc_uuid is None or self.picker_uuid is None or not self.pixel_size
+            self.mc_uuid is None or self.picker_uuid is None
         ):
             raise ValueError(
                 "In SPA mode the following must be provided: "
                 f"mc_uuid (given {self.mc_uuid}), "
-                f"picker_uuid (given {self.picker_uuid}), "
-                f"pixel_size (given {self.pixel_size})"
+                f"picker_uuid (given {self.picker_uuid})"
             )
         return self
 
@@ -329,6 +329,29 @@ def cryolo(self, rw, header: dict, message: dict):
             }
             rw.send_to("ispyb_connector", ispyb_parameters_tomo)
 
+            # Get the diameters of the particles in Angstroms for Murfey
+            try:
+                cbox_file = cif.read_file(cryolo_params.output_path)
+                cbox_block = cbox_file.find_block("cryolo")
+                cbox_sizes = (
+                    np.array(cbox_block.find_loop("_EstWidth"), dtype=float)
+                    + np.array(cbox_block.find_loop("_EstHeight"), dtype=float)
+                ) / 2
+                cryolo_particle_sizes = cbox_sizes * cryolo_params.pixel_size
+            except (FileNotFoundError, OSError, AttributeError):
+                cryolo_particle_sizes = []
+
+            # Send to murfey for extraction
+            extraction_parameters = {
+                "register": "picked_tomogram",
+                "tomogram": cryolo_params.raw_tomogram or cryolo_params.input_path,
+                "cbox_3d": cryolo_params.output_path,
+                "pixel_size": cryolo_params.pixel_size,
+                "particle_diameters": list(cryolo_particle_sizes),
+                "particle_count": len(cryolo_particle_sizes),
+            }
+            rw.send_to("murfey_feedback", extraction_parameters)
+
             self.log.info(
                 f"Done {self.job_type} {cryolo_params.experiment_type} "
                 f"for {cryolo_params.input_path}."
@@ -345,14 +368,11 @@ def cryolo(self, rw, header: dict, message: dict):
                 )
             )
             cbox_block = cbox_file.find_block("cryolo")
-            cbox_sizes = np.append(
-                np.array(cbox_block.find_loop("_EstWidth"), dtype=float),
-                np.array(cbox_block.find_loop("_EstHeight"), dtype=float),
-            )
-            cbox_confidence = np.append(
-                np.array(cbox_block.find_loop("_Confidence"), dtype=float),
-                np.array(cbox_block.find_loop("_Confidence"), dtype=float),
-            )
+            cbox_sizes = (
+                np.array(cbox_block.find_loop("_EstWidth"), dtype=float)
+                + np.array(cbox_block.find_loop("_EstHeight"), dtype=float)
+            ) / 2
+            cbox_confidence = np.array(cbox_block.find_loop("_Confidence"), dtype=float)
 
             # Select only a fraction of particles based on confidence if requested
             if (

src/cryoemservices/services/ctffind.py

@@ -5,6 +5,8 @@
 from pathlib import Path
 from typing import Any, Optional
 
+import mrcfile
+import numpy as np
 from pydantic import BaseModel, Field, ValidationError, field_validator
 from workflows.recipe import wrap_subscribe
 
@@ -352,3 +354,60 @@ def ctf_find(self, rw, header: dict, message: dict):
 
         self.log.info(f"Done {self.job_type} for {ctf_params.input_image}.")
         rw.transport.ack(header)
+
+
+def ctf(
+    num_pixels_x: int, num_pixels_y: int, pixel_size: float, defocus: float
+) -> np.array:
+    C_s: float = 2.7e7
+    wavelength: float = 0.0196
+    grid = np.meshgrid(
+        np.fft.fftfreq(num_pixels_x, pixel_size),
+        np.fft.fftfreq(num_pixels_y, pixel_size),
+    )
+    grid = np.fft.fftshift(grid)
+    ksq = grid[0] ** 2 + grid[1] ** 2
+    w = -defocus * wavelength * ksq / 2 + C_s * wavelength**3 + ksq**2 / 4
+    return np.sin(2 * np.pi * w)
+
+
+def ctf_micrograph(mic_file, output_file, defocus_u, defocus_v):
+    with mrcfile.open(mic_file) as mrc:
+        im = mrc.data
+        num_pixels_x = int(mrc.header.mx)
+        num_pixels_y = int(mrc.header.my)
+        pixel_size = mrc.header.cella.x / mrc.header.mx
+    dtype = im.dtype
+
+    fft = np.fft.fft2(im)
+    fft = np.fft.fftshift(fft)
+
+    ctf_mask = ctf(
+        num_pixels_x, num_pixels_y, pixel_size, np.sqrt(defocus_u**2 + defocus_v**2)
+    )
+    ctf_mask[ctf_mask < 0] = -1
+    ctf_mask[ctf_mask >= 0] = 1
+
+    fft = fft * ctf_mask
+
+    fft_shifted = np.fft.ifftshift(fft)
+    corrected_im = np.real(np.fft.ifft2(fft_shifted))
+    corrected_im = corrected_im.astype(dtype)
+    mrcfile.write(output_file, corrected_im, overwrite=True)
+
+
+def ctf_of_tomo_star(root_dir, tomo_star, output_dir):
+    from gemmi import cif
+
+    data = cif.read_file(str(Path(root_dir) / tomo_star))
+    mics = list(data.sole_block().find_loop("_rlnMicrographName"))
+    defocus_u = list(data.sole_block().find_loop("_rlnDefocusU"))
+    defocus_v = list(data.sole_block().find_loop("_rlnDefocusV"))
+
+    for i, mic in enumerate(mics):
+        ctf_micrograph(
+            Path(root_dir) / mic,
+            Path(root_dir) / output_dir / Path(mic).name,
+            float(defocus_u[i]),
+            float(defocus_v[i]),
+        )

src/cryoemservices/services/denoise.py

@@ -297,10 +297,12 @@ def denoise(self, rw, header: dict, message: dict):
             "relion_options": dict(denoise_params.relion_options),
         }
         cryolo_parameters = {
+            "raw_tomogram": denoise_params.volume,
             "input_path": str(denoised_full_path),
             "output_path": str(cryolo_dir / f"CBOX_3D/{denoised_full_path.stem}.cbox"),
             "experiment_type": "tomography",
             "cryolo_box_size": 40,
+            "pixel_size": str(denoise_params.relion_options.pixel_size_downscaled),
             "relion_options": dict(denoise_params.relion_options),
         }
         rw.send_to("segmentation", segmentation_parameters)