matplotlib surface plot for (x-y)^2

Question

Can anyone help me with plotting a 3D surface plot for the equation f(x,y) = (x-y)^2

The Z axis should represent the function f(x,y)

I have the below function:

def fnc(X):
    return (X[0] - X[1]) ** 2

Here X is a numpy array with first parameter as X and the second as Y. I specifically need it to be this way. So please don't suggest me to change the signature. ;)

I tried the below from this solution:

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
x = y = np.linspace(-5,5,100)
X, Y = np.meshgrid(x, y)
Z = fnc1([np.linspace(-5,5,100) , np.linspace(-6,6,100)])
ax.plot_surface(X, Y, Z)
ax.set_xlabel('X Label')
ax.set_ylabel('Y Label')
ax.set_zlabel('Z Label')
plt.show()

However i get the wrong plot.

Answer 1

The array Z to be plotted should be a 2D array, just like X and Y are 2D arrays, such that for each pair of values from X and Y you get exactly one point in Z. It therefore makes sense to use those arrays X and Y as input to your fnc function, Z = fnc([X,Y])

The complete code would look like

import numpy as np
import matplotlib.pylab as plt
from mpl_toolkits.mplot3d import Axes3D


def fnc(X):
    return (X[0] - X[1]) ** 2

fig = plt.figure()
ax = fig.add_subplot(111, projection=Axes3D.name)
x = y = np.linspace(-5,5,100)
X, Y = np.meshgrid(x, y)
Z = fnc([X,Y])
ax.plot_surface(X, Y, Z)
ax.set_xlabel('X Label')
ax.set_ylabel('Y Label')
ax.set_zlabel('Z Label')
ax.view_init(elev=15, azim=-118)
plt.show()