Enhance Fourier Corrected operator (#180)

* fix typoes (off_resonnance -> off_resonance)

* feat: add field interpolators calculation (#124)

This add "get_interpolators_from_fieldmap", supporting both real (B0 map only) and complex (Zmap = R2* + 1j * B0map) fields
in arbitrary dimension (2D and 3D). The routine supports both numpy / cupy arrays and torch tensors on CPU and GPU,
the latter requiring cupy due to limitations in torch.histogram / torch.histogramdd (pytorch/pytorch#69519).

Calculation uses time segmentation with uniform time samples and LS coefficients (using histogram).

Based on MIRT (https://github.com/JeffFessler/mirt/blob/main/mri/mri_exp_approx.m)
and its Python porting from MIRTORCH (https://github.com/guanhuaw/MIRTorch/blob/master/mirtorch/linear/mri.py)
and SigPy (https://github.com/mikgroup/sigpy/blob/main/sigpy/mri/util.py).

* feat: field interpolator estimation in MRIFourierCorrected constructor.

In addition, add off-resonance example.

* feat: set autograd_available based on base FourierOperator

* Update test_offres_exp_approx.py

Avoid torch cuda test case if cupy is not available.

* Update src/mrinufft/operators/off_resonance.py

remove deprecated get_grad

Co-authored-by: Chaithya G R <chaithyagr@gmail.com>

* address pr review

* remove phantominator

* remove duplicated example and replacing subject in example_offres

* address PR rev #2

* \!docs_build and fix on \!style

* \!docs_build

* skip field coefficient gpu test case if cupy not available

* \!docs_build fix gpuNUFFT

* Update src/mrinufft/operators/off_resonance.py

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>

* Update src/mrinufft/operators/off_resonance.py

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>

* \!docs_build gpunufft

* Update src/mrinufft/operators/off_resonance.py

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>

* Update src/mrinufft/operators/off_resonance.py

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>

* address PR rev 3

* Update pyproject.toml

* Update src/mrinufft/operators/off_resonance.py

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>

* \![docs_build]

* temp, \![docs_build] fix

* fix cufinufft, \!docs_build fix

---------

Co-authored-by: Chaithya G R <chaithyagr@gmail.com>
Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>

.github/workflows/test-ci.yml

@@ -137,6 +137,8 @@ jobs:
             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
@@ -213,7 +215,7 @@ 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 sigpy scikit-image
+          python -m pip install gpuNUFFT "cufinufft<2.3" sigpy scikit-image
 
       - name: Run examples
         shell: bash
@@ -324,7 +326,7 @@ 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
+          python -m pip install gpuNUFFT "cufinufft<2.3"
       
       - name: Build API documentation
         run: |

docs/sphinx_add_colab_link.py

@@ -73,6 +73,7 @@ def notebook_modifier(self, notebook_path, commands):
         idx = self.find_index_of_colab_link(notebook)
         code_lines = ["# Install libraries"]
         code_lines.append(commands)
+        code_lines.append("pip install brainweb-dl  # Required for data")
         dummy_notebook_content = {"cells": []}
         add_code_cell(
             dummy_notebook_content,

examples/example_offresonance.py

@@ -0,0 +1,110 @@
+"""
+======================
+Off-resonance Corrected NUFFT Operator
+======================
+
+Example of Off-resonance Corrected NUFFT trajectory operator.
+
+This examples show how to use the Off-resonance Corrected NUFFT operator to acquire 
+and reconstruct data in presence of field inhomogeneities.
+Here a spiral trajectory is used as a demonstration.
+
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from mrinufft import display_2D_trajectory
+
+plt.rcParams["image.cmap"] = "gray"
+
+# %%
+# Data Generation
+# ===============
+# For realistic 2D image we will use a slice from the brainweb dataset.
+# installable using ``pip install brainweb-dl``
+
+from brainweb_dl import get_mri
+
+mri_data = get_mri(0, "T1")
+mri_data = mri_data[::-1, ...][90]
+plt.imshow(mri_data), plt.axis("off"), plt.title("ground truth")
+
+# %%
+# Masking
+# ===============
+# Here, we generate a binary mask to exclude the background.
+# We perform a simple binary threshold; in real-world application,
+# it is advised to use other tools (e.g., FSL-BET).
+
+brain_mask = mri_data > 0.1 * mri_data.max()
+plt.imshow(brain_mask), plt.axis("off"), plt.title("brain mask")
+
+# %%
+# Field Generation
+# ===============
+# Here, we generate a radial B0 field with the same shape of
+# the input Shepp-Logan phantom
+
+from mrinufft.extras import make_b0map
+
+# generate field
+b0map, _ = make_b0map(mri_data.shape, b0range=(-200, 200), mask=brain_mask)
+plt.imshow(brain_mask * b0map, cmap="bwr", vmin=-200, vmax=200), plt.axis(
+    "off"
+), plt.colorbar(), plt.title("B0 map [Hz]")
+
+# %%
+# Generate a Spiral trajectory
+# ----------------------------
+
+from mrinufft import initialize_2D_spiral
+from mrinufft.density import voronoi
+from mrinufft.trajectories.utils import DEFAULT_RASTER_TIME
+
+samples = initialize_2D_spiral(Nc=48, Ns=600, nb_revolutions=10)
+t_read = np.arange(samples.shape[1]) * DEFAULT_RASTER_TIME * 1e-3
+t_read = np.repeat(t_read[None, ...], samples.shape[0], axis=0)
+density = voronoi(samples)
+
+display_2D_trajectory(samples)
+
+# %%
+# Setup the Operator
+# ==================
+
+from mrinufft import get_operator
+from mrinufft.operators.off_resonance import MRIFourierCorrected
+
+# Generate standard NUFFT operator
+nufft = get_operator("finufft")(
+    samples=samples,
+    shape=mri_data.shape,
+    density=density,
+)
+
+# Generate Fourier Corrected operator
+mfi_nufft = MRIFourierCorrected(
+    nufft, b0_map=b0map, readout_time=t_read, mask=brain_mask
+)
+
+# Generate K-Space
+kspace = mfi_nufft.op(mri_data)
+
+# Reconstruct without field correction
+mri_data_adj = nufft.adj_op(kspace)
+mri_data_adj = np.squeeze(abs(mri_data_adj))
+
+# Reconstruct with field correction
+mri_data_adj_mfi = mfi_nufft.adj_op(kspace)
+mri_data_adj_mfi = np.squeeze(abs(mri_data_adj_mfi))
+
+fig2, ax2 = plt.subplots(1, 2)
+ax2[0].imshow(mri_data_adj), ax2[0].axis("off"), ax2[0].set_title("w/o correction")
+ax2[1].imshow(mri_data_adj_mfi), ax2[1].axis("off"), ax2[1].set_title("with correction")
+
+plt.show()
+
+# %%
+# The blurring is significantly reduced using the Off-resonance Corrected
+# operator (right)

pyproject.toml

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

src/mrinufft/__init__.py

@@ -10,6 +10,7 @@
     get_operator,
     check_backend,
     list_backends,
+    get_interpolators_from_fieldmap,
 )
 
 from .trajectories import (
@@ -56,6 +57,7 @@
     "get_operator",
     "check_backend",
     "list_backends",
+    "get_interpolators_from_fieldmap",
     "initialize_2D_radial",
     "initialize_2D_spiral",
     "initialize_2D_fibonacci_spiral",

src/mrinufft/extras/__init__.py

@@ -1,10 +1,13 @@
 """Sensitivity map estimation methods."""
 
+from .data import make_b0map, make_t2smap
 from .smaps import low_frequency
 from .utils import get_smaps
 
 
 __all__ = [
+    "make_b0map",
+    "make_t2smap",
     "low_frequency",
     "get_smaps",
 ]

src/mrinufft/extras/data.py

@@ -0,0 +1,106 @@
+"""Field map generator module."""
+
+import numpy as np
+
+
+def make_b0map(shape, b0range=(-300, 300), mask=None):
+    """
+    Make radial B0 map.
+
+    Parameters
+    ----------
+    shape : tuple[int]
+        Matrix size. Only supports isotropic matrices.
+    b0range : tuple[float], optional
+        Frequency shift range in [Hz]. The default is (-300, 300).
+    mask : np.ndarray
+        Spatial support of the objec. If not provided,
+        build a radial mask with radius = 0.3 * shape
+
+    Returns
+    -------
+    np.ndarray
+        B0 map of shape (*shape) in [Hz],
+        with values included in (*b0range).
+    mask : np.ndarray, optional
+        Spatial support binary mask.
+
+    """
+    assert np.unique(shape).size, ValueError("Only isotropic matriex are supported.")
+    ndim = len(shape)
+    if ndim == 2:
+        radial_mask, fieldmap = _make_disk(shape)
+    elif ndim == 3:
+        radial_mask, fieldmap = _make_sphere(shape)
+    if mask is None:
+        mask = radial_mask
+
+    # build map
+    fieldmap *= mask
+    fieldmap = (b0range[1] - b0range[0]) * fieldmap / fieldmap.max() + b0range[0]  # Hz
+    fieldmap *= mask
+
+    # remove nan
+    fieldmap = np.nan_to_num(fieldmap, neginf=0.0, posinf=0.0)
+
+    return fieldmap.astype(np.float32), mask
+
+
+def make_t2smap(shape, t2svalue=15.0, mask=None):
+    """
+    Make homogeneous T2* map.
+
+    Parameters
+    ----------
+    shape : tuple[int]
+        Matrix size.
+    t2svalue : float, optional
+        Object T2* in [ms]. The default is 15.0.
+    mask : np.ndarray
+        Spatial support of the objec. If not provided,
+        build a radial mask with radius = 0.3 * shape
+
+    Returns
+    -------
+    np.ndarray
+        T2* map of shape (*shape) in [ms].
+    mask : np.ndarray, optional
+        Spatial support binary mask.
+
+    """
+    assert np.unique(shape).size, ValueError("Only isotropic matriex are supported.")
+    ndim = len(shape)
+    if ndim == 2:
+        radial_mask, fieldmap = _make_disk(shape)
+    elif ndim == 3:
+        radial_mask, fieldmap = _make_sphere(shape)
+    if mask is None:
+        mask = radial_mask
+
+    # build map
+    fieldmap = t2svalue * mask  # ms
+
+    # remove nan
+    fieldmap = np.nan_to_num(fieldmap, neginf=0.0, posinf=0.0)
+
+    return fieldmap.astype(np.float32), mask
+
+
+def _make_disk(shape, frac_radius=0.3):
+    """Make circular binary mask."""
+    ny, nx = shape
+    yy, xx = np.mgrid[:ny, :nx]
+    yy, xx = yy - ny // 2, xx - nx // 2
+    yy, xx = yy / ny, xx / nx
+    rr = (xx**2 + yy**2) ** 0.5
+    return rr < frac_radius, rr
+
+
+def _make_sphere(shape, frac_radius=0.3):
+    """Make spherical binary mask."""
+    nz, ny, nx = shape
+    zz, yy, xx = np.mgrid[:nz, :ny, :nx]
+    zz, yy, xx = zz - nz // 2, yy - ny // 2, xx - nx // 2
+    zz, yy, xx = zz / nz, yy / ny, xx / nx
+    rr = (xx**2 + yy**2 + zz**2) ** 0.5
+    return rr < frac_radius, rr

src/mrinufft/operators/__init__.py

@@ -10,7 +10,7 @@
     list_backends,
     check_backend,
 )
-from .off_resonnance import MRIFourierCorrected
+from .off_resonance import MRIFourierCorrected, get_interpolators_from_fieldmap
 from .stacked import MRIStackedNUFFT
 
 #
@@ -28,4 +28,5 @@
     "check_backend",
     "get_operator",
     "list_backends",
+    "get_interpolators_from_fieldmap",
 ]

src/mrinufft/operators/base.py

@@ -352,9 +352,9 @@ def data_consistency(self, image, obs_data):
         """
         return self.adj_op(self.op(image) - obs_data)
 
-    def with_off_resonnance_correction(self, B, C, indices):
+    def with_off_resonance_correction(self, B, C, indices):
         """Return a new operator with Off Resonnance Correction."""
-        from ..off_resonnance import MRIFourierCorrected
+        from ..off_resonance import MRIFourierCorrected
 
         return MRIFourierCorrected(self, B, C, indices)