python - How can I add a 2D colorbar, or a color wheel, to matplotlib?

Question

I am analyzing the magnetization mapping of a sample. After getting the gradient and its direction, I plotted them as an HSV (the direction from -π to π was mapped to Hue from 0 to 1, and Value was the normalized gradient) converted to RGB by img_rgb = mpl.colors.hsv_to_rgb(img_hsv).

I managed to add an HSV colorbar by using vmin and vmax, but this does not show the magnitude of the gradient:

plt.imshow(img_rgb, cmap='hsv', vmin=-180, vmax=180, extent=(0, 100, 0,100))
plt.xlabel('μm')
plt.ylabel('μm')
plt.colorbar()

My current plot:

Ideally, I would like to add a color wheel which encodes both the direction and the magnitude (maybe as something like a polar plot?). If that is not possible, adding a 2D plot which extends the current colorbar to include the gradient magnitude on the x-axis.

Subplots are obviously possible, but they seem like a kludge. Is there a better way?

Answer 1

First off, if you have two different parameters that you want to visualise simultaneously, you can do that by assigning two different channels to them (say red and green). This can be done by normalising your two 2d arrays and feeding them to imshow stacked similarly to this answer.

If you are content with a square-shaped 2d colormap, you can then get this colormap in the same way, by creating a meshgrid that you then again stack and feed to imshow:

from matplotlib import pyplot as plt
import numpy as np

##generating some  data
x,y = np.meshgrid(
    np.linspace(0,1,100),
    np.linspace(0,1,100),
)
directions = (np.sin(2*np.pi*x)*np.cos(2*np.pi*y)+1)*np.pi
magnitude = np.exp(-(x*x+y*y))


##normalize data:
def normalize(M):
    return (M-np.min(M))/(np.max(M)-np.min(M))

d_norm = normalize(directions)
m_norm = normalize(magnitude)

fig,(plot_ax, bar_ax) = plt.subplots(nrows=1,ncols=2,figsize=(8,4))

plot_ax.imshow(
    np.dstack((d_norm,m_norm, np.zeros_like(directions))),
    aspect = 'auto',
    extent = (0,100,0,100),
)

bar_ax.imshow(
    np.dstack((x, y, np.zeros_like(x))),
    extent = (
        np.min(directions),np.max(directions),
        np.min(magnitude),np.max(magnitude),
    ),
    aspect = 'auto',
    origin = 'lower',
)
bar_ax.set_xlabel('direction')
bar_ax.set_ylabel('magnitude')

plt.show()

The result looks like this:

In principle the same thing should also be doable with a polar Axes, but according to a comment in this github ticket, imshow does not support polar axes and I couldn't make imshow fill the entire disc.

EDIT:

Thanks to ImportanceOfBeingErnest and his answer to another question (the color keyword did it), here now a 2d colormap on a polar axis using pcolormesh. There were a few caveats, most notable, the colors dimension needs to be one smaller than the meshgrid in theta direction, otherwise the colormap has a spiral form:

fig= plt.figure(figsize=(8,4))
plot_ax = fig.add_subplot(121)
bar_ax = fig.add_subplot(122, projection = 'polar')

plot_ax.imshow(
    np.dstack((d_norm,m_norm, np.zeros_like(directions))),
    aspect = 'auto',
    extent = (0,100,0,100),
)

theta, R = np.meshgrid(
    np.linspace(0,2*np.pi,100),
    np.linspace(0,1,100),
)

t,r = np.meshgrid(
    np.linspace(0,1,99),
    np.linspace(0,1,100),
)    

image = np.dstack((t, r, np.zeros_like(r)))

color = image.reshape((image.shape[0]*image.shape[1],image.shape[2]))

bar_ax.pcolormesh(
    theta,R,
    np.zeros_like(R),
    color = color,
)

bar_ax.set_xticks(np.linspace(0,2*np.pi,5)[:-1])
bar_ax.set_xticklabels(
    ['{:.2}'.format(i) for i in np.linspace(np.min(directions),np.max(directions),5)[:-1]]
)
bar_ax.set_yticks(np.linspace(0,1,5))
bar_ax.set_yticklabels(
    ['{:.2}'.format(i) for i in np.linspace(np.min(magnitude),np.max(magnitude),5)]
)
bar_ax.grid('off')

plt.show()

This produces this figure: