matplotlib pyplotのquiver()でベクトル場を書くとクッソ見づらくなる問題の2つの解決法


Matplotlib.pyplotでベクトル場を書く際に定番となっているメソッドplt.quiver()はめっちゃ使いづらい。
きれいに描くコツと、上位互換としてのplt.streamplot()の利用を提案する。

前提

微分方程式Uまたは(「または」)Vが与えられている

これらの解軌道の振る舞いを知りたい

plt.quiver(X,Y,U,V,C,cmap)

初期値とそこからの差分を指定することでベクトルを書いてくれる。斜体はオプションだけどめっちゃ便利なのでおすすめなやつ。
X,Yには初期値をmgridなどで設定
U,Vにはそれら格子点から刻み幅ぶん移動した際の「差分」→微分方程式の解軌道を知りたい場合はRunge-Kuttaなどを使って数値的に解いてあげて、その際の差分の項、または刻み幅$\Delta t$(定数,X-Tプロットの場合のみ)を設定する。解軌道をRunge-Kuttaでどうせ出す場合によく使われがち。
そのままだとベクトルの大きさはそのまま矢印の大きさに反映されてしまう!この場合、描画範囲内で相対的にベクトルの大きさが異なるとこういうことになる。

plt.quiver(u,v,U,V)

いや、見づらいわ。見づらすぎて水ナスになったわね。もはや計算を間違えたのかと疑いたくなる。
非線形ダイナミクスの名著「ストロガッツ 非線形ダイナミクスとカオス」にはこうある。

「ベクトルの矢じりや長さの違いは、残念ながら図を乱雑にしてしまいがちである」(p.162)

おっしゃるとおりである。同書では、大きさを統一して方向場を示すことを提案、採用している。これにならった上で、白黒の本ではできなかったカラーマップを導入することでスピードを表現、本家超えをめざす。

まずベクトルの大きさで自身を割ることでサイズを統一。
この「大きさ」の値を第5の引数にあたえることで、cmapで指定した「カラーマップ」に紐付けるッ

plt.quiver(u,v,U/np.sqrt(pow(U,2)+pow(V,2)),V/np.sqrt(pow(U,2)+pow(V,2)),np.sqrt(pow(U,2)+pow(V,2)), cmap='jet')

はいわかりやすい。差は歴然である。
ちなみにこれは神経の膜電位を表すモデルに4次のRunge-Kutta法を刻み幅0.1で1ステップ適用したもの。

plt.streamplot(X,Y,U,V)

初期値に対応した解軌道を勝手に計算してくれる。正確さについては後述。
X,Yには初期値をmgridなどで設定するが、この際なぜかYを先に書かないと絶対にエラーがでる(ValueError: The rows of 'x' must be equal)。入れ替えればいいだけなので特に調査はしてないけど、パッとわかる方いたらコメントで教えて下さい。

U,Vがちがって、ここにはそのまま微分方程式(を表す変数)を書いてしまう。
連続した矢印が書かれるのでベクトルの大きさが矢印の大きさと対応しないのだが、オプションを使って色に反映させることができる。

例)

plt.streamplot(X,Y,U,V,color = np.sqrt(U**2+V**2)), cmap='jet'

実装例

u,v二次元の微分方程式をudot,vdotで定義、quiver用の初期値をu,v、streamplot用の初期値をx,yで与えておく(前述の通り順序を変えないとエラーになるので)。

import numpy as np
import matplotlib.pyplot as plt

#考えている微分方程式(cf;ストロガッツ非線形ダイナミクスとカオスp.172)
def udot(u,v):
  return u*(3-u-2*v)
def vdot(u,v):
  return v*(2-u-v)

#streamplot(なんと二行で済む!ヒートマップ的に大きいところは赤、小さいところは青で表示していくれるカラーマップを使用。)
y, x= np.mgrid[0:2:20j,0:3:20j]
plt.streamplot(x,y,udot(x,y),vdot(x,y),color=np.sqrt(udot(x,y)**2+vdot(x,y)**2),cmap='jet')

#つぎにquiver用にルンゲクッタを用意。
def rk4th_1step_2d(func,u,v,dt):
  k1 = func(u,v)*dt
  k2 = func(u+k1/2, v+k1/2)*dt
  k3 = func(u+k2/2, v+k2/2)*dt
  k4 = func(u+k3, v+k3)*dt
  return (k1+2*k2+2*k3+k4)/6

#quiverをプロット
u, v = np.mgrid[0:3:20j,0:2:20j]
U = rk4th_1step_2d(udot,u,v,0.1)
V = rk4th_1step_2d(vdot,u,v,0.1)
plt.quiver(u,v,U,V)

結果の比較

ルンゲクッタでは原理的に刻み幅を小さくするほど正確な値に近づくので,dt=0.1の場合とdt=0.01の場合で解軌道がどう変化するのか調べてみた。

dt=0.1の場合

dt=0.01の場合

dtで指定する刻み幅を小さくするとstreamplotが出した解軌道と平行に近づいていることがわかる。右下あたりとかがわかりやすい。streamplotは結構正確な計算を裏でやっているみたい。どんな計算をしているかはいずれちゃんと読みたい。

追記

こちらのサイトにソースコードがのっていて、主に適応刻み幅制御をつかった2次のRunge-Kutta法を利用していることがわかった。
https://matplotlib.org/3.1.3/_modules/matplotlib/streamplot.html