忘れられた備忘録

機械学習の研究をしています.強化学習,逆強化学習が専門です.

Matplotlibでヒートマップ(カラーバー付き)

matplotlib.pyplot.pcolormeshを使う

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
x = np.linspace(-4, 4, 200)
y = np.linspace(-5, 5, 200)
x_mesh, y_mesh = np.meshgrid(x, y)
z_mesh = - np.sqrt(x_mesh**2 + y_mesh**2)
plt.pcolormesh(x_mesh, y_mesh, z_mesh)
plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x11e4a1080>

f:id:yusuke135:20181011004020p:plain

ヒートマップが描けた.

ちなみに関数を3Dで表示するとこんな感じ.

from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(x_mesh, y_mesh, z_mesh)

<mpl_toolkits.mplot3d.art3d.Poly3DCollection at 0x11f147080>

f:id:yusuke135:20181011004025p:plain

参考URL: matplotlib.pyplot.pcolormesh — Matplotlib 3.0.0 documentation