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
Related
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 would like to create 3d scatter plot with colormap range from min(u), u =64 to max(u), u=100. u is a 1d array
The code works as expected, u is increasing from the center (x,y,z)=(0,0,0) but the colors is incorrect, the color gradient should range according to u, from min(u) to max(u) instead of depending on x,y,z coordinate. Also colorbar is not correct (should be from 0 to 100)
fig = plt.figure(figsize = (8,6))
ax = fig.add_subplot(111, projection='3d')
ax.set_title('normal distribution')
#add the line/data in our plot
x = 18 * np.random.normal(size =500)
y = 18 * np.random.normal(size =500)
z = 18 * np.random.normal(size =500)
u = np.linspace(64, 100, 500)
norma = mpl.colors.Normalize(min(u), max(u))
color = np.linalg.norm([x,y,z], axis=0)
track = ax.scatter(x,y,z, s=35, c = color, alpha = 1, cmap='inferno', norm = norma)
plt.colorbar(track, label='color map', shrink=0.6)
fig = plt.figure(figsize = (8,6))
ax = fig.add_subplot(111, projection='3d')
ax.set_title('normal distribution')
the above code figure
When the color map Normalise to vmin=min(u) and vmax=max(u), the color gradient is lost and colormap gradient values are spread randomly along the x,y,z axis instead of being in ordered array.
Does someone know how to fix the color gradient along the axis, while the center of u is at (0,0,0) with the correct color bar (0-100) please?
fig = plt.figure(figsize = (8,6))
ax = fig.add_subplot(111, projection='3d')
ax.set_title('normal distribution')
#add the line/data in our plot
x = 18 * np.random.normal(size =500)
y = 18 * np.random.normal(size =500)
z = 18 * np.random.normal(size =500)
u = np.linspace(100, 64, 500)
norma = mpl.colors.Normalize(vmin=0, vmax = 100)
color = np.linalg.norm([u], axis=0)
track = ax.scatter(x,y,z, s=35, c = color, alpha = 1, cmap='inferno', norm = norma)
plt.colorbar(track, label='color map', shrink=0.6)
The result of the second example
x = 18 * np.random.normal(size =500)
y = 18 * np.random.normal(size =500)
z = 18 * np.random.normal(size =500)
# collect all data in array
data = np.array([x,y,z])
# center in a given dimension is the mean of all datapoints:
# reshape to allow easy subtraction
center = np.mean(data, axis=1).reshape(3,-1)
# for each datapoint, calculate distance to center and use as color value
color = np.linalg.norm(data - center, axis=0)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
track = ax.scatter(x,y,z, s=35, c = color, alpha = 1, cmap='inferno')
plt.colorbar(track, label='color map', shrink=0.6)
I found this question which seems to answer your question about the coordinates. The answers also show how to evenly distribute coordinates if you prefer to do that.
After getting the coordinates, you can then get the distance from the center as the color value (like warped did in his answer). I adjusted the distance to reflect your specifications. This is the resulting code:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from mpl_toolkits.mplot3d import Axes3D
number_of_particles = 500
sphere_radius = 18
# create the particles
radius = sphere_radius * np.random.uniform(0.0, 1.0, number_of_particles)
theta = np.random.uniform(0., 1., number_of_particles) * 2 * np.pi
phi = np.random.uniform(0., 1., number_of_particles) * 2 * np.pi
x = radius * np.sin(theta) * np.cos(phi)
y = radius * np.sin(theta) * np.sin(phi)
z = radius * np.cos(theta)
# collect all data in array
data = np.array([x, y, z])
# for each datapoint, calculate distance to center and use as color value
color = radius
color /= sphere_radius
color = color * 36 + 64
# initialize a figure with a plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
# add the points and the colorbar
track = ax.scatter(x, y, z, s=35, c=color, alpha=1, cmap='inferno',
norm=Normalize(0, 100))
plt.colorbar(track, label='color map', shrink=0.6)
plt.show()
My result looks like this:
I am trying to create a map showing an interpolated values over a scatter plot, here my code so far
# define map extent
lllon = dd
lllat = bb
urlon = cc
urlat = aa
# Set up Basemap instance
m = Basemap(
projection = 'merc',
llcrnrlon = lllon, llcrnrlat = lllat, urcrnrlon = urlon, urcrnrlat = urlat,
resolution='h')
# transform lon / lat coordinates to map projection
newdf['LONGITUDE'], newdf['LATITUDE'] = m(*(newdf['LONGITUDE'].values, newdf['LATITUDE'].values))
# grid data
#numcols, numrows = count_col, count_row
#xi = np.linspace(dd, cc, numcols)
#yi = np.linspace(bb, aa, numrows)
#xi, yi = np.meshgrid(xi, yi)
count_row = newdf.shape[0] # gives number of row count
count_col = newdf.shape[1] # gives number of col count
xi = np.linspace(newdf['LONGITUDE'].min(), newdf['LONGITUDE'].max(), count_col)
yi = np.linspace(newdf['LATITUDE'].min(), newdf['LATITUDE'].max(), count_row)
xi, yi = np.meshgrid(xi, yi)
x, y, z = newdf['LONGITUDE'].values, newdf['LATITUDE'].values, newdf['MUD_WGHT'].values
#zi = griddata(x, y, z, xi, yi)
zi = griddata((x,y),z,(xi,yi),method='linear')
# interpolate
#x, y, z = newdf['LONGITUDE'].values, newdf['LATITUDE'].values, newdf['MUD_WGHT'].values
#zi = griddata((x,y),z,(xi,yi),method='linear')
# draw map details
m.drawmapboundary(fill_color = 'white')
m.fillcontinents(color='#C0C0C0', lake_color='#7093DB')
m.drawcountries(
linewidth=.75, linestyle='solid', color='#000073',
antialiased=True,
ax=ax, zorder=3)
m.drawparallels(
np.arange(lllat, urlat, 2.),
color = 'black', linewidth = 0.5,
labels=[True, False, False, False])
m.drawmeridians(
np.arange(lllon, urlon, 2.),
color = '0.25', linewidth = 0.5,
labels=[False, False, False, True])
# contour plot
con = m.contourf(xi, yi, zi, zorder=4, alpha=0.6, cmap='RdPu')
# scatter plot
m.scatter(
newdf['LONGITUDE'],
newdf['LATITUDE'],
color='#545454',
edgecolor='#ffffff',
alpha=.75,
s=50 * norm(newdf['MUD_WGHT'].values),
cmap='RdPu',
ax=ax,
vmin=zi.min(), vmax=zi.max(), zorder=4)
# add colour bar and title
# add colour bar, title, and scale
cbar = plt.colorbar(orientation='horizontal', fraction=.057, pad=0.05)
cbar.set_label("Mud Weight - PPG")
#m.drawmapscale(
# 24., -9., 28., -13,
# 100,
# units='km', fontsize=10,
#yoffset=None,
#barstyle='fancy', labelstyle='simple',
#fillcolor1='w', fillcolor2='#000000',
#fontcolor='#000000',
#zorder=5)
plt.title("Regional Mud Weights in The Lateral")
plt.show()
The result is the following:
How can I get the interpolated contour region to extend to the full scatter plot? I have been focusing on if this is an issue with the meshgrid, so I am not sure if the meshgrid isn't interpolating all of the data or if it is an issue with the plotting.
You need to do triangulation and use tricontour and/or tricontourf. Here is a demonstration code and sample plot.
from mpl_toolkits.basemap import Basemap
from matplotlib.tri import Triangulation
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
# data coordinates must conform to the projection
x = np.random.random((30))
y = np.random.random((30))
z = x * y
yn = 30*x
xn = 80 + 30*y
tri = Triangulation(xn,yn) #create tri mesh
fig = plt.figure(figsize=(7, 7))
m = Basemap(projection = 'cyl',
llcrnrlat = 0,
urcrnrlat = 30,
llcrnrlon = 80,
urcrnrlon = 110,
resolution = 'l')
ctf = plt.tricontourf(tri, z, cmap=cm.coolwarm, zorder=10, alpha=0.75)
#plt.tricontour(tri, z, )
plt.scatter(xn, yn, c='g', zorder=13)
m.drawparallels(np.arange(-90, 90,10), labels=[1,0,0,0])
m.drawmeridians(np.arange(-180, 180, 10), labels = [0,0,0,1])
m.drawcoastlines(linewidth=0.8, color='blue', zorder=12)
m.fillcontinents(color='#C0C0C0', lake_color='#7093DB')
cbar = plt.colorbar(ctf , orientation='horizontal', fraction=.045, pad=0.08)
plt.show()
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.
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()