[ENH] Deepar v2 final fixes

docs/source/tutorials/stallion.ipynb

@@ -41,7 +41,6 @@
    "source": [
     "import warnings\n",
     "\n",
-    "\n",
     "warnings.filterwarnings(\"ignore\")  # avoid printing out absolute paths"
    ]
   },

pytorch_forecasting/models/deepar/__deepar_pkg_v2.py

@@ -0,0 +1,117 @@
+"""
+Packages container for DeepAR model.
+"""
+
+from pytorch_forecasting.base._base_pkg import Base_pkg
+
+
+class DeepAR_pkg_v2(Base_pkg):
+    """DeepAR package container."""
+
+    _tags = {
+        "info:name": "DeepAR",
+        "info:compute": 3,
+        "authors": ["jdb78"],
+        "capability:exogenous": True,
+        "capability:multivariate": True,
+        "capability:pred_int": True,
+        "capability:flexible_history_length": True,
+        "capability:cold_start": False,
+    }
+
+    @classmethod
+    def get_cls(cls):
+        """Get model class."""
+        from pytorch_forecasting.models.deepar._deepar_v2 import DeepAR
+
+        return DeepAR
+
+    @classmethod
+    def get_datamodule_cls(cls):
+        """Get the underlying DataModule class."""
+        from pytorch_forecasting.data.data_module import (
+            EncoderDecoderTimeSeriesDataModule,
+        )
+
+        return EncoderDecoderTimeSeriesDataModule
+
+    @classmethod
+    def get_base_test_params(cls):
+        """Return testing parameter settings for the trainer."""
+        return [
+            {},
+            dict(
+                cell_type="GRU",
+                hidden_size=16,
+                rnn_layers=2,
+            ),
+        ]
+
+    @classmethod
+    def get_test_train_params(cls):
+        """Return testing parameter settings for the trainer.
+
+        Returns
+        -------
+        params : dict or list of dict, default = {}
+            Parameters to create testing instances of the class
+            Each dict are parameters to construct an "interesting" test instance, i.e.,
+            `MyClass(**params)` or `MyClass(**params[i])` creates a valid test instance.
+            `create_test_instance` uses the first (or only) dictionary in `params`
+        """
+        from pytorch_forecasting.metrics import NormalDistributionLoss
+
+        params = [
+            dict(
+                loss=NormalDistributionLoss(),
+            ),
+            dict(
+                loss=NormalDistributionLoss(),
+                cell_type="GRU",
+                hidden_size=16,
+                rnn_layers=2,
+            ),
+            dict(
+                loss=NormalDistributionLoss(),
+                hidden_size=32,
+                rnn_layers=3,
+                dropout=0.2,
+            ),
+            dict(
+                loss=NormalDistributionLoss(),
+                hidden_size=20,
+                datamodule_cfg=dict(
+                    max_encoder_length=7,
+                    max_prediction_length=5,
+                ),
+            ),
+            dict(
+                loss=NormalDistributionLoss(),
+                hidden_size=16,
+                n_validation_samples=50,
+                n_plotting_samples=25,
+            ),
+            dict(
+                loss=NormalDistributionLoss(),
+                hidden_size=10,
+                rnn_layers=1,
+                dropout=0.0,
+                datamodule_cfg=dict(
+                    max_encoder_length=3,
+                    max_prediction_length=2,
+                ),
+            ),
+        ]
+
+        default_dm_cfg = {
+            "max_encoder_length": 4,
+            "max_prediction_length": 3,
+        }
+
+        for param in params:
+            current_dm_cfg = param.get("datamodule_cfg", {})
+            merged_dm_cfg = default_dm_cfg.copy()
+            merged_dm_cfg.update(current_dm_cfg)
+            param["datamodule_cfg"] = merged_dm_cfg
+
+        return params

pytorch_forecasting/models/deepar/__init__.py

@@ -1,6 +1,8 @@
 """DeepAR: Probabilistic forecasting with autoregressive recurrent networks."""
 
+from pytorch_forecasting.models.deepar.__deepar_pkg_v2 import DeepAR_pkg_v2
 from pytorch_forecasting.models.deepar._deepar import DeepAR
 from pytorch_forecasting.models.deepar._deepar_pkg import DeepAR_pkg
+from pytorch_forecasting.models.deepar._deepar_v2 import DeepAR as DeepAR_v2
 
-__all__ = ["DeepAR", "DeepAR_pkg"]
+__all__ = ["DeepAR", "DeepAR_v2", "DeepAR_pkg", "DeepAR_pkg_v2"]