Implement scipy.stats.multivariate_normal.cdf

[Mv77]

Hi,

Some of the distributions from scipy that you have implemented (laplace, logistic, norm) support their .cdf method for evaluating the cumulative distribution function. One that does not is jax.random.multivariate_normal. It would be great if you could add support for the multivariate normal CDF. The here is the doc of scipy's multivariate_normal.

Please:

• Check for duplicate requests.
• Describe your goal, and if possible provide a code snippet with a motivating example.

[jakevdp]

Thanks for the request! We would definitely welcome a contribution of this functionality.

[sharadmv]

FWIW a general MVN CDF is pretty difficult to compute (see this TFP issue for reference). A Monte Carlo approximation (as suggested in that issue) could be a good approximation if you need something right now.

Alternatively if you know your MVN has a diagonal covariance, you can use a product of independent normal CDFs instead.

[riven314]

@sharadmv @jakevdp
I am interested in this ticket but I would like to have more clarity on how I should implement this.

From scipy source code, I see that jax.random.multivariate_normal internally triggers helper functions from scipy.stats.mvndst.f (it's not written in Python) for monte carlo sampling, and the script (mvndst.f) is an implementation of this paper.

Should I re-implement the script in JAX?

[flaviovdf]

Hi,

I have implemented this code for the two-variate case. It is based on the recent paper: " A simple approximation for the bivariate normal integral" Wen-Jen Tsay & Peng-Hsuan Ke
https://www.tandfonline.com/doi/abs/10.1080/03610918.2021.1884718

https://colab.research.google.com/drive/1w2tI1-1LWzPSdG_jE0FXzdrs6VAsJwhv?usp=sharing

Posting it here in case someone also needs this CDF. It is jax.grad friendly. With this code I have replicated the authors' paper, thus I believe it's accurate enough.

[simontheflutist]

I found the relative error in Tsay and Ke's paper unacceptable in one of my applications. I have implemented a formula from this paper which is faster than Tsay and Ke, branchless, and accurate to near machine precision if the number of quadrature points is high enough.

By the Fundamental Theorem of Calculus

$$ \Phi_2(h, k; \rho) - \Phi(h) \Phi(k) = \int_0^\rho f(r, h, k) ,\mathrm{d} r $$

where $$f(\rho, h, k) = \partial_\rho \Phi_2(h, k; \rho) $$ has an explicit formula.

In my implementation,

• gauss_legendre_on_0_x is ChatGPT's implementation of Gaussian quadrature,
• dΦ_2__dρ is $f$, and
• Φ_2_increment_quad is \Phi_2(h, k; \rho) - \Phi(h) \Phi(k)

def gauss_legendre_on_0_x(n, x):
    """
    Generate Gauss-Legendre quadrature nodes and weights for the interval [0, x].

    Parameters:
    - n: int, number of quadrature points
    - x: float, right endpoint of the interval [0, x]

    Returns:
    - nodes: ndarray, transformed quadrature nodes in [0, x]
    - weights: ndarray, transformed quadrature weights
    """
    # Get nodes and weights for [-1, 1]
    nodes, weights = np.polynomial.legendre.leggauss(n)

    # Affine transformation from [-1, 1] to [0, x]
    nodes_transformed = 0.5 * (nodes + 1) * x
    weights_transformed = 0.5 * x * weights

    return nodes_transformed, weights_transformed


def dΦ_2__dρ(ρ, h, k):
    numerator = jnp.exp(-(h**2 - 2 * ρ * h * k + k**2) / (2 * (1 - ρ**2)))
    denominator = 2 * jnp.pi * jnp.sqrt(1 - ρ**2)
    return numerator / denominator


def Φ_2_increment_quad(h, k, ρ, num_points=20):
    """Returns the difference
            Φ_2(h, k, ρ) - Φ_2(h, k, 0)
       using Gaussian quadrature.
    """
    nodes, weights = gauss_legendre_on_0_x(num_points, ρ)
    integrand_values = jax.vmap(dΦ_2__dρ, in_axes=(0, None, None))(nodes, h, k)
    return weights @ integrand_values