Merge pull request #203 from hanjinliu/relion-job

Support importing RELION 3D class job

Author: hanjinliu

cylindra/widgets/batch/main.py

@@ -189,6 +189,7 @@ def _new_projects_from_table(
                 prj.project_path = save_path
             self.constructor.projects._add(prj.project_path)
         self.save_batch_project(save_path=save_root)
+        self.show()
 
     @set_design(text=capitalize, location=constructor)
     @thread_worker

cylindra/widgets/main.py

@@ -3064,6 +3064,8 @@ def rotate_molecule_toward_spline(
             the splines.
         """
         layer = assert_layer(layer, self.parent_viewer)
+        if self.splines.count() == 0:
+            raise ValueError("No spline exists in the tomogram.")
         mole = layer.molecules
         if spline_id_column == "":
             dist_stack = []

cylindra_builtins/relion/_relion_utils.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+try:
+    import starfile
+except ImportError:
+
+    class _starfile_module:
+        def __getattr__(self, name):
+            raise ImportError(
+                "The 'starfile' package is required for RELION I/O functions. "
+                "Please install it via 'pip install starfile'."
+            )
+
+    starfile = _starfile_module()
+
+
+def get_optimisation_set_star(job_dir_path: Path) -> Path:
+    if fp := _relion_job_get_last(job_dir_path, "run_it*_optimisation_set.star"):
+        return fp
+    raise ValueError(
+        f"No optimisation set star files found in {job_dir_path}. "
+        "Please ensure at least one iteration has finished."
+    )
+
+
+def get_run_data_star(job_dir_path: Path) -> Path | None:
+    return _relion_job_get_last(job_dir_path, "run_it*_data.star")
+
+
+def _relion_job_get_last(job_dir_path: Path, pattern: str) -> Path | None:
+    path_list = sorted(job_dir_path.glob(pattern), key=lambda p: p.stem)
+    if len(path_list) == 0:
+        return None
+    return path_list[-1]
+
+
+def relion_project_path(path: Path) -> Path:
+    return path.parent.parent
+
+
+def get_job_type(job_dir: Path) -> str:
+    """Determine the type of RELION job based on the directory structure."""
+    if (job_star_path := job_dir / "job.star").exists():
+        return starfile.read(job_star_path, always_dict=True)["job"]["rlnJobTypeLabel"]
+    raise ValueError(f"{job_dir} is not a RELION job folder.")
+
+
+def shape_to_center_zyx(shape: tuple[int, int, int], scale: float) -> np.ndarray:
+    return (np.array(shape) / 2 - 1) * scale
+
+
+def strip_relion5_prefix(name: str):
+    """Strip the RELION 5.0 "rec" prefix from the name."""
+    if name.startswith("rec_"):
+        name = name[4:]
+    if "." in name:
+        name = name.split(".")[0]
+    return name

cylindra_builtins/relion/io.py

@@ -10,26 +10,14 @@
 from magicclass.types import Optional, Path
 from magicclass.widgets import ConsoleTextEdit
 
-try:
-    import starfile
-except ImportError:
-
-    class _starfile_module:
-        def __getattr__(self, name):
-            raise ImportError(
-                "The 'starfile' package is required for RELION I/O functions. "
-                "Please install it via 'pip install starfile'."
-            )
-
-    starfile = _starfile_module()
-
-
 from cylindra.const import FileFilter, nm
 from cylindra.core import ACTIVE_WIDGETS
 from cylindra.plugin import register_function
 from cylindra.widget_utils import add_molecules
 from cylindra.widgets import CylindraMainWidget
 from cylindra.widgets._annotated import MoleculesLayersType, assert_list_of_layers
+from cylindra_builtins.relion import _relion_utils
+from cylindra_builtins.relion._relion_utils import starfile
 
 if TYPE_CHECKING:
     from cylindra.components.tomogram import CylTomogram
@@ -131,7 +119,9 @@ def save_molecules(
     """
     save_path = Path(save_path)
     layers = assert_list_of_layers(layers, ui.parent_viewer)
-    tomo_name = tomo_name_override or _strip_relion5_prefix(ui.tomogram.image.name)
+    tomo_name = tomo_name_override or _relion_utils.strip_relion5_prefix(
+        ui.tomogram.image.name
+    )
     df = _mole_to_star_df(
         [layer.molecules for layer in layers],
         ui.tomogram,
@@ -261,7 +251,7 @@ def _iter_dataframe_from_path_sets(
     for path_info in path_sets:
         path_info = PathInfo(*path_info)
         prj = path_info.project_instance(missing_ok=False)
-        tomo_name = _strip_relion5_prefix(path_info.image.stem)
+        tomo_name = _relion_utils.strip_relion5_prefix(path_info.image.stem)
         img = path_info.lazy_imread()
         tomo = CylTomogram.from_image(
             img,
@@ -412,8 +402,8 @@ def _parse_relion_job(path, project_root):
     if path.name != "job.star" or not path.is_file() or not path.exists():
         raise ValueError(f"Path must be an existing RELION job.star file, got {path}")
     job_dir_path = Path(path).parent
-    rln_project_path = _relion_project_path(job_dir_path)
-    jobtype = _get_job_type(job_dir_path)
+    rln_project_path = _relion_utils.relion_project_path(job_dir_path)
+    jobtype = _relion_utils.get_job_type(job_dir_path)
     if project_root is None:
         project_root = job_dir_path / "cylindra"
     if jobtype in ("relion.reconstructtomograms", "relion.denoisetomo"):
@@ -438,36 +428,21 @@ def _parse_relion_job(path, project_root):
         paths, scales, moles, tilt_models = _parse_optimisation_star(
             opt_star_path, rln_project_path
         )
-    elif jobtype in ("relion.initialmodel.tomo", "relion.refine3d.tomo"):
-        opt_set_path_list = sorted(
-            job_dir_path.glob("run_it*_optimisation_set.star"),
-            key=lambda p: p.stem,
-        )
-        if len(opt_set_path_list) == 0:
-            raise ValueError(
-                f"No optimisation set star files found in {job_dir_path}. "
-                "Please ensure at least one iteration has finished."
-            )
-        opt_star_path = opt_set_path_list[-1]
+    elif jobtype in (
+        "relion.initialmodel.tomo",
+        "relion.class3d",
+        "relion.refine3d.tomo",
+    ):
+        opt_star_path = _relion_utils.get_optimisation_set_star(job_dir_path)
+        run_data_path = _relion_utils.get_run_data_star(job_dir_path)
         paths, scales, moles, tilt_models = _parse_optimisation_star(
-            opt_star_path, rln_project_path
+            opt_star_path, rln_project_path, run_data_path
         )
     else:
         raise ValueError(f"Job {job_dir_path.name} is not a supported RELION job.")
     return project_root, paths, scales, moles, tilt_models
 
 
-def _relion_project_path(path: Path) -> Path:
-    return path.parent.parent
-
-
-def _get_job_type(job_dir: Path) -> str:
-    """Determine the type of RELION job based on the directory structure."""
-    if (job_star_path := job_dir / "job.star").exists():
-        return starfile.read(job_star_path, always_dict=True)["job"]["rlnJobTypeLabel"]
-    raise ValueError(f"{job_dir} is not a RELION job folder.")
-
-
 class TomogramStar:
     """Object to parse tomograms.star file."""
 
@@ -518,12 +493,18 @@ def iter_tomo_shapes(self) -> Iterator[tuple[float, float, float]]:
         yield from zip(nzs, nys, nxs, strict=False)
 
 
-def _parse_optimisation_star(opt_star_path: Path, rln_project_path: Path):
+def _parse_optimisation_star(
+    opt_star_path: Path,
+    rln_project_path: Path,
+    run_data_path: Path | None = None,
+):
     paths = []
     scales = []
     molecules = []
     tilt_models = []
-    for item in _iter_from_optimisation_star(opt_star_path, rln_project_path):
+    for item in _iter_from_optimisation_star(
+        opt_star_path, rln_project_path, run_data_path
+    ):
         paths.append(rln_project_path / item.tomo_path)
         scales.append(item.scale)
         molecules.append(item.molecules)
@@ -543,6 +524,7 @@ class OptimizationSetItem:
 def _iter_from_optimisation_star(
     path: Path,
     rln_project_path: Path,
+    run_data_path: Path | None = None,
 ) -> "Iterator[OptimizationSetItem]":
     opt_star_df = starfile.read(path)
     if isinstance(opt_star_df, pd.DataFrame):
@@ -557,13 +539,16 @@ def _iter_from_optimisation_star(
     tomo_paths = list(tomostar.iter_tomo_paths(rln_project_path))
     tomo_shapes = list(tomostar.iter_tomo_shapes())
     tilt_models = list(tomostar.iter_tilt_models(rln_project_path))
-    particles_df = starfile.read(rln_project_path / particles_path)
+    if run_data_path is None:
+        particles_df = starfile.read(rln_project_path / particles_path)
+    else:
+        particles_df = starfile.read(run_data_path)
 
     if isinstance(particles_df, dict):
         particles_df = particles_df["particles"]
     assert isinstance(particles_df, pd.DataFrame)
     name_to_center_map = {
-        tomo_name: _shape_to_center_zyx(tomo_shape, sc_nm)
+        tomo_name: _relion_utils.shape_to_center_zyx(tomo_shape, sc_nm)
         for tomo_name, tomo_shape, sc_nm in zip(
             tomo_names, tomo_shapes, scale_nm, strict=False
         )
@@ -592,7 +577,7 @@ def _iter_from_optimisation_star(
 
 def _read_star(path: str, tomo: "CylTomogram") -> dict[str, Molecules]:
     fpath = Path(path)
-    center_zyx = _shape_to_center_zyx(tomo.image.shape, tomo.scale)
+    center_zyx = _relion_utils.shape_to_center_zyx(tomo.image.shape, tomo.scale)
     star = starfile.read(fpath)
     if not isinstance(star, dict):
         star = {"particles": star}  # assume particles block
@@ -648,19 +633,6 @@ def _particles_to_molecules(
     return {default_key: mole}
 
 
-def _shape_to_center_zyx(shape: tuple[int, int, int], scale: nm) -> np.ndarray:
-    return (np.array(shape) / 2 - 1) * scale
-
-
-def _strip_relion5_prefix(name: str):
-    """Strip the RELION 5.0 rec prefix from the name."""
-    if name.startswith("rec_"):
-        name = name[4:]
-    if "." in name:
-        name = name.split(".")[0]
-    return name
-
-
 def _mole_to_star_df(
     moles: list[Molecules],
     tomo: "CylTomogram",
@@ -682,7 +654,9 @@ def _mole_to_star_df(
         ROT_COLUMNS[2]: euler_angle[:, 2],
     }
     if centered:
-        centerz, centery, centerx = _shape_to_center_zyx(tomo.image.shape, scale)
+        centerz, centery, centerx = _relion_utils.shape_to_center_zyx(
+            tomo.image.shape, scale
+        )
         # Angstrom
         _pos_dict = {
             POS_CENTERED[2]: (mole.pos[:, 2] - centerx + orig.x) * 10,

cylindra_builtins/relion/tests/test_jobs/Class3D/job_class3d/job.star

@@ -0,0 +1,84 @@
+
+# version 50001
+
+data_job
+
+_rlnJobTypeLabel             relion.class3d
+_rlnJobIsContinue                       0
+_rlnJobIsTomo                           1
+
+
+# version 50001
+
+data_joboptions_values
+
+loop_
+_rlnJobOptionVariable #1
+_rlnJobOptionValue #2
+allow_coarser         No
+ctf_intact_first_peak         No
+do_apply_helical_symmetry        Yes
+  do_blush         No
+do_combine_thru_disc         No
+do_ctf_correction        Yes
+do_fast_subsets         No
+  do_helix         No
+do_local_ang_searches        Yes
+do_local_search_helical_symmetry         No
+   do_pad1         No
+do_parallel_discio        Yes
+do_preread_images         No
+  do_queue         No
+do_zero_mask        Yes
+dont_skip_align        Yes
+   fn_cont         ""
+   fn_mask         ""
+    fn_ref Refine3D/job015/run_class001.mrc
+   gpu_ids         ""
+helical_nr_asu          1
+helical_range_distance         -1
+helical_rise_inistep          0
+helical_rise_initial          0
+helical_rise_max          0
+helical_rise_min          0
+helical_tube_inner_diameter         -1
+helical_tube_outer_diameter         -1
+helical_twist_inistep          0
+helical_twist_initial          0
+helical_twist_max          0
+helical_twist_min          0
+helical_z_percentage         30
+highres_limit         -1
+in_optimisation         ""
+in_particles XXX
+in_tomograms YYY
+in_trajectories         ""
+  ini_high         60
+keep_tilt_prior_fixed        Yes
+min_dedicated         24
+nr_classes          9
+   nr_iter         25
+    nr_mpi          5
+   nr_pool         30
+nr_threads          6
+offset_range          5
+offset_step          1
+other_args         ""
+particle_diameter        230
+      qsub     sbatch
+qsubscript /public/EM/RELION/relion-slurm-gpu-4.0.csh
+ queuename    openmpi
+ range_psi         10
+ range_rot         -1
+range_tilt         15
+ref_correct_greyscale        Yes
+ relax_sym         ""
+  sampling "1.8 degrees"
+scratch_dir $RELION_SCRATCH_DIR
+sigma_angles          5
+sigma_tilt         10
+  sym_name         C1
+ tau_fudge          1
+trust_ref_size        Yes
+use_direct_entries        Yes
+   use_gpu        Yes