def quad_mesh(x, y, cut_x_edges=False, cut_y_edges=False): """ Construct quadrilateral mesh arrays from two grids. Intended for use with e.g. plt.pcolor. Parameters ---------- x : 1d array Grid for last axis of the mesh. y : 1d array Grid for first axis of the mesh. cut_x_edges, cut_y_edges : bool, optional True to truncate edge quadrilaterals at x/y grid edges. False (default) to center edge quadrilaterals at x/y grid edges. """ # Get 1d vertex vectors xvert = get_1d_vertices(x, cut_edges=cut_x_edges) yvert = get_1d_vertices(y, cut_edges=cut_y_edges) # Reshape as multidimensional vectors xvert = reshape_vector(xvert, dim=2, axis=1) yvert = reshape_vector(yvert, dim=2, axis=0) # Broadcast up to arrays xmesh = xvert * np.ones_like(yvert) ymesh = yvert * np.ones_like(xvert) return xmesh, ymesh
def __init__(self, arg, basis0, basis1, **kw): arg = Operand.cast(arg) super().__init__(arg, **kw) self.basis0 = self.domain.get_basis_object(basis0) self.basis1 = self.domain.get_basis_object(basis1) self.axis0 = self.domain.bases.index(self.basis0) self.axis1 = self.domain.bases.index(self.basis1) if self.axis0 > self.axis1: raise ValueError("Cannot evaluate specified axis order.") if self.basis0.interval != self.basis1.interval: raise ValueError("Bases must occupy same interval.") self.name = 'Diag[%s=%s]' % (self.basis0.name, self.basis1.name) # Shear array k0 = reshape_vector(self.basis0.wavenumbers, dim=self.domain.dim, axis=self.axis0) x1 = self.domain.grid(self.axis1, scales=self.domain.dealias) dx0 = x1 - self.basis1.interval[0] self.shear = np.exp(1j * k0 * dx0) # Filter mask slices = self.domain.dist.coeff_layout.slices(self.domain.dealias) k0 = reshape_vector(self.basis0.wavenumbers[slices[self.axis0]], dim=self.domain.dim, axis=self.axis0) k1 = reshape_vector(self.basis1.wavenumbers[slices[self.axis1]], dim=self.domain.dim, axis=self.axis1) self.filter_mask = (k0 == k1)
def compute_second_cumulant(domain, f, g): """ Compute second cumulant of two 2D functions under x' averaging: cfg(x, y0, y1) = <f(x'+x, y0) g(x', y1)> Parameters ---------- domain : domain object Domain f, g: functions of (x, y) Fields over which to compute second cumulant Notes ----- Assumes triple-Fourier real domain over (x, y0, y1). """ x, y0, y1 = domain.grids() # Build 2D fields as functions of (x,y0) and (x,y1) a = domain.new_field() a['g'] = f(x, y0) b = domain.new_field() b['g'] = g(x, y1) # Compute cumulant directly from Fourier modes c = domain.new_field() kx = reshape_vector(domain.bases[0].wavenumbers, dim=3, axis=0) shift = np.exp(1j * kx * domain.bases[0].interval[0]) c['c'] = a['c'][:,:,0:1] * np.roll(b['c'],-1,2)[:,0:1,::-1].conj() * shift return c
def __init__(self, arg, basis0, basis1, **kw): arg = Operand.cast(arg) super().__init__(arg, **kw) self.basis0 = self.domain.get_basis_object(basis0) self.basis1 = self.domain.get_basis_object(basis1) self.axis0 = self.domain.bases.index(self.basis0) self.axis1 = self.domain.bases.index(self.basis1) if self.axis0 > self.axis1: raise ValueError("Cannot evaluate specified axis order.") if self.basis0.interval != self.basis1.interval: raise ValueError("Bases must occupy same interval.") self.name = 'CoeffDiag[%s=%s]' % (self.basis0.name, self.basis1.name) # Filter mask slices = self.domain.dist.layouts[2].slices(scales=1) k0 = reshape_vector(self.basis0.wavenumbers[slices[self.axis0]], dim=self.domain.dim, axis=self.axis0) k1 = reshape_vector(self.basis1.wavenumbers[slices[self.axis1]], dim=self.domain.dim, axis=self.axis1) self.filter_mask = (k0 == k1)
def weight(self, axis, dimensions=2): if axis == 0 and dimensions == 2: weight = self.S.weights if axis == 1 and dimensions == 2: weight = self.dV if axis == 1 and dimensions == 3: weight = self.S.weights if axis == 2 and dimensions == 3: weight = self.dV return reshape_vector(weight, dimensions, axis)
def grid(self, axis, dimensions=2): if axis == 0 and dimensions == 2: grid = self.theta if axis == 1 and dimensions == 2: grid = self.radius if axis == 1 and dimensions == 3: grid = self.theta if axis == 2 and dimensions == 3: grid = self.radius return reshape_vector(grid, dimensions, axis)
def get_vertex_vector(r, axis, trim_edges=False): #axis = vector_axis(r) vflat = get_vertex_flat(r.ravel(), trim_edges=trim_edges) return reshape_vector(vflat, dim=2, axis=axis)