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$の関数に当てはめて書き換えました．

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))

結果↓

ダメですね．
あー，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))

例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))

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