Color matplotlib quiver field according to magnitude and direction

Question

I'm attempting to achieve the same behavior as this function in Matlab, whereby the color of each arrow corresponds to both its magnitude and direction, essentially drawing its color from a wheel. I saw this question, but it only seems to work for barbs. I also saw this answer, but quiver complains that the color array must be two-dimensional.

What is the best way to compute C for matplotlib.pyplot.quiver, taking into account both magnitude and direction?

Answer 1

Even though this is quite old now, I've come across the same problem. Based on matplotlibs quiver demo and my own answer to this post, I created the following example. The idea is to convert the angle of a vector to the color using HSV colors Hue value. The absolute value of the vector is used as the saturation and the value.

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

def vector_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))

X = np.arange(-10, 10, 1)
Y = np.arange(-10, 10, 1)
U, V = np.meshgrid(X, Y)

angles = np.arctan2(V, U)
lengths = np.sqrt(np.square(U) + np.square(V))

max_abs = np.max(lengths)
c = np.array(list(map(vector_to_rgb, angles.flatten(), lengths.flatten())))

fig, ax = plt.subplots()
q = ax.quiver(X, Y, U, V, color=c)

plt.show()

The color wheel is the following. The code for generating it is mentioned in the Edit.

---

Edit

I just noticed, that the linked matlab function "renders a vector field as a grid of unit-length arrows. The arrow direction indicates vector field direction, and the color indicates the magnitude". So my above example is not really what is in the question. Here are some modifications.

The left graph is the same as above. The right one does, what the cited matlab function does: A unit-length arrow plot with the color indicating the magnitude. The center one does not use the magnitude but only the direction in the color which might be useful too. I hope other combinations are clear from this example.

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

def vector_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))

X = np.arange(-10, 10, 1)
Y = np.arange(-10, 10, 1)
U, V = np.meshgrid(X, Y)

angles = np.arctan2(V, U)
lengths = np.sqrt(np.square(U) + np.square(V))
max_abs = np.max(lengths)

# color is direction, hue and value are magnitude
c1 = np.array(list(map(vector_to_rgb, angles.flatten(), lengths.flatten())))

ax = plt.subplot(131)
ax.set_title("Color is lenth,\nhue and value are magnitude")
q = ax.quiver(X, Y, U, V, color=c1)

# color is length only
c2 = np.array(list(map(vector_to_rgb, angles.flatten(), 
                                      np.ones_like(lengths.flatten()) * max_abs)))

ax = plt.subplot(132)
ax.set_title("Color is direction only")
q = ax.quiver(X, Y, U, V, color=c2)

# color is direction only
c3 = np.array(list(map(vector_to_rgb, 2 * np.pi * lengths.flatten() / max_abs, 
                                      max_abs * np.ones_like(lengths.flatten()))))

# create one-length vectors
U_ddash = np.ones_like(U)
V_ddash = np.zeros_like(V)
# now rotate them
U_dash = U_ddash * np.cos(angles) - V_ddash * np.sin(angles)
V_dash = U_ddash * np.sin(angles) + V_ddash * np.cos(angles)

ax = plt.subplot(133)
ax.set_title("Uniform length,\nColor is magnitude only")
q = ax.quiver(X, Y, U_dash, V_dash, color=c3)

plt.show()

---

To plot the color wheel use the following code. Note that this uses the max_abs value from above which is the maximum value that the color hue and value can reach. The vector_to_rgb() function is also re-used here.

ax = plt.subplot(236, projection='polar')

n = 200
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(vector_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([])

Answer 2

I don't know if you've since found that quiver with matplotlib 1.4.x has 3d capability. This capability is limited when attempting to colour the arrows however.

A friend and I write the following script (in half an hour or so) to plot my experiment data using hex values from a spreadsheet, for my thesis. We're going to make this more automated once we're done with the semester but the issue with passing a colour map to quiver is that it can't accept a vector form for some reason.

This link is to my git repository where the code I used, slightly neatened up by another friend, is hosted.

I hope I can save someone the time it took me.