5.4.3 heat1d-i-v3.py

単に陰解法にしたバージョン。

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

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
dt=0.001
skip=int(dt/tau)

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

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)
line,=window.plot(x,u)

def update(i):
    global n
    if n<nmax:
        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])
        n=n+skip
        line.set_ydata(u)
    else:
        sys.exit(0)
    return line,

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