app.py

## Imports
import pickle
import warnings
import streamlit as st
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import datetime

import torch
from torch.distributions import Normal
from pytorch_forecasting import (
    TimeSeriesDataSet,
    TemporalFusionTransformer,
)

## Functions 
def raw_preds_to_df(raw,quantiles = None):
    """
    raw is output of model.predict with return_index=True
    quantiles can be provided like [0.1,0.5,0.9] to get interpretable quantiles
    in the output, time_idx is the first prediction time index (one step after knowledge cutoff)
    pred_idx the index of the predicted date i.e. time_idx + h - 1
    """
    index = raw[2]
    preds = raw[0].prediction
    dec_len = preds.shape[1]
    n_quantiles = preds.shape[-1]
    preds_df = pd.DataFrame(index.values.repeat(dec_len * n_quantiles, axis=0),columns=index.columns)
    preds_df = preds_df.assign(h=np.tile(np.repeat(np.arange(1,1+dec_len),n_quantiles),len(preds_df)//(dec_len*n_quantiles)))
    preds_df = preds_df.assign(q=np.tile(np.arange(n_quantiles),len(preds_df)//n_quantiles))
    preds_df = preds_df.assign(pred=preds.flatten().numpy())
    if quantiles is not None:
        preds_df['q'] = preds_df['q'].map({i:q for i,q in enumerate(quantiles)})

    preds_df['pred_idx'] = preds_df['time_idx'] + preds_df['h'] - 1
    return preds_df

def prepare_dataset(parameters, df, rain, temperature, datepicker, mapping):
    if rain != "Default":
        df["MTXWTH_Day_precip"] = rain_mapping[rain]
    
    df["MTXWTH_Temp_min"] = df["MTXWTH_Temp_min"] + temperature
    df["MTXWTH_Temp_max"] = df["MTXWTH_Temp_max"] + temperature

    lowerbound = datepicker - datetime.timedelta(days = 35) 
    upperbound = datepicker + datetime.timedelta(days = 30) 

    df = df.loc[(df["Date"].dt.date>lowerbound) & (df["Date"].dt.date<=upperbound)]
    
    df = TimeSeriesDataSet.from_parameters(parameters, df)
    return df.to_dataloader(train=False, batch_size=256,num_workers = 0)

@st.cache_data
def predict(_model, _dataloader, datepicker):
    out = _model.predict(_dataloader, mode="raw", return_x=True, return_index=True)
    preds = raw_preds_to_df(out, quantiles = None)

    def add_dates(group):
        group["date_imputed"] = [datepicker + datetime.timedelta(days=x) for x in range(30)]
        return group
    
    preds["date_imputed"] = preds.groupby("Group").pred.transform(add_dates)
    
    return preds[["date_imputed", "Group", "pred"]]


def generate_plot(df, preds):
    fig, axs = plt.subplots(2, 2, figsize=(8, 6))
    
    df = pd.merge(df, pred, left_on='Date', right_on='date_imputed') 
    # Plot scatter plots for each group
    axs[0, 0].scatter(df.loc[df['Group'] == '4', 'Date'], df.loc[df['Group'] == '4', 'sales'], color='grey', marker='o')
    axs[0, 0].plot(df.loc[df['Group'] == '4', 'Date'], df.loc[df['Group'] == '4', 'pred'], color = 'red')
    axs[0, 0].set_title('Article Group 1')
    
    axs[0, 1].scatter(df.loc[df['Group'] == '7', 'Date'], df.loc[df['Group'] == '7', 'sales'], color='blue', marker='o')
    axs[0, 1].set_title('Article Group 2')
    
    axs[1, 0].scatter(df.loc[df['Group'] == '1', 'Date'], df.loc[df['Group'] == '1', 'sales'], color='green', marker='o')
    axs[1, 0].set_title('Article Group 3')
    
    axs[1, 1].scatter(df.loc[df['Group'] == '6', 'Date'], df.loc[df['Group'] == '6', 'sales'], color='yellow', marker='o')
    axs[1, 1].set_title('Article Group 4')
    
    # Adjust spacing between subplots
    plt.tight_layout()
    
    #for ax in axs.flat:
    #    ax.set_xlim(df['Date'].min(), df['Date'].max())
    #    ax.set_ylim(df['sales'].min(), df['sales'].max())
    
    st.pyplot(fig)

@st.cache_data
def load_data():
    with open('data/parameters.pkl', 'rb') as f:
        parameters = pickle.load(f)
    df = pd.read_pickle('data/test_data.pkl')
    df = df.loc[(df["Branch"] == "15") & (df["Group"].isin(["6","7","4","1"]))]
    return parameters, df

@st.cache_resource
def init_model():
    model = TemporalFusionTransformer.load_from_checkpoint('model/tft_check.ckpt', map_location=torch.device('cpu'))
    return model
    
def main():
    ## Initiate Data
    parameters, df = load_data()
    model = init_model()
    rain_mapping = {
        "Yes" : 1,
        "No" : 0
    }
    
    # Start App
    st.title("Temporal Fusion Transformers for Interpretable Multi-horizon Time Series Forecasting")
    
    st.markdown(body = """
        ### Abstract
        Multi-horizon forecasting often contains a complex mix of inputs – including
        static (i.e. time-invariant) covariates, known future inputs, and other exogenous
        time series that are only observed in the past – without any prior information
        on how they interact with the target. Several deep learning methods have been
        proposed, but they are typically ‘black-box’ models which do not shed light on
        how they use the full range of inputs present in practical scenarios. In this pa-
        per, we introduce the Temporal Fusion Transformer (TFT) – a novel attention-
        based architecture which combines high-performance multi-horizon forecasting
        with interpretable insights into temporal dynamics. To learn temporal rela-
        tionships at different scales, TFT uses recurrent layers for local processing and
        interpretable self-attention layers for long-term dependencies. TFT utilizes spe-
        cialized components to select relevant features and a series of gating layers to
        suppress unnecessary components, enabling high performance in a wide range of
        scenarios. On a variety of real-world datasets, we demonstrate significant per-
        formance improvements over existing benchmarks, and showcase three practical
        interpretability use cases of TFT.
    
        ### Experiments
    """)
    st.write(df.head(5))
    rain = st.radio("Rain Indicator", ('Default', 'Yes', 'No'))
    
    temperature = st.slider('Change in Temperature', min_value=-10.0, max_value=10.0, value=0.0, step=0.25)
    
    datepicker = st.date_input("Start of Forecast", datetime.date(2022, 12, 24), min_value=datetime.date(2022, 6, 26) + datetime.timedelta(days = 35), max_value=datetime.date(2023, 6, 26) - datetime.timedelta(days = 30))

    if st.button("Forecast Sales", type="primary"):
        dataloader = prepare_dataset(parameters, df, rain, temperature, datepicker, rain_mapping)
        preds = predict(model, dataloader, datepicker)
        generate_plot(df, preds)

if __name__ == '__main__':
    main()