示例#1
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
def do_test(lines, scalars, show=False, txt=""):
    viscid.logger.info('--> ' + txt)
    title = txt + '\n' + "\n".join(
        textwrap.wrap("scalars = {0}".format(scalars), width=50))

    try:
        from viscid.plot import vpyplot as vlt
        from matplotlib import pyplot as plt

        vlt.clf()
        vlt.plot_lines(lines, scalars=scalars)
        plt.title(title)
        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if show:
            vlt.show()
    except ImportError:
        pass

    try:
        from mayavi import mlab
        vlab, _ = get_mvi_fig()

        vlab.clf()
        vlab.plot_lines3d(lines, scalars=scalars)
        vlab.fancy_axes()
        mlab.text(0.05, 0.05, title)
        vlab.savefig(next_plot_fname(__file__, series='q3'))
        if show:
            vlab.show(stop=True)
    except ImportError:
        pass
示例#3
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
def run_test(_fld, _seeds, plot2d=True, plot3d=True, title='', show=False,
             **kwargs):
    lines, topo = viscid.calc_streamlines(_fld, _seeds, **kwargs)
    topo_color = viscid.topology2color(topo)

    # downsample lines for plotting
    lines = [line[:, ::8] for line in lines]

    try:
        if not plot2d:
            raise ImportError
        from viscid.plot import vpyplot as vlt
        from matplotlib import pyplot as plt
        plt.clf()

        vlt.plot2d_lines(lines, scalars=topo_color, symdir='y', marker='^')
        if title:
            plt.title(title)

        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,
                             bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))
            _global_ns['figure'] = fig

        fld_mag = np.log(viscid.magnitude(_fld))
        try:
            # note: mayavi.mlab.mesh can't take color tuples as scalars
            #       so one can't use topo_color on a mesh surface. This
            #       is a limitation of mayavi. To actually plot a specific
            #       set of colors on a mesh, one must use a texture
            mesh = vlab.mesh_from_seeds(_seeds, scalars=topo, opacity=0.6)
            mesh.actor.property.backface_culling = True
        except RuntimeError:
            pass
        vlab.plot_lines(lines, scalars=fld_mag, tube_radius=0.01,
                       cmap='viridis')
        if title:
            vlab.title(title)

        vlab.savefig(next_plot_fname(__file__, series='3d'))
        if show:
            vlab.show()

    except ImportError:
        pass
示例#5
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
示例#6
0
def _main():
    crd_system = 'gse'
    print(
        viscid.get_dipole_moment_ang(dip_tilt=45.0,
                                     dip_gsm=0.0,
                                     crd_system=crd_system))
    print(
        viscid.get_dipole_moment_ang(dip_tilt=0.0,
                                     dip_gsm=45.0,
                                     crd_system=crd_system))
    print(
        viscid.get_dipole_moment_ang(dip_tilt=45.0,
                                     dip_gsm=45.0,
                                     crd_system=crd_system))

    print("---")
    ptsNP = np.array([[+2, -2, +2], [+2, -1, +2], [+2, 1, +2], [+2, 2, +2]]).T
    ptsSP = np.array([[+2, -2, -2], [+2, -1, -2], [+2, 1, -2], [+2, 2, -2]]).T

    ptsNN = np.array([[-2, -2, +2], [-2, -1, +2], [-2, 1, +2], [-2, 2, +2]]).T
    ptsSN = np.array([[-2, -2, -2], [-2, -1, -2], [-2, 1, -2], [-2, 2, -2]]).T

    mapped_ptsNP = dipole_map(ptsNP)
    mapped_ptsNN = dipole_map(ptsNN)
    mapped_ptsSP = dipole_map(ptsSP)
    mapped_ptsSN = dipole_map(ptsSN)

    try:
        from viscid.plot import vlab
        colors1 = np.array([(0.6, 0.2, 0.2), (0.2, 0.2, 0.6), (0.6, 0.6, 0.2),
                            (0.2, 0.6, 0.6)])
        colors2 = colors1 * 0.5

        vlab.points3d(ptsNP, scale_factor=0.4, color=tuple(colors1[0]))
        vlab.points3d(ptsNN, scale_factor=0.4, color=tuple(colors1[1]))
        vlab.points3d(ptsSP, scale_factor=0.4, color=tuple(colors1[2]))
        vlab.points3d(ptsSN, scale_factor=0.4, color=tuple(colors1[3]))

        vlab.points3d(mapped_ptsNP, scale_factor=0.4, color=tuple(colors2[0]))
        vlab.points3d(mapped_ptsNN, scale_factor=0.4, color=tuple(colors2[1]))
        vlab.points3d(mapped_ptsSP, scale_factor=0.4, color=tuple(colors2[2]))
        vlab.points3d(mapped_ptsSN, scale_factor=0.4, color=tuple(colors2[3]))

        b = make_dipole()

        vlab.plot_lines(
            viscid.calc_streamlines(b, mapped_ptsNP, ibound=0.5)[0])
        vlab.plot_lines(
            viscid.calc_streamlines(b, mapped_ptsNN, ibound=0.5)[0])

        vlab.show()

    except ImportError:
        print("Mayavi not installed, no 3D plots", file=sys.stderr)
示例#7
0
def run_test(_fld, _seeds, plot2d=True, plot3d=True, title='', show=False,
             **kwargs):
    lines, topo = viscid.calc_streamlines(_fld, _seeds, **kwargs)
    topo_color = viscid.topology2color(topo)

    # downsample lines for plotting
    lines = [line[:, ::8] for line in lines]

    try:
        if not plot2d:
            raise ImportError
        from viscid.plot import vpyplot as vlt
        from matplotlib import pyplot as plt
        plt.clf()

        vlt.plot2d_lines(lines, scalars=topo_color, symdir='y', marker='^')
        if title:
            plt.title(title)

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

        fld_mag = np.log(viscid.magnitude(_fld))
        try:
            # note: mayavi.mlab.mesh can't take color tuples as scalars
            #       so one can't use topo_color on a mesh surface. This
            #       is a limitation of mayavi. To actually plot a specific
            #       set of colors on a mesh, one must use a texture
            mesh = vlab.mesh_from_seeds(_seeds, scalars=topo, opacity=0.6)
            mesh.actor.property.backface_culling = True
        except RuntimeError:
            pass
        vlab.plot_lines(lines, scalars=fld_mag, tube_radius=0.01,
                       cmap='viridis')
        if title:
            vlab.title(title)

        vlab.savefig(next_plot_fname(__file__, series='3d'))
        if show:
            vlab.show()

    except ImportError:
        pass
示例#8
0
def _main():
    crd_system = 'gse'
    print(viscid.get_dipole_moment_ang(dip_tilt=45.0, dip_gsm=0.0,
                                       crd_system=crd_system))
    print(viscid.get_dipole_moment_ang(dip_tilt=0.0, dip_gsm=45.0,
                                       crd_system=crd_system))
    print(viscid.get_dipole_moment_ang(dip_tilt=45.0, dip_gsm=45.0,
                                       crd_system=crd_system))

    print("---")
    ptsNP = np.array([[+2, -2, +2], [+2, -1, +2], [+2, 1, +2], [+2, 2, +2]]).T
    ptsSP = np.array([[+2, -2, -2], [+2, -1, -2], [+2, 1, -2], [+2, 2, -2]]).T

    ptsNN = np.array([[-2, -2, +2], [-2, -1, +2], [-2, 1, +2], [-2, 2, +2]]).T
    ptsSN = np.array([[-2, -2, -2], [-2, -1, -2], [-2, 1, -2], [-2, 2, -2]]).T

    mapped_ptsNP = dipole_map(ptsNP)
    mapped_ptsNN = dipole_map(ptsNN)
    mapped_ptsSP = dipole_map(ptsSP)
    mapped_ptsSN = dipole_map(ptsSN)

    try:
        from viscid.plot import vlab
        colors1 = np.array([(0.6, 0.2, 0.2),
                            (0.2, 0.2, 0.6),
                            (0.6, 0.6, 0.2),
                            (0.2, 0.6, 0.6)])
        colors2 = colors1 * 0.5

        vlab.points3d(ptsNP, scale_factor=0.4, color=tuple(colors1[0]))
        vlab.points3d(ptsNN, scale_factor=0.4, color=tuple(colors1[1]))
        vlab.points3d(ptsSP, scale_factor=0.4, color=tuple(colors1[2]))
        vlab.points3d(ptsSN, scale_factor=0.4, color=tuple(colors1[3]))

        vlab.points3d(mapped_ptsNP, scale_factor=0.4, color=tuple(colors2[0]))
        vlab.points3d(mapped_ptsNN, scale_factor=0.4, color=tuple(colors2[1]))
        vlab.points3d(mapped_ptsSP, scale_factor=0.4, color=tuple(colors2[2]))
        vlab.points3d(mapped_ptsSN, scale_factor=0.4, color=tuple(colors2[3]))

        b = make_dipole()

        vlab.plot_lines(viscid.calc_streamlines(b, mapped_ptsNP, ibound=0.5)[0])
        vlab.plot_lines(viscid.calc_streamlines(b, mapped_ptsNN, ibound=0.5)[0])

        vlab.show()

    except ImportError:
        print("Mayavi not installed, no 3D plots", file=sys.stderr)
示例#9
0
def _main():
    try:
        # raise ImportError
        from viscid.plot import vlab
        _HAS_MVI = True
    except ImportError:
        _HAS_MVI = False

    def _test(_p1, _p2, r1=None, r2=None, color=(0.8, 0.8, 0.8)):
        if r1 is not None:
            _p1 = r1 * np.asarray(_p1) / np.linalg.norm(_p1)
        if r2 is not None:
            _p2 = r2 * np.asarray(_p2) / np.linalg.norm(_p2)
        circ = great_circle(_p1, _p2)
        if not np.all(np.isclose(circ[:, 0], _p1)):
            print("!! great circle error P1:", _p1, ", P2:", _p2)
            print("             first_point:", circ[:, 0], "!= P1")
        if not np.all(np.isclose(circ[:, -1], _p2)):
            print("!! great circle error P1:", _p1, ", P2:", _p2)
            print("              last_point:", circ[:, -1], "!= P2")

        if _HAS_MVI:
            vlab.plot_lines([circ], tube_radius=0.02, color=color)

    print("TEST 1")
    _test([1, 0, 0], [0, 1, 0], r1=1.0, r2=1.0, color=(0.8, 0.8, 0.2))
    print("TEST 2")
    _test([1, 0, 0], [-1, 0, 0], r1=1.0, r2=1.0, color=(0.2, 0.8, 0.8))
    print("TEST 3")
    _test([1, 1, 0.01], [-1, -1, 0.01], r1=1.0, r2=1.5, color=(0.8, 0.2, 0.8))

    print("TEST 4")
    _test([-0.9947146, 1.3571029, 2.6095123],
          [-0.3371437, -1.5566425, 2.6634643],
          color=(0.8, 0.2, 0.2))
    print("TEST 5")
    _test([0.9775307, -1.3741084, 2.6030273],
          [0.3273931, 1.5570284, 2.6652965],
          color=(0.2, 0.2, 0.8))

    if _HAS_MVI:
        vlab.plot_blue_marble(r=1.0, lines=False, ntheta=64, nphi=128)
        vlab.plot_earth_3d(radius=1.01, night_only=True, opacity=0.5)
        vlab.show()

    return 0
示例#10
0
def _main():
    try:
        # raise ImportError
        from viscid.plot import vlab
        _HAS_MVI = True
    except ImportError:
        _HAS_MVI = False

    def _test(_p1, _p2, r1=None, r2=None, color=(0.8, 0.8, 0.8)):
        if r1 is not None:
            _p1 = r1 * np.asarray(_p1) / np.linalg.norm(_p1)
        if r2 is not None:
            _p2 = r2 * np.asarray(_p2) / np.linalg.norm(_p2)
        circ = great_circle(_p1, _p2)
        if not np.all(np.isclose(circ[:, 0], _p1)):
            print("!! great circle error P1:", _p1, ", P2:", _p2)
            print("             first_point:", circ[:, 0], "!= P1")
        if not np.all(np.isclose(circ[:, -1], _p2)):
            print("!! great circle error P1:", _p1, ", P2:", _p2)
            print("              last_point:", circ[:, -1], "!= P2")

        if _HAS_MVI:
            vlab.plot_lines([circ], tube_radius=0.02, color=color)

    print("TEST 1")
    _test([1, 0, 0], [0, 1, 0], r1=1.0, r2=1.0, color=(0.8, 0.8, 0.2))
    print("TEST 2")
    _test([1, 0, 0], [-1, 0, 0], r1=1.0, r2=1.0, color=(0.2, 0.8, 0.8))
    print("TEST 3")
    _test([1, 1, 0.01], [-1, -1, 0.01], r1=1.0, r2=1.5, color=(0.8, 0.2, 0.8))

    print("TEST 4")
    _test([-0.9947146, 1.3571029, 2.6095123], [-0.3371437, -1.5566425, 2.6634643],
          color=(0.8, 0.2, 0.2))
    print("TEST 5")
    _test([0.9775307, -1.3741084, 2.6030273], [0.3273931, 1.5570284, 2.6652965],
          color=(0.2, 0.2, 0.8))

    if _HAS_MVI:
        vlab.plot_blue_marble(r=1.0, lines=False, ntheta=64, nphi=128)
        vlab.plot_earth_3d(radius=1.01, night_only=True, opacity=0.5)
        vlab.show()

    return 0
示例#11
0
def do_test(lines, scalars, show=False, txt=""):
    viscid.logger.info('--> ' + txt)
    title = txt + '\n' + "\n".join(textwrap.wrap("scalars = {0}".format(scalars),
                                                 width=50))

    try:
        from matplotlib import pyplot as plt
        from viscid.plot import vpyplot as vlt

        vlt.clf()
        vlt.plot_lines(lines, scalars=scalars)
        plt.title(title)
        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if show:
            vlt.show()
    except ImportError:
        pass

    try:
        from mayavi import mlab
        from viscid.plot import vlab

        try:
            fig = _global_ns['figure']
            vlab.clf()
        except KeyError:
            fig = vlab.figure(size=[1200, 800], offscreen=not show,
                              bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))
            _global_ns['figure'] = fig

        vlab.clf()
        vlab.plot_lines3d(lines, scalars=scalars)
        vlab.fancy_axes()
        mlab.text(0.05, 0.05, title)
        vlab.savefig(next_plot_fname(__file__, series='q3'))
        if show:
            vlab.show(stop=True)
    except ImportError:
        pass
def do_test(lines, scalars, show=False, txt=""):
    viscid.logger.info('--> ' + txt)
    title = txt + '\n' + "\n".join(textwrap.wrap("scalars = {0}".format(scalars),
                                                 width=50))

    try:
        from viscid.plot import vpyplot as vlt
        from matplotlib import pyplot as plt

        vlt.clf()
        vlt.plot_lines(lines, scalars=scalars)
        plt.title(title)
        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if show:
            vlt.show()
    except ImportError:
        pass

    try:
        from mayavi import mlab
        from viscid.plot import vlab

        try:
            fig = _global_ns['figure']
            vlab.clf()
        except KeyError:
            fig = vlab.figure(size=[1200, 800], offscreen=not show,
                              bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))
            _global_ns['figure'] = fig

        vlab.clf()
        vlab.plot_lines3d(lines, scalars=scalars)
        vlab.fancy_axes()
        mlab.text(0.05, 0.05, title)
        vlab.savefig(next_plot_fname(__file__, series='q3'))
        if show:
            vlab.show(stop=True)
    except ImportError:
        pass
示例#13
0
def _get_sep_pts_bisect(fld, seed, trace_opts=None, min_depth=3, max_depth=7,
                        plot=False, perimeter_check=perimeter_check_bitwise_or,
                        make_3d=True, start_uneven=False, _base_quadrent="",
                        _uneven_mask=0, _first_recurse=True):
    if len(_base_quadrent) == max_depth:
        return [_base_quadrent]  # causes pylint to complain
    if trace_opts is None:
        trace_opts = dict()

    nx, ny = seed.uv_shape
    (xlim, ylim) = seed.uv_extent

    if _first_recurse and start_uneven:
        _uneven_mask = UNEVEN_MASK

    if _first_recurse and plot:
        from viscid.plot import vlab
        from viscid.plot import vpyplot as vlt
        vlt.clf()
        _, all_topo = viscid.calc_streamlines(fld, seed, **trace_opts)
        vlt.plot(np.bitwise_and(all_topo, 15), show=False)
        verts, arr = seed.wrap_mesh(all_topo.data)
        vlab.mesh(verts[0], verts[1], verts[2], scalars=arr, opacity=0.75)

    # quadrents and lines are indexed as follows...
    # directions are counter clackwise around the quadrent with
    # lower index (which matters for lines which are shared among
    # more than one quadrent, aka, lines 1,2,6,7). Notice that even
    # numbered lines are horizontal, like the interstate system :)
    # -<--10-----<-8---
    # |       ^       ^
    # 11  2   9   3   7
    # \/      |       |
    # --<-2-----<-6----
    # |       ^       ^
    # 3   0   1   1   5
    # \/      |       |
    # ----0->-----4->--

    # find low(left), mid(center), and high(right) crds in x and y
    low_quad = "{0}{1:x}".format(_base_quadrent, 0 | _uneven_mask)
    high_quad = "{0}{1:x}".format(_base_quadrent, 3 | _uneven_mask)
    xl, xm, yl, ym = _quadrent_limits(low_quad, xlim, ylim)
    _, xh, _, yh = _quadrent_limits(high_quad, xlim, ylim)
    segsx, segsy = [None] * 12, [None] * 12
    topo = [None] * 12
    nxm, nym = nx //2, ny // 2

    # make all the line segments
    segsx[0], segsy[0] = np.linspace(xl, xm, nxm), np.linspace(yl, yl, nxm)
    segsx[1], segsy[1] = np.linspace(xm, xm, nym), np.linspace(yl, ym, nym)
    segsx[2], segsy[2] = np.linspace(xm, xl, nxm), np.linspace(ym, ym, nxm)
    segsx[3], segsy[3] = np.linspace(xl, xl, nym), np.linspace(ym, yl, nym)

    segsx[4], segsy[4] = np.linspace(xm, xh, nxm), np.linspace(yl, yl, nxm)
    segsx[5], segsy[5] = np.linspace(xh, xh, nym), np.linspace(yl, ym, nym)
    segsx[6], segsy[6] = np.linspace(xh, xm, nxm), np.linspace(ym, ym, nxm)

    segsx[7], segsy[7] = np.linspace(xh, xh, nym), np.linspace(ym, yh, nym)
    segsx[8], segsy[8] = np.linspace(xh, xm, nxm), np.linspace(yh, yh, nxm)
    segsx[9], segsy[9] = np.linspace(xm, xm, nym), np.linspace(ym, yh, nym)

    segsx[10], segsy[10] = np.linspace(xm, xl, nxm), np.linspace(yh, yh, nxm)
    segsx[11], segsy[11] = np.linspace(xl, xl, nym), np.linspace(yh, ym, nym)

    allx = np.concatenate(segsx)
    ally = np.concatenate(segsy)

    # print("plot::", _base_quadrent, '|', _uneven_mask, '|', len(allx), len(ally))

    pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
    _, all_topo = viscid.calc_streamlines(fld, pts3d, **trace_opts)

    topo[0] = all_topo[:len(segsx[0])]
    cnt = len(topo[0])
    for i, segx in zip(count(1), segsx[1:]):
        topo[i] = all_topo[cnt:cnt + len(segx)]
        # print("??", i, cnt, cnt + len(segx), np.bitwise_and.reduce(topo[i]))
        cnt += len(topo[i])

    # assemble the lines into the four quadrents
    quad_topo = [None] * 4

    # all arrays snip off the last element since those are
    # duplicated by the next line... reversed arrays do the
    # snipping with -1:0:-1
    quad_topo[0] = np.concatenate([topo[0][:-1], topo[1][:-1],
                                   topo[2][:-1], topo[3][:-1]])

    quad_topo[1] = np.concatenate([topo[4][:-1], topo[5][:-1],
                                   topo[6][:-1], topo[1][-1:0:-1]])

    quad_topo[2] = np.concatenate([topo[2][-1:0:-1], topo[9][:-1],
                                   topo[10][:-1], topo[11][:-1]])

    quad_topo[3] = np.concatenate([topo[6][-1:0:-1], topo[7][:-1],
                                   topo[8][:-1], topo[9][-1:0:-1]])

    # now that the quad arrays are populated, decide which quadrents
    # still contain the separator (could be > 1)
    required_uneven_subquads = False
    ret = []
    for i in range(4):
        if perimeter_check(quad_topo[i]):
            next_quad = "{0}{1:x}".format(_base_quadrent, i | _uneven_mask)
            subquads = _get_sep_pts_bisect(fld, seed, trace_opts=trace_opts,
                                           min_depth=min_depth,
                                           max_depth=max_depth, plot=plot,
                                           _base_quadrent=next_quad,
                                           _uneven_mask=0,
                                           _first_recurse=False)
            ret += subquads

    if len(ret) == 0:
        perimeter = np.concatenate([topo[0][::-1], topo[4][::-1],
                                    topo[5][::-1], topo[7][::-1],
                                    topo[8][::-1], topo[10][::-1],
                                    topo[11][::-1], topo[3][::-1]])
        if _uneven_mask:
            if len(_base_quadrent) > min_depth:
                print("sep trace issue, but min depth reached: {0} > {1}"
                      "".format(len(_base_quadrent), min_depth))
                ret = [_base_quadrent]
            else:
                print("sep trace issue, the separator ended prematurely")
        elif perimeter_check(perimeter):
            ret = _get_sep_pts_bisect(fld, seed, trace_opts=trace_opts,
                                      min_depth=min_depth, max_depth=max_depth,
                                      plot=plot, _base_quadrent=_base_quadrent,
                                      _uneven_mask=UNEVEN_MASK,
                                      _first_recurse=False)
            required_uneven_subquads = True

    if plot and not required_uneven_subquads:
        from viscid.plot import vlab
        from viscid.plot import vpyplot as vlt
        from matplotlib import pyplot as plt
        _pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
        vlab.points3d(_pts3d[0], _pts3d[1], _pts3d[2],
                      all_topo.data.reshape(-1), scale_mode='none',
                      scale_factor=0.02)
        plt.scatter(allx, ally, color=np.bitwise_and(all_topo, 15),
                        vmin=0, vmax=15, marker='o', edgecolor='y', s=40)

    if _first_recurse:
        # turn quadrent strings into locations
        xc = np.empty(len(ret))
        yc = np.empty(len(ret))
        for i, r in enumerate(ret):
            xc[i], yc[i] = _quadrent_center(r, xlim, ylim)
        pts_uv = np.array([xc, yc])
        if plot:
            from viscid.plot import vlab
            from viscid.plot import vpyplot as vlt
            from matplotlib import pyplot as plt
            plt.plot(pts_uv[0], pts_uv[1], "y*", ms=20,
                         markeredgecolor='k', markeredgewidth=1.0)
            vlt.show(block=False)
            vlab.show(stop=True)
        # return seed.to_3d(seed.uv_to_local(pts_uv))
        # if pts_uv.size == 0:
        #     return None
        if make_3d:
            return seed.uv_to_3d(pts_uv)
        else:
            return pts_uv
    else:
        return ret
示例#14
0
def _main():
    global offscreen_vlab

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--notwo", dest='notwo', action="store_true")
    parser.add_argument("--nothree", dest='nothree', action="store_true")
    parser.add_argument("--show", "--plot", action="store_true")
    args = viscid.vutil.common_argparse(parser, default_verb=0)

    plot2d = not args.notwo
    plot3d = not args.nothree

    # plot2d = True
    # plot3d = True
    # args.show = True

    offscreen_vlab = not args.show

    img = np.load(os.path.join(sample_dir, "logo.npy"))
    x = np.linspace(-1, 1, img.shape[0])
    y = np.linspace(-1, 1, img.shape[1])
    z = np.linspace(-1, 1, img.shape[2])
    logo = viscid.arrays2field([x, y, z], img)

    if 1:
        viscid.logger.info('Testing Point with custom local coordinates...')
        pts = np.vstack([[-1, -0.5, 0, 0.5, 1],
                         [-1, -0.5, 0, 0.5, 1],
                         [ 0,  0.5, 1, 1.5, 2]])
        local_crds = viscid.asarray_datetime64([0, 60, 120, 180, 240],
                                               conservative=True)
        seeds = viscid.Point(pts, local_crds=local_crds)
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)

    if 1:
        viscid.logger.info('Testing Line...')
        seeds = viscid.Line([-1, -1, 0], [1, 1, 2], n=5)
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)

    if 1:
        viscid.logger.info('Testing Plane...')
        seeds = viscid.Plane([0.0, 0.0, 0.0], [1, 1, 1], [1, 0, 0], 2, 2,
                             nl=160, nm=170, NL_are_vectors=True)
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)

    if 1:
        viscid.logger.info('Testing Volume...')
        seeds = viscid.Volume([-0.8, -0.8, -0.8], [0.8, 0.8, 0.8],
                              n=[64, 64, 3])
        # note: can't make a 2d plot of the volume w/o a slice
        run_test(logo, seeds, plot2d=False, plot3d=plot3d, add_title="3d",
                 show=args.show)

    if 1:
        viscid.logger.info('Testing Volume (with ignorable dim)...')
        seeds = viscid.Volume([-0.8, -0.8, 0.0], [0.8, 0.8, 0.0],
                              n=[64, 64, 1])
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="2d",
                 show=args.show)

    if 1:
        viscid.logger.info('Testing Spherical Sphere (phi, theta)...')
        seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170,
                              pole=[-1, -1, -1], theta_phi=False)
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT",
                 show=args.show)

    if 1:
        viscid.logger.info('Testing Spherical Sphere (theta, phi)...')
        seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170,
                              pole=[-1, -1, -1], theta_phi=True)
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP",
                 show=args.show)

    if 1:
        viscid.logger.info('Testing Spherical Cap (phi, theta)...')
        seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80,
                                    pole=[-1, -1, -1], theta_phi=False)
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT",
                 view_kwargs=dict(azimuth=180, elevation=180), show=args.show)

    if 1:
        viscid.logger.info('Testing Spherical Cap (theta, phi)...')
        seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80,
                                    pole=[-1, -1, -1], theta_phi=True)
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP",
                 view_kwargs=dict(azimuth=180, elevation=180), show=args.show)

    if 1:
        viscid.logger.info('Testing Spherical Patch...')
        seeds = viscid.SphericalPatch(p0=[0, 0, 0], p1=[0, -0, -1],
                                      max_alpha=30.0, max_beta=59.9,
                                      nalpha=65, nbeta=80, r=0.5, roll=45.0)
        run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show)

    if 1:
        # this spline test is very custom
        viscid.logger.info('Testing Spline...')
        try:
            import scipy.interpolate as interpolate
        except ImportError:
            msg = "XFail: ImportError (is scipy installed?)"
            if plot2d:
                try:
                    from viscid.plot import vpyplot as vlt
                    from matplotlib import pyplot as plt
                    plt.clf()
                    plt.annotate(msg, xy=(0.3, 0.4), xycoords='axes fraction')
                    plt.savefig(next_plot_fname(__file__, series='2d'))
                    plt.savefig(next_plot_fname(__file__, series='2d'))
                    plt.savefig(next_plot_fname(__file__, series='3d'))
                    if args.show:
                        plt.show()
                except ImportError:
                    pass
        else:
            knots = np.array([[ 0.2,  0.5, 0.0], [-0.2,  0.5, 0.2],
                              [-0.2,  0.0, 0.4], [ 0.2,  0.0, 0.2],
                              [ 0.2, -0.5, 0.0], [-0.2, -0.5, 0.2]]).T
            seed_name = "Spline"
            fld = logo
            seeds = viscid.Spline(knots)
            seed_pts = seeds.get_points()
            interp_fld = viscid.interp_trilin(fld, seeds)

            if plot2d:
                try:
                    from viscid.plot import vpyplot as vlt
                    from matplotlib import pyplot as plt
                    plt.clf()
                    vlt.plot(interp_fld)
                    plt.title(seed_name)
                    plt.savefig(next_plot_fname(__file__, series='2d'))
                    if args.show:
                        plt.show()

                    plt.clf()
                    from matplotlib import rcParams
                    _ms = rcParams['lines.markersize']
                    plt.gca().scatter(knots[0, :], knots[1, :],
                                      s=(2 * _ms)**2, marker='^', color='y')
                    plt.gca().scatter(seed_pts[0, :], seed_pts[1, :],
                                      s=(1.5 * _ms)**2, marker='o', color='k')
                    vlt.plot2d_line(seed_pts, scalars=interp_fld.flat_data,
                                    symdir='z')
                    plt.title(seed_name)
                    plt.savefig(next_plot_fname(__file__, series='2d'))
                    if args.show:
                        plt.show()
                except ImportError:
                    pass
            if plot3d:
                try:
                    vlab, _ = get_mvi_fig()
                    vlab.points3d(knots[0], knots[1], knots[2],
                                  color=(1.0, 1.0, 0), scale_mode='none',
                                  scale_factor=0.04)
                    p = vlab.points3d(seed_pts[0], seed_pts[1], seed_pts[2],
                                      color=(0, 0, 0), scale_mode='none',
                                      scale_factor=0.03)
                    vlab.plot_line(seed_pts, scalars=interp_fld.flat_data,
                                   tube_radius=0.01)
                    vlab.axes(p)
                    vlab.title(seed_name)
                    vlab.mlab.roll(-90.0)
                    vlab.savefig(next_plot_fname(__file__, series='3d'))
                    if args.show:
                        vlab.show(stop=True)
                except ImportError:
                    pass

    if 1:
        viscid.logger.info('Testing RectilinearMeshPoints...')
        f = viscid.load_file(os.path.join(sample_dir, 'sample_xdmf.3d.[-1].xdmf'))
        slc = 'x=-40j:12j, y=-10j:10j, z=-10j:10j'
        b = f['b'][slc]
        z = b.get_crd('z')
        sheet_iz = np.argmin(b['x']**2, axis=2)
        sheet_pts = b['z=0:1'].get_points()
        sheet_pts[2, :] = z[sheet_iz].reshape(-1)
        isphere_mask = np.sum(sheet_pts[:2, :]**2, axis=0) < 5**2
        day_mask = sheet_pts[0:1, :] > -1.0
        sheet_pts[2, :] = np.choose(isphere_mask, [sheet_pts[2, :], 0])
        sheet_pts[2, :] = np.choose(day_mask, [sheet_pts[2, :], 0])
        nx, ny, _ = b.sshape
        sheet_seed = viscid.RectilinearMeshPoints(sheet_pts.reshape(3, nx, ny))
        vx_sheet = viscid.interp_nearest(f['vx'], sheet_seed)

        try:
            if not plot2d:
                raise ImportError
            from viscid.plot import vpyplot as vlt
            from matplotlib import pyplot as plt
            vlt.clf()
            vlt.plot(vx_sheet, symmetric=True)
            plt.savefig(next_plot_fname(__file__, series='2d'))
            if args.show:
                vlt.show()
        except ImportError:
            pass

        try:
            if not plot3d:
                raise ImportError
            vlab, _ = get_mvi_fig()
            mesh = vlab.mesh_from_seeds(sheet_seed, scalars=vx_sheet,
                                        clim=(-400, 400))
            vlab.plot_earth_3d(crd_system=b)
            vlab.view(azimuth=+90.0 + 45.0, elevation=90.0 - 25.0,
                      distance=30.0, focalpoint=(-10.0, +1.0, +1.0))

            vlab.title("RectilinearMeshPoints")
            vlab.savefig(next_plot_fname(__file__, series='3d'))
            if args.show:
                vlab.show(stop=True)

        except ImportError:
            pass

    # prevent weird xorg bad-instructions on tear down
    if 'figure' in _global_ns and _global_ns['figure'] is not None:
        from viscid.plot import vlab
        vlab.mlab.close(_global_ns['figure'])

    return 0
示例#15
0
def _get_sep_pts_bisect(fld,
                        seed,
                        trace_opts=None,
                        min_depth=3,
                        max_depth=7,
                        plot=False,
                        perimeter_check=perimeter_check_bitwise_or,
                        make_3d=True,
                        start_uneven=False,
                        _base_quadrent="",
                        _uneven_mask=0,
                        _first_recurse=True):
    if len(_base_quadrent) == max_depth:
        return [_base_quadrent]  # causes pylint to complain
    if trace_opts is None:
        trace_opts = dict()

    nx, ny = seed.uv_shape
    (xlim, ylim) = seed.uv_extent

    if _first_recurse and start_uneven:
        _uneven_mask = UNEVEN_MASK

    if _first_recurse and plot:
        from viscid.plot import vlab
        from viscid.plot import vpyplot as vlt
        vlt.clf()
        _, all_topo = viscid.calc_streamlines(fld, seed, **trace_opts)
        vlt.plot(np.bitwise_and(all_topo, 15), show=False)
        verts, arr = seed.wrap_mesh(all_topo.data)
        vlab.mesh(verts[0], verts[1], verts[2], scalars=arr, opacity=0.75)

    # quadrents and lines are indexed as follows...
    # directions are counter clackwise around the quadrent with
    # lower index (which matters for lines which are shared among
    # more than one quadrent, aka, lines 1,2,6,7). Notice that even
    # numbered lines are horizontal, like the interstate system :)
    # -<--10-----<-8---
    # |       ^       ^
    # 11  2   9   3   7
    # \/      |       |
    # --<-2-----<-6----
    # |       ^       ^
    # 3   0   1   1   5
    # \/      |       |
    # ----0->-----4->--

    # find low(left), mid(center), and high(right) crds in x and y
    low_quad = "{0}{1:x}".format(_base_quadrent, 0 | _uneven_mask)
    high_quad = "{0}{1:x}".format(_base_quadrent, 3 | _uneven_mask)
    xl, xm, yl, ym = _quadrent_limits(low_quad, xlim, ylim)
    _, xh, _, yh = _quadrent_limits(high_quad, xlim, ylim)
    segsx, segsy = [None] * 12, [None] * 12
    topo = [None] * 12
    nxm, nym = nx // 2, ny // 2

    # make all the line segments
    segsx[0], segsy[0] = np.linspace(xl, xm, nxm), np.linspace(yl, yl, nxm)
    segsx[1], segsy[1] = np.linspace(xm, xm, nym), np.linspace(yl, ym, nym)
    segsx[2], segsy[2] = np.linspace(xm, xl, nxm), np.linspace(ym, ym, nxm)
    segsx[3], segsy[3] = np.linspace(xl, xl, nym), np.linspace(ym, yl, nym)

    segsx[4], segsy[4] = np.linspace(xm, xh, nxm), np.linspace(yl, yl, nxm)
    segsx[5], segsy[5] = np.linspace(xh, xh, nym), np.linspace(yl, ym, nym)
    segsx[6], segsy[6] = np.linspace(xh, xm, nxm), np.linspace(ym, ym, nxm)

    segsx[7], segsy[7] = np.linspace(xh, xh, nym), np.linspace(ym, yh, nym)
    segsx[8], segsy[8] = np.linspace(xh, xm, nxm), np.linspace(yh, yh, nxm)
    segsx[9], segsy[9] = np.linspace(xm, xm, nym), np.linspace(ym, yh, nym)

    segsx[10], segsy[10] = np.linspace(xm, xl, nxm), np.linspace(yh, yh, nxm)
    segsx[11], segsy[11] = np.linspace(xl, xl, nym), np.linspace(yh, ym, nym)

    allx = np.concatenate(segsx)
    ally = np.concatenate(segsy)

    # print("plot::", _base_quadrent, '|', _uneven_mask, '|', len(allx), len(ally))

    pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
    _, all_topo = viscid.calc_streamlines(fld, pts3d, **trace_opts)

    topo[0] = all_topo[:len(segsx[0])]
    cnt = len(topo[0])
    for i, segx in zip(count(1), segsx[1:]):
        topo[i] = all_topo[cnt:cnt + len(segx)]
        # print("??", i, cnt, cnt + len(segx), np.bitwise_and.reduce(topo[i]))
        cnt += len(topo[i])

    # assemble the lines into the four quadrents
    quad_topo = [None] * 4

    # all arrays snip off the last element since those are
    # duplicated by the next line... reversed arrays do the
    # snipping with -1:0:-1
    quad_topo[0] = np.concatenate(
        [topo[0][:-1], topo[1][:-1], topo[2][:-1], topo[3][:-1]])

    quad_topo[1] = np.concatenate(
        [topo[4][:-1], topo[5][:-1], topo[6][:-1], topo[1][-1:0:-1]])

    quad_topo[2] = np.concatenate(
        [topo[2][-1:0:-1], topo[9][:-1], topo[10][:-1], topo[11][:-1]])

    quad_topo[3] = np.concatenate(
        [topo[6][-1:0:-1], topo[7][:-1], topo[8][:-1], topo[9][-1:0:-1]])

    # now that the quad arrays are populated, decide which quadrents
    # still contain the separator (could be > 1)
    required_uneven_subquads = False
    ret = []
    for i in range(4):
        if perimeter_check(quad_topo[i]):
            next_quad = "{0}{1:x}".format(_base_quadrent, i | _uneven_mask)
            subquads = _get_sep_pts_bisect(fld,
                                           seed,
                                           trace_opts=trace_opts,
                                           min_depth=min_depth,
                                           max_depth=max_depth,
                                           plot=plot,
                                           _base_quadrent=next_quad,
                                           _uneven_mask=0,
                                           _first_recurse=False)
            ret += subquads

    if len(ret) == 0:
        perimeter = np.concatenate([
            topo[0][::-1], topo[4][::-1], topo[5][::-1], topo[7][::-1],
            topo[8][::-1], topo[10][::-1], topo[11][::-1], topo[3][::-1]
        ])
        if _uneven_mask:
            if len(_base_quadrent) > min_depth:
                print("sep trace issue, but min depth reached: {0} > {1}"
                      "".format(len(_base_quadrent), min_depth))
                ret = [_base_quadrent]
            else:
                print("sep trace issue, the separator ended prematurely")
        elif perimeter_check(perimeter):
            ret = _get_sep_pts_bisect(fld,
                                      seed,
                                      trace_opts=trace_opts,
                                      min_depth=min_depth,
                                      max_depth=max_depth,
                                      plot=plot,
                                      _base_quadrent=_base_quadrent,
                                      _uneven_mask=UNEVEN_MASK,
                                      _first_recurse=False)
            required_uneven_subquads = True

    if plot and not required_uneven_subquads:
        from viscid.plot import vlab
        from matplotlib import pyplot as plt
        from viscid.plot import vpyplot as vlt
        _pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
        vlab.points3d(_pts3d[0],
                      _pts3d[1],
                      _pts3d[2],
                      all_topo.data.reshape(-1),
                      scale_mode='none',
                      scale_factor=0.02)
        plt.scatter(allx,
                    ally,
                    color=np.bitwise_and(all_topo, 15),
                    vmin=0,
                    vmax=15,
                    marker='o',
                    edgecolor='y',
                    s=40)

    if _first_recurse:
        # turn quadrent strings into locations
        xc = np.empty(len(ret))
        yc = np.empty(len(ret))
        for i, r in enumerate(ret):
            xc[i], yc[i] = _quadrent_center(r, xlim, ylim)
        pts_uv = np.array([xc, yc])
        if plot:
            from viscid.plot import vlab
            from matplotlib import pyplot as plt
            from viscid.plot import vpyplot as vlt
            plt.plot(pts_uv[0],
                     pts_uv[1],
                     "y*",
                     ms=20,
                     markeredgecolor='k',
                     markeredgewidth=1.0)
            vlt.show(block=False)
            vlab.show(stop=True)
        # return seed.to_3d(seed.uv_to_local(pts_uv))
        # if pts_uv.size == 0:
        #     return None
        if make_3d:
            return seed.uv_to_3d(pts_uv)
        else:
            return pts_uv
    else:
        return ret
示例#16
0
def _main():
    f = viscid.load_file("$WORK/xi_fte_001/*.3d.[4050f].xdmf")
    mp = get_mp_info(f['pp'], f['b'], f['j'], f['e_cc'], fit='mp_xloc',
                     slc="x=6.5j:10.5j, y=-4j:4j, z=-4.8j:3j", cache=False)

    y, z = mp['pp_max_xloc'].meshgrid_flat(prune=True)
    x = mp['pp_max_xloc'].data.reshape(-1)

    Y, Z = mp['pp_max_xloc'].meshgrid(prune=True)
    x2 = paraboloid(Y, Z, *mp['paraboloid'][0])

    skip = 117
    n = paraboloid_normal(Y, Z, *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]

    minvar_y = Y.reshape(-1)[::skip]
    minvar_z = Z.reshape(-1)[::skip]
    minvar_n = np.zeros([3, len(minvar_y)])
    for i in range(minvar_n.shape[0]):
        p0 = [0.0, minvar_y[i], minvar_z[i]]
        p0[0] = mp['pp_max_xloc']['y={0[0]}f, z={0[1]}f'.format(p0)]
        minvar_n[:, i] = viscid.find_minvar_lmn_around(f['b'], p0, l=2.0, n=64)[2, :]

    # 2d plots, normals don't look normal in the matplotlib projection
    if False:  # pylint: disable=using-constant-test
        from viscid.plot import vpyplot as vlt
        from matplotlib import pyplot as plt

        normals = paraboloid_normal(Y, Z, *mp['paraboloid'][0])
        p0 = np.array([x2, Y, Z]).reshape(3, -1)
        p1 = p0 + normals.reshape(3, -1)

        vlt.scatter_3d(np.vstack([x, y, z])[:, ::skip], equal=True)
        for i in range(0, p0.shape[1], skip):
            plt.gca().plot([p0[0, i], p1[0, i]],
                               [p0[1, i], p1[1, i]],
                               [p0[2, i], p1[2, i]], color='c')
        # z2 = _ellipsiod(X, Y, *popt)
        plt.gca().plot_surface(Y, Z, x2, color='r')
        vlt.show()

    # mayavi 3d plots, normals look better here
    if True:  # pylint: disable=using-constant-test
        from viscid.plot import vlab
        vlab.points3d(x[::skip], y[::skip], z[::skip], scale_factor=0.25,
                      color=(0.0, 0.0, 1.0))

        mp_width = mp['mp_width']['x=0']
        mp_sheath_edge = mp['mp_sheath_edge']['x=0']
        mp_sphere_edge = mp_sheath_edge - mp_width

        vlab.mesh(x2, Y, Z, scalars=mp_width.data)
        vlab.mesh(mp_sheath_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)
        vlab.mesh(mp_sphere_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)

        n = paraboloid_normal(Y, Z, *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]
        vlab.quiver3d(x2.reshape(-1)[::skip],
                      Y.reshape(-1)[::skip],
                      Z.reshape(-1)[::skip],
                      n[0], n[1], n[2], color=(1, 0, 0))
        vlab.quiver3d(x2.reshape(-1)[::skip],
                      Y.reshape(-1)[::skip],
                      Z.reshape(-1)[::skip],
                      minvar_n[0], minvar_n[1], minvar_n[2], color=(0, 0, 1))
        vlab.show()
示例#17
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = vutil.common_argparse(parser)

    b, e = make_arcade(8.0, N=[64, 64, 64])
    epar = viscid.project(e, b)
    epar.pretty_name = "E parallel"

    ###############
    # Calculate Xi
    seeds = viscid.Volume(xl=[-10, 0.0, -10], xh=[10, 0.0, 10], n=[64, 1, 64])
    b_lines, _ = viscid.calc_streamlines(b, seeds)

    xi_dat = viscid.integrate_along_lines(b_lines, e, reduction='dot')
    xi = seeds.wrap_field(xi_dat, name='xi', pretty_name=r"$\Xi$")

    ################################
    # Make 2D Matplotlib plot of Xi
    vlt.plot(xi,
             x=(-10, 10),
             y=(-10, 10),
             style='contourf',
             levels=256,
             lin=(2e-4, 1.5718))
    vlt.plot(xi,
             x=(-10, 10),
             y=(-10, 10),
             style='contour',
             colors='grey',
             levels=[0.5, 1.0])
    vlt.savefig(next_plot_fname(__file__))
    if args.show:
        vlt.show()

    ############################################################
    # Make 3D mayavi plot of Xi and the 'brightest' field lines
    # as well as some other field lines for context
    try:
        from viscid.plot import vlab
    except ImportError:
        xfail("Mayavi not installed")

    vlab.figure(size=[1200, 800], offscreen=not args.show)

    inds = np.argsort(xi_dat)[-64:]
    inds = np.concatenate([inds, np.arange(len(xi_dat))[::71]])
    s = vlab.plot_lines(b_lines[inds], scalars=epar, cmap='viridis')
    vlab.mesh_from_seeds(seeds, scalars=xi, cmap='inferno')
    vlab.colorbar(s, orientation='horizontal', title=epar.pretty_name)
    # vlab.streamline(b, scalars=e, seedtype='sphere', seed_resolution=4,
    #                 integration_direction='both')

    oa = vlab.orientation_axes()
    oa.marker.set_viewport(0.75, 0.75, 1.0, 1.0)
    vlab.view(roll=0,
              azimuth=90,
              elevation=25,
              distance=30.0,
              focalpoint=[0, 2, 0])

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

    try:
        vlab.mlab.close()
    except AttributeError:
        pass

    return 0
示例#18
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

    # #################################################
    # viscid.logger.info("Testing field lines on 2d field...")
    B = viscid.make_dipole(twod=True)
    line = viscid.seed.Line((0.2, 0.0, 0.0), (1.0, 0.0, 0.0), 10)
    obound0 = np.array([-4, -4, -4], dtype=B.data.dtype)
    obound1 = np.array([4, 4, 4], dtype=B.data.dtype)
    run_test(B, line, plot2d=plot2d, plot3d=plot3d, title='2D', show=args.show,
             ibound=0.07, obound0=obound0, obound1=obound1)

    #################################################
    viscid.logger.info("Testing field lines on 3d field...")
    B = viscid.make_dipole(m=[0.2, 0.3, -0.9])
    sphere = viscid.seed.Sphere((0.0, 0.0, 0.0), 2.0, ntheta=20, nphi=10)
    obound0 = np.array([-4, -4, -4], dtype=B.data.dtype)
    obound1 = np.array([4, 4, 4], dtype=B.data.dtype)
    run_test(B, sphere, plot2d=plot2d, plot3d=plot3d, title='3D', show=args.show,
             ibound=0.12, obound0=obound0, obound1=obound1, method=viscid.RK12)

    # The Remainder of this test makes sure higher order methods are indeed
    # more accurate than lower order methods... this could find a bug in
    # the integrators

    ##################################################
    # test accuracy of streamlines in an ideal dipole
    cotr = viscid.Cotr(dip_tilt=15.0, dip_gsm=21.0)  # pylint: disable=not-callable
    m = cotr.get_dipole_moment(crd_system='gse')
    seeds = viscid.seed.Sphere((0.0, 0.0, 0.0), 2.0, pole=-m, ntheta=25, nphi=25,
                               thetalim=(5, 90), philim=(5, 360), phi_endpoint=False)
    B = viscid.make_dipole(m=m, crd_system='gse', n=(256, 256, 256),
                           l=(-25, -25, -25), h=(25, 25, 25), dtype='f8')

    seeds_xyz = seeds.get_points()
    # seeds_lsp = viscid.xyz2lsrlp(seeds_xyz, cotr=cotr, crd_system=B)[(0, 3), :]
    seeds_lsp = viscid.xyz2lsrlp(seeds_xyz, cotr=cotr, crd_system=B)[(0, 3), :]

    e1_lines, e1_lsps, t_e1 = lines_and_lsps(B, seeds, method='euler1',
                                             ibound=1.0, cotr=cotr)
    rk2_lines, rk2_lsps, t_rk2 = lines_and_lsps(B, seeds, method='rk2',
                                                ibound=1.0, cotr=cotr)
    rk4_lines, rk4_lsps, t_rk4 = lines_and_lsps(B, seeds, method='rk4',
                                                ibound=1.0, cotr=cotr)
    e1a_lines, e1a_lsps, t_e1a = lines_and_lsps(B, seeds, method='euler1a',
                                                ibound=1.0, cotr=cotr)
    rk12_lines, rk12_lsps, t_rk12 = lines_and_lsps(B, seeds, method='rk12',
                                                   ibound=1.0, cotr=cotr)
    rk45_lines, rk45_lsps, t_rk45 = lines_and_lsps(B, seeds, method='rk45',
                                                   ibound=1.0, cotr=cotr)

    def _calc_rel_diff(_lsp, _ideal_lsp, _d):
        _diffs = []
        for _ilsp, _iideal in zip(_lsp, _ideal_lsp.T):
            _a = _ilsp[_d, :]
            _b = _iideal[_d]
            _diffs.append((_a - _b) / _b)
        return _diffs

    lshell_diff_e1 = _calc_rel_diff(e1_lsps, seeds_lsp, 0)
    phi_diff_e1 = _calc_rel_diff(e1_lsps, seeds_lsp, 1)

    lshell_diff_rk2 = _calc_rel_diff(rk2_lsps, seeds_lsp, 0)
    phi_diff_rk2 = _calc_rel_diff(rk2_lsps, seeds_lsp, 1)

    lshell_diff_rk4 = _calc_rel_diff(rk4_lsps, seeds_lsp, 0)
    phi_diff_rk4 = _calc_rel_diff(rk4_lsps, seeds_lsp, 1)

    lshell_diff_e1a = _calc_rel_diff(e1a_lsps, seeds_lsp, 0)
    phi_diff_e1a = _calc_rel_diff(e1a_lsps, seeds_lsp, 1)

    lshell_diff_rk12 = _calc_rel_diff(rk12_lsps, seeds_lsp, 0)
    phi_diff_rk12 = _calc_rel_diff(rk12_lsps, seeds_lsp, 1)

    lshell_diff_rk45 = _calc_rel_diff(rk45_lsps, seeds_lsp, 0)
    phi_diff_rk45 = _calc_rel_diff(rk45_lsps, seeds_lsp, 1)

    methods = ['Euler 1', 'Runge Kutta 2', 'Runge Kutta 4',
               'Euler 1 Adaptive Step', 'Runge Kutta 12 Adaptive Step',
               'Runge Kutta 45 Adaptive Step']
    wall_ts = [t_e1, t_rk2, t_rk4, t_e1a, t_rk12, t_rk45]
    all_lines = [e1_lines, rk2_lines, rk4_lines, e1a_lines, rk12_lines,
                 rk45_lines]
    all_lshell_diffs = [lshell_diff_e1, lshell_diff_rk2, lshell_diff_rk4,
                        lshell_diff_e1a, lshell_diff_rk12, lshell_diff_rk45]
    lshell_diffs = [np.abs(np.concatenate(lshell_diff_e1, axis=0)),
                    np.abs(np.concatenate(lshell_diff_rk2, axis=0)),
                    np.abs(np.concatenate(lshell_diff_rk4, axis=0)),
                    np.abs(np.concatenate(lshell_diff_e1a, axis=0)),
                    np.abs(np.concatenate(lshell_diff_rk12, axis=0)),
                    np.abs(np.concatenate(lshell_diff_rk45, axis=0))]
    phi_diffs = [np.abs(np.concatenate(phi_diff_e1, axis=0)),
                 np.abs(np.concatenate(phi_diff_rk2, axis=0)),
                 np.abs(np.concatenate(phi_diff_rk4, axis=0)),
                 np.abs(np.concatenate(phi_diff_e1a, axis=0)),
                 np.abs(np.concatenate(phi_diff_rk12, axis=0)),
                 np.abs(np.concatenate(phi_diff_rk45, axis=0))]
    npts = [len(lsd) for lsd in lshell_diffs]
    lshell_75 = [np.percentile(lsdiff, 75) for lsdiff in lshell_diffs]

    # # 3D DEBUG PLOT:: for really getting under the covers
    # vlab.clf()
    # earth1 = viscid.seed.Sphere((0.0, 0.0, 0.0), 1.0, pole=-m, ntheta=60, nphi=120,
    #                             thetalim=(15, 165), philim=(0, 360))
    # ls1 = viscid.xyz2lsrlp(earth1.get_points(), cotr=cotr, crd_system='gse')[0, :]
    # earth2 = viscid.seed.Sphere((0.0, 0.0, 0.0), 2.0, pole=-m, ntheta=60, nphi=120,
    #                             thetalim=(15, 165), philim=(0, 360))
    # ls2 = viscid.xyz2lsrlp(earth2.get_points(), cotr=cotr, crd_system='gse')[0, :]
    # earth4 = viscid.seed.Sphere((0.0, 0.0, 0.0), 4.0, pole=-m, ntheta=60, nphi=120,
    #                             thetalim=(15, 165), philim=(0, 360))
    # ls4 = viscid.xyz2lsrlp(earth4.get_points(), cotr=cotr, crd_system='gse')[0, :]
    # clim = [2.0, 6.0]
    # vlab.mesh_from_seeds(earth1, scalars=ls1, clim=clim, logscale=True)
    # vlab.mesh_from_seeds(earth2, scalars=ls2, clim=clim, logscale=True, opacity=0.5)
    # vlab.mesh_from_seeds(earth4, scalars=ls2, clim=clim, logscale=True, opacity=0.25)
    # vlab.plot3d_lines(e1_lines, scalars=[_e1_lsp[0, :] for _e1_lsp in e1_lsps],
    #                  clim=clim, logscale=True)
    # vlab.colorbar(title="L-Shell")
    # vlab.show()

    assert lshell_75[1] < lshell_75[0], "RK2 should have less error than Euler"
    assert lshell_75[2] < lshell_75[1], "RK4 should have less error than RK2"
    assert lshell_75[3] < lshell_75[0], "Euler 1a should have less error than Euler 1"
    assert lshell_75[4] < lshell_75[0], "RK 12 should have less error than Euler 1"
    assert lshell_75[5] < lshell_75[1], "RK 45 should have less error than RK2"

    try:
        if not plot2d:
            raise ImportError
        from viscid.plot import vpyplot as vlt
        from matplotlib import pyplot as plt

        # stats on error for all points on all lines
        _ = plt.figure(figsize=(15, 8))
        ax1 = vlt.subplot(121)
        v = plt.violinplot(lshell_diffs, showextrema=False, showmedians=False,
                               vert=False)
        colors = set_violin_colors(v)
        xl, xh = plt.gca().get_xlim()
        for i, txt, c in zip(count(), methods, colors):
            t_txt = ", took {0:.2e} seconds".format(wall_ts[i])
            stat_txt = format_data_range(lshell_diffs[i])
            plt.text(xl + 0.35 * (xh - xl), i + 1.15, txt + t_txt, color=c)
            plt.text(xl + 0.35 * (xh - xl), i + 0.85, stat_txt, color=c)
        ax1.get_yaxis().set_visible(False)
        plt.title('L-Shell')
        plt.xlabel('Relative Difference from Ideal (as fraction)')

        ax2 = vlt.subplot(122)
        v = plt.violinplot(phi_diffs, showextrema=False, showmedians=False,
                               vert=False)
        colors = set_violin_colors(v)
        xl, xh = plt.gca().get_xlim()
        for i, txt, c in zip(count(), methods, colors):
            t_txt = ", took {0:.2e} seconds".format(wall_ts[i])
            stat_txt = format_data_range(phi_diffs[i])
            plt.text(xl + 0.35 * (xh - xl), i + 1.15, txt + t_txt, color=c)
            plt.text(xl + 0.35 * (xh - xl), i + 0.85, stat_txt, color=c)
        ax2.get_yaxis().set_visible(False)
        plt.title('Longitude')
        plt.xlabel('Relative Difference from Ideal (as fraction)')

        vlt.auto_adjust_subplots()

        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if args.show:
            vlt.show()

        # stats for ds for all points on all lines
        _ = plt.figure(figsize=(10, 8))
        ax1 = vlt.subplot(111)

        ds = [np.concatenate([np.linalg.norm(_l[:, 1:] - _l[:, :-1], axis=0)
                              for _l in lines]) for lines in all_lines]
        v = plt.violinplot(ds, showextrema=False, showmedians=False,
                               vert=False)
        colors = set_violin_colors(v)
        xl, xh = plt.gca().get_xlim()
        for i, txt, c in zip(count(), methods, colors):
            stat_txt = format_data_range(ds[i])
            plt.annotate(txt, xy=(0.55, i / len(methods) + 0.1), color=c,
                         xycoords='axes fraction')
            plt.annotate(stat_txt, xy=(0.55, i / len(methods) + 0.04), color=c,
                         xycoords='axes fraction')
        ax1.get_yaxis().set_visible(False)
        plt.xscale('log')
        plt.title('Step Size')
        plt.xlabel('Absolute Step Size')
        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if args.show:
            vlt.show()


        # random other information
        _ = plt.figure(figsize=(13, 10))

        ## wall time for each method
        vlt.subplot(221)
        plt.scatter(range(len(methods)), wall_ts, color=colors,
                        s=150, marker='s', edgecolors='none')
        for i, meth in enumerate(methods):
            meth = meth.replace(" Adaptive Step", "\nAdaptive Step")
            plt.annotate(meth, (i, wall_ts[i]), xytext=(0, 15.0),
                             color=colors[i], horizontalalignment='center',
                             verticalalignment='bottom',
                             textcoords='offset points')
        plt.ylabel("Wall Time (s)")
        x_padding = 0.5
        plt.xlim(-x_padding, len(methods) - x_padding)
        yl, yh = np.min(wall_ts), np.max(wall_ts)
        y_padding = 0.4 * (yh - yl)
        plt.ylim(yl - y_padding, yh + y_padding)
        plt.gca().get_xaxis().set_visible(False)
        for _which in ('right', 'top'):
            plt.gca().spines[_which].set_color('none')

        ## number of points calculated for each method
        vlt.subplot(222)
        plt.scatter(range(len(methods)), npts, color=colors,
                        s=150, marker='s', edgecolors='none')
        for i, meth in enumerate(methods):
            meth = meth.replace(" Adaptive Step", "\nAdaptive Step")
            plt.annotate(meth, (i, npts[i]), xytext=(0, 15.0),
                             color=colors[i], horizontalalignment='center',
                             verticalalignment='bottom',
                             textcoords='offset points')
        plt.ylabel("Number of Streamline Points Calculated")
        x_padding = 0.5
        plt.xlim(-x_padding, len(methods) - x_padding)
        yl, yh = np.min(npts), np.max(npts)
        y_padding = 0.4 * (yh - yl)
        plt.ylim(yl - y_padding, yh + y_padding)
        plt.gca().get_xaxis().set_visible(False)
        for _which in ('right', 'top'):
            plt.gca().spines[_which].set_color('none')

        ## Wall time per segment, this should show the overhead of the method
        vlt.subplot(223)
        wall_t_per_seg = np.asarray(wall_ts) / np.asarray(npts)
        plt.scatter(range(len(methods)), wall_t_per_seg, color=colors,
                        s=150, marker='s', edgecolors='none')
        for i, meth in enumerate(methods):
            meth = meth.replace(" Adaptive Step", "\nAdaptive Step")
            plt.annotate(meth, (i, wall_t_per_seg[i]), xytext=(0, 15.0),
                             color=colors[i], horizontalalignment='center',
                             verticalalignment='bottom',
                             textcoords='offset points')
        plt.ylabel("Wall Time Per Line Segment")
        x_padding = 0.5
        plt.xlim(-x_padding, len(methods) - x_padding)
        yl, yh = np.min(wall_t_per_seg), np.max(wall_t_per_seg)
        y_padding = 0.4 * (yh - yl)
        plt.ylim(yl - y_padding, yh + y_padding)
        plt.gca().get_xaxis().set_visible(False)
        plt.gca().xaxis.set_major_formatter(viscid.plot.mpl_extra.steve_axfmt)
        for _which in ('right', 'top'):
            plt.gca().spines[_which].set_color('none')

        ## 75th percentile of l-shell error for each method
        vlt.subplot(224)
        plt.scatter(range(len(methods)), lshell_75, color=colors,
                        s=150, marker='s', edgecolors='none')
        plt.yscale('log')

        for i, meth in enumerate(methods):
            meth = meth.replace(" Adaptive Step", "\nAdaptive Step")
            plt.annotate(meth, (i, lshell_75[i]), xytext=(0, 15.0),
                             color=colors[i], horizontalalignment='center',
                             verticalalignment='bottom',
                             textcoords='offset points')
        plt.ylabel("75th Percentile of Relative L-Shell Error")
        x_padding = 0.5
        plt.xlim(-x_padding, len(methods) - x_padding)
        ymin, ymax = np.min(lshell_75), np.max(lshell_75)
        plt.ylim(0.75 * ymin, 2.5 * ymax)
        plt.gca().get_xaxis().set_visible(False)
        for _which in ('right', 'top'):
            plt.gca().spines[_which].set_color('none')

        vlt.auto_adjust_subplots(subplot_params=dict(wspace=0.25, hspace=0.15))

        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if args.show:
            vlt.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 args.show,
                              bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))
            _global_ns['figure'] = fig

        for i, method in zip(count(), methods):
            # if i in (3, 4):
            #     next_plot_fname(__file__, series='q3')
            #     print(i, "::", [line.shape[1] for line in all_lines[i]])
            #     # continue
            vlab.clf()
            _lshell_diff = [np.abs(s) for s in all_lshell_diffs[i]]
            vlab.plot3d_lines(all_lines[i], scalars=_lshell_diff)
            vlab.colorbar(title="Relative L-Shell Error (as fraction)")
            vlab.title(method, size=0.5)
            vlab.orientation_axes()
            vlab.view(azimuth=40, elevation=140, distance=80.0,
                      focalpoint=[0, 0, 0])
            vlab.savefig(next_plot_fname(__file__, series='q3'))
            if args.show:
                vlab.show()
    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
示例#19
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
示例#20
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    parser.add_argument("--interact", "-i", action="store_true")
    args = vutil.common_argparse(parser)

    f3d = viscid.load_file(os.path.join(sample_dir, 'sample_xdmf.3d.[0].xdmf'))
    f_iono = viscid.load_file(os.path.join(sample_dir, "sample_xdmf.iof.[0].xdmf"))

    b = f3d["b"]
    v = f3d["v"]
    pp = f3d["pp"]
    e = f3d["e_cc"]

    vlab.figure(size=(1280, 800), offscreen=not args.show)

    ##########################################################
    # make b a dipole inside 3.1Re and set e = 0 inside 4.0Re
    cotr = viscid.Cotr(time='1990-03-21T14:48', dip_tilt=0.0)  # pylint: disable=not-callable
    moment = cotr.get_dipole_moment(crd_system=b)
    isphere_mask = viscid.make_spherical_mask(b, rmax=3.1)
    viscid.fill_dipole(b, m=moment, mask=isphere_mask)
    e_mask = viscid.make_spherical_mask(b, rmax=4.0)
    viscid.set_in_region(e, 0.0, alpha=0.0, mask=e_mask, out=e)

    ######################################
    # plot a scalar cut plane of pressure
    pp_src = vlab.field2source(pp, center='node')
    scp = vlab.scalar_cut_plane(pp_src, plane_orientation='z_axes', opacity=0.5,
                                transparent=True, view_controls=False,
                                cmap="inferno", logscale=True)
    scp.implicit_plane.normal = [0, 0, -1]
    scp.implicit_plane.origin = [0, 0, 0]
    scp.enable_contours = True
    scp.contour.filled_contours = True
    scp.contour.number_of_contours = 64
    cbar = vlab.colorbar(scp, title=pp.name, orientation='vertical')
    cbar.scalar_bar_representation.position = (0.01, 0.13)
    cbar.scalar_bar_representation.position2 = (0.08, 0.76)

    ######################################
    # plot a vector cut plane of the flow
    vcp = vlab.vector_cut_plane(v, scalars=pp_src, plane_orientation='z_axes',
                                view_controls=False, mode='arrow',
                                cmap='Greens_r')
    vcp.implicit_plane.normal = [0, 0, -1]
    vcp.implicit_plane.origin = [0, 0, 0]

    ##############################
    # plot very faint isosurfaces
    vx_src = vlab.field2source(v['x'], center='node')
    iso = vlab.iso_surface(vx_src, contours=[0.0], opacity=0.008, cmap='Pastel1')

    ##############################################################
    # calculate B field lines && topology in Viscid and plot them
    seedsA = viscid.SphericalPatch([0, 0, 0], [2, 0, 1], 30, 15, r=5.0,
                                   nalpha=5, nbeta=5)
    seedsB = viscid.SphericalPatch([0, 0, 0], [1.9, 0, -20], 30, 15, r=5.0,
                                   nalpha=1, nbeta=5)
    seeds = np.concatenate([seedsA, seedsB], axis=1)
    b_lines, topo = viscid.calc_streamlines(b, seeds, ibound=3.5,
                                            obound0=[-25, -20, -20],
                                            obound1=[15, 20, 20], wrap=True)
    vlab.plot_lines(b_lines, scalars=viscid.topology2color(topo))

    ######################################################################
    # plot a random circle at geosynchronus orbit with scalars colored
    # by the Matplotlib viridis color map, just because we can; this is
    # a useful toy for debugging
    circle = viscid.Circle(p0=[0, 0, 0], r=6.618, n=128, endpoint=True)
    scalar = np.sin(circle.as_local_coordinates().get_crd('phi'))
    surf = vlab.plot_line(circle.get_points(), scalars=scalar, clim=0.8,
                          cmap="Spectral_r")

    ######################################################################
    # Use Mayavi (VTK) to calculate field lines using an interactive seed
    # These field lines are colored by E parallel
    epar = viscid.project(e, b)
    epar.name = "Epar"
    bsl2 = vlab.streamline(b, epar, seedtype='plane', seed_resolution=4,
                           integration_direction='both', clim=(-0.05, 0.05))

    # now tweak the VTK streamlines
    bsl2.stream_tracer.maximum_propagation = 20.
    bsl2.seed.widget.origin = [-11, -5.0, -2.0]
    bsl2.seed.widget.point1 = [-11, 5.0, -2.0]
    bsl2.seed.widget.point2 = [-11.0, -5.0, 2.0]
    bsl2.streamline_type = 'tube'
    bsl2.tube_filter.radius = 0.03
    bsl2.stop()  # this stop/start was a hack to get something to update
    bsl2.start()
    bsl2.seed.widget.enabled = False

    cbar = vlab.colorbar(bsl2, title=epar.name, label_fmt='%.3f',
                         orientation='horizontal')
    cbar.scalar_bar_representation.position = (0.15, 0.01)
    cbar.scalar_bar_representation.position2 = (0.72, 0.10)

    ###############################################################
    # Make a contour at the open-closed boundary in the ionosphere
    seeds_iono = viscid.Sphere(r=1.063, pole=-moment, ntheta=256, nphi=256,
                               thetalim=(0, 180), philim=(0, 360), crd_system=b)
    _, topo_iono = viscid.calc_streamlines(b, seeds_iono, ibound=1.0,
                                           nr_procs='all',
                                           output=viscid.OUTPUT_TOPOLOGY)
    topo_iono = np.log2(topo_iono)

    m = vlab.mesh_from_seeds(seeds_iono, scalars=topo_iono, opacity=1.0,
                             clim=(0, 3), color=(0.992, 0.445, 0.0))
    m.enable_contours = True
    m.actor.property.line_width = 4.0
    m.contour.number_of_contours = 4

    ####################################################################
    # Plot the ionosphere, note that the sample data has the ionosphere
    # at a different time, so the open-closed boundary found above
    # will not be consistant with the field aligned currents
    fac_tot = 1e9 * f_iono['fac_tot']

    m = vlab.plot_ionosphere(fac_tot, bounding_lat=30.0, vmin=-300, vmax=300,
                             opacity=0.75, rotate=cotr, crd_system=b)
    m.actor.property.backface_culling = True

    ########################################################################
    # Add some markers for earth, i.e., real earth, and dayside / nightside
    # representation
    vlab.plot_blue_marble(r=1.0, lines=False, ntheta=64, nphi=128,
                         rotate=cotr, crd_system=b)
    # now shade the night side with a transparent black hemisphere
    vlab.plot_earth_3d(radius=1.01, night_only=True, opacity=0.5, crd_system=b)

    ####################
    # Finishing Touches
    # vlab.axes(pp_src, nb_labels=5)
    oa = vlab.orientation_axes()
    oa.marker.set_viewport(0.75, 0.75, 1.0, 1.0)

    # note that resize won't work if the current figure has the
    # off_screen_rendering flag set
    # vlab.resize([1200, 800])
    vlab.view(azimuth=45, elevation=70, distance=35.0, focalpoint=[-2, 0, 0])

    ##############
    # Save Figure

    # print("saving png")
    # vlab.savefig('mayavi_msphere_sample.png')
    # print("saving x3d")
    # # x3d files can be turned into COLLADA files with meshlab, and
    # # COLLADA (.dae) files can be opened in OS X's preview
    # #
    # # IMPORTANT: for some reason, using bounding_lat in vlab.plot_ionosphere
    # #            causes a segfault when saving x3d files
    # #
    # vlab.savefig('mayavi_msphere_sample.x3d')
    # print("done")

    vlab.savefig(next_plot_fname(__file__))

    ###########################
    # Interact Programatically
    if args.interact:
        vlab.interact()

    #######################
    # Interact Graphically
    if args.show:
        vlab.show()

    try:
        vlab.mlab.close()
    except AttributeError:
        pass

    return 0
示例#21
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 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:
        viscid.logger.info('Testing RectilinearMeshPoints...')
        f = viscid.load_file(os.path.join(sample_dir, 'sample_xdmf.3d.[-1].xdmf'))
        slc = 'x=-40f:12f, y=-10f:10f, z=-10f:10f'
        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 matplotlib import pyplot as plt
            from viscid.plot import vpyplot as vlt
            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
            vlab.clf()
            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()

        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
示例#22
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    parser.add_argument("--interact", "-i", action="store_true")
    args = vutil.common_argparse(parser)

    f3d = viscid.load_file(os.path.join(sample_dir, 'sample_xdmf.3d.[0].xdmf'))
    f_iono = viscid.load_file(
        os.path.join(sample_dir, "sample_xdmf.iof.[0].xdmf"))

    b = f3d["b"]
    v = f3d["v"]
    pp = f3d["pp"]
    e = f3d["e_cc"]

    vlab.mlab.options.offscreen = not args.show
    vlab.figure(size=(1280, 800))

    ##########################################################
    # make b a dipole inside 3.1Re and set e = 0 inside 4.0Re
    cotr = viscid.Cotr(time='1990-03-21T14:48', dip_tilt=0.0)  # pylint: disable=not-callable
    moment = cotr.get_dipole_moment(crd_system=b)
    isphere_mask = viscid.make_spherical_mask(b, rmax=3.1)
    viscid.fill_dipole(b, m=moment, mask=isphere_mask)
    e_mask = viscid.make_spherical_mask(b, rmax=4.0)
    viscid.set_in_region(e, 0.0, alpha=0.0, mask=e_mask, out=e)

    ######################################
    # plot a scalar cut plane of pressure
    pp_src = vlab.field2source(pp, center='node')
    scp = vlab.scalar_cut_plane(pp_src,
                                plane_orientation='z_axes',
                                opacity=0.5,
                                transparent=True,
                                view_controls=False,
                                cmap="inferno",
                                logscale=True)
    scp.implicit_plane.normal = [0, 0, -1]
    scp.implicit_plane.origin = [0, 0, 0]
    scp.enable_contours = True
    scp.contour.filled_contours = True
    scp.contour.number_of_contours = 64
    cbar = vlab.colorbar(scp, title=pp.name, orientation='vertical')
    cbar.scalar_bar_representation.position = (0.01, 0.13)
    cbar.scalar_bar_representation.position2 = (0.08, 0.76)

    ######################################
    # plot a vector cut plane of the flow
    vcp = vlab.vector_cut_plane(v,
                                scalars=pp_src,
                                plane_orientation='z_axes',
                                view_controls=False,
                                mode='arrow',
                                cmap='Greens_r')
    vcp.implicit_plane.normal = [0, 0, -1]
    vcp.implicit_plane.origin = [0, 0, 0]

    ##############################
    # plot very faint isosurfaces
    vx_src = vlab.field2source(v['x'], center='node')
    iso = vlab.iso_surface(vx_src,
                           contours=[0.0],
                           opacity=0.008,
                           cmap='Pastel1')

    ##############################################################
    # calculate B field lines && topology in Viscid and plot them
    seedsA = viscid.SphericalPatch([0, 0, 0], [2, 0, 1],
                                   30,
                                   15,
                                   r=5.0,
                                   nalpha=5,
                                   nbeta=5)
    seedsB = viscid.SphericalPatch([0, 0, 0], [1.9, 0, -20],
                                   30,
                                   15,
                                   r=5.0,
                                   nalpha=1,
                                   nbeta=5)
    seeds = np.concatenate([seedsA, seedsB], axis=1)
    b_lines, topo = viscid.calc_streamlines(b,
                                            seeds,
                                            ibound=3.5,
                                            obound0=[-25, -20, -20],
                                            obound1=[15, 20, 20],
                                            wrap=True)
    vlab.plot_lines(b_lines, scalars=viscid.topology2color(topo))

    ######################################################################
    # plot a random circle at geosynchronus orbit with scalars colored
    # by the Matplotlib viridis color map, just because we can; this is
    # a useful toy for debugging
    circle = viscid.Circle(p0=[0, 0, 0], r=6.618, n=128, endpoint=True)
    scalar = np.sin(circle.as_local_coordinates().get_crd('phi'))
    surf = vlab.plot_line(circle.get_points(),
                          scalars=scalar,
                          clim=0.8,
                          cmap="Spectral_r")

    ######################################################################
    # Use Mayavi (VTK) to calculate field lines using an interactive seed
    # These field lines are colored by E parallel
    epar = viscid.project(e, b)
    epar.name = "Epar"
    bsl2 = vlab.streamline(b,
                           epar,
                           seedtype='plane',
                           seed_resolution=4,
                           integration_direction='both',
                           clim=(-0.05, 0.05))

    # now tweak the VTK streamlines
    bsl2.stream_tracer.maximum_propagation = 20.
    bsl2.seed.widget.origin = [-11, -5.0, -2.0]
    bsl2.seed.widget.point1 = [-11, 5.0, -2.0]
    bsl2.seed.widget.point2 = [-11.0, -5.0, 2.0]
    bsl2.streamline_type = 'tube'
    bsl2.tube_filter.radius = 0.03
    bsl2.stop()  # this stop/start was a hack to get something to update
    bsl2.start()
    bsl2.seed.widget.enabled = False

    cbar = vlab.colorbar(bsl2,
                         title=epar.name,
                         label_fmt='%.3f',
                         orientation='horizontal')
    cbar.scalar_bar_representation.position = (0.15, 0.01)
    cbar.scalar_bar_representation.position2 = (0.72, 0.10)

    ###############################################################
    # Make a contour at the open-closed boundary in the ionosphere
    seeds_iono = viscid.Sphere(r=1.063,
                               pole=-moment,
                               ntheta=256,
                               nphi=256,
                               thetalim=(0, 180),
                               philim=(0, 360),
                               crd_system=b)
    _, topo_iono = viscid.calc_streamlines(b,
                                           seeds_iono,
                                           ibound=1.0,
                                           nr_procs='all',
                                           output=viscid.OUTPUT_TOPOLOGY)
    topo_iono = np.log2(topo_iono)

    m = vlab.mesh_from_seeds(seeds_iono,
                             scalars=topo_iono,
                             opacity=1.0,
                             clim=(0, 3),
                             color=(0.992, 0.445, 0.0))
    m.enable_contours = True
    m.actor.property.line_width = 4.0
    m.contour.number_of_contours = 4

    ####################################################################
    # Plot the ionosphere, note that the sample data has the ionosphere
    # at a different time, so the open-closed boundary found above
    # will not be consistant with the field aligned currents
    fac_tot = 1e9 * f_iono['fac_tot']

    m = vlab.plot_ionosphere(fac_tot,
                             bounding_lat=30.0,
                             vmin=-300,
                             vmax=300,
                             opacity=0.75,
                             rotate=cotr,
                             crd_system=b)
    m.actor.property.backface_culling = True

    ########################################################################
    # Add some markers for earth, i.e., real earth, and dayside / nightside
    # representation
    vlab.plot_blue_marble(r=1.0,
                          lines=False,
                          ntheta=64,
                          nphi=128,
                          rotate=cotr,
                          crd_system=b)
    # now shade the night side with a transparent black hemisphere
    vlab.plot_earth_3d(radius=1.01, night_only=True, opacity=0.5, crd_system=b)

    ####################
    # Finishing Touches
    # vlab.axes(pp_src, nb_labels=5)
    oa = vlab.orientation_axes()
    oa.marker.set_viewport(0.75, 0.75, 1.0, 1.0)

    # note that resize won't work if the current figure has the
    # off_screen_rendering flag set
    # vlab.resize([1200, 800])
    vlab.view(azimuth=45, elevation=70, distance=35.0, focalpoint=[-2, 0, 0])

    ##############
    # Save Figure

    # print("saving png")
    # vlab.savefig('mayavi_msphere_sample.png')
    # print("saving x3d")
    # # x3d files can be turned into COLLADA files with meshlab, and
    # # COLLADA (.dae) files can be opened in OS X's preview
    # #
    # # IMPORTANT: for some reason, using bounding_lat in vlab.plot_ionosphere
    # #            causes a segfault when saving x3d files
    # #
    # vlab.savefig('mayavi_msphere_sample.x3d')
    # print("done")

    vlab.savefig(next_plot_fname(__file__))

    ###########################
    # Interact Programatically
    if args.interact:
        vlab.interact()

    #######################
    # Interact Graphically
    if args.show:
        vlab.show()

    try:
        vlab.mlab.close()
    except AttributeError:
        pass

    return 0
示例#23
0
def _main():
    f = viscid.load_file("$WORK/xi_fte_001/*.3d.[4050f].xdmf")
    mp = get_mp_info(f['pp'],
                     f['b'],
                     f['j'],
                     f['e_cc'],
                     fit='mp_xloc',
                     slc="x=6.5f:10.5f, y=-4f:4f, z=-4.8f:3f",
                     cache=False)

    y, z = mp['pp_max_xloc'].meshgrid_flat(prune=True)
    x = mp['pp_max_xloc'].data.reshape(-1)

    Y, Z = mp['pp_max_xloc'].meshgrid(prune=True)
    x2 = paraboloid(Y, Z, *mp['paraboloid'][0])

    skip = 117
    n = paraboloid_normal(Y, Z, *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]

    minvar_y = Y.reshape(-1)[::skip]
    minvar_z = Z.reshape(-1)[::skip]
    minvar_n = np.zeros([3, len(minvar_y)])
    for i in range(minvar_n.shape[0]):
        p0 = [0.0, minvar_y[i], minvar_z[i]]
        p0[0] = mp['pp_max_xloc']['y={0[0]}f, z={0[1]}f'.format(p0)]
        minvar_n[:, i] = viscid.find_minvar_lmn_around(f['b'], p0, l=2.0,
                                                       n=64)[2, :]

    # 2d plots, normals don't look normal in the matplotlib projection
    if False:  # pylint: disable=using-constant-test
        from matplotlib import pyplot as plt
        from viscid.plot import vpyplot as vlt

        normals = paraboloid_normal(Y, Z, *mp['paraboloid'][0])
        p0 = np.array([x2, Y, Z]).reshape(3, -1)
        p1 = p0 + normals.reshape(3, -1)

        vlt.scatter_3d(np.vstack([x, y, z])[:, ::skip], equal=True)
        for i in range(0, p0.shape[1], skip):
            plt.gca().plot([p0[0, i], p1[0, i]], [p0[1, i], p1[1, i]],
                           [p0[2, i], p1[2, i]],
                           color='c')
        # z2 = _ellipsiod(X, Y, *popt)
        plt.gca().plot_surface(Y, Z, x2, color='r')
        vlt.show()

    # mayavi 3d plots, normals look better here
    if True:  # pylint: disable=using-constant-test
        from viscid.plot import vlab
        vlab.points3d(x[::skip],
                      y[::skip],
                      z[::skip],
                      scale_factor=0.25,
                      color=(0.0, 0.0, 1.0))

        mp_width = mp['mp_width']['x=0']
        mp_sheath_edge = mp['mp_sheath_edge']['x=0']
        mp_sphere_edge = mp_sheath_edge - mp_width

        vlab.mesh(x2, Y, Z, scalars=mp_width.data)
        vlab.mesh(mp_sheath_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)
        vlab.mesh(mp_sphere_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)

        n = paraboloid_normal(Y, Z,
                              *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]
        vlab.quiver3d(x2.reshape(-1)[::skip],
                      Y.reshape(-1)[::skip],
                      Z.reshape(-1)[::skip],
                      n[0],
                      n[1],
                      n[2],
                      color=(1, 0, 0))
        vlab.quiver3d(x2.reshape(-1)[::skip],
                      Y.reshape(-1)[::skip],
                      Z.reshape(-1)[::skip],
                      minvar_n[0],
                      minvar_n[1],
                      minvar_n[2],
                      color=(0, 0, 1))
        vlab.show()
示例#24
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

    # #################################################
    # viscid.logger.info("Testing field lines on 2d field...")
    B = viscid.make_dipole(twod=True)
    line = viscid.seed.Line((0.2, 0.0, 0.0), (1.0, 0.0, 0.0), 10)
    obound0 = np.array([-4, -4, -4], dtype=B.data.dtype)
    obound1 = np.array([4, 4, 4], dtype=B.data.dtype)
    run_test(B,
             line,
             plot2d=plot2d,
             plot3d=plot3d,
             title='2D',
             show=args.show,
             ibound=0.07,
             obound0=obound0,
             obound1=obound1)

    #################################################
    viscid.logger.info("Testing field lines on 3d field...")
    B = viscid.make_dipole(m=[0.2, 0.3, -0.9])
    sphere = viscid.seed.Sphere((0.0, 0.0, 0.0), 2.0, ntheta=20, nphi=10)
    obound0 = np.array([-4, -4, -4], dtype=B.data.dtype)
    obound1 = np.array([4, 4, 4], dtype=B.data.dtype)
    run_test(B,
             sphere,
             plot2d=plot2d,
             plot3d=plot3d,
             title='3D',
             show=args.show,
             ibound=0.12,
             obound0=obound0,
             obound1=obound1,
             method=viscid.RK12)

    # The Remainder of this test makes sure higher order methods are indeed
    # more accurate than lower order methods... this could find a bug in
    # the integrators

    ##################################################
    # test accuracy of streamlines in an ideal dipole
    cotr = viscid.Cotr(dip_tilt=15.0, dip_gsm=21.0)  # pylint: disable=not-callable
    m = cotr.get_dipole_moment(crd_system='gse')
    seeds = viscid.seed.Sphere((0.0, 0.0, 0.0),
                               2.0,
                               pole=-m,
                               ntheta=25,
                               nphi=25,
                               thetalim=(5, 90),
                               philim=(5, 360),
                               phi_endpoint=False)
    B = viscid.make_dipole(m=m,
                           crd_system='gse',
                           n=(256, 256, 256),
                           l=(-25, -25, -25),
                           h=(25, 25, 25),
                           dtype='f8')

    seeds_xyz = seeds.get_points()
    # seeds_lsp = viscid.xyz2lsrlp(seeds_xyz, cotr=cotr, crd_system=B)[(0, 3), :]
    seeds_lsp = viscid.xyz2lsrlp(seeds_xyz, cotr=cotr, crd_system=B)[(0, 3), :]

    e1_lines, e1_lsps, t_e1 = lines_and_lsps(B,
                                             seeds,
                                             method='euler1',
                                             ibound=1.0,
                                             cotr=cotr)
    rk2_lines, rk2_lsps, t_rk2 = lines_and_lsps(B,
                                                seeds,
                                                method='rk2',
                                                ibound=1.0,
                                                cotr=cotr)
    rk4_lines, rk4_lsps, t_rk4 = lines_and_lsps(B,
                                                seeds,
                                                method='rk4',
                                                ibound=1.0,
                                                cotr=cotr)
    e1a_lines, e1a_lsps, t_e1a = lines_and_lsps(B,
                                                seeds,
                                                method='euler1a',
                                                ibound=1.0,
                                                cotr=cotr)
    rk12_lines, rk12_lsps, t_rk12 = lines_and_lsps(B,
                                                   seeds,
                                                   method='rk12',
                                                   ibound=1.0,
                                                   cotr=cotr)
    rk45_lines, rk45_lsps, t_rk45 = lines_and_lsps(B,
                                                   seeds,
                                                   method='rk45',
                                                   ibound=1.0,
                                                   cotr=cotr)

    def _calc_rel_diff(_lsp, _ideal_lsp, _d):
        _diffs = []
        for _ilsp, _iideal in zip(_lsp, _ideal_lsp.T):
            _a = _ilsp[_d, :]
            _b = _iideal[_d]
            _diffs.append((_a - _b) / _b)
        return _diffs

    lshell_diff_e1 = _calc_rel_diff(e1_lsps, seeds_lsp, 0)
    phi_diff_e1 = _calc_rel_diff(e1_lsps, seeds_lsp, 1)

    lshell_diff_rk2 = _calc_rel_diff(rk2_lsps, seeds_lsp, 0)
    phi_diff_rk2 = _calc_rel_diff(rk2_lsps, seeds_lsp, 1)

    lshell_diff_rk4 = _calc_rel_diff(rk4_lsps, seeds_lsp, 0)
    phi_diff_rk4 = _calc_rel_diff(rk4_lsps, seeds_lsp, 1)

    lshell_diff_e1a = _calc_rel_diff(e1a_lsps, seeds_lsp, 0)
    phi_diff_e1a = _calc_rel_diff(e1a_lsps, seeds_lsp, 1)

    lshell_diff_rk12 = _calc_rel_diff(rk12_lsps, seeds_lsp, 0)
    phi_diff_rk12 = _calc_rel_diff(rk12_lsps, seeds_lsp, 1)

    lshell_diff_rk45 = _calc_rel_diff(rk45_lsps, seeds_lsp, 0)
    phi_diff_rk45 = _calc_rel_diff(rk45_lsps, seeds_lsp, 1)

    methods = [
        'Euler 1', 'Runge Kutta 2', 'Runge Kutta 4', 'Euler 1 Adaptive Step',
        'Runge Kutta 12 Adaptive Step', 'Runge Kutta 45 Adaptive Step'
    ]
    wall_ts = [t_e1, t_rk2, t_rk4, t_e1a, t_rk12, t_rk45]
    all_lines = [
        e1_lines, rk2_lines, rk4_lines, e1a_lines, rk12_lines, rk45_lines
    ]
    all_lshell_diffs = [
        lshell_diff_e1, lshell_diff_rk2, lshell_diff_rk4, lshell_diff_e1a,
        lshell_diff_rk12, lshell_diff_rk45
    ]
    lshell_diffs = [
        np.abs(np.concatenate(lshell_diff_e1, axis=0)),
        np.abs(np.concatenate(lshell_diff_rk2, axis=0)),
        np.abs(np.concatenate(lshell_diff_rk4, axis=0)),
        np.abs(np.concatenate(lshell_diff_e1a, axis=0)),
        np.abs(np.concatenate(lshell_diff_rk12, axis=0)),
        np.abs(np.concatenate(lshell_diff_rk45, axis=0))
    ]
    phi_diffs = [
        np.abs(np.concatenate(phi_diff_e1, axis=0)),
        np.abs(np.concatenate(phi_diff_rk2, axis=0)),
        np.abs(np.concatenate(phi_diff_rk4, axis=0)),
        np.abs(np.concatenate(phi_diff_e1a, axis=0)),
        np.abs(np.concatenate(phi_diff_rk12, axis=0)),
        np.abs(np.concatenate(phi_diff_rk45, axis=0))
    ]
    npts = [len(lsd) for lsd in lshell_diffs]
    lshell_75 = [np.percentile(lsdiff, 75) for lsdiff in lshell_diffs]

    # # 3D DEBUG PLOT:: for really getting under the covers
    # vlab.clf()
    # earth1 = viscid.seed.Sphere((0.0, 0.0, 0.0), 1.0, pole=-m, ntheta=60, nphi=120,
    #                             thetalim=(15, 165), philim=(0, 360))
    # ls1 = viscid.xyz2lsrlp(earth1.get_points(), cotr=cotr, crd_system='gse')[0, :]
    # earth2 = viscid.seed.Sphere((0.0, 0.0, 0.0), 2.0, pole=-m, ntheta=60, nphi=120,
    #                             thetalim=(15, 165), philim=(0, 360))
    # ls2 = viscid.xyz2lsrlp(earth2.get_points(), cotr=cotr, crd_system='gse')[0, :]
    # earth4 = viscid.seed.Sphere((0.0, 0.0, 0.0), 4.0, pole=-m, ntheta=60, nphi=120,
    #                             thetalim=(15, 165), philim=(0, 360))
    # ls4 = viscid.xyz2lsrlp(earth4.get_points(), cotr=cotr, crd_system='gse')[0, :]
    # clim = [2.0, 6.0]
    # vlab.mesh_from_seeds(earth1, scalars=ls1, clim=clim, logscale=True)
    # vlab.mesh_from_seeds(earth2, scalars=ls2, clim=clim, logscale=True, opacity=0.5)
    # vlab.mesh_from_seeds(earth4, scalars=ls2, clim=clim, logscale=True, opacity=0.25)
    # vlab.plot3d_lines(e1_lines, scalars=[_e1_lsp[0, :] for _e1_lsp in e1_lsps],
    #                  clim=clim, logscale=True)
    # vlab.colorbar(title="L-Shell")
    # vlab.show()

    assert lshell_75[1] < lshell_75[0], "RK2 should have less error than Euler"
    assert lshell_75[2] < lshell_75[1], "RK4 should have less error than RK2"
    assert lshell_75[3] < lshell_75[
        0], "Euler 1a should have less error than Euler 1"
    assert lshell_75[4] < lshell_75[
        0], "RK 12 should have less error than Euler 1"
    assert lshell_75[5] < lshell_75[1], "RK 45 should have less error than RK2"

    try:
        if not plot2d:
            raise ImportError
        from matplotlib import pyplot as plt
        from viscid.plot import vpyplot as vlt

        # stats on error for all points on all lines
        _ = plt.figure(figsize=(15, 8))
        ax1 = vlt.subplot(121)
        v = plt.violinplot(lshell_diffs,
                           showextrema=False,
                           showmedians=False,
                           vert=False)
        colors = set_violin_colors(v)
        xl, xh = plt.gca().get_xlim()
        for i, txt, c in zip(count(), methods, colors):
            t_txt = ", took {0:.2e} seconds".format(wall_ts[i])
            stat_txt = format_data_range(lshell_diffs[i])
            plt.text(xl + 0.35 * (xh - xl), i + 1.15, txt + t_txt, color=c)
            plt.text(xl + 0.35 * (xh - xl), i + 0.85, stat_txt, color=c)
        ax1.get_yaxis().set_visible(False)
        plt.title('L-Shell')
        plt.xlabel('Relative Difference from Ideal (as fraction)')

        ax2 = vlt.subplot(122)
        v = plt.violinplot(phi_diffs,
                           showextrema=False,
                           showmedians=False,
                           vert=False)
        colors = set_violin_colors(v)
        xl, xh = plt.gca().get_xlim()
        for i, txt, c in zip(count(), methods, colors):
            t_txt = ", took {0:.2e} seconds".format(wall_ts[i])
            stat_txt = format_data_range(phi_diffs[i])
            plt.text(xl + 0.35 * (xh - xl), i + 1.15, txt + t_txt, color=c)
            plt.text(xl + 0.35 * (xh - xl), i + 0.85, stat_txt, color=c)
        ax2.get_yaxis().set_visible(False)
        plt.title('Longitude')
        plt.xlabel('Relative Difference from Ideal (as fraction)')

        vlt.auto_adjust_subplots()

        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if args.show:
            vlt.show()

        # stats for ds for all points on all lines
        _ = plt.figure(figsize=(10, 8))
        ax1 = vlt.subplot(111)

        ds = [
            np.concatenate([
                np.linalg.norm(_l[:, 1:] - _l[:, :-1], axis=0) for _l in lines
            ]) for lines in all_lines
        ]
        v = plt.violinplot(ds,
                           showextrema=False,
                           showmedians=False,
                           vert=False)
        colors = set_violin_colors(v)
        xl, xh = plt.gca().get_xlim()
        for i, txt, c in zip(count(), methods, colors):
            stat_txt = format_data_range(ds[i])
            plt.text(xl + 0.01 * (xh - xl), i + 1.15, txt, color=c)
            plt.text(xl + 0.01 * (xh - xl), i + 0.85, stat_txt, color=c)
        ax1.get_yaxis().set_visible(False)
        plt.xscale('log')
        plt.title('Step Size')
        plt.xlabel('Absolute Step Size')
        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if args.show:
            vlt.show()

        # random other information
        _ = plt.figure(figsize=(13, 10))

        ## wall time for each method
        vlt.subplot(221)
        plt.scatter(range(len(methods)),
                    wall_ts,
                    color=colors,
                    s=150,
                    marker='s',
                    edgecolors='none')
        for i, meth in enumerate(methods):
            meth = meth.replace(" Adaptive Step", "\nAdaptive Step")
            plt.annotate(meth, (i, wall_ts[i]),
                         xytext=(0, 15.0),
                         color=colors[i],
                         horizontalalignment='center',
                         verticalalignment='bottom',
                         textcoords='offset points')
        plt.ylabel("Wall Time (s)")
        x_padding = 0.5
        plt.xlim(-x_padding, len(methods) - x_padding)
        yl, yh = np.min(wall_ts), np.max(wall_ts)
        y_padding = 0.4 * (yh - yl)
        plt.ylim(yl - y_padding, yh + y_padding)
        plt.gca().get_xaxis().set_visible(False)
        for _which in ('right', 'top'):
            plt.gca().spines[_which].set_color('none')

        ## number of points calculated for each method
        vlt.subplot(222)
        plt.scatter(range(len(methods)),
                    npts,
                    color=colors,
                    s=150,
                    marker='s',
                    edgecolors='none')
        for i, meth in enumerate(methods):
            meth = meth.replace(" Adaptive Step", "\nAdaptive Step")
            plt.annotate(meth, (i, npts[i]),
                         xytext=(0, 15.0),
                         color=colors[i],
                         horizontalalignment='center',
                         verticalalignment='bottom',
                         textcoords='offset points')
        plt.ylabel("Number of Streamline Points Calculated")
        x_padding = 0.5
        plt.xlim(-x_padding, len(methods) - x_padding)
        yl, yh = np.min(npts), np.max(npts)
        y_padding = 0.4 * (yh - yl)
        plt.ylim(yl - y_padding, yh + y_padding)
        plt.gca().get_xaxis().set_visible(False)
        for _which in ('right', 'top'):
            plt.gca().spines[_which].set_color('none')

        ## Wall time per segment, this should show the overhead of the method
        vlt.subplot(223)
        wall_t_per_seg = np.asarray(wall_ts) / np.asarray(npts)
        plt.scatter(range(len(methods)),
                    wall_t_per_seg,
                    color=colors,
                    s=150,
                    marker='s',
                    edgecolors='none')
        for i, meth in enumerate(methods):
            meth = meth.replace(" Adaptive Step", "\nAdaptive Step")
            plt.annotate(meth, (i, wall_t_per_seg[i]),
                         xytext=(0, 15.0),
                         color=colors[i],
                         horizontalalignment='center',
                         verticalalignment='bottom',
                         textcoords='offset points')
        plt.ylabel("Wall Time Per Line Segment")
        x_padding = 0.5
        plt.xlim(-x_padding, len(methods) - x_padding)
        yl, yh = np.min(wall_t_per_seg), np.max(wall_t_per_seg)
        y_padding = 0.4 * (yh - yl)
        plt.ylim(yl - y_padding, yh + y_padding)
        plt.gca().get_xaxis().set_visible(False)
        plt.gca().xaxis.set_major_formatter(viscid.plot.mpl_extra.steve_axfmt)
        for _which in ('right', 'top'):
            plt.gca().spines[_which].set_color('none')

        ## 75th percentile of l-shell error for each method
        vlt.subplot(224)
        plt.scatter(range(len(methods)),
                    lshell_75,
                    color=colors,
                    s=150,
                    marker='s',
                    edgecolors='none')
        plt.yscale('log')

        for i, meth in enumerate(methods):
            meth = meth.replace(" Adaptive Step", "\nAdaptive Step")
            plt.annotate(meth, (i, lshell_75[i]),
                         xytext=(0, 15.0),
                         color=colors[i],
                         horizontalalignment='center',
                         verticalalignment='bottom',
                         textcoords='offset points')
        plt.ylabel("75th Percentile of Relative L-Shell Error")
        x_padding = 0.5
        plt.xlim(-x_padding, len(methods) - x_padding)
        ymin, ymax = np.min(lshell_75), np.max(lshell_75)
        plt.ylim(0.75 * ymin, 2.5 * ymax)
        plt.gca().get_xaxis().set_visible(False)
        for _which in ('right', 'top'):
            plt.gca().spines[_which].set_color('none')

        vlt.auto_adjust_subplots(subplot_params=dict(wspace=0.25, hspace=0.15))

        vlt.savefig(next_plot_fname(__file__, series='q2'))
        if args.show:
            vlt.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 args.show,
                              bgcolor=(1, 1, 1),
                              fgcolor=(0, 0, 0))
            _global_ns['figure'] = fig

        for i, method in zip(count(), methods):
            # if i in (3, 4):
            #     next_plot_fname(__file__, series='q3')
            #     print(i, "::", [line.shape[1] for line in all_lines[i]])
            #     # continue
            vlab.clf()
            _lshell_diff = [np.abs(s) for s in all_lshell_diffs[i]]
            vlab.plot3d_lines(all_lines[i], scalars=_lshell_diff)
            vlab.colorbar(title="Relative L-Shell Error (as fraction)")
            vlab.title(method, size=0.5)
            vlab.orientation_axes()
            vlab.view(azimuth=40,
                      elevation=140,
                      distance=80.0,
                      focalpoint=[0, 0, 0])
            vlab.savefig(next_plot_fname(__file__, series='q3'))
            if args.show:
                vlab.show()
    except ImportError:
        pass

    # prevent weird xorg bad-instructions on tear down
    if 'figure' in _global_ns and _global_ns['figure'] is not None:
        from viscid.plot import vlab
        vlab.mlab.close(_global_ns['figure'])

    return 0
def _main():
    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
示例#26
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
示例#27
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