Matplotlib动画:如何同时绘制线条?

时间:2016-01-10 10:25:48

标签: python python-3.x animation numpy matplotlib

我想使用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()

1 个答案:

答案 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]