pytme 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1.post1__cp311-cp311-macosx_15_0_arm64.whl

tme/backends/mlx_backend.py

@@ -6,7 +6,7 @@ Copyright (c) 2024 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from typing import Tuple, List, Callable
+from typing import Tuple, List
 
 import numpy as np
 
@@ -144,32 +144,6 @@ class MLXBackend(NumpyFFTWBackend):
         box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
         return arr[box]
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int] = None,
-        inv_output_shape: Tuple[int] = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        **kwargs,
-    ) -> Tuple[Callable, Callable]:
-        # Runs on mlx.core.cpu until Metal support is available
-        rfft_shape = self._format_fft_shape(fwd_shape, fwd_axes)
-        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
-        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
-
-        def rfftn(arr: MlxArray, out: MlxArray = None, s=rfft_shape, axes=fwd_axes):
-            out[:] = self._array_backend.fft.rfftn(
-                arr, s=s, axes=axes, stream=self._array_backend.cpu
-            )
-
-        def irfftn(arr: MlxArray, out: MlxArray = None, s=irfft_shape, axes=inv_axes):
-            out[:] = self._array_backend.fft.irfftn(
-                arr, s=s, axes=axes, stream=self._array_backend.cpu
-            )
-
-        return rfftn, irfftn
-
     def rfftn(self, arr, *args, **kwargs):
         return self.fft.rfftn(arr, stream=self._array_backend.cpu, **kwargs)
 

tme/backends/npfftw_backend.py

@@ -9,17 +9,23 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 import os
 from psutil import virtual_memory
 from contextlib import contextmanager
-from typing import Tuple, Dict, List, Type
+from typing import Tuple, List, Type
 
 import scipy
 import numpy as np
 from scipy.ndimage import maximum_filter, affine_transform
-from pyfftw.builders import rfftn as rfftn_builder, irfftn as irfftn_builder
-from pyfftw import zeros_aligned, simd_alignment, FFTW, next_fast_len, interfaces
+from pyfftw import (
+    zeros_aligned,
+    simd_alignment,
+    next_fast_len,
+    interfaces,
+    config,
+)
 
 from ..types import NDArray, BackendArray, shm_type
 from .matching_backend import MatchingBackend, _create_metafunction
 
+
 os.environ["MKL_NUM_THREADS"] = "1"
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["PYFFTW_NUM_THREADS"] = "1"
@@ -103,6 +109,20 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         self.solve_triangular = self._solve_triangular
         self.linalg.solve_triangular = scipy.linalg.solve_triangular
 
+        try:
+            from ._numpyfftw_utils import rfftn as rfftn_cache
+            from ._numpyfftw_utils import irfftn as irfftn_cache
+
+            self._rfftn = rfftn_cache
+            self._irfftn = irfftn_cache
+        except Exception as e:
+            print(e)
+
+        config.NUM_THREADS = 1
+        config.PLANNER_EFFORT = "FFTW_MEASURE"
+        interfaces.cache.enable()
+        interfaces.cache.set_keepalive_time(360)
+
     def _linalg_cholesky(self, arr, lower=False, *args, **kwargs):
         # Upper argument is not supported until numpy 2.0
         ret = self._array_backend.linalg.cholesky(arr, *args, **kwargs)
@@ -138,7 +158,7 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         return float
 
     def free_cache(self):
-        pass
+        interfaces.cache.disable()
 
     def transpose(self, arr: NDArray, *args, **kwargs) -> NDArray:
         return self._array_backend.transpose(arr, *args, **kwargs)
@@ -181,6 +201,9 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         sorted_indices = self.unravel_index(indices=sorted_indices, shape=arr.shape)
         return sorted_indices
 
+    def ssum(self, arr, *args, **kwargs):
+        return self.sum(self.square(arr), *args, **kwargs)
+
     def indices(self, *args, **kwargs) -> NDArray:
         return self._array_backend.indices(*args, **kwargs)
 
@@ -240,70 +263,53 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         b[tuple(bind)] = arr[tuple(aind)]
         return b
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int],
-        real_dtype: type,
-        cmpl_dtype: type,
-        fftargs: Dict = {},
-        inv_output_shape: Tuple[int] = None,
-        temp_fwd: NDArray = None,
-        temp_inv: NDArray = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-    ) -> Tuple[FFTW, FFTW]:
-        if temp_fwd is None:
-            temp_fwd = (
-                self.zeros(fwd_shape, real_dtype) if temp_fwd is None else temp_fwd
-            )
-        if temp_inv is None:
-            temp_inv = (
-                self.zeros(inv_shape, cmpl_dtype) if temp_inv is None else temp_inv
-            )
-
-        default_values = {
-            "planner_effort": "FFTW_MEASURE",
-            "auto_align_input": False,
-            "auto_contiguous": False,
-            "avoid_copy": True,
-            "overwrite_input": True,
-            "threads": 1,
-        }
-        for key in default_values:
-            if key in fftargs:
-                continue
-            fftargs[key] = default_values[key]
-
-        rfft_shape = self._format_fft_shape(temp_fwd.shape, fwd_axes)
-        _rfftn = rfftn_builder(temp_fwd, s=rfft_shape, axes=fwd_axes, **fftargs)
-        overwrite_input = fftargs.pop("overwrite_input", None)
-
-        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
-        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
-        _irfftn = irfftn_builder(temp_inv, s=irfft_shape, axes=inv_axes, **fftargs)
-
-        def _rfftn_wrapper(arr, out, *args, **kwargs):
-            return _rfftn(arr, out)
-
-        def _irfftn_wrapper(arr, out, *args, **kwargs):
-            return _irfftn(arr, out)
-
-        fftargs["overwrite_input"] = overwrite_input
-        return _rfftn_wrapper, _irfftn_wrapper
+    def _rfftn(self, arr, out=None, **kwargs):
+        ret = interfaces.numpy_fft.rfftn(arr, **kwargs)
+        if out is not None:
+            out[:] = ret
+            return out
+        return ret
 
-    @staticmethod
-    def _format_fft_shape(shape: Tuple[int], axes: Tuple[int] = None):
-        if axes is None:
-            return shape
-        axes = tuple(sorted(range(len(shape))[i] for i in axes))
-        return tuple(shape[i] for i in axes)
+    def _irfftn(self, arr, out=None, **kwargs):
+        ret = interfaces.numpy_fft.irfftn(arr, **kwargs)
+        if out is not None:
+            out[:] = ret
+            return out
+        return ret
 
-    def rfftn(self, arr: NDArray, *args, **kwargs) -> NDArray:
-        return interfaces.numpy_fft.rfftn(arr, **kwargs)
+    def rfftn(
+        self,
+        arr: NDArray,
+        out=None,
+        auto_align_input: bool = False,
+        auto_contiguous: bool = False,
+        overwrite_input: bool = True,
+        **kwargs,
+    ) -> NDArray:
+        return self._rfftn(
+            arr,
+            auto_align_input=auto_align_input,
+            auto_contiguous=auto_contiguous,
+            overwrite_input=overwrite_input,
+            **kwargs,
+        )
 
-    def irfftn(self, arr: NDArray, *args, **kwargs) -> NDArray:
-        return interfaces.numpy_fft.irfftn(arr, **kwargs)
+    def irfftn(
+        self,
+        arr: NDArray,
+        out=None,
+        auto_align_input: bool = False,
+        auto_contiguous: bool = False,
+        overwrite_input: bool = True,
+        **kwargs,
+    ) -> NDArray:
+        return self._irfftn(
+            arr,
+            auto_align_input=auto_align_input,
+            auto_contiguous=auto_contiguous,
+            overwrite_input=overwrite_input,
+            **kwargs,
+        )
 
     def extract_center(self, arr: NDArray, newshape: Tuple[int]) -> NDArray:
         new_shape = self.to_backend_array(newshape)

tme/backends/pytorch_backend.py

@@ -7,7 +7,7 @@ Copyright (c) 2023 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from typing import Tuple, Callable
+from typing import Tuple
 from contextlib import contextmanager
 from multiprocessing import shared_memory
 from multiprocessing.managers import SharedMemoryManager
@@ -273,31 +273,6 @@ class PytorchBackend(NumpyFFTWBackend):
             kwargs["device"] = self.device
         return self._array_backend.eye(*args, **kwargs)
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int],
-        inv_output_shape: Tuple[int] = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        **kwargs,
-    ) -> Tuple[Callable, Callable]:
-        rfft_shape = self._format_fft_shape(fwd_shape, fwd_axes)
-        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
-        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
-
-        def rfftn(
-            arr: TorchTensor, out: TorchTensor, s=rfft_shape, axes=fwd_axes
-        ) -> TorchTensor:
-            return self._array_backend.fft.rfftn(arr, s=s, out=out, dim=axes)
-
-        def irfftn(
-            arr: TorchTensor, out: TorchTensor = None, s=irfft_shape, axes=inv_axes
-        ) -> TorchTensor:
-            return self._array_backend.fft.irfftn(arr, s=s, out=out, dim=axes)
-
-        return rfftn, irfftn
-
     def rfftn(self, arr: NDArray, *args, **kwargs) -> NDArray:
         kwargs["dim"] = kwargs.pop("axes", None)
         return self._array_backend.fft.rfftn(arr, **kwargs)

tme/density.py

@@ -36,6 +36,7 @@ from .matching_utils import (
     array_to_memmap,
     memmap_to_array,
     minimum_enclosing_box,
+    is_gzipped,
 )
 
 __all__ = ["Density"]
@@ -331,6 +332,7 @@ class Density:
         if non_standard_crs:
             data = np.transpose(data, crs_index)
             origin = np.take(origin, crs_index)
+            sampling_rate = np.take(sampling_rate, crs_index)
 
         return data.T, origin[::-1], sampling_rate[::-1], metadata
 
@@ -2257,9 +2259,3 @@ class Density:
         coordinates = np.array(np.where(data > 0))
         weights = self.data[tuple(coordinates)]
         return align_to_axis(coordinates.T, weights=weights, axis=axis, flip=flip)
-
-
-def is_gzipped(filename: str) -> bool:
-    """Check if a file is a gzip file by reading its magic number."""
-    with open(filename, "rb") as f:
-        return f.read(2) == b"\x1f\x8b"

tme/filters/ctf.py

@@ -36,7 +36,7 @@ class CTF(ComposableFilter):
 
     #: The shape of the to-be created mask.
     shape: Tuple[int] = None
-    #: The defocus value in x direction (in units of sampling rate).
+    #: The defocus in x direction (in units of sampling rate).
     defocus_x: Tuple[float] = None
     #: The tilt angles.
     angles: Tuple[float] = None
@@ -164,7 +164,7 @@ class CTF(ComposableFilter):
         shape : tuple of int
             The shape of the CTF.
         defocus_x : tuple of float
-            The defocus value in x direction.
+            The defocus in x direction (in units of sampling rate).
         angles : tuple of float
             The tilt angles.
         opening_axis : int, optional
@@ -178,7 +178,7 @@ class CTF(ComposableFilter):
         defocus_angle : tuple of float, optional
             The defocus angle in radians, defaults to 0.
         defocus_y : tuple of float, optional
-            The defocus value in y direction, defaults to None.
+            The defocus in x direction (in units of sampling rate).
         correct_defocus_gradient : bool, optional
             Whether to correct defocus gradient, defaults to False.
         sampling_rate : tuple of float, optional
@@ -219,14 +219,12 @@ class CTF(ComposableFilter):
                 corrected_tilt_axis -= 1
 
         for index, angle in enumerate(angles):
-            defocus_x, defocus_y = defoci_x[index], defoci_y[index]
-
             correction = correct_defocus_gradient and angle is not None
             chi = create_ctf(
                 angle=angle,
                 shape=ctf_shape,
-                defocus_x=defocus_x,
-                defocus_y=defocus_y,
+                defocus_x=defoci_x[index],
+                defocus_y=defoci_y[index],
                 sampling_rate=sampling_rate,
                 acceleration_voltage=acceleration_voltage[index],
                 correct_defocus_gradient=correction,
@@ -243,12 +241,10 @@ class CTF(ComposableFilter):
             stack[index] = chi
 
         # Avoid contrast inversion
-        np.negative(stack, out=stack)
+        stack = np.negative(stack, out=stack)
         if flip_phase:
-            np.abs(stack, out=stack)
-
-        stack = be.to_backend_array(np.squeeze(stack))
-        return stack
+            stack = np.abs(stack, out=stack)
+        return be.to_backend_array(np.squeeze(stack))
 
 
 class CTFReconstructed(CTF):
@@ -281,7 +277,7 @@ class CTFReconstructed(CTF):
         shape : tuple of int
             The shape of the CTF.
         defocus_x : tuple of float
-            The defocus value in x direction.
+            The defocus in x direction in units of sampling rate.
         opening_axis : int, optional
             The axis around which the wedge is opened, defaults to 2.
         amplitude_contrast : float, optional
@@ -291,7 +287,7 @@ class CTFReconstructed(CTF):
         defocus_angle : tuple of float, optional
             The defocus angle in radians, defaults to 0.
         defocus_y : tuple of float, optional
-            The defocus value in y direction, defaults to None.
+            The defocus in y direction in units of sampling rate.
         sampling_rate : tuple of float, optional
             The sampling rate, defaults to 1.
         acceleration_voltage : float, optional
@@ -321,18 +317,15 @@ class CTFReconstructed(CTF):
             defocus_angle=defocus_angle,
             amplitude_contrast=amplitude_contrast,
         )
-        stack = shift_fourier(data=stack, shape_is_real_fourier=False)
-
         # Avoid contrast inversion
         np.negative(stack, out=stack)
         if flip_phase:
             np.abs(stack, out=stack)
 
-        stack = be.to_backend_array(np.squeeze(stack))
+        stack = shift_fourier(data=stack, shape_is_real_fourier=False)
         if return_real_fourier:
             stack = crop_real_fourier(stack)
-
-        return stack
+        return be.to_backend_array(np.squeeze(stack))
 
 
 def _from_xml(filename: str) -> Dict:
@@ -436,6 +429,9 @@ def _from_ctffind(filename: str) -> Dict:
         output[key] = np.array(output[key])
 
     output["additional_phase_shift"] = np.degrees(output["additional_phase_shift"])
+    cs = output.get("spherical_aberration")
+    if cs is not None:
+        output["spherical_aberration"] = float(cs) * 1e7
     return output
 
 
@@ -566,7 +562,7 @@ def create_ctf(
     amplitude_contrast : float, optional
         Amplitude contrast of microscope, defaults to 0.07.
     spherical_aberration : float, optional
-        Spherical aberration of microscope in Angstrom.
+        Spherical aberration of microscope in units of sampling rate.
     angle : float, optional
         Assume the created CTF is a projection over opening_axis observed at angle.
     opening_axis : int, optional
@@ -590,10 +586,14 @@ def create_ctf(
     electron_wavelength = _compute_electron_wavelength(acceleration_voltage)
     electron_wavelength /= sampling_rate
     aberration = (spherical_aberration / sampling_rate) * electron_wavelength**2
+
+    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:
         if len(shape) < 2:
             raise ValueError(f"Length of shape needs to be at least 2, got {shape}")
 
+        # Axial distance from grid center in multiples of sampling rate
         sampling = tuple(float(x) for x in np.divide(sampling_rate, shape))
         grid = fftfreqn(
             shape=shape,
@@ -619,6 +619,7 @@ def create_ctf(
         defocus_sum = np.add(defocus_x, defocus_y)
         defocus_difference = np.subtract(defocus_x, defocus_y)
 
+        # Reusing grid, but in principle pure frequencies would suffice
         angular_grid = np.arctan2(grid[1], grid[0])
         defocus_difference = np.multiply(
             defocus_difference,
@@ -627,7 +628,7 @@ def create_ctf(
         defocus_x = np.add(defocus_sum, defocus_difference)
         defocus_x *= 0.5
 
-    frequency_grid = fftfreqn(shape, sampling_rate=True, compute_euclidean_norm=True)
+    frequency_grid = fftfreqn(shape, sampling_rate=1, compute_euclidean_norm=True)
     if angle is not None and opening_axis is not None and full_shape is not None:
         frequency_grid = frequency_grid_at_angle(
             shape=full_shape,

tme/filters/wedge.py

@@ -15,7 +15,7 @@ import numpy as np
 from ..types import NDArray
 from ..backends import backend as be
 from .compose import ComposableFilter
-from ..matching_utils import centered
+from ..matching_utils import center_slice
 from ..rotations import euler_to_rotationmatrix
 from ..parser import XMLParser, StarParser, MDOCParser
 from ._utils import (
@@ -207,11 +207,10 @@ class Wedge(ComposableFilter):
             )
             sigma = np.sqrt(self.weights[index] * 4 / (8 * np.pi**2))
             sigma = -2 * np.pi**2 * sigma**2
-            np.square(frequency_grid, out=frequency_grid)
-            np.multiply(sigma, frequency_grid, out=frequency_grid)
-            np.exp(frequency_grid, out=frequency_grid)
-            np.multiply(frequency_grid, np.cos(np.radians(angle)), out=frequency_grid)
-            wedges[index] = frequency_grid
+            frequency_grid = np.square(frequency_grid, out=frequency_grid)
+            frequency_grid = np.multiply(sigma, frequency_grid, out=frequency_grid)
+            frequency_grid = np.exp(frequency_grid, out=frequency_grid)
+            wedges[index] = np.multiply(frequency_grid, np.cos(np.radians(angle)))
 
         return wedges
 
@@ -490,7 +489,11 @@ class WedgeReconstructed:
             )
             wedge_volume += plane_rotated * weights[index]
 
-        wedge_volume = centered(wedge_volume, (shape[opening_axis], shape[tilt_axis]))
+        subset = center_slice(
+            wedge_volume.shape, (shape[opening_axis], shape[tilt_axis])
+        )
+        wedge_volume = wedge_volume[subset]
+
         np.fmin(wedge_volume, np.max(weights), wedge_volume)
 
         if opening_axis > tilt_axis: