bivariate_quantiles module

Bivariate vector quantiles and ranks.

This module implements the bivariate case of the vector quantiles and vector ranks construction of Chernozhukov, Galichon, Hallin, and Henry (2017).

The main workflow is:

1. Solve for the dual weights v with :func:_solve_for_v.
2. Evaluate quantiles with :func:bivariate_quantiles_v or :func:bivariate_quantiles.
3. Read off barycentric ranks with :func:bivariate_ranks.

References

Chernozhukov, Galichon, Hallin, and Henry. "Monge-Kantorovich Depth, Quantiles, Ranks and Signs." Annals of Statistics 45(1), 2017.

bivariate_quantiles(y, tau, n_nodes=32, verbose=False)

Solve for the dual weights and evaluate bivariate quantiles.

Parameters:

Name	Type	Description	Default
y	ndarray	Observations with shape (n, 2).	required
tau	ndarray	Query points in [0, 1]^2 with shape (m, 2).	required
n_nodes	int	Number of Chebyshev nodes for the quadrature.	32
verbose	bool	Print optimisation diagnostics when True.	False

Returns:

Type	Description
ndarray	Bivariate quantiles evaluated at tau.

Source code in bs_python_utils/stats/bivariate_quantiles.py

def bivariate_quantiles(
    y: np.ndarray,
    tau: np.ndarray,
    n_nodes: int = 32,
    verbose: bool = False,
) -> np.ndarray:
    """Solve for the dual weights and evaluate bivariate quantiles.

    Args:
        y: Observations with shape ``(n, 2)``.
        tau: Query points in ``[0, 1]^2`` with shape ``(m, 2)``.
        n_nodes: Number of Chebyshev nodes for the quadrature.
        verbose: Print optimisation diagnostics when ``True``.

    Returns:
        Bivariate quantiles evaluated at ``tau``.
    """
    v, _ = _solve_for_v(y, n_nodes, verbose)
    return bivariate_quantiles_v(y, tau, v)

bivariate_quantiles_v(y, tau, v)

Evaluate vector quantiles for fixed dual weights.

Parameters:

Name	Type	Description	Default
y	ndarray	Observations with shape (n, 2).	required
tau	ndarray	Evaluation points in [0, 1]^2 with shape (m, 2).	required
v	ndarray	Dual weights solving the optimal transport problem, with length n.	required

Returns:

Type	Description
ndarray	Array of quantile locations with shape (m, 2).

Raises:

Type	Description
SystemExit	If tau does not have exactly two columns.

Source code in bs_python_utils/stats/bivariate_quantiles.py

def bivariate_quantiles_v(y: np.ndarray, tau: np.ndarray, v: np.ndarray) -> np.ndarray:
    """Evaluate vector quantiles for fixed dual weights.

    Args:
        y: Observations with shape ``(n, 2)``.
        tau: Evaluation points in ``[0, 1]^2`` with shape ``(m, 2)``.
        v: Dual weights solving the optimal transport problem, with length
            ``n``.

    Returns:
        Array of quantile locations with shape ``(m, 2)``.

    Raises:
        SystemExit: If ``tau`` does not have exactly two columns.
    """
    if tau.shape[1] != 2:
        bs_error_abort("tau must have two columns")
    q = y[np.argmax(tau @ y.T - v, axis=1), :]
    return cast(np.ndarray, q)

bivariate_ranks(y, n_nodes=32, verbose=False)

Compute barycentric ranks for each observation.

Parameters:

Name	Type	Description	Default
y	ndarray	Observations with shape (n, 2).	required
n_nodes	int	Number of Chebyshev nodes used in the quadrature.	32
verbose	bool	Print diagnostics when True.	False

Returns:

Type	Description
ndarray	Array of average ranks (shape (n, 2)) with nan for zero-mass cells.

Source code in bs_python_utils/stats/bivariate_quantiles.py

def bivariate_ranks(
    y: np.ndarray,
    n_nodes: int = 32,
    verbose: bool = False,
) -> np.ndarray:
    """Compute barycentric ranks for each observation.

    Args:
        y: Observations with shape ``(n, 2)``.
        n_nodes: Number of Chebyshev nodes used in the quadrature.
        verbose: Print diagnostics when ``True``.

    Returns:
        Array of average ranks (shape ``(n, 2)``) with ``nan`` for zero-mass cells.
    """
    d = y.shape[1]

    if d != 2:
        bs_error_abort(f"only works for 2-dimensional y, not for {d}")

    _, bivranks = _solve_for_v(y, n_nodes, verbose)
    return bivranks

bivariate_ranks_simul(x, rng, n_draws=10000, h_mult=100.0)

This computes bivariate ranks for a matrix `x using simulations

Parameters:

Name	Type	Description	Default
x	ndarray	Input matrix of shape (n, 2).	required
rng	Generator	Random number generator for reproducibility.	required
n_draws	int	Number of random draws for the simulation.	10000
h_mult	float	we set the bandwidth h as std(v_init)/h_mult, default is 100.0	100.0

Source code in bs_python_utils/stats/bivariate_quantiles.py

def bivariate_ranks_simul(
    x: np.ndarray,
    rng: np.random.Generator,
    n_draws: int = 10_000,
    h_mult: float = 100.0,
) -> np.ndarray:
    """This computes bivariate ranks for a matrix ``x` using simulations

    Args:
        x: Input matrix of shape ``(n, 2)``.
        rng: Random number generator for reproducibility.
        n_draws: Number of random draws for the simulation.
        h_mult: we set the bandwidth h as std(v_init)/h_mult, default is 100.0
    """
    n, d = check_matrix(x)
    if d != 2:
        bs_error_abort("Input matrix x must have 2 columns for bivariate ranks.")

    v_init = np.mean(x, axis=1)[:-1]
    h = np.std(v_init) / 100.0

    tau_draws = rng.uniform(low=0.0, high=1.0, size=2 * n_draws).reshape((n_draws, 2))

    def objv_and_grad(
        v: np.ndarray,
        ret: int = 0,
    ) -> float | np.ndarray | None:
        vn = np.zeros(n)
        vn[:-1] = v.copy()
        vn[-1] = -np.sum(v)
        psi_vals = tau_draws @ x.T - vn
        lse_vals = logsumexp(psi_vals / h, axis=1)
        psi_maxs = h * lse_vals
        obj_val = np.mean(psi_maxs)
        probs = np.exp(psi_vals / h - lse_vals.reshape((-1, 1)))
        mean_probs = np.mean(probs, axis=0)
        if ret == 0:
            return cast(float, obj_val)
        elif ret == 1:
            grad = -mean_probs[:-1] + mean_probs[-1]
            return cast(np.ndarray, grad)
        elif ret == 2:
            ranks = (probs.T @ tau_draws / n_draws) / mean_probs.reshape((-1, 1))
            return cast(np.ndarray, ranks)
        else:
            bs_error_abort("Invalid value for ret. Must be 0, 1, or 2.")
            return None

    def objv(v: np.ndarray, args: list) -> float:
        return cast(float, objv_and_grad(v, ret=0))

    def grad_objv(v: np.ndarray, args: list) -> np.ndarray:
        return cast(np.ndarray, objv_and_grad(v, ret=1))

    def ranks(v: np.ndarray, args: list) -> np.ndarray:
        return cast(np.ndarray, objv_and_grad(v, ret=2))

    resv = minimize_free(objv, grad_objv, v_init, args=[])
    if not resv.success:
        bs_error_abort(f"Optimization failed, message: {resv.message}")

    v_sol = resv.x
    r_x = ranks(v_sol, args=[])
    return r_x