Add PC prior for Student-T degrees of freedom

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"
+        )

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())