How can I plot a surface using a function with a single vector or array input using matplotlib?

Question

I want to plot a function R^2 -> R using numpy and matplotlib.

In most matplotlib examples, a function with two inputs is used, as here:

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)

Now, this is fine, except that I want to plot functions that I am also optimizing. Those typically use a single vector input, as here:

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

How would I go about generating a surface plot using such a function? I would like to use the same functions for both my plots and my optimization routines, since transcribing them is both wasteful and error-prone.

Answer 1

If the input x is a 3-D array representing a regular mesh, you can do, assuming a shape (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

such that out will be a 2-D array with shape (m, n) , ready to be plot with matplotlib:

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