Added density compensation and gpunufft stuff

examples/GPU_Examples/NonCartesian_Pipe_DensityCompensation.py

@@ -0,0 +1,95 @@
+"""
+Neuroimaging non-cartesian reconstruction
+=========================================
+
+Author: Chaithya G R
+
+In this tutorial we will reconstruct an MR Image directly with density
+compensation
+
+Import neuroimaging data
+------------------------
+
+We use the toy datasets available in pysap, more specifically a 2D brain slice
+and the radial acquisition scheme (non-cartesian).
+"""
+
+# %%
+# Package import
+from mri.operators import NonCartesianFFT
+from mri.operators.fourier.utils import estimate_density_compensation, \
+    convert_locations_to_mask, gridded_inverse_fourier_transform_nd
+from pysap.data import get_sample_data
+
+# Third party import
+from modopt.math.metrics import ssim
+import numpy as np
+import matplotlib.pyplot as plt
+
+# %%
+# Loading input data
+image = get_sample_data('2d-mri').data.astype(np.complex64)
+
+# Obtain MRI non-cartesian mask and estimate the density compensation
+kspace_loc = get_sample_data("mri-radial-samples").data
+density_comp = estimate_density_compensation(kspace_loc, image.shape, 'pipe', backend='gpunufft')
+
+# %%
+# View Input
+plt.subplot(1, 2, 1)
+plt.imshow(np.abs(image), cmap='gray')
+plt.title("MRI Data")
+plt.subplot(1, 2, 2)
+plt.imshow(convert_locations_to_mask(kspace_loc, image.shape), cmap='gray')
+plt.title("K-space Sampling Mask")
+plt.show()
+
+# %%
+# Generate the kspace
+# -------------------
+#
+# From the 2D brain slice and the acquisition mask, we retrospectively
+# undersample the k-space using a radial acquisition mask
+# We then reconstruct using adjoint with and without density compensation
+
+# Get the locations of the kspace samples and the associated observations
+fourier_op = NonCartesianFFT(
+    samples=kspace_loc,
+    shape=image.shape,
+)
+fourier_op_density_comp = NonCartesianFFT(
+    samples=kspace_loc,
+    shape=image.shape,
+    density_comp=density_comp
+)
+# Get the kspace data retrospectively. Note that this can be done with
+# `fourier_op_density_comp` as the forward operator is the same
+kspace_obs = fourier_op.op(image)
+
+# %%
+# Gridded solution
+grid_space = np.linspace(-0.5, 0.5, num=image.shape[0])
+grid2D = np.meshgrid(grid_space, grid_space)
+grid_soln = gridded_inverse_fourier_transform_nd(kspace_loc, kspace_obs,
+                                                 tuple(grid2D), 'linear')
+base_ssim = ssim(grid_soln, image, mask=np.abs(image)>np.mean(np.abs(image)))
+plt.imshow(np.abs(grid_soln), cmap='gray')
+plt.title('Gridded solution : SSIM = ' + str(np.around(base_ssim, 2)))
+plt.show()
+# %%
+# Simple adjoint
+# This preconditions k-space giving a result closer to inverse
+image_rec = fourier_op.adj_op(kspace_obs)
+recon_ssim = ssim(image_rec, image, mask=np.abs(image)>np.mean(np.abs(image)))
+plt.imshow(np.abs(image_rec), cmap='gray')
+plt.title('Simple NUFFT Adjoint : SSIM = ' + str(np.around(recon_ssim, 2)))
+plt.show()
+
+# %%
+# Density Compensation adjoint:
+# This preconditions k-space giving a result closer to inverse
+image_rec = fourier_op_density_comp.adj_op(kspace_obs)
+recon_ssim = ssim(image_rec, image, mask=np.abs(image)>np.mean(np.abs(image)))
+plt.imshow(np.abs(image_rec), cmap='gray')
+plt.title('Density Compensated Adjoint : SSIM = ' + str(np.around(recon_ssim, 2)))
+plt.show()

examples/GPU_Examples/NonCartesian_gpuNUFFT_SENSE_DensityCompensation.py

@@ -14,32 +14,43 @@
 and the radial acquisition scheme (non-cartesian).
 """
 
+# %%
 # Package import
-from mri.operators import NonCartesianFFT, WaveletUD2
-from mri.operators.utils import convert_locations_to_mask, \
-    gridded_inverse_fourier_transform_nd
+from mri.operators import NonCartesianFFT, WaveletN
+from mri.operators.utils import normalize_frequency_locations
 from mri.operators.fourier.utils import estimate_density_compensation
-from mri.reconstructors import SingleChannelReconstructor
+from mri.reconstructors import SelfCalibrationReconstructor
 from mri.reconstructors.utils.extract_sensitivity_maps import get_Smaps
-import pysap
 from pysap.data import get_sample_data
 
 # Third party import
 from modopt.math.metrics import ssim
 from modopt.opt.linear import Identity
 from modopt.opt.proximity import SparseThreshold
 import numpy as np
+import matplotlib.pyplot as plt
 
+# %%
 # Loading input data
-image = get_sample_data('3d-pmri')
+image = get_sample_data('3d-pmri').data.astype(np.complex64)
 cartesian = np.linalg.norm(image, axis=0)
 
 # Obtain MRI non-cartesian mask and estimate the density compensation
 radial_mask = get_sample_data("mri-radial-3d-samples")
-kspace_loc = radial_mask.data
-density_comp = estimate_density_compensation(kspace_loc, cartesian.shape)
+kspace_loc = normalize_frequency_locations(radial_mask.data)
+density_comp = estimate_density_compensation(kspace_loc, cartesian.shape, 'pipe', backend='gpunufft')
 
-#############################################################################
+# %%
+# View Input
+plt.subplot(1, 2, 1)
+plt.imshow(cartesian[..., 80], cmap='gray')
+plt.title("MRI Data")
+ax = plt.subplot(1, 2, 2, projection='3d')
+ax.scatter(*kspace_loc[::500].T, s=0.1, alpha=0.5)
+plt.title("K-space Sampling Mask")
+plt.show()
+
+# %%
 # Generate the kspace
 # -------------------
 #
@@ -54,7 +65,9 @@
     n_coils=image.shape[0],
     implementation='gpuNUFFT',
 )
-kspace_obs = fourier_op.op(image.data)
+kspace_obs = fourier_op.op(image)
+# %%
+# Obtrain the Sensitivity Maps
 Smaps, SOS = get_Smaps(
     k_space=kspace_obs,
     img_shape=fourier_op.shape,
@@ -66,6 +79,16 @@
     density_comp=density_comp,
     mode='NFFT',
 )
+# %%
+# View Input
+for i in range(9):
+    plt.subplot(3, 3, i+1)
+    plt.imshow(np.abs(Smaps[i][..., 80]), cmap='gray')
+plt.suptitle("Sensitivty Maps")
+plt.show()
+
+# %%
+# Density Compensation adjoint:
 fourier_op_sense_dc = NonCartesianFFT(
     samples=kspace_loc,
     shape=cartesian.shape,
@@ -74,13 +97,44 @@
     density_comp=density_comp,
     smaps=Smaps,
 )
-
-# Density Compensation SENSE adjoint:
 # This preconditions k-space giving a result closer to inverse
-image_rec1 = pysap.Image(data=np.abs(
-    fourier_op_sense_dc.adj_op(kspace_obs))
+image_rec = fourier_op_sense_dc.adj_op(kspace_obs)
+recon_ssim = ssim(image_rec, cartesian, mask=np.abs(image)>np.mean(np.abs(image)))
+plt.imshow(np.abs(image_rec)[..., 80], cmap='gray')
+plt.title('Density Compensated Adjoint : SSIM = ' + str(np.around(recon_ssim, 2)))
+plt.show()
+
+
+# %%
+# FISTA optimization
+# ------------------
+#
+# We now want to refine the zero order solution using a FISTA optimization.
+# The cost function is set to Proximity Cost + Gradient Cost
+
+# Setup the operators
+linear_op = WaveletN(
+    wavelet_name='sym8',
+    nb_scale=4,
+    dim=3,
+)
+regularizer_op = SparseThreshold(Identity(), 3 * 1e-9, thresh_type="soft")
+# Setup Reconstructor
+reconstructor = SelfCalibrationReconstructor(
+    fourier_op=fourier_op_sense_dc,
+    linear_op=linear_op,
+    regularizer_op=regularizer_op,
+    gradient_formulation='synthesis',
+    verbose=1,
+)
+# %%
+# Run the FISTA reconstruction and view results
+image_rec, costs, metrics = reconstructor.reconstruct(
+    kspace_data=kspace_obs,
+    optimization_alg='fista',
+    num_iterations=30,
 )
-# image_rec1.show()
-base_ssim = ssim(image_rec1, cartesian)
-print('The SSIM for simple Density '
-      'compensated SENSE Adjoint is : ' + str(base_ssim))
+recon_ssim = ssim(image_rec, cartesian)
+plt.imshow(np.abs(image_rec)[..., 80], cmap='gray')
+plt.title('Iterative Reconstruction : SSIM = ' + str(np.around(recon_ssim, 2)))
+plt.show()

examples/NonCartesian_DensityCompensation.py

@@ -0,0 +1,95 @@
+"""
+Neuroimaging non-cartesian reconstruction
+=========================================
+
+Author: Chaithya G R
+
+In this tutorial we will reconstruct an MR Image directly with density
+compensation
+
+Import neuroimaging data
+------------------------
+
+We use the toy datasets available in pysap, more specifically a 2D brain slice
+and the radial acquisition scheme (non-cartesian).
+"""
+
+# %%
+# Package import
+from mri.operators import NonCartesianFFT
+from mri.operators.fourier.utils import estimate_density_compensation, \
+    convert_locations_to_mask, gridded_inverse_fourier_transform_nd
+from pysap.data import get_sample_data
+
+# Third party import
+from modopt.math.metrics import ssim
+import numpy as np
+import matplotlib.pyplot as plt
+
+# %%
+# Loading input data
+image = get_sample_data('2d-mri').data.astype(np.complex64)
+
+# Obtain MRI non-cartesian mask and estimate the density compensation
+kspace_loc = get_sample_data("mri-radial-samples").data
+density_comp = estimate_density_compensation(kspace_loc, image.shape, 'cell_count')
+
+# %%
+# View Input
+plt.subplot(1, 2, 1)
+plt.imshow(np.abs(image), cmap='gray')
+plt.title("MRI Data")
+plt.subplot(1, 2, 2)
+plt.imshow(convert_locations_to_mask(kspace_loc, image.shape), cmap='gray')
+plt.title("K-space Sampling Mask")
+plt.show()
+
+# %%
+# Generate the kspace
+# -------------------
+#
+# From the 2D brain slice and the acquisition mask, we retrospectively
+# undersample the k-space using a radial acquisition mask
+# We then reconstruct using adjoint with and without density compensation
+
+# Get the locations of the kspace samples and the associated observations
+fourier_op = NonCartesianFFT(
+    samples=kspace_loc,
+    shape=image.shape,
+)
+fourier_op_density_comp = NonCartesianFFT(
+    samples=kspace_loc,
+    shape=image.shape,
+    density_comp=density_comp
+)
+# Get the kspace data retrospectively. Note that this can be done with
+# `fourier_op_density_comp` as the forward operator is the same
+kspace_obs = fourier_op.op(image)
+
+# %%
+# Gridded solution
+grid_space = np.linspace(-0.5, 0.5, num=image.shape[0])
+grid2D = np.meshgrid(grid_space, grid_space)
+grid_soln = gridded_inverse_fourier_transform_nd(kspace_loc, kspace_obs,
+                                                 tuple(grid2D), 'linear')
+base_ssim = ssim(grid_soln, image, mask=np.abs(image)>np.mean(np.abs(image)))
+plt.imshow(np.abs(grid_soln), cmap='gray')
+plt.title('Gridded solution : SSIM = ' + str(np.around(base_ssim, 2)))
+plt.show()
+# %%
+# Simple adjoint
+# This preconditions k-space giving a result closer to inverse
+image_rec = fourier_op.adj_op(kspace_obs)
+recon_ssim = ssim(image_rec, image, mask=np.abs(image)>np.mean(np.abs(image)))
+plt.imshow(np.abs(image_rec), cmap='gray')
+plt.title('Simple NUFFT Adjoint : SSIM = ' + str(np.around(recon_ssim, 2)))
+plt.show()
+
+# %%
+# Density Compensation adjoint:
+# This preconditions k-space giving a result closer to inverse
+image_rec = fourier_op_density_comp.adj_op(kspace_obs)
+recon_ssim = ssim(image_rec, image, mask=np.abs(image)>np.mean(np.abs(image)))
+plt.imshow(np.abs(image_rec), cmap='gray')
+plt.title('Density Compensated Adjoint : SSIM = ' + str(np.around(recon_ssim, 2)))
+plt.show()

mri/operators/fourier/non_cartesian.py

@@ -57,6 +57,7 @@ def __init__(self, samples, shape, implementation='finufft', n_coils=1,
             self.samples,
             self.shape,
             density=self.density_comp,
+            n_coils=self.n_coils,
             **self.kwargs
         )
 

mri/operators/fourier/utils/processing.py

@@ -64,6 +64,7 @@ def convert_locations_to_mask(samples_locations, img_shape):
                          "the dimension of the image shape")
     locations = np.copy(samples_locations).astype("float")
     locations += 0.5
+    locations = np.clip(locations, 0, 0.999)
     locations *= img_shape
     locations = locations.astype("int")
     mask = np.zeros(img_shape, dtype="int")

mri/reconstructors/utils/extract_sensitivity_maps.py

@@ -139,7 +139,7 @@ def get_Smaps(k_space, img_shape, samples, thresh,
               method='linear',
               window_fun=None,
               density_comp=None, n_cpu=1,
-              fourier_op_kwargs=None):
+              fourier_op_kwargs={}):
     r"""
     Get Smaps for from pMRI sample data.
 
@@ -183,7 +183,7 @@ def get_Smaps(k_space, img_shape, samples, thresh,
         use density compensated adjoint for Smap estimation
     n_cpu: int default=1
         Number of parallel jobs in case of parallel MRI
-    fourier_op_kwargs: dict, default None
+    fourier_op_kwargs: dict, default {}
         The keyword arguments given to fourier_op initialization if
         mode == 'NFFT'. If None, we choose implementation of fourier op to
         'gpuNUFFT' if library is installed.