"""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)