def div(fld):
""" first order """
vx, vy, vz = fld.component_views()
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
divcenter = "Cell"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
divcenter = "Node"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
else:
raise NotImplementedError("Can only do cell and node centered divs")
xp = crdx[2:,:,:]; xm = crdx[:-2,:,:] #pylint: disable=W0612,C0321,C0324
yp = crdy[:,2:,:]; ym = crdy[:,:-2,:] #pylint: disable=W0612,C0321,C0324
zp = crdz[:,:,2:]; zm = crdz[:,:,:-2] #pylint: disable=W0612,C0321,C0324
vxp = vx[2:,1:-1,1:-1]; vxm = vx[:-2,1:-1,1:-1] #pylint: disable=W0612,C0321,C0324,C0301
vyp = vy[1:-1,2:,1:-1]; vym = vy[1:-1,:-2,1:-1] #pylint: disable=W0612,C0321,C0324,C0301
vzp = vz[1:-1,1:-1,2:]; vzm = vz[1:-1,1:-1,:-2] #pylint: disable=W0612,C0321,C0324,C0301
# print(vxp.shape, vyp.shape, vzp.shape)
# print(vxm.shape, vym.shape, vzm.shape)
# print(xp.shape, yp.shape, zp.shape)
# print(xm.shape, ym.shape, zm.shape)
div_arr = ne.evaluate("(vxp-vxm)/(xp-xm) + (vyp-vym)/(yp-ym) + "
"(vzp-vzm)/(zp-zm)")
return field.wrap_field(div_arr, divcrds, name="div " + fld.name,
center=divcenter, time=fld.time, parents=[fld])
def _parse_file(self, filename, parent_node):
# we do minimal file parsing here for performance. we just
# make data wrappers from the templates we got from the first
# file in the group, and package them up into grids
# frame = int(re.match(self._detector, filename).group(3))
_grid = self._make_grid(parent_node, name="<GkeyllGrid>",
**self._grid_opts)
# FIXME: To get the time at load, we have to open all hdf5 files
# which defeats the purpose of making templates etc. in attempt to
# be lazy. Maybe there's a way to use frame above?
with h5py.File(filename, 'r') as f:
step = f['timeData'].attrs['vsStep']
time = f['timeData'].attrs['vsTime']
self.set_info("step", step)
self.time = time
_grid.time = time
_grid.set_crds(self._crds)
if self._fld_templates is None:
self._fld_templates = self._make_template(filename)
for i, meta in enumerate(self._fld_templates):
# FIXME: the transpose goes xyz -> zyx
h5_data = H5pyDataWrapper(self.fname, "StructGridField",
comp_dim=-1, comp_idx=i)
fld = field.wrap_field(h5_data, self._crds, meta['fld_name'],
center="cell", pretty_name=meta['pretty_name'])
_grid.add_field(fld)
_grid.set_info('field_type', self.get_info('field_type'))
return _grid
def _parse_file(self, filename, parent_node):
# we do minimal file parsing here for performance. we just
# make data wrappers from the templates we got from the first
# file in the group, and package them up into grids
# frame = int(re.match(self._detector, filename).group(3))
_grid = self._make_grid(parent_node, name="<GkeyllGrid>",
**self._grid_opts)
# FIXME: To get the time at load, we have to open all hdf5 files
# which defeats the purpose of making templates etc. in attempt to
# be lazy. Maybe there's a way to use frame above?
with h5py.File(filename, 'r') as f:
step = f['timeData'].attrs['vsStep']
time = f['timeData'].attrs['vsTime']
self.set_info("step", step)
self.time = time
_grid.time = time
_grid.set_crds(self._crds)
if self._fld_templates is None:
self._fld_templates = self._make_template(filename)
for i, meta in enumerate(self._fld_templates):
# FIXME: the transpose goes xyz -> zyx
h5_data = H5pyDataWrapper(self.fname, "StructGridField",
comp_dim=-1, comp_idx=i)
fld = field.wrap_field(h5_data, self._crds, meta['fld_name'],
center="cell", pretty_name=meta['pretty_name'])
_grid.add_field(fld)
return _grid
def _make_field(self, parent_node, fldtype, name, crds, data, **kwargs):
"""Use this instead of calling Grid(...) yourself
Args:
parent_node (Dataset, Grid, or None): Hint at parent in
the tree, needed if info is used before this object
is added to its parent
"""
fld = field.wrap_field(data, crds, name=name, fldtype=fldtype, **kwargs)
if parent_node is not None:
parent_node.prepare_child(fld)
return fld
def _evaluate_numexpr(grid, result_name, eqn, slc=None):
"""
Returns:
Field
Raises:
RuntimeError, if no numexpr, or if evaluate is not enabled
TypeError, if numexpr couldn't understand a math input
"""
if not _has_numexpr:
raise RuntimeError("Evaluate not enabled, or numexpr not installed.")
# salt symbols that don't look like math functions and look them up
# in the grid
salt = "SALT"
_symbol_re = r"\b([_A-Za-z][_a-zA-Z0-9]*)\b"
var_re = _symbol_re + r"(?!\s*\()"
flds = []
# for security
eqn = eqn.replace("__", "")
local_dict = dict()
def var_salter(symbols):
symbol = symbols.groups()[0]
salted_symbol = salt + symbol
# yes, i'm not using dict.update on purpose since grid's
# getitem might copy the array
if salted_symbol not in local_dict:
this_fld = grid.get_field(symbol, slc=slc)
local_dict[salted_symbol] = this_fld
if len(flds) == 0:
if isinstance(this_fld, field.Field):
flds.append(this_fld)
else:
raise RuntimeError(
"reduced to scalar, no need for numexpr")
return salted_symbol
salted_eqn = re.sub(var_re, var_salter, eqn)
arr = ne.evaluate(salted_eqn,
local_dict=local_dict,
global_dict={"__builtins__": {}})
# FIXME: actually detect the type of field instead of asserting it's
# a scalar... also maybe auto-detect reduction operations?
if len(flds) > 0:
ctx = dict(name=result_name, pretty_name=result_name)
return flds[0].wrap(arr, context=ctx)
else:
logger.warn("Strange input to numexpr evaluator: %s", eqn)
return field.wrap_field(arr, grid.crds, name=result_name)
def _evaluate_numexpr(grid, result_name, eqn, slc=Ellipsis):
"""
Returns:
Field
Raises:
RuntimeError, if no numexpr, or if evaluate is not enabled
TypeError, if numexpr couldn't understand a math input
"""
if not _has_numexpr:
raise RuntimeError("Evaluate not enabled, or numexpr not installed.")
# salt symbols that don't look like math functions and look them up
# in the grid
salt = "SALT"
_symbol_re = r"\b([_A-Za-z][_a-zA-Z0-9]*)\b"
var_re = _symbol_re + r"(?!\s*\()"
flds = []
# for security
eqn = eqn.replace("__", "")
local_dict = dict()
def var_salter(symbols):
symbol = symbols.groups()[0]
salted_symbol = salt + symbol
# yes, i'm not using dict.update on purpose since grid's
# getitem might copy the array
if salted_symbol not in local_dict:
this_fld = grid.get_field(symbol, slc=slc)
local_dict[salted_symbol] = this_fld
if len(flds) == 0:
if isinstance(this_fld, field.Field):
flds.append(this_fld)
else:
raise RuntimeError("reduced to scalar, no need for numexpr")
return salted_symbol
salted_eqn = re.sub(var_re, var_salter, eqn)
arr = ne.evaluate(salted_eqn, local_dict=local_dict,
global_dict={"__builtins__": {}})
# FIXME: actually detect the type of field instead of asserting it's
# a scalar... also maybe auto-detect reduction operations?
if len(flds) > 0:
ctx = dict(name=result_name, pretty_name=result_name)
return flds[0].wrap(arr, context=ctx)
else:
logger.warning("Strange input to numexpr evaluator: %s", eqn)
return field.wrap_field(arr, grid.crds, name=result_name)
def _make_field(self, parent_node, fldtype, name, crds, data, **kwargs):
"""Use this instead of calling Grid(...) yourself
Args:
parent_node (Dataset, Grid, or None): Hint at parent in
the tree, needed if info is used before this object
is added to its parent
"""
fld = field.wrap_field(data,
crds,
name=name,
fldtype=fldtype,
**kwargs)
if parent_node is not None:
parent_node.prepare_child(fld)
return fld
def _curl_np(fld, bnd=True):
"""2nd order centeral diff, 1st order @ boundaries if bnd"""
if bnd:
fld = viscid.extend_boundaries(fld, order=0, crd_order=0)
vx, vy, vz = fld.component_views()
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
curlcenter = "cell"
# curlcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
curlcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
curlcenter = "node"
# curlcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
curlcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
else:
raise NotImplementedError("Can only do cell and node centered divs")
xp, xm = crdx[2:, :, :], crdx[:-2, :, :] # pylint: disable=bad-whitespace
yp, ym = crdy[:, 2:, :], crdy[:, :-2, :] # pylint: disable=bad-whitespace
zp, zm = crdz[:, :, 2:], crdz[:, :, :-2] # pylint: disable=bad-whitespace
vxpy, vxmy = vx[1:-1, 2:, 1:-1], vx[1:-1, :-2, 1:-1] # pylint: disable=bad-whitespace
vxpz, vxmz = vx[1:-1, 1:-1, 2:], vx[1:-1, 1:-1, :-2] # pylint: disable=bad-whitespace
vypx, vymx = vy[2:, 1:-1, 1:-1], vy[:-2, 1:-1, 1:-1] # pylint: disable=bad-whitespace
vypz, vymz = vy[1:-1, 1:-1, 2:], vy[1:-1, 1:-1, :-2] # pylint: disable=bad-whitespace
vzpx, vzmx = vz[2:, 1:-1, 1:-1], vz[:-2, 1:-1, 1:-1] # pylint: disable=bad-whitespace
vzpy, vzmy = vz[1:-1, 2:, 1:-1], vz[1:-1, :-2, 1:-1] # pylint: disable=bad-whitespace
curl_x = (vzpy - vzmy) / (yp - ym) - (vypz - vymz) / (zp - zm)
curl_y = -(vzpx - vzmx) / (xp - xm) + (vxpz - vxmz) / (zp - zm)
curl_z = (vypx - vymx) / (xp - xm) - (vxpy - vxmy) / (yp - ym)
return field.wrap_field([curl_x, curl_y, curl_z],
curlcrds,
name="curl " + fld.name,
fldtype="Vector",
center=curlcenter,
time=fld.time,
parents=[fld])
def curl(fld, bnd=True):
"""2nd order centeral diff, 1st order @ boundaries if bnd"""
if bnd:
fld = viscid.extend_boundaries(fld, order=0, crd_order=0)
vx, vy, vz = fld.component_views()
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
curlcenter = "cell"
# curlcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
curlcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
curlcenter = "node"
# curlcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
curlcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
else:
raise NotImplementedError("Can only do cell and node centered divs")
xp, xm = crdx[2:, :, :], crdx[:-2, : , : ] # pylint: disable=C0326
yp, ym = crdy[ :, 2:, :], crdy[: , :-2, : ] # pylint: disable=C0326
zp, zm = crdz[ :, :, 2:], crdz[: , : , :-2] # pylint: disable=C0326
vxpy, vxmy = vx[1:-1, 2: , 1:-1], vx[1:-1, :-2, 1:-1] # pylint: disable=C0326
vxpz, vxmz = vx[1:-1, 1:-1, 2: ], vx[1:-1, 1:-1, :-2] # pylint: disable=C0326
vypx, vymx = vy[2: , 1:-1, 1:-1], vy[ :-2, 1:-1, 1:-1] # pylint: disable=C0326
vypz, vymz = vy[1:-1, 1:-1, 2: ], vy[1:-1, 1:-1, :-2] # pylint: disable=C0326
vzpx, vzmx = vz[2: , 1:-1, 1:-1], vz[ :-2, 1:-1, 1:-1] # pylint: disable=C0326
vzpy, vzmy = vz[1:-1, 2: , 1:-1], vz[1:-1, :-2, 1:-1] # pylint: disable=C0326
curl_x = ne.evaluate("(vzpy-vzmy)/(yp-ym) - (vypz-vymz)/(zp-zm)")
curl_y = ne.evaluate("-(vzpx-vzmx)/(xp-xm) + (vxpz-vxmz)/(zp-zm)")
curl_z = ne.evaluate("(vypx-vymx)/(xp-xm) - (vxpy-vxmy)/(yp-ym)")
return field.wrap_field([curl_x, curl_y, curl_z], curlcrds,
name="curl " + fld.name, fldtype="Vector",
center=curlcenter, time=fld.time,
parents=[fld])
def _div_np(fld, bnd=True):
"""2nd order centeral diff, 1st order @ boundaries if bnd"""
if fld.iscentered("Face"):
# dispatch fc div immediately since that does its own pre-processing
return viscid.div_fc(fld, bnd=bnd)
if bnd:
fld = viscid.extend_boundaries(fld, order=0, crd_order=0)
vx, vy, vz = fld.component_views()
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
divcenter = "Cell"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
divcenter = "Node"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
else:
raise NotImplementedError("Can only do cell and node centered divs")
xp, xm = crdx[2:, :, :], crdx[:-2, :, :] # pylint: disable=bad-whitespace
yp, ym = crdy[:, 2:, :], crdy[:, :-2, :] # pylint: disable=bad-whitespace
zp, zm = crdz[:, :, 2:], crdz[:, :, :-2] # pylint: disable=bad-whitespace
vxp, vxm = vx[2:, 1:-1, 1:-1], vx[:-2, 1:-1, 1:-1] # pylint: disable=bad-whitespace
vyp, vym = vy[1:-1, 2:, 1:-1], vy[1:-1, :-2, 1:-1] # pylint: disable=bad-whitespace
vzp, vzm = vz[1:-1, 1:-1, 2:], vz[1:-1, 1:-1, :-2] # pylint: disable=bad-whitespace
div_arr = ((vxp - vxm) / (xp - xm) + (vyp - vym) / (yp - ym) +
(vzp - vzm) / (zp - zm))
return field.wrap_field(div_arr,
divcrds,
name="div " + fld.name,
center=divcenter,
time=fld.time,
parents=[fld])
def curl(fld):
""" first order """
vx, vy, vz = fld.component_views()
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
curlcenter = "cell"
curlcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
curlcenter = "node"
curlcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
else:
raise NotImplementedError("Can only do cell and node centered divs")
xp = crdx[2:,:,:]; xm = crdx[:-2,:,:] #pylint: disable=W0612,C0321,C0324
yp = crdy[:,2:,:]; ym = crdy[:,:-2,:] #pylint: disable=W0612,C0321,C0324
zp = crdz[:,:,2:]; zm = crdz[:,:,:-2] #pylint: disable=W0612,C0321,C0324
# vxpx = vx[1:-1,1:-1,2:]; vxmx = vx[1:-1,1:-1,:-2] #pylint: disable=W0612,C0321,C0324,C0301
vxpy = vx[1:-1,2:,1:-1]; vxmy = vx[1:-1,:-2,1:-1] #pylint: disable=W0612,C0321,C0324,C0301
vxpz = vx[1:-1,1:-1,2:]; vxmz = vx[1:-1,1:-1,:-2] #pylint: disable=W0612,C0321,C0324,C0301
vypx = vy[2:,1:-1,1:-1]; vymx = vy[:-2,1:-1,1:-1] #pylint: disable=W0612,C0321,C0324,C0301
# vypy = vy[1:-1,2:,1:-1]; vymy = vy[1:-1,:-2,1:-1] #pylint: disable=W0612,C0321,C0324,C0301
vypz = vy[1:-1,1:-1,2:]; vymz = vy[1:-1,1:-1,:-2] #pylint: disable=W0612,C0321,C0324,C0301
vzpx = vz[2:,1:-1,1:-1]; vzmx = vz[:-2,1:-1,1:-1] #pylint: disable=W0612,C0321,C0324,C0301
vzpy = vz[1:-1,2:,1:-1]; vzmy = vz[1:-1,:-2,1:-1] #pylint: disable=W0612,C0321,C0324,C0301
# vzpz = vz[2:,1:-1,1:-1]; vzmz = vz[:-2,1:-1,1:-1] #pylint: disable=W0612,C0321,C0324,C0301
curl_x = ne.evaluate("(vzpy-vzmy)/(yp-ym) - (vypz-vymz)/(zp-zm)")
curl_y = ne.evaluate("-(vzpx-vzmx)/(xp-xm) + (vxpz-vxmz)/(zp-zm)")
curl_z = ne.evaluate("(vypx-vymx)/(xp-xm) - (vxpy-vxmy)/(yp-ym)")
return field.wrap_field([curl_x, curl_y, curl_z], curlcrds,
name="curl " + fld.name, fldtype="Vector",
center=curlcenter, time=fld.time,
parents=[fld])
def div(fld, bnd=True):
"""2nd order centeral diff, 1st order @ boundaries if bnd"""
if fld.iscentered("Face"):
# dispatch fc div immediately since that does its own pre-processing
return viscid.div_fc(fld, bnd=bnd)
if bnd:
fld = viscid.extend_boundaries(fld, order=0, crd_order=0)
vx, vy, vz = fld.component_views()
if fld.iscentered("Cell"):
crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
divcenter = "Cell"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
elif fld.iscentered("Node"):
crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
divcenter = "Node"
# divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
else:
raise NotImplementedError("Can only do cell and node centered divs")
xp, xm = crdx[2:, :, :], crdx[:-2, : , : ] # pylint: disable=bad-whitespace
yp, ym = crdy[ :, 2:, :], crdy[: , :-2, : ] # pylint: disable=bad-whitespace
zp, zm = crdz[ :, :, 2:], crdz[: , : , :-2] # pylint: disable=bad-whitespace
vxp, vxm = vx[2: , 1:-1, 1:-1], vx[ :-2, 1:-1, 1:-1] # pylint: disable=bad-whitespace
vyp, vym = vy[1:-1, 2: , 1:-1], vy[1:-1, :-2, 1:-1] # pylint: disable=bad-whitespace
vzp, vzm = vz[1:-1, 1:-1, 2: ], vz[1:-1, 1:-1, :-2] # pylint: disable=bad-whitespace
div_arr = ne.evaluate("(vxp-vxm)/(xp-xm) + (vyp-vym)/(yp-ym) + "
"(vzp-vzm)/(zp-zm)")
return field.wrap_field(div_arr, divcrds, name="div " + fld.name,
center=divcenter, time=fld.time, parents=[fld])
def calc_psi(B, reversed=False):
"""Calc Flux function (only valid in 2d)
Parameters:
B (VectorField): magnetic field, should only have two
spatial dimensions so we can infer the symmetry dimension
reversed (bool): since this integration doesn't like going
through undefined regions (like within 1 earth radius of
the origin for openggcm), you can use this to start
integrating from the opposite corner.
Returns:
ScalarField: 2-D scalar flux function
Raises:
ValueError: If B has <> 2 spatial dimensions
"""
# TODO: if this is painfully slow, i bet just putting this exact
# code in a cython module would make it a bunch faster, the problem
# being that the loops are in python instead of some broadcasting
# numpy type thing
B = B.slice_reduce(":")
if B.nr_sdims != 2:
raise ValueError("flux function implemented for 2D fields")
comps = ""
for comp in "xyz":
if comp in B.crds.axes:
comps += comp
# ex: comps = "yz", comp_inds = [1, 2]
comp_inds = [dict(x=0, y=1, z=2)[comp] for comp in comps]
# Note: what follows says y, z, but it has been generalized
# to any two directions, so hy isn't necessarily hy, but it's
# easier to see at a glance if it's correct using a specific
# example
ycc, zcc = B.get_crds(comps)
comp_views = B.component_views()
hy, hz = comp_views[comp_inds[0]], comp_views[comp_inds[1]]
dy = ycc[1:] - ycc[:-1]
dz = zcc[1:] - zcc[:-1]
ny, nz = len(ycc), len(zcc)
A = np.empty((ny, nz), dtype=B.dtype)
if reversed:
A[-1, -1] = 0.0
for i in range(ny - 2, -1, -1):
A[i, -1] = A[i + 1, -1] - dy[i] * 0.5 * (hz[i, -1] + hz[i + 1, -1])
for j in range(nz - 2, -1, -1):
A[:, j] = A[:, j + 1] + dz[j] * 0.5 * (hy[:, j + 1] + hy[:, j])
else:
A[0, 0] = 0.0
for i in range(1, ny):
A[i, 0] = A[i - 1, 0] + dy[i - 1] * 0.5 * (hz[i, 0] + hz[i - 1, 0])
for j in range(1, nz):
A[:, j] = A[:, j - 1] - dz[j - 1] * 0.5 * (hy[:, j - 1] + hy[:, j])
return field.wrap_field(A, B.crds, name="psi", center=B.center,
pretty_name=r"$\psi$", parents=[B])
def calc_psi(B, rev=False):
"""Calc Flux function (only valid in 2d)
Parameters:
B (VectorField): magnetic field, should only have two
spatial dimensions so we can infer the symmetry dimension
rev (bool): since this integration doesn't like going
through undefined regions (like within 1 earth radius of
the origin for openggcm), you can use this to start
integrating from the opposite corner.
Returns:
ScalarField: 2-D scalar flux function
Raises:
ValueError: If B has <> 2 spatial dimensions
"""
# TODO: if this is painfully slow, i bet just putting this exact
# code in a cython module would make it a bunch faster, the problem
# being that the loops are in python instead of some broadcasting
# numpy type thing
B = B.slice_reduce(":")
# try to guess if a dim of a 3D field is invariant
reduced_axes = []
if B.nr_sdims > 2:
slcs = [slice(None)] * B.nr_sdims
for i, nxi in enumerate(B.sshape):
if nxi <= 2:
slcs[i] = 0
reduced_axes.append(B.crds.axes[i])
slcs.insert(B.nr_comp, slice(None))
B = B[slcs]
# ok, so the above didn't work... just nip out the smallest dim?
if B.nr_sdims == 3:
slcs = [slice(None)] * B.nr_sdims
i = np.argmin(B.sshape)
slcs[i] = 0
reduced_axes.append(B.crds.axes[i])
logger.warning("Tried to get the flux function of a 3D field. "
"I can't do that, so I'm\njust ignoring the {0} "
"dimension".format(reduced_axes[-1]))
slcs.insert(B.nr_comp, slice(None))
B = B[slcs]
if B.nr_sdims != 2:
raise ValueError("flux function only implemented for 2D fields")
comps = ""
for comp in "xyz":
if comp in B.crds.axes:
comps += comp
# ex: comps = "yz", comp_inds = [1, 2]
comp_inds = [dict(x=0, y=1, z=2)[comp] for comp in comps]
# Note: what follows says y, z, but it has been generalized
# to any two directions, so hy isn't necessarily hy, but it's
# easier to see at a glance if it's correct using a specific
# example
ycc, zcc = B.get_crds(comps)
comp_views = B.component_views()
hy, hz = comp_views[comp_inds[0]], comp_views[comp_inds[1]]
dy = ycc[1:] - ycc[:-1]
dz = zcc[1:] - zcc[:-1]
ny, nz = len(ycc), len(zcc)
A = np.empty((ny, nz), dtype=B.dtype)
if rev:
A[-1, -1] = 0.0
for i in range(ny - 2, -1, -1):
A[i, -1] = A[i + 1, -1] - dy[i] * 0.5 * (hz[i, -1] + hz[i + 1, -1])
for j in range(nz - 2, -1, -1):
A[:, j] = A[:, j + 1] + dz[j] * 0.5 * (hy[:, j + 1] + hy[:, j])
else:
A[0, 0] = 0.0
for i in range(1, ny):
A[i, 0] = A[i - 1, 0] + dy[i - 1] * 0.5 * (hz[i, 0] + hz[i - 1, 0])
for j in range(1, nz):
A[:, j] = A[:, j - 1] - dz[j - 1] * 0.5 * (hy[:, j - 1] + hy[:, j])
psi = field.wrap_field(A, B.crds, name="psi", center=B.center,
pretty_name=r"$\psi$", parents=[B])
if reduced_axes:
slc = "..., " + ", ".join("{0}=None".format(ax) for ax in reduced_axes)
psi = psi[slc]
return psi
# globals(), locals(), "interp.prof")
# plane = seed.Plane((1., 1., 1.), (1., 1., 1.), (1., 0., 0.),
# 1.0, 1.0, 500, 500)
vol = B.crds
# print(plane.get_points())
cProfile.runctx("""interp_vals = cycalc.interp_trilin(B, vol)""",
globals(), locals(), "interp.prof")
t1 = time()
print("Total time: {0:.3e}".format(t1 - t0))
s = pstats.Stats("interp.prof")
s.strip_dirs().sort_stats("cumtime").print_stats()
# print([line.shape for line in lines])
# mpl.scatter_3d(vol.get_points(), interp_vals[:, 2], show=True)
interp_field = field.wrap_field(interp_vals, vol.as_coordinates(), # pylint: disable=undefined-variable
name="interp", fldtype="vector",
center="cell")
interp_y1 = interp_field["y=1"]
exact_y1 = B["y=1"]
bxi, byi, bzi = interp_y1.component_fields()
bxe, bye, bze = exact_y1.component_fields()
for interp, exact in zip([bxi, byi, bzi], [bxe, bye, bze]):
plt.clf()
plt.subplot(211)
mpl.plot(exact, show=False)
plt.subplot(212)
mpl.plot(interp, show=True)
##
## EOF
##