Package: torchkbnufft
torchkbnufftDocumentation|GitHub|Notebook ExamplesSimple installation from PyPI:pipinstalltorchkbnufftAbouttorchkbnufftimplements 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 atorch.autograd.Function, allowing backpropagation through NUFFT operators
for training neural networks.This package was inspired in large part by the NUFFT implementation in theMichigan Image Reconstruction Toolbox (Matlab).Operation Modes and StagesThe 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:TypeSpeedToeplitzFastestTableMediumSparse MatrixSlow (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 aKbNufftorKbNufftAdjointobject.DocumentationAn html-based documentation reference onRead the Docs.Most files are accompanied with docstrings that can be read withhelpwhile running IPython. Example:fromtorchkbnufftimportKbNuffthelp(KbNufft)Examplestorchkbnufftcan 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:Basic Example in ColabSENSE-NUFFT Example in ColabSparse Matrix Example in ColabToeplitz Example in ColabSimple Forward NUFFTBasic NUFFT Example in ColabThe following code loads a Shepp-Logan phantom and computes a single radial
spoke of k-space data:importtorchimporttorchkbnufftastkbnimportnumpyasnpfromskimage.dataimportshepp_logan_phantomx=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=64ktraj=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 datakdata=nufft_ob(x,ktraj)SENSE-NUFFTSENSE-NUFFT Example in ColabThe 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 insmaps, 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 PrecomputationSparse Matrix Example in ColabSparse 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 matricesinterp_mats=tkbn.calc_tensor_spmatrix(ktraj,im_size=im_size)# use sparse matrices in adjointimage=adjnufft_ob(kdata,ktraj,interp_mats)Sparse matrix multiplication is only implemented for real numbers in PyTorch,
which can limit their speed.Toeplitz EmbeddingToeplitz Example in ColabThe 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 kernelkernel=tkbn.calc_toeplitz_kernel(ktraj,im_size)# use FFT kernel from embedded Toeplitz matriximage=toep_ob(image,kernel)Running on the GPUAll 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 underlyingdtypeanddeviceof 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 SpeedThe 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
usedtorchkbnufftversion 1.0.0 andtorchversion 1.7.1.(n) = normal, (spm) = sparse matrix, (toep) = Toeplitz embedding, (f/b) = forward/backward combinedOperationCPU (n)CPU (spm)CPU (toep)GPU (n)GPU (spm)GPU (toep)Forward NUFFT0.820.770.058 (f/b)2.58e-027.44e-023.03e-03 (f/b)Adjoint NUFFT0.750.76N/A3.56e-027.93e-02N/AProfiling for your machine can be done by runningpipinstall-rdev-requirements.txtpythonprofile_torchkbnufft.pyOther PackagesFor users interested in NUFFT implementations for other computing platforms,
the following is a partial list of other projects:TF KB-NUFFT(KB-NUFFT 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)ReferencesFessler, J. A., & Sutton, B. P. (2003).Nonuniform fast Fourier transforms using min-max interpolation.IEEE Transactions on Signal Processing, 51(2), 560-574.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.Feichtinger, H. G., Gr, K., & Strohmer, T. (1995).Efficient numerical methods in non-uniform sampling theory.Numerische Mathematik, 69(4), 423-440.CitationIf 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}.}