在Python中为colorplot使用矩阵函数

我试图绘制一个pcolormesh，其中z是（轴）变量x和y的函数。函数z是我编写的一个模块，它使用元素依赖于x和y的矩阵，并返回最终矩阵的一个元素（0和1之间的数字）。当我试着运行这个时，我得到了一个错误

ValueError:使用序列设置数组元素

如何将z定义为在矩阵中使用x和y的函数。这是我正在运行的代码。我没有包括NewATFS模块，但它本质上只是一些矩阵乘法，并返回最终矩阵的特定元素。 状态=矩阵（[[1,0,0]，[0,0,0]，[0,0,0]]） dx，dy=0.5，1

y, x = np.mgrid[slice(1, 170 + dy, dy),
            slice(0.5, 5 + dx, dx)]
z = NewATFS.ATFS(state, 1.25, 0.01, y, 1, 0.25, 0.125, 1, x, 1, 10, True)
z = z[:-1, :-1]
z_min, z_max = -np.abs(z).max(), np.abs(z).max()

plt.pcolormesh(x, y, z, cmap='RdBu', vmin=z_min, vmax=z_max)
plt.title('pcolormesh')
plt.axis([x.min(), x.max(), y.min(), y.max()])
plt.colorbar()