qra/organized_uncertainty_analysis/bayesian/elasticity.py

"""Numerical elasticities of Bayesian x50 and s50 to scientific parameters."""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from .core import P_with, slope_at_x, theta_max, x_at_p

PARAMETER_NAMES = ("omega", "a", "b", "s", "k", "vartheta")


def _scientific_vector(theta_hat):
    return np.asarray(theta_hat[:6], float)


def _as_theta(values):
    omega, a, b, s, k, vartheta = map(float, values)
    cap = theta_max(a, b, k, s)
    if not (0 < omega < 1 and a > k > 0 and b > 0 and s > 0 and 0 < vartheta < cap):
        raise ValueError("Perturbed parameter vector violates model constraints.")
    return omega, a, b, s, k, vartheta, cap


def elasticity_x50_s50(theta_hat, relative_step=1e-4):
    """Return dimensionless local elasticities using central perturbations."""
    base = _scientific_vector(theta_hat)
    base_theta = _as_theta(base)
    x0 = x_at_p(base_theta)
    s0 = slope_at_x(base_theta, x0)
    rows = []
    for j, name in enumerate(PARAMETER_NAMES):
        step = relative_step * max(abs(base[j]), 1e-8)
        plus, minus = base.copy(), base.copy()
        plus[j] += step
        minus[j] -= step
        try:
            tp, tm = _as_theta(plus), _as_theta(minus)
            xp, xm = x_at_p(tp), x_at_p(tm)
            sp, sm = slope_at_x(tp, xp), slope_at_x(tm, xm)
            ex = (base[j] / x0) * (xp - xm) / (2.0 * step)
            es = (base[j] / s0) * (sp - sm) / (2.0 * step)
        except (ValueError, FloatingPointError):
            ex = es = np.nan
        rows.append({"Parameter": name, "Elasticity_x50": ex, "Elasticity_s50": es})
    return pd.DataFrame(rows)


def elasticity_full_trim(fit_result, relative_step=1e-4):
    frames = []
    for dataset in ("FULL", "TRIM"):
        frame = elasticity_x50_s50(fit_result[dataset]["theta"], relative_step)
        frame.insert(0, "Dataset", dataset)
        frames.append(frame)
    return pd.concat(frames, ignore_index=True)


def plot_combined_elasticity(table):
    """Plot all Bayesian parameter elasticities in one two-panel figure."""
    parameter_order = list(PARAMETER_NAMES)
    display_labels = [r"$\pi$", r"$a$", r"$b$", r"$s$", r"$k$", r"$\vartheta$"]
    x = np.arange(len(parameter_order))
    width = 0.36
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2), dpi=180)
    for offset, dataset in zip((-width / 2, width / 2), ("FULL", "TRIM")):
        block = table.set_index(["Dataset", "Parameter"]).loc[dataset].reindex(parameter_order)
        axes[0].bar(x + offset, block["Elasticity_x50"].abs(), width=width, label=dataset)
        axes[1].bar(x + offset, block["Elasticity_s50"].abs(), width=width, label=dataset)
    for label, axis in zip(("A", "B"), axes):
        axis.set_xticks(x, display_labels)
        axis.set_ylabel("Absolute elasticity")
        axis.text(0.02, 0.97, label, transform=axis.transAxes, ha="left", va="top", fontsize=15)
        axis.grid(alpha=0.25, axis="y")
    axes[0].legend(frameon=True, fontsize=9, loc="upper left")
    fig.tight_layout()
    return fig, axes