feat(Tests): added basic unit tests for walk_forward_train_test() (#22)

* feat(Tests): added basic unit tests for walk_forward_train_test()

* fix(Tests): inherit from BaseEstimator, fix index problems in walk_forward_train_test

* fix(WalkForward): predictions were mistakenly removed, oops

* fix(WalkForward): mistakenly re-assiging model

.gitignore

@@ -128,3 +128,5 @@ dmypy.json
 # Pyre type checker
 .pyre/
 lightning/lightning_logs/
+
+results.csv

run_pipeline.py

@@ -92,7 +92,7 @@ results = pd.DataFrame()
 all_assets = get_all_assets('data/')
 
 for asset in all_assets:
-    for method in ['regression', 'classification']:
+    for method in ['regression']:
         current_result = run_whole_pipeline(
             ticker_to_predict = asset,
             models = regression_models if method == 'regression' else classification_models,

tests/test_walk_forward.py

@@ -0,0 +1,41 @@
+import pytest
+import numpy as np
+import pandas as pd
+from utils.walk_forward import walk_forward_train_test
+from sklearn.base import BaseEstimator
+
+def __generate_test_data():
+    no_columns = 6
+    no_rows = 100
+    X = [[row] * no_columns for row in range(no_rows)]
+    assert X[0][0] == 0
+    assert X[1][0] == 1
+    assert X[2][0] == 2
+    assert X[3][0] == 3
+    X = pd.DataFrame(X)
+
+    y = [row+1 for row in range(no_rows)]
+    assert y[0] == 1
+    assert y[1] == 2
+    assert y[2] == 3
+    y = pd.Series(y)
+
+    return X, y
+
+
+def test_walk_forward_train_test():
+    X, y = __generate_test_data()
+
+    window_length = 10
+    class StubModel(BaseEstimator):
+
+        def fit(self, X, y):
+            assert len(X) == window_length
+            for i in range(len(X)):
+                assert X[i][0] + 1 == y[i]
+
+        def predict(self, X):
+            return np.array([X[0][0] + 1])
+
+    model = StubModel()
+    walk_forward_train_test('test', model, X, y, window_length, 10)

utils/walk_forward.py

@@ -12,8 +12,8 @@ def walk_forward_train_test(
                             retrain_every: int
                         ) -> tuple[pd.Series, pd.Series]:
                         
-    predictions = pd.Series(index=y.index)
+    predictions = pd.Series(index=y.index).rename(model_name)
-    models = pd.Series(index=y.index)
+    models = pd.Series(index=y.index).rename(model_name)
 
     train_from = window_size
     train_till = y.index[-1]
@@ -30,7 +30,7 @@ def walk_forward_train_test(
 
         if iterations_since_retrain >= retrain_every or pd.isna(models[i-1]):
             current_model = clone(model)
-            current_model.fit(X_slice.to_numpy(), y_slice)
+            current_model.fit(X_slice.to_numpy(), y_slice.to_numpy())
             iterations_since_retrain = 0
         else:
             current_model = models[i-1]