Merge pull request #1 from felixdivo/feature/multivariate-wrapper

Improve testing code for MultivariateForecastingModelWrapper

darts/models/forecasting/multivariate_forecasting_model_wrapper.py

@@ -74,7 +74,7 @@ def _predict(
     ) -> TimeSeries:
         predictions = [
             model.predict(n=n, future_covariates=future_covariates)
-            if isinstance(model, FutureCovariatesLocalForecastingModel)
+            if model.supports_future_covariates
             else model.predict(n=n)
             for model in self._trained_models
         ]

darts/tests/models/forecasting/test_multivariate_forecasting_model_wrapper.py

@@ -1,5 +1,7 @@
 import copy
 
+import pytest
+
 from darts import TimeSeries
 from darts.logging import get_logger
 from darts.models import (
@@ -55,53 +57,71 @@ class TestMultivariateForecastingModelWrapper:
     n_pred = 5
 
     univariate = tg.gaussian_timeseries(length=ts_length, mean=50)
-    multivariate = univariate.stack(univariate)
+    multivariate = univariate.stack(tg.gaussian_timeseries(length=ts_length, mean=20))
 
     future_covariates = tg.gaussian_timeseries(length=ts_length + n_pred, mean=50)
 
-    def test_fit_predict_local_models(self):
-        for model in local_models:
-            self._test_predict_with_base_model(model, None)
+    @pytest.mark.parametrize("model", local_models)
+    def test_fit_predict_local_models(self, model):
+        self._test_predict_with_base_model(model)
 
-    def test_fit_predict_local_future_covariates_models(self):
-        for model in future_covariates_models:
-            self._test_predict_with_base_model(model, self.future_covariates)
+    @pytest.mark.parametrize("model", future_covariates_models)
+    def test_fit_predict_local_future_covariates_models(self, model):
+        self._test_predict_with_base_model(model, self.future_covariates)
 
-    def test_encoders_support(self):
+    @pytest.mark.parametrize("model_object", future_covariates_models)
+    def test_encoders_support(self, model_object):
         add_encoders = {
             "position": {"future": ["relative"]},
         }
 
-        # test some models that support encoders
-        for model_object in future_covariates_models:
-            # test once with user supplied covariates, and once without
-            for fc in [self.future_covariates, None]:
-                model_params = {
-                    k: vals
-                    for k, vals in copy.deepcopy(model_object.model_params).items()
-                }
-                model_params["add_encoders"] = add_encoders
-                model = model_object.__class__(**model_params)
-
-                self._test_predict_with_base_model(model, fc)
-
-    def _test_predict_with_base_model(self, model, future_covariates):
-        series_univariate = self.univariate
-        series_multivariate = self.multivariate
-
-        combinations = [
-            series_univariate,
-            series_multivariate,
-        ]
-
-        for combination in combinations:
+        # test once with user supplied covariates, and once without
+        for fc in [self.future_covariates, None]:
+            model_params = {
+                k: vals for k, vals in copy.deepcopy(model_object.model_params).items()
+            }
+            model_params["add_encoders"] = add_encoders
+            model = model_object.__class__(**model_params)
+
+            self._test_predict_with_base_model(model, fc)
+
+    def _test_predict_with_base_model(self, model, future_covariates=None):
+        for combination in [self.univariate, self.multivariate]:
             preds = self.trained_model_predictions(
                 model, self.n_pred, combination, future_covariates
             )
             assert isinstance(preds, TimeSeries)
             assert preds.n_components == combination.n_components
 
+            # Make sure that the compound prediction is the same as the individual predictions
+            individual_preds = self.trained_individual_model_predictions(
+                model, self.n_pred, combination, future_covariates
+            )
+            for component in range(combination.n_components):
+                assert (
+                    preds.univariate_component(component) == individual_preds[component]
+                )
+
     def trained_model_predictions(self, base_model, n, series, future_covariates):
         model = MultivariateForecastingModelWrapper(base_model)
         model.fit(series, future_covariates=future_covariates)
         return model.predict(n=n, series=series, future_covariates=future_covariates)
+
+    def trained_individual_model_predictions(
+        self, base_model, n, series, future_covariates
+    ):
+        predictions = []
+        for component in range(series.n_components):
+            single_series = series.univariate_component(component)
+
+            model = base_model.untrained_model()
+            if model.supports_future_covariates:
+                model.fit(single_series, future_covariates=future_covariates)
+                predictions.append(
+                    model.predict(n=n, future_covariates=future_covariates)
+                )
+            else:
+                model.fit(single_series)
+                predictions.append(model.predict(n=n))
+
+        return predictions