pythonで行列を3次元のように２次元でプロットしたい。

python3.5を使っています。例えば３次元プロットを２次元に射影するためには次のようなコードで書けます。

import matplotlib.pyplot as plt
import numpy as np

f = lambda x,y: np.exp(-(x**2 + y**2))

x = np.linspace(-4,4,100)
y = np.linspace(-4,4,100)
X,Y = np.meshgrid(x,y)
Z = f(X,Y)
plt.figure()
#ax = plt.subplot(121, projection='3d')
#ax.plot_surface(X,Y,Z,cmap='Reds')
ax = plt.subplot(111)
cs = ax.contourf(X,Y,Z,100, cmap='Reds')
plt.colorbar(cs)
plt.show()

これを実行すると次のような２次元グラフが得られます。

これを参考にして、ある行列を与えられたときに、行列の縦と横をx軸、y軸とし、行列の各要素をｚ座標の値として扱って同じような２次元グラフを書きたいです。次は試しに書いてみたコードです。

import matplotlib.pyplot as plt
import numpy as np

n_number=21
lam_number=21
mat=np.zeros((n_number,lam_number))
#print(mat)
for i in range(0,n_number):
    for j in range(0,lam_number):
        mat[i][j]=(i-10)**2+(j-10)**2

x=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]
y=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]
print(x)

X,Y = np.meshgrid(x,y)
print(x)
f = lambda x,y: mat[x][y]
Z = f(X,Y)
print(Z.shape)


#ax = plt.subplot(121, projection='3d')
#ax.plot_surface(X,Y,Z,cmap='Reds')
ax = plt.subplot(111)
cs = ax.contourf(X,Y,Z[0][0],100, cmap='Reds')
plt.colorbar(cs)
plt.show()

これを実行すると、TypeError: Input z must be a 2D array.とエラーを吐きます。確かに、Zの形を見てみるとZ.shape=(21,21,21,21)でした。

また、

import matplotlib.pyplot as plt
import numpy as np

n_number=21
lam_number=21
mat=np.zeros((n_number,lam_number))
#print(mat)
for i in range(0,n_number):
    for j in range(0,lam_number):
        mat[i][j]=(i-10)**2+(j-10)**2

x=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]
y=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]
print(x)

X,Y = np.meshgrid(x,y)
print(x)
f = lambda x,y: mat[x][y]
Z = f(X,Y)
print(Z.shape)


plt.figure()
ax = plt.subplot(111)
cs=ax.contourf(x,y,mat[x][y],cmap="Reds")
plt.colorbar(cs)
plt.show()

としたところ、下のようにカクカクのグラフが得られました。

しかし、より滑らかな図形を得たいので、やはりX,Y = np.meshgrid(x,y)を使って書いてみたいです。