matplotlibでキャビティ流れを可視化し、gifアニメで保存する

を見て、キャビティ流れの書き方がわかったが、アニメーションしないので物足りない。

最終結果だけを表示するコードは以下。（元ネタではiPython Notebookを使ってやっている）

cavorg.py

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation

nx = 41
ny = 41
nt = 5
nit=50
c = 1
dx = 2.0/(nx-1)
dy = 2.0/(ny-1)
x = np.linspace(0,2,nx)
y = np.linspace(0,2,ny)
X,Y = np.meshgrid(x,y)

rho = 1
nu = .1
dt = .001

u = np.zeros((ny, nx))
v = np.zeros((ny, nx))
p = np.zeros((ny, nx)) 
b = np.zeros((ny, nx))


def buildUpB(b, rho, dt, u, v, dx, dy):
    
    b[1:-1,1:-1]=rho*(1/dt*((u[1:-1,2:]-u[1:-1,0:-2])/(2*dx)+(v[2:,1:-1]-v[0:-2,1:-1])/(2*dy))-\
                      ((u[1:-1,2:]-u[1:-1,0:-2])/(2*dx))**2-\
                      2*((u[2:,1:-1]-u[0:-2,1:-1])/(2*dy)*(v[1:-1,2:]-v[1:-1,0:-2])/(2*dx))-\
                      ((v[2:,1:-1]-v[0:-2,1:-1])/(2*dy))**2)

    return b
    
def presPoisson(p, dx, dy, b):
    pn = np.empty_like(p)
    pn = p.copy()
    
    for q in range(nit):
        pn = p.copy()
        p[1:-1,1:-1] = ((pn[1:-1,2:]+pn[1:-1,0:-2])*dy**2+(pn[2:,1:-1]+pn[0:-2,1:-1])*dx**2)/\
                        (2*(dx**2+dy**2)) -\
                        dx**2*dy**2/(2*(dx**2+dy**2))*b[1:-1,1:-1]

        p[-1,:] =p[-2,:] ##dp/dy = 0 at y = 2
        p[0,:] = p[1,:]  ##dp/dy = 0 at y = 0
        p[:,0]=p[:,1]    ##dp/dx = 0 at x = 0
        p[:,-1]=0        ##p = 0 at x = 2
        
    return p
    
def cavityFlow(nt, u, v, dt, dx, dy, p, rho, nu):
    un = np.empty_like(u)
    vn = np.empty_like(v)
    b = np.zeros((ny, nx))
    
    for n in range(nt):
        un = u.copy()
        vn = v.copy()
        
        b = buildUpB(b, rho, dt, u, v, dx, dy)
        p = presPoisson(p, dx, dy, b)
        
        u[1:-1,1:-1] = un[1:-1,1:-1]-\
                        un[1:-1,1:-1]*dt/dx*(un[1:-1,1:-1]-un[1:-1,0:-2])-\
                        vn[1:-1,1:-1]*dt/dy*(un[1:-1,1:-1]-un[0:-2,1:-1])-\
                        dt/(2*rho*dx)*(p[1:-1,2:]-p[1:-1,0:-2])+\
                        nu*(dt/dx**2*(un[1:-1,2:]-2*un[1:-1,1:-1]+un[1:-1,0:-2])+\
                        dt/dy**2*(un[2:,1:-1]-2*un[1:-1,1:-1]+un[0:-2,1:-1]))

        v[1:-1,1:-1] = vn[1:-1,1:-1]-\
                        un[1:-1,1:-1]*dt/dx*(vn[1:-1,1:-1]-vn[1:-1,0:-2])-\
                        vn[1:-1,1:-1]*dt/dy*(vn[1:-1,1:-1]-vn[0:-2,1:-1])-\
                        dt/(2*rho*dy)*(p[2:,1:-1]-p[0:-2,1:-1])+\
                        nu*(dt/dx**2*(vn[1:-1,2:]-2*vn[1:-1,1:-1]+vn[1:-1,0:-2])+\
                        (dt/dy**2*(vn[2:,1:-1]-2*vn[1:-1,1:-1]+vn[0:-2,1:-1])))

        u[0,:] = 0
        u[:,0] = 0
        u[:,-1] = 0
        u[-1,:] = 1    #set velocity on cavity lid equal to 1
        v[0,:] = 0
        v[-1,:]=0
        v[:,0] = 0
        v[:,-1] = 0
   
    return u, v, p
    
fig = plt.figure(figsize=(11,7), dpi=100)
u, v, p = cavityFlow(nt, u, v, dt, dx, dy, p, rho, nu)

plt.contourf(X,Y,p,alpha=0.5)
plt.colorbar()
plt.contour(X,Y,p)
plt.quiver(X[::2,::2],Y[::2,::2],u[::2,::2],v[::2,::2])
plt.xlabel('X')
plt.ylabel('Y')
plt.show()

matplotlib.animationを使うとgifアニメが作れるらしいので挑戦してみたがうまくいかなかった。理由を以下で説明する。

gifアニメは、imagemagickが入っていれば作れる。これは以下参照。
http://kiito.hatenablog.com/entry/2013/11/27/233507

アニメーションさせるための具体的な方法は以下。

アニメーションをさせる方法は二通りのやり方があるようです。
animation.ArtistAnimation
...事前に用意してあるデータを描画
animation.FuncAnimation
...随時データを更新する

from http://cflat-inc.hatenablog.com/entry/2014/03/17/214719

その他の例
http://qiita.com/HirofumiYashima/items/5864fb9384053b77cf6f
http://d.hatena.ne.jp/xef/20120629/p2
http://stackoverflow.com/questions/24166236/add-text-to-image-animated-with-matplotlib-artistanimation

最初のシミュレーションコードはcavityFlow()の中でforループを回しているので、ループごとにデータを吐き出して、ArtistAnimationで描画しようと思ったのだが、描画用データの配列にcontourfを渡すとバグが出る。これは他のところでも報告されているのだが、対応されているのかどうかわからない。

http://matplotlib.1069221.n5.nabble.com/Matplotlib-1-1-0-animation-vs-contour-plots-td18703.html
http://stackoverflow.com/questions/23070305/how-can-i-make-an-animation-with-contourf

とりあえずArtistAnimationではなくFuncAnimationでやればcontourfでもアニメになるらしいが、ArtistAnimationでもできるよという情報もある。
contourfを除いて、流線だけ描写してみたところ無事作成できたので、これはやはりcontourfが原因のバグっぽい。