Ejemplo n.º 1
0
def main():
    parser = argparse.ArgumentParser(description="Test quasi potential")
    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
    mpl.plot(xi, x=(-10, 10), y=(-10, 10), style='contourf', levels=256,
             lin=(2e-4, 1.5718))
    mpl.plot(xi, x=(-10, 10), y=(-10, 10), style='contour', colors='grey',
             levels=[0.5, 1.0])
    mpl.savefig(next_plot_fname(__file__))
    if args.show:
        mpl.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 mvi
    except ImportError:
        xfail("Mayavi not installed")

    mvi.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 = mvi.plot_lines(b_lines[inds], scalars=epar, cmap='viridis')
    mvi.mesh_from_seeds(seeds, scalars=xi, cmap='inferno')
    mvi.colorbar(s, orientation='horizontal', title=epar.pretty_name)
    # mvi.streamline(b, scalars=e, seedtype='sphere', seed_resolution=4,
    #                integration_direction='both')

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

    mvi.savefig(next_plot_fname(__file__))
    if args.show:
        mvi.show()
Ejemplo n.º 2
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 viscid.plot import mpl
        mpl.plt.clf()
        # mpl.plt.plot(seeds.get_points()[2, :], fld)
        mpl_plot_kwargs = dict()
        if interpolated_fld.is_spherical():
            mpl_plot_kwargs['hemisphere'] = 'north'
        mpl.plot(interpolated_fld, **mpl_plot_kwargs)
        mpl.plt.title(seed_name)

        mpl.plt.savefig(next_plot_fname(__file__, series='2d'))
        if show:
            mpl.plt.show()
    except ImportError:
        pass

    try:
        if not plot3d:
            raise ImportError
        from viscid.plot import mvi

        try:
            fig = _global_ns['figure']
            mvi.clf()
        except KeyError:
            fig = mvi.figure(size=[1200, 800], offscreen=not show)
            _global_ns['figure'] = fig

        try:
            mesh = mvi.mesh_from_seeds(seeds, scalars=interpolated_fld)
            mesh.actor.property.backface_culling = True
        except RuntimeError:
            pass

        pts = seeds.get_points()
        p = mvi.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data,
                         scale_mode='none', scale_factor=0.02)
        mvi.axes(p)
        mvi.title(seed_name)
        if view_kwargs:
            mvi.view(**view_kwargs)

        mvi.savefig(next_plot_fname(__file__, series='3d'))
        if show:
            mvi.show()
    except ImportError:
        pass
Ejemplo n.º 3
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 mpl

        mpl.plt.clf()

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

        mpl.plt.savefig(next_plot_fname(__file__, series="2d"))
        if show:
            mpl.plt.show()
    except ImportError:
        pass

    try:
        if not plot3d:
            raise ImportError
        from viscid.plot import mvi

        try:
            fig = _global_ns["figure"]
            mvi.clf()
        except KeyError:
            fig = mvi.figure(size=[1200, 800], offscreen=not show)
            _global_ns["figure"] = fig

        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 = mvi.mesh_from_seeds(_seeds, scalars=topo, opacity=0.6)
            mesh.actor.property.backface_culling = True
        except RuntimeError:
            pass
        mvi.plot_lines(lines, scalars=fld_mag, tube_radius=0.01, cmap="viridis")
        if title:
            mvi.title(title)

        mvi.savefig(next_plot_fname(__file__, series="3d"))
        if show:
            mvi.show()

    except ImportError:
        pass
Ejemplo n.º 4
0
def _main():
    """This _main is a faux unit-test of cotr

    It makes 3d plots using both Viscid and Mayavi
    """
    plot_mpl = True
    plot_mvi = True
    crd_system = 'gse'

    dtfmt = "%Y-%m-%d %H:%M"

    t = "2010-06-23T00:00:00.0"
    print("Dipole Moment at {0},".format(format_datetime(t, dtfmt)))
    print("    - mhd:", get_dipole_moment(t, crd_system='gse', strength=1.0))
    print("    - gse:", get_dipole_moment(t, crd_system='mhd', strength=1.0))

    if plot_mpl:
        from matplotlib import pyplot as plt
        import matplotlib.dates as mdates
        try:
            from viscid.plot import mpl
        except ImportError:
            pass

        times = linspace_datetime64("2010-01-01T00:00:00.0",
                                    "2010-06-21T00:00:00.0", n=(365//2*24))
        # times = linspace_datetime64("1952-01-01T00:00:00.0",
        #                             "1956-01-01T00:00:00.0", n=(4*365*24*2))

        t_dt = as_datetime(times).tolist()
        m_gse = np.empty((len(times), 3), dtype='f8')
        m_sm = np.empty((len(times), 3), dtype='f8')
        psi = np.empty((len(times), ), dtype='f8')
        mu = np.empty((len(times), ), dtype='f8')
        message_cadence = len(times) // 20
        for i, t in enumerate(times):
            if i % message_cadence == 0:
                print("Getting moment for: {0} {1} ({2:.0f}% complete)"
                      "".format(i, t, 100 * i / len(times)))
            c = Cotr(t)
            m_gse[i, :] = c.get_dipole_moment(crd_system='gse', strength=1)
            m_sm[i, :] = c.get_dipole_moment(crd_system='sm', strength=1)
            mu[i], psi[i] = c.get_dipole_angles()
        dip_angle = np.rad2deg(np.arccos(np.sum(m_gse * m_sm, axis=1)))

        i_smallest_diptilt = np.argmin(np.abs(mu))
        i_largest_diptilt = np.argmax(np.abs(mu))
        i_smallest_dipangle = np.argmin(np.abs(dip_angle))
        i_largest_dipangle = np.argmax(np.abs(dip_angle))

        print()
        print("Dipole info between {0} and {1},"
              "".format(format_datetime(times[0], dtfmt),
                        format_datetime(times[-1], dtfmt)))
        print("    - Smallest dipole tilt angle: {1:.03g} deg @ {0}"
              "".format(format_datetime(times[i_smallest_diptilt], dtfmt),
                        mu[i_smallest_diptilt]))
        print("    - Largest dipole tilt angle: {1:.03g} deg @ {0}"
              "".format(format_datetime(times[i_largest_diptilt], dtfmt),
                        mu[i_largest_diptilt]))
        print("    - Smallest angle between dip and GSE-Z: {1:.03g} deg @ {0}"
              "".format(format_datetime(times[i_smallest_dipangle], dtfmt),
                        dip_angle[i_smallest_dipangle]))
        print("    - Largest angle between dip and GSE-Z: {1:.03g} deg @ {0}"
              "".format(format_datetime(times[i_largest_dipangle], dtfmt),
                        dip_angle[i_largest_dipangle]))

        plt.clf()
        ax0 = plt.subplot(211)
        plt.plot(t_dt, psi, label='GSM Angle')
        plt.plot(t_dt, mu, label='DIP Tilt Angle')
        plt.gca().get_xaxis().set_visible(False)
        plt.legend(loc=0)
        plt.subplot(212, sharex=ax0)
        plt.plot(t_dt, dip_angle, label='Angle between dip and GSE-Z')
        # plt.ylim(11.62, 11.73)  # to see knee on 1954-12-14 @ 17:05
        dateFmt = mdates.DateFormatter(dtfmt)
        plt.gca().xaxis.set_major_formatter(dateFmt)
        plt.gcf().autofmt_xdate()
        plt.legend(loc=0)
        plt.subplots_adjust(left=0.15, right=0.98, top=0.97)
        plt.show()

        # m = m_gse
        # plt.clf()
        # plt.subplot(311)
        # plt.plot(times, m[:, 0], label='M$_{0}$')
        # plt.subplot(312)
        # plt.plot(times, m[:, 1], label='M$_{1}$')
        # plt.subplot(313)
        # plt.plot(times, m[:, 2], label='M$_{2}$')
        # plt.show()

    if plot_mvi:
        import os
        import viscid
        from viscid.plot import mvi

        mvi.figure(size=(768, 768), fgcolor=(0, 0, 0), bgcolor=(1, 1, 1),
                   offscreen=True)

        def _plot_time_range(times, figname):
            for i, t in enumerate(times):
                mvi.clf()
                cotr = Cotr(t)

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

                mvi.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])

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

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

                mvi.view(azimuth=0.0, elevation=90.0, distance=5.0,
                         focalpoint=[0, 0, 0])
                mvi.savefig("{0}_eq_{1:06d}.png".format(figname, i))
                mvi.view(azimuth=0.0, elevation=0.0, distance=5.0,
                         focalpoint=[0, 0, 0])
                mvi.savefig("{0}_pole_{1:06d}.png".format(figname, i))

        path = os.path.expanduser("~/Desktop/day/{0}/".format(crd_system))
        if not os.path.exists(path):
            os.makedirs(path)
        print()
        print("Writing 3D images in:", path)

        # summer solstice (earth-sun line @ tropic of cancer)
        print()
        print("Making a movie for the June solstice")
        solstice_times = linspace_datetime64("2010-06-21T00:00:00.0",
                                             "2010-06-22T00:00:00.0", n=49)
        _plot_time_range(solstice_times, path + "solstice")
        pfx = path + "solstice_eq"
        viscid.vutil.make_animation(pfx + '.mp4', pfx, framerate=23.976, yes=1)
        pfx = path + "solstice_pole"
        viscid.vutil.make_animation(pfx + '.mp4', pfx, framerate=23.976, yes=1)

        # autumnal equinox (earth-sun line @ equator)
        print()
        print("Making a movie for the September equinox")
        equinox_times = linspace_datetime64("2010-09-23T00:00:00.0",
                                            "2010-09-24T00:00:00.0", n=49)
        _plot_time_range(equinox_times, path + "equinox")
        pfx = path + "equinox_eq"
        viscid.vutil.make_animation(pfx + '.mp4', pfx, framerate=23.976, yes=1)
        pfx = path + "equinox_pole"
        viscid.vutil.make_animation(pfx + '.mp4', pfx, framerate=23.976, yes=1)

        # Watching the magnetic pole move year-by-year
        print()
        print("Making a movie to show magnetic pole motion")
        years = range(1950, 2016, 1)
        times = [as_datetime64("{0:04d}-06-21".format(y)) for y in years]
        _plot_time_range(times, path + "year")
        pfx = path + "year_eq"
        viscid.vutil.make_animation(pfx + '.mp4', pfx, yes=1, framerate=8)
        pfx = path + "year_pole"
        viscid.vutil.make_animation(pfx + '.mp4', pfx, yes=1, framerate=8)
Ejemplo n.º 5
0
def main():
    parser = argparse.ArgumentParser(description="Test calc")
    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(sample_dir + '/sample_xdmf.3d.[0].xdmf')
    f_iono = viscid.load_file(sample_dir + "/sample_xdmf.iof.[0].xdmf")

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

    mvi.figure(size=(1200, 800), offscreen=not args.show)

    ##########################################################
    # make b a dipole inside 3.1Re and set e = 0 inside 4.0Re
    cotr = viscid.Cotr(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 = mvi.field2source(pp, center='node')
    scp = mvi.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]
    cbar = mvi.colorbar(scp, title=pp.name, orientation='vertical')

    ######################################
    # plot a vector cut plane of the flow
    vcp = mvi.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
    iso = mvi.iso_surface(pp_src, contours=5, opacity=0.1, cmap=False)

    ##############################################################
    # calculate B field lines && topology in Viscid and plot them
    seeds = viscid.SphericalPatch([0, 0, 0], [2, 0, 1], 30, 15, r=5.0,
                                  nalpha=5, nbeta=5)
    b_lines, topo = viscid.calc_streamlines(b, seeds, ibound=3.5,
                                            obound0=[-25, -20, -20],
                                            obound1=[15, 20, 20], wrap=True)
    mvi.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 = mvi.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 = mvi.streamline(b, epar, seedtype='sphere', 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.center = [-11, 0, 0]
    bsl2.seed.widget.radius = 1.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 = mvi.colorbar(bsl2, title=epar.name, orientation='horizontal')
    cbar.scalar_bar_representation.position = (0.2, 0.01)
    cbar.scalar_bar_representation.position2 = (0.6, 0.14)

    ###############################################################
    # 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 = mvi.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

    ####################################################################
    # 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 = mvi.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
    mvi.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
    mvi.plot_earth_3d(radius=1.01, night_only=True, opacity=0.5, crd_system=b)

    ####################
    # Finishing Touches
    # mvi.axes(pp_src, nb_labels=5)
    oa = mvi.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
    # mvi.resize([1200, 800])
    mvi.view(azimuth=45, elevation=70, distance=35.0, focalpoint=[-2, 0, 0])

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

    # print("saving png")
    # mvi.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 mvi.plot_ionosphere
    # #            causes a segfault when saving x3d files
    # #
    # mvi.savefig('mayavi_msphere_sample.x3d')
    # print("done")

    mvi.savefig(next_plot_fname(__file__))

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

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