Merge pull request #68 from mgiammar/main

Sync minor test and linting updates upstream

Author: mgiammar

src/leopard_em/utils/fourier_slice.py

@@ -3,7 +3,6 @@
 import roma
 import torch
 from torch_fourier_slice import extract_central_slices_rfft_3d
-from torch_fourier_slice._dft_utils import _fftshift_3d, _ifftshift_2d
 from torch_grid_utils import fftfreq_grid
 
 
@@ -35,11 +34,12 @@ def _rfft_slices_to_real_projections(
     torch.Tensor
         The real-space projections.
     """
-    fourier_slices = _ifftshift_2d(fourier_slices, rfft=True)
+    # pylint: disable=not-callable
+    fourier_slices = torch.fft.fftshift(fourier_slices, dim=(-2,))
     # pylint: disable=not-callable
     projections = torch.fft.irfftn(fourier_slices, dim=(-2, -1))
-    projections = _ifftshift_2d(projections, rfft=False)
-
+    # pylint: disable=not-callable
+    projections = torch.fft.ifftshift(projections, dim=(-2, -1))
     return projections
 
 
@@ -72,10 +72,9 @@ def get_rfft_slices_from_volume(
 
     """
     shape = volume.shape
-    volume_rfft = _fftshift_3d(volume, rfft=False)
-    # pylint: disable=not-callable
-    volume_rfft = torch.fft.fftn(volume_rfft, dim=(-3, -2, -1))
-    volume_rfft = _fftshift_3d(volume_rfft, rfft=True)
+    volume_rfft = torch.fft.fftshift(volume, dim=(-3, -2, -1))  # pylint: disable=not-callable
+    volume_rfft = torch.fft.fftn(volume_rfft, dim=(-3, -2, -1))  # pylint: disable=not-callable
+    volume_rfft = torch.fft.fftshift(volume_rfft, dim=(-3, -2))  # pylint: disable=not-callable
 
     # Use roma to keep angles on same device
     rot_matrix = roma.euler_to_rotmat("ZYZ", (phi, theta, psi), degrees=degrees)

tests/self_consistency/test_backend_cross_correlate.py

@@ -22,6 +22,7 @@
 import torch
 from scipy.ndimage import gaussian_filter
 from torch_fourier_filter.ctf import calculate_ctf_2d
+from torch_fourier_filter.envelopes import b_envelope
 
 from leopard_em.backend.cross_correlation import (
     do_batched_orientation_cross_correlate,
@@ -70,7 +71,7 @@ def sample_input_data() -> dict[str, torch.Tensor]:
     template_fft = torch.fft.fftshift(template_fft, dim=(0, 1))
 
     # Generate a set of projective filters (CTFs) for the template
-    defocus_values = torch.linspace(500, 1500, NUM_DEFOCUS_VALUES)
+    defocus_values = torch.linspace(2000, 4000, NUM_DEFOCUS_VALUES)
     pixel_sizes = torch.linspace(0.8, 1.2, NUM_PIXEL_SIZES)
     cs_values = get_cs_range(1.0, pixel_sizes, 2.7)
 
@@ -83,7 +84,6 @@ def sample_input_data() -> dict[str, torch.Tensor]:
             voltage=300,  # 300 kV
             spherical_aberration=cs_val,
             amplitude_contrast=0.07,
-            b_factor=100.0,
             phase_shift=0.0,
             pixel_size=1.0,
             image_shape=template.shape[-2:],
@@ -92,7 +92,17 @@ def sample_input_data() -> dict[str, torch.Tensor]:
         )
         ctf_list.append(tmp)
 
+    # Apply a b-factor envelope to the CTFs
+    b_factor = 100.0  # arbitrary value
+    b_envelope_values = b_envelope(
+        B=b_factor,
+        image_shape=template.shape[-2:],
+        pixel_size=1.0,
+        device="cuda",
+    )
+
     projective_filters = torch.stack(ctf_list, dim=0).to(device="cuda")
+    projective_filters *= b_envelope_values[None, None]
 
     return {
         "image_dft": image_fft,

tests/utils/test_crop_extraction.py

@@ -61,23 +61,28 @@ def test_get_cropped_image_regions_numpy_fixed_positions():
 def test_get_cropped_image_regions_numpy_random_nonoverlapping():
     """Random non-overlapping pos for _get_cropped_image_regions_numpy function."""
     box_size = (5, 5)
-    num_patches = 10
+    num_patches = 32
     image_size = (256, 256)
     test_patch = get_test_patch(box_size)
     image = np.zeros(image_size, dtype=np.float32)
 
     # Generate non-overlapping positions
     positions = []
     for _ in range(num_patches):
-        while True:
+        total_failures = 0
+        while total_failures < 100:
             y = np.random.randint(0, image_size[0] - box_size[0] + 1)
             x = np.random.randint(0, image_size[1] - box_size[1] + 1)
-            if all(
-                not (y <= py < y + box_size[0] and x <= px < x + box_size[1])
+            # Check if new position overlaps with any existing position
+            overlap = any(
+                abs(y - py) < box_size[0] and abs(x - px) < box_size[1]
                 for py, px in positions
-            ):
+            )
+
+            if not overlap:
                 positions.append((y, x))
                 break
+            total_failures += 1
 
     pos_y, pos_x = zip(*positions)