From dc1951b19e572634880896851cc3c4169969de95 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Sat, 4 Jan 2020 20:47:04 +0100
Subject: [PATCH] added independent normal output

---
 pts/modules/__init__.py            |  1 +
 pts/modules/distribution_output.py | 29 +++++++++++++++++++++++++++--
 test/model/test_deepvar.py         | 21 ++++++++++++++-------
 3 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/pts/modules/__init__.py b/pts/modules/__init__.py
index 3e356b2..fd35a0b 100644
--- a/pts/modules/__init__.py
+++ b/pts/modules/__init__.py
@@ -5,6 +5,7 @@ from .distribution_output import (
     StudentTOutput,
     BetaOutput,
     NegativeBinomialOutput,
+    IndependentNormalOutput,
     LowRankMultivariateNormalOutput,
 )
 from .lambda_layer import LambdaLayer
diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py
index ee594ec..b26802a 100644
--- a/pts/modules/distribution_output.py
+++ b/pts/modules/distribution_output.py
@@ -10,7 +10,10 @@ from torch.distributions import (
     Beta,
     NegativeBinomial,
     StudentT,
+    Normal,
+    Independent,
     LowRankMultivariateNormal,
+    MultivariateNormal,
     TransformedDistribution,
     AffineTransform,
 )
@@ -85,7 +88,7 @@ class DistributionOutput(Output, ABC):
 
 class BetaOutput(DistributionOutput):
     args_dim: Dict[str, int] = {"concentration1": 1, "concentration0": 1}
-    distr_cls: type = Beta
+    distr_cls: Distribution = Beta
 
     @classmethod
     def domain_map(cls, concentration1, concentration0):
@@ -100,7 +103,6 @@ class BetaOutput(DistributionOutput):
 
 class NegativeBinomialOutput(DistributionOutput):
     args_dim: Dict[str, int] = {"mu": 1, "alpha": 1}
-    distr_cls: Distribution = NegativeBinomial
 
     @classmethod
     def domain_map(cls, mu, alpha):
@@ -173,3 +175,26 @@ class LowRankMultivariateNormalOutput(DistributionOutput):
     @property
     def event_shape(self) -> Tuple:
         return (self.dim,)
+
+
+class IndependentNormalOutput(DistributionOutput):
+    def __init__(self, dim: int) -> None:
+        self.dim = dim
+        self.args_dim = {"loc": self.dim, "scale": self.dim}
+
+    def domain_map(self, loc, scale):
+        return loc, F.softplus(scale)
+
+    @property
+    def event_shape(self) -> Tuple:
+        return (self.dim,)
+
+    def distribution(
+        self, distr_args, scale: Optional[torch.Tensor] = None
+    ) -> Distribution:
+        distr = Independent(Normal(*distr_args), 1)
+
+        if scale is None:
+            return distr
+        else:
+            return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)])
diff --git a/test/model/test_deepvar.py b/test/model/test_deepvar.py
index a8a931d..dafe665 100644
--- a/test/model/test_deepvar.py
+++ b/test/model/test_deepvar.py
@@ -16,8 +16,9 @@ import pytest
 
 from pts.dataset.artificial import constant_dataset
 from pts.modules import (
-    # MultivariateGaussianOutput,
+    IndependentNormalOutput,
     LowRankMultivariateNormalOutput,
+    # MultivariateNormalOutput,
 )
 from pts.evaluation import backtest_metrics
 from pts.model.deepvar import DeepVAREstimator
@@ -45,22 +46,28 @@ metadata = dataset.metadata
 estimator = DeepVAREstimator
 
 
-@pytest.mark.timeout(10)
+#@pytest.mark.timeout(10)
 @pytest.mark.parametrize(
-    "distr_output, num_batches_per_epoch, Estimator, " "use_marginal_transformation",
+    "distr_output, num_batches_per_epoch, Estimator, use_marginal_transformation",
     [
         (
-            LowRankMultivariateNormalOutput(dim=target_dim, rank=2),
+            IndependentNormalOutput(dim=target_dim),
             10,
             estimator,
             True,
         ),
         (
-            LowRankMultivariateNormalOutput(dim=target_dim, rank=2),
+            IndependentNormalOutput(dim=target_dim),
             10,
             estimator,
             False,
         ),
+        (
+            LowRankMultivariateNormalOutput(dim=target_dim, rank=2),
+            10,
+            estimator,
+            True,
+        ),
         (
             LowRankMultivariateNormalOutput(dim=target_dim, rank=2),
             10,
@@ -78,7 +85,7 @@ estimator = DeepVAREstimator
         #     MultivariateGaussianOutput(dim=target_dim),
         #     10,
         #     estimator,
-        #     True,
+        #     False,
         # ),
     ],
 )
@@ -90,10 +97,10 @@ def test_deepvar(
         input_size=44,
         num_cells=20,
         num_layers=1,
+        dropout_rate=0.0,
         pick_incomplete=True,
         target_dim=target_dim,
         prediction_length=metadata.prediction_length,
-        # target_dim=target_dim,
         freq=metadata.freq,
         distr_output=distr_output,
         scaling=False,