如何使用仅支持标量输入的自定义函数绘制Matplotlib等高线图？

我刚好遇到过类似的问题，当自定义函数只接受标量输入时，直接传给meshgrid生成的二维数组肯定会报错，这里有几种实用的解决办法，你可以根据自己的需求选择：

方法1：用np.vectorize包装标量函数 #

这是最快的解决方案，不需要修改你的原函数，numpy的vectorize会帮你把标量函数转换成能处理数组的版本——它本质上是在底层遍历数组的每个元素，逐个调用你的函数。

举个例子：

import numpy as np
import matplotlib.pyplot as plt

# 假设这是你的复杂标量函数（比如负对数似然计算）
def my_function(x1, x2):
    return np.log(x1**2 + x2**2 + 1)

# 把标量函数包装成向量化函数
vectorized_nllh = np.vectorize(my_function)

# 生成网格点
xlist = np.linspace(-3.0, 3.0, 100)
ylist = np.linspace(-3.0, 3.0, 100)
X, Y = np.meshgrid(xlist, ylist)

# 现在可以直接用包装后的函数计算Z了
Z = vectorized_nllh(X, Y)

# 绘制等高线图
fig, ax = plt.subplots()
cp = ax.contourf(X, Y, Z)
plt.colorbar(cp)
ax.set_title('Contour Plot with Vectorized Custom Function')
ax.set_xlabel('x (cm)')
ax.set_ylabel('y (cm)')
plt.show()

注意：np.vectorize不是真正的底层向量化（底层还是Python循环），所以如果你的网格非常大（比如1000x1000以上），速度会有点慢，但对于常规的等高线图需求完全够用。

方法2：手动遍历网格点（适合调试） #

如果你想清楚看到每个点的计算过程，或者需要在循环里加一些调试逻辑，手动遍历是个不错的选择：

import numpy as np
import matplotlib.pyplot as plt

def my_function(x1, x2):
    return np.log(x1**2 + x2**2 + 1)

xlist = np.linspace(-3.0, 3.0, 100)
ylist = np.linspace(-3.0, 3.0, 100)
X, Y = np.meshgrid(xlist, ylist)

# 初始化一个和X/Y同维度的空数组
Z = np.zeros_like(X)

# 逐个计算每个网格点的值
for i in range(X.shape[0]):
    for j in range(X.shape[1]):
        Z[i][j] = my_function(X[i][j], Y[i][j])

# 绘图部分和之前一致
fig, ax = plt.subplots()
cp = ax.contourf(X, Y, Z)
plt.colorbar(cp)
ax.set_title('Contour Plot with Manual Looping')
ax.set_xlabel('x (cm)')
ax.set_ylabel('y (cm)')
plt.show()

这种方法逻辑最直观，但Python的嵌套循环效率不高，网格点多的时候会比较卡，所以更适合小网格或者调试场景。

方法3：改造函数为真正的向量化版本（最优性能） #

如果你的函数逻辑允许，把它改成支持数组输入的向量化版本是性能最好的选择——numpy的向量化操作是用C实现的，速度比Python循环快几个数量级。

比如把原来的标量函数改成：

import numpy as np
import matplotlib.pyplot as plt

# 改造后的向量化函数，直接支持数组输入
def my_function(x1, x2):
    # 所有操作都用numpy的向量化方法，比如**、+、log都是原生支持数组的
    return np.log(x1**2 + x2**2 + 1)

xlist = np.linspace(-3.0, 3.0, 100)
ylist = np.linspace(-3.0, 3.0, 100)
X, Y = np.meshgrid(xlist, ylist)

# 直接调用函数，不需要额外处理
Z = my_function(X, Y)

# 绘图
fig, ax = plt.subplots()
cp = ax.contourf(X, Y, Z)
plt.colorbar(cp)
ax.set_title('Contour Plot with Native Vectorized Function')
ax.set_xlabel('x (cm)')
ax.set_ylabel('y (cm)')
plt.show()

如果你的原函数里有标量的条件判断（比如if x1 > 0:），可以用np.where代替实现向量化判断：

def my_function(x1, x2):
    # 向量化的条件分支逻辑
    return np.where(x1 > 0, np.log(x1**2 + x2**2), np.sqrt(x1**2 + x2**2))

这样改造后，函数就能直接处理meshgrid的数组了，速度最快。

内容的提问来源于stack exchange，提问作者hex93