Пример #1
0
def make_vector_fld():
    x, y, z = make_grid()
    f0 = viscid.empty((x, y, z), dtype=DTYPE, nr_comps=3, center='node')
    viscid.fill_dipole(f0)
    # seeds = viscid.Sphere(r=10.0, nphi=64, ntheta=32)
    seeds = viscid.Sphere(r=10.0, nphi=32, ntheta=48)
    return f0, seeds
Пример #2
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
Пример #3
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
Пример #4
0
def benchmark_streamline(precompile=True, profile=True, scale=True, plot=True):
    which = "streamline"
    print(which)
    print('-' * len(which))

    f0, seeds = make_vector_fld()

    print("Timing", which)

    sl_kwargs = dict(ibound=3.7, ds0=0.020)
    lines, _ = viscid.streamlines(f0, seeds, **sl_kwargs)
    nsegs_cython = np.sum([line.shape[1] for line in lines])
    lines = None

    ft_stats, cy_stats = dict(), dict()
    bench_output_type = viscid.OUTPUT_TOPOLOGY
    sl_kwargs.update(output=bench_output_type)

    retFT = viscid.timeit(fort_topology,
                          f0,
                          seeds,
                          timeit_repeat=6,
                          timeit_stats=ft_stats)
    _, retCY = viscid.timeit(viscid.streamlines,
                             f0,
                             seeds,
                             timeit_repeat=6,
                             timeit_stats=cy_stats,
                             **sl_kwargs)

    fort_per_seg = ft_stats['min'] / retFT.get_info('nsegs')
    cy_per_seg = cy_stats['min'] / nsegs_cython

    print("Segs Fortran", retFT.get_info('nsegs'))
    print("Segs Cython ", nsegs_cython)

    print("Fortran took {0:.3g} sec/seg".format(fort_per_seg))
    print("Cython took {0:.3g} sec/seg".format(cy_per_seg))
    print_seedup("Cython", cy_per_seg, "Fortran", fort_per_seg, prefix="@ ")

    if plot:
        from viscid.plot import mpl
        mpl.clf()
        mpl.subplot(121, projection='polar')
        mpl.plot(retCY, hemisphere='north')
        mpl.subplot(122, projection='polar')
        mpl.plot(retCY, hemisphere='south')
        mpl.show()

    if scale:
        thetas = np.logspace(np.log10(3), np.log10(144), 8).astype('i')

        cy_nsegs = [None] * len(thetas)
        fort_nsegs = [None] * len(thetas)
        cy_mintime = [None] * len(thetas)
        fort_mintime = [None] * len(thetas)

        for i, ntheta in enumerate(thetas):
            seeds = viscid.Sphere(r=10.0, ntheta=ntheta, nphi=32)
            _stats = dict()

            topo = viscid.timeit(fort_topology,
                                 f0,
                                 seeds,
                                 timeit_repeat=5,
                                 timeit_stats=_stats,
                                 timeit_quiet=True)
            fort_nsegs[i] = topo.get_info('nsegs')
            fort_mintime[i] = _stats['min']

            _, topo = viscid.timeit(viscid.calc_streamlines,
                                    f0,
                                    seeds,
                                    ibound=3.7,
                                    ds0=0.020,
                                    output=bench_output_type,
                                    timeit_repeat=5,
                                    timeit_stats=_stats,
                                    timeit_quiet=True)
            lines, _ = viscid.streamlines(f0, seeds, ibound=3.7, ds0=0.020)
            cy_nsegs[i] = np.sum([line.shape[1] for line in lines])
            cy_mintime[i] = _stats['min']

        from viscid.plot import mpl
        mpl.clf()
        mpl.plt.plot(cy_nsegs, cy_mintime, label="Cython")
        mpl.plt.plot(fort_nsegs, fort_mintime, label="Fortran")
        mpl.plt.legend(loc=0)
        mpl.plt.xlabel('Number of segments calculated')
        mpl.plt.ylabel('time to calculate')
        mpl.show()

        mpl.clf()
        cy_tperseg = np.array(cy_mintime) / np.array(cy_nsegs)
        fort_tperseg = np.array(fort_mintime) / np.array(fort_nsegs)
        mpl.plt.plot(thetas, cy_tperseg / fort_tperseg, label="over cython")
        mpl.plt.xlabel('ntheta')
        mpl.plt.ylabel('Fortran Speedup')
        mpl.show()
Пример #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 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
Пример #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
Пример #7
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