Add Legends To Altair Plots

hubverse_annotator/ui.py

@@ -17,6 +17,7 @@
 import streamlit as st
 from streamlit_shortcuts import add_shortcuts
 from utils import (
+    build_ci_specs_from_df,
     get_available_locations,
     get_initial_window_range,
     get_reference_dates,
@@ -31,6 +32,7 @@
 CHART_TITLE_FONT_SIZE = 18
 REF_DATE_STROKE_WIDTH = 2.5
 REF_DATE_STROKE_DASH = [6, 6]
+MARKER_SIZE = 65
 ROOT = pathlib.Path(__file__).resolve().parent.parent
 
 
@@ -414,11 +416,36 @@ def plotting_ui(
     # empty streamlit object (DeltaGenerator) needed for
     # plots to reload successfully with new data.
     base_chart = st.empty()
-    forecast_layer = quantile_forecast_chart(
-        forecasts_to_plot, selected_target, scale=scale, grid=show_grid
+
+    has_obs = not data_to_plot.is_empty()
+    has_fc = not forecasts_to_plot.is_empty()
+    ci_specs = build_ci_specs_from_df(forecasts_to_plot) if has_fc else {}
+
+    legend_labels = ["Observations"] if has_obs else []
+    color_range = ["limegreen"] if has_obs else []
+
+    if has_fc and ci_specs:
+        legend_labels.extend(ci_specs.keys())
+        color_range.extend([spec["color"] for spec in ci_specs.values()])
+    color_enc = alt.Color(
+        "legend_label:N",
+        title=None,
+        scale=alt.Scale(domain=legend_labels, range=color_range),
     )
     observed_layer = target_data_chart(
-        data_to_plot, selected_target, scale=scale, grid=show_grid
+        data_to_plot,
+        selected_target,
+        color_enc=color_enc,
+        scale=scale,
+        grid=show_grid,
+    )
+    forecast_layer = quantile_forecast_chart(
+        forecasts_to_plot,
+        selected_target,
+        ci_specs,
+        color_enc=color_enc,
+        scale=scale,
+        grid=show_grid,
     )
     sub_layers = [
         layer
@@ -472,6 +499,13 @@ def plotting_ui(
         .interactive()
         .resolve_scale(y="independent")
         .resolve_axis(x="independent")
+        .configure_legend(
+            orient="top",
+            direction="horizontal",
+            symbolType="circle",
+            symbolSize=MARKER_SIZE,
+            titleAnchor="middle",
+        )
     )
     base_chart.altair_chart(
         chart,

hubverse_annotator/utils.py

@@ -11,18 +11,23 @@
 import datetime
 import logging
 import pathlib
+from itertools import combinations
 from typing import Literal
 
 import altair as alt
+import colorbrewer
 import forecasttools
 import polars as pl
 import polars.selectors as cs
 import streamlit as st
+from matplotlib.colors import to_hex
 from streamlit.runtime.uploaded_file_manager import UploadedFile
 
 PLOT_WIDTH = 625
 STROKE_WIDTH = 2
-MARKER_SIZE = 55
+MARKER_SIZE = 65
+MAX_NUM_CIS = 7
+
 
 type ScaleType = Literal["linear", "log"]
 
@@ -132,6 +137,73 @@ def get_initial_window_range(
     return (start_date, end_date)
 
 
+def wide_quantile_df(forecast_table: pl.DataFrame) -> pl.DataFrame:
+    """
+    Converts long-format quantile forecast table into
+    wide format where each quantile is a column.
+    """
+    return (
+        forecast_table.filter(pl.col("output_type") == "quantile")
+        .pivot(
+            on="output_type_id",
+            index=cs.exclude("output_type_id", "value"),
+            values="value",
+        )
+        .rename({"0.5": "median"})
+    )
+
+
+def build_ci_specs_from_df(
+    forecast_table: pl.DataFrame,
+) -> dict[str, dict[str, str]]:
+    """
+    Automatically constructs a CI_SPECS-style dict for
+    altair legend creation by finding quantile columns in
+    a wide forecast table.
+
+    Parameters
+    ----------
+    forecast_table : pl.DataFrame
+        A Polars DataFrame of forecasts, with a quantile
+        forecast columns.
+
+    Returns
+    -------
+    dict[str, dict[str, str]]
+        Mapping from CI label (e.g. "95% CI") to its
+        bounds and color.
+    """
+    df_wide = wide_quantile_df(forecast_table)
+    quant_vals = sorted(
+        float(col) for col in df_wide.columns if col.startswith("0.")
+    )
+    ci_pairs = [
+        (low, high)
+        for low, high in combinations(quant_vals, 2)
+        if low < 0.5 < high and abs(low + high - 1.0) < 1e-6
+    ]
+    ci_pairs.sort(key=lambda p: p[1] - p[0], reverse=True)
+
+    labels = [f"{round((high - low) * 100)}% CI" for low, high in ci_pairs]
+
+    palette_rgb255 = list(
+        colorbrewer.Blues[max(3, min(MAX_NUM_CIS, len(labels)))]
+    )
+    palette = [(r / 255, g / 255, b / 255) for r, g, b in palette_rgb255]
+
+    specs = {
+        label: {
+            "low": f"{low:.3f}".rstrip("0").rstrip("."),
+            "high": f"{high:.3f}".rstrip("0").rstrip("."),
+            "color": to_hex(color),
+        }
+        for (low, high), label, color in zip(
+            ci_pairs, labels, palette, strict=False
+        )
+    }
+    return specs
+
+
 def is_empty_chart(chart: alt.LayerChart) -> bool:
     """
     Checks if an altair layer is empty. Primarily used for
@@ -172,6 +244,7 @@ def is_empty_chart(chart: alt.LayerChart) -> bool:
 def target_data_chart(
     observed_data_table: pl.DataFrame,
     selected_target: str,
+    color_enc: alt.Color,
     scale: ScaleType = "log",
     grid: bool = True,
 ) -> alt.Chart | alt.LayerChart:
@@ -186,6 +259,9 @@ def target_data_chart(
     selected_target : str
         The target for filtering in the forecast and or
         observed hubverse tables.
+    color_enc : alt.Color
+        An Altair color encoding used for plotting the
+        observations color and legend.
     scale : str
         The scale to use for the Y axis during plotting.
         Defaults to logarithmic.
@@ -203,7 +279,6 @@ def target_data_chart(
     x_enc = alt.X(
         "date:T",
         axis=alt.Axis(title="Date", grid=grid, ticks=True, labels=True),
-        scale=alt.Scale(type=scale),
     )
     y_enc = alt.Y(
         "observation:Q",
@@ -218,10 +293,15 @@ def target_data_chart(
     )
     obs_layer = (
         alt.Chart(observed_data_table, width=PLOT_WIDTH)
-        .mark_point(filled=True, size=MARKER_SIZE, color="limegreen")
+        .transform_calculate(legend_label="'Observations'")
+        .mark_point(
+            filled=True,
+            size=MARKER_SIZE,
+        )
         .encode(
             x=x_enc,
             y=y_enc,
+            color=color_enc,
             tooltip=[
                 alt.Tooltip("date:T", title="Date"),
                 alt.Tooltip("observation:Q", title="Value"),
@@ -234,6 +314,8 @@ def target_data_chart(
 def quantile_forecast_chart(
     forecast_table: pl.DataFrame,
     selected_target: str,
+    ci_specs,
+    color_enc: alt.Color,
     scale: ScaleType = "log",
     grid: bool = True,
 ) -> alt.LayerChart:
@@ -250,6 +332,9 @@ def quantile_forecast_chart(
     selected_target : str
         The target for filtering in the forecast and or
         observed hubverse tables.
+    color_enc : alt.Color
+        An Altair color encoding used for plotting the
+        quantile bands color and legend.
     scale : str
         The scale to use for the Y axis during plotting.
         Defaults to logarithmic.
@@ -264,24 +349,26 @@ def quantile_forecast_chart(
     """
     if forecast_table.is_empty():
         return alt.layer()
-    df_wide = (
-        forecast_table.filter(pl.col("output_type") == "quantile")
-        .pivot(
-            on="output_type_id",
-            index=cs.exclude("output_type_id", "value"),
-            values="value",
-        )
-        .rename({"0.5": "median"})
-    )
+    df_wide = wide_quantile_df(forecast_table)
     x_enc = alt.X("target_end_date:T", title="Date", axis=alt.Axis(grid=grid))
     y_enc = alt.Y(
         "median:Q",
         axis=alt.Axis(grid=grid),
         scale=alt.Scale(type=scale),
     )
-    base = alt.Chart(df_wide, width=PLOT_WIDTH).encode(x=x_enc, y=y_enc)
+    base = (
+        alt.Chart(df_wide, width=PLOT_WIDTH)
+        .transform_calculate(
+            date="toDate(datum.target_end_date)",
+            data_type="'Forecast'",
+        )
+        .encode(
+            x=x_enc,
+            y=y_enc,
+        )
+    )
 
-    def band(low: str, high: str, opacity: float) -> alt.Chart:
+    def band(low: str, high: str, label: str) -> alt.Chart:
         """
         Builds an errorband layer for a quantile.
 
@@ -293,27 +380,33 @@ def band(low: str, high: str, opacity: float) -> alt.Chart:
         high : str
             Upper-bound column name in the wide forecast
             table (e.g., "0.975").
-        opacity : float
-            Fill opacity for the band in the range
-            [0.0, 1.0].
+        label : str
+            The label in the legend for the confidence
+            interval (e.g. "97.5% CI").
 
         Returns
         -------
         alt.Chart
             An Altair layer with the filled band from
             ``low`` to ``high``, with step interpolation.
         """
-        return base.mark_errorband(opacity=opacity, interpolate="step").encode(
-            y=alt.Y(f"{low}:Q", title=f"{selected_target}"),
-            y2=f"{high}:Q",
-            fill=alt.value("steelblue"),
+        return (
+            base.transform_calculate(legend_label=f"'{label}'")
+            .mark_errorband(interpolate="step")
+            .encode(
+                y=alt.Y(f"{low}:Q", title=f"{selected_target}"),
+                y2=f"{high}:Q",
+                color=color_enc,
+                opacity=alt.value(1.0),
+            )
         )
 
     bands = [
-        band("0.025", "0.975", 0.10),
-        band("0.1", "0.9", 0.20),
-        band("0.25", "0.75", 0.30),
+        band(spec["low"], spec["high"], label)
+        for label, spec in ci_specs.items()
+        if spec["low"] in df_wide.columns and spec["high"] in df_wide.columns
     ]
+
     median = base.mark_line(
         strokeWidth=STROKE_WIDTH, interpolate="step", color="navy"
     )

pyproject.toml

@@ -18,6 +18,8 @@ dependencies = [
     "streamlit>=1.47.0",
     "polars>=0.19.37",
     "streamlit-shortcuts>=1.1.5",
+    "colorbrewer>=0.3",
+    "matplotlib>=3.10.3",
 ]
 
 [project.urls]