bitcoin-model/btcmodel/plots.py

"""Charts. One colour per model, fixed by its position in MODELS."""

from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
from matplotlib.ticker import FuncFormatter, LogLocator, PercentFormatter  # noqa: E402

from .evaluate import COVERAGES, horizon_label  # noqa: E402
from .forecast import Forecast  # noqa: E402
from .halving import CYCLE_STARTS  # noqa: E402
from .models import BASELINE, MODELS  # noqa: E402

SURFACE = "#fcfcfb"
INK = "#0b0b0b"
INK_SECONDARY = "#52514e"
MUTED = "#898781"
GRID = "#e1e0d9"
AXIS = "#c3c2b7"
SERIES = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#4a3aa7", "#e34948"]

plt.rcParams.update(
    {
        "figure.facecolor": SURFACE,
        "axes.facecolor": SURFACE,
        "savefig.facecolor": SURFACE,
        "axes.edgecolor": AXIS,
        "axes.linewidth": 0.8,
        "axes.labelcolor": INK_SECONDARY,
        "axes.titlecolor": INK,
        "axes.titlesize": 11,
        "axes.titleweight": "semibold",
        "axes.titlelocation": "left",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.grid": True,
        "grid.color": GRID,
        "grid.linewidth": 0.8,
        "xtick.color": MUTED,
        "ytick.color": MUTED,
        "xtick.labelcolor": INK_SECONDARY,
        "ytick.labelcolor": INK_SECONDARY,
        "legend.frameon": False,
        "legend.labelcolor": INK_SECONDARY,
        "lines.linewidth": 2,
        "lines.solid_capstyle": "round",
        "lines.solid_joinstyle": "round",
        "font.family": "sans-serif",
        "font.size": 9,
    }
)


def model_colors(names) -> dict[str, str]:
    """Registered models keep their MODELS slot; others follow in order of appearance."""
    names = list(dict.fromkeys(names))
    if all(n in MODELS for n in names):
        return {n: SERIES[list(MODELS).index(n)] for n in names}
    return {n: SERIES[i % len(SERIES)] for i, n in enumerate(names)}


def price_formatter(x, _=None) -> str:
    for scale, suffix in ((1e9, "B"), (1e6, "M"), (1e3, "K")):
        if x >= scale:
            return f"${x / scale:.3g}{suffix}"
    return f"${x:.3g}"


def fan_chart(history: pd.DataFrame, forecasts: dict[str, Forecast], path: Path) -> None:
    """History plus 50/80/95% intervals, one panel per model on shared axes."""
    fig, axes = plt.subplots(
        len(forecasts), 1, figsize=(10, 3.2 * len(forecasts)), sharex=True, sharey=True
    )
    axes = np.atleast_1d(axes)
    shown = history[history.index >= history.index[-1] - pd.Timedelta(days=6 * 365)]
    colors = model_colors(forecasts)
    for ax, (name, f) in zip(axes, forecasts.items(), strict=True):
        color = colors[name]
        for c in sorted(COVERAGES, reverse=True):
            lo, hi = f.interval(c)
            ax.fill_between(f.dates, np.exp(lo), np.exp(hi), color=color, alpha=0.1, lw=0)
        median = np.exp(f.quantile(0.5))
        ax.plot(shown.index, shown["close"], color=INK, lw=1.2)
        ax.plot(f.dates, median, color=color)
        ax.annotate(
            f"median {price_formatter(median[-1])}",
            (f.dates[-1], median[-1]),
            xytext=(6, 0),
            textcoords="offset points",
            va="center",
            color=INK_SECONDARY,
        )
        for start in CYCLE_STARTS:
            if shown.index[0] <= start <= f.dates[-1]:
                ax.axvline(start, color=AXIS, lw=0.8, zorder=0)
        ax.set_yscale("log")
        ax.yaxis.set_major_locator(LogLocator(base=10, subs=(1, 2, 5)))
        ax.yaxis.set_major_formatter(FuncFormatter(price_formatter))
        ax.yaxis.set_minor_formatter(FuncFormatter(lambda *_: ""))
        ax.set_title(f"{name}: forecast from {f.origin:%Y-%m-%d}")
    axes[0].fill_between([], [], color=MUTED, alpha=0.3, label="50% interval")
    axes[0].fill_between([], [], color=MUTED, alpha=0.2, label="80% interval")
    axes[0].fill_between([], [], color=MUTED, alpha=0.1, label="95% interval")
    axes[0].legend(loc="upper left")
    axes[-1].set_xlabel("vertical lines: halvings (future ones projected)", color=MUTED)
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)


def skill_chart(summary: pd.DataFrame, path: Path, title: str, baseline: str = BASELINE) -> None:
    """CRPS skill vs `baseline`, by horizon, with bootstrap intervals."""
    horizons = sorted(summary["horizon"].unique())
    x = np.arange(len(horizons))
    colors = model_colors(summary["model"])
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.axhline(0, color=colors[baseline], lw=2, label=baseline)
    for name, g in summary[summary.model != baseline].groupby("model", sort=False):
        g = g.set_index("horizon").reindex(horizons)
        color = colors[name]
        ax.fill_between(x, g.skill_lo, g.skill_hi, color=color, alpha=0.1, lw=0)
        ax.plot(x, g.skill, color=color, marker="o", ms=6, mec=SURFACE, mew=2, label=name)
        ax.annotate(
            name,
            (x[-1], g.skill.iloc[-1]),
            xytext=(8, 0),
            textcoords="offset points",
            va="center",
            color=INK_SECONDARY,
        )
    ax.set_xticks(x, [horizon_label(h) for h in horizons])
    ax.set_xlabel("forecast horizon")
    ax.set_ylabel(f"CRPS skill vs {baseline} (higher is better)")
    ax.yaxis.set_major_formatter(PercentFormatter(1.0, decimals=0))
    ax.set_title(title)
    ax.legend(loc="lower left")
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)


def calibration_chart(summary: pd.DataFrame, path: Path, title: str) -> None:
    """How often each nominal interval contained the outcome, by horizon."""
    horizons = sorted(summary["horizon"].unique())
    x = np.arange(len(horizons))
    colors = model_colors(summary["model"])
    fig, axes = plt.subplots(1, len(COVERAGES), figsize=(12, 4), sharey=True)
    for ax, c in zip(axes, COVERAGES, strict=True):
        ax.axhline(c, color=INK_SECONDARY, lw=1)
        ax.annotate(
            "target", (x[-1], c), xytext=(0, 4), textcoords="offset points", ha="right", color=MUTED
        )
        for name, g in summary.groupby("model", sort=False):
            g = g.set_index("horizon").reindex(horizons)
            ax.plot(
                x,
                g[f"cov{c:.0%}"],
                color=colors[name],
                marker="o",
                ms=6,
                mec=SURFACE,
                mew=2,
                label=name,
            )
        ax.set_xticks(x, [horizon_label(h) for h in horizons])
        ax.set_title(f"{c:.0%} interval")
        ax.set_ylim(0, 1.02)
        ax.yaxis.set_major_formatter(PercentFormatter(1.0, decimals=0))
    axes[0].set_ylabel("share of outcomes inside")
    axes[-1].legend(loc="lower left")
    fig.suptitle(title, x=0.01, ha="left", color=INK, fontsize=11, fontweight="semibold")
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)