From d4676e099bafc63631d06d83c5589f185e6b3d3b Mon Sep 17 00:00:00 2001
From: Daniel Szemerey <szemy2@gmail.com>
Date: Thu, 18 Nov 2021 10:59:06 +0100
Subject: [PATCH] feat(Pytorch): added custom model to pytorch-forecasting (#8)

* feat: Refractored and created new model. Pipeline not ready yet.

* feat: Implemented and refactored a data pipeline.

* ref: Refractored to make more sense.

* feat: Training works now with models that you can change.

* feat: Added predict function but without working instructions.

* feat: gitignore.
---
 .gitignore                                |  1 +
 {lightning => archive}/forecasting.py     |  0
 {lightning => archive}/model_lightning.py |  0
 {lightning => archive}/run_lightning.py   |  0
 lightning/create_dataset.py               | 64 +++++++++++++++
 lightning/define_model.py                 |  0
 lightning/models/built_in_models.py       | 19 +++++
 lightning/models/custom_model.py          | 98 +++++++++++++++++++++++
 lightning/options.py                      | 63 +++++++++++++++
 lightning/run_pipeline.py                 | 41 ++++++++++
 lightning/timeseries_dataset.py           | 18 -----
 lightning/train_predict.py                | 67 ++++++++++++++++
 12 files changed, 353 insertions(+), 18 deletions(-)
 rename {lightning => archive}/forecasting.py (100%)
 rename {lightning => archive}/model_lightning.py (100%)
 rename {lightning => archive}/run_lightning.py (100%)
 create mode 100644 lightning/create_dataset.py
 create mode 100644 lightning/define_model.py
 create mode 100644 lightning/models/built_in_models.py
 create mode 100644 lightning/models/custom_model.py
 create mode 100644 lightning/options.py
 create mode 100644 lightning/run_pipeline.py
 delete mode 100644 lightning/timeseries_dataset.py
 create mode 100644 lightning/train_predict.py

diff --git a/.gitignore b/.gitignore
index b6e4761..cbaaafb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -127,3 +127,4 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+lightning/lightning_logs/
diff --git a/lightning/forecasting.py b/archive/forecasting.py
similarity index 100%
rename from lightning/forecasting.py
rename to archive/forecasting.py
diff --git a/lightning/model_lightning.py b/archive/model_lightning.py
similarity index 100%
rename from lightning/model_lightning.py
rename to archive/model_lightning.py
diff --git a/lightning/run_lightning.py b/archive/run_lightning.py
similarity index 100%
rename from lightning/run_lightning.py
rename to archive/run_lightning.py
diff --git a/lightning/create_dataset.py b/lightning/create_dataset.py
new file mode 100644
index 0000000..35171bb
--- /dev/null
+++ b/lightning/create_dataset.py
@@ -0,0 +1,64 @@
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor
+
+from pytorch_forecasting import TimeSeriesDataSet, TemporalFusionTransformer
+from pytorch_forecasting.data import GroupNormalizer
+from pytorch_forecasting.metrics import QuantileLoss
+
+import sys
+sys.path.insert(0, '..')
+
+from load_data import load_files
+
+print("success")
+
+def load_format_data(data_dir):
+    print("Data Starting    ===>", end="  ")
+    # load data
+    data = load_files(data_dir, add_features=True, log_returns=False, narrow_format=True)
+    # need to treat time as an independent column
+    data = data.reset_index().rename({'index':'time'}, axis = 'columns')
+    # we need a `time_idx` column for pytorch-forecasting, so we convert the time column to a time index by encoding the dates as consecutive days from the first date.
+    data['time_idx'] = (data['time']-data['time'].min()).astype('timedelta64[D]').astype(int)+1
+    # volume needs some love before we can use it
+    data.drop(columns=['volume'], inplace=True)
+
+    data['month'] = data['month'].astype(str)
+    data['day_month'] = data['day_month'].astype(str)
+    data['day_week'] = data['day_week'].astype(str)
+    
+    
+    print("<===  Data Loaded")
+    print(data.head(3))
+    print(data.describe())
+    print("")
+
+    
+    
+    
+    return data
+
+
+def create_dataloaders(data, kwargs):
+    print("DataLoader Starting ===>", end="  ")
+    # training_cutoff = "YYYY-MM-DD"  # day for cutoff
+    # training_cutoff = data["time_idx"].max() - max_prediction_length
+    batch_size = kwargs['batch_size']
+    del kwargs['batch_size']
+
+    training_dataset = TimeSeriesDataSet(
+        data, # data[lambda x: x.date < training_cutoff],
+        **kwargs
+    )
+
+    #%%
+    # create validation and training dataset
+    validation = TimeSeriesDataSet.from_dataset(training_dataset, data, predict=True, stop_randomization=True)
+    train_dataloader = training_dataset.to_dataloader(train=True, batch_size=batch_size, num_workers=0)
+    val_dataloader = validation.to_dataloader(train=False, batch_size=batch_size, num_workers=0)
+
+
+    print("<===  DataLoader Created")
+    
+    
+    return training_dataset, train_dataloader, val_dataloader
\ No newline at end of file
diff --git a/lightning/define_model.py b/lightning/define_model.py
new file mode 100644
index 0000000..e69de29
diff --git a/lightning/models/built_in_models.py b/lightning/models/built_in_models.py
new file mode 100644
index 0000000..b2fe570
--- /dev/null
+++ b/lightning/models/built_in_models.py
@@ -0,0 +1,19 @@
+from pytorch_forecasting import TemporalFusionTransformer
+from pytorch_forecasting.metrics import QuantileLoss
+
+def create_TemporalFusionTransformer(training_dataset, model_options):
+    # create the model
+    tft = TemporalFusionTransformer.from_dataset(
+        training_dataset,
+        learning_rate=0.03,
+        hidden_size=32,
+        attention_head_size=1,
+        dropout=0.1,
+        hidden_continuous_size=16,
+        output_size=7,
+        loss=QuantileLoss(),
+        log_interval=2,
+        reduce_on_plateau_patience=4
+    )
+    print(f"Number of parameters in network: {tft.size()/1e3:.1f}k")
+    return tft
\ No newline at end of file
diff --git a/lightning/models/custom_model.py b/lightning/models/custom_model.py
new file mode 100644
index 0000000..94c3372
--- /dev/null
+++ b/lightning/models/custom_model.py
@@ -0,0 +1,98 @@
+#%%
+import warnings
+from typing import Dict
+
+warnings.filterwarnings("ignore")
+
+import torch
+from torch import nn
+
+from pytorch_forecasting.models import BaseModel
+from pytorch_forecasting import TimeSeriesDataSet
+
+
+#%%
+class FullyConnectedModule(nn.Module):
+    def __init__(self, input_size: int, output_size: int, hidden_size: int, n_hidden_layers: int):
+        super().__init__()
+
+        # input layer
+        module_list = [nn.Linear(input_size, hidden_size), nn.ReLU()]
+        # hidden layers
+        for _ in range(n_hidden_layers):
+            module_list.extend([nn.Linear(hidden_size, hidden_size), nn.ReLU()])
+        # output layer
+        module_list.append(nn.Linear(hidden_size, output_size))
+
+        self.sequential = nn.Sequential(*module_list)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x of shape: batch_size x n_timesteps_in
+        # output of shape batch_size x n_timesteps_out
+        return self.sequential(x)
+
+
+
+#%%
+class FullyConnectedModel(BaseModel):
+    def __init__(self, input_size: int, output_size: int, hidden_size: int, n_hidden_layers: int, **kwargs):
+        # saves arguments in signature to `.hparams` attribute, mandatory call - do not skip this
+        self.save_hyperparameters()
+        # pass additional arguments to BaseModel.__init__, mandatory call - do not skip this
+        super().__init__(**kwargs)
+        self.network = FullyConnectedModule(
+            input_size=self.hparams.input_size,
+            output_size=self.hparams.output_size,
+            hidden_size=self.hparams.hidden_size,
+            n_hidden_layers=self.hparams.n_hidden_layers,
+        )
+
+    def forward(self, x: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        # x is a batch generated based on the TimeSeriesDataset
+        network_input = x["encoder_cont"].squeeze(-1)
+        prediction = self.network(network_input)
+
+        # rescale predictions into target space
+        prediction = self.transform_output(prediction, target_scale=x["target_scale"])
+
+        # We need to return a dictionary that at least contains the prediction
+        # The parameter can be directly forwarded from the input.
+        # The conversion to a named tuple can be directly achieved with the `to_network_output` function.
+        return self.to_network_output(prediction=prediction)
+    
+    @classmethod
+    def from_dataset(cls, dataset: TimeSeriesDataSet, **kwargs):
+        new_kwargs = {
+            "output_size": dataset.max_prediction_length,
+            "input_size": dataset.max_encoder_length,
+        }
+        new_kwargs.update(kwargs)  # use to pass real hyperparameters and override defaults set by dataset
+        # example for dataset validation
+        assert dataset.max_prediction_length == dataset.min_prediction_length, "Decoder only supports a fixed length"
+        assert dataset.min_encoder_length == dataset.max_encoder_length, "Encoder only supports a fixed length"
+        assert (
+            len(dataset.time_varying_known_categoricals) == 0
+            and len(dataset.time_varying_known_reals) == 0
+            and len(dataset.time_varying_unknown_categoricals) == 0
+            and len(dataset.static_categoricals) == 0
+            and len(dataset.static_reals) == 0
+            and len(dataset.time_varying_unknown_reals) == 1
+            and dataset.time_varying_unknown_reals[0] == dataset.target
+        ), "Only covariate should be the target in 'time_varying_unknown_reals'"
+
+        return super().from_dataset(dataset, **new_kwargs)
+    
+    def calculate_prediction_actual_by_variable(x, train):
+        pass
+    
+
+
+
+
+# %%
+def create_FullyConnectedModel(training_dataset, kwargs):
+    model = FullyConnectedModel.from_dataset(training_dataset, **kwargs)
+    model.summarize("full")  # print model summary
+    model.hparams
+    
+    return model
diff --git a/lightning/options.py b/lightning/options.py
new file mode 100644
index 0000000..bbf6c9a
--- /dev/null
+++ b/lightning/options.py
@@ -0,0 +1,63 @@
+from pytorch_forecasting.data import GroupNormalizer
+from pytorch_forecasting.models.temporal_fusion_transformer import TemporalFusionTransformer
+
+from lightning.models.custom_model import FullyConnectedModel
+
+training_options = dict(
+    max_epochs=100,
+    gpus=0,
+    gradient_clip_val=0.1,
+    limit_train_batches=30
+)
+
+model_options_fcn = dict(
+    hidden_size=64,
+    n_hidden_layers=2,
+)
+
+model_options_tft = dict(
+    hidden_size=64,
+    n_hidden_layers=2,
+)
+
+dataset_options_tft = dict(
+    time_idx= 'time_idx',
+    target= 'returns',
+    # weight="weight",
+    group_ids=[ 'ticker' ],
+    min_encoder_length = 36, # this is the look-back window, see https://github.com/jdb78/pytorch-forecasting/issues/448
+    max_encoder_length = 36,
+    min_prediction_length = 1,
+    max_prediction_length = 1,
+    static_categoricals=[  ],
+    static_reals=[  ],
+    time_varying_known_categoricals=[ 'day_month', 'day_week', 'month' ],
+    time_varying_known_reals=[ 'vol_10', 'vol_20', 'vol_30', 'vol_60', 'mom_10', 'mom_20', 'mom_30', 'mom_60', 'mom_90'],
+    time_varying_unknown_categoricals=[  ],
+    time_varying_unknown_reals=[ 'returns' ],
+    allow_missing_timesteps=True,
+    target_normalizer=GroupNormalizer(
+        groups=['ticker'], transformation="softplus"),
+    batch_size=128
+)
+
+dataset_options_fcn = dict(
+    time_idx= 'time_idx',
+    target= 'returns',
+    # weight="weight",
+    group_ids=[ 'ticker' ],
+    min_encoder_length = 36, # this is the look-back window, see https://github.com/jdb78/pytorch-forecasting/issues/448
+    max_encoder_length = 36,
+    min_prediction_length = 1,
+    max_prediction_length = 1,
+    static_categoricals=[  ],
+    static_reals=[  ],
+    time_varying_known_categoricals=[ ],
+    time_varying_known_reals=[ ],
+    time_varying_unknown_categoricals=[  ],
+    time_varying_unknown_reals=[ 'returns' ],
+    allow_missing_timesteps=True,
+    target_normalizer=GroupNormalizer(
+        groups=['ticker'], transformation="softplus"),
+    batch_size=128
+)
diff --git a/lightning/run_pipeline.py b/lightning/run_pipeline.py
new file mode 100644
index 0000000..8d73c75
--- /dev/null
+++ b/lightning/run_pipeline.py
@@ -0,0 +1,41 @@
+
+#%%
+from create_dataset import load_format_data, create_dataloaders
+from train_predict import train_model, predict
+from models.built_in_models import create_TemporalFusionTransformer
+from models.custom_model import create_FullyConnectedModel
+from options import training_options, model_options_tft, model_options_fcn, dataset_options_tft, dataset_options_fcn
+
+import warnings
+warnings.filterwarnings("ignore")
+
+#%%
+def run_pipeline(model_name, data_dir):
+    data = load_format_data(data_dir)  
+    
+    dataset_options, model_options, _create_model = select_model(model_name)
+    training_dataset, train_dataloader, val_dataloader = create_dataloaders(data, dataset_options)
+    
+    model = _create_model( training_dataset, model_options )
+    trainer = train_model(model, train_dataloader, val_dataloader, training_options)
+    predict(trainer, model, val_dataloader)
+    
+#%%
+
+def select_model(model_name):
+    if model_name == "FullyConnectedLayer":
+        return dataset_options_fcn, model_options_fcn, create_FullyConnectedModel
+    elif model_name == "TemporalFusionTransformer":
+        return dataset_options_tft, model_options_tft, create_TemporalFusionTransformer
+    else:
+        assert False, "No such model exists."
+    
+
+#%%
+run_pipeline("FullyConnectedLayer", '../data/')
+
+# #%%
+# if __name__ == '__main__':
+#     run_pipeline("FullyConnectedLayer", '../data/')
+
+ 
\ No newline at end of file
diff --git a/lightning/timeseries_dataset.py b/lightning/timeseries_dataset.py
deleted file mode 100644
index e0d72c8..0000000
--- a/lightning/timeseries_dataset.py
+++ /dev/null
@@ -1,18 +0,0 @@
-import os
-import pandas as pd
-
-import torch
-from torch.utils.data import Dataset
-
-
-class TimeSeriesDataset(Dataset):
-    def __init__(self, file_loader_hook):
-        
-        df = file_loader_hook()
-
-    def __len__(self):
-        return len(self.img_labels)
-
-    def __getitem__(self, idx):
-        pass
-        # return image, label
\ No newline at end of file
diff --git a/lightning/train_predict.py b/lightning/train_predict.py
new file mode 100644
index 0000000..d23a76a
--- /dev/null
+++ b/lightning/train_predict.py
@@ -0,0 +1,67 @@
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor
+import torch
+from torch import nn
+
+
+import warnings
+warnings.filterwarnings("ignore")
+
+import sys
+sys.path.insert(0, '..')
+
+from load_data import load_files
+
+
+
+
+def train_model(model, train_dataloader, val_dataloader, kwargs):
+    print()
+    print("Creating Trainer ===>", end="  ") 
+    # define trainer with early stopping
+    early_stop_callback = EarlyStopping(monitor="val_loss", min_delta=1e-4, patience=1, verbose=False, mode="min")
+    lr_logger = LearningRateMonitor()
+    
+    trainer = pl.Trainer(
+        **kwargs,
+        callbacks=[lr_logger, early_stop_callback],
+    )
+    print("<===  Trainer Created")
+    print("Finding Optimal LR  ===>", end="  ")
+    # find optimal learning rate (set limit_train_batches to 1.0 and log_interval = -1)
+    res = trainer.tuner.lr_find(
+        model, train_dataloader=train_dataloader, val_dataloaders=val_dataloader, early_stop_threshold=1000.0, max_lr=0.3,
+    )
+
+    print(f"<=== suggested learning rate: {res.suggestion()}")
+    fig = res.plot(show=True, suggest=True)
+    fig.show()
+
+    print("Training the model  ===>", end="  ")
+    # fit the model
+    trainer.fit(
+        model, train_dataloader=train_dataloader, val_dataloaders=val_dataloader,
+    )
+    
+    print("<=== Training Finished")
+    return trainer
+    
+    
+def predict(trainer, model, val_dataloader):
+    pass
+    # best_model_path = trainer.checkpoint_callback.best_model_path
+    # best_model = model.load_from_checkpoint(best_model_path)
+    
+    # # calcualte mean absolute error on validation set
+    # actuals = torch.cat([y[0] for x, y in iter(val_dataloader)])
+    # predictions = best_model.predict(val_dataloader)
+    # (actuals - predictions).abs().mean()
+    
+    # raw_predictions, x = best_model.predict(val_dataloader, mode="raw", return_x=True)
+    # for idx in range(10):  # plot 10 examples
+    #     best_model.plot_prediction(x, raw_predictions, idx=idx, add_loss_to_title=True)
+    
+    # predictions, x = best_model.predict(val_dataloader, return_x=True)
+    # predictions_vs_actuals = best_model.calculate_prediction_actual_by_variable(x, predictions)
+    # best_model.plot_prediction_actual_by_variable(predictions_vs_actuals)
+    
\ No newline at end of file