5.4.4 heat1d-i-v4.py

意味は分からないけれど， Matplotlib Animation Tutorial の真似をして書き換えてみた。

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

import matplotlib
matplotlib.use('TkAgg')
import sys
import numpy as np
import matplotlib.pyplot as plot
import matplotlib.animation as animation
import trid

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
theta=0.5

Tmax=1
dt=0.001
skip=int(dt/tau)
nframes=int(Tmax/dt)
print nframes

ad = (1+2*theta*lam)*np.ones(N-1)
al = -theta*lam*np.ones(N-1)
au = -theta*lam*np.ones(N-1)
trid.trilu(N-1,al,ad,au)

fig=plot.figure()
window=fig.add_subplot(111)
ax = plot.axes(xlim=(-0.02, 1.02), ylim=(-0.02, 1.02))
line, = ax.plot([], [], lw=1)

def init():
    line.set_data([], [])
    return line,

def update(i):
    if i==0:
        line.set_data(x,u)
    else:
        for j in xrange(skip):
            u[1:N]=(1-2*(1-theta)*lam)*u[1:N]+(1-theta)*lam*(u[0:N-1]+u[2:N+1])
            trid.trisol(N-1,al,ad,au,u[1:N])
        line.set_data(x,u)
    return line,

ani=animation.FuncAnimation(fig,
                            update,
                            frames=nframes, init_func=init,
                            interval=1, blit=True, repeat=False)
plot.show()