# Source code for botorch.utils.probability.bvn

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# LICENSE file in the root directory of this source tree.

r"""
Methods for computing bivariate normal probabilities and statistics.

.. [Genz2004bvnt]
A. Genz. Numerical computation of rectangular bivariate and trivariate normal and
t probabilities. Statistics and Computing, 2004.

.. [Muthen1990moments]
B. Muthen. Moments of the censored and truncated bivariate normal distribution.
British Journal of Mathematical and Statistical Psychology, 1990.
"""

from __future__ import annotations

from math import pi as _pi
from typing import Optional, Tuple

import torch
from botorch.exceptions import UnsupportedError
from botorch.utils.probability.utils import (
case_dispatcher,
get_constants_like,
leggauss,
ndtr as Phi,
phi,
STANDARDIZED_RANGE,
)
from botorch.utils.safe_math import (
div as safe_div,
exp as safe_exp,
mul as safe_mul,
sub as safe_sub,
)
from torch import Tensor

# Some useful constants
_inf = float("inf")
_2pi = 2 * _pi
_sqrt_2pi = _2pi**0.5
_inv_2pi = 1 / _2pi

[docs]def bvn(r: Tensor, xl: Tensor, yl: Tensor, xu: Tensor, yu: Tensor) -> Tensor:
r"""A function for computing bivariate normal probabilities.

Calculates P(xl < x < xu, yl < y < yu) where x and y are bivariate normal with
unit variance and correlation coefficient r. See Section 2.4 of [Genz2004bvnt]_.

This method uses a sign flip trick to improve numerical performance. Many of bvnus
internal branches rely on evaluations Phi(-bound). For a < b < 0, the term
Phi(-a) - Phi(-b) goes to zero faster than Phi(b) - Phi(a) because
finfo(dtype).epsneg is typically much larger than finfo(dtype).tiny. In these
cases, flipping the sign can prevent situations where bvnu(...) - bvnu(...) would
otherwise be zero due to round-off error.

Args:
r: Tensor of correlation coefficients.
xl: Tensor of lower bounds for x, same shape as r.
yl: Tensor of lower bounds for y, same shape as r.
xu: Tensor of upper bounds for x, same shape as r.
yu: Tensor of upper bounds for y, same shape as r.

Returns:
Tensor of probabilities P(xl < x < xu, yl < y < yu).

"""
if not (r.shape == xl.shape == xu.shape == yl.shape == yu.shape):
raise UnsupportedError("Arguments to bvn must have the same shape.")

# Sign flip trick
_0, _1, _2 = get_constants_like(values=(0, 1, 2), ref=r)
flip_x = xl.abs() > xu  # is xl more negative than xu is positive?
flip_y = yl.abs() > yu
flip = (flip_x & (~flip_y | yu.isinf())) | (flip_y & (~flip_x | xu.isinf()))
if flip.any():  # symmetric calls to bvnu below makes swapping bounds unnecessary
sign = _1 - _2 * flip.to(dtype=r.dtype)
xl = sign * xl  # becomes -xu if flipped
xu = sign * xu  # becomes -xl
yl = sign * yl  # becomes -yu
yu = sign * yu  # becomes -yl

p = bvnu(r, xl, yl) - bvnu(r, xu, yl) - bvnu(r, xl, yu) + bvnu(r, xu, yu)
return p.clip(_0, _1)

[docs]def bvnu(r: Tensor, h: Tensor, k: Tensor) -> Tensor:
r"""Solves for P(x > h, y > k) where x and y are standard bivariate normal
random variables with correlation coefficient r. In [Genz2004bvnt]_, this is (1)

L(h, k, r) = P(x < -h, y < -k) \
= 1/(a 2\pi) \int_{h}^{\infty} \int_{k}^{\infty} f(x, y, r) dy dx,

where f(x, y, r) = e^{-1/(2a^2) (x^2 - 2rxy + y^2)} and a = (1 - r^2)^{1/2}.

[Genz2004bvnt]_ report the following integation scheme incurs a maximum of 5e-16
error when run in double precision: if |r| >= 0.925, use a 20-point quadrature
rule on a 5th order Taylor expansion; else, numerically integrate in polar
coordinates using no more than 20 quadrature points.

Args:
r: Tensor of correlation coefficients.
h: Tensor of negative upper bounds for x, same shape as r.
k: Tensor of negative upper bounds for y, same shape as r.

Returns:
A tensor of probabilities P(x > h, y > k).
"""
if not (r.shape == h.shape == k.shape):
raise UnsupportedError("Arguments to bvnu must have the same shape.")
_0, _1, lower, upper = get_constants_like((0, 1) + STANDARDIZED_RANGE, r)
x_free = h < lower
y_free = k < lower
return case_dispatcher(
out=torch.empty_like(r),
cases=(  # Special cases admitting closed-form solutions
(lambda: (h > upper) | (k > upper), lambda mask: _0),
(lambda: x_free & y_free, lambda mask: _1),
(  # For |r| >= 0.925, use a Taylor approximation
lambda: r.abs() >= get_constants_like(0.925, r),
lambda m: _bvnu_taylor(r[m], h[m], k[m]),
),
),  # For |r| < 0.925, integrate in polar coordinates.
)

def _bvnu_polar(
r: Tensor, h: Tensor, k: Tensor, num_points: Optional[int] = None
) -> Tensor:
r"""Solves for P(x > h, y > k) by integrating in polar coordinates as

L(h, k, r) = \Phi(-h)\Phi(-k) + 1/(2\pi) \int_{0}^{sin^{-1}(r)} f(t) dt \
f(t) = e^{-0.5 cos(t)^{-2} (h^2 + k^2 - 2hk sin(t))}

For details, see Section 2.2 of [Genz2004bvnt]_.
"""
if num_points is None:
mar = r.abs().max()
num_points = 6 if mar < 0.3 else 12 if mar < 0.75 else 20

_0, _1, _i2, _i2pi = get_constants_like(values=(0, 1, 0.5, _inv_2pi), ref=r)

x, w = leggauss(num_points, dtype=r.dtype, device=r.device)
x = x + _1
asin_r = _i2 * torch.asin(r)
sin_asrx = (asin_r.unsqueeze(-1) * x).sin()

_h = h.unsqueeze(-1)
_k = k.unsqueeze(-1)
vals = safe_exp(
safe_sub(safe_mul(sin_asrx, _h * _k), _i2 * (_h.square() + _k.square()))
/ (_1 - sin_asrx.square())
)
probs = Phi(-h) * Phi(-k) + _i2pi * asin_r * (vals @ w)
return probs.clip(min=_0, max=_1)  # necessary due to "safe" handling of inf

def _bvnu_taylor(r: Tensor, h: Tensor, k: Tensor, num_points: int = 20) -> Tensor:
r"""Solves for P(x > h, y > k) via Taylor expansion.

Per Section 2.3 of [Genz2004bvnt]_, the bvnu equation (1) may be rewritten as

L(h, k, r) = L(h, k, s) - s/(2\pi) \int_{0}^{a} f(x) dx \
f(x) = (1 - x^2){-1/2} e^{-0.5 ((h - sk)/ x)^2} e^{-shk/(1 + (1 - x^2)^{1/2})},

where s = sign(r) and a = sqrt(1 - r^{2}). The term L(h, k, s) is analytic.
The second integral is approximated via Taylor expansion. See Sections 2.3 and
2.4 of [Genz2004bvnt]_.
"""
_0, _1, _ni2, _i2pi, _sq2pi = get_constants_like(
values=(0, 1, -0.5, _inv_2pi, _sqrt_2pi), ref=r
)

x, w = leggauss(num_points, dtype=r.dtype, device=r.device)
x = x + _1

s = get_constants_like(2, r) * (r > _0).to(r) - _1  # sign of r where sign(0) := 1
sk = s * k
skh = sk * h
comp_r2 = _1 - r.square()

a = comp_r2.clip(min=0).sqrt()
b = safe_sub(h, sk)
b2 = b.square()
c = get_constants_like(1 / 8, r) * (get_constants_like(4, r) - skh)
d = get_constants_like(1 / 80, r) * (get_constants_like(12, r) - skh)

# ---- Solve for L(h, k, s)
int_from_0_to_s = case_dispatcher(
out=torch.empty_like(r),
)

# ---- Solve for s/(2\pi) \int_{0}^{a} f(x) dx
# Analytic part
_a0 = _ni2 * (safe_div(b2, comp_r2) + skh)
_a1 = c * get_constants_like(1 / 3, r) * (_1 - d * b2)
_a2 = _1 - b2 * _a1
abs_b = b.abs()
analytic_part = torch.subtract(  # analytic part of solution
a * (_a2 + comp_r2 * _a1 + c * d * comp_r2.square()) * safe_exp(_a0),
_sq2pi * Phi(safe_div(-abs_b, a)) * abs_b * _a2 * safe_exp(_ni2 * skh),
)

_b2 = b2.unsqueeze(-1)
_skh = skh.unsqueeze(-1)
_q0 = get_constants_like(0.25, r) * comp_r2.unsqueeze(-1) * x.square()
_q1 = (_1 - _q0).sqrt()
_q2 = _ni2 * (_b2 / _q0 + _skh)

_b2 = b2.unsqueeze(-1)
_c = c.unsqueeze(-1)
_d = d.unsqueeze(-1)
vals = (_ni2 * (_b2 / _q0 + _skh)).exp() * torch.subtract(
_1 + _c * _q0 * (_1 + get_constants_like(5, r) * _d * _q0),
safe_exp(_ni2 * _q0 / (_1 + _q1).square() * _skh) / _q1,
)
mask = _q2 > get_constants_like(-100, r)
quadrature_part = _ni2 * a * (vals @ w)

# Return P(x > h, y > k)
int_from_0_to_a = _i2pi * s * (analytic_part + quadrature_part)
return (int_from_0_to_s - int_from_0_to_a).clip(min=_0, max=_1)

[docs]def bvnmom(
r: Tensor,
xl: Tensor,
yl: Tensor,
xu: Tensor,
yu: Tensor,
p: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor]:
r"""Computes the expected values of truncated, bivariate normal random variables.

Let x and y be a pair of standard bivariate normal random variables having
correlation r. This function computes E([x,y] \| [xl,yl] < [x,y] < [xu,yu]).

Following [Muthen1990moments]_ equations (4) and (5), we have

E(x \| [xl, yl] < [x, y] < [xu, yu]) \
= Z^{-1} \phi(xl) P(yl < y < yu \| x=xl) - \phi(xu) P(yl < y < yu \| x=xu),

where Z = P([xl, yl] < [x, y] < [xu, yu]) and \phi is the standard normal PDF.

Args:
r: Tensor of correlation coefficients.
xl: Tensor of lower bounds for x, same shape as r.
xu: Tensor of upper bounds for x, same shape as r.
yl: Tensor of lower bounds for y, same shape as r.
yu: Tensor of upper bounds for y, same shape as r.
p: Tensor of probabilities P(xl < x < xu, yl < y < yu), same shape as r.

Returns:
E(x \| [xl, yl] < [x, y] < [xu, yu]) and
E(y \| [xl, yl] < [x, y] < [xu, yu]).
"""
if not (r.shape == xl.shape == xu.shape == yl.shape == yu.shape):
raise UnsupportedError("Arguments to bvn must have the same shape.")

if p is None:
p = bvn(r=r, xl=xl, xu=xu, yl=yl, yu=yu)

corr = r[..., None, None]
istd = (1 - corr.square()).rsqrt()
lower = torch.stack([xl, yl], -1)
upper = torch.stack([xu, yu], -1)
bounds = torch.stack([lower, upper], -1)
deltas = safe_mul(corr, bounds)

# Compute densities and conditional probabilities
density_at_bounds = phi(bounds)
prob_given_bounds = Phi(
safe_mul(istd, safe_sub(upper.flip(-1).unsqueeze(-1), deltas))
) - Phi(safe_mul(istd, safe_sub(lower.flip(-1).unsqueeze(-1), deltas)))

# Evaluate Muthen's formula
p_diffs = -(density_at_bounds * prob_given_bounds).diff().squeeze(-1)
moments = (1 / p).unsqueeze(-1) * (p_diffs + r.unsqueeze(-1) * p_diffs.flip(-1))
return moments.unbind(-1)