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

tme/matching_data.py

@@ -171,6 +171,7 @@ class MatchingData:
                 np.subtract(right_pad, data_voxels_right),
             )
         )
+        # The reflections are later cropped from the scores
         arr = np.pad(arr, padding, mode="reflect")
 
         if invert:
@@ -449,7 +450,7 @@ class MatchingData:
         template_shape: NDArray,
         batch_mask: NDArray = None,
         pad_fourier: bool = False,
-    ) -> Tuple[Tuple, Tuple, Tuple]:
+    ) -> Tuple[Tuple, Tuple, Tuple, Tuple]:
         """
         Determines an efficient shape for Fourier transforms considering zero-padding.
         """
@@ -467,31 +468,39 @@ class MatchingData:
 
         pad_shape = np.maximum(target_shape, template_shape)
         ret = be.compute_convolution_shapes(pad_shape, fourier_pad)
-        convolution_shape, fast_shape, fast_ft_shape = ret
+        conv_shape, fast_shape, fast_ft_shape = ret
+
+        template_mod = np.mod(template_shape, 2)
         if not pad_fourier:
             fourier_shift = 1 - np.divide(template_shape, 2).astype(int)
-            fourier_shift -= np.mod(template_shape, 2)
-            shape_diff = np.subtract(fast_shape, convolution_shape)
-            shape_diff = np.divide(shape_diff, 2).astype(int)
-            shape_diff = np.multiply(shape_diff, 1 - batch_mask)
-            np.add(fourier_shift, shape_diff, out=fourier_shift)
-
-        fourier_shift = fourier_shift.astype(int)
-
-        shape_diff = np.subtract(target_shape, template_shape)
-        shape_diff = np.multiply(shape_diff, 1 - batch_mask)
-        if np.sum(shape_diff < 0) and not pad_fourier:
-            warnings.warn(
-                "Target is larger than template and Fourier padding is turned off. "
-                "This may lead to inaccurate results. Prefer swapping template and target, "
-                "enable padding or turn off template centering."
-            )
-            fourier_shift = np.subtract(fourier_shift, np.divide(shape_diff, 2))
-            fourier_shift = fourier_shift.astype(int)
+            fourier_shift = np.subtract(fourier_shift, template_mod)
+
+        shape_diff = np.multiply(
+            np.subtract(target_shape, template_shape), 1 - batch_mask
+        )
+        if np.sum(shape_diff < 0):
+            shape_shift = np.divide(shape_diff, 2)
+            offset = np.mod(shape_diff, 2)
+            if pad_fourier:
+                offset = -np.subtract(
+                    offset,
+                    np.logical_and(np.mod(target_shape, 2) == 0, template_mod == 1),
+                )
+            else:
+                warnings.warn(
+                    "Template is larger than target and padding is turned off. Consider "
+                    "swapping them or activate padding. Correcting the shift for now."
+                )
+
+            shape_shift = np.add(shape_shift, offset)
+            fourier_shift = np.subtract(fourier_shift, shape_shift).astype(int)
 
-        return tuple(fast_shape), tuple(fast_ft_shape), tuple(fourier_shift)
+        fourier_shift = tuple(fourier_shift.astype(int))
+        return tuple(conv_shape), tuple(fast_shape), tuple(fast_ft_shape), fourier_shift
 
-    def fourier_padding(self, pad_fourier: bool = False) -> Tuple[Tuple, Tuple, Tuple]:
+    def fourier_padding(
+        self, pad_fourier: bool = False
+    ) -> Tuple[Tuple, Tuple, Tuple, Tuple]:
         """
         Computes efficient shape four Fourier transforms and potential associated shifts.
 
@@ -503,8 +512,8 @@ class MatchingData:
 
         Returns
         -------
-        Tuple[tuple of int, tuple of int, tuple of int]
-            Tuple with real and complex Fourier transform shape, and corresponding shift.
+        Tuple[tuple of int, tuple of int, tuple of int, tuple of int]
+            Tuple with convolution, forward FT, inverse FT shape and corresponding shift.
         """
         return self._fourier_padding(
             target_shape=be.to_numpy_array(self._output_target_shape),

tme/matching_exhaustive.py

@@ -73,35 +73,48 @@ def _setup_template_filter_apply_target_filter(
     if not filter_template and not filter_target:
         return template_filter
 
-    target_temp = be.topleft_pad(matching_data.target, fast_shape)
-    target_temp_ft = be.zeros(fast_ft_shape, be._complex_dtype)
-
     inv_mask = be.subtract(1, be.to_backend_array(matching_data._batch_mask))
     filter_shape = be.multiply(be.to_backend_array(fast_ft_shape), inv_mask)
     filter_shape = tuple(int(x) if x != 0 else 1 for x in filter_shape)
-
     fast_shape = be.multiply(be.to_backend_array(fast_shape), inv_mask)
     fast_shape = tuple(int(x) for x in fast_shape if x != 0)
 
+    fastt_shape, fastt_ft_shape = fast_shape, filter_shape
+    if filter_template and not pad_template_filter:
+        # FFT shape acrobatics for faster filter application
+        # _, fastt_shape, _, _ = matching_data._fourier_padding(
+        #     target_shape=be.to_numpy_array(matching_data._template.shape),
+        #     template_shape=be.to_numpy_array(
+        #         [1 for _ in matching_data._template.shape]
+        #     ),
+        #     pad_fourier=False,
+        # )
+        fastt_shape = matching_data._template.shape
+        matching_data.template = be.reverse(
+            be.topleft_pad(matching_data.template, fastt_shape)
+        )
+        matching_data.template_mask = be.reverse(
+            be.topleft_pad(matching_data.template_mask, fastt_shape)
+        )
+        matching_data._set_matching_dimension(
+            target_dims=matching_data._target_dims,
+            template_dims=matching_data._template_dims,
+        )
+        fastt_ft_shape = [int(x) for x in matching_data._output_template_shape]
+        fastt_ft_shape[-1] = fastt_ft_shape[-1] // 2 + 1
+
+    target_temp = be.topleft_pad(matching_data.target, fast_shape)
+    target_temp_ft = be.zeros(fast_ft_shape, be._complex_dtype)
     target_temp_ft = rfftn(target_temp, target_temp_ft)
     if filter_template:
-        # TODO: Pad to fast shapes and adapt _setup_template_filtering accordingly
-        template_fast_shape, template_filter_shape = fast_shape, filter_shape
-        if not pad_template_filter:
-            template_fast_shape = tuple(int(x) for x in matching_data._template.shape)
-            template_filter_shape = [
-                int(x) for x in matching_data._output_template_shape
-            ]
-            template_filter_shape[-1] = template_filter_shape[-1] // 2 + 1
-
         template_filter = matching_data.template_filter(
-            shape=template_fast_shape,
+            shape=fastt_shape,
             return_real_fourier=True,
             shape_is_real_fourier=False,
             data_rfft=target_temp_ft,
             batch_dimension=matching_data._target_dims,
         )["data"]
-        template_filter = be.reshape(template_filter, template_filter_shape)
+        template_filter = be.reshape(template_filter, fastt_ft_shape)
 
     if filter_target:
         target_filter = matching_data.target_filter(
@@ -212,9 +225,13 @@ def scan(
 
     """
     matching_data.to_backend()
-    fast_shape, fast_ft_shape, fourier_shift = matching_data.fourier_padding(
-        pad_fourier=pad_fourier
-    )
+    (
+        conv_shape,
+        fast_shape,
+        fast_ft_shape,
+        fourier_shift,
+    ) = matching_data.fourier_padding(pad_fourier=pad_fourier)
+    template_shape = matching_data.template.shape
 
     rfftn, irfftn = be.build_fft(
         fast_shape=fast_shape,
@@ -256,7 +273,8 @@ def scan(
         "fourier_shift": fourier_shift,
         "convolution_mode": convmode,
         "targetshape": matching_data.target.shape,
-        "templateshape": matching_data.template.shape,
+        "templateshape": template_shape,
+        "convolution_shape": conv_shape,
         "fast_shape": fast_shape,
         "indices": getattr(matching_data, "indices", None),
         "shared_memory_handler": shared_memory_handler,
@@ -382,7 +400,8 @@ def scan_subsets(
     The template matching procedure is determined by ``matching_setup`` and
     ``matching_score``, which are unique to each score. In the following,
     we will be using the `FLCSphericalMask` score, which is composed of
-    :py:meth:`flcSphericalMask_setup` and :py:meth:`corr_scoring`
+    :py:meth:`tme.matching_scores.flcSphericalMask_setup` and
+    :py:meth:`tme.matching_scores.corr_scoring`
 
     >>> from tme.matching_exhaustive import MATCHING_EXHAUSTIVE_REGISTER
     >>> funcs = MATCHING_EXHAUSTIVE_REGISTER.get("FLCSphericalMask")

tme/matching_scores.py

@@ -86,7 +86,7 @@ def _setup_template_filtering(
         forward_ft_shape = template_shape
         inverse_ft_shape = template_filter.shape
 
-    if rfftn is not None and irfftn is not None:
+    if (rfftn is not None and irfftn is not None) or shape_mismatch:
         rfftn, irfftn = be.build_fft(
             fast_shape=forward_ft_shape,
             fast_ft_shape=inverse_ft_shape,
@@ -109,7 +109,10 @@ def _setup_template_filtering(
 
         def _apply_filter_shape_mismatch(template, ft_temp, template_filter):
             _template[:] = template[real_subset]
-            return _apply_template_filter(_template, _ft_temp, template_filter)
+            template[real_subset] = _apply_template_filter(
+                _template, _ft_temp, template_filter
+            )
+            return template
 
         return _apply_filter_shape_mismatch
 

tme/matching_utils.py

@@ -467,6 +467,7 @@ def apply_convolution_mode(
     convolution_mode: str,
     s1: Tuple[int],
     s2: Tuple[int],
+    convolution_shape: Tuple[int] = None,
     mask_output: bool = False,
 ) -> BackendArray:
     """
@@ -490,6 +491,8 @@ def apply_convolution_mode(
         Tuple of integers corresponding to shape of convolution array 1.
     s2 : tuple of ints
         Tuple of integers corresponding to shape of convolution array 2.
+    convolution_shape : tuple of ints, optional
+        Size of the actually computed convolution. s1 + s2 - 1 by default.
     mask_output : bool, optional
         Whether to mask values outside of convolution_mode rather than
         removing them. Defaults to False.
@@ -500,7 +503,9 @@ def apply_convolution_mode(
         The array after applying the convolution mode.
     """
     # Remove padding to next fast Fourier length
-    arr = arr[tuple(slice(s1[i] + s2[i] - 1) for i in range(len(s1)))]
+    if convolution_shape is None:
+        convolution_shape = [s1[i] + s2[i] - 1 for i in range(len(s1))]
+    arr = arr[tuple(slice(0, x) for x in convolution_shape)]
 
     if convolution_mode not in ("full", "same", "valid"):
         raise ValueError("Supported convolution_mode are 'full', 'same' and 'valid'.")
@@ -640,6 +645,7 @@ def get_rotation_matrices(
         dets = np.linalg.det(ret)
         neg_dets = dets < 0
         ret[neg_dets, :, -1] *= -1
+        ret[0] = np.eye(dim, dtype = ret.dtype)
     return ret
 
 
@@ -1220,7 +1226,7 @@ def scramble_phases(
     arr: NDArray,
     noise_proportion: float = 0.5,
     seed: int = 42,
-    normalize_power: bool = True,
+    normalize_power: bool = False,
 ) -> NDArray:
     """
     Perform random phase scrambling of ``arr``.

tme/orientations.py

@@ -6,6 +6,7 @@
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 import re
+import warnings
 from collections import deque
 from dataclasses import dataclass
 from string import ascii_lowercase
@@ -301,7 +302,7 @@ class Orientations:
     def _to_relion_star(
         self,
         filename: str,
-        name_prefix: str = None,
+        name: str = None,
         ctf_image: str = None,
         sampling_rate: float = 1.0,
         subtomogram_size: int = 0,
@@ -313,8 +314,9 @@ class Orientations:
         ----------
         filename : str
             The name of the file to save the orientations.
-        name_prefix : str, optional
-            A prefix to add to the image names in the STAR file.
+        name : str or list of str, optional
+            Path to image file the orientation is in reference to. If name is a list
+            its assumed to correspond to _rlnImageName, otherwise _rlnMicrographName.
         ctf_image : str, optional
             Path to CTF or wedge mask RELION.
         sampling_rate : float, optional
@@ -352,6 +354,21 @@ class Orientations:
         optics_header = "\n".join(optics_header)
         optics_data = "\t".join(optics_data)
 
+        if name is None:
+            name = ""
+            warnings.warn(
+                "Consider specifying the name argument. A single string will be "
+                "interpreted as path to the original micrograph, a list of strings "
+                "as path to individual subsets."
+            )
+
+        name_reference = "_rlnImageName"
+        if isinstance(name, str):
+            name = [
+                name,
+            ] * self.translations.shape[0]
+            name_reference = "_rlnMicrographName"
+
         header = [
             "data_particles",
             "",
@@ -359,7 +376,7 @@ class Orientations:
             "_rlnCoordinateX",
             "_rlnCoordinateY",
             "_rlnCoordinateZ",
-            "_rlnImageName",
+            name_reference,
             "_rlnAngleRot",
             "_rlnAngleTilt",
             "_rlnAnglePsi",
@@ -371,8 +388,6 @@ class Orientations:
         ctf_image = "" if ctf_image is None else f"\t{ctf_image}"
 
         header = "\n".join(header)
-        name_prefix = "" if name_prefix is None else name_prefix
-
         with open(filename, mode="w", encoding="utf-8") as ofile:
             _ = ofile.write(f"{optics_header}\n")
             _ = ofile.write(f"{optics_data}\n")
@@ -387,9 +402,8 @@ class Orientations:
 
                 translation_string = "\t".join([str(x) for x in translation][::-1])
                 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"
+                    f"{translation_string}\t{name[index]}\t{angle_string}\t1{ctf_image}\n"
                 )
 
         return None
@@ -584,7 +598,10 @@ class Orientations:
         cls, filename: str, delimiter: str = None
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         ret = cls._parse_star(filename=filename, delimiter=delimiter)
-        ret = ret["data_particles"]
+
+        ret = ret.get("data_particles", None)
+        if ret is None:
+            raise ValueError(f"No data_particles section found in {filename}.")
 
         translation = np.vstack(
             (ret["_rlnCoordinateZ"], ret["_rlnCoordinateY"], ret["_rlnCoordinateX"])

tme/preprocessing/_utils.py

@@ -19,8 +19,8 @@ def compute_tilt_shape(shape: Tuple[int], opening_axis: int, reduce_dim: bool =
     """
     Given an opening_axis, computes the shape of the remaining dimensions.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     shape : Tuple[int]
         The shape of the input array.
     opening_axis : int
@@ -28,8 +28,8 @@ def compute_tilt_shape(shape: Tuple[int], opening_axis: int, reduce_dim: bool =
     reduce_dim : bool, optional (default=False)
         Whether to reduce the dimensionality after tilting.
 
-    Returns:
-    --------
+    Returns
+    -------
     Tuple[int]
         The shape of the array after tilting.
     """
@@ -44,13 +44,13 @@ def centered_grid(shape: Tuple[int]) -> NDArray:
     """
     Generate an integer valued grid centered around size // 2
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     shape : Tuple[int]
         The shape of the grid.
 
-    Returns:
-    --------
+    Returns
+    -------
     NDArray
         The centered grid.
     """
@@ -70,8 +70,8 @@ def frequency_grid_at_angle(
     """
     Generate a frequency grid from 0 to 1/(2 * sampling_rate) in each axis.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     shape : Tuple[int]
         The shape of the grid.
     angle : float
@@ -128,8 +128,8 @@ def fftfreqn(
     """
     Generate the n-dimensional discrete Fourier transform sample frequencies.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     shape : Tuple[int]
         The shape of the data.
     sampling_rate : float or Tuple[float]
@@ -180,8 +180,8 @@ def crop_real_fourier(data: BackendArray) -> BackendArray:
     """
     Crop the real part of a Fourier transform.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     data : BackendArray
         The Fourier transformed data.
 

tme/preprocessing/composable_filter.py

@@ -17,15 +17,16 @@ class ComposableFilter(ABC):
     @abstractmethod
     def __call__(self, *args, **kwargs) -> Dict:
         """
-        Parameters:
-        -----------
+
+        Parameters
+        ----------
         *args : tuple
             Variable length argument list.
         **kwargs : dict
             Arbitrary keyword arguments.
 
-        Returns:
-        --------
+        Returns
+        -------
         Dict
             A dictionary representing the result of the filtering operation.
         """

tme/preprocessing/compose.py

@@ -17,13 +17,13 @@ class Compose:
     This class allows composing multiple transformations together. Each transformation
     is expected to be a callable that accepts keyword arguments and returns metadata.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     transforms : Tuple[object]
         A tuple containing transformation objects.
 
-    Returns:
-    --------
+    Returns
+    -------
     Dict
         Metadata resulting from the composed transformations.
 

tme/preprocessing/frequency_filters.py

@@ -18,12 +18,10 @@ from ._utils import fftfreqn, crop_real_fourier, shift_fourier, compute_fourier_
 
 class BandPassFilter:
     """
-    This class provides methods to generate bandpass filters in Fourier space,
-    either by directly specifying the frequency cutoffs (discrete_bandpass) or
-    by using Gaussian functions (gaussian_bandpass).
+    Generate bandpass filters in Fourier space.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     lowpass : float, optional
         The lowpass cutoff, defaults to None.
     highpass : float, optional
@@ -67,8 +65,8 @@ class BandPassFilter:
         """
         Generate a bandpass filter using discrete frequency cutoffs.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         shape : tuple of int
             The shape of the bandpass filter.
         lowpass : float
@@ -84,8 +82,8 @@ class BandPassFilter:
         **kwargs : dict
             Additional keyword arguments.
 
-        Returns:
-        --------
+        Returns
+        -------
         BackendArray
             The bandpass filter in Fourier space.
         """
@@ -98,17 +96,18 @@ class BandPassFilter:
             shape_is_real_fourier=shape_is_real_fourier,
             compute_euclidean_norm=True,
         )
-
-        lowpass = 0 if lowpass is None else lowpass
-        highpass = 1e10 if highpass is None else highpass
+        grid = be.to_backend_array(grid)
+        sampling_rate = be.to_backend_array(sampling_rate)
 
         highcut = grid.max()
-        if lowpass > 0:
-            highcut = np.max(2 * sampling_rate / lowpass)
-        lowcut = np.max(2 * sampling_rate / highpass)
+        if lowpass is not None:
+            highcut = be.max(2 * sampling_rate / lowpass)
 
-        bandpass_filter = ((grid <= highcut) & (grid >= lowcut)) * 1.0
+        lowcut = 0
+        if highpass is not None:
+            lowcut = be.max(2 * sampling_rate / highpass)
 
+        bandpass_filter = ((grid <= highcut) & (grid >= lowcut)) * 1.0
         bandpass_filter = shift_fourier(
             data=bandpass_filter, shape_is_real_fourier=shape_is_real_fourier
         )
@@ -129,10 +128,10 @@ class BandPassFilter:
         **kwargs,
     ) -> BackendArray:
         """
-        Generate a bandpass filter using Gaussian functions.
+        Generate a bandpass filter using Gaussians.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         shape : tuple of int
             The shape of the bandpass filter.
         lowpass : float
@@ -148,8 +147,8 @@ class BandPassFilter:
         **kwargs : dict
             Additional keyword arguments.
 
-        Returns:
-        --------
+        Returns
+        -------
         BackendArray
             The bandpass filter in Fourier space.
         """
@@ -216,15 +215,13 @@ class BandPassFilter:
 
 class LinearWhiteningFilter:
     """
-    This class provides methods to compute the spectrum of the input data and
-    apply linear whitening to the Fourier coefficients.
+    Compute Fourier power spectrums and perform whitening.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     **kwargs : Dict, optional
         Additional keyword arguments.
 
-
     References
     ----------
     .. [1] de Teresa-Trueba, I.; Goetz, S. K.; Mattausch, A.; Stojanovska, F.; Zimmerli, C. E.;
@@ -243,10 +240,10 @@ class LinearWhiteningFilter:
         data_rfft: BackendArray, n_bins: int = None, batch_dimension: int = None
     ) -> Tuple[BackendArray, BackendArray]:
         """
-        Compute the spectrum of the input data.
+        Compute the power spectrum of the input data.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         data_rfft : BackendArray
             The Fourier transform of the input data.
         n_bins : int, optional
@@ -254,8 +251,8 @@ class LinearWhiteningFilter:
         batch_dimension : int, optional
             Batch dimension to average over.
 
-        Returns:
-        --------
+        Returns
+        -------
         bins : BackendArray
             Array containing the bin indices for the spectrum.
         radial_averages : BackendArray
@@ -330,8 +327,8 @@ class LinearWhiteningFilter:
         """
         Apply linear whitening to the data and return the result.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         data : BackendArray, optional
             The input data, defaults to None.
         data_rfft : BackendArray, optional
@@ -345,8 +342,8 @@ class LinearWhiteningFilter:
         **kwargs : Dict
             Additional keyword arguments.
 
-        Returns:
-        --------
+        Returns
+        -------
         Dict
            Filter data and associated parameters.
         """