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 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