[FIX] Code refactoring

.github/workflows/ci.yaml

@@ -40,7 +40,7 @@ jobs:
           cache-environment: true
 
       - name: Install pip requirements
-        run: pip install ./
+        run: pip install ".[dev]"
 
       - name: Tests
         run: nbdev_test --do_print --timing --n_workers 0 --flags polars

action_files/test_models/src/multivariate_models.py

@@ -10,7 +10,7 @@
 from neuralforecast.models.tsmixer import TSMixer
 from neuralforecast.models.tsmixerx import TSMixerx
 from neuralforecast.models.itransformer import iTransformer
-from neuralforecast.models.stemgnn import StemGNN
+# from neuralforecast.models.stemgnn import StemGNN
 from neuralforecast.models.mlpmultivariate import MLPMultivariate
 from neuralforecast.models.timemixer import TimeMixer
 
@@ -26,13 +26,13 @@ def main(dataset: str = 'multivariate', group: str = 'ETTm2') -> None:
     train['ds'] = pd.to_datetime(train['ds'])
 
     models = [
-        SOFTS(h=horizon, n_series=7, input_size=2 * horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500),
-        TSMixer(h=horizon, n_series=7, input_size=2 * horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500),
-        TSMixerx(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500),
-        iTransformer(h=horizon, n_series=7, input_size=2 * horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500),
-        StemGNN(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout_rate=0.0, max_steps=1000, val_check_steps=500),
-        MLPMultivariate(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), max_steps=1000, val_check_steps=500),
-        TimeMixer(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500)
+        SOFTS(h=horizon, n_series=7, input_size=2 * horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500, windows_batch_size=64, inference_windows_batch_size=64),
+        TSMixer(h=horizon, n_series=7, input_size=2 * horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500, windows_batch_size=64, inference_windows_batch_size=64),
+        TSMixerx(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500, windows_batch_size=64, inference_windows_batch_size=64),
+        iTransformer(h=horizon, n_series=7, input_size=2 * horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500, windows_batch_size=64, inference_windows_batch_size=64),
+        # StemGNN(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout_rate=0.0, max_steps=1000, val_check_steps=500, windows_batch_size=64, inference_windows_batch_size=64),
+        MLPMultivariate(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), max_steps=1000, val_check_steps=500, windows_batch_size=64, inference_windows_batch_size=64),
+        TimeMixer(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500, windows_batch_size=64, inference_windows_batch_size=64)
     ]
 
     # Models

nbs/common.base_auto.ipynb

@@ -238,7 +238,11 @@
     "        self.callbacks = callbacks\n",
     "\n",
     "        # Base Class attributes\n",
-    "        self.SAMPLING_TYPE = cls_model.SAMPLING_TYPE\n",
+    "        self.EXOGENOUS_FUTR = cls_model.EXOGENOUS_FUTR\n",
+    "        self.EXOGENOUS_HIST = cls_model.EXOGENOUS_HIST\n",
+    "        self.EXOGENOUS_STAT = cls_model.EXOGENOUS_STAT\n",
+    "        self.MULTIVARIATE = cls_model.MULTIVARIATE    \n",
+    "        self.RECURRENT = cls_model.RECURRENT    \n",
     "\n",
     "    def __repr__(self):\n",
     "        return type(self).__name__ if self.alias is None else self.alias\n",

nbs/common.model_checks.ipynb

@@ -0,0 +1,234 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| default_exp common._model_checks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
+     ]
+    }
+   ],
+   "source": [
+    "#| hide\n",
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 1. Checks for models"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This file provides a set of unit tests for all models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "import pandas as pd\n",
+    "import neuralforecast.losses.pytorch as losses\n",
+    "\n",
+    "from neuralforecast import NeuralForecast\n",
+    "from neuralforecast.utils import AirPassengersPanel, AirPassengersStatic, generate_series"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "seed = 0\n",
+    "test_size = 14\n",
+    "FREQ = \"D\"\n",
+    "\n",
+    "# 1 series, no exogenous\n",
+    "N_SERIES_1 = 1\n",
+    "df = generate_series(n_series=N_SERIES_1, seed=seed, freq=FREQ, equal_ends=True)\n",
+    "max_ds = df.ds.max() - pd.Timedelta(test_size, FREQ)\n",
+    "Y_TRAIN_DF_1 = df[df.ds < max_ds]\n",
+    "Y_TEST_DF_1 = df[df.ds >= max_ds]\n",
+    "\n",
+    "# 5 series, no exogenous\n",
+    "N_SERIES_2 = 5\n",
+    "df = generate_series(n_series=N_SERIES_2, seed=seed, freq=FREQ, equal_ends=True)\n",
+    "max_ds = df.ds.max() - pd.Timedelta(test_size, FREQ)\n",
+    "Y_TRAIN_DF_2 = df[df.ds < max_ds]\n",
+    "Y_TEST_DF_2 = df[df.ds >= max_ds]\n",
+    "\n",
+    "# 1 series, with static and temporal exogenous\n",
+    "N_SERIES_3 = 1\n",
+    "df, STATIC_3 = generate_series(n_series=N_SERIES_3, n_static_features=2, \n",
+    "                     n_temporal_features=2, seed=seed, freq=FREQ, equal_ends=True)\n",
+    "max_ds = df.ds.max() - pd.Timedelta(test_size, FREQ)\n",
+    "Y_TRAIN_DF_3 = df[df.ds < max_ds]\n",
+    "Y_TEST_DF_3 = df[df.ds >= max_ds]\n",
+    "\n",
+    "# 5 series, with static and temporal exogenous\n",
+    "N_SERIES_4 = 5\n",
+    "df, STATIC_4 = generate_series(n_series=N_SERIES_4, n_static_features=2, \n",
+    "                     n_temporal_features=2, seed=seed, freq=FREQ, equal_ends=True)\n",
+    "max_ds = df.ds.max() - pd.Timedelta(test_size, FREQ)\n",
+    "Y_TRAIN_DF_4 = df[df.ds < max_ds]\n",
+    "Y_TEST_DF_4 = df[df.ds >= max_ds]\n",
+    "\n",
+    "# Generic test for a given config for a model\n",
+    "def _run_model_tests(model_class, config):\n",
+    "    if model_class.RECURRENT:\n",
+    "        config[\"inference_input_size\"] = config[\"input_size\"]\n",
+    "\n",
+    "    # DF_1\n",
+    "    if model_class.MULTIVARIATE:\n",
+    "        config[\"n_series\"] = N_SERIES_1\n",
+    "    if isinstance(config[\"loss\"], losses.relMSE):\n",
+    "        config[\"loss\"].y_train = Y_TRAIN_DF_1[\"y\"].values   \n",
+    "    if isinstance(config[\"valid_loss\"], losses.relMSE):\n",
+    "        config[\"valid_loss\"].y_train = Y_TRAIN_DF_1[\"y\"].values   \n",
+    "\n",
+    "    model = model_class(**config)\n",
+    "    fcst = NeuralForecast(models=[model], freq=FREQ)\n",
+    "    fcst.fit(df=Y_TRAIN_DF_1, val_size=24)\n",
+    "    _ = fcst.predict(futr_df=Y_TEST_DF_1)\n",
+    "    # DF_2\n",
+    "    if model_class.MULTIVARIATE:\n",
+    "        config[\"n_series\"] = N_SERIES_2\n",
+    "    if isinstance(config[\"loss\"], losses.relMSE):\n",
+    "        config[\"loss\"].y_train = Y_TRAIN_DF_2[\"y\"].values   \n",
+    "    if isinstance(config[\"valid_loss\"], losses.relMSE):\n",
+    "        config[\"valid_loss\"].y_train = Y_TRAIN_DF_2[\"y\"].values\n",
+    "    model = model_class(**config)\n",
+    "    fcst = NeuralForecast(models=[model], freq=FREQ)\n",
+    "    fcst.fit(df=Y_TRAIN_DF_2, val_size=24)\n",
+    "    _ = fcst.predict(futr_df=Y_TEST_DF_2)\n",
+    "\n",
+    "    if model.EXOGENOUS_STAT and model.EXOGENOUS_FUTR:\n",
+    "        # DF_3\n",
+    "        if model_class.MULTIVARIATE:\n",
+    "            config[\"n_series\"] = N_SERIES_3\n",
+    "        if isinstance(config[\"loss\"], losses.relMSE):\n",
+    "            config[\"loss\"].y_train = Y_TRAIN_DF_3[\"y\"].values   \n",
+    "        if isinstance(config[\"valid_loss\"], losses.relMSE):\n",
+    "            config[\"valid_loss\"].y_train = Y_TRAIN_DF_3[\"y\"].values\n",
+    "        model = model_class(**config)\n",
+    "        fcst = NeuralForecast(models=[model], freq=FREQ)\n",
+    "        fcst.fit(df=Y_TRAIN_DF_3, static_df=STATIC_3, val_size=24)\n",
+    "        _ = fcst.predict(futr_df=Y_TEST_DF_3)\n",
+    "\n",
+    "        # DF_4\n",
+    "        if model_class.MULTIVARIATE:\n",
+    "            config[\"n_series\"] = N_SERIES_4\n",
+    "        if isinstance(config[\"loss\"], losses.relMSE):\n",
+    "            config[\"loss\"].y_train = Y_TRAIN_DF_4[\"y\"].values   \n",
+    "        if isinstance(config[\"valid_loss\"], losses.relMSE):\n",
+    "            config[\"valid_loss\"].y_train = Y_TRAIN_DF_4[\"y\"].values \n",
+    "        model = model_class(**config)\n",
+    "        fcst = NeuralForecast(models=[model], freq=FREQ)\n",
+    "        fcst.fit(df=Y_TRAIN_DF_4, static_df=STATIC_4, val_size=24)\n",
+    "        _ = fcst.predict(futr_df=Y_TEST_DF_4) \n",
+    "\n",
+    "# Tests a model against every loss function\n",
+    "def check_loss_functions(model_class):\n",
+    "    loss_list = [losses.MAE(), losses.MSE(), losses.RMSE(), losses.MAPE(), losses.SMAPE(), losses.MASE(seasonality=7), \n",
+    "              losses.QuantileLoss(q=0.5), losses.MQLoss(), losses.IQLoss(), losses.DistributionLoss(\"Normal\"), \n",
+    "              losses.DistributionLoss(\"StudentT\"), losses.DistributionLoss(\"Poisson\"), losses.DistributionLoss(\"NegativeBinomial\"), \n",
+    "              losses.DistributionLoss(\"Tweedie\", rho=1.5), losses.DistributionLoss(\"ISQF\"), losses.PMM(), losses.PMM(weighted=True), \n",
+    "              losses.GMM(), losses.GMM(weighted=True), losses.NBMM(), losses.NBMM(weighted=True), losses.HuberLoss(), \n",
+    "            losses.TukeyLoss(), losses.HuberQLoss(q=0.5), losses.HuberMQLoss()]\n",
+    "    for loss in loss_list:\n",
+    "        test_name = f\"{model_class.__name__}: checking {loss._get_name()}\"\n",
+    "        print(f\"{test_name}\")\n",
+    "        config = {'max_steps': 2,\n",
+    "            'h': 7,\n",
+    "            'input_size': 28,\n",
+    "            'loss': loss,\n",
+    "            'valid_loss': None,\n",
+    "            'enable_progress_bar': False,\n",
+    "            'enable_model_summary': False,\n",
+    "            'val_check_steps': 2}        \n",
+    "        try:\n",
+    "            _run_model_tests(model_class, config) \n",
+    "        except RuntimeError:\n",
+    "            raise Exception(f\"{test_name} failed.\")\n",
+    "        except Exception:\n",
+    "            print(f\"{test_name} skipped on raised Exception.\")\n",
+    "            pass\n",
+    "\n",
+    "# Tests a model against the AirPassengers dataset\n",
+    "def check_airpassengers(model_class):\n",
+    "    print(f\"{model_class.__name__}: checking forecast AirPassengers dataset\")\n",
+    "    Y_train_df = AirPassengersPanel[AirPassengersPanel.ds<AirPassengersPanel['ds'].values[-12]] # 132 train\n",
+    "    Y_test_df = AirPassengersPanel[AirPassengersPanel.ds>=AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 12 test\n",
+    "\n",
+    "    config = {'max_steps': 2,\n",
+    "        'h': 12,\n",
+    "        'input_size': 24,\n",
+    "        'enable_progress_bar': False,\n",
+    "        'enable_model_summary': False,\n",
+    "        'val_check_steps': 2,\n",
+    "        }\n",
+    "\n",
+    "    if model_class.MULTIVARIATE:\n",
+    "        config[\"n_series\"] = Y_train_df[\"unique_id\"].nunique()\n",
+    "    # Normal forecast\n",
+    "    fcst = NeuralForecast(models=[model_class(**config)], freq='M')\n",
+    "    fcst.fit(df=Y_train_df, static_df=AirPassengersStatic)\n",
+    "    _ = fcst.predict(futr_df=Y_test_df)   \n",
+    "\n",
+    "    # Cross-validation\n",
+    "    fcst = NeuralForecast(models=[model_class(**config)], freq='M')\n",
+    "    _ = fcst.cross_validation(df=AirPassengersPanel, static_df=AirPassengersStatic, n_windows=2, step_size=12)\n",
+    "\n",
+    "# Add unit test functions to this function\n",
+    "def check_model(model_class, checks=[\"losses\", \"airpassengers\"]):\n",
+    "    \"\"\"\n",
+    "    Check model with various tests. Options for checks are:<br>\n",
+    "    \"losses\": test the model against all loss functions<br>\n",
+    "    \"airpassengers\": test the model against the airpassengers dataset for forecasting and cross-validation<br>\n",
+    "    \n",
+    "    \"\"\"\n",
+    "    if \"losses\" in checks:\n",
+    "        check_loss_functions(model_class)   \n",
+    "    if \"airpassengers\" in checks:\n",
+    "        try:\n",
+    "            check_airpassengers(model_class)   \n",
+    "        except RuntimeError:\n",
+    "            raise Exception(f\"{model_class.__name__}: AirPassengers forecast test failed.\")\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "python3",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

nbs/common.modules.ipynb

@@ -691,6 +691,66 @@
     "            x = x + self.mean\n",
     "        return x"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "class RevINMultivariate(nn.Module):\n",
+    "    \"\"\" \n",
+    "    ReversibleInstanceNorm1d for Multivariate models\n",
+    "    \"\"\"  \n",
+    "    def __init__(self, num_features: int, eps=1e-5, affine=False, subtract_last=False, non_norm=False):\n",
+    "        super().__init__()\n",
+    "        self.num_features = num_features\n",
+    "        self.eps = eps\n",
+    "        self.affine = affine\n",
+    "        if self.affine:\n",
+    "            self._init_params()\n",
+    "\n",
+    "    def forward(self, x, mode: str):\n",
+    "        if mode == 'norm':\n",
+    "            x = self._normalize(x)\n",
+    "        elif mode == 'denorm':\n",
+    "            x = self._denormalize(x)\n",
+    "        else:\n",
+    "            raise NotImplementedError\n",
+    "        return x\n",
+    "\n",
+    "    def _init_params(self):\n",
+    "        # initialize RevIN params: (C,)\n",
+    "        self.affine_weight = nn.Parameter(torch.ones((1, 1, self.num_features)))\n",
+    "        self.affine_bias = nn.Parameter(torch.zeros((1, 1, self.num_features)))\n",
+    "\n",
+    "    def _normalize(self, x):\n",
+    "        # Batch statistics\n",
+    "        self.batch_mean = torch.mean(x, axis=1, keepdim=True).detach()\n",
+    "        self.batch_std = torch.sqrt(torch.var(x, axis=1, keepdim=True, unbiased=False) + self.eps).detach()\n",
+    "    \n",
+    "        # Instance normalization\n",
+    "        x = x - self.batch_mean\n",
+    "        x = x / self.batch_std\n",
+    "        \n",
+    "        if self.affine:\n",
+    "            x = x * self.affine_weight\n",
+    "            x = x + self.affine_bias\n",
+    "\n",
+    "        return x\n",
+    "\n",
+    "    def _denormalize(self, x):\n",
+    "        # Reverse the normalization\n",
+    "        if self.affine:\n",
+    "            x = x - self.affine_bias\n",
+    "            x = x / self.affine_weight       \n",
+    "        \n",
+    "        x = x * self.batch_std\n",
+    "        x = x + self.batch_mean       \n",
+    "\n",
+    "        return x"
+   ]
   }
  ],
  "metadata": {

nbs/common.scalers.ipynb

@@ -682,11 +682,11 @@
     "    def _init_params(self, num_features):\n",
     "        # Initialize RevIN scaler params to broadcast:\n",
     "        if self.dim==1: # [B,T,C]  [1,1,C]\n",
-    "            self.revin_bias = nn.Parameter(torch.zeros(1,1,num_features))\n",
-    "            self.revin_weight = nn.Parameter(torch.ones(1,1,num_features))\n",
+    "            self.revin_bias = nn.Parameter(torch.zeros(1, 1, num_features, 1))\n",
+    "            self.revin_weight = nn.Parameter(torch.ones(1, 1, num_features, 1))\n",
     "        elif self.dim==-1: # [B,C,T]  [1,C,1]\n",
-    "            self.revin_bias = nn.Parameter(torch.zeros(1,num_features,1))\n",
-    "            self.revin_weight = nn.Parameter(torch.ones(1,num_features,1))\n",
+    "            self.revin_bias = nn.Parameter(torch.zeros(1, num_features, 1, 1))\n",
+    "            self.revin_weight = nn.Parameter(torch.ones(1, num_features, 1, 1))\n",
     "\n",
     "    #@torch.no_grad()\n",
     "    def transform(self, x, mask):\n",
@@ -863,8 +863,8 @@
     "#| hide\n",
     "# Validate scalers\n",
     "for scaler_type in [None, 'identity', 'standard', 'robust', 'minmax', 'minmax1', 'invariant', 'revin']:\n",
-    "    x = 1.0*torch.tensor(np_x)\n",
-    "    mask = torch.tensor(np_mask)\n",
+    "    x = 1.0*torch.tensor(np_x).unsqueeze(-1)\n",
+    "    mask = torch.tensor(np_mask).unsqueeze(-1)\n",
     "    scaler = TemporalNorm(scaler_type=scaler_type, dim=1, num_features=np_x.shape[-1])\n",
     "    x_scaled = scaler.transform(x=x, mask=mask)\n",
     "    x_recovered = scaler.inverse_transform(x_scaled)\n",
@@ -987,14 +987,6 @@
     "nf = NeuralForecast(models=[model], freq='MS')\n",
     "Y_hat_df = nf.cross_validation(df=Y_df, val_size=12, n_windows=1)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b2f50bd8",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {