Fix sinogram stack handle

Author: neon60

docs/source/examples/raven_filter_example.py

@@ -5,7 +5,7 @@
 import pyfftw
 import pyfftw.interfaces.numpy_fft as fft
 import httomolibgpu
-from httomolibgpu.misc.raven_filter import raven_filter
+from httomolibgpu.prep.stripe import raven_filter
 
 import matplotlib.pyplot as plt
 import time
@@ -23,63 +23,70 @@
 
 sino_shape = sinogram.shape
 
-print("The shape of the sinogram is {}".format(cp.shape(sinogram)))
+sinogram_stack = cp.stack([sinogram] * 20, axis=1)
+
+print("The shape of the sinogram stack is {}".format(cp.shape(sinogram_stack)))
 
 # Parameters
 v0 = 2
 n = 4
 u0 = 20
 
 # Make a numpy copy
-sinogram_padded = np.pad(sinogram.get(), 20, "edge")
+sinogram_padded = np.pad(sinogram_stack.get(), [(20, 20), (0, 0), (20, 20)], "edge")
 
 start_time = time.time()
 # GPU filter
-sinogram_gpu_filter = raven_filter(sinogram, u0, n, v0)
+sinogram_gpu_filter = raven_filter(sinogram_stack, u0, n, v0)
 print("--- %s seconds ---" % (time.time() - start_time))
 
 start_time = time.time()
 
 # Size
-width1 = sino_shape[1] + 2 * 20
-height1 = sino_shape[0] + 2 * 20
+height, images, width = sinogram_padded.shape
 
 # Generate filter function
-centerx = np.ceil(width1 / 2.0) - 1.0
-centery = np.int16(np.ceil(height1 / 2.0) - 1)
+centerx = np.ceil(width / 2.0) - 1.0
+centery = np.int16(np.ceil(height / 2.0) - 1)
 row1 = centery - v0
 row2 = centery + v0 + 1
-listx = np.arange(width1) - centerx
+listx = np.arange(width) - centerx
 filtershape = 1.0 / (1.0 + np.power(listx / u0, 2 * n))
 filtershapepad2d = np.zeros((row2 - row1, filtershape.size))
 filtershapepad2d[:] = np.float64(filtershape)
 filtercomplex = filtershapepad2d + filtershapepad2d * 1j
 
 # Generate filter objects
-a = pyfftw.empty_aligned((height1, width1), dtype='complex128', n=16)
-b = pyfftw.empty_aligned((height1, width1), dtype='complex128', n=16)
-c = pyfftw.empty_aligned((height1, width1), dtype='complex128', n=16)
-d = pyfftw.empty_aligned((height1, width1), dtype='complex128', n=16)
-fft_object  = pyfftw.FFTW(a, b, axes=(0, 1))
-ifft_object = pyfftw.FFTW(c, d, axes=(0, 1), direction='FFTW_BACKWARD')
+a = pyfftw.empty_aligned((height, images, width), dtype='complex128', n=16)
+b = pyfftw.empty_aligned((height, images, width), dtype='complex128', n=16)
+c = pyfftw.empty_aligned((height, images, width), dtype='complex128', n=16)
+d = pyfftw.empty_aligned((height, images, width), dtype='complex128', n=16)
+fft_object  = pyfftw.FFTW(a, b, axes=(0, 2))
+ifft_object = pyfftw.FFTW(c, d, axes=(0, 2), direction='FFTW_BACKWARD')
+
+sino = fft.fftshift(fft_object(sinogram_padded), axes=(0, 2))
+for m in range(sino.shape[1]):
+    sino[row1:row2, m] = sino[row1:row2, m] * filtercomplex
+sino = ifft_object(fft.ifftshift(sino, axes=(0, 2)))
 
-sino = fft.fftshift(fft_object(sinogram_padded))
-sino[row1:row2] = sino[row1:row2] * filtercomplex
-sino = ifft_object(fft.ifftshift(sino))
+# Remove padding
+sino = sino[20:height-20, :, 20:width-20]
 
 print("--- %s seconds ---" % (time.time() - start_time))
 
 #subplot(r,c) provide the no. of rows and columns
 f, axarr = plt.subplots(2,2) 
 
+sino_index = 10
+
 # use the created array to output your multiple images. In this case I have stacked 4 images vertically
-axarr[0, 0].imshow(sinogram_padded)
+axarr[0, 0].imshow(sinogram_stack.get()[:, sino_index, :])
 axarr[0, 0].set_title('Original sinogram')
-axarr[0, 1].imshow(sinogram_padded - sinogram_gpu_filter.get().real)
+axarr[0, 1].imshow(sinogram_stack.get()[:, sino_index, :] - sinogram_gpu_filter.get().real[:, sino_index, :])
 axarr[0, 1].set_title('Difference of original and GPU filtered')
-axarr[1, 0].imshow(sinogram_padded - sino.real)
+axarr[1, 0].imshow(sinogram_stack.get()[:, sino_index, :] - sino.real[:, sino_index, :])
 axarr[1, 0].set_title('Difference of original and CPU filtered')
-axarr[1, 1].imshow(sinogram_gpu_filter.get().real - sino.real)
+axarr[1, 1].imshow(sinogram_gpu_filter.get().real[:, sino_index, :] - sino.real[:, sino_index, :])
 axarr[1, 1].set_title('Difference of GPU and CPU filtered')
 
 plt.show()

httomolibgpu/cuda_kernels/raven_filter.cu

@@ -1,21 +1,24 @@
 #include <cupy/complex.cuh>
 
 extern "C" __global__ void 
-raven_filter(complex<float> *input, complex<float> *output, int width1, int height1, int u0, int n, int v0) {
+raven_filter(
+  complex<float> *input,
+  complex<float> *output,
+  int width, int images, int height, 
+  int u0, int n, int v0) {
 
-  int centerx = width1 / 2;
-  int centery = height1 / 2;
+  const int px = threadIdx.x + blockIdx.x * blockDim.x;
+  const int py = threadIdx.y + blockIdx.y * blockDim.y;
+  const int pz = threadIdx.z + blockIdx.z * blockDim.z;
 
-  int px = threadIdx.x + blockIdx.x * blockDim.x;
-  int py = threadIdx.y + blockIdx.y * blockDim.y;
-
-  if (px >= width1)
-    return;
-  if (py >= height1)
+  if (px >= width || py >= images || pz >= height)
     return;
 
-  complex<float> value = input[py * width1 + px];
-  if( py >= (centery - v0) && py < (centery + v0 + 1) ) {
+  int centerx = width / 2;
+  int centerz = height / 2;
+
+  complex<float> value = input[pz * width * images + py * width + px];
+  if( pz >= (centerz - v0) && pz < (centerz + v0 + 1) ) {
 
     // +1 needed to match with CPU implementation
     float base = float(px - centerx + 1) / u0;
@@ -28,10 +31,10 @@ raven_filter(complex<float> *input, complex<float> *output, int width1, int heig
   }
 
   // ifftshifting positions
-  int xshift = (width1 + 1) / 2;
-  int yshift = (height1 + 1) / 2;
-  int outX = (px + xshift) % width1;
-  int outY = (py + yshift) % height1;
+  int xshift = (width + 1) / 2;
+  int zshift = (height + 1) / 2;
+  int outX = (px + xshift) % width;
+  int outZ = (pz + zshift) % height;
 
-  output[outY * width1 + outX] = value;
+  output[outZ * width * images + py * width + outX] = value;
 }

httomolibgpu/misc/raven_filter.py

@@ -30,7 +30,7 @@
 
 if cupy_run:
     from cupyx.scipy.ndimage import median_filter, binary_dilation, uniform_filter1d
-    from cupyx.scipy.fft import fft2, ifft2, fftshift
+    from cupyx.scipy.fft import fft2, ifft2, fftshift, ifftshift
     from httomolibgpu.cuda_kernels import load_cuda_module
 else:
     median_filter = Mock()
@@ -39,13 +39,15 @@
     fft2 = Mock()
     ifft2 = Mock()
     fftshift = Mock()
+    ifftshift = Mock()
 
 from typing import Union
 
 __all__ = [
     "remove_stripe_based_sorting",
     "remove_stripe_ti",
     "remove_all_stripe",
+    "raven_filter",
 ]
 
 
@@ -376,36 +378,36 @@ def raven_filter(
     Raven filter
     """
 
-    # Padding of the data
-    padded_data = cp.pad(sinogram, ((0, 0), (pad_y, pad_y), (pad_x, pad_x)), mode=pad_method)
-    # padded_data = cp.pad(sinogram, ((pad_y, pad_y), (pad_x, pad_x)), mode=pad_method)
+    # Padding of the sinogram
+    sinogram = cp.pad(sinogram, ((pad_y, pad_y), (0, 0), (pad_x, pad_x)), mode=pad_method)
 
-    # FFT and shift of data
-    fft_data = fft2(padded_data, axes=(-2, -1), overwrite_x=True)
-    fft_data_shifted = fftshift(fft_data)
+    # FFT and shift of sinogram
+    fft_data = fft2(sinogram, axes=(0, 2), overwrite_x=True)
+    fft_data_shifted = fftshift(fft_data, axes=(0, 2))
 
     # Setup various values for the filter
-    _, height, width = sinogram.shape
-
-    height1 = height + 2 * pad_y
-    width1 = width + 2 * pad_x
+    height, images, width = sinogram.shape
 
     # setting grid/block parameters
     block_x = 128
     block_dims = (block_x, 1, 1)
-    grid_x = (width1 + block_x - 1) // block_x
-    grid_y = height1
-    grid_dims = (grid_x, grid_y, 1)
-    params = (fft_data_shifted, fft_data, width1, height1, uvalue, nvalue, vvalue)
+    grid_x = (width + block_x - 1) // block_x
+    grid_y = images
+    grid_z = height
+    grid_dims = (grid_x, grid_y, grid_z)
+    params = (fft_data_shifted, fft_data, width, images, height, uvalue, nvalue, vvalue)
 
     raven_module = load_cuda_module("raven_filter")
     raven_filt = raven_module.get_function("raven_filter")
 
     raven_filt(grid_dims, block_dims, params)
 
     # raven_filt already doing ifftshifting
-    # fft_data = ifftshift(fft_data_shifted)
-    sinogram = ifft2(fft_data, axes=(-2, -1), overwrite_x=True)
+    # fft_data = ifftshift(fft_data_shifted, axes=(0, 2))
+    sinogram = ifft2(fft_data, axes=(0, 2), overwrite_x=True)
+
+    # Removing padding
+    sinogram = sinogram[pad_y:height-pad_y, :, pad_x:width-pad_x].real
 
     return sinogram
 

httomolibgpu/prep/stripe.py

@@ -57,7 +57,8 @@ dev = [
   "toml",
   "imageio",
   "h5py",
-  "pre-commit"
+  "pre-commit",
+  "pyfftw"
 ]