1次元拡散方程式のPython実装例

以下の1次元の拡散方程式の実装例です。$0\leq x\leq 1$, $t\geq 0$ とします。

$$\begin{align}
\frac{\partial u(x,t)}{\partial t}&=\frac{\partial^2 u(x,t)}{\partial x^2},\\
u(0,t) &= u(1,t) = 0,\\
u(x,0)&=x(1-x).
\end{align}$$

レポート課題として定番の問題なので、様々な実装例がネット検索で沢山出てきますが、Pythonらしい実装 (なるべく配列の要素を番号で指定しない書き方) が無さそうだったので作ってみました。

########### 準備と前処理 ###########

# 関連パッケージの読み込み
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal

# パラメータ設定
n_space = 100  # 空間方向の分割数
n_time  = 500  # 時間方向の更新回数
alpha   = 0.5  # 1ステップ当たりの変化の大きさ (dt / dx**2)
t_list  = [100, 200, 500]  # 図に描画する時点のリスト (0以外で)

# 0から1まで等間隔に並んだ点を生成
x = np.linspace(0, 1, n_space)  

# 初期条件の指定 x(1-x)
u = x * (1-x)

# カーネルの定義 (左からの流入、左右への流出、右からの流入)
kernel = alpha * np.array([1,-2,1])


########### メインの計算部分 ###########

# 結果を格納するリストを用意
result_list = []

# 時刻t=0と初期状態を結果リストに追加
result_list.append((0,u))

# 時間方向の繰り返し、1から開始
for t in np.arange(n_time)+1:

  # 状態を1ステップ分更新 (畳み込み計算)
  u = u + signal.convolve(u, kernel, mode='same')

  # 左端の値を0に修正
  u[0] = 0

  # 右端の値を0に修正
  u[-1] = 0

  # 図に描画する時点かどうか
  if t in t_list:

    # 現時点の時刻と状態を結果リストに追加
    result_list.append((t,u))


########### 図の表示 ###########

# 空の図を用意
plt.figure()

# 結果リストの各要素を取得
for (t,u) in result_list:

  # 各時点の折れ線グラフを追加 (横軸x, 縦軸u, 凡例のためにラベルも指定)
  plt.plot(x, u, label=f't={t}')

# 凡例を追加
plt.legend()

# 完成した図を表示
plt.show()

# ファイルに保存
plt.savefig('fig.png')