【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](https://qiita-image-store.s3.amazonaws.com/0/192073/b39da136-cc05-5839-2485-ff3c034fab4a.png) ![Fig3-2.png](https://qiita-image-store.s3.amazonaws.com/0/192073/1d4b94b5-af52-1814-5a35-b67852331464.png)
![Fig2-1.png](https://qiita-image-store.s3.amazonaws.com/0/192073/3bee4687-156d-9400-78b1-3abaa8816a50.png) ![Fig2-2.png](https://qiita-image-store.s3.amazonaws.com/0/192073/abf18d4b-e1f4-450a-873c-22aa5966999b.png)



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

![Fig1-1.png](https://qiita-image-store.s3.amazonaws.com/0/192073/488435c3-77e0-d38a-f43e-8695a42ab384.png) ![Fig1-2.png](https://qiita-image-store.s3.amazonaws.com/0/192073/3fabc2a8-76af-0525-b2f3-cb64b4f40634.png)




# 私の解答はこちら

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


```ruby: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は非常に賢いですね．

```ruby: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)$ に書きかえたいときにはどうすれば…？


```ruby:失敗例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()使ってみるか．

```ruby:失敗例2
if (np.sqrt(X**2+Y**2)<=1).any() == True:
    Z=18*np.exp(-(X**2+Y**2))
```

結果↓
![Fig4-1.png](https://qiita-image-store.s3.amazonaws.com/0/192073/0f208352-2bf3-94e6-9f40-11c40e6d097f.png)


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



そうだ，今回のパターンでは，
$z(r)=x^2+y^2=r^2$　　 $|r|\leqq1$
において，$z(r)\leqq1$になるので，その条件を満たす部分を書き換えよう．

```ruby:失敗例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個の値で書き換えないといけないのか…．
わかってきたぞ．
なら以下の式はいけるのか？

```ruby:失敗例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点書き換えていくしかないな．

```ruby:成功例
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点書き換えること！**

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

```ruby:例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](https://qiita-image-store.s3.amazonaws.com/0/192073/54158294-f85d-1159-f2a8-23f8e5149f1e.png) ![Fig5-2.png](https://qiita-image-store.s3.amazonaws.com/0/192073/bc615322-ee07-74dc-0140-940581b65502.png)


```ruby:例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](https://qiita-image-store.s3.amazonaws.com/0/192073/792281da-8456-0c92-6332-afde02833910.png) ![Fig6-2.png](https://qiita-image-store.s3.amazonaws.com/0/192073/d215f736-e2aa-7a88-782d-a30f7d52b0bb.png)



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