Exemplo n.º 1
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    ####### test 5-moment uniform grids
    gk_uniform = viscid.load_file(os.path.join(sample_dir,
                                               'sample_gkeyll_uniform_q_*.h5'))

    plt.figure(figsize=(9, 3))
    for i, grid in enumerate(gk_uniform.iter_times(":")):
        plt.subplot2grid((1, 2), (0, i))
        vlt.plot(grid['rho_i'], logscale=True, style='contourf', levels=128)
        seeds = viscid.Line((-1.2, 0, 0), (1.4, 0, 0), 8)
        b_lines, _ = viscid.calc_streamlines(grid['b'], seeds, method='euler1',
                                             max_length=20.0)
        vlt.plot2d_lines(b_lines, scalars='#000000', symdir='z', linewidth=1.0)
        plt.title(grid.format_time('.02f'))
    vlt.auto_adjust_subplots()
    plt.suptitle("Uniform Gkeyll Dataset")

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        plt.show()
    plt.clf()

    return 0
Exemplo n.º 2
0
def default_fluid_callback(i,
                           t,
                           seeds=None,
                           v_field=None,
                           anc_fields=None,
                           grid0=None,
                           grid1=None,
                           streamlines=None,
                           series_fname=None,
                           show=False,
                           **kwargs):
    from viscid.plot import vpyplot as vlt
    from matplotlib import pyplot as plt

    _, ax0 = plt.subplots(1, 1)
    vlt.plot(viscid.magnitude(v_field).atmost_nd(2), ax=ax0)
    vlt.scatter_2d(seeds, ax=ax0)
    vlt.streamplot(v_field.atmost_nd(2))
    vlt.plot_lines2d(streamlines, ax=ax0)
    plt.title('{0:8.3f}'.format(t))
    if series_fname:
        plt.savefig('{0}_{1:06d}.png'.format(series_fname, i))
    if show:
        vlt.show()
    else:
        plt.close()
Exemplo n.º 3
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    ####### test 5-moment uniform grids
    gk_uniform = viscid.load_file(os.path.join(sample_dir,
                                               'sample_gkeyll_uniform_q_*.h5'))

    _, axes = plt.subplots(1, 2, figsize=(9, 3))
    for i, grid in enumerate(gk_uniform.iter_times(":")):
        vlt.plot(grid['rho_i'], logscale=True, style='contourf', levels=128,
                 ax=axes[i])
        seeds = viscid.Line((-1.2, 0, 0), (1.4, 0, 0), 8)
        b_lines, _ = viscid.calc_streamlines(grid['b'], seeds, method='euler1',
                                             max_length=20.0)
        vlt.plot2d_lines(b_lines, scalars='#000000', symdir='z', linewidth=1.0)
        plt.title(grid.format_time('.02f'))
    vlt.auto_adjust_subplots()
    plt.suptitle("Uniform Gkeyll Dataset")

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        plt.show()
    plt.clf()

    return 0
Exemplo n.º 4
0
def run_div_test(fld, exact, title='', show=False, ignore_inexact=False):
    t0 = time()
    result_numexpr = viscid.div(fld, preferred="numexpr", only=False)
    t1 = time()
    logger.info("numexpr magnitude runtime: %g", t1 - t0)

    result_diff = viscid.diff(result_numexpr, exact)['x=1:-1, y=1:-1, z=1:-1']
    if not ignore_inexact and not (result_diff.data < 5e-5).all():
        logger.warning("numexpr result is far from the exact result")
    logger.info("min/max(abs(numexpr - exact)): %g / %g",
                np.min(result_diff.data), np.max(result_diff.data))

    planes = ["y=0j", "z=0j"]
    nrows = 2
    ncols = len(planes)
    _, axes = plt.subplots(nrows, ncols, squeeze=False)

    for i, p in enumerate(planes):
        vlt.plot(result_numexpr, p, ax=axes[0, i], show=False)
        vlt.plot(result_diff, p, ax=axes[1, i], show=False)

    plt.suptitle(title)
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 5
0
def guess_dipole_moment(b, r=2.0, strength=DEFAULT_STRENGTH, cap_angle=40,
                        cap_ntheta=121, cap_nphi=121, plot=False):
    """guess dipole moment from a B field"""
    viscid.logger.warning("guess_dipole_moment doesn't seem to do better than "
                          "1.6 degrees, you may want to use cotr instead.")
    cap = seed.SphericalCap(r=r, angle=cap_angle, ntheta=cap_ntheta,
                            nphi=cap_nphi)
    b_cap = interp_trilin(b, cap)

    # FIXME: this doesn't get closer than 1.6 deg @ (theta, mu) = (0, 7.5)
    #        so maybe the index is incorrect somehow?
    idx = np.argmax(viscid.magnitude(b_cap).data)
    pole = cap.points()[:, idx]
    # FIXME: it should be achievabe to get strength from the magimum magnitude,
    #        up to the direction
    pole = strength * pole / np.linalg.norm(pole)
    # # not sure where 0.133 comes from, this is probably not correct
    # pole *= 0.133 * np.dot(pole, b_cap.data.reshape(-1, 3)[idx, :]) * r**3

    if plot:
        from viscid.plot import vpyplot as vlt
        from matplotlib import pyplot as plt
        vlt.plot(viscid.magnitude(b_cap))
        vlt.plot(viscid.magnitude(b_cap), style='contour', levels=10,
                 colors='k', colorbar=False, ax=plt.gca())
        vlt.show()
    return pole
Exemplo n.º 6
0
def run_test_2d(f, main__file__, show=False):
    vlt.clf()
    slc = "x=-20j:12j, y=0j"
    plot_kwargs = dict(title=True, earth=True)
    vlt.subplot(141)
    vlt.plot(f['pp'], slc, logscale=True, **plot_kwargs)
    vlt.plot(np.abs(f['psi']), style='contour', logscale=True, levels=30,
             linewidths=0.8, colors='grey', linestyles='solid', colorbar=None,
             x=(-20, 12))
    vlt.subplot(142)
    vlt.plot(viscid.magnitude(f['bcc']), slc, logscale=True, **plot_kwargs)
    vlt.plot2d_quiver(f['v'][slc], step=5, color='y', pivot='mid', width=0.03,
                      scale=600)
    vlt.subplot(143)
    vlt.plot(f['jy'], slc, clim=[-0.005, 0.005], **plot_kwargs)
    vlt.streamplot(f['v'][slc], linewidth=0.3)
    vlt.subplot(144)
    vlt.plot(f['jy'], "x=7j:12j, y=0j, z=0j")

    plt.suptitle("2D File")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9, wspace=1.3))
    plt.gcf().set_size_inches(10, 4)

    vlt.savefig(next_plot_fname(main__file__))
    if show:
        vlt.show()
Exemplo n.º 7
0
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
Exemplo n.º 8
0
def run_test_iof(f, main__file__, show=False):
    vlt.clf()

    fac_tot = 1e9 * f["fac_tot"]
    plot_args = dict(projection="polar",
                     lin=[-300, 300],
                     bounding_lat=35.0,
                     drawcoastlines=True,  # for basemap only
                     title="Total FAC\n",
                     gridec='gray',
                     label_lat=True,
                     label_mlt=True,
                     colorbar=True,
                     cbar_kwargs=dict(pad=0.15)  # pad the colorbar away from the plot
                    )

    ax1 = vlt.subplot(121, projection='polar')
    vlt.plot(fac_tot, ax=ax1, hemisphere='north', **plot_args)
    ax1.annotate('(a)', xy=(0, 0), textcoords="axes fraction",
                 xytext=(-0.1, 1.0), fontsize=18)

    ax2 = vlt.subplot(122, projection='polar')
    plot_args['gridec'] = False
    vlt.plot(fac_tot, ax=ax2, hemisphere="south", style="contourf",
             levels=50, extend="both", **plot_args)
    ax2.annotate('(b)', xy=(0, 0), textcoords="axes fraction",
                 xytext=(-0.1, 1.0), fontsize=18)

    vlt.auto_adjust_subplots(subplot_params=dict())
    plt.gcf().set_size_inches(8, 4)

    plt.savefig(next_plot_fname(main__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 9
0
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
Exemplo n.º 10
0
def _main():
    import viscid
    from viscid import sample_dir
    from viscid.plot import vpyplot as vlt
    logger.setLevel(viscid.logging.DEBUG)

    f = viscid.load_file(os.path.join(sample_dir, 'vpic_sample', 'global.vpc'))
    vlt.plot(-f['n_e']['y=0j'], logscale=True, show=True)
Exemplo n.º 11
0
def run_test(show=False):
    f = viscid.load_file(os.path.join(viscid.sample_dir, "amr.xdmf"))
    plot_kwargs = dict(patchec='y')
    vlt.plot(f['f'], "z=0.0f", **plot_kwargs)

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.show()
Exemplo n.º 12
0
def run_test(fld, seeds, kind, show=False):
    plt.clf()
    vlt.plot(viscid.interp(fld, seeds, kind=kind))
    plt.title(kind)

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.show()
Exemplo n.º 13
0
def run_test(show=False):
    f = viscid.load_file(os.path.join(viscid.sample_dir, "amr.xdmf"))
    plot_kwargs = dict(patchec='y')
    vlt.plot(f['f'], "z=0.0j", **plot_kwargs)

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.show()
Exemplo n.º 14
0
def run_test(fld, seeds, kind, show=False):
    plt.clf()
    vlt.plot(viscid.interp(fld, seeds, kind=kind))
    plt.title(kind)

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.show()
Exemplo n.º 15
0
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
Exemplo n.º 16
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    f = viscid.load_file(os.path.join(viscid.sample_dir, "test.asc"))
    vlt.plot(f['c1'], show=False)
    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()

    return 0
Exemplo n.º 17
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = viscid.vutil.common_argparse(parser)
    # args.show = True

    t = viscid.linspace_datetime64('2006-06-10 12:30:00.0',
                                   '2006-06-10 12:33:00.0', 16)
    tL = viscid.as_datetime64('2006-06-10 12:31:00.0')
    tR = viscid.as_datetime64('2006-06-10 12:32:00.0')
    y = np.linspace(2 * np.pi, 4 * np.pi, 12)

    ### plots with a datetime64 axis
    f0 = viscid.ones([t, y], crd_names='ty', center='node')
    T, Y = f0.get_crds(shaped=True)
    f0.data += np.arange(T.size).reshape(T.shape)
    f0.data += np.cos(Y)

    fig = plt.figure(figsize=(10, 5))
    # 1D plot
    vlt.subplot(121)
    vlt.plot(f0[tL:tR]['y=0'], marker='^')
    plt.xlim(*viscid.as_datetime(t[[0, -1]]).tolist())
    # 2D plot
    vlt.subplot(122)
    vlt.plot(f0, x=(t[0], t[-1]))

    plt.suptitle("datetime64")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))
    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.close(fig)

    ### plots with a timedelta64 axis
    tL = tL - t[0]
    tR = tR - t[0]
    t = t - t[0]
    f0 = viscid.ones([t, y], crd_names='ty', center='node')
    T, Y = f0.get_crds(shaped=True)
    f0.data += np.arange(T.size).reshape(T.shape)
    f0.data += np.cos(Y)

    fig = plt.figure(figsize=(10, 5))
    # 1D plot
    vlt.subplot(121)
    vlt.plot(f0[tL:tR]['y=0'], marker='^')
    plt.xlim(*viscid.as_datetime(t[[0, -1]]).tolist())
    # 2D plot
    vlt.subplot(122)
    vlt.plot(f0, x=(t[0], t[-1]))

    plt.suptitle("timedelta64")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))
    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.close(fig)

    return 0
Exemplo n.º 18
0
def run_mpl_testA(show=False):
    logger.info("2D cell centered tests")

    x = np.array(np.linspace(-10, 10, 100), dtype=dtype)
    y = np.array(np.linspace(-10, 10, 120), dtype=dtype)
    z = np.array(np.linspace(-1, 1, 2), dtype=dtype)

    fld_s = viscid.empty([x, y, z], center='cell')
    Xcc, Ycc, Zcc = fld_s.get_crds_cc(shaped=True)  # pylint: disable=unused-variable
    fld_s[:, :, :] = np.sin(Xcc) + np.cos(Ycc)

    _, axes = plt.subplots(4, 1, squeeze=False)

    vlt.plot(fld_s, "y=20j", ax=axes[0, 0], show=False, plot_opts="lin_0")
    vlt.plot(fld_s,
             "x=0j:20j,y=0j:5j",
             ax=axes[1, 0],
             earth=True,
             show=False,
             plot_opts="x_-10_0,y_0_7")
    vlt.plot(fld_s, "y=0j", ax=axes[2, 0], show=False, plot_opts="lin_-1_1")
    vlt.plot(fld_s,
             "z=0j,x=-20j:0j",
             ax=axes[3, 0],
             earth=True,
             show=False,
             plot_opts="lin_-5_5")

    plt.suptitle("2d cell centered")
    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 19
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    f = viscid.load_file(os.path.join(sample_dir, 'vpic_sample', 'global.vpc'))

    # some slices that are good to check
    vlt.clf()
    vlt.plot(f['bx']['x=:32.01j'])
    plt.close()
    vlt.clf()
    vlt.plot(f['bx']['x=:33.0j'])
    plt.close()

    _, axes = vlt.subplots(2, 2, figsize=(8, 4))

    for i, ti in enumerate([0, -1]):
        f.activate_time(ti)
        vlt.plot(f['n_e']['y=0j'], symmetric=False, ax=axes[0, i])
        vlt.plot(f['bx']['y=0j'], symmetric=True, ax=axes[1, i])
        axes[0, i].set_title(f.get_grid().time)

    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.close()

    return 0
Exemplo n.º 20
0
def run_mag_test(fld, title="", show=False):
    vx, vy, vz = fld.component_views()  # pylint: disable=W0612
    vx, vy, vz = fld.component_fields()

    try:
        t0 = time()
        mag_ne = viscid.magnitude(fld, preferred="numexpr", only=False)
        t1 = time()
        logger.info("numexpr mag runtime: %g", t1 - t0)
    except viscid.verror.BackendNotFound:
        xfail("Numexpr is not installed")

    planes = ["z=0", "y=0"]
    nrows = 4
    ncols = len(planes)

    _, axes = plt.subplots(nrows, ncols, sharex=True, sharey=True, squeeze=False)

    for ind, p in enumerate(planes):
        vlt.plot(vx, p, ax=axes[0, ind], show=False)
        vlt.plot(vy, p, ax=axes[1, ind], show=False)
        vlt.plot(vz, p, ax=axes[2, ind], show=False)
        vlt.plot(mag_ne, p, ax=axes[3, ind], show=False)

    plt.suptitle(title)
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9, right=0.9))
    plt.gcf().set_size_inches(6, 7)

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 21
0
def run_mpl_testB(show=False):
    logger.info("3D node centered tests")

    x = np.array(np.linspace(-10, 10, 100), dtype=dtype)
    y = np.array(np.linspace(-10, 10, 120), dtype=dtype)
    z = np.array(np.linspace(-10, 10, 140), dtype=dtype)

    fld_s = viscid.empty([x, y, z], center='node')
    X, Y, Z = fld_s.get_crds_nc(shaped=True)  # pylint: disable=W0612
    fld_s[:, :, :] = np.sin(X) + np.cos(Y) - np.cos(Z)
    # print("shape: ", fld_s.data.shape)

    nrows = 4
    ncols = 1

    plt.subplot2grid((nrows, ncols), (0, 0))
    vlt.plot(fld_s, "z=0,x=:30", earth=True, plot_opts="lin_0")
    plt.subplot2grid((nrows, ncols), (1, 0))
    vlt.plot(fld_s, "z=0.75f,x=-4:-1,y=-3f:3f", earth=True)
    plt.subplot2grid((nrows, ncols), (2, 0))
    vlt.plot(fld_s, "x=-0.5f:,y=-3f:3f,z=0f", earth=True)
    plt.subplot2grid((nrows, ncols), (3, 0))
    vlt.plot(fld_s, "x=0.0f,y=-5.0f:5.0f", earth=True, plot_opts="log,g")

    plt.suptitle("3d node centered")
    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 22
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    f = viscid.load_file(os.path.join(sample_dir, 'vpic_sample', 'global.vpc'))

    # some slices that are good to check
    vlt.clf()
    vlt.plot(f['bx']['x=:32.01j'])
    plt.close()
    vlt.clf()
    vlt.plot(f['bx']['x=:33.0j'])
    plt.close()

    _, axes = vlt.subplots(2, 2, figsize=(8, 4))

    for i, ti in enumerate([0, -1]):
        f.activate_time(ti)
        vlt.plot(f['n_e']['y=0j'], symmetric=False, ax=axes[0, i])
        vlt.plot(f['bx']['y=0j'], symmetric=True, ax=axes[1, i])
        axes[0, i].set_title(f.get_grid().time)

    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.close()

    return 0
Exemplo n.º 23
0
def run_test_3d(f, main__file__, show=False):
    vlt.clf()
    slc = "x=-20f:12f, y=0f"
    plot_kwargs = dict(title=True, earth=True)
    vlt.subplot(141)
    vlt.plot(f['pp'], slc, logscale=True, **plot_kwargs)
    vlt.subplot(142)
    vlt.plot(viscid.magnitude(f['bcc']), slc, logscale=True, **plot_kwargs)
    vlt.plot2d_quiver(f['v'][slc],
                      step=5,
                      color='y',
                      pivot='mid',
                      width=0.03,
                      scale=600)
    vlt.subplot(143)
    vlt.plot(f['jy'], slc, clim=(-0.005, 0.005), **plot_kwargs)
    vlt.streamplot(f['v'][slc], linewidth=0.3)
    vlt.subplot(144)
    vlt.plot(f['jy'], "x=7f:12f, y=0f, z=0f")

    plt.suptitle("3D File")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9, wspace=1.3))
    plt.gcf().set_size_inches(10, 4)

    vlt.savefig(next_plot_fname(main__file__))
    if show:
        vlt.show()
Exemplo n.º 24
0
def run_mpl_testB(show=False):
    logger.info("3D node centered tests")

    x = np.array(np.linspace(-10, 10, 100), dtype=dtype)
    y = np.array(np.linspace(-10, 10, 120), dtype=dtype)
    z = np.array(np.linspace(-10, 10, 140), dtype=dtype)

    fld_s = viscid.empty([x, y, z], center='node')
    X, Y, Z = fld_s.get_crds_nc(shaped=True)  # pylint: disable=W0612
    fld_s[:, :, :] = np.sin(X) + np.cos(Y) - np.cos(Z)
    # print("shape: ", fld_s.data.shape)

    _, axes = plt.subplots(4, 1, squeeze=False)

    vlt.plot(fld_s, "z=0,x=:30", ax=axes[0, 0], earth=True, plot_opts="lin_0")
    vlt.plot(fld_s, "z=0.75j,x=-4:-1,y=-3j:3j", ax=axes[1, 0], earth=True)
    vlt.plot(fld_s, "x=-0.5j:,y=-3j:3j,z=0j", ax=axes[2, 0], earth=True)
    vlt.plot(fld_s,
             "x=0.0j,y=-5.0j:5.0j",
             ax=axes[3, 0],
             earth=True,
             plot_opts="log,g")

    plt.suptitle("3d node centered")
    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 25
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = viscid.vutil.common_argparse(parser)
    # args.show = True

    t = viscid.linspace_datetime64('2006-06-10 12:30:00.0',
                                   '2006-06-10 12:33:00.0', 16)
    tL = viscid.as_datetime64('2006-06-10 12:31:00.0')
    tR = viscid.as_datetime64('2006-06-10 12:32:00.0')
    y = np.linspace(2 * np.pi, 4 * np.pi, 12)

    ### plots with a datetime64 axis
    f0 = viscid.ones([t, y], crd_names='ty', center='node')
    T, Y = f0.get_crds(shaped=True)
    f0.data += np.arange(T.size).reshape(T.shape)
    f0.data += np.cos(Y)

    fig = plt.figure(figsize=(10, 5))
    # 1D plot
    vlt.subplot(121)
    vlt.plot(f0[tL:tR]['y=0'], marker='^')
    plt.xlim(*viscid.as_datetime(t[[0, -1]]).tolist())
    # 2D plot
    vlt.subplot(122)
    vlt.plot(f0, x=(t[0], t[-1]))

    plt.suptitle("datetime64")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))
    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.close(fig)

    ### plots with a timedelta64 axis
    tL = tL - t[0]
    tR = tR - t[0]
    t = t - t[0]
    f0 = viscid.ones([t, y], crd_names='ty', center='node')
    T, Y = f0.get_crds(shaped=True)
    f0.data += np.arange(T.size).reshape(T.shape)
    f0.data += np.cos(Y)

    fig = plt.figure(figsize=(10, 5))
    # 1D plot
    vlt.subplot(121)
    vlt.plot(f0[tL:tR]['y=0'], marker='^')
    plt.xlim(*viscid.as_datetime(t[[0, -1]]).tolist())
    # 2D plot
    vlt.subplot(122)
    vlt.plot(f0, x=(t[0], t[-1]))

    plt.suptitle("timedelta64")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))
    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.close(fig)

    return 0
Exemplo n.º 26
0
def run_mpl_testA(show=False):
    logger.info("2D cell centered tests")

    x = np.array(np.linspace(-10, 10, 100), dtype=dtype)
    y = np.array(np.linspace(-10, 10, 120), dtype=dtype)
    z = np.array(np.linspace(-1, 1, 2), dtype=dtype)

    fld_s = viscid.empty([x, y, z], center='cell')
    Xcc, Ycc, Zcc = fld_s.get_crds_cc(shaped=True)  # pylint: disable=unused-variable
    fld_s[:, :, :] = np.sin(Xcc) + np.cos(Ycc)

    _, axes = plt.subplots(4, 1, squeeze=False)

    vlt.plot(fld_s, "y=20j", ax=axes[0, 0], show=False, plot_opts="lin_0")
    vlt.plot(fld_s, "x=0j:20j,y=0j:5j", ax=axes[1, 0], earth=True, show=False,
             plot_opts="x_-10_0,y_0_7")
    vlt.plot(fld_s, "y=0j", ax=axes[2, 0], show=False, plot_opts="lin_-1_1")
    vlt.plot(fld_s, "z=0j,x=-20j:0j", ax=axes[3, 0], earth=True, show=False,
             plot_opts="lin_-5_5")

    plt.suptitle("2d cell centered")
    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 27
0
def run_mpl_testB(show=False):
    logger.info("3D node centered tests")

    x = np.array(np.linspace(-10, 10, 100), dtype=dtype)
    y = np.array(np.linspace(-10, 10, 120), dtype=dtype)
    z = np.array(np.linspace(-10, 10, 140), dtype=dtype)

    fld_s = viscid.empty([x, y, z], center='node')
    X, Y, Z = fld_s.get_crds_nc(shaped=True)  # pylint: disable=W0612
    fld_s[:, :, :] = np.sin(X) + np.cos(Y) - np.cos(Z)
    # print("shape: ", fld_s.data.shape)

    _, axes = plt.subplots(4, 1, squeeze=False)

    vlt.plot(fld_s, "z=0,x=:30", ax=axes[0, 0], earth=True, plot_opts="lin_0")
    vlt.plot(fld_s, "z=0.75j,x=-4:-1,y=-3j:3j", ax=axes[1, 0], earth=True)
    vlt.plot(fld_s, "x=-0.5j:,y=-3j:3j,z=0j", ax=axes[2, 0], earth=True)
    vlt.plot(fld_s, "x=0.0j,y=-5.0j:5.0j", ax=axes[3, 0], earth=True,
             plot_opts="log,g")

    plt.suptitle("3d node centered")
    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 28
0
def default_fluid_callback(i, t, seeds=None, v_field=None, anc_fields=None,
                           grid0=None, grid1=None, streamlines=None,
                           series_fname=None, show=False, **kwargs):
    from viscid.plot import vpyplot as vlt
    from matplotlib import pyplot as plt

    _, ax0 = plt.subplots(1, 1)
    vlt.plot(viscid.magnitude(v_field).atmost_nd(2), ax=ax0)
    vlt.scatter_2d(seeds, ax=ax0)
    vlt.streamplot(v_field.atmost_nd(2))
    vlt.plot_lines2d(streamlines, ax=ax0)
    plt.title('{0:8.3f}'.format(t))
    if series_fname:
        plt.savefig('{0}_{1:06d}.png'.format(series_fname, i))
    if show:
        vlt.show()
    else:
        plt.close()
Exemplo n.º 29
0
def compare_vectors(cc_fld, ecfc_fld, to_cc_fn, catol=1e-8, rtol=2.2e-6,
                    trim_slc=':', bnd=True, make_plots=False):
    trimmed = cc_fld[trim_slc]
    cc = to_cc_fn(ecfc_fld, bnd=bnd)
    reldiff = (cc - trimmed) / viscid.magnitude(trimmed)
    reldiff = reldiff["x=1:-1, y=1:-1, z=1:-1"]
    reldiff.name = cc.name + " - " + trimmed.name
    reldiff.pretty_name = cc.pretty_name + " - " + trimmed.pretty_name
    comp_beyond_limit = [False] * 3

    for i, d in enumerate('xyz'):
        abs_max_rel_diff = np.nanmax(np.abs(reldiff[d]))
        # trim first and last when diffing crds since they're extrapolated
        # in the ecfc field, so we already know they won't match up do dx
        max_crd_diff = np.max((trimmed.get_crd(d) - cc.get_crd(d))[1:-1])
        print("{0}{1}, max absolute relative diff: {2:.3e} ({1} crds: {3:.1e})"
              "".format(cc_fld.name, d, abs_max_rel_diff, max_crd_diff))
        if abs_max_rel_diff > rtol or abs(max_crd_diff) > catol:
            comp_beyond_limit[i] = True

        # plot differences?
        if make_plots:
            ax1 = plt.subplot(311)
            vlt.plot(cc_fld[d]['y=0j'], symmetric=True, earth=True)
            plt.subplot(312, sharex=ax1, sharey=ax1)
            vlt.plot(cc[d]['y=0j'], symmetric=True, earth=True)
            plt.subplot(313, sharex=ax1, sharey=ax1)
            vlt.plot(reldiff[d]['y=0j'], symmetric=True, earth=True)
            vlt.show()

    if any(comp_beyond_limit):
        raise RuntimeError("Tolerance exceeded on ->CC accuracy")
Exemplo n.º 30
0
def compare_vectors(orig_fld,
                    cmp_fld,
                    catol=1e-8,
                    rtol=2e-6,
                    trim_slc='x=2:-2, y=2:-2, z=2:-2',
                    make_plots=False):
    # mag_shape = list(orig_fld.sshape) + [1]
    mag = viscid.magnitude(orig_fld)  # .data.reshape(mag_shape)
    reldiff = (cmp_fld - orig_fld) / mag
    reldiff = reldiff[trim_slc]
    reldiff.name = cmp_fld.name + " - " + orig_fld.name
    reldiff.pretty_name = cmp_fld.pretty_name + " - " + orig_fld.pretty_name
    comp_beyond_limit = [False] * 3

    for i, d in enumerate('xyz'):
        abs_max_rel_diff = np.nanmax(np.abs(reldiff[d]))
        max_crd_diff = np.max(orig_fld.get_crd(d) - cmp_fld.get_crd(d))
        print("{0}{1}, max absolute relative diff: {2:.3e} ({1} crds: {3:.1e})"
              "".format(cmp_fld.name, d, abs_max_rel_diff, max_crd_diff))
        if abs_max_rel_diff > rtol or abs(max_crd_diff) > catol:
            comp_beyond_limit[i] = True

        # plot differences?
        if make_plots:
            plt.clf()
            ax1 = plt.subplot(311)
            vlt.plot(orig_fld[d]['y=0j'], symmetric=True, earth=True)
            plt.subplot(312, sharex=ax1, sharey=ax1)
            vlt.plot(cmp_fld[d]['y=0j'], symmetric=True, earth=True)
            plt.subplot(313, sharex=ax1, sharey=ax1)
            vlt.plot(reldiff[d]['y=0j'], symmetric=True, earth=True)
            vlt.show()
Exemplo n.º 31
0
def compare_vectors(orig_fld, cmp_fld, catol=1e-8, rtol=2e-6,
                          trim_slc='x=2:-2, y=2:-2, z=2:-2', make_plots=False):
    # mag_shape = list(orig_fld.sshape) + [1]
    mag = viscid.magnitude(orig_fld)  # .data.reshape(mag_shape)
    reldiff = (cmp_fld - orig_fld) / mag
    reldiff = reldiff[trim_slc]
    reldiff.name = cmp_fld.name + " - " + orig_fld.name
    reldiff.pretty_name = cmp_fld.pretty_name + " - " + orig_fld.pretty_name
    comp_beyond_limit = [False] * 3

    for i, d in enumerate('xyz'):
        abs_max_rel_diff = np.nanmax(np.abs(reldiff[d]))
        max_crd_diff = np.max(orig_fld.get_crd(d) - cmp_fld.get_crd(d))
        print("{0}{1}, max absolute relative diff: {2:.3e} ({1} crds: {3:.1e})"
              "".format(cmp_fld.name, d, abs_max_rel_diff, max_crd_diff))
        if abs_max_rel_diff > rtol or abs(max_crd_diff) > catol:
            comp_beyond_limit[i] = True

        # plot differences?
        if make_plots:
            plt.clf()
            ax1 = plt.subplot(311)
            vlt.plot(orig_fld[d]['y=0j'], symmetric=True, earth=True)
            plt.subplot(312, sharex=ax1, sharey=ax1)
            vlt.plot(cmp_fld[d]['y=0j'], symmetric=True, earth=True)
            plt.subplot(313, sharex=ax1, sharey=ax1)
            vlt.plot(reldiff[d]['y=0j'], symmetric=True, earth=True)
            vlt.show()
Exemplo n.º 32
0
def main():
    f = viscid.load_file("~/dev/work/tmedium/*.3d.[-1].xdmf")
    grid = f.get_grid()

    gslc = "x=-26f:12.5f, y=-15f:15f, z=-15f:15f"
    # gslc = "x=-12.5f:26f, y=-15f:15f, z=-15f:15f"

    b_cc = f['b_cc'][gslc]
    b_cc.name = "b_cc"
    b_fc = f['b_fc'][gslc]
    b_fc.name = "b_fc"

    e_cc = f['e_cc'][gslc]
    e_cc.name = "e_cc"
    e_ec = f['e_ec'][gslc]
    e_ec.name = "e_ec"

    pp = f['pp'][gslc]
    pp.name = 'pp'

    pargs = dict(logscale=True, earth=True)

    # vlt.clf()
    # ax1 = vlt.subplot(211)
    # vlt.plot(f['pp']['y=0f'], **pargs)
    # # vlt.plot(viscid.magnitude(f['b_cc']['y=0f']), **pargs)
    # # vlt.show()
    # vlt.subplot(212, sharex=ax1, sharey=ax1)
    # vlt.plot(viscid.magnitude(viscid.fc2cc(f['b_fc'])['y=0f']), **pargs)
    # vlt.show()

    basename = './tmediumR.3d.{0:06d}'.format(int(grid.time))
    viscid.save_fields(basename + '.h5', [b_cc, b_fc, e_cc, e_ec, pp])

    f2 = viscid.load_file(basename + ".xdmf")

    pargs = dict(logscale=True, earth=True)

    vlt.clf()
    ax1 = vlt.subplot(211)
    vlt.plot(f2['pp']['y=0f'], style='contour', levels=5, colorbar=None,
             colors='k', **pargs)
    vlt.plot(viscid.magnitude(f2['b_cc']['y=0f']), **pargs)
    vlt.subplot(212, sharex=ax1, sharey=ax1)
    vlt.plot(viscid.magnitude(viscid.fc2cc(f2['b_fc'])['y=0f']), **pargs)
    vlt.show()

    os.remove(basename + '.h5')
    os.remove(basename + '.xdmf')

    return 0
Exemplo n.º 33
0
def main():
    f = viscid.load_file("~/dev/work/tmedium/*.3d.[-1].xdmf")
    grid = f.get_grid()

    gslc = "x=-26j:12.5j, y=-15j:15j, z=-15j:15j"
    # gslc = "x=-12.5j:26j, y=-15j:15j, z=-15j:15j"

    b_cc = f['b_cc'][gslc]
    b_cc.name = "b_cc"
    b_fc = f['b_fc'][gslc]
    b_fc.name = "b_fc"

    e_cc = f['e_cc'][gslc]
    e_cc.name = "e_cc"
    e_ec = f['e_ec'][gslc]
    e_ec.name = "e_ec"

    pp = f['pp'][gslc]
    pp.name = 'pp'

    pargs = dict(logscale=True, earth=True)

    # vlt.clf()
    # ax1 = vlt.subplot(211)
    # vlt.plot(f['pp']['y=0j'], **pargs)
    # # vlt.plot(viscid.magnitude(f['b_cc']['y=0j']), **pargs)
    # # vlt.show()
    # vlt.subplot(212, sharex=ax1, sharey=ax1)
    # vlt.plot(viscid.magnitude(viscid.fc2cc(f['b_fc'])['y=0j']), **pargs)
    # vlt.show()

    basename = './tmediumR.3d.{0:06d}'.format(int(grid.time))
    viscid.save_fields(basename + '.h5', [b_cc, b_fc, e_cc, e_ec, pp])

    f2 = viscid.load_file(basename + ".xdmf")

    pargs = dict(logscale=True, earth=True)

    vlt.clf()
    ax1 = vlt.subplot(211)
    vlt.plot(f2['pp']['y=0j'], style='contour', levels=5, colorbar=None,
             colors='k', **pargs)
    vlt.plot(viscid.magnitude(f2['b_cc']['y=0j']), **pargs)
    vlt.subplot(212, sharex=ax1, sharey=ax1)
    vlt.plot(viscid.magnitude(viscid.fc2cc(f2['b_fc'])['y=0j']), **pargs)
    vlt.show()

    os.remove(basename + '.h5')
    os.remove(basename + '.xdmf')

    return 0
Exemplo n.º 34
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    ####### test binary files
    f_bin = viscid.load_file(os.path.join(sample_dir, 'ath_sample.*.bin'))

    for i, grid in enumerate(f_bin.iter_times(":")):
        plt.subplot2grid((2, 2), (0, i))
        vlt.plot(grid['bx'])
        plt.subplot2grid((2, 2), (1, i))
        vlt.plot(grid['by'])
    plt.suptitle("athena bin (binary) files")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.clf()

    ####### test ascii files
    f_tab = viscid.load_file(os.path.join(sample_dir, 'ath_sample.*.tab'))

    for i, grid in enumerate(f_tab.iter_times(":")):
        plt.subplot2grid((2, 2), (0, i))
        vlt.plot(grid['bx'])
        plt.subplot2grid((2, 2), (1, i))
        vlt.plot(grid['by'])
    plt.suptitle("athena tab (ascii) files")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.clf()

    return 0
Exemplo n.º 35
0
def run_test_3d(f, main__file__, show=False):
    vlt.clf()
    slc = "x=-20j:12j, y=0j"
    plot_kwargs = dict(title=True, earth=True)
    vlt.subplot(141)
    vlt.plot(f['pp'], slc, logscale=True, **plot_kwargs)
    vlt.subplot(142)
    vlt.plot(viscid.magnitude(f['bcc']), slc, logscale=True, **plot_kwargs)
    vlt.plot2d_quiver(f['v'][slc], step=5, color='y', pivot='mid', width=0.03,
                      scale=600)
    vlt.subplot(143)
    vlt.plot(f['jy'], slc, clim=(-0.005, 0.005), **plot_kwargs)
    vlt.streamplot(f['v'][slc], linewidth=0.3)
    vlt.subplot(144)
    vlt.plot(f['jy'], "x=7j:12j, y=0j, z=0j")

    plt.suptitle("3D File")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9, wspace=1.3))
    plt.gcf().set_size_inches(10, 4)

    vlt.savefig(next_plot_fname(main__file__))
    if show:
        vlt.show()
Exemplo n.º 36
0
def run_mpl_testA(show=False):
    logger.info("2D cell centered tests")

    x = np.array(np.linspace(-10, 10, 100), dtype=dtype)
    y = np.array(np.linspace(-10, 10, 120), dtype=dtype)
    z = np.array(np.linspace(-1, 1, 2), dtype=dtype)

    fld_s = viscid.empty([x, y, z], center='cell')
    Xcc, Ycc, Zcc = fld_s.get_crds_cc(shaped=True)  # pylint: disable=unused-variable
    fld_s[:, :, :] = np.sin(Xcc) + np.cos(Ycc)

    nrows = 4
    ncols = 1

    plt.subplot2grid((nrows, ncols), (0, 0))
    vlt.plot(fld_s, "y=20f", show=False, plot_opts="lin_0")
    plt.subplot2grid((nrows, ncols), (1, 0))
    vlt.plot(fld_s,
             "x=0f:20f,y=0f:5f",
             earth=True,
             show=False,
             plot_opts="x_-10_0,y_0_7")
    plt.subplot2grid((nrows, ncols), (2, 0))
    vlt.plot(fld_s, "y=0f", show=False, plot_opts="lin_-1_1")
    plt.subplot2grid((nrows, ncols), (3, 0))
    vlt.plot(fld_s,
             "z=0f,x=-20f:0f",
             earth=True,
             show=False,
             plot_opts="lin_-5_5")

    plt.suptitle("2d cell centered")
    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if show:
        vlt.mplshow()
Exemplo n.º 37
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    parser.add_argument("--keep", action="store_true")
    args = vutil.common_argparse(parser)

    # setup a simple force free field
    x = np.linspace(-2, 2, 20)
    y = np.linspace(-2.5, 2.5, 25)
    z = np.linspace(-3, 3, 30)
    psi = viscid.empty([x, y, z], name='psi', center='node')
    b = viscid.empty([x, y, z],
                     nr_comps=3,
                     name='b',
                     center='cell',
                     layout='interlaced')

    X, Y, Z = psi.get_crds_nc("xyz", shaped=True)
    Xcc, Ycc, Zcc = psi.get_crds_cc("xyz", shaped=True)
    psi[:, :, :] = 0.5 * (X**2 + Y**2 - Z**2)
    b['x'] = Xcc
    b['y'] = Ycc
    b['z'] = -Zcc

    # save an hdf5 file with companion xdmf file
    h5_fname = os.path.join(viscid.sample_dir, "test.h5")
    viscid.save_fields(h5_fname, [psi, b])

    # load the companion xdmf file
    xdmf_fname = h5_fname[:-3] + ".xdmf"
    f = viscid.load_file(xdmf_fname)
    plt.subplot(131)
    vlt.plot(f['psi'], "y=0")
    plt.subplot(132)
    vlt.plot(f['b'].component_fields()[0], "y=0")
    plt.subplot(133)
    vlt.plot(f['b'].component_fields()[2], "y=0")

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        plt.show()

    if not args.keep:
        os.remove(h5_fname)
        os.remove(xdmf_fname)

    return 0
Exemplo n.º 38
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    ####### test binary files
    f_bin = viscid.load_file(os.path.join(sample_dir, 'ath_sample.*.bin'))

    _, axes = plt.subplots(2, 2)
    for i, grid in enumerate(f_bin.iter_times(":")):
        vlt.plot(grid['bx'], ax=axes[0, i])
        vlt.plot(grid['by'], ax=axes[1, i])
    plt.suptitle("athena bin (binary) files")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.close()

    ####### test ascii files
    f_tab = viscid.load_file(os.path.join(sample_dir, 'ath_sample.*.tab'))

    _, axes = plt.subplots(2, 2)
    for i, grid in enumerate(f_tab.iter_times(":")):
        vlt.plot(grid['bx'], ax=axes[0, i])
        vlt.plot(grid['by'], ax=axes[1, i])
    plt.suptitle("athena tab (ascii) files")
    vlt.auto_adjust_subplots(subplot_params=dict(top=0.9))

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()
    plt.close()

    return 0
Exemplo n.º 39
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    parser.add_argument("--keep", action="store_true")
    args = vutil.common_argparse(parser)

    # setup a simple force free field
    x = np.linspace(-2, 2, 20)
    y = np.linspace(-2.5, 2.5, 25)
    z = np.linspace(-3, 3, 30)
    psi = viscid.empty([x, y, z], name='psi', center='node')
    b = viscid.empty([x, y, z], nr_comps=3, name='b', center='cell',
                     layout='interlaced')

    X, Y, Z = psi.get_crds_nc("xyz", shaped=True)
    Xcc, Ycc, Zcc = psi.get_crds_cc("xyz", shaped=True)
    psi[:, :, :] = 0.5 * (X**2 + Y**2 - Z**2)
    b['x'] = Xcc
    b['y'] = Ycc
    b['z'] = -Zcc

    # save an hdf5 file with companion xdmf file
    h5_fname = os.path.join(".", "test.h5")
    viscid.save_fields(h5_fname, [psi, b])

    # load the companion xdmf file
    xdmf_fname = h5_fname[:-3] + ".xdmf"
    f = viscid.load_file(xdmf_fname)
    plt.subplot(131)
    vlt.plot(f['psi'], "y=0")
    plt.subplot(132)
    vlt.plot(f['b'].component_fields()[0], "y=0")
    plt.subplot(133)
    vlt.plot(f['b'].component_fields()[2], "y=0")

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        plt.show()

    if not args.keep:
        os.remove(h5_fname)
        os.remove(xdmf_fname)

    return 0
Exemplo n.º 40
0
def compare_vectors(cc_fld,
                    ecfc_fld,
                    to_cc_fn,
                    catol=1e-8,
                    rtol=2.2e-6,
                    trim_slc=':',
                    bnd=True,
                    make_plots=False):
    trimmed = cc_fld[trim_slc]
    cc = to_cc_fn(ecfc_fld, bnd=bnd)
    reldiff = (cc - trimmed) / viscid.magnitude(trimmed)
    reldiff = reldiff["x=1:-1, y=1:-1, z=1:-1"]
    reldiff.name = cc.name + " - " + trimmed.name
    reldiff.pretty_name = cc.pretty_name + " - " + trimmed.pretty_name
    comp_beyond_limit = [False] * 3

    for i, d in enumerate('xyz'):
        abs_max_rel_diff = np.nanmax(np.abs(reldiff[d]))
        # trim first and last when diffing crds since they're extrapolated
        # in the ecfc field, so we already know they won't match up do dx
        max_crd_diff = np.max((trimmed.get_crd(d) - cc.get_crd(d))[1:-1])
        print("{0}{1}, max absolute relative diff: {2:.3e} ({1} crds: {3:.1e})"
              "".format(cc_fld.name, d, abs_max_rel_diff, max_crd_diff))
        if abs_max_rel_diff > rtol or abs(max_crd_diff) > catol:
            comp_beyond_limit[i] = True

        # plot differences?
        if make_plots:
            ax1 = plt.subplot(311)
            vlt.plot(cc_fld[d]['y=0f'], symmetric=True, earth=True)
            plt.subplot(312, sharex=ax1, sharey=ax1)
            vlt.plot(cc[d]['y=0f'], symmetric=True, earth=True)
            plt.subplot(313, sharex=ax1, sharey=ax1)
            vlt.plot(reldiff[d]['y=0f'], symmetric=True, earth=True)
            vlt.show()

    if any(comp_beyond_limit):
        raise RuntimeError("Tolerance exceeded on ->CC accuracy")
Exemplo n.º 41
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
Exemplo n.º 42
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prof", action="store_true")
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    b = viscid.make_dipole(l=(-5, -5, -5), h=(5, 5, 5), n=(255, 255, 127),
                           m=(0, 0, -1))
    b2 = np.sum(b * b, axis=b.nr_comp)

    if args.prof:
        print("Without boundaries")
        viscid.timeit(viscid.grad, b2, bnd=False, timeit_repeat=10,
                      timeit_print_stats=True)
        print("With boundaries")
        viscid.timeit(viscid.grad, b2, bnd=True, timeit_repeat=10,
                      timeit_print_stats=True)

    grad_b2 = viscid.grad(b2)
    grad_b2.pretty_name = r"$\nabla$ B$^2$"
    conv = viscid.convective_deriv(b)
    conv.pretty_name = r"(B $\cdot \nabla$) B"

    _ = plt.figure(figsize=(9, 4.2))

    ax1 = vlt.subplot(231)
    vlt.plot(b2['z=0j'], logscale=True)
    vlt.plot(b2['z=0j'], logscale=True, style='contour', levels=10, colors='grey')
    # vlt.plot2d_quiver(viscid.normalize(b['z=0j']), step=16, pivot='mid')
    ax2 = vlt.subplot(234)
    vlt.plot(b2['y=0j'], logscale=True)
    vlt.plot(b2['y=0j'], logscale=True, style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(b['y=0j'], preferred='numpy'),
                      step=16, pivot='mid')

    vlt.subplot(232, sharex=ax1, sharey=ax1)
    vlt.plot(1e-4 + viscid.magnitude(grad_b2['z=0j']), logscale=True)
    vlt.plot(1e-4 + viscid.magnitude(grad_b2['z=0j']), logscale=True,
             style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(grad_b2['z=0j']), step=16, pivot='mid')
    vlt.subplot(235, sharex=ax2, sharey=ax2)
    vlt.plot(1e-4 + viscid.magnitude(grad_b2['y=0j']), logscale=True)
    vlt.plot(1e-4 + viscid.magnitude(grad_b2['y=0j']), logscale=True,
             style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(grad_b2['y=0j']), step=16, pivot='mid')

    vlt.subplot(233, sharex=ax1, sharey=ax1)
    vlt.plot(viscid.magnitude(conv['z=0j']), logscale=True)
    vlt.plot(viscid.magnitude(conv['z=0j']), logscale=True,
             style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(conv['z=0j']), step=16, pivot='mid')
    vlt.subplot(236, sharex=ax2, sharey=ax2)
    vlt.plot(viscid.magnitude(conv['y=0j']), logscale=True)
    vlt.plot(viscid.magnitude(conv['y=0j']), logscale=True,
             style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(conv['y=0j']), step=16, pivot='mid')

    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()

    return 0
Exemplo n.º 43
0
def apply_labels(labels=None,
                 colors=None,
                 ax=None,
                 magnet=(0.5, 0.75),
                 magnetcoords="axes fraction",
                 padding=None,
                 paddingcoords="offset points",
                 choices="00:02:20:22",
                 n_candidates=32,
                 ignore_filling=False,
                 spacing='linear',
                 _debug=False,
                 **kwargs):
    """Apply labels directly to series in liu of a legend

    The `choices` offsets are as follows::

        ---------------------
        |  02  |  12  | 22  |
        |-------------------|
        |  01  |  XX  | 21  |
        |-------------------|
        |  00  |  10  | 20  |
        ---------------------

    Args:
        labels (sequence): Optional sequence of labels to override the
            labels already in the data series
        colors (str, sequence): color as hex string, list of hex
            strings to color each label, or an Nx4 ndarray of rgba
            values for N labels
        ax (matplotlib.axis): axis; defaults to `plt.gca()`
        magnet (tuple): prefer positions that are closer to the magnet
        magnetcoords (str): 'offset pixels', 'offset points' or 'axes fraction'
        padding (tuple): padding for text in the (x, y) directions
        paddingcoords (str): 'offset pixels', 'offset points' or 'axes fraction'
        choices (str): colon separated list of possible label positions
            relative to the data values. The positions are summarized
            above.
        alpha (float): alpha channel (opacity) of label text. Defaults
            to 1.0 to make text visible. Set to `None` to use the
            underlying alpha from the handle's color.
        n_candidates (int): number of potential label locations to
            consider for each data series.
        ignore_filling (bool): if True, then assume it's ok to place
            labels inside paths that are filled with color
        spacing (str): one of 'linear' or 'random' to specify how far
            apart candidate locations are spaced along path segments
        _debug (bool): Mark up all possible label locations
        **kwargs: passed to plt.annotate

    Returns:
        List: annotation objects
    """
    if not ax:
        ax = plt.gca()

    if isinstance(colors, (list, tuple)):
        pass

    spacing = spacing.strip().lower()
    if spacing not in ('linear', 'random'):
        raise ValueError("Spacing '{0}' not understood".format(spacing))
    rand_state = np.random.get_state() if spacing == 'random' else None
    if rand_state is not None:
        # save the RNG state to restore it later so that plotting functions
        # don't change the results of scripts that use random numbers
        np.random.seed(1)

    _xl, _xh = ax.get_xlim()
    _yl, _yh = ax.get_ylim()
    axbb0 = np.array([_xl, _yl]).reshape(1, 2)
    axbb1 = np.array([_xh, _yh]).reshape(1, 2)

    # choices:: "01:02:22" -> [(0, 1), (0, 2), (2, 2)]
    choices = [(int(c[0]), int(c[1])) for c in choices.split(':')]

    _size = kwargs.get('fontsize', kwargs.get('size', None))
    _fontproperties = kwargs.get('fontproperties', None)
    font_size_pts = text_size_points(size=_size,
                                     fontproperties=_fontproperties)

    # set the default padding equal to the font size
    if paddingcoords == 'offset pixels':
        default_padding = font_size_pts * 72 / ax.figure.dpi
    elif paddingcoords == 'offset points':
        default_padding = font_size_pts
    elif paddingcoords == 'axes fraction':
        default_padding = 0.05
    else:
        raise ValueError("Bad padding coords '{0}'".format(paddingcoords))

    # print("fontsize pt:", font_size_pts,
    #       "fontsize px:", xy_as_pixels([font_size_pts, font_size_pts],
    #                                    'offset points')[0])

    if not isinstance(padding, (list, tuple)):
        padding = [padding, padding]
    padding = [default_padding if pd is None else pd for pd in padding]
    # print("padding::", paddingcoords, padding)

    magnet_px = xy_as_pixels(magnet, magnetcoords, ax=ax)
    padding_px = xy_as_pixels(padding, paddingcoords, ax=ax)
    # print("padding px::", padding_px)

    annotations = []

    cand_map = {}
    for choice in choices:
        cand_map[choice] = np.zeros([n_candidates, 2, 2], dtype='f')

    # these paths are all the paths we can get our hands on so that the text
    # doesn't overlap them. bboxes around labels are added as we go
    paths_px = []
    # here is a list of bounding boxes around the text boxes as we add them
    bbox_paths_px = []
    is_filled = []

    ## how many vertices to avoid ?
    # artist
    # collection
    # image
    # line
    # patch
    # table
    # container

    for line in ax.lines:
        paths_px += [ax.transData.transform_path(line.get_path())]
        is_filled += [False]
    for collection in ax.collections:
        for pth in collection.get_paths():
            paths_px += [ax.transData.transform_path(pth)]
            is_filled += [collection.get_fill()]

    if ignore_filling:
        is_filled = [False] * len(is_filled)

    hands, hand_labels = ax.get_legend_handles_labels()

    colors = _cycle_colors(colors, len(hands))

    # >>> debug >>>
    if _debug:
        import viscid
        from matplotlib import patches as mpatches
        from viscid.plot import vpyplot as vlt

        _fig_width = int(ax.figure.bbox.width)
        _fig_height = int(ax.figure.bbox.height)
        fig_fld = viscid.zeros((_fig_width, _fig_height),
                               dtype='f',
                               center='node')
        _X, _Y = fig_fld.get_crds(shaped=True)

        _axXL, _axYL, _axXH, _axYH = ax.bbox.extents

        _mask = np.bitwise_and(np.bitwise_and(_X >= _axXL, _X <= _axXH),
                               np.bitwise_and(_Y >= _axYL, _Y <= _axYH))
        fig_fld.data[_mask] = 1.0

        dfig, dax = plt.subplots(1, 1, figsize=ax.figure.get_size_inches())

        vlt.plot(fig_fld, ax=dax, cmap='ocean', colorbar=None)
        for _, path in enumerate(paths_px):
            dax.plot(path.vertices[:, 0], path.vertices[:, 1])
        dfig.subplots_adjust(bottom=0.0, left=0.0, top=1.0, right=1.0)
    else:
        dfig, dax = None, None
    # <<< debug <<<

    for i, hand, label_i in zip(count(), hands, hand_labels):
        if labels and i < len(labels):
            label = labels[i]
        else:
            label = label_i

        # divine color of label
        if colors[i]:
            color = colors[i]
        else:
            try:
                color = hand.get_color()
            except AttributeError:
                color = hand.get_facecolor()[0]

        # get path vertices to determine candidate label positions
        try:
            verts = hand.get_path().vertices
        except AttributeError:
            verts = [p.vertices for p in hand.get_paths()]
            verts = np.concatenate(verts, axis=0)

        segl_dat = verts[:-1, :]
        segh_dat = verts[1:, :]

        # take out path segments that have one vertex outside the view
        _seg_mask = np.all(
            np.bitwise_and(segl_dat >= axbb0, segl_dat <= axbb1)
            & np.bitwise_and(segh_dat >= axbb0, segh_dat <= axbb1),
            axis=1)
        segl_dat = segl_dat[_seg_mask, :]
        segh_dat = segh_dat[_seg_mask, :]

        if np.prod(segl_dat.shape) == 0:
            print("no full segments are visible, skipping path", i, hand)
            continue

        segl_px = ax.transData.transform(segl_dat)
        segh_px = ax.transData.transform(segh_dat)
        seglen_px = np.linalg.norm(segh_px - segl_px, axis=1)

        # take out path segments that are 0 pixels in length
        _non0_seg_mask = seglen_px > 0
        segl_dat = segl_dat[_non0_seg_mask, :]
        segh_dat = segh_dat[_non0_seg_mask, :]
        segl_px = segl_px[_non0_seg_mask, :]
        segh_px = segh_px[_non0_seg_mask, :]
        seglen_px = seglen_px[_non0_seg_mask]

        if np.prod(segl_dat.shape) == 0:
            print("no non-0 segments are visible, skipping path", i, hand)
            continue

        # i deeply appologize for how convoluted this got, but the punchline
        # is that each line segment gets candidates proportinal to their
        # length in pixels on the figure
        s_src = np.concatenate([[0], np.cumsum(seglen_px)])
        if rand_state is not None:
            s_dest = s_src[-1] * np.sort(np.random.rand(n_candidates))
        else:
            s_dest = np.linspace(0, s_src[-1], n_candidates)
        _diff = s_dest.reshape(1, -1) - s_src.reshape(-1, 1)
        iseg = np.argmin(np.ma.masked_where(_diff <= 0, _diff), axis=0)
        frac = (s_dest - s_src[iseg]) / seglen_px[iseg]

        root_dat = (segl_dat[iseg] + frac.reshape(-1, 1) *
                    (segh_dat[iseg] - segl_dat[iseg]))
        root_px = ax.transData.transform(root_dat)

        # estimate the width and height of the label's text
        txt_size = np.array(
            estimate_text_size_px(label, fig=ax.figure, size=font_size_pts))
        txt_size = txt_size.reshape([1, 2])

        # this initial offset is needed to shift the center of the label
        # to the data point
        offset0 = -txt_size / 2

        # now we can shift the label away from the data point by an amount
        # equal to half the text width/height + the padding
        offset1 = padding_px + txt_size / 2

        for key, abs_px_arr in cand_map.items():
            ioff = np.array(key, dtype='i').reshape(1, 2) - 1
            total_offset = offset0 + ioff * offset1
            # approx lower left corner of the text box in absolute pixels
            abs_px_arr[:, :, 0] = root_px + total_offset
            # approx upper right corner of the text box in absolute pixels
            abs_px_arr[:, :, 1] = abs_px_arr[:, :, 0] + txt_size

        # candidates_abs_px[i] has root @ root_px[i % n_candidates]
        candidates_abs_px = np.concatenate([cand_map[c] for c in choices],
                                           axis=0)

        # find how many other things each candidate overlaps
        n_overlaps = np.zeros_like(candidates_abs_px[:, 0, 0])

        for k, candidate in enumerate(candidates_abs_px):
            cand_bbox = Bbox(candidate.T)

            # penalty for each time a box overlaps a path that's already
            # on the plot
            for ipth, path in enumerate(paths_px):
                if path.intersects_bbox(cand_bbox, filled=is_filled[ipth]):
                    n_overlaps[k] += 1

            # slightly larger penalty if we intersect a text box that we
            # just added to the plot
            for ipth, path in enumerate(bbox_paths_px):
                if path.intersects_bbox(cand_bbox, filled=is_filled[ipth]):
                    n_overlaps[k] += 5

            # big penalty if the candidate is out of the current view
            if not (ax.bbox.contains(*cand_bbox.min)
                    and ax.bbox.contains(*cand_bbox.max)):
                n_overlaps[k] += 100

        # sort candidates by distance between center of text box and magnet
        magnet_dist = np.linalg.norm(np.mean(candidates_abs_px, axis=-1) -
                                     magnet_px,
                                     axis=1)
        isorted = np.argsort(magnet_dist)
        magnet_dist = np.array(magnet_dist[isorted])
        candidates_abs_px = np.array(candidates_abs_px[isorted, :, :])
        n_overlaps = np.array(n_overlaps[isorted])
        root_dat = np.array(root_dat[isorted % n_candidates, :])
        root_px = np.array(root_px[isorted % n_candidates, :])

        # sort candidates so the ones with the fewest overlaps are first
        # but do it with a stable algorithm so among the best candidates,
        # choose the one closest to the magnet
        sargs = np.argsort(n_overlaps, kind='mergesort')

        # >>> debug >>>
        if dax is not None:
            for _candidate, n_overlap in zip(candidates_abs_px, n_overlaps):
                _cand_bbox = Bbox(_candidate.T)
                _x0 = _cand_bbox.get_points()[0]
                _bbox_center = np.mean(_candidate, axis=-1)
                _ray_x = [_bbox_center[0], magnet_px[0]]
                _ray_y = [_bbox_center[1], magnet_px[1]]
                dax.plot(_ray_x, _ray_y, '-', alpha=0.3, color='grey')
                _rect = mpatches.Rectangle(_x0,
                                           _cand_bbox.width,
                                           _cand_bbox.height,
                                           fill=False)
                dax.add_patch(_rect)
                plt.text(_x0[0], _x0[1], label, color='gray')
                plt.text(_x0[0], _x0[1], '{0}'.format(n_overlap))
        # <<< debug <<<

        # pick winning candidate and add its bounding box to this list of
        # paths to avoid
        winner_abs_px = candidates_abs_px[sargs[0], :, :]
        xy_root_px = root_px[sargs[0], :]
        xy_root_dat = np.array(root_dat[sargs[0], :])
        xy_txt_offset = np.array(winner_abs_px[:, 0] - xy_root_px)

        corners = Bbox(winner_abs_px.T).corners()[(0, 1, 3, 2), :]
        bbox_paths_px += [Path(corners)]

        # a = plt.annotate(label, xy=xy_root_dat, xycoords='data',
        #                  xytext=xy_txt_offset, textcoords="offset pixels",
        #                  color=color, **kwargs)
        a = ax.annotate(label,
                        xy=xy_root_dat,
                        xycoords='data',
                        xytext=xy_txt_offset,
                        textcoords="offset pixels",
                        color=color,
                        **kwargs)
        annotations.append(a)

    if rand_state is not None:
        np.random.set_state(rand_state)

    return annotations
Exemplo n.º 44
0
    # run eval
    fld = eval(salted_eqn, {"__builtins__": {}}, local_dict)  # pylint: disable=eval-used
    try:
        fld.name = result_name
        fld.pretty_name = result_name
    except AttributeError:
        pass
    return fld


if __name__ == "__main__":
    import os
    import viscid
    from viscid.plot import vpyplot as vlt
    import matplotlib.pyplot as plt
    enabled = True
    _d = os.path.dirname(viscid.__file__)
    _g = viscid.load_file(_d + "/../sample/sample.py_0.xdmf").get_grid()
    plt.subplot(211)
    _fld = evaluate(_g, "speed", "sqrt(vx**2 + vy**2 + sqrt(vz**4))")
    vlt.plot(_fld, show=False)
    plt.subplot(212)
    _fld = evaluate(_g, "speed", "sqrt(vx**2 + vy**2 + sqrt(vz**4))",
                    try_numexpr=False)
    vlt.plot(_fld, show=True)

##
## EOF
##
Exemplo n.º 45
0
def main():
    mhd_type = "C"
    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()

    do_fill_dipole = True

    gslc = "x=-21.2j:12j, y=-11j:11j, z=-11j:11j"
    b = f['b_cc'][gslc]
    b1 = f['b_fc'][gslc]
    e_cc = f['e_cc'][gslc]
    e_ec = f['e_ec'][gslc]

    if do_fill_dipole:
        mask = viscid.make_spherical_mask(b, rmax=3.5)
        viscid.fill_dipole(b, mask=mask)

        mask = viscid.make_spherical_mask(b1, rmax=3.5)
        viscid.fill_dipole(b1, mask=mask)

        mask = None

    # seeds = viscid.SphericalCap(r=1.02, ntheta=64, nphi=32, angle0=17, angle=20,
    #                             philim=(100, 260), roll=-180.0)
    # seeds = viscid.SphericalCap(r=1.02, ntheta=64, nphi=32, angle0=17, angle=20,
    #                             philim=(0, 10), roll=0.0)
    seedsN = viscid.Sphere(r=1.02, ntheta=16, nphi=16, thetalim=(15, 25),
                           philim=(0, 300), crd_system=b)
    seedsS = viscid.Sphere(r=1.02, ntheta=16, nphi=16, thetalim=(155, 165),
                           philim=(0, 300), crd_system=b)

    bl_kwargs = dict(ibound=0.9, obound0=(-20, -10, -10), obound1=(11, 10, 10))

    # blines_cc, topo_cc = viscid.streamlines(b, seeds, **bl_kwargs)
    blinesN_fc, topoN_fc = viscid.streamlines(b1, seedsN, **bl_kwargs)
    _, topoS_fc = viscid.streamlines(b1, seedsS, output=viscid.OUTPUT_TOPOLOGY,
                                     **bl_kwargs)

    if True:
        from viscid.plot import vlab
        mesh = vlab.mesh_from_seeds(seedsN, scalars=topoN_fc)
        mesh.actor.property.backface_culling = True
        # vlab.plot_lines(blines_cc, scalars="#000000", tube_radius=0.03)
        vlab.plot_lines(blinesN_fc, scalars=viscid.topology2color(topoN_fc),
                        opacity=0.7)

        vlab.plot_blue_marble(r=1.0)
        vlab.plot_earth_3d(radius=1.01, crd_system=b, night_only=True,
                           opacity=0.5)
        vlab.show()

    if True:
        vlt.subplot(121, projection='polar')
        vlt.plot(topoN_fc)
        vlt.subplot(122, projection='polar')
        vlt.plot(topoS_fc)
        vlt.show()

    return 0
Exemplo n.º 46
0
def main():
    xl, xh, nx = -1.0, 1.0, 41
    yl, yh, ny = -1.5, 1.5, 41
    zl, zh, nz = -2.0, 2.0, 41
    x = np.linspace(xl, xh, nx)
    y = np.linspace(yl, yh, ny)
    z = np.linspace(zl, zh, nz)
    crds = coordinate.wrap_crds("nonuniform_cartesian",
                                [('z', z), ('y', y), ('x', x)])
    bx = field.empty(crds, name="$B_x$", center="Node")
    by = field.empty(crds, name="$B_y$", center="Node")
    bz = field.empty(crds, name="$B_z$", center="Node")
    fld = field.empty(crds, name="B", nr_comps=3, center="Node",
                      layout="interlaced")
    X, Y, Z = crds.get_crds(shaped=True)

    x01, y01, z01 = 0.5, 0.5, 0.5
    x02, y02, z02 = 0.5, 0.5, 0.5
    x03, y03, z03 = 0.5, 0.5, 0.5

    bx[:] = 0.0 + 1.0 * (X - x01) + 1.0 * (Y - y01) + 1.0 * (Z - z01) + \
              1.0 * (X - x01) * (Y - y01) + 1.0 * (Y - y01) * (Z - z01) + \
              1.0 * (X - x01) * (Y - y01) * (Z - z01)
    by[:] = 0.0 + 1.0 * (X - x02) - 1.0 * (Y - y02) + 1.0 * (Z - z02) + \
              1.0 * (X - x02) * (Y - y02) + 1.0 * (Y - y02) * (Z - z02) - \
              1.0 * (X - x02) * (Y - y02) * (Z - z02)
    bz[:] = 0.0 + 1.0 * (X - x03) + 1.0 * (Y - y03) - 1.0 * (Z - z03) + \
              1.0 * (X - x03) * (Y - y03) + 1.0 * (Y - y03) * (Z - z03) + \
              1.0 * (X - x03) * (Y - y03) * (Z - z03)
    fld[..., 0] = bx
    fld[..., 1] = by
    fld[..., 2] = bz

    fig = mlab.figure(size=(1150, 850),
                      bgcolor=(1.0, 1.0, 1.0),
                      fgcolor=(0.0, 0.0, 0.0))
    f1_src = vlab.add_field(bx)
    f2_src = vlab.add_field(by)
    f3_src = vlab.add_field(bz)
    mlab.pipeline.iso_surface(f1_src, contours=[0.0],
                              opacity=1.0, color=(1.0, 0.0, 0.0))
    mlab.pipeline.iso_surface(f2_src, contours=[0.0],
                              opacity=1.0, color=(0.0, 1.0, 0.0))
    mlab.pipeline.iso_surface(f3_src, contours=[0.0],
                              opacity=1.0, color=(0.0, 0.0, 1.0))
    mlab.axes()
    mlab.show()

    nullpt = cycalc.interp_trilin(fld, [(0.5, 0.5, 0.5)])
    print("f(0.5, 0.5, 0.5):", nullpt)

    _, axes = plt.subplots(4, 3, sharex=True, sharey=True)
    all_roots = []
    positive_roots = []
    ix = iy = iz = 0

    for di, d in enumerate([0, -1]):
        #### XY face
        a1 = bx[iz + d, iy, ix]
        b1 = bx[iz + d, iy, ix - 1] - a1
        c1 = bx[iz + d, iy - 1, ix] - a1
        d1 = bx[iz + d, iy - 1, ix - 1] - c1 - b1 - a1

        a2 = by[iz + d, iy, ix]
        b2 = by[iz + d, iy, ix - 1] - a2
        c2 = by[iz + d, iy - 1, ix] - a2
        d2 = by[iz + d, iy - 1, ix - 1] - c2 - b2 - a2

        a3 = bz[iz + d, iy, ix]
        b3 = bz[iz + d, iy, ix - 1] - a3
        c3 = bz[iz + d, iy - 1, ix] - a3
        d3 = bz[iz + d, iy - 1, ix - 1] - c3 - b3 - a3

        roots1, roots2 = find_roots_face(a1, b1, c1, d1, a2, b2, c2, d2)

        # for rt1, rt2 in zip(roots1, roots2):
        #     print("=")
        #     print("fx", a1 + b1 * rt1 + c1 * rt2 + d1 * rt1 * rt2)
        #     print("fy", a2 + b2 * rt1 + c2 * rt2 + d2 * rt1 * rt2)
        #     print("=")

        # find f3 at the root points
        f3 = np.empty_like(roots1)
        markers = [None] * len(f3)
        for i, rt1, rt2 in zip(count(), roots1, roots2):
            f3[i] = a3 + b3 * rt1 + c3 * rt2 + d3 * rt1 * rt2
            all_roots.append((rt1, rt2, d))  # switch order here
            if f3[i] >= 0.0:
                markers[i] = 'k^'
                positive_roots.append((rt1, rt2, d))  # switch order here
            else:
                markers[i] = 'w^'

        # rescale the roots to the original domain
        roots1 = (xh - xl) * roots1 + xl
        roots2 = (yh - yl) * roots2 + yl

        xp = np.linspace(0.0, 1.0, nx)

        vlt.plot(fld['x'], "z={0}i".format(d), ax=axes[0 + 2 * di, 0],
                 plot_opts="x={0}_{1},y={2}_{3},lin_-10_10".format(xl, xh, yl, yh))
        y1 = - (a1 + b1 * xp) / (c1 + d1 * xp)
        plt.plot(x, (yh - yl) * y1 + yl, 'k')
        for i, xrt, yrt in zip(count(), roots1, roots2):
            plt.plot(xrt, yrt, markers[i])

        vlt.plot(fld['y'], "z={0}i".format(d), ax=axes[1 + 2 * di, 0],
                 plot_opts="x={0}_{1},y={2}_{3},lin_-10_10".format(xl, xh, yl, yh))
        y2 = - (a2 + b2 * xp) / (c2 + d2 * xp)
        plt.plot(x, (yh - yl) * y2 + yl, 'k')
        for xrt, yrt in zip(roots1, roots2):
            plt.plot(xrt, yrt, markers[i])

        #### YZ face
        a1 = bx[iz, iy, ix + d]
        b1 = bx[iz, iy - 1, ix + d] - a1
        c1 = bx[iz - 1, iy, ix + d] - a1
        d1 = bx[iz - 1, iy - 1, ix + d] - c1 - b1 - a1

        a2 = by[iz, iy, ix + d]
        b2 = by[iz, iy - 1, ix + d] - a2
        c2 = by[iz - 1, iy, ix + d] - a2
        d2 = by[iz - 1, iy - 1, ix + d] - c2 - b2 - a2

        a3 = bz[iz, iy, ix + d]
        b3 = bz[iz, iy - 1, ix + d] - a3
        c3 = bz[iz - 1, iy, ix + d] - a3
        d3 = bz[iz - 1, iy - 1, ix + d] - c3 - b3 - a3

        roots1, roots2 = find_roots_face(a1, b1, c1, d1, a2, b2, c2, d2)

        # for rt1, rt2 in zip(roots1, roots2):
        #     print("=")
        #     print("fx", a1 + b1 * rt1 + c1 * rt2 + d1 * rt1 * rt2)
        #     print("fy", a2 + b2 * rt1 + c2 * rt2 + d2 * rt1 * rt2)
        #     print("=")

        # find f3 at the root points
        f3 = np.empty_like(roots1)
        markers = [None] * len(f3)
        for i, rt1, rt2 in zip(count(), roots1, roots2):
            f3[i] = a3 + b3 * rt1 + c3 * rt2 + d3 * rt1 * rt2
            all_roots.append((d, rt1, rt2))  # switch order here
            if f3[i] >= 0.0:
                markers[i] = 'k^'
                positive_roots.append((d, rt1, rt2))  # switch order here
            else:
                markers[i] = 'w^'

        # rescale the roots to the original domain
        roots1 = (yh - yl) * roots1 + yl
        roots2 = (zh - zl) * roots2 + zl

        yp = np.linspace(0.0, 1.0, ny)

        # plt.subplot(121)
        vlt.plot(fld['x'], "x={0}i".format(d), ax=axes[0 + 2 * di, 1],
                 plot_opts="x={0}_{1},y={2}_{3},lin_-10_10".format(yl, yh, zl, zh))
        z1 = - (a1 + b1 * yp) / (c1 + d1 * yp)
        plt.plot(y, (zh - zl) * z1 + zl, 'k')
        for i, yrt, zrt in zip(count(), roots1, roots2):
            plt.plot(yrt, zrt, markers[i])

        # plt.subplot(122)
        vlt.plot(fld['y'], "x={0}i".format(d), ax=axes[1 + 2 * di, 1],
                 plot_opts="x={0}_{1},y={2}_{3},lin_-10_10".format(yl, yh, zl, zh))
        z1 = - (a2 + b2 * yp) / (c2 + d2 * yp)
        plt.plot(y, (zh - zl) * z1 + zl, 'k')
        for i, yrt, zrt in zip(count(), roots1, roots2):
            plt.plot(yrt, zrt, markers[i])

        #### ZX face
        a1 = bx[iz, iy + d, ix]
        b1 = bx[iz - 1, iy + d, ix] - a1
        c1 = bx[iz, iy + d, ix - 1] - a1
        d1 = bx[iz - 1, iy + d, ix - 1] - c1 - b1 - a1

        a2 = by[iz, iy + d, ix]
        b2 = by[iz - 1, iy + d, ix] - a2
        c2 = by[iz, iy + d, ix - 1] - a2
        d2 = by[iz - 1, iy + d, ix - 1] - c2 - b2 - a2

        a3 = bz[iz, iy + d, ix]
        b3 = bz[iz - 1, iy + d, ix] - a3
        c3 = bz[iz, iy + d, ix - 1] - a3
        d3 = bz[iz - 1, iy + d, ix - 1] - c3 - b3 - a3

        roots1, roots2 = find_roots_face(a1, b1, c1, d1, a2, b2, c2, d2)

        # for rt1, rt2 in zip(roots1, roots2):
        #     print("=")
        #     print("fx", a1 + b1 * rt1 + c1 * rt2 + d1 * rt1 * rt2)
        #     print("fy", a2 + b2 * rt1 + c2 * rt2 + d2 * rt1 * rt2)
        #     print("=")

        # find f3 at the root points
        f3 = np.empty_like(roots1)
        markers = [None] * len(f3)
        for i, rt1, rt2 in zip(count(), roots1, roots2):
            f3[i] = a3 + b3 * rt1 + c3 * rt2 + d3 * rt1 * rt2
            all_roots.append((rt2, d, rt1))  # switch order here
            if f3[i] >= 0.0:
                markers[i] = 'k^'
                positive_roots.append((rt2, d, rt1))  # switch order here
            else:
                markers[i] = 'w^'

        # rescale the roots to the original domain
        roots1 = (zh - zl) * roots1 + zl
        roots2 = (xh - xl) * roots2 + xl

        zp = np.linspace(0.0, 1.0, nz)

        # plt.subplot(121)
        vlt.plot(fld['x'], "y={0}i".format(d), ax=axes[0 + 2 * di, 2],
                 plot_opts="x={0}_{1},y={2}_{3},lin_-10_10".format(xl, xh, zl, zh))
        x1 = - (a1 + b1 * zp) / (c1 + d1 * zp)
        plt.plot(z, (xh - xl) * x1 + xl, 'k')
        for i, zrt, xrt in zip(count(), roots1, roots2):
            plt.plot(xrt, zrt, markers[i])

        # plt.subplot(121)
        vlt.plot(fld['y'], "y={0}i".format(d), ax=axes[1 + 2 * di, 2],
                 plot_opts="x={0}_{1},y={2}_{3},lin_-10_10".format(xl, xh, zl, zh))
        x1 = - (a2 + b2 * zp) / (c2 + d2 * zp)
        plt.plot(z, (xh - xl) * x1 + xl, 'k')
        for i, zrt, xrt in zip(count(), roots1, roots2):
            plt.plot(xrt, zrt, markers[i])

    print("all:", len(all_roots), "positive:", len(positive_roots))
    if len(all_roots) % 2 == 1:
        print("something is fishy, there are an odd number of root points "
              "on the surface of your cube, there is probably a degenerate "
              "line or surface of nulls")
    print("Null Point?", (len(positive_roots) % 2 == 1))

    plt.show()
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
Exemplo n.º 48
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)
    dtype = 'float32'

    ########################################################################
    # hard core test transpose (since this is used for mapfield transforms)
    x = np.array(np.linspace( 1, -1, 9), dtype=dtype)
    y = np.array(np.linspace(-1,  1, 9), dtype=dtype)
    z = np.array(np.linspace(-1,  1, 9), dtype=dtype)
    vI1 = viscid.empty([x, y, z], nr_comps=3, name='V', center='cell',
                       layout='interlaced')
    vI2 = viscid.empty([z, y, x], nr_comps=3, name='V', center='cell',
                       layout='interlaced', crd_names='zyx')
    vF1 = viscid.empty([x, y, z], nr_comps=3, name='V', center='cell',
                       layout='flat')
    vF2 = viscid.empty([z, y, x], nr_comps=3, name='V', center='cell',
                       layout='flat', crd_names='zyx')
    X1, Y1, Z1 = vI1.get_crds_cc(shaped=True)
    X2, Y2, Z2 = vI2.get_crds_cc(shaped=True)

    for v in (vI1, vF1):
        v['x'] = (0.5 * X1) + (0.0 * Y1) + (0.0 * Z1)
        v['y'] = (0.0 * X1) + (0.5 * Y1) + (0.5 * Z1)
        v['z'] = (0.0 * X1) + (0.5 * Y1) + (0.5 * Z1)
    for v in (vI2, vF2):
        v['z'] = (0.5 * X2) + (0.5 * Y2) + (0.0 * Z2)
        v['y'] = (0.5 * X2) + (0.5 * Y2) + (0.0 * Z2)
        v['x'] = (0.0 * X2) + (0.0 * Y2) + (0.5 * Z2)

    assert_different(vI1, vI2)

    # test some straight up transposes of both interlaced and flat fields
    assert_similar(vI1.spatial_transpose(), vI2)
    assert_similar(vI1.ST, vI2)
    assert_similar(vF1.spatial_transpose(), vF2)
    assert_similar(vI1.transpose(), vF2)
    assert_similar(vI1.T, vF2)
    assert_different(vI1.transpose(), vI2)
    assert_similar(vF1.transpose(), vI2)
    assert_different(vF1.transpose(), vF2)
    assert_similar(np.transpose(vI1), vF2)
    assert_similar(np.transpose(vF1), vI2)

    # now specify specific axes using all 3 interfaces
    assert_similar(vI1.spatial_transpose('x', 'z', 'y'), vI1)
    assert_similar(np.transpose(vI1, axes=[0, 2, 1, 3]), vI1)
    assert_similar(vI1.transpose(0, 2, 1, 3), vI1)
    assert_similar(vF1.spatial_transpose('x', 'z', 'y'), vF1)
    assert_similar(np.transpose(vF1, axes=(0, 1, 3, 2)), vF1)
    assert_similar(vF1.transpose(0, 1, 3, 2), vF1)

    # now test swapaxes since that uses
    assert_similar(vI1.swapaxes(1, 2), vI1)
    assert_similar(np.swapaxes(vI1, 1, 2), vI1)
    assert_similar(vI1.swap_crd_axes('y', 'z'), vI1)

    ##############################
    # test some other mathy stuff
    x = np.array(np.linspace(-1, 1, 2), dtype=dtype)
    y = np.array(np.linspace(-2, 2, 30), dtype=dtype)
    z = np.array(np.linspace(-5, 5, 90), dtype=dtype)
    v = viscid.empty([x, y, z], nr_comps=3, name='V', center='cell',
                     layout='interlaced')
    X, Y, Z = v.get_crds_cc(shaped=True)

    v['x'] = (0.5 * X**2) + (      Y   ) + (0.0 * Z   )
    v['y'] = (0.0 * X   ) + (0.5 * Y**2) + (0.0 * Z   )
    v['z'] = (0.0 * X   ) + (0.0 * Y   ) + (0.5 * Z**2)

    mag = viscid.magnitude(v)
    mag2 = np.sqrt(np.sum(v * v, axis=v.nr_comp))
    another = np.transpose(mag)

    plt.subplot(151)
    vlt.plot(v['x'])
    plt.subplot(152)
    vlt.plot(v['y'])
    plt.subplot(153)
    vlt.plot(mag)
    plt.subplot(154)
    vlt.plot(mag2)
    plt.subplot(155)
    vlt.plot(another)

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        plt.show()

    return 0
Exemplo n.º 49
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
Exemplo n.º 50
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)
    dtype = 'float32'

    ########################################################################
    # hard core test transpose (since this is used for mapfield transforms)
    x = np.array(np.linspace(1, -1, 9), dtype=dtype)
    y = np.array(np.linspace(-1, 1, 9), dtype=dtype)
    z = np.array(np.linspace(-1, 1, 9), dtype=dtype)
    vI1 = viscid.empty([x, y, z],
                       nr_comps=3,
                       name='V',
                       center='cell',
                       layout='interlaced')
    vI2 = viscid.empty([z, y, x],
                       nr_comps=3,
                       name='V',
                       center='cell',
                       layout='interlaced',
                       crd_names='zyx')
    vF1 = viscid.empty([x, y, z],
                       nr_comps=3,
                       name='V',
                       center='cell',
                       layout='flat')
    vF2 = viscid.empty([z, y, x],
                       nr_comps=3,
                       name='V',
                       center='cell',
                       layout='flat',
                       crd_names='zyx')
    X1, Y1, Z1 = vI1.get_crds_cc(shaped=True)
    X2, Y2, Z2 = vI2.get_crds_cc(shaped=True)

    for v in (vI1, vF1):
        v['x'] = (0.5 * X1) + (0.0 * Y1) + (0.0 * Z1)
        v['y'] = (0.0 * X1) + (0.5 * Y1) + (0.5 * Z1)
        v['z'] = (0.0 * X1) + (0.5 * Y1) + (0.5 * Z1)
    for v in (vI2, vF2):
        v['z'] = (0.5 * X2) + (0.5 * Y2) + (0.0 * Z2)
        v['y'] = (0.5 * X2) + (0.5 * Y2) + (0.0 * Z2)
        v['x'] = (0.0 * X2) + (0.0 * Y2) + (0.5 * Z2)

    assert_different(vI1, vI2)

    # test some straight up transposes of both interlaced and flat fields
    assert_similar(vI1.spatial_transpose(), vI2)
    assert_similar(vI1.ST, vI2)
    assert_similar(vF1.spatial_transpose(), vF2)
    assert_similar(vI1.transpose(), vF2)
    assert_similar(vI1.T, vF2)
    assert_different(vI1.transpose(), vI2)
    assert_similar(vF1.transpose(), vI2)
    assert_different(vF1.transpose(), vF2)
    assert_similar(np.transpose(vI1), vF2)
    assert_similar(np.transpose(vF1), vI2)

    # now specify specific axes using all 3 interfaces
    assert_similar(vI1.spatial_transpose('x', 'z', 'y'), vI1)
    assert_similar(np.transpose(vI1, axes=[0, 2, 1, 3]), vI1)
    assert_similar(vI1.transpose(0, 2, 1, 3), vI1)
    assert_similar(vF1.spatial_transpose('x', 'z', 'y'), vF1)
    assert_similar(np.transpose(vF1, axes=(0, 1, 3, 2)), vF1)
    assert_similar(vF1.transpose(0, 1, 3, 2), vF1)

    # now test swapaxes since that uses
    assert_similar(vI1.swapaxes(1, 2), vI1)
    assert_similar(np.swapaxes(vI1, 1, 2), vI1)
    assert_similar(vI1.swap_crd_axes('y', 'z'), vI1)

    ##############################
    # test some other mathy stuff
    x = np.array(np.linspace(-1, 1, 2), dtype=dtype)
    y = np.array(np.linspace(-2, 2, 30), dtype=dtype)
    z = np.array(np.linspace(-5, 5, 90), dtype=dtype)
    v = viscid.empty([x, y, z],
                     nr_comps=3,
                     name='V',
                     center='cell',
                     layout='interlaced')
    X, Y, Z = v.get_crds_cc(shaped=True)

    v['x'] = (0.5 * X**2) + (Y) + (0.0 * Z)
    v['y'] = (0.0 * X) + (0.5 * Y**2) + (0.0 * Z)
    v['z'] = (0.0 * X) + (0.0 * Y) + (0.5 * Z**2)

    mag = viscid.magnitude(v)
    mag2 = np.sqrt(np.sum(v * v, axis=v.nr_comp))
    another = np.transpose(mag)

    plt.subplot(151)
    vlt.plot(v['x'])
    plt.subplot(152)
    vlt.plot(v['y'])
    plt.subplot(153)
    vlt.plot(mag)
    plt.subplot(154)
    vlt.plot(mag2)
    plt.subplot(155)
    vlt.plot(another)

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        plt.show()

    return 0
Exemplo n.º 51
0
def _do_multiplot(tind, grid, plot_vars=None, global_popts=None, kwopts=None,
                  share_axes=False, show=False, subplot_params=None,
                  first_run_result=None, first_run=False, **kwargs):
    import matplotlib.pyplot as plt
    from viscid.plot import vpyplot as vlt

    logger.info("Plotting timestep: %d, %g", tind, grid.time)

    if plot_vars is None:
        raise ValueError("No plot_vars given to `_do_multiplot` :(")
    if kwargs:
        logger.info("Unused kwargs: {0}".format(kwargs))

    if kwopts is None:
        kwopts = {}
    transpose = kwopts.get("transpose", False)
    plot_size = kwopts.get("plot_size", None)
    dpi = kwopts.get("dpi", None)
    out_prefix = kwopts.get("out_prefix", None)
    out_format = kwopts.get("out_format", "png")
    selection = kwopts.get("selection", None)
    timeformat = kwopts.get("timeformat", ".02f")
    tighten = kwopts.get("tighten", False)
    # wicked hacky
    # subplot_params = kwopts.get("subplot_params", _subplot_params)

    # nrows = len(plot_vars)
    nrows = len([pv[0] for pv in plot_vars if not pv[0].startswith('^')])
    ncols = 1
    if transpose:
        nrows, ncols = ncols, nrows

    if nrows == 0:
        logger.warn("I have no variables to plot")
        return

    fig = plt.gcf()
    if plot_size is not None:
        fig.set_size_inches(*plot_size, forward=True)
    if dpi is not None:
        fig.set_dpi(dpi)

    shareax = None

    this_row = -1
    for i, fld_meta in enumerate(plot_vars):
        if not fld_meta[0].startswith('^'):
            this_row += 1
            same_axis = False
        else:
            same_axis = True

        fld_name_meta = fld_meta[0].lstrip('^')
        fld_name_split = fld_name_meta.split(',')
        if '=' in fld_name_split[0]:
            # if fld_name is actually an equation, assume
            # there's no slice, and commas are part of the
            # equation
            fld_name = ",".join(fld_name_split)
            fld_slc = ""
        else:
            fld_name = fld_name_split[0]
            fld_slc = ",".join(fld_name_split[1:])
        if selection is not None:
            # fld_slc += ",{0}".format(selection)
            if fld_slc != "":
                fld_slc = ",".join([fld_slc, selection])
            else:
                fld_slc = selection
        if fld_slc.strip() == "":
            fld_slc = None

        # print("fld_time:", fld.time)
        if this_row < 0:
            raise ValueError("first plot can't begin with a +")
        row = this_row
        col = 0
        if transpose:
            row, col = col, row
        if not same_axis:
            ax = plt.subplot2grid((nrows, ncols), (row, col),
                                  sharex=shareax, sharey=shareax)
        if i == 0 and share_axes:
            shareax = ax

        if "plot_opts" not in fld_meta[1]:
            fld_meta[1]["plot_opts"] = global_popts
        elif global_popts is not None:
            fld_meta[1]["plot_opts"] = "{0},{1}".format(
                fld_meta[1]["plot_opts"], global_popts)

        with grid.get_field(fld_name, slc=fld_slc) as fld:
            vlt.plot(fld, masknan=True, **fld_meta[1])
        # print("fld cache", grid[fld_meta[0]]._cache)

    if timeformat and timeformat.lower() != "none":
        plt.suptitle(grid.format_time(timeformat))

    # for adjusting subplots / tight_layout and applying the various
    # hacks to keep plots from dancing around in movies
    if not subplot_params and first_run_result:
        subplot_params = first_run_result
    if tighten:
        tighten = dict(rect=[0, 0.03, 1, 0.90])
    ret = vlt.auto_adjust_subplots(tight_layout=tighten,
                                   subplot_params=subplot_params)
    if not first_run:
        ret = None

    if out_prefix:
        plt.savefig("{0}_{1:06d}.{2}".format(out_prefix, tind + 1, out_format))
    if show:
        plt.show()
    plt.clf()

    return ret
Exemplo n.º 52
0
def main():
    mhd_type = "C"
    make_plots = 1
    test_fc = 1
    test_ec = 1
    test_div = 1
    test_interp = 1
    test_streamline = 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 = 5e-6
    elif mhd_type in ("F", "FORTRAN"):
        f = viscid.load_file("$WORK/tmedium3/*.3df.[-1]")
        catol = 1e-8
        rtol = 7e-2
    else:
        raise ValueError()

    ISLICE = slice(None)
    # ISLICE = 'y=0j:0.15j'

    # #################
    # # test out fc2cc
    if test_fc:
        b = f['b'][ISLICE]
        b1 = f['b1'][ISLICE]

        compare_vectors(b, b1, viscid.fc2cc, catol=catol, rtol=rtol,
                        make_plots=make_plots)

    #################
    # test out ec2cc
    if test_ec:
        e_cc = f['e_cc'][ISLICE]
        e_ec = f['e_ec'][ISLICE]

        if mhd_type not in ("F", "FORTRAN"):
            compare_vectors(e_cc, e_ec, viscid.ec2cc, catol=catol, rtol=rtol,
                            make_plots=make_plots)

    #################
    # test out divfc
    # Note: Relative error on Div B is meaningless b/c the coordinates
    #       are not the same up to order (dx/4) I think. You can see this
    #       since (fcdiv - divb_trimmed) is both noisy and stripy
    if test_div:
        bnd = 0

        if mhd_type not in ("F", "FORTRAN"):
            b1 = f['b1'][ISLICE]
            divb = f['divB'][ISLICE]
            if bnd:
                trimmed = divb
            else:
                trimmed = divb['x=1:-1, y=1:-1, z=1:-1']
            b1mag = viscid.magnitude(viscid.fc2cc(b1, bnd=bnd))

            divb1 = viscid.div_fc(b1, bnd=bnd)

            viscid.set_in_region(trimmed, trimmed, alpha=0.0, beta=0.0, out=trimmed,
                                 mask=viscid.make_spherical_mask(trimmed, rmax=5.0))
            viscid.set_in_region(divb1, divb1, alpha=0.0, beta=0.0, out=divb1,
                                 mask=viscid.make_spherical_mask(divb1, rmax=5.0))

            reldiff = (divb1 - trimmed) / b1mag
            reldiff = reldiff["x=1:-1, y=1:-1, z=1:-1"]
            reldiff.name = divb1.name + " - " + trimmed.name
            reldiff.pretty_name = divb1.pretty_name + " - " + trimmed.pretty_name

            abs_max_rel_diff = np.nanmax(np.abs(reldiff))
            max_crd_diff = [0.0] * 3
            for i, d in enumerate('xyz'):
                max_crd_diff[i] = np.max(trimmed.get_crd(d) - divb1.get_crd(d))
            print("divB max absolute relative diff: {0:.3e} "
                  "(crds: X: {1[0]:.3e}, Y: {1[1]:.3e}, Z: {1[2]:.3e})"
                  "".format(abs_max_rel_diff, max_crd_diff))

            # plot differences?
            if make_plots:
                ax1 = plt.subplot(311)
                vlt.plot(divb['y=0j'], symmetric=True, earth=True)
                plt.subplot(312, sharex=ax1, sharey=ax1)
                vlt.plot(divb1['y=0j'], symmetric=True, earth=True)
                plt.subplot(313, sharex=ax1, sharey=ax1)
                vlt.plot(reldiff['y=0j'], symmetric=True, earth=True)
                vlt.show()

            # Since the coordinates will be different by order dx^2 (i think),
            # there is no way to compare the divB from simulation with the
            # one we get here. However, they should be the same up to a few %, and
            # down to noise level with stripes of enhanced noise. These stripes
            # are the errors in the coordinate values (since the output only
            # gives us weird nc = averaged cc locations)
            #
            # if abs_max_rel_diff > rtol or np.any(np.abs(max_crd_diff) > catol):
            #     raise RuntimeError("Tolerance exceeded on divB calculation")

    if test_streamline:
        b_cc = f['b_cc']['x=-40j:12j, y=-15j:15j, z=-15j:15j']
        b_fc = f['b_fc']['x=-40j:12j, y=-15j:15j, z=-15j:15j']

        cotr = viscid.cotr.Cotr()
        r_mask = 3.0
        # set b_cc to dipole inside some sphere
        isphere_mask = viscid.make_spherical_mask(b_cc, rmax=r_mask)
        moment = cotr.get_dipole_moment(crd_system=b_cc)
        viscid.fill_dipole(b_cc, m=moment, mask=isphere_mask)
        # set b_fc to dipole inside some sphere
        isphere_mask = viscid.make_spherical_mask(b_fc, rmax=r_mask)
        moment = cotr.get_dipole_moment(crd_system=b_fc)
        viscid.fill_dipole(b_fc, m=moment, mask=isphere_mask)

        seeds = viscid.Volume([-10, 0, -5], [10, 0, 5], (16, 1, 3))
        sl_kwargs = dict(ibound=1.0, method=viscid.EULER1A)
        lines_cc, topo_cc = viscid.calc_streamlines(b_cc, seeds, **sl_kwargs)
        lines_fc, topo_fc = viscid.calc_streamlines(b_fc, seeds, **sl_kwargs)

        if make_plots:
            plt.figure(figsize=(10, 6))

            ax0 = plt.subplot(211)
            topo_cc_colors = viscid.topology2color(topo_cc)
            vlt.plot(f['pp']['y=0j'], logscale=True, earth=True, cmap='plasma')
            vlt.plot2d_lines(lines_cc, topo_cc_colors, symdir='y')

            ax0 = plt.subplot(212, sharex=ax0, sharey=ax0)
            topo_fc_colors = viscid.topology2color(topo_fc)
            vlt.plot(f['pp']['y=0j'], logscale=True, earth=True, cmap='plasma')
            vlt.plot2d_lines(lines_fc, topo_fc_colors, symdir='y')

            plt.xlim(-20, 10)
            plt.ylim(-10, 10)
            vlt.auto_adjust_subplots()
            vlt.show()

    if test_interp:
        # test interpolation with E . B / B
        b_cc = f['b_cc']
        b_fc = f['b_fc']
        e_cc = f['e_cc']
        e_ec = f['e_ec']

        cotr = viscid.cotr.Cotr()
        r_mask = 3.0
        # set b_cc to dipole inside some sphere
        isphere_mask = viscid.make_spherical_mask(b_cc, rmax=r_mask)
        moment = cotr.get_dipole_moment(crd_system=b_cc)
        viscid.fill_dipole(b_cc, m=moment, mask=isphere_mask)
        # set b_fc to dipole inside some sphere
        isphere_mask = viscid.make_spherical_mask(b_fc, rmax=r_mask)
        moment = cotr.get_dipole_moment(crd_system=b_fc)
        viscid.fill_dipole(b_fc, m=moment, mask=isphere_mask)
        # zero out e_cc inside some sphere
        viscid.set_in_region(e_cc, e_cc, alpha=0.0, beta=0.0, out=e_cc,
                             mask=viscid.make_spherical_mask(e_cc, rmax=r_mask))
        # zero out e_ec inside some sphere
        viscid.set_in_region(e_ec, e_ec, alpha=0.0, beta=0.0, out=e_ec,
                             mask=viscid.make_spherical_mask(e_ec, rmax=r_mask))

        tmp = viscid.empty([np.linspace(-10, 10, 64), np.linspace(-10, 10, 64),
                            np.linspace(-10, 10, 64)], center="Cell")

        b_cc_interp = viscid.interp_linear(b_cc, tmp)
        b_fc_interp = viscid.interp_linear(b_fc, tmp)
        e_cc_interp = viscid.interp_linear(e_cc, tmp)
        e_ec_interp = viscid.interp_linear(e_ec, tmp)

        epar_cc = viscid.dot(e_cc_interp, b_cc_interp) / viscid.magnitude(b_cc_interp)
        epar_ecfc = viscid.dot(e_ec_interp, b_fc_interp) / viscid.magnitude(b_fc_interp)

        if make_plots:
            # plt.figure()
            # ax0 = plt.subplot(121)
            # vlt.plot(b_cc['x']['y=0j'], clim=(-40, 40))
            # plt.subplot(122, sharex=ax0, sharey=ax0)
            # vlt.plot(b_fc['x']['y=0j'], clim=(-40, 40))
            # vlt.show()

            plt.figure(figsize=(14, 5))
            ax0 = plt.subplot(131)
            vlt.plot(epar_cc['y=0j'], symmetric=True, cbarlabel="Epar CC")
            plt.subplot(132, sharex=ax0, sharey=ax0)
            vlt.plot(epar_ecfc['y=0j'], symmetric=True, cbarlabel="Epar ECFC")
            plt.subplot(133, sharex=ax0, sharey=ax0)
            vlt.plot(((epar_cc - epar_ecfc) / epar_cc)['y=0j'], clim=(-10, 10),
                     cbarlabel="Rel Diff")
            vlt.auto_adjust_subplots()
            vlt.show()

    return 0
Exemplo n.º 53
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prof", action="store_true")
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    b = viscid.make_dipole(l=(-5, -5, -5), h=(5, 5, 5), n=(255, 255, 127),
                           m=(0, 0, -1))
    b2 = np.sum(b * b, axis=b.nr_comp)

    if args.prof:
        print("Without boundaries")
        viscid.timeit(viscid.grad, b2, bnd=False, timeit_repeat=10,
                      timeit_print_stats=True)
        print("With boundaries")
        viscid.timeit(viscid.grad, b2, bnd=True, timeit_repeat=10,
                      timeit_print_stats=True)

    grad_b2 = viscid.grad(b2)
    grad_b2.pretty_name = r"$\nabla$ B$^2$"
    conv = viscid.convective_deriv(b)
    conv.pretty_name = r"(B $\cdot \nabla$) B"

    _ = plt.figure(figsize=(9, 4.2))

    ax1 = vlt.subplot(231)
    vlt.plot(b2['z=0f'], logscale=True)
    vlt.plot(b2['z=0f'], logscale=True, style='contour', levels=10, colors='grey')
    # vlt.plot2d_quiver(viscid.normalize(b['z=0f']), step=16, pivot='mid')
    ax2 = vlt.subplot(234)
    vlt.plot(b2['y=0f'], logscale=True)
    vlt.plot(b2['y=0f'], logscale=True, style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(b['y=0f'], preferred='numpy'),
                      step=16, pivot='mid')

    vlt.subplot(232, sharex=ax1, sharey=ax1)
    vlt.plot(1e-4 + viscid.magnitude(grad_b2['z=0f']), logscale=True)
    vlt.plot(1e-4 + viscid.magnitude(grad_b2['z=0f']), logscale=True,
             style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(grad_b2['z=0f']), step=16, pivot='mid')
    vlt.subplot(235, sharex=ax2, sharey=ax2)
    vlt.plot(1e-4 + viscid.magnitude(grad_b2['y=0f']), logscale=True)
    vlt.plot(1e-4 + viscid.magnitude(grad_b2['y=0f']), logscale=True,
             style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(grad_b2['y=0f']), step=16, pivot='mid')

    vlt.subplot(233, sharex=ax1, sharey=ax1)
    vlt.plot(viscid.magnitude(conv['z=0f']), logscale=True)
    vlt.plot(viscid.magnitude(conv['z=0f']), logscale=True,
             style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(conv['z=0f']), step=16, pivot='mid')
    vlt.subplot(236, sharex=ax2, sharey=ax2)
    vlt.plot(viscid.magnitude(conv['y=0f']), logscale=True)
    vlt.plot(viscid.magnitude(conv['y=0f']), logscale=True,
             style='contour', levels=10, colors='grey')
    vlt.plot2d_quiver(viscid.normalize(conv['y=0f']), step=16, pivot='mid')

    vlt.auto_adjust_subplots()

    plt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()

    return 0