How to fill in a shape composed of multiple curves? - python

I'm attempting to fill in an area enclosed by four curves but can't seem to get the fill I want: I'm using fill_between() but that can only handle two lines at most. Here's the code I currently have:
fig,ax = plt.subplots()
x = np.arange(start = -80,stop = 80,step = 1e-3)
yArc = np.sqrt(80**2 - x**2)
yLL = -2*x - 10
yRL = 2*x - 10
yTop = 0*x + 150
ax.fill_between(x,yArc,yTop,where = yArc<yTop,color = 'red',alpha = 0.8)
plt.show()
Here's the output:
Note: In case you can't see, there is a horizontal line just about where the free-throw line is.

You can use set operations between different filters to fill the central part first,
import numpy as np
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
x = np.arange(start=-80, stop=80, step=1e-3)
yArc = np.sqrt(80 ** 2 - x ** 2)
yLL = -2 * x - 10
yRL = 2 * x - 10
yTop = 0 * x + 150
plt.plot(x, yArc, "r")
plt.plot(x, yLL, "b")
plt.plot(x, yRL, "g")
plt.plot(x, yTop, "k")
filt_1 = yArc < yTop
filt_2 = yArc > yLL
filt_3 = yArc > yRL
filt = filt_1 & filt_2 & filt_3
ax.fill_between(x, yArc, yTop, where=filt, color="orange", alpha=0.8)
plt.show()
Producing:
And finally fill the two remaining right triangles, see for example:
matplotlib: use fill_between to make coloured triangles

The most obvious way to tackle this, is to create a bottom curve as the maximum of the three curves:
import numpy as np
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
x = np.arange(start=-80, stop=80, step=1e-3)
yArc = np.sqrt(80 ** 2 - x ** 2)
yLL = -2 * x - 10
yRL = 2 * x - 10
yTop = 0 * x + 150
plt.plot(x, yArc, "r")
plt.plot(x, yLL, "b")
plt.plot(x, yRL, "g")
plt.plot(x, yTop, "k")
yBottom = np.max([yArc, yLL, yRL], axis=0)
ax.fill_between(x, yBottom, yTop, where=yBottom <= yTop, color="orange", alpha=0.8)

Related

Python3 Matplotlib to Show Solid of Revolution from Method of Disks, Method of Shells, and Method of Washers around x-axis and y-axis

I have created this topic before: How to Add another subplot to show Solid of Revolution toward x-axis?
I want to plot the solid from a curve 3 + 2x - x^2 that is revolved about:
a. the x-axis
b. the y-axis
c. the line y = -1
d. the line x = 4
this is my MWE (the problem is to find the inverse of 3 + 2x - x^2 is not something easy thus I have no idea how to make this code works):
# Compare the plot at xy axis with the solid of revolution toward x and y axis
# For function x=(y)^(3/2)
import matplotlib.pyplot as plt
import numpy as np
n = 100
fig = plt.figure(figsize=(14, 7))
ax1 = fig.add_subplot(221)
ax2 = fig.add_subplot(222, projection='3d')
ax3 = fig.add_subplot(223)
ax4 = fig.add_subplot(224, projection='3d')
y = np.linspace(0, 9, n)
x = (y) ** (3 / 2)
t = np.linspace(0, np.pi * 2, n)
xn = np.outer(x, np.cos(t))
yn = np.outer(x, np.sin(t))
zn = np.zeros_like(xn)
for i in range(len(x)):
zn[i:i + 1, :] = np.full_like(zn[0, :], y[i])
ax1.plot(x, y)
ax1.set_title("$f(x)$")
ax2.plot_surface(xn, yn, zn)
ax2.set_title("$f(x)$: Revolution around $y$")
# find the inverse of the function
x_inverse = y
y_inverse = np.power(x_inverse, 3 / 2)
xn_inverse = np.outer(x_inverse, np.cos(t))
yn_inverse = np.outer(x_inverse, np.sin(t))
zn_inverse = np.zeros_like(xn_inverse)
for i in range(len(x_inverse)):
zn_inverse[i:i + 1, :] = np.full_like(zn_inverse[0, :], y_inverse[i])
ax3.plot(x_inverse, y_inverse)
ax3.set_title("Inverse of $f(x)$")
ax4.plot_surface(xn_inverse, yn_inverse, zn_inverse)
ax4.set_title("$f(x)$: Revolution around $x$")
plt.tight_layout()
plt.show()

How to Add another subplot to show Solid of Revolution toward x-axis?

I have this code modified from the topic here:
How to produce a revolution of a 2D plot with matplotlib in Python
The plot contains a subplot in the XY plane and another subplot of the solid of revolution toward the y-axis.
I want to add another subplot that is the solid of revolution toward the x-axis + how to add a legend to each subplot (above them), so there will be 3 subplots.
This is my MWE:
# Compare the plot at xy axis with the solid of revolution
# For function x=(y-2)^(1/3)
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
n = 100
fig = plt.figure(figsize=(12,6))
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122,projection='3d')
y = np.linspace(np.pi/8, np.pi*40/5, n)
x = (y-2)**(1/3) # x = np.sin(y)
t = np.linspace(0, np.pi*2, n)
xn = np.outer(x, np.cos(t))
yn = np.outer(x, np.sin(t))
zn = np.zeros_like(xn)
for i in range(len(x)):
zn[i:i+1,:] = np.full_like(zn[0,:], y[i])
ax1.plot(x, y)
ax2.plot_surface(xn, yn, zn)
plt.show()
Option 1:
Simply reverse x and y to switch the axes of the function.
x = np.linspace(np.pi/8, np.pi*40/5, n)
y = (x-2)**(1/3)
Option 2:
It is a little complicated. You can also accomplish this by finding the inverse of the original function.
The inverse of f(x) = y = x^3 + 2 is f^{-1}(y) = (y - 2)^(1/3).
I modified the code you provided.
import matplotlib.pyplot as plt
import numpy as np
n = 100
fig = plt.figure(figsize=(14, 7))
ax1 = fig.add_subplot(221)
ax2 = fig.add_subplot(222, projection='3d')
ax3 = fig.add_subplot(223)
ax4 = fig.add_subplot(224, projection='3d')
y = np.linspace(np.pi / 8, np.pi * 40 / 5, n)
x = (y - 2) ** (1 / 3)
t = np.linspace(0, np.pi * 2, n)
xn = np.outer(x, np.cos(t))
yn = np.outer(x, np.sin(t))
zn = np.zeros_like(xn)
for i in range(len(x)):
zn[i:i + 1, :] = np.full_like(zn[0, :], y[i])
ax1.plot(x, y)
ax1.set_title("$f(x)$")
ax2.plot_surface(xn, yn, zn)
ax2.set_title("$f(x)$: Revolution around $y$")
# find the inverse of the function
x_inverse = y
y_inverse = np.power(x_inverse - 2, 1 / 3)
xn_inverse = np.outer(x_inverse, np.cos(t))
yn_inverse = np.outer(x_inverse, np.sin(t))
zn_inverse = np.zeros_like(xn_inverse)
for i in range(len(x_inverse)):
zn_inverse[i:i + 1, :] = np.full_like(zn_inverse[0, :], y_inverse[i])
ax3.plot(x_inverse, y_inverse)
ax3.set_title("Inverse of $f(x)$")
ax4.plot_surface(xn_inverse, yn_inverse, zn_inverse)
ax4.set_title("$f(x)$: Revolution around $x$")
plt.tight_layout()
plt.show()

Animating two figures with different data units using matplotlib

Trying to plot two separate animations, i.e. in different windows as separate figures. Running this code for me rightly creates two windows, but animates the data on the second figure at the same time. Closing figure 1 results in only the intended data for figure 2 being animated, removing the overlap from the data intended for figure 1. Closing figure 2 results in only the intended data for figure 1 being animated, removing the overlap from the data intended for figure 2.
Minimum code below:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
dx, dv, N, Nb, decp = 2, 1.5, 100, 12, int(1)
Pd = np.zeros([N + 1, 2 * Nb])
Vd = np.zeros([N + 1, 2 * Nb])
Pd[:, 1] = 4
Vd[:, 3] = 2
t = np.zeros(N + 1)
t[0] = 0
for i in range(0, N):
t[i + 1] = (i + 1) * 0.1
Px = []
for i in range(0, (2 * Nb)):
PX = dx * (-Nb + i) / 4
Px.append(PX)
lblx = []
for i in range(0, int((Nb / 2) + 1)):
if i == (Nb / 4):
LBL = r"$\mu_x$"
lblx.append(LBL)
else:
LBL = r"${0}\sigma_x$".format(-(Nb / 4) + i)
lblx.append(LBL)
Pv = []
for i in range(0, (2 * Nb)):
PV = dv * (-Nb + i) / 4
Pv.append(PV)
lblv = []
for i in range(0, int((Nb / 2) + 1)):
if i == (Nb / 4):
LBL = r"$\mu_v$"
lblv.append(LBL)
else:
LBL = r"${0}\sigma_v$".format(-(Nb / 4) + i)
lblv.append(LBL)
fig1 = plt.figure(figsize=(8,6))
def animatex(i):
fig1.clear()
plt.bar(Px, Pd[i, :], width = dx / 4, align = 'edge', color = 'b', \
label = 't = {} seconds'.format(round(t[i], decp)))
s_ticks = np.arange(-3 * dx, (3 + 1) * dx, dx)
plt.xticks(s_ticks, lblx)
plt.ylim(0, np.max(Pd))
plt.xlim(-3 * dx, 3 * dx)
plt.legend()
plt.draw()
anix = FuncAnimation(fig1, animatex, repeat = True, interval = 200, frames = N + 1)
fig2 = plt.figure(figsize=(8,6))
def animatev(i):
fig2.clear()
plt.bar(Pv, Vd[i, :], width = dv / 4, align = 'edge', color = 'b', \
label = 't = {} seconds'.format(round(t[i], decp)))
s_ticks = np.arange(-3 * dv, (3 + 1) * dv, dv)
plt.xticks(s_ticks, lblv)
plt.ylim(0, np.max(Vd))
plt.xlim(-3 * dv, 3 * dv)
plt.legend()
plt.draw()
aniv = FuncAnimation(fig2, animatev, repeat = True, interval = 200, frames = N + 1)
plt.show()
As is probably clear, they are two bar plots, with different vertical and horizontal dimensions. I've seen some solutions for these kinds of problems where the data shares an axis through a shared variable, but here they are not (as can be seen).
For this minimum code, the solution involves having the two bars, one in Pd and the other in Vd, being on their respective intended figures, not both on the second figure.
Let me know if there are any issues with the information here i.e. minimal code requirements not met, more information etc. and I will update.
Ignore any wayward writing style, it is not relevant.
Simplifying your code:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
dx, dv, N, Nb, decp = 2, 1.5, 10, 12, int(1)
Px = np.arange(Nb)
Pd = np.random.randn(N, Nb)
Vd = np.random.randn(N, Nb)
fig1, ax1 = plt.subplots(figsize=(8, 6))
def animatex(i):
ax1.clear()
ax1.bar(Px, Pd[i, :], width=dx / 4, align='edge', color='b')
anix = FuncAnimation(fig1, animatex, repeat=True, interval=200, frames=N)
fig2, ax2 = plt.subplots(figsize=(8, 6))
def animatev(i):
ax2.clear()
ax2.bar(Px, Vd[i, :], width = dv / 4, align='edge', color='b')
aniv = FuncAnimation(fig2, animatev, repeat=True, interval=200, frames=N)
plt.show()
works fine for me. You can add the esthetic/data details back in...

Way to contour outer edge of selected grid region in Python

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:

Find and draw regression plane to a set of points

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

Categories