pytme 0.2.0__tar.gz → 0.2.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytme
-Version: 0.2.0
+Version: 0.2.1
 Summary: Python Template Matching Engine
 Author: Valentin Maurer
 Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>

{pytme-0.2.0 → 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.0"
+version="0.2.1"
 description="Python Template Matching Engine"
 readme="README.md"
 requires-python = ">=3.11"

{pytme-0.2.0 → 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

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()

{pytme-0.2.0 → pytme-0.2.1}/scripts/match_template.py

@@ -217,7 +217,7 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
             if args.tilt_weighting not in ("angle", None):
                 raise ValueError(
                     "Tilt weighting schemes other than 'angle' or 'None' require "
-                    "a specification of electron doses."
+                    "a specification of electron doses via --tilt_angles."
                 )
 
             wedge = Wedge(
@@ -240,31 +240,58 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
         wedge.sampling_rate = template.sampling_rate
         template_filter.append(wedge)
         if not isinstance(wedge, WedgeReconstructed):
-            template_filter.append(ReconstructFromTilt(
-                reconstruction_filter = args.reconstruction_filter
-            ))
+            template_filter.append(
+                ReconstructFromTilt(
+                    reconstruction_filter=args.reconstruction_filter,
+                    interpolation_order=args.reconstruction_interpolation_order,
+                )
+            )
 
-    if args.ctf_file is not None:
+    if args.ctf_file is not None or args.defocus is not None:
         from tme.preprocessing.tilt_series import CTF
 
-        ctf = CTF.from_file(args.ctf_file)
-        n_tilts_ctfs, n_tils_angles = len(ctf.defocus_x), len(wedge.angles)
-        if n_tilts_ctfs != n_tils_angles:
-            raise ValueError(
-                f"CTF file contains {n_tilts_ctfs} micrographs, but match_template "
-                f"recieved {n_tils_angles} tilt angles. Expected one angle "
-                "per micrograph."
+        needs_reconstruction = True
+        if args.ctf_file is not None:
+            ctf = CTF.from_file(args.ctf_file)
+            n_tilts_ctfs, n_tils_angles = len(ctf.defocus_x), len(wedge.angles)
+            if n_tilts_ctfs != n_tils_angles:
+                raise ValueError(
+                    f"CTF file contains {n_tilts_ctfs} micrographs, but match_template "
+                    f"recieved {n_tils_angles} tilt angles. Expected one angle "
+                    "per micrograph."
+                )
+            ctf.angles = wedge.angles
+            ctf.opening_axis, ctf.tilt_axis = args.wedge_axes
+        else:
+            needs_reconstruction = False
+            ctf = CTF(
+                defocus_x=[args.defocus],
+                phase_shift=[args.phase_shift],
+                defocus_y=None,
+                angles=[0],
+                shape=None,
+                return_real_fourier=True,
             )
-        ctf.angles = wedge.angles
-        ctf.opening_axis, ctf.tilt_axis = args.wedge_axes
-
-        if isinstance(template_filter[-1], ReconstructFromTilt):
+        ctf.sampling_rate = template.sampling_rate
+        ctf.flip_phase = not args.no_flip_phase
+        ctf.amplitude_contrast = args.amplitude_contrast
+        ctf.spherical_aberration = args.spherical_aberration
+        ctf.acceleration_voltage = args.acceleration_voltage * 1e3
+        ctf.correct_defocus_gradient = args.correct_defocus_gradient
+
+        if not needs_reconstruction:
+            template_filter.append(ctf)
+        elif isinstance(template_filter[-1], ReconstructFromTilt):
             template_filter.insert(-1, ctf)
         else:
             template_filter.insert(0, ctf)
-            template_filter.insert(1, ReconstructFromTilt(
-                reconstruction_filter = args.reconstruction_filter
-            ))
+            template_filter.insert(
+                1,
+                ReconstructFromTilt(
+                    reconstruction_filter=args.reconstruction_filter,
+                    interpolation_order=args.reconstruction_interpolation_order,
+                ),
+            )
 
     if args.lowpass or args.highpass is not None:
         lowpass, highpass = args.lowpass, args.highpass
@@ -293,6 +320,14 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
         template_filter.append(whitening_filter)
         target_filter.append(whitening_filter)
 
+    needs_reconstruction = any(
+        [isinstance(t, ReconstructFromTilt) for t in template_filter]
+    )
+    if needs_reconstruction and args.reconstruction_filter is None:
+        warnings.warn(
+            "Consider using a --reconstruction_filter such as 'ramp' to avoid artifacts."
+        )
+
     template_filter = Compose(template_filter) if len(template_filter) else None
     target_filter = Compose(target_filter) if len(target_filter) else None
 
@@ -510,7 +545,7 @@ def parse_args():
         dest="no_pass_smooth",
         action="store_false",
         default=True,
-        help="Whether a hard edge filter should be used for --lowpass and --highpass."
+        help="Whether a hard edge filter should be used for --lowpass and --highpass.",
     )
     filter_group.add_argument(
         "--pass_format",
@@ -519,7 +554,7 @@ def parse_args():
         required=False,
         choices=["sampling_rate", "voxel", "frequency"],
         help="How values passed to --lowpass and --highpass should be interpreted. "
-        "By default, they are assumed to be in units of sampling rate, e.g. Ångstrom."
+        "By default, they are assumed to be in units of sampling rate, e.g. Ångstrom.",
     )
     filter_group.add_argument(
         "--whiten_spectrum",
@@ -561,23 +596,90 @@ def parse_args():
         "grigorieff (exposure filter as defined in Grant and Grigorieff 2015)."
         "relion and grigorieff require electron doses in --tilt_angles weights column.",
     )
-    # filter_group.add_argument(
-    #     "--ctf_file",
-    #     dest="ctf_file",
-    #     type=str,
-    #     required=False,
-    #     default=None,
-    #     help="Path to a file with CTF parameters from CTFFIND4.",
-    # )
     filter_group.add_argument(
         "--reconstruction_filter",
         dest="reconstruction_filter",
         type=str,
         required=False,
-        choices = ["ram-lak", "ramp", "shepp-logan", "cosine", "hamming"],
+        choices=["ram-lak", "ramp", "ramp-cont", "shepp-logan", "cosine", "hamming"],
         default=None,
         help="Filter applied when reconstructing (N+1)-D from N-D filters.",
     )
+    filter_group.add_argument(
+        "--reconstruction_interpolation_order",
+        dest="reconstruction_interpolation_order",
+        type=int,
+        default=1,
+        required=False,
+        help="Analogous to --interpolation_order but for reconstruction.",
+    )
+
+    ctf_group = parser.add_argument_group("Contrast Transfer Function")
+    ctf_group.add_argument(
+        "--ctf_file",
+        dest="ctf_file",
+        type=str,
+        required=False,
+        default=None,
+        help="Path to a file with CTF parameters from CTFFIND4. Each line will be "
+        "interpreted as tilt obtained at the angle specified in --tilt_angles. ",
+    )
+    ctf_group.add_argument(
+        "--defocus",
+        dest="defocus",
+        type=float,
+        required=False,
+        default=None,
+        help="Defocus in units of sampling rate (typically Ångstrom). "
+        "Superseded by --ctf_file.",
+    )
+    ctf_group.add_argument(
+        "--phase_shift",
+        dest="phase_shift",
+        type=float,
+        required=False,
+        default=0,
+        help="Phase shift in degrees. Superseded by --ctf_file.",
+    )
+    ctf_group.add_argument(
+        "--acceleration_voltage",
+        dest="acceleration_voltage",
+        type=float,
+        required=False,
+        default=300,
+        help="Acceleration voltage in kV, defaults to 300.",
+    )
+    ctf_group.add_argument(
+        "--spherical_aberration",
+        dest="spherical_aberration",
+        type=float,
+        required=False,
+        default=2.7e7,
+        help="Spherical aberration in units of sampling rate (typically Ångstrom).",
+    )
+    ctf_group.add_argument(
+        "--amplitude_contrast",
+        dest="amplitude_contrast",
+        type=float,
+        required=False,
+        default=0.07,
+        help="Amplitude contrast, defaults to 0.07.",
+    )
+    ctf_group.add_argument(
+        "--no_flip_phase",
+        dest="no_flip_phase",
+        action="store_false",
+        required=False,
+        help="Whether the phase of the computed CTF should not be flipped.",
+    )
+    ctf_group.add_argument(
+        "--correct_defocus_gradient",
+        dest="correct_defocus_gradient",
+        action="store_true",
+        required=False,
+        help="[Experimental] Whether to compute a more accurate 3D CTF incorporating "
+        "defocus gradients.",
+    )
 
     performance_group = parser.add_argument_group("Performance")
     performance_group.add_argument(
@@ -655,12 +757,11 @@ def parse_args():
     )
 
     args = parser.parse_args()
+    args.version = __version__
 
     if args.interpolation_order < 0:
         args.interpolation_order = None
 
-    args.ctf_file = None
-
     if args.temp_directory is None:
         default = abspath(".")
         if os.environ.get("TMPDIR", None) is not None:
@@ -725,12 +826,13 @@ def main():
             sampling_rate=target.sampling_rate,
         )
 
-    if not np.allclose(target.sampling_rate, template.sampling_rate):
-        print(
-            f"Resampling template to {target.sampling_rate}. "
-            "Consider providing a template with the same sampling rate as the target."
-        )
-        template = template.resample(target.sampling_rate, order=3)
+    if target.sampling_rate.size == template.sampling_rate.size:
+        if not np.allclose(target.sampling_rate, template.sampling_rate):
+            print(
+                f"Resampling template to {target.sampling_rate}. "
+                "Consider providing a template with the same sampling rate as the target."
+            )
+            template = template.resample(target.sampling_rate, order=3)
 
     template_mask = load_and_validate_mask(
         mask_target=template, mask_path=args.template_mask
@@ -863,31 +965,46 @@ def main():
     if args.memory is None:
         args.memory = int(args.memory_scaling * available_memory)
 
-    target_padding = np.zeros_like(template.shape)
-    if args.pad_target_edges:
-        target_padding = template.shape
+    callback_class = MaxScoreOverRotations
+    if args.peak_calling:
+        callback_class = PeakCallerMaximumFilter
+
+    matching_data = MatchingData(
+        target=target,
+        template=template.data,
+        target_mask=target_mask,
+        template_mask=template_mask,
+        invert_target=args.invert_target_contrast,
+        rotations=parse_rotation_logic(args=args, ndim=template.data.ndim),
+    )
+
+    template_filter, target_filter = setup_filter(args, template, target)
+    matching_data.template_filter = template_filter
+    matching_data.target_filter = target_filter
 
-    template_box = template.shape
+    template_box = matching_data._output_template_shape
     if not args.pad_fourier:
         template_box = np.ones(len(template_box), dtype=int)
 
-    callback_class = MaxScoreOverRotations
-    if args.peak_calling:
-        callback_class = PeakCallerMaximumFilter
+    target_padding = np.zeros(
+        (backend.size(matching_data._output_template_shape)), dtype=int
+    )
+    if args.pad_target_edges:
+        target_padding = matching_data._output_template_shape
 
     splits, schedule = compute_parallelization_schedule(
         shape1=target.shape,
-        shape2=template_box,
-        shape1_padding=target_padding,
+        shape2=tuple(int(x) for x in template_box),
+        shape1_padding=tuple(int(x) for x in target_padding),
         max_cores=args.cores,
         max_ram=args.memory,
         split_only_outer=args.use_gpu,
         matching_method=args.score,
         analyzer_method=callback_class.__name__,
         backend=backend._backend_name,
-        float_nbytes=backend.datatype_bytes(backend._default_dtype),
+        float_nbytes=backend.datatype_bytes(backend._float_dtype),
         complex_nbytes=backend.datatype_bytes(backend._complex_dtype),
-        integer_nbytes=backend.datatype_bytes(backend._default_dtype_int),
+        integer_nbytes=backend.datatype_bytes(backend._int_dtype),
     )
 
     if splits is None:
@@ -898,20 +1015,6 @@ def main():
         exit(-1)
 
     matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[args.score]
-    matching_data = MatchingData(target=target, template=template.data)
-    matching_data.rotations = parse_rotation_logic(args=args, ndim=target.data.ndim)
-
-    template_filter, target_filter = setup_filter(args, template, target)
-    matching_data.template_filter = template_filter
-    matching_data.target_filter = target_filter
-
-    matching_data.template_filter = template_filter
-    matching_data._invert_target = args.invert_target_contrast
-    if target_mask is not None:
-        matching_data.target_mask = target_mask
-    if template_mask is not None:
-        matching_data.template_mask = template_mask.data
-
     n_splits = np.prod(list(splits.values()))
     target_split = ", ".join(
         [":".join([str(x) for x in axis]) for axis in splits.items()]
@@ -945,13 +1048,23 @@ def main():
         "Lowpass": args.lowpass,
         "Highpass": args.highpass,
         "Smooth Pass": args.no_pass_smooth,
-        "Pass Format" : args.pass_format,
+        "Pass Format": args.pass_format,
         "Spectral Whitening": args.whiten_spectrum,
         "Wedge Axes": args.wedge_axes,
         "Tilt Angles": args.tilt_angles,
         "Tilt Weighting": args.tilt_weighting,
-        "CTF": args.ctf_file,
+        "Reconstruction Filter": args.reconstruction_filter,
     }
+    if args.ctf_file is not None or args.defocus is not None:
+        filter_args["CTF File"] = args.ctf_file
+        filter_args["Defocus"] = args.defocus
+        filter_args["Phase Shift"] = args.phase_shift
+        filter_args["No Flip Phase"] = args.no_flip_phase
+        filter_args["Acceleration Voltage"] = args.acceleration_voltage
+        filter_args["Spherical Aberration"] = args.spherical_aberration
+        filter_args["Amplitude Contrast"] = args.amplitude_contrast
+        filter_args["Correct Defocus"] = args.correct_defocus_gradient
+
     filter_args = {k: v for k, v in filter_args.items() if v is not None}
     if len(filter_args):
         print_block(
@@ -1000,15 +1113,16 @@ def main():
             candidates[0] *= target_mask.data
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", category=UserWarning)
+            nbytes = backend.datatype_bytes(backend._float_dtype)
+            dtype = np.float32 if nbytes == 4 else np.float16
+            rot_dim = matching_data.rotations.shape[1]
             candidates[3] = {
                 x: euler_from_rotationmatrix(
-                    np.frombuffer(i, dtype=matching_data.rotations.dtype).reshape(
-                        candidates[0].ndim, candidates[0].ndim
-                    )
+                    np.frombuffer(i, dtype=dtype).reshape(rot_dim, rot_dim)
                 )
                 for i, x in candidates[3].items()
             }
-    candidates.append((target.origin, template.origin, target.sampling_rate, args))
+    candidates.append((target.origin, template.origin, template.sampling_rate, args))
     write_pickle(data=candidates, filename=args.output)
 
     runtime = time() - start