pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__cp311-cp311-macosx_15_0_arm64.whl

scripts/eval.py

@@ -0,0 +1,93 @@
+#!python3
+""" Apply tme.preprocessor.Preprocessor methods to an input file based
+    on a provided yaml configuration obtaiend from preprocessor_gui.py.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+import yaml
+import argparse
+import textwrap
+from tme import Preprocessor, Density
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description=textwrap.dedent(
+            """
+        Apply preprocessing to an input file based on a provided YAML configuration.
+
+        Expected YAML file format:
+        ```yaml
+        <method_name>:
+            <parameter1>: <value1>
+            <parameter2>: <value2>
+            ...
+        ```
+        """
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "-i",
+        "--input_file",
+        type=str,
+        required=True,
+        help="Path to the input data file in CCP4/MRC format.",
+    )
+    parser.add_argument(
+        "-y",
+        "--yaml_file",
+        type=str,
+        required=True,
+        help="Path to the YAML configuration file.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output_file",
+        type=str,
+        required=True,
+        help="Path to output file in CPP4/MRC format..",
+    )
+    parser.add_argument(
+        "--compress", action="store_true", help="Compress the output file using gzip."
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+def main():
+    args = parse_args()
+    with open(args.yaml_file, "r") as f:
+        preprocess_settings = yaml.safe_load(f)
+
+    if len(preprocess_settings) > 1:
+        raise NotImplementedError(
+            "Multiple preprocessing methods specified. "
+            "The script currently supports one method at a time."
+        )
+
+    method_name = list(preprocess_settings.keys())[0]
+    if not hasattr(Preprocessor, method_name):
+        raise ValueError(f"Method {method_name} does not exist in Preprocessor.")
+
+    density = Density.from_file(args.input_file)
+    output = density.empty
+
+    method_params = preprocess_settings[method_name]
+    preprocessor = Preprocessor()
+    method = getattr(preprocessor, method_name, None)
+    if not method:
+        raise ValueError(
+            f"{method} does not exist in dge.preprocessor.Preprocessor class."
+        )
+
+    output.data = method(template=density.data, **method_params)
+    output.to_file(args.output_file, gzip=args.compress)
+
+
+if __name__ == "__main__":
+    main()

scripts/extract_candidates.py

@@ -1,19 +1,22 @@
 #!python3
-""" Prepare orientations stack for refinement.
+"""Prepare orientations stack for refinement.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 import argparse
-from os.path import splitext
+from os import unlink
+from os.path import splitext, basename
 
 import numpy as np
+from collections import defaultdict
 
+from tme.parser import StarParser
 from tme import Density, Orientations
-from tme.matching_utils import (
-    generate_tempfile_name,
-    rotation_aligning_vectors,
+from tme.matching_utils import generate_tempfile_name
+from tme.rotations import (
+    align_vectors,
     euler_from_rotationmatrix,
     euler_to_rotationmatrix,
 )
@@ -25,7 +28,7 @@ class ProgressBar:
     """
 
     def __init__(self, message: str, nchars: int, total: int):
-        self._size = nchars - len(message) - (len(str(total))+2) * 2
+        self._size = nchars - len(message) - (len(str(total)) + 2) * 2
         self._message = message
         self._total = total
 
@@ -45,20 +48,14 @@ def parse_args():
     )
 
     io_group = parser.add_argument_group("Input / Output")
-    io_group.add_argument(
-        "--target",
-        required=True,
-        type=str,
-        help="Extract candidates from this target.",
-    )
     io_group.add_argument(
         "--orientations",
         required=True,
         type=str,
-        help="Path to file generated by postprocess.py using output_format orientations.",
+        help="Star file with picks and micrograph names.",
     )
     io_group.add_argument(
-        "--orientations_sampling",
+        "--orientations-scaling",
         required=False,
         type=float,
         default=1.0,
@@ -67,54 +64,59 @@ def parse_args():
     )
     io_group.add_argument(
         "-o",
-        "--output_file",
+        "--output-prefix",
         required=True,
         type=str,
-        help="Path to write output H5 file.",
+        help="Output prefix to use.",
     )
 
     alignment_group = parser.add_argument_group("Alignment")
     alignment_group.add_argument(
-        "--align_orientations",
+        "--align-orientations",
         action="store_true",
         required=False,
         help="Whether to align extracted orientations based on their angles. Allows "
         "for efficient subsequent sampling of cone angles.",
     )
     alignment_group.add_argument(
-        "--angles_are_vector",
+        "--angles-are-vector",
         action="store_true",
         required=False,
         help="Considers euler_z euler_y, euler_x as vector that will be rotated to align "
         "with the z-axis (1,0,0). Only considered when --align_orientations is set.",
     )
     alignment_group.add_argument(
-        "--interpolation_order",
-        dest="interpolation_order",
+        "--interpolation-order",
         required=False,
         type=int,
         default=1,
         help="Interpolation order for alignment, less than zero is no interpolation.",
     )
+    alignment_group.add_argument(
+        "--split-by-micrograph",
+        action="store_true",
+        required=False,
+        help="Create separate output files for each micrograph."
+    )
 
     extraction_group = parser.add_argument_group("Extraction")
     extraction_group.add_argument(
-        "--box_size",
-        required=False,
+        "--box-size",
+        required=True,
         type=int,
-        help="Box size for extraction, defaults to two times the template.",
+        help="Box size for extraction.",
     )
     extraction_group.add_argument(
-        "--translation_uncertainty",
+        "--translation-uncertainty",
         required=False,
         type=int,
         help="Sets box size for extraction to template box plus this value.",
     )
     extraction_group.add_argument(
-        "--keep_out_of_box",
+        "--drop-out-of-box",
         action="store_true",
         required=False,
-        help="Whether to keep orientations that fall outside the box. If the "
+        help="Whether to drop orientations that fall outside the box. If the "
         "orientations are sensible, it is safe to pass this flag.",
     )
 
@@ -125,100 +127,117 @@ def parse_args():
 
 def main():
     args = parse_args()
+
+    data = StarParser(args.orientations, delimiter="\t")
+    key = list(data.keys())[0]
+
+    index_map = defaultdict(list)
+    for index, value in enumerate(data[key]["_rlnMicrographName"]):
+        index_map[value].append(index)
+
     orientations = Orientations.from_file(args.orientations)
     orientations.translations = np.divide(
-        orientations.translations, args.orientations_sampling
+        orientations.translations, args.orientations_scaling
     )
 
-    target = Density.from_file(args.target, use_memmap=True)
-
     box_size = np.array(args.box_size)
-    box_size = np.repeat(box_size, target.data.ndim // box_size.size).astype(int)
-
+    box_size = np.repeat(box_size, 3 // box_size.size).astype(int)
     extraction_shape = np.copy(box_size)
-    if args.align_orientations:
-        extraction_shape[:] = int(np.linalg.norm(box_size) + 1)
-
-    orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
-        target_shape=target.shape,
-        extraction_shape=extraction_shape,
-        drop_out_of_box=not args.keep_out_of_box,
-        return_orientations=True,
-    )
 
     if args.align_orientations:
+        extraction_shape[:] = int(np.linalg.norm(box_size) + 1)
         for index in range(orientations.rotations.shape[0]):
             rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
             rotation_matrix = np.linalg.inv(rotation_matrix)
             if args.angles_are_vector:
-                rotation_matrix = rotation_aligning_vectors(
+                rotation_matrix = align_vectors(
                     orientations.rotations[index], target_vector=(1, 0, 0)
                 )
             orientations.rotations[index] = euler_from_rotationmatrix(rotation_matrix)
 
-    filename = generate_tempfile_name()
-    output_dtype = target.data.dtype
-    if args.align_orientations is not None:
-        output_dtype = np.float32
-
-    target.data = target.data.astype(output_dtype)
-
-    dens = Density(
-        np.memmap(
-            filename,
-            mode="w+",
-            shape=(len(obs_slices), *box_size),
-            dtype=output_dtype,
-        ),
-        sampling_rate=(1, *target.sampling_rate),
-        origin=(0, *target.origin),
-    )
-    dens.data[:] = target.metadata["mean"]
-
-    data_subset = np.zeros(extraction_shape, dtype=target.data.dtype)
-    pbar = ProgressBar(message="Aligning ", nchars=80, total=len(obs_slices))
-    for index, (obs_slice, cand_slice) in enumerate(zip(obs_slices, cand_slices)):
-        pbar.update(index + 1)
-
-        data_subset.fill(0)
-        data_subset[cand_slice] = target.data[obs_slice]
-        target_subset = Density(
-            data_subset,
-            sampling_rate=target.sampling_rate,
-            origin=target.origin,
+    ret_orientations, ret_dens, ix = [], [], 0
+    n_particles = orientations.translations.shape[0]
+    pbar = ProgressBar(message="Processing ", nchars=80, total=n_particles)
+    for target_path, indices in index_map.items():
+
+        target = Density.from_file(target_path, use_memmap=True)
+
+        subset = orientations[indices]
+        subset, cand_slices, obs_slices = subset.get_extraction_slices(
+            target_shape=target.shape,
+            extraction_shape=extraction_shape,
+            drop_out_of_box=args.drop_out_of_box,
+            return_orientations=True,
         )
 
-        if args.align_orientations:
-            rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
-            target_subset = target_subset.rigid_transform(
-                rotation_matrix=rotation_matrix,
-                use_geometric_center=True,
-                order=args.interpolation_order,
+        dens = Density(
+            np.memmap(
+                generate_tempfile_name(),
+                mode="w+",
+                shape=(subset.translations.shape[0], *box_size),
+                dtype=np.float32,
+            ),
+            sampling_rate = (1, *target.sampling_rate),
+            metadata = {"batch_dimension" : (0,), "path" : target_path}
+        )
+
+        data_subset = np.zeros(extraction_shape, dtype=target.data.dtype)
+        for index, (obs_slice, cand_slice) in enumerate(zip(obs_slices, cand_slices)):
+            pbar.update(ix + 1)
+
+            data_subset.fill(0)
+            data_subset[cand_slice] = target.data[obs_slice]
+            target_subset = Density(
+                data_subset,
+                sampling_rate=target.sampling_rate,
+                origin=target.origin,
             )
-        target_subset.pad(box_size, center=True)
 
-        # target_value = target.data[tuple(orientations.translations[index].astype(int))]
-        # center = np.divide(target_subset.data.shape, 2).astype(int)
-        # print(np.where(target_subset.data == target_value), center)
-        # print(target_subset.data[tuple(center.astype(int))],
-        # target_value,
-        # target_subset.data[tuple(center.astype(int))] == target_value
-        # )
+            if args.align_orientations:
+                rotation_matrix = euler_to_rotationmatrix(subset.rotations[index])
+                target_subset = target_subset.rigid_transform(
+                    rotation_matrix=rotation_matrix,
+                    use_geometric_center=True,
+                    order=args.interpolation_order,
+                )
+            target_subset.pad(box_size, center=True)
+            dens.data[index] = target_subset.data.astype(np.float32)
+            ix += 1
 
-        dens.data[index] = target_subset.data
-    print("")
+        ret_dens.append(dens)
+        ret_orientations.append(subset)
 
-    target_meta = {
-        k: v for k, v in target.metadata.items() if k in ("mean", "max", "min", "std")
-    }
-    dens.metadata.update(target_meta)
-    dens.metadata["batch_dimension"] = (0,)
-    dens.metadata["normals"] = orientations.rotations
+    if not len(ret_dens):
+        exit("Found no valid particles.")
 
-    dens.to_file(args.output_file)
-    orientations.to_file(
-        f"{splitext(args.output_file)[0]}_aligned.tsv", file_format="text"
-    )
+    print("")
+    if not args.split_by_micrograph:
+        ret_orientations = [Orientations(
+            translations=np.concatenate([x.translations for x in ret_orientations]),
+            rotations=np.concatenate([x.rotations for x in ret_orientations]),
+            scores=np.concatenate([x.scores for x in ret_orientations]),
+            details=np.concatenate([x.details for x in ret_orientations]),
+        )]
+        dens_data = Density(
+            np.concatenate([x.data for x in ret_dens]),
+            sampling_rate=ret_dens[0].sampling_rate
+        )
+        _ = [unlink(x.data.filename) for x in ret_dens]
+        dens_data.metadata.update({"batch_dimension" : (0, )})
+        ret_dens = [dens_data]
+
+    for orientation, dens in zip(ret_orientations, ret_dens):
+        fname = args.output_prefix
+        if args.split_by_micrograph:
+            target = splitext(basename(dens.metadata["path"]))[0]
+            fname = f"{args.output_prefix}_{target}"
+
+        dens.to_file(f"{fname}.h5")
+        orientation.to_file(f"{fname}_aligned.star")
+        try:
+            unlink(dens.data.filename)
+        except Exception:
+            continue
 
 if __name__ == "__main__":
     main()