qra/organized_uncertainty_analysis/bayesian/data_fit_gof.py

"""Data preparation, constrained MAP fitting, prediction, and fit diagnostics."""

import numpy as np

from .core import (
    load_xy,
    make_trimmed_dataset,
    sigmoid,
    softplus,
    theta_max,
    unpack,
    dE_full,
    P_with,
    make_priors,
    init_phi,
    neg_post,
    fit_bayes,
    x_at_p,
    slope_at_x,
)


def goodness_of_fit(X, y, theta_hat, threshold=0.5):
    """Return conditional Bernoulli fit measures for a fitted CB model."""
    X = np.asarray(X, float).ravel()
    y = np.asarray(y, int).ravel()
    p = np.clip(np.asarray(P_with(theta_hat, X), float), 1e-12, 1.0 - 1e-12)
    pred = (p >= threshold).astype(int)
    log_loss = -float(np.mean(y * np.log(p) + (1 - y) * np.log1p(-p)))
    return {
        "n": int(y.size),
        "events": int(y.sum()),
        "log_loss": log_loss,
        "brier_score": float(np.mean((y - p) ** 2)),
        "accuracy": float(np.mean(pred == y)),
        "sensitivity": float(np.mean(pred[y == 1] == 1)) if np.any(y == 1) else np.nan,
        "specificity": float(np.mean(pred[y == 0] == 0)) if np.any(y == 0) else np.nan,
    }


def fit_full_trim(percentile=95, **fit_kwargs):
    """Load, trim, and fit the FULL and TRIM constrained Bayesian models."""
    X, y = load_xy(perc=percentile)
    ds = make_trimmed_dataset(X, y)
    result = {"data": ds}
    for key, x_key, y_key in (
        ("FULL", "X_orig", "y_orig"),
        ("TRIM", "X_trim", "y_trim"),
    ):
        theta, optimizer = fit_bayes(ds[x_key], ds[y_key], **fit_kwargs)
        result[key] = {
            "theta": theta,
            "optimizer": optimizer,
            "gof": goodness_of_fit(ds[x_key], ds[y_key], theta),
        }
    return result