pytme 0.2.2__cp311-cp311-macosx_14_0_arm64.whl → 0.2.3__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:
@@ -467,29 +468,35 @@ 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)
+            fourier_shift = np.subtract(fourier_shift, template_mod)
 
-        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:
+        shape_diff = np.multiply(
+            np.subtract(target_shape, template_shape), 1 - batch_mask
+        )
+        if np.sum(shape_diff < 0):
             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."
+                "Template is larger than target and padding is turned off. Consider "
+                "swapping them or activate padding. Correcting the shift for now."
             )
-            fourier_shift = np.subtract(fourier_shift, np.divide(shape_diff, 2))
-            fourier_shift = fourier_shift.astype(int)
 
-        return tuple(fast_shape), tuple(fast_ft_shape), tuple(fourier_shift)
+            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),
+                )
+
+            shape_shift = np.add(shape_shift, offset)
+            fourier_shift = np.subtract(fourier_shift, shape_shift).astype(int)
+
+        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]:
         """

tme/matching_exhaustive.py

@@ -73,35 +73,47 @@ 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,
+        )
+        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 +224,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 +272,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,

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'.")
@@ -1220,7 +1225,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/composable_filter.py

@@ -17,15 +17,18 @@ 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/tilt_series.py

@@ -673,6 +673,9 @@ class WedgeReconstructed:
             omit_negative_frequencies=return_real_fourier,
         )
 
+        if not weight_wedge:
+            ret = (ret > 0) * 1.0
+
         return ret
 
 
@@ -790,16 +793,6 @@ class CTF:
     def __post_init__(self):
         self.defocus_angle = np.radians(self.defocus_angle)
 
-        kwargs = {
-            "defocus_x": self.defocus_x,
-            "defocus_y": self.defocus_y,
-            "spherical_aberration": self.spherical_aberration,
-        }
-        kwargs = {k: v for k, v in kwargs.items() if v is not None}
-        self._update_parameters(
-            electron_wavelength=self._compute_electron_wavelength(), **kwargs
-        )
-
     def _compute_electron_wavelength(self, acceleration_voltage: int = None):
         """Computes the wavelength of an electron in angstrom."""
 
@@ -818,28 +811,10 @@ class CTF:
         electron_wavelength = np.divide(
             planck_constant * light_velocity, np.sqrt(denominator)
         )
+        # Convert to Ångstrom
         electron_wavelength *= 1e10
         return electron_wavelength
 
-    def _update_parameters(self, **kwargs):
-        """Update multiple parameters of the CTF instance."""
-        voxel_based = [
-            "electron_wavelength",
-            "spherical_aberration",
-            "defocus_x",
-            "defocus_y",
-        ]
-        if "sampling_rate" in kwargs:
-            self.sampling_rate = kwargs["sampling_rate"]
-
-        if "acceleration_voltage" in kwargs:
-            kwargs["electron_wavelength"] = self._compute_electron_wavelength()
-
-        for key, value in kwargs.items():
-            if key in voxel_based and value is not None:
-                value = np.divide(value, np.max(self.sampling_rate))
-            setattr(self, key, value)
-
     def __call__(self, **kwargs) -> NDArray:
         func_args = vars(self).copy()
         func_args.update(kwargs)
@@ -865,7 +840,6 @@ class CTF:
         shape: Tuple[int],
         defocus_x: Tuple[float],
         angles: Tuple[float],
-        electron_wavelength: float = None,
         opening_axis: int = None,
         tilt_axis: int = None,
         amplitude_contrast: float = 0.07,
@@ -891,8 +865,6 @@ class CTF:
             The defocus value in x direction.
         angles : tuple of float
             The tilt angles.
-        electron_wavelength : float, optional
-            The electron wavelength, defaults to None.
         opening_axis : int, optional
             The axis around which the wedge is opened, defaults to None.
         tilt_axis : int, optional
@@ -939,9 +911,13 @@ class CTF:
         correct_defocus_gradient &= tilt_axis is not None
         correct_defocus_gradient &= opening_axis is not None
 
+        spherical_aberration /= sampling_rate
+        electron_wavelength = self._compute_electron_wavelength() / sampling_rate
         for index, angle in enumerate(angles):
             defocus_x, defocus_y = defoci_x[index], defoci_y[index]
 
+            defocus_x = defocus_x / sampling_rate if defocus_x is not None else None
+            defocus_y = defocus_y / sampling_rate if defocus_y is not None else None
             if correct_defocus_gradient or defocus_y is not None:
                 grid = fftfreqn(
                     shape=shape,
@@ -978,6 +954,7 @@ class CTF:
                 angle=angle,
                 sampling_rate=1,
             )
+            frequency_mask = frequency_grid < 0.5
             np.square(frequency_grid, out=frequency_grid)
 
             electron_aberration = spherical_aberration * electron_wavelength**2
@@ -993,6 +970,7 @@ class CTF:
                 )
             )
             np.sin(-chi, out=chi)
+            np.multiply(chi, frequency_mask, out=chi)
             stack[index] = chi
 
         # Avoid contrast inversion