I was able to simulate a mass-spring system under damped oscillations. However, I wanted to add a subplot of position vs time and another subplot velocity vs position (phase path) so that I will be having three animations. How can I add them? The source code that I used is shown below
Update: I tried adding the first subplot position vs time but I cannot get the desired curve for it.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
#Constants
w = np.pi #angular frequency
b = np.pi * 0.2 #damping parameter
#Function that implements rk4 integration
def rk4(t, x, f, dt):
dx1 = f(t, x)*dt
dx2 = f(t+0.5*dt, x+0.5*dx1)*dt
dx3 = f(t+0.5*dt, x+0.5*dx2)*dt
dx4 = f(t+dt, x+dx3)*dt
return x+dx1/6.0+dx2/3.0+dx3/3.0+dx4/6.0
#Function that returns dX/dt for the linearly damped oscillator
def dXdt(t, X):
x = X[0]
vx = X[1]
ax = -2*b*vx - w**2*x
return np.array([vx, ax])
#Initialize Variables
x0 = 5.0 #Initial x position
vx0 = 1.0 #Initial x Velocity
#Setting time array for graph visualization
dt = 0.05
N = 200
x = np.zeros(N)
vx = np.zeros(N)
y = []
# integrate equations of motion using rk4;
# X is a vector that contains the positions and velocities being integrated
X = np.array([x0, vx0])
for i in range(N):
x[i] = X[0]
vx[i] = X[1]
y.append(0)
# update the vector X to the next time step
X = rk4(i*dt, X, dXdt, dt)
fig, (ax1, ax2) = plt.subplots(2,1, figsize=(8,6))
fig.suptitle(r' Damped Oscillation with $\beta$$\approx$' + str(round(b,2)) + r' and $\omega$$\approx$'
+ str(round(w,2)), fontsize=16)
line1, = ax1.plot([], [], lw=10,c="blue",ls="-",ms=50,marker="s",mfc="gray",fillstyle="none",mec="black",markevery=2)
line2, = ax2.plot([], [], lw=2, color='r')
time_template = '\nTime = %.1fs'
time_text = ax1.text(0.1, 0.9, '', transform=ax1.transAxes)
for ax in [ax1, ax2]:
ax1.set_xlim(1.2*min(x), 1.2*max(x))
ax2.set_ylim(1.2*min(x), 1.2*max(x),1)
ax2.set_xlim(0, N*dt)
ax1.set_yticklabels([])
def init():
line1.set_data([], [])
line2.set_data([], [])
time_text.set_text('')
return line1, line2, time_text
def animate(i):
thisx1 = [x[i],0]
thisy1 = [y[i],0]
thisx2 = [i*dt,0]
thisy2 = [x[i],0]
line1.set_data(thisx1, thisy1)
line2.set_data(thisx2, thisy2)
time_text.set_text(time_template % (i*dt))
return line1, line2, time_text
ani = animation.FuncAnimation(fig, animate, np.arange(1, N),
interval=50, blit=True, init_func=init,repeat=False)
After some minors changes to your initial code, the most noteworthy being:
thisx1 = [x[i],0]
thisy1 = [y[i],0]
thisx2 = [i*dt,0]
thisy2 = [x[i],0]
line1.set_data(thisx1, thisy1)
line2.set_data(thisx2, thisy2)
# should be written like this
line1.set_data([x[i],0], [y[i],0])
line2.set_data(t[:i], x[:i])
line3.set_data(x[:i], vx[:i])
The working version, with phase space plot in green, is as follows:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
#Constants
w = np.pi #angular frequency
b = np.pi * 0.2 #damping parameter
#Function that implements rk4 integration
def rk4(t, x, f, dt):
dx1 = f(t, x)*dt
dx2 = f(t+0.5*dt, x+0.5*dx1)*dt
dx3 = f(t+0.5*dt, x+0.5*dx2)*dt
dx4 = f(t+dt, x+dx3)*dt
return x+dx1/6.0+dx2/3.0+dx3/3.0+dx4/6.0
#Function that returns dX/dt for the linearly damped oscillator
def dXdt(t, X):
x = X[0]
vx = X[1]
ax = -2*b*vx - w**2*x
return np.array([vx, ax])
#Initialize Variables
x0 = 5.0 #Initial x position
vx0 = 1.0 #Initial x Velocity
#Setting time array for graph visualization
dt = 0.05
N = 200
t = np.linspace(0,N*dt,N,endpoint=False)
x = np.zeros(N)
vx = np.zeros(N)
y = np.zeros(N)
# integrate equations of motion using rk4;
# X is a vector that contains the positions and velocities being integrated
X = np.array([x0, vx0])
for i in range(N):
x[i] = X[0]
vx[i] = X[1]
# update the vector X to the next time step
X = rk4(i*dt, X, dXdt, dt)
fig, (ax1, ax2, ax3) = plt.subplots(3,1, figsize=(8,12))
fig.suptitle(r' Damped Oscillation with $\beta$$\approx$' + str(round(b,2)) + r' and $\omega$$\approx$'
+ str(round(w,2)), fontsize=16)
line1, = ax1.plot([], [], lw=10,c="blue",ls="-",ms=50,marker="s",mfc="gray",fillstyle="none",mec="black",markevery=2)
line2, = ax2.plot([], [], lw=1, color='r')
line3, = ax3.plot([], [], lw=1, color='g')
time_template = '\nTime = %.1fs'
time_text = ax1.text(0.1, 0.9, '', transform=ax1.transAxes)
ax1.set_xlim(1.2*min(x), 1.2*max(x))
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax2.set_ylim(1.2*min(x), 1.2*max(x),1)
ax2.set_xlim(0, N*dt)
ax2.set_xlabel('t')
ax2.set_ylabel('x')
ax3.set_xlim(1.2*min(x), 1.2*max(x),1)
ax3.set_ylim(1.2*min(vx), 1.2*max(vx),1)
ax3.set_xlabel('x')
ax3.set_ylabel('vx')
def init():
line1.set_data([], [])
line2.set_data([], [])
line3.set_data([], [])
time_text.set_text('')
return line1, line2, line3, time_text
def animate(i):
line1.set_data([x[i],0], [y[i],0])
line2.set_data(t[:i], x[:i])
line3.set_data(x[:i], vx[:i])
time_text.set_text(time_template % (i*dt))
return line1, line2, line3, time_text
ani = animation.FuncAnimation(fig, animate, np.arange(1, N),
interval=50, blit=True, init_func=init,repeat=False)
ani.save('anim.gif')
and gives:
Related
Recently I have managed to do the animation. But now I can't make the tangent line went as I want (not even displayed yet). The formula of the tangent line is y=(2/r)(sqrt(1-((r^2)/4))-1)x +r. The formula is obtained from 2 circles equation (C1 and C2). C1(blue) : x^2+y^2=r^2, and C2(green) : (x-1)^2+y^2=1. My goal is to obtain this kind of animation as and my current animation goes like this .
How should the code looks like when the animation looks like the reference one (the first one)? Any comments and answers would be very helpful for me as a beginner, I appreciate it, Thank you.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
fig, ax = plt.subplots(1)
line, = ax.plot([], [], lw=2)
line2, = ax.plot([], [], lw=2)
ax.set_xlim(-5,5)
ax.set_ylim(-5,5)
# theta goes from 0 to 2pi
theta = np.linspace(0, 2*np.pi, 100)
# the radius of the circle
r = np.sqrt(1)
r2 = np.sqrt(4)
# compute x1 and x2
x1 = 1+r*np.cos(theta)
y1 = r*np.sin(theta)
x2 = r2*np.cos(theta)
y2 = r2*np.sin(theta)
# Move left y-axis and bottim x-axis to centre, passing through (0,0)
ax.spines['left'].set_position('center')
ax.spines['bottom'].set_position('center')
# Eliminate upper and right axes
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
# Show ticks in the left and lower axes only
ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_ticks_position('left')
def init():
line.set_data([], [])
return line,
def init2():
line.set_data([], [])
return line2,
def animate(i):
x2 = np.sqrt(i)*np.cos(theta)
y2 = np.sqrt(i)*np.sin(theta)
line.set_data(x2, y2)
return line,
def animate2(i):
x3 = np.linspace(0,r2**2,100)
y3 = ((2/r2)*(np.sqrt(1-(r2**2)/4)-1)*x3)+r2
line.set_data(x3, y3)
return line,
# create the figure
ax.plot(x1,y1)
ax.set_aspect(1)
plt.grid()
anim = animation.FuncAnimation(fig, animate, init_func=init,
interval=1000, blit=False,\
frames=np.hstack([range(0),
range(4)[::-1]]))
anim2 = animation.FuncAnimation(fig, animate2, init_func=init2,
interval=1000, blit=False)
plt.show()
f = r"D:/UNPAR/Semester 2/Pemrograman Komputer/Project/SK.gif"
writergif = animation.PillowWriter(fps=30)
anim.save(f, writer=writergif)
The animation functions need to be combined into one. We will combine them into an initialization function and an animation function.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
fig, ax = plt.subplots(1)
line, = ax.plot([], [], lw=2, color='green')
line2, = ax.plot([], [], lw=2, color='red')
ax.set_xlim(-5,5)
ax.set_ylim(-5,5)
# theta goes from 0 to 2pi
theta = np.linspace(0, 2*np.pi, 100)
# the radius of the circle
r = np.sqrt(1)
r2 = np.sqrt(4)
# compute x1 and x2
x1 = 1+r*np.cos(theta)
y1 = r*np.sin(theta)
x2 = r2*np.cos(theta)
y2 = r2*np.sin(theta)
# create the figure
ax.plot(x1,y1)
ax.set_aspect(1)
plt.grid()
def init():
line.set_data([], [])
line2.set_data([], [])
return line,line2
def animate(i):
x2 = np.sqrt(i)*np.cos(theta)
y2 = np.sqrt(i)*np.sin(theta)
print(max(y2))
line.set_data(x2, y2)
x3 = [0, 4-(0.1*i)]
y3 = [max(y2), 0]
line2.set_data(x3, y3)
return line,line2
anim = animation.FuncAnimation(fig, animate, init_func=init, interval=1000, blit=False,frames=np.arange(10,0,-1))
plt.show()
My Python code here:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.projections import PolarAxes
import mpl_toolkits.axisartist.grid_finder as gf
import mpl_toolkits.axisartist.floating_axes as fa
class TaylorDiagram(object):
def __init__(self, STD ,fig=None, rect=111, label='_'):
self.STD = STD
tr = PolarAxes.PolarTransform()
# Correlation labels
rlocs = np.concatenate(((np.arange(11.0) / 10.0), [0.95, 0.99]))
tlocs = np.arccos(rlocs) # Conversion to polar angles
gl1 = gf.FixedLocator(tlocs) # Positions
tf1 = gf.DictFormatter(dict(zip(tlocs, map(str, rlocs))))
# Standard deviation axis extent
self.smin = 0
self.smax = 1.6 * self.STD
gh = fa.GridHelperCurveLinear(tr,extremes=(0,(np.pi/2),self.smin,self.smax),grid_locator1=gl1,tick_formatter1=tf1,)
if fig is None:
fig = plt.figure()
ax = fa.FloatingSubplot(fig, rect, grid_helper=gh)
fig.add_subplot(ax)
# Angle axis
ax.axis['top'].set_axis_direction('bottom')
ax.axis['top'].label.set_text("Correlation coefficient")
ax.axis['top'].toggle(ticklabels=True, label=True)
ax.axis['top'].major_ticklabels.set_axis_direction('top')
ax.axis['top'].label.set_axis_direction('top')
# X axis
ax.axis['left'].set_axis_direction('bottom')
ax.axis['left'].label.set_text("Standard deviation")
ax.axis['left'].toggle(ticklabels=True, label=True)
ax.axis['left'].major_ticklabels.set_axis_direction('bottom')
ax.axis['left'].label.set_axis_direction('bottom')
# Y axis
ax.axis['right'].set_axis_direction('top')
ax.axis['right'].label.set_text("Standard deviation")
ax.axis['right'].toggle(ticklabels=True, label=True)
ax.axis['right'].major_ticklabels.set_axis_direction('left')
ax.axis['right'].label.set_axis_direction('top')
# Useless
ax.axis['bottom'].set_visible(False)
# Contours along standard deviations
ax.grid()
self._ax = ax # Graphical axes
self.ax = ax.get_aux_axes(tr) # Polar coordinates
# Add reference point and STD contour
l , = self.ax.plot([0], self.STD, 'k*', ls='', ms=12, label=label)
l1 , = self.ax.plot([0], self.STD, 'k*', ls='', ms=12, label=label)
t = np.linspace(0, (np.pi / 2.0))
t1 = np.linspace(0, (np.pi / 2.0))
r = np.zeros_like(t) + self.STD
r1 = np.zeros_like(t) + self.STD
self.ax.plot(t, r, 'k--', label='_')
# Collect sample points for latter use (e.g. legend)
self.samplePoints = [l]
self.samplePoints = [l1]
def add_sample(self,STD,r,*args,**kwargs):
l,= self.ax.plot(np.arccos(r), STD, *args, **kwargs) # (theta, radius)
self.samplePoints.append(l)
return l
def add_sample(self,STD,r1,*args,**kwargs):
l1,= self.ax.plot(np.arccos(r1), STD, *args, **kwargs) # (theta, radius)
self.samplePoints.append(l1)
return l1
def add_contours(self,levels=5,**kwargs):
rs, ts = np.meshgrid(np.linspace(self.smin, self.smax), np.linspace(0, (np.pi / 2.0)))
RMSE=np.sqrt(np.power(self.STD, 2) + np.power(rs, 2) - (2.0 * self.STD * rs *np.cos(ts)))
contours = self.ax.contour(ts, rs, RMSE, levels, **kwargs)
return contours
def srl(obsSTD, s, s1, r, r1, l, l1, fname):
fig=plt.figure(figsize=(8,8))
dia=TaylorDiagram(obsSTD, fig=fig, rect=111, label='ref')
plt.clabel(dia.add_contours(colors='#808080'), inline=1, fontsize=10)
srlc = zip(s, r, l)
srlc1 = zip(s1, r1, l1)
for i in srlc:
dia.add_sample(i[0], i[1], label=i[2], marker='o',mec = 'red', mfc = 'none', mew=1.6)
for i in srlc1:
dia.add_sample(i[0], i[1], label=i[2], marker='^', mec = 'blue', mfc = 'none', mew=1.6)
spl = [p.get_label() for p in dia.samplePoints]
fig.legend(dia.samplePoints, spl, numpoints=1, prop=dict(size='small'), loc=[0.7,0.3])
plt.show()
obsSTD = 1
s = [1.275412605,1.391302157,1.424314937]
s1 = [0.980035327,0.997244197, 1.003002031]
r = [0.572272,0.533529,0.477572]
r1 = [0.82,0.72,0.8]
l = ['A', 'B','C']
l1 = ['A', 'B','C']
fname = 'TaylorDiagram.jpg'
srl(obsSTD, s, s1, r, r1, l,l1, fname)
And give a result that: taylor_dig that the red is before processing, the blue for after one. How could I label the position of A, B, C models as 1,2,3 (enumerate) with the position of before and after processing as well as legends without repeating?
I am actually working on a Python Project where I need to animate variations over time of a concentration, to do so I need the values of my concentration at time t-1 to calculate the concentration at time t.
I would like the "animate" function to use the new values of n0 and n1 when being called. Presently it only uses the same values of n0 and n1 as I declared them in my initial conditions and do not take into account the changes it makes to n0 and n1 during the first plot.
Here is a sample of my code (K, L and M are constant floats):
# Initial conditions
n0 = f(X)
n1 = n0.copy()
n1[1:-1] = K * (n0[0:-2] - 2 * n0[1:-1] + n0[2:]) - L * (n0[2:] - n0[:-2]) + M * n0[1:-1]
n0 = n1.copy()
# Animation
fig = plt.figure()
ax = plt.axes(xlim=(-0.5, 1.3), ylim=(min(f(X)), max(f(X))), xlabel="x", ylabel="n", title="n en fonction de x")
line, = ax.plot([], [], lw=2)
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)
def init():
time_text.set_text('')
line.set_data([],[])
return line, time_text
def animate(i,n0,n1):
n1[1:-1] = K * (n0[0:-2] - 2 * n0[1:-1] + n0[2:]) - L * (n0[2:] - n0[:-2]) + M * n0[1:-1]
n0 = n1.copy()
t = i * dt
time_text.set_text('time = %f ms' % (t * 1000))
y = n1
line.set_data(X, y)
return line, time_text
ani = animation.FuncAnimation(fig, animate, init_func=init, frames=nt, blit=True, interval=1, repeat=False, fargs=(n0, n1))
plt.show()
Hope you can help me because I am a new user of animation.Funcanimation and I do not understand everything about it.
A rather crude way to achieve the desired result is to use global variable
Fmin2 = 0
Fmin1 = 1
fig,ax = plt.subplots(figsize=(4,4))
ax.set_xlim(0,200)
ax.set_ylim(0,200)
line, = ax.plot([], [], 'ro')
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)
def init():
global Fmin2,Fmin1
Fmin2 = 0
Fmin1 = 1
time_text.set_text('')
line.set_data([],[])
return line, time_text
def animate(i):
global Fmin2,Fmin1
Fn = Fmin1+Fmin2
Fmin2 = Fmin1
Fmin1 = Fn
time_text.set_text('i={:d}'.format(i))
line.set_data([Fmin1],[Fn])
return line, time_text
ani = animation.FuncAnimation(fig, animate, init_func=init, frames=10, blit=True)
plt.show()
A cleaner way would be to integrate your animation in a class, and have n0 and n1 be members of that class that are updated in the animate(self, i) function.
I've been playing with the animation module from matplotlib and I realized I couldn't efficiently make a sine wave loop between two limits (in this case between -180° and 180°).
Like this...
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
fig = plt.figure()
ax = plt.axes(xlim=(0, 0.04), ylim=(-1.5, 1.5))
# initialize moving plots
line1, = ax.plot([], [], linewidth=2, label='sine')
line2, = ax.plot([], [], label='cosine')
ax.legend()
ax.grid()
def animate(i):
step = np.pi/30
# loop by hand...
if i < 30:
phase = i*step
elif 30 <= i < 90:
phase = -i*step
elif 90 <= i < 150:
phase = i*step
elif 150 <= i < 210:
phase = -i*step
else:
phase = i*step
x = np.linspace(0, 0.04, 1000)
y1 = np.sin( 2*np.pi*50*x - phase )
y2 = 0.5*np.cos( 2*np.pi*50*x + phase )
line1.set_data(x, y1)
line2.set_data(x, y2)
print('i:',i) # debug i
return line1, line2
anim = animation.FuncAnimation(fig, animate, interval=250, blit=True)
plt.show()
The reason is because I'm using the i variable, that is used for the frames count and only increases with time. Is there a way to loop indefinitely without writing if conditions until the end of time?
From this answer I found that is posible to refresh the data from the plot, and I've manage to make it loop almost like I wanted.
Adapted example... (workaround not complete)
import matplotlib.pyplot as plt
import numpy as np
def Yvalue(t, phase):
"""Function to plot"""
w = 2*np.pi*50
return np.sin(w*t + phase)
plt.ion() # You probably won't need this if you're embedding things in a tkinter plot...
step = np.pi/30 # steps for phase shifting
t = np.linspace(0, 0.04) # x values
y1 = Yvalue(t, 0) # y values
# starts figure
fig = plt.figure()
ax = plt.axes(xlim=(0, 0.04), ylim=(-1.5, 1.5))
# Returns a tuple of line objects, thus the comma
line1, = ax.plot(t, y1, linewidth=2, label='sine')
# static plot (cosine)
ax.plot(t, np.cos(2*np.pi*50*t), label='cosine static')
ax.legend()
ax.grid()
# initial values
phase = 0
direction = 1 # 1: shifting plot to left; 0: shifting plot to right
UpperLimit = np.pi
LowerLimit = -np.pi
# magic begins...
for something in range(210):
# while 1:
if direction and phase < UpperLimit:
phase += step
direction = 1
else:
phase -= step
direction = 0
# condition that helps to return to left shifting
if phase < LowerLimit:
direction = 1
line1.set_ydata( Yvalue(t, phase) )
fig.canvas.draw()
The problem with is that it doesn't allow me to close the window like it would be with the animation module. Therefore the program must be killed manually when changing the for loop by the while loop.
You would usually not use the animating function itself to calculate its animating parameter. Instead you would provide that parameter as argument to it using the frames argument.
In this case you would want the animating function to take the phase as argument. To create the phase, which is a kind of sawtooth function you can use numpy like
a = np.linspace(0,1, 30, endpoint=False)
phase = np.concatenate((a, 1-a, -a, a-1))*np.pi
Complete example:
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
fig = plt.figure()
ax = plt.axes(xlim=(0, 0.04), ylim=(-1.5, 1.5))
line1, = ax.plot([], [], linewidth=2, label='sine')
line2, = ax.plot([], [], label='cosine')
ax.legend()
ax.grid()
x = np.linspace(0, 0.04, 1000)
a = np.linspace(0,1, 30, endpoint=False)
phase = np.concatenate((a, 1-a, -a, a-1))*np.pi
def animate(phase):
y1 = np.sin( 2*np.pi*50*x - phase )
y2 = 0.5*np.cos( 2*np.pi*50*x + phase )
line1.set_data(x, y1)
line2.set_data(x, y2)
return line1, line2
anim = animation.FuncAnimation(fig, animate, frames=phase, interval=50, blit=True)
plt.show()
I don't know if I understand your problem because I don't see problem to use second method (used in for loop) inside animate
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
fig = plt.figure()
ax = plt.axes(xlim=(0, 0.04), ylim=(-1.5, 1.5))
# initialize moving plots
line1, = ax.plot([], [], linewidth=2, label='sine')
line2, = ax.plot([], [], label='cosine')
ax.legend()
ax.grid()
# -------------------------------------------------
def func1(t, phase):
"""Function to plot"""
w = 2*np.pi*50
return np.sin( w*t + phase)
def func2(t, phase):
"""Function to plot"""
w = 2*np.pi*50
return np.sin( w*t - phase)
# -------------------------------------------------
t = np.linspace(0, 0.04)
step = np.pi/30
UpperLimit = np.pi
LowerLimit = -np.pi
direction = 1
phase = 0
def animate(i):
global direction
global phase
if direction:
phase += step
if phase >= UpperLimit:
direction = 0
else:
phase -= step
if phase < LowerLimit:
direction = 1
line1.set_data(t, func1(t, phase))
line2.set_data(t, func2(t, phase))
return line1, line2
anim = animation.FuncAnimation(fig, animate, interval=250, blit=True)
plt.show()
Or even without variable direction
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
fig = plt.figure()
ax = plt.axes(xlim=(0, 0.04), ylim=(-1.5, 1.5))
# initialize moving plots
line1, = ax.plot([], [], linewidth=2, label='sine')
line2, = ax.plot([], [], label='cosine')
ax.legend()
ax.grid()
# -------------------------------------------------
def func1(t, phase):
"""Function to plot"""
w = 2*np.pi*50
return np.sin( w*t + phase)
def func2(t, phase):
"""Function to plot"""
w = 2*np.pi*50
return np.sin( w*t - phase)
# -------------------------------------------------
t = np.linspace(0, 0.04)
step = np.pi/30
UpperLimit = np.pi
LowerLimit = -np.pi
phase = 0
def animate(i):
global phase
global step
phase += step
if phase >= UpperLimit or phase <= LowerLimit:
step = -step
line1.set_data(t, func1(t, phase))
line2.set_data(t, func2(t, phase))
return line1, line2
anim = animation.FuncAnimation(fig, animate, interval=250, blit=True)
plt.show()
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()