diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index cffea7f94e..cfe4dbe52d 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -89,6 +89,9 @@ def polyagamma_cdf(*args, **kwargs): logpow, normal_lccdf, normal_lcdf, + pc_prior_studentt_kld_dist, + pc_prior_studentt_logp, + pc_prior_studentt_kld_dist_inv_op, zvalue, ) from pymc.distributions.distribution import DIST_PARAMETER_TYPES, Continuous @@ -3973,3 +3976,69 @@ def logcdf(value, h, z): h > 0, msg="h > 0", ) + + +class PCPriorStudentT_dof_RV(RandomVariable): + name = "pc_prior_studentt_dof" + ndim_supp = 0 + ndims_params = [0] + dtype = "floatX" + _print_name = ("PCTDoF", "\\operatorname{PCPriorStudentT_dof}") + + @classmethod + def rng_fn(cls, rng, lam, size=None) -> np.ndarray: + return pc_prior_studentt_kld_dist_inv_op.spline( + rng.exponential(scale=1.0 / lam, size=size) + ) + + +pc_prior_studentt_dof = PCPriorStudentT_dof_RV() + + +class PCPriorStudentT_dof(PositiveContinuous): + + rv_op = pc_prior_studentt_dof + + @classmethod + def dist( + cls, + alpha: Optional[DIST_PARAMETER_TYPES] = None, + U: Optional[DIST_PARAMETER_TYPES] = None, + lam: Optional[DIST_PARAMETER_TYPES] = None, + *args, + **kwargs + ): + lam = cls.get_lam(alpha, U, lam) + return super().dist([lam], *args, **kwargs) + + def moment(rv, size, lam): + mean = pc_prior_studentt_kld_dist_inv_op(1.0 / lam) + if not rv_size_is_none(size): + mean = pt.full(size, mean) + return mean + + + @classmethod + def get_lam(cls, alpha=None, U=None, lam=None): + if (alpha is not None) and (U is not None): + return -np.log(alpha) / pc_prior_studentt_kld_dist(U) + elif (lam is not None): + return lam + else: + raise ValueError( + "Incompatible parameterization. Either use alpha and U, or lam to specify the " + "distribution." + ) + + def logp(value, lam): + res = pc_prior_studentt_logp(value, lam) + res = pt.switch( + pt.lt(value, 2 + 1e-6), # 2 + 1e-6 smallest value for nu + -np.inf, + res, + ) + return check_parameters( + res, + lam > 0, + msg="lam > 0" + ) diff --git a/pymc/distributions/dist_math.py b/pymc/distributions/dist_math.py index 7a72f27ebd..eea237cae2 100644 --- a/pymc/distributions/dist_math.py +++ b/pymc/distributions/dist_math.py @@ -508,3 +508,60 @@ def incomplete_beta(a, b, value): stacklevel=2, ) return pt.betainc(a, b, value) + + +def tri_gamma_approx(x): + """ Approximation to the trigamma function (derivative of digamma). This function is in + PyTensor, but doesn't currently have a derivative implemented. Until that's in this is + a series expansion from Wikipedia: https://en.wikipedia.org/wiki/Trigamma_function + """ + return ( + 1 / x + + (1 / (2 * x**2)) + + (1 / (6 * x**3)) + - (1 / (30 * x**5)) + + (1 / (42 * x**7)) + - (1 / (30 * x**9)) + + (5 / (66 * x**11)) + - (691 / (2730 * x**13)) + + (7 / (6 * x**15)) + ) + + +def pc_prior_studentt_kld_dist(nu): + """ Calculates the distance in terms of sqrt(2KLD) between the base model, Gaussian, and the + more complex model, Student t. + """ + return pt.sqrt( + 1 + pt.log(2 * pt.reciprocal(nu - 2)) + + 2 * pt.gammaln((nu + 1) / 2) + - 2 * pt.gammaln(nu / 2) + - (nu + 1) * (pt.digamma((nu + 1) / 2) - pt.digamma(nu / 2)) + ) + + +def pc_prior_studentt_logp(nu, lam): + """ Log-probability for the PC prior of the Student-t degrees of freedom parameter. + """ + return ( + pt.log(lam) + + pt.log((1 / (nu - 2)) + ((nu + 1) / 2) + * (tri_gamma_approx((nu + 1) / 2) - tri_gamma_approx(nu / 2))) + - pt.log(4 * pc_prior_studentt_kld_dist(nu)) + - lam * pc_prior_studentt_kld_dist(nu) + + pt.log(2) + ) + + +def _make_pct_inv_func(): + """ Numerically calculate the inverse of the PC Student-t distance function. + """ + nu = np.concatenate(( + np.linspace(2.0 + 1e-6, 2.4, 2000), + np.linspace(2.4 + 1e-4, 4000, 10000) + )) + return UnivariateSpline( + pc_prior_studentt_kld_dist(nu).eval()[::-1], nu[::-1], ext=3, k=3, s=0, + ) + +pc_prior_studentt_kld_dist_inv_op = SplineWrapper(_make_pct_inv_func())