pytme 0.1.1__tar.gz → 0.1.3__tar.gz

{pytme-0.1.1 → pytme-0.1.3}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.1
 Name: pytme
-Version: 0.1.1
+Version: 0.1.3
 Summary: Python Template Matching Engine
 Author: Valentin Maurer
 Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 License: Proprietary
-Project-URL: Homepage, https://git.embl.de/vmaurer/pytme
+Project-URL: Homepage, https://github.com/KosinskiLab/pyTME
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.11

{pytme-0.1.1 → pytme-0.1.3}/pyproject.toml

@@ -7,7 +7,7 @@ name="pytme"
 authors = [
   { name = "Valentin Maurer", email = "valentin.maurer@embl-hamburg.de" },
 ]
-version="0.1.1"
+version="0.1.3"
 description="Python Template Matching Engine"
 readme="README.md"
 requires-python = ">=3.11"
@@ -38,7 +38,7 @@ preprocess = "tme.scripts:preprocess"
 postprocess = "tme.scripts:postprocess"
 
 [project.urls]
-"Homepage" = "https://git.embl.de/vmaurer/pytme"
+"Homepage" = "https://github.com/KosinskiLab/pyTME"
 
 [tool.setuptools]
 include-package-data = true

{pytme-0.1.1 → pytme-0.1.3}/pytme.egg-info/SOURCES.txt

@@ -21,6 +21,7 @@ tme/matching_optimization.py
 tme/matching_utils.py
 tme/parser.py
 tme/preprocessor.py
+tme/scoring.py
 tme/structure.py
 tme/types.py
 tme/backends/__init__.py

{pytme-0.1.1 → pytme-0.1.3}/scripts/match_template.py

@@ -43,7 +43,7 @@ def print_block(name: str, data: dict, label_width=20) -> None:
     """Prints a formatted block of information."""
     print(f"\n> {name}")
     for key, value in data.items():
-        formatted_value = str(value)  # Convert non-string values to string
+        formatted_value = str(value)
         print(f"  - {key + ':':<{label_width}} {formatted_value}")
 
 
@@ -426,7 +426,7 @@ def main():
 
     if not np.allclose(target.sampling_rate, template.sampling_rate):
         print(
-            f"Resampling template to {target.sampling_rate}."
+            f"Resampling template to {target.sampling_rate}. "
             "Consider providing a template with the same sampling rate as the target."
         )
         template = template.resample(target.sampling_rate, order=3)
@@ -506,7 +506,7 @@ def main():
                 -tilt_start, tilt_stop + args.tilt_step, args.tilt_step
             )
             angles = np.zeros((template.data.ndim, tilt_angles.size))
-            angles[1, :] = tilt_angles
+            angles[2, :] = tilt_angles
             template_filter["wedge_mask"] = {
                 "tilt_angles": angles,
                 "sigma": args.wedge_smooth,
@@ -516,6 +516,7 @@ def main():
                 "start_tilt": tilt_start,
                 "stop_tilt": tilt_stop,
                 "tilt_axis": 1,
+                "infinite_plane": True,
                 "sigma": args.wedge_smooth,
             }
 
@@ -597,6 +598,10 @@ def main():
     if not args.pad_fourier:
         template_box = np.ones(len(template_box), dtype=int)
 
+    callback_class = MaxScoreOverRotations
+    if args.peak_calling:
+        callback_class = PeakCallerMaximumFilter
+
     splits, schedule = compute_parallelization_schedule(
         shape1=target.shape,
         shape2=template_box,
@@ -605,7 +610,7 @@ def main():
         max_ram=args.ram,
         split_only_outer=args.use_gpu,
         matching_method=args.score,
-        analyzer_method="MaxScoreOverRotations",
+        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),
@@ -627,16 +632,12 @@ def main():
     }
 
     matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[args.score]
-    callback_class = MaxScoreOverRotations
-    if args.peak_calling:
-        callback_class = PeakCallerMaximumFilter
-
     matching_data = MatchingData(target=target, template=template.data)
     matching_data.rotations = get_rotation_matrices(
         angular_sampling=args.angular_sampling, dim=target.data.ndim
     )
-    matching_data.template_filter = template_filter
 
+    matching_data.template_filter = template_filter
     if target_mask is not None:
         matching_data.target_mask = target_mask
     if template_mask is not None:

{pytme-0.1.1 → pytme-0.1.3}/scripts/postprocess.py

@@ -100,9 +100,12 @@ def main():
                 orientations.append((translation, angles, score, detail))
         else:
             candidates = data
-            for translation, rotation, score, detail in zip(*candidates):
-                angles = euler_from_rotationmatrix(rotation)
-                orientations.append((translation, angles, score, detail))
+            translation, rotation, score, detail, *_ = data
+            for i in range(translation.shape[0]):
+                angles = euler_from_rotationmatrix(rotation[i])
+                orientations.append(
+                    (np.array(translation[i]), angles, score[i], detail[i])
+                )
     else:
         with open(args.orientations, mode="r", encoding="utf-8") as infile:
             data = [x.strip().split("\t") for x in infile.read().split("\n")]

{pytme-0.1.1 → pytme-0.1.3}/scripts/preprocessor_gui.py

@@ -108,6 +108,31 @@ def mean(
     return preprocessor.mean_filter(template=template, width=width)
 
 
+def wedge(
+    template: NDArray,
+    tilt_start: float,
+    tilt_stop: float,
+    gaussian_sigma: float,
+    tilt_axis: int = 1,
+    infinite_plane: bool = True,
+    extrude_plane: bool = True,
+):
+    template_ft = np.fft.rfftn(template)
+    wedge_mask = preprocessor.continuous_wedge_mask(
+        start_tilt=tilt_start,
+        stop_tilt=tilt_stop,
+        tilt_axis=tilt_axis,
+        shape=template.shape,
+        sigma=gaussian_sigma,
+        omit_negative_frequencies=True,
+        infinite_plane=infinite_plane,
+        extrude_plane=extrude_plane,
+    )
+    np.multiply(template_ft, wedge_mask, out=template_ft)
+    template = np.real(np.fft.irfftn(template_ft))
+    return template
+
+
 def widgets_from_function(function: Callable, exclude_params: List = ["self"]):
     """
     Creates list of magicui widgets by inspecting function typing ann
@@ -166,7 +191,8 @@ WRAPPED_FUNCTIONS = {
     "ntree_filter": ntree,
     "local_gaussian_filter": local_gaussian_filter,
     "difference_of_gaussian_filter": difference_of_gaussian_filter,
-    "mean_filter" : mean,
+    "mean_filter": mean,
+    "continuous_wedge_mask": wedge,
 }
 
 EXCLUDED_FUNCTIONS = [
@@ -178,7 +204,7 @@ EXCLUDED_FUNCTIONS = [
     "interpolate_box",
     "molmap",
     "local_gaussian_alignment_filter",
-    "continuous_wedge_mask",
+    # "continuous_wedge_mask",
     "wedge_mask",
     "bandpass_mask",
 ]
@@ -372,19 +398,22 @@ def wedge_mask(
     tilt_stop: float,
     gaussian_sigma: float,
     tilt_axis: int = 1,
+    omit_negative_frequencies: bool = True,
+    extrude_plane: bool = True,
+    infinite_plane: bool = True,
 ):
-    template_ft = np.fft.fftn(template)
     wedge_mask = preprocessor.continuous_wedge_mask(
         start_tilt=tilt_start,
         stop_tilt=tilt_stop,
         tilt_axis=tilt_axis,
-        shape=template_ft.shape,
+        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)
-    template = np.real(np.fft.ifftn(template_ft))
-    return template
+    wedge_mask = np.fft.fftshift(wedge_mask)
+    return wedge_mask
 
 
 class MaskWidget(widgets.Container):
@@ -410,34 +439,73 @@ class MaskWidget(widgets.Container):
         )
         self.method_dropdown.changed.connect(self._on_method_changed)
 
-        self.adapt_button = widgets.PushButton(
-            text="Adapt to current layer", enabled=False
-        )
+        self.adapt_button = widgets.PushButton(text="Adapt to layer", enabled=False)
         self.adapt_button.changed.connect(self._update_initial_values)
 
         self.viewer.layers.selection.events.active.connect(
             self._update_action_button_state
         )
 
+        self.align_button = widgets.PushButton(text="Align to axis", enabled=False)
+        self.align_button.changed.connect(self._align_with_axis)
+        self.density_field = widgets.Label()
+        # self.density_field.value = f"Positive Density in Mask: {0:.2f}%"
+
         self.append(self.method_dropdown)
         self.append(self.adapt_button)
+        self.append(self.align_button)
         self.append(self.action_button)
+        self.append(self.density_field)
 
         # Create GUI for initially selected filtering method
         self._on_method_changed(None)
 
     def _update_action_button_state(self, event):
+        self.align_button.enabled = bool(self.viewer.layers.selection.active)
         self.action_button.enabled = bool(self.viewer.layers.selection.active)
         self.adapt_button.enabled = bool(self.viewer.layers.selection.active)
 
+    def _align_with_axis(self):
+        active_layer = self.viewer.layers.selection.active
+
+        if active_layer.metadata.get("is_aligned", False):
+            return
+
+        coords = np.array(np.where(active_layer.data > 0)).T
+        centered_coords = coords - np.mean(coords, axis=0)
+        cov_matrix = np.cov(centered_coords, rowvar=False)
+
+        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
+        principal_eigenvector = eigenvectors[:, np.argmax(eigenvalues)]
+
+        rotation_axis = np.cross(principal_eigenvector, [1, 0, 0])
+        rotation_angle = np.arccos(np.dot(principal_eigenvector, [1, 0, 0]))
+        k = rotation_axis / np.linalg.norm(rotation_axis)
+        K = np.array([[0, -k[2], k[1]], [k[2], 0, -k[0]], [-k[1], k[0], 0]])
+        rotation_matrix = np.eye(3)
+        rotation_matrix += np.sin(rotation_angle) * K
+        rotation_matrix += (1 - np.cos(rotation_angle)) * np.dot(K, K)
+
+        rotated_data = Density.rotate_array(
+            arr=active_layer.data,
+            rotation_matrix=rotation_matrix,
+            use_geometric_center=False,
+        )
+        eps = np.finfo(rotated_data.dtype).eps
+        rotated_data[rotated_data < eps] = 0
+
+        active_layer.metadata["is_aligned"] = True
+        active_layer.data = rotated_data
+
     def _update_initial_values(self, event=None):
         active_layer = self.viewer.layers.selection.active
-        center_of_mass = Density.center_of_mass(np.abs(active_layer.data))
+        center_of_mass = Density.center_of_mass(np.abs(active_layer.data), 0)
         coordinates = np.array(np.where(active_layer.data > 0))
         coordinates_min = coordinates.min(axis=1)
         coordinates_max = coordinates.max(axis=1)
         coordinates_heights = coordinates_max - coordinates_min
         coordinate_radius = np.divide(coordinates_heights, 2)
+        center_of_mass = coordinate_radius + coordinates_min
 
         defaults = dict(zip(["center_x", "center_y", "center_z"], center_of_mass))
         defaults.update(
@@ -465,7 +533,7 @@ class MaskWidget(widgets.Container):
         widgets = widgets_from_function(function)
         for widget in widgets:
             self.action_widgets.append(widget)
-            self.insert(-2, widget)
+            self.insert(1, widget)
 
     def _action(self):
         function = self.methods.get(self.method_dropdown.value)
@@ -479,7 +547,8 @@ class MaskWidget(widgets.Container):
             selected_layer = self.viewer.layers[new_layer_name]
 
         processed_data = processed_data.astype(np.float32)
-        mask = selected_layer.metadata.get("mask", False)
+        metadata = selected_layer.metadata
+        mask = metadata.get("mask", False)
         if mask == self.method_dropdown.value:
             selected_layer.data = processed_data
         else:
@@ -490,8 +559,18 @@ class MaskWidget(widgets.Container):
             metadata = selected_layer.metadata.copy()
             metadata["filter_parameters"] = {self.method_dropdown.value: kwargs.copy()}
             metadata["mask"] = self.method_dropdown.value
+            metadata["origin_layer"] = selected_layer.name
             new_layer.metadata = metadata
 
+        origin_layer = metadata["origin_layer"]
+        if origin_layer in self.viewer.layers:
+            origin_layer = self.viewer.layers[origin_layer]
+            if np.allclose(origin_layer.data.shape, processed_data.shape):
+                in_mask = np.sum(np.fmax(origin_layer.data * processed_data, 0))
+                in_mask /= np.sum(np.fmax(origin_layer.data, 0))
+                in_mask *= 100
+                self.density_field.value = f"Positive Density in Mask: {in_mask:.2f}%"
+
 
 class ExportWidget(widgets.Container):
     def __init__(self, viewer):

pytme-0.1.3/tme/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.1.3"

{pytme-0.1.1 → pytme-0.1.3}/tme/analyzer.py

@@ -121,7 +121,7 @@ class PeakCaller(ABC):
         fourier_shift = kwargs.get(
             "fourier_shift", backend.zeros(peak_positions.shape[1], dtype=int)
         )
-        if np.any(fourier_shift) != 0:
+        if backend.sum(fourier_shift != 0) != 0:
             peak_positions = backend.mod(
                 backend.add(peak_positions, fourier_shift), score_space.shape
             )
@@ -197,6 +197,7 @@ class PeakCaller(ABC):
             if len(candidate) == 0:
                 continue
             peak_positions, rotations, peak_scores, peak_details = candidate
+            kwargs["translation_offset"] = backend.zeros(peak_positions.shape[1])
             base._update(
                 peak_positions=backend.to_backend_array(peak_positions),
                 peak_details=backend.to_backend_array(peak_details),
@@ -237,7 +238,6 @@ class PeakCaller(ABC):
         translation_offset = backend.astype(translation_offset, peak_positions.dtype)
 
         backend.add(peak_positions, translation_offset, out=peak_positions)
-
         if not len(self.peak_list):
             self.peak_list = [peak_positions, rotations, peak_scores, peak_details]
             return None

{pytme-0.1.1 → pytme-0.1.3}/tme/backends/pytorch_backend.py

@@ -491,8 +491,11 @@ class PytorchBackend(NumpyFFTWBackend):
             Operates as a context manager, yielding None and providing
             the set GPU context for enclosed operations.
         """
-        with self._array_backend.cuda.device(device_index):
-            yield
+        if self.device == "cuda":
+            with self._array_backend.cuda.device(device_index):
+                yield
+        else:
+            yield None
 
     def device_count(self) -> int:
         """
@@ -505,7 +508,7 @@ class PytorchBackend(NumpyFFTWBackend):
         """
         return self._array_backend.cuda.device_count()
 
-    def reverse(arr: TorchTensor) -> TorchTensor:
+    def reverse(self, arr: TorchTensor) -> TorchTensor:
         """
         Reverse the order of elements in a tensor along all its axes.
 
@@ -519,4 +522,4 @@ class PytorchBackend(NumpyFFTWBackend):
         TorchTensor
             Reversed tensor.
         """
-        return self._array_backend.flip(arr, [i for i in range(arr.dim())])
+        return self._array_backend.flip(arr, [i for i in range(arr.ndim)])

{pytme-0.1.1 → pytme-0.1.3}/tme/density.py

@@ -46,10 +46,10 @@ class Density:
     ----------
     data : NDArray
         Electron density data.
-    origin : NDArray
-        Origin of the coordinate system.
-    sampling_rate : NDArray
-        Sampling rate along data axis.
+    origin : NDArray, optional
+        Origin of the coordinate system. Defaults to zero.
+    sampling_rate : NDArray, optional
+        Sampling rate along data axis. Defaults to one.
     metadata : dict, optional
         Dictionary with metadata information, empty by default.
 
@@ -62,16 +62,18 @@ class Density:
     --------
     >>> import numpy as np
     >>> data = np.random.rand(50,50,50)
-    >>> Density(data = data, origin = (0, 0, 0), sampling_rate = (0, 0, 0))
+    >>> Density(data = data, origin = (0, 0, 0), sampling_rate = (1, 1, 1))
     """
 
     def __init__(
         self,
         data: NDArray,
-        origin: NDArray,
-        sampling_rate: NDArray,
+        origin: NDArray = None,
+        sampling_rate: NDArray = None,
         metadata: Dict = {},
     ):
+        origin = 0 if origin is None else origin
+        sampling_rate = 1 if sampling_rate is None else sampling_rate
         sampling_rate, origin = np.asarray(sampling_rate), np.asarray(origin)
         sampling_rate = np.repeat(sampling_rate, data.ndim // sampling_rate.size)
 
@@ -127,7 +129,7 @@ class Density:
 
         Examples
         --------
-        >>> density = Density.from_file("/path/to/mrc")
+        >>> density = Density.from_file("/path/to/file")
 
         Notes
         -----
@@ -548,8 +550,9 @@ class Density:
             Which weight should be given to individual atoms. For valid values
             see :py:meth:`Structure.from_file`.
         chain : str, optional
-            Which chain of the protein should be considered. Default value
-            corresponds to using all chains.
+            The chain identifier. If multiple chains should be selected they need
+            to be a comma separated string, e.g. 'A,B,CE'. If chain None,
+            all chains are returned. Default is None.
         filter_by_elements : set, optional
             Set of atomic elements to keep. Default is all atoms.
         filter_by_residues : set, optional
@@ -876,7 +879,6 @@ class Density:
         box_start = np.array([b.start for b in box])
         box_stop = np.array([b.stop for b in box])
         left_pad = -np.minimum(box_start, np.zeros(len(box), dtype=int))
-        has_extension = box_start < 0
 
         right_pad = box_stop - box_start * (box_start > 0)
         right_pad -= np.array(self.shape, dtype=int)
@@ -905,8 +907,6 @@ class Density:
         self.data = self.data[crop_box].copy()
 
         # In case the box is larger than the current map
-        before_shape = self.data.shape
-        after_shape = [b.stop - b.start for b in box]
         self.data = self._pad_slice(box)
 
         # Adjust the origin
@@ -934,10 +934,19 @@ class Density:
         tuple
             A tuple containing slice objects representing the box.
 
+        Raises
+        ------
+        ValueError
+            If the cutoff is larger than or equal to the maximum density value.
+
         See Also
         --------
         :py:meth:`Density.adjust_box`
         """
+        if cutoff >= self.data.max():
+            raise ValueError(
+                f"Cutoff exceeds data range ({cutoff} >= {self.data.max()})."
+            )
         starts, stops = [], []
         for axis in range(self.data.ndim):
             projected_max = np.max(

{pytme-0.1.1 → pytme-0.1.3}/tme/matching_data.py

@@ -116,7 +116,6 @@ class MatchingData:
         ).astype(int)
 
         ret_shape = np.add(slice_shape, padding)
-
         arr_start = np.subtract(slice_start, data_voxels_left)
         arr_stop = np.add(slice_stop, data_voxels_right)
         arr_slice = tuple(slice(*pos) for pos in zip(arr_start, arr_stop))
@@ -138,7 +137,6 @@ class MatchingData:
                     arr.filename, mode="r", shape=arr.shape, dtype=arr.dtype
                 )
             arr = np.asarray(arr[*arr_mesh])
-
         ret = np.full(
             shape=np.add(slice_shape, padding), fill_value=arr.mean(), dtype=arr.dtype
         )
@@ -188,8 +186,8 @@ class MatchingData:
             template_pad = np.zeros(len(self.target.shape), dtype=int)
 
         indices = compute_full_convolution_index(
-            outer_shape=self.target.shape,
-            inner_shape=self.template.shape,
+            outer_shape=self._target.shape,
+            inner_shape=self._template.shape,
             outer_split=target_slice,
             inner_split=template_slice,
         )

{pytme-0.1.1 → pytme-0.1.3}/tme/matching_exhaustive.py

@@ -1136,7 +1136,6 @@ def mcc_scoring(
             mask_overlap, axis=axes, keepdims=True
         )
         temp[mask_overlap < number_px_threshold] = 0.0
-
         convolution_mode = kwargs.get("convolution_mode", "full")
         score = apply_convolution_mode(
             temp, convolution_mode=convolution_mode, s1=targetshape, s2=templateshape
@@ -1165,7 +1164,6 @@ def device_memory_handler(func: Callable):
                     return_value = func(shared_memory_handler=smh, *args, **kwargs)
         except Exception as e:
             print(e)
-            return None
             last_type, last_value, last_traceback = sys.exc_info()
         finally:
             handle_traceback(last_type, last_value, last_traceback)

{pytme-0.1.1 → pytme-0.1.3}/tme/matching_memory.py

@@ -327,7 +327,7 @@ def estimate_ram_usage(
     """
     if matching_method not in MATCHING_MEMORY_REGISTRY:
         raise ValueError(
-            f"Supported fitters are {','.join(MATCHING_MEMORY_REGISTRY.keys())}"
+            f"Supported options are {','.join(MATCHING_MEMORY_REGISTRY.keys())}"
         )
 
     convolution_shape, fast_shape, ft_shape = _compute_convolution_shapes(

{pytme-0.1.1 → pytme-0.1.3}/tme/matching_optimization.py

@@ -246,6 +246,7 @@ class CrossCorrelation(MatchCoordinatesToDensity):
     .. math::
 
         \\text{score} = \\text{target_weights} \\cdot \\text{template_weights}
+
     """
 
     def __init__(self, **kwargs):
@@ -298,6 +299,7 @@ class LaplaceCrossCorrelation(CrossCorrelation):
 
         \\text{score} = \\nabla^{2} \\text{target_weights} \\cdot
                         \\nabla^{2} \\text{template_weights}
+
     """
 
     def __init__(self, **kwargs):
@@ -617,6 +619,7 @@ class Chamfer(MatchCoordinatesToCoordinates):
         -------
         float
             The negative of the Chamfer distance score.
+
         """
         dist, _ = self.target_tree.query(self.template_coordinates_rotated.T)
         score = np.mean(dist)
@@ -638,6 +641,7 @@ class MutualInformation(MatchCoordinatesToDensity):
     .. [1]  Daven Vasishtan and Maya Topf, "Scoring functions for cryoEM density
             fitting", Journal of Structural Biology, vol. 174, no. 2,
             pp. 333--343, 2011. DOI: https://doi.org/10.1016/j.jsb.2011.01.012
+
     """
 
     def __init__(self, **kwargs):
@@ -776,6 +780,7 @@ class NormalVectorScore(MatchCoordinatesToCoordinates):
     .. [1]  Daven Vasishtan and Maya Topf, "Scoring functions for cryoEM density
             fitting", Journal of Structural Biology, vol. 174, no. 2,
             pp. 333--343, 2011. DOI: https://doi.org/10.1016/j.jsb.2011.01.012
+
     """
 
     def __init__(self, **kwargs):

{pytme-0.1.1 → pytme-0.1.3}/tme/matching_utils.py

@@ -762,6 +762,8 @@ def euler_to_rotationmatrix(angles: Tuple[float]) -> NDArray:
     NDArray
         The generated rotation matrix.
     """
+    if len(angles) == 1:
+        angles = (angles, 0, 0)
     rotation_matrix = (
         Rotation.from_euler("zyx", angles, degrees=True).as_matrix().astype(np.float32)
     )
@@ -775,13 +777,17 @@ def euler_from_rotationmatrix(rotation_matrix: NDArray) -> Tuple:
     Parameters
     ----------
     rotation_matrix : NDArray
-        A 3 x 3 rotation matrix in z y x form.
+        A 2 x 2 or 3 x 3 rotation matrix in z y x form.
 
     Returns
     -------
     Tuple
         The generate euler angles in degrees
     """
+    if rotation_matrix.shape[0] == 2:
+        temp_matrix = np.eye(3)
+        temp_matrix[:2, :2] = rotation_matrix
+        rotation_matrix = temp_matrix
     euler_angles = (
         Rotation.from_matrix(rotation_matrix)
         .as_euler("zyx", degrees=True)