I want to fit a plane to some data points and draw it. My current code is this:
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
points = [(1.1,2.1,8.1),
(3.2,4.2,8.0),
(5.3,1.3,8.2),
(3.4,2.4,8.3),
(1.5,4.5,8.0)]
xs, ys, zs = zip(*points)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(xs, ys, zs)
point = np.array([0.0, 0.0, 8.1])
normal = np.array([0.0, 0.0, 1.0])
d = -point.dot(normal)
xx, yy = np.meshgrid([-5,10], [-5,10])
z = (-normal[0] * xx - normal[1] * yy - d) * 1. /normal[2]
ax.plot_surface(xx, yy, z, alpha=0.2, color=[0,1,0])
ax.set_xlim(-10,10)
ax.set_ylim(-10,10)
ax.set_zlim( 0,10)
plt.show()
which results in the following:
As you can see at the moment I create the plane manually. How can I calculate it? I guess it is possible with scipy.optimize.minimize somehow. The kind of error function is not that important to me at the moment. I think least squares (vertical point-plane-distance) would be fine. It would be cool if one of you could show me how to do it.
Oh, the idea just came to my mind. It's quite easy. :-)
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import scipy.optimize
import functools
def plane(x, y, params):
a = params[0]
b = params[1]
c = params[2]
z = a*x + b*y + c
return z
def error(params, points):
result = 0
for (x,y,z) in points:
plane_z = plane(x, y, params)
diff = abs(plane_z - z)
result += diff**2
return result
def cross(a, b):
return [a[1]*b[2] - a[2]*b[1],
a[2]*b[0] - a[0]*b[2],
a[0]*b[1] - a[1]*b[0]]
points = [(1.1,2.1,8.1),
(3.2,4.2,8.0),
(5.3,1.3,8.2),
(3.4,2.4,8.3),
(1.5,4.5,8.0)]
fun = functools.partial(error, points=points)
params0 = [0, 0, 0]
res = scipy.optimize.minimize(fun, params0)
a = res.x[0]
b = res.x[1]
c = res.x[2]
xs, ys, zs = zip(*points)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(xs, ys, zs)
point = np.array([0.0, 0.0, c])
normal = np.array(cross([1,0,a], [0,1,b]))
d = -point.dot(normal)
xx, yy = np.meshgrid([-5,10], [-5,10])
z = (-normal[0] * xx - normal[1] * yy - d) * 1. /normal[2]
ax.plot_surface(xx, yy, z, alpha=0.2, color=[0,1,0])
ax.set_xlim(-10,10)
ax.set_ylim(-10,10)
ax.set_zlim( 0,10)
plt.show()
Sorry for asking unnecessarily.
Another way is with a straight forward least squares solution.
The equation for a plane is: ax + by + c = z. So set up matrices like this with all your data:
x_0 y_0 1
A = x_1 y_1 1
...
x_n y_n 1
And
a
x = b
c
And
z_0
B = z_1
...
z_n
In other words: Ax = B. Now solve for x which are your coefficients. But since (I assume) you have more than 3 points, the system is over-determined so you need to use the left pseudo inverse. So the answer is:
a
b = (A^T A)^-1 A^T B
c
And here is some simple Python code with an example:
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
N_POINTS = 10
TARGET_X_SLOPE = 2
TARGET_y_SLOPE = 3
TARGET_OFFSET = 5
EXTENTS = 5
NOISE = 5
# create random data
xs = [np.random.uniform(2*EXTENTS)-EXTENTS for i in range(N_POINTS)]
ys = [np.random.uniform(2*EXTENTS)-EXTENTS for i in range(N_POINTS)]
zs = []
for i in range(N_POINTS):
zs.append(xs[i]*TARGET_X_SLOPE + \
ys[i]*TARGET_y_SLOPE + \
TARGET_OFFSET + np.random.normal(scale=NOISE))
# plot raw data
plt.figure()
ax = plt.subplot(111, projection='3d')
ax.scatter(xs, ys, zs, color='b')
# do fit
tmp_A = []
tmp_b = []
for i in range(len(xs)):
tmp_A.append([xs[i], ys[i], 1])
tmp_b.append(zs[i])
b = np.matrix(tmp_b).T
A = np.matrix(tmp_A)
fit = (A.T * A).I * A.T * b
errors = b - A * fit
residual = np.linalg.norm(errors)
print "solution:"
print "%f x + %f y + %f = z" % (fit[0], fit[1], fit[2])
print "errors:"
print errors
print "residual:"
print residual
# plot plane
xlim = ax.get_xlim()
ylim = ax.get_ylim()
X,Y = np.meshgrid(np.arange(xlim[0], xlim[1]),
np.arange(ylim[0], ylim[1]))
Z = np.zeros(X.shape)
for r in range(X.shape[0]):
for c in range(X.shape[1]):
Z[r,c] = fit[0] * X[r,c] + fit[1] * Y[r,c] + fit[2]
ax.plot_wireframe(X,Y,Z, color='k')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
plt.show()
Thanks #Ben for sharing! Since np.matrix is deprecated, I edited your code so it works with np arrays
import matplotlib.pyplot as plt
import numpy as np
from numpy.linalg import inv
# Pass the function array of points, shape (3, X)
def plane_from_points(points):
# Create this matrix correctly without transposing it later?
A = np.array([
points[0,:],
points[1,:],
np.ones(points.shape[1])
]).T
b = np.array([points[2, :]]).T
# fit = (A.T * A).I * A.T * b
fit = np.dot(np.dot(inv(np.dot(A.T, A)), A.T), b)
# errors = b - np.dot(A, fit)
# residual = np.linalg.norm(errors)
return fit
N_POINTS = 10
TARGET_X_SLOPE = 2
TARGET_y_SLOPE = 3
TARGET_OFFSET = 5
EXTENTS = 5
NOISE = 3
# create random data
xs = [np.random.uniform(2*EXTENTS)-EXTENTS for i in range(N_POINTS)]
ys = [np.random.uniform(2*EXTENTS)-EXTENTS for i in range(N_POINTS)]
zs = []
for i in range(N_POINTS):
zs.append(xs[i]*TARGET_X_SLOPE + \
ys[i]*TARGET_y_SLOPE + \
TARGET_OFFSET + np.random.normal(scale=NOISE))
points = np.array([xs, ys, zs])
# plot raw data
plt.figure()
ax = plt.subplot(111, projection='3d')
ax.scatter(xs, ys, zs, color='b')
fit = plane_from_points(points)
# plot plane
xlim = ax.get_xlim()
ylim = ax.get_ylim()
X,Y = np.meshgrid(np.arange(xlim[0], xlim[1]),
np.arange(ylim[0], ylim[1]))
Z = np.zeros(X.shape)
for r in range(X.shape[0]):
for c in range(X.shape[1]):
Z[r,c] = fit[0] * X[r,c] + fit[1] * Y[r,c] + fit[2]
ax.plot_wireframe(X,Y,Z, color='k')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
plt.show()
I'm surprised nobody has mentioned lsq_linear. There you can more or less directly plug in the data points and get the plane coefficients out:
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
points = np.array([[1.1,2.1,8.1],
[3.2,4.2,8.0],
[5.3,1.3,8.2],
[3.4,2.4,8.3],
[1.5,4.5,8.0]])
A = np.hstack((points[:,:2], np.ones((len(xs),1))))
b = points[:,2]
res = scipy.optimize.lsq_linear(A, b)
assert res.success
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(xs, ys, zs)
XnY = np.linspace(-5,10,10)
X, Y = np.meshgrid(XnY, XnY)
Z = res.x[0] * X + res.x[1] * Y + res.x[2]
surf = ax.plot_surface(X, Y, Z, alpha=0.2, color=[0,1,0])
ax.set_xlim(-5,10)
ax.set_ylim(-5,10)
ax.set_zlim( 0,10)
plt.show()
Related
I have 4 arrays x, y, z and T of length n and I want to plot a 3D curve using matplotlib. The (x, y, z) are the points positions and T is the value of each point (which is plotted as color), like the temperature of each point. How can I do it?
Example code:
import numpy as np
from matplotlib import pyplot as plt
fig = plt.figure()
ax = fig.gca(projection='3d')
n = 100
cmap = plt.get_cmap("bwr")
theta = np.linspace(-4 * np.pi, 4 * np.pi, n)
z = np.linspace(-2, 2, n)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)
T = (2*np.random.rand(n) - 1) # All the values are in [-1, 1]
What I found over the internet:
It's possible to use cmap with scatter like shown in the docs and in this stackoverflow question
ax = plt.gca()
ax.scatter(x, y, z, cmap=cmap, c=T)
The problem is that scatter is a set of points, not a curve.
In this stackoverflow question the solution was divide in n-1 intervals and each interval we use a different color like
t = (T - np.min(T))/(np.max(T)-np.min(T)) # Normalize
for i in range(n-1):
plt.plot(x[i:i+2], y[i:i+2], z[i:i+2], c=cmap(t[i])
The problem is that each segment has only one color, but it should be an gradient. The last value is not even used.
Useful links:
Matplotlib - Colormaps
Matplotlib - Tutorial 3D
This is a case where you probably need to use Line3DCollection. This is the recipe:
create segments from your array of coordinates.
create a Line3DCollection object.
add that collection to the axis.
set the axis limits.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Line3DCollection
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize
def get_segments(x, y, z):
"""Convert lists of coordinates to a list of segments to be used
with Matplotlib's Line3DCollection.
"""
points = np.ma.array((x, y, z)).T.reshape(-1, 1, 3)
return np.ma.concatenate([points[:-1], points[1:]], axis=1)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
n = 100
cmap = plt.get_cmap("bwr")
theta = np.linspace(-4 * np.pi, 4 * np.pi, n)
z = np.linspace(-2, 2, n)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)
T = np.cos(theta)
segments = get_segments(x, y, z)
c = Line3DCollection(segments, cmap=cmap, array=T)
ax.add_collection(c)
fig.colorbar(c)
ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.min(), y.max())
ax.set_zlim(z.min(), z.max())
plt.show()
I have a list with x,y,z, r coordinates (cartesians). I need to plot some circles in a polar plot, but I don't know how to do it with cartesians coordinates.
I am trying to do it with this line
circle1 = plt.Circle((x[i], y[i]), r[i], transform=ax3.transData._b, color = 'r', alpha=0.5, fill=False)
but this doesn't seem to work because I obtain the circles too far away from the center of the origin.
any help?
data1 = pd.read_csv('Uchuu_lightcone_0_11.9_voids.txt', sep='\s+', header=None)
data1 = pd.DataFrame(data1)
x = data1[0]
y = data1[1]
r = data1[3]
z = data1[2]
azvoids, elvoids, rvoids = cart2sph(x,y,z)
d = ax3.scatter(azvoids, rvoids, s=3, c='red', alpha=1, marker='.')
for i in range(len(x)):
if elvoids[i] > 35 and elvoids[i] < 45:
circle1 = plt.Circle((x[i], y[i]), r[i], transform=ax3.transData._b, color = 'r', alpha=0.5, fill=False)
ax3.add_artist(circle1)
# The cart2sph function is
def cart2sph(x,y,z):
""" x, y, z : ndarray coordinates
ceval: backend to use:
- eval : pure Numpy
- numexpr.evaluate: Numexpr """
azimuth = arctan2(y,x)*180/math.pi
xy2 = x**2 + y**2
elevation = arctan2(z, sqrt(xy2))*180/math.pi
r = sqrt(xy2 + z**2)
return azimuth, elevation, r
You should use azvoids and rvoids to plot the center of the circle since you use those to show tham in the scatter plot
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.transforms as transforms
def cart2sph(x, y, z):
""" x, y, z : ndarray coordinates
ceval: backend to use:
- eval : pure Numpy
- numexpr.evaluate: Numexpr """
azimuth = np.arctan2(y, x) * 180 / math.pi
xy2 = x ** 2 + y ** 2
elevation = np.arctan2(z, np.sqrt(xy2)) * 180 / math.pi
r = np.sqrt(xy2 + z ** 2)
return azimuth, elevation, r
#
# data1 = pd.read_csv('Uchuu_lightcone_0_11.9_voids.txt', sep='\s+', header=None)
# data1 = pd.DataFrame(data1)
N=100
x = (np.random.rand(N)-0.5)*100
y = (np.random.rand(N)-0.5)*100
z = (np.random.rand(N)-0.5)*100
r = np.random.rand(N)*10
azvoids, elvoids, rvoids = cart2sph(x, y, z)
fig = plt.figure()
ax3 = plt.subplot(111 )
d = plt.scatter(azvoids, elvoids, s=3 , c='red', alpha=1, marker='.' )
for i in range(len(x)):
if elvoids[i] > 35 and elvoids[i] < 45:
# circle1 = plt.Circle((azvoids[i], elvoids[i]), rvoids[i], color='r', alpha=0.5, fill=False)
x, y = ax3.transData.transform((azvoids[i], elvoids[i]))
trans = (fig.dpi_scale_trans +
transforms.ScaledTranslation(azvoids[i], elvoids[i], ax3.transData))
circle1 = plt.Circle((azvoids[i], elvoids[i]), rvoids[i] , color='r', alpha=0.5, fill=None)
ax3.add_artist(circle1)
plt.axis('equal')
plt.show()
# The cart2sph function is
I'm trying to place a plane on the surface of a sphere, although I think the math is correct, the resulting figure displays the plane at some point else.
Here is the code to compute and visualize it;
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
plt.style.use('dark_background')
# point and the unit vector
point = np.array([ 349370.39178182, 5570903.05977037, 3085958.36621096])
unit_vector = point/np.linalg.norm(point)
# the math
print(f'plane equation:\nAx + By + Cz + D = 0')
D = np.sum(unit_vector * point)
print(f'D=-(point * unit_vector) = {D:.2f}')
print(f'plane equation:\n{unit_vector[0]:1.4f}x + {unit_vector[1]:1.4f}y + {unit_vector[2]:1.4f}z + {D:.1f} = 0')
print(f'{-1*unit_vector[2]:1.4f}z = {unit_vector[0]:1.4f}x + {unit_vector[1]:1.4f}y + {D:.1f}')
print(f'z = ({unit_vector[0]:1.4f}x + {unit_vector[1]:1.4f}y + {D:.1f}) / {-1*unit_vector[2]:1.4f}')
x = np.linspace(-3e6,+3e6,100)
y = np.linspace(-3e6,+3e6,100)
X,Y = np.meshgrid(x,y)
Z = (0.05477656*X +0.87344241*Y + 6378100.0)/-0.48383662
# plotting stuff
def set_axis_equal_scale(ax, ticks_off=True):
xl = ax.set_xlim()
yl = ax.set_ylim()
zl = ax.set_zlim()
maxx=max(max(xl), max(yl), max(zl))
minn=min(min(xl), min(yl), min(zl))
ax.set_xlim(minn, maxx)
ax.set_ylim(minn, maxx)
ax.set_zlim(minn, maxx)
if ticks_off:
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
fig = plt.figure(figsize=(5,5))
ax = fig.add_subplot(111, projection='3d')
# plot wireframe
radius = 6.3781e6 # in units m
uu, vv = np.mgrid[0:2*np.pi:200j, 0:np.pi:100j]
xE = radius * np.cos(uu)*np.sin(vv)
yE = radius * np.sin(uu)*np.sin(vv)
zE = radius * np.cos(vv)
ax.plot_wireframe(xE,yE,zE, color='w', alpha=0.1)
ax.scatter(point[0], point[1], point[2], s=500, color='r')
ax.plot([0,point[0]], [0,point[1]], [0,point[2]], color='w', lw=2)
surf = ax.plot_surface(X, Y, Z)
ax.scatter(0,0,0, marker='o', s=900, color='b')
ax.view_init(25, -190)
ax.axis('off')
set_axis_equal_scale(ax)
I expect the plane to be on where the red marker is and perpendicular to the white line connecting the center and the red marker.
[here is the image][1]
[1]: https://i.stack.imgur.com/LHPzW.png
I have an equation as followed:
y = x^T * A * x + b^T * x + c
where x, b, c are vectors in n space and A is a nxn matrix.
I can plot a linear equation in matplotlib, but not sure how a matrix equation can be (if possible) shown also in a 3d plot.
I tried with following code, A is given matrix and w, c and b are column vectors. X and Y are mesh and Z is the solution.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
# if using a Jupyter notebook, include:
%matplotlib inline
fig = plt.figure(figsize=(10,6))
ax1 = fig.add_subplot(111, projection='3d')
n = 50
i = -5.0
j = 5.0
A = np.random.randint(i, j, size=(n, n))
w = np.random.randint(i, j, size=(n, 1))
c = b = np.random.randint(i, j, size=(n, 1))
X,Y = np.meshgrid(n,n)
Z = w.T*A*w + b.T*w + c
mycmap = plt.get_cmap('gist_earth')
surf1 = ax1.plot_surface(X, A, Z, cmap=mycmap)
fig.colorbar(surf1, ax=ax1, shrink=0.5, aspect=10)
plt.show()
The resulting plot does not seem to be a satisfied plot.
There are two problems in your code: 1) meshgrid was being used incorrectly (it needs two arrays, not two ints); 2) in the surface plot, you were using X, A, Z instead of X, Y, Z -- X, A, Z will work, and might make sense, but I'm guessing it wasn't your intention.
Here's a working solution:
fig = plt.figure(figsize=(10,6))
ax1 = fig.add_subplot(111, projection='3d')
n = 10
i = -5.0
j = 5.0
A = np.random.randint(i, j, size=(n, n))
w = np.random.randint(i, j, size=(n, 1))
c = b = np.random.randint(i, j, size=(n, 1))
X,Y = np.meshgrid(np.arange(n),np.arange(n))
Z = w.T*A*w + b.T*w + c
mycmap = plt.get_cmap('gist_earth')
surf1 = ax1.plot_surface(X, Y, Z, cmap=mycmap)
fig.colorbar(surf1, ax=ax1, shrink=0.5, aspect=10)
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()
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.