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.

using a single vector is no problem, but are your input points x[0, :] and x[1, :] distributed in a regular mesh? — Saullo G. P. Castro
– Saullo G. P. Castro, Commented Nov 19, 2014 at 18:28
@SaulloCastro: Yes, the points are in a regular mesh (we would use numpy.linspace() to generate them). — Stephen Bosch
– Stephen Bosch, Commented Nov 20, 2014 at 15:15

Saullo G. P. Castro · Accepted Answer · 2014-11-20 15:20:33Z

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)

answered Nov 20, 2014 at 15:20

Saullo G. P. Castro

59.4k28 gold badges191 silver badges244 bronze badges

Sign up to request clarification or add additional context in comments.

7 Comments

Stephen Bosch Over a year ago

Ah, ok. I was wondering if there was some way of doing it without having to reassign; that is, with the vector being called with indices in the expression, but this is also a good answer, so I will give it a point for now and wait a bit before declaring it the accepted answer.

Saullo G. P. Castro Over a year ago

@StephenBosch reassigning in this case is not copying the arrays, just creating shortcuts through new variables...

Stephen Bosch Over a year ago

Right, they are just pointers, that is easy to forget. Still, it turns out this isn't working. If I use it in the code above, I still end up with a zi that has the wrong dimensions (100,), which is the same result I get when I use the single vector input example from my question. I probably have to change other parts of my code.

Saullo G. P. Castro Over a year ago

@StephenBosch if the shape of x is (2, m, n), z should be (m, n)... check if you are not passing yi as the input sign, which would give a wrong result...

Stephen Bosch Over a year ago

Thanks for the tips, I found a couple of oversights (the sign parameter being one of them). I have it so that I get usable plots now. Unfortunately, this function does not work with scipy.optimize.minimize(), it gives IndexError: too many indices at x1 = x[0,:]. scipy.optimize.minimize() must be passing (n,) shaped vectors to the function.

|

Collectives™ on Stack Overflow

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

1 Answer 1

7 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

7 Comments

Your Answer

Sign up or log in

Post as a guest