如何使用帶有使用matplotlib的單個矢量或數組輸入的函數來繪制表面？

Question

我想使用numpy和matplotlib繪制函數R ^ 2->R。

在大多數matplotlib示例中，使用帶有兩個輸入的函數，如下所示：

import numpy as np
import matplotlib.pyplot as mplot
import matplotlib.cm as cm
from mpl_toolkits.mplot3d import Axes3D as m3d

def f(x,y,sign=1.0):
    out = (np.sin(x - y/8)**2 + np.sin(x + y/8)**2)/(np.sqrt((x - 8.6998)**2 + (y - 6.7665)**2) + 1)
    return out

x = np.linspace(-5,5,num=100)
y = x
xi, yi = np.meshgrid(x,y)
zi = f(xi,yi)
fig = mplot.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(xi,yi,zi,cmap=cm.coolwarm)

現在，這很好，除了我要繪制同時也在優化的函數。 這些通常使用單個向量輸入，如下所示：

def f(x,sign=1.0):
    # Objective function
    out = sign*(x[0]**3 + 3*x[0]**2 - 2*x[0]*x[1] + 3*x[0] + x[1]**3 + 3*x[1]**2 + 3*x[1])
    return out

我將如何使用這樣的函數生成表面圖？ 我想對繪圖和優化例程使用相同的函數，因為轉錄它們既浪費又容易出錯。

Answer 1

如果輸入x是表示規則網格的3-D數組，則可以假設形狀為(2, m, n) ：

def f(x, sign=1.0):
    x1 = x[0, :]
    x2 = x[1, :]
    # Objective function
    out = sign*(x1**3 + 3*x1**2 - 2*x1*x2 + 3*x1 + x2**3 + 3*x2**2 + 3*x2)
    return out

使得out將是一個2-d陣列形狀(m, n)准備與matplotlib情節：

ax.plot_surface(x[0, :], x[1, :], f(x), cmap=cm.coolwarm)