Merge pull request #110 from Lucaslab-Berkeley/jd_ctf_aberrations

CTF aberrations now included

Author: jdickerson95

pyproject.toml

@@ -48,15 +48,16 @@ dependencies = [
     "types-PyYAML",
     "roma",
     "tqdm",
-    "torch-fourier-slice>=v0.2.0",
+    "torch-fourier-slice>=v0.4.0",
     "torch-fourier-filter>=v0.2.6",
     "torch-so3>=v0.2.0",
     "ttsim3d>=v0.4.0",
     "lmfit",
     "zenodo-get",
     "torch-fourier-shift",
     "torch-motion-correction>=0.0.4",
-    "torch-grid-utils>=v0.0.9"
+    "torch-grid-utils>=v0.0.9",
+    "torch-ctf"
 ]
 
 [tool.hatch.metadata]
@@ -156,7 +157,10 @@ pretty = true
 [tool.pytest.ini_options]
 minversion = "7.0"
 testpaths = ["tests"]
-filterwarnings = ["error"]
+filterwarnings = [
+    "error",
+    "ignore::FutureWarning",
+]
 addopts = "-m 'not slow'"               # Skip slow tests on default
 markers = ["slow: marks test as slow"]
 

src/leopard_em/analysis/__init__.py

@@ -6,7 +6,10 @@
     match_template_peaks_to_dict,
 )
 from .pvalue_metric import extract_peaks_and_statistics_p_value
-from .zscore_metric import extract_peaks_and_statistics_zscore, gaussian_noise_zscore_cutoff
+from .zscore_metric import (
+    extract_peaks_and_statistics_zscore,
+    gaussian_noise_zscore_cutoff,
+)
 
 __all__ = [
     "MatchTemplatePeaks",

src/leopard_em/backend/core_differentiable_refine.py

@@ -68,6 +68,7 @@ def core_differentiable_refine(
     device: torch.device | list[torch.device],
     batch_size: int = 32,
     num_cuda_streams: int = 1,
+    transform_matrix: torch.Tensor | None = None,
 ) -> dict[torch.Tensor, torch.Tensor]:
     """Core function to refine orientations and defoci of a set of particles.
 
@@ -112,6 +113,9 @@ def core_differentiable_refine(
         The number of cross-correlations to process in one batch, defaults to 32.
     num_cuda_streams : int, optional
         Number of CUDA streams to use for parallel processing. Defaults to 1.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
 
     Returns
     -------
@@ -143,6 +147,7 @@ def core_differentiable_refine(
         batch_size=batch_size,
         devices=device,
         num_cuda_streams=num_cuda_streams,
+        transform_matrix=transform_matrix,
     )
 
     results = {}
@@ -230,6 +235,7 @@ def _core_refine_template_single_gpu(
     batch_size: int,
     device: torch.device,
     num_cuda_streams: int = 1,
+    transform_matrix: torch.Tensor | None = None,
 ) -> None:
     """Run refine template on a subset of particles on a single GPU.
 
@@ -273,6 +279,9 @@ def _core_refine_template_single_gpu(
         Torch device to run this process on.
     num_cuda_streams : int, optional
         Number of CUDA streams to use for parallel processing. Defaults to 1.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
     """
     streams = [torch.cuda.Stream(device=device) for _ in range(num_cuda_streams)]
 
@@ -293,6 +302,8 @@ def _core_refine_template_single_gpu(
     corr_mean = corr_mean.to(device)
     corr_std = corr_std.to(device)
     projective_filters = projective_filters.to(device)
+    if transform_matrix is not None:
+        transform_matrix = transform_matrix.to(device)
 
     ########################################
     ### Setup constants and progress bar ###
@@ -345,6 +356,7 @@ def _core_refine_template_single_gpu(
                 projective_filter=projective_filters[i],
                 batch_size=batch_size,
                 device_id=device_id,
+                transform_matrix=transform_matrix,
             )
             refined_statistics.append(refined_stats)
 
@@ -457,6 +469,7 @@ def _core_refine_template_single_thread(
     projective_filter: torch.Tensor,
     batch_size: int = 32,
     device_id: int = 0,
+    transform_matrix: torch.Tensor | None = None,
 ) -> dict[str, float | int]:
     """Run the single-threaded core refine template function.
 
@@ -497,6 +510,9 @@ def _core_refine_template_single_thread(
         The number of orientations to cross-correlate at once. Default is 32.
     device_id : int, optional
         The ID of the device/process. Default is 0.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
 
     Returns
     -------
@@ -600,13 +616,15 @@ def _core_refine_template_single_thread(
                 rotation_matrices=rot_matrix_batch,
                 projective_filters=combined_projective_filter,
                 requires_grad=True,
+                transform_matrix=transform_matrix,
             )
         else:
             cross_correlation = do_batched_orientation_cross_correlate_cpu(
                 image_dft=particle_image_dft,
                 template_dft=template_dft,
                 rotation_matrices=rot_matrix_batch,
                 projective_filters=combined_projective_filter,
+                transform_matrix=transform_matrix,
             )
 
         cross_correlation = cross_correlation[..., :crop_h, :crop_w]  # valid crop

src/leopard_em/backend/core_match_template.py

@@ -155,6 +155,7 @@ def core_match_template(
     orientation_batch_size: int = 1,
     num_cuda_streams: int = 1,
     backend: str = "streamed",
+    transform_matrix: torch.Tensor | None = None,
 ) -> dict[str, torch.Tensor]:
     """Core function for performing the whole-orientation search.
 
@@ -209,6 +210,9 @@ def core_match_template(
     backend : str, optional
         The backend to use for computation. Defaults to 'streamed'.
         Must be 'streamed' or 'batched'.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
 
     Returns
     -------
@@ -259,6 +263,9 @@ def core_match_template(
     defocus_values = defocus_values.cpu()
     pixel_values = pixel_values.cpu()
     euler_angles = euler_angles.cpu()
+    # Move transform_matrix to CPU if it's not None
+    if transform_matrix is not None:
+        transform_matrix = transform_matrix.cpu()
 
     ##############################################################
     ### Pre-multiply the whitening filter with the CTF filters ###
@@ -308,6 +315,7 @@ def core_match_template(
             "num_cuda_streams": num_cuda_streams,
             "backend": backend,
             "device": d,
+            "transform_matrix": transform_matrix,
         }
 
         kwargs_per_device.append(kwargs)
@@ -371,6 +379,7 @@ def _core_match_template_single_gpu(
     num_cuda_streams: int,
     backend: str,
     device: torch.device,
+    transform_matrix: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Single-GPU call for template matching.
 
@@ -410,6 +419,9 @@ def _core_match_template_single_gpu(
         Defaults to 'streamed'. Must be 'streamed' or 'batched'.
     device : torch.device
         Device to run the computation on. All tensors must be allocated on this device.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
 
     Returns
     -------
@@ -435,6 +447,9 @@ def _core_match_template_single_gpu(
     template_dft = template_dft.to(device)
     euler_angles = euler_angles.to(device)
     projective_filters = projective_filters.to(device)
+    # Move transform_matrix to device if it's not None
+    if transform_matrix is not None:
+        transform_matrix = transform_matrix.to(device)
 
     num_orientations = euler_angles.shape[0]
     num_defocus = defocus_values.shape[0]
@@ -510,6 +525,7 @@ def _core_match_template_single_gpu(
                         template_dft=template_dft,
                         rotation_matrices=rot_matrix,
                         projective_filters=projective_filters,
+                        transform_matrix=transform_matrix,
                     )
                 else:
                     cross_correlation = do_streamed_orientation_cross_correlate(
@@ -518,6 +534,7 @@ def _core_match_template_single_gpu(
                         rotation_matrices=rot_matrix,
                         projective_filters=projective_filters,
                         streams=streams,
+                        transform_matrix=transform_matrix,
                     )
 
                 # Update the tracked statistics

src/leopard_em/backend/core_refine_template.py

@@ -9,7 +9,7 @@
 import roma
 import torch
 import tqdm
-from torch_fourier_slice import extract_central_slices_rfft_3d
+from torch_fourier_slice import extract_central_slices_rfft_3d, transform_slice_2d
 
 from leopard_em.backend.cross_correlation import (
     do_batched_orientation_cross_correlate,
@@ -61,6 +61,7 @@ def core_refine_template(
     device: torch.device | list[torch.device],
     batch_size: int = 32,
     num_cuda_streams: int = 1,
+    transform_matrix: torch.Tensor | None = None,
 ) -> dict[str, torch.Tensor]:
     """Core function to refine orientations and defoci of a set of particles.
 
@@ -105,6 +106,9 @@ def core_refine_template(
         The number of cross-correlations to process in one batch, defaults to 32.
     num_cuda_streams : int, optional
         Number of CUDA streams to use for parallel processing. Defaults to 1.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
 
     Returns
     -------
@@ -136,6 +140,7 @@ def core_refine_template(
         batch_size=batch_size,
         devices=device,
         num_cuda_streams=num_cuda_streams,
+        transform_matrix=transform_matrix,
     )
 
     results = run_multiprocess_jobs(
@@ -229,6 +234,7 @@ def construct_multi_gpu_refine_template_kwargs(
     batch_size: int,
     devices: list[torch.device],
     num_cuda_streams: int,
+    transform_matrix: torch.Tensor | None = None,
 ) -> list[dict]:
     """Split particle stack between requested devices.
 
@@ -266,6 +272,9 @@ def construct_multi_gpu_refine_template_kwargs(
         List of devices to split across.
     num_cuda_streams : int
         Number of CUDA streams to use per device.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
 
     Returns
     -------
@@ -321,6 +330,7 @@ def construct_multi_gpu_refine_template_kwargs(
             "batch_size": batch_size,
             "num_cuda_streams": num_cuda_streams,
             "device": device,
+            "transform_matrix": transform_matrix,
         }
 
         kwargs_per_device.append(kwargs)
@@ -350,6 +360,7 @@ def _core_refine_template_single_gpu(
     batch_size: int,
     device: torch.device,
     num_cuda_streams: int = 1,
+    transform_matrix: torch.Tensor | None = None,
 ) -> None:
     """Run refine template on a subset of particles on a single GPU.
 
@@ -393,6 +404,9 @@ def _core_refine_template_single_gpu(
         Torch device to run this process on.
     num_cuda_streams : int, optional
         Number of CUDA streams to use for parallel processing. Defaults to 1.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
     """
     streams = [torch.cuda.Stream(device=device) for _ in range(num_cuda_streams)]
 
@@ -464,6 +478,7 @@ def _core_refine_template_single_gpu(
                 corr_std=corr_std[i],
                 projective_filter=projective_filters[i],
                 batch_size=batch_size,
+                transform_matrix=transform_matrix,
                 device_id=device_id,
             )
             refined_statistics.append(refined_stats)
@@ -565,6 +580,7 @@ def _core_refine_template_single_thread(
     projective_filter: torch.Tensor,
     batch_size: int = 32,
     device_id: int = 0,
+    transform_matrix: torch.Tensor | None = None,
 ) -> dict[str, float | int]:
     """Run the single-threaded core refine template function.
 
@@ -605,6 +621,9 @@ def _core_refine_template_single_thread(
         The number of orientations to cross-correlate at once. Default is 32.
     device_id : int, optional
         The ID of the device/process. Default is 0.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
 
     Returns
     -------
@@ -693,13 +712,15 @@ def _core_refine_template_single_thread(
                 template_dft=template_dft,
                 rotation_matrices=rot_matrix_batch,
                 projective_filters=combined_projective_filter,
+                transform_matrix=transform_matrix,
             )
         else:
             cross_correlation = do_batched_orientation_cross_correlate_cpu(
                 image_dft=particle_image_dft,
                 template_dft=template_dft,
                 rotation_matrices=rot_matrix_batch,
                 projective_filters=combined_projective_filter,
+                transform_matrix=transform_matrix,
             )
 
         cross_correlation = cross_correlation[..., :crop_h, :crop_w]  # valid crop
@@ -765,6 +786,7 @@ def cross_correlate_particle_stack(
     projective_filters: torch.Tensor,  # (N, h, w)
     mode: Literal["valid", "same"] = "valid",
     batch_size: int = 1024,
+    transform_matrix: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Cross-correlate a stack of particle images against a template.
 
@@ -793,6 +815,9 @@ def cross_correlate_particle_stack(
         The number of particle images to cross-correlate at once. Default is 1024.
         Larger sizes will consume more memory. If -1, then the entire stack will be
         cross-correlated at once.
+    transform_matrix : torch.Tensor | None, optional
+        Anisotropic magnification transform matrix of shape (2, 2). If None,
+        no magnification transform is applied. Default is None.
 
     Returns
     -------
@@ -839,9 +864,19 @@ def cross_correlate_particle_stack(
         # Extract the Fourier slice and apply the projective filters
         fourier_slice = extract_central_slices_rfft_3d(
             volume_rfft=template_dft,
-            image_shape=(template_h,) * 3,
             rotation_matrices=batch_rotation_matrices,
         )
+        # Apply anisotropic magnification transform if provided
+        # pylint: disable=duplicate-code
+        if transform_matrix is not None:
+            rfft_shape = (template_h, template_w)
+            stack_shape = (batch_rotation_matrices.shape[0],)
+            fourier_slice = transform_slice_2d(
+                projection_image_dfts=fourier_slice,
+                rfft_shape=rfft_shape,
+                stack_shape=stack_shape,
+                transform_matrix=transform_matrix,
+            )
         fourier_slice = torch.fft.ifftshift(fourier_slice, dim=(-2,))
         fourier_slice[..., 0, 0] = 0 + 0j  # zero out the DC component (mean zero)
         fourier_slice *= -1  # flip contrast