7.3.5 heat1d-i-v2.py

#!/usr/bin/env python3
# heat1d-i-v2.py
# http://d.hatena.ne.jp/Megumi221/20080306/1204770689 を参考にした
# http://www.scipy.org/Cookbook/Matplotlib/Animations

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

def f(x):
    return min(x,1-x)

N=50
x=np.linspace(0.0, 1.0, N+1)
# vf=np.vectorize(f)
# u=vf(x)
u=np.vectorize(f)(x)

h=1.0/N
lam=0.5  # lambda は予約語で使えない？
tau=lam*h*h
theta=0.5
Tmax=1.0
nmax=int(Tmax/tau)
dt=0.001
skip=int(dt/tau)

plt.ion()
line,=plt.plot(x,u)

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)

for n in range(1,nmax):
    b=(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,b)
    u[1:N]=b
    if n%skip == 0:
        line.set_ydata(u)
        plt.pause(0.0001)

もう少しケチれるかな。b を消して余計なコピーを無くせる。

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,al,ad,au,u[1:N])