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

scripts/preprocess.py

@@ -1,92 +1,147 @@
 #!python3
-""" Apply tme.preprocessor.Preprocessor methods to an input file based
-    on a provided yaml configuration obtaiend from preprocessor_gui.py.
+""" Preprocessing routines for template matching.
 
     Copyright (c) 2023 European Molecular Biology Laboratory
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
-import yaml
+import warnings
 import argparse
-import textwrap
-from tme import Preprocessor, Density
+import numpy as np
+
+from tme import Density, Structure
+from tme.backends import backend as be
+from tme.preprocessing.frequency_filters import BandPassFilter
 
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description=textwrap.dedent(
-            """
-        Apply preprocessing to an input file based on a provided YAML configuration.
-
-        Expected YAML file format:
-        ```yaml
-        <method_name>:
-            <parameter1>: <value1>
-            <parameter2>: <value2>
-            ...
-        ```
-        """
-        ),
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description="Perform template matching preprocessing.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser.add_argument(
-        "-i",
-        "--input_file",
+
+    io_group = parser.add_argument_group("Input / Output")
+    io_group.add_argument(
+        "-m",
+        "--data",
+        dest="data",
         type=str,
         required=True,
-        help="Path to the input data file in CCP4/MRC format.",
+        help="Path to a file in PDB/MMCIF, CCP4/MRC, EM, H5 or a format supported by "
+        "tme.density.Density.from_file "
+        "https://kosinskilab.github.io/pyTME/reference/api/tme.density.Density.from_file.html",
     )
-    parser.add_argument(
-        "-y",
-        "--yaml_file",
+    io_group.add_argument(
+        "-o",
+        "--output",
+        dest="output",
         type=str,
         required=True,
-        help="Path to the YAML configuration file.",
+        help="Path the output should be written to.",
     )
-    parser.add_argument(
-        "-o",
-        "--output_file",
-        type=str,
+
+    box_group = parser.add_argument_group("Box")
+    box_group.add_argument(
+        "--box_size",
+        dest="box_size",
+        type=int,
         required=True,
-        help="Path to output file in CPP4/MRC format..",
+        help="Box size of the output",
     )
-    parser.add_argument(
-        "--compress", action="store_true", help="Compress the output file using gzip."
+    box_group.add_argument(
+        "--sampling_rate",
+        dest="sampling_rate",
+        type=float,
+        required=True,
+        help="Sampling rate of the output file.",
     )
 
+    modulation_group = parser.add_argument_group("Modulation")
+    modulation_group.add_argument(
+        "--invert_contrast",
+        dest="invert_contrast",
+        action="store_true",
+        required=False,
+        help="Inverts the template contrast.",
+    )
+    modulation_group.add_argument(
+        "--lowpass",
+        dest="lowpass",
+        type=float,
+        required=False,
+        default=None,
+        help="Lowpass filter the template to the given resolution. Nyquist by default. "
+        "A value of 0 disables the filter.",
+    )
+    modulation_group.add_argument(
+        "--no_centering",
+        dest="no_centering",
+        action="store_true",
+        help="Assumes the template is already centered and omits centering.",
+    )
+    modulation_group.add_argument(
+        "--backend",
+        dest="backend",
+        type=str,
+        default=None,
+        choices=be.available_backends(),
+        help="Determines more suitable box size for the given compute backend.",
+    )
     args = parser.parse_args()
-
     return args
 
 
 def main():
     args = parse_args()
-    with open(args.yaml_file, "r") as f:
-        preprocess_settings = yaml.safe_load(f)
-
-    if len(preprocess_settings) > 1:
-        raise NotImplementedError(
-            "Multiple preprocessing methods specified. "
-            "The script currently supports one method at a time."
-        )
-
-    method_name = list(preprocess_settings.keys())[0]
-    if not hasattr(Preprocessor, method_name):
-        raise ValueError(f"Method {method_name} does not exist in Preprocessor.")
-
-    density = Density.from_file(args.input_file)
-    output = density.empty
-
-    method_params = preprocess_settings[method_name]
-    preprocessor = Preprocessor()
-    method = getattr(preprocessor, method_name, None)
-    if not method:
-        raise ValueError(
-            f"{method} does not exist in dge.preprocessor.Preprocessor class."
-        )
-
-    output.data = method(template=density.data, **method_params)
-    output.to_file(args.output_file, gzip=args.compress)
+
+    try:
+        data = Structure.from_file(args.data)
+        data = Density.from_structure(data, sampling_rate=args.sampling_rate)
+    except NotImplementedError:
+        data = Density.from_file(args.data)
+
+    if not args.no_centering:
+        data, _ = data.centered(0)
+
+    for name in be.available_backends():
+        be.change_backend(name, device="cpu")
+        box = be.compute_convolution_shapes([args.box_size], [1])[1][0]
+        if box != args.box_size and args.backend is None:
+            print(f"Consider --box_size {box} instead of {args.box_size} for {name}.")
+
+    if args.backend is not None:
+        be.change_backend(name, device="cpu")
+        box = be.compute_convolution_shapes([args.box_size], [1])[1][0]
+        if box != args.box_size:
+            print(f"Changed --box_size from {args.box_size} to {box}.")
+        args.box_size = box
+
+    data.pad(
+        np.multiply(args.box_size, np.divide(args.sampling_rate, data.sampling_rate)),
+        center=True,
+    )
+
+    bpf_mask = 1
+    lowpass = 2 * args.sampling_rate if args.lowpass is None else args.lowpass
+    if args.lowpass != 0:
+        bpf_mask = BandPassFilter(
+            lowpass=lowpass,
+            highpass=None,
+            use_gaussian=True,
+            return_real_fourier=True,
+            shape_is_real_fourier=False,
+        )(shape=data.shape)["data"]
+
+    data_ft = np.fft.rfftn(data.data, s=data.shape)
+    data_ft = np.multiply(data_ft, bpf_mask, out=data_ft)
+    data.data = np.fft.irfftn(data_ft, s=data.shape).real
+
+    data = data.resample(args.sampling_rate, method="spline", order=3)
+
+    if args.invert_contrast:
+        data.data = data.data * -1
+
+    data.to_file(args.output)
 
 
 if __name__ == "__main__":

scripts/preprocessor_gui.py

@@ -132,14 +132,6 @@ def local_gaussian_filter(
     )
 
 
-def ntree(
-    template: NDArray,
-    sigma_range: Tuple[float, float],
-    **kwargs: dict,
-) -> NDArray:
-    return preprocessor.ntree_filter(template=template, sigma_range=sigma_range)
-
-
 def mean(
     template: NDArray,
     width: int,
@@ -155,9 +147,7 @@ def wedge(
     tilt_step: float = 0,
     opening_axis: int = 0,
     tilt_axis: int = 1,
-    gaussian_sigma: float = 0,
     omit_negative_frequencies: bool = True,
-    extrude_plane: bool = True,
     infinite_plane: bool = True,
 ) -> NDArray:
     template_ft = np.fft.rfftn(template)
@@ -169,9 +159,7 @@ def wedge(
             tilt_axis=tilt_axis,
             opening_axis=opening_axis,
             shape=template.shape,
-            sigma=gaussian_sigma,
             omit_negative_frequencies=omit_negative_frequencies,
-            extrude_plane=extrude_plane,
             infinite_plane=infinite_plane,
         )
         np.multiply(template_ft, wedge_mask, out=template_ft)
@@ -185,7 +173,6 @@ def wedge(
         tilt_step=tilt_step,
         opening_axis=opening_axis,
         shape=template.shape,
-        sigma=gaussian_sigma,
         omit_negative_frequencies=omit_negative_frequencies,
     )
     np.multiply(template_ft, wedge_mask, out=template_ft)
@@ -197,6 +184,10 @@ def compute_power_spectrum(template: NDArray) -> NDArray:
     return np.fft.fftshift(np.log(np.abs(np.fft.fftn(template))))
 
 
+def invert_contrast(template: NDArray) -> NDArray:
+    return template * -1
+
+
 def widgets_from_function(function: Callable, exclude_params: List = ["self"]):
     """
     Creates list of magicui widgets by inspecting function typing ann
@@ -252,13 +243,13 @@ WRAPPED_FUNCTIONS = {
     "gaussian_filter": gaussian_filter,
     "bandpass_filter": bandpass_filter,
     "edge_gaussian_filter": edge_gaussian_filter,
-    "ntree_filter": ntree,
     "local_gaussian_filter": local_gaussian_filter,
     "difference_of_gaussian_filter": difference_of_gaussian_filter,
     "mean_filter": mean,
     "wedge_filter": wedge,
     "power_spectrum": compute_power_spectrum,
     "ctf": ctf_filter,
+    "invert_contrast": invert_contrast,
 }
 
 EXCLUDED_FUNCTIONS = [
@@ -487,10 +478,9 @@ def wedge_mask(
     tilt_step: float = 0,
     opening_axis: int = 0,
     tilt_axis: int = 2,
-    gaussian_sigma: float = 0,
     omit_negative_frequencies: bool = False,
-    extrude_plane: bool = True,
-    infinite_plane: bool = True,
+    infinite_plane: bool = False,
+    weight_angle: bool = False,
     **kwargs,
 ) -> NDArray:
     if tilt_step <= 0:
@@ -500,22 +490,23 @@ def wedge_mask(
             tilt_axis=tilt_axis,
             opening_axis=opening_axis,
             shape=template.shape,
-            sigma=gaussian_sigma,
             omit_negative_frequencies=omit_negative_frequencies,
-            extrude_plane=extrude_plane,
             infinite_plane=infinite_plane,
         )
         wedge_mask = np.fft.fftshift(wedge_mask)
         return wedge_mask
 
+    weights = None
+    tilt_angles = np.arange(-tilt_start, tilt_stop + tilt_step, tilt_step)
+    if weight_angle:
+        weights = np.cos(np.radians(tilt_angles))
+
     wedge_mask = preprocessor.step_wedge_mask(
-        start_tilt=tilt_start,
-        stop_tilt=tilt_stop,
+        tilt_angles=tilt_angles,
         tilt_axis=tilt_axis,
-        tilt_step=tilt_step,
         opening_axis=opening_axis,
         shape=template.shape,
-        sigma=gaussian_sigma,
+        weights=weights,
         omit_negative_frequencies=omit_negative_frequencies,
     )
 
@@ -634,6 +625,7 @@ class MaskWidget(widgets.Container):
 
         data = active_layer.data.copy()
         cutoff = np.quantile(data, self.percentile_range_edit.value / 100)
+        cutoff = max(cutoff, np.finfo(np.float32).resolution)
         data[data < cutoff] = 0
 
         center_of_mass = Density.center_of_mass(np.abs(data), 0)

tests/data/README.md

@@ -0,0 +1,2 @@
+Test data was derived from EMD-8249
+https://www.ebi.ac.uk/pdbe/entry/emdb/EMD-8249