Class Image line projection method

src/aspire/image/image.py

@@ -9,7 +9,7 @@
 from scipy.linalg import lstsq
 
 import aspire.volume
-from aspire.nufft import anufft
+from aspire.nufft import anufft, nufft
 from aspire.numeric import fft, xp
 from aspire.utils import FourierRingCorrelation, anorm, crop_pad_2d, grid_2d
 from aspire.volume import SymmetryGroup
@@ -186,6 +186,42 @@ def __init__(self, data, dtype=None):
         self.__array_interface__ = self._data.__array_interface__
         self.__array__ = self._data
 
+    def project(self, angles):
+        """
+        Computes the Radon Transform on an Image Stack using
+        Non-Uniform Fast Fourier Transforms. This method projects the
+        Image stack along different angles and returns the Radon
+        Transform.
+
+        :param angles: A 1-D Numpy Array of angles in Radians.
+            This is used to compute the Radon Transform at different angles.
+        :return: Radon transform of the Image Stack.
+        :rtype: Ndarray (stack size, number of angles, image resolution)
+        """
+        # number of points to sample on radial line in polar grid
+        n_points = self.resolution
+        original_stack = self.stack_shape
+
+        # 2-D grid
+        radial_idx = np.fft.rfftfreq(n_points) * np.pi * 2
+        n_real_points = len(radial_idx)
+        n_angles = len(angles)
+
+        pts = np.empty((2, n_angles, n_real_points), dtype=self.dtype)
+        pts[0] = radial_idx[np.newaxis, :] * np.sin(angles)[:, np.newaxis]
+        pts[1] = radial_idx[np.newaxis, :] * np.cos(angles)[:, np.newaxis]
+        pts = pts.reshape(2, n_real_points * n_angles)
+
+        # compute the polar nufft (NUFFT)
+        image_ft = nufft(self.stack_reshape(-1)._data, pts).reshape(
+            self.n_images, n_angles, n_real_points
+        )
+
+        # Radon transform, output: (stack size, angles, points)
+        image_rt = np.fft.fftshift(np.fft.irfft(image_ft, n=n_points, axis=-1), axes=-1)
+        image_rt = image_rt.reshape(*original_stack, n_angles, n_points)
+        return image_rt
+
     @property
     def res(self):
         warn(

tests/test_sinogram.py

@@ -0,0 +1,138 @@
+import numpy as np
+import pytest
+from skimage import data
+from skimage.transform import radon
+
+from aspire.image import Image
+from aspire.utils import grid_2d
+
+# Relative tolerance comparing line projections to scikit
+# The same tolerance will be used in all scikit comparisons
+SK_TOL = 0.005
+
+IMG_SIZES = [
+    511,
+    512,
+]
+
+DTYPES = [
+    np.float32,
+    np.float64,
+]
+
+ANGLES = [
+    1,
+    50,
+    pytest.param(90, marks=pytest.mark.expensive),
+    pytest.param(117, marks=pytest.mark.expensive),
+    pytest.param(180, marks=pytest.mark.expensive),
+    pytest.param(360, marks=pytest.mark.expensive),
+]
+
+
+@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}", scope="module")
+def dtype(request):
+    """
+    Dtypes for image.
+    """
+    return request.param
+
+
+@pytest.fixture(params=IMG_SIZES, ids=lambda x: f"px={x}", scope="module")
+def img_size(request):
+    """
+    Image size.
+    """
+    return request.param
+
+
+@pytest.fixture(params=ANGLES, ids=lambda x: f"n_angles={x}", scope="module")
+def num_ang(request):
+    """
+    Number of angles in radon transform.
+    """
+    return request.param
+
+
+@pytest.fixture
+def masked_image(dtype, img_size):
+    """
+    Creates a masked image fixture using camera data from Scikit-Image.
+    """
+    g = grid_2d(img_size, normalized=True, shifted=True)
+    mask = g["r"] < 1
+
+    image = data.camera().astype(dtype)
+    image = image[:img_size, :img_size]
+    return Image(image * mask)
+
+
+# Image.project and compare results to skimage.radon
+def test_image_project(masked_image, num_ang):
+    """
+    Test Image.project on a single stack of images. Compares project method output with skimage project.
+    """
+    ny = masked_image.resolution
+    angles = np.linspace(0, 360, num_ang, endpoint=False)
+    rads = angles / 180 * np.pi
+    s = masked_image.project(rads)
+    assert s.shape == (1, len(angles), ny)
+
+    # sci-kit image `radon` reference
+    #
+    # Note, Image.project's angles are wrt projection line (ie
+    # grid), while sk's radon are wrt the image. To correspond the
+    # rotations are inverted.  This was the convention prefered by
+    # the original author of this method.
+    #
+    # Note, transpose sk output to match (angles, points)
+    reference_sinogram = radon(masked_image._data[0], theta=angles[::-1]).T
+    assert reference_sinogram.shape == (len(angles), ny), "Incorrect Shape"
+
+    # compare project method on ski-image reference
+    nrms = np.sqrt(np.mean((s[0] - reference_sinogram) ** 2, axis=-1)) / np.linalg.norm(
+        reference_sinogram, axis=-1
+    )
+
+    np.testing.assert_array_less(nrms, SK_TOL, "Error in image projections.")
+
+
+def test_multidim(num_ang):
+    """
+    Test Image.project on stacks of images. Extension of test_image_project but for multi-dimensional stacks.
+    """
+
+    L = 512  # pixels
+    n = 3
+    m = 2
+
+    # Generate a mask
+    g = grid_2d(L, normalized=True, shifted=True)
+    mask = g["r"] < 1
+
+    # Generate images
+    imgs = Image(np.random.random((m, n, L, L))) * mask
+
+    # Generate line project angles
+    angles = np.linspace(0, 360, num_ang, endpoint=False)
+    rads = angles / 180.0 * np.pi
+    s = imgs.project(rads)
+
+    # Compare
+    reference_sinograms = np.empty((m, n, num_ang, L))
+    for i in range(m):
+        for j in range(n):
+            img = imgs[i, j]
+            # Compute the singleton case, and compare with stack.
+            single_sinogram = img.project(rads)
+
+            # These should be allclose up to determinism in the FFT and NUFFT.
+            np.testing.assert_allclose(s[i, j : j + 1], single_sinogram)
+
+            # Next individually compute sk's radon transform for each image.
+            reference_sinograms[i, j] = radon(img._data[0], theta=angles[::-1]).T
+
+    _nrms = np.sqrt(np.mean((s - reference_sinograms) ** 2, axis=-1)) / np.linalg.norm(
+        reference_sinograms, axis=-1
+    )
+    np.testing.assert_array_less(_nrms, SK_TOL, "Error in image projections.")