def test_type2(shape=(16, 16, 16), M=4096, tol=1e-3):
kxyz = utils.gen_nu_pts(M)
fk = utils.gen_uniform_data(shape)
kxyz_gpu = gpuarray.to_gpu(kxyz)
fk_gpu = gpuarray.to_gpu(fk)
c_gpu = gpuarray.GPUArray(shape=(M,), dtype=np.complex64)
plan = cufinufft.plan(2, shape, -1, tol)
cufinufft.set_nu_pts(plan, M, kxyz_gpu[0].gpudata, kxyz_gpu[1].gpudata,
kxyz_gpu[2].gpudata)
cufinufft.execute(plan, c_gpu.gpudata, fk_gpu.gpudata)
cufinufft.destroy(plan)
c = c_gpu.get()
ind = M // 2
c_est = c[ind]
c_target = utils.direct_type2(fk, kxyz[:, ind])
type2_rel_err = np.abs(c_target - c_est) / np.abs(c_target)
print('Type 2 relative error:', type2_rel_err)
def test_opts(shape=(8, 8, 8), M=32, tol=1e-3):
dtype = np.float32
complex_dtype = utils._complex_dtype(dtype)
dim = len(shape)
k = utils.gen_nu_pts(M, dim=dim).astype(dtype)
c = utils.gen_nonuniform_data(M).astype(complex_dtype)
k_gpu = gpuarray.to_gpu(k)
c_gpu = gpuarray.to_gpu(c)
fk_gpu = gpuarray.GPUArray(shape, dtype=complex_dtype)
plan = cufinufft(1,
shape,
eps=tol,
dtype=dtype,
gpu_sort=False,
gpu_maxsubprobsize=10)
plan.set_pts(k_gpu[0], k_gpu[1], k_gpu[2])
plan.execute(c_gpu, fk_gpu)
fk = fk_gpu.get()
ind = int(0.1789 * np.prod(shape))
fk_est = fk.ravel()[ind]
fk_target = utils.direct_type1(c, k, shape, ind)
type1_rel_err = np.abs(fk_target - fk_est) / np.abs(fk_target)
assert type1_rel_err < 0.01
def test_type1(shape=(16, 16, 16), M=4096, tol=1e-3):
kxyz = utils.gen_nu_pts(M)
c = utils.gen_nonuniform_data(M)
kxyz_gpu = gpuarray.to_gpu(kxyz)
c_gpu = gpuarray.to_gpu(c)
fk_gpu = gpuarray.GPUArray(shape, dtype=np.complex64)
plan = cufinufft.plan(1, shape, 1, tol)
cufinufft.set_nu_pts(plan, M, kxyz_gpu[0].gpudata, kxyz_gpu[1].gpudata,
kxyz_gpu[2].gpudata)
cufinufft.execute(plan, c_gpu.gpudata, fk_gpu.gpudata)
fk = fk_gpu.get()
ind = int(0.1789 * np.prod(shape))
fk_est = fk.ravel()[ind]
fk_target = utils.direct_type1(c, kxyz, shape, ind)
type1_rel_err = np.abs(fk_target - fk_est) / np.abs(fk_target)
print('Type 1 relative error:', type1_rel_err)
def _test_type2(dtype, shape=(16, 16, 16), M=4096, tol=1e-3):
complex_dtype = utils._complex_dtype(dtype)
k = utils.gen_nu_pts(M).astype(dtype)
fk = utils.gen_uniform_data(shape).astype(complex_dtype)
k_gpu = gpuarray.to_gpu(k)
fk_gpu = gpuarray.to_gpu(fk)
c_gpu = gpuarray.GPUArray(shape=(M, ), dtype=complex_dtype)
plan = cufinufft(2, shape, eps=tol, dtype=dtype)
plan.set_pts(k_gpu[0], k_gpu[1], k_gpu[2])
plan.execute(c_gpu, fk_gpu)
c = c_gpu.get()
ind = M // 2
c_est = c[ind]
c_target = utils.direct_type2(fk, k[:, ind])
type2_rel_err = np.abs(c_target - c_est) / np.abs(c_target)
print('Type 2 relative error:', type2_rel_err)
assert type2_rel_err < 0.01
def _test_type1(dtype, shape=(16, 16, 16), M=4096, tol=1e-3):
complex_dtype = utils._complex_dtype(dtype)
dim = len(shape)
k = utils.gen_nu_pts(M, dim=dim).astype(dtype)
c = utils.gen_nonuniform_data(M).astype(complex_dtype)
k_gpu = gpuarray.to_gpu(k)
c_gpu = gpuarray.to_gpu(c)
fk_gpu = gpuarray.GPUArray(shape, dtype=complex_dtype)
plan = cufinufft(1, shape, eps=tol, dtype=dtype)
plan.set_pts(k_gpu[0], k_gpu[1], k_gpu[2])
plan.execute(c_gpu, fk_gpu)
fk = fk_gpu.get()
ind = int(0.1789 * np.prod(shape))
fk_est = fk.ravel()[ind]
fk_target = utils.direct_type1(c, k, shape, ind)
type1_rel_err = np.abs(fk_target - fk_est) / np.abs(fk_target)
print('Type 1 relative error:', type1_rel_err)
assert type1_rel_err < 0.01
def test_exec_raises_on_dtype():
dtype = np.float32
complex_dtype = np.complex64
M = 4096
tol = 1e-3
shape = (16, 16, 16)
dim = len(shape)
kxyz = utils.gen_nu_pts(M, dim=dim).astype(dtype)
c = utils.gen_nonuniform_data(M).astype(complex_dtype)
c_gpu = gpuarray.to_gpu(c)
# Using c.real gives us wrong dtype here...
c_gpu_wrong_dtype = gpuarray.to_gpu(c.real)
kxyz_gpu = gpuarray.to_gpu(kxyz)
fk_gpu = gpuarray.GPUArray(shape, dtype=complex_dtype)
# Here we'll intentionally contruct an incorrect array dtype.
fk_gpu_wrong_dtype = gpuarray.GPUArray(shape, dtype=np.complex128)
plan = cufinufft(1, shape, eps=tol, dtype=dtype)
plan.set_pts(kxyz_gpu[0], kxyz_gpu[1], kxyz_gpu[2])
with pytest.raises(TypeError):
plan.execute(c_gpu, fk_gpu_wrong_dtype)
with pytest.raises(TypeError):
plan.execute(c_gpu_wrong_dtype, fk_gpu)
def test_set_nu_raises_on_dtype():
dtype = np.float32
M = 4096
tol = 1e-3
shape = (16, 16, 16)
dim = len(shape)
kxyz = utils.gen_nu_pts(M, dim=dim).astype(dtype)
kxyz_gpu = gpuarray.to_gpu(kxyz)
# Here we'll intentionally contruct an incorrect array dtype.
kxyz_gpu_wrong_type = gpuarray.to_gpu(kxyz.astype(np.float64))
plan = cufinufft(1, shape, eps=tol, dtype=dtype)
with pytest.raises(TypeError):
plan.set_pts(kxyz_gpu_wrong_type[0], kxyz_gpu[1], kxyz_gpu[2])
with pytest.raises(TypeError):
plan.set_pts(kxyz_gpu[0], kxyz_gpu_wrong_type[1], kxyz_gpu[2])
with pytest.raises(TypeError):
plan.set_pts(kxyz_gpu[0], kxyz_gpu[1], kxyz_gpu_wrong_type[2])
with pytest.raises(TypeError):
plan.set_pts(kxyz_gpu_wrong_type[0], kxyz_gpu_wrong_type[1],
kxyz_gpu_wrong_type[2])
def test_multi_type1(dtype=np.float32, shape=(16, 16, 16), M=4096, tol=1e-3):
complex_dtype = utils._complex_dtype(dtype)
drv.init()
dev_count = drv.Device.count()
if dev_count == 1:
pytest.skip()
devs = [drv.Device(dev_id) for dev_id in range(dev_count)]
dim = len(shape)
errs = []
for dev_id, dev in enumerate(devs):
ctx = dev.make_context()
k = utils.gen_nu_pts(M, dim=dim).astype(dtype)
c = utils.gen_nonuniform_data(M).astype(complex_dtype)
k_gpu = gpuarray.to_gpu(k)
c_gpu = gpuarray.to_gpu(c)
fk_gpu = gpuarray.GPUArray(shape, dtype=complex_dtype)
plan = cufinufft(1, shape, eps=tol, dtype=dtype, gpu_device_id=dev_id)
plan.set_pts(k_gpu[0], k_gpu[1], k_gpu[2])
plan.execute(c_gpu, fk_gpu)
fk = fk_gpu.get()
ind = int(0.1789 * np.prod(shape))
fk_est = fk.ravel()[ind]
fk_target = utils.direct_type1(c, k, shape, ind)
type1_rel_err = np.abs(fk_target - fk_est) / np.abs(fk_target)
print(f'Type 1 relative error (GPU {dev_id}):', type1_rel_err)
ctx.pop()
errs.append(type1_rel_err)
assert all(err < 0.01 for err in errs)
def test_set_pts_raises_on_size():
dtype = np.float32
M = 8
tol = 1e-3
shape = (16, 16, 16)
dim = len(shape)
kxyz = utils.gen_nu_pts(M, dim=dim).astype(dtype)
kxyz_gpu = gpuarray.to_gpu(kxyz)
plan = cufinufft(1, shape, eps=tol, dtype=dtype)
with pytest.raises(TypeError) as err:
plan.set_pts(kxyz_gpu[0], kxyz_gpu[1][:4])
assert 'kx and ky must be equal' in err.value.args[0]
with pytest.raises(TypeError) as err:
plan.set_pts(kxyz_gpu[0], kxyz_gpu[1], kxyz_gpu[2][:4])
assert 'kx and kz must be equal' in err.value.args[0]