This PR refactors the repository, rewriting almost the entire package to use native complex tensors in PyTorch. The benefit of doing this is roughly a factor-of-2 speedup in tests on my personal laptop. The updated documentation can be viewed at https://torchkbnufft.readthedocs.io/en/refactor/index.html.
Since so much has changed, v1.0.0 is very incompatible with v0.3.4, but I think the new design choices are much better, so this is worth it.
A summary of changes follows:
Native support for PyTorch complex tensors. We get a speed improvement for doing this, so v1.0.0 will require the user to be on PyTorch 1.7 or greater. This also required a change to the API. Users wanting to input real tensors will need the last dimension to be of size 2, rather than the dimension after the channel dimension. This almost completely removed the need for permutations in the package and brings us in line complex operations in other parts of the PyTorch ecosystem. For speed, the backend always converts the tensors to complex numbers where possible.
The backend has been substantially rewritten to a higher code quality, adding type annotations and compiling performant-critical functions with torch.jit.script to get rid of the Python GIL. This ultimately gave more than a four-fold increase in CPU compute time and a 2-fold increase in GPU compute time.
A much improved density compensation function, calc_density_compensation_function, thanks to a contribution of @chaithyagr on the suggestion of @zaccharieramzi.
Simplified utility functions for calc_toeplitz_kernel and calc_tensor_spmatrix.
The documentation has been completely rewritten, upgrading to the Read the Docs template, an improved table of contents, adding mathematical descriptions of core operators, and having a dedicated basic usage section.
Dedicated SENSE-NUFFT operators have been removed. Wrapping these with torch.autograd.Function didn't give us any benefits, so there's no need to have them. Users will now pass their sensitivity maps into the forward function of KbNufft and KbNufftAdjoint directly.
Testing:
The test suite has also been completely redone. Tests are largely done on an integration basis by using the core modules, KbInterp, KbInterpAdjoint, KbNufft, KbNufftAdjoint, and ToepNufft. We also created some small datasets to test accuracy and assert that the new package is computationally equivalent to the old one. Lastly, we now enforce black code formatting and lint with both flake8 and mypy.
This PR refactors the repository, rewriting almost the entire package to use native complex tensors in PyTorch. The benefit of doing this is roughly a factor-of-2 speedup in tests on my personal laptop. The updated documentation can be viewed at https://torchkbnufft.readthedocs.io/en/refactor/index.html.
Since so much has changed, v1.0.0 is very incompatible with v0.3.4, but I think the new design choices are much better, so this is worth it.
A summary of changes follows:
torch.jit.scriptto get rid of the Python GIL. This ultimately gave more than a four-fold increase in CPU compute time and a 2-fold increase in GPU compute time.calc_density_compensation_function, thanks to a contribution of @chaithyagr on the suggestion of @zaccharieramzi.calc_toeplitz_kernelandcalc_tensor_spmatrix.torch.autograd.Functiondidn't give us any benefits, so there's no need to have them. Users will now pass their sensitivity maps into theforwardfunction ofKbNufftandKbNufftAdjointdirectly.Testing:
The test suite has also been completely redone. Tests are largely done on an integration basis by using the core modules,
KbInterp,KbInterpAdjoint,KbNufft,KbNufftAdjoint, andToepNufft. We also created some small datasets to test accuracy and assert that the new package is computationally equivalent to the old one. Lastly, we now enforceblackcode formatting and lint with bothflake8andmypy.