16/07/2024

An ugly method to get matplotlib to add transparency to streamplots



Well, they said in the API docs: "This container will probably change in the future to allow changes to the colormap, alpha, etc. for both lines and arrows, but these changes should be backward compatible." but it looks like there's still a lot broken in the backend (see e.g. this GitHub issue).

So here's a very ugly workaround, based off one of the official streamplot examples
We plot some kind of velocity field, with the absolute speed in a background colour map, and alpha, width and colour mapping of the streamlines set by the speed. To get the alpha, we use an intermediate greyscale colourmap, and retroactively feed those grey values into the colour map and alpha values we want to actually use.  (TL;DR: I don't understand all the details, but it works)

from matplotlib import pyplot as plt
from matplotlib import cm
import numpy as np

fig,ax=plt.subplots(1)

w = 3
Y, X = np.mgrid[-w:w:100j, -w:w:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U**2 + V**2)
alph = (speed-speed.min())/(speed.max()-speed.min())
lw=4*alph
speed = np.sqrt(U**2 + V**2)
im=ax.imshow(speed,extent=[-w,w,-w,w],cmap=cm.bone_r)
stream=ax.streamplot(X, Y, U, V, density=0.6,linewidth=lw,color=alph,cmap=cm.gray)

#populate the colours list via draw()
fig.canvas.draw()

cols=stream.lines.get_colors()
alphas=cols[:,0] #copy alphas
cols=cm.afmhot(cols[:,0])#apply colour map
cols[:,3]=alphas #apply alphas
stream.lines.set_colors(cols)

#since the 'arrows' collection of the streamplot doesn't work,
#we access the arrow props via the axes' patch list.
for p in ax.patches:
    c=p.get_ec()#copy edge colour
    col=list(cm.afmhot(c[0]))[:3]+[c[0]]
    p.set_ec(col)
    p.set_fc(col)

plt.colorbar(im,label='speed')
plt.savefig('myfigure.png')
plt.show()

No comments: