Refactor project

pyproject.toml

@@ -46,7 +46,7 @@ dependencies = [
 # "extras" (e.g. for `pip install .[test]`)
 [project.optional-dependencies]
 # add dependencies used for testing here
-test = ["pytest", "pytest-cov"]
+test = ["pytest", "pytest-cov", "mrcfile"]
 # add anything else you like to have in your dev environment here
 dev = [
     "ipython",

src/torch_fourier_slice/grids/central_slice_grid.py

@@ -1,10 +1,11 @@
 import einops
 import torch
-
+from functools import lru_cache
 from .fftfreq_grid import _construct_fftfreq_grid_2d
 from ..dft_utils import rfft_shape, fftshift_2d
 
 
+@lru_cache(1) #Alternativelly, we can have an argument that needs to be propagated to extract_central_slices_rfft_3d
 def central_slice_fftfreq_grid(
     volume_shape: tuple[int, int, int],
     rfft: bool,

src/torch_fourier_slice/project.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn.functional as F
+from typing import List
 
 from .grids import fftfreq_grid
 from .slice_extraction import extract_central_slices_rfft_3d
@@ -13,6 +14,51 @@ def project_3d_to_2d(
 ) -> torch.Tensor:
     """Project a cubic volume by sampling a central slice through its DFT.
 
+    Parameters
+    ----------
+    volume: torch.Tensor
+        `(d, d, d)` volume.
+    rotation_matrices: torch.Tensor
+        `(..., 3, 3)` array of rotation matrices for insert of `images`.
+        Rotation matrices left-multiply column vectors containing xyz coordinates.
+    pad: bool
+        Whether to pad the volume 2x with zeros to increase sampling rate in the DFT.
+    fftfreq_max: float | None
+        Maximum frequency (cycles per pixel) included in the projection.
+
+    Returns
+    -------
+    projections: torch.Tensor
+        `(..., d, d)` array of projection images.
+    """
+
+    dft, pad_length = _compute_dft3d_for_project(volume, pad)
+
+    # make projections by taking central slices
+    projections = extract_central_slices_rfft_3d(
+        volume_rfft=dft,
+        image_shape=volume.shape,
+        rotation_matrices=rotation_matrices,
+        fftfreq_max=fftfreq_max
+    )  # (..., h, w) rfft stack
+
+    # transform back to real space
+    projections = torch.fft.ifftshift(projections, dim=(-2,))  # ifftshift of 2D rfft
+    projections = torch.fft.irfftn(projections, dim=(-2, -1))
+    projections = torch.fft.ifftshift(projections, dim=(-2, -1))  # recenter 2D image in real space
+
+    # unpad if required
+    if pad is True:
+        projections = projections[..., pad_length:-pad_length, pad_length:-pad_length]
+    return torch.real(projections)
+
+
+def _compute_dft3d_for_project(    
+    volume: torch.Tensor,
+    pad: bool = True,
+) -> (torch.Tensor, int):
+    """Project a cubic volume by sampling a central slice through its DFT.
+
     Parameters
     ----------
     volume: torch.Tensor
@@ -31,6 +77,7 @@ def project_3d_to_2d(
         `(..., d, d)` array of projection images.
     """
     # padding
+    pad_length = 0
     if pad is True:
         pad_length = volume.shape[-1] // 2
         volume = F.pad(volume, pad=[pad_length] * 6, mode='constant', value=0)
@@ -50,20 +97,8 @@ def project_3d_to_2d(
     dft = torch.fft.rfftn(dft, dim=(-3, -2, -1))
     dft = torch.fft.fftshift(dft, dim=(-3, -2,))  # actual fftshift of 3D rfft
 
-    # make projections by taking central slices
-    projections = extract_central_slices_rfft_3d(
-        volume_rfft=dft,
-        image_shape=volume.shape,
-        rotation_matrices=rotation_matrices,
-        fftfreq_max=fftfreq_max
-    )  # (..., h, w) rfft stack
 
-    # transform back to real space
-    projections = torch.fft.ifftshift(projections, dim=(-2,))  # ifftshift of 2D rfft
-    projections = torch.fft.irfftn(projections, dim=(-2, -1))
-    projections = torch.fft.ifftshift(projections, dim=(-2, -1))  # recenter 2D image in real space
+    return dft, pad_length
+
+
 
-    # unpad if required
-    if pad is True:
-        projections = projections[..., pad_length:-pad_length, pad_length:-pad_length]
-    return torch.real(projections)

tests/test_torch_fourier_slice.py

@@ -2,3 +2,80 @@
 
 def test_something():
     pass
+
+
+
+def test_central_slice():
+
+    import os
+    import mrcfile
+    import requests
+    import tempfile
+    import torch
+    from scipy.spatial.transform import Rotation as R
+    from torch_fourier_slice import project_3d_to_2d, backproject_2d_to_3d
+
+    tmpdir= tempfile.gettempdir()
+    fname = os.path.join(tmpdir, "emd_17129.map.gz")
+    if not os.path.isfile(fname):
+        response = requests.get("https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-17129/map/emd_17129.map.gz", stream=True)
+        if response.status_code == 200:
+            # Open a file in write-binary mode
+
+            with open(fname, "wb") as f:
+                # Write the content of the response to the file in chunks
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            print("Download completed successfully!")
+        else:
+            print(f"Failed to download the file. Status code: {response.status_code}")
+            raise RuntimeError()
+    volume = torch.as_tensor(mrcfile.read(fname), dtype=torch.float32)
+
+    rot_mats =  torch.as_tensor(R.from_euler("ZYZ", [[0,0,0],[0,0,90], [0,90,0]], degrees=True
+                                            ).as_matrix(), dtype=torch.float32)
+    projs = project_3d_to_2d(
+                    volume=volume,
+                    rotation_matrices=rot_mats,
+                    pad=False,
+                    fftfreq_max=None)
+
+
+    affine_mats = torch.zeros(rot_mats.shape[0], 3, 4)
+    affine_mats[:,:3,:3] = rot_mats
+#    affine_mats[:,:3,-1] += 1./volume.shape[-1] #TODO: It seems that the projections may be off by half a pixel
+
+
+    volume = volume[None,None,...].expand(rot_mats.shape[0], -1, -1, -1, -1)
+    rot_vols = torch.nn.functional.grid_sample(volume, torch.nn.functional.affine_grid(affine_mats, size=volume.shape), align_corners=False)
+    projs_sum = rot_vols.sum(2).squeeze(1)
+
+    for i in range(projs.shape[0]):
+        assert torch.isclose(projs[i], projs_sum[i], atol=1e-1).all(), f"Error, disagreement in projections {i}"
+
+    diff = torch.abs(projs - projs_sum)
+    print(diff.mean(-1).mean(-1))
+
+#    from matplotlib import pyplot as plt
+#    f, axes = plt.subplots(3,3)
+#    axes[0,0].imshow(projs[0])
+#    axes[0,1].imshow(projs[1])
+#    axes[0,2].imshow(projs[2])
+#    
+#    axes[1,0].imshow(diff[0])
+#    axes[1,1].imshow(diff[1])
+#    axes[1,2].imshow(diff[2])
+#    
+#    axes[2,0].imshow(projs_sum[0])
+#    axes[2,1].imshow(projs_sum[1])
+#    axes[2,2].imshow(projs_sum[2])
+#    plt.show()
+
+"""
+ PYTHONPATH=../torch-image-lerp/src:$PYTHONPATH  python -m src.torch_fourier_slice.project
+"""
+if __name__ == "__main__":
+    test_central_slice()
+    """
+PYTHONPATH=../torch-image-lerp/src:src/:$PYTHONPATH python tests/test_torch_fourier_slice.py
+    """