Add support for fiNUFFT spread interpolate (#224)

* Added

* Fix

* Remove bymistake add

* Fix

* Fixed lint

* Lint

* Added refbackend

* Fix NDFT

* feat: use finufft as ref backend.

* feat(tests): move ndft vs nufft tests to own file.

* Add support for pipe

* \!docs_build try to run cufinufft tests

* \!docs_build fix style

* Added next235 for stability

* Fix lint

* Fix CUPY

* WIP

* Updates

* fix back learn examples

* move tto flatiron

* fix black

* Move to test on GPU

* Update pyproject toml and use it in test-ci, to prevent duplication of dependencies and actually test them!

* Make CI build shorter

* Test run to run

* \!docs_build Added

* adding density normalization

* Start support for finufft spread and interpolate (pipe)

* Add support for density comp

* Add support for normlize back

* Added a bunch of extra codes

* PEP fixes

* Update siemens.py

* Added fixes

* add [docs]

* Fixes and updates on the locatuions

* Update the codes to be in sync with cufinufft / finufft 2.4.0

* Merged to master

* [style]

* Fix toml ffile [docs]

* Fix toml ffile [docs]

* Update testbatch stuff

* Update the tests

* Fixed tests againnnn

* All set

* Add finufft

* Updates

* Fixes

* More fixes

* make tests less strict again

* remove bymistake

* Remove strictness, everything works, hopefully [docs]

* [docs] setup cufinufft also

* [docs] fix style

* [docs] more updates, fixes

* pywavelets

* [docs] fixes further

* [docs] try again

* [docs] final comments

* [docs] method class use

---------

Co-authored-by: chaithyagr <[email protected]>
Co-authored-by: Pierre-antoine Comby <[email protected]>
Co-authored-by: Asma TANABENE <[email protected]>

Author: 4 people

.github/workflows/test-ci.yml

@@ -51,36 +51,9 @@ jobs:
           ${{ env.create_venv }}
           ${{ env.activate_venv }}
           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: |
-          ${{ env.activate_venv }}
-          python -m pip install "pynfft2>=1.4.3"
-
-      - name: Install pynufft
-        if: ${{ matrix.backend == 'pynufft-cpu' || env.ref_backend == 'pynufft-cpu' }}
-        run: |
-          ${{ env.activate_venv }}
-          python -m pip install pynufft
-
-      - name: Install finufft
-        if: ${{ matrix.backend == 'finufft'  || env.ref_backend == 'finufft'}}
-        shell: bash
-        run: |
-          ${{ env.activate_venv }}
-          python -m pip install finufft
-
-      - name: Install Sigpy
-        if: ${{ matrix.backend == 'sigpy'  || env.ref_backend == 'sigpy'}}
-        shell: bash
-        run: |
-          ${{ env.activate_venv }}
-          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: |
@@ -92,13 +65,6 @@ jobs:
           make
           echo $PWD >> $GITHUB_PATH
 
-      - name: Install torchkbnufft-cpu
-        if: ${{ matrix.backend == 'torchkbnufft-cpu'  || env.ref_backend == 'torchkbnufft-cpu'}}
-        run: |
-          ${{ env.activate_venv }}
-          python -m pip install torchkbnufft
-
-
       - name: Run Tests
         shell: bash
         run: |
@@ -127,33 +93,22 @@ jobs:
       - name: Install mri-nufft and finufft
         shell: bash
         run: |
-          ${{ env.create_venv }}
-          ${{ env.activate_venv }}
-          python -m pip install --upgrade pip wheel
-          python -m pip install -e .[test]
-          python -m pip install cupy-cuda12x finufft "numpy<2.0"
+          cd $RUNNER_WORKSPACE
+          python --version
+          python -m venv venv
+          source $RUNNER_WORKSPACE/venv/bin/activate
+          pip install --upgrade pip wheel
+          pip install -e mri-nufft[test,finufft]
 
-      - name: Install torch with CUDA 12.1
-        shell: bash
-        if: ${{ matrix.backend != 'tensorflow'}}
-        run: |
-          ${{ env.activate_venv }}
-          ${{ env.setup_cuda }}
-          python -m pip install torch
 
       - name: Install backend
         shell: bash
         run: |
-          ${{ env.activate_venv }}
-          if [[ ${{ matrix.backend }} == "torchkbnufft-gpu" ]]; then
-            python -m pip install torchkbnufft
-          elif [[ ${{ matrix.backend }} == "tensorflow" ]]; then
-            python -m 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
-            python -m pip install "cufinufft<2.3"
-          else
-            python -m pip install ${{ matrix.backend }}
-          fi
+          source $RUNNER_WORKSPACE/venv/bin/activate
+          export CUDA_BIN_PATH=/usr/local/cuda-12.4/
+          export PATH=/usr/local/cuda-12.4/bin/:${PATH}
+          export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64/:${LD_LIBRARY_PATH}
+          pip install -e .[${{ matrix.backend }},autodiff]
 
       - name: Run Tests
         shell: bash
@@ -270,13 +225,14 @@ jobs:
           python -m pip install -e .[extra,test,dev]
           python -m pip install finufft pooch brainweb-dl torch fastmri
 
-      - name: Install GPU related interfaces
+      - name: Install Python deps
+        shell: bash
         run: |
           ${{ env.activate_venv }}
           ${{ env.setup_cuda }}
-          pip install cupy-cuda12x torch
-          python -m pip install gpuNUFFT "cufinufft<2.3" sigpy scikit-image fastmri
-
+          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: |
@@ -377,23 +333,22 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: "3.10"
-
+        
+      - name: Point to CUDA 12.4
+        run: |
+          export CUDA_BIN_PATH=/usr/local/cuda-12.4/
+          export PATH=/usr/local/cuda-12.4/bin/:${PATH}
+          export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64/:${LD_LIBRARY_PATH}
+      
       - name: Install dependencies
         shell: bash -l {0}
         run: |
           ${{ env.create_venv }}
           ${{ env.activate_venv }}
           python -m pip install --upgrade pip
-          python -m pip install .[doc,extra]
-          python -m pip install finufft
-
-      - name: Install GPU related interfaces
-        run: |
-          ${{ env.activate_venv }}
-          ${{ env.setup_cuda }}
-          pip install cupy-cuda12x torch
-          python -m pip install gpuNUFFT "cufinufft<2.3" sigpy scikit-image fastmri
-
+          python -m pip install .[doc,finufft,autodiff,gpunufft,cufinufft,sigpy,extra] fastmri
+      
+          
       - name: Build API documentation
         run: |
           ${{ env.activate_venv }}

docs/binder/environment.yml

@@ -10,3 +10,4 @@ dependencies:
       - joblib
       - brainweb-dl
       - torch
+      - pywavelets

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(np.squeeze(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,
@@ -103,7 +103,7 @@ def forward(self, x):
             self.trajectory.detach().numpy(),
             self.img_size,
             kspace.detach(),
-            backend="gpunufft",
+            backend="cufinufft",
             density=self.sense_op.density,
             blurr_factor=20,
         )

examples/example_learn_samples_multires.py

@@ -83,6 +83,7 @@ def __init__(
         self.operator = get_operator("finufft", wrt_data=True, wrt_traj=True)(
             sample_points,
             shape=img_size,
+            density=True,
             squeeze_dims=False,
         )
         self.img_size = img_size

pyproject.toml

@@ -12,11 +12,23 @@ 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.4.0b1", "cupy-cuda12x"]
+tensorflow = ["tensorflow-mri==0.21.0", "tensorflow-probability==0.17.0", "tensorflow-io==0.27.0", "matplotlib==3.7"]
+finufft = ["finufft==2.4.0rc1"]
+sigpy = ["sigpy"]
 pynfft = ["pynfft3"]
+
 pynufft = ["pynufft"]
+pynufft-cpu = ["pynufft"]
+pynufft-gpu = ["pynufft"]
+
+io = ["pymapvbvd"]
+smaps = ["scikit-image"]
 extra = ["pymapvbvd", "scikit-image", "scikit-learn", "pywavelets"]
 autodiff = ["torch"]
 

src/mrinufft/operators/base.py

@@ -14,12 +14,24 @@
 import numpy as np
 from numpy.typing import NDArray
 
-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.
 DTYPE_R2C = {"float32": "complex64", "float64": "complex128"}
 
@@ -309,8 +321,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

@@ -39,11 +39,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."""
 
@@ -64,7 +59,9 @@ 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)
@@ -278,7 +275,7 @@ def samples(self, new_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, samples=self._samples)
         self.compute_density(self._density_method)
 
     @FourierOperatorBase.density.setter
@@ -850,3 +847,55 @@ 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(
+                "cufinufft is not available, cannot estimate the density compensation"
+            )
+        grid_op = cls(
+            samples=kspace_loc,
+            shape=volume_shape,
+            upsampfac=osf,
+            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 = cls(samples=kspace_loc, shape=volume_shape, **kwargs)
+            test_im = cp.ones(volume_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()