from scipy.interpolate import bisplrep, bisplev
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.axes3d import Axes3D
from matplotlib import cm

#%% Define function over 50x50 grid
xgrid = np.linspace(-1, 1, 50)
ygrid = xgrid
x, y = np.meshgrid(xgrid, ygrid)

z = x*y

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(x,
                y,
                z,
                rstride=2, cstride=2,
                cmap=cm.jet,
                alpha=0.7,
                linewidth=0.25)
plt.title('Sparsely sampled function')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

#%% Interpolate function over new 70x70 grid
xgrid2 = np.linspace(-1, 1, 70)
ygrid2 = xgrid2
xnew,ynew = np.meshgrid(xgrid2, ygrid2, indexing='ij')

tck = bisplrep(x, y, z, s=0.1) # Build the spline
znew = bisplev(xnew[:,0], ynew[0,:], tck) # Evaluate the spline

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(xnew,
                ynew,
                znew,
                rstride=2, cstride=2,
                cmap=cm.jet,
                alpha=0.7,
                linewidth=0.25)
plt.title('Interpolated function')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

plt.show()