Ejemplo n.º 1
0
def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="",
             view_kwargs=None, show=False):
    interpolated_fld = viscid.interp_trilin(fld, seeds)
    seed_name = seeds.__class__.__name__
    if add_title:
        seed_name += " " + add_title

    try:
        if not plot2d:
            raise ImportError
        from viscid.plot import mpl
        mpl.plt.clf()
        # mpl.plt.plot(seeds.get_points()[2, :], fld)
        mpl_plot_kwargs = dict()
        if interpolated_fld.is_spherical():
            mpl_plot_kwargs['hemisphere'] = 'north'
        mpl.plot(interpolated_fld, **mpl_plot_kwargs)
        mpl.plt.title(seed_name)

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

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

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

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

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

        mvi.savefig(next_plot_fname(__file__, series='3d'))
        if show:
            mvi.show()
    except ImportError:
        pass
Ejemplo n.º 2
0
def _main():
    f = viscid.load_file("$WORK/xi_fte_001/*.3d.[4050f].xdmf")
    mp = get_mp_info(f['pp'], f['b'], f['j'], f['e_cc'], fit='mp_xloc',
                     slc="x=6.5f:10.5f, y=-4f:4f, z=-4.8f:3f", cache=False)

    y, z = mp['pp_max_xloc'].meshgrid_flat(prune=True)
    x = mp['pp_max_xloc'].data.reshape(-1)

    Y, Z = mp['pp_max_xloc'].meshgrid(prune=True)
    x2 = paraboloid(Y, Z, *mp['paraboloid'][0])

    skip = 117
    n = paraboloid_normal(Y, Z, *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]

    minvar_y = Y.reshape(-1)[::skip]
    minvar_z = Z.reshape(-1)[::skip]
    minvar_n = np.zeros([3, len(minvar_y)])
    for i in range(minvar_n.shape[0]):
        p0 = [0.0, minvar_y[i], minvar_z[i]]
        p0[0] = mp['pp_max_xloc']['y={0[0]}f, z={0[1]}f'.format(p0)]
        minvar_n[:, i] = viscid.find_minvar_lmn_around(f['b'], p0, l=2.0, n=64)[2, :]

    # 2d plots, normals don't look normal in the matplotlib projection
    if False:  # pylint: disable=using-constant-test
        from viscid.plot import mpl

        normals = paraboloid_normal(Y, Z, *mp['paraboloid'][0])
        p0 = np.array([x2, Y, Z]).reshape(3, -1)
        p1 = p0 + normals.reshape(3, -1)

        mpl.scatter_3d(np.vstack([x, y, z])[:, ::skip], equal=True)
        for i in range(0, p0.shape[1], skip):
            mpl.plt.gca().plot([p0[0, i], p1[0, i]],
                               [p0[1, i], p1[1, i]],
                               [p0[2, i], p1[2, i]], color='c')
        # z2 = _ellipsiod(X, Y, *popt)
        mpl.plt.gca().plot_surface(Y, Z, x2, color='r')
        mpl.show()

    # mayavi 3d plots, normals look better here
    if True:  # pylint: disable=using-constant-test
        from viscid.plot import mvi
        mvi.points3d(x[::skip], y[::skip], z[::skip], scale_factor=0.25,
                     color=(0.0, 0.0, 1.0))

        mp_width = mp['mp_width']['x=0']
        mp_sheath_edge = mp['mp_sheath_edge']['x=0']
        mp_sphere_edge = mp_sheath_edge - mp_width

        mvi.mesh(x2, Y, Z, scalars=mp_width.data)
        mvi.mesh(mp_sheath_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)
        mvi.mesh(mp_sphere_edge.data, Y, Z, opacity=0.75, color=(0.75, ) * 3)

        n = paraboloid_normal(Y, Z, *mp['paraboloid'][0]).reshape(3, -1)[:, ::skip]
        mvi.quiver3d(x2.reshape(-1)[::skip],
                     Y.reshape(-1)[::skip],
                     Z.reshape(-1)[::skip],
                     n[0], n[1], n[2], color=(1, 0, 0))
        mvi.quiver3d(x2.reshape(-1)[::skip],
                     Y.reshape(-1)[::skip],
                     Z.reshape(-1)[::skip],
                     minvar_n[0], minvar_n[1], minvar_n[2], color=(0, 0, 1))
        mvi.show()
Ejemplo n.º 3
0
def _get_sep_pts_bisect(
    fld,
    seed,
    trace_opts=None,
    min_depth=3,
    max_depth=7,
    plot=False,
    perimeter_check=perimeter_check_bitwise_or,
    make_3d=True,
    start_uneven=False,
    _base_quadrent="",
    _uneven_mask=0,
    _first_recurse=True,
):
    if len(_base_quadrent) == max_depth:
        return [_base_quadrent]  # causes pylint to complain
    if trace_opts is None:
        trace_opts = dict()

    nx, ny = seed.uv_shape
    (xlim, ylim) = seed.uv_extent

    if _first_recurse and start_uneven:
        _uneven_mask = UNEVEN_MASK

    if _first_recurse and plot:
        from viscid.plot import mvi
        from viscid.plot import mpl

        mpl.clf()
        _, all_topo = viscid.calc_streamlines(fld, seed, **trace_opts)
        mpl.plot(np.bitwise_and(all_topo, 15), show=False)
        verts, arr = seed.wrap_mesh(all_topo.data)
        mvi.mesh(verts[0], verts[1], verts[2], scalars=arr, opacity=0.75)

    # quadrents and lines are indexed as follows...
    # directions are counter clackwise around the quadrent with
    # lower index (which matters for lines which are shared among
    # more than one quadrent, aka, lines 1,2,6,7). Notice that even
    # numbered lines are horizontal, like the interstate system :)
    # -<--10-----<-8---
    # |       ^       ^
    # 11  2   9   3   7
    # \/      |       |
    # --<-2-----<-6----
    # |       ^       ^
    # 3   0   1   1   5
    # \/      |       |
    # ----0->-----4->--

    # find low(left), mid(center), and high(right) crds in x and y
    low_quad = "{0}{1:x}".format(_base_quadrent, 0 | _uneven_mask)
    high_quad = "{0}{1:x}".format(_base_quadrent, 3 | _uneven_mask)
    xl, xm, yl, ym = _quadrent_limits(low_quad, xlim, ylim)
    _, xh, _, yh = _quadrent_limits(high_quad, xlim, ylim)
    segsx, segsy = [None] * 12, [None] * 12
    topo = [None] * 12
    nxm, nym = nx // 2, ny // 2

    # make all the line segments
    segsx[0], segsy[0] = np.linspace(xl, xm, nxm), np.linspace(yl, yl, nxm)
    segsx[1], segsy[1] = np.linspace(xm, xm, nym), np.linspace(yl, ym, nym)
    segsx[2], segsy[2] = np.linspace(xm, xl, nxm), np.linspace(ym, ym, nxm)
    segsx[3], segsy[3] = np.linspace(xl, xl, nym), np.linspace(ym, yl, nym)

    segsx[4], segsy[4] = np.linspace(xm, xh, nxm), np.linspace(yl, yl, nxm)
    segsx[5], segsy[5] = np.linspace(xh, xh, nym), np.linspace(yl, ym, nym)
    segsx[6], segsy[6] = np.linspace(xh, xm, nxm), np.linspace(ym, ym, nxm)

    segsx[7], segsy[7] = np.linspace(xh, xh, nym), np.linspace(ym, yh, nym)
    segsx[8], segsy[8] = np.linspace(xh, xm, nxm), np.linspace(yh, yh, nxm)
    segsx[9], segsy[9] = np.linspace(xm, xm, nym), np.linspace(ym, yh, nym)

    segsx[10], segsy[10] = np.linspace(xm, xl, nxm), np.linspace(yh, yh, nxm)
    segsx[11], segsy[11] = np.linspace(xl, xl, nym), np.linspace(yh, ym, nym)

    allx = np.concatenate(segsx)
    ally = np.concatenate(segsy)

    # print("plot::", _base_quadrent, '|', _uneven_mask, '|', len(allx), len(ally))

    pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
    _, all_topo = viscid.calc_streamlines(fld, pts3d, **trace_opts)

    topo[0] = all_topo[: len(segsx[0])]
    cnt = len(topo[0])
    for i, segx in zip(count(1), segsx[1:]):
        topo[i] = all_topo[cnt : cnt + len(segx)]
        # print("??", i, cnt, cnt + len(segx), np.bitwise_and.reduce(topo[i]))
        cnt += len(topo[i])

    # assemble the lines into the four quadrents
    quad_topo = [None] * 4

    # all arrays snip off the last element since those are
    # duplicated by the next line... reversed arrays do the
    # snipping with -1:0:-1
    quad_topo[0] = np.concatenate([topo[0][:-1], topo[1][:-1], topo[2][:-1], topo[3][:-1]])

    quad_topo[1] = np.concatenate([topo[4][:-1], topo[5][:-1], topo[6][:-1], topo[1][-1:0:-1]])

    quad_topo[2] = np.concatenate([topo[2][-1:0:-1], topo[9][:-1], topo[10][:-1], topo[11][:-1]])

    quad_topo[3] = np.concatenate([topo[6][-1:0:-1], topo[7][:-1], topo[8][:-1], topo[9][-1:0:-1]])

    # now that the quad arrays are populated, decide which quadrents
    # still contain the separator (could be > 1)
    required_uneven_subquads = False
    ret = []
    for i in range(4):
        if perimeter_check(quad_topo[i]):
            next_quad = "{0}{1:x}".format(_base_quadrent, i | _uneven_mask)
            subquads = _get_sep_pts_bisect(
                fld,
                seed,
                trace_opts=trace_opts,
                min_depth=min_depth,
                max_depth=max_depth,
                plot=plot,
                _base_quadrent=next_quad,
                _uneven_mask=0,
                _first_recurse=False,
            )
            ret += subquads

    if len(ret) == 0:
        perimeter = np.concatenate(
            [
                topo[0][::-1],
                topo[4][::-1],
                topo[5][::-1],
                topo[7][::-1],
                topo[8][::-1],
                topo[10][::-1],
                topo[11][::-1],
                topo[3][::-1],
            ]
        )
        if _uneven_mask:
            if len(_base_quadrent) > min_depth:
                print("sep trace issue, but min depth reached: {0} > {1}" "".format(len(_base_quadrent), min_depth))
                ret = [_base_quadrent]
            else:
                print("sep trace issue, the separator ended prematurely")
        elif perimeter_check(perimeter):
            ret = _get_sep_pts_bisect(
                fld,
                seed,
                trace_opts=trace_opts,
                min_depth=min_depth,
                max_depth=max_depth,
                plot=plot,
                _base_quadrent=_base_quadrent,
                _uneven_mask=UNEVEN_MASK,
                _first_recurse=False,
            )
            required_uneven_subquads = True

    if plot and not required_uneven_subquads:
        from viscid.plot import mvi
        from viscid.plot import mpl

        _pts3d = seed.to_3d(seed.uv_to_local(np.array([allx, ally])))
        mvi.points3d(_pts3d[0], _pts3d[1], _pts3d[2], all_topo.data.reshape(-1), scale_mode="none", scale_factor=0.02)
        mpl.plt.scatter(
            allx, ally, color=np.bitwise_and(all_topo, 15), vmin=0, vmax=15, marker="o", edgecolor="y", s=40
        )

    if _first_recurse:
        # turn quadrent strings into locations
        xc = np.empty(len(ret))
        yc = np.empty(len(ret))
        for i, r in enumerate(ret):
            xc[i], yc[i] = _quadrent_center(r, xlim, ylim)
        pts_uv = np.array([xc, yc])
        if plot:
            from viscid.plot import mvi
            from viscid.plot import mpl

            mpl.plt.plot(pts_uv[0], pts_uv[1], "y*", ms=20, markeredgecolor="k", markeredgewidth=1.0)
            mpl.show(block=False)
            mvi.show(stop=True)
        # return seed.to_3d(seed.uv_to_local(pts_uv))
        # if pts_uv.size == 0:
        #     return None
        if make_3d:
            return seed.uv_to_3d(pts_uv)
        else:
            return pts_uv
    else:
        return ret
Ejemplo n.º 4
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 mvi

            mvi.plot_earth_3d(crd_system="gse")
            mvi.points3d(nulls2[0], nulls2[1], nulls2[2], color=(0, 0, 0), scale_factor=1.0)
            mvi.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