Styjun

I'm looking for help to draw a 3D cone using matplotlib.
My goal is to draw a HSL cone, then base on the vertex coordinats i will select the color.

from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

theta1 = np.linspace(0, 2*np.pi, 100)
r1 = np.linspace(-2, 0, 100)
t1, R1 = np.meshgrid(theta1, r1)

X1 = R1*np.cos(t1)
Y1 = R1*np.sin(t1)
Z1 = 5+R1*2.5

theta2 = np.linspace(0, 2*np.pi, 100)
r2 = np.linspace(0, 2, 100)
t2, R2 = np.meshgrid(theta2, r2)

X2 = R2*np.cos(t2)
Y2 = R2*np.sin(t2)
Z2 = -5+R2*2.5


ax.set_xlabel('x axis')
ax.set_ylabel('y axis')
ax.set_zlabel('z axis')
# ax.set_xlim(-2.5, 2.5)
# ax.set_ylim(-2.5, 2.5)
# ax.set_zlim(0, 5)
ax.set_aspect('equal')

ax.plot_surface(X1, Y1, Z1, alpha=0.8, color="blue")
ax.plot_surface(X2, Y2, Z2, alpha=0.8, color="blue")
# ax.plot_surface(X, Y, Z, alpha=0.8)
#fig. savefig ("Cone.png", dpi=100, transparent = False)

plt.show()

HSL CONE

My cone

So my question now is how to define color of each element.

i have found a solution, maybe it will be usefull for others.

from mpl_toolkits.mplot3d import Axes3D

from matplotlib import cm

import matplotlib.pyplot as plt

import numpy as np
import colorsys


from matplotlib.tri import Triangulation

from mpl_toolkits.mplot3d.art3d import Poly3DCollection


n_angles = 80

n_radii = 20


# An array of radii

# Does not include radius r=0, this is to eliminate duplicate points

radii = np.linspace(0.0, 0.5, n_radii)


# An array of angles

angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)


# Repeat all angles for each radius

angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1)


# Convert polar (radii, angles) coords to cartesian (x, y) coords

# (0, 0) is added here. There are no duplicate points in the (x, y) plane


x = np.append(0, (radii*np.cos(angles)).flatten())

y = np.append(0, (radii*np.sin(angles)).flatten())


# Pringle surface

z = 1+-np.sqrt(x**2+y**2)*2


print(x.shape, y.shape, angles.shape, radii.shape, z.shape)

# NOTE: This assumes that there is a nice projection of the surface into the x/y-plane!
tri = Triangulation(x, y)

triangle_vertices = np.array([np.array([[x[T[0]], y[T[0]], z[T[0]]],

 [x[T[1]], y[T[1]], z[T[1]]],

 [x[T[2]], y[T[2]], z[T[2]]]]) for T in tri.triangles])


x2 = np.append(0, (radii*np.cos(angles)).flatten())

y2 = np.append(0, (radii*np.sin(angles)).flatten())


# Pringle surface
z2 = -1+np.sqrt(x**2+y**2)*2


# NOTE: This assumes that there is a nice projection of the surface into the x/y-plane!
tri2 = Triangulation(x2, y2)

triangle_vertices2 = np.array([np.array([[x2[T[0]], y2[T[0]], z2[T[0]]],

 [x2[T[1]], y2[T[1]], z2[T[1]]],

 [x2[T[2]], y2[T[2]], z2[T[2]]]]) for T in tri2.triangles])

triangle_vertices = np.concatenate([triangle_vertices, triangle_vertices2])

midpoints = np.average(triangle_vertices, axis=1)


def find_color_for_point(pt):

 c_x, c_y, c_z = pt

 angle = np.arctan2(c_x, c_y)*180/np.pi

 if (angle < 0):
 angle = angle + 360

 if c_z < 0:

 l = 0.5 - abs(c_z)/2
 #l=0
 if c_z == 0:
 l = 0.5
 if c_z > 0:
 l = (1 - (1-c_z)/2)

 if c_z > 0.97:

 l = (1 - (1-c_z)/2)

 col = colorsys.hls_to_rgb(angle/360, l, 1)

 return col


facecolors = [find_color_for_point(pt) for pt in midpoints] # smooth gradient
# facecolors = [np.random.random(3) for pt in midpoints] # random colors


coll = Poly3DCollection(
 triangle_vertices, facecolors=facecolors, edgecolors=None)


fig = plt.figure()

ax = fig.gca(projection='3d')

ax.add_collection(coll)

ax.set_xlim(-1, 1)

ax.set_ylim(-1, 1)

ax.set_zlim(-1, 1)

ax.elev = 50


plt.show()