import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
import matplotlib as mpl
%matplotlib inline

from scipy.interpolate import (
    LinearNDInterpolator, RectBivariateSpline,
    RegularGridInterpolator)
from scipy.ndimage import map_coordinates
from dolointerpolation import MultilinearInterpolator

# Tweak how images are plotted with imshow
mpl.rcParams['image.interpolation'] = 'none' # no interpolation
mpl.rcParams['image.origin'] = 'lower' # origin at lower left corner
mpl.rcParams['image.cmap'] = 'RdBu_r'

def f_2d(x,y):
    '''a function with 2D input to interpolate on [0,1]'''
    twopi = 2*pi
    return np.exp(-x)*np.cos(x*2*pi)*np.sin(y*2*pi)

# quick check :
f_2d(0,0.25)

def f_3d(x,y,z):
    '''a function with 3D input to interpolate on [0,1]'''
    twopi = 2*pi
    return np.sin(x*2*pi)*np.sin(y*2*pi)*np.sin(z*2*pi)

# quick check :
f_3d(0.25,0.25,0.25)

Ndata = 50
xgrid = np.linspace(0,1, Ndata)
ygrid = np.linspace(0,1, Ndata+1) # use a slighly different size to check differences
zgrid = np.linspace(0,1, Ndata+2)

f_2d_grid = f_2d(xgrid.reshape(-1,1), ygrid)

plt.imshow(f_2d_grid.T)
plt.title(u'image of a 2D function ({}² pts)'.format(Ndata));

f_2d_grid.shape

f_3d_grid = f_3d(xgrid.reshape(-1,1,1), ygrid.reshape(1,-1,1), zgrid)
f_3d_grid.shape

# Define the grid to interpolate on :
Ninterp = 1000
xinterp = np.linspace(0,1, Ninterp)
yinterp = np.linspace(0,1, Ninterp+1) # use a slighly different size to check differences
zinterp = np.linspace(0,1, 5) # small dimension to avoid size explosion

# Build data for the interpolator
points_x, points_y = np.broadcast_arrays(xgrid.reshape(-1,1), ygrid)
points = np.vstack((points_x.flatten(),
                    points_y.flatten())).T
values = f_2d_grid.flatten()

# Build
%timeit f_2d_interp = LinearNDInterpolator(points, values)

f_2d_interp = LinearNDInterpolator(points, values)

# Evaluate
%timeit f_2d_interp(xinterp.reshape(-1,1), yinterp)

# Display
plt.imshow(f_2d_interp(xinterp.reshape(-1,1), yinterp).T)
plt.title(u'interpolation of a 2D function ({}² pts)'.format(Ninterp));

# Build data for the interpolator
points_x, points_y, points_z = np.broadcast_arrays(xgrid.reshape(-1,1,1), ygrid.reshape(1,-1,1), zgrid)
points = np.vstack((points_x.flatten(),
                    points_y.flatten(),
                    points_z.flatten()
                  )).T
values = f_3d_grid.flatten()
points.shape

%time f_3d_interp = LinearNDInterpolator(points, values) 

f_3d_interp = LinearNDInterpolator(points, values)

# Evaluate:
%time f_3d_interp(xinterp.reshape(-1,1), np.linspace(0,1,5), 0.25);

f_3d_interp_grid = f_3d_interp(xinterp.reshape(-1,1,1), 0.25, 0.25)

plt.plot(f_3d_interp_grid.flatten());

%%timeit # Build
f_2d_interp = RectBivariateSpline(xgrid, ygrid, f_2d_grid)

f_2d_interp = RectBivariateSpline(xgrid, ygrid, f_2d_grid)

%%timeit # Evaluate
f_2d_interp(xinterp, yinterp)

# Display
plt.imshow(f_2d_interp(xinterp, yinterp).T)
plt.title(u'interpolation of a 2D function ({}² pts)'.format(Ninterp));

# Prepare the coordinates to evaluate the array on :
points_x, points_y = np.broadcast_arrays(xinterp.reshape(-1,1), yinterp)
coord = np.vstack((points_x.flatten()*(len(xgrid)-1) , # a weird formula !
                   points_y.flatten()*(len(ygrid)-1)))
coord.shape

%%timeit # Build and Evaluate
f_2d_interp = map_coordinates(f_2d_grid, coord, order=1)
# Reshape
f_2d_interp = f_2d_interp.reshape(len(xinterp), len(yinterp))

# Display 
f_2d_interp = map_coordinates(f_2d_grid, coord, order=1).reshape(len(xinterp), len(yinterp))

plt.imshow(f_2d_interp.T)
plt.title(u'interpolation of a 2D function ({}² pts)'.format(Ninterp));

# Prepare the coordinates to evaluate the array on :
points_x, points_y, points_z = np.broadcast_arrays(xinterp.reshape(-1,1,1), yinterp.reshape(1,-1,1), zinterp)
coord = np.vstack((points_x.flatten()*(len(xgrid)-1), # a weird formula !
                   points_y.flatten()*(len(ygrid)-1),
                   points_z.flatten()*(len(zgrid)-1)
                   ))
coord.shape

%%timeit # Build and Evaluate
f_3d_interp = map_coordinates(f_3d_grid, coord, order=1)
# Reshape
f_3d_interp = f_3d_interp.reshape(len(xinterp), len(yinterp), len(zinterp))

f_3d_interp = map_coordinates(f_3d_grid, coord, order=1)
f_3d_interp = f_3d_interp.reshape(len(xinterp), len(yinterp), len(zinterp))
f_3d_interp.shape

n_z = 1
plt.imshow(f_3d_interp[:,:,n_z])
plt.title('f_3d(x,y, z={:.2f})'.format(zinterp[n_z]))
plt.colorbar();

smin = [xgrid[0], ygrid[0]]
smax = [xgrid[-1], ygrid[-1]]
orders = [len(xgrid), len(ygrid)]

print(smin, smax, orders)

%%timeit
minterp = MultilinearInterpolator(smin,smax,orders)
minterp.set_values(np.atleast_2d(f_2d_grid.flatten()))

minterp.grid.shape

f_2d_grid.shape

# Prepare the coordinates to evaluate the array on :
points_x, points_y = np.broadcast_arrays(xinterp.reshape(-1,1), yinterp)
coord = np.vstack((points_x.flatten(),
                   points_y.flatten()))
coord.shape

%timeit f_2d_interp = minterp(coord).reshape(len(xinterp), len(yinterp))

# Display 
f_2d_interp = minterp(coord).reshape(len(xinterp), len(yinterp))

plt.imshow(f_2d_interp.T)
plt.title(u'interpolation of a 2D function ({}² pts)'.format(Ninterp));
plt.colorbar();

smin = [xgrid[0], ygrid[0], zgrid[0]]
smax = [xgrid[-1], ygrid[-1], zgrid[-1]]
orders = [len(xgrid), len(ygrid), len(zgrid)]

print(smin, smax, orders)

minterp = MultilinearInterpolator(smin,smax,orders)
minterp.set_values(np.atleast_2d(f_3d_grid.flatten()))

# Prepare the coordinates to evaluate the array on :
points_x, points_y, points_z = np.broadcast_arrays(xinterp.reshape(-1,1,1), yinterp.reshape(1,-1,1), zinterp)
coord = np.vstack((points_x.flatten(), # a weird formula !
                   points_y.flatten(),
                   points_z.flatten()
                   ))
coord.shape

%timeit minterp(coord)

f_3d_interp = minterp(coord)
f_3d_interp = f_3d_interp.reshape(len(xinterp), len(yinterp), len(zinterp))
f_3d_interp.shape

n_z = 1
plt.imshow(f_3d_interp[:,:,n_z])
plt.title('f_3d(x,y, z={:.2f})'.format(zinterp[n_z]))
plt.colorbar();

print('{:.1f} Mpts/s'.format(5e6 / 0.105 /1e6) )

%%timeit # Instanciate the interpolator
f_2d_interp = RegularGridInterpolator((xgrid, ygrid), f_2d_grid)

f_2d_interp = RegularGridInterpolator((xgrid, ygrid), f_2d_grid)

# Prepare the coordinates to evaluate the array on :
points_x, points_y = np.broadcast_arrays(xinterp.reshape(-1,1), yinterp)
coord = np.vstack((points_x.flatten(),
                   points_y.flatten()))
coord.shape

f_2d_interp

%%timeit # Interpolate
f_2d_interp_res = f_2d_interp(coord.T).reshape(len(xinterp), len(yinterp))

# Display 

f_2d_interp_res = f_2d_interp(coord.T).reshape(len(xinterp), len(yinterp))

plt.imshow(f_2d_interp_res.T)
plt.title(u'interpolation of a 2D function ({}² pts)'.format(Ninterp));
plt.colorbar();

f_3d_interp = RegularGridInterpolator((xgrid, ygrid, zgrid), f_3d_grid)

# Prepare the coordinates to evaluate the array on :
points_x, points_y, points_z = np.broadcast_arrays(xinterp.reshape(-1,1,1), yinterp.reshape(1,-1,1), zinterp)
coord = np.vstack((points_x.flatten(), # a weird formula !
                   points_y.flatten(),
                   points_z.flatten()
                   ))
coord.shape

%%timeit # Interpolate
f_3d_interp_res = f_3d_interp(coord.T).reshape(len(xinterp), len(yinterp), len(zinterp))