3D plot of the CONE using matplotlib

Question

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.

Answer 1

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

Answer 2

Inspired from Jake Vanderplas with Python Data Science Handbook , when you are drawing some 3-D plot whose base is a circle, it is likely that you would try:

# Actually not sure about the math here though:
u, v = np.mgrid[0:2*np.pi:100j, 0:np.pi:20j]
x = np.cos(u)*np.sin(v)
y = np.sin(u)*np.sin(v)

and then think about the z-axis. Since viewing from the z-axis the cone is just a circle, so the relationships between z and x and y is clear, which is simply: z = np.sqrt(x ** 2 + y ** 2). Then you can draw the cone based on the codes below:

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

def f(x, y):
    return np.sqrt(x ** 2 + y ** 2)

fig = plt.figure()
ax = plt.axes(projection='3d')
# Can manipulate with 100j and 80j values to make your cone looks different
u, v = np.mgrid[0:2*np.pi:100j, 0:np.pi:80j]
x = np.cos(u)*np.sin(v)
y = np.sin(u)*np.sin(v)
z = f(x, y)

ax.plot_surface(x, y, z, cmap=cm.coolwarm)

# Some other effects you may want to try based on your needs:
# ax.plot_surface(x, y, -z, cmap=cm.coolwarm)
# ax.scatter3D(x, y, z, color="b")
# ax.plot_wireframe(x, y, z, color="b")
# ax.plot_wireframe(x, y, -z, color="r")

# Can set your view from different angles. 
ax.view_init(azim=15, elev=15)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
plt.show()
ax.set_ylabel("y")
ax.set_zlabel("z")
plt.show()

And from my side, the cone looks like:

and hope it helps.