"""
Result containers for BayesianMMM.
This module contains the dataclasses used to return results from
model fitting, prediction, and contribution analysis.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
import numpy as np
import pandas as pd
if TYPE_CHECKING:
import arviz as az
import pymc as pm
from ..data_loader import PanelDataset
from ..estimands.spec import EstimandResult
[docs]
@dataclass
class MMMResults:
"""Container for fitted model results."""
trace: az.InferenceData
model: pm.Model
panel: PanelDataset
channel_contributions: pd.DataFrame | None = None
diagnostics: dict = field(default_factory=dict)
y_mean: float = 0.0
y_std: float = 1.0
# Declarative estimands realized from the posterior at fit time, keyed by
# estimand name (wildcard-channel estimands expand to "{name}:{channel}").
# Populated best-effort when the model declares estimands; otherwise empty
# (use ``model.evaluate_estimands()`` on demand). See mmm_framework.estimands.
estimands: "dict[str, EstimandResult]" = field(default_factory=dict)
# Approximate fits (MAP / ADVI / Pathfinder) set this to True. The posterior
# is a fast approximation: convergence diagnostics (R-hat/ESS) do not apply
# and the uncertainty is not calibrated. Use for model checking, not for
# final inference. ``diagnostics["fit_method"]`` records which method ran.
approximate: bool = False
@property
def converged(self) -> bool | None:
"""MCMC convergence verdict (R-hat / ESS / divergences).
``True``/``False`` for NUTS fits; ``None`` when not assessable (an
approximate MAP/ADVI fit, or no usable diagnostics). ``None`` is NOT
"converged" -- surface it as "N/A". Do not act on intervals/ROI from a
fit where this is ``False``.
"""
from ..diagnostics.convergence import is_converged
return is_converged(self.diagnostics)
@property
def convergence_flags(self) -> list[str]:
"""Which convergence checks failed: subset of ``{divergences, rhat, ess}``."""
from ..diagnostics.convergence import convergence_flags
return convergence_flags(self.diagnostics)
[docs]
def summary(self, var_names: list[str] | None = None) -> pd.DataFrame:
"""Get posterior summary statistics."""
from ..utils.arviz_compat import summary as az_summary
return az_summary(self.trace, var_names=var_names)
[docs]
def plot_trace(self, var_names: list[str] | None = None, **kwargs):
"""Plot trace diagnostics."""
import arviz as az
return az.plot_trace(self.trace, var_names=var_names, **kwargs)
[docs]
def plot_posterior(self, var_names: list[str] | None = None, **kwargs):
"""Plot posterior distributions."""
import arviz as az
return az.plot_posterior(self.trace, var_names=var_names, **kwargs)
[docs]
@dataclass
class PredictionResults:
"""Container for prediction results."""
posterior_predictive: az.InferenceData
y_pred_mean: np.ndarray
y_pred_std: np.ndarray
y_pred_hdi_low: np.ndarray
y_pred_hdi_high: np.ndarray
y_pred_samples: np.ndarray # Shape: (n_samples, n_obs)
@property
def n_samples(self) -> int:
return self.y_pred_samples.shape[0]
@property
def n_obs(self) -> int:
return self.y_pred_samples.shape[1]
[docs]
@dataclass
class ContributionResults:
"""
Container for channel contribution results.
These results are computed via counterfactual analysis:
contribution = prediction(all channels) - prediction(channel zeroed out)
Attributes
----------
channel_contributions : pd.DataFrame
Per-observation contributions, shape (n_obs, n_channels).
Index matches the original panel index.
total_contributions : pd.Series
Total contribution by channel (summed over time).
contribution_pct : pd.Series
Contribution as percentage of total.
baseline_prediction : np.ndarray
Prediction with all channels present.
counterfactual_predictions : dict[str, np.ndarray]
Counterfactual predictions with each channel zeroed out.
time_period : tuple[int, int] | None
Time period used for calculation, if any.
contribution_hdi_low : pd.Series | None
Lower HDI for total contributions (if uncertainty computed).
contribution_hdi_high : pd.Series | None
Upper HDI for total contributions (if uncertainty computed).
"""
channel_contributions: pd.DataFrame
total_contributions: pd.Series
contribution_pct: pd.Series
baseline_prediction: np.ndarray
counterfactual_predictions: dict[str, np.ndarray]
time_period: tuple[int, int] | None = None
contribution_hdi_low: pd.Series | None = None
contribution_hdi_high: pd.Series | None = None
[docs]
def summary(self) -> pd.DataFrame:
"""Get summary DataFrame."""
data = {
"Channel": self.total_contributions.index,
"Total Contribution": self.total_contributions.values,
"Contribution %": self.contribution_pct.values,
}
if self.contribution_hdi_low is not None:
data["HDI 3%"] = self.contribution_hdi_low.values
data["HDI 97%"] = self.contribution_hdi_high.values
return pd.DataFrame(data)
[docs]
@dataclass
class ComponentDecomposition:
"""Container for full component decomposition results."""
# Component contributions (original scale, per observation)
intercept: np.ndarray
trend: np.ndarray
seasonality: np.ndarray
media_total: np.ndarray
media_by_channel: pd.DataFrame
controls_total: np.ndarray
controls_by_var: pd.DataFrame | None
geo_effects: np.ndarray | None
product_effects: np.ndarray | None
# Aggregated totals
total_intercept: float
total_trend: float
total_seasonality: float
total_media: float
total_controls: float
total_geo: float | None
total_product: float | None
# Scaling parameters for reference
y_mean: float
y_std: float
# Holiday / event effects (#143), original scale — separate from the smooth
# ``seasonality``. ``None`` when no event block was fit.
events: np.ndarray | None = None
total_events: float | None = None
# Cross-channel synergy / interaction (#142), original scale — the joint
# effect not captured by the additive per-channel contributions. ``None``
# when no interaction terms were fit.
interactions: np.ndarray | None = None
total_interactions: float | None = None
# Price & promotion levers (#138), original scale — the combined price-
# elasticity + promo-lift contribution. ``None`` when no lever was fit.
levers: np.ndarray | None = None
total_levers: float | None = None
[docs]
def summary(self) -> pd.DataFrame:
"""Get summary of component contributions."""
components = {
"Base (Intercept)": self.total_intercept,
"Trend": self.total_trend,
"Seasonality": self.total_seasonality,
"Media (Total)": self.total_media,
"Controls (Total)": self.total_controls,
}
if self.total_geo is not None:
components["Geo Effects"] = self.total_geo
if self.total_product is not None:
components["Product Effects"] = self.total_product
if self.total_events is not None:
components["Events / Holidays"] = self.total_events
if self.total_interactions is not None:
components["Synergy / Interactions"] = self.total_interactions
if self.total_levers is not None:
components["Price & Promotion"] = self.total_levers
total = sum(components.values())
df = pd.DataFrame(
{
"Component": list(components.keys()),
"Total Contribution": list(components.values()),
"Contribution %": [
v / total * 100 if total != 0 else 0 for v in components.values()
],
}
)
return df
[docs]
def controls_summary(self) -> pd.DataFrame | None:
"""Get detailed control variable breakdown."""
if self.controls_by_var is None:
return None
totals = self.controls_by_var.sum()
total_controls = totals.sum()
return pd.DataFrame(
{
"Variable": totals.index,
"Total Contribution": totals.values,
"Share of Controls %": (
(totals / total_controls * 100).values
if total_controls != 0
else [0] * len(totals)
),
}
)
__all__ = [
"MMMResults",
"PredictionResults",
"ContributionResults",
"ComponentDecomposition",
]