split train predict + fixes

src/model.py

@@ -7,6 +7,7 @@ from darts.dataprocessing.transformers import Scaler
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import os
 
 def load_ILINetDataset():
     """
@@ -41,7 +42,7 @@ def load_ILINetDataset():
 
 def preprocess_data(ilidata):
     ilitotal = ilidata['ILITOTAL']
-    covariates = ilidata.drop_columns(col_names='ILITOTAL')
+    covariates = ilidata.drop_columns(col_names='ILITOTAL')  # TODO perform predictions with covariates time series
 
     # Apply logaritmic transformation
     log_ili = ilitotal.map(np.log)
@@ -69,14 +70,27 @@ def scale_train(train_ili):
     return train_ili_scaled, scaler
 
 
-def train_exponential_smoothing(train_ili_scaled, val_ili):
-    exp_sm = ExponentialSmoothing()
-    exp_sm.fit(train_ili_scaled)
-    exp_preds = exp_sm.predict(len(val_ili))
+def train_(train_ili_scaled, val_ili, save_model_path='./models'):
+    model = ExponentialSmoothing()
+    model.fit(train_ili_scaled)
+    
+    model_name = str(model).split('(')[0]
+    path = f'{save_model_path}/{model_name}.pkl'
+    if not os.path.exists(path=path):
+        os.makedirs(path)
+        model.save(path)
+    else:
+        model = ExponentialSmoothing.load(path)
 
     # print('Exponential Smoothing mape:', mape(val_ili, scaler.inverse_transform(exp_preds)))
 
-    return exp_preds
+    return model
+
+
+def predict(model, val_ili):
+    preds = model.predict(len(val_ili))
+
+    return preds
 
 
 def inverse_scale_predictions(scaled_preds, scaler):
@@ -85,11 +99,13 @@ def inverse_scale_predictions(scaled_preds, scaler):
     return preds
 
 
-def plot_results(train_ili, val_ili, exp_preds):
+def plot_results(train_ili, val_ili, preds):
+    fig, _ = plt.subplots()
     train_ili.plot(label='train')
     val_ili.plot(label='val')
-    exp_preds.plot(label='preds')
-    plt.savefig('./charts/results.png')
+    preds.plot(label='preds')
+    
+    return fig