A robust, easytodeploy nonuniform Fast Fourier Transform in PyTorch.
Project description
Torch KBNUFFT
API  GitHub  Notebook Examples
Simple installation from PyPI:
pip install torchkbnufft
About
Torch KBNUFFT implements a nonuniform Fast Fourier Transform with KaiserBessel gridding in PyTorch. The implementation is completely in Python, facilitating robustness and flexible deployment in humanreadable code. NUFFT functions are each wrapped as a torch.autograd.Function
, allowing backpropagation through NUFFT operators for training neural networks.
Documentation
Most files are accompanied with docstrings that can be read with help
while running IPython. Example:
from torchkbnufft import KbNufft help(KbNufft)
Behavior can also be inferred by inspecting the source code here. An htmlbased API reference is here.
Examples
Simple Forward NUFFT
Jupyter notebook examples are in the notebooks/
folder. The following minimalist code loads a SheppLogan phantom and computes a single radial spoke of kspace data:
import torch import numpy as np from torchkbnufft import KbNufft from skimage.data import shepp_logan_phantom x = shepp_logan_phantom().astype(np.complex) im_size = x.shape x = np.stack((np.real(x), np.imag(x))) # convert to tensor, unsqueeze batch and coil dimension # output size: (1, 1, 2, ny, nx) x = torch.tensor(x).unsqueeze(0).unsqueeze(0) klength = 64 ktraj = np.stack( (np.zeros(64), np.linspace(np.pi, np.pi, klength)) ) # convert to tensor, unsqueeze batch dimension # output size: (1, 2, klength) ktraj = torch.tensor(ktraj).unsqueeze(0) nufft_ob = KbNufft(im_size=im_size) # outputs a (1, 1, 2, klength) vector of kspace data kdata = nufft_ob(x, ktraj)
A detailed example of basic NUFFT usage is included in notebooks/Basic Example.ipynb
.
SENSENUFFT
The package also includes utilities for working with SENSENUFFT operators. The above code can be modified to replace the nufft_ob
with the following sensenufft_ob
:
from torchkbnufft import MriSenseNufft smap = torch.rand(1, 8, 2, 400, 400) sensenufft_ob = MriSenseNufft(im_size=im_size, smap=smap) sense_data = sensenufft_ob(x, ktraj)
Application of the object in place of nufft_ob
above would first multiply by the sensitivity coils in smap
, then compute a 64length radial spoke for each coil. All operations are broadcast across coils, which minimizes interaction with the Python interpreter and maximizes computation speed.
A detailed example of SENSENUFFT usage is included in notebooks/SENSE Example.ipynb
.
Sparse Matrix Precomputation
Previously, sparse matrixbased interpolation was the fastest operation mode. As of v0.2.0, this is no longer the case on the GPU. Nonetheless, sparse matrices can be faster than normal operation mode in certain situations (e.g., when using many coils and slices). The following code calculates sparse interpolation matrices and uses them to compute a single radial spoke of kspace data:
from torchkbnufft import AdjKbNufft from torchkbnufft.nufft.sparse_interp_mat import precomp_sparse_mats adjnufft_ob = AdjKbNufft(im_size=im_size) # precompute the sparse interpolation matrices real_mat, imag_mat = precomp_sparse_mats(ktraj, adjnufft_ob) interp_mats = { 'real_interp_mats': real_mat, 'imag_interp_mats': imag_mat } # use sparse matrices in adjoint image = adjnufft_ob(kdata, ktraj, interp_mats)
A detailed example of sparse matrix precomputation usage is included in notebooks/Sparse Matrix Example.ipynb
.
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)
Similar to programming lowlevel code, 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
TorchKbNufft is first and foremost designed to be lightweight with minimal dependencies outside of PyTorch. The following computation times in seconds were observed on a workstation with a Xeon E51620 CPU and an Nvidia GTX 1080 GPU for a 15coil, 405spoke 2D radial problem.
Operation  CPU (normal)  CPU (sparse matrix)  GPU (normal)  GPU (sparse matrix) 

Forward NUFFT  3.23  3.38  9.34e02  9.34e02 
Adjoint NUFFT  4.48  1.04  1.03e01  1.15e01 
Profiling for your machine can be done by running
python profile_torchkbnufft.py
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.
Citation
If you want to cite the package, you can use the following:
@misc{Muckley2019, author = {Muckley, M.J. et al.}, title = {Torch KBNUFFT}, year = {2019}, publisher = {GitHub}, journal = {GitHub repository}, howpublished = {\url{https://github.com/mmuckley/torchkbnufft}} }
Project details
Release history Release notifications
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Filename, size  File type  Python version  Upload date  Hashes 

Filename, size torchkbnufft0.2.1py3noneany.whl (31.5 kB)  File type Wheel  Python version py3  Upload date  Hashes View hashes 
Filename, size torchkbnufft0.2.1.tar.gz (21.2 kB)  File type Source  Python version None  Upload date  Hashes View hashes 
Hashes for torchkbnufft0.2.1py3noneany.whl
Algorithm  Hash digest  

SHA256  65b5174d26ee5bbcbb4402a4fe9c42129c34f4fe33c00b06dfe507183945f693 

MD5  a35656f5d0a03794482fceb9fd3da597 

BLAKE2256  8d9b81f6158b31c4038ff27935b40312e1f7294b7a33a66bbcbd4584f92c1a53 