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 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 _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 _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 _get_sep_pts_bisect(fld,
seed,
trace_opts=None,
min_depth=3,
max_depth=7,
plot=False,
perimeter_check=perimeter_check_bitwise_or,
make_3d=True,
start_uneven=False,
_base_quadrent="",
_uneven_mask=0,
_first_recurse=True):
if len(_base_quadrent) == max_depth:
return [_base_quadrent] # causes pylint to complain
if trace_opts is None:
trace_opts = dict()
nx, ny = seed.uv_shape
(xlim, ylim) = seed.uv_extent
if _first_recurse and start_uneven:
_uneven_mask = UNEVEN_MASK
if _first_recurse and plot:
from viscid.plot import vlab
from viscid.plot import vpyplot as vlt
vlt.clf()
_, all_topo = viscid.calc_streamlines(fld, seed, **trace_opts)
vlt.plot(np.bitwise_and(all_topo, 15), show=False)
verts, arr = seed.wrap_mesh(all_topo.data)
vlab.mesh(verts[0], verts[1], verts[2], scalars=arr, opacity=0.75)
# quadrents and lines are indexed as follows...
# directions are counter clackwise around the quadrent with
# lower index (which matters for lines which are shared among
# more than one quadrent, aka, lines 1,2,6,7). Notice that even
# numbered lines are horizontal, like the interstate system :)
# -<--10-----<-8---
# | ^ ^
# 11 2 9 3 7
# \/ | |
# --<-2-----<-6----
# | ^ ^
# 3 0 1 1 5
# \/ | |
# ----0->-----4->--
# find low(left), mid(center), and high(right) crds in x and y
low_quad = "{0}{1:x}".format(_base_quadrent, 0 | _uneven_mask)
high_quad = "{0}{1:x}".format(_base_quadrent, 3 | _uneven_mask)
xl, xm, yl, ym = _quadrent_limits(low_quad, xlim, ylim)
_, xh, _, yh = _quadrent_limits(high_quad, xlim, ylim)
segsx, segsy = [None] * 12, [None] * 12
topo = [None] * 12
nxm, nym = nx // 2, ny // 2
# make all the line segments
segsx[0], segsy[0] = np.linspace(xl, xm, nxm), np.linspace(yl, yl, nxm)
segsx[1], segsy[1] = np.linspace(xm, xm, nym), np.linspace(yl, ym, nym)
segsx[2], segsy[2] = np.linspace(xm, xl, nxm), np.linspace(ym, ym, nxm)
segsx[3], segsy[3] = np.linspace(xl, xl, nym), np.linspace(ym, yl, nym)
segsx[4], segsy[4] = np.linspace(xm, xh, nxm), np.linspace(yl, yl, nxm)
segsx[5], segsy[5] = np.linspace(xh, xh, nym), np.linspace(yl, ym, nym)
segsx[6], segsy[6] = np.linspace(xh, xm, nxm), np.linspace(ym, ym, nxm)
segsx[7], segsy[7] = np.linspace(xh, xh, nym), np.linspace(ym, yh, nym)
segsx[8], segsy[8] = np.linspace(xh, xm, nxm), np.linspace(yh, yh, nxm)
segsx[9], segsy[9] = np.linspace(xm, xm, nym), np.linspace(ym, yh, nym)
segsx[10], segsy[10] = np.linspace(xm, xl, nxm), np.linspace(yh, yh, nxm)
segsx[11], segsy[11] = np.linspace(xl, xl, nym), np.linspace(yh, ym, nym)
allx = np.concatenate(segsx)
ally = np.concatenate(segsy)
# print("plot::", _base_quadrent, '|', _uneven_mask, '|', len(allx), len(ally))
pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
_, all_topo = viscid.calc_streamlines(fld, pts3d, **trace_opts)
topo[0] = all_topo[:len(segsx[0])]
cnt = len(topo[0])
for i, segx in zip(count(1), segsx[1:]):
topo[i] = all_topo[cnt:cnt + len(segx)]
# print("??", i, cnt, cnt + len(segx), np.bitwise_and.reduce(topo[i]))
cnt += len(topo[i])
# assemble the lines into the four quadrents
quad_topo = [None] * 4
# all arrays snip off the last element since those are
# duplicated by the next line... reversed arrays do the
# snipping with -1:0:-1
quad_topo[0] = np.concatenate(
[topo[0][:-1], topo[1][:-1], topo[2][:-1], topo[3][:-1]])
quad_topo[1] = np.concatenate(
[topo[4][:-1], topo[5][:-1], topo[6][:-1], topo[1][-1:0:-1]])
quad_topo[2] = np.concatenate(
[topo[2][-1:0:-1], topo[9][:-1], topo[10][:-1], topo[11][:-1]])
quad_topo[3] = np.concatenate(
[topo[6][-1:0:-1], topo[7][:-1], topo[8][:-1], topo[9][-1:0:-1]])
# now that the quad arrays are populated, decide which quadrents
# still contain the separator (could be > 1)
required_uneven_subquads = False
ret = []
for i in range(4):
if perimeter_check(quad_topo[i]):
next_quad = "{0}{1:x}".format(_base_quadrent, i | _uneven_mask)
subquads = _get_sep_pts_bisect(fld,
seed,
trace_opts=trace_opts,
min_depth=min_depth,
max_depth=max_depth,
plot=plot,
_base_quadrent=next_quad,
_uneven_mask=0,
_first_recurse=False)
ret += subquads
if len(ret) == 0:
perimeter = np.concatenate([
topo[0][::-1], topo[4][::-1], topo[5][::-1], topo[7][::-1],
topo[8][::-1], topo[10][::-1], topo[11][::-1], topo[3][::-1]
])
if _uneven_mask:
if len(_base_quadrent) > min_depth:
print("sep trace issue, but min depth reached: {0} > {1}"
"".format(len(_base_quadrent), min_depth))
ret = [_base_quadrent]
else:
print("sep trace issue, the separator ended prematurely")
elif perimeter_check(perimeter):
ret = _get_sep_pts_bisect(fld,
seed,
trace_opts=trace_opts,
min_depth=min_depth,
max_depth=max_depth,
plot=plot,
_base_quadrent=_base_quadrent,
_uneven_mask=UNEVEN_MASK,
_first_recurse=False)
required_uneven_subquads = True
if plot and not required_uneven_subquads:
from viscid.plot import vlab
from matplotlib import pyplot as plt
from viscid.plot import vpyplot as vlt
_pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
vlab.points3d(_pts3d[0],
_pts3d[1],
_pts3d[2],
all_topo.data.reshape(-1),
scale_mode='none',
scale_factor=0.02)
plt.scatter(allx,
ally,
color=np.bitwise_and(all_topo, 15),
vmin=0,
vmax=15,
marker='o',
edgecolor='y',
s=40)
if _first_recurse:
# turn quadrent strings into locations
xc = np.empty(len(ret))
yc = np.empty(len(ret))
for i, r in enumerate(ret):
xc[i], yc[i] = _quadrent_center(r, xlim, ylim)
pts_uv = np.array([xc, yc])
if plot:
from viscid.plot import vlab
from matplotlib import pyplot as plt
from viscid.plot import vpyplot as vlt
plt.plot(pts_uv[0],
pts_uv[1],
"y*",
ms=20,
markeredgecolor='k',
markeredgewidth=1.0)
vlt.show(block=False)
vlab.show(stop=True)
# return seed.to_3d(seed.uv_to_local(pts_uv))
# if pts_uv.size == 0:
# return None
if make_3d:
return seed.uv_to_3d(pts_uv)
else:
return pts_uv
else:
return ret
def trace_separator(grid,
b_slcstr="x=-25f:15f, y=-30f:30f, z=-15f:15f",
r=1.0,
plot=False,
trace_opts=None,
cache=True,
cache_dir=None):
"""Trace a separator line from most dawnward null
**Still in testing** Uses the bisection algorithm.
Args:
grid (Grid): A grid that has a "b" field
b_slcstr (str): Some valid slice for B field
r (float): spatial step of separator line
plot (bool): make debugging plots
trace_opts (dict): passed to streamline function
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
Raises:
IOError: Description
Returns:
tuple: (separator_lines, nulls)
- **separator_lines** (list): list of M 3xN ndarrays that
represent M separator lines with N points
- **nulls** (ndarray): 3xN array of N null points
"""
if not cache_dir:
cache_dir = grid.find_info("_viscid_dirname", "./")
run_name = grid.find_info("run")
sep_fname = "{0}/{1}.sep.{2:06.0f}".format(cache_dir, run_name, grid.time)
try:
if isinstance(cache,
string_types) and cache.strip().lower() == "force":
raise IOError()
with np.load(sep_fname + ".npz") as dat:
sep_iter = (f for f in dat.files if f.startswith("arr_"))
_it = sorted(sep_iter, key=lambda s: int(s[len("arr_"):]))
seps = [dat[n] for n in _it]
nulls = dat['nulls']
except IOError:
_b = grid['b'][b_slcstr]
_, nulls = viscid.find_nulls(_b['x=-30f:15f'], ibound=5.0)
# get most dawnward null, nulls2 is all nulls except p0
nullind = np.argmin(nulls[1, :])
p0 = nulls[:, nullind]
nulls2 = np.concatenate([nulls[:, :nullind], nulls[:, (nullind + 1):]],
axis=1)
if plot:
from viscid.plot import vlab
vlab.plot_earth_3d(crd_system='gse')
vlab.points3d(nulls2[0],
nulls2[1],
nulls2[2],
color=(0, 0, 0),
scale_factor=1.0)
vlab.points3d(nulls[0, nullind],
nulls[1, nullind],
nulls[2, nullind],
color=(1, 1, 1),
scale_factor=1.0)
seed = viscid.Sphere(p0=p0,
r=r,
ntheta=30,
nphi=60,
theta_endpoint=True,
phi_endpoint=True)
p1 = viscid.get_sep_pts_bisect(_b,
seed,
max_depth=12,
plot=plot,
trace_opts=trace_opts)
# print("p1 shape", p1.shape)
# if p1.shape[1] > 2:
# raise RuntimeError("Invalid B field, should be no branch @ null")
seps = []
sep_stubs = []
for i in range(p1.shape[1]):
sep_stubs.append([p0, p1[:, i]])
# print("??", sep_stubs)
while sep_stubs:
sep = sep_stubs.pop(0)
# print("!!! new stub")
for i in count():
# print("::", i)
seed = viscid.SphericalPatch(p0=sep[-1],
p1=sep[-1] - sep[-2],
r=r,
nalpha=240,
nbeta=240)
pn = viscid.get_sep_pts_bisect(_b,
seed,
max_depth=8,
plot=plot,
trace_opts=trace_opts)
if pn.shape[1] == 0:
# print("END: pn.shape[1] == 0")
break
# print("++", nulls2.shape, pn.shape)
closest_null_dist = np.min(
np.linalg.norm(nulls2 - pn[:, :1], axis=0))
# print("closest_null_dist:", closest_null_dist)
if closest_null_dist < 1.01 * r:
# print("END: within 1.01 of a null")
break
# print("??", pn)
for j in range(1, pn.shape[1]):
# print("inserting new stub")
sep_stubs.insert(0, [sep[-1], pn[:, j]])
sep.append(pn[:, 0])
# print("sep", sep)
seps.append(np.stack(sep, axis=1))
if cache:
np.savez_compressed(sep_fname, *seps, nulls=nulls)
return seps, nulls
def _main():
f = viscid.load_file("$WORK/xi_fte_001/*.3d.[4050f].xdmf")
mp = get_mp_info(f['pp'], f['b'], f['j'], f['e_cc'], fit='mp_xloc',
slc="x=6.5j:10.5j, y=-4j:4j, z=-4.8j:3j", cache=False)
y, z = mp['pp_max_xloc'].meshgrid_flat(prune=True)
x = mp['pp_max_xloc'].data.reshape(-1)
Y, Z = mp['pp_max_xloc'].meshgrid(prune=True)
x2 = paraboloid(Y, Z, *mp['paraboloid'][0])
skip = 117
n = paraboloid_normal(Y, Z, *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]
minvar_y = Y.reshape(-1)[::skip]
minvar_z = Z.reshape(-1)[::skip]
minvar_n = np.zeros([3, len(minvar_y)])
for i in range(minvar_n.shape[0]):
p0 = [0.0, minvar_y[i], minvar_z[i]]
p0[0] = mp['pp_max_xloc']['y={0[0]}f, z={0[1]}f'.format(p0)]
minvar_n[:, i] = viscid.find_minvar_lmn_around(f['b'], p0, l=2.0, n=64)[2, :]
# 2d plots, normals don't look normal in the matplotlib projection
if False: # pylint: disable=using-constant-test
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
normals = paraboloid_normal(Y, Z, *mp['paraboloid'][0])
p0 = np.array([x2, Y, Z]).reshape(3, -1)
p1 = p0 + normals.reshape(3, -1)
vlt.scatter_3d(np.vstack([x, y, z])[:, ::skip], equal=True)
for i in range(0, p0.shape[1], skip):
plt.gca().plot([p0[0, i], p1[0, i]],
[p0[1, i], p1[1, i]],
[p0[2, i], p1[2, i]], color='c')
# z2 = _ellipsiod(X, Y, *popt)
plt.gca().plot_surface(Y, Z, x2, color='r')
vlt.show()
# mayavi 3d plots, normals look better here
if True: # pylint: disable=using-constant-test
from viscid.plot import vlab
vlab.points3d(x[::skip], y[::skip], z[::skip], scale_factor=0.25,
color=(0.0, 0.0, 1.0))
mp_width = mp['mp_width']['x=0']
mp_sheath_edge = mp['mp_sheath_edge']['x=0']
mp_sphere_edge = mp_sheath_edge - mp_width
vlab.mesh(x2, Y, Z, scalars=mp_width.data)
vlab.mesh(mp_sheath_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)
vlab.mesh(mp_sphere_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)
n = paraboloid_normal(Y, Z, *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]
vlab.quiver3d(x2.reshape(-1)[::skip],
Y.reshape(-1)[::skip],
Z.reshape(-1)[::skip],
n[0], n[1], n[2], color=(1, 0, 0))
vlab.quiver3d(x2.reshape(-1)[::skip],
Y.reshape(-1)[::skip],
Z.reshape(-1)[::skip],
minvar_n[0], minvar_n[1], minvar_n[2], color=(0, 0, 1))
vlab.show()
def _main():
crd_system = 'gse'
print(viscid.get_dipole_moment_ang(dip_tilt=45.0, dip_gsm=0.0,
crd_system=crd_system))
print(viscid.get_dipole_moment_ang(dip_tilt=0.0, dip_gsm=45.0,
crd_system=crd_system))
print(viscid.get_dipole_moment_ang(dip_tilt=45.0, dip_gsm=45.0,
crd_system=crd_system))
print("---")
ptsNP = np.array([[+2, -2, +2], [+2, -1, +2], [+2, 1, +2], [+2, 2, +2]]).T
ptsSP = np.array([[+2, -2, -2], [+2, -1, -2], [+2, 1, -2], [+2, 2, -2]]).T
ptsNN = np.array([[-2, -2, +2], [-2, -1, +2], [-2, 1, +2], [-2, 2, +2]]).T
ptsSN = np.array([[-2, -2, -2], [-2, -1, -2], [-2, 1, -2], [-2, 2, -2]]).T
mapped_ptsNP = dipole_map(ptsNP)
mapped_ptsNN = dipole_map(ptsNN)
mapped_ptsSP = dipole_map(ptsSP)
mapped_ptsSN = dipole_map(ptsSN)
try:
from viscid.plot import vlab
colors1 = np.array([(0.6, 0.2, 0.2),
(0.2, 0.2, 0.6),
(0.6, 0.6, 0.2),
(0.2, 0.6, 0.6)])
colors2 = colors1 * 0.5
vlab.points3d(ptsNP, scale_factor=0.4, color=tuple(colors1[0]))
vlab.points3d(ptsNN, scale_factor=0.4, color=tuple(colors1[1]))
vlab.points3d(ptsSP, scale_factor=0.4, color=tuple(colors1[2]))
vlab.points3d(ptsSN, scale_factor=0.4, color=tuple(colors1[3]))
vlab.points3d(mapped_ptsNP, scale_factor=0.4, color=tuple(colors2[0]))
vlab.points3d(mapped_ptsNN, scale_factor=0.4, color=tuple(colors2[1]))
vlab.points3d(mapped_ptsSP, scale_factor=0.4, color=tuple(colors2[2]))
vlab.points3d(mapped_ptsSN, scale_factor=0.4, color=tuple(colors2[3]))
b = make_dipole()
vlab.plot_lines(viscid.calc_streamlines(b, mapped_ptsNP, ibound=0.5)[0])
vlab.plot_lines(viscid.calc_streamlines(b, mapped_ptsNN, ibound=0.5)[0])
vlab.show()
except ImportError:
print("Mayavi not installed, no 3D plots", file=sys.stderr)
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():
f = viscid.load_file("$WORK/xi_fte_001/*.3d.[4050f].xdmf")
mp = get_mp_info(f['pp'],
f['b'],
f['j'],
f['e_cc'],
fit='mp_xloc',
slc="x=6.5f:10.5f, y=-4f:4f, z=-4.8f:3f",
cache=False)
y, z = mp['pp_max_xloc'].meshgrid_flat(prune=True)
x = mp['pp_max_xloc'].data.reshape(-1)
Y, Z = mp['pp_max_xloc'].meshgrid(prune=True)
x2 = paraboloid(Y, Z, *mp['paraboloid'][0])
skip = 117
n = paraboloid_normal(Y, Z, *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]
minvar_y = Y.reshape(-1)[::skip]
minvar_z = Z.reshape(-1)[::skip]
minvar_n = np.zeros([3, len(minvar_y)])
for i in range(minvar_n.shape[0]):
p0 = [0.0, minvar_y[i], minvar_z[i]]
p0[0] = mp['pp_max_xloc']['y={0[0]}f, z={0[1]}f'.format(p0)]
minvar_n[:, i] = viscid.find_minvar_lmn_around(f['b'], p0, l=2.0,
n=64)[2, :]
# 2d plots, normals don't look normal in the matplotlib projection
if False: # pylint: disable=using-constant-test
from matplotlib import pyplot as plt
from viscid.plot import vpyplot as vlt
normals = paraboloid_normal(Y, Z, *mp['paraboloid'][0])
p0 = np.array([x2, Y, Z]).reshape(3, -1)
p1 = p0 + normals.reshape(3, -1)
vlt.scatter_3d(np.vstack([x, y, z])[:, ::skip], equal=True)
for i in range(0, p0.shape[1], skip):
plt.gca().plot([p0[0, i], p1[0, i]], [p0[1, i], p1[1, i]],
[p0[2, i], p1[2, i]],
color='c')
# z2 = _ellipsiod(X, Y, *popt)
plt.gca().plot_surface(Y, Z, x2, color='r')
vlt.show()
# mayavi 3d plots, normals look better here
if True: # pylint: disable=using-constant-test
from viscid.plot import vlab
vlab.points3d(x[::skip],
y[::skip],
z[::skip],
scale_factor=0.25,
color=(0.0, 0.0, 1.0))
mp_width = mp['mp_width']['x=0']
mp_sheath_edge = mp['mp_sheath_edge']['x=0']
mp_sphere_edge = mp_sheath_edge - mp_width
vlab.mesh(x2, Y, Z, scalars=mp_width.data)
vlab.mesh(mp_sheath_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)
vlab.mesh(mp_sphere_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)
n = paraboloid_normal(Y, Z,
*mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]
vlab.quiver3d(x2.reshape(-1)[::skip],
Y.reshape(-1)[::skip],
Z.reshape(-1)[::skip],
n[0],
n[1],
n[2],
color=(1, 0, 0))
vlab.quiver3d(x2.reshape(-1)[::skip],
Y.reshape(-1)[::skip],
Z.reshape(-1)[::skip],
minvar_n[0],
minvar_n[1],
minvar_n[2],
color=(0, 0, 1))
vlab.show()
def _get_sep_pts_bisect(fld, seed, trace_opts=None, min_depth=3, max_depth=7,
plot=False, perimeter_check=perimeter_check_bitwise_or,
make_3d=True, start_uneven=False, _base_quadrent="",
_uneven_mask=0, _first_recurse=True):
if len(_base_quadrent) == max_depth:
return [_base_quadrent] # causes pylint to complain
if trace_opts is None:
trace_opts = dict()
nx, ny = seed.uv_shape
(xlim, ylim) = seed.uv_extent
if _first_recurse and start_uneven:
_uneven_mask = UNEVEN_MASK
if _first_recurse and plot:
from viscid.plot import vlab
from viscid.plot import vpyplot as vlt
vlt.clf()
_, all_topo = viscid.calc_streamlines(fld, seed, **trace_opts)
vlt.plot(np.bitwise_and(all_topo, 15), show=False)
verts, arr = seed.wrap_mesh(all_topo.data)
vlab.mesh(verts[0], verts[1], verts[2], scalars=arr, opacity=0.75)
# quadrents and lines are indexed as follows...
# directions are counter clackwise around the quadrent with
# lower index (which matters for lines which are shared among
# more than one quadrent, aka, lines 1,2,6,7). Notice that even
# numbered lines are horizontal, like the interstate system :)
# -<--10-----<-8---
# | ^ ^
# 11 2 9 3 7
# \/ | |
# --<-2-----<-6----
# | ^ ^
# 3 0 1 1 5
# \/ | |
# ----0->-----4->--
# find low(left), mid(center), and high(right) crds in x and y
low_quad = "{0}{1:x}".format(_base_quadrent, 0 | _uneven_mask)
high_quad = "{0}{1:x}".format(_base_quadrent, 3 | _uneven_mask)
xl, xm, yl, ym = _quadrent_limits(low_quad, xlim, ylim)
_, xh, _, yh = _quadrent_limits(high_quad, xlim, ylim)
segsx, segsy = [None] * 12, [None] * 12
topo = [None] * 12
nxm, nym = nx //2, ny // 2
# make all the line segments
segsx[0], segsy[0] = np.linspace(xl, xm, nxm), np.linspace(yl, yl, nxm)
segsx[1], segsy[1] = np.linspace(xm, xm, nym), np.linspace(yl, ym, nym)
segsx[2], segsy[2] = np.linspace(xm, xl, nxm), np.linspace(ym, ym, nxm)
segsx[3], segsy[3] = np.linspace(xl, xl, nym), np.linspace(ym, yl, nym)
segsx[4], segsy[4] = np.linspace(xm, xh, nxm), np.linspace(yl, yl, nxm)
segsx[5], segsy[5] = np.linspace(xh, xh, nym), np.linspace(yl, ym, nym)
segsx[6], segsy[6] = np.linspace(xh, xm, nxm), np.linspace(ym, ym, nxm)
segsx[7], segsy[7] = np.linspace(xh, xh, nym), np.linspace(ym, yh, nym)
segsx[8], segsy[8] = np.linspace(xh, xm, nxm), np.linspace(yh, yh, nxm)
segsx[9], segsy[9] = np.linspace(xm, xm, nym), np.linspace(ym, yh, nym)
segsx[10], segsy[10] = np.linspace(xm, xl, nxm), np.linspace(yh, yh, nxm)
segsx[11], segsy[11] = np.linspace(xl, xl, nym), np.linspace(yh, ym, nym)
allx = np.concatenate(segsx)
ally = np.concatenate(segsy)
# print("plot::", _base_quadrent, '|', _uneven_mask, '|', len(allx), len(ally))
pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
_, all_topo = viscid.calc_streamlines(fld, pts3d, **trace_opts)
topo[0] = all_topo[:len(segsx[0])]
cnt = len(topo[0])
for i, segx in zip(count(1), segsx[1:]):
topo[i] = all_topo[cnt:cnt + len(segx)]
# print("??", i, cnt, cnt + len(segx), np.bitwise_and.reduce(topo[i]))
cnt += len(topo[i])
# assemble the lines into the four quadrents
quad_topo = [None] * 4
# all arrays snip off the last element since those are
# duplicated by the next line... reversed arrays do the
# snipping with -1:0:-1
quad_topo[0] = np.concatenate([topo[0][:-1], topo[1][:-1],
topo[2][:-1], topo[3][:-1]])
quad_topo[1] = np.concatenate([topo[4][:-1], topo[5][:-1],
topo[6][:-1], topo[1][-1:0:-1]])
quad_topo[2] = np.concatenate([topo[2][-1:0:-1], topo[9][:-1],
topo[10][:-1], topo[11][:-1]])
quad_topo[3] = np.concatenate([topo[6][-1:0:-1], topo[7][:-1],
topo[8][:-1], topo[9][-1:0:-1]])
# now that the quad arrays are populated, decide which quadrents
# still contain the separator (could be > 1)
required_uneven_subquads = False
ret = []
for i in range(4):
if perimeter_check(quad_topo[i]):
next_quad = "{0}{1:x}".format(_base_quadrent, i | _uneven_mask)
subquads = _get_sep_pts_bisect(fld, seed, trace_opts=trace_opts,
min_depth=min_depth,
max_depth=max_depth, plot=plot,
_base_quadrent=next_quad,
_uneven_mask=0,
_first_recurse=False)
ret += subquads
if len(ret) == 0:
perimeter = np.concatenate([topo[0][::-1], topo[4][::-1],
topo[5][::-1], topo[7][::-1],
topo[8][::-1], topo[10][::-1],
topo[11][::-1], topo[3][::-1]])
if _uneven_mask:
if len(_base_quadrent) > min_depth:
print("sep trace issue, but min depth reached: {0} > {1}"
"".format(len(_base_quadrent), min_depth))
ret = [_base_quadrent]
else:
print("sep trace issue, the separator ended prematurely")
elif perimeter_check(perimeter):
ret = _get_sep_pts_bisect(fld, seed, trace_opts=trace_opts,
min_depth=min_depth, max_depth=max_depth,
plot=plot, _base_quadrent=_base_quadrent,
_uneven_mask=UNEVEN_MASK,
_first_recurse=False)
required_uneven_subquads = True
if plot and not required_uneven_subquads:
from viscid.plot import vlab
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
_pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
vlab.points3d(_pts3d[0], _pts3d[1], _pts3d[2],
all_topo.data.reshape(-1), scale_mode='none',
scale_factor=0.02)
plt.scatter(allx, ally, color=np.bitwise_and(all_topo, 15),
vmin=0, vmax=15, marker='o', edgecolor='y', s=40)
if _first_recurse:
# turn quadrent strings into locations
xc = np.empty(len(ret))
yc = np.empty(len(ret))
for i, r in enumerate(ret):
xc[i], yc[i] = _quadrent_center(r, xlim, ylim)
pts_uv = np.array([xc, yc])
if plot:
from viscid.plot import vlab
from viscid.plot import vpyplot as vlt
from matplotlib import pyplot as plt
plt.plot(pts_uv[0], pts_uv[1], "y*", ms=20,
markeredgecolor='k', markeredgewidth=1.0)
vlt.show(block=False)
vlab.show(stop=True)
# return seed.to_3d(seed.uv_to_local(pts_uv))
# if pts_uv.size == 0:
# return None
if make_3d:
return seed.uv_to_3d(pts_uv)
else:
return pts_uv
else:
return ret
def trace_separator(grid, b_slcstr="x=-25j:15j, y=-30j:30j, z=-15j:15j",
r=1.0, plot=False, trace_opts=None, cache=True,
cache_dir=None):
"""Trace a separator line from most dawnward null
**Still in testing** Uses the bisection algorithm.
Args:
grid (Grid): A grid that has a "b" field
b_slcstr (str): Some valid slice for B field
r (float): spatial step of separator line
plot (bool): make debugging plots
trace_opts (dict): passed to streamline function
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
Raises:
IOError: Description
Returns:
tuple: (separator_lines, nulls)
- **separator_lines** (list): list of M 3xN ndarrays that
represent M separator lines with N points
- **nulls** (ndarray): 3xN array of N null points
"""
if not cache_dir:
cache_dir = grid.find_info("_viscid_dirname", "./")
run_name = grid.find_info("run")
sep_fname = "{0}/{1}.sep.{2:06.0f}".format(cache_dir, run_name, grid.time)
try:
if isinstance(cache, string_types) and cache.strip().lower() == "force":
raise IOError()
with np.load(sep_fname + ".npz") as dat:
sep_iter = (f for f in dat.files if f.startswith("arr_"))
_it = sorted(sep_iter, key=lambda s: int(s[len("arr_"):]))
seps = [dat[n] for n in _it]
nulls = dat['nulls']
except IOError:
_b = grid['b'][b_slcstr]
_, nulls = viscid.find_nulls(_b['x=-30j:15j'], ibound=5.0)
# get most dawnward null, nulls2 is all nulls except p0
nullind = np.argmin(nulls[1, :])
p0 = nulls[:, nullind]
nulls2 = np.concatenate([nulls[:, :nullind], nulls[:, (nullind + 1):]],
axis=1)
if plot:
from viscid.plot import vlab
vlab.plot_earth_3d(crd_system='gse')
vlab.points3d(nulls2[0], nulls2[1], nulls2[2],
color=(0, 0, 0), scale_factor=1.0)
vlab.points3d(nulls[0, nullind], nulls[1, nullind], nulls[2, nullind],
color=(1, 1, 1), scale_factor=1.0)
seed = viscid.Sphere(p0=p0, r=r, ntheta=30, nphi=60,
theta_endpoint=True, phi_endpoint=True)
p1 = viscid.get_sep_pts_bisect(_b, seed, max_depth=12, plot=plot,
trace_opts=trace_opts)
# print("p1 shape", p1.shape)
# if p1.shape[1] > 2:
# raise RuntimeError("Invalid B field, should be no branch @ null")
seps = []
sep_stubs = []
for i in range(p1.shape[1]):
sep_stubs.append([p0, p1[:, i]])
# print("??", sep_stubs)
while sep_stubs:
sep = sep_stubs.pop(0)
# print("!!! new stub")
for i in count():
# print("::", i)
seed = viscid.SphericalPatch(p0=sep[-1], p1=sep[-1] - sep[-2],
r=r, nalpha=240, nbeta=240)
pn = viscid.get_sep_pts_bisect(_b, seed, max_depth=8, plot=plot,
trace_opts=trace_opts)
if pn.shape[1] == 0:
# print("END: pn.shape[1] == 0")
break
# print("++", nulls2.shape, pn.shape)
closest_null_dist = np.min(np.linalg.norm(nulls2 - pn[:, :1], axis=0))
# print("closest_null_dist:", closest_null_dist)
if closest_null_dist < 1.01 * r:
# print("END: within 1.01 of a null")
break
# print("??", pn)
for j in range(1, pn.shape[1]):
# print("inserting new stub")
sep_stubs.insert(0, [sep[-1], pn[:, j]])
sep.append(pn[:, 0])
# print("sep", sep)
seps.append(np.stack(sep, axis=1))
if cache:
np.savez_compressed(sep_fname, *seps, nulls=nulls)
return seps, nulls
