%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

# Use seaborn for plotting defaults.
# This can be safely commented-out if seaborn is not installed
import seaborn; seaborn.set()

# Install the package
# source at http://github.com/jakevdp/supersmoother
# or use ``pip install supersmoother``
from supersmoother import SuperSmoother, LinearSmoother

def make_test_set(N=200, rseed_x=None, rseed_y=None):
    """Code to generate the test set from Friedman 1984"""
    rng_x = np.random.RandomState(rseed_x)
    rng_y = np.random.RandomState(rseed_y)
    x = rng_x.rand(N)
    dy = x
    y = np.sin(2 * np.pi * (1 - x) ** 2) + dy * rng_y.randn(N)
    return x, y, dy

# Generate and visualize the data
t, y, dy = make_test_set(rseed_x=0, rseed_y=1)
plt.errorbar(t, y, dy, fmt='o', alpha=0.3);

# fit the supersmoother model
model = SuperSmoother()
model.fit(t, y, dy)

# find the smoothed fit to the data
tfit = np.linspace(0, 1, 1000)
yfit = model.predict(tfit)

# Show the smoothed model of the data
plt.errorbar(t, y, dy, fmt='o', alpha=0.3)
plt.plot(tfit, yfit, '-k');

plt.errorbar(t, y, dy, fmt='o', alpha=0.3)

for smooth in model.primary_smooths:
    plt.plot(tfit, smooth.predict(tfit),
             label='span = {0:.2f}'.format(smooth.span))
plt.legend();

t = np.linspace(0, 1, 1000)
plt.plot(t, model.span(t))
plt.xlabel('t')
plt.ylabel('smoothed span value');

N = 1000
span = span2 = 0

tfit = np.linspace(0, 1, 100)

for rseed in np.arange(N):
    t, y, dy = make_test_set(rseed_x=0, rseed_y=rseed)
    model = SuperSmoother().fit(t, y, dy)
    span += model.span(tfit)
    span2 += model.span(tfit) ** 2
    
mean = span / N
std = np.sqrt(span2 / N - mean ** 2)
plt.plot(tfit, mean)
plt.fill_between(tfit, mean - std, mean + std, alpha=0.3)
plt.xlabel('t')
plt.ylabel('resulting span');

rng = np.random.RandomState(0)
t = rng.rand(200)
dy = 0.5
y = np.sin(5 * np.pi * t ** 2) + dy * rng.randn(200)

plt.errorbar(t, y, dy, fmt='o', alpha=0.3)

for alpha in [0, 8, 10]:
    smoother = SuperSmoother(alpha=alpha)
    smoother.fit(t, y, dy)
    plt.plot(tfit, smoother.predict(tfit),
             label='alpha = {0}'.format(alpha))
plt.legend(loc=2);