def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="",
view_kwargs=None, show=False, scatter_mpl=False, mesh_mvi=True):
interpolated_fld = viscid.interp_trilin(fld, seeds)
seed_name = seeds.__class__.__name__
if add_title:
seed_name += " " + add_title
try:
if not plot2d:
raise ImportError
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
plt.clf()
# plt.plot(seeds.get_points()[2, :], fld)
mpl_plot_kwargs = dict()
if interpolated_fld.is_spherical():
mpl_plot_kwargs['hemisphere'] = 'north'
vlt.plot(interpolated_fld, **mpl_plot_kwargs)
plt.title(seed_name)
plt.savefig(next_plot_fname(__file__, series='2d'))
if show:
plt.show()
if scatter_mpl:
plt.clf()
vlt.plot2d_line(seeds.get_points(), fld, symdir='z', marker='o')
plt.savefig(next_plot_fname(__file__, series='2d'))
if show:
plt.show()
except ImportError:
pass
try:
if not plot3d:
raise ImportError
from viscid.plot import vlab
_ = get_mvi_fig(offscreen=not show)
try:
if mesh_mvi:
mesh = vlab.mesh_from_seeds(seeds, scalars=interpolated_fld)
mesh.actor.property.backface_culling = True
except RuntimeError:
pass
pts = seeds.get_points()
p = vlab.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data,
scale_mode='none', scale_factor=0.02)
vlab.axes(p)
vlab.title(seed_name)
if view_kwargs:
vlab.view(**view_kwargs)
vlab.savefig(next_plot_fname(__file__, series='3d'))
if show:
vlab.show(stop=True)
except ImportError:
pass
def minvar(B, p1, p2, n=40):
"""Find minimum variance eigenvectors of `B` between `p1` and `p2`
Do minimum variance analysis (MVA) of vector field `B` between
two points, i.e., `p1` and `p2` are like two spacecraft locations
on either side of the boundary for the MVA.
Args:
B (:py:class:`viscid.field.VectorField`): Vector field for MVA
p1 (sequence): xyz point on one side of the boundary
p2 (sequence): xyz point on the other side of the boundary
n (int): number of points to sample B for the MVA
Returns:
(evals, evecs)
All are ordered toward increasing eigenvalue magnitude. `evecs`
is a 2d array where the vectors are columns (2nd axis), i.e.,
the min variance eigenvector is evecs[0, :]
"""
p1 = np.array(p1).reshape((3,))
p2 = np.array(p2).reshape((3,))
line = viscid.Line(p1, p2, n)
b = viscid.interp_trilin(B, line)
return minvar_series(b.data)
def minvar(B, p1, p2, n=40):
"""Find minimum variance eigenvectors of `B` between `p1` and `p2`
Do minimum variance analysis (MVA) of vector field `B` between
two points, i.e., `p1` and `p2` are like two spacecraft locations
on either side of the boundary for the MVA.
Args:
B (:py:class:`viscid.field.VectorField`): Vector field for MVA
p1 (sequence): xyz point on one side of the boundary
p2 (sequence): xyz point on the other side of the boundary
n (int): number of points to sample B for the MVA
Returns:
(evals, evecs)
All are ordered toward increasing eigenvalue magnitude. `evecs`
is a 2d array where the vectors are columns (2nd axis), i.e.,
the min variance eigenvector is evecs[0, :]
"""
p1 = np.array(p1).reshape((3, ))
p2 = np.array(p2).reshape((3, ))
line = viscid.Line(p1, p2, n)
b = viscid.interp_trilin(B, line)
return minvar_series(b.data)
def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="",
view_kwargs=None, show=False, scatter_mpl=False, mesh_mvi=True):
interpolated_fld = viscid.interp_trilin(fld, seeds)
seed_name = seeds.__class__.__name__
if add_title:
seed_name += " " + add_title
try:
if not plot2d:
raise ImportError
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
plt.clf()
# plt.plot(seeds.get_points()[2, :], fld)
mpl_plot_kwargs = dict()
if interpolated_fld.is_spherical():
mpl_plot_kwargs['hemisphere'] = 'north'
vlt.plot(interpolated_fld, **mpl_plot_kwargs)
plt.title(seed_name)
plt.savefig(next_plot_fname(__file__, series='2d'))
if show:
plt.show()
if scatter_mpl:
plt.clf()
vlt.plot2d_line(seeds.get_points(), fld, symdir='z', marker='o')
plt.savefig(next_plot_fname(__file__, series='2d'))
if show:
plt.show()
except ImportError:
pass
try:
if not plot3d:
raise ImportError
vlab, _ = get_mvi_fig()
try:
if mesh_mvi:
mesh = vlab.mesh_from_seeds(seeds, scalars=interpolated_fld)
mesh.actor.property.backface_culling = True
except RuntimeError:
pass
pts = seeds.get_points()
p = vlab.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data,
scale_mode='none', scale_factor=0.02)
vlab.axes(p)
vlab.title(seed_name)
if view_kwargs:
vlab.view(**view_kwargs)
vlab.savefig(next_plot_fname(__file__, series='3d'))
if show:
vlab.show(stop=True)
except ImportError:
pass
def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="",
view_kwargs=None, show=False):
interpolated_fld = viscid.interp_trilin(fld, seeds)
seed_name = seeds.__class__.__name__
if add_title:
seed_name += " " + add_title
try:
if not plot2d:
raise ImportError
from matplotlib import pyplot as plt
from viscid.plot import vpyplot as vlt
plt.clf()
# plt.plot(seeds.get_points()[2, :], fld)
mpl_plot_kwargs = dict()
if interpolated_fld.is_spherical():
mpl_plot_kwargs['hemisphere'] = 'north'
vlt.plot(interpolated_fld, **mpl_plot_kwargs)
plt.title(seed_name)
plt.savefig(next_plot_fname(__file__, series='2d'))
if show:
plt.show()
except ImportError:
pass
try:
if not plot3d:
raise ImportError
from viscid.plot import vlab
try:
fig = _global_ns['figure']
vlab.clf()
except KeyError:
fig = vlab.figure(size=[1200, 800], offscreen=not show)
_global_ns['figure'] = fig
try:
mesh = vlab.mesh_from_seeds(seeds, scalars=interpolated_fld)
mesh.actor.property.backface_culling = True
except RuntimeError:
pass
pts = seeds.get_points()
p = vlab.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data,
scale_mode='none', scale_factor=0.02)
vlab.axes(p)
vlab.title(seed_name)
if view_kwargs:
vlab.view(**view_kwargs)
vlab.savefig(next_plot_fname(__file__, series='3d'))
if show:
vlab.show()
except ImportError:
pass
def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="",
view_kwargs=None, show=False):
interpolated_fld = viscid.interp_trilin(fld, seeds)
seed_name = seeds.__class__.__name__
if add_title:
seed_name += " " + add_title
try:
if not plot2d:
raise ImportError
from viscid.plot import mpl
mpl.plt.clf()
# mpl.plt.plot(seeds.get_points()[2, :], fld)
mpl_plot_kwargs = dict()
if interpolated_fld.is_spherical():
mpl_plot_kwargs['hemisphere'] = 'north'
mpl.plot(interpolated_fld, **mpl_plot_kwargs)
mpl.plt.title(seed_name)
mpl.plt.savefig(next_plot_fname(__file__, series='2d'))
if show:
mpl.plt.show()
except ImportError:
pass
try:
if not plot3d:
raise ImportError
from viscid.plot import mvi
try:
fig = _global_ns['figure']
mvi.clf()
except KeyError:
fig = mvi.figure(size=[1200, 800], offscreen=not show)
_global_ns['figure'] = fig
try:
mesh = mvi.mesh_from_seeds(seeds, scalars=interpolated_fld)
mesh.actor.property.backface_culling = True
except RuntimeError:
pass
pts = seeds.get_points()
p = mvi.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data,
scale_mode='none', scale_factor=0.02)
mvi.axes(p)
mvi.title(seed_name)
if view_kwargs:
mvi.view(**view_kwargs)
mvi.savefig(next_plot_fname(__file__, series='3d'))
if show:
mvi.show()
except ImportError:
pass
def project_along_line(line, fld):
"""Project a Vector Field Parallel to a streamline
Args:
line (ndarray): 3xN of points along the
fld (VectorField): Field to interpolate and project onto the
line
"""
fld_on_verts = viscid.interp_trilin(fld, line)
dsvec = line[:, 2:] - line[:, :-2]
dsvec = np.concatenate([dsvec[:, 0:1], dsvec, dsvec[:, -2:-1]], axis=1)
dsvec = dsvec / np.linalg.norm(dsvec, axis=0)
return np.sum(fld_on_verts * dsvec.T, axis=1)
def minvar(B, p1, p2, n=40):
"""Find minimum variance eigenvectors of `B` between `p1` and `p2`
Do minimum variance analysis (MVA) of vector field `B` between
two points, i.e., `p1` and `p2` are like two spacecraft locations
on either side of the boundary for the MVA.
Args:
B (:py:class:`viscid.field.VectorField`): Vector field for MVA
p1 (sequence): xyz point on one side of the boundary
p2 (sequence): xyz point on the other side of the boundary
n (int): number of points to sample B for the MVA
Returns:
(evals, evecs)
All are ordered toward increasing eigenvalue magnitude. `evecs`
is a 2d array where the vectors are columns (2nd axis), i.e.,
the min variance eigenvector is evecs[0, :]
"""
p1 = np.array(p1).reshape((3,))
p2 = np.array(p2).reshape((3,))
line = viscid.Line(p1, p2, n)
b = viscid.interp_trilin(B, line)
m = np.empty((3, 3))
for i in range(3):
for j in range(i+1):
bibj = np.average(b[:, i] * b[:, j])
m[i, j] = bibj - (np.average(b[:, i]) * np.average(b[:, j]))
m[j, i] = bibj - (np.average(b[:, i]) * np.average(b[:, j]))
evals, evecs = np.linalg.eig(m)
eval_mags = np.abs(evals)
sorted_inds = np.argsort(eval_mags)
evec_min = evecs[:, sorted_inds[0]].flatten()
evec_int = evecs[:, sorted_inds[1]].flatten()
evec_max = evecs[:, sorted_inds[2]].flatten()
eval_min = eval_mags[sorted_inds[0]]
eval_int = eval_mags[sorted_inds[1]]
warn_thresh = 0.05
if (eval_int - eval_min) / eval_min < warn_thresh:
viscid.logger.warn("Minvar says minimum and intermediate eigenvalues "
"are too close together: {0:g} - {1:g} < {2:g}%"
"".format(eval_int, eval_min, warn_thresh))
return evals[sorted_inds], np.array([evec_min, evec_int, evec_max]).T
def integrate_along_lines(lines, fld, reduction="dot", mask_func=None):
"""Integrate the value of fld along a list of lines
Args:
lines (list): list of 3xN ndarrays, N needs not be the same for
all lines
fld (Field): Field to interpolate / integrate
reduction (str): If fld is a vector field, what quantity to
integrate. Can be "dot" to dot the vectors with ds along
the line, or "mag" to integrate the magnitude.
Returns:
ndarray with shape (len(lines), )
"""
arr = np.zeros((len(lines),), dtype=fld.dtype)
cum_n = np.cumsum([0] + [line.shape[1] for line in lines])
all_verts = np.concatenate(lines, axis=1)
fld_on_verts = viscid.interp_trilin(fld, all_verts).data
for i, start, stop in izip(count(), cum_n[:-1], cum_n[1:]):
ds = np.linalg.norm(lines[i][:, 1:] - lines[i][:, :-1], axis=0)
if len(fld_on_verts.shape) > 1:
reduction = reduction.strip().lower()
if reduction == "dot":
dsvec = lines[i][:, 1:] - lines[i][:, :-1]
dsvec = dsvec / np.linalg.norm(dsvec, axis=0)
values = 0.5 * (fld_on_verts[start:stop - 1, :] +
fld_on_verts[start + 1:stop, :])
values = values * dsvec.T
if mask_func is not None:
values = np.ma.masked_where(mask_func(values), values)
values = np.sum(values, axis=1)
elif reduction in ["mag", "magnitude", "norm"]:
mag = np.linalg.norm(fld_on_verts[start:stop], axis=1)
values = 0.5 * (mag[start:stop - 1] + mag[start + 1:stop])
else:
raise ValueError("Unknown reduction: {0}".format(reduction))
else:
values = 0.5 * (fld_on_verts[start:stop - 1] +
fld_on_verts[start + 1:stop])
arr[i] = np.sum(values * ds)
return arr
def minvar(B, p1, p2, n=40):
"""Find minimum variance eigenvectors of `B` between `p1` and `p2`
Do minimum variance analysis (MVA) of vector field `B` between
two points, i.e., `p1` and `p2` are like two spacecraft locations
on either side of the boundary for the MVA.
Args:
B (:py:class:`viscid.field.VectorField`): Vector field for MVA
p1 (sequence): xyz point on one side of the boundary
p2 (sequence): xyz point on the other side of the boundary
n (int): number of points to sample B for the MVA
Returns:
(evals, evecs)
All are ordered toward increasing eigenvalue magnitude. `evecs`
is a 2d array where the vectors are columns (2nd axis), i.e.,
the min variance eigenvector is evecs[0, :]
"""
p1 = np.array(p1).reshape((3, ))
p2 = np.array(p2).reshape((3, ))
line = viscid.Line(p1, p2, n)
b = viscid.interp_trilin(B, line)
m = np.empty((3, 3))
for i in range(3):
for j in range(i + 1):
bibj = np.average(b[:, i] * b[:, j])
m[i, j] = bibj - (np.average(b[:, i]) * np.average(b[:, j]))
m[j, i] = bibj - (np.average(b[:, i]) * np.average(b[:, j]))
evals, evecs = np.linalg.eig(m)
eval_mags = np.abs(evals)
sorted_inds = np.argsort(eval_mags)
evec_min = evecs[:, sorted_inds[0]].flatten()
evec_int = evecs[:, sorted_inds[1]].flatten()
evec_max = evecs[:, sorted_inds[2]].flatten()
eval_min = eval_mags[sorted_inds[0]]
eval_int = eval_mags[sorted_inds[1]]
warn_thresh = 0.05
if (eval_int - eval_min) / eval_min < warn_thresh:
viscid.logger.warn("Minvar says minimum and intermediate eigenvalues "
"are too close together: {0:g} - {1:g} < {2:g}%"
"".format(eval_int, eval_min, warn_thresh))
return evals[sorted_inds], np.array([evec_min, evec_int, evec_max]).T
def integrate_along_lines(lines, fld):
"""Integrate the value of fld along a list of lines
Args:
lines (list): list of 3xN ndarrays, N needs not be the same for
all lines
fld (Field): Field to interpolate / integrate
Returns:
ndarray with shape (len(lines), )
"""
arr = np.zeros((len(lines),), dtype=fld.dtype)
cum_n = np.cumsum([0] + [line.shape[1] for line in lines])
all_verts = np.concatenate(lines, axis=1)
fld_on_verts = viscid.interp_trilin(fld, all_verts)
for i, start, stop in izip(count(), cum_n[:-1], cum_n[1:]):
ds = np.linalg.norm(lines[i][:, 1:] - lines[i][:, :-1], axis=0)
values = 0.5 * (fld_on_verts[start:stop - 1] +
fld_on_verts[start + 1:stop])
arr[i] = np.sum(values * ds)
return arr
def prepare_lines(lines, scalars=None, do_connections=False, other=None):
"""Concatenate and standardize a list of lines
Args:
lines (list): Must be a list of 3xN or 4xN ndarrays of xyz(s)
data for N points along the line. N need not be the same
for all lines. Can alse be 6xN such that lines[:][3:, :]
are interpreted as rgb colors
scalars (sequence): can be one of::
- single hex color (ex, `#FD7400`)
- sequence of Nlines hex colors
- single rgb(a) tuple
- sequence of Nlines rgb(a) sequences
- sequence of N values that go with each vertex and will
be mapped with a colormap
- sequence of Nlines values that go each line and will
be mapped with a colormap
- Field. If `np.prod(fld.shape) in (N, nlines)` then the
field is interpreted as a simple sequence of values
(ie, the topology result from calc_streamlines for
coloring each line). Otherwise, the field is
interpolated onto each vertex.
do_connections (bool): Whether or not to make connections array
other (dict): a dictionary of other arrays that should be
reshaped and the like the same way scalars is
Returns:
(vertices, scalars, connections, other)
* vertices (ndarray): 3xN array of N xyz points. N is the sum
of the lengths of all the lines
* scalars (ndarray): N array of scalars, 3xN array of uint8
rgb values, 4xN array of uint8 rgba values, or None
* connections (ndarray): Nx2 array of ints (indices along
axis 1 of vertices) describing the forward and backward
connectedness of the lines, or None
* other (dict): a dict of N length arrays
Raises:
ValueError: If rgb data is not in a valid range or the shape
of scalars is not understood
"""
nlines = len(lines)
npts = [line.shape[1] for line in lines]
N = np.sum(npts)
first_idx = np.cumsum([0] + npts[:-1])
vertices = [np.asarray(line) for line in lines]
vertices = np.concatenate(lines, axis=1)
if vertices.dtype.kind not in 'fc':
vertices = np.asarray(vertices, dtype='f')
if vertices.shape[0] > 3:
if scalars is not None:
viscid.logger.warning("Overriding line scalars with scalars kwarg")
else:
scalars = vertices[3:, :]
vertices = vertices[:3, :]
if scalars is not None:
scalars_are_strings = False
if isinstance(scalars, viscid.field.Field):
if np.prod(scalars.shape) in (nlines, N):
scalars = np.asarray(scalars).reshape(-1)
else:
scalars = viscid.interp_trilin(scalars, vertices)
if scalars.size != N:
raise ValueError("Scalars was not a scalar field")
elif isinstance(scalars, (list, tuple, viscid.string_types, np.ndarray)):
# string types need some extra massaging
if any(isinstance(s, viscid.string_types) for s in scalars):
assert all(isinstance(s, viscid.string_types) for s in scalars)
scalars_are_strings = True
scalars = _string_colors_as_hex(scalars)
elif isinstance(scalars, np.ndarray):
scalars = scalars
else:
for i, si in enumerate(scalars):
if not isinstance(si, np.ndarray):
scalars[i] = np.asarray(si)
scalars[i] = np.atleast_2d(scalars[i])
try:
scalars = np.concatenate(scalars, axis=0)
except ValueError:
scalars = np.concatenate(scalars, axis=1)
scalars = np.atleast_2d(scalars)
if scalars.dtype == np.dtype('object'):
raise RuntimeError("Scalars must be numbers, tuples of numbers "
"that indicate rgb(a), or hex strings - they "
"must not be python objects")
if scalars.shape == (1, 1):
scalars = scalars.repeat(N, axis=1)
elif scalars.shape == (1, nlines) or scalars.shape == (nlines, 1):
# one scalar for each line, so broadcast it
scalars = scalars.reshape(nlines, 1)
scalars = [scalars[i].repeat(ni) for i, ni in enumerate(npts)]
scalars = np.concatenate(scalars, axis=0).reshape(1, N)
elif scalars.shape == (N, 1) or scalars.shape == (1, N):
# catch these so they're not interpreted as colors if
# nlines == 1 and N == 3; ie. 1 line with 3 points
scalars = scalars.reshape(1, N)
elif scalars.shape in [(3, nlines), (nlines, 3), (4, nlines), (nlines, 4)]:
# one rgb(a) color for each line, so broadcast it
if (scalars.shape in [(3, nlines), (4, nlines)] and
scalars.shape not in [(3, 3), (4, 4)]):
# the guard against shapes (3, 3) and (4, 4) mean that
# these square shapes are assumed Nlines x {3,4}
scalars = scalars.T
nccomps = scalars.shape[1] # 3 for rgb, 4 for rgba
colors = []
for i, ni in enumerate(npts):
c = scalars[i].reshape(nccomps, 1).repeat(ni, axis=1)
colors.append(c)
scalars = np.concatenate(colors, axis=1)
elif scalars.shape in [(3, N), (N, 3), (4, N), (N, 4)]:
# one rgb(a) color for each vertex
if (scalars.shape in [(3, N), (4, N)] and
scalars.shape not in [(3, 3), (4, 4)]):
# the guard against shapes (3, 3) and (4, 4) mean that
# these square shapes are assumed N x {3,4}
scalars = scalars.T
elif scalars.shape in [(1, 3), (3, 1), (1, 4), (4, 1)]:
# interpret a single rgb(a) color, and repeat/broadcast it
scalars = scalars.reshape([-1, 1]).repeat(N, axis=1)
else:
scalars = scalars.reshape(-1, N)
# scalars now has shape (1, N), (3, N), or (4, N)
if scalars_are_strings:
# translate hex colors (#ff00ff) into rgb(a) values
scalars = np.char.lstrip(scalars, '#')
strlens = np.char.str_len(scalars)
min_strlen, max_strlen = np.min(strlens), np.max(strlens)
if min_strlen == max_strlen == 8:
# 32-bit rgba (two hex chars per channel)
scalars = _hexchar2int(scalars.astype('S8')).reshape(-1, 4).T
elif min_strlen == max_strlen == 6:
# 24-bit rgb (two hex chars per channel)
scalars = _hexchar2int(scalars.astype('S6')).reshape(-1, 3).T
else:
raise NotImplementedError("This should never happen as "
"scalars as colors should already "
"be preprocessed appropriately")
elif scalars.shape[0] == 1:
# normal scalars, cast them down to a single dimension
scalars = scalars.reshape(-1)
elif scalars.shape[0] in (3, 4):
# The scalars encode rgb data, standardize the result to a
# 3xN or 4xN ndarray of 1 byte unsigned ints [0..255]
if np.all(scalars >= 0) and np.all(scalars <= 1):
scalars = (255 * scalars).round().astype('u1')
elif np.all(scalars >= 0) and np.all(scalars < 256):
scalars = scalars.round().astype('u1')
else:
raise ValueError("Rgb data should be in range [0, 1] or "
"[0, 255], range given is [{0}, {1}]"
"".format(np.min(scalars), np.max(scalars)))
else:
raise ValueError("Scalars should either be a number, or set of "
"rgb values, shape is {0}".format(scalars.shape))
# scalars should now have shape (N, ) or be a uint8 array with shape
# (3, N) or (4, N) encoding an rgb(a) color for each point [0..255]
# ... done with scalars...
# broadcast / reshape additional arrays given in other
if other:
for key, arr in other.items():
if arr is None:
pass
elif arr.shape == (1, nlines) or arr.shape == (nlines, 1):
arr = arr.reshape(nlines, 1)
arr = [arr[i].repeat(ni) for i, ni in enumerate(npts)]
other[key] = np.concatenate(arr, axis=0).reshape(1, N)
else:
try:
other[key] = arr.reshape(-1, N)
except ValueError:
viscid.logger.warning("Unknown dimension, dropping array {0}"
"".format(key))
if do_connections:
connections = [None] * nlines
for i, ni in enumerate(npts):
# i0 is the index of the first point of the i'th line in lines
i0 = first_idx[i]
connections[i] = np.vstack([np.arange(i0, i0 + ni - 1.5),
np.arange(i0 + 1, i0 + ni - 0.5)]).T
connections = np.concatenate(connections, axis=0).astype('i')
else:
connections = None
return vertices, scalars, connections, other
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.as_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
def prepare_lines(lines, scalars=None, do_connections=False, other=None):
"""Concatenate and standardize a list of lines
Args:
lines (list): Must be a list of 3xN or 4xN ndarrays of xyz(s)
data for N points along the line. N need not be the same
for all lines. Can alse be 6xN such that lines[:][3:, :]
are interpreted as rgb colors
scalars (sequence): can be one of::
- single hex color (ex, `#FD7400`)
- sequence of Nlines hex colors
- single rgb(a) tuple
- sequence of Nlines rgb(a) sequences
- sequence of N values that go with each vertex and will
be mapped with a colormap
- sequence of Nlines values that go each line and will
be mapped with a colormap
- Field. If `np.prod(fld.shape) in (N, nlines)` then the
field is interpreted as a simple sequence of values
(ie, the topology result from calc_streamlines for
coloring each line). Otherwise, the field is
interpolated onto each vertex.
do_connections (bool): Whether or not to make connections array
other (dict): a dictionary of other arrays that should be
reshaped and the like the same way scalars is
Returns:
(vertices, scalars, connections, other)
* vertices (ndarray): 3xN array of N xyz points. N is the sum
of the lengths of all the lines
* scalars (ndarray): N array of scalars, 3xN array of uint8
rgb values, 4xN array of uint8 rgba values, or None
* connections (ndarray): Nx2 array of ints (indices along
axis 1 of vertices) describing the forward and backward
connectedness of the lines, or None
* other (dict): a dict of N length arrays
Raises:
ValueError: If rgb data is not in a valid range or the shape
of scalars is not understood
"""
nlines = len(lines)
npts = [line.shape[1] for line in lines]
N = np.sum(npts)
first_idx = np.cumsum([0] + npts[:-1])
vertices = [np.asarray(line) for line in lines]
vertices = np.concatenate(lines, axis=1)
if vertices.dtype.kind not in 'fc':
vertices = np.asarray(vertices, dtype='f')
if vertices.shape[0] > 3:
if scalars is not None:
viscid.logger.warn("Overriding line scalars with scalars kwarg")
else:
scalars = vertices[3:, :]
vertices = vertices[:3, :]
if scalars is not None:
scalars_are_strings = False
if isinstance(scalars, viscid.field.Field):
if np.prod(scalars.shape) in (nlines, N):
scalars = np.asarray(scalars).reshape(-1)
else:
scalars = viscid.interp_trilin(scalars, vertices)
if scalars.size != N:
raise ValueError("Scalars was not a scalar field")
elif isinstance(scalars,
(list, tuple, viscid.string_types, np.ndarray)):
# string types need some extra massaging
if any(isinstance(s, viscid.string_types) for s in scalars):
assert all(isinstance(s, viscid.string_types) for s in scalars)
scalars_are_strings = True
scalars = _string_colors_as_hex(scalars)
elif isinstance(scalars, np.ndarray):
scalars = scalars
else:
for i, si in enumerate(scalars):
if not isinstance(si, np.ndarray):
scalars[i] = np.asarray(si)
scalars[i] = np.atleast_2d(scalars[i])
try:
scalars = np.concatenate(scalars, axis=0)
except ValueError:
scalars = np.concatenate(scalars, axis=1)
scalars = np.atleast_2d(scalars)
if scalars.dtype == np.dtype('object'):
raise RuntimeError("Scalars must be numbers, tuples of numbers "
"that indicate rgb(a), or hex strings - they "
"must not be python objects")
if scalars.shape == (1, 1):
scalars = scalars.repeat(N, axis=1)
elif scalars.shape == (1, nlines) or scalars.shape == (nlines, 1):
# one scalar for each line, so broadcast it
scalars = scalars.reshape(nlines, 1)
scalars = [scalars[i].repeat(ni) for i, ni in enumerate(npts)]
scalars = np.concatenate(scalars, axis=0).reshape(1, N)
elif scalars.shape == (N, 1) or scalars.shape == (1, N):
# catch these so they're not interpreted as colors if
# nlines == 1 and N == 3; ie. 1 line with 3 points
scalars = scalars.reshape(1, N)
elif scalars.shape in [(3, nlines), (nlines, 3), (4, nlines),
(nlines, 4)]:
# one rgb(a) color for each line, so broadcast it
if (scalars.shape in [(3, nlines), (4, nlines)]
and scalars.shape not in [(3, 3), (4, 4)]):
# the guard against shapes (3, 3) and (4, 4) mean that
# these square shapes are assumed Nlines x {3,4}
scalars = scalars.T
nccomps = scalars.shape[1] # 3 for rgb, 4 for rgba
colors = []
for i, ni in enumerate(npts):
c = scalars[i].reshape(nccomps, 1).repeat(ni, axis=1)
colors.append(c)
scalars = np.concatenate(colors, axis=1)
elif scalars.shape in [(3, N), (N, 3), (4, N), (N, 4)]:
# one rgb(a) color for each vertex
if (scalars.shape in [(3, N), (4, N)]
and scalars.shape not in [(3, 3), (4, 4)]):
# the guard against shapes (3, 3) and (4, 4) mean that
# these square shapes are assumed N x {3,4}
scalars = scalars.T
elif scalars.shape in [(1, 3), (3, 1), (1, 4), (4, 1)]:
# interpret a single rgb(a) color, and repeat/broadcast it
scalars = scalars.reshape([-1, 1]).repeat(N, axis=1)
else:
scalars = scalars.reshape(-1, N)
# scalars now has shape (1, N), (3, N), or (4, N)
if scalars_are_strings:
# translate hex colors (#ff00ff) into rgb(a) values
scalars = np.char.lstrip(scalars, '#')
strlens = np.char.str_len(scalars)
min_strlen, max_strlen = np.min(strlens), np.max(strlens)
if min_strlen == max_strlen == 8:
# 32-bit rgba (two hex chars per channel)
scalars = _hexchar2int(scalars.astype('S8')).reshape(-1, 4).T
elif min_strlen == max_strlen == 6:
# 24-bit rgb (two hex chars per channel)
scalars = _hexchar2int(scalars.astype('S6')).reshape(-1, 3).T
else:
raise NotImplementedError("This should never happen as "
"scalars as colors should already "
"be preprocessed appropriately")
elif scalars.shape[0] == 1:
# normal scalars, cast them down to a single dimension
scalars = scalars.reshape(-1)
elif scalars.shape[0] in (3, 4):
# The scalars encode rgb data, standardize the result to a
# 3xN or 4xN ndarray of 1 byte unsigned ints [0..255]
if np.all(scalars >= 0) and np.all(scalars <= 1):
scalars = (255 * scalars).round().astype('u1')
elif np.all(scalars >= 0) and np.all(scalars < 256):
scalars = scalars.round().astype('u1')
else:
raise ValueError("Rgb data should be in range [0, 1] or "
"[0, 255], range given is [{0}, {1}]"
"".format(np.min(scalars), np.max(scalars)))
else:
raise ValueError("Scalars should either be a number, or set of "
"rgb values, shape is {0}".format(scalars.shape))
# scalars should now have shape (N, ) or be a uint8 array with shape
# (3, N) or (4, N) encoding an rgb(a) color for each point [0..255]
# ... done with scalars...
# broadcast / reshape additional arrays given in other
if other:
for key, arr in other.items():
if arr is None:
pass
elif arr.shape == (1, nlines) or arr.shape == (nlines, 1):
arr = arr.reshape(nlines, 1)
arr = [arr[i].repeat(ni) for i, ni in enumerate(npts)]
other[key] = np.concatenate(arr, axis=0).reshape(1, N)
else:
try:
other[key] = arr.reshape(-1, N)
except ValueError:
viscid.logger.warn("Unknown dimension, dropping array {0}"
"".format(key))
if do_connections:
connections = [None] * nlines
for i, ni in enumerate(npts):
# i0 is the index of the first point of the i'th line in lines
i0 = first_idx[i]
connections[i] = np.vstack([
np.arange(i0, i0 + ni - 1.5),
np.arange(i0 + 1, i0 + ni - 0.5)
]).T
connections = np.concatenate(connections, axis=0).astype('i')
else:
connections = None
return vertices, scalars, connections, other
def _main():
global offscreen_vlab
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--notwo", dest='notwo', action="store_true")
parser.add_argument("--nothree", dest='nothree', action="store_true")
parser.add_argument("--show", "--plot", action="store_true")
args = viscid.vutil.common_argparse(parser, default_verb=0)
plot2d = not args.notwo
plot3d = not args.nothree
# plot2d = True
# plot3d = True
# args.show = True
offscreen_vlab = not args.show
img = np.load(os.path.join(sample_dir, "logo.npy"))
x = np.linspace(-1, 1, img.shape[0])
y = np.linspace(-1, 1, img.shape[1])
z = np.linspace(-1, 1, img.shape[2])
logo = viscid.arrays2field([x, y, z], img)
if 1:
viscid.logger.info('Testing Point with custom local coordinates...')
pts = np.vstack([[-1, -0.5, 0, 0.5, 1],
[-1, -0.5, 0, 0.5, 1],
[ 0, 0.5, 1, 1.5, 2]])
local_crds = viscid.asarray_datetime64([0, 60, 120, 180, 240],
conservative=True)
seeds = viscid.Point(pts, local_crds=local_crds)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)
if 1:
viscid.logger.info('Testing Line...')
seeds = viscid.Line([-1, -1, 0], [1, 1, 2], n=5)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)
if 1:
viscid.logger.info('Testing Plane...')
seeds = viscid.Plane([0.0, 0.0, 0.0], [1, 1, 1], [1, 0, 0], 2, 2,
nl=160, nm=170, NL_are_vectors=True)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)
if 1:
viscid.logger.info('Testing Volume...')
seeds = viscid.Volume([-0.8, -0.8, -0.8], [0.8, 0.8, 0.8],
n=[64, 64, 3])
# note: can't make a 2d plot of the volume w/o a slice
run_test(logo, seeds, plot2d=False, plot3d=plot3d, add_title="3d",
show=args.show)
if 1:
viscid.logger.info('Testing Volume (with ignorable dim)...')
seeds = viscid.Volume([-0.8, -0.8, 0.0], [0.8, 0.8, 0.0],
n=[64, 64, 1])
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="2d",
show=args.show)
if 1:
viscid.logger.info('Testing Spherical Sphere (phi, theta)...')
seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170,
pole=[-1, -1, -1], theta_phi=False)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT",
show=args.show)
if 1:
viscid.logger.info('Testing Spherical Sphere (theta, phi)...')
seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170,
pole=[-1, -1, -1], theta_phi=True)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP",
show=args.show)
if 1:
viscid.logger.info('Testing Spherical Cap (phi, theta)...')
seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80,
pole=[-1, -1, -1], theta_phi=False)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT",
view_kwargs=dict(azimuth=180, elevation=180), show=args.show)
if 1:
viscid.logger.info('Testing Spherical Cap (theta, phi)...')
seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80,
pole=[-1, -1, -1], theta_phi=True)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP",
view_kwargs=dict(azimuth=180, elevation=180), show=args.show)
if 1:
viscid.logger.info('Testing Spherical Patch...')
seeds = viscid.SphericalPatch(p0=[0, 0, 0], p1=[0, -0, -1],
max_alpha=30.0, max_beta=59.9,
nalpha=65, nbeta=80, r=0.5, roll=45.0)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)
if 1:
# this spline test is very custom
viscid.logger.info('Testing Spline...')
try:
import scipy.interpolate as interpolate
except ImportError:
msg = "XFail: ImportError (is scipy installed?)"
if plot2d:
try:
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
plt.clf()
plt.annotate(msg, xy=(0.3, 0.4), xycoords='axes fraction')
plt.savefig(next_plot_fname(__file__, series='2d'))
plt.savefig(next_plot_fname(__file__, series='2d'))
plt.savefig(next_plot_fname(__file__, series='3d'))
if args.show:
plt.show()
except ImportError:
pass
else:
knots = np.array([[ 0.2, 0.5, 0.0], [-0.2, 0.5, 0.2],
[-0.2, 0.0, 0.4], [ 0.2, 0.0, 0.2],
[ 0.2, -0.5, 0.0], [-0.2, -0.5, 0.2]]).T
seed_name = "Spline"
fld = logo
seeds = viscid.Spline(knots)
seed_pts = seeds.get_points()
interp_fld = viscid.interp_trilin(fld, seeds)
if plot2d:
try:
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
plt.clf()
vlt.plot(interp_fld)
plt.title(seed_name)
plt.savefig(next_plot_fname(__file__, series='2d'))
if args.show:
plt.show()
plt.clf()
from matplotlib import rcParams
_ms = rcParams['lines.markersize']
plt.gca().scatter(knots[0, :], knots[1, :],
s=(2 * _ms)**2, marker='^', color='y')
plt.gca().scatter(seed_pts[0, :], seed_pts[1, :],
s=(1.5 * _ms)**2, marker='o', color='k')
vlt.plot2d_line(seed_pts, scalars=interp_fld.flat_data,
symdir='z')
plt.title(seed_name)
plt.savefig(next_plot_fname(__file__, series='2d'))
if args.show:
plt.show()
except ImportError:
pass
if plot3d:
try:
vlab, _ = get_mvi_fig()
vlab.points3d(knots[0], knots[1], knots[2],
color=(1.0, 1.0, 0), scale_mode='none',
scale_factor=0.04)
p = vlab.points3d(seed_pts[0], seed_pts[1], seed_pts[2],
color=(0, 0, 0), scale_mode='none',
scale_factor=0.03)
vlab.plot_line(seed_pts, scalars=interp_fld.flat_data,
tube_radius=0.01)
vlab.axes(p)
vlab.title(seed_name)
vlab.mlab.roll(-90.0)
vlab.savefig(next_plot_fname(__file__, series='3d'))
if args.show:
vlab.show(stop=True)
except ImportError:
pass
if 1:
viscid.logger.info('Testing RectilinearMeshPoints...')
f = viscid.load_file(os.path.join(sample_dir, 'sample_xdmf.3d.[-1].xdmf'))
slc = 'x=-40j:12j, y=-10j:10j, z=-10j:10j'
b = f['b'][slc]
z = b.get_crd('z')
sheet_iz = np.argmin(b['x']**2, axis=2)
sheet_pts = b['z=0:1'].get_points()
sheet_pts[2, :] = z[sheet_iz].reshape(-1)
isphere_mask = np.sum(sheet_pts[:2, :]**2, axis=0) < 5**2
day_mask = sheet_pts[0:1, :] > -1.0
sheet_pts[2, :] = np.choose(isphere_mask, [sheet_pts[2, :], 0])
sheet_pts[2, :] = np.choose(day_mask, [sheet_pts[2, :], 0])
nx, ny, _ = b.sshape
sheet_seed = viscid.RectilinearMeshPoints(sheet_pts.reshape(3, nx, ny))
vx_sheet = viscid.interp_nearest(f['vx'], sheet_seed)
try:
if not plot2d:
raise ImportError
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
vlt.clf()
vlt.plot(vx_sheet, symmetric=True)
plt.savefig(next_plot_fname(__file__, series='2d'))
if args.show:
vlt.show()
except ImportError:
pass
try:
if not plot3d:
raise ImportError
vlab, _ = get_mvi_fig()
mesh = vlab.mesh_from_seeds(sheet_seed, scalars=vx_sheet,
clim=(-400, 400))
vlab.plot_earth_3d(crd_system=b)
vlab.view(azimuth=+90.0 + 45.0, elevation=90.0 - 25.0,
distance=30.0, focalpoint=(-10.0, +1.0, +1.0))
vlab.title("RectilinearMeshPoints")
vlab.savefig(next_plot_fname(__file__, series='3d'))
if args.show:
vlab.show(stop=True)
except ImportError:
pass
# prevent weird xorg bad-instructions on tear down
if 'figure' in _global_ns and _global_ns['figure'] is not None:
from viscid.plot import vlab
vlab.mlab.close(_global_ns['figure'])
return 0
def main():
mhd_type = "C3"
make_plots = 1
mhd_type = mhd_type.upper()
if mhd_type.startswith("C"):
if mhd_type in ("C",):
f = viscid.load_file("$WORK/tmedium/*.3d.[-1].xdmf")
elif mhd_type in ("C2", "C3"):
f = viscid.load_file("$WORK/tmedium2/*.3d.[-1].xdmf")
else:
raise ValueError()
catol = 1e-8
rtol = 2e-6
elif mhd_type in ("F", "FORTRAN"):
f = viscid.load_file("$WORK/tmedium3/*.3df.[-1]")
catol = 1e-8
rtol = 7e-2
else:
raise ValueError()
b = f['b_cc']
b1 = f['b_fc']
e_cc = f['e_cc']
e_ec = f['e_ec']
# divb = f['divB']
# viscid.interact()
if True:
bD = viscid.empty_like(b)
bD.data = np.array(b.data)
b1D = viscid.empty_like(b1)
b1D.data = np.array(b1.data)
mask5 = viscid.make_spherical_mask(bD, rmax=3.5)
mask1_5 = viscid.make_spherical_mask(bD, rmax=1.5)
viscid.fill_dipole(bD, mask=mask5)
viscid.set_in_region(bD, bD, 0.0, 0.0, mask=mask1_5, out=bD)
# compare_vectors(_b, bD, make_plots=True)
mask5 = viscid.make_spherical_mask(b1D, rmax=3.5)
mask1_5 = viscid.make_spherical_mask(b1D, rmax=1.5)
viscid.fill_dipole(b1D, mask=mask5)
viscid.set_in_region(b1D, b1D, 0.0, 0.0, mask=mask1_5, out=b1D)
compare_vectors(bD["x=1:-1, y=1:-1, z=1:-1"], b1D.as_cell_centered(),
make_plots=True)
# plt.clf()
# dkwargs = dict(symmetric=True, earth=True, clim=(-1e2, 1e2))
# ax1 = plt.subplot(311)
# vlt.plot(viscid.div(b1)['y=0j'], **dkwargs)
# plt.subplot(312, sharex=ax1, sharey=ax1)
# vlt.plot(viscid.div(b)['y=0j'], **dkwargs)
# plt.subplot(313, sharex=ax1, sharey=ax1)
# vlt.plot(viscid.div(b1D)['y=0j'], **dkwargs)
# vlt.show()
bD = b1D = mask5 = mask1_5 = None
# straight up interpolate b1 to cc crds and compare with b
if True:
b1_cc = viscid.interp_trilin(b1, b).as_flat()
viscid.set_in_region(b, b, alpha=0.0, beta=0.0, out=b,
mask=viscid.make_spherical_mask(b, rmax=5.0))
viscid.set_in_region(b1_cc, b1_cc, alpha=0.0, beta=0.0, out=b1_cc,
mask=viscid.make_spherical_mask(b1_cc, rmax=5.0))
compare_vectors(b, b1_cc, make_plots=True)
# make div?
if True:
# make seeds for 1.5x supersampling b1
n = 128
seeds = viscid.Volume((5.1, -0.02, -5.0), (12.0, 0.02, 5.0), (n, 3, n))
# do interpolation onto new seeds
b2 = viscid.interp_trilin(b1, seeds)
div_b = viscid.div(b)
div_b1 = viscid.div(b1)
div_b2 = viscid.div(b2)
viscid.set_in_region(div_b, div_b, alpha=0.0, beta=0.0, out=div_b,
mask=viscid.make_spherical_mask(div_b, rmax=5.0))
viscid.set_in_region(div_b1, div_b1, alpha=0.0, beta=0.0, out=div_b1,
mask=viscid.make_spherical_mask(div_b1, rmax=5.0))
viscid.set_in_region(div_b2, div_b2, alpha=0.0, beta=0.0, out=div_b2,
mask=viscid.make_spherical_mask(div_b2, rmax=5.0))
viscid.set_in_region(divb, divb, alpha=0.0, beta=0.0, out=divb,
mask=viscid.make_spherical_mask(divb, rmax=5.0))
plt.clf()
ax1 = vlt.subplot(311)
vlt.plot(div_b['y=0j'], symmetric=True, earth=True)
vlt.subplot(312, sharex=ax1, sharey=ax1)
# vlt.plot(div_b1['y=0j'], symmetric=True, earth=True)
vlt.plot(div_b2['y=0j'], symmetric=True, earth=True)
vlt.subplot(313, sharex=ax1, sharey=ax1)
vlt.plot(divb['y=0j'], symmetric=True, earth=True)
vlt.show()
return 0
def get_mp_info(pp, b, j, e, cache=True, cache_dir=None,
slc="x=5.5j:11.0j, y=-4.0j:4.0j, z=-3.6j:3.6j",
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.as_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", list(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
def main():
mhd_type = "C3"
make_plots = 1
mhd_type = mhd_type.upper()
if mhd_type.startswith("C"):
if mhd_type in ("C", ):
f = viscid.load_file("$WORK/tmedium/*.3d.[-1].xdmf")
elif mhd_type in ("C2", "C3"):
f = viscid.load_file("$WORK/tmedium2/*.3d.[-1].xdmf")
else:
raise ValueError()
catol = 1e-8
rtol = 2e-6
elif mhd_type in ("F", "FORTRAN"):
f = viscid.load_file("$WORK/tmedium3/*.3df.[-1]")
catol = 1e-8
rtol = 7e-2
else:
raise ValueError()
b = f['b_cc']
b1 = f['b_fc']
e_cc = f['e_cc']
e_ec = f['e_ec']
# divb = f['divB']
# viscid.interact()
if True:
bD = viscid.empty_like(b)
bD.data = np.array(b.data)
b1D = viscid.empty_like(b1)
b1D.data = np.array(b1.data)
mask5 = viscid.make_spherical_mask(bD, rmax=3.5)
mask1_5 = viscid.make_spherical_mask(bD, rmax=1.5)
viscid.fill_dipole(bD, mask=mask5)
viscid.set_in_region(bD, bD, 0.0, 0.0, mask=mask1_5, out=bD)
# compare_vectors(_b, bD, make_plots=True)
mask5 = viscid.make_spherical_mask(b1D, rmax=3.5)
mask1_5 = viscid.make_spherical_mask(b1D, rmax=1.5)
viscid.fill_dipole(b1D, mask=mask5)
viscid.set_in_region(b1D, b1D, 0.0, 0.0, mask=mask1_5, out=b1D)
compare_vectors(bD["x=1:-1, y=1:-1, z=1:-1"],
b1D.as_cell_centered(),
make_plots=True)
# plt.clf()
# dkwargs = dict(symmetric=True, earth=True, clim=(-1e2, 1e2))
# ax1 = plt.subplot(311)
# vlt.plot(viscid.div(b1)['y=0j'], **dkwargs)
# plt.subplot(312, sharex=ax1, sharey=ax1)
# vlt.plot(viscid.div(b)['y=0j'], **dkwargs)
# plt.subplot(313, sharex=ax1, sharey=ax1)
# vlt.plot(viscid.div(b1D)['y=0j'], **dkwargs)
# vlt.show()
bD = b1D = mask5 = mask1_5 = None
# straight up interpolate b1 to cc crds and compare with b
if True:
b1_cc = viscid.interp_trilin(b1, b).as_flat()
viscid.set_in_region(b,
b,
alpha=0.0,
beta=0.0,
out=b,
mask=viscid.make_spherical_mask(b, rmax=5.0))
viscid.set_in_region(b1_cc,
b1_cc,
alpha=0.0,
beta=0.0,
out=b1_cc,
mask=viscid.make_spherical_mask(b1_cc, rmax=5.0))
compare_vectors(b, b1_cc, make_plots=True)
# make div?
if True:
# make seeds for 1.5x supersampling b1
n = 128
seeds = viscid.Volume((5.1, -0.02, -5.0), (12.0, 0.02, 5.0), (n, 3, n))
# do interpolation onto new seeds
b2 = viscid.interp_trilin(b1, seeds)
div_b = viscid.div(b)
div_b1 = viscid.div(b1)
div_b2 = viscid.div(b2)
viscid.set_in_region(div_b,
div_b,
alpha=0.0,
beta=0.0,
out=div_b,
mask=viscid.make_spherical_mask(div_b, rmax=5.0))
viscid.set_in_region(div_b1,
div_b1,
alpha=0.0,
beta=0.0,
out=div_b1,
mask=viscid.make_spherical_mask(div_b1, rmax=5.0))
viscid.set_in_region(div_b2,
div_b2,
alpha=0.0,
beta=0.0,
out=div_b2,
mask=viscid.make_spherical_mask(div_b2, rmax=5.0))
viscid.set_in_region(divb,
divb,
alpha=0.0,
beta=0.0,
out=divb,
mask=viscid.make_spherical_mask(divb, rmax=5.0))
plt.clf()
ax1 = vlt.subplot(311)
vlt.plot(div_b['y=0j'], symmetric=True, earth=True)
vlt.subplot(312, sharex=ax1, sharey=ax1)
# vlt.plot(div_b1['y=0j'], symmetric=True, earth=True)
vlt.plot(div_b2['y=0j'], symmetric=True, earth=True)
vlt.subplot(313, sharex=ax1, sharey=ax1)
vlt.plot(divb['y=0j'], symmetric=True, earth=True)
vlt.show()
return 0
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 _main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--notwo", dest='notwo', action="store_true")
parser.add_argument("--nothree", dest='nothree', action="store_true")
parser.add_argument("--show", "--plot", action="store_true")
args = viscid.vutil.common_argparse(parser, default_verb=0)
plot2d = not args.notwo
plot3d = not args.nothree
# plot2d = True
# plot3d = True
# args.show = True
img = np.load(os.path.join(sample_dir, "logo.npy"))
x = np.linspace(-1, 1, img.shape[0])
y = np.linspace(-1, 1, img.shape[1])
z = np.linspace(-1, 1, img.shape[2])
logo = viscid.arrays2field([x, y, z], img)
if 1:
viscid.logger.info('Testing Point with custom local coordinates...')
pts = np.vstack([[-1, -0.5, 0, 0.5, 1],
[-1, -0.5, 0, 0.5, 1],
[ 0, 0.5, 1, 1.5, 2]])
local_crds = viscid.asarray_datetime64([0, 60, 120, 180, 240],
conservative=True)
seeds = viscid.Point(pts, local_crds=local_crds)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)
if 1:
viscid.logger.info('Testing Line...')
seeds = viscid.Line([-1, -1, 0], [1, 1, 2], n=5)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)
if 1:
viscid.logger.info('Testing Plane...')
seeds = viscid.Plane([0.0, 0.0, 0.0], [1, 1, 1], [1, 0, 0], 2, 2,
nl=160, nm=170, NL_are_vectors=True)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)
if 1:
viscid.logger.info('Testing Volume...')
seeds = viscid.Volume([-0.8, -0.8, -0.8], [0.8, 0.8, 0.8],
n=[64, 64, 3])
# note: can't make a 2d plot of the volume w/o a slice
run_test(logo, seeds, plot2d=False, plot3d=plot3d, add_title="3d",
show=args.show)
if 1:
viscid.logger.info('Testing Volume (with ignorable dim)...')
seeds = viscid.Volume([-0.8, -0.8, 0.0], [0.8, 0.8, 0.0],
n=[64, 64, 1])
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="2d",
show=args.show)
if 1:
viscid.logger.info('Testing Spherical Sphere (phi, theta)...')
seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170,
pole=[-1, -1, -1], theta_phi=False)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT",
show=args.show)
if 1:
viscid.logger.info('Testing Spherical Sphere (theta, phi)...')
seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170,
pole=[-1, -1, -1], theta_phi=True)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP",
show=args.show)
if 1:
viscid.logger.info('Testing Spherical Cap (phi, theta)...')
seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80,
pole=[-1, -1, -1], theta_phi=False)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT",
view_kwargs=dict(azimuth=180, elevation=180), show=args.show)
if 1:
viscid.logger.info('Testing Spherical Cap (theta, phi)...')
seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80,
pole=[-1, -1, -1], theta_phi=True)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP",
view_kwargs=dict(azimuth=180, elevation=180), show=args.show)
if 1:
viscid.logger.info('Testing Spherical Patch...')
seeds = viscid.SphericalPatch(p0=[0, 0, 0], p1=[0, -0, -1],
max_alpha=30.0, max_beta=59.9,
nalpha=65, nbeta=80, r=0.5, roll=45.0)
run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)
if 1:
# this spline test is very custom
viscid.logger.info('Testing Spline...')
try:
import scipy.interpolate as interpolate
except ImportError:
msg = "XFail: ImportError (is scipy installed?)"
if plot2d:
try:
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
plt.clf()
plt.annotate(msg, xy=(0.3, 0.4), xycoords='axes fraction')
plt.savefig(next_plot_fname(__file__, series='2d'))
plt.savefig(next_plot_fname(__file__, series='2d'))
plt.savefig(next_plot_fname(__file__, series='3d'))
if args.show:
plt.show()
except ImportError:
pass
else:
knots = np.array([[ 0.2, 0.5, 0.0], [-0.2, 0.5, 0.2],
[-0.2, 0.0, 0.4], [ 0.2, 0.0, 0.2],
[ 0.2, -0.5, 0.0], [-0.2, -0.5, 0.2]]).T
seed_name = "Spline"
fld = logo
seeds = viscid.Spline(knots)
seed_pts = seeds.get_points()
interp_fld = viscid.interp_trilin(fld, seeds)
if plot2d:
try:
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
plt.clf()
vlt.plot(interp_fld)
plt.title(seed_name)
plt.savefig(next_plot_fname(__file__, series='2d'))
if args.show:
plt.show()
plt.clf()
from matplotlib import rcParams
_ms = rcParams['lines.markersize']
plt.gca().scatter(knots[0, :], knots[1, :],
s=(2 * _ms)**2, marker='^', color='y')
plt.gca().scatter(seed_pts[0, :], seed_pts[1, :],
s=(1.5 * _ms)**2, marker='o', color='k')
vlt.plot2d_line(seed_pts, scalars=interp_fld.flat_data,
symdir='z')
plt.title(seed_name)
plt.savefig(next_plot_fname(__file__, series='2d'))
if args.show:
plt.show()
except ImportError:
pass
if plot3d:
try:
from viscid.plot import vlab
_ = get_mvi_fig(offscreen=not args.show)
vlab.points3d(knots[0], knots[1], knots[2],
color=(1.0, 1.0, 0), scale_mode='none',
scale_factor=0.04)
p = vlab.points3d(seed_pts[0], seed_pts[1], seed_pts[2],
color=(0, 0, 0), scale_mode='none',
scale_factor=0.03)
vlab.plot_line(seed_pts, scalars=interp_fld.flat_data,
tube_radius=0.01)
vlab.axes(p)
vlab.title(seed_name)
vlab.mlab.roll(-90.0)
vlab.savefig(next_plot_fname(__file__, series='3d'))
if args.show:
vlab.show(stop=True)
except ImportError:
pass
if 1:
viscid.logger.info('Testing RectilinearMeshPoints...')
f = viscid.load_file(os.path.join(sample_dir, 'sample_xdmf.3d.[-1].xdmf'))
slc = 'x=-40j:12j, y=-10j:10j, z=-10j:10j'
b = f['b'][slc]
z = b.get_crd('z')
sheet_iz = np.argmin(b['x']**2, axis=2)
sheet_pts = b['z=0:1'].get_points()
sheet_pts[2, :] = z[sheet_iz].reshape(-1)
isphere_mask = np.sum(sheet_pts[:2, :]**2, axis=0) < 5**2
day_mask = sheet_pts[0:1, :] > -1.0
sheet_pts[2, :] = np.choose(isphere_mask, [sheet_pts[2, :], 0])
sheet_pts[2, :] = np.choose(day_mask, [sheet_pts[2, :], 0])
nx, ny, _ = b.sshape
sheet_seed = viscid.RectilinearMeshPoints(sheet_pts.reshape(3, nx, ny))
vx_sheet = viscid.interp_nearest(f['vx'], sheet_seed)
try:
if not plot2d:
raise ImportError
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
vlt.clf()
vlt.plot(vx_sheet, symmetric=True)
plt.savefig(next_plot_fname(__file__, series='2d'))
if args.show:
vlt.show()
except ImportError:
pass
try:
if not plot3d:
raise ImportError
from viscid.plot import vlab
_ = get_mvi_fig(offscreen=not args.show)
mesh = vlab.mesh_from_seeds(sheet_seed, scalars=vx_sheet,
clim=(-400, 400))
vlab.plot_earth_3d(crd_system=b)
vlab.view(azimuth=+90.0 + 45.0, elevation=90.0 - 25.0,
distance=30.0, focalpoint=(-10.0, +1.0, +1.0))
vlab.title("RectilinearMeshPoints")
vlab.savefig(next_plot_fname(__file__, series='3d'))
if args.show:
vlab.show(stop=True)
except ImportError:
pass
# prevent weird xorg bad-instructions on tear down
if 'figure' in _global_ns and _global_ns['figure'] is not None:
from viscid.plot import vlab
vlab.mlab.close(_global_ns['figure'])
return 0
def prepare_lines(lines, scalars=None, do_connections=False, other=None):
"""Concatenate and standardize a list of lines
Args:
lines (list): Must be a list of 3xN or 4xN ndarrays of xyz(s)
data for N points along the line. N need not be the same
for all lines. Can alse be 6xN such that lines[:][3:, :]
are interpreted as rgb colors
scalars (ndarray, list): Can have shape 1xN for a single scalar
or 3xN for an rgb color for each point. If the shape is
1xNlines, the scalar is broadcast so the whole line gets
the same value, and likewise for 3xNlines and rgb colors.
Can also be a list of hex color (#ffffff) strings.
Otherwise, scalars is reshaped to -1xN.
do_connections (bool): Whether or not to make connections array
other (dict): a dictionary of other arrays that should be
reshaped and the like the same way scalars is
Returns:
(vertices, scalars, connections, other)
* vertices (ndarray): 3xN array of N xyz points. N is the sum
of the lengths of all the lines
* scalars (ndarray): N array of scalars, 3xN array of uint8
rgb values, or None
* connections (ndarray): Nx2 array of ints (indices along
axis 1 of vertices) describing the forward and backward
connectedness of the lines, or None
* other (dict): a dict of N length arrays
Raises:
ValueError: If rgb data is not in a valid range or the shape
of scalars is not understood
"""
nlines = len(lines)
npts = [line.shape[1] for line in lines]
N = np.sum(npts)
first_idx = np.cumsum([0] + npts[:-1])
vertices = [np.asarray(line) for line in lines]
vertices = np.concatenate(lines, axis=1)
if vertices.shape[0] > 3:
if scalars is not None:
viscid.logger.warn("Overriding line scalars with scalars kwarg")
else:
scalars = vertices[3:, :]
vertices = vertices[:3, :]
if scalars is not None:
if isinstance(scalars, viscid.field.Field):
scalars = viscid.interp_trilin(scalars, vertices)
if scalars.size != N:
raise ValueError("Scalars was not a scalar field")
scalars = np.atleast_2d(scalars)
if scalars.shape == (1, 1):
scalars = scalars.repeat(N, axis=1)
elif scalars.shape == (1, nlines) or scalars.shape == (nlines, 1):
# one scalar for each line, so broadcast it
scalars = scalars.reshape(nlines, 1)
scalars = [scalars[i].repeat(ni) for i, ni in enumerate(npts)]
scalars = np.concatenate(scalars, axis=0).reshape(1, N)
elif scalars.shape == (N, 1) or scalars.shape == (1, N):
# catch these so they're not interpreted as colors if
# nlines == 1 and N == 3; ie. 1 line with 3 points
scalars = scalars.reshape(1, N)
elif scalars.shape == (3, nlines) or scalars.shape == (nlines, 3):
# one rgb color for each line, so broadcast it
if scalars.shape == (3, nlines):
scalars = scalars.T
colors = []
for i, ni in enumerate(npts):
c = scalars[i].reshape(3, 1).repeat(ni, axis=1)
colors.append(c)
scalars = np.concatenate(colors, axis=1)
else:
scalars = scalars.reshape(-1, N)
if scalars.dtype.kind == 'S':
# translate hex colors (#ff00ff) into rgb values
scalars = np.char.lstrip(scalars, '#').astype('S6')
scalars = np.char.zfill(scalars, 6)
scalars = np.frombuffer(np.char.decode(scalars, 'hex'), dtype='u1')
scalars = scalars.reshape(-1, 3).T
elif scalars.shape[0] == 1:
# normal scalars
scalars = scalars.reshape(-1)
elif scalars.shape[0] == 3:
# The scalars encode rgb data, standardize the result to a
# 3xN ndarray of 1 byte unsigned ints (chars)
if np.all(scalars >= 0) and np.all(scalars <= 1):
scalars = (255 * scalars).round().astype('u1')
elif np.all(scalars >= 0) and np.all(scalars < 256):
scalars = scalars.round().astype('u1')
else:
raise ValueError("Rgb data should be in range [0, 1] or "
"[0, 255], range given is [{0}, {1}]"
"".format(np.min(scalars), np.max(scalars)))
else:
raise ValueError("Scalars should either be a number, or set of "
"rgb values, shape is {0}".format(scalars.shape))
# broadcast / reshape additional arrays given in other
if other:
for key, arr in other.items():
if arr is None:
pass
elif arr.shape == (1, nlines) or arr.shape == (nlines, 1):
arr = arr.reshape(nlines, 1)
arr = [arr[i].repeat(ni) for i, ni in enumerate(npts)]
other[key] = np.concatenate(arr, axis=0).reshape(1, N)
else:
try:
other[key] = arr.reshape(-1, N)
except ValueError:
viscid.logger.warn("Unknown dimension, dropping array {0}"
"".format(key))
if do_connections:
connections = [None] * nlines
for i, ni in enumerate(npts):
# i0 is the index of the first point of the i'th line in lines
i0 = first_idx[i]
connections[i] = np.vstack([np.arange(i0, i0 + ni - 1.5),
np.arange(i0 + 1, i0 + ni - 0.5)]).T
connections = np.concatenate(connections, axis=0).astype('i')
else:
connections = None
return vertices, scalars, connections, other
