Gram Operator with Toeplitz Embedding

This example demonstrates how the auto-adjoint Gram operator of the NUFFT operator can be efficiently implemented using Toeplitz embedding.

This approach leverages the convolutional structure of the Gram operator to reduce computational complexity, making it suitable for large-scale imaging problems.

import matplotlib.pyplot as plt
import numpy as np

from mrinufft.operators import (
    get_operator,
    compute_toeplitz_kernel,
    apply_toeplitz_kernel,
)
from mrinufft.display import display_2D_trajectory
from mrinufft.trajectories import initialize_2D_spiral
from mrinufft.density import voronoi

plt.rcParams["image.cmap"] = "gray"

/volatile/github-ci-mind-inria/gpu_mind_runner/_work/mri-nufft/mri-nufft/.venv/lib/python3.12/site-packages/cupyx/jit/_interface.py:247: FutureWarning: cupyx.jit.rawkernel is experimental. The interface can change in the future.
  cupy._util.experimental('cupyx.jit.rawkernel')

Data preparation

Image loading

For realistic a 2D image we will use the BrainWeb dataset, installable using pip install brainweb-dl.

from brainweb_dl import get_mri

mri_data = get_mri(0, "T1").astype(np.complex64)
# toeplitz embedding requires the image to be cropped to an even size.
mri_data = np.flip(mri_data, axis=(0, 1, 2))[90, :-1, :-1]
mri_data = np.ascontiguousarray(mri_data)

plt.imshow(abs(mri_data))
plt.axis("off")
plt.title("Groundtruth")
plt.show()

Trajectory generation

samples = initialize_2D_spiral(Nc=24, Ns=5000, nb_revolutions=10)
density = voronoi(samples)

display_2D_trajectory(samples)
plt.show()

Operator setup

nufft = get_operator("finufft")(samples, mri_data.shape, density=density)

Naive Gram operator

gram_naive = nufft.adj_op(nufft.op(mri_data))

gram_optim = nufft.gram_op(mri_data)  # will compute the kernel internally once.

relative_squared_error = (np.abs(gram_naive - gram_optim) / np.abs(gram_naive)) ** 2
plt.figure(figsize=(12, 4))
plt.subplot(131)
plt.imshow(np.abs(gram_naive), vmin=0, vmax=np.percentile(np.abs(gram_naive), 99))
plt.title("Gram Naive")
plt.axis("off")
plt.subplot(132)
plt.imshow(np.abs(gram_optim), vmin=0, vmax=np.percentile(np.abs(gram_optim), 99))
plt.title("Gram Optimized")
plt.axis("off")
plt.subplot(133)
plt.imshow(
    relative_squared_error,
    vmin=0,
    vmax=np.percentile(relative_squared_error, 99),
)
plt.title("Relative Squared Error")
plt.axis("off")
plt.colorbar(fraction=0.046, pad=0.04)
plt.show()

Comparing the timings

from time import perf_counter

naive_times = []
for _ in range(100):
    tic = perf_counter()
    gram_naive = nufft.adj_op(nufft.op(mri_data))
    toc = perf_counter()
    naive_times.append(toc - tic)

setup_times = []
apply_times = []

for _ in range(10):
    tic = perf_counter()
    toeplitz_kernel = nufft.compute_toeplitz_kernel()
    toc = perf_counter()
    setup_times.append(toc - tic)

padded_array = np.zeros_like(toeplitz_kernel)
for _ in range(100):
    tic = perf_counter()
    gram_optim = apply_toeplitz_kernel(
        mri_data, nufft._toeplitz_kernel, padded_array=padded_array
    )
    toc = perf_counter()
    apply_times.append(toc - tic)

print(
    f"Naive Gram time: {np.mean(naive_times)*1e3:.2f} ms ± {np.std(naive_times)*1e3:.2f} ms"
)
print(
    f"Toeplitz setup time: {np.mean(setup_times)*1e3:.2f} ms ± {np.std(setup_times)*1e3:.2f} ms"
)
print(
    f"Toeplitz apply time: {np.mean(apply_times)*1e3:.2f} ms ± {np.std(apply_times)*1e3:.2f} ms"
)

/volatile/github-ci-mind-inria/gpu_mind_runner/_work/mri-nufft/mri-nufft/src/mrinufft/operators/toeplitz.py:263: ComplexWarning: Casting complex values to real discards the imaginary part
  padded_array[tl_corner] = image
Naive Gram time: 8.64 ms ± 1.04 ms
Toeplitz setup time: 43.78 ms ± 3.29 ms
Toeplitz apply time: 4.22 ms ± 0.37 ms

Going Further

The Toeplitz embedding technique only speeds up the application of the Gram operator for the NUFFT steps, and the expected speedup depends on the problem size and the trajectory. For more details on the implication for the different backends refer to the Benchmark