CODEWITHSUNDEEP
pythonpython-3.xmatplotlibmatplotlib-3d
Job Stancil
Job Stancil

Reputation: 143

How to get same color for arrow shaft and arrow head with 3d quiver

The program below plots 3d arrows. I set the color for arrows depending on elevation (z coordinate).

WHAT WORKS: Arrow shafts get the correct color

PROBLEM: Arrow heads (whole or part) get a different color. Ugly.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # noqa: F401 unused import
import matplotlib.cm as cm

n = 5
n3 = n*n*n
u, v, w = np.arange(n3), np.arange(n3), np.arange(n3)  # Positions
x, y, z = np.arange(n3), np.arange(n3), np.arange(n3)  # Vector components
colors = [(0, 0, 255)]*n3
mid = n/2
for i in range(n):
    for j in range(n):
        for k in range(n):
            ii = i*n*n + j*n + k
            u[ii], v[ii], w[ii] = -(j-mid), i-mid, 0   # Whorl
            x[ii], y[ii], z[ii] = i, j, k
            colors[ii] = (1, 0, 0)
            if abs(k-(n-1)/2) < 1.1: colors[ii] = (0, 1, 0)
            if abs(k-(n-1)/2) < 0.1: colors[ii] = (0, 0, 1)

figure = plt.figure()
axes = figure.gca(projection='3d')  # gca = get/create current axes

q = axes.quiver(x, y, z, u, v, w, color=colors, length = 0.5, normalize=True)
plt.show()  # Finally display the plot

Output: Output from the code above

Upvotes: 2

Views: 2298

Answers (2)

athantas
athantas

Reputation: 53

That's correct, but for me, it only worked after I added the same colors list cs, to the edgecolor variable; that is: ax.quiver(...., color=cs, edgecolor=cs)

Upvotes: 0

hsl
hsl

Reputation: 681

I think the problem is you didn't specify the color of the arrowheads, so matplotlib is cycling through your colors variable to color the arrowheads.

After your for loop, use the code below instead to fix the color of the arrowheads. Each vector is composed of 3 elements: a line for the body fo the vector plus two lines for the arrowhead. So if you have n vectors, you'll need to provide n*3 colors to specify the colors properly.

Let's say you're plotting 2 vectors. Here's how the colors for the vector elements should be sorted in the array of colors: line1col, line2col, l1arrow1col, l1arrow2col, l2arrow1col, l2arrow2col, l3arrow1col, l3arrow2col.

# fix arrowhead colors
cs = colors[:]
for c in colors:
    cs.append(c)
    cs.append(c)

figure = plt.figure(figsize=(13, 13))
axes = figure.gca(projection="3d")  # gca = get/create current axes
q = axes.quiver(x, y, z, u, v, w, color=cs, length=0.5, normalize=True)
plt.show()  # Finally display the plot

Here's a simpler example with 3 vectors that illustrates how to color the vectors.

fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.quiver(
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0],
    [-3, 0, 0],
    [-4, 4, -1],
    [-2, 2, 1],
    color=plt.cm.viridis(
        [200, 50, 100, 200, 200, 50, 50, 100, 100]
    ),  # line1, line2, line3, l1arrow1, l1arrow2, l2arrow1, l2arrow2, l3arrow1, l3arrow2
    arrow_length_ratio=0.3,
)
ax.set(ylim=(-5, 5), xlim=(-5, 5), zlim=(-5, 5))
plt.show()

Upvotes: 3

Related Questions