pytme 0.2.0b0__tar.gz → 0.2.1__tar.gz

{pytme-0.2.0b0 → pytme-0.2.1}/MANIFEST.in

@@ -60,7 +60,6 @@ global-exclude .DS_Store
 global-exclude .gitconfig
 global-exclude *.pickle
 prune tme/temp
-prune tme/preprocessing
 global-exclude doc/compileApiReference.sh
 prune public
 global-exclude tme/matching_backend.py
@@ -70,4 +69,4 @@ global-exclude tme/scoring.py
 global-exclude tme/package.py
 global-exclude tme/transforms.py
 global-exclude tme/tests/test_packaging.py
-global-exclude scripts/match_template_devel.py
+global-exclude scripts/match_template_devel.py

{pytme-0.2.0b0 → pytme-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytme
-Version: 0.2.0b0
+Version: 0.2.1
 Summary: Python Template Matching Engine
 Author: Valentin Maurer
 Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>
@@ -22,7 +22,7 @@ Requires-Dist: scikit_learn>=1.2.1
 Requires-Dist: scipy>=1.9.1
 Requires-Dist: pybind11
 Requires-Dist: psutil
-Requires-Dist: tifffile[all]
+Requires-Dist: tifffile
 Requires-Dist: h5py
 
 # Python Template Matching Engine (PyTME)

{pytme-0.2.0b0 → pytme-0.2.1}/pyproject.toml

@@ -7,7 +7,7 @@ name="pytme"
 authors = [
   { name = "Valentin Maurer", email = "valentin.maurer@embl-hamburg.de" },
 ]
-version="0.2.0b"
+version="0.2.1"
 description="Python Template Matching Engine"
 readme="README.md"
 requires-python = ">=3.11"
@@ -23,7 +23,7 @@ dependencies=[
     "scipy>=1.9.1",
     "pybind11",
     "psutil",
-    "tifffile[all]",
+    "tifffile",
     "h5py"
 ]
 license = {text = "Proprietary"}

{pytme-0.2.0b0 → pytme-0.2.1}/pytme.egg-info/SOURCES.txt

@@ -31,6 +31,7 @@ tme/structure.py
 tme/types.py
 tme/backends/__init__.py
 tme/backends/cupy_backend.py
+tme/backends/jax_backend.py
 tme/backends/matching_backend.py
 tme/backends/mlx_backend.py
 tme/backends/npfftw_backend.py
@@ -57,4 +58,10 @@ tme/data/c48u8649.npy
 tme/data/c600v.npy
 tme/data/c600vc.npy
 tme/data/metadata.yaml
-tme/data/quat_to_numpy.py
+tme/data/quat_to_numpy.py
+tme/preprocessing/__init__.py
+tme/preprocessing/_utils.py
+tme/preprocessing/composable_filter.py
+tme/preprocessing/compose.py
+tme/preprocessing/frequency_filters.py
+tme/preprocessing/tilt_series.py

pytme-0.2.1/scripts/extract_candidates.py

@@ -0,0 +1,250 @@
+#!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 (
+    load_pickle,
+    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)
+        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."
+    )
+
+    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=True,
+        type=str,
+        help="Path to output HDF5 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(
+        "--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()
+
+    data_present = args.target is not None and args.template is not None
+    if args.input_file is None and not data_present:
+        raise ValueError(
+            "Either --input_file or --target and --template need to be specified."
+        )
+    elif args.input_file is not None and data_present:
+        raise ValueError(
+            "Please specific either --input_file or --target and --template."
+        )
+
+    return args
+
+
+def main():
+    args = parse_args()
+    orientations = Orientations.from_file(args.orientations)
+
+    if args.input_file is not None:
+        data = load_pickle(args.input_file)
+        target_origin, _, sampling_rate, cli_args = data[-1]
+        args.target, args.template  = cli_args.target, cli_args.template
+
+    target = Density.from_file(args.target, use_memmap=True)
+
+    try:
+        template = Density.from_file(args.template)
+    except Exception:
+        template = Density.from_structure(args.template, sampling_rate = target.sampling_rate)
+
+    box_size = args.box_size
+    if box_size is None:
+        box_size = np.multiply(template.shape, 2)
+    box_size = np.array(box_size)
+    box_size = np.repeat(box_size, template.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:
+        orientations.rotations = orientations.rotations.astype(np.float32)
+        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"]
+
+    print(target.data.shape)
+    # There appears to be an isseu with the stack creation. Trace this further
+    data_subset = np.zeros(extraction_shape, dtype = target.data.dtype)
+    pbar = ProgressBar(message = "Orientation ", 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 ) + np.mod(target_subset.shape, 2)
+        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.to_file(args.output_file)
+    orientations.to_file(
+        f"{splitext(args.output_file)[0]}_aligned.tsv",
+        file_format="text"
+    )
+    orientations.to_file(
+        f"{splitext(args.output_file)[0]}_aligned.star",
+        file_format="relion"
+    )
+
+if __name__ == "__main__":
+    main()