【Python】meshgridの値を書き換えるときにドハマりしたので解決方法を【備忘録】

meshgridを使ったカラーマップを作成するとき,値を書き換える際にドはまりした内容を備忘録も兼ねて投稿.


やりたいこと

meshgridを使った強度マップで

①距離$r$に依存する関数を

②距離$r$の値で場合分けしたい

①については,$r=\sqrt{x^2+y^2}$でいいのですが,問題は②.

今回は,状況がわかりやすいように以下のような2つの関数を組み合わせたグラフを書くことを目的とします.

$z(r)=18\exp(-r^2)$    $|r|\leqq1$

$z(r)=x^2+y^2=r^2$   $|r|>1$

Fig3-1.png Fig3-2.png

Fig2-1.png Fig2-2.png

こんな感じの図を作ろうと試行錯誤しました.

Fig1-1.png Fig1-2.png


私の解答はこちら

for文を2回使って位置$(x,y)$での値を$r$の関数に当てはめて書き換えました.


code

import numpy as np

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from mpl_toolkits.mplot3d import axes3d

plt.rcParams['figure.figsize'] = [6, 5]

num = 301
a = 1
x = np.linspace(-3*a, 3*a ,num)
y = np.linspace(-3*a, 3*a, num)
X,Y = np.meshgrid(x, y)

Z = []
Z = X**2+Y**2

for m1 in range(num):
for m2 in range(num):
if np.sqrt(X[m1][m2]**2+Y[m1][m2]**2)<=a:
Z[m1][m2] = 18*np.exp(-(X[m1][m2]**2+Y[m1][m2]**2))

fig = plt.figure()
ax = plt.axes()
plt.pcolor(X, Y, Z, cmap='jet')
c = patches.Circle(xy=(0, 0), radius=a, fill=False, ec='white', lw=3, ls='--')
ax.add_patch(c)
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar()
#plt.savefig("Fig1-1")

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap="jet")
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
#plt.savefig("Fig1-2")

plt.show()



どこにドハマりしたのか

Numpyのmeshgridは非常に賢いですね.


code

X,Y = np.meshgrid(x, y)

Z = X**2+Y**2

print(Z)
[[ 18. 17.8804 17.7616 ..., 17.7616 17.8804 18. ]
[ 17.8804 17.7608 17.642 ..., 17.642 17.7608 17.8804]
[ 17.7616 17.642 17.5232 ..., 17.5232 17.642 17.7616]
...,
[ 17.7616 17.642 17.5232 ..., 17.5232 17.642 17.7616]
[ 17.8804 17.7608 17.642 ..., 17.642 17.7608 17.8804]
[ 18. 17.8804 17.7616 ..., 17.7616 17.8804 18. ]]

print(Z.shape)
(301, 301)


これだけで,$z(r)=x^2+y^2$ の値を得られます.

meshgridで作成した$X$は$(x,y)$の成分を持っているのに,$x$だけを2乗にしてくれます.

$Y$も同様に$y$だけを2乗にしてくれます.

さて,$|r|\leqq1$ の範囲で $z(r)=18\exp(-r^2)$ に書きかえたいときにはどうすれば…?


失敗例1

if np.sqrt(X**2+Y**2)<=1:

Z=18*np.exp(-(X**2+Y**2))

"""
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
複数の要素を持つ配列の真理値はあいまい.a.any() か a.all() を使ってね
"""


なるほど,確かにあいまいだと言われたらそうだな,a.any()使ってみるか.


失敗例2

if (np.sqrt(X**2+Y**2)<=1).any() == True:

Z=18*np.exp(-(X**2+Y**2))

結果↓

Fig4-1.png

ダメですね.

あー,1個でもTrueなら,すべての$z(r)$が書き換えられちゃうのか.

そうだ,今回のパターンでは,

$z(r)=x^2+y^2=r^2$   $|r|\leqq1$

において,$z(r)\leqq1$になるので,その条件を満たす部分を書き換えよう.


失敗例3

Z[Z<=1]=18*np.exp(-(X**2+Y**2)) #これはダメ

"""
ValueError: boolean index array should have 1 dimension
ブールインデックス配列は1次元じゃないとだめ
"""

#Z[Z<=1]=18 ただしこれは成立する


あーそうか.等号不等号を使ってるから1個の値で書き換えないといけないのか….

わかってきたぞ.

なら以下の式はいけるのか?


失敗例4

for m1 in range(num):

for m2 in range(num):
if np.sqrt(X[m1]**2+Y[m2]**2)<=1:
Z[m1,m2] = 18*np.exp(-(X[m1]**2+Y[m2]**2))
"""
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
複数の要素を持つ配列の真理値はあいまい.a.any() か a.all() を使ってね
"""


またかよ!でも,確かにこの式でも2次元の配列になっちゃてる.

1点1点書き換えていくしかないな.


成功例

for m1 in range(num):

for m2 in range(num):
if np.sqrt(X[m1][m2]**2+Y[m1][m2]**2)<=1:
Z[m1][m2] = 18*np.exp(-(X[m1][m2]**2+Y[m1][m2]**2))

最終的に,$X(x,y)$ および $Y(x,y)$ の値を使って,該当する $Z(x,y)$ の値を書き換えました.


結論

meshgridでできた値を書き換えたいとき,

つまり「複数の要素を持つ配列を書き換えたい」ときは,1点1点書き換えること!

この方法なら,条件式を変えれば応用もききそうです.


例1

for m1 in range(num):

for m2 in range(num):
if np.sqrt(X[m1][m2]**2)<=1:
Z[m1][m2] = 18*np.exp(-(X[m1][m2]**2+Y[m1][m2]**2))

Fig5-1.png Fig5-2.png


例2

for m1 in range(num):

for m2 in range(num):
if np.sqrt(X[m1][m2]**2)<=a or np.sqrt(Y[m1][m2]**2)<=a:
Z[m1][m2] = 18*np.exp(-(X[m1][m2]**2+Y[m1][m2]**2))

Fig6-1.png Fig6-2.png

私と同じようにドハマりする人が出ないよう,この記事が役に立てば幸いです.

以上です.