diff --git a/gwpopulation/utils.py b/gwpopulation/utils.py index b8a909af..8d1b0b53 100644 --- a/gwpopulation/utils.py +++ b/gwpopulation/utils.py @@ -183,6 +183,51 @@ def logsubexp(log_p, log_q): return xp.nan_to_num(xp.exp(log_pdf)) * (xx >= low) * (xx <= high) +@apply_conditions(dict(aa=(gt, 0), bb=(gt, 0), scale=(gt, 0))) +def skewt(xx, aa, bb, loc=0, scale=1): + r""" + Jones and Faddy skew-t distribution (implementation based on :code:`scipy`). + + .. math:: + + z &= \frac{x - \mu}{\sigma} \\ + p(x) &= \frac{1}{\sigma C_{a,b}} + \left(1 + \frac{z}{\sqrt{a + b + z^2}}\right)^{a + \frac{1}{2}} + \left(1 - \frac{z}{\sqrt{a + b + z^2}}\right)^{b + \frac{1}{2}} \\ + C_{a,b} &= {2^{a + b - 1} B(a, b) \sqrt{a + b}} + + Parameters + ---------- + xx: float, array-like + The abscissa values (:math:`x`) + aa: float + The first shape parameter (:math:`a`) + bb: float + The second shape parameter (:math:`b`) + loc: float + The location parameter (:math:`\mu`) + scale: float + The scale parameter (:math:`\sigma`) + + Returns + ------- + prob: float, array-like + The distribution evaluated at `xx` + """ + zz = (xx - loc) / scale + denom = xp.sqrt(aa + bb + zz**2) + log_c = ( + (aa + bb - 1) * np.log(2) + + scs.betaln(aa, bb) + + xp.log(aa + bb) / 2 + + xp.log(scale) + ) + log_pdf = ( + (aa + 0.5) * xp.log1p(zz / denom) + (bb + 0.5) * xp.log1p(-zz / denom) - log_c + ) + return xp.nan_to_num(xp.exp(log_pdf)) + + def unnormalized_2d_gaussian(xx, yy, mu_x, mu_y, sigma_x, sigma_y, covariance): r""" Compute the probability distribution for a correlated 2-dimensional Gaussian diff --git a/test/utils_test.py b/test/utils_test.py index 56a7c66c..7e4e8d75 100644 --- a/test/utils_test.py +++ b/test/utils_test.py @@ -127,7 +127,23 @@ def test_truncnorm_matches_scipy(backend): assert max(abs(gwpop_vals - scipy_vals)) < 1e-3 -def test_matches_scipy(backend): +def test_skewt_matches_scipy(backend): + from scipy.stats import jf_skew_t + + gwpopulation.set_backend(backend) + xp = gwpopulation.utils.xp + xx = xp.linspace(-2, 2, 1000) + for ii in range(N_TEST): + mu = np.random.uniform(-10, 10) + sigma = np.random.uniform(0, 5) + aa = np.random.uniform(0, 100) + bb = np.random.uniform(0, 100) + gwpop_vals = utils.to_numpy(utils.skewt(xx, aa=aa, bb=bb, loc=mu, scale=sigma)) + scipy_vals = jf_skew_t(loc=mu, scale=sigma, a=aa, b=bb).pdf(utils.to_numpy(xx)) + assert max(abs(gwpop_vals - scipy_vals)) < 1e-3 + + +def test_vonmises_matches_scipy(backend): gwpopulation.set_backend(backend) xp = gwpopulation.utils.xp xx = xp.linspace(0, 2 * np.pi, 1000)