Repository: awslabs/gluonts

Issue:
Implement Poisson distribution
The Poisson distribution is currently missing.


Discussion:
@lostella Can I have a go at this?
Sure!
@michiboo you can take inspiration from other distributions to see how they are implemented, as well as they're companion "output" layers. See e.g. [negative binomial](src/gluonts/distribution/neg_binomial.py), which also has integer support.
Hi I am unsure what to put in log_prob function, 
I found this formula on the internet but not sure if it is correct
https://mc-stan.org/docs/2_21/functions-reference/poisson-distribution-log-parameterization.html
@michiboo you should be implementing the [probability mass function](https://en.wikipedia.org/wiki/Poisson_distribution#Probability_mass_function)

Resolving patch:
diff --git a/src/gluonts/distribution/__init__.py b/src/gluonts/distribution/__init__.py
--- a/src/gluonts/distribution/__init__.py
+++ b/src/gluonts/distribution/__init__.py
@@ -49,6 +49,7 @@
 from .uniform import Uniform, UniformOutput
 
 from .gamma import Gamma, GammaOutput
+from .poisson import Poisson, PoissonOutput
 
 __all__ = [
     "Distribution",
@@ -77,6 +78,8 @@
     "BinnedOutput",
     "PiecewiseLinear",
     "PiecewiseLinearOutput",
+    "Poisson",
+    "PoissonOutput",
     "TransformedPiecewiseLinear",
     "TransformedDistribution",
     "TransformedDistributionOutput",
diff --git a/src/gluonts/distribution/poisson.py b/src/gluonts/distribution/poisson.py
new file mode 100644
--- /dev/null
+++ b/src/gluonts/distribution/poisson.py
@@ -0,0 +1,115 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+
+# Standard library imports
+from typing import Dict, Optional, Tuple, List
+
+# Third-party imports
+import numpy as np
+
+# First-party imports
+from gluonts.model.common import Tensor
+from gluonts.core.component import validated
+
+# Relative imports
+from .distribution import Distribution, _sample_multiple, getF, softplus
+from .distribution_output import DistributionOutput
+
+
+class Poisson(Distribution):
+    r"""
+    Poisson distribution, i.e. the distribution of the number of
+    successes in a specified region.
+
+    Parameters
+    ----------
+    rate
+        Tensor containing the means, of shape `(*batch_shape, *event_shape)`.
+    F
+    """
+
+    is_reparameterizable = False
+
+    @validated()
+    def __init__(self, rate: Tensor, F=None) -> None:
+        self.rate = rate
+        self.F = F if F else getF(rate)
+
+    @property
+    def batch_shape(self) -> Tuple:
+        return self.rate.shape
+
+    @property
+    def event_shape(self) -> Tuple:
+        return ()
+
+    @property
+    def event_dim(self) -> int:
+        return 0
+
+    def log_prob(self, x: Tensor) -> Tensor:
+        F = self.F
+        ll = x * F.log(self.rate) - F.gammaln(x + 1.0) - self.rate
+        return ll
+
+    @property
+    def mean(self) -> Tensor:
+        return self.rate
+
+    @property
+    def stddev(self) -> Tensor:
+        return self.F.sqrt(self.rate)
+
+    def sample(
+        self, num_samples: Optional[int] = None, dtype=np.float32
+    ) -> Tensor:
+        def s(rate: Tensor) -> Tensor:
+            return self.F.random.poisson(lam=rate, dtype=dtype)
+
+        return _sample_multiple(s, rate=self.rate, num_samples=num_samples)
+
+    @property
+    def args(self) -> List:
+        return [self.rate]
+
+
+class PoissonOutput(DistributionOutput):
+    args_dim: Dict[str, int] = {"rate": 1}
+    distr_cls: type = Poisson
+
+    @classmethod
+    def domain_map(cls, F, rate):
+        rate = softplus(F, rate) + 1e-8
+        return rate.squeeze(axis=-1)
+
+    # Overwrites the parent class method.
+    # We cannot scale using the affine transformation since Poisson should return integers.
+    # Instead we scale the parameters.
+    def distribution(
+        self,
+        distr_args,
+        loc: Optional[Tensor] = None,
+        scale: Optional[Tensor] = None,
+    ) -> Poisson:
+        rate = distr_args
+        assert loc is None
+        if scale is None:
+            return Poisson(rate)
+        else:
+            F = getF(rate)
+            rate = F.broadcast_mul(rate, scale)
+            return Poisson(rate, F)
+
+    @property
+    def event_shape(self) -> Tuple:
+        return ()


Test patch:
diff --git a/test/distribution/test_distribution_inference.py b/test/distribution/test_distribution_inference.py
--- a/test/distribution/test_distribution_inference.py
+++ b/test/distribution/test_distribution_inference.py
@@ -52,6 +52,8 @@
     LaplaceOutput,
     Gaussian,
     GaussianOutput,
+    Poisson,
+    PoissonOutput,
     PiecewiseLinear,
     PiecewiseLinearOutput,
     Binned,
@@ -805,3 +807,33 @@ def test_binned_likelihood(
     assert all(
         mx.nd.abs(mx.nd.array(bin_prob_hat) - bin_prob) < TOL * bin_prob
     ), f"bin_prob did not match: bin_prob = {bin_prob}, bin_prob_hat = {bin_prob_hat}"
+
+
+@pytest.mark.parametrize("rate", [1.0])
+@pytest.mark.parametrize("hybridize", [True, False])
+def test_poisson_likelihood(rate: float, hybridize: bool) -> None:
+    """
+    Test to check that maximizing the likelihood recovers the parameters
+    """
+
+    # generate samples
+    rates = mx.nd.zeros(NUM_SAMPLES) + rate
+
+    distr = Poisson(rates)
+    samples = distr.sample()
+
+    init_biases = [inv_softplus(rate - START_TOL_MULTIPLE * TOL * rate)]
+
+    rate_hat = maximum_likelihood_estimate_sgd(
+        PoissonOutput(),
+        samples,
+        init_biases=init_biases,
+        hybridize=hybridize,
+        learning_rate=PositiveFloat(0.05),
+        num_epochs=PositiveInt(20),
+    )
+
+    print("rate:", rate_hat)
+    assert (
+        np.abs(rate_hat[0] - rate) < TOL * rate
+    ), f"mu did not match: rate = {rate}, rate_hat = {rate_hat}"
diff --git a/test/distribution/test_distribution_methods.py b/test/distribution/test_distribution_methods.py
--- a/test/distribution/test_distribution_methods.py
+++ b/test/distribution/test_distribution_methods.py
@@ -27,6 +27,7 @@
     Beta,
     MultivariateGaussian,
     PiecewiseLinear,
+    Poisson,
     Binned,
     TransformedDistribution,
 )
@@ -83,6 +84,7 @@
             ).repeat(axis=0, repeats=2),
         },
     ),
+    (Poisson, {"rate": mx.nd.array([1000.0, 1.0])}),
 ]
 
 test_output = {
@@ -126,6 +128,11 @@
         "stddev": mx.nd.array([1.377416, 1.377416]),
         "variance": mx.nd.array([1.8972749, 1.8972749]),
     },
+    "Poisson": {
+        "mean": mx.nd.array([1000.0, 1.0]),
+        "stddev": mx.nd.array([31.622776, 1.0]),
+        "variance": mx.nd.array([1000.0, 1.0]),
+    },
 }
 
 # TODO: implement stddev methods for MultivariateGaussian and LowrankMultivariateGaussian
@@ -137,6 +144,7 @@
     NegativeBinomial,
     Uniform,
     Binned,
+    Poisson,
 ]
 
 
diff --git a/test/distribution/test_distribution_output_shapes.py b/test/distribution/test_distribution_output_shapes.py
--- a/test/distribution/test_distribution_output_shapes.py
+++ b/test/distribution/test_distribution_output_shapes.py
@@ -29,6 +29,7 @@
     LowrankMultivariateGaussianOutput,
     NegativeBinomialOutput,
     PiecewiseLinearOutput,
+    PoissonOutput,
     StudentTOutput,
     UniformOutput,
     DirichletOutput,
@@ -156,6 +157,14 @@
             (3, 4),
             (5,),
         ),
+        (
+            PoissonOutput(),
+            mx.nd.random.normal(shape=(3, 4, 5, 6)),
+            [None],
+            [None, mx.nd.ones(shape=(3, 4, 5))],
+            (3, 4, 5),
+            (),
+        ),
     ],
 )
 def test_distribution_output_shapes(
diff --git a/test/distribution/test_distribution_sampling.py b/test/distribution/test_distribution_sampling.py
--- a/test/distribution/test_distribution_sampling.py
+++ b/test/distribution/test_distribution_sampling.py
@@ -27,6 +27,7 @@
     Gamma,
     Beta,
     MultivariateGaussian,
+    Poisson,
     PiecewiseLinear,
     Binned,
     TransformedDistribution,
@@ -88,6 +89,7 @@
             ).repeat(axis=0, repeats=2),
         },
     ),
+    (Poisson, {"rate": mx.nd.array([1000.0, 0])}),
 ]
 
 
diff --git a/test/distribution/test_distribution_shapes.py b/test/distribution/test_distribution_shapes.py
--- a/test/distribution/test_distribution_shapes.py
+++ b/test/distribution/test_distribution_shapes.py
@@ -27,6 +27,7 @@
     MultivariateGaussian,
     NegativeBinomial,
     PiecewiseLinear,
+    Poisson,
     StudentT,
     Uniform,
     TransformedDistribution,
@@ -104,6 +105,7 @@
             (3, 4, 5),
             (),
         ),
+        (Poisson(rate=mx.nd.ones(shape=(3, 4, 5))), (3, 4, 5), ()),
         (
             Uniform(
                 low=-mx.nd.ones(shape=(3, 4, 5)),