add testcases

arnet/ar_net.py

@@ -1,12 +1,14 @@
 from dataclasses import dataclass, field
 import logging
+import os
 import pandas as pd
 import matplotlib.pyplot as plt
 
 from fastai.data.core import DataLoaders
 from fastai.tabular.core import TabularPandas, TabDataLoader
 from fastai.tabular.learner import tabular_learner, TabularLearner
 from fastai.data.transforms import Normalize
+from fastai.learner import load_learner
 
 from arnet import utils, utils_data, plotting, fastai_mods
 
@@ -184,3 +186,29 @@ def plot_weights(self, **kwargs):
             ar=self.ar_params,
             **kwargs,
         )
+
+    def plot_fitted_obs(self, num_obs=100, **kwargs):
+        preds, y = self.learn.get_preds()
+        if num_obs is not None:
+            y = y[0:num_obs]
+            preds = preds[0:num_obs]
+        plotting.plot_prediction_sample(preds, y, **kwargs)
+
+    def plot_errors(self, **kwargs):
+        preds, y = self.learn.get_preds()
+        plotting.plot_error_scatter(preds, y, **kwargs)
+
+    def save_model(self, results_path="results", model_name=None):
+        # self.learn.freeze()
+        if model_name is None:
+            model_name = "ar{}_sparse_{:.3f}_ahead_{}_epoch_{}.pkl".format(
+                self.ar_order, self.sparsity, self.n_forecasts, self.n_epoch
+            )
+        self.learn.export(fname=os.path.join(results_path, model_name))
+        return self
+
+    def load_model(self, results_path="results", model_name=None, cpu=True):
+        self.learn = load_learner(fname=os.path.join(results_path, model_name), cpu=cpu)
+        # can unfreeze the model and fine_tune
+        self.learn.unfreeze()
+        return self

arnet/plotting.py

@@ -37,11 +37,11 @@ def plot_weights(ar_val, weights, ar=None, model_name="AR-Net", save=False, save
         plt.show()
 
 
-def plot_prediction_sample(predicted, actual, num_obs=100, model_name="AR-Net", save=False, savedir="results"):
+def plot_prediction_sample(predicted, actual, model_name="AR-Net", save=False, savedir="results"):
     fig2 = plt.figure()
     fig2.set_size_inches(10, 6)
-    plt.plot(actual[0:num_obs])
-    plt.plot(predicted[0:num_obs])
+    plt.plot(actual)
+    plt.plot(predicted)
     plt.legend(["Actual Time-Series", "{}-Prediction".format(model_name)])
     if save:
         if not os.path.exists(savedir):

tests/test_integration.py

@@ -11,6 +11,8 @@
 
 warnings.filterwarnings("ignore", message=".*nonzero.*", category=UserWarning)
 
+from fastai.learner import load_learner
+
 ## lazy imports ala fastai2 style (needed for nice print functionality)
 # from fastai.basics import *
 # from fastai.tabular.all import *
@@ -28,33 +30,57 @@
 
 
 class IntegrationTests(unittest.TestCase):
-    verbose = False
     plot = False
     save = False
 
     def test_random_data(self):
         df = pd.DataFrame({"x": [random.gauss(0.0, 1.0) for i in range(1000)]})
-        m = arnet.ARNet(
-            ar_order=3,
-            n_epoch=3,
-        )
+        m = arnet.ARNet(ar_order=3, n_epoch=3)
         m = m.tabularize(df)
         m = m.make_datasets()
         m = m.create_learner(
             sparsity=0.3,
         )
         m = m.find_lr(plot=False)
-        m = m.fit(cycles=3, plot=False)
+        m = m.fit(cycles=2, plot=False)
         log.info("coeff of random data: {}".format(m.coeff))
 
-    def test_ar_data(self):
-        self.save = True
-        if self.save:
-            if not os.path.exists(results_path):
-                os.makedirs(results_path)
+    def test_plot(self):
+        if not os.path.exists(results_path):
+            os.makedirs(results_path)
+
+        df = pd.DataFrame({"x": [random.gauss(0.0, 1.0) for i in range(1000)]})
+        m = arnet.ARNet(ar_order=3, n_epoch=3)
+        m = m.fit_with_defaults(series=df)
+        if self.plot:
+            m.learn.recorder.plot_loss()
+        m.plot_weights(save=True, savedir=results_path)
+        m.plot_fitted_obs(num_obs=100, save=True, savedir=results_path)
+        m.plot_errors(save=True, savedir=results_path)
+
+        shutil.rmtree(results_path)
+
+    def test_save_load(self):
+        if not os.path.exists(results_path):
+            os.makedirs(results_path)
 
+        df = pd.DataFrame({"x": [random.gauss(0.0, 1.0) for i in range(1000)]})
+        m = arnet.ARNet(ar_order=3, n_epoch=3)
+        m = m.fit_with_defaults(series=df)
+
+        # Optional:save and create inference learner
+        model_name = "ar{}_sparse_{:.3f}_ahead_{}_epoch_{}.pkl".format(m.ar_order, m.sparsity, m.n_forecasts, m.n_epoch)
+        m = m.save_model(results_path=results_path, model_name=model_name)
+        # can be loaded like this
+        m = m.load_model(results_path, model_name)
+        # can unfreeze the model and fine_tune
+        m.learn.fit_one_cycle(2, 0.0001)
+
+        shutil.rmtree(results_path)
+
+    def test_ar_data(self):
         data_name = "ar_3_ma_0_noise_0.100_len_10000"
-        df, data_config = arnet.load_from_file(data_path, data_name, load_config=True, verbose=self.verbose)
+        df, data_config = arnet.load_from_file(data_path, data_name, load_config=True, verbose=False)
         df = df[:1000]
 
         # Hyperparameters
@@ -75,19 +101,7 @@ def test_ar_data(self):
         log.info("ar params: {}".format(arnet.nice_print_list(ar_params)))
         log.info("model weights: {}".format(arnet.nice_print_list(m.coeff)))
 
-        # should be [0.20, 0.30, -0.50, ...]
-
-        preds, y = m.learn.get_preds()
-        if self.plot or self.save:
-            if self.plot:
-                m.learn.recorder.plot_loss()
-            arnet.plot_weights(
-                ar_val=len(ar_params[0]), weights=m.coeff[0], ar=ar_params[0], save=not self.plot, savedir=results_path
-            )
-            arnet.plot_prediction_sample(preds, y, num_obs=100, save=not self.plot, savedir=results_path)
-            arnet.plot_error_scatter(preds, y, save=not self.plot, savedir=results_path)
-
-        # if self.save:
+        # if save:
         #     # Optional:save and create inference learner
         #     learn.freeze()
         #     model_name = "ar{}_sparse_{:.3f}_ahead_{}_epoch_{}.pkl".format(ar_order, sparsity, n_forecasts, n_epoch)
@@ -97,17 +111,3 @@ def test_ar_data(self):
         # # can unfreeze the model and fine_tune
         # learn.unfreeze()
         # learn.fit_one_cycle(1, lr_at_min / 100)
-        #
-        # coeff2 = arnet.coeff_from_model(learn.model)
-        # log.info("ar params", arnet.nice_print_list(ar_params))
-        # log.info("model weights", arnet.nice_print_list(coeff))
-        # log.info("model weights2", arnet.nice_print_list(coeff2))
-
-        # if self.plot or self.save:
-        #     if self.plot:
-        #         learn.recorder.plot_loss()
-        #     arnet.plot_weights(
-        #         ar_val=len(ar_params[0]), weights=coeff2[0], ar=ar_params[0], save=not self.plot, savedir=results_path
-        #     )
-        # if self.save:
-        #     shutil.rmtree(results_path)

tests/test_legacy.py

@@ -28,7 +28,7 @@
 EPOCHS = 2
 
 
-class IntegrationTests(unittest.TestCase):
+class LegacyTests(unittest.TestCase):
     verbose = False
     plot = False
     save = False
@@ -114,7 +114,7 @@ def test_legacy_ar(self):
 
     def test_legacy_random(self):
         df = pd.DataFrame({"x": [random.gauss(0.0, 1.0) for i in range(1000)]})
-        learn = arnet.init_ar_learner(
+        learn = init_ar_learner(
             series=df,
             ar_order=3,
             n_forecasts=1,
@@ -129,7 +129,7 @@ def test_legacy_random(self):
         if self.plot:
             plt.show()
         print("lr at minimum: {}; steeptes lr: {}".format(lr_at_min, lr_steep))
-        learn.fit_one_cycle(n_epoch=n_epoch, lr_max=lr_at_min / 10)
+        learn.fit_one_cycle(n_epoch=EPOCHS, lr_max=lr_at_min / 10)
         if self.plot:
             learn.recorder.plot_loss()
             plt.show()