pytme 0.1.6__cp311-cp311-macosx_14_0_arm64.whl → 0.1.8__cp311-cp311-macosx_14_0_arm64.whl

{pytme-0.1.6.data → pytme-0.1.8.data}/scripts/estimate_ram_usage.py

@@ -88,7 +88,7 @@ def main():
         shape2=template_box,
         matching_method=args.score,
         ncores=args.ncores,
-        analyzer_method="MaxScoreOverRotations"
+        analyzer_method="MaxScoreOverRotations",
     )
     print(result)
 

{pytme-0.1.6.data → pytme-0.1.8.data}/scripts/match_template.py

@@ -252,7 +252,7 @@ def parse_args():
         type=check_positive,
         default=40.0,
         help="Angular sampling rate for template matching. "
-        "A lower number yields more rotations.",
+        "A lower number yields more rotations. Values >= 180 sample only the identity.",
     )
     parser.add_argument(
         "-p",
@@ -380,8 +380,8 @@ def parse_args():
         type=float,
         required=False,
         default=None,
-        help="Step size between tilts, e.g. '5'. When set a more accurate"
-        " wedge mask will be computed.",
+        help="Step size between tilts, e.g. '5'. When set the wedge mask"
+        " reflects the individual tilts, otherwise a continuous mask is used.",
     )
     parser.add_argument(
         "--wedge_axes",
@@ -488,14 +488,16 @@ def main():
         if target_mask:
             args.target_mask = generate_tempfile_name(suffix=".mrc")
             target_mask.to_file(args.target_mask)
-            print_block(
-                name="Target Mask",
-                data={
-                    "Initial Shape": initial_shape,
-                    "Sampling Rate": tuple(np.round(target_mask.sampling_rate, 2)),
-                    "Final Shape": target_mask.shape,
-                },
-            )
+
+    if target_mask:
+        print_block(
+            name="Target Mask",
+            data={
+                "Initial Shape": initial_shape,
+                "Sampling Rate": tuple(np.round(target_mask.sampling_rate, 2)),
+                "Final Shape": target_mask.shape,
+            },
+        )
 
     initial_shape = template.shape
     _ = crop_data(data=template, data_mask=template_mask, cutoff=args.cutoff_template)
@@ -537,15 +539,14 @@ def main():
         opening_axis, tilt_axis = [int(x) for x in args.wedge_axes.split(",")]
 
         if args.tilt_step is not None:
-            tilt_angles = np.arange(
-                -tilt_start, tilt_stop + args.tilt_step, args.tilt_step
-            )
-            angles = np.zeros((template.data.ndim, tilt_angles.size))
-            angles[tilt_axis, :] = tilt_angles
-            template_filter["wedge_mask"] = {
-                "tilt_angles": angles,
+            template_filter["step_wedge_mask"] = {
+                "start_tilt": tilt_start,
+                "stop_tilt": tilt_stop,
+                "tilt_step": args.tilt_step,
                 "sigma": args.wedge_smooth,
-                "opening_axes": opening_axis,
+                "opening_axis": opening_axis,
+                "tilt_axis": tilt_axis,
+                "omit_negative_frequencies": True,
             }
         else:
             template_filter["continuous_wedge_mask"] = {
@@ -555,6 +556,7 @@ def main():
                 "opening_axis": opening_axis,
                 "infinite_plane": True,
                 "sigma": args.wedge_smooth,
+                "omit_negative_frequencies": True,
             }
 
     if template_mask is None:
@@ -577,8 +579,9 @@ def main():
     template_mask = template_mask.rigid_transform(
         rotation_matrix=np.eye(template_mask.data.ndim),
         translation=-translation,
+        order=1,
     )
-
+    template_mask.origin = template.origin.copy()
     print_block(
         name="Template Mask",
         data={
@@ -590,7 +593,9 @@ def main():
     print("\n" + "-" * 80)
 
     if args.scramble_phases:
-        template.data = scramble_phases(template.data, noise_proportion=1.0)
+        template.data = scramble_phases(
+            template.data, noise_proportion=1.0, normalize_power=True
+        )
 
     available_memory = backend.get_available_memory()
     if args.use_gpu:
@@ -668,7 +673,7 @@ def main():
         exit(-1)
 
     analyzer_args = {
-        "score_threshold": 0.0,
+        "score_threshold": 0,
         "number_of_peaks": 1000,
         "convolution_mode": "valid",
         "use_memmap": args.use_memmap,

{pytme-0.1.6.data → pytme-0.1.8.data}/scripts/postprocess.py

@@ -5,6 +5,8 @@
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
+from os import getcwd
+from os.path import join
 import argparse
 from sys import exit
 from typing import List, Tuple
@@ -13,6 +15,7 @@ from dataclasses import dataclass
 
 import numpy as np
 from scipy.spatial.transform import Rotation
+from numpy.typing import NDArray
 
 from tme import Density, Structure
 from tme.analyzer import (
@@ -26,6 +29,7 @@ from tme.matching_utils import (
     load_pickle,
     euler_to_rotationmatrix,
     euler_from_rotationmatrix,
+    centered_mask,
 )
 
 PEAK_CALLERS = {
@@ -52,16 +56,33 @@ def parse_args():
         help="Prefix for the output file name. Extension depends on output_format.",
     )
     parser.add_argument(
-        "--number_of_peaks", type=int, default=1000, help="Number of peaks to consider."
+        "--number_of_peaks",
+        type=int,
+        default=1000,
+        help="Number of peaks to consider. Note, this is the number of called peaks "
+        ", subject to min_distance and min_boundary_distance filtering. Therefore, the "
+        "returned number of peaks will be at most equal to number_of_peaks. "
+        "Ignored when --orientations is provided.",
     )
     parser.add_argument(
-        "--min_distance", type=int, default=5, help="Minimum distance between peaks."
+        "--min_distance",
+        type=int,
+        default=5,
+        help="Minimum distance between peaks. Ignored when --orientations is provided.",
     )
     parser.add_argument(
         "--min_boundary_distance",
         type=int,
         default=0,
-        help="Minimum distance from target boundaries.",
+        help="Minimum distance from target boundaries. Ignored when --orientations "
+        "is provided.",
+    )
+    parser.add_argument(
+        "--mask_edges",
+        action="store_true",
+        default=False,
+        help="Whether to mask edges of the input score array according to the template shape."
+        "Uses twice the value of --min_boundary_distance if boht are provided.",
     )
     parser.add_argument(
         "--wedge_mask",
@@ -73,7 +94,8 @@ def parse_args():
         "--peak_caller",
         choices=list(PEAK_CALLERS.keys()),
         default="PeakCallerScipy",
-        help="Peak caller to use for analysis. Ignored if input_file contains peaks.",
+        help="Peak caller to use for analysis. Ignored if input_file contains peaks or when "
+        "--orientations is provided.",
     )
     parser.add_argument(
         "--orientations",
@@ -205,8 +227,12 @@ class Orientations:
         return None
 
     def _to_relion_star(
-        self, filename: str, name_prefix: str = None, ctf_image: str = None,
-        sampling_rate : float = None, subtomogram_size : int = None
+        self,
+        filename: str,
+        name_prefix: str = None,
+        ctf_image: str = None,
+        sampling_rate: float = 1.0,
+        subtomogram_size: int = 0,
     ) -> None:
         """
         Save orientations in RELION's STAR file format.
@@ -249,12 +275,11 @@ class Orientations:
             "300.000000",
             str(int(subtomogram_size)),
             "3",
-            str(float(sampling_rate))
+            str(float(sampling_rate)),
         ]
         optics_header = "\n".join(optics_header)
         optics_data = "\t".join(optics_data)
 
-
         header = [
             "data_particles",
             "",
@@ -283,13 +308,14 @@ class Orientations:
             _ = ofile.write("\n# version 30001\n")
             _ = ofile.write(f"{header}\n")
 
+            # pyTME uses a zyx data layout
             for index, (translation, rotation, score, detail) in enumerate(self):
                 rotation = Rotation.from_euler("zyx", rotation, degrees=True)
-                rotation = rotation.as_euler(seq="zyz", degrees=True)
+                rotation = rotation.as_euler(seq="xyx", degrees=True)
 
                 translation_string = "\t".join([str(x) for x in translation][::-1])
-                angle_string = "\t".join([str(x) for x in rotation[::-1]])
-                name = f"{name_prefix}{index}.mrc"
+                angle_string = "\t".join([str(x) for x in rotation])
+                name = f"{name_prefix}_{index}.mrc"
                 _ = ofile.write(
                     f"{translation_string}\t{name}\t{angle_string}\t1{ctf_image}\n"
                 )
@@ -386,11 +412,25 @@ class Orientations:
         return translation, rotation, score, detail
 
 
+def load_template(filepath: str, sampling_rate: NDArray) -> "Density":
+    try:
+        template = Density.from_file(filepath)
+        template, _ = template.centered(0)
+        center_of_mass = template.center_of_mass(template.data)
+    except ValueError:
+        template = Structure.from_file(filepath)
+        center_of_mass = template.center_of_mass()[::-1]
+        template = Density.from_structure(template, sampling_rate=sampling_rate)
+
+    return template, center_of_mass
+
+
 def main():
     args = parse_args()
     data = load_pickle(args.input_file)
 
     meta = data[-1]
+    target_origin, _, sampling_rate, cli_args = meta
 
     if args.orientations is not None:
         orientations = Orientations.from_file(
@@ -402,13 +442,29 @@ def main():
         # Output is MaxScoreOverRotations
         if data[0].ndim == data[2].ndim:
             scores, offset, rotation_array, rotation_mapping, meta = data
+            if args.mask_edges:
+                template, center_of_mass = load_template(
+                    cli_args.template, sampling_rate=sampling_rate
+                )
+                if not cli_args.no_centering:
+                    template, *_ = template.centered(0)
+                mask_size = template.shape
+                if args.min_boundary_distance > 0:
+                    mask_size = 2 * args.min_boundary_distance
+                scores = centered_mask(scores, np.subtract(scores.shape, mask_size) + 1)
+
             peak_caller = PEAK_CALLERS[args.peak_caller](
                 number_of_peaks=args.number_of_peaks,
                 min_distance=args.min_distance,
                 min_boundary_distance=args.min_boundary_distance,
             )
             peak_caller(scores, rotation_matrix=np.eye(3))
-            candidates = peak_caller.merge([tuple(peak_caller)])
+            candidates = peak_caller.merge(
+                candidates=[tuple(peak_caller)],
+                number_of_peaks=args.number_of_peaks,
+                min_distance=args.min_distance,
+                min_boundary_distance=args.min_boundary_distance,
+            )
             if len(candidates) == 0:
                 exit(
                     "Found no peaks. Try reducing min_distance or min_boundary_distance."
@@ -436,18 +492,11 @@ def main():
         orientations.to_file(filename=f"{args.output_prefix}.tsv", file_format="text")
         exit(0)
 
-    target_origin, _, sampling_rate, cli_args = meta
-
-    template_is_density, index = True, 0
     _, template_extension = splitext(cli_args.template)
-    try:
-        template = Density.from_file(cli_args.template)
-        template, _ = template.centered(0)
-        center_of_mass = template.center_of_mass(template.data)
-    except ValueError:
-        template_is_density = False
-        template = Structure.from_file(cli_args.template)
-        center_of_mass = template.center_of_mass()[::-1]
+    template, center_of_mass = load_template(
+        filepath=cli_args.template, sampling_rate=sampling_rate
+    )
+    template_is_density, index = isinstance(template, Density), 0
 
     if args.output_format == "relion":
         new_shape = np.add(template.shape, np.mod(template.shape, 2))
@@ -457,10 +506,6 @@ def main():
 
     if args.output_format in ("extraction", "relion"):
         target = Density.from_file(cli_args.target)
-        if isinstance(template, Structure):
-            template = Density.from_structure(
-                template, sampling_rate=target.sampling_rate
-            )
 
         if not np.all(np.divide(target.shape, template.shape) > 2):
             print(
@@ -489,14 +534,15 @@ def main():
         )
 
         orientations = orientations[keep_peaks]
+        working_directory = getcwd()
         if args.output_format == "relion":
             orientations.to_file(
                 filename=f"{args.output_prefix}.star",
                 file_format="relion",
-                name_prefix=args.output_prefix,
+                name_prefix=join(working_directory, args.output_prefix),
                 ctf_image=args.wedge_mask,
-                sampling_rate = target.sampling_rate.max(),
-                subtomogram_size = template.shape[0]
+                sampling_rate=target.sampling_rate.max(),
+                subtomogram_size=template.shape[0],
             )
 
         peaks = peaks[keep_peaks,]
@@ -543,7 +589,9 @@ def main():
                 origin=candidate_starts[index] * sampling_rate,
             )
             # out_density.data = out_density.data * template_mask.data
-            out_density.to_file(f"{args.output_prefix}{index}.mrc")
+            out_density.to_file(
+                join(working_directory, f"{args.output_prefix}_{index}.mrc")
+            )
 
         exit(0)
 
@@ -566,7 +614,10 @@ def main():
                 translation=translation[::-1],
                 rotation_matrix=rotation_matrix[::-1, ::-1],
             )
-        transformed_template.to_file(f"{args.output_prefix}{index}{template_extension}")
+        # template_extension should contain the extension '.'
+        transformed_template.to_file(
+            f"{args.output_prefix}_{index}{template_extension}"
+        )
         index += 1