Different colours for arrows in quiver plot

Question

I am plotting an arrow graph and my code uses an external file as follows:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from pylab import rcParams

data=np.loadtxt(r'data.dat')

x = data[:,0] 
y = data[:,1] 
u = data[:,2] 
v = data[:,3] 


plt.quiver(x, y, u, v, angles='xy', scale_units='xy', scale=1, pivot='mid',color='g')

The data file basically looks like :

1 1 0 0
0 1 1 0
0 1 1 0
1 1 0 1

Is there a way to plot this with different colours for the different arrow directions?

Ps.: I have got a lot more arrows in my data file in a not very logical sentence like the one I am using as example.

Colleen

Answer 1

This probably do the trick:

plt.quiver(x, y, u, v, np.arctan2(v, u), angles='xy', scale_units='xy', scale=1, pivot='mid',color='g')

Note that the fifth's argument of plt.quiver is a color.

---

UPD. If you want to control the colors, you have to use colormaps. Here are a couple of examples:

Use colormap with colors parameter:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import Normalize

%matplotlib inline

ph = np.linspace(0, 2*np.pi, 13)
x = np.cos(ph)
y = np.sin(ph)
u = np.cos(ph)
v = np.sin(ph)
colors = arctan2(u, v)

norm = Normalize()
norm.autoscale(colors)
# we need to normalize our colors array to match it colormap domain
# which is [0, 1]

colormap = cm.inferno
# pick your colormap here, refer to 
# http://matplotlib.org/examples/color/colormaps_reference.html
# and
# http://matplotlib.org/users/colormaps.html
# for details
plt.figure(figsize=(6, 6))
plt.xlim(-2, 2)
plt.ylim(-2, 2)
plt.quiver(x, y, u, v, color=colormap(norm(colors)),  angles='xy', 
           scale_units='xy', scale=1, pivot='mid')

You can also stick with fifth argument like in my first example (which works in a bit different way comparing with colors) and change default colormap to control the colors.

plt.rcParams['image.cmap'] = 'Paired'

plt.figure(figsize=(6, 6))
plt.xlim(-2, 2)
plt.ylim(-2, 2)

plt.quiver(x, y, u, v, np.arctan2(v, u), angles='xy', scale_units='xy', scale=1, pivot='mid')

You can also create your own colormaps, see e.g. here.

Actin

Yeah, that did what I wanted, but how does that np.(arctan2(u,v) actually works? I am not sure what it is doing there – Colleen

Usually the y coordinate (v in this case) is passed as the first argument to arctan2, but that shouldn't matter much. @Mac, arctangent computes the angle (in radians) between the arrow vector and the positive x-axis between -pi and pi. – Gilles

As you need the color to depend on an angle of an arrow, you have to calculate the angle somehow. To this end I used function atan2 (numpy.arctan2) which uses coordinates of a vector to calculate an angle (in radians). See en.wikipedia.org/wiki/Atan2 for details. – Actin

Thanks guys! That was very helpful. If its not stressing too much, is there a way to chose the colours of the arrow, rather than getting them automatically? – Colleen

If I understand you right, you are interested in a way to control the actual colors used by matplotlib. To do so, you need to use colormap (which maps numbers to colors). I expanded an answer to provide detauls. (Note that if the answer answers you question, you can 'accept' it by clicking on green checkmark.) – Actin

@Colleen Yes of course you can choose the individual colors, passing to the keyword argument color= a Numpy array with shape (NPOINTS,4), where each individual row, corresponding to an arrow, consists of four floating number values 0 ≤ v_i ≤ 1 indicating red, green, blue and alpha of the corresponding arrow. I think there are also other possibilities but such an array is what colormap(norm(colors)) returns. – Mendicity

@Mendicity sorry, but I didn't quite catch your answer. So if I see color = 'blue', for example, how does that makes only the arrows pointing up blue, and how do I set other colours to the other arrows? Sorry again. – Colleen