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

user_8448554 0
Viewed 3323

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:
enter image description here

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?

matplotlib python
2 Answers

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.

Resulting image

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()
user_59343160
Related