pytme 0.2.9__cp311-cp311-macosx_15_0_arm64.whl → 0.3.0__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

@@ -0,0 +1,224 @@
+#!python3
+""" Prepare orientations stack for refinement.
+
+    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 (
+    generate_tempfile_name,
+    rotation_aligning_vectors,
+    euler_from_rotationmatrix,
+    euler_to_rotationmatrix,
+)
+
+
+class ProgressBar:
+    """
+    ASCII progress bar.
+    """
+
+    def __init__(self, message: str, nchars: int, total: int):
+        self._size = nchars - len(message) - (len(str(total))+2) * 2
+        self._message = message
+        self._total = total
+
+    def update(self, cur):
+        x = int(cur * self._size / self._total)
+        print(
+            "%s[%s%s] %i/%i\r"
+            % (self._message, "#" * x, "." * (self._size - x), cur, self._total),
+            end="",
+        )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Extract matching candidates for further refinement.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    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.",
+    )
+    io_group.add_argument(
+        "--orientations_sampling",
+        required=False,
+        type=float,
+        default=1.0,
+        help="Factor to map candidate coordinates onto the target. Only relevant if "
+        "target sampling rate differs from candidate orientation sampling rate.",
+    )
+    io_group.add_argument(
+        "-o",
+        "--output_file",
+        required=True,
+        type=str,
+        help="Path to write output H5 file.",
+    )
+
+    alignment_group = parser.add_argument_group("Alignment")
+    alignment_group.add_argument(
+        "--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",
+        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",
+        required=False,
+        type=int,
+        default=1,
+        help="Interpolation order for alignment, less than zero is no interpolation.",
+    )
+
+    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.",
+    )
+    extraction_group.add_argument(
+        "--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",
+        action="store_true",
+        required=False,
+        help="Whether to keep orientations that fall outside the box. If the "
+        "orientations are sensible, it is safe to pass this flag.",
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+def main():
+    args = parse_args()
+    orientations = Orientations.from_file(args.orientations)
+    orientations.translations = np.divide(
+        orientations.translations, args.orientations_sampling
+    )
+
+    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)
+
+    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:
+        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(
+                    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,
+        )
+
+        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,
+            )
+        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
+        # )
+
+        dens.data[index] = target_subset.data
+    print("")
+
+    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
+
+    dens.to_file(args.output_file)
+    orientations.to_file(
+        f"{splitext(args.output_file)[0]}_aligned.tsv", file_format="text"
+    )
+
+if __name__ == "__main__":
+    main()