qra/organized_uncertainty_analysis/logistic/elasticity.py

"""Elasticity calculations for logistic x50 and s50."""

import numpy as np
import matplotlib.pyplot as plt

from .core import (
    elasticity_x50_s50,
    elasticity_table_4panels,
    plot_x50_values,
    plot_s50_values,
    plot_x50_theta1_elasticity,
    plot_s50_theta1_elasticity,
)


def plot_combined_elasticity(table):
    """Plot all logistic coefficient elasticities in one two-panel figure."""
    labels = table["Panel"].tolist()
    x = np.arange(2)
    width = 0.17
    offsets = np.linspace(-1.5, 1.5, len(labels)) * width
    x50_values = table[["E_x50_theta0", "E_x50_theta1"]].abs().to_numpy(float)
    s50_values = table[["E_s50_theta0", "E_s50_theta1"]].abs().to_numpy(float)
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.2), dpi=180)
    for index, label in enumerate(labels):
        axes[0].bar(x + offsets[index], x50_values[index], width=width, label=label)
        axes[1].bar(x + offsets[index], s50_values[index], width=width, label=label)
    axes[0].set_xticks(x, [r"$|\mathcal{E}_{\theta_0}(x_{50})|$", r"$|\mathcal{E}_{\theta_1}(x_{50})|$"])
    axes[1].set_xticks(x, [r"$|\mathcal{E}_{\theta_0}(s_{50})|$", r"$|\mathcal{E}_{\theta_1}(s_{50})|$"])
    axes[0].set_ylabel("Absolute elasticity")
    for label, axis in zip(("A", "B"), axes):
        axis.text(0.03, 0.95, label, transform=axis.transAxes, ha="left", va="top", fontsize=15)
        axis.grid(alpha=0.25, axis="y")
    axes[1].legend(frameon=True, fontsize=9, loc="upper left")
    fig.tight_layout()
    return fig, axes