def test_pencil():
from itertools import product
comm = MPI.COMM_WORLD
dims = (2, 3)
sizes = (7, 8, 9)
types = 'fdFD' #'hilfdgFDG'
for typecode in types:
for dim in dims:
for shape in product(*([sizes] * dim)):
axes = list(range(dim))
for axis1, axis2, axis3 in product(axes, axes, axes):
if axis1 == axis2: continue
if axis2 == axis3: continue
axis3 -= len(shape)
#if comm.rank == 0:
# print(shape, axis1, axis2, axis3, typecode)
for pdim in [None] + list(range(1, dim - 1)):
subcomm = Subcomm(comm, pdim)
pencil0 = Pencil(subcomm, shape)
pencilA = pencil0.pencil(axis1)
pencilB = pencilA.pencil(axis2)
pencilC = pencilB.pencil(axis3)
trans1 = Pencil.transfer(pencilA, pencilB, typecode)
trans2 = Pencil.transfer(pencilB, pencilC, typecode)
X = np.random.random(pencilA.subshape).astype(typecode)
A = np.empty(pencilA.subshape, dtype=typecode)
B = np.empty(pencilB.subshape, dtype=typecode)
C = np.empty(pencilC.subshape, dtype=typecode)
A[...] = X
B.fill(0)
trans1.forward(A, B)
C.fill(0)
trans2.forward(B, C)
B.fill(0)
trans2.backward(C, B)
A.fill(0)
trans1.backward(B, A)
assert np.allclose(A, X)
trans1.destroy()
trans2.destroy()
subcomm.destroy()
def test_3Darray():
N = (8, 8, 8)
for subcomm in ((0, 0, 1), (0, 1, 0), (1, 0, 0), (0, 1, 1), (1, 0, 1),
(1, 1, 0), None, Subcomm(comm, (0, 0, 1))):
for rank in (0, 1, 2):
M = (3, ) * rank + N
alignment = None
if subcomm is None and rank == 1:
alignment = 2
a = DistArray(M,
subcomm=subcomm,
val=1,
rank=rank,
alignment=alignment)
assert a.rank == rank
assert a.global_shape == M
_ = a.substart
_ = a.subcomm
z = a.commsizes
_ = a.pencil
assert np.prod(np.array(z)) == comm.Get_size()
if rank > 0:
a0 = a[0]
assert isinstance(a0, DistArray)
assert a0.rank == rank - 1
if rank == 2:
a0 = a[0, 1]
assert isinstance(a0, DistArray)
assert a0.rank == 0
aa = a.v
assert isinstance(aa, np.ndarray)
try:
k = a.get((0, ) * rank + (0, 0, slice(None)))
if comm.Get_rank() == 0:
assert len(k) == N[2]
assert np.sum(k) == N[2]
k = a.get((0, ) * rank + (slice(None), 0, 0))
if comm.Get_rank() == 0:
assert len(k) == N[0]
assert np.sum(k) == N[0]
except ModuleNotFoundError:
pass
_ = a.local_slice()
newaxis = (a.alignment + 1) % 3
_ = a.get_pencil_and_transfer(newaxis)
a[:] = MPI.COMM_WORLD.Get_rank()
b = a.redistribute(newaxis)
a = b.redistribute(out=a)
s0 = MPI.COMM_WORLD.reduce(np.linalg.norm(a)**2)
s1 = MPI.COMM_WORLD.reduce(np.linalg.norm(b)**2)
if MPI.COMM_WORLD.Get_rank() == 0:
assert abs(s0 - s1) < 1e-1
def __init__(self, decomp, queue, grid_shape, dtype, **kwargs):
self.decomp = decomp
self.grid_shape = grid_shape
self.proc_shape = decomp.proc_shape
self.dtype = np.dtype(dtype)
self.is_real = self.dtype.kind == "f"
from pystella.fourier import get_complex_dtype_with_matching_prec
self.cdtype = get_complex_dtype_with_matching_prec(self.dtype)
from pystella.fourier import get_real_dtype_with_matching_prec
self.rdtype = get_real_dtype_with_matching_prec(self.dtype)
if self.proc_shape[0] > 1 and self.proc_shape[1] == 1:
slab = True
else:
slab = False
from mpi4py_fft.pencil import Subcomm
default_kwargs = dict(
axes=([0], [1], [2]),
threads=16,
backend="fftw",
collapse=True,
)
default_kwargs.update(kwargs)
comm = decomp.comm if slab else Subcomm(decomp.comm, self.proc_shape)
from mpi4py_fft import PFFT
self.fft = PFFT(comm,
grid_shape,
dtype=dtype,
slab=slab,
**default_kwargs)
self.fx = self.fft.forward.input_array
self.fk = self.fft.forward.output_array
slc = self.fft.local_slice(True)
self.sub_k = get_sliced_momenta(grid_shape, self.dtype, slc, queue)
def __init__(self,
comm,
bases,
axes=None,
dtype=None,
slab=False,
collapse_fourier=False,
**kw):
# Note do not call __init__ of super
self.comm = comm
self.bases = bases
shape = list(self.global_shape())
assert shape
assert min(shape) > 0
if axes is not None:
axes = list(axes) if np.ndim(axes) else [axes]
for i, axis in enumerate(axes):
if axis < 0:
axes[i] = axis + len(shape)
else:
axes = list(range(len(shape)))
assert min(axes) >= 0
assert max(axes) < len(shape)
assert 0 < len(axes) <= len(shape)
assert sorted(axes) == sorted(set(axes))
if dtype is None:
dtype = np.complex if isinstance(bases[axes[-1]],
C2CBasis) else np.float
else:
if isinstance(bases[axes[-1]], C2CBasis):
assert np.dtype(dtype).char in 'FDG'
elif isinstance(bases[axes[-1]], R2CBasis):
assert np.dtype(dtype).char in 'fdg'
dtype = np.dtype(dtype)
assert dtype.char in 'fdgFDG'
if isinstance(comm, Subcomm):
assert slab is False
assert len(comm) == len(shape)
assert comm[axes[-1]].Get_size() == 1
self.subcomm = comm
else:
if slab:
dims = [1] * len(shape)
dims[axes[0]] = comm.Get_size()
else:
dims = [0] * len(shape)
dims[axes[-1]] = 1
self.subcomm = Subcomm(comm, dims)
self.axes = tuple((axis, ) for axis in axes)
self.xfftn = []
self.transfer = []
self.pencil = [None, None]
# At this points axes is a tuple of tuples of length one.
# Try to collapse some Fourier transforms into one.
if np.any([abs(base.padding_factor - 1.0) > 1e-6 for base in bases]):
collapse_fourier = False
if collapse_fourier:
F = lambda ax: bases[ax].family() == 'fourier' and self.subcomm[
ax].Get_size() == 1
axis = self.axes[-1][-1]
groups = [list(self.axes[-1])]
F0 = F(axis)
for ax in reversed(self.axes[:-1]):
axis = ax[-1]
if F0 and F(axis):
groups[0].insert(0, axis)
else:
groups.insert(0, list(ax))
F0 = F(axis)
self.axes = groups
self.axes = tuple(map(tuple, self.axes))
for base in self.bases:
base.tensorproductspace = self
# Configure all transforms
axes = self.axes[-1]
pencil = Pencil(self.subcomm, shape, axes[-1])
self.xfftn.append(self.bases[axes[-1]])
self.xfftn[-1].plan(pencil.subshape, axes, dtype, kw)
self.pencil[0] = pencilA = pencil
if not shape[axes[-1]] == self.xfftn[-1].forward.output_array.shape[
axes[-1]]:
dtype = self.xfftn[-1].forward.output_array.dtype
shape[axes[-1]] = self.xfftn[-1].forward.output_array.shape[
axes[-1]]
pencilA = Pencil(self.subcomm, shape, axes[-1])
for i, axes in enumerate(reversed(self.axes[:-1])):
pencilB = pencilA.pencil(axes[-1])
transAB = pencilA.transfer(pencilB, dtype)
xfftn = self.bases[axes[-1]]
xfftn.plan(pencilB.subshape, axes, dtype, kw)
self.xfftn.append(xfftn)
self.transfer.append(transAB)
pencilA = pencilB
if not shape[axes[-1]] == xfftn.forward.output_array.shape[
axes[-1]]:
dtype = xfftn.forward.output_array.dtype
shape[axes[-1]] = xfftn.forward.output_array.shape[axes[-1]]
pencilA = Pencil(pencilB.subcomm, shape, axes[-1])
self.pencil[1] = pencilA
self.forward = Transform([o.forward for o in self.xfftn],
[o.forward for o in self.transfer],
self.pencil)
self.backward = Transform([o.backward for o in self.xfftn[::-1]],
[o.backward for o in self.transfer[::-1]],
self.pencil[::-1])
self.scalar_product = Transform([o.scalar_product for o in self.xfftn],
[o.forward for o in self.transfer],
self.pencil)
for i, base in enumerate(bases):
base.axis = i
if base.boundary_condition(
) == 'Dirichlet' and not base.family() in ('laguerre', 'hermite'):
base.bc.set_tensor_bcs(base, self)
if family == 'jacobi':
a = 0
b = 0
x, y, z = symbols("x,y,z", real=True)
ue = (cos(4 * x) + sin(2 * y) +
sin(4 * z)) * (1 - y**2) + a * (1 - y) / 2. + b * (1 + y) / 2.
# Size of discretization
N = int(sys.argv[-2])
N = [N, N + 1, N + 2]
#N = (14, 15, 16)
SD = FunctionSpace(N[1], family=family, bc=(a, b))
K1 = FunctionSpace(N[0], family='F', dtype='D')
K2 = FunctionSpace(N[2], family='F', dtype='d')
subcomms = Subcomm(MPI.COMM_WORLD, [0, 0, 1])
T = TensorProductSpace(subcomms, (K1, SD, K2), axes=(1, 0, 2))
u = TrialFunction(T)
v = TestFunction(T)
# Get manufactured right hand side
fe = div(grad(u)).tosympy(basis=ue)
K = T.local_wavenumbers()
# Get f on quad points
fj = Array(T, buffer=fe)
# Compute right hand side of Poisson equation
f_hat = inner(v, fj)
def __init__(self, comm, bases, axes=None, dtype=None, slab=False, **kw):
self.comm = comm
self.bases = bases
shape = self.shape()
assert shape
assert min(shape) > 0
if axes is not None:
axes = list(axes) if np.ndim(axes) else [axes]
for i, axis in enumerate(axes):
if axis < 0:
axes[i] = axis + len(shape)
else:
axes = list(range(len(shape)))
assert min(axes) >= 0
assert max(axes) < len(shape)
assert 0 < len(axes) <= len(shape)
assert sorted(axes) == sorted(set(axes))
if dtype is None:
dtype = np.complex if isinstance(bases[axes[-1]],
C2CBasis) else np.float
else:
if isinstance(bases[axes[-1]], C2CBasis):
assert np.dtype(dtype).char in 'FDG'
elif isinstance(bases[axes[-1]], R2CBasis):
assert np.dtype(dtype).char in 'fdg'
dtype = self.dtype = np.dtype(dtype)
assert dtype.char in 'fdgFDG'
if isinstance(comm, Subcomm):
assert slab is False
assert len(comm) == len(shape)
assert comm[axes[-1]].Get_size() == 1
self.subcomm = comm
else:
if slab:
dims = [1] * len(shape)
dims[axes[0]] = comm.Get_size()
else:
dims = [0] * len(shape)
dims[axes[-1]] = 1
self.subcomm = Subcomm(comm, dims)
collapse = False # kw.pop('collapse', True)
if collapse:
groups = [[]]
for axis in reversed(axes):
if self.subcomm[axis].Get_size() == 1:
groups[0].insert(0, axis)
else:
groups.insert(0, [axis])
self.axes = tuple(map(tuple, groups))
else:
self.axes = tuple((axis, ) for axis in axes)
self.xfftn = []
self.transfer = []
self.pencil = [None, None]
axes = self.axes[-1]
pencil = Pencil(self.subcomm, shape, axes[-1])
self.xfftn.append(self.bases[axes[-1]])
self.xfftn[-1].plan(pencil.subshape, axes, dtype, kw)
self.pencil[0] = pencilA = pencil
if not shape[axes[-1]] == self.xfftn[-1].forward.output_array.shape[
axes[-1]]:
dtype = self.xfftn[-1].forward.output_array.dtype
shape[axes[-1]] = self.xfftn[-1].forward.output_array.shape[
axes[-1]]
pencilA = Pencil(self.subcomm, shape, axes[-1])
for i, axes in enumerate(reversed(self.axes[:-1])):
pencilB = pencilA.pencil(axes[-1])
transAB = pencilA.transfer(pencilB, dtype)
xfftn = self.bases[axes[-1]]
xfftn.plan(pencilB.subshape, axes, dtype, kw)
self.xfftn.append(xfftn)
self.transfer.append(transAB)
pencilA = pencilB
if not shape[axes[-1]] == xfftn.forward.output_array.shape[
axes[-1]]:
dtype = xfftn.forward.output_array.dtype
shape[axes[-1]] = xfftn.forward.output_array.shape[axes[-1]]
pencilA = Pencil(pencilB.subcomm, shape, axes[-1])
self.pencil[1] = pencilA
self.forward = Transform([o.forward for o in self.xfftn],
[o.forward for o in self.transfer],
self.pencil)
self.backward = Transform([o.backward for o in self.xfftn[::-1]],
[o.backward for o in self.transfer[::-1]],
self.pencil[::-1])
self.scalar_product = Transform([o.scalar_product for o in self.xfftn],
[o.forward for o in self.transfer],
self.pencil)
for base in self.bases:
if isinstance(base, (legendre.bases.ShenDirichletBasis,
chebyshev.bases.ShenDirichletBasis)):
base.bc.set_tensor_bcs(self)
class TensorProductSpace(object):
"""Class for multidimensional tensorproductspaces.
The tensorproductspaces are created as Cartesian products from a set of 1D
bases. The 1D bases are subclassed instances of the :class:`.SpectralBase`
class.
Parameters
----------
comm : MPI communicator
bases : list
List of 1D bases
axes : tuple of ints, optional
A tuple containing the order of which to perform transforms.
Last item is transformed first. Defaults to range(len(bases))
dtype : data-type, optional
Type of input data in real physical space. If not provided it
will be inferred from the bases.
slab : bool, optional
Use 1D slab decomposition instead of default pencil.
kw : dict, optional
Dictionary that can be used to plan transforms. Input to method
`plan` for the bases.
"""
def __init__(self, comm, bases, axes=None, dtype=None, slab=False, **kw):
self.comm = comm
self.bases = bases
shape = self.shape()
assert shape
assert min(shape) > 0
if axes is not None:
axes = list(axes) if np.ndim(axes) else [axes]
for i, axis in enumerate(axes):
if axis < 0:
axes[i] = axis + len(shape)
else:
axes = list(range(len(shape)))
assert min(axes) >= 0
assert max(axes) < len(shape)
assert 0 < len(axes) <= len(shape)
assert sorted(axes) == sorted(set(axes))
if dtype is None:
dtype = np.complex if isinstance(bases[axes[-1]],
C2CBasis) else np.float
else:
if isinstance(bases[axes[-1]], C2CBasis):
assert np.dtype(dtype).char in 'FDG'
elif isinstance(bases[axes[-1]], R2CBasis):
assert np.dtype(dtype).char in 'fdg'
dtype = self.dtype = np.dtype(dtype)
assert dtype.char in 'fdgFDG'
if isinstance(comm, Subcomm):
assert slab is False
assert len(comm) == len(shape)
assert comm[axes[-1]].Get_size() == 1
self.subcomm = comm
else:
if slab:
dims = [1] * len(shape)
dims[axes[0]] = comm.Get_size()
else:
dims = [0] * len(shape)
dims[axes[-1]] = 1
self.subcomm = Subcomm(comm, dims)
collapse = False # kw.pop('collapse', True)
if collapse:
groups = [[]]
for axis in reversed(axes):
if self.subcomm[axis].Get_size() == 1:
groups[0].insert(0, axis)
else:
groups.insert(0, [axis])
self.axes = tuple(map(tuple, groups))
else:
self.axes = tuple((axis, ) for axis in axes)
self.xfftn = []
self.transfer = []
self.pencil = [None, None]
axes = self.axes[-1]
pencil = Pencil(self.subcomm, shape, axes[-1])
self.xfftn.append(self.bases[axes[-1]])
self.xfftn[-1].plan(pencil.subshape, axes, dtype, kw)
self.pencil[0] = pencilA = pencil
if not shape[axes[-1]] == self.xfftn[-1].forward.output_array.shape[
axes[-1]]:
dtype = self.xfftn[-1].forward.output_array.dtype
shape[axes[-1]] = self.xfftn[-1].forward.output_array.shape[
axes[-1]]
pencilA = Pencil(self.subcomm, shape, axes[-1])
for i, axes in enumerate(reversed(self.axes[:-1])):
pencilB = pencilA.pencil(axes[-1])
transAB = pencilA.transfer(pencilB, dtype)
xfftn = self.bases[axes[-1]]
xfftn.plan(pencilB.subshape, axes, dtype, kw)
self.xfftn.append(xfftn)
self.transfer.append(transAB)
pencilA = pencilB
if not shape[axes[-1]] == xfftn.forward.output_array.shape[
axes[-1]]:
dtype = xfftn.forward.output_array.dtype
shape[axes[-1]] = xfftn.forward.output_array.shape[axes[-1]]
pencilA = Pencil(pencilB.subcomm, shape, axes[-1])
self.pencil[1] = pencilA
self.forward = Transform([o.forward for o in self.xfftn],
[o.forward for o in self.transfer],
self.pencil)
self.backward = Transform([o.backward for o in self.xfftn[::-1]],
[o.backward for o in self.transfer[::-1]],
self.pencil[::-1])
self.scalar_product = Transform([o.scalar_product for o in self.xfftn],
[o.forward for o in self.transfer],
self.pencil)
for base in self.bases:
if isinstance(base, (legendre.bases.ShenDirichletBasis,
chebyshev.bases.ShenDirichletBasis)):
base.bc.set_tensor_bcs(self)
def convolve(self, a_hat, b_hat, ab_hat):
"""Convolution of a_hat and b_hat
Parameters
----------
a_hat : array
Input array of shape and type as output array from
self.forward
b_hat : array
Input array of shape and type as output array from
self.forward
ab_hat : array
Return array of same type and shape as a_hat and b_hat
Note
----
The return array ab_hat is truncated to the shape of a_hat and b_hat.
Also note that self should have bases with padding for this method to give
a convolution without aliasing. The padding is specified when creating
instances of bases for the TensorProductSpace.
FIXME Efficiency due to allocation
"""
a = self.backward.output_array.copy()
b = self.backward.output_array.copy()
a = self.backward(a_hat, a)
b = self.backward(b_hat, b)
ab_hat = self.forward(a * b, ab_hat)
return ab_hat
def eval(self, points, coefficients, output_array=None, cython=True):
"""Evaluate Function at points, given expansion coefficients
Parameters
----------
points : float or array of floats
coefficients : array
Expansion coefficients
output_array : array, optional
Return array, function values at points
cython : bool, optional
Whether to use optimized cython implementation or not
"""
if cython:
return self._eval_cython(points, coefficients, output_array)
else:
return self._eval_python(points, coefficients, output_array)
def _eval_python(self,
points,
coefficients,
output_array=None): # pragma : no cover
"""Evaluate Function at points, given expansion coefficients
Parameters
----------
points : float or array of floats
coefficients : array
Expansion coefficients
output_array : array, optional
Return array, function values at points
"""
shape = list(self.local_shape())
out = coefficients
for base in reversed(self):
axis = base.axis
shape[axis] = len(points)
out2 = np.zeros(shape, dtype=out.dtype)
x = base.map_reference_domain(points[..., axis])
out2 = self.vandermonde_evaluate_local_expansion(
base, x, out, out2)
out = out2
# Get the 'diagonals' of out, that is, for 2 points get out[i, i] for i = (0, 1), and same for larger numer of points
out = np.array(
[out[tuple([s] * len(shape))] for s in range(len(points))],
dtype=self.forward.input_array.dtype)
out = np.atleast_1d(np.squeeze(out))
out = self.comm.allreduce(out)
if not output_array is None:
output_array[:] = out
return output_array
return out
def _eval_cython(self, points, coefficients, output_array=None):
"""Evaluate Function at points, given expansion coefficients
Parameters
----------
points : float or array of floats
coefficients : array
Expansion coefficients
output_array : array, optional
Return array, function values at points
"""
out = coefficients
P = []
r2c = -1
last_conj_index = -1
sl = -1
for base in self:
axis = base.axis
x = base.map_reference_domain(points[..., axis])
V = base.vandermonde(x)
D = base.get_vandermonde_basis(V)
P.append(D[..., self.local_slice()[axis]])
if isinstance(base, R2CBasis):
r2c = axis
M = base.N // 2 + 1
if base.N % 2 == 0:
last_conj_index = M - 1
else:
last_conj_index = M
sl = self.local_slice()[axis].start
out = np.zeros(len(points), dtype=self.forward.input_array.dtype)
if len(self) == 2:
out = shenfun.optimization.evaluate.evaluate_2D(out,
coefficients,
P,
r2c=r2c,
M=last_conj_index,
start=sl)
elif len(self) == 3:
out = shenfun.optimization.evaluate.evaluate_3D(out,
coefficients,
P,
r2c=r2c,
M=last_conj_index,
start=sl)
out = np.atleast_1d(out)
out = self.comm.allreduce(out)
if not output_array is None:
output_array[:] = out
return output_array
return out
def vandermonde_evaluate_local_expansion(self, base, points, input_array,
output_array):
"""Evaluate expansion at certain points, possibly different from
the quadrature points
Parameters
----------
base : SpectralBase
The base class
points : float or array of floats
coefficients : array
Expansion coefficients
output_array : array
Return array, function values at points
"""
V = base.vandermonde(points)
P = base.get_vandermonde_basis(V)
P = P[..., self.local_slice()[base.axis]]
last_conj_index = -1
sl = -1
if isinstance(base, R2CBasis):
M = base.N // 2 + 1
if base.N % 2 == 0:
last_conj_index = M - 1
else:
last_conj_index = M
sl = self.local_slice()[base.axis].start
base.vandermonde_evaluate_local_expansion(P, input_array,
output_array,
last_conj_index, sl)
else:
base.vandermonde_evaluate_local_expansion(P, input_array,
output_array)
return output_array
def destroy(self):
"""Destructor"""
self.subcomm.destroy()
for trans in self.transfer:
trans.destroy()
def wavenumbers(self, scaled=False, eliminate_highest_freq=False):
"""Return list of wavenumbers of TensorProductSpace
Parameters
----------
scaled : bool, optional
Scale wavenumbers with size of box
eliminate_highest_freq : bool, optional
Set Nyquist frequency to zero for evenly
shaped axes
"""
K = []
N = self.shape()
for axis, base in enumerate(self):
K.append(
base.wavenumbers(
N,
axis,
scaled=scaled,
eliminate_highest_freq=eliminate_highest_freq))
return K
def local_wavenumbers(self,
broadcast=False,
scaled=False,
eliminate_highest_freq=False):
"""Return list of local wavenumbers of TensorProductSpace
Parameters
----------
broadcast : bool, optional
Broadcast returned wavenumber arrays to actual
dimensions of TensorProductSpace
scaled : bool, optional
Scale wavenumbers with size of box
eliminate_highest_freq : bool, optional
Set Nyquist frequency to zero for evenly
shaped axes
"""
k = self.wavenumbers(scaled=scaled,
eliminate_highest_freq=eliminate_highest_freq)
lk = []
for axis, (n, s) in enumerate(zip(k, self.local_slice(True))):
ss = [slice(None)] * len(k)
ss[axis] = s
lk.append(n[ss])
if broadcast is True:
return [np.broadcast_to(m, self.local_shape(True)) for m in lk]
return lk
def mesh(self):
"""Return list of 1D physical mesh for each dimension of
TensorProductSpace
"""
X = []
N = self.shape(False)
for axis, base in enumerate(self):
X.append(base.mesh(N, axis))
return X
def local_mesh(self, broadcast=False):
"""Return list of 1D physical mesh for each dimension of
TensorProductSpace
Parameters
----------
broadcast : bool, optional
Broadcast each 1D mesh to real shape of
TensorProductSpace
"""
m = self.mesh()
lm = []
for axis, (n, s) in enumerate(zip(m, self.local_slice(False))):
ss = [slice(None)] * len(m)
ss[axis] = s
lm.append(n[ss])
if broadcast is True:
return [np.broadcast_to(m, self.local_shape(False)) for m in lm]
return lm
def shape(self, spectral=False):
"""Return shape of TensorProductSpace in physical space
Parameters
----------
spectral : bool, optional
If True then return shape of spectral space, i.e.,
the input to a backward transfer. If False then return
shape of physical space, i.e., the input to a
forward transfer.
"""
if spectral == False:
return [
int(np.round(base.N * base.padding_factor)) for base in self
]
return self.spectral_shape()
def spectral_shape(self):
"""Return shape of TensorProductSpace in spectral space
Note
----
Spectral space corresponds to the result of a forward transfer
"""
return [base.spectral_shape() for base in self]
def __iter__(self):
return iter(self.bases)
def local_shape(self, spectral=True):
"""Return local shape of TensorProductSpace
Parameters
----------
spectral : bool, optional
If True then return local shape of spectral space, i.e.,
the input to a backward transfer. If False then return
local shape of physical space, i.e., the input to a
forward transfer.
"""
if not spectral:
return self.forward.input_pencil.subshape
else:
return self.backward.input_pencil.subshape
def local_slice(self, spectral=True):
"""Return the local view into the global data
Parameters
----------
spectral : bool, optional
If True then return local slice of spectral pace, i.e.,
the input to a backward transfer. If False then return
local slice of physical space, i.e., the input to a
forward transfer.
"""
if spectral is not True:
ip = self.forward.input_pencil
s = [
slice(start, start + shape)
for start, shape in zip(ip.substart, ip.subshape)
]
else:
ip = self.backward.input_pencil
s = [
slice(start, start + shape)
for start, shape in zip(ip.substart, ip.subshape)
]
return s
def rank(self):
"""Return rank of TensorProductSpace"""
return 1
def ndim(self):
"""Return dimension of TensorProductSpace"""
return len(self.bases)
def __len__(self):
"""Return dimension of TensorProductSpace"""
return len(self.bases)
def num_components(self):
"""Return number of spaces in TensorProductSpace"""
return 1
def __getitem__(self, i):
"""Return instance of base i
Parameters
----------
i : int
"""
return self.bases[i]
def test_mpifft():
from itertools import product
comm = MPI.COMM_WORLD
dims = (2, 3, 4,)
sizes = (12, 13)
assert MPI.COMM_WORLD.Get_size() < 8, "due to sizes"
types = ''
for t in 'fdg':
if fftw.get_fftw_lib(t):
types += t+t.upper()
grids = {2: (None,),
3: ((-1,), None),
4: ((-1,), None)}
for typecode in types:
for dim in dims:
for shape in product(*([sizes]*dim)):
if dim < 3:
n = min(shape)
if typecode in 'fdg':
n //= 2
n += 1
if n < comm.size:
continue
for grid in grids[dim]:
padding = False
for collapse in (True, False):
for backend in backends:
transforms = None
if dim < 3:
allaxes = [None, (-1,), (-2,),
(-1, -2,), (-2, -1),
(-1, 0), (0, -1),
((0,), (1,))]
elif dim < 4:
allaxes = [None, ((0,), (1, 2)),
((0,), (-2, -1))]
elif dim > 3:
allaxes = [None, ((0,), (1,), (2,), (3,)),
((0,), (1, 2, 3)),
((0,), (1,), (2, 3))]
dctn = functools.partial(fftw.dctn, type=3)
idctn = functools.partial(fftw.idctn, type=3)
if not typecode in 'FDG':
if backend == 'pyfftw':
transforms = {(3,): (pyfftw.builders.rfftn, pyfftw.builders.irfftn),
(2, 3): (pyfftw.builders.rfftn, pyfftw.builders.irfftn),
(1, 2, 3): (pyfftw.builders.rfftn, pyfftw.builders.irfftn),
(0, 1, 2, 3): (pyfftw.builders.rfftn, pyfftw.builders.irfftn)}
else:
transforms = {(3,): (dctn, idctn),
(2, 3): (dctn, idctn),
(1, 2, 3): (dctn, idctn),
(0, 1, 2, 3): (dctn, idctn)}
for axes in allaxes:
# Test also the slab is number interface
_grid = grid
if grid is not None:
ax = -1
if axes is not None:
ax = axes[-1] if isinstance(axes[-1], int) else axes[-1][-1]
_slab = (ax+1) % len(shape)
_grid = [1]*(_slab+1)
_grid[_slab] = 0
_comm = comm
# Test also the comm is Subcomm interfaces
# For PFFT the Subcomm needs to be as long as shape
if len(shape) > 2 and axes is None and grid is None:
_dims = [0] * len(shape)
_dims[-1] = 1 # distribute all but last axis (axes is None)
_comm = comm
if random_true_or_false(comm) == 1:
# then test Subcomm with a MPI.CART argument
_dims = MPI.Compute_dims(comm.Get_size(), _dims)
_comm = comm.Create_cart(_dims)
_dims = None
_comm = Subcomm(_comm, _dims)
#print(typecode, shape, axes, collapse, _grid)
fft = PFFT(_comm, shape, axes=axes, dtype=typecode,
padding=padding, grid=_grid, collapse=collapse,
backend=backend, transforms=transforms)
#if comm.rank == 0:
# grid_ = [c.size for c in fft.subcomm]
# print('grid:{} shape:{} typecode:{} backend:{} axes:{}'
# .format(grid_, shape, typecode, backend, axes))
assert fft.dtype(True) == fft.forward.output_array.dtype
assert fft.dtype(False) == fft.forward.input_array.dtype
assert len(fft.axes) == len(fft.xfftn)
assert len(fft.axes) == len(fft.transfer) + 1
assert (fft.forward.input_pencil.subshape ==
fft.forward.input_array.shape)
assert (fft.forward.output_pencil.subshape ==
fft.forward.output_array.shape)
assert (fft.backward.input_pencil.subshape ==
fft.backward.input_array.shape)
assert (fft.backward.output_pencil.subshape ==
fft.backward.output_array.shape)
assert np.alltrue(np.array(fft.global_shape(True)) == np.array(fft.forward.output_pencil.shape))
assert np.alltrue(np.array(fft.global_shape(False)) == np.array(fft.forward.input_pencil.shape))
ax = -1 if axes is None else axes[-1] if isinstance(axes[-1], int) else axes[-1][-1]
assert fft.forward.input_pencil.substart[ax] == 0
assert fft.backward.output_pencil.substart[ax] == 0
ax = 0 if axes is None else axes[0] if isinstance(axes[0], int) else axes[0][0]
assert fft.forward.output_pencil.substart[ax] == 0
assert fft.backward.input_pencil.substart[ax] == 0
assert fft.dimensions == len(shape)
U = random_like(fft.forward.input_array)
if random_true_or_false(comm) == 1:
F = fft.forward(U)
V = fft.backward(F)
assert allclose(V, U)
else:
fft.forward.input_array[...] = U
fft.forward()
fft.backward()
V = fft.backward.output_array
assert allclose(V, U)
fft.destroy()
padding = [1.5]*len(shape)
for backend in backends:
if dim < 3:
allaxes = [None, (-1,), (-2,),
(-1, -2,), (-2, -1),
(-1, 0), (0, -1),
((0,), (1,))]
elif dim < 4:
allaxes = [None, ((0,), (1,), (2,)),
((0,), (-2,), (-1,))]
elif dim > 3:
allaxes = [None, (0, 1, -2, -1),
((0,), (1,), (2,), (3,))]
for axes in allaxes:
_grid = grid
if grid is not None:
ax = -1
if axes is not None:
ax = axes[-1] if isinstance(axes[-1], int) else axes[-1][-1]
_slab = (ax+1) % len(shape)
_grid = [1]*(_slab+1)
_grid[_slab] = 0
fft = PFFT(comm, shape, axes=axes, dtype=typecode,
padding=padding, grid=_grid, backend=backend)
#if comm.rank == 0:
# grid = [c.size for c in fft.subcomm]
# print('grid:{} shape:{} typecode:{} backend:{} axes:{}'
# .format(grid, shape, typecode, backend, axes))
assert len(fft.axes) == len(fft.xfftn)
assert len(fft.axes) == len(fft.transfer) + 1
assert (fft.forward.input_pencil.subshape ==
fft.forward.input_array.shape)
assert (fft.forward.output_pencil.subshape ==
fft.forward.output_array.shape)
assert (fft.backward.input_pencil.subshape ==
fft.backward.input_array.shape)
assert (fft.backward.output_pencil.subshape ==
fft.backward.output_array.shape)
ax = -1 if axes is None else axes[-1] if isinstance(axes[-1], int) else axes[-1][-1]
assert fft.forward.input_pencil.substart[ax] == 0
assert fft.backward.output_pencil.substart[ax] == 0
ax = 0 if axes is None else axes[0] if isinstance(axes[0], int) else axes[0][0]
assert fft.forward.output_pencil.substart[ax] == 0
assert fft.backward.input_pencil.substart[ax] == 0
U = random_like(fft.forward.input_array)
F = fft.forward(U)
if random_true_or_false(comm) == 1:
Fc = F.copy()
V = fft.backward(F)
F = fft.forward(V)
assert allclose(F, Fc)
else:
fft.backward.input_array[...] = F
fft.backward()
fft.forward()
V = fft.forward.output_array
assert allclose(F, V)
# Test normalization on backward transform instead of default
fft.backward.input_array[...] = F
fft.backward(normalize=True)
fft.forward(normalize=False)
V = fft.forward.output_array
assert allclose(F, V)
fft.destroy()