Add density compensation functions

sigpy/mri/__init__.py

@@ -2,20 +2,23 @@
 
 It provides convenient simulation and sampling functions,
 such as the poisson-disc sampling function. It also
-provides functions to compute preconditioners.
+provides functions to compute preconditioners,
+and density compensation factors.
 
 """
 from sigpy.mri import app, linop
 
-from sigpy.mri import bloch, precond, samp, sim, util
+from sigpy.mri import bloch, dcf, precond, samp, sim, util
 from sigpy.mri.bloch import *  # noqa
+from sigpy.mri.dcf import *  # noqa
 from sigpy.mri.precond import *  # noqa
 from sigpy.mri.samp import *  # noqa
 from sigpy.mri.sim import *  # noqa
 from sigpy.mri.util import *  # noqa
 
 __all__ = ['app', 'linop']
 __all__.extend(bloch.__all__)
+__all__.extend(dcf.__all__)
 __all__.extend(precond.__all__)
 __all__.extend(samp.__all__)
 __all__.extend(sim.__all__)

sigpy/mri/dcf.py

@@ -0,0 +1,65 @@
+"""Density compensation functions.
+
+"""
+import sigpy as sp
+from tqdm.auto import tqdm
+
+
+__all__ = ['pipe_menon_dcf']
+
+
+def pipe_menon_dcf(coord, device=sp.cpu_device, max_iter=30,
+                   n=128, beta=8, width=4, show_pbar=True):
+    r"""Compute Pipe Menon density compensation factor.
+
+    Perform the following iteration:
+
+    .. math::
+
+        w = \frac{w}{|G^H G w|}
+
+    with :math:`G` as the gridding operator.
+
+    Args:
+        coord (array): k-space coordinates.
+        device (Device): computing device.
+        max_iter (int): number of iterations.
+        n (int): Kaiser-Bessel sampling numbers for gridding operator.
+        beta (float): Kaiser-Bessel kernel parameter.
+        width (float): Kaiser-Bessel kernel width.
+        show_pbar (bool): show progress bar.
+
+    Returns:
+        array: density compensation factor.
+
+    References:
+        Pipe, James G., and Padmanabhan Menon.
+        Sampling Density Compensation in MRI:
+        Rationale and an Iterative Numerical Solution.
+        Magnetic Resonance in Medicine 41, no. 1 (1999): 179–86.
+
+
+    """
+    device = sp.Device(device)
+    xp = device.xp
+
+    with device:
+        w = xp.ones(coord.shape[:-1], dtype=coord.dtype)
+        img_shape = sp.estimate_shape(coord)
+
+        # Get kernel
+        x = xp.arange(n, dtype=coord.dtype) / n
+        kernel = xp.i0(beta * (1 - x**2)**0.5).astype(coord.dtype)
+        kernel /= kernel.max()
+
+        G = sp.linop.Gridding(img_shape, coord, width, kernel)
+        with tqdm(total=max_iter, disable=not show_pbar) as pbar:
+            for it in range(max_iter):
+                GHGw = G.H * G * w
+                w /= xp.abs(GHGw)
+                resid = xp.abs(GHGw - 1).max().item()
+
+                pbar.set_postfix(resid='{0:.2E}'.format(resid))
+                pbar.update()
+
+    return w

tests/mri/test_dcf.py

@@ -0,0 +1,28 @@
+import unittest
+import numpy as np
+import sigpy as sp
+import numpy.testing as npt
+
+from sigpy.mri import dcf, samp
+
+if __name__ == '__main__':
+    unittest.main()
+
+
+class TestApp(unittest.TestCase):
+
+    def shepp_logan_setup(self):
+        img_shape = [16, 16]
+        coord_shape = [int(16 * np.pi), 16, 2]
+
+        img = sp.shepp_logan(img_shape)
+        coord = samp.radial(coord_shape, img_shape)
+        ksp = sp.nufft(img, coord)
+        return img, coord, ksp
+
+    def test_shepp_logan_dcf(self):
+        img, coord, ksp = self.shepp_logan_setup()
+        pm_dcf = dcf.pipe_menon_dcf(coord, show_pbar=False)
+        img_dcf = sp.nufft_adjoint(ksp * pm_dcf, coord, oshape=img.shape)
+        img_dcf /= np.abs(img_dcf).max()
+        npt.assert_allclose(img, img_dcf, atol=1, rtol=1e-1)