from typing import Dict, List
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots
[docs]class ForecastVisualizer:
r"""A tool for quickly visualizing forecast results from the holdout
dataframe of a forecasting job.
.. code-block:: python
import kumoai
# Retrieve job results from a training training job. Note
# that the job ID passed here must be in a completed state:
job_result = kumoai.TrainingJob("trainingjob-...").result()
# Read the holdout table as a Pandas DataFrame:
holdout_df = job_result.holdout_df()
# Pass holdout table to ForecastVisualizer and visualize results
holdout_forecast = kumoai.utils.ForecastVisualizer(holdout_df)
holdout_forecast.visualize()
"""
[docs] def __init__(self, holdout_df: pd.DataFrame) -> None:
# Sort the holdout dataframe and extract unique entities:
self.forecast = holdout_df.sort_values(['ENTITY', 'TIMESTAMP'])
self.entities = holdout_df['ENTITY'].unique().tolist()
self.fig = self._initialize_subplot()
self.buttons: List[Dict] = []
self.plot_config = {
'target': {
'color': 'blue',
'name': 'TARGET'
},
'prediction': {
'color': 'red',
'name': 'TARGET_PRED'
},
'residuals': {
'color': 'green',
'name': 'Residuals'
},
'residuals_time': {
'color': 'orange',
'name': 'Residuals Over Time'
}
}
@staticmethod
def _initialize_subplot() -> go.Figure:
r"""Initializes the subplot structure with three rows:
Row 1: Line plot of actual forecast vs predicted
Row 2: Line plot of residuals overtime
Row 3: Histogram distribution of residuals
"""
return make_subplots(
rows=3,
cols=1,
specs=[[{
"type": "scatter"
}], [{
"type": "scatter"
}], [{
"type": "xy"
}]],
row_heights=[0.5, 0.25, 0.25],
subplot_titles=('Forecast vs Actual', 'Residuals Over Time',
'Residuals Distribution'),
)
def _create_time_series_trace(
self,
data: pd.Series,
entity: str,
trace_type: str,
visibility: bool,
) -> go.Scatter:
r"""Create a time series trace for either target or prediction"""
config = self.plot_config[trace_type]
return go.Scatter(
x=data["TIMESTAMP"],
y=data[config['name']],
name=f"{entity} - {config['name']}",
mode="lines",
line=dict(color=config['color']),
visible=visibility,
opacity=0.75,
)
def _create_residuals_time_trace(
self,
data: pd.Series,
entity: str,
visibility: bool,
) -> go.Scatter:
r"""Create a time series trace for residuals over time"""
residuals = data["TARGET"] - data["TARGET_PRED"]
return go.Scatter(
x=data["TIMESTAMP"],
y=residuals,
name=f"{entity} - Residuals Over Time",
mode="lines+markers",
line=dict(color=self.plot_config['residuals_time']['color']),
visible=visibility,
opacity=0.75,
)
def _create_residuals_hist_trace(
self,
data: pd.Series,
entity: str,
visibility: bool,
) -> go.Histogram:
r"""Create a histogram trace for residuals distribution."""
residuals = data["TARGET"] - data["TARGET_PRED"]
return go.Histogram(
x=residuals,
name=f"{entity} - Residuals Distribution",
marker=dict(color=self.plot_config['residuals']['color']),
visible=visibility,
opacity=0.75,
nbinsx=30,
)
def _create_button(self, index: int, entity: str) -> None:
r"""Create visibility toggle button for an entity."""
# target, prediction, residuals time, and residuals hist:
num_traces_per_entity = 4
total_traces = len(self.entities) * num_traces_per_entity
button = dict(label=entity, method="update", args=[{
"visible": [False] * total_traces
}])
# Set visibility for the entity's traces:
base_index = index * num_traces_per_entity
for i in range(num_traces_per_entity):
button["args"][0]["visible"][base_index + i] = True # type: ignore
self.buttons.append(button)
def _create_traces(self) -> None:
"""Create all traces for the visualization."""
for i, entity in enumerate(self.entities):
entity_data = self.forecast.loc[self.forecast.ENTITY == entity]
# First entity's traces are visible by default:
visibility = (i == 0)
# Create traces
trace_target = self._create_time_series_trace(
entity_data, entity, 'target', visibility)
trace_pred = self._create_time_series_trace(
entity_data, entity, 'prediction', visibility)
trace_residuals_time = self._create_residuals_time_trace(
entity_data, entity, visibility)
trace_residuals_hist = self._create_residuals_hist_trace(
entity_data, entity, visibility)
# Add traces to appropriate subplots
self.fig.add_trace(trace_target, row=1, col=1)
self.fig.add_trace(trace_pred, row=1, col=1)
self.fig.add_trace(trace_residuals_time, row=2, col=1)
self.fig.add_trace(trace_residuals_hist, row=3, col=1)
self._create_button(i, entity)
def _update_layout(self) -> None:
r"""Update the figure layout with all necessary configurations."""
self.fig.update_layout(
updatemenus=[
dict(active=0, buttons=self.buttons, direction="down", pad={
"r": 10,
"t": 10
}, showactive=True, x=1, xanchor="left", y=1.07, yanchor="top")
],
title="Forecast Results by Department",
height=1000, # Increased height to accommodate third plot
width=1300,
showlegend=True,
hovermode='x unified')
# Update axis labels and add zero reference line for residuals
self.fig.update_xaxes(title_text="Timestamp", row=1, col=1)
self.fig.update_xaxes(title_text="Timestamp", row=2, col=1)
self.fig.update_xaxes(title_text="Residual Value", row=3, col=1)
self.fig.update_yaxes(title_text="Patient Volume", row=1, col=1)
self.fig.update_yaxes(title_text="Residual Value", row=2, col=1)
self.fig.update_yaxes(title_text="Frequency", row=3, col=1)
# Add zero reference line for residuals time series
self.fig.add_hline(
y=0,
line_dash="dash",
line_color="gray",
row=2,
col=1,
)
[docs] def visualize(self) -> None:
r"""Generate and display the complete visualization."""
self._create_traces()
self._update_layout()
self.fig.show()
