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:

Answer 2

I am having a similar problem when trying to visualize the radial and absolute components of a surface gradient.

I am converting the absolute value of the gradient plus the angle to a color via hsv (using the hue as the angle and the saturation and the value as the absolute value). This is the same as in magnetization plots as any vector field can be used replacing the gradient. The following function illustrates the idea. The full code is provided in the end of the answer.

import matplotlib.colors

# gradabs is the absolute gradient value, 
# gradang is the angle direction, z the vector field
# the gradient was calculated of

max_abs = np.max(gradabs) 

def grad_to_rgb(angle, absolute):
    """Get the rgb value for the given `angle` and the `absolute` value

    Parameters
    ----------
    angle : float
        The angle in radians
    absolute : float
        The absolute value of the gradient

    Returns
    -------
    array_like
        The rgb value as a tuple with values [0..1]
    """
    global max_abs

    # normalize angle
    angle = angle % (2 * np.pi)
    if angle < 0:
        angle += 2 * np.pi

    return matplotlib.colors.hsv_to_rgb((angle / 2 / np.pi, 
                                         absolute / max_abs, 
                                         absolute / max_abs))

# convert to colors via hsv
grad = np.array(list(map(grad_to_rgb, gradang.flatten(), gradabs.flatten())))

# reshape
grad = grad.reshape(tuple(list(z.shape) + [3]))

The resulting graph is the following.

---

The full example code on showing the gradient field of a surface:

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

r = np.linspace(0, np.pi, num=100)
x, y = np.meshgrid(r, r)
z = np.sin(y) * np.cos(x)

fig = plt.figure()

ax = fig.add_subplot(1, 3, 1, projection='3d')
ax.plot_surface(x, y, z)
# ax.imshow(z)
ax.set_title("Surface")

ax = fig.add_subplot(1, 3, 2)
ax.set_title("Gradient")

# create gradient
grad_y, grad_x = np.gradient(z)

# calculate length
gradabs = np.sqrt(np.square(grad_x) + np.square(grad_y))
max_abs = np.max(gradabs)

# calculate angle component
gradang = np.arctan2(grad_y, grad_x)

def grad_to_rgb(angle, absolute):
    """Get the rgb value for the given `angle` and the `absolute` value

    Parameters
    ----------
    angle : float
        The angle in radians
    absolute : float
        The absolute value of the gradient
    
    Returns
    -------
    array_like
        The rgb value as a tuple with values [0..1]
    """
    global max_abs

    # normalize angle
    angle = angle % (2 * np.pi)
    if angle < 0:
        angle += 2 * np.pi

    return matplotlib.colors.hsv_to_rgb((angle / 2 / np.pi, 
                                         absolute / max_abs, 
                                         absolute / max_abs))

# convert to colors via hsv
grad = np.array(list(map(grad_to_rgb, gradang.flatten(), gradabs.flatten())))

# reshape
grad = grad.reshape(tuple(list(z.shape) + [3]))

ax.imshow(grad)

n = 5
gx, gy = np.meshgrid(np.arange(z.shape[0] / n), np.arange(z.shape[1] / n))
ax.quiver(gx * n, gy * n, grad_x[::n, ::n], grad_y[::n, ::n])

# plot color wheel
# Generate a figure with a polar projection, inspired by
# https://stackoverflow.com/a/48253413/5934316
ax = fig.add_subplot(1, 3, 3, projection='polar')

n = 200  # the number of secants for the mesh
t = np.linspace(0, 2 * np.pi, n)
r = np.linspace(0, max_abs, n)
rg, tg = np.meshgrid(r, t)
c = np.array(list(map(grad_to_rgb, tg.T.flatten(), rg.T.flatten())))
cv = c.reshape((n, n, 3))

m = ax.pcolormesh(t, r, cv[:,:,1], color=c, shading='auto')
m.set_array(None)
ax.set_yticklabels([])

plt.show()