Streamplot 範例

作者：王一哲
日期：2019/12/2

前言

我為了寫〈均勻帶電圓環的電場〉找到了 matplotlib 官網上關於 streamplot 指令的範例，範例中利用 numpy 的指令直接產生繪圖資料，但是我需要從 csv 檔中讀取由 VPython 產生的資料，因此我需要改寫範例的程式碼。

streamplot 指令範例繪圖成果

範例

import numpy as np
import matplotlib.pyplot as plt

Y, X = np.mgrid[-3:3:100j, -3:3:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U*U + V*V)

plt.streamplot(X, Y, U, V, color=U, linewidth=2, cmap=plt.cm.autumn)
plt.colorbar()
plt.savefig('streamplot.png')
plt.show()

原來的範例中底下還有另外兩張圖，但我只是要測試指令及存取資料的方法，因此只留下一張圖。範例程式的第4行使用了 numpy.mgrid 產生 X、Y 資料，其語法為

Y, X = np.mgrid[y軸最小值:y軸最大值:y軸點數+j, x軸最小值:x軸最大值:x軸點數+j]
Y, X = np.mgrid[y軸最小值:y軸最大值:y軸間隔, x軸最小值:x軸最大值:x軸間隔]

因此第4行的功能是將

$-3\le x\le 3$ 及

$-3\le y\le 3$ 各分為100個點，再將這些數值交錯產生兩個2維陣列，其中X為

[[-3.         -2.93939394 -2.87878788 ...  2.87878788  2.93939394
   3.        ]
 [-3.         -2.93939394 -2.87878788 ...  2.87878788  2.93939394
   3.        ]
 [-3.         -2.93939394 -2.87878788 ...  2.87878788  2.93939394
   3.        ]
 ...
 [-3.         -2.93939394 -2.87878788 ...  2.87878788  2.93939394
   3.        ]
 [-3.         -2.93939394 -2.87878788 ...  2.87878788  2.93939394
   3.        ]
 [-3.         -2.93939394 -2.87878788 ...  2.87878788  2.93939394
   3.        ]]

Y為

[[-3.         -3.         -3.         ... -3.         -3.
  -3.        ]
 [-2.93939394 -2.93939394 -2.93939394 ... -2.93939394 -2.93939394
  -2.93939394]
 [-2.87878788 -2.87878788 -2.87878788 ... -2.87878788 -2.87878788
  -2.87878788]
 ...
 [ 2.87878788  2.87878788  2.87878788 ...  2.87878788  2.87878788
   2.87878788]
 [ 2.93939394  2.93939394  2.93939394 ...  2.93939394  2.93939394
   2.93939394]
 [ 3.          3.          3.         ...  3.          3.
   3.        ]]

也可以改用 numpy.meshgrid 達成同樣的效果，將第4行改為以下的程式碼效果相同。

x = np.linspace(-3, 3, 100)
y = np.linspace(-3, 3, 100)
X, Y = np.meshgrid(x, y)

產生資料並存成 csv 檔

import numpy as np

Y, X = np.mgrid[-3:3:100j, -3:3:100j]
U = -1 - X**2 + Y
V = 1 + X - Y**2
speed = np.sqrt(U*U + V*V)

rows = len(X)
cols = len(X[0])

X = X.reshape(1, rows*cols)
Y = X.reshape(1, rows*cols)
U = U.reshape(1, rows*cols)
V = V.reshape(1, rows*cols)
speed = speed.reshape(1, rows*cols)
data = np.concatenate((X, Y, U, V, speed), axis=0).transpose()
np.savetxt('data.csv', data, delimiter=',')

1. 第8、9行：從陣列 X 取得列數 rows 及欄數 cols。
2. 第11 ~ 15行：用 numpy.reshape 將陣列 X、Y、U、V、speed 轉成1列、
   $rows\times cols$ 欄的2維陣列。
3. 第16行：用 numpy.concatenate 將5個陣列接起來再轉置，將資料存在2維陣例 data 當中，這樣各欄的資料依序是 X、Y、U、V、speed。其中選項 axis=0 是預設值，理論上可以省略。
4. 第17行：用 numpy.savetxt 將陣列 data 的資料用逗號分隔儲存到文字檔 data.csv。
5. 以下是陣列 data 的前10列資料。

[[ -3.          -3.         -13.         -11.          17.02938637]
 [ -2.93939394  -2.93939394 -12.64003673 -10.93939394  16.71648493]
 [ -2.87878788  -2.87878788 -12.28741965 -10.87878788  16.41123723]
 [ -2.81818182  -2.81818182 -11.94214876 -10.81818182  16.11359596]
 [ -2.75757576  -2.75757576 -11.60422406 -10.75757576  15.82350948]
 [ -2.6969697   -2.6969697  -11.27364555 -10.6969697   15.54092161]
 [ -2.63636364  -2.63636364 -10.95041322 -10.63636364  15.26577155]
 [ -2.57575758  -2.57575758 -10.63452709 -10.57575758  14.99799369]
 [ -2.51515152  -2.51515152 -10.32598714 -10.51515152  14.73751749]
 [ -2.45454545  -2.45454545 -10.02479339 -10.45454545  14.48426744]]

從 csv 檔讀取資料並繪圖

import numpy as np
import matplotlib.pyplot as plt

X, Y, U, V, speed = np.loadtxt('data.csv', delimiter=',', unpack=True)

num = int(np.sqrt(len(X)))
X = X.reshape(num, num)[0, :]
Y = Y.reshape(num, num)[0, :]
U = U.reshape(num, num)
V = V.reshape(num, num)
speed = speed.reshape(num, num)

plt.streamplot(X, Y, U, V, color=U, linewidth=2, cmap=plt.cm.autumn)
plt.colorbar()
plt.savefig('streamplotTest.png')
plt.show()

1. 第6行：從陣列 X 的長度推算原來的列數及欄數。
2. 第7、8行：用 numpy.reshape 將陣列 X、Y 轉成2維陣列，再取出其中第0列的資料存回 X、Y 當中轉成1維陣列，這是為了符合 streamplot 輸入資料的格式。
3. 第9 ~ 11行：用 numpy.reshape 將陣列 U、V、speed 轉成2維陣列。
4. 第13行：用 streamplot 指令繪圖，繪讀成果與範例一模一樣。

從 csv 檔讀取資料繪圖成果

結語

1. 用 numpy.mgrid 或 numpy.meshgrid 產生資料很方便，但這兩種的差異實在很微妙，產生的資料排序方式也很難想像。
2. 為了將資料依照欄位存進 csv 檔，還要先改變陣列的格式、連接再轉置，這個方法看起來很怪，但是它至少會成功，我目前還想不更好的方法。
3. 最重要的一點，我找到從 csv 檔讀取資料並繪圖的方法，以後就可以用別的工具產生資料並儲存成 csv 檔，甚至可以處理真實的實驗數據。

參考資料

1. streamplot_demo_features.py：https://matplotlib.org/1.3.1/examples/images_contours_and_fields/streamplot_demo_features.html
2. numpy.mgrid：https://docs.scipy.org/doc/numpy/reference/generated/numpy.mgrid.html
3. numpy.meshgrid：https://docs.scipy.org/doc/numpy/reference/generated/numpy.meshgrid.html
4. CSDN numpy中mgrid与meshgrid的区别：https://blog.csdn.net/tymatlab/article/details/79027162
5. matplotlib.streamplot：https://docs.scipy.org/doc/numpy/reference/generated/numpy.loadtxt.html
6. numpy.concatenate：https://docs.scipy.org/doc/numpy/reference/generated/numpy.concatenate.html
7. numpy.savetxt：https://docs.scipy.org/doc/numpy/reference/generated/numpy.savetxt.html
8. StackOverflow plot streamlines with matplotlib from file：https://stackoverflow.com/questions/55207810/plot-streamlines-with-matplotlib-from-file

---

tags:Python