Merge pull request #111 from Lucaslab-Berkeley/jd_update_extract

fix: Update breaking change in torch-fourier-slice

Author: jdickerson95

pyproject.toml

@@ -48,7 +48,7 @@ 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",
@@ -153,7 +153,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/backend/core_refine_template.py

@@ -839,7 +839,6 @@ 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,
         )
         fourier_slice = torch.fft.ifftshift(fourier_slice, dim=(-2,))

src/leopard_em/backend/cross_correlation.py

@@ -71,7 +71,6 @@ def do_streamed_orientation_cross_correlate(
     # Do a batched Fourier slice extraction for all the orientations at once.
     fourier_slices = extract_central_slices_rfft_3d(
         volume_rfft=template_dft,
-        image_shape=(projection_shape_real[0],) * 3,
         rotation_matrices=rotation_matrices,
     )
     fourier_slices = torch.fft.ifftshift(fourier_slices, dim=(-2,))
@@ -202,7 +201,6 @@ def do_batched_orientation_cross_correlate(
     # Extract central slice(s) from the template volume
     fourier_slice = extract_central_slices_rfft_3d(
         volume_rfft=template_dft,
-        image_shape=(projection_shape_real[0],) * 3,  # NOTE: requires cubic template
         rotation_matrices=rotation_matrices,
     )
     fourier_slice = torch.fft.ifftshift(fourier_slice, dim=(-2,))
@@ -279,7 +277,6 @@ def do_batched_orientation_cross_correlate_cpu(
     # Extract central slice(s) from the template volume
     fourier_slice = extract_central_slices_rfft_3d(
         volume_rfft=template_dft,
-        image_shape=(projection_shape_real[0],) * 3,  # NOTE: requires cubic template
         rotation_matrices=rotation_matrices,
     )
     fourier_slice = torch.fft.ifftshift(fourier_slice, dim=(-2,))

src/leopard_em/utils/fourier_slice.py

@@ -71,7 +71,6 @@ def get_rfft_slices_from_volume(
         The Fourier slices of the volume.
 
     """
-    shape = volume.shape
     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
@@ -82,7 +81,6 @@ def get_rfft_slices_from_volume(
     # Use torch_fourier_slice to take the Fourier slice
     fourier_slices = extract_central_slices_rfft_3d(
         volume_rfft=volume_rfft,
-        image_shape=shape,
         rotation_matrices=rot_matrix,
     )