Parametrize Binomial and Categorical distributions via logit_p

pymc/distributions/discrete.py

@@ -114,7 +114,15 @@ class Binomial(Discrete):
     rv_op = binomial
 
     @classmethod
-    def dist(cls, n, p, *args, **kwargs):
+    def dist(cls, n, p=None, logit_p=None, *args, **kwargs):
+        if p is not None and logit_p is not None:
+            raise ValueError("Incompatible parametrization. Can't specify both p and logit_p.")
+        elif p is None and logit_p is None:
+            raise ValueError("Incompatible parametrization. Must specify either p or logit_p.")
+
+        if logit_p is not None:
+            p = at.sigmoid(logit_p)
+
         n = at.as_tensor_variable(intX(n))
         p = at.as_tensor_variable(floatX(p))
         return super().dist([n, p], **kwargs)
@@ -1252,7 +1260,15 @@ class Categorical(Discrete):
     rv_op = categorical
 
     @classmethod
-    def dist(cls, p, **kwargs):
+    def dist(cls, p=None, logit_p=None, **kwargs):
+        if p is not None and logit_p is not None:
+            raise ValueError("Incompatible parametrization. Can't specify both p and logit_p.")
+        elif p is None and logit_p is None:
+            raise ValueError("Incompatible parametrization. Must specify either p or logit_p.")
+
+        if logit_p is not None:
+            p = pm.math.softmax(logit_p, axis=-1)
+
         if isinstance(p, np.ndarray) or isinstance(p, list):
             if (np.asarray(p) < 0).any():
                 raise ValueError(f"Negative `p` parameters are not valid, got: {p}")

pymc/tests/test_distributions_random.py

@@ -1005,6 +1005,25 @@ class TestBinomial(BaseTestDistributionRandom):
     checks_to_run = ["check_pymc_params_match_rv_op"]
 
 
+class TestLogitBinomial(BaseTestDistributionRandom):
+    pymc_dist = pm.Binomial
+    pymc_dist_params = {"n": 100, "logit_p": 2.197224577}
+    expected_rv_op_params = {"n": 100, "p": 0.9}
+    tests_to_run = ["check_pymc_params_match_rv_op"]
+
+    @pytest.mark.parametrize(
+        "n, p, logit_p, expected",
+        [
+            (5, None, None, "Must specify either p or logit_p."),
+            (5, 0.5, 0.5, "Can't specify both p and logit_p."),
+        ],
+    )
+    def test_binomial_init_fail(self, n, p, logit_p, expected):
+        with pm.Model() as model:
+            with pytest.raises(ValueError, match=f"Incompatible parametrization. {expected}"):
+                pm.Binomial("x", n=n, p=p, logit_p=logit_p)
+
+
 class TestNegativeBinomial(BaseTestDistributionRandom):
     pymc_dist = pm.NegativeBinomial
     pymc_dist_params = {"n": 100, "p": 0.33}
@@ -1410,6 +1429,28 @@ class TestCategorical(BaseTestDistributionRandom):
     ]
 
 
+class TestLogitCategorical(BaseTestDistributionRandom):
+    pymc_dist = pm.Categorical
+    pymc_dist_params = {"logit_p": np.array([-0.944461608841, 0.489548225319, -2.197224577336])}
+    expected_rv_op_params = {"p": np.array([0.28, 0.62, 0.10])}
+    tests_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_rv_size",
+    ]
+
+    @pytest.mark.parametrize(
+        "p, logit_p, expected",
+        [
+            (None, None, "Must specify either p or logit_p."),
+            (0.5, 0.5, "Can't specify both p and logit_p."),
+        ],
+    )
+    def test_categorical_init_fail(self, p, logit_p, expected):
+        with pm.Model() as model:
+            with pytest.raises(ValueError, match=f"Incompatible parametrization. {expected}"):
+                pm.Categorical("x", p=p, logit_p=logit_p)
+
+
 class TestGeometric(BaseTestDistributionRandom):
     pymc_dist = pm.Geometric
     pymc_dist_params = {"p": 0.9}