#!/usr/bin/env python
# coding: utf-8

# There are faster ways of doing spline interpolation, but a quick and dirty example of doing so in a python-stratify context follows:

# In[1]:


# Taken from https://docs.scipy.org/doc/scipy-0.18.1/reference/generated/scipy.interpolate.CubicSpline.html#scipy.interpolate.CubicSpline
import numpy as np
import scipy.interpolate

x = np.arange(10)
y = np.sin(x)
cs = scipy.interpolate.CubicSpline(x, y)
xs = np.arange(-0.5, 9.6, 0.1)

ys = cs(xs)

import matplotlib.pyplot as plt
plt.plot(x, y, 'ob', label='real')
plt.plot(xs, ys, color='red', label='scipy.interpolate')
plt.show()


# In[2]:


import numpy as np
import stratify


# In[3]:


class SplineInterp(stratify.PyFuncInterpolator):
    def column_prep(self, z_target, z_src, fz_src, increasing):
        self.spline = scipy.interpolate.CubicSpline(z_src, fz_src, axis=-1)
    
    def interp_kernel(self, index, z_src, fz_src, level, output_array):
        output_array[:] = self.spline(level)
        

class SplineExtrap(stratify.PyFuncExtrapolator):
    def column_prep(self, z_target, z_src, fz_src, increasing):
        self.spline = scipy.interpolate.CubicSpline(z_src, fz_src, axis=-1)
    
    def extrap_kernel(self, direction, z_src, fz_src, level, output_array):
        output_array[:] = self.spline(level)


# In[4]:


result = stratify.interpolate(xs, x, y,
                              interpolation=SplineInterp(),
                              extrapolation=SplineExtrap())


# In[5]:


import matplotlib.pyplot as plt
plt.plot(x, y, 'ob', label='actual')
plt.plot(xs, result, color='red', label='stratify')
plt.show()


# Just to demonstrate that we are paying a price for dropping back into python at each iteration of the target levels, let's take a look at the execution time.

# In[6]:


get_ipython().run_cell_magic('timeit', '', 'result = stratify.interpolate(xs, x, y,\n                              interpolation=SplineInterp(),\n                              extrapolation=SplineExtrap())\n')


# In[7]:


get_ipython().run_cell_magic('timeit', '', 'cs = scipy.interpolate.CubicSpline(x, y)\nys = cs(xs)\n')