Make Gamma distribution multivariate

cuqi/distribution/_gamma.py

@@ -7,17 +7,21 @@ class Gamma(Distribution):
     """
     Represents a multivariate Gamma distribution characterized by shape and rate parameters of independent random variables x_i. Each is distributed according to the PDF function
 
-    f(x; shape, rate) = rate^shape * x^(shape-1) * exp(-rate * x) / Gamma(shape)
+    f(x_i; shape, rate) = rate^shape * x_i^(shape-1) * exp(-rate * x_i) / Gamma(shape)
 
     where `shape` and `rate` are the parameters of the distribution, and Gamma is the Gamma function.
 
+    In case shape and/or rate are arrays, the pdf looks like
+
+    f(x_i; shape_i, rate_i) = rate_i^shape_i * x_i^(shape_i-1) * exp(-rate_i * x_i) / Gamma(shape_i)
+
     Parameters
     ----------
     shape : float or array_like, optional
         The shape parameter of the Gamma distribution. Must be positive.
 
     rate : float or array_like, optional
-        The rate parameter of the Gamma distribution. Must be positive. The inverse of the scale parameter.
+        The rate parameter of the Gamma distribution. Must be positive.
 
     Examples
     --------
@@ -27,16 +31,35 @@ class Gamma(Distribution):
         import cuqi
         import matplotlib.pyplot as plt
 
-        # Create a Gamma distribution instance
+        # Create a multivariate Gamma distribution with the same shape and rate parameters
         shape = 1
         rate = 1e-4
+        gamma_dist = cuqi.distribution.Gamma(shape=shape, rate=rate, geometry=10)
+
+        # Generate samples
+        samples = gamma_dist.sample(10000)
+
+        # Plot histogram of samples for index 0
+        samples.hist_chain(0, bins=70)
+
+
+    .. code-block:: python
+
+        import numpy as np
+        import cuqi
+        import matplotlib.pyplot as plt
+
+        # Create a multivariate Gamma distribution with different shape and rate parameters
+        shape = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        rate = [1e-4, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2, 1e3, 1e4, 1e5]
         gamma_dist = cuqi.distribution.Gamma(shape=shape, rate=rate)
 
         # Generate samples
         samples = gamma_dist.sample(10000)
 
-        # Plot histogram of samples
+        # Plot histogram of samples for index 0
         samples.hist_chain(0, bins=70)
+
     """
     def __init__(self, shape=None, rate=None, is_symmetric=False, **kwargs):
         # Init from abstract distribution class
@@ -47,6 +70,7 @@ def __init__(self, shape=None, rate=None, is_symmetric=False, **kwargs):
 
     @property
     def shape(self):
+        """ Shape parameter of the Gamma distribution. Must be positive. """
         return self._shape
 
     @shape.setter
@@ -55,6 +79,7 @@ def shape(self, value):
 
     @property
     def rate(self):
+        """ Rate parameter of the Gamma distribution. Must be positive. """
         return self._rate
 
     @rate.setter
@@ -63,6 +88,7 @@ def rate(self, value):
 
     @property
     def scale(self):
+        """ Scale parameter of the Gamma distribution. Must be positive. This is the inverse of the rate parameter. """
         return 1/self.rate
 
     def logpdf(self, x):

tests/test_distribution.py

@@ -375,6 +375,32 @@ def my_pdf( x, a, location, scale):
 
     else:
         raise ValueError
+
+@pytest.mark.parametrize("shape", [1, 2, 3, 5])
+@pytest.mark.parametrize("rate", [1e-4, 1e-3, 1e-2, 1e-1, 1, 1e4, 1e5])
+@pytest.mark.parametrize("value", [1, 2, 3, 4, 5])
+def test_Gamma_pdf(shape, rate, value):
+    G = cuqi.distribution.Gamma(shape, rate)
+    assert np.isclose(G.pdf(value), scipy_stats.gamma(shape, scale=1/rate).pdf(value))
+
+@pytest.mark.parametrize("shape", [1, 2, 3, 5])
+@pytest.mark.parametrize("rate", [1e-4, 1e-3, 1e-2, 1e-1, 1, 1e4, 1e5])
+@pytest.mark.parametrize("value", [1, 2, 3, 4, 5])
+def test_Gamma_cdf(shape, rate, value):
+    G = cuqi.distribution.Gamma(shape, rate)
+    assert np.isclose(G.cdf(value), scipy_stats.gamma(shape, scale=1/rate).cdf(value))
+
+@pytest.mark.parametrize("shape", [1, 2, 3, 5])
+@pytest.mark.parametrize("rate", [1e-4, 1e-3, 1e-2, 1e-1, 1, 1e4, 1e5])
+def test_Gamma_sample(shape, rate):
+    rng = np.random.RandomState(3)
+    G = cuqi.distribution.Gamma(shape, rate)
+    cuqi_samples = G.sample(3, rng=rng)
+
+    rng2 = np.random.RandomState(3)
+    np_samples = rng2.gamma(shape=shape, scale=1/rate, size=(3, 1)).T
+
+    assert np.allclose(cuqi_samples.samples, np_samples)
 
 @pytest.mark.parametrize("location", [-1, -2, -3, 0, 1, 2, 3])
 @pytest.mark.parametrize("scale", [1e-3, 1e-1, 1e0, 1e1, 1e3])