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

data = np.array([
    [0, 0, 1], 
    [0, 1, 1],
    [0, 2, 1],
    [1, 0, 2],
    [1, 1, 3],
    [1, 2, 2],
    [2, 0, 2],
    [2, 1, 2],
    [2, 2, 2],
])

print('data shape:', data.shape)
n = data.shape[0]   # number of points
M = int(np.sqrt(n)) # square M*M = n

X = data[:,0]
Y = data[:,1]
Z = data[:,2]

X = X.reshape(M, M)
Y = Y.reshape(M, M)
Z = Z.reshape(M, M)

print('sX:', X.shape)
print('X:', X)
print('sY:', Y.shape)
print('Y:', Y)
print('sZ:', Z.shape)
print('Z:', Z)

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

surf = ax.plot_surface(X, Y, Z)

plt.show()