diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml index 4720af6ae..9984fb3ad 100644 --- a/.github/workflows/test-ci.yml +++ b/.github/workflows/test-ci.yml @@ -45,29 +45,13 @@ 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 }}] + if [[ ${{ matrix.backend }} == "finufft" ]]; then + # Install torch for testing autodiff + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + fi + + - name: Install BART if needed if: ${{ matrix.backend == 'bart' || env.ref_backend == 'bart'}} shell: bash run: | @@ -79,11 +63,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: | @@ -116,32 +95,17 @@ jobs: python -m venv venv 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" + pip install -e mri-nufft[test,${{ env.ref_backend }}] - - 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 run: | source $RUNNER_WORKSPACE/venv/bin/activate - 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 + 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 --no-binary cufinufft -e .[${{ matrix.backend }},autodiff] - name: Run Tests shell: bash @@ -163,10 +127,10 @@ jobs: shell: bash run: | cd $RUNNER_WORKSPACE - ls -al - rm -rf finufft - rm -rf gpuNUFFT - rm -rf venv + # ls -al + # rm -rf finufft + # rm -rf gpuNUFFT + # rm -rf venv get-commit-message: runs-on: ubuntu-latest @@ -209,14 +173,18 @@ 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: Point to CUDA 12.4 #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 + 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 Python deps + shell: bash + run: | + python -m pip install --upgrade pip + python -m pip install --no-binary cufinufft -e .[test,dev,extra,finufft,cufinufft,gpuNUFFT,sigpy,autodiff,doc] + - name: Run examples shell: bash run: | @@ -313,21 +281,19 @@ jobs: 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: | python -m pip install --upgrade pip - python -m pip install .[doc] - python -m pip install finufft - - - name: Install GPU related interfaces - 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" + python -m pip install --no-binary cufinufft .[doc,finufft,autodiff,gpunufft,cufinufft] + - name: Build API documentation run: | python -m sphinx docs docs_build diff --git a/examples/GPU/example_density.py b/examples/GPU/example_density.py index 03c0235a3..24e517437 100644 --- a/examples/GPU/example_density.py +++ b/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) diff --git a/examples/GPU/example_learn_samples.py b/examples/GPU/example_learn_samples.py index c0e03a046..81f12ed4f 100644 --- a/examples/GPU/example_learn_samples.py +++ b/examples/GPU/example_learn_samples.py @@ -24,7 +24,7 @@ # .. colab-link:: # :needs_gpu: 1 # -# !pip install mri-nufft[gpunufft] scikit-image +# !pip install mri-nufft[cufinufft] scikit-image import time import joblib @@ -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, diff --git a/examples/GPU/example_learn_samples_multicoil.py b/examples/GPU/example_learn_samples_multicoil.py index da8198ecc..c61d4282f 100644 --- a/examples/GPU/example_learn_samples_multicoil.py +++ b/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, diff --git a/pyproject.toml b/pyproject.toml index 1d594e0c0..fccd0cc0f 100644 --- a/pyproject.toml +++ b/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 = ["mri-nufft[torchkbnufft]"] +torchkbnufft-gpu = ["mri-nufft[torchkbnufft]"] + +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 = ["mri-nufft[pynufft]"] +pynufft-gpu = ["mri-nufft[pynufft]"] + extra = ["pymapvbvd", "scikit-image", "scikit-learn", "pywavelets"] autodiff = ["torch"] diff --git a/src/mrinufft/density/utils.py b/src/mrinufft/density/utils.py index 0b42ac918..e1992fedb 100644 --- a/src/mrinufft/density/utils.py +++ b/src/mrinufft/density/utils.py @@ -51,3 +51,19 @@ def normalize_weights(weights): """ inv_weights = np.sum(weights) / weights return inv_weights / (np.sum(inv_weights)) + + +def normalize_density(kspace_loc, shape, density, backend, **kwargs): + """Normalize the density to ensure that the reconstruction is stable.""" + from mrinufft import get_operator + + xp = np + if backend == "cufinufft": + import cupy as cp + + xp = cp + test_op = get_operator(backend)(samples=kspace_loc, shape=shape, **kwargs) + test_im = xp.ones(shape, dtype=test_op.cpx_dtype) + test_im_recon = test_op.adj_op(density * test_op.op(test_im)) + density /= xp.mean(xp.abs(test_im_recon)) + return density diff --git a/src/mrinufft/operators/base.py b/src/mrinufft/operators/base.py index f66c20126..5a129016e 100644 --- a/src/mrinufft/operators/base.py +++ b/src/mrinufft/operators/base.py @@ -13,7 +13,12 @@ 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 @@ -21,6 +26,8 @@ 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 diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py index 9f3c360f4..77cd96fb6 100644 --- a/src/mrinufft/operators/interfaces/cufinufft.py +++ b/src/mrinufft/operators/interfaces/cufinufft.py @@ -18,6 +18,7 @@ pin_memory, sizeof_fmt, ) +from mrinufft.density.utils import normalize_density CUFINUFFT_AVAILABLE = CUPY_AVAILABLE try: @@ -40,9 +41,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 +87,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 +112,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 +300,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 +837,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 +876,60 @@ 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: + density_comp = normalize_density( + kspace_loc=kspace_loc, + shape=original_shape, + density=density_comp, + backend=cls.backend, + **kwargs, + ) + return density_comp.squeeze() diff --git a/src/mrinufft/operators/interfaces/gpunufft.py b/src/mrinufft/operators/interfaces/gpunufft.py index b5d65af66..02ef22f7d 100644 --- a/src/mrinufft/operators/interfaces/gpunufft.py +++ b/src/mrinufft/operators/interfaces/gpunufft.py @@ -6,6 +6,7 @@ from ..base import FourierOperatorBase, with_numpy_cupy from mrinufft._utils import proper_trajectory, get_array_module, auto_cast from mrinufft.operators.interfaces.utils import is_cuda_array, is_host_array +from mrinufft.density.utils import normalize_density GPUNUFFT_AVAILABLE = True try: @@ -596,6 +597,7 @@ def pipe( raise ValueError( "gpuNUFFT is not available, cannot " "estimate the density compensation" ) + original_shape = volume_shape volume_shape = (np.array(volume_shape) * osf).astype(int) grid_op = MRIGpuNUFFT( samples=kspace_loc, @@ -607,11 +609,13 @@ def pipe( max_iter=num_iterations ) if normalize: - spike = np.zeros(volume_shape) - mid_loc = tuple(v // 2 for v in volume_shape) - spike[mid_loc] = 1 - psf = grid_op.adj_op(grid_op.op(spike)) - density_comp /= np.linalg.norm(psf) + density_comp = normalize_density( + kspace_loc=kspace_loc, + shape=original_shape, + density=density_comp, + backend=cls.backend, + **kwargs, + ) return density_comp.squeeze() def get_lipschitz_cst(self, max_iter=10, tolerance=1e-5, **kwargs): diff --git a/tests/operators/test_density_for_op.py b/tests/operators/test_density_for_op.py index 3e24cecf1..30bd3708c 100644 --- a/tests/operators/test_density_for_op.py +++ b/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)