rocket-fft extends Numba by scipy.fft and numpy.fft
Project description
rocket-fft
Rocket-FFT makes Numba aware of numpy.fft
and scipy.fft
. Rocket-FFT takes its name from the PocketFFT Fast Fourier Transformation library that powers it, and Numba's goal of making your scientific Python code blazingly fast - like a rocket 🚀.
Rocket-FFT has been tested against both the SciPy and Numpy test suites, plus some additional typing tests. Therefore, it is considered safe to use, but the author still welcomes bug reports to help improve the project.
Getting started
The easiest way to get Rocket-FFT is to:
pip install rocket-fft
Alternatively, you can build it yourself:
git clone https://github.com/styfenschaer/rocket-fft.git
pip install ./rocket-fft
The latter will require a C++ compiler to be installed on your system.
Rocket-FFT uses setuptools
entry points to register itself as an extension of Numba, so there is no need for additional imports. Once installed successfully, the following will work:
import numba as nb
import numpy as np
@nb.njit
def jit_fft(x):
return np.fft.fft(x)
a = np.array([3, 1, 4, 1, 5, 9, 2, 6])
jit_fft(a)
Supported Numpy functions
The whole numpy.fft
module is supported, which contains all the functions listed below:
-
numpy.fft.fft
-
numpy.fft.ifft
-
numpy.fft.fft2
* -
numpy.fft.ifft2
* -
numpy.fft.fftn
* -
numpy.fft.ifftn
* -
numpy.fft.rfft
-
numpy.fft.irfft
-
numpy.fft.rfft2
-
numpy.fft.irfft2
-
numpy.fft.rfftn
-
numpy.fft.irfftn
-
numpy.fft.hfft
-
numpy.fft.ihfft
-
numpy.fft.fftfreq
-
numpy.fft.rfftfreq
-
numpy.fft.fftshift
-
numpy.fft.ifftshift
*Rocket-FFT follows SciPy's approach of not allowing duplicate axes
Supported SciPy functions
Most of the scipy.fft
module is supported as well, including:
-
scipy.fft.fft
-
scipy.fft.ifft
-
scipy.fft.fft2
-
scipy.fft.ifft2
-
scipy.fft.fftn
-
scipy.fft.ifftn
-
scipy.fft.rfft
-
scipy.fft.irfft
-
scipy.fft.rfft2
-
scipy.fft.irfft2
-
scipy.fft.rfftn
-
scipy.fft.irfftn
-
scipy.fft.hfft
-
scipy.fft.ihfft
-
scipy.fft.hfft2
-
scipy.fft.ihfft2
-
scipy.fft.hfftn
-
scipy.fft.ihfftn
-
scipy.fft.dct
-
scipy.fft.dct2
-
scipy.fft.dctn
-
scipy.fft.idctn
-
scipy.fft.dst
-
scipy.fft.idst
-
scipy.fft.dstn
-
scipy.fft.idstn
-
scipy.fft.fht
-
scipy.fft.ifht
-
scipy.fft.fftshift
-
scipy.fft.ifftshift
-
scipy.fft.fftfreq
-
scipy.fft.ifftfreq
-
scipy.fft.fhtoffset
-
scipy.fft.next_fast_len
Type conversion
By default, Rocket-FFT follows SciPy's approach to type conversion. However, if you would like to use Numpy's approach instead, you can do so by calling the numpy_like
function from the rocket_fft
namespace:
from rocket_fft import numpy_like, scipy_like
numpy_like()
You can switch back to the SciPy approach by calling the scipy_like
function. Keep in mind that these changes must be made before the internal functions of Rocket-FFT are compiled, as the type conversion rule is frozen upon the first compilation.
Compilation time
Rocket-FFT employs multiple strategies to achieve both flexibility in function signatures similar to scipy.fft
and numpy.fft
, while minimizing compilation times.
One of the key approaches is to generate and assemble specialized functions based on the provided type information prior to compilation. This avoids compiling code that will never be executed, and Rocket-FFT can thus limit compilation times to a few hundred milliseconds in most cases. Calls with default arguments are handled specially, resulting in the fastest compilation times.
Limitations on Linux and MacOS
Rocket-FFT uses a C interface to the PocketFFT C++ library. On Linux and MacOS, the C interface functions are accessed via the Python built-in ctypes
module due to an unresolved issue with LLVM. This fallback solution results in all function not being cachable, which means that numba.njit(cache=True)
does not work. On Windows the C interface is accessed differently and therefore all functions can be cached.
Low-level interface
If you don't need the convenience of the flexible signatures provided by SciPy and Numpy, you can use the low-level interface to the PocketFFT library instead. Using the low-level interface can significantly reduce compilation times and has slightly lower overhead. Additionally, it offers some additional functions that are not available through the Scipy and Numpy interfaces. It is an option to consider if you prioritize speed on small transforms and have a good understanding of how to use FFT. You can import the low-level interface functions from the rocket_fft
namespace:
from rocket_fft import c2c, dct, ...
The low-level interface provides the following functions:
from numpy import complex64, complex128, float32, float64, int64, ndarray
# Note that single and double precision can't be mixed
complex_array = ndarray[complex64] | ndarray[complex128]
real_array = ndarray[float32] | ndarray[float64]
def c2c(ain: complex_array, aout: complex_array, axes: ndarray[int64], forward: bool, fct: float64, nthreads: int64) -> None: ...
def r2c(ain: real_array, aout: complex_array, axes: ndarray[int64], forward: bool, fct: float64, nthreads: int64) -> None: ...
def c2r(ain: complex_array, aout: real_array, axes: ndarray[int64], forward: bool, fct: float64, nthreads: int64) -> None: ...
# Like c2c but on real-valued input array and thus using symmetry
def c2c_sym(ain: real_array, aout: complex_array, axes: ndarray[int64], forward: bool, fct: float64, nthreads: int64) -> None: ...
def dst(ain: real_array, aout: real_array, axes: ndarray[int64], type: int64, fct: float64, ortho: bool, nthreads: int64) -> None: ...
def dct(ain: real_array, aout: real_array, axes: ndarray[int64], type: int64, fct: float64, ortho: bool, nthreads: int64) -> None: ...
def separable_hartley(ain: real_array, aout: real_array, axes: ndarray[int64], fct: float64, nthreads: int64) -> None: ...
def genuine_hartley(ain: real_array, aout: real_array, axes: ndarray[int64], fct: float64, nthreads: int64) -> None: ...
def fftpack(ain: real_array, aout: real_array, axes: ndarray[int64], real2hermitian: bool, forward: bool, fct: float64, nthreads: int64) -> None: ...
def good_size(target: int64, real: bool) -> int64: ...
It's important to note that the low-level interface does not provide the same level of safety and convenience as SciPy and Numpy. There is no safety net, and it is up to the user to ensure proper usage. You may even need to refer to the original PocketFFT C++ implementation to understand how to use the functions properly.
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 Distributions
Hashes for rocket_fft-0.0.2-cp310-cp310-win_amd64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 0efc2afdb1cb216758d92fa390982c91e661186e0c2ac8f5bcba9438dc04fc4a |
|
MD5 | 5690d0949ba1cc15f683b0ea3f9c2627 |
|
BLAKE2b-256 | d27a140d9c8ccbe86c3baca090c12f4fdc3576d4a049bcf2a00d53bd74087799 |
Hashes for rocket_fft-0.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 44bdf148b5cedd82c6a8e3baff4349f9240e7e5804a2ea1e55114f0ec2604eaa |
|
MD5 | 83bf985b60785f9325aebf89d44e3db9 |
|
BLAKE2b-256 | adba1b70298b079accf4499d5c90e23eeef92f20d1298cf2ad2c3390de8da1f7 |
Hashes for rocket_fft-0.0.2-cp310-cp310-macosx_10_9_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 3841feaa9e967cccfa2f64885033a11ee451f0958552d9d273eac8854ffb2efb |
|
MD5 | 47294c5a3f3b10c2bd2c5ad3d40555d2 |
|
BLAKE2b-256 | 3d879b25b7229bd7217464c2a057b4643c78f6fcded7969a54cfae5f0b10952f |
Hashes for rocket_fft-0.0.2-cp39-cp39-win_amd64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 7ef9d6c07cee448748b46c276af08de5ee142e9aa32bb1150c7fdd7bb365b5b8 |
|
MD5 | 665669667424e2bc002a6742c81e9b7a |
|
BLAKE2b-256 | 20bfa2d98c1e16d98a7d547a705da30e6773f0a2cb1d59280832e8919bdc17e9 |
Hashes for rocket_fft-0.0.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 2730132c04b2bab423e09cd36dc08712620236e70a75f5e29690c79055a84342 |
|
MD5 | d6a72eddd5b0be6cc291c65f02e8f4db |
|
BLAKE2b-256 | f9ecad12c5d03ce1a00534cb3d190aa0ee628b46dc33dc9f2ab6da9d18352f79 |
Hashes for rocket_fft-0.0.2-cp39-cp39-macosx_10_9_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 2c9363a3be2d7d4e0fe073f168c953f1df4320b63b0905f2edc70375a365bb58 |
|
MD5 | f6ce26e922163b26f566e18dad626892 |
|
BLAKE2b-256 | 04171d7f3762fcb13eb9358caf644549879b767146a0f301421d9851fcc6f7a7 |
Hashes for rocket_fft-0.0.2-cp38-cp38-win_amd64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | fcaa6d3ab476736d7537ce0fb69863838539aea53c5195c65bbb00486ec3526a |
|
MD5 | 7d3d5c85cc04055fc92a04df5eb881e5 |
|
BLAKE2b-256 | a2f9e987375b0d8885990ccc0c9ad5fd5a2e19edf75a82bf21ca94c0a0cac315 |
Hashes for rocket_fft-0.0.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | bdad9366d9538cf842984400dd2770d125c743c605ce637d677b6b727e41c1ef |
|
MD5 | b0dcfa42c074b25db603aab0f649e51f |
|
BLAKE2b-256 | 80a92dc87ea4af1e4a00fa587108acaaa682135a15fd29434ef6c46dae413316 |
Hashes for rocket_fft-0.0.2-cp38-cp38-macosx_10_9_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | a8e41ad81762668d336d4abeea89b6ac1a4e3f4b8dc474620d438f0ec7cbdd66 |
|
MD5 | a191a24460e13272b6b417606bf66eb7 |
|
BLAKE2b-256 | e230492acdac97313826ccfeac2874bbd6f31b7be403377970e83a4f674a91ab |