pytme 0.2.0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.1__cp311-cp311-macosx_14_0_arm64.whl

scripts/refine_matches.py

@@ -1,52 +1,121 @@
 #!python3
-""" CLI to refine template matching candidates.
+""" Iterative template matching parameter tuning.
 
-    Copyright (c) 2023-2024 European Molecular Biology Laboratory
+    Copyright (c) 2024 European Molecular Biology Laboratory
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 import argparse
-from time import time
+import subprocess
+from sys import exit
+from shutil import copyfile
+from typing import Tuple, List, Dict
 
 import numpy as np
-from numpy.typing import NDArray
-
-from tme.backends import backend
-from tme import Density, Structure
-from tme.matching_data import MatchingData
-from tme.analyzer import MaxScoreOverRotations, MaxScoreOverTranslations
-from tme.matching_utils import (
-    load_pickle,
-    get_rotation_matrices,
-    compute_parallelization_schedule,
-)
-from tme.matching_exhaustive import scan_subsets, MATCHING_EXHAUSTIVE_REGISTER
-
-from postprocess import Orientations
-from match_template import load_and_validate_mask
+from scipy import optimize
+
+from tme import Orientations, Density
+from tme.matching_utils import generate_tempfile_name, load_pickle
+from tme.matching_exhaustive import MATCHING_EXHAUSTIVE_REGISTER
+
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description="Refine Template Matching Orientations."
+        description="Refine template matching candidates using deep matching.",
     )
-    parser.add_argument(
-        "--input_file",
+    io_group = parser.add_argument_group("Input / Output")
+    io_group.add_argument(
+        "--orientations",
         required=True,
         type=str,
-        help="Path to the output of match_template.py.",
+        help="Path to an orientations file in a supported format. See "
+        "https://kosinskilab.github.io/pyTME/reference/api/tme.orientations.Orientations.from_file.html"
+        " for available options."
     )
-    parser.add_argument(
-        "--orientations",
+    io_group.add_argument(
+        "--output_prefix",
         required=True,
         type=str,
-        help="Path to orientations from postprocess.py.",
+        help="Path to write output to."
+    )
+    io_group.add_argument(
+        "--iterations",
+        required=False,
+        default=0,
+        type=int,
+        help="Number of refinement iterations to perform.",
+    )
+    io_group.add_argument(
+        "--verbose",
+        action="store_true",
+        default=False,
+        help="More verbose and more files written to disk.",
+    )
+    matching_group = parser.add_argument_group("Template Matching")
+    matching_group.add_argument(
+        "--input_file",
+        required=False,
+        type=str,
+        help="Path to the output of match_template.py.",
+    )
+    matching_group.add_argument(
+        "-m",
+        "--target",
+        dest="target",
+        type=str,
+        required=False,
+        help="Path to a target in CCP4/MRC, EM, H5 or another format supported by "
+        "tme.density.Density.from_file "
+        "https://kosinskilab.github.io/pyTME/reference/api/tme.density.Density.from_file.html",
     )
-    parser.add_argument(
-        "--output_file",
+    matching_group.add_argument(
+        "--target_mask",
+        dest="target_mask",
+        type=str,
+        required=False,
+        help="Path to a mask for the target in a supported format (see target).",
+    )
+    matching_group.add_argument(
+        "-i",
+        "--template",
+        dest="template",
+        type=str,
+        required=False,
+        help="Path to a template in PDB/MMCIF or other supported formats (see target).",
+    )
+    matching_group.add_argument(
+        "--template_mask",
+        dest="template_mask",
+        type=str,
+        required=False,
+        help="Path to a mask for the template in a supported format (see target).",
+    )
+    matching_group.add_argument(
+        "--invert_target_contrast",
+        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 "
+        "negative values, e.g. tomograms.",
+    )
+    matching_group.add_argument(
+        "--angular_sampling",
+        dest="angular_sampling",
         required=True,
-        help="Path to the refined output orientations.",
+        default=None,
+        help="Angular sampling rate using optimized rotational sets."
+        "A lower number yields more rotations. Values >= 180 sample only the identity.",
+    )
+    matching_group.add_argument(
+        "-s",
+        dest="score",
+        type=str,
+        default="FLCSphericalMask",
+        choices=list(MATCHING_EXHAUSTIVE_REGISTER.keys()),
+        help="Template matching scoring function.",
     )
-    parser.add_argument(
+    matching_group.add_argument(
         "-n",
         dest="cores",
         required=False,
@@ -54,165 +123,300 @@ def parse_args():
         default=4,
         help="Number of cores used for template matching.",
     )
+    matching_group.add_argument(
+        "--use_gpu",
+        dest="use_gpu",
+        action="store_true",
+        default=False,
+        help="Whether to perform computations on the GPU.",
+    )
+    matching_group.add_argument(
+        "--no_centering",
+        dest="no_centering",
+        action="store_true",
+        help="Assumes the template is already centered and omits centering.",
+    )
+    matching_group.add_argument(
+        "--no_edge_padding",
+        dest="no_edge_padding",
+        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.",
+    )
+    matching_group.add_argument(
+        "--no_fourier_padding",
+        dest="no_fourier_padding",
+        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.",
+    )
+
+    peak_group = parser.add_argument_group("Peak Calling")
+    peak_group.add_argument(
+        "--number_of_peaks",
+        type=int,
+        default=100,
+        required=False,
+        help="Upper limit of peaks to call, subject to filtering parameters. Default 1000. "
+        "If minimum_score is provided all peaks scoring higher will be reported.",
+    )
+    extraction_group = parser.add_argument_group("Extraction")
+    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.",
+    )
+
+    optimization_group = parser.add_argument_group("Optimization")
+    optimization_group.add_argument(
+        "--lowpass",
+        dest="lowpass",
+        action="store_true",
+        default=False,
+        help="Optimize template matching lowpass filter cutoff.",
+    )
+    optimization_group.add_argument(
+        "--highpass",
+        dest="highpass",
+        action="store_true",
+        default=False,
+        help="Optimize template matching highpass filter cutoff.",
+    )
 
     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 load_template(filepath: str, sampling_rate: NDArray) -> "Density":
-    try:
-        template = Density.from_file(filepath)
-    except ValueError:
-        template = Structure.from_file(filepath)
-        template = Density.from_structure(template, sampling_rate=sampling_rate)
+def argdict_to_command(input_args: Dict, executable: str) -> List:
+    ret = []
+    for key, value in input_args.items():
+        if value is None:
+            continue
+        elif isinstance(value, bool):
+            if value:
+                ret.append(key)
+        else:
+            ret.extend([key, value])
+
+    ret = [str(x) for x in ret]
+    ret.insert(0, executable)
+    return " ".join(ret)
+
+
+def create_stacking_argdict(args) -> Dict:
+    arg_dict = {
+        "--input_file": args.input_file,
+        "--orientations": args.orientations,
+        "--output_file": args.candidate_stack_path,
+        "--keep_out_of_box": args.keep_out_of_box,
+    }
+    return arg_dict
+
+
+def create_matching_argdict(args) -> Dict:
+    arg_dict = {
+        "--target": args.target,
+        "--template": args.template,
+        "--template_mask": args.template_mask,
+        "--target_mask": args.target_mask,
+        "-o": args.match_template_path,
+        "-a": args.angular_sampling,
+        "-s": args.score,
+        "--no_fourier_padding": True,
+        "--no_edge_padding": True,
+        "--no_centering": args.no_centering,
+        "-n": args.cores,
+        "--invert_target_contrast": args.invert_target_contrast,
+        "--use_gpu": args.use_gpu,
+    }
+    return arg_dict
+
+
+def create_postprocessing_argdict(args) -> Dict:
+    arg_dict = {
+        "--input_file": args.match_template_path,
+        "--target_mask": args.target_mask,
+        "--output_prefix": args.new_orientations_path,
+        "--peak_caller": "PeakCallerMaximumFilter",
+        "--number_of_peaks": args.number_of_peaks,
+        "--output_format": "orientations",
+        "--mask_edges" : True,
+    }
+    return arg_dict
+
+def update_orientations(old, new, args):
+    stack_shape = Density.from_file(
+        args.candidate_stack_path, use_memmap = True
+    ).shape
+
+    stack_center = np.add(
+        np.divide(stack_shape, 2).astype(int),
+        np.mod(stack_shape, 2)
+    )
 
-    return template
+    print(old.translations[new.translations[:, 0].astype(int)])
+    print(new.translations[:, 1:], stack_center, stack_shape)
 
+    new.translations[:, 1:] = np.add(
+        old.translations[new.translations[:, 0].astype(int)],
+        np.subtract(new.translations, stack_center)[:, 1:]
+    )
+    print(new.translations)
+    # The effect of --align_orientations should be handled herer
+    return new
 
-def main():
-    args = parse_args()
-    meta = load_pickle(args.input_file)[-1]
-    target_origin, _, sampling_rate, cli_args = meta
 
-    orientations = Orientations.from_file(
-        filename=args.orientations, file_format="text"
-    )
 
-    template = load_template(cli_args.template, sampling_rate)
-    template_mask = load_and_validate_mask(
-        mask_target=template, mask_path=cli_args.template_mask
-    )
 
-    if not cli_args.no_centering:
-        template, translation = template.centered(0)
+class DeepMatcher:
+    def __init__(self, args):
+        self.orientations = Orientations.from_file(args.orientations)
 
-    if template_mask is None:
-        template_mask = template.empty
-        if not cli_args.no_centering:
-            enclosing_box = template.minimum_enclosing_box(
-                0, use_geometric_center=False
-            )
-            template_mask.adjust_box(enclosing_box)
+        match_template_args = create_matching_argdict(args)
+        match_template_args["--target"] = args.candidate_stack_path
+        self.match_template_args = match_template_args
 
-        template_mask.data[:] = 1
-        translation = np.zeros_like(translation)
+        self.filter_parameters = {}
+        if args.lowpass:
+            self.filter_parameters["--lowpass"] = 0
+        if args.highpass:
+            self.filter_parameters["--highpass"] = 200
 
-    template_mask.pad(template.shape, center=False)
-    origin_translation = np.divide(
-        np.subtract(template.origin, template_mask.origin), template.sampling_rate
-    )
-    translation = np.add(translation, origin_translation)
+        self.postprocess_args = create_postprocessing_argdict(args)
+        self.postprocess_args["--number_of_peaks"] = 1
 
-    template_mask = template_mask.rigid_transform(
-        rotation_matrix=np.eye(template_mask.data.ndim),
-        translation=-translation,
-        order=1,
-    )
-    template_mask.origin = template.origin.copy()
 
-    target = Density.from_file(cli_args.target)
-    peaks = orientations.translations.astype(int)
-    half_shape = np.divide(template.shape, 2).astype(int)
-    observation_starts = np.subtract(peaks, half_shape)
-    observation_stops = np.add(peaks, half_shape) + np.mod(template.shape, 2).astype(
-        int
-    )
+    def get_initial_values(self) -> Tuple[float]:
+        ret = tuple(float(x) for x in self.filter_parameters.values())
+        return ret
 
-    pruned_starts = np.maximum(observation_starts, 0)
-    pruned_stops = np.minimum(observation_stops, target.shape)
+    def format_parameters(self, parameter_values: Tuple[float]) -> Dict:
+        ret = {}
+        for value, key in zip(parameter_values, ret.keys()):
+            ret[key] = value
+            if isinstance(self.filter_parameters[key], bool):
+                ret[key] = value > 0.5
+        return ret
 
-    keep_peaks = (
-        np.sum(
-            np.multiply(
-                observation_starts == pruned_starts, observation_stops == pruned_stops
-            ),
-            axis=1,
+    def __call__(self, x: Tuple[float]):
+        filter_parameters = self.format_parameters(x)
+        self.match_template_args.update(filter_parameters)
+        match_template = argdict_to_command(
+            self.match_template_args,
+            executable="python3 /Users/vmaurer/src/pytme/scripts/match_template_filters.py"
         )
-        == observation_starts.shape[1]
-    )
-    observation_starts = observation_starts[keep_peaks]
-    observation_stops = observation_stops[keep_peaks]
+        _ = subprocess.run(match_template, capture_output=True, shell = True)
 
-    observation_slices = [
-        tuple(slice(s, e) for s, e in zip(start_row, stop_row))
-        for start_row, stop_row in zip(observation_starts, observation_stops)
-    ]
+        # Assume we get a new peak for each input in the same order
+        postprocess = argdict_to_command(
+            self.postprocess_args,
+            executable="python3 /Users/vmaurer/src/pytme/scripts/postprocess.py"
+        )
+        _ = subprocess.run(postprocess, capture_output=True, shell = True)
 
-    matching_data = MatchingData(target=target, template=template)
-    matching_data.rotations = np.eye(template.data.ndim).reshape(1, 3, 3)
+        orientations_new = Orientations.from_file(
+            f"{self.postprocess_args['--output_prefix']}.tsv"
+        )
+        orientations_new = orientations_new[np.argsort(orientations_new.translations[:, 0])]
 
-    target_pad = matching_data.target_padding(pad_target=True)
-    out_shape = np.add(target_pad, template.shape).astype(int)
+        label, dist = self.orientations.details, orientations_new.scores
+        loss = np.add(
+            (1 - label) * np.square(dist),
+            label * np.square(np.fmax(.5 - dist, 0.0))
+        )
+        return np.sum(loss)
 
-    observations = np.zeros((len(observation_slices), *out_shape))
 
 
-    for idx, obs_slice in enumerate(observation_slices):
-        subset = matching_data.subset_by_slice(
-            target_slice=obs_slice,
-            target_pad=target_pad,
-            invert_target=cli_args.invert_target_contrast,
+def main():
+    args = parse_args()
+
+    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
+
+    args.candidate_stack_path = generate_tempfile_name(suffix = ".h5")
+    args.new_orientations_path = generate_tempfile_name()
+    args.match_template_path = generate_tempfile_name()
+
+    match_deep = DeepMatcher(args)
+    initial_values = match_deep.get_initial_values()
+
+    # Do a single pass over the data
+    if len(initial_values) == 0:
+        create_image_stack = create_stacking_argdict(args)
+        create_image_stack = argdict_to_command(
+            create_image_stack,
+            executable="python3 /Users/vmaurer/src/pytme/scripts/extract_candidates.py"
         )
-        xd = template.copy()
-        xd.pad(subset.target.shape, center = True)
-        # observations[idx] = subset.target
-        observations[idx] = xd.data
-
-
-    matching_data = MatchingData(target=observations, template=template)
-    matching_data._set_batch_dimension(target_dims=0, template_dims=None)
-    matching_data.rotations = get_rotation_matrices(15)
-    if template_mask is not None:
-        matching_data.template_mask = template_mask.data
-
-    template_box = np.ones(len(matching_data._output_template_shape), dtype=int)
-    target_padding = np.zeros_like(matching_data._output_target_shape)
-
-    scoring_method = "FLC"
-    callback_class = MaxScoreOverRotations
-    splits, schedule = compute_parallelization_schedule(
-        shape1=matching_data._output_target_shape,
-        shape2=template_box,
-        shape1_padding=target_padding,
-        max_cores=args.cores,
-        max_ram=backend.get_available_memory(),
-        split_only_outer=False,
-        matching_method=scoring_method,
-        analyzer_method=callback_class.__name__,
-        backend=backend._backend_name,
-        float_nbytes=backend.datatype_bytes(backend._default_dtype),
-        complex_nbytes=backend.datatype_bytes(backend._complex_dtype),
-        integer_nbytes=backend.datatype_bytes(backend._default_dtype_int),
-    )
-
-    matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[scoring_method]
-
-    start = time()
-    candidates = scan_subsets(
-        matching_data=matching_data,
-        matching_score=matching_score,
-        matching_setup=matching_setup,
-        callback_class=callback_class,
-        callback_class_args={
-            # "score_space_shape" : (
-                # matching_data.rotations.shape[0], observations.shape[0]
-            # ),
-            # "score_space_dtype" : backend._default_dtype,
-            # "template_shape" : (matching_data.rotations.shape[0], *matching_data._template.shape)
-        },
-        target_splits=splits,
-        job_schedule=schedule,
-        pad_target_edges=False,
-        pad_fourier=False,
-        interpolation_order=cli_args.interpolation_order,
-    )
-    print(candidates[0].max())
-    Density(candidates[0][0]).to_file("scores.mrc")
-
-    runtime = time() - start
-    print(f"\nRuntime real: {runtime:.3f}s user: {(runtime * args.cores):.3f}s.")
+        _ = subprocess.run(create_image_stack, capture_output=True, shell = True)
+        if args.verbose:
+            copyfile(args.candidate_stack_path, f"{args.output_prefix}_stack.h5")
+        _ = match_deep(x = ())
+        if args.verbose:
+            copyfile(args.match_template_path, f"{args.output_prefix}_stack.pickle")
+        orientations = Orientations.from_file(f"{args.new_orientations_path}.tsv")
+        orientations = update_orientations(
+            new = orientations,
+            old = Orientations.from_file(args.orientations),
+            args= args
+        )
+        orientations.to_file(f"{args.output_prefix}.tsv")
+        exit(0)
+
+    for current_iteration in range(args.iterations):
+        create_image_stack = create_stacking_argdict(args)
+        create_image_stack = argdict_to_command(
+            create_image_stack,
+            executable="python3 /Users/vmaurer/src/pytme/scripts/extract_candidates.py"
+        )
+        _ = subprocess.run(create_image_stack, capture_output=True, shell = True)
+
+        # Perhaps we need a different optimizer here to use sensible steps for each parameter
+        match_deep = DeepMatcher(args)
+        parameters = optimize.minimize(
+            x0=match_deep.get_initial_values(),
+            fun=match_deep,
+            options = {"maxiter" : 2}
+        )
+        parameter_dict = match_deep.format_parameters(parameters)
 
+        match_template = create_matching_argdict(args)
+        match_template.update(parameter_dict)
+        match_template = argdict_to_command(
+            match_template,
+            executable="python3 /Users/vmaurer/src/pytme/scripts/match_template_filters.py"
+        )
+        _ = subprocess.run(match_template, capture_output=True, shell = True)
+
+        # Some form of labelling is necessary for these matches
+        # 1. All of them are true positives
+        # 2. All of them are true positives up to a certain threshold
+        # 3. Kernel fitting
+        # 4. Perhaps also sensible to include a certain percentage of low scores as true negatives
+        postprocess = create_postprocessing_argdict(args)
+        postprocess = argdict_to_command(postprocess, executable="postprocess.py")
+        _ = subprocess.run(postprocess, capture_output=True, shell = True)
+        args.orientations = f"{args.new_orientations_path}.tsv"
+        orientations = Orientations.from_file(args.orientations)
+        orientations.to_file(f"{args.output_prefix}_{current_iteration}.tsv")
 
 if __name__ == "__main__":
     main()

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.0"
+__version__ = "0.2.1"