Python pyplot.streamplot方法代码示例

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def test_colormap():
    X, Y, U, V = velocity_field()
    plt.streamplot(X, Y, U, V, color=U, density=0.6, linewidth=2,
                   cmap=plt.cm.autumn)
    plt.colorbar() 

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def test_linewidth():
    X, Y, U, V = velocity_field()
    speed = np.sqrt(U*U + V*V)
    lw = 5*speed/speed.max()
    df = 25. / 30.   # Compatibility factor for old test image
    plt.streamplot(X, Y, U, V, density=[0.5 * df, 1. * df], color='k',
                   linewidth=lw) 

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def test_masks_and_nans():
    X, Y, U, V = velocity_field()
    mask = np.zeros(U.shape, dtype=bool)
    mask[40:60, 40:60] = 1
    U = np.ma.array(U, mask=mask)
    U[:20, :20] = np.nan
    with np.errstate(invalid='ignore'):
        plt.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues) 

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def test_startpoints():
    X, Y, U, V = velocity_field()
    start_x = np.linspace(X.min(), X.max(), 10)
    start_y = np.linspace(Y.min(), Y.max(), 10)
    start_points = np.column_stack([start_x, start_y])
    plt.streamplot(X, Y, U, V, start_points=start_points)
    plt.plot(start_x, start_y, 'ok') 

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def test_masks_and_nans():
    X, Y, U, V = velocity_field()
    mask = np.zeros(U.shape, dtype=bool)
    mask[40:60, 40:60] = 1
    U[:20, :20] = np.nan
    U = np.ma.array(U, mask=mask)
    with np.errstate(invalid='ignore'):
        plt.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues) 

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def test_maxlength():
    x, y, U, V = swirl_velocity_field()
    plt.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]],
                   linewidth=2, density=2) 

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def plt_stream(transform, ax, npts=200, title="Density streamflow", device="cpu"):
    side = np.linspace(LOW, HIGH, npts)
    xx, yy = np.meshgrid(side, side)
    x = np.hstack([xx.reshape(-1, 1), yy.reshape(-1, 1)])
    with torch.no_grad():
        logqx, z = transform(torch.tensor(x).float().to(device))
    d_z_x = -(x - z.cpu().numpy())[:, 0].reshape(xx.shape)
    d_z_y = -(x - z.cpu().numpy())[:, 1].reshape(xx.shape)
    plt.streamplot(xx, yy, d_z_x, d_z_y, color=(d_z_y**2 + d_z_x**2)/2, cmap='autumn') 

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def plot_vec_field(fld, field="field", fig=None, ax=None):  # pragma: no cover
    """
    Plot a spatial random vector field.

    Parameters
    ----------
    fld : :class:`Field`
        The given field class instance.
    field : :class:`str`, optional
        Field that should be plotted. Default: "field"
    fig : :class:`Figure` or :any:`None`, optional
        Figure to plot the axes on. If `None`, a new one will be created.
        Default: `None`
    ax : :class:`Axes` or :any:`None`, optional
        Axes to plot on. If `None`, a new one will be added to the figure.
        Default: `None`
    """
    if fld.mesh_type != "structured":
        raise RuntimeError(
            "Only structured vector fields are supported"
            + " for plotting. Please create one on a structured grid."
        )
    plot_field = getattr(fld, field)
    assert not (fld.pos is None or plot_field is None)

    norm = np.sqrt(plot_field[0, :].T ** 2 + plot_field[1, :].T ** 2)

    fig, ax = _get_fig_ax(fig, ax)
    title = "Field 2D " + fld.mesh_type + ": " + str(plot_field.shape)
    x, y, __ = pos2xyz(fld.pos, max_dim=2)

    sp = plt.streamplot(
        x,
        y,
        plot_field[0, :].T,
        plot_field[1, :].T,
        color=norm,
        linewidth=norm / 2,
    )
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_title(title)
    fig.colorbar(sp.lines)
    fig.show()
    return ax 

# 需要导入模块: from matplotlib import pyplot [as 别名]
# 或者: from matplotlib.pyplot import streamplot [as 别名]
def integrate_streamline(
    X, Y, U, V, integration_direction, init_states, interpolation_num=250, average=True
):
    """use streamline's integrator to alleviate stacking of the solve_ivp. Need to update with the correct time."""
    import matplotlib.pyplot as plt

    n_cell = init_states.shape[0]

    res = np.zeros((n_cell * interpolation_num, 2))
    j = -1  # this index will become 0 when the first trajectory found

    for i in tqdm(range(n_cell), "integration with streamline"):
        strm = plt.streamplot(
            X,
            Y,
            U,
            V,
            start_points=init_states[i, None],
            integration_direction=integration_direction,
            density=100,
        )
        strm_res = np.array(strm.lines.get_segments()).reshape((-1, 2))

        if len(strm_res) == 0:
            continue
        else:
            j += 1
        t = np.arange(strm_res.shape[0])
        t_linspace = np.linspace(t[0], t[-1], interpolation_num)
        f = interpolate.interp1d(t, strm_res.T)

        cur_rng = np.arange(j * interpolation_num, (j + 1) * interpolation_num)
        res[cur_rng, :] = f(t_linspace).T

    res = res[: cur_rng[-1], :]  # remove all empty trajectories
    n_cell = int(res.shape[0] / interpolation_num)

    if n_cell > 1 and average:
        t_len = len(t_linspace)
        avg = np.zeros((t_len, 2))

        for i in range(t_len):
            cur_rng = np.arange(n_cell) * t_len + i
            avg[i, :] = np.mean(res[cur_rng, :], 0)

        res = avg

    plt.close()

    return t_linspace, res


# ---------------------------------------------------------------------------------------------------
# fate related