예제 #1
0
 def __cotr__(self):
     if self.has_info("dipoletime"):
         return viscid.Cotr(time=self.find_info('dipoletime'),
                            notilt1967=True)
     elif self.has_info("basetime"):
         return viscid.Cotr(time=self.find_info('basetime'),
                            notilt1967=True)
     else:
         return NotImplemented
예제 #2
0
def check():
    """Runtime check compiled modules"""
    import os
    import sys

    import numpy as np
    import viscid

    ret = 0

    check_version()
    print()

    #####################################################
    # run streamline calculation (checks cython modules)
    try:
        cotr = viscid.Cotr(dip_tilt=15.0, dip_gsm=21.0)  # pylint: disable=not-callable
        m = cotr.get_dipole_moment(crd_system='gse')
        seeds = viscid.seed.Sphere((0.0, 0.0, 0.0),
                                   2.0,
                                   pole=-m,
                                   ntheta=25,
                                   nphi=25,
                                   thetalim=(5, 90),
                                   philim=(5, 360),
                                   phi_endpoint=False)
        B = viscid.make_dipole(m=m,
                               crd_system='gse',
                               n=(32, 32, 32),
                               l=(-25, -25, -25),
                               h=(25, 25, 25),
                               dtype='f8')
        lines, _ = viscid.calc_streamlines(B, seeds, ibound=1.0)
        for line in lines:
            if np.any(np.isnan(line)):
                raise ValueError("NaN in line")
        print("Cython module ran successfully")
    except Exception as e:
        print("Cython module has runtime errors.")
        print(str(e))
        ret |= (1 << 0)
    print()

    ####################################
    # load a jrrle file (checks fortran)
    try:
        f3d = viscid.load_file(
            os.path.join(viscid.sample_dir, 'sample_jrrle.3df.*'))
        _ = np.array(f3d['pp'].data)
        print("Fortran module ran successfully")
    except Exception as e:
        print("Fortran module has runtime errors.")
        print(str(e))
        ret |= (1 << 1)
    print()

    return ret
예제 #3
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--notwo", dest='notwo', action="store_true")
    parser.add_argument("--nothree", dest='nothree', action="store_true")
    parser.add_argument("--show", "--plot", action="store_true")
    args = viscid.vutil.common_argparse(parser, default_verb=0)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        vlt.auto_adjust_subplots()

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

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

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

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

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

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

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

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

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

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

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

    except ImportError:
        pass

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

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

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

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

    return 0
예제 #4
0
def _main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--show", "--plot", action="store_true")
    args = viscid.vutil.common_argparse(parser, default_verb=0)

    viscid.logger.setLevel(viscid.logging.DEBUG)
    args.show = False

    cotr = viscid.Cotr(dip_tilt=20.0, dip_gsm=15.0)  # pylint: disable=not-callable
    b = viscid.make_dipole(m=cotr.get_dipole_moment(), n=(32, 32, 32))

    seeds = viscid.Circle(n=5, r=1.5, pole=[0, 0, 1])
    lines, topo = viscid.calc_streamlines(b, seeds, ibound=1.4, method='rk45')

    for i in range(2):
        # make sure this works for lines with 0, 1, 2, 3 vertices
        if i == 1:
            lines[1] = lines[2][:, :0]
            lines[2] = lines[2][:, :1]
            lines[3] = lines[3][:, :2]
            lines[4] = lines[4][:, :3]

        viscid.logger.debug('---')
        viscid.logger.debug('{0}'.format(len(lines)))
        for line in lines:
            viscid.logger.debug('line shape: {0}'.format(line.shape))
        viscid.logger.debug('---')

        do_test(lines,
                scalars=None,
                txt='given None',
                show=args.show)

        do_test(lines,
                scalars='#ff0000',
                txt='given a single 24bit rgb hex color',
                show=args.show)

        do_test(lines,
                scalars='#ff000066',
                txt='given a single 32bit rgba hex color',
                show=args.show)

        do_test(lines,
                scalars='#f00',
                txt='given a single 12bit rgb hex color',
                show=args.show)

        do_test(lines,
                scalars='#f006',
                txt='given a single 16bit rgba hex color',
                show=args.show)

        do_test(lines,
                scalars=['#ff0000', '#cc0000', '#aa0000', '#880000', '#660000'],
                txt='given a list of Nlines 24bit rgb hex colors',
                show=args.show)

        do_test(lines,
                scalars=['#ff000066', '#cc000066', '#aa000066', '#88000066',
                         '#66000066'],
                txt='given a list of Nlines 32bit rgba hex colors',
                show=args.show)

        do_test(lines,
                scalars=['#f00', '#c00', '#a00', '#800', '#600'],
                txt='given a list of Nlines 12bit rgb hex colors',
                show=args.show)

        do_test(lines,
                scalars=['#f00a', '#c009', '#a008', '#8007', '#6006'],
                txt='given a list of Nlines 16bit rgba hex colors',
                show=args.show)

        do_test(lines,
                scalars=[0.8, 0.0, 0.2],
                txt='given a single rgb [0..1] color',
                show=args.show)

        do_test(lines,
                scalars=[0.8, 0.0, 0.2, 0.8],
                txt='given a single rgba [0..1] color',
                show=args.show)

        do_test(lines,
                scalars=[(0.8, 0.0, 0.2), (0.7, 0.0, 0.3), (0.6, 0.0, 0.4),
                         (0.5, 0.0, 0.5), (0.4, 0.0, 0.6)],
                txt='given a list of Nlines rgb [0..1] tuples',
                show=args.show)

        do_test(lines,
                scalars=[(0.8, 0.0, 0.2, 1.0), (0.7, 0.0, 0.3, 0.9),
                         (0.6, 0.0, 0.4, 0.8), (0.5, 0.0, 0.5, 0.7),
                         (0.4, 0.0, 0.6, 0.6)],
                txt='given a list of Nlines rgba [0..1] tuples',
                show=args.show)

        do_test(lines,
                scalars=[250, 0, 250],
                txt='given a single rgb [0..255] color',
                show=args.show)

        do_test(lines,
                scalars=[250, 0, 250, 190],
                txt='given a single rgba [0..255] color',
                show=args.show)

        do_test(lines,
                scalars=[(204, 0, 51), (179, 0, 77), (153, 0, 102),
                         (127, 0, 127), (0.4, 0, 102)],
                txt='given a list of Nlines rgb [0..255] tuples',
                show=args.show)

        do_test(lines,
                scalars=[(204, 0, 51, 255), (179, 0, 77, 230),
                         (153, 0, 102, 204), (127, 0, 127, 179),
                         (102, 0, 102, 153)],
                txt='given a list of Nlines rgba [0..255] tuples',
                show=args.show)

        do_test(lines,
                scalars=['#ff000088', 'blue', 'lavenderblush', 'c', '#4f4'],
                txt='given a mix of color hex/html color names',
                show=args.show)

        do_test(lines,
                scalars=topo,
                txt='scalars == topo value',
                show=args.show)

        do_test(lines,
                scalars=viscid.topology2color(topo),
                txt='scalars == topo2color value',
                show=args.show)

        do_test(lines,
                scalars=np.log(viscid.magnitude(b)),
                txt='given bmag',
                show=args.show)

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

    return 0
def _main():
    f = viscid.load_file('~/dev/work/xi_fte_001/*.3d.*.xdmf')
    time_slice = ':'
    times = np.array([grid.time for grid in f.iter_times(time_slice)])

    # XYZ coordinates of virtual satelites in warped "plasma sheet coords"
    x_sat_psc = np.linspace(-30, 0, 31)  # X (GSE == PSC)
    y_sat_psc = np.linspace(-10, 10, 21)  # Y (GSE == PSC)
    z_sat_psc = np.linspace(-2, 2, 5)  # Z in PSC (z=0 is the plasma sheet)

    # the GSE z location of the virtual satelites in the warped plasma sheet
    # coordinates, so sat_z_gse_ts['x=5j, y=1j, z=0j'] would give the
    # plasma sheet location at x=5.0, y=1.0
    # These fields depend on time because the plasma sheet moves in time
    sat_z_gse_ts = viscid.zeros([times, x_sat_psc, y_sat_psc, z_sat_psc],
                                crd_names='txyz', center='node',
                                name='PlasmaSheetZ_GSE')
    vx_ts = viscid.zeros_like(sat_z_gse_ts)
    bz_ts = viscid.zeros_like(sat_z_gse_ts)

    for itime, grid in enumerate(f.iter_times(time_slice)):
        print("Processing time slice", itime, grid.time)

        gse_slice = 'x=-35j:0j, y=-15j:15j, z=-6j:6j'
        bx = grid['bx'][gse_slice]
        bx_argmin = np.argmin(bx**2, axis=2)
        z_gse = bx.get_crd('z')
        # ps_zloc_gse is the plasma sheet z location along the GGCM grid x/y
        ps_z_gse = viscid.zeros_like(bx[:, :, 0:1])
        ps_z_gse[...] = z_gse[bx_argmin]

        # Note: Here you could apply a gaussian filter to
        #       ps_z_gse[:, :, 0].data in order to smooth the surface
        #       if desired. Scipy / Scikit-Image have some functions
        #       that do this

        # ok, we found the plasma sheet z GSE location on the actual GGCM
        # grid, but we just want a subset of that grid for our virtual
        # satelites, so just interpolate the ps z location to our subset
        ps_z_gse_subset = viscid.interp_trilin(ps_z_gse,
                                               sat_z_gse_ts[itime, :, :, 0:1],
                                               wrap=True)
        # now we know the plasma sheet z location in GSE, and how far
        # apart we want the satelites in z, so put those two things together
        # to get a bunch of satelite locations
        sat_z_gse_ts[itime] = ps_z_gse_subset.data + z_sat_psc.reshape(1, 1, -1)

        # make a seed generator that we can use to fill the vx and bz
        # time series for this instant in time
        sat_loc_gse = sat_z_gse_ts[itime].get_points()
        sat_loc_gse[2, :] = sat_z_gse_ts[itime].data.reshape(-1)

        # slicing the field before doing the interpolation makes this
        # faster for hdf5 data, but probably for other data too
        vx_ts[itime] = viscid.interp_trilin(grid['vx'][gse_slice],
                                            sat_loc_gse,
                                            wrap=False
                                            ).reshape(vx_ts.shape[1:])
        bz_ts[itime] = viscid.interp_trilin(grid['bz'][gse_slice],
                                            sat_loc_gse,
                                            wrap=False
                                            ).reshape(bz_ts.shape[1:])

        # 2d plots of the plasma sheet z location to make sure we did the
        # interpolation correctly
        if False:  # pylint: disable=using-constant-test
            from viscid.plot import vpyplot as vlt
            fig, (ax0, ax1) = vlt.subplots(2, 1)  # pylint: disable=unused-variable
            vlt.plot(ps_z_gse, ax=ax0, clim=(-5, 5))
            vlt.plot(ps_z_gse_subset, ax=ax1, clim=(-5, 5))
            vlt.auto_adjust_subplots()
            vlt.show()

        # make a 3d plot of the plasma sheet surface to verify that it
        # makes sense
        if True:  # pylint: disable=using-constant-test
            from viscid.plot import vlab
            fig = vlab.figure(size=(1280, 800), bgcolor=(1, 1, 1),
                              fgcolor=(0, 0, 0))
            vlab.clf()
            # plot the plasma sheet coloured by vx
            # Note: points closer to x = 0 are unsightly since the plasma
            #       sheet criteria starts to fall apart on the flanks, so
            #       just remove the first few rows
            ps_z_gse_tail = ps_z_gse['x=:-2.25j']
            ps_mesh_shape = [3, ps_z_gse_tail.shape[0], ps_z_gse_tail.shape[1]]
            ps_pts = ps_z_gse_tail.get_points().reshape(ps_mesh_shape)
            ps_pts[2, :, :] = ps_z_gse_tail[:, :, 0]
            plasma_sheet = viscid.RectilinearMeshPoints(ps_pts)
            ps_vx = viscid.interp_trilin(grid['vx'][gse_slice], plasma_sheet)
            _ = vlab.mesh_from_seeds(plasma_sheet, scalars=ps_vx)
            vx_clim = (-1400, 1400)
            vx_cmap = 'viridis'
            vlab.colorbar(title='Vx', clim=vx_clim, cmap=vx_cmap,
                          nb_labels=5)
            # plot satelite locations as dots colored by Vx with the same
            # limits and color as the plasma sheet mesh
            sat3d = vlab.points3d(sat_loc_gse[0], sat_loc_gse[1], sat_loc_gse[2],
                                  vx_ts[itime].data.reshape(-1),
                                  scale_mode='none', scale_factor=0.2)
            vlab.apply_cmap(sat3d, clim=vx_clim, cmap=vx_cmap)

            # plot Earth for reference
            cotr = viscid.Cotr(dip_tilt=0.0)  # pylint: disable=not-callable
            vlab.plot_blue_marble(r=1.0, lines=False, ntheta=64, nphi=128,
                                  rotate=cotr, crd_system='mhd')
            vlab.plot_earth_3d(radius=1.01, night_only=True, opacity=0.5,
                               crd_system='gse')
            vlab.view(azimuth=45, elevation=70, distance=35.0,
                      focalpoint=[-9, 3, -1])
            vlab.savefig('plasma_sheet_3d_{0:02d}.png'.format(itime))
            vlab.show()
            try:
                vlab.mlab.close(fig)
            except TypeError:
                pass  # this happens if the figure is already closed

    # now do what we will with the time series... this is not a good
    # presentation of this data, but you get the idea
    from viscid.plot import vpyplot as vlt
    fig, axes = vlt.subplots(4, 4, figsize=(12, 12))
    for ax_row, yloc in zip(axes, np.linspace(-5, 5, len(axes))[::-1]):
        for ax, xloc in zip(ax_row, np.linspace(4, 7, len(ax_row))):
            vlt.plot(vx_ts['x={0}j, y={1}j, z=0j'.format(xloc, yloc)], ax=ax)
            ax.set_ylabel('')
            vlt.plt.title('x = {0:g}, y = {1:g}'.format(xloc, yloc))
    vlt.plt.suptitle('Vx [km/s]')
    vlt.auto_adjust_subplots()
    vlt.show()

    return 0
예제 #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