Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add nu parametrization to beta distribution #6344

Merged
merged 10 commits into from
Feb 3, 2023
21 changes: 15 additions & 6 deletions pymc/distributions/continuous.py
Original file line number Diff line number Diff line change
Expand Up @@ -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::
Expand All @@ -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
Expand All @@ -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
ricardoV94 marked this conversation as resolved.
Show resolved Hide resolved
Alternative "sample size" of a Beta distribution (``nu`` > 0).

Notes
-----
Expand All @@ -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))

Expand All @@ -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
Expand Down
10 changes: 10 additions & 0 deletions pymc/tests/distributions/test_continuous.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down