# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # KTRX Example

# %load_ext autoreload
# %autoreload 2

# +
import pandas as pd
import numpy as np
from orbit.models.ktrlite import KTRLiteMAP
from orbit.models.ktrx import KTRXFull, KTRXAggregated

from orbit.estimators.pyro_estimator import PyroEstimatorVI
from orbit.estimators.stan_estimator import StanEstimatorMAP

from orbit.utils.features import make_fourier_series_df, make_fourier_series
from orbit.diagnostics.plot import plot_predicted_data, plot_predicted_components, plot_ktr_lev_knots
from orbit.diagnostics.metrics import smape
from orbit.utils.dataset import load_iclaims, load_electricity_demand

import matplotlib.pyplot as plt
plt.style.use("fivethirtyeight")

# %matplotlib inline
# -

pd.set_option('display.float_format', lambda x: '%.5f' % x)

# ## Data

# +
df = load_iclaims()
date_col = 'week'
response_col = 'claims'

print(df.shape)
df.head()
# -

print(f'starts with {df[date_col].min()}\nends with {df[date_col].max()}\nshape: {df.shape}')

test_size = 52
train_df = df[:-test_size]
test_df = df[-test_size:]

# ## KTRLite

# +
# level_knot_dates = pd.date_range(start='1981-01-01', end='1990-12-31', periods=21)
# level_knot_dates
# -

ktrlite = KTRLiteMAP(
    response_col=response_col,
    date_col=date_col,
    # seasonality
    seasonality=[52],
    seasonality_fs_order=[3],
    level_knot_scale=.1,
    span_level=.05,
    # level_knot_dates=level_knot_dates,
    # date_freq='D',
    span_coefficients=.3,
    estimator_type=StanEstimatorMAP,
    n_bootstrap_draws=1e4,
)

ktrlite.fit(train_df)

predicted_df = ktrlite.predict(df=test_df, decompose=True)
predicted_df.head()

'{:.2%}'.format(smape(predicted_df['prediction'].values, test_df[response_col].values))

_ = plot_predicted_data(training_actual_df=train_df, predicted_df=predicted_df, 
                        date_col=date_col, actual_col=response_col, 
                        test_actual_df=test_df)

_ = plot_predicted_components(predicted_df=predicted_df, date_col=date_col, 
                              plot_components=['trend', 'seasonality_52'])

# get the knots
lev_knots = ktrlite._aggregate_posteriors['map']['lev_knot']
decomp_df = ktrlite.predict(df=train_df, decompose=True)

_ = plot_ktr_lev_knots(train_df, decomp_df, date_col, response_col, ktrlite._level_knot_dates, lev_knots)

# ## KTRX

# prepare the input
level_knot_dates = ktrlite._level_knot_dates
level_knots = ktrlite._aggregate_posteriors['map']['lev_knot'][0]
seasonal_knots_input = {
    '_seas_coef_knot_dates': ktrlite._coef_knot_dates,
    '_sea_coef_knot': ktrlite._aggregate_posteriors['map']['coef_knot'],
    '_seasonality': ktrlite._seasonality,
    '_seasonality_fs_order': ktrlite._seasonality_fs_order,
}

# +
ktrx = KTRXAggregated(
    response_col=response_col,
    date_col=date_col,
    level_knot_dates=level_knot_dates,
    level_knots=level_knots,
    seasonal_knots_input=seasonal_knots_input,
    regressor_col=['trend.unemploy', 'trend.filling', 'trend.job'],
    level_knot_scale=.1, #.01
    span_coefficients=0.3,
    rho_coefficients=0.15,
    # pyro optimization parameters
    seed=8888,
    num_steps=1000,
    num_sample=1000,
    learning_rate=0.1,
    #learning_rate_total_decay=0.05,
    verbose=True,
    message=100,
    aggregate_method="median",
    estimator_type=PyroEstimatorVI,
)

ktrx.fit(train_df)

# +
# ktrx = KTRXFull(
#     response_col=response_col,
#     date_col=date_col,
#     level_knot_dates=level_knot_dates,
#     level_knots=level_knots,
#     seasonal_knots_input=seasonal_knots_input,
#     regressor_col=['trend.unemploy', 'trend.filling', 'trend.job'],
#     level_knot_scale=.1, #.01
#     span_coefficients=0.3,
#     rho_coefficients=0.15,
#     # pyro optimization parameters
#     seed=8888,
#     num_steps=1000,
#     num_sample=1000,
#     learning_rate=0.1,
#     #learning_rate_total_decay=0.05,
#     verbose=True,
#     message=100,
#     estimator_type=PyroEstimatorVI,
# )

# ktrx.fit(train_df)
# -

coef_df = ktrx.get_regression_coefs(coefficient_method="smooth") # include_ci=True
print(coef_df.head())
print(coef_df.tail())

coef_df = ktrx.get_regression_coefs(coefficient_method="empirical") # include_ci=True
print(coef_df.head())
print(coef_df.tail())

coef_df = ktrx.get_regression_coefs(coefficient_method="smooth" ,
                                    date_array=['2017-06-11', '2017-06-18', '2017-06-25',
                                                '2017-07-02', '2017-07-09', '2017-07-16']) 
coef_df

coef_df = ktrx.get_regression_coefs(coefficient_method="empirical" ,
                                    date_array=['2017-06-11', '2017-06-18', '2017-06-25',
                                                '2017-07-02', '2017-07-09', '2017-07-16']) 
coef_df

_ = ktrx.plot_regression_coefs(with_knot=True, coefficient_method='smooth', ncol=3, figsize=(10, 3))

predicted_df = ktrx.predict(df=test_df, coefficient_method="smooth", decompose=False)
predicted_df.head()

predicted_df = ktrx.predict(df=test_df, coefficient_method="empirical", decompose=False)
predicted_df.head()

'{:.2%}'.format(smape(predicted_df['prediction'].values, test_df[response_col].values))

_ = plot_predicted_data(training_actual_df=train_df, predicted_df=predicted_df, 
                    date_col=date_col, actual_col=response_col, 
                    test_actual_df=test_df)

_ = plot_predicted_components(predicted_df=predicted_df, date_col=date_col, 
                          plot_components=['trend', 
                                           'seasonality_input',
                                           'regression'])

# ## Try a different data

# +
# from 2000-01-01 to 2008-12-31
df = load_electricity_demand()
df['electricity'] = np.log(df['electricity'])

date_col = 'date'
response_col = 'electricity'

print(df.shape)
df.head()
# -

test_size = 365
train_df = df[:-test_size]
test_df = df[-test_size:]

# +
ktrlite = KTRLiteMAP(
    response_col=response_col,
    date_col=date_col,
    # seasonality
    seasonality=[7, 365.25],
    seasonality_fs_order=[2, 5],
    level_knot_scale=.1,
    span_level=.05,
    # level_knot_dates=level_knot_dates,
    # date_freq='D',
    span_coefficients=.3,
    estimator_type=StanEstimatorMAP,
    n_bootstrap_draws=1e4,
)

ktrlite.fit(train_df)
# -

# prepare the input
level_knot_dates = ktrlite._level_knot_dates
level_knots = ktrlite._aggregate_posteriors['map']['lev_knot'][0]
seasonal_knots_input = {
    '_seas_coef_knot_dates': ktrlite._coef_knot_dates,
    '_sea_coef_knot': ktrlite._aggregate_posteriors['map']['coef_knot'],
    '_seasonality': ktrlite._seasonality,
    '_seasonality_fs_order': ktrlite._seasonality_fs_order,
}

# +
ktrx = KTRXFull(
    response_col=response_col,
    date_col=date_col,
    level_knot_dates=level_knot_dates,
    level_knots=level_knots,
    seasonal_knots_input=seasonal_knots_input,
    level_knot_scale=.1, #.01
    span_coefficients=0.3,
    rho_coefficients=0.15,
    # pyro optimization parameters
    seed=8888,
    num_steps=1000,
    num_sample=1000,
    learning_rate=0.1,
    #learning_rate_total_decay=0.05,
    verbose=True,
    message=100,
    estimator_type=PyroEstimatorVI,
)

ktrx.fit(train_df)
# -

predicted_df = ktrx.predict(df=test_df, decompose=True)
predicted_df.head()

'{:.2%}'.format(smape(predicted_df['prediction'].values, test_df[response_col].values))

_ = plot_predicted_data(training_actual_df=train_df, predicted_df=predicted_df, 
                    date_col=date_col, actual_col=response_col, 
                    test_actual_df=test_df)



StarcoderdataJupyterScriptsDedupFiltered