Add nu parametrization to beta distribution (#6344)

* Add (mu, nu) parametrization

* add test beta logp

* Add relationship between alpha, beta and nu

* change test structure

* remove white space accidently added in docstring

* remove sample size from docstring

---------

Co-authored-by: Michael Osthege <[email protected]>

Author: alekracicot

pymc/distributions/continuous.py

@@ -1062,7 +1062,7 @@ class Beta(UnitContinuous):
     ========  ==============================================================
 
     Beta distribution can be parameterized either in terms of alpha and
-    beta or mean and standard deviation. The link between the two
+    beta, mean and standard deviation or mean and sample size. The link between the three
     parametrizations is given by
 
     .. math::
@@ -1072,6 +1072,9 @@ class Beta(UnitContinuous):
 
        \text{where } \kappa = \frac{\mu(1-\mu)}{\sigma^2} - 1
 
+       \alpha = \mu * \nu
+       \beta = (1 - \mu) * \nu
+
     Parameters
     ----------
     alpha : tensor_like of float, optional
@@ -1081,7 +1084,9 @@ class Beta(UnitContinuous):
     mu : tensor_like of float, optional
         Alternative mean (0 < ``mu`` < 1).
     sigma : tensor_like of float, optional
-        Alternative standard deviation (1 < ``sigma`` < sqrt(``mu`` * (1 - ``mu``))).
+        Alternative standard deviation (0 < ``sigma`` < sqrt(``mu`` * (1 - ``mu``))).
+    nu : tensor_like of float, optional
+        Alternative "sample size" of a Beta distribution (``nu`` > 0).
 
     Notes
     -----
@@ -1092,8 +1097,8 @@ class Beta(UnitContinuous):
     rv_op = pytensor.tensor.random.beta
 
     @classmethod
-    def dist(cls, alpha=None, beta=None, mu=None, sigma=None, *args, **kwargs):
-        alpha, beta = cls.get_alpha_beta(alpha, beta, mu, sigma)
+    def dist(cls, alpha=None, beta=None, mu=None, sigma=None, nu=None, *args, **kwargs):
+        alpha, beta = cls.get_alpha_beta(alpha, beta, mu, sigma, nu)
         alpha = at.as_tensor_variable(floatX(alpha))
         beta = at.as_tensor_variable(floatX(beta))
 
@@ -1106,17 +1111,21 @@ def moment(rv, size, alpha, beta):
         return mean
 
     @classmethod
-    def get_alpha_beta(self, alpha=None, beta=None, mu=None, sigma=None):
+    def get_alpha_beta(self, alpha=None, beta=None, mu=None, sigma=None, nu=None):
         if (alpha is not None) and (beta is not None):
             pass
         elif (mu is not None) and (sigma is not None):
             kappa = mu * (1 - mu) / sigma**2 - 1
             alpha = mu * kappa
             beta = (1 - mu) * kappa
+        elif (mu is not None) and (nu is not None):
+            alpha = mu * nu
+            beta = (1 - mu) * nu
         else:
             raise ValueError(
                 "Incompatible parameterization. Either use alpha "
-                "and beta, or mu and sigma to specify distribution."
+                "and beta, mu and sigma or mu and nu to specify "
+                "distribution."
             )
 
         return alpha, beta

pymc/tests/distributions/test_continuous.py

@@ -2026,6 +2026,16 @@ class TestBetaMuSigma(BaseTestDistributionRandom):
     checks_to_run = ["check_pymc_params_match_rv_op"]
 
 
+class TestBetaMuNu(BaseTestDistributionRandom):
+    pymc_dist = pm.Beta
+    pymc_dist_params = {"mu": 0.5, "nu": 3}
+    expected_alpha, expected_beta = pm.Beta.get_alpha_beta(
+        mu=pymc_dist_params["mu"], nu=pymc_dist_params["nu"]
+    )
+    expected_rv_op_params = {"alpha": expected_alpha, "beta": expected_beta}
+    checks_to_run = ["check_pymc_params_match_rv_op"]
+
+
 class TestExponential(BaseTestDistributionRandom):
     pymc_dist = pm.Exponential
     pymc_dist_params = {"lam": 10.0}