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)