Python port of Wetzels & Wagenmakers (2012) R code for calculating Bayes factors for (partial) correlation values

correlation_bf.py

"""
Calculate the Bayes factor for (partial) correlation values using the JZS prior.

Ported from R code in Wetzels & Wagenmakers (2012) "A default Bayesian hypothesis test for correlations and partial
correlations". Psychon Bull Rev. 19:1057–1064 doi:10.3758/s13423-012-0295-x
"""

from math import sqrt, exp, gamma
from numpy import inf
from scipy import integrate


def _jzs_integrand(g, r, n, p):
    return (
        (1 + g) ** ((n - p - 1) / 2)
        * (1 + (1 - r ** 2) * g) ** (-(n - 1) / 2)
        * g ** (-3 / 2)
        * exp(-n / (2 * g)))


def jzs_cor_bf(r, n):
    """
    r: correlation value
    n: sample size
    
    From MacLean et al. (2010):   jzs_cor_bf(r=-0.3616346, n=54) = 3.85
    From Kanai et al. (in press): jzs_cor_bf(r= 0.4844199, n=40) = 17.87
    """
    bf10 = sqrt(n / 2) / gamma(1 / 2) * integrate.quad(func=_jzs_integrand, a=0, b=inf, args=(r, n, 1))[0]
    return bf10


def jzs_parcor_bf(r0, r1, n, p0, p1):
    """
    r0: the correlation from the null model (all covariates except covariate of interest)
    r1: correlation from the full model (all covariates)
    n: sample size
    p0: number of covariates (null model)
    p1: number of covariates (partial model)
    
    From Lleras et al. (in press): jzs_parcor_bf(r0=sqrt(0.6084), r1=sqrt(0.6084408), n=40, p0=1, p1=2) = 0.13
    """
    # BF21 = BF20/BF10
    bf21 = (
        integrate.quad(func=_jzs_integrand, a=0, b=inf, args=(r1, n, p1))[0] /
        integrate.quad(func=_jzs_integrand, a=0, b=inf, args=(r0, n, p0))[0]
    )
    return bf21