Add support for density compensation (#109)

* Add support for density compensation

* Fix Documentation and PEP8

* Update to utils.py and README

* Update tests for sensitivity extraction

* remove warning for pyNUFFT and make it an error

* Update based on comments from Zac

* Revert a single change to prevent unwanted changes

* Final changes

* Make den_c None in case not used

* Fix Doc

Co-authored-by: chaithyagr <chaithyagr@gitlab.com>

README.rst

@@ -38,7 +38,7 @@ Special Installations
 
 For faster NUFFT operation, pysap-mri uses gpuNUFFT, to run the NUFFT on GPU. To install gpuNUFFT, please use:
 
-``pip install git+https://github.com/andyschwarzl/gpuNUFFT``
+``pip install gpuNUFFT``
 
 We are still in the process of merging this and getting a pip release for gpuNUFFT. Note, that you can still use CPU
 NFFT without installing the package.

examples/GPU_Examples/NonCartesian_gpuNUFFT_DensityCompensation.py

@@ -0,0 +1,78 @@
+"""
+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, WaveletUD2
+from mri.operators.utils import convert_locations_to_mask, \
+    gridded_inverse_fourier_transform_nd
+from mri.operators.fourier.utils import estimate_density_compensation
+from mri.reconstructors import SingleChannelReconstructor
+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
+
+# Loading input data
+image = get_sample_data('2d-mri')
+
+# Obtain MRI non-cartesian mask and estimate the density compensation
+radial_mask = get_sample_data("mri-radial-samples")
+kspace_loc = radial_mask.data
+density_comp = estimate_density_compensation(kspace_loc, image.shape)
+
+#############################################################################
+# 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,
+    implementation='gpuNUFFT',
+)
+fourier_op_density_comp = NonCartesianFFT(
+    samples=kspace_loc,
+    shape=image.shape,
+    implementation='gpuNUFFT',
+    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.data)
+
+# Simple adjoint
+image_rec0 = pysap.Image(data=np.abs(fourier_op.adj_op(kspace_obs)))
+# image_rec0.show()
+base_ssim = ssim(image_rec0, image)
+print('The SSIM from Adjoint is : ' + str(base_ssim))
+
+# Density Compensation adjoint:
+# This preconditions k-space giving a result closer to inverse
+image_rec1 = pysap.Image(data=np.abs(
+    fourier_op_density_comp.adj_op(kspace_obs))
+)
+# image_rec1.show()
+new_ssim = ssim(image_rec1, image)
+print('The SSIM from Density '
+      'compensated Adjoint is : ' + str(new_ssim))

examples/GPU_Examples/NonCartesian_gpuNUFFT_SENSE_DensityCompensation.py

@@ -0,0 +1,86 @@
+"""
+Neuroimaging non-cartesian reconstruction
+=========================================
+
+Author: Chaithya G R
+
+In this tutorial we will reconstruct an MR Image directly with density
+compensation and SENSE from gpuNUFFT
+
+Import neuroimaging data
+------------------------
+
+We use the toy datasets available in pysap, more specifically a 3D orange data
+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.fourier.utils import estimate_density_compensation
+from mri.reconstructors import SingleChannelReconstructor
+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
+
+# Loading input data
+image = get_sample_data('3d-pmri')
+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)
+
+#############################################################################
+# Generate the kspace
+# -------------------
+#
+# From the 3D orange slice and 3D radial acquisition mask, we retrospectively
+# undersample the k-space
+# 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=cartesian.shape,
+    n_coils=image.shape[0],
+    implementation='gpuNUFFT',
+)
+kspace_obs = fourier_op.op(image.data)
+Smaps, SOS = get_Smaps(
+    k_space=kspace_obs,
+    img_shape=fourier_op.shape,
+    samples=kspace_loc,
+    thresh=(0.05, 0.05, 0.05),  # The cutoff threshold in each kspace
+                                # direction between 0 and kspace_max (0.5)
+    min_samples=kspace_loc.min(axis=0),
+    max_samples=kspace_loc.max(axis=0),
+    density_comp=density_comp,
+    mode='NFFT',
+)
+fourier_op_sense_dc = NonCartesianFFT(
+    samples=kspace_loc,
+    shape=cartesian.shape,
+    implementation='gpuNUFFT',
+    n_coils=image.shape[0],
+    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_rec1.show()
+base_ssim = ssim(image_rec1, cartesian)
+print('The SSIM for simple Density '
+      'compensated SENSE Adjoint is : ' + str(base_ssim))

mri/operators/fourier/non_cartesian.py

@@ -27,12 +27,13 @@
     warnings.warn("pynfft python package has not been found. If needed use "
                   "the master release.")
     pass
+pynufft_available = False
 try:
     from pynufft import NUFFT_hsa, NUFFT_cpu
 except Exception:
-    warnings.warn("pynufft python package has not been found. If needed use "
-                  "the master release. Till then you cannot use NUFFT on GPU")
     pass
+else:
+    pynufft_available = True
 
 gpunufft_available = False
 try:
@@ -236,6 +237,10 @@ def __init__(self, samples, shape, platform='cuda', Kd=None, Jd=None,
         """
         if (n_coils < 1) or (type(n_coils) is not int):
             raise ValueError('The number of coils should be an integer >= 1')
+        if not pynufft_available:
+            raise ValueError('PyNUFFT Package is not installed, please '
+                             'consider using `gpuNUFFT` or install the '
+                             'PyNUFFT package')
         self.nb_coils = n_coils
         self.shape = shape
         self.platform = platform
@@ -464,14 +469,17 @@ def __init__(self, samples, shape, n_coils=1, density_comp=None,
             balance_workload
         )
 
-    def op(self, image):
+    def op(self, image, interpolate_data=False):
         """ This method calculates the masked non-cartesian Fourier transform
         of a 2D / 3D image.
 
         Parameters
         ----------
         image: np.ndarray
             input array with the same shape as shape.
+        interpolate_data: bool, default False
+            if set to True, the image is just apodized and interpolated to
+            kspace locations. This is used for density estimation.
 
         Returns
         -------
@@ -484,31 +492,36 @@ def op(self, image):
         if self.n_coils > 1 and not self.uses_sense:
             coeff = self.operator.op(np.asarray(
                 [np.reshape(image_ch.T, image_ch.size) for image_ch in image]
-            ).T)
+            ).T, interpolate_data)
         else:
-            coeff = self.operator.op(np.reshape(image.T, image.size))
+            coeff = self.operator.op(
+                np.reshape(image.T, image.size),
+                interpolate_data
+            )
             # Data is always returned as num_channels X coeff_array,
             # so for single channel, we extract single array
             if not self.uses_sense:
                 coeff = coeff[0]
         return coeff
 
-    def adj_op(self, coeff):
+    def adj_op(self, coeff, grid_data=False):
         """ This method calculates adjoint of non-uniform Fourier
         transform of a 1-D coefficients array.
 
         Parameters
         ----------
         coeff: np.ndarray
             masked non-uniform Fourier transform 1D data.
-
+        grid_data: bool, default False
+            if True, the kspace data is gridded and returned,
+            this is used for density compensation
         Returns
         -------
         np.ndarray
             adjoint operator of Non Uniform Fourier transform of the
             input coefficients.
         """
-        image = self.operator.adj_op(coeff)
+        image = self.operator.adj_op(coeff, grid_data)
         if self.n_coils > 1 and not self.uses_sense:
             image = np.asarray(
                 [image_ch.T for image_ch in image]
@@ -523,7 +536,7 @@ def adj_op(self, coeff):
 class NonCartesianFFT(OperatorBase):
     """This class wraps around different implementation algorithms for NFFT"""
     def __init__(self, samples, shape, implementation='cpu', n_coils=1,
-                 **kwargs):
+                 density_comp=None, **kwargs):
         """ Initialize the class.
 
         Parameters
@@ -548,9 +561,11 @@ def __init__(self, samples, shape, implementation='cpu', n_coils=1,
         self.samples = samples
         self.n_coils = n_coils
         if implementation == 'cpu':
+            self.density_comp = density_comp
             self.implementation = NFFT(samples=samples, shape=shape,
                                        n_coils=self.n_coils)
         elif implementation == 'cuda' or implementation == 'opencl':
+            self.density_comp = density_comp
             self.implementation = NUFFT(samples=samples, shape=shape,
                                         platform=implementation,
                                         n_coils=self.n_coils)
@@ -559,14 +574,19 @@ def __init__(self, samples, shape, implementation='cpu', n_coils=1,
                 raise ValueError('gpuNUFFT library is not installed, '
                                  'please refer to README'
                                  'or use cpu for implementation')
-            self.implementation = gpuNUFFT(samples=samples, shape=shape,
-                                           n_coils=self.n_coils, **kwargs)
+            self.implementation = gpuNUFFT(
+                samples=samples,
+                shape=shape,
+                n_coils=self.n_coils,
+                density_comp=density_comp,
+                **kwargs
+            )
         else:
             raise ValueError('Bad implementation ' + implementation +
                              ' chosen. Please choose between "cpu" | "cuda" |'
                              '"opencl" | "gpuNUFFT"')
 
-    def op(self, data):
+    def op(self, data, *args):
         """ This method calculates the masked non-cartesian Fourier transform
         of an image.
 
@@ -579,9 +599,9 @@ def op(self, data):
         -------
             masked Fourier transform of the input image.
         """
-        return self.implementation.op(data)
+        return self.implementation.op(data, *args)
 
-    def adj_op(self, coeffs):
+    def adj_op(self, coeffs, *args):
         """ This method calculates inverse masked non-uniform Fourier
         transform of a 1-D coefficients array.
 
@@ -594,7 +614,14 @@ def adj_op(self, coeffs):
         -------
             inverse discrete Fourier transform of the input coefficients.
         """
-        return self.implementation.adj_op(coeffs)
+        if not isinstance(self.implementation, gpuNUFFT) and \
+                self.density_comp is not None:
+            return self.implementation.adj_op(
+                coeffs * self.density_comp,
+                *args
+            )
+        else:
+            return self.implementation.adj_op(coeffs, *args)
 
 
 class Stacked3DNFFT(OperatorBase):

mri/operators/fourier/utils.py

@@ -308,3 +308,36 @@ def check_if_fourier_op_uses_sense(fourier_op):
         return fourier_op.implementation.uses_sense
     else:
         return False
+
+
+def estimate_density_compensation(kspace_loc, volume_shape, num_iterations=10):
+    """ Utils function to obtain the density compensator for a
+    given set of kspace locations.
+
+    Parameters
+    ----------
+    kspace_loc: np.ndarray
+        the kspace locations
+    volume_shape: np.ndarray
+        the volume shape
+    num_iterations: int default 10
+        the number of iterations for density estimation
+    """
+    from .non_cartesian import NonCartesianFFT
+    from .non_cartesian import gpunufft_available
+    if gpunufft_available is False:
+        raise ValueError("gpuNUFFT is not available, cannot "
+                         "estimate the density compensation")
+    grid_op = NonCartesianFFT(
+        samples=kspace_loc,
+        shape=volume_shape,
+        implementation='gpuNUFFT',
+        osf=1,
+    )
+    density_comp = np.ones(kspace_loc.shape[0])
+    for _ in range(num_iterations):
+        density_comp = (
+                density_comp /
+                np.abs(grid_op.op(grid_op.adj_op(density_comp, True), True))
+        )
+    return density_comp