Pythonで任意の定義域を持つヒートマップの作り方

　$f=(x,y)$的なデータを図示するにあたって、ヒートマップは強力な手段です。

　Pythonではseabornが有名ですが、xとyの定義域がそれぞれ1,2,...に限定されてしまいます。どうせなら任意の定義域でやりたいですよね。scikitlearnのサンプルコードを参考にしたやり方を解説します。

不格好だけどあらゆる関数に対応できるやつ

import numpy as np
import matplotlib.pyplot as plt
x = np.arange(-2,2,0.1)
y = np.arange(-2,2,0.1)
xx,yy = np.meshgrid(x,y)
z = np.zeros((len(x),len(y)))
for i in range(len(x)):
    for j in range(len(y)):
        v = np.array([xx[i][j],yy[i][j]])
        r = np.linalg.norm(v)
        z[i][j] = np.exp(-r)
plt.figure(figsize=(5,5))
plt.contourf(x,y,z,cmap='coolwarm')
plt.show()

　for文の中身が実質的な$f(x,y)$になっていて、上のサンプルコードでは例として、

$$
f(x,y) = e^{\frac{-1}{\sqrt{x^2+y^2}}}
$$

　を実装しています。煩雑なので3行に分けていますが、やっていることは同じです。

　これはあらゆる関数に対応可能ですが、for文を使っているのがいかにも不格好です。何より、行列を一回かければ大丈夫なパターンや、各種機械学習の予測モデル（だいたいはデータ数 x 入力ベクトル）の形に対応していないので不便です。

データ数 x 入力ベクトルの形に直したやつ

import numpy as np
import matplotlib.pyplot as plt

def model(xy):
    ar = np.array([-2,1])
    return np.matmul(xy,ar)

x = np.arange(-2,2,0.1)
y = np.arange(-2,2,0.1)
xx,yy = np.meshgrid(x,y)
xx = xx.reshape(-1,1)
yy = yy.reshape(-1,1)
xy = np.concatenate([xx,yy],1)
z = model(xy)
z = z.reshape(-1,len(x))
plt.figure(figsize=(5,5))
plt.contourf(x,y,z,cmap='coolwarm')
plt.show()

　x,yの形を整形しています。図にすると以下のような感じです。

　この形式であれば、例えばlda.precdict()やmodel()（PyTorch等）の関数に対応できます。

おまけ

　contourfではなくpcolormedhにするとグリッド表示になります。見た目はキレイですが、データの傾向はむしろわかりづらくなる気がするので、個人的にはcontourfの方が好みです。

import numpy as np
import matplotlib.pyplot as plt
x = np.arange(-2,2,0.01) #解像度を0.01に
y = np.arange(-2,2,0.01) #解像度を0.01に
xx,yy = np.meshgrid(x,y)
z = np.zeros((len(x),len(y)))
for i in range(len(x)):
    for j in range(len(y)):
        v = np.array([xx[i][j],yy[i][j]])
        r = np.linalg.norm(v)
        z[i][j] = np.exp(-r)
plt.figure(figsize=(5,5))
plt.pcolormesh(x,y,z,cmap='coolwarm') #ここを変更
plt.show()