pytme 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.2__cp311-cp311-macosx_14_0_arm64.whl

scripts/postprocess.py

@@ -8,9 +8,8 @@
 import argparse
 from sys import exit
 from os import getcwd
-from os.path import join, abspath
-from typing import List
-from os.path import splitext
+from typing import List, Tuple
+from os.path import join, abspath, splitext
 
 import numpy as np
 from numpy.typing import NDArray
@@ -26,9 +25,11 @@ from tme.analyzer import (
 )
 from tme.matching_utils import (
     load_pickle,
+    centered_mask,
     euler_to_rotationmatrix,
     euler_from_rotationmatrix,
 )
+from tme.matching_optimization import create_score_object, optimize_match
 
 PEAK_CALLERS = {
     "PeakCallerSort": PeakCallerSort,
@@ -40,9 +41,7 @@ PEAK_CALLERS = {
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Peak Calling for Template Matching Outputs"
-    )
+    parser = argparse.ArgumentParser(description="Analyze Template Matching Outputs")
 
     input_group = parser.add_argument_group("Input")
     output_group = parser.add_argument_group("Output")
@@ -55,6 +54,13 @@ def parse_args():
         nargs="+",
         help="Path to the output of match_template.py.",
     )
+    input_group.add_argument(
+        "--background_file",
+        required=False,
+        nargs="+",
+        help="Path to an output of match_template.py used for normalization. "
+        "For instance from --scramble_phases or a different template.",
+    )
     input_group.add_argument(
         "--target_mask",
         required=False,
@@ -86,7 +92,7 @@ def parse_args():
             "average",
         ],
         default="orientations",
-        help="Available output formats:"
+        help="Available output formats: "
         "orientations (translation, rotation, and score), "
         "alignment (aligned template to target based on orientations), "
         "extraction (extract regions around peaks from targets, i.e. subtomograms), "
@@ -181,6 +187,13 @@ def parse_args():
         required=False,
         help="Number of accepted false-positives picks to determine minimum score.",
     )
+    additional_group.add_argument(
+        "--local_optimization",
+        action="store_true",
+        required=False,
+        help="[Experimental] Perform local optimization of candidates. Useful when the "
+        "number of identified candidats is small (< 10).",
+    )
 
     args = parser.parse_args()
 
@@ -195,38 +208,53 @@ def parse_args():
 
     if args.minimum_score is not None or args.n_false_positives is not None:
         args.number_of_peaks = np.iinfo(np.int64).max
-    else:
+    elif args.number_of_peaks is None:
         args.number_of_peaks = 1000
 
+    if args.background_file is None:
+        args.background_file = [None]
+    if len(args.background_file) == 1:
+        args.background_file = args.background_file * len(args.input_file)
+    elif len(args.background_file) not in (0, len(args.input_file)):
+        raise ValueError(
+            "--background_file needs to be specified once or for each --input_file."
+        )
+
     return args
 
 
-def load_template(filepath: str, sampling_rate: NDArray, center: bool = True):
+def load_template(
+    filepath: str,
+    sampling_rate: NDArray,
+    centering: bool = True,
+    target_shape: Tuple[int] = None,
+):
     try:
         template = Density.from_file(filepath)
-        center_of_mass = template.center_of_mass(template.data)
+        center = np.divide(np.subtract(template.shape, 1), 2)
         template_is_density = True
-    except ValueError:
+    except Exception:
         template = Structure.from_file(filepath)
-        center_of_mass = template.center_of_mass()[::-1]
+        center = template.center_of_mass()[::-1]
         template = Density.from_structure(template, sampling_rate=sampling_rate)
         template_is_density = False
 
-    translation = np.zeros_like(center_of_mass)
-    if center:
+    translation = np.zeros_like(center)
+    if centering and template_is_density:
         template, translation = template.centered(0)
+        center = np.divide(np.subtract(template.shape, 1), 2)
 
-    return template, center_of_mass, translation, template_is_density
+    return template, center, translation, template_is_density
 
 
-def merge_outputs(data, filepaths: List[str], args):
-    if len(filepaths) == 0:
+def merge_outputs(data, foreground_paths: List[str], background_paths: List[str], args):
+    if len(foreground_paths) == 0:
         return data, 1
 
     if data[0].ndim != data[2].ndim:
         return data, 1
 
-    from tme.matching_exhaustive import _normalize_under_mask
+    from tme.matching_exhaustive import normalize_under_mask
 
     def _norm_scores(data, args):
         target_origin, _, sampling_rate, cli_args = data[-1]
@@ -235,7 +263,7 @@ def merge_outputs(data, filepaths: List[str], args):
         ret = load_template(
             filepath=cli_args.template,
             sampling_rate=sampling_rate,
-            center=not cli_args.no_centering,
+            centering=not cli_args.no_centering,
         )
         template, center_of_mass, translation, template_is_density = ret
 
@@ -256,13 +284,16 @@ def merge_outputs(data, filepaths: List[str], args):
             mask.shape, np.multiply(args.min_boundary_distance, 2)
         ).astype(int)
         mask[cropped_shape] = 0
-        _normalize_under_mask(template=data[0], mask=mask, mask_intensity=mask.sum())
+        normalize_under_mask(template=data[0], mask=mask, mask_intensity=mask.sum())
         return data[0]
 
     entities = np.zeros_like(data[0])
     data[0] = _norm_scores(data=data, args=args)
-    for index, filepath in enumerate(filepaths):
-        new_scores = _norm_scores(data=load_pickle(filepath), args=args)
+    for index, filepath in enumerate(foreground_paths):
+        new_scores = _norm_scores(
+            data=load_match_template_output(filepath, background_paths[index]),
+            args=args,
+        )
         indices = new_scores > data[0]
         entities[indices] = index + 1
         data[0][indices] = new_scores[indices]
@@ -270,9 +301,18 @@ def merge_outputs(data, filepaths: List[str], args):
     return data, entities
 
 
+def load_match_template_output(foreground_path, background_path):
+    data = load_pickle(foreground_path)
+    if background_path is not None:
+        data_background = load_pickle(background_path)
+        data[0] = (data[0] - data_background[0]) / (1 - data_background[0])
+        np.fmax(data[0], 0, out=data[0])
+    return data
+
+
 def main():
     args = parse_args()
-    data = load_pickle(args.input_file[0])
+    data = load_match_template_output(args.input_file[0], args.background_file[0])
 
     target_origin, _, sampling_rate, cli_args = data[-1]
 
@@ -280,7 +320,7 @@ def main():
     ret = load_template(
         filepath=cli_args.template,
         sampling_rate=sampling_rate,
-        center=not cli_args.no_centering,
+        centering=not cli_args.no_centering,
     )
     template, center_of_mass, translation, template_is_density = ret
 
@@ -310,7 +350,14 @@ def main():
         max_shape = np.max(template.shape)
         args.min_boundary_distance = np.ceil(np.divide(max_shape, 2))
 
-    # data, entities = merge_outputs(data=data, filepaths=args.input_file[1:], args=args)
+    entities = None
+    if len(args.input_file) > 1:
+        data, entities = merge_outputs(
+            data=data,
+            foreground_paths=args.input_file,
+            background_paths=args.background_file,
+            args=args,
+        )
 
     orientations = args.orientations
     if orientations is None:
@@ -323,57 +370,69 @@ def main():
                 target_mask = Density.from_file(args.target_mask)
                 scores = scores * target_mask.data
 
-            if args.n_false_positives is not None:
-                args.n_false_positives = max(args.n_false_positives, 1)
-                cropped_shape = np.subtract(
-                    scores.shape, np.multiply(args.min_boundary_distance, 2)
-                ).astype(int)
+            cropped_shape = np.subtract(
+                scores.shape, np.multiply(args.min_boundary_distance, 2)
+            ).astype(int)
 
-                cropped_shape = tuple(
+            if args.min_boundary_distance > 0:
+                scores = centered_mask(scores, new_shape=cropped_shape)
+
+            if args.n_false_positives is not None:
+                # Rickgauer et al. 2017
+                cropped_slice = tuple(
                     slice(
                         int(args.min_boundary_distance),
                         int(x - args.min_boundary_distance),
                     )
                     for x in scores.shape
                 )
-                # Rickgauer et al. 2017
-                n_correlations = np.size(scores[cropped_shape]) * len(rotation_mapping)
+                args.n_false_positives = max(args.n_false_positives, 1)
+                n_correlations = np.size(scores[cropped_slice]) * len(rotation_mapping)
                 minimum_score = np.multiply(
                     erfcinv(2 * args.n_false_positives / n_correlations),
-                    np.sqrt(2) * np.std(scores[cropped_shape]),
+                    np.sqrt(2) * np.std(scores[cropped_slice]),
                 )
                 print(f"Determined minimum score cutoff: {minimum_score}.")
                 minimum_score = max(minimum_score, 0)
                 args.minimum_score = minimum_score
 
-            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,
-            )
-            if args.minimum_score is not None:
-                args.number_of_peaks = np.inf
+            args.batch_dims = None
+            if hasattr(cli_args, "target_batch"):
+                args.batch_dims = cli_args.target_batch
+
+            peak_caller_kwargs = {
+                "number_of_peaks": args.number_of_peaks,
+                "min_distance": args.min_distance,
+                "min_boundary_distance": args.min_boundary_distance,
+                "batch_dims": args.batch_dims,
+                "minimum_score": args.minimum_score,
+                "maximum_score": args.maximum_score,
+            }
 
+            peak_caller = PEAK_CALLERS[args.peak_caller](**peak_caller_kwargs)
             peak_caller(
                 scores,
-                rotation_matrix=np.eye(3),
+                rotation_matrix=np.eye(template.data.ndim),
                 mask=template.data,
                 rotation_mapping=rotation_mapping,
                 rotation_array=rotation_array,
-                minimum_score=args.minimum_score,
             )
             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,
+                candidates=[tuple(peak_caller)], **peak_caller_kwargs
             )
             if len(candidates) == 0:
-                print("Found no peaks. Consider changing peak calling parameters.")
+                candidates = [[], [], [], []]
+                print("Found no peaks, consider changing peak calling parameters.")
                 exit(-1)
 
             for translation, _, score, detail in zip(*candidates):
-                rotations.append(rotation_mapping[rotation_array[tuple(translation)]])
+                rotation_index = rotation_array[tuple(translation)]
+                rotation = rotation_mapping.get(
+                    rotation_index, np.zeros(template.data.ndim, int)
+                )
+                if rotation.ndim == 2:
+                    rotation = euler_from_rotationmatrix(rotation)
+                rotations.append(rotation)
 
         else:
             candidates = data
@@ -381,8 +440,13 @@ def main():
             for i in range(translation.shape[0]):
                 rotations.append(euler_from_rotationmatrix(rotation[i]))
 
-        rotations = np.vstack(rotations).astype(float)
+        if len(rotations):
+            rotations = np.vstack(rotations).astype(float)
         translations, scores, details = candidates[0], candidates[2], candidates[3]
+
+        if entities is not None:
+            details = entities[tuple(translations.T)]
+
         orientations = Orientations(
             translations=translations,
             rotations=rotations,
@@ -390,14 +454,55 @@ def main():
             details=details,
         )
 
-    if args.minimum_score is not None:
+    if args.minimum_score is not None and len(orientations.scores):
         keep = orientations.scores >= args.minimum_score
         orientations = orientations[keep]
 
-    if args.maximum_score is not None:
+    if args.maximum_score is not None and len(orientations.scores):
         keep = orientations.scores <= args.maximum_score
         orientations = orientations[keep]
 
+    if args.peak_oversampling > 1:
+        peak_caller = peak_caller = PEAK_CALLERS[args.peak_caller]()
+        if data[0].ndim != data[2].ndim:
+            print(
+                "Input pickle does not contain template matching scores."
+                " Cannot oversample peaks."
+            )
+            exit(-1)
+        orientations.translations = peak_caller.oversample_peaks(
+            scores=data[0],
+            peak_positions=orientations.translations,
+            oversampling_factor=args.peak_oversampling,
+        )
+
+    if args.local_optimization:
+        target = Density.from_file(cli_args.target)
+        orientations.translations = orientations.translations.astype(np.float32)
+        orientations.rotations = orientations.rotations.astype(np.float32)
+        for index, (translation, angles, *_) in enumerate(orientations):
+            score_object = create_score_object(
+                score="FLC",
+                target=target.data.copy(),
+                template=template.data.copy(),
+                template_mask=template_mask.data.copy(),
+            )
+
+            center = np.divide(template.shape, 2)
+            init_translation = np.subtract(translation, center)
+            bounds_translation = tuple((x - 5, x + 5) for x in init_translation)
+
+            translation, rotation_matrix, score = optimize_match(
+                score_object=score_object,
+                optimization_method="basinhopping",
+                bounds_translation=bounds_translation,
+                maxiter=3,
+                x0=[*init_translation, *angles],
+            )
+            orientations.translations[index] = np.add(translation, center)
+            orientations.rotations[index] = angles
+            orientations.scores[index] = score * -1
+
     if args.output_format == "orientations":
         orientations.to_file(filename=f"{args.output_prefix}.tsv", file_format="text")
         exit(0)
@@ -506,7 +611,7 @@ def main():
             return_orientations=True,
         )
         out = np.zeros_like(template.data)
-        out = np.zeros(np.multiply(template.shape, 2).astype(int))
+        # out = np.zeros(np.multiply(template.shape, 2).astype(int))
         for index in range(len(cand_slices)):
             from scipy.spatial.transform import Rotation
 
@@ -515,7 +620,6 @@ def main():
             )
             rotation_matrix = rotation.inv().as_matrix()
 
-            # rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
             subset = Density(target.data[obs_slices[index]])
             subset = subset.rigid_transform(rotation_matrix=rotation_matrix, order=1)
 
@@ -526,35 +630,30 @@ def main():
         ret.to_file(f"{args.output_prefix}_average.mrc")
         exit(0)
 
-    if args.peak_oversampling > 1:
-        peak_caller = peak_caller = PEAK_CALLERS[args.peak_caller]()
-        if data[0].ndim != data[2].ndim:
-            print(
-                "Input pickle does not contain template matching scores."
-                " Cannot oversample peaks."
-            )
-            exit(-1)
-        orientations.translations = peak_caller.oversample_peaks(
-            score_space=data[0],
-            translations=orientations.translations,
-            oversampling_factor=args.oversampling_factor,
-        )
+    template, center, *_ = load_template(
+        filepath=cli_args.template,
+        sampling_rate=sampling_rate,
+        centering=not cli_args.no_centering,
+        target_shape=target.shape,
+    )
 
     for index, (translation, angles, *_) in enumerate(orientations):
         rotation_matrix = euler_to_rotationmatrix(angles)
         if template_is_density:
-            translation = np.subtract(translation, center_of_mass)
+            translation = np.subtract(translation, center)
             transformed_template = template.rigid_transform(
                 rotation_matrix=rotation_matrix
             )
-            new_origin = np.add(target_origin / sampling_rate, translation)
-            transformed_template.origin = np.multiply(new_origin, sampling_rate)
+            transformed_template.origin = np.add(
+                target_origin, np.multiply(translation, sampling_rate)
+            )
+
         else:
             template = Structure.from_file(cli_args.template)
             new_center_of_mass = np.add(
                 np.multiply(translation, sampling_rate), target_origin
             )
-            translation = np.subtract(new_center_of_mass, center_of_mass)
+            translation = np.subtract(new_center_of_mass, center)
             transformed_template = template.rigid_transform(
                 translation=translation[::-1],
                 rotation_matrix=rotation_matrix[::-1, ::-1],