3D Surface Plot where z is a function that takes a vector formed from x and y

Question

I am trying to produce a 3D surface plot where X and Y are values between -50 and 50, and Z is calculated by a function depending on X and Y.

This function takes a vector as a parameter in the form of an np array. The vector's first row is a value from X and the second a value from Y. All combinations of X and Y should produce a Z value, hence the meshgrid.

Here is my implementation, for ZI am currently creating a vector where the first row is the entire dataset of X, and the second the entire dataset of Y. This is of course incorrect.

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np
import matplotlib.pyplot as plt

def myFunction(v):
    return v.dot(np.array([1, 2]))

fig = plt.figure()
ax = fig.gca(projection='3d')

# Make data.
X = np.linspace(-50,50, 100)
Y = np.linspace(-50,50, 100)
X, Y = np.meshgrid(X, Y)
Z = myFunction(np.array([X, Y])) # <-- Here is the problem

# Plot the surface.
surf = ax.plot_surface(X, Y, Z, cmap=cm.Greens,
                       linewidth=0, antialiased=False)


ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

I hope I have made sense, Thanks

Answer 1

You would probably like to supply an array with all x values in the frst column and all y values in the second column to the function. That would ensure to have the dimensions match for the dot product. The result can then be reshaped to the shape of the mesh.

Z = myFunction(np.array([X.flatten(), Y.flatten()]).T).reshape(X.shape)

Complete example:

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np
import matplotlib.pyplot as plt

def myFunction(v):
    return v.dot(np.array([1, 2]))

fig = plt.figure()
ax = fig.gca(projection='3d')

# Make data.
X = np.linspace(-50,50, 100)
Y = np.linspace(-50,50, 100)
X, Y = np.meshgrid(X, Y)
Z = myFunction(np.array([X.flatten(), Y.flatten()]).T).reshape(X.shape)

# Plot the surface.
surf = ax.plot_surface(X, Y, Z, cmap=cm.Greens,
                       linewidth=0, antialiased=False)


ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()