pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__cp311-cp311-macosx_15_0_arm64.whl

tme/filters/ctf.py

@@ -7,7 +7,6 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import re
-import warnings
 from typing import Tuple, Dict
 from dataclasses import dataclass
 
@@ -16,13 +15,14 @@ import numpy as np
 from ..types import NDArray
 from ..backends import backend as be
 from .compose import ComposableFilter
-from ..parser import StarParser, XMLParser
+from ..parser import StarParser, XMLParser, MDOCParser
 from ._utils import (
     frequency_grid_at_angle,
     compute_tilt_shape,
     crop_real_fourier,
     fftfreqn,
     shift_fourier,
+    pad_to_length,
 )
 
 __all__ = ["CTF", "CTFReconstructed", "create_ctf"]
@@ -34,39 +34,39 @@ class CTF(ComposableFilter):
     Generate a per-tilt contrast transfer function mask.
     """
 
-    #: The shape of the to-be reconstructed volume.
+    #: The shape of the to-be created mask.
     shape: Tuple[int] = None
-    #: The defocus value in x direction.
-    defocus_x: float = None
+    #: The defocus in x direction (in units of sampling rate).
+    defocus_x: Tuple[float] = None
     #: The tilt angles.
     angles: Tuple[float] = None
-    #: The microscope projection axis, defaults to None.
+    #: The microscope projection axis, defaults to 2 (z).
     opening_axis: int = 2
-    #: The axis along which the tilt is applied, defaults to 2 (z).
+    #: The axis along which the tilt is applied, defaults to 0 (x).
     tilt_axis: int = 0
-    #: Whether to correct defocus gradient, defaults to 0 (x).
+    #: Whether to correct defocus gradient, defaults False.
     correct_defocus_gradient: bool = False
-    #: The sampling rate, defaults to 1 Angstrom / Voxel.
+    #: The sampling rate, defaults to 1 Ångstrom / voxel.
     sampling_rate: Tuple[float] = 1
     #: The acceleration voltage in Volts, defaults to 300e3.
-    acceleration_voltage: float = 300e3
-    #: The spherical aberration coefficient, defaults to 2.7e7.
-    spherical_aberration: float = 2.7e7
+    acceleration_voltage: Tuple[float] = 300e3
+    #: The spherical aberration, defaults to 2.7e7 (in units of sampling rate).
+    spherical_aberration: Tuple[float] = 2.7e7
     #: The amplitude contrast, defaults to 0.07.
-    amplitude_contrast: float = 0.07
+    amplitude_contrast: Tuple[float] = 0.07
     #: The phase shift in degrees, defaults to 0.
-    phase_shift: float = 0
+    phase_shift: Tuple[float] = 0
     #: The defocus angle in degrees, defaults to 0.
-    defocus_angle: float = 0
-    #: The defocus value in y direction, defaults to None.
-    defocus_y: float = None
-    #: Whether the returned CTF should be phase-flipped.
+    defocus_angle: Tuple[float] = 0
+    #: The defocus value in y direction, defaults to None (in units of sampling rate).
+    defocus_y: Tuple[float] = None
+    #: Whether the returned CTF should be phase-flipped, defaults to True.
     flip_phase: bool = True
-    #: Whether to return a format compliant with rfft. Only relevant for single angles.
+    #: Whether to return a ctf mask for rfft (for :py:class:`CTFReconstructed`).
     return_real_fourier: bool = False
 
     @classmethod
-    def from_file(cls, filename: str) -> "CTF":
+    def from_file(cls, filename: str, **kwargs) -> "CTF":
         """
         Initialize :py:class:`CTF` from file.
 
@@ -80,36 +80,50 @@ class CTF(ComposableFilter):
             +-------+---------------------------------------------------------+
             | .xml  | WARP/M XML file                                         |
             +-------+---------------------------------------------------------+
+            | .mdoc | SerialEM file                                           |
+            +-------+---------------------------------------------------------+
             | .*    | CTFFIND4 file                                           |
             +-------+---------------------------------------------------------+
+        **kwargs : optional
+            Overwrite fields that cannot be extracted from input file.
         """
         func = _from_ctffind
         if filename.lower().endswith("star"):
-            func = _from_gctf
+            func = _from_star
         elif filename.lower().endswith("xml"):
             func = _from_xml
+        elif filename.lower().endswith("mdoc"):
+            func = _from_mdoc
 
         data = func(filename=filename)
 
         # Pixel size needs to be overwritten by pixel size the ctf is generated for
-        return cls(
-            shape=None,
-            angles=data.get("angles", None),
-            defocus_x=data["defocus_1"],
-            sampling_rate=data["pixel_size"],
-            acceleration_voltage=data["acceleration_voltage"],
-            spherical_aberration=data["spherical_aberration"],
-            amplitude_contrast=data["amplitude_contrast"],
-            phase_shift=data["additional_phase_shift"],
-            defocus_angle=data["azimuth_astigmatism"],
-            defocus_y=data["defocus_2"],
-        )
+        init_kwargs = {
+            "shape": None,
+            "angles": data.get("angles", None),
+            "defocus_x": data["defocus_1"],
+            "sampling_rate": data["pixel_size"],
+            "acceleration_voltage": np.multiply(data["acceleration_voltage"], 1e3),
+            "spherical_aberration": data.get("spherical_aberration"),
+            "amplitude_contrast": data.get("amplitude_contrast"),
+            "phase_shift": data.get("additional_phase_shift"),
+            "defocus_angle": data.get("azimuth_astigmatism"),
+            "defocus_y": data["defocus_2"],
+        }
+        for k, v in kwargs.items():
+            if k in init_kwargs and init_kwargs.get(k) is None:
+                init_kwargs[k] = v
+        init_kwargs = {k: v for k, v in init_kwargs.items() if v is not None}
+        return cls(**init_kwargs)
 
     def __post_init__(self):
         self.defocus_angle = np.radians(self.defocus_angle)
         self.phase_shift = np.radians(self.phase_shift)
 
     def __call__(self, **kwargs) -> NDArray:
+        """
+        Returns a CTF stack of chosen parameters with DC component in the center.
+        """
         func_args = vars(self).copy()
         func_args.update(kwargs)
 
@@ -122,11 +136,6 @@ class CTF(ComposableFilter):
             "is_multiplicative_filter": True,
         }
 
-    @staticmethod
-    def _pad_to_length(arr, length: int):
-        ret = np.atleast_1d(arr)
-        return np.repeat(ret, length // ret.size)
-
     def weight(
         self,
         shape: Tuple[int],
@@ -134,7 +143,7 @@ class CTF(ComposableFilter):
         angles: Tuple[float],
         opening_axis: int = 2,
         tilt_axis: int = 0,
-        amplitude_contrast: float = 0.07,
+        amplitude_contrast: Tuple[float] = 0.07,
         phase_shift: Tuple[float] = 0,
         defocus_angle: Tuple[float] = 0,
         defocus_y: Tuple[float] = None,
@@ -155,21 +164,21 @@ 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
             The axis around which the wedge is opened, defaults to 2.
         tilt_axis : int, optional
             The axis along which the tilt is applied, defaults to 0.
-        amplitude_contrast : float, optional
+        amplitude_contrast : tuple of float, optional
             The amplitude contrast, defaults to 0.07.
         phase_shift : tuple of float, optional
             The phase shift in radians, defaults to 0.
         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
@@ -179,7 +188,7 @@ class CTF(ComposableFilter):
         spherical_aberration : float, optional
             The spherical aberration coefficient, defaults to 2.7e3.
         flip_phase : bool, optional
-            Whether the returned CTF should be phase-flipped.
+            Whether the returned CTF should be phase-flipped, defaults to True.
         **kwargs : Dict
             Additional keyword arguments.
 
@@ -189,12 +198,13 @@ class CTF(ComposableFilter):
             A stack containing the CTF weight.
         """
         angles = np.atleast_1d(angles)
-        defoci_x = self._pad_to_length(defocus_x, angles.size)
-        defoci_y = self._pad_to_length(defocus_y, angles.size)
-        phase_shift = self._pad_to_length(phase_shift, angles.size)
-        defocus_angle = self._pad_to_length(defocus_angle, angles.size)
-        spherical_aberration = self._pad_to_length(spherical_aberration, angles.size)
-        amplitude_contrast = self._pad_to_length(amplitude_contrast, angles.size)
+        defoci_x = pad_to_length(defocus_x, angles.size)
+        defoci_y = pad_to_length(defocus_y, angles.size)
+        phase_shift = pad_to_length(phase_shift, angles.size)
+        defocus_angle = pad_to_length(defocus_angle, angles.size)
+        spherical_aberration = pad_to_length(spherical_aberration, angles.size)
+        amplitude_contrast = pad_to_length(amplitude_contrast, angles.size)
+        acceleration_voltage = pad_to_length(acceleration_voltage, angles.size)
 
         sampling_rate = np.max(sampling_rate)
         ctf_shape = compute_tilt_shape(
@@ -209,16 +219,14 @@ 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=self.acceleration_voltage,
+                acceleration_voltage=acceleration_voltage[index],
                 correct_defocus_gradient=correction,
                 spherical_aberration=spherical_aberration[index],
                 cutoff_frequency=cutoff_frequency,
@@ -233,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):
@@ -271,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
@@ -281,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
@@ -311,21 +317,18 @@ 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 be.to_backend_array(np.squeeze(stack))
 
-        return stack
 
-
-def _from_xml(filename: str):
+def _from_xml(filename: str) -> Dict:
     data = XMLParser(filename)
 
     params = {
@@ -353,6 +356,7 @@ def _from_xml(filename: str):
         params["PhaseShift"] = [
             ctf_phase[i]["@attributes"]["Value"] for i in range(len(ctf_phase))
         ]
+        params["PhaseShift"] = np.degrees(params["PhaseShift"])
         ctf_ast = data["GridCTFDefocusAngle"]["Node"]
         params["DefocusAngle"] = [
             ctf_ast[i]["@attributes"]["Value"] for i in range(len(ctf_ast))
@@ -384,7 +388,7 @@ def _from_xml(filename: str):
     return {k: params[v] for k, v in mapping.items()}
 
 
-def _from_ctffind(filename: str):
+def _from_ctffind(filename: str) -> Dict:
     parameter_regex = {
         "pixel_size": r"Pixel size: ([0-9.]+) Angstroms",
         "acceleration_voltage": r"acceleration voltage: ([0-9.]+) keV",
@@ -425,18 +429,64 @@ def _from_ctffind(filename: str):
         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
 
 
-def _from_gctf(filename: str):
+def _from_star(filename: str) -> Dict:
     parser = StarParser(filename)
-    ctf_data = parser["data_"]
+
+    if "data_stopgap_wedgelist" in parser:
+        key = "data_stopgap_wedgelist"
+        mapping = {
+            "angles": ("_tilt_angle", float, 1),
+            "defocus_1": ("_defocus", float, 1e4),
+            "defocus_2": (None, float, 1e4),
+            "pixel_size": ("_pixelsize", float, 1),
+            "acceleration_voltage": ("_voltage", float, 1),
+            "spherical_aberration": ("_cs", float, 1e7),
+            "amplitude_contrast": ("_amp_contrast", float, 1),
+            "additional_phase_shift": (None, float, 1),
+            "azimuth_astigmatism": (None, float, 1),
+        }
+    else:
+        key = "data_"
+        mapping = {
+            "defocus_1": ("_rlnDefocusU", float, 1),
+            "defocus_2": ("_rlnDefocusV", float, 1),
+            "pixel_size": ("_rlnDetectorPixelSize", float, 1),
+            "acceleration_voltage": ("_rlnVoltage", float, 1),
+            "spherical_aberration": ("_rlnSphericalAberration", float, 1),
+            "amplitude_contrast": ("_rlnAmplitudeContrast", float, 1),
+            "additional_phase_shift": (None, float, 1),
+            "azimuth_astigmatism": ("_rlnDefocusAngle", float, 1),
+        }
+
+    output = {}
+    ctf_data = parser[key]
+    for out_key, (key, key_dtype, scale) in mapping.items():
+        key_value = ctf_data.get(key)
+        if key_value is not None:
+            try:
+                key_value = [key_dtype(x) * scale for x in key_value]
+            except Exception:
+                pass
+        output[out_key] = key_value
+    return output
+
+
+def _from_mdoc(filename: str) -> Dict:
+    parser = MDOCParser(filename)
 
     mapping = {
-        "defocus_1": ("_rlnDefocusU", float),
-        "defocus_2": ("_rlnDefocusV", float),
+        "angles": ("TiltAngle", float),
+        "defocus_1": ("Defocus", float),
+        "acceleration_voltage": ("Voltage", float),
+        # These will be None, but on purpose
         "pixel_size": ("_rlnDetectorPixelSize", float),
-        "acceleration_voltage": ("_rlnVoltage", float),
+        "defocus_2": ("Defocus2", float),
         "spherical_aberration": ("_rlnSphericalAberration", float),
         "amplitude_contrast": ("_rlnAmplitudeContrast", float),
         "additional_phase_shift": (None, float),
@@ -444,14 +494,10 @@ def _from_gctf(filename: str):
     }
     output = {}
     for out_key, (key, key_dtype) in mapping.items():
-        if key not in ctf_data and key is not None:
-            warnings.warn(f"ctf_data is missing key {key}.")
-
-        key_value = ctf_data.get(key, [0])
-        output[out_key] = [key_dtype(x) for x in key_value]
+        output[out_key] = parser.get(key, None)
 
-    longest_key = max(map(len, output.values()))
-    output = {k: v * longest_key if len(v) == 1 else v for k, v in output.items()}
+    # Adjust convention and convert to Angstrom
+    output["defocus_1"] = np.multiply(output["defocus_1"], -1e4)
     return output
 
 
@@ -516,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
@@ -540,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,
@@ -569,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,
@@ -577,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,
@@ -589,10 +640,10 @@ def create_ctf(
     frequency_mask = frequency_grid < cutoff_frequency
 
     # k^2*π*λ(dx - 0.5 * sph_abb * λ^2 * k^2) + phase_shift + ampl_contrast_term)
-    np.square(frequency_grid, out=frequency_grid)
+    frequency_grid = np.square(frequency_grid, out=frequency_grid)
     chi = defocus_x - 0.5 * aberration * frequency_grid
-    np.multiply(chi, np.pi * electron_wavelength, out=chi)
-    np.multiply(chi, frequency_grid, out=chi)
+    chi = np.multiply(chi, np.pi * electron_wavelength, out=chi)
+    chi = np.multiply(chi, frequency_grid, out=chi)
     chi += phase_shift
     chi += np.arctan(
         np.divide(
@@ -600,6 +651,6 @@ def create_ctf(
             np.sqrt(1 - np.square(amplitude_contrast)),
         )
     )
-    np.sin(-chi, out=chi)
-    np.multiply(chi, frequency_mask, out=chi)
+    chi = np.sin(-chi, out=chi)
+    chi = np.multiply(chi, frequency_mask, out=chi)
     return chi

tme/filters/reconstruction.py

@@ -11,14 +11,14 @@ from dataclasses import dataclass
 
 import numpy as np
 
-from ..types import NDArray
 from ..backends import backend as be
+from ..types import NDArray, BackendArray
 
 from .compose import ComposableFilter
 from ..rotations import euler_to_rotationmatrix
 from ._utils import crop_real_fourier, shift_fourier, create_reconstruction_filter
 
-__all__ = ["ReconstructFromTilt"]
+__all__ = ["ReconstructFromTilt", "ShiftFourier"]
 
 
 @dataclass
@@ -52,7 +52,9 @@ class ReconstructFromTilt(ComposableFilter):
         shape : tuple of int
             The shape of the reconstruction volume.
         data : BackendArray
-            D-dimensional image stack with shape (n, ...)
+            D-dimensional image stack with shape (n, ...). The data is assumed to be
+            a Fourier transform of the stack you are trying to reconstruct with
+            DC component in the center.
         angles : tuple of float
             Angle of each individual tilt.
         return_real_fourier : bool, optional
@@ -87,13 +89,14 @@ class ReconstructFromTilt(ComposableFilter):
 
         ret = self.reconstruct(**func_args)
 
+        ret = shift_fourier(data=ret, shape_is_real_fourier=False)
         if return_real_fourier:
             ret = crop_real_fourier(ret)
 
         return {
             "data": ret,
             "shape": func_args["shape"],
-            "shape_is_real_fourier": return_real_fourier,
+            "return_real_fourier": return_real_fourier,
             "is_multiplicative_filter": False,
         }
 
@@ -114,7 +117,7 @@ class ReconstructFromTilt(ComposableFilter):
         Parameters
         ----------
         data : NDArray
-            The tilt series data.
+            The Fourier transform of tilt series data.
         shape : tuple of int
             Shape of the reconstruction.
         angles : tuple of float
@@ -138,9 +141,9 @@ class ReconstructFromTilt(ComposableFilter):
             return data
 
         data = be.to_backend_array(data)
-        volume_temp = be.zeros(shape, dtype=be._float_dtype)
-        volume_temp_rotated = be.zeros(shape, dtype=be._float_dtype)
-        volume = be.zeros(shape, dtype=be._float_dtype)
+        volume_temp = be.zeros(shape, dtype=data.dtype)
+        volume_temp_rotated = be.zeros(shape, dtype=data.dtype)
+        volume = be.zeros(shape, dtype=data.dtype)
 
         slices = tuple(slice(a // 2, (a // 2) + 1) for a in shape)
         subset = tuple(
@@ -187,4 +190,30 @@ class ReconstructFromTilt(ComposableFilter):
             )
             volume = be.add(volume, volume_temp_rotated, out=volume)
 
-        return shift_fourier(data=volume, shape_is_real_fourier=False)
+        return volume
+
+
+class ShiftFourier(ComposableFilter):
+    def __call__(
+        self,
+        data: BackendArray,
+        shape_is_real_fourier: bool = False,
+        return_real_fourier: bool = True,
+        **kwargs,
+    ):
+        ret = []
+        for index in range(data.shape[0]):
+            mask = be.to_numpy_array(data[index])
+
+            mask = shift_fourier(data=mask, shape_is_real_fourier=shape_is_real_fourier)
+            if return_real_fourier:
+                mask = crop_real_fourier(mask)
+            ret.append(mask[None])
+        ret = np.concatenate(ret, axis=0)
+
+        return {
+            "data": ret,
+            "shape": kwargs.get("shape"),
+            "return_real_fourier": return_real_fourier,
+            "is_multiplicative_filter": False,
+        }