The program calculates the reflection coefficient of the multilayer system, depending on the thickness of 50 layers (d1, d2). If I take any two numbers (d1,d2), it works correct. But I need to get Wireframe plots, where d1, d2 takes meaning in some range, I get an error: "ValueError: input must be a square array" in line 13. How can I fix it?
from math import pi
import numpy as np
import matplotlib.pyplot as plt
def R(n1, n2, d1, d2, lamda):
phy1 = (-2*pi*n1*d1/lamda)
phy2 = (-2*pi*n2*d2/lamda)
DPD1 = 0.5*np.array([[2*np.cos(phy1), 2j*np.sin(phy1)/n1], [n1*2j*np.sin(phy1), 2*np.cos(phy1) ]])
DPD2 = 0.5*np.array([[2*np.cos(phy2), 2j*np.sin(phy2)/n2], [n2*2j*np.sin(phy2), 2*np. cos(phy2) ]])
D0 = 0.5 * np.array([[1, 1], [1, -1]])
DS = np.array([[1, 1], [n1, -n1]])
DPD = np.dot(DPD1, DPD2)
DPD = np.linalg.matrix_power(DPD, 50)
M = np.dot(D0, DPD)
M = np.dot(M, DS)
return(abs(M[1,0]/M[0,0])**2)
x = np.arange(0, 10, 1)
y = np.arange(0, 10, 1)
X, Y = np.meshgrid(x, y)
Z = R(0.99910053+0.00183184j, 0.92373900+0.00644652j, X, Y, 13.5)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_wireframe(X, Y, Z, antialiased=True)
I don't know how to do it without for loop. Here is my solution:
from math import pi
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
def R(n1, n2, d1, d2, lamda):
phy1 = (-2*pi*n1*d1/lamda)
phy2 = (-2*pi*n2*d2/lamda)
DPD1 = 0.5*np.array([[2*np.cos(phy1), 2j*np.sin(phy1)/n1], [n1*2j*np.sin(phy1), 2*np.cos(phy1) ]])
DPD2 = 0.5*np.array([[2*np.cos(phy2), 2j*np.sin(phy2)/n2], [n2*2j*np.sin(phy2), 2*np. cos(phy2) ]])
DPD = np.dot(DPD1, DPD2)
DPD = np.linalg.matrix_power(DPD, 50)
D0 = 0.5 * np.array([[1, 1], [1, -1]])
DS = np.array([[1, 1], [n1, -n1]])
M = np.dot(D0, DPD)
M = np.dot(M, DS)
return(abs(M[1,0]/M[0,0])**2)
x = np.arange(0, 10, 1)
y = np.arange(0, 10, 1)
X, Y = np.meshgrid(x, y)
Z = np.zeros(X.shape).ravel()
for i, (x, y) in enumerate(zip(X.ravel(), Y.ravel())):
Z[i] = R(0.99910053+0.00183184j, 0.92373900+0.00644652j, x, y, 13.5)
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_wireframe(X, Y, Z.reshape(X.shape), antialiased=True)
plt.show()
Related
I have the following code:
import numpy as np
import matplotlib.pyplot as plt
x = np.linspace(-np.pi/2, np.pi/2, 30)
y = np.linspace(-np.pi/2, np.pi/2, 30)
x,y = np.meshgrid(x,y)
z = np.sin(x**2+y**2)[:-1,:-1]
fig,ax = plt.subplots()
ax.pcolormesh(x,y,z)
Which gives this image:
Now lets say I want to highlight the edge certain grid boxes:
highlight = (z > 0.9)
I could use the contour function, but this would result in a "smoothed" contour. I just want to highlight the edge of a region, following the edge of the grid boxes.
The closest I've come is adding something like this:
highlight = np.ma.masked_less(highlight, 1)
ax.pcolormesh(x, y, highlight, facecolor = 'None', edgecolors = 'w')
Which gives this plot:
Which is close, but what I really want is for only the outer and inner edges of that "donut" to be highlighted.
So essentially I am looking for some hybrid of the contour and pcolormesh functions - something that follows the contour of some value, but follows grid bins in "steps" rather than connecting point-to-point. Does that make sense?
Side note: In the pcolormesh arguments, I have edgecolors = 'w', but the edges still come out to be blue. Whats going on there?
EDIT:
JohanC's initial answer using add_iso_line() works for the question as posed. However, the actual data I'm using is a very irregular x,y grid, which cannot be converted to 1D (as is required for add_iso_line().
I am using data which has been converted from polar coordinates (rho, phi) to cartesian (x,y). The 2D solution posed by JohanC does not appear to work for the following case:
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
def pol2cart(rho, phi):
x = rho * np.cos(phi)
y = rho * np.sin(phi)
return(x, y)
phi = np.linspace(0,2*np.pi,30)
rho = np.linspace(0,2,30)
pp, rr = np.meshgrid(phi,rho)
xx,yy = pol2cart(rr, pp)
z = np.sin(xx**2 + yy**2)
scale = 5
zz = ndimage.zoom(z, scale, order=0)
fig,ax = plt.subplots()
ax.pcolormesh(xx,yy,z[:-1, :-1])
xlim = ax.get_xlim()
ylim = ax.get_ylim()
xmin, xmax = xx.min(), xx.max()
ymin, ymax = yy.min(), yy.max()
ax.contour(np.linspace(xmin,xmax, zz.shape[1]) + (xmax-xmin)/z.shape[1]/2,
np.linspace(ymin,ymax, zz.shape[0]) + (ymax-ymin)/z.shape[0]/2,
np.where(zz < 0.9, 0, 1), levels=[0.5], colors='red')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
This post shows a way to draw such lines. As it is not straightforward to adapt to the current pcolormesh, the following code demonstrates a possible adaption.
Note that the 2d versions of x and y have been renamed, as the 1d versions are needed for the line segments.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
x = np.linspace(-np.pi / 2, np.pi / 2, 30)
y = np.linspace(-np.pi / 2, np.pi / 2, 30)
xx, yy = np.meshgrid(x, y)
z = np.sin(xx ** 2 + yy ** 2)[:-1, :-1]
fig, ax = plt.subplots()
ax.pcolormesh(x, y, z)
def add_iso_line(ax, value, color):
v = np.diff(z > value, axis=1)
h = np.diff(z > value, axis=0)
l = np.argwhere(v.T)
vlines = np.array(list(zip(np.stack((x[l[:, 0] + 1], y[l[:, 1]])).T,
np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
l = np.argwhere(h.T)
hlines = np.array(list(zip(np.stack((x[l[:, 0]], y[l[:, 1] + 1])).T,
np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
lines = np.vstack((vlines, hlines))
ax.add_collection(LineCollection(lines, lw=1, colors=color))
add_iso_line(ax, 0.9, 'r')
plt.show()
Here is an adaption of the second answer, which can work with only 2d arrays:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from scipy import ndimage
x = np.linspace(-np.pi / 2, np.pi / 2, 30)
y = np.linspace(-np.pi / 2, np.pi / 2, 30)
x, y = np.meshgrid(x, y)
z = np.sin(x ** 2 + y ** 2)
scale = 5
zz = ndimage.zoom(z, scale, order=0)
fig, ax = plt.subplots()
ax.pcolormesh(x, y, z[:-1, :-1] )
xlim = ax.get_xlim()
ylim = ax.get_ylim()
xmin, xmax = x.min(), x.max()
ymin, ymax = y.min(), y.max()
ax.contour(np.linspace(xmin,xmax, zz.shape[1]) + (xmax-xmin)/z.shape[1]/2,
np.linspace(ymin,ymax, zz.shape[0]) + (ymax-ymin)/z.shape[0]/2,
np.where(zz < 0.9, 0, 1), levels=[0.5], colors='red')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
plt.show()
I'll try to refactor add_iso_line method in order to make it more clear an open for optimisations. So, at first, there comes a must-do part:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
x = np.linspace(-np.pi/2, np.pi/2, 30)
y = np.linspace(-np.pi/2, np.pi/2, 30)
x, y = np.meshgrid(x,y)
z = np.sin(x**2+y**2)[:-1,:-1]
fig, ax = plt.subplots()
ax.pcolormesh(x,y,z)
xlim, ylim = ax.get_xlim(), ax.get_ylim()
highlight = (z > 0.9)
Now highlight is a binary array that looks like this:
After that we can extract indexes of True cells, look for False neighbourhoods and identify positions of 'red' lines. I'm not comfortable enough with doing it in a vectorised manner (like here in add_iso_line method) so just using simple loop:
lines = []
cells = zip(*np.where(highlight))
for x, y in cells:
if x == 0 or highlight[x - 1, y] == 0: lines.append(([x, y], [x, y + 1]))
if x == highlight.shape[0] or highlight[x + 1, y] == 0: lines.append(([x + 1, y], [x + 1, y + 1]))
if y == 0 or highlight[x, y - 1] == 0: lines.append(([x, y], [x + 1, y]))
if y == highlight.shape[1] or highlight[x, y + 1] == 0: lines.append(([x, y + 1], [x + 1, y + 1]))
And, finally, I resize and center coordinates of lines in order to fit with pcolormesh:
lines = (np.array(lines) / highlight.shape - [0.5, 0.5]) * [xlim[1] - xlim[0], ylim[1] - ylim[0]]
ax.add_collection(LineCollection(lines, colors='r'))
plt.show()
In conclusion, this is very similar to JohanC solution and, in general, slower. Fortunately, we can reduce amount of cells significantly, extracting contours only using python-opencv package:
import cv2
highlight = highlight.astype(np.uint8)
contours, hierarchy = cv2.findContours(highlight, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
cells = np.vstack(contours).squeeze()
This is an illustration of cells being checked:
Say I have a grid
xGrid = np.linspace(0.1, 1, 10)
yGrid = np.linspace(5, 10, 5)
and some data on that grid:
X, Y = np.meshgrid(xGrid, yGrid, indexing='ij')
Z = X*Y + 1
I could easily now plot Z(x, y). Now, there is a transformation t(x, y):
T = X+1+Y/2
and I would like to plot Z(t(x, y), y) instead. To do that, I need to project my Z data onto the t(x,y)-y plane. What'd be the best way of doing that?
Since I ultimately want to plot the data and not do any further work with it, direct methods of doing this in matplotlib (but actually drawing onto the correct new coordinates, not just relabeling the ticks) are also accepted.
You can use interpolation to compute the values in the projection, for example with scipy.interpolate.RectBivariateSpline:
import numpy as np
import scipy.interpolate
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
xGrid = np.linspace(0.1, 1, 10)
yGrid = np.linspace(5, 10, 5)
X, Y = np.meshgrid(xGrid, yGrid, indexing='ij')
Z = X * Y + 1
T = X + 1 + Y / 2
# Interpolate values
interp = scipy.interpolate.RectBivariateSpline(xGrid, yGrid, Z)
Zt = interp.ev(T.ravel(), Y.ravel()).reshape(Z.shape)
# Plot
fig = plt.figure(figsize=(8, 10))
ax1 = fig.add_subplot(211, projection='3d')
ax1.set_title('Original')
ax1.plot_surface(X, Y, Z)
ax2 = fig.add_subplot(212, projection='3d')
ax2.set_title('Projected')
ax2.plot_surface(T, Y, Zt)
fig.tight_layout()
Output:
If I understand your problem you could use pcolormesh that can be used for non regular meshes
In [8]: import numpy as np
...: import matplotlib.pyplot as plt
...: from matplotlib.collections import PatchCollection, QuadMesh
...: from matplotlib.patches import Rectangle
...:
...: np.random.seed(2018)
...: xGrid = np.linspace(0.1, 1, 10)
...: yGrid = np.linspace(5, 10, 6)
...: X, Y = np.meshgrid(xGrid, yGrid, indexing='ij')
...: Z = X*Y + 1
...: T = X+1+Y/2
...: Zt = T*Y + 1
...: plt.pcolormesh(T, Y, Zt)
...: plt.colorbar()
Out[8]: <matplotlib.colorbar.Colorbar at 0x7fda83cd4ef0>
that produces
If the bands are too ugly use plt.pcolormesh(T, Y, Zt, shading='gouraud')
How can I create a multivariate skew normal function, where then by inputting x and y points we can create a surface diagram in 3d (x,y and z coordinates)
I wrote a blog post about this, but here is complete working code:
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import (multivariate_normal as mvn,
norm)
class multivariate_skewnorm:
def __init__(self, a, cov=None):
self.dim = len(a)
self.a = np.asarray(a)
self.mean = np.zeros(self.dim)
self.cov = np.eye(self.dim) if cov is None else np.asarray(cov)
def pdf(self, x):
return np.exp(self.logpdf(x))
def logpdf(self, x):
x = mvn._process_quantiles(x, self.dim)
pdf = mvn(self.mean, self.cov).logpdf(x)
cdf = norm(0, 1).logcdf(np.dot(x, self.a))
return np.log(2) + pdf + cdf
xx = np.linspace(-2, 2, 100)
yy = np.linspace(-2, 2, 100)
X, Y = np.meshgrid(xx, yy)
pos = np.dstack((X, Y))
fig = plt.figure(figsize=(10, 10), dpi=150)
axes = [
fig.add_subplot(1, 3, 1, projection='3d'),
fig.add_subplot(1, 3, 2, projection='3d'),
fig.add_subplot(1, 3, 3, projection='3d')
]
for a, ax in zip([[0, 0], [5, 1], [1, 5]], axes):
Z = multivariate_skewnorm(a=a).pdf(pos)
ax.plot_surface(X, Y, Z, cmap=cm.viridis)
ax.set_title(r'$\alpha$ = %s, cov = $\mathbf{I}$' % str(a), fontsize=18)
That code will generate this figure:
You can add direction to multivariate normal distribution by adding a sigma covariance matrix:
import numpy as np
from scipy.stats import multivariate_normal
mu = [20,20] # center of distribution.
sigma_size_top, sigma_size_bot = np.random.uniform(5, 20, size=2)
cov_max = np.sqrt(sigma_size_top * sigma_size_bot) * 0.9 # Cov max can't be larger than sqrt of the other elements
sigma_cov = np.random.uniform(-cov_max, cov_max)
sigma = np.array([[sigma_size_top, sigma_cov],[sigma_cov, sigma_size_bot]])
And then pass it in to your multivariate_normal:
dist = multivariate_normal(mu, sigma)
Put this into a 2D mapping by:
x = np.linspace(0, 40, 41)
y = x.copy()
xx, yy = np.meshgrid(x, y)
pos = np.empty(xx.shape + (2,))
pos[:, :, 0] = xx
pos[:, :, 1] = yy
my_map = dist.pdf(pos)
You'll then have a skewed multivariate normal distribution on a matrix. I suggest you scale this matrix as the values will be small.
I am trying to color each individual face of a cylinder, however I am not sure how to go about it, I have tried the following:
for i in range(10):
col.append([])
for i in range(10):
for j in range(20):
col[i].append(plt.cm.Blues(0.4))
ax.plot_surface(X, Y, Z,facecolors = col,edgecolor = "red")
I want each face to be assigned its own color, so I would think I would supply an array of colors for each of the faces in a 2d array.
But this gives an error:
in plot_surface
colset.append(fcolors[rs][cs])
IndexError: list index out of range
Here is the full code:
import numpy as np
from matplotlib import cm
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.linalg import norm
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
origin = np.array([0, 0, 0])
#axis and radius
p0 = np.array([1, 3, 2])
p1 = np.array([8, 5, 9])
R = 5
#vector in direction of axis
v = p1 - p0
#find magnitude of vector
mag = norm(v)
#unit vector in direction of axis
v = v / mag
#make some vector not in the same direction as v
not_v = np.array([1, 0, 0])
if (v == not_v).all():
not_v = np.array([0, 1, 0])
#make vector perpendicular to v
n1 = np.cross(v, not_v)
#normalize n1
n1 /= norm(n1)
#make unit vector perpendicular to v and n1
n2 = np.cross(v, n1)
#surface ranges over t from 0 to length of axis and 0 to 2*pi
t = np.linspace(0, mag, 200)
theta = np.linspace(0, 2 * np.pi, 100)
#use meshgrid to make 2d arrays
t, theta = np.meshgrid(t, theta)
#generate coordinates for surface
X, Y, Z = [p0[i] + v[i] * t + R * np.sin(theta) * n1[i] + R * np.cos(theta) * n2[i] for i in [0, 1, 2]]
col = []
for i in range(10):
col.append([])
for i in range(10):
for j in range(20):
col[i].append(plt.cm.Blues(0.4))
ax.plot_surface(X, Y, Z,facecolors = col,edgecolor = "red")
#plot axis
ax.plot(*zip(p0, p1), color = 'red')
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.set_zlim(0, 10)
plt.axis('off')
ax.axes.get_xaxis().set_visible(False)
ax.axes.get_yaxis().set_visible(False)
plt.show()
Your Z array is of size 100x200, yet you are only specifying 10x20 colors. A quicker way to make col (with the right dimensions) might be something like:
col1 = plt.cm.Blues(np.linspace(0,1,200)) # linear gradient along the t-axis
col1 = np.repeat(col1[np.newaxis,:, :], 100, axis=0) # expand over the theta-axis
col2 = plt.cm.Blues(np.linspace(0,1,100)) # linear gradient along the theta-axis
col2 = np.repeat(col2[:, np.newaxis, :], 200, axis=1) # expand over the t-axis
ax=plt.subplot(121, projection='3d')
ax.plot_surface(X, Y, Z, facecolors=col1)
ax=plt.subplot(122, projection='3d')
ax.plot_surface(X, Y, Z, facecolors=col2)
Which produces:
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()