Add logp and logcdf methods

pymc3/distributions/continuous.py

@@ -23,10 +23,10 @@
 
 from aesara.assert_op import Assert
 from aesara.tensor.random.basic import (
+    RandomVariable,
     beta,
     cauchy,
     exponential,
-    RandomVariable,
     gamma,
     halfcauchy,
     halfnormal,
@@ -36,7 +36,7 @@
 )
 
 try:
-    from polyagamma import random_polyagamma
+    from polyagamma import polyagamma_logcdf, polyagamma_logpdf, random_polyagamma
 except ImportError:  # pragma: no cover
 
     def random_polyagamma(*args, **kwargs):
@@ -4271,32 +4271,35 @@ class PolyaGamma(PositiveContinuous):
     (exponential tilting parameter). The pdf of this distribution is
 
     .. math::
+
        f(x \mid h, z) = cosh^h(\frac{z}{2})e^{-\frac{1}{2}xz^2}f(x \mid h, 0),
 
-    where :math:`f(x \mid h, 0)` is the pdf of a :math:`PG(h, 0)` variable. It
-    is equivalent in distribution to an infinite convolution of Gamma
-    distributions,
+    where :math:`f(x \mid h, 0)` is the pdf of a :math:`PG(h, 0)` variable.
+    Notice that the pdf of this distribution is expressed as an alternating-sign
+    sum of inverse-Gaussian densities.
 
     .. math::
+
         X = \Sigma_{k=1}^{\infty}\frac{Ga(h, 1)}{d_k},
 
     where :math:`d_k = 2(k - 0.5)^2\pi^2 + z^2/2`, :math:`Ga(h, 1)` is a gamma
     random variable with shape  parameter ``h`` and scale parameter ``1``.
 
     .. plot::
+
         import matplotlib.pyplot as plt
         import numpy as np
-        from polyagamma import random_polyagamma
-        import seaborn as sns
+        from polyagamma import polyagamma_pdf
         plt.style.use('seaborn-darkgrid')
+        x = np.linspace(0.01, 5, 500);x.sort()
         hs = [1., 5., 10., 15.]
-        zs = [0., 1., 5., 10.]
-        data = {}
+        zs = [0.] * 4
         for h, z in zip(hs, zs):
-            data[f'h = {h}, z = {z}'] = random_polyagamma(h, z, size=25000)
-        sns.kdeplot(data=data)
+            pdf = polyagamma_pdf(x, h=h, z=z)
+            plt.plot(x, pdf, label=r'$h$ = {}, $z$ = {}'.format(h, z))
         plt.xlabel('x', fontsize=12)
         plt.ylabel('f(x)', fontsize=12)
+        plt.legend(loc=1)
         plt.show()
 
     ========  =============================
@@ -4342,11 +4345,46 @@ class PolyaGamma(PositiveContinuous):
 
     @classmethod
     def dist(cls, h=1.0, z=0.0, rng=None, size=None, **kwargs):
-        hh = aet.as_tensor_variable(floatX(h))
-        zz = aet.as_tensor_variable(floatX(z))
+        # really not sure whether h & z are required to be tensors? The
+        # distribution uses numpy arrays or scalars for parameter values.
+        hh = aet.as_tensor_variable(h)
+        zz = aet.as_tensor_variable(z)
 
         msg = f"The variable {hh} specified for PolyaGamma has non-positive "
         msg += "values, making it unsuitable for this parameter."
         Assert(msg)(hh, aet.all(aet.gt(hh, 0.00001)))
 
         return super().dist([h, z], size=size, rng=rng, **kwargs)
+
+    def logp(value, h, z):
+        """
+        Calculate log-probability of Normal distribution at specified value.
+
+        Parameters
+        ----------
+        value: numeric
+            Value(s) for which log-probability is calculated. If the log probabilities for multiple
+            values are desired the values must be provided in a numpy array.
+
+        Returns
+        -------
+        TensorVariable
+        """
+        return aet.as_tensor_variable(polyagamma_logpdf(value, h, z))
+
+    def logcdf(value, h, z):
+        """
+        Compute the log of the cumulative distribution function for Normal distribution
+        at the specified value.
+
+        Parameters
+        ----------
+        value: numeric or np.ndarray or `TensorVariable`
+            Value(s) for which log CDF is calculated. If the log CDF for multiple
+            values are desired the values must be provided in a numpy array.
+
+        Returns
+        -------
+        TensorVariable
+        """
+        return aet.as_tensor_variable(polyagamma_logcdf(value, h, z))