[matplotlib] カラーマップの上にグリッドに合わせた境界線を引く

困っていたこと

matplotlibでカラーマップを描き、その上に、ある条件を満たす領域の境界線をグリッドに合わせて描きたい。

「ある条件＝boolの２次元配列」として、 ax.contour で levels=[0.5] を指定するとうまく行きそうだが、グリッドに沿わずに斜めの境界線になってしまう。

↑の境界線を格子に沿ったカクカクにしたい。

理想図

失敗図の書き方

失敗図は次のように描いている。

データ

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(-1, 1, 10)
y = np.linspace(-1, 1, 10)

extent = (x.min(), x.max(), y.min(), y.max())

X, Y = np.meshgrid(x, y)
Z = np.exp(-((X-0.5)**2 + Y**2))
condition = Z > 0.8

失敗図のプロット（斜めの境界線になる）

fig, ax = plt.subplots(dpi=200)
ax.imshow(Z, vmin=-1, extent=extent)
ax.contour(condition, levels=[0.5], extent=extent)
plt.show()

解決策

途方に暮れていたところ、 matplotlib imshow contour horizontal でググったら以下のstackoverflowが出てきた。
How can I have straight contourlines in matplotlib?

発想は、extentを固定し、二次元配列(bool)の解像度のみを上げていき、斜めが目立たないようにする というものだった。天才かよ。

成功図のプロット(まっすぐの境界線になる)

np.kronを使えば、二次元配列(bool)の解像度を上げつつ、extentで端点を固定した時に全体的な値は変化しないようにできる。

fac = 10
condition = np.kron(Z > 0.8, np.ones((fac, fac)))

結論

ググり方は大事。horizontalというワードを思いついたら正解にたどり着けた。