具有多个粒子的 FuncAnimation

Question

以下是我已经研究了几个星期的模型，随着我学习如何编码（编码新手），它逐渐增加了复杂性。如果不清楚，我正在尝试创建一个具有 2 种不同密度的粒子模型，这些粒子根据斯托克斯沉降速度（作为浓度的函数）进行沉降。我无法让动画与第二个粒子一起工作（工作代码在底部有一个动画粒子）。我已经分解了两个不同粒子的变量以尝试调试代码，但没有任何运气确定我做错了什么。

任何帮助将不胜感激！

import numpy as np
import matplotlib.pyplot as plt
from pylab import *
from matplotlib.animation import FuncAnimation
import random
import pdb
from scipy import spatial
from scipy.spatial import KDTree

n = 100
n2 = 50
sigma = 0.01
sigma2 = 0.01

pp = 2.56 # pp = particle density (Sphene=3.53) #(Feldspar=2.56) 
#(g/cm^3)

pp2 = 3.53

pf = 2.7 # pf = fluid density(g/cm^3)

pf2 = 2.7

g = 9.8 # g = gravity (m/s^2)

g2 = 9.8

r = 0.003 # r = radius of sphere (meter)

r2 = 0.0002

mu = 0.53 # mu = dynamic viscosity of fluid (log10Poise)

mu2 = 0.53

rp = 0.01 #radius around particle to check for nearest neighbor

rp2 = 0.001

dt = 0.008

dt2 = 0.008

fig, ax = plt.subplots()

az = plt.axes(xlim=(-.5, .5), ylim=(-1, 0))

xdata, ydata = [0.0], [0.0]

xdata2, ydata2 = [0.0], [0.0]

ln, = plt.plot([], [], marker= 'o', 
markerfacecolor='r',markeredgecolor='k', linestyle='None', 
animated=True)

ln2, = plt.plot([], [], marker= 's', 
markerfacecolor='b',markeredgecolor='k', linestyle='None', 
animated=True)

#pdb.set_trace()
def v_stokes(pp,pf,g,r,mu):
    top=2*(pp-pf)*g*(r**2)
    bottom=9*mu
    ws=top/bottom
    return ws

def v_stokes2(pp2,pf2,g2,r2,mu2):
    top2=2*(pp2-pf2)*g2*(r2**2)
    bottom2=9*mu2
    ws2=top2/bottom2
    return ws2


def init():
    ax.set_xlim(  -2, 2)
    ax.set_ylim(-10, 0)
    return ln, ln2,

def concentration(xdata, ydata, rp):
    coords = list(zip(xdata, ydata))
    tree = spatial.KDTree(coords)
    test = np.column_stack([xdata, ydata])
    nnl = tree.query_ball_point(test, rp) #nearest neighbors as a list 
    #(had tree in here before test but shape was wrong)
    #pdb.set.trace()
    nnt = np.zeros(len(nnl)) #nearest neighbors total
    for i in range(len(nnt)): 
        nnt[i] = len(nnl[i])
    return nnt

def concentration2(xdata2, ydata2, rp2):
    coords2 = list(zip(xdata2, ydata2))
    tree2 = spatial.KDTree(coords2)
    test2 = np.column_stack([xdata2, ydata2])
    nnl2 = tree2.query_ball_point(test2, rp2)  
    nnt2 = np.zeros(len(nnl2)) #nearest neighbors total
    for i in range(len(nnt2)): 
        nnt2[i] = len(nnl2[i])
    return nnt2

#y0 = []
#y1 = []
#y2 = []
#y3 = []
#y4 = []
def update(frame):

    global xdata
    global ydata
    global concentration, v_stokes, pp, pf, g, r, mu, rp, dt, n, sigma

    xdata = xdata + np.random.normal(0, sigma, n)

    wss = v_stokes(pp,pf,g,r,mu)

    if frame == 0.0:
        ydata = np.zeros(len(xdata)) #makes the ydata length = xdata at 
        #time 0 print(ydata)
        rp = 0.003

    if frame > 10:
        rp = 0.008

    cp = concentration(xdata, ydata, rp)
    if np.any(cp == 1):
        cp = cp-1

    if frame > 0.0:
        #y0.append(ydata)
        #y1.append(wss)
        #y2.append(cp)
        #y3.append(dt)
        #y4.append(xdata)
        ydata = ydata + (wss*(1-cp)**3) - dt # [0] 
        for v in ydata:
            if v < -1.001:
                ydata = -1
    ln.set_data(xdata, ydata)
    return ln,


def update2(frame2):

    global xdata2
    global ydata2
    global concentration2, v_stokes2, pp2, pf2, g2, r2, mu2, rp2, dt2, 
    n2, sigma2

    xdata2 = xdata2 + np.random.normal(0, sigma2, n2)

    wss2 = v_stokes2(pp2,pf2,g2,r2,mu2)

    if frame2 == 0.0:
        ydata2 = np.zeros(len(xdata2)) #makes the ydata length = xdata 
        #at time 0 print(ydata)
        rp2 = 0.003

    if frame2 > 10:
        rp2 = 0.008

    cp2 = concentration2(xdata2, ydata2, rp2)
    if np.any(cp2 == 1):
        cp2 = cp2-1

    if frame2 > 0.0:
        #y5.append(ydata2)
        #y6.append(wss2)
        #y7.append(cp2)
        #y8.append(dt2)
        #y9.append(xdata2)

        ydata2 = ydata2 + (wss2*(1-cp2)**3) - dt2 # [0] 
        for v2 in ydata2:
            if v2 < -1.001:
                ydata2 = -1

     ln2.set_data(xdata2, ydata2)
     return ln2,

def update_all(i):
    l1 = update(i)
    l2 = update2(i)
    return l1, l2,

ani = FuncAnimation(fig, update_all, frames=range(0,200),
                    init_func=init, blit=True, interval=100, repeat = 
False)
# change frames=range(0:1000) to change the number of frames 

plt.show()

以下是带有一个动画粒子的原始工作代码：

import numpy as np
import matplotlib.pyplot as plt
from pylab import *
from matplotlib.animation import FuncAnimation
import random
import pdb
from scipy import spatial
from scipy.spatial import KDTree

n=100
sigma= 0.01


#m = np.random.uniform(n)

pp = 2.56 # pp = particle density (Sphene=3.53) #(Feldspar=2.56) 
#(g/cm^3)

pf = 2.7 # pf = fluid density(g/cm^3)

g = 9.8 # g = gravity (m/s^2)

r = 0.003 # r = radius of sphere (meter)

mu = 0.53 # mu = dynamic viscosity of fluid (log10Poise)

rp = 0.01 #radius around particle to check for nearest neighbor

dt =0.008

fig, ax = plt.subplots()
az = plt.axes(xlim=(-.5, .5), ylim=(-1, 0))
xdata, ydata = [0.0], [0.0]
ln, = plt.plot([], [], marker= 'o', 
markerfacecolor='r',markeredgecolor='k', linestyle='None', 
animated=True)

#pdb.set_trace()
def v_stokes(pp,pf,g,r,mu):
    top=2*(pp-pf)*g*(r**2)
    bottom=9*mu
    ws=top/bottom
    return ws


def init():
    ax.set_xlim(  -2, 2)
    ax.set_ylim(-10, 0)
    return ln,

def concentration(xdata, ydata, rp):
    coords = list(zip(xdata, ydata))
    tree = spatial.KDTree(coords)
    test = np.column_stack([xdata, ydata])
    nnl = tree.query_ball_point(test, rp) #nearest neighbors as a list 
    #(had tree in here before test but shape was wrong)
    #pdb.set.trace()
    nnt = np.zeros(len(nnl)) #nearest neighbors total
    for i in range(len(nnt)): 
        nnt[i] = len(nnl[i])
    return nnt

y0 = []
y1 = []
y2 = []
y3 = []
y4 = []
def update(frame):

    global xdata
    global ydata
    global concentration, v_stokes, pp, pf, g, r, mu, rp, dt, n

    xdata = xdata + np.random.normal(0, sigma, n)

    wss = v_stokes(pp,pf,g,r,mu)
    #print(wss)

    if frame == 0.0:
        ydata = np.zeros(len(xdata)) #makes the ydata length = xdata at 
        #time 0 print(ydata)
        rp = 0.003

    if frame > 10:
        rp = 0.008

    if frame > 30:
        rp = 0.01

    if frame > 50: #notice that the particles move up instead of just 
    #slowing down
        rp = 0.013

    cp = concentration(xdata, ydata, rp)
    #print(cp)
    if np.any(cp == 1):
    cp = cp-1
    #print(xdata)
    print(cp)
    if frame > 0.0:
        y0.append(ydata)
        y1.append(wss)
        y2.append(cp)
        y3.append(dt)
        y4.append(xdata)
        ydata = ydata + (wss*(1-cp)**3) - dt # [0] 
        for v in ydata:
            if v < -1.001:
                ydata = -1
        #if np.all(ydata) > 0:
            #print(ydata)
            #print(frame)
    #print(ydata[0:5])
    #ydata = ydata + (wss*(1-cp))
    ln.set_data(xdata, ydata)
    return ln,

#print(test)
ani = FuncAnimation(fig, update, frames=range(0,200),
                init_func=init, blit=True, interval=100, repeat = 
False)
# change frames=range(0:1000) to change the number of frames 

plt.show()

Answer 1

matplotlib.animation.FuncAnimation(fig, func, ...) 期望它的第二个参数是一个可调用的（例如一个函数）来更新绘图。

func：可调用
在每一帧调用的函数。第一个参数将是帧中的下一个值。任何额外的位置参数都可以通过 fargs 参数提供。

如果您有两个函数要更新，则需要将它们合并为一个，

def update_all(i):
    l1 = update(i)
    l2 = update2(i)
    return l1,l2

ani = FuncAnimation(fig, update_all, frames=range(0,200), init_func=init, ...)

请注意，您需要像往常一样定义plots，ln, = plt.plot，并且您需要在两个更新函数中返回线对象本身。

完整的运行代码：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy import spatial

n=100
sigma= 0.01
pp = 2.56 # pp = particle density (Sphene=3.53) #(Feldspar=2.56) 
#(g/cm^3)
pp2 = 3.53
pf = 2.7 # pf = fluid density(g/cm^3)
pf2 = 2.7
g = 9.8 # g = gravity (m/s^2)
g2 = 9.8
r = 0.003 # r = radius of sphere (meter)
r2 = 0.0002
mu = 0.53 # mu = dynamic viscosity of fluid (log10Poise)
mu2 = 0.53
rp = 0.01 #radius around particle to check for nearest neighbor
rp2 = 0.001
dt = 0.008
dt2 = 0.008

fig, ax = plt.subplots()

az = plt.axes(xlim=(-.5, .5), ylim=(-1, 0))

xdata, ydata = [0.0], [0.0]
xdata2, ydata2 = [0.0], [0.0]

ln, = plt.plot([], [], marker= 'o', 
               markerfacecolor='r',markeredgecolor='k', linestyle='None', 
               animated=True)

ln2, = plt.plot([], [], marker= 's', 
                markerfacecolor='b',markeredgecolor='k', linestyle='None', 
                animated=True)

#pdb.set_trace()
def v_stokes(pp,pf,g,r,mu):
    top=2*(pp-pf)*g*(r**2)
    bottom=9*mu
    ws=top/bottom
    return ws

def v_stokes2(pp2,pf2,g2,r2,mu2):
    top2=2*(pp2-pf2)*g2*(r2**2)
    bottom2=9*mu2
    ws2=top2/bottom2
    return ws2


def init():
    ax.set_xlim(  -2, 2)
    ax.set_ylim(-10, 0)
    return ln, ln2

def concentration(xdata, ydata, rp):
    coords = list(zip(xdata, ydata))
    tree = spatial.KDTree(coords)
    test = np.column_stack([xdata, ydata])
    nnl = tree.query_ball_point(test, rp)  
    nnt = np.zeros(len(nnl)) #nearest neighbors total
    for i in range(len(nnt)): 
        nnt[i] = len(nnl[i])
    return nnt

def concentration2(xdata2, ydata2, rp2):
    coords2 = list(zip(xdata2, ydata2))
    tree2 = spatial.KDTree(coords2)
    test2 = np.column_stack([xdata2, ydata2])
    nnl2 = tree2.query_ball_point(test2, rp2) 
    nnt2 = np.zeros(len(nnl2)) #nearest neighbors total
    for i in range(len(nnt2)): 
        nnt2[i] = len(nnl2[i])
    return nnt2

y0 = [];y1 = [];y2 = [];y3 = [];y4 = []

def update(frame):

    global xdata
    global ydata
    global concentration, v_stokes, pp, pf, g, r, mu, rp, dt, n

    xdata = xdata + np.random.normal(0, sigma, n)

    wss = v_stokes(pp,pf,g,r,mu)

    if frame == 0.0:
        ydata = np.zeros(len(xdata)) #makes the ydata length = xdata at 
        #time 0 print(ydata)
        rp = 0.003

    if frame > 10:
        rp = 0.008

    if frame > 30:
        rp = 0.01

    if frame > 50: #notice that the particles move up instead of just 
    #slowing down
        rp = 0.013

    cp = concentration(xdata, ydata, rp)
    if np.any(cp == 1):
        cp = cp-1
    #print(cp)
    if frame > 0.0:
        y0.append(ydata)
        y1.append(wss)
        y2.append(cp)
        y3.append(dt)
        y4.append(xdata)
        ydata = ydata + (wss*(1-cp)**3) - dt # [0] 
        for j,v in enumerate(ydata):
            if v < -1.001:
                ydata[j] = -1
    ln.set_data(xdata, ydata)
    return ln


def update2(frame2):

    global xdata2
    global ydata2
    global concentration2, v_stokes, pp2, pf2, g2, r2, mu2, rp2, dt2, n

    xdata2 = xdata2 + np.random.normal(0, sigma, n)

    wss2 = v_stokes2(pp2,pf2,g2,r2,mu2)

    if frame2 == 0.0:
        ydata2 = np.zeros(len(xdata2)) #makes the ydata length = xdata 
        #at time 0 print(ydata)
        rp2 = 0.003

    if frame2 > 10:
        rp2 = 0.008

    if frame2 > 30:
        rp2 = 0.01

    if frame2 > 50: #notice that the particles move up instead of just 
    #slowing down
        rp2 = 0.013

    cp2 = concentration2(xdata2, ydata2, rp2)
    if np.any(cp2 == 1):
        cp2 = cp2-1

    if frame2 > 0.0:
        y0.append(ydata2)
        y1.append(wss2)
        y2.append(cp2)
        y3.append(dt2)
        y4.append(xdata2)
        ydata2 = ydata2 + (wss2*(1-cp2)**3) - dt2 # [0] 
        for j,v in enumerate(ydata2):
            if v < -1.001:
                ydata2[j] = -1

    ln2.set_data(xdata2, ydata2)
    return ln2

def update_all(i):
    l1 = update(i)
    l2 = update2(i)
    return l1,l2

ani = FuncAnimation(fig, update_all, frames=range(0,200),
                    init_func=init, blit=True, interval=100, repeat = False)

plt.show()

由于这两个粒子类的代码似乎大多是多余的，我建议通过将其放入一个类中来简化它。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy import spatial

n=100
sigma= 0.01
pp = 2.56 # pp = particle density (Sphene=3.53) #(Feldspar=2.56) 
#(g/cm^3)
pp2 = 3.53
pf = 2.7 # pf = fluid density(g/cm^3)
pf2 = 2.7
g = 9.8 # g = gravity (m/s^2)
g2 = 9.8
r = 0.003 # r = radius of sphere (meter)
r2 = 0.0002
mu = 0.53 # mu = dynamic viscosity of fluid (log10Poise)
mu2 = 0.53
rp = 0.01 #radius around particle to check for nearest neighbor
rp2 = 0.001
dt = 0.008
dt2 = 0.008


class Particles():
    def __init__(self,ax, xdata,
                 pp, pf, g, r, mu, rp, dt, n,
                 marker="o",mfc="r"):
        self.xdata=xdata
        self.ydata = np.zeros(len(self.xdata))
        self.pp=pp;self.pf=pf;self.g=g;self.r=r;self.mu=mu
        self.rp=rp;self.dt=dt;self.n=n
        self.ln, = ax.plot([], [], marker= marker, 
               markerfacecolor=mfc,markeredgecolor='k', linestyle='None', 
               animated=True)
           
    def v_stokes(self, pp,pf,g,r,mu):
        top=2*(pp-pf)*g*(r**2)
        bottom=9*mu
        return top/bottom
    
    def concentration(self,xdata, ydata, rp):
        coords = list(zip(xdata, ydata))
        tree = spatial.KDTree(coords)
        test = np.column_stack([xdata, ydata])
        nnl = tree.query_ball_point(test, rp)  
        nnt = np.zeros(len(nnl)) #nearest neighbors total
        for i in range(len(nnt)): 
            nnt[i] = len(nnl[i])
        return nnt
    
    def update(self, frame):
        self.xdata += np.random.normal(0, sigma, n)
        if frame == 0:
            self.ydata = np.zeros(len(self.xdata))
    
        wss = self.v_stokes(self.pp,self.pf,self.g,self.r,self.mu)

        cp = self.concentration(self.xdata, self.ydata, self.rp)
        if np.any(cp == 1):
            cp = cp-1

        if frame > 0.0:
            self.ydata += (wss*(1-cp)**3) - dt 
            for j,v in enumerate(self.ydata):
                if v < -1.001:
                    self.ydata[j] = -1
        self.ln.set_data(self.xdata, self.ydata)
        return self.ln

fig, ax = plt.subplots()

xdata, xdata2 = [0.0], [0.0]

p1 = Particles(ax,xdata, pp, pf, g, r, mu, rp, dt, n,
                 marker="o",mfc="r")
p2 = Particles(ax,xdata2, pp2, pf2, g2, r2, mu2, rp2, dt2, n,
                 marker="s",mfc="b")

def init():
    ax.set_xlim(-2, 2)
    ax.set_ylim(-1, 1)
    return p1.ln, p2.ln

def update(i):
    l1 = p1.update(i)
    l2 = p2.update(i)
    return l1, l2

ani = FuncAnimation(fig, update, frames=range(0,200),
                    init_func=init, blit=True, interval=100, repeat = False)

plt.show()