Matplotlib：streamplot的线宽键？

Question

I am using streamplot to plot streamlines of wind, with the line width set by the wind speed. I don't want to use colour as this is to be overlaid on a filled contour plot of a different field.

Is there a way to add some kind of key or legend to indicate the magnitude associated with a particular line thickness, similar to quiverkey for quiver plots?

Answer 1

Here's an example for how you could go about making a legend yourself, using the LineCollection that is returned from streamplot . It modifies the example from the matplotlib gallery, here .

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

# Some fake data
Y, X = np.mgrid[-3:3:100j, -3:3:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U*U + V*V)

# Create you figure
fig = plt.figure()

# Create axes, ax for your plot, and lx for the legend
gs = gridspec.GridSpec(2, 2, height_ratios=(1,2), width_ratios=(4,1))
ax = fig.add_subplot(gs[:, 0])
lx = fig.add_subplot(gs[0, 1])

def speedToLW(speed):
    ''' 
    Function to convert windspeed into a sensible linewidth
    This will need to change depending on your data
    '''
    return 0.5 + speed / 5.

def LWToSpeed(lw):
    ''' The inverse of speedToLW, to get the speed back from the linewidth '''
    return (lw - 0.5) * 5.

def makeStreamLegend(strm, lx, convertFunc, nlines=5, color='k', fmt='{:g}'):

    ''' Make a legend for a streamplot on a separate axes instance '''

    # Get the linewidths from the streamplot LineCollection
    lws = np.array(strm.lines.get_linewidths())

    # Turn off axes lines and ticks, and set axis limits
    lx.axis('off')
    lx.set_xlim(0, 1)
    lx.set_ylim(0, 1)

    # Loop over the desired number of lines in the legend
    for i, y in enumerate(np.linspace(0.1, 0.9, nlines)):

        # This linewidth
        lw = lws.min()+float(i) * lws.ptp()/float(nlines-1)

        # Plot a line in the legend, of the correct length
        lx.axhline(y, 0.1, 0.4, c=color, lw=lw)

        # Add a text label, after converting the lw back to a speed
        lx.text(0.5, y, fmt.format(convertFunc(lw)), va='center')

# Make the stream plot
strm = ax.streamplot(X, Y, U, V, color='k', linewidth=speedToLW(speed))

# Add a legend, with 5 lines
makeStreamLegend(strm, lx, LWToSpeed, nlines=5, fmt='{:6.3f}')

plt.show()