repository·main·Indexed 42 days ago
https://github.com/mmuckley/torchkbnufftA PyTorch implementation of the Non-Uniform Fast Fourier Transform (NUFFT) using Kaiser-Bessel gridding. Designed for deep learning integration, all functions are wrapped as torch.autograd.Function to support backpropagation. It offers three operation modes: Table interpolation (standard), Toeplitz-embedded FFTs (fastest), and Sparse matrix-based interpolation. Supports GPU acceleration, batched k-space trajectories as of version 1.1.0, and sensitivity maps for coil-aware operations.
The package provides utility functions for common tasks in NUFFT-based reconstruction and signal processing:
calc_density_compensation_function: Calculates density compensation weights to account for non-uniform sampling density.calc_tensor_spmatrix: Calculates a tensor sparse matrix.calc_toeplitz_kernel: Calculates a Toeplitz kernel, often used in iterative reconstruction algorithms.The primary workhorse modules for applying Non-Uniform Fast Fourier Transform (NUFFT) operations are found in the torchkbnufft package. These modules handle the core transformations between Cartesian and non-Cartesian grids.
Key modules include:
KbNufft: Performs the forward NUFFT operation.KbNufftAdjoint: Performs the adjoint (backward) NUFFT operation.KbInterp: Interpolation-based NUFFT.KbInterpAdjoint: Adjoint interpolation-based NUFFT.ToepNufft: Toeplitz-based NUFFT.The package provides three main classes of NUFFT operation modes. When choosing a mode, consider the trade-off between speed and accuracy. It is generally recommended to start with Table interpolation and then experiment with other modes based on your specific requirements.
| Type | Speed | Notes |
|---|---|---|
| Toeplitz | Fastest | Uses embedded FFTs; excellent for gradient descent. |
| Table | Medium | The standard starting point. |
| Sparse Matrix | Slow | Not recommended for general use; can be more accurate than Table mode. |
The performance of torchkbnufft is primarily limited by advanced indexing and multiplications. While multiplications scale well, advanced indexing is limited by PyTorch.
What scales well:
What does NOT scale well:
If you encounter performance issues in 3D or with long trajectories, try lowering precision, reducing the oversampling ratio, or using fewer interpolation neighbors.
The package provides three major classes of NUFFT (Non-uniform Fast Fourier Transform) interpolation modes. Choosing the right mode depends on your specific problem requirements for speed and memory:
Recommendation: Start with Table interpolation and experiment with other modes if performance needs to be optimized.
torchkbnufft.KbNufft or torchkbnufft.KbNufftAdjoint objects. This allows the operator to account for coil sensitivities during the forward or adjoint transform.To use the basic functionality demonstrated in the examples, install torchkbnufft along with its required dependencies: matplotlib, numpy, torch, and scikit-image.
!pip install torchkbnufft matplotlib numpy torch scikit-imageThe NUFFT objects act as callable modules.
nufft_ob(image, ktraj) to obtain k-space data.adjnufft_ob(kdata, ktraj) to reconstruct the image.Note: The input image tensor should typically have dimensions for batch, coil, and spatial resolution (e.g., [batch, coil, ny, nx]).
# Forward NUFFT: Image -> K-space
kdata = nufft_ob(image, ktraj)
# Adjoint NUFFT: K-space -> Image
image_reconstructed = adjnufft_ob(kdata, ktraj)The number of interpolation neighbors significantly impacts speed. The default is 6 neighbors in each dimension. You can reduce this by setting the npoints parameter during KbNufft initialization. You can provide a single integer for all dimensions or a tuple to specify different neighbor counts for different dimensions (e.g., reducing neighbors only in the z-dimension).
# Use 4 neighbors in all dimensions
forw_ob = tkbn.KbNufft(im_size=im_size, numpoints=4)
# Use specific neighbor counts per dimension (e.g., 4 in Z, 6 in Y, 6 in X)
forw_ob = tkbn.KbNufft(im_size=im_size, numpoints=(4, 6, 6))When performing the adjoint NUFFT, applying a density compensation function (DCF) helps mitigate blurring caused by non-uniform sampling. You can calculate the DCF using tkbn.calc_density_compensation_function and multiply it by your k-space data before passing it to the adjoint operator.
# 1. Calculate density compensation function
dcomp = tkbn.calc_density_compensation_function(ktraj=ktraj, im_size=im_size)
# 2. Apply DCF to k-space data and then perform adjoint NUFFT
image_sharp = adjnufft_ob(kdata * dcomp, ktraj)The primary way to use torchkbnufft is through PyTorch modules. You initialize a module (like KbNufft) with your imaging setup properties (e.g., im_size). The module handles Kaiser-Bessel kernel calculation and interpolation parameters. You then pass your image data and trajectory tensors to the module. NUFFT operations are wrapped in torch.autograd.Function to support backpropagation and training.
import torch
import torchkbnufft as tkbn
import numpy as np
from skimage.data import shepp_logan_phantom
x = shepp_logan_phantom().astype(complex)
im_size = x.shape
# Prepare image tensor: (batch, coil, ny, nx)
x = torch.tensor(x).unsqueeze(0).unsqueeze(0).to(torch.complex64)
klength = 64
ktraj = np.stack((np.zeros(64), np.linspace(-np.pi, np.pi, klength)))
# Prepare trajectory tensor: (dim, klength)
ktraj = torch.tensor(ktraj).to(torch.float)
nufft_ob = tkbn.KbNufft(im_size=im_size)
# Compute k-space data
kdata = nufft_ob(x, ktraj)To use torchkbnufft with the SENSE example functionality, ensure you have torch, matplotlib, numpy, scikit-image, and torchkbnufft installed. If you are using the mrisensesim utility from the repository, you may need to clone the repository or ensure the module is in your path.
!pip install torchkbnufft matplotlib numpy torch scikit-image