Source code for mmm_framework.reporting.extractors.bayesian
"""
BayesianMMMExtractor - Extract data from mmm-framework's BayesianMMM class.
This is the primary extractor for the core MMM framework.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
import numpy as np
from loguru import logger
from ..helpers import _safe_get_column
from .base import DataExtractor
from .mixins import (
AggregationMixin,
EstimandPPCMixin,
GeoExtractionMixin,
ProductExtractionMixin,
)
from .bundle import MMMDataBundle
if TYPE_CHECKING:
pass
[docs]
class BayesianMMMExtractor(
DataExtractor,
AggregationMixin,
GeoExtractionMixin,
ProductExtractionMixin,
EstimandPPCMixin,
):
"""
Extract data from mmm-framework's BayesianMMM class.
Inherits shared utilities from DataExtractor and AggregationMixin
for HDI computation, fit statistics, and data aggregation.
Parameters
----------
mmm : BayesianMMM
Fitted BayesianMMM instance
panel : PanelDataset
Panel data used for fitting
results : MMMResults, optional
Fit results if available
ci_prob : float
Credible interval probability (default 0.8)
"""
[docs]
def __init__(
self,
mmm: Any,
panel: Any | None = None,
results: Any | None = None,
ci_prob: float = 0.8,
):
logger.debug("Initializing BayesianMMMExtractor")
self.mmm = mmm
self.panel = panel or getattr(mmm, "panel", None)
self.results = results or getattr(mmm, "_results", None)
self._ci_prob = ci_prob
@property
def ci_prob(self) -> float:
"""Credible interval probability."""
return self._ci_prob
[docs]
def extract(self) -> MMMDataBundle:
"""Extract all available data from BayesianMMM."""
bundle = MMMDataBundle()
logger.debug("Extracting data from BayesianMMM model")
# Extract basic info
bundle.channel_names = self._get_channel_names()
pooled = getattr(self.mmm, "_pooled_channels", None)
bundle.pooled_channels = sorted(pooled) if pooled else None
bundle.time_varying_betas = self._extract_time_varying_betas()
bundle.reach_frequency = self._extract_reach_frequency()
bundle.dates = self._get_dates()
# Actual values
bundle.actual = self._get_actual()
# Extract predictions if model is fitted
if self.results is not None or getattr(self.mmm, "_trace", None) is not None:
logger.debug(
"Model appears to be fitted, extracting predictions and diagnostics"
)
bundle.predicted = self._get_predictions()
bundle.fit_statistics = self._compute_fit_statistics(
bundle.actual, bundle.predicted
)
bundle.diagnostics = self._merge_fit_provenance(
self._extract_diagnostics(getattr(self.mmm, "_trace", None))
)
self.stamp_inference_family(bundle, bundle.diagnostics)
# ROI and contributions
bundle.channel_roi = self._compute_channel_roi()
bundle.component_totals = self._get_component_totals()
bundle.component_time_series = self._get_component_time_series()
# Summary metrics
bundle.total_revenue = (
float(bundle.actual.sum()) if bundle.actual is not None else None
)
bundle.marketing_attributed_revenue = self._compute_marketing_attribution()
bundle.blended_roi = self._compute_blended_roi()
bundle.marketing_contribution_pct = (
self._compute_marketing_contribution_pct(bundle.total_revenue)
)
# Saturation and adstock
bundle.saturation_curves = self._get_saturation_curves()
bundle.adstock_curves = self._get_adstock_curves()
bundle.current_spend = self._get_current_spend()
# Trace data for diagnostics
bundle.trace_data, bundle.trace_parameters = self._get_trace_data()
# Prior/posterior. A frequentist fit has NO prior — the transforms
# came from an out-of-sample search and the coefficients from a
# penalized solve — so prior-vs-posterior contraction has nothing to
# measure. Skipped rather than filled with the fallback synthetic
# priors, which would render a "prior-dominated" chip about a prior
# that does not exist.
if not bundle.is_frequentist:
bundle.prior_samples, bundle.posterior_samples = (
self._get_prior_posterior()
)
# Estimand results (mean + CI) and posterior-predictive goodness-of-fit.
# Both are best-effort so a report never fails on them.
bundle = self._extract_estimands(bundle)
if not bundle.is_frequentist:
# A posterior-predictive p-value is undefined without a
# posterior; leaving the field empty is what gates the section.
bundle = self._extract_posterior_predictive(bundle)
# Short-term vs long-term / brand split (issue #106).
bundle = self._extract_long_term(bundle)
# CFO one-pager: P&L rollup + spend-cut revenue-at-risk (issue #108).
bundle = self._extract_cfo(bundle)
# Evidence tier + identifiability gate on every channel number
# (issue #102) — runs after channel_roi + estimands are populated so
# it can stamp them.
bundle = self._extract_channel_evidence(bundle)
bundle = self._extract_aggregated_fit_data(bundle)
bundle = self._extract_aggregated_decomposition(bundle)
bundle = self._extract_geo_level_fit_data(bundle)
bundle = self._extract_geo_level_decomposition(bundle)
# Product-level data
bundle = self._extract_product_level_fit_data(bundle)
bundle = self._extract_product_level_decomposition(bundle)
# Model specification
bundle.model_specification = self._get_model_specification()
# Causal assumptions (identification strategy, designated confounders,
# and the unobserved-confounding robustness-value table). This makes the
# always-on CausalAssumptionsSection come alive with model-specific data.
bundle = self._extract_causal_assumptions(bundle)
# Latent-structure (hybrid models): an MMM that ALSO estimates a latent
# factor (e.g. LatentFactorMMM) gets a factor-loadings section ALONGSIDE
# the channel/ROI sections. Best-effort; never blocks report generation.
bundle = self._extract_latent_structure(bundle)
return bundle
def _extract_latent_structure(self, bundle: MMMDataBundle) -> MMMDataBundle:
"""Populate the latent-structure bundle fields when this MMM also exposes a
factor/class summary (duck-typed ``factor_loadings_summary`` /
``class_profile_summary``). Leaves ``bundle.model_kind == "mmm"`` so every
MMM section stays on; the factor-analysis section is gated separately on
the presence of these fields (see ``generator._initialize_sections``)."""
try:
from ...garden.contract import has_latent_structure
if not has_latent_structure(self.mmm):
return bundle
# Reuse the family-agnostic latent extractor's table/estimands logic.
from .factor_analysis import FactorAnalysisExtractor
fx = FactorAnalysisExtractor(self.mmm, ci_prob=self.ci_prob)
bundle.factor_loadings = fx._table("factor_loadings_summary")
bundle.cfa_fit_indices = fx._estimands()
bundle.latent_section_title = (
getattr(self.mmm, "LATENT_SECTION_TITLE", None) or "Latent Factor"
)
bundle.latent_table_title = "Factor loadings"
bundle.latent_estimands_title = "Latent estimands"
except Exception: # noqa: BLE001 — reporting must never hard-fail
logger.debug("latent-structure extraction skipped", exc_info=True)
return bundle
def _extract_causal_assumptions(self, bundle: MMMDataBundle) -> MMMDataBundle:
"""Populate ``bundle.causal_assumptions`` for the CausalAssumptionsSection.
Surfaces what the fitted model actually rests on: the confounders it was
told to adjust for (the *resolved* causal roles it enforced, P1), a plain
identification-strategy statement, and the per-channel robustness value to
unobserved confounding (P0-2). Degrades gracefully -- any failure leaves
``causal_assumptions`` as ``None`` so the section still renders its
always-on no-unobserved-confounding caveat.
"""
try:
from ...config import CausalControlRole
assumptions: dict[str, Any] = {}
# Designated confounders = controls the model resolved to CONFOUNDER.
# Read the RESOLVED roles the model enforced (guaranteed to match the
# fit), not a re-derivation from a DAG the model may not have.
control_names = list(getattr(self.mmm, "control_names", []) or [])
roles = list(getattr(self.mmm, "_control_causal_roles", []) or [])
# Only read roles when they align 1:1 with the control names; a
# mismatch (e.g. a corrupted/hand-mutated model) would let zip()
# silently truncate and under-report confounders, so skip instead.
if len(control_names) == len(roles):
confounders = [
name
for name, role in zip(control_names, roles)
if role == CausalControlRole.CONFOUNDER
]
else:
confounders = []
if confounders:
assumptions["assumed_confounders"] = confounders
assumptions["identification_strategy"] = (
"Effects are identified by back-door adjustment for the "
f"designated confounder(s): {', '.join(confounders)}. SUTVA "
"and <em>no unobserved</em> confounding are assumed; the "
"estimate is causal only if every common cause of spend and "
"the KPI is among the adjusted variables."
)
else:
# Honest framing: nothing was designated, so the number leans
# entirely on the (usually false) no-unobserved-confounding
# assumption rather than on any adjustment.
assumptions["identification_strategy"] = (
"No confounders were explicitly designated, so the reported "
"effects rest entirely on the no-unobserved-confounding "
"assumption (see caveat) rather than on any back-door "
"adjustment. Anchor high-stakes channels with a geo-lift / "
"incrementality experiment (mmm_framework.calibration)."
)
# Robustness to unobserved confounding (P0-2) -- needs a fitted trace.
if getattr(self.mmm, "_trace", None) is not None:
try:
from ...validation.sensitivity_unobserved import (
UnobservedConfoundingAnalysis,
)
sensitivity = UnobservedConfoundingAnalysis(self.mmm).run()
rob = sensitivity.to_dict()
if rob.get("channels"):
assumptions["robustness"] = rob
except Exception as exc: # pragma: no cover - defensive
logger.debug("Robustness extraction skipped: %s", exc)
bundle.causal_assumptions = assumptions or None
except Exception as exc: # pragma: no cover - defensive
logger.warning("Causal-assumptions extraction failed: %s", exc)
bundle.causal_assumptions = None
return bundle
# -------------------------------------------------------------------------
# Basic data extraction
# -------------------------------------------------------------------------
def _get_channel_names(self) -> list[str]:
"""Get channel names from model or panel."""
if hasattr(self.mmm, "channel_names"):
logger.debug("Retrieving channel names from model")
return list(self.mmm.channel_names)
if self.panel is not None and hasattr(self.panel, "channel_names"):
logger.debug("Retrieving channel names from panel data")
return list(self.panel.channel_names)
logger.debug("Channel names not found")
return []
def _extract_time_varying_betas(self) -> dict | None:
"""TVP (#137): per-period coefficient trajectory for each time-varying
channel (``beta_tv_<ch>`` deterministic) — median + 90% band."""
trace = getattr(self.mmm, "_trace", None)
if trace is None or not hasattr(trace, "posterior"):
return None
post = trace.posterior
out: dict[str, dict] = {}
for ch in self._get_channel_names():
name = f"beta_tv_{ch}"
if name not in post:
continue
arr = post[name].values.reshape(-1, post[name].shape[-1])
out[ch] = {
"median": np.median(arr, axis=0).tolist(),
"lower": np.percentile(arr, 5, axis=0).tolist(),
"upper": np.percentile(arr, 95, axis=0).tolist(),
}
return out or None
def _extract_reach_frequency(self) -> dict | None:
"""Reach & frequency (#141): per-channel effective-frequency insight —
the ``effective_frequency_<ch>`` scalar (raw frequency at 90% of the
response asymptote / half-saturation) with a 90% credible band and the
curve shape, for the reach-vs-frequency tradeoff insight."""
trace = getattr(self.mmm, "_trace", None)
rf = getattr(self.mmm, "_reach_freq", None)
if trace is None or not hasattr(trace, "posterior") or not rf:
return None
post = trace.posterior
out: dict[str, dict] = {}
for ch, entry in rf.items():
cfg = entry[0]
name = f"effective_frequency_{ch}"
if name not in post:
continue
arr = np.asarray(post[name].values).reshape(-1)
out[ch] = {
"effective_frequency": float(np.median(arr)),
"lower": float(np.percentile(arr, 5)),
"upper": float(np.percentile(arr, 95)),
"response": getattr(cfg.response, "value", str(cfg.response)),
"mean_frequency": float(entry[2]),
}
return out or None
def _get_dates(self) -> np.ndarray | None:
"""Get date index."""
if self.panel is not None:
logger.debug("Retrieving dates from panel data")
if hasattr(self.panel, "dates"):
logger.debug("Using 'dates' attribute from panel")
return np.array(self.panel.dates)
if hasattr(self.panel, "index"):
logger.debug("Using 'index' attribute from panel")
return np.array(self.panel.index)
logger.debug("Dates not found")
return None
def _get_unique_periods(self) -> list | None:
"""Get unique period labels in order."""
if hasattr(self.panel, "coords") and hasattr(self.panel.coords, "periods"):
periods = self.panel.coords.periods
if hasattr(periods[0], "strftime"):
return [p.strftime("%Y-%m-%d") for p in periods]
return list(periods)
if hasattr(self, "_dates") and self._dates is not None:
return list(self._dates)
return None
def _get_period_labels_per_obs(self) -> list | None:
"""Get period label for each observation."""
if hasattr(self.panel, "time_idx") and self.panel.time_idx is not None:
unique_periods = self._get_unique_periods()
if unique_periods is None:
return None
time_idx = self.panel.time_idx
return [unique_periods[int(t)] for t in time_idx]
# Fallback: try to get from panel index
if hasattr(self.panel, "y") and hasattr(self.panel.y, "index"):
idx = self.panel.y.index
if hasattr(idx, "get_level_values"):
# MultiIndex - get period level (try both cases)
for level_name in ["period", "Period"]:
try:
periods = idx.get_level_values(level_name)
if hasattr(periods[0], "strftime"):
return [p.strftime("%Y-%m-%d") for p in periods]
return list(periods)
except KeyError:
continue
# Try first level if period not found by name
try:
periods = idx.get_level_values(0)
if hasattr(periods[0], "strftime"):
return [p.strftime("%Y-%m-%d") for p in periods]
return list(periods)
except Exception as e:
logger.debug(f"period formatting failed: {e}")
else:
# Simple index (period only)
if hasattr(idx[0], "strftime"):
return [p.strftime("%Y-%m-%d") for p in idx]
return list(idx)
return None
def _get_actual(self) -> np.ndarray | None:
"""Get actual KPI values."""
if hasattr(self.mmm, "y"):
logger.debug("Retrieving actual values from model")
y = self.mmm.y
# Unstandardize if needed
if hasattr(self.mmm, "y_mean") and hasattr(self.mmm, "y_std"):
logger.debug("Unstandardizing actual values")
return y * self.mmm.y_std + self.mmm.y_mean
return np.array(y)
if self.panel is not None and hasattr(self.panel, "y"):
logger.debug("Retrieving actual values from panel data")
return np.array(self.panel.y)
logger.debug("Actual values not found")
return None
def _get_predictions(self) -> dict[str, np.ndarray] | None:
"""Get posterior predictive mean and CI."""
try:
if hasattr(self.mmm, "predict"):
logger.debug("Generating predictions using model's predict method")
pred = self.mmm.predict()
return {
"mean": np.array(pred.y_pred_mean),
"lower": np.array(pred.y_pred_hdi_low),
"upper": np.array(pred.y_pred_hdi_high),
}
# Try to get from trace
trace = getattr(self.mmm, "_trace", None)
if trace is not None and hasattr(trace, "posterior_predictive"):
logger.debug("Extracting predictions from model trace")
pp = trace.posterior_predictive
if "y_obs" in pp:
samples = pp["y_obs"].values.reshape(-1, pp["y_obs"].shape[-1])
mean = samples.mean(axis=0)
lower, upper = np.percentile(samples, [10, 90], axis=0)
return {"mean": mean, "lower": lower, "upper": upper}
except Exception as e:
logger.warning(f"Error extracting predictions: {e}")
pass
logger.debug("Predictions not found")
return None
# -------------------------------------------------------------------------
# Estimand results + posterior-predictive goodness-of-fit
#
# The orchestration (``_extract_estimands`` / ``_extract_posterior_predictive``
# / ``_ppc_bayes_p`` / ``_downsample_rows``) lives in ``EstimandPPCMixin``;
# this extractor only supplies the model hook + the predict-based PPC arrays.
# -------------------------------------------------------------------------
def _estimand_model(self) -> Any:
return self.mmm
def _ppc_arrays(self):
"""Aligned ``(observed, y_rep, pred_mean, pred_lower, pred_upper)`` in
original KPI scale at the observation level, or ``(None,)*5`` if the model
cannot produce posterior-predictive draws. ``y_rep`` is ``(n_draws, n_obs)``.
A core ``BayesianMMM`` resamples via ``predict`` (original scale).
"""
none = (None, None, None, None, None)
observed = self._get_actual()
if observed is None:
return none
observed = np.asarray(observed, dtype=float).ravel()
predict = getattr(self.mmm, "predict", None)
if not callable(predict):
return none
try:
pred = predict(return_original_scale=True, hdi_prob=self.ci_prob)
except TypeError:
pred = predict()
y_rep = getattr(pred, "y_pred_samples", None)
pred_mean = getattr(pred, "y_pred_mean", None)
if y_rep is None or pred_mean is None:
return none
y_rep = np.asarray(y_rep, dtype=float)
pred_mean = np.asarray(pred_mean, dtype=float).ravel()
if y_rep.ndim != 2 or y_rep.shape[1] != observed.shape[0]:
return none
# Non-finite values (a pathological / non-converged fit) would silently
# propagate NaN through coverage / p-values and serialize as literal NaN
# in the report. Bail out of the whole PPC instead — the section no-ops.
if not (
np.all(np.isfinite(observed))
and np.all(np.isfinite(y_rep))
and np.all(np.isfinite(pred_mean))
):
return none
lo = getattr(pred, "y_pred_hdi_low", None)
hi = getattr(pred, "y_pred_hdi_high", None)
pred_lower = np.asarray(lo, dtype=float).ravel() if lo is not None else None
pred_upper = np.asarray(hi, dtype=float).ravel() if hi is not None else None
# Drop interval bounds that are non-finite rather than fail the whole view.
if pred_lower is not None and not np.all(np.isfinite(pred_lower)):
pred_lower = None
if pred_upper is not None and not np.all(np.isfinite(pred_upper)):
pred_upper = None
return observed, y_rep, pred_mean, pred_lower, pred_upper
# -------------------------------------------------------------------------
# Aggregated fit data extraction
# -------------------------------------------------------------------------
def _extract_aggregated_fit_data(self, bundle: MMMDataBundle) -> MMMDataBundle:
"""
Extract properly aggregated model fit data for multi-geo models.
This ensures bundle.dates, bundle.actual, and bundle.predicted are
aggregated to period-level (summed over geo and product).
IMPORTANT: Uncertainty bounds must be computed from aggregated samples,
not by summing bounds directly.
"""
import pandas as pd
if not hasattr(self, "panel") or self.panel is None:
return bundle
# Get unique periods
unique_periods = self._get_unique_periods()
if unique_periods is None:
return bundle
# Check if we have multi-dimensional data
geo_idx = self._get_geo_indices()
has_geo = geo_idx is not None and len(set(geo_idx)) > 1
if not has_geo:
# Single geo - just ensure dates are unique periods
bundle.dates = unique_periods
return bundle
try:
# Get period labels for each observation
periods = self._get_period_labels_per_obs()
if periods is None:
return bundle
# Get observed values (original scale)
y_obs = self._get_actual_original_scale()
if y_obs is None:
return bundle
# Aggregate observed values by period
df_obs = pd.DataFrame(
{
"period": periods,
"actual": y_obs,
}
)
agg_obs = df_obs.groupby("period")["actual"].sum()
agg_obs = agg_obs.reindex(unique_periods).fillna(0)
# Update bundle with aggregated observed data
bundle.dates = unique_periods
bundle.actual = agg_obs.values
# For predictions, we need to aggregate samples properly to get valid bounds
bundle.predicted = self._aggregate_predictions_with_uncertainty(
periods, unique_periods
)
if bundle.predicted is not None:
# Recompute fit statistics on aggregated data
bundle.fit_statistics = self._compute_fit_statistics(
bundle.actual, bundle.predicted
)
except Exception as e:
import logging
import traceback
logging.warning(f"Failed to extract aggregated fit data: {e}")
logging.warning(traceback.format_exc())
return bundle
def _aggregate_predictions_with_uncertainty(
self,
periods: list,
unique_periods: list,
) -> dict[str, np.ndarray] | None:
"""
Aggregate predictions while properly propagating uncertainty.
The key insight: we must sum the posterior samples first,
THEN compute percentiles. Summing percentiles directly gives
invalid (too wide) bounds.
"""
import logging
try:
# Method 1: Use model's predict() method (most reliable)
y_samples_orig = None
if hasattr(self.mmm, "predict"):
try:
logging.info("Attempting to get predictions via model.predict()")
pred_results = self.mmm.predict(
return_original_scale=True, hdi_prob=self.ci_prob
)
y_samples_orig = pred_results.y_pred_samples # (n_samples, n_obs)
logging.info(
f"predict() successful, samples shape: {y_samples_orig.shape}, mean: {y_samples_orig.mean():.2f}"
)
except Exception as e:
logging.warning(f"predict() failed: {e}")
y_samples_orig = None
# Method 2: Try to reconstruct from trace components
if y_samples_orig is None:
logging.info("Attempting to reconstruct predictions from trace")
y_samples_orig = self._reconstruct_predictions_from_trace()
if y_samples_orig is not None:
logging.info(
f"Reconstruction successful, samples shape: {y_samples_orig.shape}, mean: {y_samples_orig.mean():.2f}"
)
if y_samples_orig is None:
logging.warning("Could not obtain prediction samples")
return None
n_samples, n_obs = y_samples_orig.shape
# Create period index for aggregation
period_to_idx = {p: i for i, p in enumerate(unique_periods)}
obs_period_idx = np.array([period_to_idx[p] for p in periods])
n_periods = len(unique_periods)
# Aggregate samples by period (sum over geo/product for each sample)
# Result: (n_samples, n_periods)
y_samples_agg = np.zeros((n_samples, n_periods))
for t in range(n_periods):
mask = obs_period_idx == t
if mask.any():
# Sum across all observations in this period (across geos/products)
y_samples_agg[:, t] = y_samples_orig[:, mask].sum(axis=1)
# Now compute statistics on the aggregated samples
y_pred_mean = y_samples_agg.mean(axis=0)
alpha = (1 - self.ci_prob) / 2
y_pred_lower = np.percentile(y_samples_agg, alpha * 100, axis=0)
y_pred_upper = np.percentile(y_samples_agg, (1 - alpha) * 100, axis=0)
return {
"mean": y_pred_mean,
"lower": y_pred_lower,
"upper": y_pred_upper,
}
except Exception as e:
import traceback
logging.warning(f"Failed to aggregate predictions with uncertainty: {e}")
logging.warning(traceback.format_exc())
return None
def _reconstruct_predictions_from_trace(self) -> np.ndarray | None:
"""
Reconstruct prediction samples by summing component samples.
This is a fallback when predict() isn't available.
Returns shape (n_samples, n_obs) in original scale.
"""
import logging
trace = getattr(self.mmm, "_trace", None)
if trace is None or not hasattr(trace, "posterior"):
logging.warning("No trace or posterior found")
return None
try:
posterior = trace.posterior
y_std = self.mmm.y_std if hasattr(self.mmm, "y_std") else 1
y_mean = self.mmm.y_mean if hasattr(self.mmm, "y_mean") else 0
logging.info(f"Reconstruction: y_std={y_std:.4f}, y_mean={y_mean:.4f}")
# Method 1: Try to get y_obs from posterior_predictive and transform
if (
hasattr(trace, "posterior_predictive")
and "y_obs" in trace.posterior_predictive
):
y_ppc = trace.posterior_predictive["y_obs"].values
n_chains, n_draws, n_obs = y_ppc.shape
y_samples = y_ppc.reshape(n_chains * n_draws, n_obs)
y_samples_orig = y_samples * y_std + y_mean
logging.info(
f"Using y_obs from posterior_predictive: shape {y_samples_orig.shape}, mean={y_samples_orig.mean():.2f}"
)
return y_samples_orig
# Method 2: Reconstruct mu from component samples
if "intercept" in posterior:
intercept_samples = posterior["intercept"].values
n_chains, n_draws = intercept_samples.shape[:2]
logging.info(
f"Intercept samples shape: {intercept_samples.shape}, mean: {intercept_samples.mean():.4f}"
)
else:
logging.warning("No intercept found in posterior")
return None
n_obs = len(self.panel.y) if self.panel is not None else None
if n_obs is None:
return None
n_samples = n_chains * n_draws
# Start with intercept (broadcast to all observations)
intercept_flat = intercept_samples.reshape(n_samples, -1)
if intercept_flat.shape[1] == 1:
mu_samples = np.broadcast_to(intercept_flat, (n_samples, n_obs)).copy()
else:
mu_samples = intercept_flat
logging.info(f"After intercept: mu mean={mu_samples.mean():.4f}")
# Add trend if present
if "trend_component" in posterior:
trend = posterior["trend_component"].values.reshape(n_samples, -1)
if trend.shape[1] == n_obs:
mu_samples = mu_samples + trend
logging.info(f"After trend: mu mean={mu_samples.mean():.4f}")
# Add seasonality if present
if "seasonality_component" in posterior:
seas = posterior["seasonality_component"].values.reshape(n_samples, -1)
if seas.shape[1] == n_obs:
mu_samples = mu_samples + seas
logging.info(f"After seasonality: mu mean={mu_samples.mean():.4f}")
# Add geo effect if present (hierarchical geo model)
if "geo_offset" in posterior and "geo_sigma" in posterior:
geo_idx = self._get_geo_indices()
if geo_idx is not None:
geo_sigma = posterior["geo_sigma"].values.reshape(n_samples, 1)
geo_offset = posterior["geo_offset"].values.reshape(n_samples, -1)
geo_effect = geo_sigma * geo_offset # (n_samples, n_geos)
geo_contrib = geo_effect[:, geo_idx] # (n_samples, n_obs)
mu_samples = mu_samples + geo_contrib
logging.info(f"After geo effect: mu mean={mu_samples.mean():.4f}")
# Add product effect if present (hierarchical product model)
if "product_offset" in posterior and "product_sigma" in posterior:
product_idx = self._get_product_indices()
if product_idx is not None:
product_sigma = posterior["product_sigma"].values.reshape(
n_samples, 1
)
product_offset = posterior["product_offset"].values.reshape(
n_samples, -1
)
product_effect = product_sigma * product_offset
product_contrib = product_effect[:, product_idx]
mu_samples = mu_samples + product_contrib
logging.info(
f"After product effect: mu mean={mu_samples.mean():.4f}"
)
# Add media contribution
if "media_total" in posterior:
media = posterior["media_total"].values.reshape(n_samples, -1)
if media.shape[1] == n_obs:
mu_samples = mu_samples + media
logging.info(f"After media: mu mean={mu_samples.mean():.4f}")
# Add control contribution if present
if "control_contributions" in posterior:
ctrl = posterior["control_contributions"].values
# Sum over control dimension if needed
if ctrl.ndim > 3:
ctrl = ctrl.sum(axis=-1) # Sum over controls
ctrl = ctrl.reshape(n_samples, -1)
if ctrl.shape[1] == n_obs:
mu_samples = mu_samples + ctrl
logging.info(f"After controls: mu mean={mu_samples.mean():.4f}")
# Otherwise, compute from beta_controls and X_controls
elif (
"beta_controls" in posterior
and hasattr(self.mmm, "X_controls")
and self.mmm.X_controls is not None
):
beta_samples = posterior["beta_controls"].values
beta_samples = beta_samples.reshape(n_samples, -1)
X_controls = self.mmm.X_controls
ctrl = np.einsum("oc,sc->so", X_controls, beta_samples)
mu_samples = mu_samples + ctrl
logging.info(
f"After controls (computed from beta): mu mean={mu_samples.mean():.4f}"
)
# Convert to original scale
logging.info(f"Final mu (standardized) mean={mu_samples.mean():.4f}")
y_samples_orig = mu_samples * y_std + y_mean
logging.info(f"Final y (original scale) mean={y_samples_orig.mean():.4f}")
return y_samples_orig
except Exception as e:
import traceback
logging.warning(f"Failed to reconstruct predictions: {e}")
logging.warning(traceback.format_exc())
return None
def _extract_aggregated_decomposition(
self,
bundle: MMMDataBundle,
) -> MMMDataBundle:
"""
Extract properly aggregated decomposition for multi-geo models.
Ensures bundle.component_time_series and bundle.component_totals
are aggregated to period-level.
"""
import pandas as pd
if not hasattr(self, "panel") or self.panel is None:
return bundle
# Check if we have multi-dimensional data
geo_idx = self._get_geo_indices()
has_geo = geo_idx is not None and len(set(geo_idx)) > 1
if not has_geo:
return bundle
try:
unique_periods = self._get_unique_periods()
periods = self._get_period_labels_per_obs()
if unique_periods is None or periods is None:
return bundle
# Get decomposition components at observation level
components = self._get_decomposition_components_obs_level()
if components is None:
return bundle
bundle.component_time_series = {}
bundle.component_totals = {}
for comp_name, comp_values in components.items():
df = pd.DataFrame(
{
"period": periods,
"value": comp_values,
}
)
agg = df.groupby("period")["value"].sum()
agg = agg.reindex(unique_periods).fillna(0)
bundle.component_time_series[comp_name] = agg.values
bundle.component_totals[comp_name] = float(agg.sum())
except Exception as e:
import logging
import traceback
logging.warning(f"Failed to extract aggregated decomposition: {e}")
logging.warning(traceback.format_exc())
return bundle
# -------------------------------------------------------------------------
# Geo-level extraction
# -------------------------------------------------------------------------
def _extract_geo_level_fit_data(self, bundle: MMMDataBundle) -> MMMDataBundle:
"""
Extract geo-level model fit data.
Aggregation strategy:
- For each geo: sum over products (if any) for each period
- Result is a time series of length n_periods for each geo
IMPORTANT: Uncertainty must be computed from aggregated samples per geo.
Populates:
- bundle.actual_by_geo
- bundle.predicted_by_geo
- bundle.fit_statistics_by_geo
"""
import pandas as pd
# Check if we have geo-level data
if not hasattr(self, "panel") or self.panel is None:
return bundle
# Get geo info
geo_names = self._get_geo_names()
if geo_names is None or len(geo_names) <= 1:
return bundle
bundle.geo_names = geo_names
try:
# Get period labels (unique periods in order)
unique_periods = self._get_unique_periods()
if unique_periods is None:
return bundle
# Get period index for each observation
periods = self._get_period_labels_per_obs()
if periods is None:
return bundle
# Get geo index for each observation
geo_idx = self._get_geo_indices()
if geo_idx is None:
return bundle
# Map geo indices to names
geo_labels = [geo_names[int(g)] for g in geo_idx]
# Get observed values (original scale)
y_obs = self._get_actual_original_scale()
if y_obs is None:
return bundle
# Build DataFrame for observed data aggregation
df_obs = pd.DataFrame(
{
"period": periods,
"geo": geo_labels,
"actual": y_obs,
}
)
# Initialize geo-level storage
bundle.actual_by_geo = {}
bundle.predicted_by_geo = {}
bundle.fit_statistics_by_geo = {}
# Get posterior samples for proper uncertainty propagation
y_samples_orig = self._get_posterior_samples_original_scale()
# Aggregate observed and predicted by geo
for geo in geo_names:
geo_df = df_obs[df_obs["geo"] == geo]
if len(geo_df) == 0:
continue
# Aggregate observed values for this geo
agg_obs = geo_df.groupby("period")["actual"].sum()
agg_obs = agg_obs.reindex(unique_periods).fillna(0)
y_obs_geo = agg_obs.values
bundle.actual_by_geo[geo] = y_obs_geo
# Aggregate predictions with proper uncertainty for this geo
if y_samples_orig is not None:
geo_mask = np.array(geo_labels) == geo
pred_geo = self._aggregate_samples_by_period(
y_samples_orig[:, geo_mask],
[p for p, g in zip(periods, geo_labels) if g == geo],
unique_periods,
)
if pred_geo is not None:
bundle.predicted_by_geo[geo] = pred_geo
bundle.fit_statistics_by_geo[geo] = (
self._compute_fit_statistics(y_obs_geo, pred_geo)
)
except Exception as e:
import logging
import traceback
logging.warning(f"Failed to extract geo-level fit data: {e}")
logging.warning(traceback.format_exc())
return bundle
def _get_posterior_samples_original_scale(self) -> np.ndarray | None:
"""Get posterior samples in original scale, shape (n_samples, n_obs)."""
try:
# Method 1: Use model's predict() method (most reliable)
if hasattr(self.mmm, "predict"):
try:
pred_results = self.mmm.predict(
return_original_scale=True, hdi_prob=self.ci_prob
)
return pred_results.y_pred_samples # (n_samples, n_obs)
except Exception as e:
logger.debug(f"posterior-predictive sampling failed: {e}")
# Method 2: Reconstruct from trace components
y_samples_orig = self._reconstruct_predictions_from_trace()
if y_samples_orig is not None:
return y_samples_orig
# Method 3: Try to get from trace directly (legacy)
trace = getattr(self.mmm, "_trace", None)
if trace is None:
return None
y_samples = None
if hasattr(trace, "posterior_predictive"):
pp = trace.posterior_predictive
for var_name in ["y", "y_obs", "likelihood"]:
if var_name in pp:
y_samples = pp[var_name].values
break
if y_samples is None:
return None
# Reshape to (n_samples, n_obs)
n_chains, n_draws = y_samples.shape[:2]
n_obs = y_samples.shape[-1]
y_samples = y_samples.reshape(n_chains * n_draws, n_obs)
# Transform to original scale
y_mean = self.mmm.y_mean if hasattr(self.mmm, "y_mean") else 0
y_std = self.mmm.y_std if hasattr(self.mmm, "y_std") else 1
return y_samples * y_std + y_mean
except Exception:
return None
def _extract_geo_level_decomposition(
self,
bundle: MMMDataBundle,
) -> MMMDataBundle:
"""
Extract geo-level decomposition data.
Aggregation strategy:
- For each geo and component: sum over products (if any) for each period
- Result is a time series of length n_periods for each geo-component pair
Populates:
- bundle.component_time_series_by_geo
- bundle.component_totals_by_geo
"""
import pandas as pd
if bundle.geo_names is None or len(bundle.geo_names) <= 1:
return bundle
try:
geo_names = bundle.geo_names
# Get period and geo info for each observation
unique_periods = self._get_unique_periods()
periods = self._get_period_labels_per_obs()
geo_idx = self._get_geo_indices()
if periods is None or geo_idx is None or unique_periods is None:
return bundle
geo_labels = [geo_names[int(g)] for g in geo_idx]
# Get decomposition components at observation level
components = self._get_decomposition_components_obs_level()
if components is None:
return bundle
bundle.component_time_series_by_geo = {}
bundle.component_totals_by_geo = {}
for geo in geo_names:
bundle.component_time_series_by_geo[geo] = {}
bundle.component_totals_by_geo[geo] = {}
for comp_name, comp_values in components.items():
# Build DataFrame for this component
df = pd.DataFrame(
{
"period": periods,
"geo": geo_labels,
"value": comp_values,
}
)
# Filter to this geo and aggregate by period
geo_df = df[df["geo"] == geo]
if len(geo_df) == 0:
continue
agg = geo_df.groupby("period")["value"].sum()
agg = agg.reindex(unique_periods).fillna(0)
comp_geo = agg.values
bundle.component_time_series_by_geo[geo][comp_name] = comp_geo
bundle.component_totals_by_geo[geo][comp_name] = float(
comp_geo.sum()
)
except Exception as e:
import logging
import traceback
logging.warning(f"Failed to extract geo-level decomposition: {e}")
logging.warning(traceback.format_exc())
return bundle
# -------------------------------------------------------------------------
# Product-level extraction
# -------------------------------------------------------------------------
def _extract_product_level_fit_data(
self,
bundle: MMMDataBundle,
) -> MMMDataBundle:
"""
Extract product-level model fit data.
Populates:
- bundle.product_names
- bundle.actual_by_product
- bundle.predicted_by_product
- bundle.fit_statistics_by_product
"""
# Check if we have product-level data
if not hasattr(self, "panel") or self.panel is None:
return bundle
mmm = self.mmm
# Get product info
product_names = self._get_product_names()
if product_names is None or len(product_names) <= 1:
return bundle
bundle.product_names = product_names
try:
# Get indices
product_idx = self._get_product_indices()
time_idx = self._get_time_indices()
if product_idx is None or time_idx is None:
return bundle
# Get observed values (original scale)
y_obs = self._get_actual_original_scale()
if y_obs is None:
return bundle
# Get predictions
y_pred_mean, y_pred_lower, y_pred_upper = (
self._get_predictions_original_scale()
)
if y_pred_mean is None:
return bundle
# Get period info
n_periods = (
mmm.n_periods if hasattr(mmm, "n_periods") else len(bundle.dates)
)
# Initialize product-level storage
bundle.actual_by_product = {}
bundle.predicted_by_product = {}
bundle.fit_statistics_by_product = {}
# Aggregate by product (sum over geos within each period)
for p_idx, product in enumerate(product_names):
# Get mask for this product
product_mask = product_idx == p_idx
# Aggregate observed values for this product over time
y_obs_prod = self._aggregate_by_period_with_indices(
y_obs, time_idx, product_mask, n_periods
)
y_pred_mean_prod = self._aggregate_by_period_with_indices(
y_pred_mean, time_idx, product_mask, n_periods
)
y_pred_lower_prod = self._aggregate_by_period_with_indices(
y_pred_lower, time_idx, product_mask, n_periods
)
y_pred_upper_prod = self._aggregate_by_period_with_indices(
y_pred_upper, time_idx, product_mask, n_periods
)
bundle.actual_by_product[product] = y_obs_prod
bundle.predicted_by_product[product] = {
"mean": y_pred_mean_prod,
"lower": y_pred_lower_prod,
"upper": y_pred_upper_prod,
}
# Compute fit statistics for this product
bundle.fit_statistics_by_product[product] = (
self._compute_fit_statistics(
y_obs_prod,
{
"mean": y_pred_mean_prod,
"lower": y_pred_lower_prod,
"upper": y_pred_upper_prod,
},
)
)
except Exception as e:
import logging
logging.warning(f"Failed to extract product-level fit data: {e}")
return bundle
def _extract_product_level_decomposition(
self,
bundle: MMMDataBundle,
) -> MMMDataBundle:
"""
Extract product-level decomposition data.
Populates:
- bundle.component_time_series_by_product
- bundle.component_totals_by_product
"""
if bundle.product_names is None or len(bundle.product_names) <= 1:
return bundle
try:
mmm = self.mmm
product_names = bundle.product_names
product_idx = self._get_product_indices()
time_idx = self._get_time_indices()
if product_idx is None or time_idx is None:
return bundle
n_periods = (
mmm.n_periods if hasattr(mmm, "n_periods") else len(bundle.dates)
)
# Get decomposition components at observation level
components = self._get_decomposition_components_obs_level()
if components is None:
return bundle
bundle.component_time_series_by_product = {}
bundle.component_totals_by_product = {}
for p_idx, product in enumerate(product_names):
product_mask = product_idx == p_idx
bundle.component_time_series_by_product[product] = {}
bundle.component_totals_by_product[product] = {}
for comp_name, comp_values in components.items():
# Aggregate this component for this product over time
comp_prod = self._aggregate_by_period_with_indices(
comp_values, time_idx, product_mask, n_periods
)
bundle.component_time_series_by_product[product][
comp_name
] = comp_prod
bundle.component_totals_by_product[product][comp_name] = float(
comp_prod.sum()
)
except Exception as e:
import logging
logging.warning(f"Failed to extract product-level decomposition: {e}")
return bundle
# -------------------------------------------------------------------------
# Scale transformation helpers
# -------------------------------------------------------------------------
def _get_actual_original_scale(self) -> np.ndarray | None:
"""Get observed values in original scale."""
if self.panel is None:
return None
y_standardized = self.panel.y.values.flatten()
y_mean = self.mmm.y_mean if hasattr(self.mmm, "y_mean") else 0
y_std = self.mmm.y_std if hasattr(self.mmm, "y_std") else 1
return y_standardized * y_std + y_mean
def _get_predictions_original_scale(
self,
) -> tuple[np.ndarray, np.ndarray, np.ndarray] | tuple[None, None, None]:
"""Get predictions in original scale with uncertainty."""
trace = getattr(self.mmm, "_trace", None)
if trace is None:
return None, None, None
try:
if (
hasattr(trace, "posterior_predictive")
and "y" in trace.posterior_predictive
):
y_samples = trace.posterior_predictive["y"].values
y_samples = y_samples.reshape(-1, y_samples.shape[-1])
elif hasattr(trace, "posterior") and "mu" in trace.posterior:
y_samples = trace.posterior["mu"].values
y_samples = y_samples.reshape(-1, y_samples.shape[-1])
else:
return None, None, None
# Transform to original scale
y_mean = self.mmm.y_mean if hasattr(self.mmm, "y_mean") else 0
y_std = self.mmm.y_std if hasattr(self.mmm, "y_std") else 1
y_samples_orig = y_samples * y_std + y_mean
y_pred_mean = y_samples_orig.mean(axis=0)
alpha = (1 - self.ci_prob) / 2
y_pred_lower = np.percentile(y_samples_orig, alpha * 100, axis=0)
y_pred_upper = np.percentile(y_samples_orig, (1 - alpha) * 100, axis=0)
return y_pred_mean, y_pred_lower, y_pred_upper
except Exception:
return None, None, None
# -------------------------------------------------------------------------
# Index extraction helpers (from GeoExtractionMixin and ProductExtractionMixin)
# -------------------------------------------------------------------------
def _get_geo_names(self) -> list[str] | None:
"""Get geography names from panel or model."""
if hasattr(self.panel, "coords") and hasattr(self.panel.coords, "geographies"):
geographies = self.panel.coords.geographies
if geographies is not None:
return list(geographies)
if hasattr(self.mmm, "geo_names"):
geo_names = self.mmm.geo_names
if geo_names is not None:
return list(geo_names)
return None
def _get_geo_indices(self) -> np.ndarray | None:
"""Get geo index for each observation."""
if hasattr(self.mmm, "geo_idx") and self.mmm.geo_idx is not None:
return np.array(self.mmm.geo_idx)
if hasattr(self.panel, "geo_idx") and self.panel.geo_idx is not None:
return np.array(self.panel.geo_idx)
# Try to extract from panel's MultiIndex
if self.panel is not None:
geo_names = self._get_geo_names()
if geo_names is not None and len(geo_names) > 1:
# Try to get from y.index if it's a MultiIndex
if hasattr(self.panel, "y") and hasattr(self.panel.y, "index"):
idx = self.panel.y.index
if hasattr(idx, "get_level_values"):
# Try both cases for geography level name
for level_name in ["geography", "Geography"]:
try:
geo_values = idx.get_level_values(level_name)
# Convert geo names to indices
geo_to_idx = {g: i for i, g in enumerate(geo_names)}
return np.array(
[geo_to_idx.get(str(g), 0) for g in geo_values]
)
except KeyError:
continue
return None
def _get_product_names(self) -> list[str] | None:
"""Get product names from panel or model."""
if hasattr(self.panel, "coords") and hasattr(self.panel.coords, "products"):
products = self.panel.coords.products
if products is not None:
return list(products)
if hasattr(self.mmm, "product_names"):
return list(self.mmm.product_names)
return None
def _get_product_indices(self) -> np.ndarray | None:
"""Get product index for each observation."""
if hasattr(self.mmm, "product_idx"):
return np.array(self.mmm.product_idx)
if hasattr(self.panel, "product_idx"):
return np.array(self.panel.product_idx)
return None
def _get_time_indices(self) -> np.ndarray | None:
"""Get time index for each observation."""
if hasattr(self.mmm, "time_idx") and self.mmm.time_idx is not None:
return np.array(self.mmm.time_idx)
# Try to extract from panel's MultiIndex
if self.panel is not None:
unique_periods = self._get_unique_periods()
if unique_periods is not None and len(unique_periods) > 0:
if hasattr(self.panel, "y") and hasattr(self.panel.y, "index"):
idx = self.panel.y.index
if hasattr(idx, "get_level_values"):
# Try both cases for period level name
for level_name in ["period", "Period"]:
try:
period_values = idx.get_level_values(level_name)
# Convert periods to indices
period_to_idx = {
str(p): i for i, p in enumerate(unique_periods)
}
return np.array(
[
period_to_idx.get(str(p), 0)
for p in period_values
]
)
except KeyError:
continue
# Try first level if period not found by name
try:
period_values = idx.get_level_values(0)
period_to_idx = {
str(p): i for i, p in enumerate(unique_periods)
}
return np.array(
[period_to_idx.get(str(p), 0) for p in period_values]
)
except Exception as e:
logger.debug(f"period index mapping failed: {e}")
else:
# Simple index - assume it's period only
period_to_idx = {
str(p): i for i, p in enumerate(unique_periods)
}
return np.array([period_to_idx.get(str(p), 0) for p in idx])
# Fallback: construct from panel coords
if self.panel is not None and hasattr(self.panel, "coords"):
n_obs = len(self.panel.y)
n_periods = self.panel.coords.n_periods
n_geos = self.panel.coords.n_geos
n_products = self.panel.coords.n_products
# Assumes data ordered as periods innermost
if n_obs == n_periods * n_geos * n_products:
return np.tile(np.arange(n_periods), n_geos * n_products)
return None
def _aggregate_by_period(
self,
values: np.ndarray,
n_periods: int,
geo_mask: np.ndarray,
) -> np.ndarray:
"""
Aggregate values by period for a specific geo.
If there are multiple products, this sums over products within each period.
"""
# Simple case: one observation per period per geo
n_obs = len(values)
if n_obs == n_periods:
return values
# Multiple products case: need to reshape and sum
n_products = n_obs // n_periods
if n_products * n_periods == n_obs:
return values.reshape(n_products, n_periods).sum(axis=0)
# Fallback: return as-is (may need custom handling)
return values
def _aggregate_by_period_with_indices(
self,
values: np.ndarray,
time_idx: np.ndarray,
dim_mask: np.ndarray,
n_periods: int,
) -> np.ndarray:
"""
Aggregate values by period using explicit time indices.
This is a more robust aggregation method that uses the time_idx array
to properly group observations by period, regardless of data ordering.
"""
result = np.zeros(n_periods)
filtered_values = values[dim_mask]
filtered_time_idx = time_idx[dim_mask]
for t in range(n_periods):
time_mask = filtered_time_idx == t
if time_mask.any():
result[t] = filtered_values[time_mask].sum()
return result
# -------------------------------------------------------------------------
# Decomposition extraction
# -------------------------------------------------------------------------
def _get_decomposition_components_obs_level(self) -> dict[str, np.ndarray] | None:
"""
Get decomposition components at observation level (not yet aggregated).
Returns dict mapping component name to array of shape (n_obs,).
"""
trace = getattr(self.mmm, "_trace", None)
if trace is None:
return None
components = {}
posterior = trace.posterior
y_mean = self.mmm.y_mean if hasattr(self.mmm, "y_mean") else 0
y_std = self.mmm.y_std if hasattr(self.mmm, "y_std") else 1
n_obs = len(self.panel.y.values.flatten())
try:
# Intercept/Baseline
if "intercept" in posterior:
intercept = float(posterior["intercept"].values.mean())
# Baseline in original scale includes both the intercept contribution
# and the mean offset (y_mean) to ensure decomposition sums to predictions
baseline_per_obs = intercept * y_std + y_mean
components["Baseline"] = np.full(n_obs, baseline_per_obs)
# Trend
if "trend" in posterior:
trend_samples = posterior["trend"].values
trend_mean = trend_samples.mean(axis=(0, 1))
if len(trend_mean) == n_obs:
components["Trend"] = trend_mean * y_std
# Seasonality
if "seasonality" in posterior:
seas_samples = posterior["seasonality"].values
seas_mean = seas_samples.mean(axis=(0, 1))
if len(seas_mean) == n_obs:
components["Seasonality"] = seas_mean * y_std
# Media channels
channel_names = self._get_channel_names()
# First try the channel_contributions array (most common for geo models)
if "channel_contributions" in posterior:
contrib_da = posterior["channel_contributions"]
contrib_vals = contrib_da.values # (chains, draws, obs, channels)
contrib_mean = contrib_vals.mean(axis=(0, 1)) # (obs, channels)
# Check shape and extract per-channel
if contrib_mean.ndim == 2:
if contrib_mean.shape[0] == n_obs and contrib_mean.shape[1] == len(
channel_names
):
# Shape is (obs, channels)
for i, ch in enumerate(channel_names):
components[ch] = contrib_mean[:, i] * y_std
elif contrib_mean.shape[1] == n_obs and contrib_mean.shape[
0
] == len(channel_names):
# Shape is (channels, obs)
for i, ch in enumerate(channel_names):
components[ch] = contrib_mean[i, :] * y_std
else:
# Fall back to individual channel contribution variables
for ch in channel_names:
contrib_key = f"channel_contribution_{ch}"
if contrib_key in posterior:
contrib = posterior[contrib_key].values.mean(axis=(0, 1))
if len(contrib) == n_obs:
components[ch] = contrib * y_std
# Controls
if hasattr(self.mmm, "control_names"):
for ctrl in self.mmm.control_names:
ctrl_key = f"control_contribution_{ctrl}"
if ctrl_key in posterior:
ctrl_contrib = posterior[ctrl_key].values.mean(axis=(0, 1))
if len(ctrl_contrib) == n_obs:
components[f"Control: {ctrl}"] = ctrl_contrib * y_std
return components if components else None
except Exception:
return None
# -------------------------------------------------------------------------
# ROI and contribution extraction
# -------------------------------------------------------------------------
def _compute_channel_roi(self) -> dict[str, dict[str, float]] | None:
"""Compute channel ROI with uncertainty."""
try:
from mmm_framework.reporting.helpers.measurement import (
resolve_channel_divisor,
)
channels = self._get_channel_names()
if not channels:
return None
trace = getattr(self.mmm, "_trace", None)
if trace is None or not hasattr(trace, "posterior"):
return None
posterior = trace.posterior
y_std = getattr(self.mmm, "y_std", 1.0)
n_obs = getattr(self.mmm, "n_obs", 52)
roi_results = {}
for ch in channels:
# Measurement-aware divisor: dollar spend for spend / cpm / cpc /
# spend_column channels (-> ROI), or per-1,000-impression volume
# for cost-less impression channels (-> efficiency).
resolved = resolve_channel_divisor(self.mmm, ch)
spend = resolved.total
meta = resolved.meta
if not resolved.found or spend <= 0:
continue
# Try to get contribution samples
contrib_samples = None
# Method 1: Direct contribution variable
contrib_names = [
f"contribution_{ch}",
f"channel_contribution_{ch}",
]
for contrib_name in contrib_names:
if contrib_name in posterior:
vals = posterior[contrib_name].values
# Flatten chains and draws, sum over time if needed
if vals.ndim > 2:
contrib_samples = (
vals.reshape(-1, *vals.shape[2:]).sum(axis=-1) * y_std
)
else:
contrib_samples = vals.flatten() * y_std * n_obs
break
# Method 2: From channel_contributions array
if contrib_samples is None and "channel_contributions" in posterior:
try:
ch_idx = channels.index(ch)
contrib_da = posterior["channel_contributions"]
vals = contrib_da.values # numpy first!
# Flatten chains/draws, extract channel, sum over time
flat = vals.reshape(-1, *vals.shape[2:])
if flat.ndim == 2: # (samples, time)
contrib_samples = flat.sum(axis=-1) * y_std
elif flat.ndim == 3: # (samples, time, channels)
contrib_samples = flat[:, :, ch_idx].sum(axis=-1) * y_std
except Exception as e:
logger.debug(
f"Could not extract from channel_contributions: {e}"
)
# Method 3: Estimate from beta
if contrib_samples is None:
for beta_name in [f"beta_{ch}", f"beta_media_{ch}"]:
if beta_name in posterior:
beta_vals = posterior[beta_name].values.flatten()
# Rough estimate
contrib_samples = beta_vals * y_std * n_obs * 0.5
break
if contrib_samples is not None and len(contrib_samples) > 0:
roi_samples = contrib_samples / spend
roi_results[ch] = {
"mean": float(np.mean(roi_samples)),
"lower": float(
np.percentile(roi_samples, (1 - self.ci_prob) / 2 * 100)
),
"upper": float(
np.percentile(roi_samples, (1 + self.ci_prob) / 2 * 100)
),
# Measurement metadata so the section/chart render ROI vs
# efficiency (and the right break-even reference).
"reference": meta.reference,
"metric_label": meta.roi_label,
"value_units": meta.value_units,
"is_monetary": meta.is_monetary,
"measurement_unit": meta.unit.value,
"cost_basis": meta.cost_basis,
}
return roi_results if roi_results else None
except Exception as e:
logger.warning(f"Error computing channel ROI: {e}")
import traceback
logger.debug(traceback.format_exc())
return None
def _components_from_decomposition(self) -> dict[str, np.ndarray] | None:
"""Build component time series (original KPI scale) from the model's
canonical decomposition, which handles inverse-scaling correctly
(the intercept carries the ``y_mean`` location shift)."""
if not hasattr(self.mmm, "compute_component_decomposition"):
return None
try:
decomp = self.mmm.compute_component_decomposition()
except Exception as e:
logger.debug(f"compute_component_decomposition unavailable: {e}")
return None
components: dict[str, np.ndarray] = {
"Baseline": np.asarray(decomp.intercept, dtype=float)
}
if np.any(decomp.trend):
components["Trend"] = np.asarray(decomp.trend, dtype=float)
if np.any(decomp.seasonality):
components["Seasonality"] = np.asarray(decomp.seasonality, dtype=float)
if decomp.events is not None and np.any(decomp.events):
# #143: holiday / event contributions, separate from Fourier seasonality.
components["Events"] = np.asarray(decomp.events, dtype=float)
if decomp.interactions is not None and np.any(decomp.interactions):
# #142: cross-channel synergy / interaction contributions.
components["Synergy"] = np.asarray(decomp.interactions, dtype=float)
if decomp.levers is not None and np.any(decomp.levers):
# #138: price elasticity + promo lift contributions.
components["Price & Promotion"] = np.asarray(decomp.levers, dtype=float)
if decomp.geo_effects is not None and np.any(decomp.geo_effects):
components["Geo"] = np.asarray(decomp.geo_effects, dtype=float)
if decomp.product_effects is not None and np.any(decomp.product_effects):
components["Product"] = np.asarray(decomp.product_effects, dtype=float)
for ch in decomp.media_by_channel.columns:
components[str(ch)] = decomp.media_by_channel[ch].to_numpy(dtype=float)
if decomp.controls_by_var is not None:
for ctrl in decomp.controls_by_var.columns:
components[f"Control: {ctrl}"] = decomp.controls_by_var[ctrl].to_numpy(
dtype=float
)
elif np.any(decomp.controls_total):
components["Controls"] = np.asarray(decomp.controls_total, dtype=float)
return components
def _get_component_totals(self) -> dict[str, float] | None:
"""Get total contribution by component (original KPI scale)."""
try:
logger.debug("Computing component totals")
# Preferred path: the model's canonical decomposition
components = self._components_from_decomposition()
if components:
return {name: float(vals.sum()) for name, vals in components.items()}
trace = getattr(self.mmm, "_trace", None)
if trace is None or not hasattr(trace, "posterior"):
return None
posterior = trace.posterior
y_mean = getattr(self.mmm, "y_mean", 0.0)
y_std = getattr(self.mmm, "y_std", 1.0)
n_obs = getattr(self.mmm, "n_obs", 52)
totals = {}
# Baseline/Intercept — carries the y_mean location shift so the
# decomposition sums to predictions in original units
if "intercept" in posterior:
intercept = (
float(posterior["intercept"].values.mean()) * y_std + y_mean
) * n_obs
totals["Baseline"] = intercept
# Media channels
channels = self._get_channel_names()
current_spend = self._get_current_spend()
for ch in channels:
# Try to get contribution from trace
contrib_names = [
f"contribution_{ch}",
f"channel_contribution_{ch}",
f"media_contribution_{ch}",
]
found = False
for contrib_name in contrib_names:
if contrib_name in posterior:
# Sum over time
contrib_vals = posterior[contrib_name].values
total_contrib = float(contrib_vals.mean()) * y_std
if contrib_vals.ndim > 2:
# Has time dimension - sum it
total_contrib = (
float(contrib_vals.sum(axis=-1).mean()) * y_std
)
totals[ch] = total_contrib
found = True
break
if not found:
# Estimate from beta
for beta_name in [f"beta_{ch}", f"beta_media_{ch}"]:
if beta_name in posterior:
beta_val = float(posterior[beta_name].values.mean())
# Try to get media sum
if current_spend and ch in current_spend:
# Rough estimate
totals[ch] = (
beta_val * y_std * n_obs * 0.5
) # Approximate
else:
totals[ch] = beta_val * y_std * n_obs
break
return totals if totals else None
except Exception as e:
logger.warning(f"Error computing component totals: {e}")
import traceback
logger.debug(traceback.format_exc())
return None
def _get_component_time_series(self) -> dict[str, np.ndarray] | None:
"""Get component time series for decomposition chart (original KPI scale)."""
try:
logger.debug("Computing component time series")
# Preferred path: the model's canonical decomposition
components = self._components_from_decomposition()
if components:
return components
trace = getattr(self.mmm, "_trace", None)
if trace is None or not hasattr(trace, "posterior"):
return None
posterior = trace.posterior
y_mean = getattr(self.mmm, "y_mean", 0.0)
y_std = getattr(self.mmm, "y_std", 1.0)
n_obs = getattr(self.mmm, "n_obs", 52)
components = {}
# Baseline - constant; carries the y_mean location shift so the
# decomposition sums to predictions in original units
if "intercept" in posterior:
intercept = float(posterior["intercept"].values.mean()) * y_std + y_mean
components["Baseline"] = np.full(n_obs, intercept)
# Trend
trend_names = ["trend", "trend_contribution", "trend_component"]
for trend_name in trend_names:
if trend_name in posterior:
trend_vals = posterior[trend_name].values
# Average over chains and draws
trend_mean = trend_vals.mean(axis=(0, 1))
if len(trend_mean) == n_obs:
components["Trend"] = trend_mean * y_std
break
# Seasonality
seas_names = ["seasonality", "seasonality_component", "seasonal_effect"]
for seas_name in seas_names:
if seas_name in posterior:
seas_vals = posterior[seas_name].values
seas_mean = seas_vals.mean(axis=(0, 1))
if len(seas_mean) == n_obs:
components["Seasonality"] = seas_mean * y_std
break
# Channel contributions
channels = self._get_channel_names()
# Try channel_contributions array
if "channel_contributions" in posterior:
contrib_da = posterior["channel_contributions"]
contrib_vals = contrib_da.values # Get numpy first!
# Average over chains and draws
if contrib_vals.ndim >= 3:
contrib_mean = contrib_vals.mean(axis=(0, 1))
if contrib_mean.ndim == 2 and contrib_mean.shape[0] == n_obs:
# Shape is (time, channels)
for i, ch in enumerate(channels):
if i < contrib_mean.shape[1]:
components[ch] = contrib_mean[:, i] * y_std
elif contrib_mean.ndim == 2 and contrib_mean.shape[1] == n_obs:
# Shape is (channels, time)
for i, ch in enumerate(channels):
if i < contrib_mean.shape[0]:
components[ch] = contrib_mean[i, :] * y_std
else:
# Try individual contribution variables
for ch in channels:
contrib_names = [
f"contribution_{ch}",
f"channel_contribution_{ch}",
f"media_contribution_{ch}",
]
for contrib_name in contrib_names:
if contrib_name in posterior:
contrib_vals = posterior[contrib_name].values
contrib_mean = contrib_vals.mean(axis=(0, 1))
if len(contrib_mean) == n_obs:
components[ch] = contrib_mean * y_std
break
return components if components else None
except Exception as e:
logger.warning(f"Error computing component time series: {e}")
import traceback
logger.debug(traceback.format_exc())
return None
def _compute_marketing_attribution(self) -> dict[str, float] | None:
"""Compute total marketing-attributed revenue with uncertainty."""
try:
# First, try using the model's compute_contributions method
if hasattr(self.mmm, "compute_counterfactual_contributions"):
contrib_results = self.mmm.compute_counterfactual_contributions(
compute_uncertainty=True,
hdi_prob=self.ci_prob,
)
total = float(contrib_results.total_contributions.sum())
# Get uncertainty from HDI if available
if contrib_results.contribution_hdi_low is not None:
lower = float(contrib_results.contribution_hdi_low.sum())
upper = float(contrib_results.contribution_hdi_high.sum())
else:
# Rough estimate: +/- 15% of total
lower = total * 0.85
upper = total * 1.15
return {"mean": total, "lower": lower, "upper": upper}
logger.debug("Computing total marketing-attributed revenue")
trace = getattr(self.mmm, "_trace", None)
logger.debug("Accessing model trace for attribution computation")
if trace is None:
logger.debug("Model trace not found for attribution computation")
return None
channels = self._get_channel_names()
# Sum contributions across channels
total_samples = None
for ch in channels:
contrib_name = f"contribution_{ch}"
if hasattr(trace, "posterior") and contrib_name in trace.posterior:
logger.debug(f"Adding contribution samples for channel: {ch}")
ch_samples = (
trace.posterior[contrib_name].values.sum(axis=-1).flatten()
)
if total_samples is None:
logger.debug("Initializing total samples for attribution")
total_samples = ch_samples
else:
logger.debug("Accumulating total samples for attribution")
total_samples = total_samples + ch_samples
if total_samples is not None:
logger.debug("Computing mean and HDI for total marketing attribution")
mean = float(total_samples.mean())
lower, upper = self._compute_hdi(total_samples, self.ci_prob)
return {"mean": mean, "lower": lower, "upper": upper}
except Exception as e:
logger.warning(f"Error computing marketing attribution: {e}")
pass
return None
def _compute_blended_roi(self) -> dict[str, float] | None:
"""Compute blended marketing ROI with uncertainty.
A single blended ROI is only meaningful when every channel is measured
in dollars (spend, or impressions/clicks costed via cpm/cpc/spend_column
-> derived dollar spend). If any channel reports per-volume *efficiency*
(no cost), the units are non-comparable, so the blended rollup is
suppressed and only per-channel numbers are shown (the chosen policy for
mixed portfolios)."""
try:
from mmm_framework.reporting.helpers.measurement import (
resolve_channel_divisor,
)
attribution = self._compute_marketing_attribution()
if attribution is None:
return None
total_spend = 0.0
for ch in self._get_channel_names() or []:
resolved = resolve_channel_divisor(self.mmm, ch)
if not resolved.found:
continue
if not resolved.meta.is_monetary:
# Mixed/efficiency portfolio -- a blended $ ROI is meaningless.
return None
total_spend += resolved.total
if total_spend == 0:
return None
# Simple division (would need full posterior for proper uncertainty)
mean_roi = attribution["mean"] / total_spend
lower_roi = attribution["lower"] / total_spend
upper_roi = attribution["upper"] / total_spend
return {"mean": mean_roi, "lower": lower_roi, "upper": upper_roi}
except Exception:
return None
def _compute_marketing_contribution_pct(
self, total_revenue: float | None
) -> dict[str, float] | None:
"""Compute marketing contribution as percentage of total."""
if total_revenue is None or total_revenue == 0:
return None
attribution = self._compute_marketing_attribution()
if attribution is None:
return None
return {
"mean": attribution["mean"] / total_revenue,
"lower": attribution["lower"] / total_revenue,
"upper": attribution["upper"] / total_revenue,
}
# -------------------------------------------------------------------------
# Saturation and adstock extraction
# -------------------------------------------------------------------------
def _get_saturation_curves(self) -> dict[str, dict[str, np.ndarray]] | None:
"""Get saturation curve data for each channel."""
try:
# If model has a method for this, use it
if hasattr(self.mmm, "compute_saturation_curves"):
result = self.mmm.compute_saturation_curves()
if result:
return result
channels = self._get_channel_names()
if not channels:
return None
# Get current spend for setting spend range
current_spend = self._get_current_spend()
trace = getattr(self.mmm, "_trace", None)
if trace is None or not hasattr(trace, "posterior"):
logger.debug("No trace/posterior found for saturation curves")
return None
posterior = trace.posterior
curves = {}
for ch in channels:
# Determine spend range
max_spend = 1e6 # Default
if current_spend and ch in current_spend:
max_spend = current_spend[ch] * 2 # 2x current spend
spend_range = np.linspace(0, max_spend, 100)
# Try Hill saturation parameters (sat_half_/sat_slope_ are the
# core BayesianMMM's per-channel Hill RVs)
kappa_names = [
f"sat_half_{ch}",
f"kappa_{ch}",
f"K_{ch}",
f"sat_K_{ch}",
]
slope_names = [
f"sat_slope_{ch}",
f"slope_{ch}",
f"S_{ch}",
f"sat_S_{ch}",
f"n_{ch}",
]
kappa_val = None
slope_val = None
kappa_matched = None
for k_name in kappa_names:
if k_name in posterior:
kappa_val = float(posterior[k_name].values.mean())
kappa_matched = k_name
break
for s_name in slope_names:
if s_name in posterior:
slope_val = float(posterior[s_name].values.mean())
break
def _beta_scaled(response: np.ndarray) -> np.ndarray:
for beta_name in [f"beta_{ch}", f"beta_media_{ch}"]:
if beta_name in posterior:
beta_val = float(posterior[beta_name].values.mean())
return response * beta_val
return response
if kappa_val is not None and slope_val is not None:
# The core model's sat_half_ is on the normalized (~[0, 1])
# spend scale, so evaluate the curve on normalized spend
# (mirroring the logistic branch below). Other naming
# conventions keep the raw-spend evaluation.
if kappa_matched is not None and kappa_matched.startswith(
"sat_half_"
):
x_grid = spend_range / max(spend_range.max(), 1)
else:
x_grid = spend_range
x_grid = np.clip(x_grid, 1e-9, None)
# Hill function
response = x_grid**slope_val / (
kappa_val**slope_val + x_grid**slope_val
)
curves[ch] = {
"spend": spend_range,
"response": _beta_scaled(response),
}
logger.debug(f"Generated Hill saturation curve for {ch}")
continue
if kappa_val is not None and slope_val is None:
# sat_half_ without a slope RV: the core model's
# michaelis_menten or tanh saturation. The two are
# indistinguishable from RV names alone, so consult the
# channel's configured saturation type.
sat_type = None
try:
sat_type = self.mmm._get_saturation_config(ch).type.value
except Exception:
sat_type = None
x_grid = spend_range / max(spend_range.max(), 1)
if sat_type == "michaelis_menten":
response = x_grid / (x_grid + kappa_val)
elif sat_type == "tanh":
response = np.tanh(x_grid / kappa_val)
else:
logger.warning(
f"Saturation curve for {ch}: found {kappa_matched} but "
"cannot determine the saturation form; skipping."
)
continue
curves[ch] = {
"spend": spend_range,
"response": _beta_scaled(response),
}
logger.debug(f"Generated {sat_type} saturation curve for {ch}")
continue
# Root / power saturation: x**k on normalized spend.
for exp_name in [f"sat_exponent_{ch}", f"saturation_exponent_{ch}"]:
if exp_name in posterior:
k_val = float(posterior[exp_name].values.mean())
x_grid = np.clip(
spend_range / max(spend_range.max(), 1), 1e-9, None
)
curves[ch] = {
"spend": spend_range,
"response": _beta_scaled(x_grid**k_val),
}
logger.debug(f"Generated root saturation curve for {ch}")
break
if ch in curves:
continue
# Logistic saturation (the core default): 1 - exp(-lam * x)
for lam_name in [f"sat_lam_{ch}", f"saturation_lam_{ch}", f"lam_{ch}"]:
if lam_name in posterior:
lam_val = float(posterior[lam_name].values.mean())
response = 1 - np.exp(
-lam_val * spend_range / max(spend_range.max(), 1)
)
curves[ch] = {
"spend": spend_range,
"response": _beta_scaled(response),
}
logger.debug(f"Generated logistic saturation curve for {ch}")
break
return curves if curves else None
except Exception as e:
logger.warning(f"Error getting saturation curves: {e}")
import traceback
logger.debug(traceback.format_exc())
return None
def _get_adstock_curves(self) -> dict[str, np.ndarray] | None:
"""Get adstock decay weights for each channel."""
try:
if hasattr(self.mmm, "compute_adstock_curves"):
result = self.mmm.compute_adstock_curves()
if result:
return result
channels = self._get_channel_names()
if not channels:
return None
trace = getattr(self.mmm, "_trace", None)
if trace is None or not hasattr(trace, "posterior"):
return None
posterior = trace.posterior
curves = {}
from ...transforms.adstock import adstock_weights
for ch in channels:
# Default l_max from the channel's adstock config if resolvable
l_max = getattr(self.mmm, "adstock_lmax", 8)
if hasattr(self.mmm, "model_config"):
l_max = getattr(self.mmm.model_config, "adstock_lmax", l_max)
try:
l_max = self.mmm._get_adstock_config(ch).l_max
except Exception as e:
logger.debug(f"adstock l_max lookup failed for {ch}: {e}")
def _mean(name: str) -> float | None:
if name in posterior:
return float(posterior[name].values.mean())
return None
# Weibull kernel (shape + scale)
shape_val = _mean(f"adstock_shape_{ch}")
scale_val = _mean(f"adstock_scale_{ch}")
if shape_val is not None and scale_val is not None:
curves[ch] = adstock_weights(
"weibull", l_max, shape=shape_val, scale=scale_val
)
logger.debug(
f"Generated weibull adstock curve for {ch}: "
f"shape={shape_val:.3f} scale={scale_val:.3f}"
)
continue
# Geometric / delayed kernel (alpha [+ theta])
alpha_val = None
for alpha_name in [
f"adstock_alpha_{ch}",
f"alpha_{ch}",
f"adstock_{ch}",
f"decay_{ch}",
]:
alpha_val = _mean(alpha_name)
if alpha_val is not None:
break
if alpha_val is not None and 0 < alpha_val < 1:
theta_val = _mean(f"adstock_theta_{ch}")
if theta_val is not None:
curves[ch] = adstock_weights(
"delayed", l_max, alpha=alpha_val, theta=theta_val
)
logger.debug(
f"Generated delayed adstock curve for {ch}: "
f"alpha={alpha_val:.3f} theta={theta_val:.2f}"
)
else:
curves[ch] = adstock_weights(
"geometric", l_max, alpha=alpha_val
)
logger.debug(
f"Generated adstock curve for {ch}: alpha={alpha_val:.3f}"
)
return curves if curves else None
except Exception as e:
logger.warning(f"Error getting adstock curves: {e}")
return None
def _get_current_spend(self) -> dict[str, float] | None:
"""Get current spend levels by channel."""
try:
channels = self._get_channel_names()
if not channels:
logger.debug("No channel names found for spend extraction")
return None
# Try to get X_media from panel
X_media = None
if self.panel is not None and hasattr(self.panel, "X_media"):
X_media = self.panel.X_media
elif hasattr(self.mmm, "X_media_raw"):
X_media = self.mmm.X_media_raw
elif hasattr(self.mmm, "X_media"):
X_media = self.mmm.X_media
if X_media is None:
logger.debug("No X_media found for spend extraction")
return None
spend = {}
for i, ch in enumerate(channels):
col_data = _safe_get_column(X_media, i, ch)
if col_data is not None:
spend[ch] = float(col_data.sum())
else:
logger.debug(f"Could not extract spend for channel {ch}")
return spend if spend else None
except Exception as e:
logger.warning(f"Error getting current spend: {e}")
return None
# -------------------------------------------------------------------------
# Diagnostic data extraction
# -------------------------------------------------------------------------
def _get_trace_data(
self,
) -> tuple[dict[str, np.ndarray] | None, list[str] | None]:
"""Get trace data for diagnostic plots."""
try:
trace = getattr(self.mmm, "_trace", None)
if trace is None or not hasattr(trace, "posterior"):
return None, None
# Get key parameters
params = []
data = {}
for var_name in trace.posterior.data_vars:
if any(
prefix in var_name
for prefix in ["beta", "alpha", "sigma", "intercept"]
):
values = trace.posterior[var_name].values
if values.ndim >= 2:
# Shape: (chain, draw, ...)
data[var_name] = values.reshape(values.shape[0], -1)[
:, : values.shape[1]
]
params.append(var_name)
return data, params
except Exception:
return None, None
def _get_prior_posterior(
self,
) -> tuple[dict[str, np.ndarray] | None, dict[str, np.ndarray] | None]:
"""Get prior and posterior samples for comparison."""
try:
trace = getattr(self.mmm, "_trace", None)
if trace is None:
logger.debug("No trace found for prior/posterior extraction")
return None, None
posterior_samples = {}
prior_samples = {}
# Get posterior samples
if hasattr(trace, "posterior"):
# Select key parameters to include
key_prefixes = [
"beta",
"sigma",
"intercept",
"alpha",
"adstock",
"sat_lam",
"sat_half",
"kappa",
"slope",
]
for var_name in trace.posterior.data_vars:
# Only include key parameters (not all deterministics)
if any(prefix in var_name for prefix in key_prefixes):
try:
samples = trace.posterior[var_name].values.flatten()
# Subsample if too many
if len(samples) > 2000:
idx = np.random.choice(
len(samples), 2000, replace=False
)
samples = samples[idx]
posterior_samples[var_name] = samples
except Exception as e:
logger.debug(
f"Could not extract posterior for {var_name}: {e}"
)
if not posterior_samples:
logger.debug("No posterior samples extracted")
return None, None
# Try to get prior samples from trace
if hasattr(trace, "prior"):
for var_name in trace.prior.data_vars:
if var_name in posterior_samples:
try:
samples = trace.prior[var_name].values.flatten()
if len(samples) > 2000:
idx = np.random.choice(
len(samples), 2000, replace=False
)
samples = samples[idx]
prior_samples[var_name] = samples
except Exception as e:
logger.debug(f"Could not extract prior for {var_name}: {e}")
# If no prior samples in trace, try to generate from model
if not prior_samples:
prior_samples = self._generate_prior_samples(posterior_samples.keys())
logger.debug(
f"Extracted {len(posterior_samples)} posterior params, {len(prior_samples)} prior params"
)
return (
prior_samples if prior_samples else None,
posterior_samples if posterior_samples else None,
)
except Exception as e:
logger.warning(f"Error in _get_prior_posterior: {e}")
import traceback
logger.debug(traceback.format_exc())
return None, None
def _generate_prior_samples(
self, param_names: list[str], n_samples: int = 1000
) -> dict[str, np.ndarray]:
"""
Generate prior samples from model specification.
This is used when sample_prior_predictive wasn't called during fitting.
"""
prior_samples = {}
try:
model = getattr(self.mmm, "model", None) or getattr(
self.mmm, "_model", None
)
if model is None:
logger.debug("No PyMC model found for prior sampling")
return {}
# Sample from prior
from ...utils import arviz_compat
with model:
prior_trace = arviz_compat.sample_prior_predictive(n_samples, 42)
if hasattr(prior_trace, "prior"):
prior_data = prior_trace.prior
for param in param_names:
if param in prior_data:
try:
samples = prior_data[param].values.flatten()
if len(samples) > n_samples:
idx = np.random.choice(
len(samples), n_samples, replace=False
)
samples = samples[idx]
prior_samples[param] = samples
except Exception as e:
logger.debug(
f"prior sample extraction failed for {param}: {e}"
)
logger.debug(f"Generated {len(prior_samples)} prior sample sets")
except Exception as e:
logger.debug(f"Could not generate prior samples: {e}")
# Fall back to generating from known prior distributions
prior_samples = self._generate_fallback_priors(param_names, n_samples)
return prior_samples
def _generate_fallback_priors(
self, param_names: list[str], n_samples: int = 1000
) -> dict[str, np.ndarray]:
"""
Generate fallback prior samples based on common MMM prior conventions.
Used when we can't sample from the actual model.
"""
prior_samples = {}
for param in param_names:
try:
if param.startswith("beta_") or param.startswith("beta_media_"):
# Beta coefficients: HalfNormal(sigma=2)
prior_samples[param] = np.abs(np.random.normal(0, 2, n_samples))
elif param == "sigma":
# Noise scale: HalfNormal(sigma=1)
prior_samples[param] = np.abs(np.random.normal(0, 1, n_samples))
elif param == "intercept":
# Intercept: Normal(0, 5)
prior_samples[param] = np.random.normal(0, 5, n_samples)
elif param.startswith("alpha_") or param.startswith("adstock_"):
# Adstock decay: Beta(1, 3) - favors lower values
prior_samples[param] = np.random.beta(1, 3, n_samples)
elif param.startswith("sat_half_"):
# Core-model half-saturation: Beta(2, 2)
prior_samples[param] = np.random.beta(2, 2, n_samples)
elif param.startswith("sat_slope_"):
# Core-model Hill slope: HalfNormal(sigma=1.5)
prior_samples[param] = np.abs(np.random.normal(0, 1.5, n_samples))
elif param.startswith("sat_lam_") or param.startswith("saturation_"):
# Saturation rate: Gamma(2, 1)
prior_samples[param] = np.random.gamma(2, 1, n_samples)
elif param.startswith("kappa_") or param.startswith("K_"):
# Hill half-saturation: Gamma(2, 1)
prior_samples[param] = np.random.gamma(2, 1, n_samples)
elif param.startswith("slope_") or param.startswith("S_"):
# Hill slope: Gamma(3, 1)
prior_samples[param] = np.random.gamma(3, 1, n_samples)
except Exception as e:
logger.debug(f"prior-sample synthesis failed: {e}")
return prior_samples
def _get_model_specification(self) -> dict[str, Any] | None:
"""Get model specification details."""
spec = {}
try:
# Get from model config
if hasattr(self.mmm, "model_config"):
config = self.mmm.model_config
spec["chains"] = getattr(config, "chains", 4)
spec["draws"] = getattr(config, "draws", 2000)
spec["tune"] = getattr(config, "tune", 1000)
# Get trend type
if hasattr(self.mmm, "trend_config"):
trend = self.mmm.trend_config
spec["baseline"] = (
f"{getattr(trend, 'type', 'Linear').value} trend + Fourier seasonality"
)
spec["likelihood"] = "Normal with estimated scale"
spec["media_effects"] = "Hill saturation × Geometric adstock"
spec["priors"] = "Weakly informative"
return spec
except Exception:
return None
__all__ = ["BayesianMMMExtractor"]