pytme 0.1.9__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0__cp311-cp311-macosx_14_0_arm64.whl

scripts/extract_candidates.py

@@ -0,0 +1,218 @@
+#!python3
+""" CLI to simplify analysing the output of match_template.py.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+import argparse
+from os.path import splitext
+
+import numpy as np
+
+from tme import Density, Orientations
+from tme.matching_utils import (
+    create_mask,
+    generate_tempfile_name,
+    rotation_aligning_vectors,
+    euler_from_rotationmatrix,
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Extract matching candidates for further refinement."
+    )
+
+    io_group = parser.add_argument_group("Input / Output")
+    io_group.add_argument(
+        "--input_file",
+        required=False,
+        type=str,
+        help="Path to the output of match_template.py.",
+    )
+    io_group.add_argument(
+        "--orientations",
+        required=True,
+        type=str,
+        help="Path to file generated by postprocess.py using output_format orientations.",
+    )
+    io_group.add_argument(
+        "--target",
+        required=False,
+        type=str,
+        help="Extract candidates from this target, can be at different sampling rate.",
+    )
+    io_group.add_argument(
+        "--template",
+        required=False,
+        type=str,
+        help="Extract candidates from this target, can be at different sampling rate.",
+    )
+    io_group.add_argument(
+        "-o",
+        "--output_file",
+        required=False,
+        type=str,
+        help="Path to output HDF5 file.",
+    )
+
+    extraction_group = parser.add_argument_group("Extraction")
+    extraction_group.add_argument(
+        "--box_size",
+        required=False,
+        type=int,
+        help="Box size for extraction, defaults to two times the template box.",
+    )
+    extraction_group.add_argument(
+        "--alignment_vector",
+        required=False,
+        type=str,
+        help="Vector to align extracted coordinate to e.g. '1,0,0'. If provided, euler_z"
+        "euler_y and euler_x are assumed to be a vector that will be rotated to match "
+        "alignment_vector.",
+    )
+    extraction_group.add_argument(
+        "--interpolation_order",
+        dest="interpolation_order",
+        required=False,
+        type=int,
+        default=1,
+        help="Spline interpolation used for template rotations. If less than zero "
+        "no interpolation is performed.",
+    )
+    extraction_group.add_argument(
+        "--translation_uncertainty",
+        dest="translation_uncertainty",
+        required=False,
+        type=int,
+        default=None,
+        help="Creates a cented spherical target mask with given radius in voxel.",
+    )
+    extraction_group.add_argument(
+        "--mask_path",
+        dest="mask_path",
+        required=False,
+        type=str,
+        default="target_mask.h5",
+        help="Path to write spherical mask to, defaults to target_mask.h5.",
+    )
+    args = parser.parse_args()
+
+    if args.alignment_vector is not None:
+        args.alignment_vector = [float(x) for x in args.alignment_vector.split(",")]
+        if len(args.alignment_vector) != 3:
+            raise ValueError("Only three-dimensional vectors are supported")
+
+    return args
+
+
+def main():
+    args = parse_args()
+    orientations = Orientations.from_file(args.orientations)
+
+    target = Density.from_file(args.target, use_memmap=True)
+    template = Density.from_file(args.template)
+
+    box_size = template.shape
+    if args.box_size is not None:
+        args.box_size = args.box_size // 2
+        box_size = np.full(template.data.ndim, fill_value=args.box_size, dtype=int)
+
+    subtomo_shape = np.multiply(box_size, 2).astype(int)
+    extraction_scaling = 3 if args.alignment_vector is not None else 2
+    extraction_shape = np.multiply(box_size, extraction_scaling).astype(int)
+    orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
+        target_shape=target.shape,
+        extraction_shape=extraction_shape,
+        drop_out_of_box=True,
+        return_orientations=True,
+    )
+
+    filename = generate_tempfile_name()
+    output_dtype = target.data.dtype
+    if args.alignment_vector is not None:
+        output_dtype = np.float32
+
+    if args.translation_uncertainty is not None:
+        dens = Density(
+            np.memmap(
+                args.mask_path,
+                mode="w+",
+                shape=(len(obs_slices), *subtomo_shape),
+                dtype=output_dtype,
+            ),
+            sampling_rate=(1, *target.sampling_rate),
+            origin=(0, *target.origin),
+        )
+        dens.data[:] = 0
+        mask = create_mask(
+            mask_type="ellipse",
+            center=template.shape,
+            radius=args.translation_uncertainty,
+            shape=subtomo_shape,
+        )
+        dens.data[:] = mask
+        dens.to_file(args.mask_path)
+
+    target.data = target.data.astype(output_dtype)
+
+    dens = Density(
+        np.memmap(
+            filename,
+            mode="w+",
+            shape=(len(obs_slices), *subtomo_shape),
+            dtype=output_dtype,
+        ),
+        sampling_rate=(1, *target.sampling_rate),
+        origin=(0, *target.origin),
+    )
+    dens.data[:] = target.metadata["mean"]
+
+    target_meta = {
+        k: v for k, v in target.metadata.items() if k in ("mean", "max", "min", "std")
+    }
+
+    total_slices = len(obs_slices)
+    magnitude = len(str(total_slices))
+    for index, obs_slice in enumerate(obs_slices):
+        print(f"Processing {index:{magnitude}} / {total_slices:{magnitude}}", end="\r")
+
+        target_subset = Density(
+            target.data[obs_slice],
+            sampling_rate=target.sampling_rate,
+            origin=target.origin,
+        )
+
+        rotation_matrix = np.eye(target_subset.data.ndim)
+        if args.alignment_vector is not None:
+            normal = orientations.rotations[index]
+            # To visualize only normals were picked we would need to do the opposite
+            # However, this allows us to align everything to the template axis
+            rotation_matrix = rotation_aligning_vectors(
+                normal, target_vector=args.alignment_vector
+            )
+
+            target_subset = target_subset.rigid_transform(
+                rotation_matrix=rotation_matrix,
+                use_geometric_center=True,
+                order=args.interpolation_order,
+            )
+        target_subset.pad(subtomo_shape, center=True)
+        dens.data[index] = target_subset.data
+        orientations.rotations[index] = euler_from_rotationmatrix(rotation_matrix)
+
+    dens.metadata.update(target_meta)
+    dens.to_file(args.output_file)
+    orientations.to_file(
+        f"{splitext(args.output_file)}_aligned.tsv",
+        file_format="text"
+    )
+    orientations.to_file(
+        f"{splitext(args.output_file)}_aligned.star",
+        file_format="relion"
+    )
+
+
+if __name__ == "__main__":
+    main()