def test_2d_on_2d_r2c(comm):
if comm.size == 1:
procmesh = pfft.ProcMesh(np=[1, 1], comm=comm)
else:
procmesh = pfft.ProcMesh(np=[2, 2], comm=comm)
N = (8, 8)
data = numpy.arange(numpy.prod(N), dtype='f8').reshape(N)
correct = numpy.fft.rfftn(data.copy())
result = numpy.zeros_like(correct)
partition = pfft.Partition(
pfft.Type.PFFT_R2C,
N,
procmesh,
flags=pfft.Flags.PFFT_ESTIMATE
| pfft.Flags.PFFT_TRANSPOSED_OUT
| pfft.Flags.PFFT_DESTROY_INPUT
# | pfft.Flags.PADDED_R2C # doesn't work yet
)
buffer1 = pfft.LocalBuffer(partition)
buffer2 = pfft.LocalBuffer(partition)
plan = pfft.Plan(partition, pfft.Direction.PFFT_FORWARD, buffer1, buffer2)
buffer1.view_input()[:] = data[partition.local_i_slice]
plan.execute(buffer1, buffer2)
result[partition.local_o_slice] = buffer2.view_output()
result = comm.allreduce(result)
assert_almost_equal(correct, result)
def test_correct_multi(comm):
procmesh = pfft.ProcMesh(np=[
comm.size,
], comm=comm)
N = (2, 3)
data = numpy.arange(numpy.prod(N), dtype='complex128').reshape(N)
correct = numpy.fft.fftn(data)
result = numpy.zeros_like(data)
partition = pfft.Partition(pfft.Type.PFFT_C2C,
N,
procmesh,
flags=pfft.Flags.PFFT_ESTIMATE)
buffer1 = pfft.LocalBuffer(partition)
buffer2 = pfft.LocalBuffer(partition)
plan = pfft.Plan(partition, pfft.Direction.PFFT_FORWARD, buffer1, buffer2)
buffer1.view_input()[:] = data[partition.local_i_slice]
plan.execute(buffer1, buffer2)
result[partition.local_o_slice] = buffer2.view_output()
result = comm.allreduce(result)
assert_almost_equal(correct, result)
def test_plan_backward(comm):
procmesh = pfft.ProcMesh(np=[1], comm=comm)
partition = pfft.Partition(pfft.Type.PFFT_R2C, [2, 2],
procmesh,
flags=pfft.Flags.PFFT_ESTIMATE
| pfft.Flags.PFFT_TRANSPOSED_OUT)
buffer1 = pfft.LocalBuffer(partition)
buffer2 = pfft.LocalBuffer(partition)
plan = pfft.Plan(partition, pfft.Direction.PFFT_FORWARD, buffer1, buffer2)
assert plan.flags & pfft.Flags.PFFT_TRANSPOSED_OUT
assert plan.type == pfft.Type.PFFT_R2C
plan = pfft.Plan(partition, pfft.Direction.PFFT_BACKWARD, buffer1, buffer2)
assert plan.flags & pfft.Flags.PFFT_TRANSPOSED_IN
assert plan.type == pfft.Type.PFFT_C2R
def test_correct_single(comm):
procmesh = pfft.ProcMesh(np=[1], comm=comm)
partition = pfft.Partition(pfft.Type.PFFT_C2C, [2, 2],
procmesh,
flags=pfft.Flags.PFFT_ESTIMATE)
buffer1 = pfft.LocalBuffer(partition)
buffer2 = pfft.LocalBuffer(partition)
plan = pfft.Plan(partition, pfft.Direction.PFFT_FORWARD, buffer1, buffer2)
buffer1.view_input()[:] = numpy.arange(4).reshape(2, 2)
correct = numpy.fft.fftn(buffer1.view_input())
plan.execute(buffer1, buffer2)
assert_array_equal(correct, buffer2.view_output())
def __init__(self, BoxSize, Nmesh, paintbrush='cic', comm=None, np=None, verbose=False, dtype='f8'):
""" create a PM object. """
# this weird sequence to intialize comm is because
# we want to be compatible with None comm == MPI.COMM_WORLD
# while not relying on pfft's full mpi4py compatibility
# (passing None through to pfft)
if comm is None:
self.comm = MPI.COMM_WORLD
else:
self.comm = comm
if np is None:
np = pfft.split_size_2d(self.comm.size)
dtype = numpy.dtype(dtype)
if dtype == numpy.dtype('f8'):
forward = pfft.Type.PFFT_R2C
backward = pfft.Type.PFFT_C2R
elif dtype == numpy.dtype('f4'):
forward = pfft.Type.PFFTF_R2C
backward = pfft.Type.PFFTF_C2R
else:
raise ValueError("dtype must be f8 or f4")
self.procmesh = pfft.ProcMesh(np, comm=comm)
self.Nmesh = Nmesh
self.BoxSize = numpy.empty(3, dtype='f8')
self.BoxSize[:] = BoxSize
self.partition = pfft.Partition(forward,
[Nmesh, Nmesh, Nmesh],
self.procmesh,
pfft.Flags.PFFT_TRANSPOSED_OUT | pfft.Flags.PFFT_DESTROY_INPUT)
buffer = pfft.LocalBuffer(self.partition)
self.real = buffer.view_input()
self.real[:] = 0
self.complex = buffer.view_output()
self.T = Timers(self.comm)
with self.T['Plan']:
self.forward = pfft.Plan(self.partition, pfft.Direction.PFFT_FORWARD,
self.real.base, self.complex.base, forward,
pfft.Flags.PFFT_ESTIMATE | pfft.Flags.PFFT_TRANSPOSED_OUT | pfft.Flags.PFFT_DESTROY_INPUT)
self.backward = pfft.Plan(self.partition, pfft.Direction.PFFT_BACKWARD,
self.complex.base, self.real.base, backward,
pfft.Flags.PFFT_ESTIMATE | pfft.Flags.PFFT_TRANSPOSED_IN | pfft.Flags.PFFT_DESTROY_INPUT)
self.domain = domain.GridND(self.partition.i_edges, comm=self.comm)
self.verbose = verbose
self.stack = []
k = []
x = []
w = []
r = []
for d in range(self.partition.ndim):
t = numpy.ones(self.partition.ndim, dtype='intp')
s = numpy.ones(self.partition.ndim, dtype='intp')
t[d] = self.partition.local_ni[d]
s[d] = self.partition.local_no[d]
wi = numpy.arange(s[d], dtype='f4') + self.partition.local_o_start[d]
ri = numpy.arange(t[d], dtype='f4') + self.partition.local_i_start[d]
wi[wi >= self.Nmesh // 2] -= self.Nmesh
ri[ri >= self.Nmesh // 2] -= self.Nmesh
wi *= (2 * numpy.pi / self.Nmesh)
ki = wi * self.Nmesh / self.BoxSize[d]
xi = ri * self.BoxSize[d] / self.Nmesh
w.append(wi.reshape(s))
r.append(ri.reshape(t))
k.append(ki.reshape(s))
x.append(xi.reshape(t))
self.w = w
self.r = r
self.k = k
self.x = x
# set the painter
self.paintbrush = paintbrush.lower()
if paintbrush == 'cic':
self.painter = cic.paint
elif paintbrush == 'tsc':
self.painter = tsc.paint
else:
raise ValueError("valid `painter` values are: ['cic', 'tsc']")
partition = pfft.Partition(
pfft.Type.PFFT_C2C, [8, 8], procmesh,
pfft.Flags.PFFT_TRANSPOSED_OUT | pfft.Flags.PFFT_DESTROY_INPUT)
for irank in range(4):
MPI.COMM_WORLD.barrier()
if irank != procmesh.rank:
continue
print 'My rank is', procmesh.this
print 'local_i_start', partition.local_i_start
print 'local_o_start', partition.local_o_start
print 'i_edges', partition.i_edges
print 'o_edges', partition.o_edges
buffer = pfft.LocalBuffer(partition)
plan = pfft.Plan(partition, pfft.Direction.PFFT_FORWARD, buffer)
iplan = pfft.Plan(
partition,
pfft.Direction.PFFT_BACKWARD,
buffer,
flags=pfft.Flags.PFFT_TRANSPOSED_OUT | pfft.Flags.PFFT_DESTROY_INPUT,
)
input = buffer.view_input()
print input.base
# now lets fill the input array in a funny way
# a[i, j] = i * 10 + j
# we will do a tranform roundtrip and print this out
indices = numpy.array(numpy.indices(input.shape))
indices += partition.local_i_start[:, None, None]
i, j = indices
def __init__(self, Nmesh, BoxSize=1.0, comm=None, np=None, dtype='f8', plan_method='estimate', resampler='cic'):
""" create a PM object. """
if comm is None:
comm = MPI.COMM_WORLD
self.comm = comm
if np is None:
if len(Nmesh) >= 3:
np = pfft.split_size_2d(self.comm.size)
elif len(Nmesh) == 2:
np = [self.comm.size]
elif len(Nmesh) == 1:
np = []
dtype = numpy.dtype(dtype)
self.dtype = dtype
if dtype == numpy.dtype('f8'):
forward = pfft.Type.PFFT_R2C
backward = pfft.Type.PFFT_C2R
elif dtype == numpy.dtype('f4'):
forward = pfft.Type.PFFTF_R2C
backward = pfft.Type.PFFTF_C2R
else:
raise ValueError("dtype must be f8 or f4")
self.procmesh = pfft.ProcMesh(np, comm=comm)
self.Nmesh = numpy.array(Nmesh, dtype='i8')
self.ndim = len(self.Nmesh)
self.BoxSize = numpy.empty(len(Nmesh), dtype='f8')
self.BoxSize[:] = BoxSize
self.partition = pfft.Partition(forward,
self.Nmesh,
self.procmesh,
pfft.Flags.PFFT_TRANSPOSED_OUT | pfft.Flags.PFFT_PADDED_R2C)
bufferin = pfft.LocalBuffer(self.partition)
bufferout = pfft.LocalBuffer(self.partition)
plan_method = {
"estimate": pfft.Flags.PFFT_ESTIMATE,
"measure": pfft.Flags.PFFT_MEASURE,
"exhaustive": pfft.Flags.PFFT_EXHAUSTIVE,
} [plan_method]
self.forward = pfft.Plan(self.partition, pfft.Direction.PFFT_FORWARD,
bufferin, bufferout, forward,
plan_method | pfft.Flags.PFFT_TRANSPOSED_OUT | pfft.Flags.PFFT_TUNE | pfft.Flags.PFFT_PADDED_R2C)
self.backward = pfft.Plan(self.partition, pfft.Direction.PFFT_BACKWARD,
bufferout, bufferin, backward,
plan_method | pfft.Flags.PFFT_TRANSPOSED_IN | pfft.Flags.PFFT_TUNE | pfft.Flags.PFFT_PADDED_C2R)
self.ipforward = pfft.Plan(self.partition, pfft.Direction.PFFT_FORWARD,
bufferin, bufferin, forward,
plan_method | pfft.Flags.PFFT_TRANSPOSED_OUT | pfft.Flags.PFFT_TUNE | pfft.Flags.PFFT_PADDED_R2C)
self.ipbackward = pfft.Plan(self.partition, pfft.Direction.PFFT_BACKWARD,
bufferout, bufferout, backward,
plan_method | pfft.Flags.PFFT_TRANSPOSED_IN | pfft.Flags.PFFT_TUNE | pfft.Flags.PFFT_PADDED_C2R)
self.domain = domain.GridND(self.partition.i_edges, comm=self.comm)
k = []
x = []
w = []
r = []
o_ind = []
i_ind = []
for d in range(self.partition.ndim):
t = numpy.ones(self.partition.ndim, dtype='intp')
s = numpy.ones(self.partition.ndim, dtype='intp')
t[d] = self.partition.local_i_shape[d]
s[d] = self.partition.local_o_shape[d]
i_indi = numpy.arange(t[d], dtype='intp') + self.partition.local_i_start[d]
o_indi = numpy.arange(s[d], dtype='intp') + self.partition.local_o_start[d]
wi = numpy.arange(s[d], dtype='f4') + self.partition.local_o_start[d]
ri = numpy.arange(t[d], dtype='f4') + self.partition.local_i_start[d]
wi[wi >= self.Nmesh[d] // 2] -= self.Nmesh[d]
ri[ri >= self.Nmesh[d] // 2] -= self.Nmesh[d]
wi *= (2 * numpy.pi / self.Nmesh[d])
ki = wi * self.Nmesh[d] / self.BoxSize[d]
xi = ri * self.BoxSize[d] / self.Nmesh[d]
o_ind.append(o_indi.reshape(s))
i_ind.append(i_indi.reshape(t))
w.append(wi.reshape(s))
r.append(ri.reshape(t))
k.append(ki.reshape(s))
x.append(xi.reshape(t))
self.i_ind = i_ind
self.o_ind = o_ind
self.w = w
self.r = r
self.k = k
self.x = x
# Transform from simulation unit to local grid unit.
self.affine = Affine(self.partition.ndim,
translate=-self.partition.local_i_start,
scale=1.0 * self.Nmesh / self.BoxSize,
period = self.Nmesh)
# Transform from global grid unit to local grid unit.
self.affine_grid = Affine(self.partition.ndim,
translate=-self.partition.local_i_start,
scale=1.0,
period = self.Nmesh)
self.resampler = FindResampler(resampler)
def main(comm):
Nmesh = [8, 8]
if len(Nmesh) == 3:
procmesh = pfft.ProcMesh(pfft.split_size_2d(comm.size), comm=comm)
else:
procmesh = pfft.ProcMesh((comm.size, ), comm=comm)
partition = pfft.Partition(
pfft.Type.R2C, Nmesh, procmesh, pfft.Flags.PADDED_R2C
| pfft.Flags.PFFT_TRANSPOSED_OUT | pfft.Flags.DESTROY_INPUT)
# generate the coordinate support.
k = [None] * partition.ndim
x = [None] * partition.ndim
for d in range(partition.ndim):
k[d] = numpy.arange(partition.no[d])[partition.local_o_slice[d]]
k[d][k[d] >= partition.n[d] // 2] -= partition.n[d]
# set to the right numpy broadcast shape
k[d] = k[d].reshape(
[-1 if i == d else 1 for i in range(partition.ndim)])
x[d] = numpy.arange(partition.ni[d])[partition.local_i_slice[d]]
# set to the right numpy broadcast shape
x[d] = x[d].reshape(
[-1 if i == d else 1 for i in range(partition.ndim)])
# allocate memory
buffer1 = pfft.LocalBuffer(partition)
phi_disp = buffer1.view_input()
buffer2 = pfft.LocalBuffer(partition)
phi_spec = buffer2.view_output()
# forward plan
disp_to_spec_inplace = pfft.Plan(
partition,
pfft.Direction.PFFT_FORWARD,
buffer2,
buffer2,
# the two lines below not needed after version 0.1.21
# type=pfft.Type.R2C,
# flags=pfft.Flags.TRANSPOSED_OUT | pfft.Flags.DESTROY_INPUT | pfft.Flags.PADDED_R2C
)
buffer3 = pfft.LocalBuffer(partition)
grad_spec = buffer3.view_output()
buffer4 = pfft.LocalBuffer(partition)
grad_disp = buffer4.view_input()
# backward plan
spec_to_disp = pfft.Plan(
partition,
pfft.Direction.PFFT_BACKWARD,
buffer3,
buffer4,
# the two lines below not needed after version 0.1.21
# type=pfft.Type.C2R,
# flags=pfft.Flags.TRANSPOSED_IN | pfft.Flags.DESTROY_INPUT | pfft.Flags.PADDED_C2R
)
# to do : fill in initial value
dx = x[0] - Nmesh[0] * 0.5 + 0.5
dy = x[1] - Nmesh[1] * 0.5 + 0.5
phi_disp[...] = dx**2 + dx * dy + dy**2
cprint('phi =', gather(partition, phi_disp).round(2), comm=comm)
# copy in to the buffer for inplace transform
# this preserves value of phi_disp
phi_spec.base.view_input()[...] = phi_disp
disp_to_spec_inplace.execute(phi_spec.base, phi_spec.base)
all_grad_disp = numpy.zeros([partition.ndim] + list(phi_disp.shape),
dtype=grad_disp.dtype)
# cprint('phi_k =', gather(partition, phi_spec, mode='output').round(2), comm=comm)
for d in range(partition.ndim):
grad_spec[...] = phi_spec[...] * (k[d] * 1j)
spec_to_disp.execute(grad_spec.base, grad_disp.base)
# copy the gradient along d th direction
all_grad_disp[d] = grad_disp
# now do your thing.
for d in range(partition.ndim):
cprint('dim =',
gather(partition, all_grad_disp[d]).round(2),
comm=comm)