極座標形式の2次元データをMatplotlibで描写する

f:id:Ibetoge:20201219144408p:plain — 右の図を作ろう

f:id:Ibetoge:20201219144020p:plain — 右の図を作ろう

数値解析を極座標で行った後，2次元の極座標形式データを描写したい時がありました．
例えば，
${\displaystyle \begin{eqnarray} \left\{ \begin{array}{l} r=x^2+y^2 \\ \theta=\arccos\left(\frac{x}{\sqrt{x^2+y^2}}\right) \end{array} \right. \end{eqnarray} }$
のような座標でのデータを描写したい時があります．つまり

double r[imax];
double theta[jmax];
double T[jmax][imax];

のような配列データを格納したcsvファイルやらが手元にあります．それらをpythonを用いてmatplotlibで描写しましょう．
( $r,~\theta$ 方向にそれぞれ ${\rm imax,~jmax}$ 分割してるような空間を配列r,thetaによって表現していて，
そのグリッド点で物理量T[jmax][imax]を定義しています．)
上の配列はそれぞれr.csv，theta.csv，T.csvに入っていることにしておきます．
今回は
${\displaystyle \begin{equation} T=\mathrm{e}^{-r^2} \end{equation} }$
として，Tを描写します．

流れとしては，

極座標配列をcsvから読み込む
x-y平面のデータに変換(スプライン補完)
描写

という感じです．
まずは必要なモジュールたちをインポートします．

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from scipy import interpolate

scipyはスプライン補完に使い，mpl_toolkits.mplot3dと matplotlib.pyplotは描写，numpyは色々なことに使います．

次にnumpyでcsvファイルを読み込んで配列に格納しましょう．

r = np.loadtxt("r.csv", delimiter=',')
theta = np.loadtxt("theta.csv", delimiter=',')
T_polar = np.loadtxt("T.csv", delimiter=',')

描写したいx-y平面とx-y平面でのTを定義します．今回は配列r,thetaで表せる部分をすべて含めて，グリッド幅は配列rと同じにしておきます．

size_r = r.shape[0]
x = np.linspace(-r[-1], r[-1], 2*size_r)
y = np.linspace(-r[-1], r[-1], 2*size_r)
T_cartesian = np.empty((2*size_r, 2*size_r))

スプライン補完でx-y平面でのTを計算

spl_T = interpolate.interp2d(theta, r, T_polar, kind='cubic')
for ix in range(0, 2*size_r):
    for iy in range(0, 2*size_r):
        r_tmp = (x[ix]**2+y[iy]**2)**0.5
        costh = x[ix]/(x[ix]**2+y[iy]**2)**0.5
        th_tmp = np.arccos(costh)
        T_cartesian[ix, iy] = spl_T(th_tmp, r_tmp)

最後に描写

x, y = np.meshgrid(x, y, indexing='ij')
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.contourf(x, y, T_cartesian)
plt.xlabel('x')
plt.ylabel('y')
plt.show()

これで完成．

割と一般化できるやり方なので，これで任意の座標系から任意の座標系へ変換して描写できるようになります．

ちなみに今回はカラーマップの等値線図を書きましたが

ax = fig.add_subplot(1, 1, 1)
ax.contourf(x, y, T_cartesian)

の部分を変えるといろいろな図にできます．

グラフの種類

等値線図

ax = fig.add_subplot(1, 1, 1)
ax.contour(x, y, T_cartesian)

f:id:Ibetoge:20201219150159p:plain — 等値線図

3次元ワイヤーフレーム

ax = fig.add_subplot(1, 1, 1, projection="3d")
ax.plot_wireframe(x, y, T_cartesian)

f:id:Ibetoge:20201219150408p:plain — ワイヤー

３次元曲面

ax = fig.add_subplot(1, 1, 1, projection="3d")
ax.plot_surface(x, y, T_cartesian)

f:id:Ibetoge:20201219150501p:plain — 三次元曲面

コード

最後にcsv作るように書いたコードとプロットするのに使ったコードを載せときます．

CSV作るよう

#include <cmath>
#include <fstream>
#include <iostream>

using namespace std;

int main() {
  int imax = 200;
  int jmax = 128;
  double r[imax], theta[jmax], T[jmax][imax];
  double pi;
  pi = acos(-1.0);

  for (int i = 0; i < imax; i++) {
    r[i] = (i + 0.50) / imax;
  }
  for (int j = 0; j < jmax; j++) {
    theta[j] = (j + 0.50) / jmax * 2.0 * pi;
  }

  for (int i = 0; i < imax; i++) {
    for (int j = 0; j < jmax; j++) {
      T[j][i] = exp(-1.0 * pow(r[i], 2));
    }
  }
  ofstream ofsr("r.csv");
  for (int i = 0; i < imax - 1; i++) {
    ofsr << r[i] << ",";
  }
  ofsr << r[imax - 1] << endl;

  ofstream ofst("theta.csv");
  for (int j = 0; j < jmax - 1; j++) {
    ofst << theta[j] << ",";
  }
  ofst << theta[jmax - 1] << endl;

  ofstream ofsv("T.csv");
  for (int i = 0; i < imax; i++) {
    for (int j = 0; j < jmax; j++) {
      if (j == jmax - 1) {
        ofsv << T[j][i] << endl;
      } else {
        ofsv << T[j][i] << ",";
      }
    }
  }
}

プロットする

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from scipy import interpolate

# csvを読み込む
r = np.loadtxt("r.csv", delimiter=',')
theta = np.loadtxt("theta.csv", delimiter=',')
T_polar = np.loadtxt("T.csv", delimiter=',')

# x-y平面を定義してスプライン補完
size_r = r.shape[0]
x = np.linspace(-r[-1], r[-1], 2*size_r)
y = np.linspace(-r[-1], r[-1], 2*size_r)
T_cartesian = np.empty((2*size_r, 2*size_r))

spl_T = interpolate.interp2d(theta, r, T_polar, kind='cubic')
for ix in range(0, 2*size_r):
    for iy in range(0, 2*size_r):
        r_tmp = (x[ix]**2+y[iy]**2)**0.5
        costh = x[ix]/(x[ix]**2+y[iy]**2)**0.5
        th_tmp = np.arccos(costh)
        T_cartesian[ix, iy] = spl_T(th_tmp, r_tmp)

# matplotlibで描写する
x, y = np.meshgrid(x, y, indexing='ij')
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.contourf(x, y, T_cartesian)
plt.xlabel('x')
plt.ylabel('y')
plt.savefig('cartesian.png')

# 極座標での描写
theta, r = np.meshgrid(theta, r, indexing='ij')
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.contourf(theta, r, T_polar)
plt.xlabel(r'$\theta$')
plt.ylabel('r')
plt.savefig('polar.png')

My Memo

私のメモ