Gadgets with mri-nufft Data: Subspace and B0 ORC

This example uses a common data pipeline for both PyGROG gadgets:

1. BrainWeb phantom from brainweb-dl.

2. Trajectory + k-space simulation + reference adjoint from mri-nufft.

3. GROG gridding to feed SparseFFT.

4. Comparison against mri-nufft reference formulations.

import matplotlib.pyplot as plt
import numpy as np
import torch

from mrinufft import get_operator, initialize_2D_spiral
from mrinufft.density import voronoi
from mrinufft.extras import fse_simulation, get_brainweb_map, make_b0map
from mrinufft.operators.off_resonance import MRIFourierCorrected
from mrinufft.trajectories.utils import Acquisition
from pygrog.calib import GrogInterpolator
from pygrog.gadgets import (
    OffResonanceGadget,
    SubspaceGadget,
)
from pygrog.operator import SparseFFT



m0, t1, t2 = get_brainweb_map(0)

Downloading phantom_1.0mm_normal_crisp: 0.00B [00:00, ?B/s]


Downloading phantom_1.0mm_normal_crisp: 1.00kB [00:00, 7.67kB/s]


Downloading phantom_1.0mm_normal_crisp: 193kB [00:00, 1.00MB/s]


Downloading phantom_1.0mm_normal_crisp: 481kB [00:00, 1.66MB/s]


Downloading phantom_1.0mm_normal_crisp: 689kB [00:00, 1.76MB/s]


                                                               /home/docs/checkouts/readthedocs.org/user_builds/pygrog/envs/latest/lib/python3.12/site-packages/mrinufft/_utils.py:67: UserWarning: Samples will be rescaled to [-pi, pi), assuming they were in [-0.5, 0.5)
  warnings.warn(

Gadget 1: SubspaceSparseFFT for Multi-Echo FSE

Learn a low-rank basis from a dictionary of FSE signal evolutions, then project GROG-interpolated k-space onto that basis to extract subspace coefficients.

/home/docs/checkouts/readthedocs.org/user_builds/pygrog/envs/latest/lib/python3.12/site-packages/mrinufft/_utils.py:67: UserWarning: Samples will be rescaled to [-pi, pi), assuming they were in [-0.5, 0.5)
  warnings.warn(

PyGROG Subspace gadget API call

Build a GROG plan with temporal frames as the leading stack axis, interpolate all frames, and project onto the learned subspace basis via the public decorator/wrapper API.

n_shots, n_read = samples.shape[:2]
coords_sub = np.broadcast_to(
    coords[np.newaxis, np.newaxis],  # (1, 1, n_shots, n_read, 2)
    (etl, 1, n_shots, n_read, 2),  # (T, 1, n_shots, n_read, 2)
).copy()
grog_sub = GrogInterpolator(
    shape=shape, coords=coords_sub, kernel_width=2, oversamp=1.25, image_shape=shape
)
grog_sub.calc_interp_table(calib_cart, lamda=0.01, precision=1)
base_op_sub = SparseFFT(plan=grog_sub.plan, smaps=smaps)

# kspace_frames: (T, C, n_shots*n_read) -> (T, C, 1, n_shots, n_read) -> (1, C, T, 1, n_shots, n_read)
kspace_sub = kspace_frames.reshape(
    etl, n_coils, 1, n_shots, n_read
).transpose(  # (T, C, 1, n_shots, n_read)
    1, 0, 2, 3, 4
)[np.newaxis]  # (1, C, T, 1, n_shots, n_read)
sparse_sub = grog_sub.interpolate(kspace_sub).reshape(
    1, n_coils, *grog_sub.plan.natural_shape
)

# Learn a rank-K temporal basis Phi from the training dictionary (K, T).
basis_torch = torch.as_tensor(_estimate_basis(train.T, rank), dtype=torch.complex64)

# Preferred public construction via wrapper class.
sub_op = SubspaceGadget(base_op_sub, basis_torch, encoding_axis=-5)
# Equivalent decorator-style wrapping (same behavior):
# sub_op = with_subspace(base_op_sub, basis_torch, encoding_axis=-5)

coeff_pygrog = np.asarray(sub_op.adjoint(sparse_sub))
# Drop leading B=1 batch axis -> (rank, H, W)
coeff_pygrog = coeff_pygrog[0]

Display all coefficients: one row per coefficient.

fig, axes = plt.subplots(rank, 3, figsize=(11, 3.5 * rank), squeeze=False)

Gadget 2: Off-resonance gadget

Model and correct for B0 field inhomogeneity on GROG-gridded data. Create a field map and per-readout timing basis, then apply ORC to reconstruct sharp images despite off-resonance blurring.

/home/docs/checkouts/readthedocs.org/user_builds/pygrog/envs/latest/lib/python3.12/site-packages/mrinufft/_array_compat.py:122: UserWarning: The input is CPU array but not C-contiguous.
  warnings.warn("The input is CPU array but not C-contiguous. ")

PyGROG off-resonance API call

Apply OffResonanceGadget to correct off-resonance blurring. The wrapper automatically constructs a low-rank temporal basis for the correction term.

n_shots, n_read = samples.shape[:2]

base_op_orc = SparseFFT(plan=grog.plan)  # no smaps
sqrt_w_orc = np.asarray(grog.plan.pre_weights)  # (n_shots, n_read, kw)

orc_pygrog = OffResonanceGadget(
    base_op_orc,
    field_map=b0_map,
    readout_time=readout_time,
    mask=brain_mask,
    n_components=-1,
    method="svd",
)

# GROG-interpolate the off-resonance k-space and pre-weight.
sparse_off = (
    grog.interpolate(
        kspace_off.reshape(n_coils, n_shots, n_read),
        ret_image=False,
    )
    * sqrt_w_orc[np.newaxis]
)

# Adjoint returns (n_coils, *image_shape) — combine with smaps
result_orc = orc_pygrog.adjoint(sparse_off)  # (n_coils, *image_shape)
image_pygrog_orc = np.abs((np.asarray(result_orc) * smaps.conj()).sum(0))

# Compose gadgets by double-decoration (recommended over instantiating
# internal sparse classes directly):
# op_joint = with_offresonance(
#     with_subspace(base_op_sub, basis_torch, encoding_axis=-5),
#     b0_map=b0_map,
#     readout_time=readout_time,
#     mask=brain_mask,
#     L=-1,
#     method="svd",
# )

truth_orc = _normalize_unit(np.abs(image_orc))