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