我想使用Python 3,NumPy和Matplotlib来模拟二维质量和弹簧网络。现在我尝试使用2个弹簧检查系统,但是动画存在问题(请参阅下面附带代码中的块ANIMATION
)。使用我的代码,我只能显示一个或另一个弹簧(将line[0]
更改为line [1]
),但我不知道如何同时绘制两条线。我将非常感谢您的帮助。
这是我的代码:
from scipy.linalg import norm
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
# INITIAL CONDITIONS
#==================
# equilibrium distance
a = 1.
# stiffness
c = 1.
# mass
m = 1.
# initial positions
r = np.array([[-1., 0.],[0.,0.],[1., 0.]])
v = np.array([[0., 0.],[0., 0.],[0., 0.]])
# calculation steps
Time = 1000
# integration step
dt = 0.01
def Linforce(r1, r2, c , a):
"""returns linear force (vector) between two particles with a stiffness 'c' and an equilibrium distance 'a'"""
dr = r2 - r1
abs_dr = norm(dr) #calculates norm of the vector
return c * dr/abs_dr * (abs_dr - a)
#def W(force, m):
# """returns the acceleration for given force field and mass"""
# return force / m
def Verlet(r1, r2, v1, v2, dt, force):
"""Verlet integration. Return new positions and velocities from current values, time step and force.
Parameters:
r is a numpy array giving the current position vector
v is a numpy array giving the current velocity vector
dt is a float value giving the length of the integration time step
force is a force
field
Works with arrays of any dimension as long as they're all the same.
"""
w_old = 1/m * Linforce(r1, r2, c , a) # initial acceleration
r1_new = r1 + v1*dt + w_old *dt**2/2
r2_new = r2 + v2*dt - w_old *dt**2/2 #3rd Newton's law: w2 = -w1
w_new = 1/m * Linforce(r1_new, r2_new, c , a) # new acceleration
v1_new = v1 + (w_old + w_new)/2 * dt
v2_new = v2 - (w_old + w_new)/2 * dt
return (r1_new, r2_new, v1_new, v2_new)
time = 0
coords = np.array(r) #array of initial coordinates
# INTEGRATION
#==================
while time < Time:
for i in np.arange(1, len(r)):
for j in np.arange(0, i):
if norm(r[i]-r[j]) < 2.1*a: #distance of interaction
r[i], r[j], v[i], v[j] = Verlet(r[i], r[j], v[i], v[j], dt, Linforce)
coords =np.concatenate((coords,r), axis = 1) #Join a sequence of coordinates at each integration step
time+=1
x = np.zeros((len(r), Time+1)) #Array of x coordinates for all particles and time steps (including zero step => Time+1)
y = np.zeros((len(r), Time+1)) #Array of y coordinates for all particles and time steps
for n in np.arange(len(r)):
x[n] = np.array(coords[n,0::2]) # array of x coordinates of n-th particle in time
y[n] = np.array(coords[n,1::2]) # array of y coordinates of n-th particle in time
# ANIMATION
#==================
fig = plt.figure()
ax = fig.add_subplot(111, autoscale_on=False, xlim=(-2, 2), ylim=(-2, 2))
ax.grid()
lines = []
line1, = ax.plot([], [], 'o-', lw=2)
line2, = ax.plot([], [], 'o-', lw=2)
time_template = 'time = %.1fs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
def init():
lines.append(line1)
lines.append(line2) # we need 2 lines in this example
time_text.set_text('')
return lines[1], time_text #change line[1] to line[0] to draw another spring
#return lines, time_text # DOES NOT WORK
def animate(i):
time_text.set_text(time_template%(i*dt))
thisx1 = [x[1][i], x[0][i]]
thisy1 = [y[1][i], y[0][i]]
thisx2 = [x[2][i], x[1][i]]
thisy2 = [y[2][i], y[1][i]]
lines[0].set_data(thisx1, thisy1)
lines[1].set_data(thisx2, thisy2)
return lines[1], time_text #change line[1] to line[0] to draw another spring
#return lines, time_text # DOES NOT WORK
ani = animation.FuncAnimation(fig, animate, np.arange(1, Time),
interval=25, blit=True, init_func=init)
plt.show()
答案 0 :(得分:0)
你可以将time_text附加到行:
def animate(i):
time_text.set_text(time_template%(i*dt))
thisx1 = [x[1][i], x[0][i]]
thisy1 = [y[1][i], y[0][i]]
thisx2 = [x[2][i], x[1][i]]
thisy2 = [y[2][i], y[1][i]]
lines[0].set_data(thisx1, thisy1)
lines[1].set_data(thisx2, thisy2)
return lines + [time_text]