5.4.2 heat1d-v3.py

とりあえず FuncAnimation() を使って動いた最初のプログラム。

#!/opt/local/bin/python2
# -*- coding: utf-8 -*-
# heat1d-v3.py

import sys
import numpy as np
import matplotlib.pyplot as plot
import matplotlib.animation as animation

def f(x):
    return min(x,1-x)
N=50
x=np.linspace(0.0, 1.0, N+1)
u=np.vectorize(f)(x)

h=1.0/N
lam=0.5
tau=lam*h*h
dt=0.001
skip=int(dt/tau)

Tmax=1
nmax=int(Tmax/tau)
n=0

fig=plot.figure()
window=fig.add_subplot(111)
line,=window.plot(x,u)

def update(i):
    global n
    if n<nmax:
        for i in range(skip):
            u[1:N]=(1-2*lam)*u[1:N]+lam*(u[0:N-1]+u[2:N+1])
        n=n+skip
        line.set_ydata(u)
    else:
        sys.exit(0)
    return line,

ani=animation.FuncAnimation(fig, update, interval=1)
plot.show()