A highlevel, easytodeploy nonuniform Fast Fourier Transform in PyTorch.
Project description
torchkbnufft
Documentation  GitHub  Notebook Examples
Simple installation from PyPI:
pip install torchkbnufft
About
torchkbnufft
implements a nonuniform Fast Fourier Transform
[1, 2] with KaiserBessel gridding in PyTorch. The
implementation is completely in Python, facilitating flexible deployment in
readable code with no compilation. NUFFT functions are each wrapped as a
torch.autograd.Function
, allowing backpropagation through NUFFT operators
for training neural networks.
This package was inspired in large part by the NUFFT implementation in the Michigan Image Reconstruction Toolbox (Matlab).
Operation Modes and Stages
The package has three major classes of NUFFT operation mode: tablebased NUFFT interpolation, sparse matrixbased NUFFT interpolation, and forward/backward operators with Toeplitzembedded FFTs [3]. Roughly, computation speed follows:
Type  Speed 

Toeplitz  Fastest 
Table  Medium 
Sparse Matrix  Slow (not recommended) 
It is generally best to start with Table interpolation and then experiment with the other modes for your problem.
Sensitivity maps can be incorporated by passing them into a KbNufft
or
KbNufftAdjoint
object.
Documentation
An htmlbased documentation reference on Read the Docs.
Most files are accompanied with docstrings that can be read with help
while running IPython. Example:
from torchkbnufft import KbNufft
help(KbNufft)
Examples
torchkbnufft
can be used for ND NUFFT transformations. The examples here
start with a simple 2D NUFFT, then expand it to SENSE (a task with multiple,
parallel 2D NUFFTs).
The last two examples demonstrate NUFFTs based on sparse matrix multiplications (which can be useful for highdimensional cases) and Toeplitz NUFFTs (which are an extremely fast forwardbackward NUFFT technique).
All examples have associated notebooks that you can run in Google Colab:
 Basic Example in Colab
 SENSENUFFT Example in Colab
 Sparse Matrix Example in Colab
 Toeplitz Example in Colab
Simple Forward NUFFT
The following code loads a SheppLogan phantom and computes a single radial spoke of kspace data:
import torch
import torchkbnufft as tkbn
import numpy as np
from skimage.data import shepp_logan_phantom
x = shepp_logan_phantom().astype(np.complex)
im_size = x.shape
# convert to tensor, unsqueeze batch and coil dimension
# output size: (1, 1, 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))
)
# convert to tensor, unsqueeze batch dimension
# output size: (2, klength)
ktraj = torch.tensor(ktraj, dtype=torch.float)
nufft_ob = tkbn.KbNufft(im_size=im_size)
# outputs a (1, 1, klength) vector of kspace data
kdata = nufft_ob(x, ktraj)
SENSENUFFT
The package also includes utilities for working with SENSENUFFT operators. The above code can be modified to include sensitivity maps.
smaps = torch.rand(1, 8, 400, 400) + 1j * torch.rand(1, 8, 400, 400)
sense_data = nufft_ob(x, ktraj, smaps=smaps.to(x))
This code first multiplies by the sensitivity coils in smaps
, then
computes a 64length radial spoke for each coil. All operations are broadcast
across coils, which minimizes interaction with the Python interpreter, helping
computation speed.
Sparse Matrix Precomputation
Sparse Matrix Example in Colab
Sparse matrices are an alternative to table interpolation. Their speed can vary, but they are a bit more accurate than standard table mode. The following code calculates sparse interpolation matrices and uses them to compute a single radial spoke of kspace data:
adjnufft_ob = tkbn.KbNufftAdjoint(im_size=im_size)
# precompute the sparse interpolation matrices
interp_mats = tkbn.calc_tensor_spmatrix(
ktraj,
im_size=im_size
)
# use sparse matrices in adjoint
image = adjnufft_ob(kdata, ktraj, interp_mats)
Sparse matrix multiplication is only implemented for real numbers in PyTorch, which can limit their speed.
Toeplitz Embedding
The package includes routines for calculating embedded Toeplitz kernels and using them as FFT filters for the forward/backward NUFFT operations [3]. This is very useful for gradient descent algorithms that must use the forward/backward ops in calculating the gradient. The following code shows an example:
toep_ob = tkbn.ToepNufft()
# precompute the embedded Toeplitz FFT kernel
kernel = tkbn.calc_toeplitz_kernel(ktraj, im_size)
# use FFT kernel from embedded Toeplitz matrix
image = toep_ob(image, kernel)
Running on the GPU
All of the examples included in this repository can be run on the GPU by sending the NUFFT object and data to the GPU prior to the function call, e.g.,
adjnufft_ob = adjnufft_ob.to(torch.device('cuda'))
kdata = kdata.to(torch.device('cuda'))
ktraj = ktraj.to(torch.device('cuda'))
image = adjnufft_ob(kdata, ktraj)
PyTorch will throw errors if the underlying dtype
and device
of all
objects are not matching. Be sure to make sure your data and NUFFT objects are
on the right device and in the right format to avoid these errors.
Computation Speed
The following computation times in seconds were observed on a workstation with
a Xeon E52698 CPU and an Nvidia Quadro GP100 GPU for a 15coil, 405spoke 2D
radial problem. CPU computations were limited to 8 threads and done with 64bit
floats, whereas GPU computations were done with 32bit floats. The benchmark
used torchkbnufft
version 1.0.0 and torch
version 1.7.1.
(n) = normal, (spm) = sparse matrix, (toep) = Toeplitz embedding, (f/b) = forward/backward combined
Operation  CPU (n)  CPU (spm)  CPU (toep)  GPU (n)  GPU (spm)  GPU (toep) 

Forward NUFFT  0.82  0.77  0.058 (f/b)  2.58e02  7.44e02  3.03e03 (f/b) 
Adjoint NUFFT  0.75  0.76  N/A  3.56e02  7.93e02  N/A 
Profiling for your machine can be done by running
pip install r devrequirements.txt
python profile_torchkbnufft.py
Other Packages
For users interested in NUFFT implementations for other computing platforms, the following is a partial list of other projects:
 TF KBNUFFT (KBNUFFT for TensorFlow)
 SigPy (for Numpy arrays, Numba (for CPU) and CuPy (for GPU) backends)
 FINUFFT (for MATLAB, Python, Julia, C, etc., very efficient)
 NFFT (for Julia)
 PyNUFFT (for Numpy, also has PyCUDA/PyOpenCL backends)
References

Fessler, J. A., & Sutton, B. P. (2003). Nonuniform fast Fourier transforms using minmax interpolation. IEEE Transactions on Signal Processing, 51(2), 560574.

Beatty, P. J., Nishimura, D. G., & Pauly, J. M. (2005). Rapid gridding reconstruction with a minimal oversampling ratio. IEEE Transactions on Medical Imaging, 24(6), 799808.

Feichtinger, H. G., Gr, K., & Strohmer, T. (1995). Efficient numerical methods in nonuniform sampling theory. Numerische Mathematik, 69(4), 423440.
Citation
If you use the package, please cite:
@conference{muckley:20:tah,
author = {M. J. Muckley and R. Stern and T. Murrell and F. Knoll},
title = {{TorchKbNufft}: A HighLevel, HardwareAgnostic NonUniform Fast {Fourier} Transform},
booktitle = {ISMRM Workshop on Data Sampling \& Image Reconstruction},
year = 2020,
note = {Source code available at https://github.com/mmuckley/torchkbnufft}.
}
Project details
Release history Release notifications  RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for torchkbnufft1.4.0py3noneany.whl
Algorithm  Hash digest  

SHA256  81a52e181df0e84f415bda8650a40d4fdaad6a71aa9ec6b4feb17e3f73c62a94 

MD5  8edbe1a184ed33ced92c4c2c90b0fe21 

BLAKE2b256  3a89de9abd7b4a889ebc59578f227de5656e0b354d140884f49fff00017eb33b 