core.py 文件源码

def colorline(x, y, cmap=None, cm_range=(0, 0.7), **kwargs):
    """Colorline plots a trajectory of (x,y) points with a colormap"""

    # plt.plot(x, y, '-k', zorder=1)
    # plt.scatter(x, y, s=40, c=plt.cm.RdBu(np.linspace(0,1,40)), zorder=2, edgecolor='k')

    assert len(cm_range)==2, "cm_range must have (min, max)"
    assert len(x) == len(y), "x and y must have the same number of elements!"

    ax = kwargs.get('ax', plt.gca())
    lw = kwargs.get('lw', 2)
    if cmap is None:
        cmap=plt.cm.Blues_r

    t = np.linspace(cm_range[0], cm_range[1], len(x))

    points = np.array([x, y]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)

    lc = LineCollection(segments, cmap=cmap, norm=plt.Normalize(0, 1),
                        zorder=50)
    lc.set_array(t)
    lc.set_linewidth(lw)

    ax.add_collection(lc)

    return lc