Add a Theano Polya-Gamma random variable type

Author: brandonwillard

.travis.yml

@@ -10,6 +10,7 @@ python:
   # - "pypy3"
 
 install:
+  - pip install Cython
   - pip install -r requirements.txt
 
 script:

requirements.txt

@@ -2,6 +2,7 @@ setuptools>=45.2.0
 six>=1.14.0
 -e ./
 sympy>=1.3
+pypolyagamma @ git+https://github.com/slinderman/pypolyagamma.git@b5883e661123862ca07d29ab14369fae85bdbc27#egg=pypolyagamma-1.2.2
 coveralls
 pydocstyle>=3.0.0
 pytest>=5.0.0

symbolic_pymc/theano/random_variables.py

@@ -5,6 +5,8 @@
 
 from functools import partial
 
+from pypolyagamma import PyPolyaGamma
+
 from .ops import RandomVariable, param_supp_shape_fn
 
 
@@ -204,7 +206,7 @@ def __init__(self):
             "invgamma",
             theano.config.floatX,
             0,
-            [0, 0, 0],
+            [0, 0],
             lambda rng, shape, rate, size: scipy.stats.invgamma.rvs(
                 shape, scale=rate, size=size, random_state=rng
             ),
@@ -343,6 +345,46 @@ def make_node(self, pvals, size=None, rng=None, name=None):
 CategoricalRV = CategoricalRVType()
 
 
+class PolyaGammaRVType(RandomVariable):
+    """Polya-Gamma random variable.
+
+    XXX: This doesn't really use the given RNG, due to the narrowness of the
+    sampler package's implementation.
+    """
+
+    print_name = ("PG", "\\operatorname{PG}")
+
+    def __init__(self):
+        super().__init__(
+            "polya-gamma", theano.config.floatX, 0, [0, 0], self._smpl_fn, inplace=True,
+        )
+
+    def make_node(self, b, c, size=None, rng=None, name=None):
+        return super().make_node(b, c, size=size, rng=rng, name=name)
+
+    @classmethod
+    def _smpl_fn(cls, rng, b, c, size):
+        pg = PyPolyaGamma(rng.randint(2 ** 16))
+
+        if not size and b.shape == c.shape == ():
+            return pg.pgdraw(b, c)
+        else:
+            b, c = np.broadcast_arrays(b, c)
+            out_shape = b.shape + tuple(size or ())
+            smpl_val = np.empty(out_shape, dtype="double")
+            b = np.tile(b, tuple(size or ()) + (1,))
+            c = np.tile(c, tuple(size or ()) + (1,))
+            pg.pgdrawv(
+                np.asarray(b.flat).astype("double", copy=True),
+                np.asarray(c.flat).astype("double", copy=True),
+                np.asarray(smpl_val.flat),
+            )
+            return smpl_val
+
+
+PolyaGammaRV = PolyaGammaRVType()
+
+
 class Observed(tt.Op):
     """An `Op` that represents an observed random variable.
 

tests/theano/test_rv.py

@@ -2,7 +2,9 @@
 
 import theano.tensor as tt
 
-from symbolic_pymc.theano.random_variables import NormalRV, MvNormalRV
+from pytest import importorskip
+
+from symbolic_pymc.theano.random_variables import NormalRV, MvNormalRV, PolyaGammaRV
 
 
 def rv_numpy_tester(rv, *params, size=None):
@@ -68,3 +70,31 @@ def test_mvnormalrv():
     # Looks like NumPy doesn't support that (and it's probably better off for
     # it).
     # rv_numpy_tester(MvNormalRV, [[0, 1, 2], [4, 5, 6]], np.diag([1, 1, 1]))
+
+
+def test_polyagammarv():
+
+    _ = importorskip("pypolyagamma")
+
+    # Sampled values should be scalars
+    pg_rv = PolyaGammaRV(1.1, -10.5)
+    assert pg_rv.eval().shape == ()
+
+    pg_rv = PolyaGammaRV(1.1, -10.5, size=[1])
+    assert pg_rv.eval().shape == (1,)
+
+    pg_rv = PolyaGammaRV(1.1, -10.5, size=[2, 3])
+    bcast_smpl = pg_rv.eval()
+    assert bcast_smpl.shape == (2, 3)
+    # Make sure they're not all equal
+    assert np.all(np.abs(np.diff(bcast_smpl.flat)) > 0.0)
+
+    pg_rv = PolyaGammaRV(np.r_[1.1, 3], -10.5)
+    bcast_smpl = pg_rv.eval()
+    assert bcast_smpl.shape == (2,)
+    assert np.all(np.abs(np.diff(bcast_smpl.flat)) > 0.0)
+
+    pg_rv = PolyaGammaRV(np.r_[1.1, 3], -10.5, size=(2, 3))
+    bcast_smpl = pg_rv.eval()
+    assert bcast_smpl.shape == (2, 2, 3)
+    assert np.all(np.abs(np.diff(bcast_smpl.flat)) > 0.0)