class hwconv:
def __init__(self,
Nx,
Ny,
padx=1.5,
pady=1.5,
num_in=6,
num_out=2,
fmultin=mult_in,
fmultout=mult_out62,
comm=MPI.COMM_WORLD):
self.comm = comm
self.num_in = num_in
self.num_out = num_out
self.nar = max(num_in, num_out)
self.ffti = PFFT(comm,
shape=(self.num_in, Nx, Ny),
axes=(1, 2),
grid=[1, -1, 1],
padding=[1, 1.5, 1.5],
collapse=False)
self.ffto = PFFT(comm,
shape=(self.num_out, Nx, Ny),
axes=(1, 2),
grid=[1, -1, 1],
padding=[1, 1.5, 1.5],
collapse=False)
self.datk = newDistArray(self.ffti, forward_output=True)
self.dat = newDistArray(self.ffti, forward_output=False)
lkx = np.r_[0:int(Nx / 2), -int(Nx / 2):0]
lky = np.r_[0:int(Ny / 2 + 1)]
self.kx = DistArray((Nx, int(Ny / 2 + 1)),
subcomm=(1, 0),
dtype=float,
alignment=0)
self.ky = DistArray((Nx, int(Ny / 2 + 1)),
subcomm=(1, 0),
dtype=float,
alignment=0)
self.kx[:], self.ky[:] = np.meshgrid(lkx[self.kx.local_slice()[0]],
lky[self.ky.local_slice()[1]],
indexing='ij')
self.ksqr = self.kx**2 + self.ky**2
self.fmultin = fmultin
self.fmultout = fmultout
def convolve(self, u):
hermitian_symmetrize(u)
if (u.local_slice()[2].stop == u.global_shape[2]):
u[:, :, -1] = 0
u[:, int(Nx / 2), :] = 0
self.fmultin(u, self.datk, self.kx, self.ky, self.ksqr)
self.ffti.backward(self.datk, self.dat)
self.fmultout(self.dat)
self.ffto.forward(self.dat[:self.num_out, ],
self.datk[:self.num_out, ])
if (self.datk.local_slice()[2].stop == self.datk.global_shape[2]):
self.datk[:, :, -1] = 0
self.datk[:, int(Nx / 2), :] = 0
return self.datk[:self.num_out, ]
def __init__(self,
Nx,
Ny,
padx=1.5,
pady=1.5,
num_in=6,
num_out=2,
fmultin=mult_in,
fmultout=mult_out62,
comm=MPI.COMM_WORLD):
self.comm = comm
self.num_in = num_in
self.num_out = num_out
self.nar = max(num_in, num_out)
self.ffti = PFFT(comm,
shape=(self.num_in, Nx, Ny),
axes=(1, 2),
grid=[1, -1, 1],
padding=[1, 1.5, 1.5],
collapse=False)
self.ffto = PFFT(comm,
shape=(self.num_out, Nx, Ny),
axes=(1, 2),
grid=[1, -1, 1],
padding=[1, 1.5, 1.5],
collapse=False)
self.datk = newDistArray(self.ffti, forward_output=True)
self.dat = newDistArray(self.ffti, forward_output=False)
lkx = np.r_[0:int(Nx / 2), -int(Nx / 2):0]
lky = np.r_[0:int(Ny / 2 + 1)]
self.kx = DistArray((Nx, int(Ny / 2 + 1)),
subcomm=(1, 0),
dtype=float,
alignment=0)
self.ky = DistArray((Nx, int(Ny / 2 + 1)),
subcomm=(1, 0),
dtype=float,
alignment=0)
self.kx[:], self.ky[:] = np.meshgrid(lkx[self.kx.local_slice()[0]],
lky[self.ky.local_slice()[1]],
indexing='ij')
self.ksqr = self.kx**2 + self.ky**2
self.fmultin = fmultin
self.fmultout = fmultout
def __new__(cls, space, val=0, buffer=None):
if hasattr(space, 'points_and_weights'): # 1D case
if cls.__name__ == 'Function':
dtype = space.forward.output_array.dtype
shape = space.forward.output_array.shape
elif cls.__name__ == 'Array':
dtype = space.forward.input_array.dtype
shape = space.forward.input_array.shape
if not space.num_components() == 1:
shape = (space.num_components(), ) + shape
obj = DistArray.__new__(cls,
shape,
buffer=buffer,
dtype=dtype,
rank=space.rank)
obj._space = space
if buffer is None and isinstance(val, Number):
obj[:] = val
return obj
if cls.__name__ == 'Function':
forward_output = True
p0 = space.pencil[1]
dtype = space.forward.output_array.dtype
elif cls.__name__ == 'Array':
forward_output = False
p0 = space.pencil[0]
dtype = space.forward.input_array.dtype
global_shape = space.global_shape(forward_output)
obj = DistArray.__new__(cls,
global_shape,
subcomm=p0.subcomm,
val=val,
dtype=dtype,
buffer=buffer,
alignment=p0.axis,
rank=space.rank)
obj._space = space
return obj
def test_2Darray():
N = (8, 8)
for subcomm in ((0, 1), (1, 0), None, Subcomm(comm, (0, 1))):
for rank in (0, 1, 2):
M = (2, ) * rank + N
alignment = None
if subcomm is None and rank == 1:
alignment = 1
a = DistArray(M,
subcomm=subcomm,
val=1,
rank=rank,
alignment=alignment)
assert a.rank == rank
assert a.global_shape == M
_ = a.substart
c = 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
aa = a.v
assert isinstance(aa, np.ndarray)
try:
k = a.get((0, ) * rank + (0, slice(None)))
if comm.Get_rank() == 0:
assert len(k) == N[1]
assert np.sum(k) == N[1]
k = a.get((0, ) * rank + (slice(None), 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) % 2
_ = a.get_pencil_and_transfer(newaxis)
a[:] = MPI.COMM_WORLD.Get_rank()
b = a.redistribute(newaxis)
a = b.redistribute(out=a)
a = b.redistribute(a.alignment, 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
c = a.redistribute(a.alignment)
assert c is a
def __new__(cls, space, val=0, buffer=None):
if hasattr(space, 'points_and_weights'): # 1D case
if cls.__name__ == 'Function':
dtype = space.forward.output_array.dtype
shape = space.forward.output_array.shape
elif cls.__name__ == 'Array':
dtype = space.forward.input_array.dtype
shape = space.forward.input_array.shape
if not space.num_components() == 1:
shape = (space.num_components(),) + shape
if hasattr(buffer, 'free_symbols'):
# Evaluate sympy function on entire mesh
import sympy
x, y, z = sympy.symbols("x,y,z")
sym0 = [sym for sym in (x, y, z) if sym in buffer.free_symbols]
buffer = sympy.lambdify(sym0, buffer)
buffer = buffer(space.mesh())
if cls.__name__ == 'Function':
buf = np.empty_like(space.forward.output_array)
buf = space.forward(buffer, buf)
buffer = buf
obj = DistArray.__new__(cls, shape, buffer=buffer, dtype=dtype,
rank=space.rank)
obj._space = space
obj._offset = 0
if buffer is None and isinstance(val, Number):
obj[:] = val
return obj
if cls.__name__ == 'Function':
forward_output = True
p0 = space.forward.output_pencil
dtype = space.forward.output_array.dtype
elif cls.__name__ == 'Array':
forward_output = False
p0 = space.backward.output_pencil
dtype = space.forward.input_array.dtype
# Evaluate sympy function on entire mesh
if hasattr(buffer, 'free_symbols'):
import sympy
x, y, z, r, s, t = sympy.symbols("x,y,z,r,s,t")
sym0 = [sym for sym in (x, y, z, r, s, t) if sym in buffer.free_symbols]
#sym0 = tuple(buffer.free_symbols)
buffer = sympy.lambdify(sym0, buffer)(*space.local_mesh())
if cls.__name__ == 'Function':
buf = np.empty_like(space.forward.output_array)
buf = space.forward(buffer, buf)
buffer = buf
global_shape = space.global_shape(forward_output)
obj = DistArray.__new__(cls, global_shape,
subcomm=p0.subcomm, val=val, dtype=dtype,
buffer=buffer, alignment=p0.axis,
rank=space.rank)
obj._space = space
obj._offset = 0
return obj
def refine(self, N, output_array=None):
"""Return self with new number of quadrature points
Parameters
----------
N : number or sequence of numbers
The new number of quadrature points
Note
----
If N is smaller than for self, then a truncated array
is returned. If N is greater than before, then the
returned array is padded with zeros.
"""
from shenfun.fourier.bases import R2CBasis
from shenfun import VectorTensorProductSpace
if self.ndim == 1:
assert isinstance(N, Number)
space = self.function_space()
if output_array is None:
refined_basis = space.get_refined(N)
output_array = Function(refined_basis)
output_array = self.assign(output_array)
return output_array
space = self.function_space()
if isinstance(space, VectorTensorProductSpace):
if output_array is None:
output_array = [None]*len(self)
for i, array in enumerate(self):
output_array[i] = array.refine(N, output_array=output_array[i])
if isinstance(output_array, list):
T = output_array[0].function_space()
VT = VectorTensorProductSpace(T)
output_array = np.array(output_array)
output_array = Function(VT, buffer=output_array)
return output_array
axes = [bx for ax in space.axes for bx in ax]
base = space.bases[axes[0]]
global_shape = list(self.global_shape) # Global shape in spectral space
factor = N[axes[0]]/self.function_space().bases[axes[0]].N
if isinstance(base, R2CBasis):
global_shape[axes[0]] = int((2*global_shape[axes[0]]-2)*factor)//2+1
else:
global_shape[axes[0]] = int(global_shape[axes[0]]*factor)
c1 = DistArray(global_shape,
subcomm=self.pencil.subcomm,
dtype=self.dtype,
alignment=self.alignment)
if self.global_shape[axes[0]] <= global_shape[axes[0]]:
base._padding_backward(self, c1)
else:
base._truncation_forward(self, c1)
for ax in axes[1:]:
c0 = c1.redistribute(ax)
factor = N[ax]/self.function_space().bases[ax].N
# Get a new padded array
base = space.bases[ax]
if isinstance(base, R2CBasis):
global_shape[ax] = int(base.N*factor)//2+1
else:
global_shape[ax] = int(global_shape[ax]*factor)
c1 = DistArray(global_shape,
subcomm=c0.pencil.subcomm,
dtype=c0.dtype,
alignment=ax)
# Copy from c0 to d0
if self.global_shape[ax] <= global_shape[ax]:
base._padding_backward(c0, c1)
else:
base._truncation_forward(c0, c1)
# Reverse transfer to get the same distribution as u_hat
for ax in reversed(axes[:-1]):
c1 = c1.redistribute(ax)
if output_array is None:
refined_space = space.get_refined(N)
output_array = Function(refined_space, buffer=c1)
else:
output_array[:] = c1
return output_array
u[:, int(Nx / 2), :] = 0
self.fmultin(u, self.datk, self.kx, self.ky, self.ksqr)
self.ffti.backward(self.datk, self.dat)
self.fmultout(self.dat)
self.ffto.forward(self.dat[:self.num_out, ],
self.datk[:self.num_out, ])
if (self.datk.local_slice()[2].stop == self.datk.global_shape[2]):
self.datk[:, :, -1] = 0
self.datk[:, int(Nx / 2), :] = 0
return self.datk[:self.num_out, ]
Nx, Ny = 8, 8
h = hwconv(Nx, Ny)
uk = DistArray((2, Nx, int(Ny / 2 + 1)),
subcomm=(1, 1, 0),
dtype=complex,
alignment=1)
vk = DistArray((2, Nx, int(Ny / 2 + 1)),
subcomm=(1, 1, 0),
dtype=complex,
alignment=1)
kx = DistArray((Nx, int(Ny / 2 + 1)), subcomm=(1, 0), dtype=float, alignment=0)
ky = DistArray((Nx, int(Ny / 2 + 1)), subcomm=(1, 0), dtype=float, alignment=0)
inds = np.r_[int(Nx / 2):Nx, 0:int(Nx / 2)]
uk[0, :, :] = np.array(
[[inds[l - 1] + 1j * m for m in np.r_[uk.local_slice()[2]]]
for l in np.r_[uk.local_slice()[1]]])
uk[1, :, :] = np.array(
[[2 * inds[l - 1] + 1j * (m + 1) for m in np.r_[uk.local_slice()[2]]]
for l in np.r_[uk.local_slice()[1]]])
vk[:] = h.convolve(uk)
def __new__(cls, space, val=0, buffer=None):
if hasattr(space, 'points_and_weights'): # 1D case
if cls.__name__ == 'Function':
dtype = space.forward.output_array.dtype
shape = space.forward.output_array.shape
elif cls.__name__ == 'Array':
dtype = space.forward.input_array.dtype
shape = space.forward.input_array.shape
if not space.num_components() == 1:
shape = (space.num_components(), ) + shape
if hasattr(buffer, 'free_symbols'):
# Evaluate sympy function on entire mesh
x = buffer.free_symbols.pop()
buffer = sp.lambdify(x, buffer)
buf = buffer(space.mesh()).astype(
space.forward.input_array.dtype)
buffer = Array(space)
buffer[:] = buf
if cls.__name__ == 'Function':
buf = Function(space)
buf = buffer.forward(buf)
buffer = buf
obj = DistArray.__new__(cls,
shape,
buffer=buffer,
dtype=dtype,
rank=space.is_composite_space)
obj._space = space
obj._offset = 0
if buffer is None and isinstance(val, Number):
obj[:] = val
return obj
if cls.__name__ == 'Function':
forward_output = True
p0 = space.forward.output_pencil
dtype = space.forward.output_array.dtype
elif cls.__name__ == 'Array':
forward_output = False
p0 = space.backward.output_pencil
dtype = space.forward.input_array.dtype
# Evaluate sympy function on entire mesh
if hasattr(buffer, 'free_symbols'):
sym0 = buffer.free_symbols
mesh = space.local_mesh(True)
m = []
for sym in sym0:
j = 'xyzrs'.index(str(sym))
m.append(mesh[j])
buf = sp.lambdify(sym0,
buffer,
modules=['numpy', {
'cot': cot,
'Ynm': Ynm
}])(*m).astype(space.forward.input_array.dtype)
buffer = Array(space)
buffer[:] = buf
if cls.__name__ == 'Function':
buf = Function(space)
buf = buffer.forward(buf)
buffer = buf
global_shape = space.global_shape(forward_output)
obj = DistArray.__new__(cls,
global_shape,
subcomm=p0.subcomm,
val=val,
dtype=dtype,
buffer=buffer,
alignment=p0.axis,
rank=space.is_composite_space)
obj._space = space
obj._offset = 0
return obj
def test_1Darray():
N = (8, )
z = DistArray(N, val=2)
assert z[0] == 2
assert z.shape == N