def div_fc(fc, force_numpy=False, bnd=True):
"""Calculate cell centered divergence of face centered field"""
fc = fc.atleast_3d()
fc = make_ecfc_field_leading(fc)
# FIXME: maybe it's possible to do the boundary correctly here without
# just faking it with a 0 order hold before returning
s0, sm, sp = _prep_slices(fc)
s0vec = list(s0)
s0vec.insert(fc.nr_comp, slice(0, 1))
div_cc = viscid.zeros_like(fc[s0vec], center="cell")
x, y, z = fc.get_crds_nc('xyz', shaped=True)
if True:
# x, y, z = x[1:, :, :], y[:, 1:, :], z[:, :, 1:]
x, y, z = x[:-1, :, :], y[:, :-1, :], z[:, :, :-1]
else:
raise NotImplementedError()
# x, y, z = fc.get_crds_cc('xyz', shaped=True)
xm, xp = x[sm[0], :, :], x[sp[0], :, :]
ym, yp = y[:, sm[1], :], y[:, sp[1], :]
zm, zp = z[:, :, sm[2]], z[:, :, sp[2]]
fcx0, fcx1 = (fc['x', sm[0], s0[1], s0[2]].data, fc['x', sp[0], s0[1],
s0[2]].data)
fcy0, fcy1 = (fc['y', s0[0], sm[1], s0[2]].data, fc['y', s0[0], sp[1],
s0[2]].data)
fcz0, fcz1 = (fc['z', s0[0], s0[1], sm[2]].data, fc['z', s0[0], s0[1],
sp[2]].data)
# xp, yp, zp = xm + 1.0, ym + 1.0, zm + 1.0
if _HAS_NUMEXPR and not force_numpy:
div_cc[:, :, :] = ne.evaluate("((fcx1 - fcx0) / (xp - xm)) + "
"((fcy1 - fcy0) / (yp - ym)) + "
"((fcz1 - fcz0) / (zp - zm))")
else:
div_cc[:, :, :] = (((fcx1 - fcx0) / (xp - xm)) +
((fcy1 - fcy0) / (yp - ym)) + ((fcz1 - fcz0) /
(zp - zm)))
if bnd:
# FIXME: this is really just faking the bnd so there aren't shape
# errors when doing math with the result
div_cc = viscid.extend_boundaries(div_cc,
nl=1,
nh=1,
order=0,
crd_order=1)
div_cc.name = "div " + fc.name
div_cc.pretty_name = "Div " + fc.pretty_name
return div_cc
def ec2cc(ec, force_numpy=False, bnd=True):
"""Average an edge centered field to cell centers"""
ec = ec.atleast_3d()
ec = make_ecfc_field_leading(ec)
s0, sm, sp = _prep_slices(ec)
s0vec = list(s0)
s0vec.insert(ec.nr_comp, slice(None))
cc = viscid.zeros_like(ec[s0vec])
cc.center = "Cell"
ecx0, ecx1, ecx2, ecx3 = (ec['x', s0[0], sm[1],
sm[2]].data, ec['x', s0[0], sm[1],
sp[2]].data,
ec['x', s0[0], sp[1],
sm[2]].data, ec['x', s0[0], sp[1],
sp[2]].data)
ecy0, ecy1, ecy2, ecy3 = (ec['y', sm[0], s0[1],
sm[2]].data, ec['y', sm[0], s0[1],
sp[2]].data,
ec['y', sp[0], s0[1],
sm[2]].data, ec['y', sp[0], s0[1],
sp[2]].data)
ecz0, ecz1, ecz2, ecz3 = (ec['z', sm[0], sm[1],
s0[2]].data, ec['z', sm[0], sp[1],
s0[2]].data,
ec['z', sp[0], sm[1],
s0[2]].data, ec['z', sp[0], sp[1],
s0[2]].data)
if _HAS_NUMEXPR and not force_numpy:
quarter = np.array([0.25], dtype=ec.dtype)[0] # pylint: disable=unused-variable
s = "quarter * (a + b + c + d)"
a, b, c, d = ecx0, ecx1, ecx2, ecx3 # pylint: disable=unused-variable
cc['x'] = ne.evaluate(s)
a, b, c, d = ecy0, ecy1, ecy2, ecy3
cc['y'] = ne.evaluate(s)
a, b, c, d = ecz0, ecz1, ecz2, ecz3
cc['z'] = ne.evaluate(s)
else:
cc['x'] = 0.25 * (ecx0 + ecx1 + ecx2 + ecx3)
cc['y'] = 0.25 * (ecy0 + ecy1 + ecy2 + ecy3)
cc['z'] = 0.25 * (ecz0 + ecz1 + ecz2 + ecz3)
if bnd:
# FIXME: this is really just faking the bnd so there aren't shape
# errors when doing math with the result
cc = viscid.extend_boundaries(cc, nl=1, nh=1, order=0, crd_order=1)
cc.name = ec.name
cc.pretty_name = ec.pretty_name
return cc
def div_fc(fc, force_numpy=False, bnd=True):
"""Calculate cell centered divergence of face centered field"""
fc = fc.atleast_3d()
fc = make_ecfc_field_leading(fc)
# FIXME: maybe it's possible to do the boundary correctly here without
# just faking it with a 0 order hold before returning
s0, sm, sp = _prep_slices(fc)
s0vec = list(s0)
s0vec.insert(fc.nr_comp, slice(0, 1))
div_cc = viscid.zeros_like(fc[s0vec], center="cell")
x, y, z = fc.get_crds_nc('xyz', shaped=True)
if True:
# x, y, z = x[1:, :, :], y[:, 1:, :], z[:, :, 1:]
x, y, z = x[:-1, :, :], y[:, :-1, :], z[:, :, :-1]
else:
raise NotImplementedError()
# x, y, z = fc.get_crds_cc('xyz', shaped=True)
xm, xp = x[sm[0], :, :], x[sp[0], :, :]
ym, yp = y[:, sm[1], :], y[:, sp[1], :]
zm, zp = z[:, :, sm[2]], z[:, :, sp[2]]
fcx0, fcx1 = (fc['x', sm[0], s0[1], s0[2]].data,
fc['x', sp[0], s0[1], s0[2]].data)
fcy0, fcy1 = (fc['y', s0[0], sm[1], s0[2]].data,
fc['y', s0[0], sp[1], s0[2]].data)
fcz0, fcz1 = (fc['z', s0[0], s0[1], sm[2]].data,
fc['z', s0[0], s0[1], sp[2]].data)
# xp, yp, zp = xm + 1.0, ym + 1.0, zm + 1.0
if _HAS_NUMEXPR and not force_numpy:
div_cc[:, :, :] = ne.evaluate("((fcx1 - fcx0) / (xp - xm)) + "
"((fcy1 - fcy0) / (yp - ym)) + "
"((fcz1 - fcz0) / (zp - zm))")
else:
div_cc[:, :, :] = (((fcx1 - fcx0) / (xp - xm)) +
((fcy1 - fcy0) / (yp - ym)) +
((fcz1 - fcz0) / (zp - zm)))
if bnd:
# FIXME: this is really just faking the bnd so there aren't shape
# errors when doing math with the result
div_cc = viscid.extend_boundaries(div_cc, nl=1, nh=1, order=0,
crd_order=1)
div_cc.name = "div " + fc.name
div_cc.pretty_name = "Div " + fc.pretty_name
return div_cc
def ec2cc(ec, force_numpy=False, bnd=True):
"""Average an edge centered field to cell centers"""
ec = ec.atleast_3d()
ec = make_ecfc_field_leading(ec)
s0, sm, sp = _prep_slices(ec)
s0vec = list(s0)
s0vec.insert(ec.nr_comp, slice(None))
cc = viscid.zeros_like(ec[s0vec])
cc.center = "Cell"
ecx0, ecx1, ecx2, ecx3 = (ec['x', s0[0], sm[1], sm[2]].data,
ec['x', s0[0], sm[1], sp[2]].data,
ec['x', s0[0], sp[1], sm[2]].data,
ec['x', s0[0], sp[1], sp[2]].data)
ecy0, ecy1, ecy2, ecy3 = (ec['y', sm[0], s0[1], sm[2]].data,
ec['y', sm[0], s0[1], sp[2]].data,
ec['y', sp[0], s0[1], sm[2]].data,
ec['y', sp[0], s0[1], sp[2]].data)
ecz0, ecz1, ecz2, ecz3 = (ec['z', sm[0], sm[1], s0[2]].data,
ec['z', sm[0], sp[1], s0[2]].data,
ec['z', sp[0], sm[1], s0[2]].data,
ec['z', sp[0], sp[1], s0[2]].data)
if _HAS_NUMEXPR and not force_numpy:
quarter = np.array([0.25], dtype=ec.dtype)[0] # pylint: disable=unused-variable
s = "quarter * (a + b + c + d)"
a, b, c, d = ecx0, ecx1, ecx2, ecx3 # pylint: disable=unused-variable
cc['x'] = ne.evaluate(s)
a, b, c, d = ecy0, ecy1, ecy2, ecy3
cc['y'] = ne.evaluate(s)
a, b, c, d = ecz0, ecz1, ecz2, ecz3
cc['z'] = ne.evaluate(s)
else:
cc['x'] = 0.25 * (ecx0 + ecx1 + ecx2 + ecx3)
cc['y'] = 0.25 * (ecy0 + ecy1 + ecy2 + ecy3)
cc['z'] = 0.25 * (ecz0 + ecz1 + ecz2 + ecz3)
if bnd:
# FIXME: this is really just faking the bnd so there aren't shape
# errors when doing math with the result
cc = viscid.extend_boundaries(cc, nl=1, nh=1, order=0, crd_order=1)
cc.name = ec.name
cc.pretty_name = ec.pretty_name
return cc
def make_arcade(eps, xl=(-10.0, 0.0, -10.0), xh=(10.0, 20.0, 10.0),
L=(5, 5, 5), N=(32, 32, 32), layout='interlaced'):
xl, xh = np.asarray(xl), np.asarray(xh)
x = np.linspace(xl[0], xh[0], N[0])
y = np.linspace(xl[1], xh[1], N[1])
z = np.linspace(xl[2], xh[2], N[2])
b = viscid.zeros([x, y, z], nr_comps=3, layout=layout)
e = viscid.zeros_like(b)
X, Y, Z = b.get_crds('xyz', shaped=True)
Y2 = Y**2 / L[1]**2
Z2 = Z**2 / L[2]**2
b['x'] = -1 - (eps * ((1 - Y2) / (1 + Y2)) * (1 / (1 + Z2)))
b['y'] = X
b['z'] = 0.2
e['z'] = Y / ((1 + Y2) * (1 + Z2))
return b, e
def fc2cc(fc, force_numpy=False, bnd=True):
"""Average a face centered field to cell centers"""
fc = fc.atleast_3d()
fc = make_ecfc_field_leading(fc)
s0, sm, sp = _prep_slices(fc)
s0vec = list(s0)
s0vec.insert(fc.nr_comp, slice(None))
cc = viscid.zeros_like(fc[s0vec])
cc.center = "Cell"
fcx0, fcx1 = (fc['x', sm[0], s0[1], s0[2]].data,
fc['x', sp[0], s0[1], s0[2]].data)
fcy0, fcy1 = (fc['y', s0[0], sm[1], s0[2]].data,
fc['y', s0[0], sp[1], s0[2]].data)
fcz0, fcz1 = (fc['z', s0[0], s0[1], sm[2]].data,
fc['z', s0[0], s0[1], sp[2]].data)
if _HAS_NUMEXPR and not force_numpy:
half = np.array([0.5], dtype=fc.dtype)[0] # pylint: disable=unused-variable
s = "half * (a + b)"
a, b = fcx0, fcx1 # pylint: disable=unused-variable
cc['x'] = ne.evaluate(s)
a, b = fcy0, fcy1
cc['y'] = ne.evaluate(s)
a, b = fcz0, fcz1
cc['z'] = ne.evaluate(s)
else:
cc['x'] = 0.5 * (fcx0 + fcx1)
cc['y'] = 0.5 * (fcy0 + fcy1)
cc['z'] = 0.5 * (fcz0 + fcz1)
if bnd:
# FIXME: this is really just faking the bnd so there aren't shape
# errors when doing math with the result
cc = viscid.extend_boundaries(cc, nl=1, nh=1, order=0, crd_order=1)
cc.name = fc.name
cc.pretty_name = fc.pretty_name
return cc
def fc2cc(fc, force_numpy=False, bnd=True):
"""Average a face centered field to cell centers"""
fc = fc.atleast_3d()
fc = make_ecfc_field_leading(fc)
s0, sm, sp = _prep_slices(fc)
s0vec = list(s0)
s0vec.insert(fc.nr_comp, slice(None))
cc = viscid.zeros_like(fc[s0vec])
cc.center = "Cell"
fcx0, fcx1 = (fc['x', sm[0], s0[1], s0[2]].data, fc['x', sp[0], s0[1],
s0[2]].data)
fcy0, fcy1 = (fc['y', s0[0], sm[1], s0[2]].data, fc['y', s0[0], sp[1],
s0[2]].data)
fcz0, fcz1 = (fc['z', s0[0], s0[1], sm[2]].data, fc['z', s0[0], s0[1],
sp[2]].data)
if _HAS_NUMEXPR and not force_numpy:
half = np.array([0.5], dtype=fc.dtype)[0] # pylint: disable=unused-variable
s = "half * (a + b)"
a, b = fcx0, fcx1 # pylint: disable=unused-variable
cc['x'] = ne.evaluate(s)
a, b = fcy0, fcy1
cc['y'] = ne.evaluate(s)
a, b = fcz0, fcz1
cc['z'] = ne.evaluate(s)
else:
cc['x'] = 0.5 * (fcx0 + fcx1)
cc['y'] = 0.5 * (fcy0 + fcy1)
cc['z'] = 0.5 * (fcz0 + fcz1)
if bnd:
# FIXME: this is really just faking the bnd so there aren't shape
# errors when doing math with the result
cc = viscid.extend_boundaries(cc, nl=1, nh=1, order=0, crd_order=1)
cc.name = fc.name
cc.pretty_name = fc.pretty_name
return cc
def make_arcade(eps,
xl=(-10.0, 0.0, -10.0),
xh=(10.0, 20.0, 10.0),
L=(5, 5, 5),
N=(32, 32, 32),
layout='interlaced'):
xl, xh = np.asarray(xl), np.asarray(xh)
x = np.linspace(xl[0], xh[0], N[0])
y = np.linspace(xl[1], xh[1], N[1])
z = np.linspace(xl[2], xh[2], N[2])
b = viscid.zeros([x, y, z], nr_comps=3, layout=layout)
e = viscid.zeros_like(b)
X, Y, Z = b.get_crds('xyz', shaped=True)
Y2 = Y**2 / L[1]**2
Z2 = Z**2 / L[2]**2
b['x'] = -1 - (eps * ((1 - Y2) / (1 + Y2)) * (1 / (1 + Z2)))
b['y'] = X
b['z'] = 0.2
e['z'] = Y / ((1 + Y2) * (1 + Z2))
return b, e
def make_ecfc_field_leading(fc, trim_leading=True):
"""Standardize staggering on edge / face centered fields"""
if fc.find_info(PREPROCESSED_KEY, False):
return fc
elif fc.find_info(PREPROCESSED_KEY, None) is None:
# this is for non-ggcm fields
fc.set_info(PREPROCESSED_KEY, True)
fc.set_info(STAGGER_KEY, [STAGGER_LEADING] * 3)
return fc
else:
# NOTE: I deeply apologize for the variable names and logic here,
# it turns out to be super confusing to deal with staggering and
# flipping of 2 components (mhd->gse), etc. The important thing is
# this logic was tested to work with step_c, step_c3, and jrrle
# output in both mhd and gse crds
center = fc.center.lower()
if center not in ('face', 'edge'):
raise ValueError("This only makes sense for Edge / Face centered "
"fields")
SA = [slice(None, -1), slice(1, None), slice(None, None)]
SB = [SA[1], SA[0], SA[2]]
# Trailing offsets/Staggering
_sd0T = 1
_sdT = 0
# Leading Offsets/Staggering
if trim_leading:
# this makes everything symmetric since trailing offset fields
# will be sliced by one value, so you can slice the first and
# last values from a CC field and add it to an fc2cc field
# regardless of the fc field's representation.
_sd0L = 1
_sdL = 1
else:
_sd0L = 2
_sdL = 2
# sd: how to slice the data for vector component xyz
# sc: how to slice the xyz NC coordinates
sd = [[None] * 3, [None] * 3, [None] * 3]
sc = [None] * 3
# staggering = fc.find_info(STAGGER_KEY, [STAGGER_TRAILING] * 3)
staggering = fc.find_info(STAGGER_KEY, [STAGGER_LEADING] * 3)
flipping = fc.find_info(FLIP_KEY, [False] * 3)
for i, n in enumerate(fc.sshape):
if staggering[i] == STAGGER_TRAILING:
_sd0, _sd = _sd0T, _sdT
else:
_sd0, _sd = _sd0L, _sdL
if n == 1:
sc[i] = slice(None, None)
for j in range(3):
sd[i][j] = slice(None, None)
elif n > 1:
sc[i] = slice(1, None)
for j in range(3):
cnd = i == j if center == 'face' else i != j
S = SB if flipping[j] and cnd else SA
if cnd:
sd[i][j] = S[_sd]
else:
sd[i][j] = S[_sd0]
else:
raise RuntimeError("I shouldn't be here")
fc.resolve() # if translating -> gse, do it once, not 4 times
s0vec = list(sc)
s0vec.insert(fc.nr_comp, slice(None))
prepared_fc = viscid.zeros_like(fc[s0vec])
prepared_fc['x'] = fc['x', sd[0][0], sd[0][1], sd[0][2]]
prepared_fc['y'] = fc['y', sd[1][0], sd[1][1], sd[1][2]]
prepared_fc['z'] = fc['z', sd[2][0], sd[2][1], sd[2][2]]
prepared_fc.set_info(PREPROCESSED_KEY, True)
prepared_fc.set_info(STAGGER_KEY, [STAGGER_LEADING] * 3)
prepared_fc.name = fc.name
prepared_fc.pretty_name = fc.pretty_name
return prepared_fc
def _main():
f = viscid.load_file('~/dev/work/xi_fte_001/*.3d.*.xdmf')
time_slice = ':'
times = np.array([grid.time for grid in f.iter_times(time_slice)])
# XYZ coordinates of virtual satelites in warped "plasma sheet coords"
x_sat_psc = np.linspace(-30, 0, 31) # X (GSE == PSC)
y_sat_psc = np.linspace(-10, 10, 21) # Y (GSE == PSC)
z_sat_psc = np.linspace(-2, 2, 5) # Z in PSC (z=0 is the plasma sheet)
# the GSE z location of the virtual satelites in the warped plasma sheet
# coordinates, so sat_z_gse_ts['x=5j, y=1j, z=0j'] would give the
# plasma sheet location at x=5.0, y=1.0
# These fields depend on time because the plasma sheet moves in time
sat_z_gse_ts = viscid.zeros([times, x_sat_psc, y_sat_psc, z_sat_psc],
crd_names='txyz', center='node',
name='PlasmaSheetZ_GSE')
vx_ts = viscid.zeros_like(sat_z_gse_ts)
bz_ts = viscid.zeros_like(sat_z_gse_ts)
for itime, grid in enumerate(f.iter_times(time_slice)):
print("Processing time slice", itime, grid.time)
gse_slice = 'x=-35j:0j, y=-15j:15j, z=-6j:6j'
bx = grid['bx'][gse_slice]
bx_argmin = np.argmin(bx**2, axis=2)
z_gse = bx.get_crd('z')
# ps_zloc_gse is the plasma sheet z location along the GGCM grid x/y
ps_z_gse = viscid.zeros_like(bx[:, :, 0:1])
ps_z_gse[...] = z_gse[bx_argmin]
# Note: Here you could apply a gaussian filter to
# ps_z_gse[:, :, 0].data in order to smooth the surface
# if desired. Scipy / Scikit-Image have some functions
# that do this
# ok, we found the plasma sheet z GSE location on the actual GGCM
# grid, but we just want a subset of that grid for our virtual
# satelites, so just interpolate the ps z location to our subset
ps_z_gse_subset = viscid.interp_trilin(ps_z_gse,
sat_z_gse_ts[itime, :, :, 0:1],
wrap=True)
# now we know the plasma sheet z location in GSE, and how far
# apart we want the satelites in z, so put those two things together
# to get a bunch of satelite locations
sat_z_gse_ts[itime] = ps_z_gse_subset.data + z_sat_psc.reshape(1, 1, -1)
# make a seed generator that we can use to fill the vx and bz
# time series for this instant in time
sat_loc_gse = sat_z_gse_ts[itime].get_points()
sat_loc_gse[2, :] = sat_z_gse_ts[itime].data.reshape(-1)
# slicing the field before doing the interpolation makes this
# faster for hdf5 data, but probably for other data too
vx_ts[itime] = viscid.interp_trilin(grid['vx'][gse_slice],
sat_loc_gse,
wrap=False
).reshape(vx_ts.shape[1:])
bz_ts[itime] = viscid.interp_trilin(grid['bz'][gse_slice],
sat_loc_gse,
wrap=False
).reshape(bz_ts.shape[1:])
# 2d plots of the plasma sheet z location to make sure we did the
# interpolation correctly
if False: # pylint: disable=using-constant-test
from viscid.plot import vpyplot as vlt
fig, (ax0, ax1) = vlt.subplots(2, 1) # pylint: disable=unused-variable
vlt.plot(ps_z_gse, ax=ax0, clim=(-5, 5))
vlt.plot(ps_z_gse_subset, ax=ax1, clim=(-5, 5))
vlt.auto_adjust_subplots()
vlt.show()
# make a 3d plot of the plasma sheet surface to verify that it
# makes sense
if True: # pylint: disable=using-constant-test
from viscid.plot import vlab
fig = vlab.figure(size=(1280, 800), bgcolor=(1, 1, 1),
fgcolor=(0, 0, 0))
vlab.clf()
# plot the plasma sheet coloured by vx
# Note: points closer to x = 0 are unsightly since the plasma
# sheet criteria starts to fall apart on the flanks, so
# just remove the first few rows
ps_z_gse_tail = ps_z_gse['x=:-2.25j']
ps_mesh_shape = [3, ps_z_gse_tail.shape[0], ps_z_gse_tail.shape[1]]
ps_pts = ps_z_gse_tail.get_points().reshape(ps_mesh_shape)
ps_pts[2, :, :] = ps_z_gse_tail[:, :, 0]
plasma_sheet = viscid.RectilinearMeshPoints(ps_pts)
ps_vx = viscid.interp_trilin(grid['vx'][gse_slice], plasma_sheet)
_ = vlab.mesh_from_seeds(plasma_sheet, scalars=ps_vx)
vx_clim = (-1400, 1400)
vx_cmap = 'viridis'
vlab.colorbar(title='Vx', clim=vx_clim, cmap=vx_cmap,
nb_labels=5)
# plot satelite locations as dots colored by Vx with the same
# limits and color as the plasma sheet mesh
sat3d = vlab.points3d(sat_loc_gse[0], sat_loc_gse[1], sat_loc_gse[2],
vx_ts[itime].data.reshape(-1),
scale_mode='none', scale_factor=0.2)
vlab.apply_cmap(sat3d, clim=vx_clim, cmap=vx_cmap)
# plot Earth for reference
cotr = viscid.Cotr(dip_tilt=0.0) # pylint: disable=not-callable
vlab.plot_blue_marble(r=1.0, lines=False, ntheta=64, nphi=128,
rotate=cotr, crd_system='mhd')
vlab.plot_earth_3d(radius=1.01, night_only=True, opacity=0.5,
crd_system='gse')
vlab.view(azimuth=45, elevation=70, distance=35.0,
focalpoint=[-9, 3, -1])
vlab.savefig('plasma_sheet_3d_{0:02d}.png'.format(itime))
vlab.show()
try:
vlab.mlab.close(fig)
except TypeError:
pass # this happens if the figure is already closed
# now do what we will with the time series... this is not a good
# presentation of this data, but you get the idea
from viscid.plot import vpyplot as vlt
fig, axes = vlt.subplots(4, 4, figsize=(12, 12))
for ax_row, yloc in zip(axes, np.linspace(-5, 5, len(axes))[::-1]):
for ax, xloc in zip(ax_row, np.linspace(4, 7, len(ax_row))):
vlt.plot(vx_ts['x={0}j, y={1}j, z=0j'.format(xloc, yloc)], ax=ax)
ax.set_ylabel('')
vlt.plt.title('x = {0:g}, y = {1:g}'.format(xloc, yloc))
vlt.plt.suptitle('Vx [km/s]')
vlt.auto_adjust_subplots()
vlt.show()
return 0
def make_ecfc_field_leading(fc, trim_leading=True):
"""Standardize staggering on edge / face centered fields"""
if fc.find_info(PREPROCESSED_KEY, False):
return fc
elif fc.find_info(PREPROCESSED_KEY, None) is None:
# this is for non-ggcm fields
fc.set_info(PREPROCESSED_KEY, True)
fc.set_info(STAGGER_KEY, [STAGGER_LEADING] * 3)
return fc
else:
# NOTE: I deeply apologize for the variable names and logic here,
# it turns out to be super confusing to deal with staggering and
# flipping of 2 components (mhd->gse), etc. The important thing is
# this logic was tested to work with step_c, step_c3, and jrrle
# output in both mhd and gse crds
center = fc.center.lower()
if center not in ('face', 'edge'):
raise ValueError("This only makes sense for Edge / Face centered "
"fields")
SA = [slice(None, -1), slice(1, None), slice(None, None)]
SB = [SA[1], SA[0], SA[2]]
# Trailing offsets/Staggering
_sd0T = 1
_sdT = 0
# Leading Offsets/Staggering
if trim_leading:
# this makes everything symmetric since trailing offset fields
# will be sliced by one value, so you can slice the first and
# last values from a CC field and add it to an fc2cc field
# regardless of the fc field's representation.
_sd0L = 1
_sdL = 1
else:
_sd0L = 2
_sdL = 2
# sd: how to slice the data for vector component xyz
# sc: how to slice the xyz NC coordinates
sd = [[None] * 3, [None] * 3, [None] * 3]
sc = [None] * 3
# staggering = fc.find_info(STAGGER_KEY, [STAGGER_TRAILING] * 3)
staggering = fc.find_info(STAGGER_KEY, [STAGGER_LEADING] * 3)
flipping = fc.find_info(FLIP_KEY, [False] * 3)
for i, n in enumerate(fc.sshape):
if staggering[i] == STAGGER_TRAILING:
_sd0, _sd = _sd0T, _sdT
else:
_sd0, _sd = _sd0L, _sdL
if n == 1:
sc[i] = slice(None, None)
for j in range(3):
sd[i][j] = slice(None, None)
elif n > 1:
sc[i] = slice(1, None)
for j in range(3):
cnd = i == j if center == 'face' else i != j
S = SB if flipping[j] and cnd else SA
if cnd:
sd[i][j] = S[_sd]
else:
sd[i][j] = S[_sd0]
else:
raise RuntimeError("I shouldn't be here")
fc.resolve() # if translating -> gse, do it once, not 4 times
s0vec = list(sc)
s0vec.insert(fc.nr_comp, slice(None))
prepared_fc = viscid.zeros_like(fc[s0vec])
prepared_fc['x'] = fc['x', sd[0][0], sd[0][1], sd[0][2]]
prepared_fc['y'] = fc['y', sd[1][0], sd[1][1], sd[1][2]]
prepared_fc['z'] = fc['z', sd[2][0], sd[2][1], sd[2][2]]
prepared_fc.set_info(PREPROCESSED_KEY, True)
prepared_fc.set_info(STAGGER_KEY, [STAGGER_LEADING] * 3)
prepared_fc.name = fc.name
prepared_fc.pretty_name = fc.pretty_name
return prepared_fc
