def plot_earth_3d(figure=None, daycol=(1, 1, 1), nightcol=(0, 0, 0),
radius=1.0, res=24, crd_system="gse", night_only=False,
**kwargs):
"""Plot a black and white sphere (Earth) showing sunward direction
Parameters:
figure (mayavi.core.scene.Scene): specific figure, or None for
:py:func:`mayavi.mlab.gcf`
daycol (tuple, optional): color of dayside (RGB)
nightcol (tuple, optional): color of nightside (RGB)
res (optional): rosolution of teh sphere
crd_system (str, other): One of ('mhd', 'gse'), or anything
that returns from :py:func:`viscid.get_crd_system`.
Returns:
Tuple (day, night) as vtk sources
"""
if figure is None:
figure = mlab.gcf()
crd_system = viscid.get_crd_system(crd_system)
if crd_system == "mhd":
theta_dusk, theta_dawn = 270, 90
elif crd_system == "gse":
theta_dusk, theta_dawn = 90, 270
else:
# use GSE convention?
theta_dusk, theta_dawn = 90, 270
night = BuiltinSurface(source='sphere', name='night')
night.data_source.set(center=(0, 0, 0), radius=radius,
start_theta=theta_dusk, end_theta=theta_dawn,
theta_resolution=res, phi_resolution=res)
mod = mlab.pipeline.surface(night, color=nightcol, figure=figure, **kwargs)
mod.actor.property.backface_culling = True
if not night_only:
day = BuiltinSurface(source='sphere', name='day')
day.data_source.set(center=(0, 0, 0), radius=radius,
start_theta=theta_dawn, end_theta=theta_dusk,
theta_resolution=res, phi_resolution=res)
mod = mlab.pipeline.surface(day, color=daycol, figure=figure, **kwargs)
mod.actor.property.backface_culling = True
else:
day = None
return day, night
def make_dipole(m=(0, 0, -DEFAULT_STRENGTH), strength=None, l=None, h=None,
n=None, twod=False, dtype='f8', nonuniform=False,
crd_system='gse', name='b'):
"""Generate a dipole field with magnetic moment m [x, y, z]"""
if l is None:
l = [-5] * 3
if h is None:
h = [5] * 3
if n is None:
n = [256] * 3
x = np.array(np.linspace(l[0], h[0], n[0]), dtype=dtype)
y = np.array(np.linspace(l[1], h[1], n[1]), dtype=dtype)
z = np.array(np.linspace(l[2], h[2], n[2]), dtype=dtype)
if twod:
y = np.array(np.linspace(-0.1, 0.1, 2), dtype=dtype)
if nonuniform:
z += 0.01 * ((h[2] - l[2]) / n[2]) * np.sin(np.linspace(0, np.pi, n[2]))
B = field.empty([x, y, z], nr_comps=3, name=name, center='cell',
layout='interlaced', dtype=dtype)
B.set_info('crd_system', viscid.get_crd_system(crd_system))
return fill_dipole(B, m=m, strength=strength)
def save_fields(cls, fname, flds, **kwargs):
""" save some fields using the format given by the class """
# FIXME: this is only good for writing cartesian rectilnear flds
# FIXME: axes are renamed if flds[0] is 1D or 2D
assert len(flds) > 0
fname = os.path.expanduser(os.path.expandvars(fname))
# FIXME: all coordinates are saved as non-uniform, the proper
# way to do this is to have let coordinate format its own
# hdf5 / xdmf / numpy binary output
clist = flds[0].crds.get_clist(full_arrays=True)
crd_arrs = [np.array([0.0])] * 3
crd_names = ["x", "y", "z"]
for i, c in enumerate(clist):
crd_arrs[i] = c[1]
crd_shape = [len(arr) for arr in crd_arrs]
time = flds[0].time
# write arrays to the hdf5 file
with h5py.File(fname, 'w') as f:
for axis_name, arr in zip(crd_names, crd_arrs):
loc = cls._CRDS_GROUP + '/' + axis_name
f[loc] = arr
for fld in flds:
loc = cls._FLD_GROUPS[fld.center.lower()] + '/' + fld.name
# xdmf files use kji ordering
f[loc] = fld.data.T
# big bad openggcm time_str hack to put basetime into hdf5 file
for fld in flds:
try:
tfmt = "%Y:%m:%d:%H:%M:%S.%f"
sec_td = viscid.as_timedelta64(fld.time, 's')
dtime = viscid.as_datetime(fld.basetime + sec_td).strftime(tfmt)
epoch = viscid.readers.openggcm.GGCM_EPOCH
ts = viscid.as_timedelta(fld.basetime - epoch).total_seconds()
ts += fld.time
timestr = "time= {0} {1:.16e} {2} 300c".format(fld.time, ts, dtime)
f.create_group('openggcm')
f['openggcm'].attrs['time_str'] = np.string_(timestr)
break
except viscid.NoBasetimeError:
pass
# now write an xdmf file
xdmf_fname = os.path.splitext(fname)[0] + ".xdmf"
relh5fname = "./" + os.path.basename(fname)
with open(xdmf_fname, 'w') as f:
xloc = cls._CRDS_GROUP + '/' + crd_names[0]
yloc = cls._CRDS_GROUP + '/' + crd_names[1]
zloc = cls._CRDS_GROUP + '/' + crd_names[2]
dim_str = " ".join([str(l) for l in crd_shape][::-1])
f.write(cls._XDMF_TEMPLATE_BEGIN.format(time=time))
s = cls._XDMF_TEMPLATE_RECTILINEAR_GRID_BEGIN.format(
grid_name="vgrid", crd_dims=dim_str, h5fname=relh5fname,
xdim=crd_shape[0], ydim=crd_shape[1], zdim=crd_shape[2],
xloc=xloc, yloc=yloc, zloc=zloc)
f.write(s)
for fld in flds:
_crd_system = viscid.get_crd_system(fld, None)
if _crd_system:
f.write(cls._XDMF_INFO_TEMPLATE.format(name="crd_system",
value=_crd_system))
break
for fld in flds:
fld = fld.as_flat().T
dt = fld.dtype.name.rstrip("0123456789").title()
precision = fld.dtype.itemsize
fld_dim_str = " ".join([str(l) for l in fld.shape])
loc = cls._FLD_GROUPS[fld.center.lower()] + '/' + fld.name
s = cls._XDMF_TEMPLATE_ATTRIBUTE.format(
fld_name=fld.name,
fld_type=fld.fldtype, center=fld.center.title(),
dtype=dt, precision=precision, fld_dims=fld_dim_str,
h5fname=relh5fname, fld_loc=loc)
f.write(s)
f.write(cls._XDMF_TEMPLATE_GRID_END)
f.write(cls._XDMF_TEMPLATE_END)
def __init__(self, p0=(0, 0, 0), r=0.0, pole=(0, 0, 1), ntheta=20, nphi=20,
thetalim=(0, 180.0), philim=(0, 360.0), roll=0.0, crd_system=None,
theta_endpoint='auto', phi_endpoint='auto', pole_is_vector=True,
theta_phi=False, cache=False, dtype=None):
"""Make seeds on the surface of a sphere
Note:
There is some funny business about the meaning of phi=0 and
`crd_system`. By default, this seed generator is agnostic
to coordinate systems and phi=0 always means the +x axis.
If crd_system is 'gse', 'mhd', or an object whose find_info
method returns a 'crd_system', then phi=0 means midnight.
This is important when specifying a phi range or plotting
on a matplotlib polar plot.
Args:
p0 (list, tuple, or ndarray): Origin of sphere as (x, y, z)
r (float): Radius of sphere; or calculated from pole if 0
pole (list, tuple, or ndarray): Vector pointing
in the direction of the north pole of the sphere.
Defaults to (0, 0, 1).
ntheta (int): Number of points in theta
nphi (int): Number of points in phi
thetalim (list): min and max theta (in degrees)
philim (list): min and max phi (in degrees)
roll (float): Roll the seeds around the pole by this angle
in deg
crd_system (str, Field): a crd system ('gse', 'mhd') or an
object that has 'crd_system' info such that phi=0
means midnight instead of the +x axis.
theta_endpoint (str): this is a bit of a hack to keep from
having redundant seeds at poles. You probably just want
auto here
phi_endpoint (bool): if true, then let phi inclue upper
value. This is false by default since 0 == 2pi.
pole_is_vector (bool): Whether pole is a vector or a
vector
theta_phi (bool): If True, the uv and local representations
are ordered (theta, phi), otherwise (phi, theta)
Raises:
ValueError: if thetalim or philim don't have 2 values each
"""
super(Sphere, self).__init__(cache=cache, dtype=dtype)
self.p0 = np.asarray(p0, dtype=self.dtype)
if pole_is_vector:
self.pole = np.asarray(pole, dtype=self.dtype)
else:
if pole is None:
pole = p0 + np.asarray([0, 0, 1], dtype=self.dtype)
else:
pole = np.asarray(pole, dtype=self.dtype)
self.pole = pole - p0
if not r:
r = np.linalg.norm(self.pole)
self.pole = self.pole / np.linalg.norm(self.pole)
if not len(thetalim) == len(philim) == 2:
raise ValueError("thetalim and philim must have both min and max")
try:
roll = float(roll)
except TypeError:
pass
# square away crd system
if crd_system:
if hasattr(crd_system, 'find_info'):
crd_system = viscid.get_crd_system(crd_system, 'none')
else:
crd_system = 'none'
if crd_system.strip().lower() == 'gse':
crd_system_roll = -180.0
else:
crd_system_roll = 0.0
self.r = r
self.ntheta = ntheta
self.nphi = nphi
self.thetalim = np.deg2rad(thetalim)
self.philim = np.deg2rad(philim)
self.theta_endpoint = theta_endpoint
self.phi_endpoint = phi_endpoint
self.theta_phi = theta_phi
self.roll = roll
self.crd_system_roll = crd_system_roll
def get_mp_info(pp, b, j, e, cache=True, cache_dir=None,
slc="x=5.5f:11.0f, y=-4.0f:4.0f, z=-3.6f:3.6f",
fit="mp_xloc", fit_p0=(9.0, 0.0, 0.0, 1.0, -1.0, -1.0)):
"""Get info about m-pause as flattened fields
Notes:
The first thing this function does is mask locations where
the GSE-y current density < 1e-4. This masks out the bow
shock and current free regions. This works for southward IMF,
but it is not very general.
Parameters:
pp (ScalarcField): pressure
b (VectorField): magnetic field
j (VectorField): current density
e (VectorField, None): electric field (same centering as b). If
None, then the info that requires E will be filled with NaN
cache (bool, str): Save to and load from cache, if "force",
then don't load from cache if it exists, but do save a
cache at the end
cache_dir (str): Directory for cache, if None, same directory
as that file to which the grid belongs
slc (str): slice that gives a box that contains the m-pause
fit (str): to which resulting field should the paraboloid be fit,
defaults to mp_xloc, but pp_max_xloc might be useful in some
circumstances
fit_p0 (tuple): Initial guess vector for paraboloid fit
Returns:
dict: Unless otherwise noted, the entiries are 2D (y-z) fields
- **mp_xloc** location of minimum abs(Bz), this works
better than max of J^2 for FTEs
- **mp_sheath_edge** location where Jy > 0.1 * Jy when
coming in from the sheath side
- **mp_sphere_edge** location where Jy > 0.1 * Jy when
coming in from the sphere side
- **mp_width** difference between m-sheath edge and
msphere edge
- **mp_shear** magnetic shear taken 6 grid points into
the m-sheath / m-sphere
- **pp_max** max pp
- **pp_max_xloc** location of max pp
- **epar_max** max e parallel
- **epar_max_xloc** location of max e parallel
- **paraboloid** numpy.recarray of paraboloid fit. The
parameters are given in the 0th element, and
the 1st element contains the 1-sigma values for the fit
Raises:
RuntimeError: if using MHD crds instead of GSE crds
"""
if not cache_dir:
cache_dir = pp.find_info("_viscid_dirname", "./")
run_name = pp.find_info("run", None)
if cache and run_name:
t = pp.time
mp_fname = "{0}/{1}.mpause.{2:06.0f}".format(cache_dir, run_name, t)
else:
mp_fname = ""
try:
force = cache.strip().lower() == "force"
except AttributeError:
force = False
try:
if force or not mp_fname or not os.path.isfile(mp_fname + ".xdmf"):
raise IOError()
mp_info = {}
with viscid.load_file(mp_fname + ".xdmf") as dat:
fld_names = ["mp_xloc", "mp_sheath_edge", "mp_sphere_edge",
"mp_width", "mp_shear", "pp_max", "pp_max_xloc",
"epar_max", "epar_max_xloc"]
for fld_name in fld_names:
mp_info[fld_name] = dat[fld_name]["x=0"]
except (IOError, KeyError):
mp_info = {}
crd_system = viscid.get_crd_system(b, None)
if crd_system != 'gse':
raise RuntimeError("get_mp_info can't work in MHD crds, "
"switch to GSE please")
if j.nr_patches == 1:
pp_block = pp[slc]
b_block = b[slc]
j_block = j[slc]
if e is None:
e_block = np.nan * viscid.empty_like(j_block)
else:
e_block = e[slc]
else:
# interpolate an amr grid so we can proceed
obnd = pp.get_slice_extent(slc)
dx = np.min(pp.skeleton.L / pp.skeleton.n, axis=0)
nx = np.ceil((obnd[1] - obnd[0]) / dx)
vol = viscid.seed.Volume(obnd[0], obnd[1], nx, cache=True)
pp_block = vol.wrap_field(viscid.interp_trilin(pp, vol),
name="P").as_cell_centered()
b_block = vol.wrap_field(viscid.interp_trilin(b, vol),
name="B").as_cell_centered()
j_block = vol.wrap_field(viscid.interp_trilin(j, vol),
name="J").as_cell_centered()
if e is None:
e_block = np.nan * viscid.empty_like(j_block)
else:
e_block = vol.wrap_field(viscid.interp_trilin(e, vol),
name="E").as_cell_centered()
# jsq = viscid.dot(j_block, j_block)
bsq = viscid.dot(b_block, b_block)
# extract ndarrays and mask out bow shock / current free regions
maskval = 1e-4
jy_mask = j_block['y'].data < maskval
masked_bsq = 1.0 * bsq
masked_bsq.data = np.ma.masked_where(jy_mask, bsq)
xcc = j_block.get_crd_cc('x')
nx = len(xcc)
mp_xloc = np.argmin(masked_bsq, axis=0) # indices
mp_xloc = mp_xloc.wrap(xcc[mp_xloc.data]) # location
pp_max = np.max(pp_block, axis=0)
pp_max_xloc = np.argmax(pp_block, axis=0) # indices
pp_max_xloc = pp_max_xloc.wrap(xcc[pp_max_xloc.data]) # location
epar = viscid.project(e_block, b_block)
epar_max = np.max(epar, axis=0)
epar_max_xloc = np.argmax(epar, axis=0) # indices
epar_max_xloc = pp_max_xloc.wrap(xcc[epar_max_xloc.data]) # location
_ret = find_mp_edges(j_block, 0.1, 0.1, maskval=maskval)
sheath_edge, msphere_edge, mp_width, sheath_ind, sphere_ind = _ret
# extract b and b**2 at sheath + 6 grid points and sphere - 6 grid pointns
# clipping cases where things go outside the block. clipped ponints are
# set to nan
step = 6
# extract b
if b_block.layout == "flat":
comp_axis = 0
ic, _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape])
ix = np.clip(sheath_ind + step, 0, nx - 1)
b_sheath = b_block.data[ic, ix, iy, iz]
ix = np.clip(sheath_ind - step, 0, nx - 1)
b_sphere = b_block.data[ic, ix, iy, iz]
elif b_block.layout == "interlaced":
comp_axis = 3
_, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape[:-1]])
ix = np.clip(sheath_ind + step, 0, nx - 1)
b_sheath = b_block.data[ix, iy, iz]
ix = np.clip(sheath_ind - step, 0, nx - 1)
b_sphere = b_block.data[ix, iy, iz]
# extract b**2
bmag_sheath = np.sqrt(np.sum(b_sheath**2, axis=comp_axis))
bmag_sphere = np.sqrt(np.sum(b_sphere**2, axis=comp_axis))
costheta = (np.sum(b_sheath * b_sphere, axis=comp_axis) /
(bmag_sphere * bmag_sheath))
costheta = np.where((sheath_ind + step < nx) & (sphere_ind - step >= 0),
costheta, np.nan)
mp_shear = mp_width.wrap((180.0 / np.pi) * np.arccos(costheta))
# don't bother with pretty name since it's not written to file
# plane_crds = b_block.crds.slice_keep('x=0', cc=True)
# fld_kwargs = dict(center="Cell", time=b.time)
mp_width.name = "mp_width"
mp_xloc.name = "mp_xloc"
sheath_edge.name = "mp_sheath_edge"
msphere_edge.name = "mp_sphere_edge"
mp_shear.name = "mp_shear"
pp_max.name = "pp_max"
pp_max_xloc.name = "pp_max_xloc"
epar_max.name = "epar_max"
epar_max_xloc.name = "epar_max_xloc"
mp_info = {}
mp_info["mp_width"] = mp_width
mp_info["mp_xloc"] = mp_xloc
mp_info["mp_sheath_edge"] = sheath_edge
mp_info["mp_sphere_edge"] = msphere_edge
mp_info["mp_shear"] = mp_shear
mp_info["pp_max"] = pp_max
mp_info["pp_max_xloc"] = pp_max_xloc
mp_info["epar_max"] = epar_max
mp_info["epar_max_xloc"] = epar_max_xloc
# cache new fields to disk
if mp_fname:
viscid.save_fields(mp_fname + ".h5", mp_info.values())
try:
_paraboloid_params = fit_paraboloid(mp_info[fit], p0=fit_p0)
mp_info["paraboloid"] = _paraboloid_params
except ImportError as _exception:
try:
msg = _exception.message
except AttributeError:
msg = _exception.msg
mp_info["paraboloid"] = viscid.DeferredImportError(msg)
mp_info["mp_width"].pretty_name = "Magnetopause Width"
mp_info["mp_xloc"].pretty_name = "Magnetopause $X_{gse}$ Location"
mp_info["mp_sheath_edge"].pretty_name = "Magnetosheath Edge"
mp_info["mp_sphere_edge"].pretty_name = "Magnetosphere Edge"
mp_info["mp_shear"].pretty_name = "Magnetic Shear"
mp_info["pp_max"].pretty_name = "Max Pressure"
mp_info["pp_max_xloc"].pretty_name = "Max Pressure Location"
mp_info["epar_max"].pretty_name = "Max E Parallel"
mp_info["epar_max_xloc"].pretty_name = "Max E Parallel Location"
return mp_info
