Пример #1
0
        def _plot_time_range(times, figname):
            for i, t in enumerate(times):
                vlab.clf()
                cotr = Cotr(t)

                vlab.plot_blue_marble(r=1.0, rotate=t, crd_system=crd_system,
                                      nphi=256, ntheta=128, res=4, lines=True)

                vlab.plot_earth_3d(radius=1.005, crd_system=crd_system,
                                   night_only=True, opacity=0.5)

                mag_north = cotr.transform('sm', crd_system, [0, 0, 1.0])

                vlab.mlab.points3d(*mag_north, scale_factor=0.05, mode='sphere',
                                   color=(0.992, 0.455, 0.0), resolution=32)
                vlab.orientation_axes(line_width=4.0)

                vlab.mlab.text(0.325, 0.95, viscid.format_datetime(t))

                vlab.view(azimuth=0.0, elevation=90.0, distance=5.0,
                          focalpoint=[0, 0, 0])
                vlab.savefig("{0}_eq_{1:06d}.png".format(figname, i))
                vlab.view(azimuth=0.0, elevation=0.0, distance=5.0,
                          focalpoint=[0, 0, 0])
                vlab.savefig("{0}_pole_{1:06d}.png".format(figname, i))
Пример #2
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
Пример #3
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
Пример #4
0
        def _plot_time_range(times, figname):
            for i, t in enumerate(times):
                vlab.clf()
                cotr = Cotr(t)

                vlab.plot_blue_marble(r=1.0,
                                      rotate=t,
                                      crd_system=crd_system,
                                      nphi=256,
                                      ntheta=128,
                                      res=4,
                                      lines=True)

                vlab.plot_earth_3d(radius=1.005,
                                   crd_system=crd_system,
                                   night_only=True,
                                   opacity=0.5)

                mag_north = cotr.transform('sm', crd_system, [0, 0, 1.0])

                vlab.mlab.points3d(*mag_north,
                                   scale_factor=0.05,
                                   mode='sphere',
                                   color=(0.992, 0.455, 0.0),
                                   resolution=32)
                vlab.orientation_axes(line_width=4.0)

                vlab.mlab.text(0.325, 0.95, viscid.format_datetime(t))

                vlab.view(azimuth=0.0,
                          elevation=90.0,
                          distance=5.0,
                          focalpoint=[0, 0, 0])
                vlab.savefig("{0}_eq_{1:06d}.png".format(figname, i))
                vlab.view(azimuth=0.0,
                          elevation=0.0,
                          distance=5.0,
                          focalpoint=[0, 0, 0])
                vlab.savefig("{0}_pole_{1:06d}.png".format(figname, i))
Пример #5
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
Пример #6
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
Пример #7
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
Пример #8
0
def trace_separator(grid,
                    b_slcstr="x=-25f:15f, y=-30f:30f, z=-15f:15f",
                    r=1.0,
                    plot=False,
                    trace_opts=None,
                    cache=True,
                    cache_dir=None):
    """Trace a separator line from most dawnward null

    **Still in testing** Uses the bisection algorithm.

    Args:
        grid (Grid): A grid that has a "b" field
        b_slcstr (str): Some valid slice for B field
        r (float): spatial step of separator line
        plot (bool): make debugging plots
        trace_opts (dict): passed to streamline function
        cache (bool, str): Save to and load from cache, if "force",
            then don't load from cache if it exists, but do save a
            cache at the end
        cache_dir (str): Directory for cache, if None, same directory
            as that file to which the grid belongs

    Raises:
        IOError: Description

    Returns:
        tuple: (separator_lines, nulls)

          - **separator_lines** (list): list of M 3xN ndarrays that
            represent M separator lines with N points
          - **nulls** (ndarray): 3xN array of N null points
    """
    if not cache_dir:
        cache_dir = grid.find_info("_viscid_dirname", "./")
    run_name = grid.find_info("run")
    sep_fname = "{0}/{1}.sep.{2:06.0f}".format(cache_dir, run_name, grid.time)

    try:
        if isinstance(cache,
                      string_types) and cache.strip().lower() == "force":
            raise IOError()
        with np.load(sep_fname + ".npz") as dat:
            sep_iter = (f for f in dat.files if f.startswith("arr_"))
            _it = sorted(sep_iter, key=lambda s: int(s[len("arr_"):]))
            seps = [dat[n] for n in _it]
            nulls = dat['nulls']
    except IOError:
        _b = grid['b'][b_slcstr]

        _, nulls = viscid.find_nulls(_b['x=-30f:15f'], ibound=5.0)

        # get most dawnward null, nulls2 is all nulls except p0
        nullind = np.argmin(nulls[1, :])
        p0 = nulls[:, nullind]
        nulls2 = np.concatenate([nulls[:, :nullind], nulls[:, (nullind + 1):]],
                                axis=1)

        if plot:
            from viscid.plot import vlab
            vlab.plot_earth_3d(crd_system='gse')
            vlab.points3d(nulls2[0],
                          nulls2[1],
                          nulls2[2],
                          color=(0, 0, 0),
                          scale_factor=1.0)
            vlab.points3d(nulls[0, nullind],
                          nulls[1, nullind],
                          nulls[2, nullind],
                          color=(1, 1, 1),
                          scale_factor=1.0)

        seed = viscid.Sphere(p0=p0,
                             r=r,
                             ntheta=30,
                             nphi=60,
                             theta_endpoint=True,
                             phi_endpoint=True)
        p1 = viscid.get_sep_pts_bisect(_b,
                                       seed,
                                       max_depth=12,
                                       plot=plot,
                                       trace_opts=trace_opts)
        # print("p1 shape", p1.shape)

        # if p1.shape[1] > 2:
        #     raise RuntimeError("Invalid B field, should be no branch @ null")

        seps = []
        sep_stubs = []
        for i in range(p1.shape[1]):
            sep_stubs.append([p0, p1[:, i]])

        # print("??", sep_stubs)

        while sep_stubs:
            sep = sep_stubs.pop(0)
            # print("!!! new stub")

            for i in count():
                # print("::", i)
                seed = viscid.SphericalPatch(p0=sep[-1],
                                             p1=sep[-1] - sep[-2],
                                             r=r,
                                             nalpha=240,
                                             nbeta=240)
                pn = viscid.get_sep_pts_bisect(_b,
                                               seed,
                                               max_depth=8,
                                               plot=plot,
                                               trace_opts=trace_opts)
                if pn.shape[1] == 0:
                    # print("END: pn.shape[1] == 0")
                    break
                # print("++", nulls2.shape, pn.shape)
                closest_null_dist = np.min(
                    np.linalg.norm(nulls2 - pn[:, :1], axis=0))
                # print("closest_null_dist:", closest_null_dist)
                if closest_null_dist < 1.01 * r:
                    # print("END: within 1.01 of a null")
                    break

                # print("??", pn)
                for j in range(1, pn.shape[1]):
                    # print("inserting new stub")
                    sep_stubs.insert(0, [sep[-1], pn[:, j]])
                sep.append(pn[:, 0])

            # print("sep", sep)
            seps.append(np.stack(sep, axis=1))

        if cache:
            np.savez_compressed(sep_fname, *seps, nulls=nulls)
    return seps, nulls
Пример #9
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
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
Пример #11
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
Пример #12
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
Пример #13
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
Пример #14
0
def trace_separator(grid, b_slcstr="x=-25j:15j, y=-30j:30j, z=-15j:15j",
                    r=1.0, plot=False, trace_opts=None, cache=True,
                    cache_dir=None):
    """Trace a separator line from most dawnward null

    **Still in testing** Uses the bisection algorithm.

    Args:
        grid (Grid): A grid that has a "b" field
        b_slcstr (str): Some valid slice for B field
        r (float): spatial step of separator line
        plot (bool): make debugging plots
        trace_opts (dict): passed to streamline function
        cache (bool, str): Save to and load from cache, if "force",
            then don't load from cache if it exists, but do save a
            cache at the end
        cache_dir (str): Directory for cache, if None, same directory
            as that file to which the grid belongs

    Raises:
        IOError: Description

    Returns:
        tuple: (separator_lines, nulls)

          - **separator_lines** (list): list of M 3xN ndarrays that
            represent M separator lines with N points
          - **nulls** (ndarray): 3xN array of N null points
    """
    if not cache_dir:
        cache_dir = grid.find_info("_viscid_dirname", "./")
    run_name = grid.find_info("run")
    sep_fname = "{0}/{1}.sep.{2:06.0f}".format(cache_dir, run_name, grid.time)

    try:
        if isinstance(cache, string_types) and cache.strip().lower() == "force":
            raise IOError()
        with np.load(sep_fname + ".npz") as dat:
            sep_iter = (f for f in dat.files if f.startswith("arr_"))
            _it = sorted(sep_iter, key=lambda s: int(s[len("arr_"):]))
            seps = [dat[n] for n in _it]
            nulls = dat['nulls']
    except IOError:
        _b = grid['b'][b_slcstr]

        _, nulls = viscid.find_nulls(_b['x=-30j:15j'], ibound=5.0)

        # get most dawnward null, nulls2 is all nulls except p0
        nullind = np.argmin(nulls[1, :])
        p0 = nulls[:, nullind]
        nulls2 = np.concatenate([nulls[:, :nullind], nulls[:, (nullind + 1):]],
                                axis=1)

        if plot:
            from viscid.plot import vlab
            vlab.plot_earth_3d(crd_system='gse')
            vlab.points3d(nulls2[0], nulls2[1], nulls2[2],
                          color=(0, 0, 0), scale_factor=1.0)
            vlab.points3d(nulls[0, nullind], nulls[1, nullind], nulls[2, nullind],
                          color=(1, 1, 1), scale_factor=1.0)

        seed = viscid.Sphere(p0=p0, r=r, ntheta=30, nphi=60,
                             theta_endpoint=True, phi_endpoint=True)
        p1 = viscid.get_sep_pts_bisect(_b, seed, max_depth=12, plot=plot,
                                       trace_opts=trace_opts)
        # print("p1 shape", p1.shape)

        # if p1.shape[1] > 2:
        #     raise RuntimeError("Invalid B field, should be no branch @ null")

        seps = []
        sep_stubs = []
        for i in range(p1.shape[1]):
            sep_stubs.append([p0, p1[:, i]])

        # print("??", sep_stubs)

        while sep_stubs:
            sep = sep_stubs.pop(0)
            # print("!!! new stub")

            for i in count():
                # print("::", i)
                seed = viscid.SphericalPatch(p0=sep[-1], p1=sep[-1] - sep[-2],
                                             r=r, nalpha=240, nbeta=240)
                pn = viscid.get_sep_pts_bisect(_b, seed, max_depth=8, plot=plot,
                                               trace_opts=trace_opts)
                if pn.shape[1] == 0:
                    # print("END: pn.shape[1] == 0")
                    break
                # print("++", nulls2.shape, pn.shape)
                closest_null_dist = np.min(np.linalg.norm(nulls2 - pn[:, :1], axis=0))
                # print("closest_null_dist:", closest_null_dist)
                if closest_null_dist < 1.01 * r:
                    # print("END: within 1.01 of a null")
                    break

                # print("??", pn)
                for j in range(1, pn.shape[1]):
                    # print("inserting new stub")
                    sep_stubs.insert(0, [sep[-1], pn[:, j]])
                sep.append(pn[:, 0])

            # print("sep", sep)
            seps.append(np.stack(sep, axis=1))

        if cache:
            np.savez_compressed(sep_fname, *seps, nulls=nulls)
    return seps, nulls