diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index 8441858470..b657164dfe 100644 --- a/pymc/distributions/continuous.py +++ b/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 diff --git a/pymc/tests/distributions/test_continuous.py b/pymc/tests/distributions/test_continuous.py index b69aea9df1..dd8d3b3922 100644 --- a/pymc/tests/distributions/test_continuous.py +++ b/pymc/tests/distributions/test_continuous.py @@ -2020,6 +2020,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}