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 _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
Ejemplo n.º 3
0
def get_mp_info(pp, b, j, e, cache=True, cache_dir=None,
                slc="x=5.5j:11.0j, y=-4.0j:4.0j, z=-3.6j:3.6j",
                fit="mp_xloc", fit_p0=(9.0, 0.0, 0.0, 1.0, -1.0, -1.0)):
    """Get info about m-pause as flattened fields

    Notes:
        The first thing this function does is mask locations where
        the GSE-y current density < 1e-4. This masks out the bow
        shock and current free regions. This works for southward IMF,
        but it is not very general.

    Parameters:
        pp (ScalarcField): pressure
        b (VectorField): magnetic field
        j (VectorField): current density
        e (VectorField, None): electric field (same centering as b). If
            None, then the info that requires E will be filled with NaN
        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
        slc (str): slice that gives a box that contains the m-pause
        fit (str): to which resulting field should the paraboloid be fit,
            defaults to mp_xloc, but pp_max_xloc might be useful in some
            circumstances
        fit_p0 (tuple): Initial guess vector for paraboloid fit

    Returns:
        dict: Unless otherwise noted, the entiries are 2D (y-z) fields

          - **mp_xloc** location of minimum abs(Bz), this works
            better than max of J^2 for FTEs
          - **mp_sheath_edge** location where Jy > 0.1 * Jy when
            coming in from the sheath side
          - **mp_sphere_edge** location where Jy > 0.1 * Jy when
            coming in from the sphere side
          - **mp_width** difference between m-sheath edge and
            msphere edge
          - **mp_shear** magnetic shear taken 6 grid points into
            the m-sheath / m-sphere
          - **pp_max** max pp
          - **pp_max_xloc** location of max pp
          - **epar_max** max e parallel
          - **epar_max_xloc** location of max e parallel
          - **paraboloid** numpy.recarray of paraboloid fit. The
            parameters are given in the 0th element, and
            the 1st element contains the 1-sigma values for the fit

    Raises:
        RuntimeError: if using MHD crds instead of GSE crds
    """
    if not cache_dir:
        cache_dir = pp.find_info("_viscid_dirname", "./")
    run_name = pp.find_info("run", None)
    if cache and run_name:
        t = pp.time
        mp_fname = "{0}/{1}.mpause.{2:06.0f}".format(cache_dir, run_name, t)
    else:
        mp_fname = ""

    try:
        force = cache.strip().lower() == "force"
    except AttributeError:
        force = False

    try:
        if force or not mp_fname or not os.path.isfile(mp_fname + ".xdmf"):
            raise IOError()

        mp_info = {}
        with viscid.load_file(mp_fname + ".xdmf") as dat:
            fld_names = ["mp_xloc", "mp_sheath_edge", "mp_sphere_edge",
                         "mp_width", "mp_shear", "pp_max", "pp_max_xloc",
                         "epar_max", "epar_max_xloc"]
            for fld_name in fld_names:
                mp_info[fld_name] = dat[fld_name]["x=0"]

    except (IOError, KeyError):
        mp_info = {}

        crd_system = viscid.as_crd_system(b, None)
        if crd_system != 'gse':
            raise RuntimeError("get_mp_info can't work in MHD crds, "
                               "switch to GSE please")

        if j.nr_patches == 1:
            pp_block = pp[slc]
            b_block = b[slc]
            j_block = j[slc]
            if e is None:
                e_block = np.nan * viscid.empty_like(j_block)
            else:
                e_block = e[slc]
        else:
            # interpolate an amr grid so we can proceed
            obnd = pp.get_slice_extent(slc)
            dx = np.min(pp.skeleton.L / pp.skeleton.n, axis=0)
            nx = np.ceil((obnd[1] - obnd[0]) / dx)
            vol = viscid.seed.Volume(obnd[0], obnd[1], nx, cache=True)
            pp_block = vol.wrap_field(viscid.interp_trilin(pp, vol),
                                      name="P").as_cell_centered()
            b_block = vol.wrap_field(viscid.interp_trilin(b, vol),
                                     name="B").as_cell_centered()
            j_block = vol.wrap_field(viscid.interp_trilin(j, vol),
                                     name="J").as_cell_centered()
            if e is None:
                e_block = np.nan * viscid.empty_like(j_block)
            else:
                e_block = vol.wrap_field(viscid.interp_trilin(e, vol),
                                         name="E").as_cell_centered()

        # jsq = viscid.dot(j_block, j_block)
        bsq = viscid.dot(b_block, b_block)

        # extract ndarrays and mask out bow shock / current free regions
        maskval = 1e-4
        jy_mask = j_block['y'].data < maskval
        masked_bsq = 1.0 * bsq
        masked_bsq.data = np.ma.masked_where(jy_mask, bsq)

        xcc = j_block.get_crd_cc('x')
        nx = len(xcc)

        mp_xloc = np.argmin(masked_bsq, axis=0)  # indices
        mp_xloc = mp_xloc.wrap(xcc[mp_xloc.data])  # location

        pp_max = np.max(pp_block, axis=0)
        pp_max_xloc = np.argmax(pp_block, axis=0)  # indices
        pp_max_xloc = pp_max_xloc.wrap(xcc[pp_max_xloc.data])  # location

        epar = viscid.project(e_block, b_block)
        epar_max = np.max(epar, axis=0)
        epar_max_xloc = np.argmax(epar, axis=0)  # indices
        epar_max_xloc = pp_max_xloc.wrap(xcc[epar_max_xloc.data])  # location

        _ret = find_mp_edges(j_block, 0.1, 0.1, maskval=maskval)
        sheath_edge, msphere_edge, mp_width, sheath_ind, sphere_ind = _ret

        # extract b and b**2 at sheath + 6 grid points and sphere - 6 grid pointns
        # clipping cases where things go outside the block. clipped ponints are
        # set to nan
        step = 6
        # extract b
        if b_block.layout == "flat":
            comp_axis = 0
            ic, _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape])
            ix = np.clip(sheath_ind + step, 0, nx - 1)
            b_sheath = b_block.data[ic, ix, iy, iz]
            ix = np.clip(sheath_ind - step, 0, nx - 1)
            b_sphere = b_block.data[ic, ix, iy, iz]
        elif b_block.layout == "interlaced":
            comp_axis = 3
            _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape[:-1]])
            ix = np.clip(sheath_ind + step, 0, nx - 1)
            b_sheath = b_block.data[ix, iy, iz]
            ix = np.clip(sheath_ind - step, 0, nx - 1)
            b_sphere = b_block.data[ix, iy, iz]
        # extract b**2
        bmag_sheath = np.sqrt(np.sum(b_sheath**2, axis=comp_axis))
        bmag_sphere = np.sqrt(np.sum(b_sphere**2, axis=comp_axis))
        costheta = (np.sum(b_sheath * b_sphere, axis=comp_axis) /
                    (bmag_sphere * bmag_sheath))
        costheta = np.where((sheath_ind + step < nx) & (sphere_ind - step >= 0),
                            costheta, np.nan)
        mp_shear = mp_width.wrap((180.0 / np.pi) * np.arccos(costheta))

        # don't bother with pretty name since it's not written to file
        # plane_crds = b_block.crds.slice_keep('x=0', cc=True)
        # fld_kwargs = dict(center="Cell", time=b.time)
        mp_width.name = "mp_width"
        mp_xloc.name = "mp_xloc"
        sheath_edge.name = "mp_sheath_edge"
        msphere_edge.name = "mp_sphere_edge"
        mp_shear.name = "mp_shear"
        pp_max.name = "pp_max"
        pp_max_xloc.name = "pp_max_xloc"
        epar_max.name = "epar_max"
        epar_max_xloc.name = "epar_max_xloc"

        mp_info = {}
        mp_info["mp_width"] = mp_width
        mp_info["mp_xloc"] = mp_xloc
        mp_info["mp_sheath_edge"] = sheath_edge
        mp_info["mp_sphere_edge"] = msphere_edge
        mp_info["mp_shear"] = mp_shear
        mp_info["pp_max"] = pp_max
        mp_info["pp_max_xloc"] = pp_max_xloc
        mp_info["epar_max"] = epar_max
        mp_info["epar_max_xloc"] = epar_max_xloc

        # cache new fields to disk
        if mp_fname:
            viscid.save_fields(mp_fname + ".h5", list(mp_info.values()))

    try:
        _paraboloid_params = fit_paraboloid(mp_info[fit], p0=fit_p0)
        mp_info["paraboloid"] = _paraboloid_params
    except ImportError as _exception:
        try:
            msg = _exception.message
        except AttributeError:
            msg = _exception.msg
        mp_info["paraboloid"] = viscid.DeferredImportError(msg)

    mp_info["mp_width"].pretty_name = "Magnetopause Width"
    mp_info["mp_xloc"].pretty_name = "Magnetopause $X_{gse}$ Location"
    mp_info["mp_sheath_edge"].pretty_name = "Magnetosheath Edge"
    mp_info["mp_sphere_edge"].pretty_name = "Magnetosphere Edge"
    mp_info["mp_shear"].pretty_name = "Magnetic Shear"
    mp_info["pp_max"].pretty_name = "Max Pressure"
    mp_info["pp_max_xloc"].pretty_name = "Max Pressure Location"
    mp_info["epar_max"].pretty_name = "Max E Parallel"
    mp_info["epar_max_xloc"].pretty_name = "Max E Parallel Location"

    return mp_info
Ejemplo n.º 4
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
Ejemplo n.º 5
0
def get_mp_info(pp,
                b,
                j,
                e,
                cache=True,
                cache_dir=None,
                slc="x=5.5f:11.0f, y=-4.0f:4.0f, z=-3.6f:3.6f",
                fit="mp_xloc",
                fit_p0=(9.0, 0.0, 0.0, 1.0, -1.0, -1.0)):
    """Get info about m-pause as flattened fields

    Notes:
        The first thing this function does is mask locations where
        the GSE-y current density < 1e-4. This masks out the bow
        shock and current free regions. This works for southward IMF,
        but it is not very general.

    Parameters:
        pp (ScalarcField): pressure
        b (VectorField): magnetic field
        j (VectorField): current density
        e (VectorField, None): electric field (same centering as b). If
            None, then the info that requires E will be filled with NaN
        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
        slc (str): slice that gives a box that contains the m-pause
        fit (str): to which resulting field should the paraboloid be fit,
            defaults to mp_xloc, but pp_max_xloc might be useful in some
            circumstances
        fit_p0 (tuple): Initial guess vector for paraboloid fit

    Returns:
        dict: Unless otherwise noted, the entiries are 2D (y-z) fields

          - **mp_xloc** location of minimum abs(Bz), this works
            better than max of J^2 for FTEs
          - **mp_sheath_edge** location where Jy > 0.1 * Jy when
            coming in from the sheath side
          - **mp_sphere_edge** location where Jy > 0.1 * Jy when
            coming in from the sphere side
          - **mp_width** difference between m-sheath edge and
            msphere edge
          - **mp_shear** magnetic shear taken 6 grid points into
            the m-sheath / m-sphere
          - **pp_max** max pp
          - **pp_max_xloc** location of max pp
          - **epar_max** max e parallel
          - **epar_max_xloc** location of max e parallel
          - **paraboloid** numpy.recarray of paraboloid fit. The
            parameters are given in the 0th element, and
            the 1st element contains the 1-sigma values for the fit

    Raises:
        RuntimeError: if using MHD crds instead of GSE crds
    """
    if not cache_dir:
        cache_dir = pp.find_info("_viscid_dirname", "./")
    run_name = pp.find_info("run", None)
    if cache and run_name:
        t = pp.time
        mp_fname = "{0}/{1}.mpause.{2:06.0f}".format(cache_dir, run_name, t)
    else:
        mp_fname = ""

    try:
        force = cache.strip().lower() == "force"
    except AttributeError:
        force = False

    try:
        if force or not mp_fname or not os.path.isfile(mp_fname + ".xdmf"):
            raise IOError()

        mp_info = {}
        with viscid.load_file(mp_fname + ".xdmf") as dat:
            fld_names = [
                "mp_xloc", "mp_sheath_edge", "mp_sphere_edge", "mp_width",
                "mp_shear", "pp_max", "pp_max_xloc", "epar_max",
                "epar_max_xloc"
            ]
            for fld_name in fld_names:
                mp_info[fld_name] = dat[fld_name]["x=0"]

    except (IOError, KeyError):
        mp_info = {}

        crd_system = viscid.as_crd_system(b, None)
        if crd_system != 'gse':
            raise RuntimeError("get_mp_info can't work in MHD crds, "
                               "switch to GSE please")

        if j.nr_patches == 1:
            pp_block = pp[slc]
            b_block = b[slc]
            j_block = j[slc]
            if e is None:
                e_block = np.nan * viscid.empty_like(j_block)
            else:
                e_block = e[slc]
        else:
            # interpolate an amr grid so we can proceed
            obnd = pp.get_slice_extent(slc)
            dx = np.min(pp.skeleton.L / pp.skeleton.n, axis=0)
            nx = np.ceil((obnd[1] - obnd[0]) / dx)
            vol = viscid.seed.Volume(obnd[0], obnd[1], nx, cache=True)
            pp_block = vol.wrap_field(viscid.interp_trilin(pp, vol),
                                      name="P").as_cell_centered()
            b_block = vol.wrap_field(viscid.interp_trilin(b, vol),
                                     name="B").as_cell_centered()
            j_block = vol.wrap_field(viscid.interp_trilin(j, vol),
                                     name="J").as_cell_centered()
            if e is None:
                e_block = np.nan * viscid.empty_like(j_block)
            else:
                e_block = vol.wrap_field(viscid.interp_trilin(e, vol),
                                         name="E").as_cell_centered()

        # jsq = viscid.dot(j_block, j_block)
        bsq = viscid.dot(b_block, b_block)

        # extract ndarrays and mask out bow shock / current free regions
        maskval = 1e-4
        jy_mask = j_block['y'].data < maskval
        masked_bsq = 1.0 * bsq
        masked_bsq.data = np.ma.masked_where(jy_mask, bsq)

        xcc = j_block.get_crd_cc('x')
        nx = len(xcc)

        mp_xloc = np.argmin(masked_bsq, axis=0)  # indices
        mp_xloc = mp_xloc.wrap(xcc[mp_xloc.data])  # location

        pp_max = np.max(pp_block, axis=0)
        pp_max_xloc = np.argmax(pp_block, axis=0)  # indices
        pp_max_xloc = pp_max_xloc.wrap(xcc[pp_max_xloc.data])  # location

        epar = viscid.project(e_block, b_block)
        epar_max = np.max(epar, axis=0)
        epar_max_xloc = np.argmax(epar, axis=0)  # indices
        epar_max_xloc = pp_max_xloc.wrap(xcc[epar_max_xloc.data])  # location

        _ret = find_mp_edges(j_block, 0.1, 0.1, maskval=maskval)
        sheath_edge, msphere_edge, mp_width, sheath_ind, sphere_ind = _ret

        # extract b and b**2 at sheath + 6 grid points and sphere - 6 grid pointns
        # clipping cases where things go outside the block. clipped ponints are
        # set to nan
        step = 6
        # extract b
        if b_block.layout == "flat":
            comp_axis = 0
            ic, _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape])
            ix = np.clip(sheath_ind + step, 0, nx - 1)
            b_sheath = b_block.data[ic, ix, iy, iz]
            ix = np.clip(sheath_ind - step, 0, nx - 1)
            b_sphere = b_block.data[ic, ix, iy, iz]
        elif b_block.layout == "interlaced":
            comp_axis = 3
            _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape[:-1]])
            ix = np.clip(sheath_ind + step, 0, nx - 1)
            b_sheath = b_block.data[ix, iy, iz]
            ix = np.clip(sheath_ind - step, 0, nx - 1)
            b_sphere = b_block.data[ix, iy, iz]
        # extract b**2
        bmag_sheath = np.sqrt(np.sum(b_sheath**2, axis=comp_axis))
        bmag_sphere = np.sqrt(np.sum(b_sphere**2, axis=comp_axis))
        costheta = (np.sum(b_sheath * b_sphere, axis=comp_axis) /
                    (bmag_sphere * bmag_sheath))
        costheta = np.where(
            (sheath_ind + step < nx) & (sphere_ind - step >= 0), costheta,
            np.nan)
        mp_shear = mp_width.wrap((180.0 / np.pi) * np.arccos(costheta))

        # don't bother with pretty name since it's not written to file
        # plane_crds = b_block.crds.slice_keep('x=0', cc=True)
        # fld_kwargs = dict(center="Cell", time=b.time)
        mp_width.name = "mp_width"
        mp_xloc.name = "mp_xloc"
        sheath_edge.name = "mp_sheath_edge"
        msphere_edge.name = "mp_sphere_edge"
        mp_shear.name = "mp_shear"
        pp_max.name = "pp_max"
        pp_max_xloc.name = "pp_max_xloc"
        epar_max.name = "epar_max"
        epar_max_xloc.name = "epar_max_xloc"

        mp_info = {}
        mp_info["mp_width"] = mp_width
        mp_info["mp_xloc"] = mp_xloc
        mp_info["mp_sheath_edge"] = sheath_edge
        mp_info["mp_sphere_edge"] = msphere_edge
        mp_info["mp_shear"] = mp_shear
        mp_info["pp_max"] = pp_max
        mp_info["pp_max_xloc"] = pp_max_xloc
        mp_info["epar_max"] = epar_max
        mp_info["epar_max_xloc"] = epar_max_xloc

        # cache new fields to disk
        if mp_fname:
            viscid.save_fields(mp_fname + ".h5", mp_info.values())

    try:
        _paraboloid_params = fit_paraboloid(mp_info[fit], p0=fit_p0)
        mp_info["paraboloid"] = _paraboloid_params
    except ImportError as _exception:
        try:
            msg = _exception.message
        except AttributeError:
            msg = _exception.msg
        mp_info["paraboloid"] = viscid.DeferredImportError(msg)

    mp_info["mp_width"].pretty_name = "Magnetopause Width"
    mp_info["mp_xloc"].pretty_name = "Magnetopause $X_{gse}$ Location"
    mp_info["mp_sheath_edge"].pretty_name = "Magnetosheath Edge"
    mp_info["mp_sphere_edge"].pretty_name = "Magnetosphere Edge"
    mp_info["mp_shear"].pretty_name = "Magnetic Shear"
    mp_info["pp_max"].pretty_name = "Max Pressure"
    mp_info["pp_max_xloc"].pretty_name = "Max Pressure Location"
    mp_info["epar_max"].pretty_name = "Max E Parallel"
    mp_info["epar_max_xloc"].pretty_name = "Max E Parallel Location"

    return mp_info
Ejemplo n.º 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.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