Skip to main content

A high-level, easy-to-deploy non-uniform Fast Fourier Transform in PyTorch.

Project description

torchkbnufft

LICENSE CI Badge Documentation Status Open In Colab

Documentation | GitHub | Notebook Examples

Simple installation from PyPI:

pip install torchkbnufft

About

torchkbnufft implements a non-uniform Fast Fourier Transform [1, 2] with Kaiser-Bessel 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: table-based NUFFT interpolation, sparse matrix-based NUFFT interpolation, and forward/backward operators with Toeplitz-embedded 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 html-based 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 N-D 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 high-dimensional cases) and Toeplitz NUFFTs (which are an extremely fast forward-backward NUFFT technique).

All examples have associated notebooks that you can run in Google Colab:

Simple Forward NUFFT

Basic NUFFT Example in Colab

The following code loads a Shepp-Logan phantom and computes a single radial spoke of k-space data:

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
# 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 k-space data
kdata = nufft_ob(x, ktraj)

SENSE-NUFFT

SENSE-NUFFT Example in Colab

The package also includes utilities for working with SENSE-NUFFT 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 64-length 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 k-space 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

Toeplitz Example in Colab

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 E5-2698 CPU and an Nvidia Quadro GP100 GPU for a 15-coil, 405-spoke 2D radial problem. CPU computations were limited to 8 threads and done with 64-bit floats, whereas GPU computations were done with 32-bit 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.58e-02 7.44e-02 3.03e-03 (f/b)
Adjoint NUFFT 0.75 0.76 N/A 3.56e-02 7.93e-02 N/A

Profiling for your machine can be done by running

pip install -r dev-requirements.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:

  1. TF KB-NUFFT (KB-NUFFT for TensorFlow)
  2. SigPy (for Numpy arrays, Numba (for CPU) and CuPy (for GPU) backends)
  3. FINUFFT (for MATLAB, Python, Julia, C, etc., very efficient)
  4. NFFT (for Julia)
  5. PyNUFFT (for Numpy, also has PyCUDA/PyOpenCL backends)

References

  1. Fessler, J. A., & Sutton, B. P. (2003). Nonuniform fast Fourier transforms using min-max interpolation. IEEE Transactions on Signal Processing, 51(2), 560-574.

  2. 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), 799-808.

  3. Feichtinger, H. G., Gr, K., & Strohmer, T. (1995). Efficient numerical methods in non-uniform sampling theory. Numerische Mathematik, 69(4), 423-440.

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 High-Level, Hardware-Agnostic Non-Uniform 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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

torchkbnufft-1.5.2.tar.gz (8.0 MB view details)

Uploaded Source

Built Distribution

torchkbnufft-1.5.2-py3-none-any.whl (37.4 kB view details)

Uploaded Python 3

File details

Details for the file torchkbnufft-1.5.2.tar.gz.

File metadata

  • Download URL: torchkbnufft-1.5.2.tar.gz
  • Upload date:
  • Size: 8.0 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.12.7

File hashes

Hashes for torchkbnufft-1.5.2.tar.gz
Algorithm Hash digest
SHA256 f10aebc6518960aed6ed3c3f16a8f3bb5c9f89e3491c5a275d13297d43037008
MD5 d784fd875ee02b7cab20e8ae91847efc
BLAKE2b-256 2df9879b1a91035841a082a16762f96e15c14d45e145c964869c03fa1a528b97

See more details on using hashes here.

File details

Details for the file torchkbnufft-1.5.2-py3-none-any.whl.

File metadata

  • Download URL: torchkbnufft-1.5.2-py3-none-any.whl
  • Upload date:
  • Size: 37.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.12.7

File hashes

Hashes for torchkbnufft-1.5.2-py3-none-any.whl
Algorithm Hash digest
SHA256 2d5d08858d631c98db44e19d11d2a5fe2c7ae0da6bcc408f663eddac71f180a7
MD5 05feb2331ea37f59dae9bd4cf5a608c1
BLAKE2b-256 2aa31598e9e32db6bd72664ea1b196630f4d3803ed27f83c9ae37665847320d8

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page