Add support for density compensation estimation with cufinufft

.github/workflows/test-ci.yml

@@ -138,7 +138,8 @@ jobs:
           elif [[ ${{ matrix.backend }} == "tensorflow" ]]; then
             pip install tensorflow-mri==0.21.0 tensorflow-probability==0.17.0 tensorflow-io==0.27.0 matplotlib==3.7
           elif [[ ${{ matrix.backend }} == "cufinufft" ]]; then
-            pip install "cufinufft<2.3" 
+            git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly
+            pip install finufft/python/cufinufft
           else
             pip install ${{ matrix.backend }}
           fi
@@ -215,7 +216,9 @@ jobs:
           export PATH=/usr/local/cuda-12.1/bin/:${PATH}
           export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH}
           pip install cupy-cuda12x torch
-          python -m pip install gpuNUFFT "cufinufft<2.3" sigpy scikit-image
+          python -m pip install gpuNUFFT sigpy scikit-image
+          git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly
+          pip install finufft/python/cufinufft
 
       - name: Run examples
         shell: bash
@@ -326,8 +329,10 @@ jobs:
           export PATH=/usr/local/cuda-12.1/bin/:${PATH}
           export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH}
           pip install cupy-cuda12x torch
-          python -m pip install gpuNUFFT "cufinufft<2.3"
-
+          python -m pip install gpuNUFFT 
+          git clone https://github.com/chaithyagr/finufft --branch fix_spreadinterponly
+          pip install finufft/python/cufinufft
+
       - name: Build API documentation
         run: |
           python -m sphinx docs docs_build

examples/GPU/example_density.py

@@ -142,7 +142,7 @@
 #    If this method is widely used in the literature, there exists no convergence guarantees for it.
 #
 # .. note::
-#    The Pipe method is currently only implemented for gpuNUFFT.
+#    The Pipe method is currently only implemented for gpuNUFFT and cufinufft backend.
 
 # %%
 flat_traj = traj.reshape(-1, 2)
@@ -158,3 +158,21 @@
 axs[2].imshow(abs(adjoint_manual))
 axs[2].set_title("Pipe density compensation")
 print(nufft.density)
+
+# %%
+# We can also do density compensation using cufinufft backend
+
+# %%
+flat_traj = traj.reshape(-1, 2)
+nufft = get_operator("cufinufft")(
+    traj, shape=mri_2D.shape, density={"name": "pipe", "osf": 2}
+)
+adjoint_manual = nufft.adj_op(kspace)
+fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+axs[0].imshow(abs(mri_2D))
+axs[0].set_title("Ground Truth")
+axs[1].imshow(abs(adjoint))
+axs[1].set_title("no density compensation")
+axs[2].imshow(abs(adjoint_manual))
+axs[2].set_title("Pipe density compensation")
+print(nufft.density)

examples/GPU/example_learn_samples_multicoil.py

@@ -75,7 +75,7 @@ def __init__(self, inital_trajectory, n_coils, img_size=(256, 256)):
             squeeze_dims=False,
         )
         # A simple density compensated adjoint SENSE operator with sensitivity maps `smaps`.
-        self.sense_op = get_operator("gpunufft", wrt_data=True, wrt_traj=True)(
+        self.sense_op = get_operator("cufinufft", wrt_data=True, wrt_traj=True)(
             sample_points,
             shape=img_size,
             density=True,

pyproject.toml

@@ -13,8 +13,8 @@ dynamic = ["version"]
 
 gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"]
 torchkbnufft = ["torchkbnufft", "cupy-cuda12x"]
-cufinufft = ["cufinufft<2.3", "cupy-cuda12x"]
-finufft = ["finufft"]
+cufinufft = ["cufinufft>=2.3", "cupy-cuda12x"]
+finufft = ["finufft>=2.3"]
 pynfft = ["pynfft2>=1.4.3", "numpy>=2.0.0"]
 pynufft = ["pynufft"]
 io = ["pymapvbvd"]

src/mrinufft/operators/interfaces/cufinufft.py

@@ -40,11 +40,6 @@
 DTYPE_R2C = {"float32": "complex64", "float64": "complex128"}
 
 
-def _error_check(ier, msg):
-    if ier != 0:
-        raise RuntimeError(msg)
-
-
 class RawCufinufftPlan:
     """Light wrapper around the guru interface of finufft."""
 
@@ -836,3 +831,51 @@ def toggle_grad_traj(self):
         if self.uses_sense:
             self.smaps = self.smaps.conj()
         self.raw_op.toggle_grad_traj()
+
+    @classmethod
+    def pipe(
+        cls,
+        kspace_loc,
+        volume_shape,
+        num_iterations=10,
+        osf=2,
+        normalize=True,
+        **kwargs,
+    ):
+        """Compute the density compensation weights 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
+        osf: float or int
+            The oversampling factor the volume shape
+        normalize: bool
+            Whether to normalize the density compensation.
+            We normalize such that the energy of PSF = 1
+        """
+        if CUFINUFFT_AVAILABLE is False:
+            raise ValueError(
+                "gpuNUFFT is not available, cannot " "estimate the density compensation"
+            )
+        volume_shape = np.array([int(osf * i) for i in volume_shape])
+        grid_op = MRICufiNUFFT(
+            samples=kspace_loc,
+            shape=volume_shape,
+            upsampfac=1,
+            gpu_spreadinterponly=1,
+            gpu_kerevalmeth=0,
+            **kwargs,
+        )
+        density_comp = cp.ones(kspace_loc.shape[0], dtype=grid_op.cpx_dtype)
+        for _ in range(num_iterations):
+            density_comp /= cp.abs(
+                grid_op.op(
+                    grid_op.adj_op(density_comp.astype(grid_op.cpx_dtype))
+                ).squeeze()
+            )
+        return density_comp.squeeze()

tests/operators/test_density_for_op.py

@@ -25,21 +25,28 @@ def radial_distance(traj, shape):
         CasesTrajectories.case_nyquist_radial3D,
     ],
 )
-@parametrize(backend=["gpunufft", "tensorflow"])
+@parametrize(backend=["gpunufft", "tensorflow", "cufinufft"])
 def test_pipe(backend, traj, shape, osf):
     """Test the pipe method."""
     distance = radial_distance(traj, shape)
     if osf != 2 and backend == "tensorflow":
         pytest.skip("OSF < 2 not supported for tensorflow.")
-    result = pipe(traj, shape, backend, osf=osf, num_iterations=10)
+    result = pipe(traj, shape, backend=backend, osf=osf, num_iterations=10)
+    if backend == "cufinufft":
+        result = result.get()
     result = result / np.mean(result)
     distance = distance / np.mean(distance)
-    if backend == "tensorflow":
-        # If tensorflow, we dont perfectly estimate, but we still want to ensure
-        # we can get density
-        assert_correlate(result, distance, slope=1, slope_err=None, r_value_err=0.5)
-    elif osf != 2:
-        # If OSF < 2, we dont perfectly estimate
-        assert_correlate(result, distance, slope=1, slope_err=None, r_value_err=0.2)
-    else:
-        assert_correlate(result, distance, slope=1, slope_err=0.1, r_value_err=0.1)
+    r_err = 0.2
+    slope_err = None
+    if osf == 2:
+        r_err = 0.1
+        slope_err = 0.1
+    if backend == "cufinufft":
+        r_err *= 2
+        slope_err = slope_err * 2 if slope_err is not None else None
+        if len(shape) == 3:
+            r_err *= 2
+            slope_err = slope_err * 2 if slope_err is not None else None
+    elif backend == "tensorflow":
+        r_err = 0.5
+    assert_correlate(result, distance, slope=1, slope_err=slope_err, r_value_err=r_err)