Spaces:
Running on Zero
Running on Zero
File size: 7,020 Bytes
ba6c090 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """Plot observed history and forecast quantiles."""
from __future__ import annotations
from collections.abc import Sequence
from pathlib import Path
from typing import TYPE_CHECKING
import torch
from torch import Tensor
if TYPE_CHECKING:
from matplotlib.figure import Figure
def plot_forecast(
history: Tensor,
prediction: Tensor,
quantiles: Tensor,
*,
timestamps: Sequence | Tensor | None = None,
ylabel: str = "Value",
output: str | Path | None = None,
) -> Figure:
"""Plot ``[variables, time]`` history and ``[variables, quantiles, horizon]`` predictions.
``quantiles`` is the model's probability grid. Each variable gets one panel,
showing the median and available nested intervals between 0.1 and 0.9.
``timestamps`` may supply a shared date axis covering history plus horizon;
otherwise the forecast starts at step zero. Returns a Matplotlib figure and
optionally saves it to ``output``.
"""
if any(not isinstance(value, Tensor) for value in (history, prediction, quantiles)):
raise ValueError("history, prediction and quantiles must be PyTorch tensors")
if any(value.is_complex() for value in (history, prediction, quantiles)):
raise ValueError("plot inputs must be real")
if history.ndim != 2 or min(history.shape) < 1:
raise ValueError("history must be a nonempty [variables, time] tensor")
if prediction.ndim != 3 or min(prediction.shape) < 1 or prediction.shape[0] != history.shape[0]:
raise ValueError("prediction must have shape [variables, quantiles, horizon]")
if (
quantiles.ndim != 1
or len(quantiles) != prediction.shape[1]
or not bool(((quantiles > 0) & (quantiles < 1)).all())
or not bool((quantiles.diff() > 0).all())
):
raise ValueError("quantiles must be a matching, strictly increasing probability grid")
if bool(history.isinf().any()) or not bool(prediction.isfinite().all()):
raise ValueError("history must be finite or NaN and predictions must be finite")
history, prediction, quantiles = (
value.detach().cpu().float() for value in (history, prediction, quantiles)
)
levels = torch.arange(1, 10) / 10
columns = (quantiles[:, None] - levels).abs().argmin(dim=0)
available = torch.isclose(quantiles[columns], levels, atol=1e-6, rtol=0)
if not available[4]:
raise ValueError("quantiles must include the median (0.5)")
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
from matplotlib import font_manager
count, length = history.shape
horizon = prediction.shape[-1]
if timestamps is None:
times = torch.arange(-length, horizon).numpy()
else:
import pandas as pd
if isinstance(timestamps, Tensor):
timestamps = timestamps.detach().cpu().numpy()
times = pd.DatetimeIndex(pd.to_datetime(timestamps))
if len(times) != length + horizon or times.hasnans or not times.is_unique:
raise ValueError("timestamps must cover history plus horizon without missing dates")
if not times.is_monotonic_increasing:
raise ValueError("timestamps must be strictly increasing")
past, future = times[:length], times[length:]
background, panel, teal = "#173446", "#1d3d50", "#549a8b"
white, muted, grid = "#fafbfb", "#b6c3c9", "#2d4d5d"
fonts = {font.name for font in font_manager.fontManager.ttflist}
with plt.rc_context(
{
"font.family": "Aptos" if "Aptos" in fonts else "DejaVu Sans",
"font.size": 14,
"axes.labelsize": 16,
"text.color": white,
"axes.labelcolor": muted,
"xtick.color": muted,
"ytick.color": muted,
"savefig.dpi": 180,
}
):
figure, panels = plt.subplots(count, 1, figsize=(13.5, 5.8 * count), squeeze=False)
figure.set_facecolor(background)
figure.subplots_adjust(left=0.09, right=0.98, bottom=0.15 / count, top=0.97, hspace=0.3)
for index, axes in enumerate(panels[:, 0]):
axes.set_facecolor(panel)
axes.axvspan(future[0], future[-1], color="#214452", linewidth=0)
axes.axvline(future[0], color=teal, linewidth=1.2, linestyle=":")
axes.grid(color=grid, linewidth=0.7)
axes.set_axisbelow(True)
axes.spines[["top", "right", "left"]].set_visible(False)
axes.spines["bottom"].set_color(grid)
axes.tick_params(length=0, pad=8)
axes.set_xlim(past[0], future[-1])
axes.plot(past, history[index], color=white, linewidth=1.7, label="History")
for lo, hi, alpha in [(0, 8, 0.22), (1, 7, 0.20), (2, 6, 0.18), (3, 5, 0.16)]:
if available[lo] and available[hi]:
axes.fill_between(
future,
prediction[index, columns[lo]],
prediction[index, columns[hi]],
color=teal,
alpha=alpha,
linewidth=0,
label="10–90% quantiles" if lo == 0 else None,
)
median = prediction[index, columns[4]]
axes.plot(
future, median, color="#76c6b3", linewidth=2.4, label="Median forecast", zorder=4
)
axes.plot(
[past[-1], future[0]],
[float(history[index, -1]), float(median[0])],
color="#76c6b3",
linewidth=1.5,
linestyle=":",
)
axes.set_ylabel(ylabel)
if count > 1:
axes.set_title(f"Series {index + 1}", color=white)
low, high = axes.get_ylim()
if low < 0 < high:
axes.axhline(0, color=muted, alpha=0.3, linewidth=0.8)
if timestamps is not None:
if times[-1] - times[0] <= pd.Timedelta(days=14):
axes.xaxis.set_major_locator(mdates.DayLocator())
axes.xaxis.set_major_formatter(mdates.DateFormatter("%a %d %b"))
else:
locator = mdates.AutoDateLocator()
axes.xaxis.set_major_locator(locator)
axes.xaxis.set_major_formatter(mdates.ConciseDateFormatter(locator))
axes.legend(
loc="upper left",
frameon=True,
facecolor=background,
edgecolor=grid,
framealpha=0.95,
fontsize=15,
labelcolor=white,
borderpad=0.9,
)
if timestamps is None:
panels[-1, 0].set_xlabel("Time step relative to forecast origin")
if output is not None:
path = Path(output).expanduser()
path.parent.mkdir(parents=True, exist_ok=True)
figure.savefig(path, facecolor=figure.get_facecolor())
return figure
|