Merge pull request #20 from mgiammar/main

Resolve a circular import error

Author: mgiammar

src/leopard_em/backend/core_refine_template.py

@@ -14,7 +14,7 @@
 )
 from leopard_em.backend.utils import normalize_template_projection
 from leopard_em.utils.cross_correlation import handle_correlation_mode
-from leopard_em.utils.pre_processing import calculate_ctf_filter_stack
+from leopard_em.utils.filter_preprocessing import calculate_ctf_filter_stack
 
 # This is assuming the Euler angles are in the ZYZ intrinsic format
 # AND that the angles are ordered in (phi, theta, psi)

src/leopard_em/pydantic_models/match_template_manager.py

@@ -18,7 +18,7 @@
 from leopard_em.pydantic_models.orientation_search import OrientationSearchConfig
 from leopard_em.pydantic_models.types import BaseModel2DTM, ExcludedTensor
 from leopard_em.utils.data_io import load_mrc_image, load_mrc_volume
-from leopard_em.utils.pre_processing import calculate_ctf_filter_stack
+from leopard_em.utils.filter_preprocessing import calculate_ctf_filter_stack
 
 
 class MatchTemplateManager(BaseModel2DTM):

src/leopard_em/pydantic_models/particle_stack.py

@@ -219,8 +219,8 @@ def construct_cropped_statistic_stack(
 
             # with reference to the exact pixel of the statistic (top-left)
             # need to account for relative extracted box size
-            pos_y = self._df.loc[indexes, "pos_y"]
-            pos_x = self._df.loc[indexes, "pos_x"]
+            pos_y = self._df.loc[indexes, "pos_y"].to_numpy()
+            pos_x = self._df.loc[indexes, "pos_x"].to_numpy()
             pos_y = torch.tensor(pos_y)
             pos_x = torch.tensor(pos_x)
             pos_y -= (H - h) // 2

src/leopard_em/utils/__init__.py

@@ -0,0 +1,31 @@
+"""Utilities submodule for various data and pre- and post-processing tasks."""
+
+from .cross_correlation import handle_correlation_mode
+from .data_io import (
+    load_mrc_image,
+    load_mrc_volume,
+    read_mrc_to_numpy,
+    read_mrc_to_tensor,
+    write_mrc_from_numpy,
+    write_mrc_from_tensor,
+)
+from .filter_preprocessing import (
+    Cs_to_pixel_size,
+    calculate_ctf_filter_stack,
+    get_Cs_range,
+)
+from .particle_stack import get_cropped_image_regions
+
+__all__ = [
+    "handle_correlation_mode",
+    "read_mrc_to_numpy",
+    "read_mrc_to_tensor",
+    "write_mrc_from_numpy",
+    "write_mrc_from_tensor",
+    "load_mrc_image",
+    "load_mrc_volume",
+    "get_cropped_image_regions",
+    "calculate_ctf_filter_stack",
+    "get_Cs_range",
+    "Cs_to_pixel_size",
+]

src/leopard_em/utils/pre_processing.py …leopard_em/utils/filter_preprocessing.pysrc/leopard_em/utils/pre_processing.py renamed to src/leopard_em/utils/filter_preprocessing.py src/leopard_em/utils/pre_processing.py renamed to src/leopard_em/utils/filter_preprocessing.py

@@ -1,10 +1,59 @@
-"""Helper functions for pre-processing data for 2DTM."""
+"""Helper functions for the CTF filter preprocessing."""
 
 import einops
 import torch
 from torch_fourier_filter.ctf import calculate_ctf_2d
 
-from leopard_em.pydantic_models import WhiteningFilterConfig
+
+def get_Cs_range(
+    pixel_size: float,
+    pixel_size_offsets: torch.Tensor,
+    Cs: float = 2.7,
+) -> torch.Tensor:
+    """Get the Cs values for a  range of pixel sizes.
+
+    Parameters
+    ----------
+    pixel_size : float
+        The nominal pixel size.
+    pixel_size_offsets : torch.Tensor
+        The pixel size offsets.
+    Cs : float, optional
+        The Cs value, by default 2.7.
+
+    Returns
+    -------
+    torch.Tensor
+        The Cs values for the range of pixel sizes.
+    """
+    pixel_sizes = pixel_size + pixel_size_offsets
+    Cs_values = Cs / torch.pow(pixel_sizes / pixel_size, 4)
+    return Cs_values
+
+
+def Cs_to_pixel_size(
+    Cs_vals: torch.Tensor,
+    nominal_pixel_size: float,
+    nominal_Cs: float = 2.7,
+) -> torch.Tensor:
+    """Convert Cs values to pixel sizes.
+
+    Parameters
+    ----------
+    Cs_vals : torch.Tensor
+        The Cs values.
+    nominal_pixel_size : float
+        The nominal pixel size.
+    nominal_Cs : float, optional
+        The nominal Cs value, by default 2.7.
+
+    Returns
+    -------
+    torch.Tensor
+        The pixel sizes.
+    """
+    pixel_size = torch.pow(nominal_Cs / Cs_vals, 0.25) * nominal_pixel_size
+    return pixel_size
 
 
 def calculate_ctf_filter_stack(
@@ -98,112 +147,3 @@ def calculate_ctf_filter_stack(
     # The CTF will have a shape of (n_Cs n_defoc, nx, ny)
     # These will catch any potential errors
     return ctf
-
-
-def do_image_preprocessing(
-    image_rfft: torch.Tensor,
-    wf_config: WhiteningFilterConfig,
-) -> torch.Tensor:
-    """Pre-processes the input image before running the algorithm.
-
-    1. Zero central pixel (0, 0)
-    2. Calculate a whitening filter
-    3. Do element-wise multiplication with the whitening filter
-    4. Zero central pixel again (superfluous, but following cisTEM)
-    5. Normalize (x /= sqrt(sum(abs(x)**2)); pixelwise)
-
-    Parameters
-    ----------
-    image_rfft : torch.Tensor
-        The input image, RFFT'd and unshifted.
-    wf_config : WhiteningFilterConfig
-        The configuration for the whitening filter.
-
-    Returns
-    -------
-    torch.Tensor
-        The pre-processed image.
-
-    """
-    H, W = image_rfft.shape
-    W = (W - 1) * 2  # Account for RFFT
-    npix_real = H * W
-
-    # Zero out the constant term
-    image_rfft[0, 0] = 0 + 0j
-
-    wf_image = wf_config.calculate_whitening_filter(
-        ref_img_rfft=image_rfft,
-        output_shape=image_rfft.shape,
-    )
-    image_rfft *= wf_image
-    image_rfft[0, 0] = 0 + 0j  # superfluous, but following cisTEM
-
-    # NOTE: Extra indexing happening with squared_sum so that Hermitian pairs are
-    # counted, but we skip the first column of the RFFT which should not be duplicated.
-    squared_image_rfft = torch.abs(image_rfft) ** 2
-    squared_sum = squared_image_rfft.sum() + squared_image_rfft[:, 1:].sum()
-    image_rfft /= torch.sqrt(squared_sum)
-
-    # # real-space image will now have mean=0 and variance=1
-    # image_rfft *= npix_real  # NOTE: This would set the variance to 1 exactly, but...
-
-    # NOTE: We add on extra division by sqrt(num_pixels) so the cross-correlograms
-    # are roughly normalized to have mean 0 and variance 1.
-    # We do this here since Fourier transform is linear, and we don't have to multiply
-    # the cross correlation at each iteration. This *will not* make the image
-    # have variance 1.
-    image_rfft *= npix_real**0.5
-
-    return image_rfft
-
-
-def get_Cs_range(
-    pixel_size: float,
-    pixel_size_offsets: torch.Tensor,
-    Cs: float = 2.7,
-) -> torch.Tensor:
-    """Get the Cs values for a  range of pixel sizes.
-
-    Parameters
-    ----------
-    pixel_size : float
-        The nominal pixel size.
-    pixel_size_offsets : torch.Tensor
-        The pixel size offsets.
-    Cs : float, optional
-        The Cs value, by default 2.7.
-
-    Returns
-    -------
-    torch.Tensor
-        The Cs values for the range of pixel sizes.
-    """
-    pixel_sizes = pixel_size + pixel_size_offsets
-    Cs_values = Cs / torch.pow(pixel_sizes / pixel_size, 4)
-    return Cs_values
-
-
-def Cs_to_pixel_size(
-    Cs_vals: torch.Tensor,
-    nominal_pixel_size: float,
-    nominal_Cs: float = 2.7,
-) -> torch.Tensor:
-    """Convert Cs values to pixel sizes.
-
-    Parameters
-    ----------
-    Cs_vals : torch.Tensor
-        The Cs values.
-    nominal_pixel_size : float
-        The nominal pixel size.
-    nominal_Cs : float, optional
-        The nominal Cs value, by default 2.7.
-
-    Returns
-    -------
-    torch.Tensor
-        The pixel sizes.
-    """
-    pixel_size = torch.pow(nominal_Cs / Cs_vals, 0.25) * nominal_pixel_size
-    return pixel_size