8.2 陰解法

#!/opt/local/bin/python3.5
# -*- coding: utf-8 -*-
# heat1d-i.py
# http://d.hatena.ne.jp/Megumi221/20080306/1204770689 を参考にした

import numpy as np
import matplotlib.pyplot as plot

def trilu(n, al, ad, au):
    for i in range(0,n-1):
        al[i+1] = al[i+1] / ad[i]
        ad[i+1] = ad[i+1] - au[i] * al[i+1]

def trisol(n, al, ad, au, b):
    nm1 = n-1
    for i in range(0,nm1):
        b[i+1] = b[i+1] - b[i] * al[i+1]
    b[nm1] = b[nm1] / ad[nm1]
    for i in range(n-2,-1,-1):
        b[i] = (b[i] - au[i] * b[i+1]) / ad[i]

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  # lambda は予約語で使えない？
tau=lam*h*h
theta=0.5
Tmax=1.0
nmax=int(Tmax/tau)
dt=0.001
skip=int(dt/tau)

plot.ion()
line,=plot.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)
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])
    trisol(N-1,al,ad,au,b)
    u[1:N]=b
    if n%skip == 0:
        line.set_ydata(u)
        plot.pause(0.001)