I have 3D plot in quiver, I need to pass colour list of strings to specify the colours of arrows. How can this be done?
e.g colour[i,j,k]= 'red'or 'blue'
from matplotlib.cm import colors
ax = plt.figure(figsize=(15,15)).add_subplot(projection='3d')
ax.quiver(x, y, z, u, v, w, color = colour, length=0.1, normalize=True)
You need to be careful how to create the color list for quivers. Read the comment in the code below.
import numpy as np
import matplotlib.pyplot as plt
from itertools import chain
x, y = np.mgrid[-5:5:5j, -5:5:5j]
z = np.zeros_like(x)
u, v, w = x, y, np.ones_like(x)
# each quiver is composed of three lines:
# * one line for the stem.
# * two lines for the head.
# Let's say you have N quivers. To apply the same color
# to each quiver (stem + head), the `colors` list must be:
# colors = ["c_stem1", ..., "cstem_N", "chead_1", "chead_1", ..., "chead_N", "chead_N"]
get_color = lambda i: ["r", "g", "b"][i % 3]
stem_colors = [get_color(i) for i, t in enumerate(u.ravel())]
head_colors = [(c, c) for c in stem_colors]
colors = stem_colors + list(chain(*head_colors))
fig = plt.figure()
ax = fig.add_subplot(projection="3d")
ax.quiver(x, y, z, u, v, w, colors=colors, normalize=True)
ax.set_zlim(-1, 2)
plt.show()
