Implement TSMixer Model (#2293)

unit8co · Apr 8, 2024 · 0d5c722 · 0d5c722
1 parent 0cdb4a5
commit 0d5c722
Show file tree

Hide file tree

Showing 16 changed files with 2,307 additions and 5 deletions.
diff --git a/.github/workflows/merge.yml b/.github/workflows/merge.yml
@@ -87,7 +87,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        example-name: [00-quickstart.ipynb, 01-multi-time-series-and-covariates.ipynb, 02-data-processing.ipynb, 03-FFT-examples.ipynb, 04-RNN-examples.ipynb, 05-TCN-examples.ipynb, 06-Transformer-examples.ipynb, 07-NBEATS-examples.ipynb, 08-DeepAR-examples.ipynb, 09-DeepTCN-examples.ipynb, 10-Kalman-filter-examples.ipynb, 11-GP-filter-examples.ipynb, 12-Dynamic-Time-Warping-example.ipynb, 13-TFT-examples.ipynb, 15-static-covariates.ipynb, 16-hierarchical-reconciliation.ipynb, 18-TiDE-examples.ipynb, 19-EnsembleModel-examples.ipynb, 20-RegressionModel-examples.ipynb]
+        example-name: [00-quickstart.ipynb, 01-multi-time-series-and-covariates.ipynb, 02-data-processing.ipynb, 03-FFT-examples.ipynb, 04-RNN-examples.ipynb, 05-TCN-examples.ipynb, 06-Transformer-examples.ipynb, 07-NBEATS-examples.ipynb, 08-DeepAR-examples.ipynb, 09-DeepTCN-examples.ipynb, 10-Kalman-filter-examples.ipynb, 11-GP-filter-examples.ipynb, 12-Dynamic-Time-Warping-example.ipynb, 13-TFT-examples.ipynb, 15-static-covariates.ipynb, 16-hierarchical-reconciliation.ipynb, 18-TiDE-examples.ipynb, 19-EnsembleModel-examples.ipynb, 20-RegressionModel-examples.ipynb, 21-TSMixer-examples.ipynb]
     steps:
       - name: "1. Clone repository"
         uses: actions/checkout@v2

diff --git a/.gitignore b/.gitignore
@@ -16,6 +16,7 @@ runs/
 htmlcov
 coverage.xml
 .darts
+darts_logs/
 docs_env
 .DS_Store
 .gradle

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,14 +9,15 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 ### For users of the library:
 **Improved**
-- 🚀🚀🚀 Improvements to metrics, historical forecasts, backtest, and residuals through major refactor. The refactor includes optimization of multiple process and improvemenets to consistency, reliability, and the documentation. Some of these necessary changes come at the cost of breaking changes. [#2284](https://github.com/unit8co/darts/pull/2284) by [Dennis Bader](https://github.com/dennisbader).
+- 🚀🚀 New forecasting model: `TSMixerModel`  as proposed in [this paper](https://arxiv.org/abs/2303.06053). An MLP based model that combines temporal, static and cross-sectional feature information using stacked mixing layers. [#1807](https://https://github.com/unit8co/darts/pull/001), by [Dennis Bader](https://github.com/dennisbader) and [Cristof Rojas](https://github.com/cristof-r).
+- 🚀🚀 Improvements to metrics, historical forecasts, backtest, and residuals through major refactor. The refactor includes optimization of multiple process and improvemenets to consistency, reliability, and the documentation. Some of these necessary changes come at the cost of breaking changes. [#2284](https://github.com/unit8co/darts/pull/2284) by [Dennis Bader](https://github.com/dennisbader).
   - Metrics:
-    - Optimized all metrics, which now run >20 times faster than before for univariate series, and >>20 times for multivariate series. This boosts direct metric computations as well as backtesting and residuals computation!
+    - Optimized all metrics, which now run **> n * 20 times faster** than before for series with `n` components/columns. This boosts direct metric computations as well as backtesting and residuals computation!
     - Added new metrics:
       - Time aggregated metric `merr()` (Mean Error)
       - Time aggregated scaled metrics  `rmsse()`, and `msse()`: The (Root) Mean Squared Scaled Error.
       - "Per time step" metrics that return a metric score per time step: `err()` (Error), `ae()` (Absolute Error), `se()` (Squared Error), `sle()` (Squared Log Error), `ase()` (Absolute Scaled Error), `sse` (Squared Scaled Error), `ape()` (Absolute Percentage Error), `sape()` (symmetric Absolute Percentage Error), `arre()` (Absolute Ranged Relative Error), `ql` (Quantile Loss)
-    - All scaled metrics now accept `insample` series that can be overlapping into `pred_series` (before that had to end exactly one step before `pred_series`).  Darts will handle the correct time extraction for you.  
+    - All scaled metrics now accept `insample` series that can be overlapping into `pred_series` (before they had to end exactly one step before `pred_series`).  Darts will handle the correct time extraction for you.  
     - Improvements to the documentation:
       - Added a summary list of all metrics to the [metrics documentation page](https://unit8co.github.io/darts/generated_api/darts.metrics.html)
       - Standardized the documentation of each metric (added formula, improved return documentation, ...) 

diff --git a/README.md b/README.md
@@ -255,6 +255,7 @@ on bringing more models and features.
 | [DLinearModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.dlinear.html#darts.models.forecasting.dlinear.DLinearModel)                                                                                                                                         | [DLinear paper](https://arxiv.org/pdf/2205.13504.pdf)                                                                                                                                                                             | 🟩 🟩                                                        | 🟩 🟩 🟩                                                                 | 🟩 🟩                                                                    | 🟩                                        |
 | [NLinearModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.nlinear.html#darts.models.forecasting.nlinear.NLinearModel)                                                                                                                                         | [NLinear paper](https://arxiv.org/pdf/2205.13504.pdf)                                                                                                                                                                             | 🟩 🟩                                                        | 🟩 🟩 🟩                                                                 | 🟩 🟩                                                                    | 🟩                                        |
 | [TiDEModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tide_model.html#darts.models.forecasting.tide_model.TiDEModel)                                                                                                                                         | [TiDE paper](https://arxiv.org/pdf/2304.08424.pdf)                                                                                                                                                                                | 🟩 🟩                                                        | 🟩 🟩 🟩                                                                 | 🟩 🟩                                                                    | 🟩                                        |
+| [TSMixerModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.tsmixer_model.html#darts.models.forecasting.tsmixer_model.TSMixerModel)                                                                                                                             | [TSMixer paper](https://arxiv.org/pdf/2303.06053.pdf), [PyTorch Implementation](https://github.com/ditschuk/pytorch-tsmixer)                                                                                                      | 🟩 🟩                                                        | 🟩 🟩 🟩                                                                 | 🟩 🟩                                                                    | 🟩                                        |
 | **Ensemble Models**<br/>([GlobalForecastingModel](https://unit8co.github.io/darts/userguide/covariates.html#global-forecasting-models-gfms)): Model support is dependent on ensembled forecasting models and the ensemble model itself                                                    |                                                                                                                                                                                                                                   |                                                              |                                                                          |                                                                          |                                           |
 | [NaiveEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.baselines.html#darts.models.forecasting.baselines.NaiveEnsembleModel)                                                                                                                         |                                                                                                                                                                                                                                   | 🟩 🟩                                                        | 🟩 🟩 🟩                                                                 | 🟩 🟩                                                                    | 🟩                                        |
 | [RegressionEnsembleModel](https://unit8co.github.io/darts/generated_api/darts.models.forecasting.regression_ensemble_model.html#darts.models.forecasting.regression_ensemble_model.RegressionEnsembleModel)                                                                               |                                                                                                                                                                                                                                   | 🟩 🟩                                                        | 🟩 🟩 🟩                                                                 | 🟩 🟩                                                                    | 🟩                                        |

diff --git a/darts/models/__init__.py b/darts/models/__init__.py
@@ -51,6 +51,7 @@
     from darts.models.forecasting.tft_model import TFTModel
     from darts.models.forecasting.tide_model import TiDEModel
     from darts.models.forecasting.transformer_model import TransformerModel
+    from darts.models.forecasting.tsmixer_model import TSMixerModel
 except ModuleNotFoundError:
     logger.warning(
         "Support for Torch based models not available. "

diff --git a/darts/models/forecasting/__init__.py b/darts/models/forecasting/__init__.py
@@ -46,6 +46,7 @@
     - :class:`~darts.models.forecasting.dlinear.DLinearModel`
     - :class:`~darts.models.forecasting.nlinear.NLinearModel`
     - :class:`~darts.models.forecasting.tide_model.TiDEModel`
+    - :class:`~darts.models.forecasting.tsmixer_model.TSMixerModel`
 Ensemble Models (`GlobalForecastingModel <https://unit8co.github.io/darts/userguide/covariates.html#global-forecasting-models-gfms>`_)
     - :class:`~darts.models.forecasting.baselines.NaiveEnsembleModel`
     - :class:`~darts.models.forecasting.regression_ensemble_model.RegressionEnsembleModel`