From 8ca012cd9d96ebb4488eb291a3401d3584189171 Mon Sep 17 00:00:00 2001
From: chaithyagr <chaithyagr@gitlab.com>
Date: Fri, 26 Apr 2024 14:45:12 +0200
Subject: [PATCH 01/44] Added

---
 .../operators/interfaces/nudft_numpy.py       |  40 +++--
 tests/test_autodiff.py                        | 145 ++++++++++++++++++
 tests/test_ndft.py                            |  59 +++----
 3 files changed, 205 insertions(+), 39 deletions(-)
 create mode 100644 tests/test_autodiff.py

diff --git a/src/mrinufft/operators/interfaces/nudft_numpy.py b/src/mrinufft/operators/interfaces/nudft_numpy.py
index 31a9fa18..9d26d02f 100644
--- a/src/mrinufft/operators/interfaces/nudft_numpy.py
+++ b/src/mrinufft/operators/interfaces/nudft_numpy.py
@@ -8,6 +8,7 @@
 from ..base import FourierOperatorCPU
 
 
+<<<<<<< Updated upstream
 def get_fourier_matrix(ktraj, shape):
     """Get the NDFT Fourier Matrix."""
     n = np.prod(shape)
@@ -17,47 +18,64 @@ def get_fourier_matrix(ktraj, shape):
     grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (ndim, np.prod(shape)))
     traj_grid = ktraj @ grid_r
     matrix = np.exp(-2j * np.pi * traj_grid)
+=======
+def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False):
+    """Get the NDFT Fourier Matrix."""
+    n = np.prod(shape)
+    ndim = len(shape)
+    matrix = np.zeros((len(ktraj), n), dtype=dtype)
+    r = [np.linspace(-s/2, s/2-1, s) for s in shape]
+    grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (ndim, np.prod(shape)))
+    traj_grid = ktraj @ grid_r
+    matrix = np.exp(-2j * np.pi * traj_grid, dtype=dtype)
+    if normalize:
+        matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)))
+>>>>>>> Stashed changes
     return matrix
 
 
-def implicit_type2_ndft(ktraj, image, shape):
+def implicit_type2_ndft(ktraj, image, shape, normalize=False):
     """Compute the NDFT using the implicit type 2 (image -> kspace) algorithm."""
-    r = [np.arange(s) for s in shape]
+    r = [np.linspace(-s/2, s/2-1, s) for s in shape]
     grid_r = np.reshape(
         np.meshgrid(*r, indexing="ij"), (len(shape), np.prod(image.shape))
     )
     res = np.zeros(len(ktraj), dtype=image.dtype)
     for j in range(np.prod(image.shape)):
         res += image[j] * np.exp(-2j * np.pi * ktraj @ grid_r[:, j])
+    if normalize:
+        matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)))
     return res
 
 
-def implicit_type1_ndft(ktraj, coeffs, shape):
+def implicit_type1_ndft(ktraj, coeffs, shape, normalize=False):
     """Compute the NDFT using the implicit type 1 (kspace -> image) algorithm."""
-    r = [np.arange(s) for s in shape]
+    r = [np.linspace(-s/2, s/2-1, s) for s in shape]
     grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (len(shape), np.prod(shape)))
     res = np.zeros(np.prod(shape), dtype=coeffs.dtype)
     for i in range(len(ktraj)):
         res += coeffs[i] * np.exp(2j * np.pi * ktraj[i] @ grid_r)
+    if normalize:
+        matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)))
     return res
 
 
-def get_implicit_matrix(ktraj, shape):
+def get_implicit_matrix(ktraj, shape, normalize=False):
     """Get the NDFT Fourier Matrix as implicit operator.
 
     This is more memory efficient than the explicit matrix.
     """
     return sp.sparse.linalg.LinearOperator(
         (len(ktraj), np.prod(shape)),
-        matvec=lambda x: implicit_type2_ndft(ktraj, x, shape),
-        rmatvec=lambda x: implicit_type1_ndft(ktraj, x, shape),
+        matvec=lambda x: implicit_type2_ndft(ktraj, x, shape, normalize),
+        rmatvec=lambda x: implicit_type1_ndft(ktraj, x, shape, normalize),
     )
 
 
 class RawNDFT:
     """Implementation of the NUDFT using numpy."""
 
-    def __init__(self, samples, shape, explicit_matrix=True):
+    def __init__(self, samples, shape, explicit_matrix=True, normalize=False):
         self.samples = samples
         self.shape = shape
         self.n_samples = len(samples)
@@ -65,13 +83,13 @@ def __init__(self, samples, shape, explicit_matrix=True):
         if explicit_matrix:
             try:
                 self._fourier_matrix = sp.sparse.linalg.aslinearoperator(
-                    get_fourier_matrix(self.samples, self.shape)
+                    get_fourier_matrix(self.samples, self.shape, normalize=normalize)
                 )
             except MemoryError:
                 warnings.warn("Not enough memory, using an implicit definition anyway")
-                self._fourier_matrix = get_implicit_matrix(self.samples, self.shape)
+                self._fourier_matrix = get_implicit_matrix(self.samples, self.shape, normalize)
         else:
-            self._fourier_matrix = get_implicit_matrix(self.samples, self.shape)
+            self._fourier_matrix = get_implicit_matrix(self.samples, self.shape, normalize)
 
     def op(self, coeffs, image):
         """Compute the forward NUDFT."""
diff --git a/tests/test_autodiff.py b/tests/test_autodiff.py
new file mode 100644
index 00000000..01de2f91
--- /dev/null
+++ b/tests/test_autodiff.py
@@ -0,0 +1,145 @@
+"""Test the autodiff functionnality."""
+
+import numpy as np
+from mrinufft.operators.interfaces.nudft_numpy import get_fourier_matrix
+import pytest
+from pytest_cases import parametrize_with_cases, parametrize, fixture
+from case_trajectories import CasesTrajectories
+from mrinufft.operators import get_operator
+
+
+from helpers import (
+    kspace_from_op,
+    image_from_op,
+    to_interface,
+)
+
+
+TORCH_AVAILABLE = True
+try:
+    import torch
+    import torch.testing as tt
+except ImportError:
+    TORCH_AVAILABLE = False
+
+
+@fixture(scope="module")
+@parametrize(backend=["cufinufft", "finufft"])
+@parametrize(autograd=["data"])
+@parametrize_with_cases(
+    "kspace_loc, shape",
+    cases=[
+        CasesTrajectories.case_grid2D,
+        CasesTrajectories.case_nyquist_radial2D,
+    ],  # 2D cases only for reduced memory footprint.
+)
+def operator(kspace_loc, shape, backend, autograd):
+    """Create NUFFT operator with autodiff capabilities."""
+    kspace_loc = kspace_loc.astype(np.float32)
+
+    nufft = get_operator(backend_name=backend, autograd=autograd)(
+        samples=kspace_loc,
+        shape=shape,
+        smaps=None,
+    )
+
+    return nufft
+
+
+@fixture(scope="module")
+def ndft_matrix(operator):
+    """Get the NDFT matrix from the operator."""
+    return get_fourier_matrix(operator.samples, operator.shape, normalize=True)
+
+
+@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"])
+@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed")
+def test_adjoint_and_grad(operator, ndft_matrix, interface):
+    """Test the adjoint and gradient of the operator."""
+    if operator.backend == "finufft" and "gpu" in interface:
+        pytest.skip("GPU not supported for finufft backend")
+    ndft_matrix_torch = to_interface(ndft_matrix, interface=interface)
+    ksp_data = to_interface(kspace_from_op(operator), interface=interface)
+    img_data = to_interface(image_from_op(operator), interface=interface)
+    ksp_data.requires_grad = True
+
+    with torch.autograd.set_detect_anomaly(True):
+        adj_data = operator.adj_op(ksp_data).reshape(img_data.shape)
+        adj_data_ndft = (ndft_matrix_torch.conj().T @ ksp_data.flatten()).reshape(
+            adj_data.shape
+        )
+        loss_nufft = torch.mean(torch.abs(adj_data) ** 2)
+        loss_ndft = torch.mean(torch.abs(adj_data_ndft) ** 2)
+
+    # Check if nufft and ndft are close in the backprop
+    grad_ndft_kdata = torch.autograd.grad(loss_ndft, ksp_data, retain_graph=True)[0]
+    grad_nufft_kdata = torch.autograd.grad(loss_nufft, ksp_data, retain_graph=True)[0]
+    tt.assert_close(grad_ndft_kdata, grad_nufft_kdata, rtol=1, atol=1)
+
+
+@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"])
+@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed")
+def test_adjoint_and_gradauto(operator, ndft_matrix, interface):
+    """Test the adjoint and gradient of the operator using autograd gradcheck."""
+    if operator.backend == "finufft" and "gpu" in interface:
+        pytest.skip("GPU not supported for finufft backend")
+
+    ksp_data = to_interface(kspace_from_op(operator), interface=interface)
+    ksp_data = torch.ones(ksp_data.shape, requires_grad=True, dtype=ksp_data.dtype)
+    print(ksp_data.shape)
+    # todo: tighten the tolerance
+    assert torch.autograd.gradcheck(
+        operator.adjoint,
+        ksp_data,
+        eps=1e-10,
+        rtol=1,
+        atol=1,
+        nondet_tol=1,
+        raise_exception=True,
+    )
+
+
+@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"])
+@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed")
+def test_forward_and_grad(operator, ndft_matrix, interface):
+    """Test the adjoint and gradient of the operator."""
+    if operator.backend == "finufft" and "gpu" in interface:
+        pytest.skip("GPU not supported for finufft backend")
+
+    ndft_matrix_torch = to_interface(ndft_matrix, interface=interface)
+    ksp_data_ref = to_interface(kspace_from_op(operator), interface=interface)
+    img_data = to_interface(image_from_op(operator), interface=interface)
+    img_data.requires_grad = True
+
+    with torch.autograd.set_detect_anomaly(True):
+        ksp_data = operator.op(img_data).reshape(ksp_data_ref.shape)
+        ksp_data_ndft = (ndft_matrix_torch @ img_data.flatten()).reshape(ksp_data.shape)
+        loss_nufft = torch.mean(torch.abs(ksp_data - ksp_data_ref) ** 2)
+        loss_ndft = torch.mean(torch.abs(ksp_data_ndft - ksp_data_ref) ** 2)
+
+    # Check if nufft and ndft are close in the backprop
+    grad_ndft_kdata = torch.autograd.grad(loss_ndft, img_data, retain_graph=True)[0]
+    grad_nufft_kdata = torch.autograd.grad(loss_nufft, img_data, retain_graph=True)[0]
+    assert torch.allclose(grad_ndft_kdata, grad_nufft_kdata, atol=6e-3)
+
+
+@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"])
+@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed")
+def test_forward_and_gradauto(operator, ndft_matrix, interface):
+    """Test the forward and gradient of the operator using autograd gradcheck."""
+    if operator.backend == "finufft" and "gpu" in interface:
+        pytest.skip("GPU not supported for finufft backend")
+
+    img_data = to_interface(image_from_op(operator), interface=interface)
+    img_data = torch.ones(img_data.shape, requires_grad=True, dtype=img_data.dtype)
+    print(img_data.shape)
+    # todo: tighten the tolerance
+    assert torch.autograd.gradcheck(
+        operator.adjoint,
+        img_data,
+        eps=1e-10,
+        rtol=1,
+        atol=1,
+        nondet_tol=1,
+        raise_exception=True,
+    )
diff --git a/tests/test_ndft.py b/tests/test_ndft.py
index 7ae595a8..5bcb3e8c 100644
--- a/tests/test_ndft.py
+++ b/tests/test_ndft.py
@@ -13,32 +13,9 @@
 
 from case_trajectories import CasesTrajectories, case_grid1D
 from helpers import assert_almost_allclose
+from mrinufft import get_operator
 
 
-@parametrize_with_cases(
-    "kspace_grid, shape",
-    cases=[
-        case_grid1D,
-        CasesTrajectories.case_grid2D,
-    ],  # 3D is ignored (too much possibility for the reordering)
-)
-def test_ndft_grid_matrix(kspace_grid, shape):
-    """Test that  the ndft matrix is a good matrix for doing fft."""
-    ndft_matrix = get_fourier_matrix(kspace_grid, shape)
-    # Create a random image
-    fft_matrix = [None] * len(shape)
-    for i in range(len(shape)):
-        fft_matrix[i] = sp.fft.fft(np.eye(shape[i]))
-    fft_mat = fft_matrix[0]
-    if len(shape) == 2:
-        fft_mat = fft_matrix[0].flatten()[:, None] @ fft_matrix[1].flatten()[None, :]
-        fft_mat = (
-            fft_mat.reshape(shape * 2)
-            .transpose(2, 0, 1, 3)
-            .reshape((np.prod(shape),) * 2)
-        )
-    assert np.allclose(ndft_matrix, fft_mat)
-
 
 @parametrize_with_cases(
     "kspace, shape",
@@ -56,7 +33,7 @@ def test_ndft_implicit2(kspace, shape):
     linop_coef = implicit_type2_ndft(kspace, random_image.flatten(), shape)
     matrix_coef = matrix @ random_image.flatten()
 
-    assert np.allclose(linop_coef, matrix_coef)
+    assert_almost_allclose(linop_coef, matrix_coef, atol=1e-4, rtol=1e-4, mismatch=5)
 
 
 @parametrize_with_cases(
@@ -76,7 +53,32 @@ def test_ndft_implicit1(kspace, shape):
     linop_coef = implicit_type1_ndft(kspace, random_kspace.flatten(), shape)
     matrix_coef = matrix.conj().T @ random_kspace.flatten()
 
-    assert np.allclose(linop_coef, matrix_coef)
+    assert_almost_allclose(linop_coef, matrix_coef, atol=1e-4, rtol=1e-4, mismatch=5)
+
+@parametrize_with_cases(
+    "kspace, shape",
+    cases=[
+        CasesTrajectories.case_random2D,
+        CasesTrajectories.case_grid2D,
+        CasesTrajectories.case_grid3D,
+    ],
+)
+def test_ndft_nufft(kspace, shape):
+    "Test that NDFT matches NUFFT"
+    ndft_op = RawNDFT(kspace, shape, normalize=True)
+    random_kspace = 1j * np.random.randn(len(kspace))
+    random_kspace += np.random.randn(len(kspace))
+    random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape)
+    operator = get_operator("pynfft")(kspace, shape)  # FIXME: @PAC, we need to get ref here
+    nufft_k = operator.op(random_image)
+    nufft_i = operator.adj_op(random_kspace)
+    
+    ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype)
+    ndft_i = np.empty(shape, dtype=random_kspace.dtype)
+    ndft_op.op(ndft_k, random_image)
+    ndft_op.adj_op(random_kspace, ndft_i)
+    assert_almost_allclose(nufft_k, ndft_k, atol=1e-4, rtol=1e-4, mismatch=5)
+    assert_almost_allclose(nufft_i, ndft_i, atol=1e-4, rtol=1e-4, mismatch=5)
 
 
 @parametrize_with_cases(
@@ -98,6 +100,7 @@ def test_ndft_fft(kspace_grid, shape):
     kspace = kspace.reshape(img.shape)
     if len(shape) >= 2:
         kspace = kspace.swapaxes(0, 1)
-    kspace_fft = sp.fft.fftn(img)
+    kspace_fft = sp.fft.fftn(sp.fft.fftshift(img))
+
+    assert_almost_allclose(kspace, kspace_fft, atol=1e-4, rtol=1e-4, mismatch=5)
 
-    assert_almost_allclose(kspace, kspace_fft, atol=1e-5, rtol=1e-5, mismatch=5)

From 093edfbe3def12e22c4aec3657f70377c238525b Mon Sep 17 00:00:00 2001
From: chaithyagr <chaithyagr@gitlab.com>
Date: Fri, 26 Apr 2024 14:46:50 +0200
Subject: [PATCH 02/44] Fix

---
 src/mrinufft/operators/interfaces/nudft_numpy.py | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/src/mrinufft/operators/interfaces/nudft_numpy.py b/src/mrinufft/operators/interfaces/nudft_numpy.py
index 9d26d02f..68690cc1 100644
--- a/src/mrinufft/operators/interfaces/nudft_numpy.py
+++ b/src/mrinufft/operators/interfaces/nudft_numpy.py
@@ -8,17 +8,6 @@
 from ..base import FourierOperatorCPU
 
 
-<<<<<<< Updated upstream
-def get_fourier_matrix(ktraj, shape):
-    """Get the NDFT Fourier Matrix."""
-    n = np.prod(shape)
-    ndim = len(shape)
-    matrix = np.zeros((len(ktraj), n), dtype=complex)
-    r = [np.arange(shape[i]) for i in range(ndim)]
-    grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (ndim, np.prod(shape)))
-    traj_grid = ktraj @ grid_r
-    matrix = np.exp(-2j * np.pi * traj_grid)
-=======
 def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False):
     """Get the NDFT Fourier Matrix."""
     n = np.prod(shape)
@@ -30,7 +19,6 @@ def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False):
     matrix = np.exp(-2j * np.pi * traj_grid, dtype=dtype)
     if normalize:
         matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)))
->>>>>>> Stashed changes
     return matrix
 
 

From 74c1ecd6d0a07dcc141156aa506be9d1967d07e5 Mon Sep 17 00:00:00 2001
From: chaithyagr <chaithyagr@gitlab.com>
Date: Fri, 26 Apr 2024 14:48:12 +0200
Subject: [PATCH 03/44] Remove bymistake add

---
 tests/test_autodiff.py | 145 -----------------------------------------
 1 file changed, 145 deletions(-)
 delete mode 100644 tests/test_autodiff.py

diff --git a/tests/test_autodiff.py b/tests/test_autodiff.py
deleted file mode 100644
index 01de2f91..00000000
--- a/tests/test_autodiff.py
+++ /dev/null
@@ -1,145 +0,0 @@
-"""Test the autodiff functionnality."""
-
-import numpy as np
-from mrinufft.operators.interfaces.nudft_numpy import get_fourier_matrix
-import pytest
-from pytest_cases import parametrize_with_cases, parametrize, fixture
-from case_trajectories import CasesTrajectories
-from mrinufft.operators import get_operator
-
-
-from helpers import (
-    kspace_from_op,
-    image_from_op,
-    to_interface,
-)
-
-
-TORCH_AVAILABLE = True
-try:
-    import torch
-    import torch.testing as tt
-except ImportError:
-    TORCH_AVAILABLE = False
-
-
-@fixture(scope="module")
-@parametrize(backend=["cufinufft", "finufft"])
-@parametrize(autograd=["data"])
-@parametrize_with_cases(
-    "kspace_loc, shape",
-    cases=[
-        CasesTrajectories.case_grid2D,
-        CasesTrajectories.case_nyquist_radial2D,
-    ],  # 2D cases only for reduced memory footprint.
-)
-def operator(kspace_loc, shape, backend, autograd):
-    """Create NUFFT operator with autodiff capabilities."""
-    kspace_loc = kspace_loc.astype(np.float32)
-
-    nufft = get_operator(backend_name=backend, autograd=autograd)(
-        samples=kspace_loc,
-        shape=shape,
-        smaps=None,
-    )
-
-    return nufft
-
-
-@fixture(scope="module")
-def ndft_matrix(operator):
-    """Get the NDFT matrix from the operator."""
-    return get_fourier_matrix(operator.samples, operator.shape, normalize=True)
-
-
-@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"])
-@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed")
-def test_adjoint_and_grad(operator, ndft_matrix, interface):
-    """Test the adjoint and gradient of the operator."""
-    if operator.backend == "finufft" and "gpu" in interface:
-        pytest.skip("GPU not supported for finufft backend")
-    ndft_matrix_torch = to_interface(ndft_matrix, interface=interface)
-    ksp_data = to_interface(kspace_from_op(operator), interface=interface)
-    img_data = to_interface(image_from_op(operator), interface=interface)
-    ksp_data.requires_grad = True
-
-    with torch.autograd.set_detect_anomaly(True):
-        adj_data = operator.adj_op(ksp_data).reshape(img_data.shape)
-        adj_data_ndft = (ndft_matrix_torch.conj().T @ ksp_data.flatten()).reshape(
-            adj_data.shape
-        )
-        loss_nufft = torch.mean(torch.abs(adj_data) ** 2)
-        loss_ndft = torch.mean(torch.abs(adj_data_ndft) ** 2)
-
-    # Check if nufft and ndft are close in the backprop
-    grad_ndft_kdata = torch.autograd.grad(loss_ndft, ksp_data, retain_graph=True)[0]
-    grad_nufft_kdata = torch.autograd.grad(loss_nufft, ksp_data, retain_graph=True)[0]
-    tt.assert_close(grad_ndft_kdata, grad_nufft_kdata, rtol=1, atol=1)
-
-
-@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"])
-@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed")
-def test_adjoint_and_gradauto(operator, ndft_matrix, interface):
-    """Test the adjoint and gradient of the operator using autograd gradcheck."""
-    if operator.backend == "finufft" and "gpu" in interface:
-        pytest.skip("GPU not supported for finufft backend")
-
-    ksp_data = to_interface(kspace_from_op(operator), interface=interface)
-    ksp_data = torch.ones(ksp_data.shape, requires_grad=True, dtype=ksp_data.dtype)
-    print(ksp_data.shape)
-    # todo: tighten the tolerance
-    assert torch.autograd.gradcheck(
-        operator.adjoint,
-        ksp_data,
-        eps=1e-10,
-        rtol=1,
-        atol=1,
-        nondet_tol=1,
-        raise_exception=True,
-    )
-
-
-@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"])
-@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed")
-def test_forward_and_grad(operator, ndft_matrix, interface):
-    """Test the adjoint and gradient of the operator."""
-    if operator.backend == "finufft" and "gpu" in interface:
-        pytest.skip("GPU not supported for finufft backend")
-
-    ndft_matrix_torch = to_interface(ndft_matrix, interface=interface)
-    ksp_data_ref = to_interface(kspace_from_op(operator), interface=interface)
-    img_data = to_interface(image_from_op(operator), interface=interface)
-    img_data.requires_grad = True
-
-    with torch.autograd.set_detect_anomaly(True):
-        ksp_data = operator.op(img_data).reshape(ksp_data_ref.shape)
-        ksp_data_ndft = (ndft_matrix_torch @ img_data.flatten()).reshape(ksp_data.shape)
-        loss_nufft = torch.mean(torch.abs(ksp_data - ksp_data_ref) ** 2)
-        loss_ndft = torch.mean(torch.abs(ksp_data_ndft - ksp_data_ref) ** 2)
-
-    # Check if nufft and ndft are close in the backprop
-    grad_ndft_kdata = torch.autograd.grad(loss_ndft, img_data, retain_graph=True)[0]
-    grad_nufft_kdata = torch.autograd.grad(loss_nufft, img_data, retain_graph=True)[0]
-    assert torch.allclose(grad_ndft_kdata, grad_nufft_kdata, atol=6e-3)
-
-
-@pytest.mark.parametrize("interface", ["torch-gpu", "torch-cpu"])
-@pytest.mark.skipif(not TORCH_AVAILABLE, reason="Pytorch is not installed")
-def test_forward_and_gradauto(operator, ndft_matrix, interface):
-    """Test the forward and gradient of the operator using autograd gradcheck."""
-    if operator.backend == "finufft" and "gpu" in interface:
-        pytest.skip("GPU not supported for finufft backend")
-
-    img_data = to_interface(image_from_op(operator), interface=interface)
-    img_data = torch.ones(img_data.shape, requires_grad=True, dtype=img_data.dtype)
-    print(img_data.shape)
-    # todo: tighten the tolerance
-    assert torch.autograd.gradcheck(
-        operator.adjoint,
-        img_data,
-        eps=1e-10,
-        rtol=1,
-        atol=1,
-        nondet_tol=1,
-        raise_exception=True,
-    )

From 0250aa8d3753bad0191cdc5f42cd1c112f589f44 Mon Sep 17 00:00:00 2001
From: chaithyagr <chaithyagr@gitlab.com>
Date: Fri, 26 Apr 2024 15:38:27 +0200
Subject: [PATCH 04/44] Fix

---
 tests/test_ndft.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_ndft.py b/tests/test_ndft.py
index 5bcb3e8c..3d972ff5 100644
--- a/tests/test_ndft.py
+++ b/tests/test_ndft.py
@@ -64,7 +64,7 @@ def test_ndft_implicit1(kspace, shape):
     ],
 )
 def test_ndft_nufft(kspace, shape):
-    "Test that NDFT matches NUFFT"
+    """Test that NDFT matches NUFFT"""
     ndft_op = RawNDFT(kspace, shape, normalize=True)
     random_kspace = 1j * np.random.randn(len(kspace))
     random_kspace += np.random.randn(len(kspace))

From 060a8bdd125d2140b82dfaa6a492ec78682a80bf Mon Sep 17 00:00:00 2001
From: chaithyagr <chaithyagr@gitlab.com>
Date: Fri, 26 Apr 2024 15:39:16 +0200
Subject: [PATCH 05/44] Fixed lint

---
 .../operators/interfaces/nudft_numpy.py       | 20 +++++++++++--------
 tests/test_ndft.py                            |  9 +++++----
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/mrinufft/operators/interfaces/nudft_numpy.py b/src/mrinufft/operators/interfaces/nudft_numpy.py
index 68690cc1..3e8e81aa 100644
--- a/src/mrinufft/operators/interfaces/nudft_numpy.py
+++ b/src/mrinufft/operators/interfaces/nudft_numpy.py
@@ -13,18 +13,18 @@ def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False):
     n = np.prod(shape)
     ndim = len(shape)
     matrix = np.zeros((len(ktraj), n), dtype=dtype)
-    r = [np.linspace(-s/2, s/2-1, s) for s in shape]
+    r = [np.linspace(-s / 2, s / 2 - 1, s) for s in shape]
     grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (ndim, np.prod(shape)))
     traj_grid = ktraj @ grid_r
     matrix = np.exp(-2j * np.pi * traj_grid, dtype=dtype)
     if normalize:
-        matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)))
+        matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))
     return matrix
 
 
 def implicit_type2_ndft(ktraj, image, shape, normalize=False):
     """Compute the NDFT using the implicit type 2 (image -> kspace) algorithm."""
-    r = [np.linspace(-s/2, s/2-1, s) for s in shape]
+    r = [np.linspace(-s / 2, s / 2 - 1, s) for s in shape]
     grid_r = np.reshape(
         np.meshgrid(*r, indexing="ij"), (len(shape), np.prod(image.shape))
     )
@@ -32,19 +32,19 @@ def implicit_type2_ndft(ktraj, image, shape, normalize=False):
     for j in range(np.prod(image.shape)):
         res += image[j] * np.exp(-2j * np.pi * ktraj @ grid_r[:, j])
     if normalize:
-        matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)))
+        matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))
     return res
 
 
 def implicit_type1_ndft(ktraj, coeffs, shape, normalize=False):
     """Compute the NDFT using the implicit type 1 (kspace -> image) algorithm."""
-    r = [np.linspace(-s/2, s/2-1, s) for s in shape]
+    r = [np.linspace(-s / 2, s / 2 - 1, s) for s in shape]
     grid_r = np.reshape(np.meshgrid(*r, indexing="ij"), (len(shape), np.prod(shape)))
     res = np.zeros(np.prod(shape), dtype=coeffs.dtype)
     for i in range(len(ktraj)):
         res += coeffs[i] * np.exp(2j * np.pi * ktraj[i] @ grid_r)
     if normalize:
-        matrix /= (np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape)))
+        matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))
     return res
 
 
@@ -75,9 +75,13 @@ def __init__(self, samples, shape, explicit_matrix=True, normalize=False):
                 )
             except MemoryError:
                 warnings.warn("Not enough memory, using an implicit definition anyway")
-                self._fourier_matrix = get_implicit_matrix(self.samples, self.shape, normalize)
+                self._fourier_matrix = get_implicit_matrix(
+                    self.samples, self.shape, normalize
+                )
         else:
-            self._fourier_matrix = get_implicit_matrix(self.samples, self.shape, normalize)
+            self._fourier_matrix = get_implicit_matrix(
+                self.samples, self.shape, normalize
+            )
 
     def op(self, coeffs, image):
         """Compute the forward NUDFT."""
diff --git a/tests/test_ndft.py b/tests/test_ndft.py
index 3d972ff5..7f90d14e 100644
--- a/tests/test_ndft.py
+++ b/tests/test_ndft.py
@@ -16,7 +16,6 @@
 from mrinufft import get_operator
 
 
-
 @parametrize_with_cases(
     "kspace, shape",
     cases=[
@@ -55,6 +54,7 @@ def test_ndft_implicit1(kspace, shape):
 
     assert_almost_allclose(linop_coef, matrix_coef, atol=1e-4, rtol=1e-4, mismatch=5)
 
+
 @parametrize_with_cases(
     "kspace, shape",
     cases=[
@@ -69,10 +69,12 @@ def test_ndft_nufft(kspace, shape):
     random_kspace = 1j * np.random.randn(len(kspace))
     random_kspace += np.random.randn(len(kspace))
     random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape)
-    operator = get_operator("pynfft")(kspace, shape)  # FIXME: @PAC, we need to get ref here
+    operator = get_operator("pynfft")(
+        kspace, shape
+    )  # FIXME: @PAC, we need to get ref here
     nufft_k = operator.op(random_image)
     nufft_i = operator.adj_op(random_kspace)
-    
+
     ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype)
     ndft_i = np.empty(shape, dtype=random_kspace.dtype)
     ndft_op.op(ndft_k, random_image)
@@ -103,4 +105,3 @@ def test_ndft_fft(kspace_grid, shape):
     kspace_fft = sp.fft.fftn(sp.fft.fftshift(img))
 
     assert_almost_allclose(kspace, kspace_fft, atol=1e-4, rtol=1e-4, mismatch=5)
-

From aecb844c74ae53ae67deb852204bc9e647ac28fd Mon Sep 17 00:00:00 2001
From: chaithyagr <chaithyagr@gitlab.com>
Date: Fri, 26 Apr 2024 15:40:50 +0200
Subject: [PATCH 06/44] Lint

---
 src/mrinufft/operators/interfaces/nudft_numpy.py | 4 ++--
 tests/test_ndft.py                               | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/mrinufft/operators/interfaces/nudft_numpy.py b/src/mrinufft/operators/interfaces/nudft_numpy.py
index 3e8e81aa..bcc6c033 100644
--- a/src/mrinufft/operators/interfaces/nudft_numpy.py
+++ b/src/mrinufft/operators/interfaces/nudft_numpy.py
@@ -32,7 +32,7 @@ def implicit_type2_ndft(ktraj, image, shape, normalize=False):
     for j in range(np.prod(image.shape)):
         res += image[j] * np.exp(-2j * np.pi * ktraj @ grid_r[:, j])
     if normalize:
-        matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))
+        res /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))
     return res
 
 
@@ -44,7 +44,7 @@ def implicit_type1_ndft(ktraj, coeffs, shape, normalize=False):
     for i in range(len(ktraj)):
         res += coeffs[i] * np.exp(2j * np.pi * ktraj[i] @ grid_r)
     if normalize:
-        matrix /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))
+        res /= np.sqrt(np.prod(shape)) * np.power(np.sqrt(2), len(shape))
     return res
 
 
diff --git a/tests/test_ndft.py b/tests/test_ndft.py
index 7f90d14e..fa66b8b2 100644
--- a/tests/test_ndft.py
+++ b/tests/test_ndft.py
@@ -64,7 +64,7 @@ def test_ndft_implicit1(kspace, shape):
     ],
 )
 def test_ndft_nufft(kspace, shape):
-    """Test that NDFT matches NUFFT"""
+    """Test that NDFT matches NUFFT."""
     ndft_op = RawNDFT(kspace, shape, normalize=True)
     random_kspace = 1j * np.random.randn(len(kspace))
     random_kspace += np.random.randn(len(kspace))

From 3130bc1c5f443294a2f71dcae30178bb8357d392 Mon Sep 17 00:00:00 2001
From: chaithyagr <chaithyagr@gitlab.com>
Date: Fri, 26 Apr 2024 17:15:00 +0200
Subject: [PATCH 07/44] Added refbackend

---
 tests/test_ndft.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/test_ndft.py b/tests/test_ndft.py
index fa66b8b2..57aedfa6 100644
--- a/tests/test_ndft.py
+++ b/tests/test_ndft.py
@@ -63,18 +63,18 @@ def test_ndft_implicit1(kspace, shape):
         CasesTrajectories.case_grid3D,
     ],
 )
-def test_ndft_nufft(kspace, shape):
+def test_ndft_nufft(kspace, shape, request):
     """Test that NDFT matches NUFFT."""
     ndft_op = RawNDFT(kspace, shape, normalize=True)
     random_kspace = 1j * np.random.randn(len(kspace))
     random_kspace += np.random.randn(len(kspace))
     random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape)
-    operator = get_operator("pynfft")(
+    operator = get_operator(request.config.getoption("ref"))(
         kspace, shape
-    )  # FIXME: @PAC, we need to get ref here
+    )
     nufft_k = operator.op(random_image)
     nufft_i = operator.adj_op(random_kspace)
-
+    
     ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype)
     ndft_i = np.empty(shape, dtype=random_kspace.dtype)
     ndft_op.op(ndft_k, random_image)

From bc014b8973e3b355f17365a9eb933cc57b92fb4b Mon Sep 17 00:00:00 2001
From: chaithyagr <chaithyagr@gitlab.com>
Date: Fri, 26 Apr 2024 17:17:48 +0200
Subject: [PATCH 08/44] Fix NDFT

---
 tests/test_ndft.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/tests/test_ndft.py b/tests/test_ndft.py
index 57aedfa6..7a157d34 100644
--- a/tests/test_ndft.py
+++ b/tests/test_ndft.py
@@ -69,12 +69,10 @@ def test_ndft_nufft(kspace, shape, request):
     random_kspace = 1j * np.random.randn(len(kspace))
     random_kspace += np.random.randn(len(kspace))
     random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape)
-    operator = get_operator(request.config.getoption("ref"))(
-        kspace, shape
-    )
+    operator = get_operator(request.config.getoption("ref"))(kspace, shape)
     nufft_k = operator.op(random_image)
     nufft_i = operator.adj_op(random_kspace)
-    
+
     ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype)
     ndft_i = np.empty(shape, dtype=random_kspace.dtype)
     ndft_op.op(ndft_k, random_image)

From 0cc73c41cf743ea19ffa053f0cd43b54f43f192e Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Mon, 29 Apr 2024 10:48:25 +0200
Subject: [PATCH 09/44] feat: use finufft as ref backend.

---
 tests/conftest.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 69598fdb..4e89f0ed 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -15,7 +15,7 @@ def pytest_addoption(parser):
     )
     parser.addoption(
         "--ref",
-        default="pynfft",
+        default="finufft",
         help="Reference backend on which the tests are performed.",
     )
 

From 21e090f21803e9e57cc721010661d30449ce1a0b Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Mon, 29 Apr 2024 10:49:37 +0200
Subject: [PATCH 10/44] feat(tests): move ndft vs nufft tests to own file.

---
 tests/operators/test_operator_ref.py | 74 ++++++++++++++++++++++++++++
 tests/test_ndft.py                   | 26 ----------
 2 files changed, 74 insertions(+), 26 deletions(-)
 create mode 100644 tests/operators/test_operator_ref.py

diff --git a/tests/operators/test_operator_ref.py b/tests/operators/test_operator_ref.py
new file mode 100644
index 00000000..b51e1633
--- /dev/null
+++ b/tests/operators/test_operator_ref.py
@@ -0,0 +1,74 @@
+"""Tests for the reference backend."""
+
+from pytest_cases import parametrize_with_cases, fixture
+from case_trajectories import CasesTrajectories
+
+from mrinufft import get_operator
+from mrinufft.operators.interfaces.nudft_numpy import MRInumpy
+from helpers import assert_almost_allclose, kspace_from_op, image_from_op
+
+
+@fixture(scope="session", autouse=True)
+def ref_backend(request):
+    """Get the reference backend from the CLI."""
+    return request.config.getoption("ref")
+
+
+@fixture(scope="module")
+@parametrize_with_cases(
+    "kspace, shape",
+    cases=[
+        CasesTrajectories.case_random2D,
+        CasesTrajectories.case_grid2D,
+        CasesTrajectories.case_grid3D,
+    ],
+)
+def ref_operator(request, ref_backend, kspace, shape):
+    """Generate a NFFT operator, matching the property of the first operator."""
+    return get_operator(ref_backend)(kspace, shape)
+
+
+@fixture(scope="module")
+def ndft_operator(ref_operator):
+    """Get a NDFT operator matching the reference operator."""
+    return MRInumpy(ref_operator.samples, ref_operator.shape)
+
+
+@fixture(scope="module")
+def image_data(ref_operator):
+    """Generate a random image. Remains constant for the module."""
+    return image_from_op(ref_operator)
+
+
+@fixture(scope="module")
+def kspace_data(ref_operator):
+    """Generate a random kspace. Remains constant for the module."""
+    return kspace_from_op(ref_operator)
+
+
+def test_ref_nufft_forward(ref_operator, ndft_operator, image_data):
+    """Test that the reference nufft matches the NDFT."""
+    nufft_ksp = ref_operator.op(image_data)
+    ndft_ksp = ndft_operator.op(image_data)
+
+    assert_almost_allclose(
+        nufft_ksp,
+        ndft_ksp,
+        atol=1e-4,
+        rtol=1e-4,
+        mismatch=5,
+    )
+
+
+def test_ref_nufft_adjoint(ref_operator, ndft_operator, kspace_data):
+    """Test that the reference nufft matches the NDFT adjoint."""
+    nufft_img = ref_operator.adj_op(kspace_data)
+    ndft_img = ndft_operator.adj_op(kspace_data)
+
+    assert_almost_allclose(
+        nufft_img,
+        ndft_img,
+        atol=1e-4,
+        rtol=1e-4,
+        mismatch=5,
+    )
diff --git a/tests/test_ndft.py b/tests/test_ndft.py
index 7a157d34..cd21622e 100644
--- a/tests/test_ndft.py
+++ b/tests/test_ndft.py
@@ -55,32 +55,6 @@ def test_ndft_implicit1(kspace, shape):
     assert_almost_allclose(linop_coef, matrix_coef, atol=1e-4, rtol=1e-4, mismatch=5)
 
 
-@parametrize_with_cases(
-    "kspace, shape",
-    cases=[
-        CasesTrajectories.case_random2D,
-        CasesTrajectories.case_grid2D,
-        CasesTrajectories.case_grid3D,
-    ],
-)
-def test_ndft_nufft(kspace, shape, request):
-    """Test that NDFT matches NUFFT."""
-    ndft_op = RawNDFT(kspace, shape, normalize=True)
-    random_kspace = 1j * np.random.randn(len(kspace))
-    random_kspace += np.random.randn(len(kspace))
-    random_image = np.random.randn(*shape) + 1j * np.random.randn(*shape)
-    operator = get_operator(request.config.getoption("ref"))(kspace, shape)
-    nufft_k = operator.op(random_image)
-    nufft_i = operator.adj_op(random_kspace)
-
-    ndft_k = np.empty(ndft_op.n_samples, dtype=random_image.dtype)
-    ndft_i = np.empty(shape, dtype=random_kspace.dtype)
-    ndft_op.op(ndft_k, random_image)
-    ndft_op.adj_op(random_kspace, ndft_i)
-    assert_almost_allclose(nufft_k, ndft_k, atol=1e-4, rtol=1e-4, mismatch=5)
-    assert_almost_allclose(nufft_i, ndft_i, atol=1e-4, rtol=1e-4, mismatch=5)
-
-
 @parametrize_with_cases(
     "kspace_grid, shape",
     cases=[

From 96a8287ea9141532da49ca6ff892fda55f364d5f Mon Sep 17 00:00:00 2001
From: mcencini <matteo.cencini@gmail.com>
Date: Thu, 8 Aug 2024 10:49:53 +0200
Subject: [PATCH 11/44] fix typoes (off_resonnance -> off_resonance)

---
 src/mrinufft/operators/__init__.py                            | 2 +-
 src/mrinufft/operators/base.py                                | 4 ++--
 .../operators/{off_resonnance.py => off_resonance.py}         | 0
 3 files changed, 3 insertions(+), 3 deletions(-)
 rename src/mrinufft/operators/{off_resonnance.py => off_resonance.py} (100%)

diff --git a/src/mrinufft/operators/__init__.py b/src/mrinufft/operators/__init__.py
index 6cb0bdab..37decfac 100644
--- a/src/mrinufft/operators/__init__.py
+++ b/src/mrinufft/operators/__init__.py
@@ -10,7 +10,7 @@
     list_backends,
     check_backend,
 )
-from .off_resonnance import MRIFourierCorrected
+from .off_resonance import MRIFourierCorrected
 from .stacked import MRIStackedNUFFT
 
 #
diff --git a/src/mrinufft/operators/base.py b/src/mrinufft/operators/base.py
index ff4c2c46..aa39956c 100644
--- a/src/mrinufft/operators/base.py
+++ b/src/mrinufft/operators/base.py
@@ -315,9 +315,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)
 
diff --git a/src/mrinufft/operators/off_resonnance.py b/src/mrinufft/operators/off_resonance.py
similarity index 100%
rename from src/mrinufft/operators/off_resonnance.py
rename to src/mrinufft/operators/off_resonance.py

From f8d5dc850a13963eef2119283796a1b81a915ee6 Mon Sep 17 00:00:00 2001
From: mcencini <matteo.cencini@gmail.com>
Date: Thu, 8 Aug 2024 18:24:44 +0200
Subject: [PATCH 12/44] 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 (https://github.com/pytorch/pytorch/issues/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).
---
 src/mrinufft/__init__.py                  |   2 +
 src/mrinufft/operators/__init__.py        |   3 +-
 src/mrinufft/operators/off_resonance.py   | 192 +++++++++++++++--
 tests/operators/test_offres_exp_approx.py | 241 ++++++++++++++++++++++
 4 files changed, 423 insertions(+), 15 deletions(-)
 create mode 100644 tests/operators/test_offres_exp_approx.py

diff --git a/src/mrinufft/__init__.py b/src/mrinufft/__init__.py
index a2d1bd90..7f34aad6 100644
--- a/src/mrinufft/__init__.py
+++ b/src/mrinufft/__init__.py
@@ -10,6 +10,7 @@
     get_operator,
     check_backend,
     list_backends,
+    get_interpolators_from_fieldmap,
 )
 
 from .trajectories import (
@@ -54,6 +55,7 @@
     "get_operator",
     "check_backend",
     "list_backends",
+    "get_interpolators_from_fieldmap",
     "initialize_2D_radial",
     "initialize_2D_spiral",
     "initialize_2D_fibonacci_spiral",
diff --git a/src/mrinufft/operators/__init__.py b/src/mrinufft/operators/__init__.py
index 37decfac..c642d47d 100644
--- a/src/mrinufft/operators/__init__.py
+++ b/src/mrinufft/operators/__init__.py
@@ -10,7 +10,7 @@
     list_backends,
     check_backend,
 )
-from .off_resonance 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",
 ]
diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 0c6f452c..a099be8f 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -4,8 +4,11 @@
 https://github.com/CEA-COSMIC/pysap-mri/blob/master/mri/operators/fourier/orc_wrapper.py
 """
 
+import math
 import numpy as np
 
+from .._utils import get_array_module
+
 from .base import FourierOperatorBase
 from .interfaces.utils import is_cuda_array
 
@@ -16,6 +19,171 @@
     CUPY_AVAILABLE = False
 
 
+TORCH_AVAILABLE = True
+try:
+    import torch
+except ImportError:
+    TORCH_AVAILABLE = False
+
+
+def get_interpolators_from_fieldmap(
+    fieldmap, t, n_time_segments=6, n_bins=(40, 10), mask=None
+):
+    r"""Create B and C matrices to approximate ``exp(-2j*pi*fieldmap*t)``.
+
+    Here, B has shape ``(n_time_segments, len(t))`` and C has shape ``(n_time_segments, *fieldmap.shape)``.
+
+    From Sigpy: https://github.com/mikgroup/sigpy and MIRT (mri_exp_approx.m): https://web.eecs.umich.edu/~fessler/code/
+
+    Parameters
+    ----------
+    fieldmap : np.ndarray or GPUarray
+        Rate map defined as ``fieldmap = R2*_map + 1j * B0_map``.
+        ``*_map`` and ``t`` should have reciprocal units.
+        If ``zmap`` is real, assume ``zmap = B0_map``.
+        Expected shape is ``(nz, ny, nx)``.
+    t : np.ndarray or GPUarray
+        Readout time in ``[s]`` of shape ``(npts,)``.
+    n_time_segments : int, optional
+        Number of time segments. The default is ``6``.
+    n_bins : int | Sequence[int] optional
+        Number of histogram bins to use for ``(B0, T2*)``. The default is ``(40, 10)``
+        If it is a scalar, assume ``n_bins = (n_bins, 10)``. For real fieldmap (B0 only),
+        ``n_bins[1]`` is ignored.
+    mask : np.ndarray or GPUarray, optional
+        Boolean mask to avoid histogram of background values.
+        The default is ``None`` (use the whole map).
+
+    Returns
+    -------
+    B : np.ndarray or GPUarray
+        Temporal interpolator of shape ``(n_time_segments, len(t))``.
+    C : np.ndarray or GPUarray
+        Off-resonance phase map at each time segment center of shape
+        ``(n_time_segments, *fieldmap.shape)``.
+    """
+    # get backend and device
+    xp = get_array_module(fieldmap)
+
+    if xp.__name__ == "torch":
+        is_torch = True
+        if fieldmap.device.type == "cpu":
+            xp = np
+            fieldmap = fieldmap.numpy(force=True)
+        else:
+            assert CUPY_AVAILABLE, "GPU computation requires Cupy!"
+            xp = cp
+            fieldmap = cp.from_dlpack(fieldmap)
+    else:
+        is_torch = False
+
+    # move t to backend
+    fieldmap = xp.asarray(fieldmap, dtype=xp.complex64)
+    t = xp.asarray(t, dtype=xp.float32)
+
+    # default
+    if isinstance(n_bins, (list, tuple)) is False:
+        n_bins = (n_bins, 10)
+
+    # transform to list
+    n_bins = list(n_bins)
+
+    # get field map
+    if xp.isreal(fieldmap).all().item():
+        r2star = None
+        b0 = fieldmap
+        fieldmap = 0.0 + 1j * b0
+    else:
+        r2star = fieldmap.real
+        b0 = fieldmap.imag
+
+    # default mask
+    if mask is None:
+        mask = xp.ones_like(fieldmap, dtype=bool)
+
+    # Hz to radians / s
+    fieldmap = 2 * math.pi * fieldmap
+
+    # create histograms
+    if r2star is not None:
+        z = fieldmap[mask].ravel()
+        z = xp.stack((z.imag, z.real), axis=1)
+        hk, ze = xp.histogramdd(z, bins=n_bins)
+        ze = list(ze)
+
+        # get bin centers
+        zc = [e[1:] - (e[1] - e[0]) / 2 for e in ze]
+
+        # complexify
+        zk = _outer_sum(1j * zc[0], zc[1])  # [K1 K2]
+        zk = zk.T
+        hk = hk.T
+    else:
+        z = fieldmap[mask].ravel()
+        hk, ze = xp.histogram(z.imag, bins=n_bins[0])
+
+        # get bin centers
+        zc = ze[1:] - (ze[1] - ze[0]) / 2
+
+        # complexify
+        zk = 0 + 1j * zc  # [K 1]
+
+    # flatten histogram values and centers
+    hk = hk.ravel()
+    zk = zk.ravel()
+
+    # generate time for each segment
+    tl = (
+        xp.linspace(0, n_time_segments, n_time_segments, dtype=xp.float32)
+        / n_time_segments
+        * t[-1]
+    )  # time seg centers in [s]
+
+    # prepare for basis calculation
+    ch = xp.exp(-tl[:, None, ...] @ zk[None, ...])
+    w = xp.diag(hk**0.5)
+    p = xp.linalg.pinv(w @ _transpose(ch)) @ w
+
+    # actual temporal basis calculation
+    B = p @ xp.exp(-zk[:, None, ...] * t[None, ...])
+    B = B.astype(xp.complex64)
+
+    # get spatial coeffs
+    C = xp.exp(-tl * fieldmap[..., None])
+    C = C[None, ...].swapaxes(0, -1)[
+        ..., 0
+    ]  # (..., n_time_segments) -> (n_time_segments, ...)
+
+    # clean-up of spatial coeffs
+    C = xp.nan_to_num(C, nan=0.0, posinf=0.0, neginf=0.0)
+
+    # back to torch if required
+    if is_torch:
+        if xp.__name__ == "cupy":
+            B = torch.from_dlpack(B)
+            C = torch.from_dlpack(C)
+        else:
+            B = torch.from_numpy(B)
+            C = torch.from_numpy(C)
+
+    return B, C
+
+
+def _outer_sum(xx, yy):
+    xx = xx[:, None, ...]  # add a singleton dimension at axis 1
+    yy = yy[None, ...]  # add a singleton dimension at axis 0
+    ss = xx + yy  # compute the outer sum
+    return ss
+
+
+def _transpose(input):
+    xp = get_array_module(input)
+    if xp.__name__ == "torch":
+        return input.mT
+    else:
+        return input.T
+
+
 class MRIFourierCorrected(FourierOperatorBase):
     """Fourier Operator with B0 Inhomogeneities compensation.
 
@@ -33,7 +201,7 @@ class MRIFourierCorrected(FourierOperatorBase):
         the backend to use for computations. Either 'cpu' or 'gpu'.
     """
 
-    def __init__(self, fourier_op, B, C, indices, backend="cpu"):
+    def __init__(self, fourier_op, B, C, mask, backend="cpu"):
         if backend == "gpu" and not CUPY_AVAILABLE:
             raise RuntimeError("Cupy is required for gpu computations.")
         if backend == "gpu":
@@ -52,13 +220,11 @@ def __init__(self, fourier_op, B, C, indices, backend="cpu"):
         self.shape = fourier_op.shape
         self.smaps = fourier_op.smaps
         self.n_interpolators = len(C)
-        self.B = self.xp.array(B)
-        self.B = self.xp.tile(self.B, (self._fourier_op.n_samples // len(B), 1))
-        self.C = self.xp.array(C)
-        self.indices = indices
+        self.B = self.xp.asarray(B)
+        self.C = self.xp.asarray(C)
 
     def op(self, data, *args):
-        """Compute Forward Operation with off-resonnances effect.
+        """Compute Forward Operation with off-resonance effect.
 
         Parameters
         ----------
@@ -70,19 +236,17 @@ def op(self, data, *args):
         numpy.ndarray or cupy.ndarray
             masked distorded N-D k-space
         """
-        y = self.xp.zeros((self.n_coils, self.n_samples), dtype=np.complex64)
+        y = 0.0
         data_d = self.xp.asarray(data)
         for idx in range(self.n_interpolators):
-            y += self.B[..., idx] * self._fourier_op.op(
-                self.C[idx, self.indices] * data_d, *args
-            )
+            y += self.B[idx] * self._fourier_op.op(self.C[idx] * data_d, *args)
         if self.xp.__name__ == "cupy" and is_cuda_array(data):
             return y
         return y.get()
 
     def adj_op(self, coeffs, *args):
         """
-        Compute Adjoint Operation with off-resonnance effect.
+        Compute Adjoint Operation with off-resonance effect.
 
         Parameters
         ----------
@@ -93,11 +257,11 @@ def adj_op(self, coeffs, *args):
         -------
             inverse Fourier transform of the distorded input k-space.
         """
-        y = self.xp.zeros(self.shape, dtype=np.complex64)
+        y = 0.0
         coeffs_d = self.xp.array(coeffs)
         for idx in range(self.n_interpolators):
-            y += cp.conj(self.C[idx, self.indices]) * self._fourier_op.adj_op(
-                cp.conj(self.B[..., idx]) * coeffs_d, *args
+            y += cp.conj(self.C[idx]) * self._fourier_op.adj_op(
+                cp.conj(self.B[idx]) * coeffs_d, *args
             )
         if self.xp.__name__ == "cupy" and is_cuda_array(coeffs):
             return y
diff --git a/tests/operators/test_offres_exp_approx.py b/tests/operators/test_offres_exp_approx.py
new file mode 100644
index 00000000..9d8a24cd
--- /dev/null
+++ b/tests/operators/test_offres_exp_approx.py
@@ -0,0 +1,241 @@
+"""Test off-resonance spatial coefficient and temporal interpolator estimation."""
+
+import math
+
+import numpy as np
+import numpy.testing as npt
+
+import pytest
+
+import mrinufft
+from mrinufft._utils import get_array_module
+
+CUPY_AVAILABLE = True
+try:
+    import cupy as cp
+except ImportError:
+    CUPY_AVAILABLE = False
+
+TORCH_AVAILABLE = True
+try:
+    import torch
+except ImportError:
+    TORCH_AVAILABLE = False
+
+# parameter combinations (ndims, backend, device)
+params = [(2, "cpu", None), (3, "cpu", None)]
+
+if CUPY_AVAILABLE:
+    params.extend([(2, "gpu", None), (3, "gpu", None)])
+
+if TORCH_AVAILABLE:
+    params.extend([(2, "torch", "cpu"), (3, "torch", "cpu")])
+
+if TORCH_AVAILABLE and torch.cuda.is_available():
+    params.extend([(2, "torch", "cuda"), (3, "torch", "cuda")])
+
+
+def calculate_true_offresonance_term(fieldmap, t):
+    """Calculate non-approximate off-resonance modulation term."""
+    xp = get_array_module(fieldmap)
+    t = xp.asarray(t)
+    if xp.__name__ == "torch":
+        t = t.to(device=fieldmap.device)
+    arg = t * fieldmap[..., None]
+    arg = arg[None, ...].swapaxes(0, -1)[..., 0]
+    return xp.exp(-arg)
+
+
+def calculate_approx_offresonance_term(B, C):
+    """Calculate approximate off-resonance modulation term."""
+    field_term = 0.0
+    for n in range(B.shape[0]):
+        tmp = B[n] * C[n][..., None]
+        tmp = tmp[None, ...].swapaxes(0, -1)[..., 0]
+        field_term += tmp
+
+    return field_term
+
+
+def make_b0map(ndim, npix=64):
+    """Make B0 map (units: Hz)."""
+    # generate support
+    shape = ndim * [npix]
+    if ndim == 2:
+        mask, fieldmap = _make_disk(shape)
+    elif ndim == 3:
+        mask, fieldmap = _make_sphere(shape)
+
+    # mask map
+    fieldmap *= mask
+
+    # rescale
+    fieldmap = 600 * fieldmap / fieldmap.max() - 300  # Hz
+
+    # mask map
+    fieldmap *= mask
+
+    return fieldmap.astype(np.float32), mask
+
+
+def make_t2smap(ndim, npix=64):
+    """Make T2* map (units: ms)."""
+    shape = ndim * [npix]
+    if ndim == 2:
+        mask, _ = _make_disk(shape)
+    elif ndim == 3:
+        mask, _ = _make_sphere(shape)
+
+    # rescale
+    fieldmap = 15.0 * mask  # ms
+
+    return fieldmap.astype(np.float32), mask
+
+
+def _make_disk(shape, frac_radius=0.3):
+    """Make circular binary mask."""
+    # calculate grid
+    ny, nx = shape
+    yy, xx = np.mgrid[:ny, :nx]
+    yy, xx = yy - ny // 2, xx - nx // 2
+    yy, xx = yy / ny, xx / nx
+
+    # radius
+    rr = (xx**2 + yy**2) ** 0.5
+    return rr < frac_radius, rr
+
+
+def _make_sphere(shape, frac_radius=0.3):
+    """Make spherical binary mask."""
+    # calculate grid
+    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
+
+    # radius
+    rr = (xx**2 + yy**2 + zz**2) ** 0.5
+    return rr < frac_radius, rr
+
+
+def get_backend(backend):
+    if backend == "cpu":
+        return np
+    elif backend == "gpu":
+        return cp
+    elif backend == "torch":
+        return torch
+
+
+def assert_allclose(actual, expected, atol, rtol):
+    """Backend agnostic assertion."""
+    xp = get_array_module(actual)
+    if xp.__name__ == "numpy":
+        npt.assert_allclose(actual, expected, atol, rtol)
+    if xp.__name__ == "cupy":
+        npt.assert_allclose(actual.get(), expected.get(), atol, rtol)
+    if xp.__name__ == "torch":
+        npt.assert_allclose(
+            actual.numpy(force=True), expected.numpy(force=True), atol, rtol
+        )
+
+
+@pytest.mark.parametrize("ndim, module, device", params)
+def test_b0map_coeff(ndim, module, device):
+    """Test exponential approximation for B0 field only."""
+    # get module
+    xp = get_backend(module)
+
+    # generate readout times
+    tread = xp.linspace(0.0, 5e-3, 501, dtype=xp.float32)  # (50,), 5 ms
+
+    # generate map
+    b0map, mask = make_b0map(ndim)
+    b0map, mask = xp.asarray(b0map), xp.asarray(mask)
+
+    # cast to device if torch tensor
+    if xp.__name__ == "torch":
+        b0map = b0map.to(device)
+
+    # generate coefficients
+    B, C = mrinufft.get_interpolators_from_fieldmap(
+        b0map, tread, mask=mask, n_time_segments=100
+    )
+
+    # correct backend
+    xpb, xpc = get_array_module(B), get_array_module(C)
+    assert xpb.__name__ == xp.__name__
+    assert xpc.__name__ == xp.__name__
+
+    if xp.__name__ == "torch":
+        assert B.device == b0map.device
+        assert C.device == b0map.device
+
+    # correct dtype
+    assert B.dtype == xp.complex64
+    assert C.dtype == xp.complex64
+
+    # correct shape
+    npt.assert_allclose(B.shape, (100, 501))
+    npt.assert_allclose(C.shape, (100, *b0map.shape))
+
+    # correct approximation
+    expected = calculate_true_offresonance_term(0 + 2 * math.pi * 1j * b0map, tread)
+    actual = calculate_approx_offresonance_term(B, C)
+    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3)
+
+
+@pytest.mark.parametrize("ndim, module, device", params)
+def test_zmap_coeff(ndim, module, device):
+    """Test exponential approximation for complex Zmap = R2* + 1j * B0 field."""
+    # get module
+    xp = get_backend(module)
+
+    # generate readout times
+    tread = xp.linspace(0.0, 5e-3, 501, dtype=xp.float32)  # (50,), 5 ms
+
+    # generate B0 map
+    b0map, mask = make_b0map(ndim)
+    b0map, mask = xp.asarray(b0map), xp.asarray(mask)
+
+    # generate T2* map
+    t2smap, mask = make_t2smap(ndim)
+    t2smap, mask = xp.asarray(t2smap), xp.asarray(mask)
+    t2smap *= 1e-3  # ms -> s
+    t2smap[t2smap == 0.0] = 1.0
+    r2smap = 1.0 / t2smap  # Hz
+    r2smap = mask * r2smap
+
+    # calculate complex fieldmap
+    zmap = r2smap + 1j * b0map
+
+    # cast to device if torch tensor
+    if xp.__name__ == "torch":
+        zmap = zmap.to(device)
+
+    # generate coefficients
+    B, C = mrinufft.get_interpolators_from_fieldmap(
+        zmap, tread, mask=mask, n_time_segments=100
+    )
+
+    # correct backend
+    xpb, xpc = get_array_module(B), get_array_module(C)
+    assert xpb.__name__ == xp.__name__
+    assert xpc.__name__ == xp.__name__
+
+    if xp.__name__ == "torch":
+        assert B.device == zmap.device
+        assert C.device == zmap.device
+
+    # correct dtype
+    assert B.dtype == xp.complex64
+    assert C.dtype == xp.complex64
+
+    # correct shape
+    npt.assert_allclose(B.shape, (100, 501))
+    npt.assert_allclose(C.shape, (100, *zmap.shape))
+
+    # correct approximation
+    expected = calculate_true_offresonance_term(2 * math.pi * zmap, tread)
+    actual = calculate_approx_offresonance_term(B, C)
+    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3)

From 63641869ba525f0dcf970b41d0be4c5a18b84168 Mon Sep 17 00:00:00 2001
From: mcencini <matteo.cencini@gmail.com>
Date: Fri, 9 Aug 2024 17:51:59 +0200
Subject: [PATCH 13/44] feat: field interpolator estimation in
 MRIFourierCorrected constructor.

In addition, add off-resonance example.
---
 .github/workflows/test-ci.yml           |   2 +-
 examples/example_offresonance.py        | 130 ++++++++++++++++++++++++
 src/mrinufft/operators/off_resonance.py |  99 ++++++++++++------
 3 files changed, 200 insertions(+), 31 deletions(-)
 create mode 100644 examples/example_offresonance.py

diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml
index d71ac1bd..fbfc652f 100644
--- a/.github/workflows/test-ci.yml
+++ b/.github/workflows/test-ci.yml
@@ -191,7 +191,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install -e .[test,dev]
-          python -m pip install finufft pooch brainweb-dl torch
+          python -m pip install finufft pooch brainweb-dl torch phantominator
         
       - name: Install GPU related interfaces
         run: |
diff --git a/examples/example_offresonance.py b/examples/example_offresonance.py
new file mode 100644
index 00000000..be78dfeb
--- /dev/null
+++ b/examples/example_offresonance.py
@@ -0,0 +1,130 @@
+"""
+======================
+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
+# ===============
+# As example, 2D shepp-logan will be used.
+# installable using ``pip install phantominator``
+
+from phantominator import ct_shepp_logan
+
+mri_data = ct_shepp_logan(256)
+plt.imshow(mri_data), plt.axis("off")
+
+# %%
+# Field Generation
+# ===============
+# Here, we generate a radial B0 field with the same shape of
+# the input Shepp-Logan phantom
+
+
+def make_b0map(obj, b0range=(-300, 300)):
+    """Make radial B0 field.
+
+    Parameters
+    ----------
+    obj : np.ndarray
+        Input object of shape (ny, nx).
+    b0range : tuple, optional
+        B0 field range in [Hz]. The default is (-300, 300).
+
+    Returns
+    -------
+    b0map : np.ndarray
+        Field inhomogeneities map of shape (ny, nx)
+    """
+    # calculate grid
+    ny, nx = obj.shape
+    yy, xx = np.mgrid[:ny, :nx]
+    yy, xx = yy - ny // 2, xx - nx // 2
+    yy, xx = yy / ny, xx / nx
+
+    # radius
+    rr = (xx**2 + yy**2) ** 0.5
+
+    # mask
+    mask = (obj != 0).astype(np.float32)
+    b0map = mask * rr
+
+    # rescale
+    b0map = (b0range[1] - b0range[0]) * b0map / b0map.max() + b0range[0]  # Hz
+
+    return b0map * mask
+
+
+# generate field
+b0map = make_b0map(mri_data)
+
+# %%
+# Generate a Spiral trajectory
+# ----------------------------
+
+from mrinufft import initialize_2D_spiral
+from mrinufft.density import voronoi
+
+samples = initialize_2D_spiral(Nc=16, Ns=2000, nb_revolutions=10)
+t_read = np.arange(samples.shape[1]) * 4e-6  # 1us raster time
+t_read = np.repeat(t_read[None, ...], samples.shape[0], axis=0)
+density = voronoi(samples)
+
+display_2D_trajectory(samples)
+
+# %%
+# Setup the Operator
+# ==================
+
+import mrinufft
+from mrinufft.operators.off_resonance import MRIFourierCorrected
+
+# Generate standard NUFFT operator
+NufftOperator = mrinufft.get_operator("finufft")
+nufft = NufftOperator(
+    samples=samples,
+    shape=mri_data.shape,
+    density=density,
+)
+
+# Generate Fourier Corrected operator
+mfi_nufft = MRIFourierCorrected(nufft, fieldmap=b0map, t=t_read, mask=mri_data != 0)
+
+# Generate K-Space
+kspace = mfi_nufft.op(mri_data)
+print(kspace.shape)
+
+# Reconstruct without field correction
+mri_data_adj = nufft.adj_op(kspace)
+mri_data_adj = np.squeeze(abs(mri_data_adj))
+print(mri_data_adj.shape)
+
+# Reconstruct with field correction
+mri_data_adj_mfi = mfi_nufft.adj_op(kspace)
+mri_data_adj_mfi = np.squeeze(abs(mri_data_adj_mfi))
+print(mri_data_adj_mfi.shape)
+
+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)
diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index a099be8f..72328447 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -31,9 +31,11 @@ def get_interpolators_from_fieldmap(
 ):
     r"""Create B and C matrices to approximate ``exp(-2j*pi*fieldmap*t)``.
 
-    Here, B has shape ``(n_time_segments, len(t))`` and C has shape ``(n_time_segments, *fieldmap.shape)``.
+    Here, B has shape ``(n_time_segments, len(t))``
+    and C has shape ``(n_time_segments, *fieldmap.shape)``.
 
-    From Sigpy: https://github.com/mikgroup/sigpy and MIRT (mri_exp_approx.m): https://web.eecs.umich.edu/~fessler/code/
+    From Sigpy: https://github.com/mikgroup/sigpy
+    and MIRT (mri_exp_approx.m): https://web.eecs.umich.edu/~fessler/code/
 
     Parameters
     ----------
@@ -43,13 +45,13 @@ def get_interpolators_from_fieldmap(
         If ``zmap`` is real, assume ``zmap = B0_map``.
         Expected shape is ``(nz, ny, nx)``.
     t : np.ndarray or GPUarray
-        Readout time in ``[s]`` of shape ``(npts,)``.
+        Readout time in ``[s]`` of shape ``(nshots, npts)`` or ``(nshots * npts,)``.
     n_time_segments : int, optional
         Number of time segments. The default is ``6``.
     n_bins : int | Sequence[int] optional
         Number of histogram bins to use for ``(B0, T2*)``. The default is ``(40, 10)``
-        If it is a scalar, assume ``n_bins = (n_bins, 10)``. For real fieldmap (B0 only),
-        ``n_bins[1]`` is ignored.
+        If it is a scalar, assume ``n_bins = (n_bins, 10)``.
+        For real fieldmap (B0 only), ``n_bins[1]`` is ignored.
     mask : np.ndarray or GPUarray, optional
         Boolean mask to avoid histogram of background values.
         The default is ``None`` (use the whole map).
@@ -79,7 +81,7 @@ def get_interpolators_from_fieldmap(
 
     # move t to backend
     fieldmap = xp.asarray(fieldmap, dtype=xp.complex64)
-    t = xp.asarray(t, dtype=xp.float32)
+    t = xp.asarray(t, dtype=xp.float32).ravel()
 
     # default
     if isinstance(n_bins, (list, tuple)) is False:
@@ -133,10 +135,8 @@ def get_interpolators_from_fieldmap(
     zk = zk.ravel()
 
     # generate time for each segment
-    tl = (
-        xp.linspace(0, n_time_segments, n_time_segments, dtype=xp.float32)
-        / n_time_segments
-        * t[-1]
+    tl = xp.linspace(
+        t.min(), t.max(), n_time_segments, dtype=xp.float32
     )  # time seg centers in [s]
 
     # prepare for basis calculation
@@ -194,16 +194,48 @@ class MRIFourierCorrected(FourierOperatorBase):
     ----------
     fourier_op: object of class FourierBase
         the fourier operator to wrap
-    B: numpy.ndarray
-    C: numpy.ndarray
-    indices: numpy.ndarray
-    backend: str, default 'cpu'
-        the backend to use for computations. Either 'cpu' or 'gpu'.
+    fieldmap : np.ndarray or GPUarray, optional
+        Rate map defined as ``fieldmap = R2*_map + 1j * B0_map``.
+        ``*_map`` and ``t`` should have reciprocal units.
+        If ``zmap`` is real, assume ``zmap = B0_map``.
+        Expected shape is ``(nz, ny, nx)``.
+    t : np.ndarray or GPUarray, optional
+        Readout time in ``[s]`` of shape ``(nshots, npts)`` or ``(nshots * npts,)``.
+    n_time_segments : int, optional
+        Number of time segments. The default is ``6``.
+    n_bins : int | Sequence[int] optional
+        Number of histogram bins to use for ``(B0, T2*)``. The default is ``(40, 10)``
+        If it is a scalar, assume ``n_bins = (n_bins, 10)``.
+        For real fieldmap (B0 only), ``n_bins[1]`` is ignored.
+    mask : np.ndarray or GPUarray, optional
+        Boolean mask to avoid histogram of background values.
+        The default is ``None`` (use the whole map).
+    B : np.ndarray or GPUarray, optional
+        Temporal interpolator of shape ``(n_time_segments, len(t))``.
+    C : np.ndarray or GPUarray, optional
+        Off-resonance phase map at each time segment center of shape
+        ``(n_time_segments, *fieldmap.shape)``.
+    backend: str, optional
+        The backend to use for computations. Either 'cpu', 'gpu' or 'torch'.
+        The default is `cpu`.
     """
 
-    def __init__(self, fourier_op, B, C, mask, backend="cpu"):
+    def __init__(
+        self,
+        fourier_op,
+        fieldmap=None,
+        t=None,
+        n_time_segments=6,
+        n_bins=(40, 10),
+        mask=None,
+        B=None,
+        C=None,
+        backend="cpu",
+    ):
         if backend == "gpu" and not CUPY_AVAILABLE:
             raise RuntimeError("Cupy is required for gpu computations.")
+        if backend == "torch":
+            self.xp = torch
         if backend == "gpu":
             self.xp = cp
         elif backend == "cpu":
@@ -212,16 +244,25 @@ def __init__(self, fourier_op, B, C, mask, backend="cpu"):
             raise ValueError("Unsupported backend.")
         self._fourier_op = fourier_op
 
-        if not fourier_op.uses_sense:
-            raise ValueError("please use smaps.")
+        # if not fourier_op.uses_sense:
+        # raise ValueError("please use smaps.")
 
-        self.n_samples = fourier_op.n_samples
+        # self.n_samples = fourier_op.n_samples
         self.n_coils = fourier_op.n_coils
         self.shape = fourier_op.shape
         self.smaps = fourier_op.smaps
-        self.n_interpolators = len(C)
-        self.B = self.xp.asarray(B)
-        self.C = self.xp.asarray(C)
+
+        if B is not None and C is not None:
+            self.B = self.xp.asarray(B)
+            self.C = self.xp.asarray(C)
+        else:
+            fieldmap = self.xp.asarray(fieldmap)
+            self.B, self.C = get_interpolators_from_fieldmap(
+                fieldmap, t, n_time_segments, n_bins, mask
+            )
+        if self.B is None or self.C is None:
+            raise ValueError("Please either provide fieldmap and t or B and C")
+        self.n_interpolators = len(self.C)
 
     def op(self, data, *args):
         """Compute Forward Operation with off-resonance effect.
@@ -240,9 +281,8 @@ def op(self, data, *args):
         data_d = self.xp.asarray(data)
         for idx in range(self.n_interpolators):
             y += self.B[idx] * self._fourier_op.op(self.C[idx] * data_d, *args)
-        if self.xp.__name__ == "cupy" and is_cuda_array(data):
-            return y
-        return y.get()
+
+        return y
 
     def adj_op(self, coeffs, *args):
         """
@@ -260,12 +300,11 @@ def adj_op(self, coeffs, *args):
         y = 0.0
         coeffs_d = self.xp.array(coeffs)
         for idx in range(self.n_interpolators):
-            y += cp.conj(self.C[idx]) * self._fourier_op.adj_op(
-                cp.conj(self.B[idx]) * coeffs_d, *args
+            y += self.xp.conj(self.C[idx]) * self._fourier_op.adj_op(
+                self.xp.conj(self.B[idx]) * coeffs_d, *args
             )
-        if self.xp.__name__ == "cupy" and is_cuda_array(coeffs):
-            return y
-        return y.get()
+
+        return y
 
     def get_grad(self, image_data, obs_data):
         """Compute the data consistency error.

From 36563b548e8987b3e32dbcb5ec8b9552604a8a85 Mon Sep 17 00:00:00 2001
From: mcencini <matteo.cencini@gmail.com>
Date: Fri, 9 Aug 2024 17:57:51 +0200
Subject: [PATCH 14/44] feat: set autograd_available based on base
 FourierOperator

---
 src/mrinufft/operators/off_resonance.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 72328447..60965566 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -251,6 +251,7 @@ def __init__(
         self.n_coils = fourier_op.n_coils
         self.shape = fourier_op.shape
         self.smaps = fourier_op.smaps
+        self.autograd_available = fourier_op.autograd_available
 
         if B is not None and C is not None:
             self.B = self.xp.asarray(B)

From 39105c67cc80d01f06f755f0b7211be85770120c Mon Sep 17 00:00:00 2001
From: Matteo Cencini <83717049+mcencini@users.noreply.github.com>
Date: Fri, 9 Aug 2024 22:33:25 +0200
Subject: [PATCH 15/44] Update test_offres_exp_approx.py

Avoid torch cuda test case if cupy is not available.
---
 tests/operators/test_offres_exp_approx.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/operators/test_offres_exp_approx.py b/tests/operators/test_offres_exp_approx.py
index 9d8a24cd..f9679ccc 100644
--- a/tests/operators/test_offres_exp_approx.py
+++ b/tests/operators/test_offres_exp_approx.py
@@ -31,7 +31,7 @@
 if TORCH_AVAILABLE:
     params.extend([(2, "torch", "cpu"), (3, "torch", "cpu")])
 
-if TORCH_AVAILABLE and torch.cuda.is_available():
+if TORCH_AVAILABLE and CUPY_AVAILABLE:
     params.extend([(2, "torch", "cuda"), (3, "torch", "cuda")])
 
 

From f3ad9ac142387e55bbe4307ac39ccf23d5b2e900 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <83717049+mcencini@users.noreply.github.com>
Date: Fri, 30 Aug 2024 16:20:03 +0200
Subject: [PATCH 16/44] Update src/mrinufft/operators/off_resonance.py

remove deprecated get_grad

Co-authored-by: Chaithya G R <chaithyagr@gmail.com>
---
 src/mrinufft/operators/off_resonance.py | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 60965566..d4f6fd24 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -307,19 +307,3 @@ def adj_op(self, coeffs, *args):
 
         return y
 
-    def get_grad(self, image_data, obs_data):
-        """Compute the data consistency error.
-
-        Parameters
-        ----------
-        image_data: numpy.ndarray or cupy.ndarray
-            N-D input image
-        obs_data: numpy.ndarray or cupy.ndarray
-            N-D observed k-space
-
-        Returns
-        -------
-        numpy.ndarray or cupy.ndarray
-            data consistency error in image space.
-        """
-        return self.adj_op(self.op(image_data) - obs_data)

From f4b910877da23545c7ce236aed19c3f70664be98 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <matteo.cencini@gmail.com>
Date: Fri, 30 Aug 2024 18:35:35 +0200
Subject: [PATCH 17/44] address pr review

---
 examples/example_offresonance.py           |  27 +--
 src/mrinufft/operators/off_resonance.py    |  69 +++---
 tests/operators/_test_offres_exp_approx.py | 236 +++++++++++++++++++++
 tests/operators/test_offres_exp_approx.py  | 198 ++++++-----------
 4 files changed, 345 insertions(+), 185 deletions(-)
 create mode 100644 tests/operators/_test_offres_exp_approx.py

diff --git a/examples/example_offresonance.py b/examples/example_offresonance.py
index be78dfeb..796214ba 100644
--- a/examples/example_offresonance.py
+++ b/examples/example_offresonance.py
@@ -21,13 +21,14 @@
 # %%
 # Data Generation
 # ===============
-# As example, 2D shepp-logan will be used.
-# installable using ``pip install phantominator``
+# For realistic 2D image we will use a slice from the brainweb dataset.
+# installable using ``pip install brainweb-dl``
 
-from phantominator import ct_shepp_logan
+from brainweb_dl import get_mri
 
-mri_data = ct_shepp_logan(256)
-plt.imshow(mri_data), plt.axis("off")
+mri_data = get_mri(44, "T1")
+mri_data = mri_data[::-1, ...][90]
+plt.imshow(mri_data), plt.axis("off"), plt.title("ground truth")
 
 # %%
 # Field Generation
@@ -79,9 +80,10 @@ def make_b0map(obj, b0range=(-300, 300)):
 
 from mrinufft import initialize_2D_spiral
 from mrinufft.density import voronoi
+from mrinufft.trajectories.utils import DEFAULT_RASTER_TIME
 
-samples = initialize_2D_spiral(Nc=16, Ns=2000, nb_revolutions=10)
-t_read = np.arange(samples.shape[1]) * 4e-6  # 1us 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)
 
@@ -91,11 +93,11 @@ def make_b0map(obj, b0range=(-300, 300)):
 # Setup the Operator
 # ==================
 
-import mrinufft
+from mrinufft import get_operator
 from mrinufft.operators.off_resonance import MRIFourierCorrected
 
 # Generate standard NUFFT operator
-NufftOperator = mrinufft.get_operator("finufft")
+NufftOperator = get_operator("finufft")
 nufft = NufftOperator(
     samples=samples,
     shape=mri_data.shape,
@@ -103,21 +105,20 @@ def make_b0map(obj, b0range=(-300, 300)):
 )
 
 # Generate Fourier Corrected operator
-mfi_nufft = MRIFourierCorrected(nufft, fieldmap=b0map, t=t_read, mask=mri_data != 0)
+mfi_nufft = MRIFourierCorrected(
+    nufft, fieldmap=b0map, readout_time=t_read, mask=mri_data != 0
+)
 
 # Generate K-Space
 kspace = mfi_nufft.op(mri_data)
-print(kspace.shape)
 
 # Reconstruct without field correction
 mri_data_adj = nufft.adj_op(kspace)
 mri_data_adj = np.squeeze(abs(mri_data_adj))
-print(mri_data_adj.shape)
 
 # Reconstruct with field correction
 mri_data_adj_mfi = mfi_nufft.adj_op(kspace)
 mri_data_adj_mfi = np.squeeze(abs(mri_data_adj_mfi))
-print(mri_data_adj_mfi.shape)
 
 fig2, ax2 = plt.subplots(1, 2)
 ax2[0].imshow(mri_data_adj), ax2[0].axis("off"), ax2[0].set_title("w/o correction")
diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index d4f6fd24..55f717ed 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -9,25 +9,17 @@
 
 from .._utils import get_array_module
 
-from .base import FourierOperatorBase
-from .interfaces.utils import is_cuda_array
+from .base import FourierOperatorBase, CUPY_AVAILABLE, AUTOGRAD_AVAILABLE
 
-CUPY_AVAILABLE = True
-try:
+if CUPY_AVAILABLE:
     import cupy as cp
-except ImportError:
-    CUPY_AVAILABLE = False
 
-
-TORCH_AVAILABLE = True
-try:
+if AUTOGRAD_AVAILABLE:
     import torch
-except ImportError:
-    TORCH_AVAILABLE = False
 
 
 def get_interpolators_from_fieldmap(
-    fieldmap, t, n_time_segments=6, n_bins=(40, 10), mask=None
+    fieldmap, readout_time, n_time_segments=6, n_bins=(40, 10), mask=None
 ):
     r"""Create B and C matrices to approximate ``exp(-2j*pi*fieldmap*t)``.
 
@@ -44,7 +36,7 @@ def get_interpolators_from_fieldmap(
         ``*_map`` and ``t`` should have reciprocal units.
         If ``zmap`` is real, assume ``zmap = B0_map``.
         Expected shape is ``(nz, ny, nx)``.
-    t : np.ndarray or GPUarray
+    readout_time : np.ndarray or GPUarray
         Readout time in ``[s]`` of shape ``(nshots, npts)`` or ``(nshots * npts,)``.
     n_time_segments : int, optional
         Number of time segments. The default is ``6``.
@@ -64,31 +56,43 @@ def get_interpolators_from_fieldmap(
         Off-resonance phase map at each time segment center of shape
         ``(n_time_segments, *fieldmap.shape)``.
     """
+    # default
+    if isinstance(n_bins, (list, tuple)) is False:
+        n_bins = (n_bins, 10)
+
+    # transform to list
+    n_bins = list(n_bins)
+
     # get backend and device
     xp = get_array_module(fieldmap)
 
+    # enforce complex field
+    fieldmap = xp.asarray(fieldmap, dtype=xp.complex64)
+
+    # cast arrays to backend
     if xp.__name__ == "torch":
         is_torch = True
         if fieldmap.device.type == "cpu":
             xp = np
             fieldmap = fieldmap.numpy(force=True)
+            if mask is not None:
+                mask = mask.numpy(force=True)
+            readout_time = readout_time.numpy(force=True)
         else:
             assert CUPY_AVAILABLE, "GPU computation requires Cupy!"
             xp = cp
             fieldmap = cp.from_dlpack(fieldmap)
+            if mask is not None:
+                mask = cp.from_dlpack(mask)
+            readout_time = cp.from_dlpack(readout_time)
     else:
         is_torch = False
 
-    # move t to backend
-    fieldmap = xp.asarray(fieldmap, dtype=xp.complex64)
-    t = xp.asarray(t, dtype=xp.float32).ravel()
-
-    # default
-    if isinstance(n_bins, (list, tuple)) is False:
-        n_bins = (n_bins, 10)
-
-    # transform to list
-    n_bins = list(n_bins)
+    readout_time = xp.asarray(readout_time, dtype=xp.float32).ravel()
+    if mask is None:
+        mask = xp.ones_like(fieldmap, dtype=bool)
+    else:
+        mask = xp.asarray(mask, dtype=bool)
 
     # get field map
     if xp.isreal(fieldmap).all().item():
@@ -99,10 +103,6 @@ def get_interpolators_from_fieldmap(
         r2star = fieldmap.real
         b0 = fieldmap.imag
 
-    # default mask
-    if mask is None:
-        mask = xp.ones_like(fieldmap, dtype=bool)
-
     # Hz to radians / s
     fieldmap = 2 * math.pi * fieldmap
 
@@ -136,7 +136,7 @@ def get_interpolators_from_fieldmap(
 
     # generate time for each segment
     tl = xp.linspace(
-        t.min(), t.max(), n_time_segments, dtype=xp.float32
+        readout_time.min(), readout_time.max(), n_time_segments, dtype=xp.float32
     )  # time seg centers in [s]
 
     # prepare for basis calculation
@@ -145,7 +145,7 @@ def get_interpolators_from_fieldmap(
     p = xp.linalg.pinv(w @ _transpose(ch)) @ w
 
     # actual temporal basis calculation
-    B = p @ xp.exp(-zk[:, None, ...] * t[None, ...])
+    B = p @ xp.exp(-zk[:, None, ...] * readout_time[None, ...])
     B = B.astype(xp.complex64)
 
     # get spatial coeffs
@@ -199,7 +199,7 @@ class MRIFourierCorrected(FourierOperatorBase):
         ``*_map`` and ``t`` should have reciprocal units.
         If ``zmap`` is real, assume ``zmap = B0_map``.
         Expected shape is ``(nz, ny, nx)``.
-    t : np.ndarray or GPUarray, optional
+    readout_time : np.ndarray or GPUarray, optional
         Readout time in ``[s]`` of shape ``(nshots, npts)`` or ``(nshots * npts,)``.
     n_time_segments : int, optional
         Number of time segments. The default is ``6``.
@@ -224,7 +224,7 @@ def __init__(
         self,
         fourier_op,
         fieldmap=None,
-        t=None,
+        readout_time=None,
         n_time_segments=6,
         n_bins=(40, 10),
         mask=None,
@@ -244,10 +244,6 @@ def __init__(
             raise ValueError("Unsupported backend.")
         self._fourier_op = fourier_op
 
-        # if not fourier_op.uses_sense:
-        # raise ValueError("please use smaps.")
-
-        # self.n_samples = fourier_op.n_samples
         self.n_coils = fourier_op.n_coils
         self.shape = fourier_op.shape
         self.smaps = fourier_op.smaps
@@ -259,7 +255,7 @@ def __init__(
         else:
             fieldmap = self.xp.asarray(fieldmap)
             self.B, self.C = get_interpolators_from_fieldmap(
-                fieldmap, t, n_time_segments, n_bins, mask
+                fieldmap, readout_time, n_time_segments, n_bins, mask
             )
         if self.B is None or self.C is None:
             raise ValueError("Please either provide fieldmap and t or B and C")
@@ -306,4 +302,3 @@ def adj_op(self, coeffs, *args):
             )
 
         return y
-
diff --git a/tests/operators/_test_offres_exp_approx.py b/tests/operators/_test_offres_exp_approx.py
new file mode 100644
index 00000000..81a13c3f
--- /dev/null
+++ b/tests/operators/_test_offres_exp_approx.py
@@ -0,0 +1,236 @@
+"""Test off-resonance spatial coefficient and temporal interpolator estimation."""
+
+import math
+
+import numpy as np
+import numpy.testing as npt
+
+import pytest
+
+import mrinufft
+from mrinufft.operators.base import CUPY_AVAILABLE, AUTOGRAD_AVAILABLE
+from mrinufft._utils import get_array_module
+
+if CUPY_AVAILABLE:
+    import cupy as cp
+
+if AUTOGRAD_AVAILABLE:
+    import torch
+
+# parameter combinations (ndims, backend, device)
+params = [(2, "cpu", None), (3, "cpu", None)]
+
+if CUPY_AVAILABLE:
+    params.extend([(2, "gpu", None), (3, "gpu", None)])
+
+if AUTOGRAD_AVAILABLE:
+    params.extend([(2, "torch", "cpu"), (3, "torch", "cpu")])
+
+if AUTOGRAD_AVAILABLE and CUPY_AVAILABLE:
+    params.extend([(2, "torch", "cuda"), (3, "torch", "cuda")])
+
+
+def calculate_true_offresonance_term(fieldmap, t):
+    """Calculate non-approximate off-resonance modulation term."""
+    xp = get_array_module(fieldmap)
+    t = xp.asarray(t)
+    if xp.__name__ == "torch":
+        t = t.to(device=fieldmap.device)
+    arg = t * fieldmap[..., None]
+    arg = arg[None, ...].swapaxes(0, -1)[..., 0]
+    return xp.exp(-arg)
+
+
+def calculate_approx_offresonance_term(B, C):
+    """Calculate approximate off-resonance modulation term."""
+    field_term = 0.0
+    for n in range(B.shape[0]):
+        tmp = B[n] * C[n][..., None]
+        tmp = tmp[None, ...].swapaxes(0, -1)[..., 0]
+        field_term += tmp
+
+    return field_term
+
+
+def make_b0map(ndim, npix=64):
+    """Make B0 map (units: Hz)."""
+    # generate support
+    shape = ndim * [npix]
+    if ndim == 2:
+        mask, fieldmap = _make_disk(shape)
+    elif ndim == 3:
+        mask, fieldmap = _make_sphere(shape)
+
+    # mask map
+    fieldmap *= mask
+
+    # rescale
+    fieldmap = 600 * fieldmap / fieldmap.max() - 300  # Hz
+
+    # mask map
+    fieldmap *= mask
+
+    return fieldmap.astype(np.float32), mask
+
+
+def make_t2smap(ndim, npix=64):
+    """Make T2* map (units: ms)."""
+    shape = ndim * [npix]
+    if ndim == 2:
+        mask, _ = _make_disk(shape)
+    elif ndim == 3:
+        mask, _ = _make_sphere(shape)
+
+    # rescale
+    fieldmap = 15.0 * mask  # ms
+
+    return fieldmap.astype(np.float32), mask
+
+
+def _make_disk(shape, frac_radius=0.3):
+    """Make circular binary mask."""
+    # calculate grid
+    ny, nx = shape
+    yy, xx = np.mgrid[:ny, :nx]
+    yy, xx = yy - ny // 2, xx - nx // 2
+    yy, xx = yy / ny, xx / nx
+
+    # radius
+    rr = (xx**2 + yy**2) ** 0.5
+    return rr < frac_radius, rr
+
+
+def _make_sphere(shape, frac_radius=0.3):
+    """Make spherical binary mask."""
+    # calculate grid
+    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
+
+    # radius
+    rr = (xx**2 + yy**2 + zz**2) ** 0.5
+    return rr < frac_radius, rr
+
+
+def get_backend(backend):
+    if backend == "cpu":
+        return np
+    elif backend == "gpu":
+        return cp
+    elif backend == "torch":
+        return torch
+
+
+def assert_allclose(actual, expected, atol, rtol):
+    """Backend agnostic assertion."""
+    xp = get_array_module(actual)
+    if xp.__name__ == "numpy":
+        npt.assert_allclose(actual, expected, atol, rtol)
+    if xp.__name__ == "cupy":
+        npt.assert_allclose(actual.get(), expected.get(), atol, rtol)
+    if xp.__name__ == "torch":
+        npt.assert_allclose(
+            actual.numpy(force=True), expected.numpy(force=True), atol, rtol
+        )
+
+
+@pytest.mark.parametrize("ndim, module, device", params)
+def test_b0map_coeff(ndim, module, device):
+    """Test exponential approximation for B0 field only."""
+    # get module
+    xp = get_backend(module)
+
+    # generate readout times
+    tread = xp.linspace(0.0, 5e-3, 501, dtype=xp.float32)  # (50,), 5 ms
+
+    # generate map
+    b0map, mask = make_b0map(ndim)
+    b0map, mask = xp.asarray(b0map), xp.asarray(mask)
+
+    # cast to device if torch tensor
+    if xp.__name__ == "torch":
+        b0map = b0map.to(device)
+
+    # generate coefficients
+    B, C = mrinufft.get_interpolators_from_fieldmap(
+        b0map, tread, mask=mask, n_time_segments=100
+    )
+
+    # correct backend
+    xpb, xpc = get_array_module(B), get_array_module(C)
+    assert xpb.__name__ == xp.__name__
+    assert xpc.__name__ == xp.__name__
+
+    if xp.__name__ == "torch":
+        assert B.device == b0map.device
+        assert C.device == b0map.device
+
+    # correct dtype
+    assert B.dtype == xp.complex64
+    assert C.dtype == xp.complex64
+
+    # correct shape
+    npt.assert_allclose(B.shape, (100, 501))
+    npt.assert_allclose(C.shape, (100, *b0map.shape))
+
+    # correct approximation
+    expected = calculate_true_offresonance_term(0 + 2 * math.pi * 1j * b0map, tread)
+    actual = calculate_approx_offresonance_term(B, C)
+    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3)
+
+
+@pytest.mark.parametrize("ndim, module, device", params)
+def test_zmap_coeff(ndim, module, device):
+    """Test exponential approximation for complex Zmap = R2* + 1j * B0 field."""
+    # get module
+    xp = get_backend(module)
+
+    # generate readout times
+    tread = xp.linspace(0.0, 5e-3, 501, dtype=xp.float32)  # (50,), 5 ms
+
+    # generate B0 map
+    b0map, mask = make_b0map(ndim)
+    b0map, mask = xp.asarray(b0map), xp.asarray(mask)
+
+    # generate T2* map
+    t2smap, mask = make_t2smap(ndim)
+    t2smap, mask = xp.asarray(t2smap), xp.asarray(mask)
+    t2smap *= 1e-3  # ms -> s
+    t2smap[t2smap == 0.0] = 1.0
+    r2smap = 1.0 / t2smap  # Hz
+    r2smap = mask * r2smap
+
+    # calculate complex fieldmap
+    zmap = r2smap + 1j * b0map
+
+    # cast to device if torch tensor
+    if xp.__name__ == "torch":
+        zmap = zmap.to(device)
+
+    # generate coefficients
+    B, C = mrinufft.get_interpolators_from_fieldmap(
+        zmap, tread, mask=mask, n_time_segments=100
+    )
+
+    # correct backend
+    xpb, xpc = get_array_module(B), get_array_module(C)
+    assert xpb.__name__ == xp.__name__
+    assert xpc.__name__ == xp.__name__
+
+    if xp.__name__ == "torch":
+        assert B.device == zmap.device
+        assert C.device == zmap.device
+
+    # correct dtype
+    assert B.dtype == xp.complex64
+    assert C.dtype == xp.complex64
+
+    # correct shape
+    npt.assert_allclose(B.shape, (100, 501))
+    npt.assert_allclose(C.shape, (100, *zmap.shape))
+
+    # correct approximation
+    expected = calculate_true_offresonance_term(2 * math.pi * zmap, tread)
+    actual = calculate_approx_offresonance_term(B, C)
+    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3)
diff --git a/tests/operators/test_offres_exp_approx.py b/tests/operators/test_offres_exp_approx.py
index f9679ccc..d1cb0915 100644
--- a/tests/operators/test_offres_exp_approx.py
+++ b/tests/operators/test_offres_exp_approx.py
@@ -8,31 +8,10 @@
 import pytest
 
 import mrinufft
+from mrinufft.operators.base import CUPY_AVAILABLE, AUTOGRAD_AVAILABLE
 from mrinufft._utils import get_array_module
 
-CUPY_AVAILABLE = True
-try:
-    import cupy as cp
-except ImportError:
-    CUPY_AVAILABLE = False
-
-TORCH_AVAILABLE = True
-try:
-    import torch
-except ImportError:
-    TORCH_AVAILABLE = False
-
-# parameter combinations (ndims, backend, device)
-params = [(2, "cpu", None), (3, "cpu", None)]
-
-if CUPY_AVAILABLE:
-    params.extend([(2, "gpu", None), (3, "gpu", None)])
-
-if TORCH_AVAILABLE:
-    params.extend([(2, "torch", "cpu"), (3, "torch", "cpu")])
-
-if TORCH_AVAILABLE and CUPY_AVAILABLE:
-    params.extend([(2, "torch", "cuda"), (3, "torch", "cuda")])
+from helpers import to_interface, from_interface
 
 
 def calculate_true_offresonance_term(fieldmap, t):
@@ -53,189 +32,138 @@ def calculate_approx_offresonance_term(B, C):
         tmp = B[n] * C[n][..., None]
         tmp = tmp[None, ...].swapaxes(0, -1)[..., 0]
         field_term += tmp
-
     return field_term
 
 
-def make_b0map(ndim, npix=64):
+def _make_b0map(ndim, npix=64):
     """Make B0 map (units: Hz)."""
-    # generate support
     shape = ndim * [npix]
     if ndim == 2:
         mask, fieldmap = _make_disk(shape)
     elif ndim == 3:
         mask, fieldmap = _make_sphere(shape)
-
-    # mask map
     fieldmap *= mask
-
-    # rescale
     fieldmap = 600 * fieldmap / fieldmap.max() - 300  # Hz
-
-    # mask map
     fieldmap *= mask
-
     return fieldmap.astype(np.float32), mask
 
 
-def make_t2smap(ndim, npix=64):
+def _make_t2smap(ndim, npix=64):
     """Make T2* map (units: ms)."""
     shape = ndim * [npix]
     if ndim == 2:
         mask, _ = _make_disk(shape)
     elif ndim == 3:
         mask, _ = _make_sphere(shape)
-
-    # rescale
     fieldmap = 15.0 * mask  # ms
-
     return fieldmap.astype(np.float32), mask
 
 
 def _make_disk(shape, frac_radius=0.3):
     """Make circular binary mask."""
-    # calculate grid
     ny, nx = shape
     yy, xx = np.mgrid[:ny, :nx]
     yy, xx = yy - ny // 2, xx - nx // 2
     yy, xx = yy / ny, xx / nx
-
-    # radius
     rr = (xx**2 + yy**2) ** 0.5
     return rr < frac_radius, rr
 
 
 def _make_sphere(shape, frac_radius=0.3):
     """Make spherical binary mask."""
-    # calculate grid
     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
-
-    # radius
     rr = (xx**2 + yy**2 + zz**2) ** 0.5
     return rr < frac_radius, rr
 
 
-def get_backend(backend):
-    if backend == "cpu":
-        return np
-    elif backend == "gpu":
-        return cp
-    elif backend == "torch":
-        return torch
+# Parameter combinations (ndims, backend, device)
+params = [(2, "cpu", None), (3, "cpu", None)]
 
+if CUPY_AVAILABLE:
+    params.extend([(2, "gpu", None), (3, "gpu", None)])
 
-def assert_allclose(actual, expected, atol, rtol):
-    """Backend agnostic assertion."""
-    xp = get_array_module(actual)
-    if xp.__name__ == "numpy":
-        npt.assert_allclose(actual, expected, atol, rtol)
-    if xp.__name__ == "cupy":
-        npt.assert_allclose(actual.get(), expected.get(), atol, rtol)
-    if xp.__name__ == "torch":
-        npt.assert_allclose(
-            actual.numpy(force=True), expected.numpy(force=True), atol, rtol
-        )
+if AUTOGRAD_AVAILABLE:
+    params.extend([(2, "torch", "cpu"), (3, "torch", "cpu")])
 
+if AUTOGRAD_AVAILABLE and CUPY_AVAILABLE:
+    params.extend([(2, "torch", "gpu"), (3, "torch", "gpu")])
 
-@pytest.mark.parametrize("ndim, module, device", params)
-def test_b0map_coeff(ndim, module, device):
-    """Test exponential approximation for B0 field only."""
-    # get module
-    xp = get_backend(module)
 
-    # generate readout times
-    tread = xp.linspace(0.0, 5e-3, 501, dtype=xp.float32)  # (50,), 5 ms
+@pytest.fixture(scope="module", params=params)
+def map_fixture(request):
+    """Fixture to generate B0 and T2* maps based on dimension and backend."""
+    ndim, module, device = request.param
+    interface = module if "torch" not in module else f"{module}-{device}"
 
-    # generate map
-    b0map, mask = make_b0map(ndim)
-    b0map, mask = xp.asarray(b0map), xp.asarray(mask)
+    # Generate maps
+    b0map, mask = _make_b0map(ndim)
+    t2smap, _ = _make_t2smap(ndim)
 
-    # cast to device if torch tensor
-    if xp.__name__ == "torch":
-        b0map = b0map.to(device)
+    # Convert maps to the appropriate interface
+    b0map = to_interface(b0map, interface)
+    mask = to_interface(mask, interface)
+    t2smap = to_interface(t2smap, interface)
 
-    # generate coefficients
-    B, C = mrinufft.get_interpolators_from_fieldmap(
-        b0map, tread, mask=mask, n_time_segments=100
-    )
+    return ndim, b0map, t2smap, mask, interface
 
-    # correct backend
-    xpb, xpc = get_array_module(B), get_array_module(C)
-    assert xpb.__name__ == xp.__name__
-    assert xpc.__name__ == xp.__name__
 
-    if xp.__name__ == "torch":
-        assert B.device == b0map.device
-        assert C.device == b0map.device
+def assert_allclose(actual, expected, atol, rtol, interface):
+    """Backend agnostic assertion using from_interface helper."""
+    actual_np = from_interface(actual, interface)
+    expected_np = from_interface(expected, interface)
+    npt.assert_allclose(actual_np, expected_np, atol=atol, rtol=rtol)
+
+
+def test_b0map_coeff(map_fixture):
+    """Test exponential approximation for B0 field only."""
+    ndim, b0map, _, mask, interface = map_fixture
+
+    # Generate readout times
+    tread = to_interface(np.linspace(0.0, 5e-3, 501, dtype=np.float32), interface)
 
-    # correct dtype
-    assert B.dtype == xp.complex64
-    assert C.dtype == xp.complex64
+    # Generate coefficients
+    B, C = mrinufft.get_interpolators_from_fieldmap(
+        b0map, tread, mask=mask, n_time_segments=100
+    )
 
-    # correct shape
-    npt.assert_allclose(B.shape, (100, 501))
-    npt.assert_allclose(C.shape, (100, *b0map.shape))
+    # Assert properties
+    assert B.shape == (100, 501)
+    assert C.shape == (100, *b0map.shape)
 
-    # correct approximation
+    # Correct approximation
     expected = calculate_true_offresonance_term(0 + 2 * math.pi * 1j * b0map, tread)
     actual = calculate_approx_offresonance_term(B, C)
-    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3)
+    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3, interface=interface)
 
 
-@pytest.mark.parametrize("ndim, module, device", params)
-def test_zmap_coeff(ndim, module, device):
+def test_zmap_coeff(map_fixture):
     """Test exponential approximation for complex Zmap = R2* + 1j * B0 field."""
-    # get module
-    xp = get_backend(module)
-
-    # generate readout times
-    tread = xp.linspace(0.0, 5e-3, 501, dtype=xp.float32)  # (50,), 5 ms
-
-    # generate B0 map
-    b0map, mask = make_b0map(ndim)
-    b0map, mask = xp.asarray(b0map), xp.asarray(mask)
-
-    # generate T2* map
-    t2smap, mask = make_t2smap(ndim)
-    t2smap, mask = xp.asarray(t2smap), xp.asarray(mask)
-    t2smap *= 1e-3  # ms -> s
-    t2smap[t2smap == 0.0] = 1.0
-    r2smap = 1.0 / t2smap  # Hz
+    ndim, b0map, t2smap, mask, interface = map_fixture
+
+    # Generate readout times
+    tread = to_interface(np.linspace(0.0, 5e-3, 501, dtype=np.float32), interface)
+
+    # Convert T2* map to R2* map
+    t2smap = t2smap * 1e-3  # ms -> s
+    r2smap = 1.0 / (t2smap + 1e-9)  # Hz
     r2smap = mask * r2smap
 
-    # calculate complex fieldmap
+    # Calculate complex fieldmap (Zmap)
     zmap = r2smap + 1j * b0map
 
-    # cast to device if torch tensor
-    if xp.__name__ == "torch":
-        zmap = zmap.to(device)
-
-    # generate coefficients
+    # Generate coefficients
     B, C = mrinufft.get_interpolators_from_fieldmap(
         zmap, tread, mask=mask, n_time_segments=100
     )
 
-    # correct backend
-    xpb, xpc = get_array_module(B), get_array_module(C)
-    assert xpb.__name__ == xp.__name__
-    assert xpc.__name__ == xp.__name__
-
-    if xp.__name__ == "torch":
-        assert B.device == zmap.device
-        assert C.device == zmap.device
-
-    # correct dtype
-    assert B.dtype == xp.complex64
-    assert C.dtype == xp.complex64
-
-    # correct shape
-    npt.assert_allclose(B.shape, (100, 501))
-    npt.assert_allclose(C.shape, (100, *zmap.shape))
+    # Assert properties
+    assert B.shape == (100, 501)
+    assert C.shape == (100, *zmap.shape)
 
-    # correct approximation
+    # Correct approximation
     expected = calculate_true_offresonance_term(2 * math.pi * zmap, tread)
     actual = calculate_approx_offresonance_term(B, C)
-    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3)
+    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3, interface=interface)

From 45ebb00a19a5e22d0b2263be21e7a6317d9087ab Mon Sep 17 00:00:00 2001
From: Matteo Cencini <matteo.cencini@gmail.com>
Date: Fri, 30 Aug 2024 18:41:40 +0200
Subject: [PATCH 18/44] remove phantominator

---
 .github/workflows/test-ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml
index fb1f5d1d..e6843a95 100644
--- a/.github/workflows/test-ci.yml
+++ b/.github/workflows/test-ci.yml
@@ -205,7 +205,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install -e .[test,dev]
-          python -m pip install finufft pooch brainweb-dl torch phantominator
+          python -m pip install finufft pooch brainweb-dl torch
         
       - name: Install GPU related interfaces
         run: |

From 308a27b7bf32da8827554597ebdcdb12d8c2c5b3 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <matteo.cencini@gmail.com>
Date: Mon, 2 Sep 2024 10:09:02 +0200
Subject: [PATCH 19/44] remove duplicated example and replacing subject in
 example_offres

---
 examples/example_offresonance.py           |   2 +-
 tests/operators/_test_offres_exp_approx.py | 236 ---------------------
 2 files changed, 1 insertion(+), 237 deletions(-)
 delete mode 100644 tests/operators/_test_offres_exp_approx.py

diff --git a/examples/example_offresonance.py b/examples/example_offresonance.py
index 796214ba..88436646 100644
--- a/examples/example_offresonance.py
+++ b/examples/example_offresonance.py
@@ -26,7 +26,7 @@
 
 from brainweb_dl import get_mri
 
-mri_data = get_mri(44, "T1")
+mri_data = get_mri(0, "T1")
 mri_data = mri_data[::-1, ...][90]
 plt.imshow(mri_data), plt.axis("off"), plt.title("ground truth")
 
diff --git a/tests/operators/_test_offres_exp_approx.py b/tests/operators/_test_offres_exp_approx.py
deleted file mode 100644
index 81a13c3f..00000000
--- a/tests/operators/_test_offres_exp_approx.py
+++ /dev/null
@@ -1,236 +0,0 @@
-"""Test off-resonance spatial coefficient and temporal interpolator estimation."""
-
-import math
-
-import numpy as np
-import numpy.testing as npt
-
-import pytest
-
-import mrinufft
-from mrinufft.operators.base import CUPY_AVAILABLE, AUTOGRAD_AVAILABLE
-from mrinufft._utils import get_array_module
-
-if CUPY_AVAILABLE:
-    import cupy as cp
-
-if AUTOGRAD_AVAILABLE:
-    import torch
-
-# parameter combinations (ndims, backend, device)
-params = [(2, "cpu", None), (3, "cpu", None)]
-
-if CUPY_AVAILABLE:
-    params.extend([(2, "gpu", None), (3, "gpu", None)])
-
-if AUTOGRAD_AVAILABLE:
-    params.extend([(2, "torch", "cpu"), (3, "torch", "cpu")])
-
-if AUTOGRAD_AVAILABLE and CUPY_AVAILABLE:
-    params.extend([(2, "torch", "cuda"), (3, "torch", "cuda")])
-
-
-def calculate_true_offresonance_term(fieldmap, t):
-    """Calculate non-approximate off-resonance modulation term."""
-    xp = get_array_module(fieldmap)
-    t = xp.asarray(t)
-    if xp.__name__ == "torch":
-        t = t.to(device=fieldmap.device)
-    arg = t * fieldmap[..., None]
-    arg = arg[None, ...].swapaxes(0, -1)[..., 0]
-    return xp.exp(-arg)
-
-
-def calculate_approx_offresonance_term(B, C):
-    """Calculate approximate off-resonance modulation term."""
-    field_term = 0.0
-    for n in range(B.shape[0]):
-        tmp = B[n] * C[n][..., None]
-        tmp = tmp[None, ...].swapaxes(0, -1)[..., 0]
-        field_term += tmp
-
-    return field_term
-
-
-def make_b0map(ndim, npix=64):
-    """Make B0 map (units: Hz)."""
-    # generate support
-    shape = ndim * [npix]
-    if ndim == 2:
-        mask, fieldmap = _make_disk(shape)
-    elif ndim == 3:
-        mask, fieldmap = _make_sphere(shape)
-
-    # mask map
-    fieldmap *= mask
-
-    # rescale
-    fieldmap = 600 * fieldmap / fieldmap.max() - 300  # Hz
-
-    # mask map
-    fieldmap *= mask
-
-    return fieldmap.astype(np.float32), mask
-
-
-def make_t2smap(ndim, npix=64):
-    """Make T2* map (units: ms)."""
-    shape = ndim * [npix]
-    if ndim == 2:
-        mask, _ = _make_disk(shape)
-    elif ndim == 3:
-        mask, _ = _make_sphere(shape)
-
-    # rescale
-    fieldmap = 15.0 * mask  # ms
-
-    return fieldmap.astype(np.float32), mask
-
-
-def _make_disk(shape, frac_radius=0.3):
-    """Make circular binary mask."""
-    # calculate grid
-    ny, nx = shape
-    yy, xx = np.mgrid[:ny, :nx]
-    yy, xx = yy - ny // 2, xx - nx // 2
-    yy, xx = yy / ny, xx / nx
-
-    # radius
-    rr = (xx**2 + yy**2) ** 0.5
-    return rr < frac_radius, rr
-
-
-def _make_sphere(shape, frac_radius=0.3):
-    """Make spherical binary mask."""
-    # calculate grid
-    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
-
-    # radius
-    rr = (xx**2 + yy**2 + zz**2) ** 0.5
-    return rr < frac_radius, rr
-
-
-def get_backend(backend):
-    if backend == "cpu":
-        return np
-    elif backend == "gpu":
-        return cp
-    elif backend == "torch":
-        return torch
-
-
-def assert_allclose(actual, expected, atol, rtol):
-    """Backend agnostic assertion."""
-    xp = get_array_module(actual)
-    if xp.__name__ == "numpy":
-        npt.assert_allclose(actual, expected, atol, rtol)
-    if xp.__name__ == "cupy":
-        npt.assert_allclose(actual.get(), expected.get(), atol, rtol)
-    if xp.__name__ == "torch":
-        npt.assert_allclose(
-            actual.numpy(force=True), expected.numpy(force=True), atol, rtol
-        )
-
-
-@pytest.mark.parametrize("ndim, module, device", params)
-def test_b0map_coeff(ndim, module, device):
-    """Test exponential approximation for B0 field only."""
-    # get module
-    xp = get_backend(module)
-
-    # generate readout times
-    tread = xp.linspace(0.0, 5e-3, 501, dtype=xp.float32)  # (50,), 5 ms
-
-    # generate map
-    b0map, mask = make_b0map(ndim)
-    b0map, mask = xp.asarray(b0map), xp.asarray(mask)
-
-    # cast to device if torch tensor
-    if xp.__name__ == "torch":
-        b0map = b0map.to(device)
-
-    # generate coefficients
-    B, C = mrinufft.get_interpolators_from_fieldmap(
-        b0map, tread, mask=mask, n_time_segments=100
-    )
-
-    # correct backend
-    xpb, xpc = get_array_module(B), get_array_module(C)
-    assert xpb.__name__ == xp.__name__
-    assert xpc.__name__ == xp.__name__
-
-    if xp.__name__ == "torch":
-        assert B.device == b0map.device
-        assert C.device == b0map.device
-
-    # correct dtype
-    assert B.dtype == xp.complex64
-    assert C.dtype == xp.complex64
-
-    # correct shape
-    npt.assert_allclose(B.shape, (100, 501))
-    npt.assert_allclose(C.shape, (100, *b0map.shape))
-
-    # correct approximation
-    expected = calculate_true_offresonance_term(0 + 2 * math.pi * 1j * b0map, tread)
-    actual = calculate_approx_offresonance_term(B, C)
-    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3)
-
-
-@pytest.mark.parametrize("ndim, module, device", params)
-def test_zmap_coeff(ndim, module, device):
-    """Test exponential approximation for complex Zmap = R2* + 1j * B0 field."""
-    # get module
-    xp = get_backend(module)
-
-    # generate readout times
-    tread = xp.linspace(0.0, 5e-3, 501, dtype=xp.float32)  # (50,), 5 ms
-
-    # generate B0 map
-    b0map, mask = make_b0map(ndim)
-    b0map, mask = xp.asarray(b0map), xp.asarray(mask)
-
-    # generate T2* map
-    t2smap, mask = make_t2smap(ndim)
-    t2smap, mask = xp.asarray(t2smap), xp.asarray(mask)
-    t2smap *= 1e-3  # ms -> s
-    t2smap[t2smap == 0.0] = 1.0
-    r2smap = 1.0 / t2smap  # Hz
-    r2smap = mask * r2smap
-
-    # calculate complex fieldmap
-    zmap = r2smap + 1j * b0map
-
-    # cast to device if torch tensor
-    if xp.__name__ == "torch":
-        zmap = zmap.to(device)
-
-    # generate coefficients
-    B, C = mrinufft.get_interpolators_from_fieldmap(
-        zmap, tread, mask=mask, n_time_segments=100
-    )
-
-    # correct backend
-    xpb, xpc = get_array_module(B), get_array_module(C)
-    assert xpb.__name__ == xp.__name__
-    assert xpc.__name__ == xp.__name__
-
-    if xp.__name__ == "torch":
-        assert B.device == zmap.device
-        assert C.device == zmap.device
-
-    # correct dtype
-    assert B.dtype == xp.complex64
-    assert C.dtype == xp.complex64
-
-    # correct shape
-    npt.assert_allclose(B.shape, (100, 501))
-    npt.assert_allclose(C.shape, (100, *zmap.shape))
-
-    # correct approximation
-    expected = calculate_true_offresonance_term(2 * math.pi * zmap, tread)
-    actual = calculate_approx_offresonance_term(B, C)
-    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3)

From f6a5cb73bf071a1042b58afdcf455403276f8d01 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <matteo.cencini@gmail.com>
Date: Mon, 2 Sep 2024 15:52:36 +0200
Subject: [PATCH 20/44] address PR rev #2

---
 examples/example_offresonance.py          |  52 +++-----
 src/mrinufft/extras/__init__.py           |   3 +
 src/mrinufft/extras/data.py               | 106 +++++++++++++++++
 src/mrinufft/operators/off_resonance.py   |  90 +++++++++-----
 tests/case_fieldmaps.py                   |  57 +++++++++
 tests/helpers/__init__.py                 |   3 +-
 tests/helpers/asserts.py                  |   8 ++
 tests/operators/test_offres_exp_approx.py | 139 ++++------------------
 8 files changed, 276 insertions(+), 182 deletions(-)
 create mode 100644 src/mrinufft/extras/data.py
 create mode 100644 tests/case_fieldmaps.py

diff --git a/examples/example_offresonance.py b/examples/example_offresonance.py
index 88436646..c5a8ef36 100644
--- a/examples/example_offresonance.py
+++ b/examples/example_offresonance.py
@@ -30,49 +30,29 @@
 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
 
-
-def make_b0map(obj, b0range=(-300, 300)):
-    """Make radial B0 field.
-
-    Parameters
-    ----------
-    obj : np.ndarray
-        Input object of shape (ny, nx).
-    b0range : tuple, optional
-        B0 field range in [Hz]. The default is (-300, 300).
-
-    Returns
-    -------
-    b0map : np.ndarray
-        Field inhomogeneities map of shape (ny, nx)
-    """
-    # calculate grid
-    ny, nx = obj.shape
-    yy, xx = np.mgrid[:ny, :nx]
-    yy, xx = yy - ny // 2, xx - nx // 2
-    yy, xx = yy / ny, xx / nx
-
-    # radius
-    rr = (xx**2 + yy**2) ** 0.5
-
-    # mask
-    mask = (obj != 0).astype(np.float32)
-    b0map = mask * rr
-
-    # rescale
-    b0map = (b0range[1] - b0range[0]) * b0map / b0map.max() + b0range[0]  # Hz
-
-    return b0map * mask
-
+from mrinufft.extras import make_b0map
 
 # generate field
-b0map = make_b0map(mri_data)
+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
@@ -106,7 +86,7 @@ def make_b0map(obj, b0range=(-300, 300)):
 
 # Generate Fourier Corrected operator
 mfi_nufft = MRIFourierCorrected(
-    nufft, fieldmap=b0map, readout_time=t_read, mask=mri_data != 0
+    nufft, fieldmap=b0map, readout_time=t_read, mask=brain_mask
 )
 
 # Generate K-Space
diff --git a/src/mrinufft/extras/__init__.py b/src/mrinufft/extras/__init__.py
index 6ba1d4a4..3c08cb88 100644
--- a/src/mrinufft/extras/__init__.py
+++ b/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",
 ]
diff --git a/src/mrinufft/extras/data.py b/src/mrinufft/extras/data.py
new file mode 100644
index 00000000..5dfcf281
--- /dev/null
+++ b/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
diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 55f717ed..3c5cc833 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -10,6 +10,7 @@
 from .._utils import get_array_module
 
 from .base import FourierOperatorBase, CUPY_AVAILABLE, AUTOGRAD_AVAILABLE
+from .interfaces.utils import is_cuda_array
 
 if CUPY_AVAILABLE:
     import cupy as cp
@@ -31,30 +32,35 @@ def get_interpolators_from_fieldmap(
 
     Parameters
     ----------
-    fieldmap : np.ndarray or GPUarray
+    fieldmap : np.ndarray
         Rate map defined as ``fieldmap = R2*_map + 1j * B0_map``.
         ``*_map`` and ``t`` should have reciprocal units.
         If ``zmap`` is real, assume ``zmap = B0_map``.
+        Also supports Cupy arrays and Torch tensors.
         Expected shape is ``(nz, ny, nx)``.
-    readout_time : np.ndarray or GPUarray
+    readout_time : np.ndarray
         Readout time in ``[s]`` of shape ``(nshots, npts)`` or ``(nshots * npts,)``.
+        Also supports Cupy arrays and Torch tensors.
     n_time_segments : int, optional
         Number of time segments. The default is ``6``.
     n_bins : int | Sequence[int] optional
         Number of histogram bins to use for ``(B0, T2*)``. The default is ``(40, 10)``
         If it is a scalar, assume ``n_bins = (n_bins, 10)``.
         For real fieldmap (B0 only), ``n_bins[1]`` is ignored.
-    mask : np.ndarray or GPUarray, optional
+    mask : np.ndarray, optional
         Boolean mask to avoid histogram of background values.
         The default is ``None`` (use the whole map).
+        Also supports Cupy arrays and Torch tensors.
 
     Returns
     -------
-    B : np.ndarray or GPUarray
+    B : np.ndarray
         Temporal interpolator of shape ``(n_time_segments, len(t))``.
-    C : np.ndarray or GPUarray
+        Array module is the same as input fieldmap.
+    C : np.ndarray
         Off-resonance phase map at each time segment center of shape
         ``(n_time_segments, *fieldmap.shape)``.
+        Array module is the same as input fieldmap.
     """
     # default
     if isinstance(n_bins, (list, tuple)) is False:
@@ -69,25 +75,24 @@ def get_interpolators_from_fieldmap(
     # enforce complex field
     fieldmap = xp.asarray(fieldmap, dtype=xp.complex64)
 
-    # cast arrays to backend
+    # cast arrays to fieldmap backend
     if xp.__name__ == "torch":
         is_torch = True
-        if fieldmap.device.type == "cpu":
-            xp = np
-            fieldmap = fieldmap.numpy(force=True)
-            if mask is not None:
-                mask = mask.numpy(force=True)
-            readout_time = readout_time.numpy(force=True)
-        else:
-            assert CUPY_AVAILABLE, "GPU computation requires Cupy!"
-            xp = cp
-            fieldmap = cp.from_dlpack(fieldmap)
-            if mask is not None:
-                mask = cp.from_dlpack(mask)
-            readout_time = cp.from_dlpack(readout_time)
     else:
         is_torch = False
 
+    if is_cuda_array(fieldmap):
+        assert CUPY_AVAILABLE, "GPU computation requires Cupy!"
+        xp = cp
+        fieldmap = _to_cupy(fieldmap)
+        readout_time = _to_cupy(readout_time)
+        mask = _to_cupy(mask)
+    else:
+        xp = np
+        fieldmap = _to_numpy(fieldmap)
+        readout_time = _to_numpy(readout_time)
+        mask = _to_numpy(mask)
+
     readout_time = xp.asarray(readout_time, dtype=xp.float32).ravel()
     if mask is None:
         mask = xp.ones_like(fieldmap, dtype=bool)
@@ -142,7 +147,7 @@ def get_interpolators_from_fieldmap(
     # prepare for basis calculation
     ch = xp.exp(-tl[:, None, ...] @ zk[None, ...])
     w = xp.diag(hk**0.5)
-    p = xp.linalg.pinv(w @ _transpose(ch)) @ w
+    p = xp.linalg.pinv(w @ ch.T) @ w
 
     # actual temporal basis calculation
     B = p @ xp.exp(-zk[:, None, ...] * readout_time[None, ...])
@@ -159,12 +164,8 @@ def get_interpolators_from_fieldmap(
 
     # back to torch if required
     if is_torch:
-        if xp.__name__ == "cupy":
-            B = torch.from_dlpack(B)
-            C = torch.from_dlpack(C)
-        else:
-            B = torch.from_numpy(B)
-            C = torch.from_numpy(C)
+        B = _to_torch(B)
+        C = _to_torch(C)
 
     return B, C
 
@@ -176,12 +177,34 @@ def _outer_sum(xx, yy):
     return ss
 
 
-def _transpose(input):
+# TODO: /* refactor with_* decorators
+def _to_numpy(input):
     xp = get_array_module(input)
+
     if xp.__name__ == "torch":
-        return input.mT
+        return input.numpy(force=True)
+    elif xp.__name__ == "cupy":
+        return input.get()
     else:
-        return input.T
+        return input
+
+
+def _to_cupy(input):
+    return cp.asarray(input)
+
+
+def _to_torch(input):
+    xp = get_array_module(input)
+
+    if xp.__name__ == "numpy":
+        return torch.from_numpy(input)
+    elif xp.__name__ == "cupy":
+        return torch.from_dlpack(input)
+    else:
+        return input
+
+
+# */
 
 
 class MRIFourierCorrected(FourierOperatorBase):
@@ -194,13 +217,15 @@ class MRIFourierCorrected(FourierOperatorBase):
     ----------
     fourier_op: object of class FourierBase
         the fourier operator to wrap
-    fieldmap : np.ndarray or GPUarray, optional
+    fieldmap : np.ndarray, optional
         Rate map defined as ``fieldmap = R2*_map + 1j * B0_map``.
         ``*_map`` and ``t`` should have reciprocal units.
         If ``zmap`` is real, assume ``zmap = B0_map``.
+        Also supports Cupy arrays and Torch tensors.
         Expected shape is ``(nz, ny, nx)``.
     readout_time : np.ndarray or GPUarray, optional
         Readout time in ``[s]`` of shape ``(nshots, npts)`` or ``(nshots * npts,)``.
+        Also supports Cupy arrays and Torch tensors.
     n_time_segments : int, optional
         Number of time segments. The default is ``6``.
     n_bins : int | Sequence[int] optional
@@ -210,9 +235,10 @@ class MRIFourierCorrected(FourierOperatorBase):
     mask : np.ndarray or GPUarray, optional
         Boolean mask to avoid histogram of background values.
         The default is ``None`` (use the whole map).
-    B : np.ndarray or GPUarray, optional
+        Also supports Cupy arrays and Torch tensors.
+    B : np.ndarray, optional
         Temporal interpolator of shape ``(n_time_segments, len(t))``.
-    C : np.ndarray or GPUarray, optional
+    C : np.ndarray, optional
         Off-resonance phase map at each time segment center of shape
         ``(n_time_segments, *fieldmap.shape)``.
     backend: str, optional
diff --git a/tests/case_fieldmaps.py b/tests/case_fieldmaps.py
new file mode 100644
index 00000000..cd36dff6
--- /dev/null
+++ b/tests/case_fieldmaps.py
@@ -0,0 +1,57 @@
+"""Fieldmap cases we want to test."""
+
+from mrinufft.extras import make_b0map, make_t2smap
+
+
+class CasesB0maps:
+    """B0 field maps cases we want to test.
+
+    Each case return a field map and the binary spatial support of the object.
+    """
+
+    def case_real2D(self, N=64, b0range=(-300, 300)):
+        """Create a real (B0 only) 2D field map."""
+        return make_b0map(2 * [N])
+
+    def case_real3D(self, N=64, b0range=(-300, 300)):
+        """Create a real (B0 only) 3D field map."""
+        return make_b0map(3 * [N])
+
+
+class CasesZmaps:
+    """Complex zmap field maps cases we want to test.
+
+    Each case return a field map and the binary spatial support of the object.
+    """
+
+    def case_complex2D(self, N=64, b0range=(-300, 300), t2svalue=15.0):
+        """Create a complex (R2* + 1j * B0) 2D field map."""
+        # Generate real and imaginary parts
+        t2smap, _ = make_t2smap(2 * [N])
+        b0map, mask = make_b0map(2 * [N])
+
+        # Convert T2* map to R2* map
+        t2smap = t2smap * 1e-3  # ms -> s
+        r2smap = 1.0 / (t2smap + 1e-9)  # Hz
+        r2smap = mask * r2smap
+
+        # Calculate complex fieldmap (Zmap)
+        zmap = r2smap + 1j * b0map
+
+        return zmap, mask
+
+    def case_complex3D(self, N=64, b0range=(-300, 300), t2svalue=15.0):
+        """Create a complex (R2* + 1j * B0) 3D field map."""
+        # Generate real and imaginary parts
+        t2smap, _ = make_t2smap(3 * [N])
+        b0map, mask = make_b0map(3 * [N])
+
+        # Convert T2* map to R2* map
+        t2smap = t2smap * 1e-3  # ms -> s
+        r2smap = 1.0 / (t2smap + 1e-9)  # Hz
+        r2smap = mask * r2smap
+
+        # Calculate complex fieldmap (Zmap)
+        zmap = r2smap + 1j * b0map
+
+        return zmap, mask
diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py
index 6afe23f0..a8132187 100644
--- a/tests/helpers/__init__.py
+++ b/tests/helpers/__init__.py
@@ -1,6 +1,6 @@
 """Helper functions for testing the operators."""
 
-from .asserts import assert_almost_allclose, assert_correlate
+from .asserts import assert_almost_allclose, assert_correlate, assert_allclose
 from .factories import (
     kspace_from_op,
     image_from_op,
@@ -14,6 +14,7 @@
 __all__ = [
     "assert_almost_allclose",
     "assert_correlate",
+    "assert_allclose"
     "kspace_from_op",
     "image_from_op",
     "to_interface",
diff --git a/tests/helpers/asserts.py b/tests/helpers/asserts.py
index ba3ed07f..df7ea6ba 100644
--- a/tests/helpers/asserts.py
+++ b/tests/helpers/asserts.py
@@ -4,6 +4,7 @@
 import numpy.testing as npt
 import scipy as sp
 
+from .factories import from_interface
 
 def assert_almost_allclose(a, b, rtol, atol, mismatch, equal_nan=False):
     """Assert allclose with a tolerance on the number of mismatched elements.
@@ -64,3 +65,10 @@ def assert_correlate(a, b, slope=1.0, slope_err=1e-3, r_value_err=1e-3):
             f"intercept={intercept}, stderr={stderr}, "
             f"intercept_stderr={intercept_stderr}"
         )
+
+
+def assert_allclose(actual, expected, atol, rtol, interface):
+    """Backend agnostic assertion using from_interface helper."""
+    actual_np = from_interface(actual, interface)
+    expected_np = from_interface(expected, interface)
+    npt.assert_allclose(actual_np, expected_np, atol=atol, rtol=rtol)
\ No newline at end of file
diff --git a/tests/operators/test_offres_exp_approx.py b/tests/operators/test_offres_exp_approx.py
index d1cb0915..b516806c 100644
--- a/tests/operators/test_offres_exp_approx.py
+++ b/tests/operators/test_offres_exp_approx.py
@@ -3,23 +3,23 @@
 import math
 
 import numpy as np
-import numpy.testing as npt
 
-import pytest
+from pytest_cases import parametrize_with_cases
 
 import mrinufft
-from mrinufft.operators.base import CUPY_AVAILABLE, AUTOGRAD_AVAILABLE
 from mrinufft._utils import get_array_module
 
-from helpers import to_interface, from_interface
+from helpers import to_interface, assert_allclose
+from helpers.factories import _param_array_interface
+from case_fieldmaps import CasesB0maps, CasesZmaps
 
 
-def calculate_true_offresonance_term(fieldmap, t):
+def calculate_true_offresonance_term(fieldmap, t, array_interface):
     """Calculate non-approximate off-resonance modulation term."""
+    fieldmap = to_interface(fieldmap, array_interface)
+    t = to_interface(t, array_interface)
+
     xp = get_array_module(fieldmap)
-    t = xp.asarray(t)
-    if xp.__name__ == "torch":
-        t = t.to(device=fieldmap.device)
     arg = t * fieldmap[..., None]
     arg = arg[None, ...].swapaxes(0, -1)[..., 0]
     return xp.exp(-arg)
@@ -35,98 +35,16 @@ def calculate_approx_offresonance_term(B, C):
     return field_term
 
 
-def _make_b0map(ndim, npix=64):
-    """Make B0 map (units: Hz)."""
-    shape = ndim * [npix]
-    if ndim == 2:
-        mask, fieldmap = _make_disk(shape)
-    elif ndim == 3:
-        mask, fieldmap = _make_sphere(shape)
-    fieldmap *= mask
-    fieldmap = 600 * fieldmap / fieldmap.max() - 300  # Hz
-    fieldmap *= mask
-    return fieldmap.astype(np.float32), mask
-
-
-def _make_t2smap(ndim, npix=64):
-    """Make T2* map (units: ms)."""
-    shape = ndim * [npix]
-    if ndim == 2:
-        mask, _ = _make_disk(shape)
-    elif ndim == 3:
-        mask, _ = _make_sphere(shape)
-    fieldmap = 15.0 * mask  # ms
-    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
-
-
-# Parameter combinations (ndims, backend, device)
-params = [(2, "cpu", None), (3, "cpu", None)]
-
-if CUPY_AVAILABLE:
-    params.extend([(2, "gpu", None), (3, "gpu", None)])
-
-if AUTOGRAD_AVAILABLE:
-    params.extend([(2, "torch", "cpu"), (3, "torch", "cpu")])
-
-if AUTOGRAD_AVAILABLE and CUPY_AVAILABLE:
-    params.extend([(2, "torch", "gpu"), (3, "torch", "gpu")])
-
-
-@pytest.fixture(scope="module", params=params)
-def map_fixture(request):
-    """Fixture to generate B0 and T2* maps based on dimension and backend."""
-    ndim, module, device = request.param
-    interface = module if "torch" not in module else f"{module}-{device}"
-
-    # Generate maps
-    b0map, mask = _make_b0map(ndim)
-    t2smap, _ = _make_t2smap(ndim)
-
-    # Convert maps to the appropriate interface
-    b0map = to_interface(b0map, interface)
-    mask = to_interface(mask, interface)
-    t2smap = to_interface(t2smap, interface)
-
-    return ndim, b0map, t2smap, mask, interface
-
-
-def assert_allclose(actual, expected, atol, rtol, interface):
-    """Backend agnostic assertion using from_interface helper."""
-    actual_np = from_interface(actual, interface)
-    expected_np = from_interface(expected, interface)
-    npt.assert_allclose(actual_np, expected_np, atol=atol, rtol=rtol)
-
-
-def test_b0map_coeff(map_fixture):
+@_param_array_interface
+@parametrize_with_cases("b0map, mask", cases=CasesB0maps)
+def test_b0map_coeff(b0map, mask, array_interface):
     """Test exponential approximation for B0 field only."""
-    ndim, b0map, _, mask, interface = map_fixture
-
     # Generate readout times
-    tread = to_interface(np.linspace(0.0, 5e-3, 501, dtype=np.float32), interface)
+    tread = np.linspace(0.0, 5e-3, 501, dtype=np.float32)
 
     # Generate coefficients
     B, C = mrinufft.get_interpolators_from_fieldmap(
-        b0map, tread, mask=mask, n_time_segments=100
+        to_interface(b0map, array_interface), tread, mask=mask, n_time_segments=100
     )
 
     # Assert properties
@@ -134,29 +52,22 @@ def test_b0map_coeff(map_fixture):
     assert C.shape == (100, *b0map.shape)
 
     # Correct approximation
-    expected = calculate_true_offresonance_term(0 + 2 * math.pi * 1j * b0map, tread)
+    expected = calculate_true_offresonance_term(
+        0 + 2 * math.pi * 1j * b0map, tread, array_interface
+    )
     actual = calculate_approx_offresonance_term(B, C)
-    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3, interface=interface)
-
+    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3, interface=array_interface)
 
-def test_zmap_coeff(map_fixture):
-    """Test exponential approximation for complex Zmap = R2* + 1j * B0 field."""
-    ndim, b0map, t2smap, mask, interface = map_fixture
 
+@_param_array_interface
+@parametrize_with_cases("zmap, mask", cases=CasesZmaps)
+def test_zmap_coeff(zmap, mask, array_interface):
     # Generate readout times
-    tread = to_interface(np.linspace(0.0, 5e-3, 501, dtype=np.float32), interface)
-
-    # Convert T2* map to R2* map
-    t2smap = t2smap * 1e-3  # ms -> s
-    r2smap = 1.0 / (t2smap + 1e-9)  # Hz
-    r2smap = mask * r2smap
-
-    # Calculate complex fieldmap (Zmap)
-    zmap = r2smap + 1j * b0map
+    tread = np.linspace(0.0, 5e-3, 501, dtype=np.float32)
 
     # Generate coefficients
     B, C = mrinufft.get_interpolators_from_fieldmap(
-        zmap, tread, mask=mask, n_time_segments=100
+        to_interface(zmap, array_interface), tread, mask=mask, n_time_segments=100
     )
 
     # Assert properties
@@ -164,6 +75,8 @@ def test_zmap_coeff(map_fixture):
     assert C.shape == (100, *zmap.shape)
 
     # Correct approximation
-    expected = calculate_true_offresonance_term(2 * math.pi * zmap, tread)
+    expected = calculate_true_offresonance_term(
+        2 * math.pi * zmap, tread, array_interface
+    )
     actual = calculate_approx_offresonance_term(B, C)
-    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3, interface=interface)
+    assert_allclose(actual, expected, atol=1e-3, rtol=1e-3, interface=array_interface)

From ed44862579f414bc9ac9d4aa9a9a663d66defd80 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Tue, 3 Sep 2024 09:17:22 +0200
Subject: [PATCH 21/44] \!docs_build and fix on \!style

---
 pyproject.toml            | 2 +-
 tests/helpers/__init__.py | 3 +--
 tests/helpers/asserts.py  | 3 ++-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 143284e7..03442100 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,7 +11,7 @@ dynamic = ["version"]
 
 [project.optional-dependencies]
 
-gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"]
+gpunufft = ["gpuNUFFT>=0.10.0", "cupy-cuda12x"]
 torchkbnufft = ["torchkbnufft", "cupy-cuda12x"]
 cufinufft = ["cufinufft", "cupy-cuda12x"]
 finufft = ["finufft"]
diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py
index a8132187..555cd224 100644
--- a/tests/helpers/__init__.py
+++ b/tests/helpers/__init__.py
@@ -14,8 +14,7 @@
 __all__ = [
     "assert_almost_allclose",
     "assert_correlate",
-    "assert_allclose"
-    "kspace_from_op",
+    "assert_allclose" "kspace_from_op",
     "image_from_op",
     "to_interface",
     "from_interface",
diff --git a/tests/helpers/asserts.py b/tests/helpers/asserts.py
index df7ea6ba..c74a015f 100644
--- a/tests/helpers/asserts.py
+++ b/tests/helpers/asserts.py
@@ -6,6 +6,7 @@
 
 from .factories import from_interface
 
+
 def assert_almost_allclose(a, b, rtol, atol, mismatch, equal_nan=False):
     """Assert allclose with a tolerance on the number of mismatched elements.
 
@@ -71,4 +72,4 @@ def assert_allclose(actual, expected, atol, rtol, interface):
     """Backend agnostic assertion using from_interface helper."""
     actual_np = from_interface(actual, interface)
     expected_np = from_interface(expected, interface)
-    npt.assert_allclose(actual_np, expected_np, atol=atol, rtol=rtol)
\ No newline at end of file
+    npt.assert_allclose(actual_np, expected_np, atol=atol, rtol=rtol)

From 70c2cf82db5c54227b549f758ab83275b1190aaa Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Tue, 3 Sep 2024 09:20:20 +0200
Subject: [PATCH 22/44] \!docs_build

---
 docs/sphinx_add_colab_link.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/sphinx_add_colab_link.py b/docs/sphinx_add_colab_link.py
index 26f6f564..c247dd75 100644
--- a/docs/sphinx_add_colab_link.py
+++ b/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,

From 45a9364421702b49a6e082420fdf3fa33b674166 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <matteo.cencini@gmail.com>
Date: Tue, 3 Sep 2024 14:16:58 +0200
Subject: [PATCH 23/44] skip field coefficient gpu test case if cupy not
 available

---
 tests/operators/test_offres_exp_approx.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/operators/test_offres_exp_approx.py b/tests/operators/test_offres_exp_approx.py
index b516806c..7c11976d 100644
--- a/tests/operators/test_offres_exp_approx.py
+++ b/tests/operators/test_offres_exp_approx.py
@@ -4,10 +4,14 @@
 
 import numpy as np
 
+import pytest
 from pytest_cases import parametrize_with_cases
 
+
 import mrinufft
 from mrinufft._utils import get_array_module
+from mrinufft.operators.base import CUPY_AVAILABLE
+
 
 from helpers import to_interface, assert_allclose
 from helpers.factories import _param_array_interface
@@ -39,6 +43,9 @@ def calculate_approx_offresonance_term(B, C):
 @parametrize_with_cases("b0map, mask", cases=CasesB0maps)
 def test_b0map_coeff(b0map, mask, array_interface):
     """Test exponential approximation for B0 field only."""
+    if array_interface == "torch-gpu" and not CUPY_AVAILABLE:
+        pytest.skip("GPU computations requires cupy")
+
     # Generate readout times
     tread = np.linspace(0.0, 5e-3, 501, dtype=np.float32)
 
@@ -62,6 +69,10 @@ def test_b0map_coeff(b0map, mask, array_interface):
 @_param_array_interface
 @parametrize_with_cases("zmap, mask", cases=CasesZmaps)
 def test_zmap_coeff(zmap, mask, array_interface):
+    """Test exponential approximation for complex Z = R2* + 1j *B0 field."""
+    if array_interface == "torch-gpu" and CUPY_AVAILABLE is False:
+        pytest.skip("GPU computations requires cupy")
+
     # Generate readout times
     tread = np.linspace(0.0, 5e-3, 501, dtype=np.float32)
 

From 1bbe2c303cff55bd8fcc506f309c440d1e429855 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Tue, 3 Sep 2024 16:58:56 +0200
Subject: [PATCH 24/44] \!docs_build fix gpuNUFFT

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 03442100..d560f8ca 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,7 +11,7 @@ dynamic = ["version"]
 
 [project.optional-dependencies]
 
-gpunufft = ["gpuNUFFT>=0.10.0", "cupy-cuda12x"]
+gpunufft = ["gpuNUFFT<0.10.0", "cupy-cuda12x"]
 torchkbnufft = ["torchkbnufft", "cupy-cuda12x"]
 cufinufft = ["cufinufft", "cupy-cuda12x"]
 finufft = ["finufft"]

From 216b7adfd3f8e1c1b4cd48227795308ee3909873 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <83717049+mcencini@users.noreply.github.com>
Date: Thu, 5 Sep 2024 10:41:29 +0200
Subject: [PATCH 25/44] Update src/mrinufft/operators/off_resonance.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>
---
 src/mrinufft/operators/off_resonance.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 3c5cc833..4ff468bf 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -76,10 +76,7 @@ def get_interpolators_from_fieldmap(
     fieldmap = xp.asarray(fieldmap, dtype=xp.complex64)
 
     # cast arrays to fieldmap backend
-    if xp.__name__ == "torch":
-        is_torch = True
-    else:
-        is_torch = False
+    is_torch = xp.__name__ == "torch"
 
     if is_cuda_array(fieldmap):
         assert CUPY_AVAILABLE, "GPU computation requires Cupy!"

From 0ac0aa0c1d3a1d318306ca1845107bde67dbbb13 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <83717049+mcencini@users.noreply.github.com>
Date: Thu, 5 Sep 2024 10:50:24 +0200
Subject: [PATCH 26/44] Update src/mrinufft/operators/off_resonance.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>
---
 src/mrinufft/operators/off_resonance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 4ff468bf..85d9f294 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -208,7 +208,7 @@ class MRIFourierCorrected(FourierOperatorBase):
     """Fourier Operator with B0 Inhomogeneities compensation.
 
     This is a wrapper around the Fourier Operator to compensate for the
-    B0 inhomogeneities  in the  k-space.
+    B0 inhomogeneities in  the k-space.
 
     Parameters
     ----------

From 9b04b510b7c48eb6c251faf82eda85d4004897f9 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Thu, 5 Sep 2024 10:51:26 +0200
Subject: [PATCH 27/44] \!docs_build gpunufft

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index d560f8ca..03442100 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,7 +11,7 @@ dynamic = ["version"]
 
 [project.optional-dependencies]
 
-gpunufft = ["gpuNUFFT<0.10.0", "cupy-cuda12x"]
+gpunufft = ["gpuNUFFT>=0.10.0", "cupy-cuda12x"]
 torchkbnufft = ["torchkbnufft", "cupy-cuda12x"]
 cufinufft = ["cufinufft", "cupy-cuda12x"]
 finufft = ["finufft"]

From a35fc350ab685b1f28d31c648632015eb9978b99 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <83717049+mcencini@users.noreply.github.com>
Date: Thu, 5 Sep 2024 10:51:47 +0200
Subject: [PATCH 28/44] Update src/mrinufft/operators/off_resonance.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>
---
 src/mrinufft/operators/off_resonance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 85d9f294..3d43f478 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -213,7 +213,7 @@ class MRIFourierCorrected(FourierOperatorBase):
     Parameters
     ----------
     fourier_op: object of class FourierBase
-        the fourier operator to wrap
+        the Fourier operator to wrap
     fieldmap : np.ndarray, optional
         Rate map defined as ``fieldmap = R2*_map + 1j * B0_map``.
         ``*_map`` and ``t`` should have reciprocal units.

From f3ed9c119a9b30b5c3cf8c845871b818038c40e1 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <83717049+mcencini@users.noreply.github.com>
Date: Thu, 5 Sep 2024 11:12:39 +0200
Subject: [PATCH 29/44] Update src/mrinufft/operators/off_resonance.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>
---
 src/mrinufft/operators/off_resonance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 3d43f478..cf04844c 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -221,7 +221,7 @@ class MRIFourierCorrected(FourierOperatorBase):
         Also supports Cupy arrays and Torch tensors.
         Expected shape is ``(nz, ny, nx)``.
     readout_time : np.ndarray or GPUarray, optional
-        Readout time in ``[s]`` of shape ``(nshots, npts)`` or ``(nshots * npts,)``.
+        Readout time in ``[s]`` of shape ``(n_shots, n_pts)`` or ``(n_shots * n_pts,)``.
         Also supports Cupy arrays and Torch tensors.
     n_time_segments : int, optional
         Number of time segments. The default is ``6``.

From c57f3b82e43b53b69a5f41be956db3fedbba73e4 Mon Sep 17 00:00:00 2001
From: Matteo Cencini <matteo.cencini@gmail.com>
Date: Thu, 5 Sep 2024 15:23:40 +0200
Subject: [PATCH 30/44] address PR rev 3

---
 examples/example_offresonance.py          |   5 +-
 src/mrinufft/operators/off_resonance.py   | 277 +++++++++++++++-------
 tests/operators/test_offres_exp_approx.py |  21 +-
 3 files changed, 209 insertions(+), 94 deletions(-)

diff --git a/examples/example_offresonance.py b/examples/example_offresonance.py
index c5a8ef36..a16c547c 100644
--- a/examples/example_offresonance.py
+++ b/examples/example_offresonance.py
@@ -77,8 +77,7 @@
 from mrinufft.operators.off_resonance import MRIFourierCorrected
 
 # Generate standard NUFFT operator
-NufftOperator = get_operator("finufft")
-nufft = NufftOperator(
+nufft = get_operator("finufft")(
     samples=samples,
     shape=mri_data.shape,
     density=density,
@@ -86,7 +85,7 @@
 
 # Generate Fourier Corrected operator
 mfi_nufft = MRIFourierCorrected(
-    nufft, fieldmap=b0map, readout_time=t_read, mask=brain_mask
+    nufft, b0_map=b0map, readout_time=t_read, mask=brain_mask
 )
 
 # Generate K-Space
diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index cf04844c..8f21b271 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -20,26 +20,27 @@
 
 
 def get_interpolators_from_fieldmap(
-    fieldmap, readout_time, n_time_segments=6, n_bins=(40, 10), mask=None
+    b0_map, readout_time, n_time_segments=6, n_bins=(40, 10), mask=None, r2star_map=None
 ):
-    r"""Create B and C matrices to approximate ``exp(-2j*pi*fieldmap*t)``.
+    r"""Approximate ``exp(-2j*pi*fieldmap*readout_time) ≈ Σ B_n(t)C_n(r)``.
 
-    Here, B has shape ``(n_time_segments, len(t))``
-    and C has shape ``(n_time_segments, *fieldmap.shape)``.
+    Here, B_n(t) are n_time_segments temporal coefficients and C_n(r)
+    are n_time_segments temporal spatial coefficients.
+
+    The matrix B has shape ``(n_time_segments, len(readout_time))``
+    and C has shape ``(n_time_segments, *b0_map.shape)``.
 
     From Sigpy: https://github.com/mikgroup/sigpy
     and MIRT (mri_exp_approx.m): https://web.eecs.umich.edu/~fessler/code/
 
     Parameters
     ----------
-    fieldmap : np.ndarray
-        Rate map defined as ``fieldmap = R2*_map + 1j * B0_map``.
-        ``*_map`` and ``t`` should have reciprocal units.
-        If ``zmap`` is real, assume ``zmap = B0_map``.
+    b0_map : np.ndarray
+        Static field inhomogeneities map.
+        ``b0_map`` and ``readout_time`` should have reciprocal units.
         Also supports Cupy arrays and Torch tensors.
-        Expected shape is ``(nz, ny, nx)``.
     readout_time : np.ndarray
-        Readout time in ``[s]`` of shape ``(nshots, npts)`` or ``(nshots * npts,)``.
+        Readout time in ``[s]`` of shape ``(n_shots, n_pts)`` or ``(n_shots * n_pts,)``.
         Also supports Cupy arrays and Torch tensors.
     n_time_segments : int, optional
         Number of time segments. The default is ``6``.
@@ -48,69 +49,76 @@ def get_interpolators_from_fieldmap(
         If it is a scalar, assume ``n_bins = (n_bins, 10)``.
         For real fieldmap (B0 only), ``n_bins[1]`` is ignored.
     mask : np.ndarray, optional
-        Boolean mask to avoid histogram of background values.
+        Boolean mask of the region of interest
+        (e.g., corresponding to the imaged object).
+        This is used to exclude the background fieldmap values
+        from histogram computation. Must have same shape as ``b0_map``.
         The default is ``None`` (use the whole map).
         Also supports Cupy arrays and Torch tensors.
+    r2star_map : np.ndarray, optional
+        Effective transverse relaxation map (R2*).
+        ``r2star_map`` and ``readout_time`` should have reciprocal units.
+        Must have same shape as ``b0_map``.
+        The default is ``None`` (purely imaginary field).
+        Also supports Cupy arrays and Torch tensors.
+
+    Notes
+    -----
+    The total field map used to calculate the field coefficients is
+    ``field_map = R2*_map + 1j * B0_map``. If R2* is not provided,
+    the field is purely immaginary: ``field_map = 1j * B0_map``.
 
     Returns
     -------
     B : np.ndarray
         Temporal interpolator of shape ``(n_time_segments, len(t))``.
-        Array module is the same as input fieldmap.
-    C : np.ndarray
-        Off-resonance phase map at each time segment center of shape
-        ``(n_time_segments, *fieldmap.shape)``.
-        Array module is the same as input fieldmap.
+        Array module is the same as input field_map.
+    tl : np.ndarray
+        Time segment centers of shape ``(n_time_segments,)``.
+        Array module is the same as input field_map.
+
     """
     # default
     if isinstance(n_bins, (list, tuple)) is False:
         n_bins = (n_bins, 10)
-
-    # transform to list
     n_bins = list(n_bins)
 
     # get backend and device
-    xp = get_array_module(fieldmap)
-
-    # enforce complex field
-    fieldmap = xp.asarray(fieldmap, dtype=xp.complex64)
+    xp = get_array_module(b0_map)
 
     # cast arrays to fieldmap backend
     is_torch = xp.__name__ == "torch"
 
-    if is_cuda_array(fieldmap):
+    if is_cuda_array(b0_map):
         assert CUPY_AVAILABLE, "GPU computation requires Cupy!"
         xp = cp
-        fieldmap = _to_cupy(fieldmap)
+        b0_map = _to_cupy(b0_map)
         readout_time = _to_cupy(readout_time)
         mask = _to_cupy(mask)
+        r2star_map = _to_cupy(r2star_map)
     else:
         xp = np
-        fieldmap = _to_numpy(fieldmap)
+        b0_map = _to_numpy(b0_map)
         readout_time = _to_numpy(readout_time)
         mask = _to_numpy(mask)
+        r2star_map = _to_numpy(r2star_map)
 
     readout_time = xp.asarray(readout_time, dtype=xp.float32).ravel()
     if mask is None:
-        mask = xp.ones_like(fieldmap, dtype=bool)
+        mask = xp.ones_like(b0_map, dtype=bool)
     else:
         mask = xp.asarray(mask, dtype=bool)
 
-    # get field map
-    if xp.isreal(fieldmap).all().item():
-        r2star = None
-        b0 = fieldmap
-        fieldmap = 0.0 + 1j * b0
-    else:
-        r2star = fieldmap.real
-        b0 = fieldmap.imag
-
     # Hz to radians / s
-    fieldmap = 2 * math.pi * fieldmap
+    field_map = _get_complex_fieldmap(b0_map, r2star_map)
+
+    # enforce precision
+    field_map = xp.asarray(field_map, dtype=xp.complex64)
 
     # create histograms
-    if r2star is not None:
-        z = fieldmap[mask].ravel()
+    z = field_map[mask].ravel()
+
+    if r2star_map is not None:
         z = xp.stack((z.imag, z.real), axis=1)
         hk, ze = xp.histogramdd(z, bins=n_bins)
         ze = list(ze)
@@ -123,14 +131,13 @@ def get_interpolators_from_fieldmap(
         zk = zk.T
         hk = hk.T
     else:
-        z = fieldmap[mask].ravel()
         hk, ze = xp.histogram(z.imag, bins=n_bins[0])
 
         # get bin centers
         zc = ze[1:] - (ze[1] - ze[0]) / 2
 
         # complexify
-        zk = 0 + 1j * zc  # [K 1]
+        zk = 1j * zc  # [K 1]
 
     # flatten histogram values and centers
     hk = hk.ravel()
@@ -150,21 +157,12 @@ def get_interpolators_from_fieldmap(
     B = p @ xp.exp(-zk[:, None, ...] * readout_time[None, ...])
     B = B.astype(xp.complex64)
 
-    # get spatial coeffs
-    C = xp.exp(-tl * fieldmap[..., None])
-    C = C[None, ...].swapaxes(0, -1)[
-        ..., 0
-    ]  # (..., n_time_segments) -> (n_time_segments, ...)
-
-    # clean-up of spatial coeffs
-    C = xp.nan_to_num(C, nan=0.0, posinf=0.0, neginf=0.0)
-
     # back to torch if required
     if is_torch:
         B = _to_torch(B)
-        C = _to_torch(C)
+        tl = _to_torch(tl)
 
-    return B, C
+    return B, tl
 
 
 def _outer_sum(xx, yy):
@@ -176,6 +174,8 @@ def _outer_sum(xx, yy):
 
 # TODO: /* refactor with_* decorators
 def _to_numpy(input):
+    if input is None:
+        return input
     xp = get_array_module(input)
 
     if xp.__name__ == "torch":
@@ -187,6 +187,8 @@ def _to_numpy(input):
 
 
 def _to_cupy(input):
+    if input is None:
+        return input
     return cp.asarray(input)
 
 
@@ -212,15 +214,11 @@ class MRIFourierCorrected(FourierOperatorBase):
 
     Parameters
     ----------
-    fourier_op: object of class FourierBase
-        the Fourier operator to wrap
-    fieldmap : np.ndarray, optional
-        Rate map defined as ``fieldmap = R2*_map + 1j * B0_map``.
-        ``*_map`` and ``t`` should have reciprocal units.
-        If ``zmap`` is real, assume ``zmap = B0_map``.
+    b0_map : np.ndarray
+        Static field inhomogeneities map.
+        ``b0_map`` and ``readout_time`` should have reciprocal units.
         Also supports Cupy arrays and Torch tensors.
-        Expected shape is ``(nz, ny, nx)``.
-    readout_time : np.ndarray or GPUarray, optional
+    readout_time : np.ndarray
         Readout time in ``[s]`` of shape ``(n_shots, n_pts)`` or ``(n_shots * n_pts,)``.
         Also supports Cupy arrays and Torch tensors.
     n_time_segments : int, optional
@@ -229,30 +227,47 @@ class MRIFourierCorrected(FourierOperatorBase):
         Number of histogram bins to use for ``(B0, T2*)``. The default is ``(40, 10)``
         If it is a scalar, assume ``n_bins = (n_bins, 10)``.
         For real fieldmap (B0 only), ``n_bins[1]`` is ignored.
-    mask : np.ndarray or GPUarray, optional
-        Boolean mask to avoid histogram of background values.
+    mask : np.ndarray, optional
+        Boolean mask of the region of interest
+        (e.g., corresponding to the imaged object).
+        This is used to exclude the background fieldmap values
+        from histogram computation.
         The default is ``None`` (use the whole map).
         Also supports Cupy arrays and Torch tensors.
     B : np.ndarray, optional
-        Temporal interpolator of shape ``(n_time_segments, len(t))``.
-    C : np.ndarray, optional
-        Off-resonance phase map at each time segment center of shape
-        ``(n_time_segments, *fieldmap.shape)``.
+        Temporal interpolator of shape ``(n_time_segments, len(readout_time))``.
+    tl : np.ndarray, optional
+        Time segment centers of shape ``(n_time_segments,)``.
+        Also supports Cupy arrays and Torch tensors.
+    r2star_map : np.ndarray, optional
+        Effective transverse relaxation map (R2*).
+        ``r2star_map`` and ``readout_time`` should have reciprocal units.
+        Must have same shape as ``b0_map``.
+        The default is ``None`` (purely imaginary field).
+        Also supports Cupy arrays and Torch tensors.
     backend: str, optional
         The backend to use for computations. Either 'cpu', 'gpu' or 'torch'.
         The default is `cpu`.
+
+    Notes
+    -----
+    The total field map used to calculate the field coefficients is
+    ``field_map = R2*_map + 1j * B0_map``. If R2* is not provided,
+    the field is purely immaginary: ``field_map = 1j * B0_map``.
+
     """
 
     def __init__(
         self,
         fourier_op,
-        fieldmap=None,
+        b0_map=None,
         readout_time=None,
         n_time_segments=6,
         n_bins=(40, 10),
         mask=None,
+        r2star_map=None,
         B=None,
-        C=None,
+        tl=None,
         backend="cpu",
     ):
         if backend == "gpu" and not CUPY_AVAILABLE:
@@ -272,35 +287,57 @@ def __init__(
         self.smaps = fourier_op.smaps
         self.autograd_available = fourier_op.autograd_available
 
-        if B is not None and C is not None:
+        if B is not None and tl is not None:
             self.B = self.xp.asarray(B)
-            self.C = self.xp.asarray(C)
+            self.tl = self.xp.asarray(tl)
         else:
-            fieldmap = self.xp.asarray(fieldmap)
-            self.B, self.C = get_interpolators_from_fieldmap(
-                fieldmap, readout_time, n_time_segments, n_bins, mask
+            b0_map = self.xp.asarray(b0_map)
+            self.B, self.tl = get_interpolators_from_fieldmap(
+                b0_map,
+                readout_time,
+                n_time_segments,
+                n_bins,
+                mask,
+                r2star_map,
             )
-        if self.B is None or self.C is None:
-            raise ValueError("Please either provide fieldmap and t or B and C")
-        self.n_interpolators = len(self.C)
+        if self.B is None or self.tl is None:
+            raise ValueError("Please either provide fieldmap and t or B and tl")
+        self.n_interpolators = self.B.shape[0]
+
+        # create spatial interpolator
+        field_map = _get_complex_fieldmap(b0_map, r2star_map)
+        if is_cuda_array(b0_map):
+            self.C = None
+            self.field_map = field_map
+        else:
+            self.C = _get_spatial_coefficients(field_map, self.tl)
+            self.field_map = None
 
     def op(self, data, *args):
         """Compute Forward Operation with off-resonance effect.
 
         Parameters
         ----------
-        x: numpy.ndarray or cupy.ndarray
-            N-D input image
+        x: numpy.ndarray
+            N-D input image.
+            Also supports Cupy arrays and Torch tensors.
 
         Returns
         -------
-        numpy.ndarray or cupy.ndarray
-            masked distorded N-D k-space
+        numpy.ndarray
+            Masked distorted N-D k-space.
+            Array module is the same as input data.
+
         """
         y = 0.0
         data_d = self.xp.asarray(data)
-        for idx in range(self.n_interpolators):
-            y += self.B[idx] * self._fourier_op.op(self.C[idx] * data_d, *args)
+        if self.C is not None:
+            for idx in range(self.n_interpolators):
+                y += self.B[idx] * self._fourier_op.op(self.C[idx] * data_d, *args)
+        else:
+            for idx in range(self.n_interpolators):
+                C = self.xp.exp(-self.field_map * self.tl[idx].item())
+                y += self.B[idx] * self._fourier_op.op(C * data_d, *args)
 
         return y
 
@@ -310,18 +347,82 @@ def adj_op(self, coeffs, *args):
 
         Parameters
         ----------
-        x: numpy.ndarray or cupy.ndarray
-            masked distorded N-D k-space
+        x: numpy.ndarray
+            Masked distorted N-D k-space.
+            Also supports Cupy arrays and Torch tensors.
+
 
         Returns
         -------
-            inverse Fourier transform of the distorded input k-space.
+        numpy.ndarray
+            Inverse Fourier transform of the distorted input k-space.
+            Array module is the same as input coeffs.
+
         """
         y = 0.0
         coeffs_d = self.xp.array(coeffs)
-        for idx in range(self.n_interpolators):
-            y += self.xp.conj(self.C[idx]) * self._fourier_op.adj_op(
-                self.xp.conj(self.B[idx]) * coeffs_d, *args
-            )
+        if self.C is not None:
+            for idx in range(self.n_interpolators):
+                y += self.xp.conj(self.C[idx]) * self._fourier_op.adj_op(
+                    self.xp.conj(self.B[idx]) * coeffs_d, *args
+                )
+        else:
+            for idx in range(self.n_interpolators):
+                C = self.xp.exp(-self.field_map * self.tl[idx].item())
+                y += self.xp.conj(C) * self._fourier_op.adj_op(
+                    self.xp.conj(self.B[idx]) * coeffs_d, *args
+                )
 
         return y
+
+    @staticmethod
+    def get_spatial_coefficients(field_map, tl):
+        """Compute spatial coefficients for field approximation.
+
+        Parameters
+        ----------
+        field_map : np.ndarray
+            Total field map used to calculate the field coefficients is
+            ``field_map = R2*_map + 1j * B0_map``.
+            Also supports Cupy arrays and Torch tensors.
+        tl : np.ndarray
+            Time segment centers of shape ``(n_time_segments,)``.
+            Also supports Cupy arrays and Torch tensors.
+
+        Returns
+        -------
+        C : np.ndarray
+            Off-resonance phase map at each time segment center of shape
+            ``(n_time_segments, *field_map.shape)``.
+            Array module is the same as input field_map.
+
+        """
+        return _get_spatial_coefficients(field_map, tl)
+
+
+def _get_complex_fieldmap(b0_map, r2star_map=None):
+    xp = get_array_module(b0_map)
+
+    if r2star_map is not None:
+        r2star_map = xp.asarray(r2star_map, dtype=xp.float32)
+        field_map = 2 * math.pi * (r2star_map + 1j * b0_map)
+    else:
+        field_map = 2 * math.pi * 1j * b0_map
+
+    return field_map
+
+
+def _get_spatial_coefficients(field_map, tl):
+    xp = get_array_module(field_map)
+
+    # get spatial coeffs
+    C = xp.exp(-tl * field_map[..., None])
+    C = C[None, ...].swapaxes(0, -1)[
+        ..., 0
+    ]  # (..., n_time_segments) -> (n_time_segments, ...)
+    C = xp.asarray(C, dtype=xp.complex64)
+
+    # clean-up of spatial coeffs
+    C = xp.nan_to_num(C, nan=0.0, posinf=0.0, neginf=0.0)
+
+    return C
diff --git a/tests/operators/test_offres_exp_approx.py b/tests/operators/test_offres_exp_approx.py
index 7c11976d..dfcc03f4 100644
--- a/tests/operators/test_offres_exp_approx.py
+++ b/tests/operators/test_offres_exp_approx.py
@@ -11,6 +11,7 @@
 import mrinufft
 from mrinufft._utils import get_array_module
 from mrinufft.operators.base import CUPY_AVAILABLE
+from mrinufft.operators.off_resonance import MRIFourierCorrected
 
 
 from helpers import to_interface, assert_allclose
@@ -50,10 +51,15 @@ def test_b0map_coeff(b0map, mask, array_interface):
     tread = np.linspace(0.0, 5e-3, 501, dtype=np.float32)
 
     # Generate coefficients
-    B, C = mrinufft.get_interpolators_from_fieldmap(
+    B, tl = mrinufft.get_interpolators_from_fieldmap(
         to_interface(b0map, array_interface), tread, mask=mask, n_time_segments=100
     )
 
+    # Calculate spatial coefficients
+    C = MRIFourierCorrected.get_spatial_coefficients(
+        to_interface(2 * math.pi * 1j * b0map, array_interface), tl
+    )
+
     # Assert properties
     assert B.shape == (100, 501)
     assert C.shape == (100, *b0map.shape)
@@ -77,8 +83,17 @@ def test_zmap_coeff(zmap, mask, array_interface):
     tread = np.linspace(0.0, 5e-3, 501, dtype=np.float32)
 
     # Generate coefficients
-    B, C = mrinufft.get_interpolators_from_fieldmap(
-        to_interface(zmap, array_interface), tread, mask=mask, n_time_segments=100
+    B, tl = mrinufft.get_interpolators_from_fieldmap(
+        to_interface(zmap.imag, array_interface),
+        tread,
+        mask=mask,
+        r2star_map=to_interface(zmap.real, array_interface),
+        n_time_segments=100,
+    )
+
+    # Calculate spatial coefficients
+    C = MRIFourierCorrected.get_spatial_coefficients(
+        to_interface(2 * math.pi * zmap, array_interface), tl
     )
 
     # Assert properties

From e46095d9c97a32147c6af94fb1e997817e762f31 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Fri, 6 Sep 2024 11:40:08 +0200
Subject: [PATCH 31/44] Update pyproject.toml

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 03442100..143284e7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,7 +11,7 @@ dynamic = ["version"]
 
 [project.optional-dependencies]
 
-gpunufft = ["gpuNUFFT>=0.10.0", "cupy-cuda12x"]
+gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"]
 torchkbnufft = ["torchkbnufft", "cupy-cuda12x"]
 cufinufft = ["cufinufft", "cupy-cuda12x"]
 finufft = ["finufft"]

From 5f3221049a41222bd86056a385a879d4e66bea1a Mon Sep 17 00:00:00 2001
From: Matteo Cencini <83717049+mcencini@users.noreply.github.com>
Date: Thu, 12 Sep 2024 19:18:11 +0200
Subject: [PATCH 32/44] Update src/mrinufft/operators/off_resonance.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>
---
 src/mrinufft/operators/off_resonance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mrinufft/operators/off_resonance.py b/src/mrinufft/operators/off_resonance.py
index 8f21b271..61c6a6e9 100644
--- a/src/mrinufft/operators/off_resonance.py
+++ b/src/mrinufft/operators/off_resonance.py
@@ -79,7 +79,7 @@ def get_interpolators_from_fieldmap(
 
     """
     # default
-    if isinstance(n_bins, (list, tuple)) is False:
+    if not isinstance(n_bins, (list, tuple)):
         n_bins = (n_bins, 10)
     n_bins = list(n_bins)
 

From 211a6aa79307dc52ea6660f4765eec3d4514452c Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Fri, 13 Sep 2024 14:56:37 +0200
Subject: [PATCH 33/44] \![docs_build]

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 143284e7..bfe9bcd7 100644
--- a/pyproject.toml
+++ b/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"]

From abd597035fe857cc5c5b6cb733f2d79592819d9a Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Fri, 13 Sep 2024 16:25:45 +0200
Subject: [PATCH 34/44] Add 3D trajectory viewer

---
 src/mrinufft/trajectories/display3D.py | 77 ++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)
 create mode 100644 src/mrinufft/trajectories/display3D.py

diff --git a/src/mrinufft/trajectories/display3D.py b/src/mrinufft/trajectories/display3D.py
new file mode 100644
index 00000000..c6ba3973
--- /dev/null
+++ b/src/mrinufft/trajectories/display3D.py
@@ -0,0 +1,77 @@
+from mrinufft import get_operator, get_density
+import numpy as np
+from typing import Tuple
+from mrinufft.io import read_trajectory
+
+
+def get_gridded_trajectory(
+    shots: np.ndarray,
+    shape: Tuple,
+    osf: int = 1,
+    grid_type: str = "density",
+    turbo_factor: int = 176,
+    backend: str = "gpunufft",
+):
+    """
+    Compute the gridded trajectory for MRI reconstruction.
+
+    Parameters
+    ----------
+    shots : ndarray
+        The input array of shape (N, M), where N is the number of shots and M is the
+        number of samples per shot.
+    shape : tuple
+        The desired shape of the gridded trajectory.
+    osf : int, optional
+        The oversampling factor for the gridded trajectory. Default is 1.
+    grid_type : str, optional
+        The type of gridded trajectory to compute. Default is "density".
+        It can be one of the following:
+            "density" : Get the sampling density in closest number of samples per voxel.
+            Helps understand suboptimal sampling.
+            "time" : Get the sampling in time, this is helpful to view and understand 
+            off-resonance effects.
+            "inversion" : Relative inversion time at the sampling location. Needs 
+            turbo_factor to be set.
+            "holes": Show the k-space holes within a elliosoid of the k-space.
+    turbo_factor : int, optional
+        The turbo factor when sampling is with inversion. Default is 176.
+    backend : str, optional
+        The backend to use for gridding. Default is "gpunufft".
+        Note that "gpunufft" is anyway used to get the `pipe` density internally.
+
+    Returns
+    -------
+    ndarray
+        The gridded trajectory of shape `shape`.
+    """
+    samples = shots.reshape(-1, shots.shape[-1])
+    dcomp = get_density("pipe")(samples, shape)
+    grid_op = get_operator(backend)(
+        samples, [sh * osf for sh in shape], density=dcomp, upsampfac=1
+    )
+    gridded_ones = grid_op.raw_op.adj_op(np.ones(samples.shape[0]), None, True)
+    if grid_type == "density":
+        return np.abs(gridded_ones).squeeze()
+    elif grid_type == "time":
+        data = grid_op.raw_op.adj_op(
+            np.tile(np.linspace(1, 10, shots.shape[1]), (shots.shape[0],)),
+            None,
+            True,
+        )
+    elif grid_type == "inversion":
+        data = grid_op.raw_op.adj_op(
+            np.repeat(
+                np.linspace(1, 10, turbo_factor), samples.shape[0] // turbo_factor + 1
+            )[: samples.shape[0]],
+            None,
+            True,
+        )
+    elif grid_type == "holes":
+        data = np.abs(gridded_ones).squeeze() == 0
+        data[
+            np.linalg.norm(
+                np.meshgrid(*[np.linspace(-1, 1, sh) for sh in shape], indexing="ij")
+            )
+        ] = 0
+    return np.squeeze(np.abs(data) / np.abs(gridded_ones))
\ No newline at end of file

From a97be0511e569c0b9f6735719f2728829c5a119d Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Mon, 16 Sep 2024 10:52:07 +0200
Subject: [PATCH 35/44] Add supoport for grads and slew

---
 src/mrinufft/trajectories/display3D.py | 37 +++++++++++++++++++++++---
 1 file changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/mrinufft/trajectories/display3D.py b/src/mrinufft/trajectories/display3D.py
index c6ba3973..0860e166 100644
--- a/src/mrinufft/trajectories/display3D.py
+++ b/src/mrinufft/trajectories/display3D.py
@@ -1,4 +1,10 @@
 from mrinufft import get_operator, get_density
+from mrinufft.trajectories.utils import (
+    convert_trajectory_to_gradients,
+    convert_gradients_to_slew_rates,
+    KMAX,
+    DEFAULT_RASTER_TIME,
+)
 import numpy as np
 from typing import Tuple
 from mrinufft.io import read_trajectory
@@ -11,6 +17,7 @@ def get_gridded_trajectory(
     grid_type: str = "density",
     turbo_factor: int = 176,
     backend: str = "gpunufft",
+    traj_params: dict = None,
 ):
     """
     Compute the gridded trajectory for MRI reconstruction.
@@ -29,16 +36,21 @@ def get_gridded_trajectory(
         It can be one of the following:
             "density" : Get the sampling density in closest number of samples per voxel.
             Helps understand suboptimal sampling.
-            "time" : Get the sampling in time, this is helpful to view and understand 
+            "time" : Get the sampling in time, this is helpful to view and understand
             off-resonance effects.
-            "inversion" : Relative inversion time at the sampling location. Needs 
+            "inversion" : Relative inversion time at the sampling location. Needs
             turbo_factor to be set.
             "holes": Show the k-space holes within a elliosoid of the k-space.
+            "gradients": Show the gradient strengths of the k-space trajectory.
+            "slew": Show the slew rate of the k-space trajectory.
     turbo_factor : int, optional
         The turbo factor when sampling is with inversion. Default is 176.
     backend : str, optional
         The backend to use for gridding. Default is "gpunufft".
         Note that "gpunufft" is anyway used to get the `pipe` density internally.
+    traj_params : dict, optional
+        The trajectory parameters. Default is None.
+        This is only needed when `grid_type` is "gradients" or "slew".
 
     Returns
     -------
@@ -74,4 +86,23 @@ def get_gridded_trajectory(
                 np.meshgrid(*[np.linspace(-1, 1, sh) for sh in shape], indexing="ij")
             )
         ] = 0
-    return np.squeeze(np.abs(data) / np.abs(gridded_ones))
\ No newline at end of file
+    elif grid_type in ["gradients", "slew"]:
+        gradients, initial_position = convert_trajectory_to_gradients(
+            shots,
+            norm_factor=KMAX,
+            resolution=np.asarray(traj_params["FOV"])
+            / np.asarray(traj_params["img_size"]),
+            raster_time=DEFAULT_RASTER_TIME,
+            gamma=traj_params["gamma"],
+        )
+        if grid_type == "gradients":
+            data = np.hstack(
+                [gradients, np.zeros((gradients.shape[0], 1, gradients.shape[2]))]
+            )
+        else:
+            slews, _ = convert_gradients_to_slew_rates(gradients, DEFAULT_RASTER_TIME)
+            data = np.hstack([slews, np.zeros((slews.shape[0], 2, slews.shape[2]))])
+        data = grid_op.raw_op.adj_op(
+            np.linalg.norm(data, axis=-1).flatten(), None, True
+        )
+    return np.squeeze(np.abs(data))  # / np.abs(gridded_ones))

From 2aff0f4ec9b791771120c8709dddfced22e4f202 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Fri, 27 Sep 2024 13:10:48 +0200
Subject: [PATCH 36/44] Fix ruff

---
 src/mrinufft/trajectories/display3D.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/mrinufft/trajectories/display3D.py b/src/mrinufft/trajectories/display3D.py
index 0860e166..e2561357 100644
--- a/src/mrinufft/trajectories/display3D.py
+++ b/src/mrinufft/trajectories/display3D.py
@@ -1,3 +1,5 @@
+"""Utils for displaying 3D trajectories."""
+
 from mrinufft import get_operator, get_density
 from mrinufft.trajectories.utils import (
     convert_trajectory_to_gradients,

From b6598a3e815fdbc9552e4188ba63d439e615e40e Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Fri, 8 Nov 2024 16:39:22 +0100
Subject: [PATCH 37/44] added example

---
 src/mrinufft/trajectories/display3D.py | 39 ++++++++++++++++----------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/src/mrinufft/trajectories/display3D.py b/src/mrinufft/trajectories/display3D.py
index e2561357..c6c515af 100644
--- a/src/mrinufft/trajectories/display3D.py
+++ b/src/mrinufft/trajectories/display3D.py
@@ -15,11 +15,13 @@
 def get_gridded_trajectory(
     shots: np.ndarray,
     shape: Tuple,
-    osf: int = 1,
     grid_type: str = "density",
-    turbo_factor: int = 176,
+    osf: int = 1,
     backend: str = "gpunufft",
     traj_params: dict = None,
+    turbo_factor: int = 176,
+    elliptical_samp: bool = True,
+    threshold: float = 1e-3,
 ):
     """
     Compute the gridded trajectory for MRI reconstruction.
@@ -31,8 +33,6 @@ def get_gridded_trajectory(
         number of samples per shot.
     shape : tuple
         The desired shape of the gridded trajectory.
-    osf : int, optional
-        The oversampling factor for the gridded trajectory. Default is 1.
     grid_type : str, optional
         The type of gridded trajectory to compute. Default is "density".
         It can be one of the following:
@@ -45,15 +45,22 @@ def get_gridded_trajectory(
             "holes": Show the k-space holes within a elliosoid of the k-space.
             "gradients": Show the gradient strengths of the k-space trajectory.
             "slew": Show the slew rate of the k-space trajectory.
-    turbo_factor : int, optional
-        The turbo factor when sampling is with inversion. Default is 176.
+    osf : int, optional
+        The oversampling factor for the gridded trajectory. Default is 1.
     backend : str, optional
         The backend to use for gridding. Default is "gpunufft".
         Note that "gpunufft" is anyway used to get the `pipe` density internally.
     traj_params : dict, optional
         The trajectory parameters. Default is None.
-        This is only needed when `grid_type` is "gradients" or "slew".
-
+        This is only needed when `grid_type` is "gradients" or "slew".    
+    turbo_factor : int, optional
+        The turbo factor when sampling is with inversion. Default is 176.
+    elliptical_samp : bool, optional
+        Whether to use elliptical sampling. Default is True.
+        This is useful while analyzing the k-space holes.
+    threshold: float, optional
+        The threshold for the k-space holes. Default is 1e-3.
+        
     Returns
     -------
     ndarray
@@ -82,12 +89,14 @@ def get_gridded_trajectory(
             True,
         )
     elif grid_type == "holes":
-        data = np.abs(gridded_ones).squeeze() == 0
-        data[
-            np.linalg.norm(
-                np.meshgrid(*[np.linspace(-1, 1, sh) for sh in shape], indexing="ij")
-            )
-        ] = 0
+        data = np.abs(gridded_ones).squeeze() < threshold
+        if elliptical_samp:
+            data[
+                np.linalg.norm(
+                    np.meshgrid(*[np.linspace(-1, 1, sh) for sh in shape], indexing="ij"),
+                    axis=0,
+                ) > 1
+            ] = 0
     elif grid_type in ["gradients", "slew"]:
         gradients, initial_position = convert_trajectory_to_gradients(
             shots,
@@ -107,4 +116,4 @@ def get_gridded_trajectory(
         data = grid_op.raw_op.adj_op(
             np.linalg.norm(data, axis=-1).flatten(), None, True
         )
-    return np.squeeze(np.abs(data))  # / np.abs(gridded_ones))
+    return np.squeeze(np.abs(data))  

From dec680bf81a88efea27ba5cabb6413ffc2362ebf Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Fri, 8 Nov 2024 16:40:35 +0100
Subject: [PATCH 38/44] \!docs_build finalizing

---
 src/mrinufft/trajectories/display3D.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/mrinufft/trajectories/display3D.py b/src/mrinufft/trajectories/display3D.py
index c6c515af..2c10bb72 100644
--- a/src/mrinufft/trajectories/display3D.py
+++ b/src/mrinufft/trajectories/display3D.py
@@ -52,7 +52,7 @@ def get_gridded_trajectory(
         Note that "gpunufft" is anyway used to get the `pipe` density internally.
     traj_params : dict, optional
         The trajectory parameters. Default is None.
-        This is only needed when `grid_type` is "gradients" or "slew".    
+        This is only needed when `grid_type` is "gradients" or "slew".
     turbo_factor : int, optional
         The turbo factor when sampling is with inversion. Default is 176.
     elliptical_samp : bool, optional
@@ -60,7 +60,7 @@ def get_gridded_trajectory(
         This is useful while analyzing the k-space holes.
     threshold: float, optional
         The threshold for the k-space holes. Default is 1e-3.
-        
+
     Returns
     -------
     ndarray
@@ -93,9 +93,12 @@ def get_gridded_trajectory(
         if elliptical_samp:
             data[
                 np.linalg.norm(
-                    np.meshgrid(*[np.linspace(-1, 1, sh) for sh in shape], indexing="ij"),
+                    np.meshgrid(
+                        *[np.linspace(-1, 1, sh) for sh in shape], indexing="ij"
+                    ),
                     axis=0,
-                ) > 1
+                )
+                > 1
             ] = 0
     elif grid_type in ["gradients", "slew"]:
         gradients, initial_position = convert_trajectory_to_gradients(
@@ -116,4 +119,4 @@ def get_gridded_trajectory(
         data = grid_op.raw_op.adj_op(
             np.linalg.norm(data, axis=-1).flatten(), None, True
         )
-    return np.squeeze(np.abs(data))  
+    return np.squeeze(np.abs(data))

From 27541583092025b7cc96159dd106b0214aee62e9 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Wed, 13 Nov 2024 09:01:08 +0100
Subject: [PATCH 39/44] \!docs_build added new example

---
 examples/GPU/example_3d_trajectory_display.py | 99 +++++++++++++++++++
 1 file changed, 99 insertions(+)
 create mode 100644 examples/GPU/example_3d_trajectory_display.py

diff --git a/examples/GPU/example_3d_trajectory_display.py b/examples/GPU/example_3d_trajectory_display.py
new file mode 100644
index 00000000..03d5c16e
--- /dev/null
+++ b/examples/GPU/example_3d_trajectory_display.py
@@ -0,0 +1,99 @@
+"""
+======================
+3D Trajectory Display
+======================
+
+In this example, we show some tools available through the `mrinufft` package to display 3D trajectories.
+This is useful to understand the sampling pattern of the k-space data, and to visualize the trajectory, see sampling time, gradient strength, slew rates etc.
+Also, another key useful feature is that it enables us to see the density of sampling pattern in the k-space, and help analyze k-space holes, which can help debug artifacts in reconstructions.
+"""
+
+# %%
+# Imports
+from mrinufft.trajectories.display3D import get_gridded_trajectory
+import mrinufft.trajectories.trajectory3D as mtt
+from mrinufft.trajectories.utils import Gammas
+import matplotlib.pyplot as plt
+from matplotlib import gridspec
+import numpy as np
+
+
+# %%
+# Utility function to plot mid-plane slices for 3D volumes
+def plot_slices(axs, volume, title=""):
+    def set_labels(ax, axis_num=None):
+        ax.set_xticks([0, 32, 64])
+        ax.set_yticks([0, 32, 64])
+        ax.set_xticklabels([r"$-\pi$", 0, r"$\pi$"])
+        ax.set_yticklabels([r"$-\pi$", 0, r"$\pi$"])
+        if axis_num is not None:
+            ax.set_xlabel(r"$k_" + "zxy"[axis_num] + r"$")
+            ax.set_ylabel(r"$k_" + "yzx"[axis_num] + r"$")
+
+    for i in range(3):
+        volume = np.rollaxis(volume, i, 0)
+        axs[i].imshow(volume[volume.shape[0] // 2])
+        axs[i].set_title(
+            ((title + f"\n") if i == 0 else "") + r"$k_{" + "xyz"[i] + r"}=0$"
+        )
+        set_labels(axs[i], i)
+
+
+def create_grid(grid_type, title="", **kwargs):
+    fig, axs = plt.subplots(3, 3, figsize=(10, 10))
+    for i, (name, traj) in enumerate(trajectories.items()):
+        grid = get_gridded_trajectory(
+            traj, (64, 64, 64), grid_type=grid_type, traj_params=traj_params, **kwargs
+        )
+        plot_slices(axs[:, i], grid, title=name)
+    plt.tight_layout()
+    plt.suptitle(title)
+    plt.show()
+
+
+# %%
+# Create a bunch of sample trajectories
+trajectories = {
+    "Radial": mtt.initialize_3D_phyllotaxis_radial(64 * 8, 64),
+    "FLORETs": mtt.initialize_3D_floret(64 * 8, 64, nb_revolutions=2),
+    "Yarn Ball": mtt.initialize_3D_seiffert_spiral(64 * 8, 64),
+}
+traj_params = {
+    "FOV": (0.23, 0.23, 0.23),
+    "img_size": (64, 64, 64),
+    "gamma": Gammas.HYDROGEN,
+}
+
+# %%
+# Display the density of the trajectories, along the 3 mid-planes. For this, make `grid_type="density"`.
+create_grid("density", "Sampling Density")
+
+
+# %%
+# Display the sample time of the trajectories, along the 3 mid-planes. For this, make `grid_type="time"`.
+# This helps in obtaining relative sampling times of the k-space sampling pattern, which helps debug off-resonance issues
+create_grid("time", "Sampling Time")
+
+# %%
+# Display the inversion time of the trajectories, along the 3 mid-planes. For this, make `grid_type="inversion"`.
+# This helps in obtaining the inversion time when particular region of k-space is sampled, assuming the trajectories are time ordered.
+# This helps understand any issues for imaging involving inversion recovery.
+# The argument `turbo_factor` can be used to tell what is the number of echoes between 2 inversion pulses.
+create_grid("inversion", "Inversion Time", turbo_factor=64)
+
+# %%
+# Display the k-space holes in the trajectories, along the 3 mid-planes. For this, make `grid_type="holes"`.
+# This helps in understanding the k-space holes, and can help debug artifacts in reconstructions.
+create_grid("holes", "K-space Holes", threshold=1e-2)
+
+# %%
+# Display the gradient strength of the trajectories, along the 3 mid-planes. For this, make `grid_type="gradients"`.
+# This helps in understanding the gradient strength applied at specific k-space region.
+# This can also be used as a surrogate to k-space "velocity", i.e. how fast does trajectory pass through a given region in k-space
+create_grid("gradients", "Gradient Strength")
+
+# %%
+# Display the slew rates of the trajectories, along the 3 mid-planes. For this, make `grid_type="slew"`.
+# This helps in understanding the slew rates applied at specific k-space region.
+# This can also be used as a surrogate to k-space "acceleration", i.e. how fast does trajectory change in a given region in k-space
+create_grid("slew", "Slew Rates")

From 96d985037c00b104d575ee7e7e12202747d1d994 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Mon, 25 Nov 2024 09:18:11 +0100
Subject: [PATCH 40/44] Apply suggestions from code review
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Guillaume Daval-Frérot <gdavalfrerot@gmail.com>
---
 examples/GPU/example_3d_trajectory_display.py | 12 ++++++------
 src/mrinufft/trajectories/display3D.py        |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/examples/GPU/example_3d_trajectory_display.py b/examples/GPU/example_3d_trajectory_display.py
index 03d5c16e..d7cf46d5 100644
--- a/examples/GPU/example_3d_trajectory_display.py
+++ b/examples/GPU/example_3d_trajectory_display.py
@@ -1,11 +1,11 @@
 """
 ======================
-3D Trajectory Display
+3D Trajectory display
 ======================
 
-In this example, we show some tools available through the `mrinufft` package to display 3D trajectories.
-This is useful to understand the sampling pattern of the k-space data, and to visualize the trajectory, see sampling time, gradient strength, slew rates etc.
-Also, another key useful feature is that it enables us to see the density of sampling pattern in the k-space, and help analyze k-space holes, which can help debug artifacts in reconstructions.
+In this example, we show some tools available to display 3D trajectories.
+It can be used to understand the k-space sampling patterns, visualize the trajectories, see the sampling times, gradient strengths, slew rates etc.
+Another key feature is to display the sampling density in k-space, for example to check for k-space holes or irregularities in the learning-based trajectories that would lead to artifacts in the images.
 """
 
 # %%
@@ -70,8 +70,8 @@ def create_grid(grid_type, title="", **kwargs):
 
 
 # %%
-# Display the sample time of the trajectories, along the 3 mid-planes. For this, make `grid_type="time"`.
-# This helps in obtaining relative sampling times of the k-space sampling pattern, which helps debug off-resonance issues
+# Display the sampling times over the trajectories, along the 3 mid-planes. For this, make `grid_type="time"`.
+# It helps to check the sampling times over the k-space trajectories, which can be responsible for excessive off-resonance artifacts.
 create_grid("time", "Sampling Time")
 
 # %%
diff --git a/src/mrinufft/trajectories/display3D.py b/src/mrinufft/trajectories/display3D.py
index 2c10bb72..48863adb 100644
--- a/src/mrinufft/trajectories/display3D.py
+++ b/src/mrinufft/trajectories/display3D.py
@@ -42,7 +42,7 @@ def get_gridded_trajectory(
             off-resonance effects.
             "inversion" : Relative inversion time at the sampling location. Needs
             turbo_factor to be set.
-            "holes": Show the k-space holes within a elliosoid of the k-space.
+            "holes": Show the k-space holes within a ellipsoid of the k-space.
             "gradients": Show the gradient strengths of the k-space trajectory.
             "slew": Show the slew rate of the k-space trajectory.
     osf : int, optional

From 2594b41be166559563e182d16f831281c920453e Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Tue, 26 Nov 2024 11:53:35 +0100
Subject: [PATCH 41/44] Handle a bunch of comments

---
 src/mrinufft/trajectories/display3D.py | 46 ++++++++++++++++++--------
 1 file changed, 32 insertions(+), 14 deletions(-)

diff --git a/src/mrinufft/trajectories/display3D.py b/src/mrinufft/trajectories/display3D.py
index 48863adb..2f8c54c9 100644
--- a/src/mrinufft/trajectories/display3D.py
+++ b/src/mrinufft/trajectories/display3D.py
@@ -25,26 +25,35 @@ def get_gridded_trajectory(
 ):
     """
     Compute the gridded trajectory for MRI reconstruction.
-
+    
+    This function helps in gridding a k-space sampling trajectory to a desired shape.
+    The gridding process can be carried out to reflect the sampling density, 
+    sampling time, inversion time, k-space holes, gradient strengths, or slew rates.
+    Please check `grid_type` parameter to know the benefits of each type of gridding.
+    
     Parameters
     ----------
     shots : ndarray
-        The input array of shape (N, M), where N is the number of shots and M is the
-        number of samples per shot.
+        The input array of shape (N, M, D), where N is the number of shots and M is the
+        number of samples per shot and D is the dimension of the trajectory (usually 3)
     shape : tuple
         The desired shape of the gridded trajectory.
     grid_type : str, optional
         The type of gridded trajectory to compute. Default is "density".
         It can be one of the following:
-            "density" : Get the sampling density in closest number of samples per voxel.
-            Helps understand suboptimal sampling.
-            "time" : Get the sampling in time, this is helpful to view and understand
-            off-resonance effects.
-            "inversion" : Relative inversion time at the sampling location. Needs
-            turbo_factor to be set.
-            "holes": Show the k-space holes within a ellipsoid of the k-space.
-            "gradients": Show the gradient strengths of the k-space trajectory.
-            "slew": Show the slew rate of the k-space trajectory.
+        "density" : Get the sampling density in closest number of samples per voxel.
+            Helps understand suboptimal sampling, by showcasing regions with strong 
+            oversampling.
+        "time" : Showcases when the k-space data is acquired in time.
+            This is helpful to view and understand off-resonance effects.
+            Generally, lower off-resonance effects occur when the sampling trajectory 
+            has smoother k-space sampling time over the k-space.
+        "inversion" : Relative inversion time at the sampling location. Needs
+            turbo_factor to be set. This is useful for analyzing the exact inversion 
+            time when the k-space is acquired, for sequences like MP(2)RAGE.
+        "holes": Show the k-space holes within a ellipsoid of the k-space.    
+        "gradients": Show the gradient strengths of the k-space trajectory.
+        "slew": Show the slew rate of the k-space trajectory.
     osf : int, optional
         The oversampling factor for the gridded trajectory. Default is 1.
     backend : str, optional
@@ -53,13 +62,20 @@ def get_gridded_trajectory(
     traj_params : dict, optional
         The trajectory parameters. Default is None.
         This is only needed when `grid_type` is "gradients" or "slew".
+        The parameters needed include `img_size`, `FOV`, and `gamma` of the sequence.
+        Generally these values are stored in the header of the trajectory file.
     turbo_factor : int, optional
-        The turbo factor when sampling is with inversion. Default is 176.
+        The turbo factor when sampling is with inversion. Default is 176, which is
+        the default turbo factor for MPRAGE acquisitions at 1mm whole 
+        brain acquisitions.
     elliptical_samp : bool, optional
         Whether to use elliptical sampling. Default is True.
-        This is useful while analyzing the k-space holes.
+        This is useful while analyzing the k-space holes, especially if the k-space
+        trajectory is expected to be elliptical sampling of k-space 
+        (i.e. ellipsoid over cuboid).
     threshold: float, optional
         The threshold for the k-space holes. Default is 1e-3.
+        This value is set heuristically to visualize the k-space hole.
 
     Returns
     -------
@@ -91,6 +107,8 @@ def get_gridded_trajectory(
     elif grid_type == "holes":
         data = np.abs(gridded_ones).squeeze() < threshold
         if elliptical_samp:
+            # If the trajectory uses elliptical sampling, ignore the k-space holes 
+            # outside the ellipsoid.
             data[
                 np.linalg.norm(
                     np.meshgrid(

From 4a07125cd951a24c4a8c6650d029b62af406c260 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Tue, 26 Nov 2024 12:49:47 +0100
Subject: [PATCH 42/44] Examples

---
 examples/GPU/example_3d_trajectory_display.py | 44 +++++++++++--------
 1 file changed, 26 insertions(+), 18 deletions(-)

diff --git a/examples/GPU/example_3d_trajectory_display.py b/examples/GPU/example_3d_trajectory_display.py
index d7cf46d5..3d365ac2 100644
--- a/examples/GPU/example_3d_trajectory_display.py
+++ b/examples/GPU/example_3d_trajectory_display.py
@@ -39,16 +39,14 @@ def set_labels(ax, axis_num=None):
         set_labels(axs[i], i)
 
 
-def create_grid(grid_type, title="", **kwargs):
+def create_grid(grid_type, trajectories, **kwargs):
     fig, axs = plt.subplots(3, 3, figsize=(10, 10))
     for i, (name, traj) in enumerate(trajectories.items()):
         grid = get_gridded_trajectory(
-            traj, (64, 64, 64), grid_type=grid_type, traj_params=traj_params, **kwargs
+            traj, traj_params['img_size'], grid_type=grid_type, traj_params=traj_params, **kwargs
         )
         plot_slices(axs[:, i], grid, title=name)
     plt.tight_layout()
-    plt.suptitle(title)
-    plt.show()
 
 
 # %%
@@ -56,7 +54,7 @@ def create_grid(grid_type, title="", **kwargs):
 trajectories = {
     "Radial": mtt.initialize_3D_phyllotaxis_radial(64 * 8, 64),
     "FLORETs": mtt.initialize_3D_floret(64 * 8, 64, nb_revolutions=2),
-    "Yarn Ball": mtt.initialize_3D_seiffert_spiral(64 * 8, 64),
+    "Seiffert Spirals": mtt.initialize_3D_seiffert_spiral(64 * 8, 64),
 }
 traj_params = {
     "FOV": (0.23, 0.23, 0.23),
@@ -66,34 +64,44 @@ def create_grid(grid_type, title="", **kwargs):
 
 # %%
 # Display the density of the trajectories, along the 3 mid-planes. For this, make `grid_type="density"`.
-create_grid("density", "Sampling Density")
+create_grid("density", trajectories)
+plt.suptitle("Sampling Density")
+plt.show()
 
 
 # %%
-# Display the sampling times over the trajectories, along the 3 mid-planes. For this, make `grid_type="time"`.
+# Display the sampling times over the trajectories. For this, make `grid_type="time"`.
 # It helps to check the sampling times over the k-space trajectories, which can be responsible for excessive off-resonance artifacts.
-create_grid("time", "Sampling Time")
+create_grid("time", trajectories)
+plt.suptitle("Sampling Time")
+plt.show()
 
 # %%
-# Display the inversion time of the trajectories, along the 3 mid-planes. For this, make `grid_type="inversion"`.
+# Display the inversion time of the trajectories. For this, make `grid_type="inversion"`.
 # This helps in obtaining the inversion time when particular region of k-space is sampled, assuming the trajectories are time ordered.
 # This helps understand any issues for imaging involving inversion recovery.
 # The argument `turbo_factor` can be used to tell what is the number of echoes between 2 inversion pulses.
-create_grid("inversion", "Inversion Time", turbo_factor=64)
-
+create_grid("inversion", trajectories, turbo_factor=64)
+plt.suptitle("Inversion Time")
+plt.show()
 # %%
-# Display the k-space holes in the trajectories, along the 3 mid-planes. For this, make `grid_type="holes"`.
+# Display the k-space holes in the trajectories. For this, make `grid_type="holes"`.
 # This helps in understanding the k-space holes, and can help debug artifacts in reconstructions.
-create_grid("holes", "K-space Holes", threshold=1e-2)
-
+create_grid("holes", trajectories, threshold=1e-2)
+plt.suptitle("K-space Holes")
+plt.show()
 # %%
-# Display the gradient strength of the trajectories, along the 3 mid-planes. For this, make `grid_type="gradients"`.
+# Display the gradient strength of the trajectories. For this, make `grid_type="gradients"`.
 # This helps in understanding the gradient strength applied at specific k-space region.
 # This can also be used as a surrogate to k-space "velocity", i.e. how fast does trajectory pass through a given region in k-space
-create_grid("gradients", "Gradient Strength")
+create_grid("gradients", trajectories)
+plt.suptitle("Gradient Strength")
+plt.show()
 
 # %%
-# Display the slew rates of the trajectories, along the 3 mid-planes. For this, make `grid_type="slew"`.
+# Display the slew rates of the trajectories. For this, make `grid_type="slew"`.
 # This helps in understanding the slew rates applied at specific k-space region.
 # This can also be used as a surrogate to k-space "acceleration", i.e. how fast does trajectory change in a given region in k-space
-create_grid("slew", "Slew Rates")
+create_grid("slew", trajectories)
+plt.suptitle("Slew Rates")
+plt.show()

From dad26ce61b1d6bdcf5a2f5bb4b12d9653567a595 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Tue, 26 Nov 2024 12:49:56 +0100
Subject: [PATCH 43/44] PEP

---
 examples/GPU/example_3d_trajectory_display.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/examples/GPU/example_3d_trajectory_display.py b/examples/GPU/example_3d_trajectory_display.py
index 3d365ac2..641d55c4 100644
--- a/examples/GPU/example_3d_trajectory_display.py
+++ b/examples/GPU/example_3d_trajectory_display.py
@@ -43,7 +43,11 @@ def create_grid(grid_type, trajectories, **kwargs):
     fig, axs = plt.subplots(3, 3, figsize=(10, 10))
     for i, (name, traj) in enumerate(trajectories.items()):
         grid = get_gridded_trajectory(
-            traj, traj_params['img_size'], grid_type=grid_type, traj_params=traj_params, **kwargs
+            traj,
+            traj_params["img_size"],
+            grid_type=grid_type,
+            traj_params=traj_params,
+            **kwargs,
         )
         plot_slices(axs[:, i], grid, title=name)
     plt.tight_layout()

From 98cb154dcafb66d7b43ff8611f5dd47876562ff8 Mon Sep 17 00:00:00 2001
From: Chaithya G R <chaithyagr@gmail.com>
Date: Tue, 26 Nov 2024 12:50:06 +0100
Subject: [PATCH 44/44] PEP

---
 src/mrinufft/trajectories/display3D.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/mrinufft/trajectories/display3D.py b/src/mrinufft/trajectories/display3D.py
index 2f8c54c9..fda76007 100644
--- a/src/mrinufft/trajectories/display3D.py
+++ b/src/mrinufft/trajectories/display3D.py
@@ -25,12 +25,12 @@ def get_gridded_trajectory(
 ):
     """
     Compute the gridded trajectory for MRI reconstruction.
-    
+
     This function helps in gridding a k-space sampling trajectory to a desired shape.
-    The gridding process can be carried out to reflect the sampling density, 
+    The gridding process can be carried out to reflect the sampling density,
     sampling time, inversion time, k-space holes, gradient strengths, or slew rates.
     Please check `grid_type` parameter to know the benefits of each type of gridding.
-    
+
     Parameters
     ----------
     shots : ndarray
@@ -42,16 +42,16 @@ def get_gridded_trajectory(
         The type of gridded trajectory to compute. Default is "density".
         It can be one of the following:
         "density" : Get the sampling density in closest number of samples per voxel.
-            Helps understand suboptimal sampling, by showcasing regions with strong 
+            Helps understand suboptimal sampling, by showcasing regions with strong
             oversampling.
         "time" : Showcases when the k-space data is acquired in time.
             This is helpful to view and understand off-resonance effects.
-            Generally, lower off-resonance effects occur when the sampling trajectory 
+            Generally, lower off-resonance effects occur when the sampling trajectory
             has smoother k-space sampling time over the k-space.
         "inversion" : Relative inversion time at the sampling location. Needs
-            turbo_factor to be set. This is useful for analyzing the exact inversion 
+            turbo_factor to be set. This is useful for analyzing the exact inversion
             time when the k-space is acquired, for sequences like MP(2)RAGE.
-        "holes": Show the k-space holes within a ellipsoid of the k-space.    
+        "holes": Show the k-space holes within a ellipsoid of the k-space.
         "gradients": Show the gradient strengths of the k-space trajectory.
         "slew": Show the slew rate of the k-space trajectory.
     osf : int, optional
@@ -66,12 +66,12 @@ def get_gridded_trajectory(
         Generally these values are stored in the header of the trajectory file.
     turbo_factor : int, optional
         The turbo factor when sampling is with inversion. Default is 176, which is
-        the default turbo factor for MPRAGE acquisitions at 1mm whole 
+        the default turbo factor for MPRAGE acquisitions at 1mm whole
         brain acquisitions.
     elliptical_samp : bool, optional
         Whether to use elliptical sampling. Default is True.
         This is useful while analyzing the k-space holes, especially if the k-space
-        trajectory is expected to be elliptical sampling of k-space 
+        trajectory is expected to be elliptical sampling of k-space
         (i.e. ellipsoid over cuboid).
     threshold: float, optional
         The threshold for the k-space holes. Default is 1e-3.
@@ -107,7 +107,7 @@ def get_gridded_trajectory(
     elif grid_type == "holes":
         data = np.abs(gridded_ones).squeeze() < threshold
         if elliptical_samp:
-            # If the trajectory uses elliptical sampling, ignore the k-space holes 
+            # If the trajectory uses elliptical sampling, ignore the k-space holes
             # outside the ellipsoid.
             data[
                 np.linalg.norm(