Python interface (GPU)¶
As mentioned in the Python GPU installation instructions, the easiest way to install the Python interface for cuFINUFFT is to run
pip install cufinufft.
Assuming cuFINUFFT has been installed, we will now consider how to calculate a 1D type 1 transform.
To manage the GPU and transfer to and from host and device, we will use the
Consequently, we start with a few import statements.
import numpy as np import pycuda.autoinit from pycuda.gpuarray import GPUArray, to_gpu from cufinufft import cufinufft
We then proceed to setting up a few parameters.
# number of nonuniform points M = 100000 # grid size N = 200000 # generate positions for the nonuniform points and the coefficients x = 2 * np.pi * np.random.uniform(size=M) c = (np.random.standard_normal(size=M) + 1J * np.random.standard_normal(size=M))
Now that the data is prepared, we need to set up a cuFINUFFT plan that can be executed on that data.
# create plan plan = cufinufft(1, (N,), dtype=np.float64) # set the nonuniform points plan.set_pts(to_gpu(x))
The cuFINUFFT interface relies on the user to supply a preallocated output array, so we use
pycuda.GPUArray for this:
# allocate output array f_gpu = GPUArray((N,), dtype=np.complex128)
With everything set up, we are now ready to execute the plan.
# execute the plan plan.execute(to_gpu(c), f_gpu) # move results off the GPU f = f_gpu.get()
The last line is needed here to copy the data from the device (GPU) to the host (CPU).
f is a size-
N array containing the NUDFT of
c at the points
Other possible calculations are possible by supplying different options during plan creation. See the full API documentation below for more information.