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$
こんな感じの図を作ろうと試行錯誤しました．
私の解答はこちら
for文を2回使って位置$(x,y)$での値を$r$の関数に当てはめて書き換えました．
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は非常に賢いですね．
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)$ に書きかえたいときにはどうすれば…？
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()使ってみるか．
if (np.sqrt(X**2+Y**2)<=1).any() == True:
Z=18*np.exp(-(X**2+Y**2))
結果↓
ダメですね．
あー，1個でもTrueなら，すべての$z(r)$が書き換えられちゃうのか．
そうだ，今回のパターンでは，
$z(r)=x^2+y^2=r^2$ $|r|\leqq1$
において，$z(r)\leqq1$になるので，その条件を満たす部分を書き換えよう．
Z[Z<=1]=18*np.exp(-(X**2+Y**2)) #これはダメ
"""
ValueError: boolean index array should have 1 dimension
ブールインデックス配列は1次元じゃないとだめ
"""
#Z[Z<=1]=18 ただしこれは成立する
あーそうか．等号不等号を使ってるから1個の値で書き換えないといけないのか…．
わかってきたぞ．
なら以下の式はいけるのか？
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点書き換えること！
この方法なら，条件式を変えれば応用もききそうです．
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))
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))
私と同じようにドハマりする人が出ないよう，この記事が役に立てば幸いです．
以上です．