pytme 0.2.2__cp311-cp311-macosx_14_0_arm64.whl → 0.2.4__cp311-cp311-macosx_14_0_arm64.whl

scripts/match_template.py

@@ -1,5 +1,5 @@
 #!python3
-""" CLI interface for basic pyTME template matching functions.
+""" CLI for basic pyTME template matching functions.
 
     Copyright (c) 2023 European Molecular Biology Laboratory
 
@@ -22,7 +22,6 @@ from tme.matching_utils import (
     get_rotations_around_vector,
     compute_parallelization_schedule,
     scramble_phases,
-    generate_tempfile_name,
     write_pickle,
 )
 from tme.matching_exhaustive import scan_subsets, MATCHING_EXHAUSTIVE_REGISTER
@@ -99,7 +98,9 @@ def load_and_validate_mask(mask_target: "Density", mask_path: str, **kwargs):
                 f"Expected shape of {mask_path} was {mask_target.shape},"
                 f" got f{mask.shape}"
             )
-        if not np.allclose(mask.sampling_rate, mask_target.sampling_rate):
+        if not np.allclose(
+            np.round(mask.sampling_rate, 2), np.round(mask_target.sampling_rate, 2)
+        ):
             raise ValueError(
                 f"Expected sampling_rate of {mask_path} was {mask_target.sampling_rate}"
                 f", got f{mask.sampling_rate}"
@@ -107,50 +108,6 @@ def load_and_validate_mask(mask_target: "Density", mask_path: str, **kwargs):
     return mask
 
 
-def crop_data(data: Density, cutoff: float, data_mask: Density = None) -> bool:
-    """
-    Crop the provided data and mask to a smaller box based on a cutoff value.
-
-    Parameters
-    ----------
-    data : Density
-        The data that should be cropped.
-    cutoff : float
-        The threshold value to determine which parts of the data should be kept.
-    data_mask : Density, optional
-        A mask for the data that should be cropped.
-
-    Returns
-    -------
-    bool
-        Returns True if the data was adjusted (cropped), otherwise returns False.
-
-    Notes
-    -----
-    Cropping is performed in place.
-    """
-    if cutoff is None:
-        return False
-
-    box = data.trim_box(cutoff=cutoff)
-    box_mask = box
-    if data_mask is not None:
-        box_mask = data_mask.trim_box(cutoff=cutoff)
-    box = tuple(
-        slice(min(arr.start, mask.start), max(arr.stop, mask.stop))
-        for arr, mask in zip(box, box_mask)
-    )
-    if box == tuple(slice(0, x) for x in data.shape):
-        return False
-
-    data.adjust_box(box)
-
-    if data_mask:
-        data_mask.adjust_box(box)
-
-    return True
-
-
 def parse_rotation_logic(args, ndim):
     if args.angular_sampling is not None:
         rotations = get_rotation_matrices(
@@ -175,8 +132,52 @@ def parse_rotation_logic(args, ndim):
     return rotations
 
 
-# TODO: Think about whether wedge mask should also be added to target
-# For now leave it at the cost of incorrect upper bound on the scores
+def compute_schedule(
+    args,
+    target: Density,
+    matching_data: MatchingData,
+    callback_class,
+    pad_edges: bool = False,
+):
+    # User requested target padding
+    if args.pad_edges is True:
+        pad_edges = True
+    template_box = matching_data._output_template_shape
+    if not args.pad_fourier:
+        template_box = tuple(0 for _ in range(len(template_box)))
+
+    target_padding = tuple(0 for _ in range(len(template_box)))
+    if pad_edges:
+        target_padding = matching_data._output_template_shape
+
+    splits, schedule = compute_parallelization_schedule(
+        shape1=target.shape,
+        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=be._backend_name,
+        float_nbytes=be.datatype_bytes(be._float_dtype),
+        complex_nbytes=be.datatype_bytes(be._complex_dtype),
+        integer_nbytes=be.datatype_bytes(be._int_dtype),
+    )
+
+    if splits is None:
+        print(
+            "Found no suitable parallelization schedule. Consider increasing"
+            " available RAM or decreasing number of cores."
+        )
+        exit(-1)
+    n_splits = np.prod(list(splits.values()))
+    if pad_edges is False and n_splits > 1:
+        args.pad_edges = True
+        return compute_schedule(args, target, matching_data, callback_class, True)
+    return splits, schedule
+
+
 def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Compose]:
     from tme.preprocessing import LinearWhiteningFilter, BandPassFilter
     from tme.preprocessing.tilt_series import (
@@ -185,8 +186,10 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
         ReconstructFromTilt,
     )
 
+    needs_reconstruction = False
     template_filter, target_filter = [], []
     if args.tilt_angles is not None:
+        needs_reconstruction = args.tilt_weighting is not None
         try:
             wedge = Wedge.from_file(args.tilt_angles)
             wedge.weight_type = args.tilt_weighting
@@ -231,19 +234,25 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
                     angles=tilt_angles,
                     weight_wedge=args.tilt_weighting == "angle",
                     create_continuous_wedge=create_continuous_wedge,
+                    reconstruction_filter=args.reconstruction_filter,
+                )
+                wedge_target = WedgeReconstructed(
+                    angles=(np.abs(np.min(tilt_angles)), np.abs(np.max(tilt_angles))),
+                    weight_wedge=False,
+                    create_continuous_wedge=True,
                 )
+                target_filter.append(wedge_target)
 
         wedge.opening_axis = args.wedge_axes[0]
         wedge.tilt_axis = args.wedge_axes[1]
         wedge.sampling_rate = template.sampling_rate
         template_filter.append(wedge)
         if not isinstance(wedge, WedgeReconstructed):
-            template_filter.append(
-                ReconstructFromTilt(
-                    reconstruction_filter=args.reconstruction_filter,
-                    interpolation_order=args.reconstruction_interpolation_order,
-                )
+            reconstruction_filter = ReconstructFromTilt(
+                reconstruction_filter=args.reconstruction_filter,
+                interpolation_order=args.reconstruction_interpolation_order,
             )
+            template_filter.append(reconstruction_filter)
 
     if args.ctf_file is not None or args.defocus is not None:
         from tme.preprocessing.tilt_series import CTF
@@ -317,18 +326,13 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
             sampling_rate=template.sampling_rate,
         )
         template_filter.append(bandpass)
-
-        if not args.no_filter_target:
-            target_filter.append(bandpass)
+        target_filter.append(bandpass)
 
     if args.whiten_spectrum:
         whitening_filter = LinearWhiteningFilter()
         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."
@@ -336,6 +340,8 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
 
     template_filter = Compose(template_filter) if len(template_filter) else None
     target_filter = Compose(target_filter) if len(target_filter) else None
+    if args.no_filter_target:
+        target_filter = None
 
     return template_filter, target_filter
 
@@ -393,8 +399,7 @@ def parse_args():
         dest="invert_target_contrast",
         action="store_true",
         default=False,
-        help="Invert the target's contrast and rescale linearly between zero and one. "
-        "This option is intended for targets where templates to-be-matched have "
+        help="Invert the target's contrast for cases where templates to-be-matched have "
         "negative values, e.g. tomograms.",
     )
     io_group.add_argument(
@@ -696,22 +701,6 @@ def parse_args():
     )
 
     performance_group = parser.add_argument_group("Performance")
-    performance_group.add_argument(
-        "--cutoff_target",
-        dest="cutoff_target",
-        type=float,
-        required=False,
-        default=None,
-        help="Target contour level (used for cropping).",
-    )
-    performance_group.add_argument(
-        "--cutoff_template",
-        dest="cutoff_template",
-        type=float,
-        required=False,
-        default=None,
-        help="Template contour level (used for cropping).",
-    )
     performance_group.add_argument(
         "--no_centering",
         dest="no_centering",
@@ -719,30 +708,28 @@ def parse_args():
         help="Assumes the template is already centered and omits centering.",
     )
     performance_group.add_argument(
-        "--no_edge_padding",
-        dest="no_edge_padding",
+        "--pad_edges",
+        dest="pad_edges",
         action="store_true",
         default=False,
-        help="Whether to not pad the edges of the target. Can be set if the target"
-        " has a well defined bounding box, e.g. a masked reconstruction.",
+        help="Whether to pad the edges of the target. Useful if the target does not "
+        "a well-defined bounding box. Defaults to True if splitting is required.",
     )
     performance_group.add_argument(
-        "--no_fourier_padding",
-        dest="no_fourier_padding",
+        "--pad_fourier",
+        dest="pad_fourier",
         action="store_true",
         default=False,
         help="Whether input arrays should not be zero-padded to full convolution shape "
-        "for numerical stability. When working with very large targets, e.g. tomograms, "
-        "it is safe to use this flag and benefit from the performance gain.",
+        "for numerical stability. Typically only useful when working with small data.",
     )
     performance_group.add_argument(
-        "--no_filter_padding",
-        dest="no_filter_padding",
+        "--pad_filter",
+        dest="pad_filter",
         action="store_true",
         default=False,
-        help="Omits padding of optional template filters. Particularly effective when "
-        "the target is much larger than the template. However, for fast osciliating "
-        "filters setting this flag can introduce aliasing effects.",
+        help="Pads the filter to the shape of the target. Particularly useful for fast "
+        "oscilating filters to avoid aliasing effects.",
     )
     performance_group.add_argument(
         "--interpolation_order",
@@ -805,9 +792,6 @@ def parse_args():
 
     os.environ["TMPDIR"] = args.temp_directory
 
-    args.pad_target_edges = not args.no_edge_padding
-    args.pad_fourier = not args.no_fourier_padding
-
     if args.score not in MATCHING_EXHAUSTIVE_REGISTER:
         raise ValueError(
             f"score has to be one of {', '.join(MATCHING_EXHAUSTIVE_REGISTER.keys())}"
@@ -862,7 +846,9 @@ def main():
         )
 
     if target.sampling_rate.size == template.sampling_rate.size:
-        if not np.allclose(target.sampling_rate, template.sampling_rate):
+        if not np.allclose(
+            np.round(target.sampling_rate, 2), np.round(template.sampling_rate, 2)
+        ):
             print(
                 f"Resampling template to {target.sampling_rate}. "
                 "Consider providing a template with the same sampling rate as the target."
@@ -877,9 +863,6 @@ def main():
     )
 
     initial_shape = target.shape
-    is_cropped = crop_data(
-        data=target, data_mask=target_mask, cutoff=args.cutoff_target
-    )
     print_block(
         name="Target",
         data={
@@ -888,13 +871,6 @@ def main():
             "Final Shape": target.shape,
         },
     )
-    if is_cropped:
-        args.target = generate_tempfile_name(suffix=".mrc")
-        target.to_file(args.target)
-
-        if target_mask:
-            args.target_mask = generate_tempfile_name(suffix=".mrc")
-            target_mask.to_file(args.target_mask)
 
     if target_mask:
         print_block(
@@ -907,8 +883,6 @@ def main():
         )
 
     initial_shape = template.shape
-    _ = crop_data(data=template, data_mask=template_mask, cutoff=args.cutoff_template)
-
     translation = np.zeros(len(template.shape), dtype=np.float32)
     if not args.no_centering:
         template, translation = template.centered(0)
@@ -956,7 +930,7 @@ def main():
 
     if args.scramble_phases:
         template.data = scramble_phases(
-            template.data, noise_proportion=1.0, normalize_power=True
+            template.data, noise_proportion=1.0, normalize_power=False
         )
 
     # Determine suitable backend for the selected operation
@@ -980,10 +954,7 @@ def main():
             available_backends.remove("jax")
         if args.use_gpu and "pytorch" in available_backends:
             available_backends = ("pytorch",)
-            if args.interpolation_order == 3:
-                raise NotImplementedError(
-                    "Pytorch does not support --interpolation_order 3, 1 is supported."
-                )
+
     # dim_match = len(template.shape) == len(target.shape) <= 3
     # if dim_match and args.use_gpu and "jax" in available_backends:
     #     args.interpolation_order = 1
@@ -1008,6 +979,12 @@ def main():
             )
         break
 
+    if pref == "pytorch" and args.interpolation_order == 3:
+        warnings.warn(
+            "Pytorch does not support --interpolation_order 3, setting it to 1."
+        )
+        args.interpolation_order = 1
+
     available_memory = be.get_available_memory() * be.device_count()
     if args.memory is None:
         args.memory = int(args.memory_scaling * available_memory)
@@ -1025,41 +1002,13 @@ def main():
         rotations=parse_rotation_logic(args=args, ndim=template.data.ndim),
     )
 
+    matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[args.score]
     matching_data.template_filter, matching_data.target_filter = setup_filter(
         args, template, target
     )
 
-    template_box = matching_data._output_template_shape
-    if not args.pad_fourier:
-        template_box = tuple(0 for _ in range(len(template_box)))
-
-    target_padding = tuple(0 for _ in range(len(template_box)))
-    if args.pad_target_edges:
-        target_padding = matching_data._output_template_shape
-
-    splits, schedule = compute_parallelization_schedule(
-        shape1=target.shape,
-        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=be._backend_name,
-        float_nbytes=be.datatype_bytes(be._float_dtype),
-        complex_nbytes=be.datatype_bytes(be._complex_dtype),
-        integer_nbytes=be.datatype_bytes(be._int_dtype),
-    )
-
-    if splits is None:
-        print(
-            "Found no suitable parallelization schedule. Consider increasing"
-            " available RAM or decreasing number of cores."
-        )
-        exit(-1)
+    splits, schedule = compute_schedule(args, target, matching_data, callback_class)
 
-    matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[args.score]
     n_splits = np.prod(list(splits.values()))
     target_split = ", ".join(
         [":".join([str(x) for x in axis]) for axis in splits.items()]
@@ -1071,8 +1020,8 @@ def main():
         "Center Template": not args.no_centering,
         "Scramble Template": args.scramble_phases,
         "Invert Contrast": args.invert_target_contrast,
-        "Extend Fourier Grid": not args.no_fourier_padding,
-        "Extend Target Edges": not args.no_edge_padding,
+        "Extend Fourier Grid": args.pad_fourier,
+        "Extend Target Edges": args.pad_edges,
         "Interpolation Order": args.interpolation_order,
         "Setup Function": f"{get_func_fullname(matching_setup)}",
         "Scoring Function": f"{get_func_fullname(matching_score)}",
@@ -1108,6 +1057,7 @@ def main():
         "Tilt Angles": args.tilt_angles,
         "Tilt Weighting": args.tilt_weighting,
         "Reconstruction Filter": args.reconstruction_filter,
+        "Extend Filter Grid": args.pad_filter,
     }
     if args.ctf_file is not None or args.defocus is not None:
         filter_args["CTF File"] = args.ctf_file
@@ -1130,8 +1080,7 @@ def main():
     analyzer_args = {
         "score_threshold": args.score_threshold,
         "number_of_peaks": args.number_of_peaks,
-        "min_distance": max(template.shape) // 2,
-        "min_boundary_distance": max(template.shape) // 2,
+        "min_distance": max(template.shape) // 3,
         "use_memmap": args.use_memmap,
     }
     print_block(
@@ -1156,9 +1105,9 @@ def main():
         callback_class=callback_class,
         callback_class_args=analyzer_args,
         target_splits=splits,
-        pad_target_edges=args.pad_target_edges,
+        pad_target_edges=args.pad_edges,
         pad_fourier=args.pad_fourier,
-        pad_template_filter=not args.no_filter_padding,
+        pad_template_filter=args.pad_filter,
         interpolation_order=args.interpolation_order,
     )
 

scripts/postprocess.py

@@ -509,19 +509,7 @@ def main():
 
     target = Density.from_file(cli_args.target)
     if args.invert_target_contrast:
-        if args.output_format == "relion":
-            target.data = target.data * -1
-            target.data = np.divide(
-                np.subtract(target.data, target.data.mean()), target.data.std()
-            )
-        else:
-            target.data = (
-                -np.divide(
-                    np.subtract(target.data, target.data.min()),
-                    np.subtract(target.data.max(), target.data.min()),
-                )
-                + 1
-            )
+        target.data = target.data * -1
 
     if args.output_format in ("extraction", "relion"):
         if not np.all(np.divide(target.shape, template.shape) > 2):
@@ -546,10 +534,14 @@ def main():
 
         working_directory = getcwd()
         if args.output_format == "relion":
+            name = [
+                join(working_directory, f"{args.output_prefix}_{index}.mrc")
+                for index in range(len(cand_slices))
+            ]
             orientations.to_file(
                 filename=f"{args.output_prefix}.star",
                 file_format="relion",
-                name_prefix=join(working_directory, args.output_prefix),
+                name=name,
                 ctf_image=args.wedge_mask,
                 sampling_rate=target.sampling_rate.max(),
                 subtomogram_size=extraction_shape[0],
@@ -606,24 +598,22 @@ def main():
     if args.output_format == "average":
         orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
             target_shape=target.shape,
-            extraction_shape=np.multiply(template.shape, 2),
+            extraction_shape=template.shape,
             drop_out_of_box=True,
             return_orientations=True,
         )
         out = np.zeros_like(template.data)
-        # out = np.zeros(np.multiply(template.shape, 2).astype(int))
         for index in range(len(cand_slices)):
-            from scipy.spatial.transform import Rotation
-
-            rotation = Rotation.from_euler(
-                angles=orientations.rotations[index], seq="zyx", degrees=True
-            )
-            rotation_matrix = rotation.inv().as_matrix()
-
             subset = Density(target.data[obs_slices[index]])
-            subset = subset.rigid_transform(rotation_matrix=rotation_matrix, order=1)
+            rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
 
+            subset = subset.rigid_transform(
+                rotation_matrix=np.linalg.inv(rotation_matrix),
+                order=1,
+                use_geometric_center=True,
+            )
             np.add(out, subset.data, out=out)
+
         out /= len(cand_slices)
         ret = Density(out, sampling_rate=template.sampling_rate, origin=0)
         ret.pad(template.shape, center=True)
@@ -637,17 +627,18 @@ def main():
         target_shape=target.shape,
     )
 
+    # Template is larger than target
     for index, (translation, angles, *_) in enumerate(orientations):
         rotation_matrix = euler_to_rotationmatrix(angles)
         if template_is_density:
-            translation = np.subtract(translation, center)
             transformed_template = template.rigid_transform(
-                rotation_matrix=rotation_matrix
-            )
-            transformed_template.origin = np.add(
-                target_origin, np.multiply(translation, sampling_rate)
+                rotation_matrix=rotation_matrix, use_geometric_center=True
             )
 
+            # Just adapting the coordinate system not the in-box position
+            shift = np.multiply(np.subtract(translation, center), sampling_rate)
+            transformed_template.origin = np.add(target_origin, shift)
+
         else:
             template = Structure.from_file(cli_args.template)
             new_center_of_mass = np.add(