util.py 文件源码

def colorline(x, y, z=None, cmap=plt.get_cmap('Spectral_r'), cmin=None, cmax=None, lw=3):
    '''
    Plot a colored line with coordinates x and y
    Optionally specify colors in the array z
    Optionally specify a colormap, a norm function and a line width
    '''

    cmap = copy(cmap)
    cmap.set_over('k')
    cmap.set_under('k')

    # Default colors equally spaced on [0,1]:
    if z is None:
        z = np.linspace(0.0, 1.0, len(x))

    # Special case if a single number:
    if not hasattr(z, "__iter__"):  # to check for numerical input -- this is a hack
        z = np.array([z])

    z = np.asarray(z)

    segments = make_segments(x, y)
    return LineCollection(segments, array=z, cmap=cmap, norm=plt.Normalize(vmin=cmin, vmax=cmax), linewidth=lw)