cross_model/scripts/sigma_contour_utils.py

"""HDR-style contour levels for 2D probability maps on a grid."""

from __future__ import annotations

import numpy as np


def compute_sigma_levels(
    posterior_norm: np.ndarray,
    credibility_mass: tuple[float, ...] | list[float],
) -> list[float]:
    """
    Highest-density containment: find density thresholds such that descending sort
    of mass covers ``credibility_mass[j]`` of total probability.

    Returned levels are ascending (suitable order for matplotlib ``contour``).
    """
    p = np.asarray(posterior_norm, dtype=np.float64).ravel()
    s = p.sum()
    if s <= 0:
        return [0.0 for _ in credibility_mass]
    ps = np.sort((p / s).flatten())[::-1]
    cdf = np.cumsum(ps)
    out: list[float] = []
    for cred in credibility_mass:
        j = int(np.searchsorted(cdf, cred, side="left"))
        j = min(max(j, 0), len(ps) - 1)
        out.append(float(ps[j]))
    return sorted(out)