f:id:Kurene:20210613222103g:plain

PyQtGraphで美しいリアルタイムプロットができることを最近知ったので、実装例として2次元波動方程式の数値シミュレーション（有限差分法）をリアルタイム実行・プロットしてみました。

実装のための数式導出と、PyQtGraphを使ったソースコードを合わせて紹介します。

2次元波動方程式の有限差分法の更新式導出

以下の2次元波動方程式の初期値・境界値問題の解を有限差分法によって近似的に解くことを考えます。

$\frac{\partial^2 u}{\partial t^2}=c^2\left(\frac{\partial^2 u}{\partial x^2}+\frac{\partial^2 u}{\partial y^2}\right)$

この $u$ を $xy$ 平面上の格子点での値として求めていきます。ただし、 $xy$ 平面のそれぞれの方向の格子間隔（微小変化）を $\Delta x = \Delta y = \Delta$ ）とします．

右辺，左辺の項をそれぞれテイラー展開し2階中心差分を導出します。詳しくは、最下部の参考ページをご参照ください。

それぞれ、以下のような形になります。

$\frac{\partial^2 u}{\partial t_{k}^2}=\frac{1}{\Delta t^2}\left(u_{i,j}^{k+1}+u_{i,j}^{k-1}-2u_{i,j}^{k}\right) \\ \frac{\partial^2 u}{\partial x_{i}^2}=\frac{1}{\Delta^2}\left(u_{i+1,j}^{k}+u_{i-1,j}^{k}-2u_{i,j}^{k}\right) \\ \frac{\partial^2 u}{\partial y_{j}^2}=\frac{1}{\Delta^2}\left(u_{i,j+1}^{k}+u_{i,j-1}^{k}-2u_{i,j}^{k}\right)$

ここで， $i, j$ は $x, y$ それぞれの格子点のインデックス， $k$ は時間ステップを表します。これらの数式を右辺・左辺に代入し，次の時間ステップの変位 $u_{i,j}^{k+1}$ に関して整理します。

$u_{i,j}^{k+1}=2u_{i,j}^{k}-u_{i,j}^{k-1}+\frac{c^2\Delta t^2}{\Delta^2}\left( u_{i+1,j}^{k}+u_{i-1,j}^{k}+u_{i,j+1}^{k}+u_{i,j-1}^{k}-4u_{i,j}^{k}\right )$

$c, \Delta t, \Delta$ は定数であるため，以上の更新式を時間ステップごとに全ての $i,j$ に対して実行すればOKです。

初期値と境界値条件について

詳しくは、最下部の参考ページをご参照ください。

ソースコードでは初期値として適当な位置にガウスパルスを与えました。

境界条件は固定端反射としてます。

PyQtGraphでの3Dサーフェイスプロット

3Dサーフェイスプロットにおける実装の流れは，以下のようになります

GLViewWidgetを生成・属性設定
GLSurfacePlotItemを生成・属性設定
リアルタイムでの定期更新のための関数や処理を記述

ここで，GLSurfacePlotItemはGLMeshItemの子クラスなので、属性等についてはGLMeshItemを調べると分かりやすいと思います。

PyQtGraphでのリアルタイムプロットの基本的な実装の流れはこちらが参考になります。

ソースコード

# -*- coding: utf-8 -*-
import numpy as np
from numba import jit

from pyqtgraph.Qt import QtCore, QtGui
import pyqtgraph as pg
import pyqtgraph.opengl as gl


# =====================================
#PyQtGraphのサブルーチン（3Dサーフェイスプロット）
app = QtGui.QApplication([])
w = gl.GLViewWidget()  # GLViewWidget生成
w.setFixedSize(500, 500)
w.setWindowTitle("FDM")
w.setCameraPosition(distance=50)
w.show()


# =====================================
# 差分法に関連するパラメタ定義
# =====================================
n_x, n_y = 128, 128  # 格子点数（境界以外）
xy_init  = 1.0, 4.0      # ガウス関数の中心座標
rad = 3                     # ガウス関数の偏差
x_range  = -10, 10 
y_range  = -10, 10
x        = np.linspace(x_range[0], x_range[1], n_x+2)
y        = np.linspace(y_range[0], y_range[1], n_y+2)
u        = np.zeros((3, n_x+2, n_y+2))
delta_xy = 0.1
delta_t  = 0.05
c        = 1
coef     = (c * delta_t / delta_xy) ** 2


x_grid, y_grid     = np.meshgrid(y, x)
u[1] = 6  * np.exp(-((x_grid-xy_init[0])**2.0)*rad**2) \
          * np.exp(-((y_grid-xy_init[1])**2.0)*rad**2)

# =====================================
# 3Dサーフェイスプロット用のProtItem生成と属性設定
# =====================================
p = gl.GLSurfacePlotItem(x=x, y=y, z=u[0],
                shader='heightColor',
                computeNormals=False,
                smooth=True)
p.shader()['colorMap'] = np.array([0.0, 0.0, 2.0, 0.4, 1.5, 2.0, 1.5, 1.5, 2.0]) 
w.addItem(p)

# =====================================
# リアルタイムプロット用サブルーチンやuの更新式など
# =====================================
@jit
def update_u(u, n_x, n_y, t, coef):
    for i in range(1, n_x+1):
        for j in range(1, n_y+1):
            u[t+1,i,j] =   u[t,i+1,j] + u[t,i-1,j] \
                         + u[t,i,j+1] + u[t,i,j-1] \
                         - 4.0 * u[t,i,j]
            u[t+1,i,j] *= coef
            u[t+1,i,j] += 2.0 * u[t,i,j] - u[t-1,i,j]

def update():
    global p, u, n_x, n_y, coef
    t = 1
    update_u(u, n_x, n_y, t, coef)
    p.setData(z=u[t+1])
    u[t-1] = u[t]
    u[t]   = u[t+1]
    u[t+1] *= 0.0

fps = 60
timer = QtCore.QTimer()
timer.timeout.connect(update)
timer.start(1.0 / fps * 1000)

if __name__ == '__main__':
    import sys
    if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
        QtGui.QApplication.instance().exec_()