Add support for density compensation estimation with cufinufft

.github/workflows/test-ci.yml

@@ -45,29 +45,9 @@ jobs:
         shell: bash
         run: |
           python -m pip install --upgrade pip
-          python -m pip install -e .[test]
-
-      - name: Install pynfft
-        if: ${{ matrix.backend == 'pynfft' || env.ref_backend == 'pynfft' }}
-        shell: bash
-        run: |
-          python -m pip install "pynfft2>=1.4.3"
-
-      - name: Install pynufft
-        if: ${{ matrix.backend == 'pynufft-cpu' || env.ref_backend == 'pynufft-cpu' }}
-        run: python -m pip install pynufft
-
-      - name: Install finufft
-        if: ${{ matrix.backend == 'finufft'  || env.ref_backend == 'finufft'}}
-        shell: bash
-        run: python -m pip install finufft
-
-      - name: Install Sigpy
-        if: ${{ matrix.backend == 'sigpy'  || env.ref_backend == 'sigpy'}}
-        shell: bash
-        run: python -m pip install sigpy
-
-      - name: Install BART
+          python -m pip install -e .[test,${{ env.ref_backend }},${{ matrix.backend }}]
+
+      - name: Install BART if needed
         if: ${{ matrix.backend == 'bart'  || env.ref_backend == 'bart'}}
         shell: bash
         run: |
@@ -79,11 +59,6 @@ jobs:
           make
           echo $PWD >> $GITHUB_PATH
 
-      - name: Install torchkbnufft-cpu
-        if: ${{ matrix.backend == 'torchkbnufft-cpu'  || env.ref_backend == 'torchkbnufft-cpu'}}
-        run: python -m pip install torchkbnufft
-
-
       - name: Run Tests
         shell: bash
         run: |
@@ -117,14 +92,7 @@ jobs:
           source $RUNNER_WORKSPACE/venv/bin/activate
           pip install --upgrade pip wheel
           pip install -e mri-nufft[test]
-          pip install cupy-cuda12x finufft "numpy<2.0"
 
-      - name: Install torch with CUDA 12.1
-        shell: bash
-        if: ${{ matrix.backend != 'tensorflow'}}
-        run: |
-          source $RUNNER_WORKSPACE/venv/bin/activate
-          pip install torch
 
       - name: Install backend
         shell: bash
@@ -133,15 +101,7 @@ jobs:
           export CUDA_BIN_PATH=/usr/local/cuda-12.1/
           export PATH=/usr/local/cuda-12.1/bin/:${PATH}
           export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH}
-          if [[ ${{ matrix.backend }} == "torchkbnufft-gpu" ]]; then
-            pip install torchkbnufft
-          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" 
-          else
-            pip install ${{ matrix.backend }}
-          fi
+          pip install -e .[${{ matrix.backend }},autodiff]
 
       - name: Run Tests
         shell: bash
@@ -202,21 +162,18 @@ jobs:
           path: ~/.cache/brainweb
           key: ${{ runner.os }}-Brainweb
 
-      - name: Install Python deps
-        shell: bash
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install -e .[test,dev]
-          python -m pip install finufft pooch brainweb-dl torch fastmri
-
-      - name: Install GPU related interfaces
+      - name: Point to CUDA 12.1 #TODO: This can be combined from other jobs
         run: |
           export CUDA_BIN_PATH=/usr/local/cuda-12.1/
           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 fastmri
 
+      - name: Install Python deps
+        shell: bash
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e .[test,dev,finufft,cufinufft,gpuNUFFT,sigpy,smaps,autodiff,doc]
+
       - name: Run examples
         shell: bash
         run: |
@@ -313,21 +270,19 @@ jobs:
         with:
           python-version: "3.10"
 
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install .[doc]
-          python -m pip install finufft
-
-      - name: Install GPU related interfaces
+      - name: Point to CUDA 12.1
         run: |
           export CUDA_BIN_PATH=/usr/local/cuda-12.1/
           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"
 
+      - name: Install dependencies
+        shell: bash -l {0}
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install .[doc,finufft,autodiff,gpunufft,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.py

@@ -54,7 +54,7 @@ def __init__(self, inital_trajectory):
             data=torch.Tensor(inital_trajectory),
             requires_grad=True,
         )
-        self.operator = get_operator("gpunufft", wrt_data=True, wrt_traj=True)(
+        self.operator = get_operator("cufinufft", wrt_data=True, wrt_traj=True)(
             self.trajectory.detach().cpu().numpy(),
             shape=(256, 256),
             density=True,

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

@@ -12,11 +12,21 @@ dynamic = ["version"]
 [project.optional-dependencies]
 
 gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"]
+
 torchkbnufft = ["torchkbnufft", "cupy-cuda12x"]
-cufinufft = ["cufinufft<2.3", "cupy-cuda12x"]
-finufft = ["finufft"]
+torchkbnufft-cpu = ["torchkbnufft", "cupy-cuda12x"]
+torchkbnufft-gpu = ["torchkbnufft", "cupy-cuda12x"]
+
+cufinufft = ["cufinufft>=2.3.1", "cupy-cuda12x"]
+tensorflow = ["tensorflow-mri==0.21.0", "tensorflow-probability==0.17.0", "tensorflow-io==0.27.0", "matplotlib==3.7"]
+finufft = ["finufft>=2.3"]
+sigpy = ["sigpy"]
 pynfft = ["pynfft2>=1.4.3; python_version < '3.12'", "numpy>=2.0.0; python_version < '3.12'"]
+
 pynufft = ["pynufft"]
+pynufft-cpu = ["pynufft"]
+pynufft-gpu = ["pynufft"]
+
 io = ["pymapvbvd"]
 smaps = ["scikit-image"]
 autodiff = ["torch"]

src/mrinufft/operators/base.py

@@ -13,14 +13,21 @@
 
 import numpy as np
 
-from mrinufft._array_compat import with_numpy, with_numpy_cupy, AUTOGRAD_AVAILABLE
+from mrinufft._array_compat import (
+    with_numpy,
+    with_numpy_cupy,
+    AUTOGRAD_AVAILABLE,
+    CUPY_AVAILABLE,
+)
 from mrinufft._utils import auto_cast, power_method
 from mrinufft.density import get_density
 from mrinufft.extras import get_smaps
 from mrinufft.operators.interfaces.utils import is_cuda_array, is_host_array
 
 if AUTOGRAD_AVAILABLE:
     from mrinufft.operators.autodiff import MRINufftAutoGrad
+if CUPY_AVAILABLE:
+    import cupy as cp
 
 
 # Mapping between numpy float and complex types.
@@ -304,8 +311,10 @@ def compute_density(self, method=None):
             if `backend` is `tensorflow`, `gpunufft` or `torchkbnufft-cpu`
                 or `torchkbnufft-gpu`.
         """
-        if isinstance(method, np.ndarray):
-            self._density = method
+        if isinstance(method, np.ndarray) or (
+            CUPY_AVAILABLE and isinstance(method, cp.ndarray)
+        ):
+            self.density = method
             return None
         if not method:
             self._density = None

src/mrinufft/operators/interfaces/cufinufft.py

@@ -40,9 +40,30 @@
 DTYPE_R2C = {"float32": "complex64", "float64": "complex128"}
 
 
-def _error_check(ier, msg):
-    if ier != 0:
-        raise RuntimeError(msg)
+def _next235beven(n, b):
+    """Find the next even integer not less than n.
+
+    This function finds the next even integer not less than n, with prime factors no
+    larger than 5, and is a multiple of b (where b is a number that only
+    has prime factors 2, 3, and 5).
+    It is used in particular with `pipe` density compensation estimation.
+    """
+    if n <= 2:
+        return 2
+    if n % 2 == 1:
+        n += 1  # make it even
+    nplus = n - 2  # to cancel out the +=2 at start of loop
+    numdiv = 2  # a dummy that is >1
+    while numdiv > 1 or nplus % b != 0:
+        nplus += 2  # stays even
+        numdiv = nplus
+        while numdiv % 2 == 0:
+            numdiv //= 2  # remove all factors of 2, 3, 5...
+        while numdiv % 3 == 0:
+            numdiv //= 3
+        while numdiv % 5 == 0:
+            numdiv //= 5
+    return nplus
 
 
 class RawCufinufftPlan:
@@ -65,10 +86,12 @@ def __init__(
         # and type 2 with 2.
         self.plans = [None, None, None]
         self.grad_plan = None
-
+        self._kx = cp.array(samples[:, 0], copy=False)
+        self._ky = cp.array(samples[:, 1], copy=False)
+        self._kz = cp.array(samples[:, 2], copy=False) if self.ndim == 3 else None
         for i in [1, 2]:
             self._make_plan(i, **kwargs)
-            self._set_pts(i, samples)
+            self._set_pts(i)
 
     @property
     def dtype(self):
@@ -88,13 +111,15 @@ def _make_plan(self, typ, **kwargs):
             **kwargs,
         )
 
-    def _set_pts(self, typ, samples):
+    def _set_kxyz(self, samples):
+        self._kx.set(samples[:, 0])
+        self._ky.set(samples[:, 1])
+        if self.ndim == 3:
+            self._kz.set(samples[:, 2])
+
+    def _set_pts(self, typ):
         plan = self.grad_plan if typ == "grad" else self.plans[typ]
-        plan.setpts(
-            cp.array(samples[:, 0], copy=False),
-            cp.array(samples[:, 1], copy=False),
-            cp.array(samples[:, 2], copy=False) if self.ndim == 3 else None,
-        )
+        plan.setpts(self._kx, self._ky, self._kz)
 
     def _destroy_plan(self, typ):
         if self.plans[typ] is not None:
@@ -274,10 +299,11 @@ def samples(self, samples):
         self._samples = np.asfortranarray(
             proper_trajectory(samples, normalize="pi").astype(np.float32, copy=False)
         )
+        self.raw_op._set_kxyz(self._samples)
         for typ in [1, 2, "grad"]:
             if typ == "grad" and not self._grad_wrt_traj:
                 continue
-            self.raw_op._set_pts(typ, self._samples)
+            self.raw_op._set_pts(typ)
         self.compute_density(self._density_method)
 
     @FourierOperatorBase.density.setter
@@ -810,7 +836,7 @@ def _make_plan_grad(self, **kwargs):
             isign=1,
             **kwargs,
         )
-        self.raw_op._set_pts(typ="grad", samples=self.samples)
+        self.raw_op._set_pts(typ="grad")
 
     def get_lipschitz_cst(self, max_iter=10, **kwargs):
         """Return the Lipschitz constant of the operator.
@@ -849,3 +875,57 @@ 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"
+            )
+        original_shape = volume_shape
+        volume_shape = np.array([_next235beven(int(osf * i), 1) 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()
+            )
+        if normalize:
+            test_op = MRICufiNUFFT(samples=kspace_loc, shape=original_shape, **kwargs)
+            test_im = cp.ones(original_shape, dtype=test_op.cpx_dtype)
+            test_im_recon = test_op.adj_op(density_comp * test_op.op(test_im))
+            density_comp /= cp.mean(cp.abs(test_im_recon))
+        return density_comp.squeeze()