Add more basic PyMC3 distributions and conversion dispatches

Author: brandonwillard

setup.py

@@ -32,7 +32,7 @@ def get_long_description():
     author_email=AUTHOR_EMAIL,
     url=URL,
     install_requires=[
-        "scipy>=1.2.0",
+        "scipy>=1.4.0",
         "Theano>=1.0.4",
         "tf-estimator-nightly==2.1.0.dev2020012309",
         "tf-nightly==2.2.0.dev20200201",

symbolic_pymc/relations/theano/conjugates.py

@@ -32,24 +32,24 @@ def _create_normal_normal_goals():
     #
     # Create the pattern/form of the prior normal distribution
     #
-    beta_name_lv = var("beta_name")
-    beta_size_lv = var("beta_size")
-    beta_rng_lv = var("beta_rng")
-    a_lv = var("a")
-    R_lv = var("R")
+    beta_name_lv = var()
+    beta_size_lv = var()
+    beta_rng_lv = var()
+    a_lv = var()
+    R_lv = var()
     beta_prior_mt = mt.MvNormalRV(a_lv, R_lv, size=beta_size_lv, rng=beta_rng_lv, name=beta_name_lv)
 
-    y_name_lv = var("y_name")
-    y_size_lv = var("y_size")
-    y_rng_lv = var("y_rng")
-    F_t_lv = var("f")
-    V_lv = var("V")
+    y_name_lv = var()
+    y_size_lv = var()
+    y_rng_lv = var()
+    F_t_lv = var()
+    V_lv = var()
     E_y_mt = mt.dot(F_t_lv, beta_prior_mt)
     Y_mt = mt.MvNormalRV(E_y_mt, V_lv, size=y_size_lv, rng=y_rng_lv, name=y_name_lv)
 
     # The variable specifying the fixed sample value of the random variable
     # given by `Y_mt`
-    obs_sample_mt = var("obs_sample")
+    obs_sample_mt = var()
 
     Y_obs_mt = mt.observed(obs_sample_mt, Y_mt)
 
@@ -80,21 +80,21 @@ def _create_normal_normal_goals():
 def _create_normal_wishart_goals():  # pragma: no cover
     """TODO."""
     # Create the pattern/form of the prior normal distribution
-    Sigma_name_lv = var("Sigma_name")
-    Sigma_size_lv = var("Sigma_size")
-    Sigma_rng_lv = var("Sigma_rng")
-    V_lv = var("V")
-    n_lv = var("n")
+    Sigma_name_lv = var()
+    Sigma_size_lv = var()
+    Sigma_rng_lv = var()
+    V_lv = var()
+    n_lv = var()
     Sigma_prior_mt = mt.WishartRV(V_lv, n_lv, Sigma_size_lv, Sigma_rng_lv, name=Sigma_name_lv)
 
-    y_name_lv = var("y_name")
-    y_size_lv = var("y_size")
-    y_rng_lv = var("y_rng")
-    V_lv = var("V")
-    f_mt = var("f")
+    y_name_lv = var()
+    y_size_lv = var()
+    y_rng_lv = var()
+    V_lv = var()
+    f_mt = var()
     Y_mt = mt.MvNormalRV(f_mt, V_lv, y_size_lv, y_rng_lv, name=y_name_lv)
 
-    y_mt = var("y")
+    y_mt = var()
     Y_obs_mt = mt.observed(y_mt, Y_mt)
 
     n_post_mt = etuple(mt.add, n_lv, etuple(mt.Shape, Y_obs_mt))

symbolic_pymc/theano/pymc3.py

@@ -37,6 +37,22 @@
     CauchyRVType,
     HalfCauchyRV,
     HalfCauchyRVType,
+    BetaRV,
+    BetaRVType,
+    BinomialRV,
+    BinomialRVType,
+    PoissonRV,
+    PoissonRVType,
+    DirichletRV,
+    DirichletRVType,
+    BernoulliRV,
+    BernoulliRVType,
+    BetaBinomialRV,
+    BetaBinomialRVType,
+    CategoricalRV,
+    CategoricalRVType,
+    MultinomialRV,
+    MultinomialRVType,
 )
 from .opt import FunctionGraph
 from .ops import RandomVariable
@@ -197,6 +213,110 @@ def _convert_rv_to_dist_HalfCauchy(op, rv):
     return pm.HalfCauchy, params
 
 
+@convert_dist_to_rv.register(pm.Beta, object)
+def convert_dist_to_rv_Beta(dist, rng):
+    size = dist.shape.astype(int)[BetaRV.ndim_supp :]
+    res = BetaRV(dist.alpha, dist.beta, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(BetaRVType, Apply)
+def _convert_rv_to_dist_Beta(op, rv):
+    params = {"alpha": rv.inputs[0], "beta": rv.inputs[1]}
+    return pm.Beta, params
+
+
+@convert_dist_to_rv.register(pm.Binomial, object)
+def convert_dist_to_rv_Binomial(dist, rng):
+    size = dist.shape.astype(int)[BinomialRV.ndim_supp :]
+    res = BinomialRV(dist.n, dist.p, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(BinomialRVType, Apply)
+def _convert_rv_to_dist_Binomial(op, rv):
+    params = {"n": rv.inputs[0], "p": rv.inputs[1]}
+    return pm.Binomial, params
+
+
+@convert_dist_to_rv.register(pm.Poisson, object)
+def convert_dist_to_rv_Poisson(dist, rng):
+    size = dist.shape.astype(int)[PoissonRV.ndim_supp :]
+    res = PoissonRV(dist.mu, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(PoissonRVType, Apply)
+def _convert_rv_to_dist_Poisson(op, rv):
+    params = {"mu": rv.inputs[0]}
+    return pm.Poisson, params
+
+
+@convert_dist_to_rv.register(pm.Dirichlet, object)
+def convert_dist_to_rv_Dirichlet(dist, rng):
+    size = dist.shape.astype(int)[DirichletRV.ndim_supp :]
+    res = DirichletRV(dist.a, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(DirichletRVType, Apply)
+def _convert_rv_to_dist_Dirichlet(op, rv):
+    params = {"a": rv.inputs[0]}
+    return pm.Dirichlet, params
+
+
+@convert_dist_to_rv.register(pm.Bernoulli, object)
+def convert_dist_to_rv_Bernoulli(dist, rng):
+    size = dist.shape.astype(int)[BernoulliRV.ndim_supp :]
+    res = BernoulliRV(dist.p, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(BernoulliRVType, Apply)
+def _convert_rv_to_dist_Bernoulli(op, rv):
+    params = {"p": rv.inputs[0]}
+    return pm.Bernoulli, params
+
+
+@convert_dist_to_rv.register(pm.BetaBinomial, object)
+def convert_dist_to_rv_BetaBinomial(dist, rng):
+    size = dist.shape.astype(int)[BetaBinomialRV.ndim_supp :]
+    res = BetaBinomialRV(dist.n, dist.alpha, dist.beta, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(BetaBinomialRVType, Apply)
+def _convert_rv_to_dist_BetaBinomial(op, rv):
+    params = {"n": rv.inputs[0], "alpha": rv.inputs[1], "beta": rv.inputs[2]}
+    return pm.BetaBinomial, params
+
+
+@convert_dist_to_rv.register(pm.Categorical, object)
+def convert_dist_to_rv_Categorical(dist, rng):
+    size = dist.shape.astype(int)[CategoricalRV.ndim_supp :]
+    res = CategoricalRV(dist.p, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(CategoricalRVType, Apply)
+def _convert_rv_to_dist_Categorical(op, rv):
+    params = {"p": rv.inputs[0]}
+    return pm.Categorical, params
+
+
+@convert_dist_to_rv.register(pm.Multinomial, object)
+def convert_dist_to_rv_Multinomial(dist, rng):
+    size = dist.shape.astype(int)[MultinomialRV.ndim_supp :]
+    res = MultinomialRV(dist.n, dist.p, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(MultinomialRVType, Apply)
+def _convert_rv_to_dist_Multinomial(op, rv):
+    params = {"n": rv.inputs[0], "p": rv.inputs[1]}
+    return pm.Multinomial, params
+
+
 # TODO: More RV conversions!
 
 
@@ -207,9 +327,17 @@ def pymc3_var_to_rv(pm_var, rand_state=None):
     new_rv.name = pm_var.name
 
     if isinstance(pm_var, pm.model.ObservedRV):
-        obs = tt.as_tensor_variable(pm_var.observations)
+        obs = pm_var.observations
+        # For some reason, the observations can be float when the RV's dtype is
+        # not.
+        if obs.dtype != pm_var.dtype:
+            obs = obs.astype(pm_var.dtype)
+
+        obs = tt.as_tensor_variable(obs)
+
         if getattr(obs, "cached", False):
             obs = obs.clone()
+
         new_rv = observed(obs, new_rv)
 
     # Let's attempt to fix the PyMC3 broadcastable dims "oracle" issue,

symbolic_pymc/theano/random_variables.py

@@ -21,6 +21,19 @@ def make_node(self, lower, upper, size=None, rng=None, name=None):
 UniformRV = UniformRVType()
 
 
+class BetaRVType(RandomVariable):
+    print_name = ("Beta", "\\operatorname{Beta}")
+
+    def __init__(self):
+        super().__init__("beta", theano.config.floatX, 0, [0, 0], "beta", inplace=True)
+
+    def make_node(self, alpha, beta, size=None, rng=None, name=None):
+        return super().make_node(alpha, beta, size=size, rng=rng, name=name)
+
+
+BetaRV = BetaRVType()
+
+
 class NormalRVType(RandomVariable):
     print_name = ("N", "\\operatorname{N}")
 
@@ -209,6 +222,61 @@ def make_node(self, b, loc=0.0, scale=1.0, size=None, rng=None, name=None):
 TruncExponentialRV = TruncExponentialRVType()
 
 
+class BernoulliRVType(RandomVariable):
+    print_name = ("Bern", "\\operatorname{Bern}")
+
+    def __init__(self):
+        super().__init__(
+            "bernoulli",
+            "int64",
+            0,
+            [0],
+            lambda rng, *args: scipy.stats.bernoulli.rvs(args[0], size=args[1], random_state=rng),
+            inplace=True,
+        )
+
+    def make_node(self, p, size=None, rng=None, name=None):
+        return super().make_node(p, size=size, rng=rng, name=name)
+
+
+BernoulliRV = BernoulliRVType()
+
+
+class BinomialRVType(RandomVariable):
+    print_name = ("Binom", "\\operatorname{Binom}")
+
+    def __init__(self):
+        super().__init__("binomial", "int64", 0, [0, 0], "binomial", inplace=True)
+
+    def make_node(self, n, p, size=None, rng=None, name=None):
+        return super().make_node(n, p, size=size, rng=rng, name=name)
+
+
+BinomialRV = BinomialRVType()
+
+
+class BetaBinomialRVType(RandomVariable):
+    print_name = ("BetaBinom", "\\operatorname{BetaBinom}")
+
+    def __init__(self):
+        super().__init__(
+            "beta_binomial",
+            "int64",
+            0,
+            [0, 0, 0],
+            lambda rng, *args: scipy.stats.betabinom.rvs(
+                *args[:-1], size=args[-1], random_state=rng
+            ),
+            inplace=True,
+        )
+
+    def make_node(self, n, a, b, size=None, rng=None, name=None):
+        return super().make_node(n, a, b, size=size, rng=rng, name=name)
+
+
+BetaBinomialRV = BetaBinomialRVType()
+
+
 # Support shape is determined by the first dimension in the *second* parameter
 # (i.e.  the probabilities vector)
 class MultinomialRVType(RandomVariable):
@@ -232,6 +300,28 @@ def make_node(self, n, pvals, size=None, rng=None, name=None):
 MultinomialRV = MultinomialRVType()
 
 
+class CategoricalRVType(RandomVariable):
+    print_name = ("Cat", "\\operatorname{Cat}")
+
+    def __init__(self):
+        super().__init__(
+            "categorical",
+            "int64",
+            0,
+            [1],
+            lambda rng, *args: scipy.stats.rv_discrete(values=(range(len(args[0])), args[0])).rvs(
+                size=args[1], random_state=rng
+            ),
+            inplace=True,
+        )
+
+    def make_node(self, pvals, size=None, rng=None, name=None):
+        return super().make_node(pvals, size=size, rng=rng, name=name)
+
+
+CategoricalRV = CategoricalRVType()
+
+
 class Observed(tt.Op):
     """An `Op` that represents an observed random variable.
 

tests/theano/test_pymc3.py

@@ -22,7 +22,7 @@
 
 
 @pytest.mark.usefixtures("run_with_theano")
-def test_pymc_convert_dists():
+def test_pymc3_convert_dists():
     """Just a basic check that all PyMC3 RVs will convert to and from Theano RVs."""
     tt.config.compute_test_value = "ignore"
     theano.config.cxx = ""
@@ -36,6 +36,15 @@ def test_pymc_convert_dists():
         gamma_rv = pm.Gamma("gamma_rv", 1.0, 1.0, observed=1.0)
         invgamma_rv = pm.InverseGamma("invgamma_rv", 1.0, 1.0, observed=1.0)
         exp_rv = pm.Exponential("exp_rv", 1.0, observed=1.0)
+        halfnormal_rv = pm.HalfNormal("halfnormal_rv", 1.0, observed=1.0)
+        beta_rv = pm.Beta("beta_rv", 2.0, 2.0, observed=1.0)
+        binomial_rv = pm.Binomial("binomial_rv", 10, 0.5, observed=5)
+        dirichlet_rv = pm.Dirichlet("dirichlet_rv", np.r_[0.1, 0.1], observed=np.r_[0.1, 0.1])
+        poisson_rv = pm.Poisson("poisson_rv", 10, observed=5)
+        bernoulli_rv = pm.Bernoulli("bernoulli_rv", 0.5, observed=0)
+        betabinomial_rv = pm.BetaBinomial("betabinomial_rv", 0.1, 0.1, 10, observed=5)
+        categorical_rv = pm.Categorical("categorical_rv", np.r_[0.5, 0.5], observed=1)
+        multinomial_rv = pm.Multinomial("multinomial_rv", 5, np.r_[0.5, 0.5], observed=np.r_[2])
 
     # Convert to a Theano `FunctionGraph`
     fgraph = model_graph(model)
@@ -53,7 +62,7 @@ def test_pymc_convert_dists():
 
 
 @pytest.mark.usefixtures("run_with_theano")
-def test_pymc_normal_model():
+def test_pymc3_normal_model():
     """Conduct a more in-depth test of PyMC3/Theano conversions for a specific model."""
     tt.config.compute_test_value = "ignore"
 
@@ -199,7 +208,7 @@ def _check_model(model):
 
 
 @pytest.mark.usefixtures("run_with_theano")
-def test_pymc_broadcastable():
+def test_pymc3_broadcastable():
     """Test PyMC3 to Theano conversion amid array broadcasting."""
     tt.config.compute_test_value = "ignore"