15

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.

Mac
  • 991
  • 3
  • 11
  • 22

1 Answers1

30

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. colored arrows

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')

different colormap

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')

Paired colormap

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

Community
  • 1
  • 1
Ilya V. Schurov
  • 7,687
  • 2
  • 40
  • 78
  • 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 – Mac Oct 13 '16 at 17:30
  • 2
    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. – wflynny Oct 13 '16 at 17:55
  • 2
    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 https://en.wikipedia.org/wiki/Atan2 for details. – Ilya V. Schurov Oct 13 '16 at 17:56
  • 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? – Mac Oct 14 '16 at 11:27
  • 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.) – Ilya V. Schurov Oct 14 '16 at 12:46
  • @Mac 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. – gboffi May 31 '17 at 13:57
  • @gboffi 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. – Mac Jun 28 '17 at 17:16