diff --git a/conjugate/distributions.py b/conjugate/distributions.py index 81d357b..7708f17 100644 --- a/conjugate/distributions.py +++ b/conjugate/distributions.py @@ -42,6 +42,7 @@ import numpy as np from scipy import stats, __version__ as scipy_version +from scipy.special import gammaln from conjugate._typing import NUMERIC from conjugate.plot import ( @@ -671,3 +672,119 @@ class CompoundGamma(ContinuousPlotDistMixin, SliceMixin): alpha: NUMERIC beta: NUMERIC lam: NUMERIC + + +@dataclass +class GammaKnownRateProportional: + """Gamma known rate proportional distribution. + + Args: + a: prod of observations + b: number of observations + c: number of observations + + """ + + a: NUMERIC + b: NUMERIC + c: NUMERIC + + def approx_log_likelihood( + self, alpha: NUMERIC, beta: NUMERIC, ln=np.log, gammaln=gammaln + ): + """Approximate log likelihood. + + Args: + alpha: shape parameter + beta: known rate parameter + ln: log function + gammaln: log gamma function + + Returns: + log likelihood up to a constant + + """ + return ( + (alpha - 1) * ln(self.a) + + alpha * self.c * ln(beta) + - self.b * gammaln(alpha) + ) + + +@dataclass +class GammaProportional: + """Gamma proportional distribution. + + Args: + p: product of r observations + q: sum of s observations + r: number of observations for p + s: number of observations for q + + """ + + p: NUMERIC + q: NUMERIC + r: NUMERIC + s: NUMERIC + + def approx_log_likelihood( + self, alpha: NUMERIC, beta: NUMERIC, ln=np.log, gammaln=gammaln + ): + """Approximate log likelihood. + + Args: + alpha: shape parameter + beta: rate parameter + ln: log function + gammaln: log gamma function + + Returns: + log likelihood up to a constant + + """ + return ( + (alpha - 1) * ln(self.p) + - self.q * beta + - self.r * gammaln(alpha) + + self.s * alpha * ln(beta) + ) + + +@dataclass +class BetaProportional: + """Beta proportional distribution. + + Args: + p: product of observations + q: product of complements + k: number of observations + + """ + + p: NUMERIC + q: NUMERIC + k: NUMERIC + + def approx_log_likelihood( + self, alpha: NUMERIC, beta: NUMERIC, ln=np.log, gammaln=gammaln + ): + """Approximate log likelihood. + + Args: + alpha: shape parameter + beta: shape parameter + ln: log function + gammaln: log gamma function + + Returns: + log likelihood up to a constant + + """ + return ( + self.k * gammaln(alpha + beta) + + alpha * ln(self.p) + + beta * ln(self.q) + - self.k * gammaln(alpha) + - self.k * gammaln(beta) + ) diff --git a/conjugate/models.py b/conjugate/models.py index 06c59a0..8d241d2 100644 --- a/conjugate/models.py +++ b/conjugate/models.py @@ -9,10 +9,13 @@ from conjugate.distributions import ( Beta, + BetaProportional, CompoundGamma, Dirichlet, DirichletMultinomial, Gamma, + GammaKnownRateProportional, + GammaProportional, NegativeBinomial, BetaNegativeBinomial, BetaBinomial, @@ -1164,3 +1167,84 @@ def pareto_gamma( beta_post = gamma_prior.beta + ln_x_total - n * ln(x_m) return Gamma(alpha=alpha_post, beta=beta_post) + + +def gamma( + x_total: NUMERIC, + x_prod: NUMERIC, + n: NUMERIC, + proportional_prior: GammaProportional, +) -> GammaProportional: + """Posterior distribution for a gamma likelihood. + + Inference on alpha and beta + + Args: + x_total: sum of all outcomes + x_prod: product of all outcomes + n: total number of samples in x_total and x_prod + proportional_prior: GammaProportional prior + + Returns: + GammaProportional posterior distribution + + """ + p_post = proportional_prior.p * x_prod + q_post = proportional_prior.q + x_total + r_post = proportional_prior.r + n + s_post = proportional_prior.s + n + + return GammaProportional(p=p_post, q=q_post, r=r_post, s=s_post) + + +def gamma_known_rate( + x_prod: NUMERIC, + n: NUMERIC, + beta: NUMERIC, + proportional_prior: GammaKnownRateProportional, +) -> GammaKnownRateProportional: + """Posterior distribution for a gamma likelihood. + + The rate beta is assumed to be known. + + Args: + x_prod: product of all outcomes + n: total number of samples in x_prod + beta: known rate parameter + + Returns: + GammaKnownRateProportional posterior distribution + + """ + a_post = proportional_prior.a * x_prod + b_post = proportional_prior.b + n + c_post = proportional_prior.c + n + + return GammaKnownRateProportional(a=a_post, b=b_post, c=c_post) + + +def beta( + x_prod: NUMERIC, + one_minus_x_prod: NUMERIC, + n: NUMERIC, + proportional_prior: BetaProportional, +) -> BetaProportional: + """Posterior distribution for a Beta likelihood. + + Inference on alpha and beta + + Args: + x_prod: product of all outcomes + one_minus_x_prod: product of all (1 - outcomes) + n: total number of samples in x_prod and one_minus_x_prod + proportional_prior: BetaProportional prior + + Returns: + BetaProportional posterior distribution + + """ + p_post = proportional_prior.p * x_prod + q_post = proportional_prior.q * one_minus_x_prod + k_post = proportional_prior.k + n + + return BetaProportional(p=p_post, q=q_post, k=k_post)