"""
Mixin classes providing shared extraction utilities.
These mixins provide reusable methods for data aggregation and
dimension-specific extraction (geo, product, time).
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
import numpy as np
from loguru import logger
if TYPE_CHECKING:
from .bundle import MMMDataBundle
def _finite(x: Any) -> float | None:
"""Coerce to a finite float, or ``None`` (for missing / non-finite values)."""
if x is None:
return None
try:
v = float(x)
except (TypeError, ValueError):
return None
return v if np.isfinite(v) else None
[docs]
class AggregationMixin:
"""
Mixin providing data aggregation utilities for extractors.
Provides methods for aggregating data by period, geography, and product
while properly propagating uncertainty through sample aggregation.
"""
def _aggregate_by_period_simple(
self,
values: np.ndarray,
periods: list,
unique_periods: list,
) -> np.ndarray:
"""
Aggregate values by period using simple summation.
Parameters
----------
values : np.ndarray
Values to aggregate, shape (n_obs,).
periods : list
Period label for each observation.
unique_periods : list
Unique periods in order.
Returns
-------
np.ndarray
Aggregated values, shape (n_periods,).
"""
period_to_idx = {p: i for i, p in enumerate(unique_periods)}
n_periods = len(unique_periods)
result = np.zeros(n_periods)
for i, (val, period) in enumerate(zip(values, periods)):
if period in period_to_idx:
result[period_to_idx[period]] += val
return result
def _aggregate_samples_by_period(
self,
samples: np.ndarray,
periods: list,
unique_periods: list,
ci_prob: float = 0.8,
) -> dict[str, np.ndarray] | None:
"""
Aggregate samples by period while preserving uncertainty.
This method properly propagates uncertainty by aggregating
samples first, then computing statistics on aggregated values.
Parameters
----------
samples : np.ndarray
Sample array of shape (n_samples, n_obs).
periods : list
Period label for each observation.
unique_periods : list
Unique periods in order.
ci_prob : float
Credible interval probability.
Returns
-------
dict[str, np.ndarray] or None
Dictionary with "mean", "lower", "upper" arrays.
"""
if samples is None or len(periods) == 0:
return None
try:
period_to_idx = {p: i for i, p in enumerate(unique_periods)}
obs_period_idx = np.array([period_to_idx.get(p, -1) for p in periods])
n_samples = samples.shape[0]
n_periods = len(unique_periods)
# Aggregate samples by period
samples_agg = np.zeros((n_samples, n_periods))
for t in range(n_periods):
mask = obs_period_idx == t
if mask.any():
samples_agg[:, t] = samples[:, mask].sum(axis=1)
# Compute statistics
alpha = (1 - ci_prob) / 2
return {
"mean": samples_agg.mean(axis=0),
"lower": np.percentile(samples_agg, alpha * 100, axis=0),
"upper": np.percentile(samples_agg, (1 - alpha) * 100, axis=0),
}
except Exception:
return None
def _aggregate_by_group(
self,
values: np.ndarray,
group_idx: np.ndarray,
n_groups: int,
) -> np.ndarray:
"""
Aggregate values by group index.
Parameters
----------
values : np.ndarray
Values to aggregate.
group_idx : np.ndarray
Group index for each value.
n_groups : int
Number of groups.
Returns
-------
np.ndarray
Aggregated values per group.
"""
result = np.zeros(n_groups)
for i, val in enumerate(values):
if 0 <= group_idx[i] < n_groups:
result[group_idx[i]] += val
return result
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.
Parameters
----------
values : np.ndarray
Values to aggregate.
n_periods : int
Number of periods.
geo_mask : np.ndarray
Boolean mask for the geo (unused in simple case).
Returns
-------
np.ndarray
Aggregated values by 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
# Assumes data is ordered: periods are innermost, then geos, then products
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.
Parameters
----------
values : np.ndarray
Values to aggregate, shape (n_obs,)
time_idx : np.ndarray
Time period index for each observation, shape (n_obs,)
dim_mask : np.ndarray
Boolean mask for filtering (e.g., specific geo or product), shape (n_obs,)
n_periods : int
Number of unique time periods
Returns
-------
np.ndarray
Aggregated values by period, shape (n_periods,)
"""
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
class EstimandPPCMixin:
"""Shared extraction of estimand results (mean + CI) and the
posterior-predictive goodness-of-fit summary.
Both ``BayesianMMMExtractor`` and ``ExtendedMMMExtractor`` inherit this. Two
hooks adapt it per model family:
- :meth:`_estimand_model` — the model object exposing ``evaluate_estimands``
(defaults to ``self.mmm`` or ``self.model``).
- :meth:`_ppc_arrays` — returns aligned, original-scale
``(observed, y_rep, pred_mean, pred_lower, pred_upper)`` posterior-predictive
arrays, or ``(None,)*5`` when the model cannot produce them. Each extractor
implements this differently (a core MMM resamples via ``predict``; an
extended model samples its PyMC graph).
Both ``_extract_*`` methods are best-effort: any failure leaves the bundle
field empty so the corresponding section simply no-ops.
"""
#: Nominal interval widths sampled for the calibration curve.
_PPC_NOMINAL_LEVELS = (0.5, 0.6, 0.7, 0.8, 0.9, 0.95)
#: Number of replicate density curves kept for the overlay chart.
_PPC_OVERLAY_CURVES = 40
# -- hooks ----------------------------------------------------------------
def _estimand_model(self) -> Any:
"""The model object exposing ``evaluate_estimands`` (family-specific)."""
return getattr(self, "mmm", None) or getattr(self, "model", None)
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``. Subclasses
override; the default produces nothing."""
return (None, None, None, None, None)
# -- estimands ------------------------------------------------------------
def _extract_estimands(self, bundle: "MMMDataBundle") -> "MMMDataBundle":
"""Realize the model's declared / default estimands as mean + CI.
Prefers estimands already realized on the fit results (computed at fit
time when the model declares them); otherwise evaluates them on demand via
``model.evaluate_estimands()``. A model that does not implement the
estimand engine (e.g. an extended model without ``evaluate_estimands``)
leaves ``bundle.estimands`` empty.
"""
try:
results = getattr(self.results, "estimands", None) if self.results else None
if not results:
model = self._estimand_model()
fn = getattr(model, "evaluate_estimands", None)
if not callable(fn):
return bundle
results = fn()
if not results:
return bundle
out: dict[str, dict[str, Any]] = {}
for key, r in results.items():
mean = _finite(getattr(r, "mean", None))
if getattr(r, "status", "ok") != "ok" or mean is None:
continue
entry: dict[str, Any] = {
"mean": mean,
"lower": _finite(getattr(r, "hdi_low", None)),
"upper": _finite(getattr(r, "hdi_high", None)),
"kind": getattr(r, "kind", "") or "",
"units": getattr(r, "units", "") or "",
"hdi_prob": float(getattr(r, "hdi_prob", 0.94) or 0.94),
}
extra = getattr(r, "extra", None) or {}
for k in ("contribution_pct", "prob_positive", "prob_profitable"):
val = _finite(extra.get(k)) if isinstance(extra, dict) else None
if val is not None:
entry[k] = val
# Impression-level measurement metadata (ROI vs efficiency, and
# the break-even reference) so the section renders + judges the
# estimand correctly. Absent for ordinary spend estimands.
if isinstance(extra, dict):
if "metric_reference" in extra:
entry["reference"] = _finite(extra.get("metric_reference"))
for k in (
"metric_label",
"metric_is_monetary",
"measurement_unit",
"cost_basis",
):
if extra.get(k) is not None:
entry[k] = extra.get(k)
out[str(key)] = entry
if out:
bundle.estimands = out
except Exception: # noqa: BLE001 — reporting must never hard-fail
logger.debug("estimands extraction skipped", exc_info=True)
return bundle
# -- short-term vs long-term / brand effect (issue #106) ------------------
def _extract_long_term(self, bundle: "MMMDataBundle") -> "MMMDataBundle":
"""Populate ``bundle.long_term`` — the immediate-vs-carryover split the
weekly model measures, plus whether a brand funnel exists.
Best-effort: reads the already-extracted ``adstock_curves`` +
``component_totals``. Even when nothing is estimable it returns a payload
so the section can render the "only short-term is measured" caveat. A
long-term multiplier scenario is added only when the report config
requests one (``config.long_term_multiplier``)."""
try:
from ..helpers.longterm import (
build_long_term_facts,
estimated_long_term_split,
)
channels = list(bundle.channel_names or [])
if not channels:
return bundle
model = self._estimand_model()
has_funnel = bool(
getattr(model, "mediator_names", None)
or getattr(model, "__garden_model_kind__", "") == "structural"
)
# A model that estimates a slow brand-equity stock (e.g.
# LongTermBrandMMM) yields a genuine long-term split — an ESTIMATE, not
# the assumption-driven scenario (issue #122). None for a plain MMM.
estimated = estimated_long_term_split(model)
# Caveat-only otherwise (no multiplier); the LongTermSection applies the
# optional long-term-multiplier scenario from the report config.
# Model-derived carryover half-life per channel. `build_long_term_facts`
# has always accepted `half_lives=`, but no caller ever passed it, so
# the classic report's half-life column was permanently empty. Read
# from the fitted kernel via the canonical per-draw reader, which is
# family-aware (geometric / delayed / weibull) — the old
# log(0.5)/log(alpha) shortcut is meaningless for a humped kernel.
half_lives = self._carryover_half_lives(model, channels)
facts = build_long_term_facts(
channels,
bundle.adstock_curves,
contribution=bundle.component_totals,
multiplier=None,
has_structural_funnel=has_funnel,
estimated=estimated,
half_lives=half_lives,
)
if facts is not None:
bundle.long_term = facts
except Exception: # noqa: BLE001 — reporting must never hard-fail
logger.debug("long-term extraction skipped", exc_info=True)
return bundle
@staticmethod
def _carryover_half_lives(model, channels: list[str]) -> dict[str, float]:
"""Posterior-mean carryover half-life per channel, best-effort.
Uses the canonical per-draw reader, so a delayed or Weibull channel gets
its true cumulative-50% horizon rather than a decay constant that only
means anything for a geometric kernel. Channels the reader cannot
resolve are simply absent from the mapping — ``build_long_term_facts``
skips a missing or non-finite entry rather than inventing one.
"""
if model is None or not channels:
return {}
try:
from ...transforms.carryover import (
carryover_half_life,
posterior_carryover_kernels,
)
except Exception: # noqa: BLE001
return {}
out: dict[str, float] = {}
try:
kernels = posterior_carryover_kernels(model, list(channels))
except Exception: # noqa: BLE001
return {}
for ch, k in (kernels or {}).items():
try:
if k is None or getattr(k, "status", "ok") != "ok":
continue
hl = carryover_half_life(k.kernel)
mean_hl = float(np.nanmean(hl))
if np.isfinite(mean_hl):
out[str(ch)] = mean_hl
except Exception: # noqa: BLE001
continue
return out
# -- evidence tier + identifiability gate (issue #102) --------------------
#: Prior draws used for the report-time prior→posterior contraction check
#: (kept modest so evidence extraction stays cheap; it is best-effort).
_EVIDENCE_PRIOR_SAMPLES = 400
def _channel_collinearity_matrix(
self, bundle: "MMMDataBundle"
) -> "np.ndarray | None":
"""An ``(n_obs, n_channels)`` design matrix for the collinearity check.
Prefers the model's raw media matrix (``X_media_raw`` / ``X_media``) —
the channels' input design the likelihood sees — and falls back to the
per-channel contribution time series already on the bundle. Returns
``None`` when neither aligns to the channel list.
"""
channels = list(bundle.channel_names or [])
n_ch = len(channels)
if n_ch < 2:
return None
model = self._estimand_model()
for attr in ("X_media_raw", "X_media"):
X = getattr(model, attr, None)
if X is None:
continue
try:
X = np.asarray(X, dtype=float)
except (TypeError, ValueError):
continue
if X.ndim == 2 and X.shape[1] == n_ch and X.shape[0] >= 3:
return X
# Fallback: contribution time series from the bundle.
cts = bundle.component_time_series or {}
cols = []
for ch in channels:
v = cts.get(ch)
if v is None:
return None
arr = np.asarray(v, dtype=float).ravel()
cols.append(arr)
if cols and len({c.shape[0] for c in cols}) == 1 and cols[0].shape[0] >= 3:
return np.column_stack(cols)
return None
def _extract_channel_evidence(self, bundle: "MMMDataBundle") -> "MMMDataBundle":
"""Attach a per-channel evidence tier + identifiability flag (issue #102).
Three orthogonal signals are folded into one :class:`ChannelEvidence`:
* *experiment-validated* — the channel's effect was calibrated against a
randomized experiment folded into THIS fit (``model.experiments``).
* *prior-dominated* — the prior→posterior contraction (best-effort
``model.compute_parameter_learning``) says the data barely moved the
channel's coefficient off its prior.
* *not separately identified* — the channel is collinear with another
over the media design (variance inflation).
The result is set on ``bundle.channel_evidence`` and stamped into each
``channel_roi[ch]["evidence"]`` and per-channel ``estimands`` entry so
every renderer reads one source of truth. Best-effort — any failure
leaves the bundle unannotated and the reports simply omit the chip.
"""
try:
from ..evidence import evidence_for_model
channels = list(bundle.channel_names or [])
if not channels:
return bundle
model = self._estimand_model()
# One gathering path shared with the fit-time dashboard snapshot
# (evidence_for_model) so the report and the live Performance page
# render the SAME tier (issue #124). The bundle's contribution-series
# collinearity fallback is passed through for models that don't expose
# their raw media design.
evidence = evidence_for_model(
model,
channels,
collinearity_matrix=self._channel_collinearity_matrix(bundle),
prior_samples=self._EVIDENCE_PRIOR_SAMPLES,
)
evd = {ch: e.to_dict() for ch, e in evidence.items()}
bundle.channel_evidence = evd
# Stamp onto channel_roi rows (the primary render source).
if bundle.channel_roi:
for ch, row in bundle.channel_roi.items():
if isinstance(row, dict) and ch in evd:
row["evidence"] = evd[ch]
# Stamp onto per-channel estimands ("{name}:{channel}").
if bundle.estimands:
for key, entry in bundle.estimands.items():
if not isinstance(entry, dict):
continue
ch = key.split(":", 1)[1] if ":" in key else None
if ch and ch in evd:
entry["evidence"] = evd[ch]
except Exception: # noqa: BLE001 — reporting must never hard-fail
logger.debug("channel evidence extraction skipped", exc_info=True)
return bundle
# -- CFO one-pager (P&L rollup + spend-cut risk) --------------------------
#: Posterior draws for the CFO facts (thinned — a board number, not a chart).
_CFO_MAX_DRAWS = 200
def _extract_cfo(self, bundle: "MMMDataBundle") -> "MMMDataBundle":
"""Best-effort CFO one-pager facts (issue #108): marketing's incremental
contribution + spend-cut revenue-at-risk with uncertainty. Revenue-scale —
a margin (for profit-at-risk) is applied on the agent-tool path when one is
configured. MMM-only (needs the channel response surface)."""
try:
model = self._estimand_model()
if model is None or not hasattr(model, "sample_channel_contributions"):
return bundle
if (
getattr(model, "X_media_raw", None) is None
or getattr(model, "y_raw", None) is None
):
return bundle
from ..helpers.cfo import cfo_facts
bundle.cfo = cfo_facts(model, max_draws=self._CFO_MAX_DRAWS)
except Exception: # noqa: BLE001 — reporting must never hard-fail
logger.debug("cfo extraction skipped", exc_info=True)
return bundle
# -- posterior-predictive goodness-of-fit ---------------------------------
def _extract_posterior_predictive(self, bundle: "MMMDataBundle") -> "MMMDataBundle":
"""Compute the posterior-predictive goodness-of-fit summary.
Pulls aligned posterior-predictive replicates (original KPI scale,
observation level) via :meth:`_ppc_arrays`, then derives the compact,
chart-ready summary the PosteriorPredictiveSection consumes: per-obs mean
+ interval, a down-sampled replicate cloud, an interval-calibration curve,
and posterior-predictive p-values for summary statistics. Best-effort.
"""
try:
observed, y_rep, pred_mean, pred_lower, pred_upper = self._ppc_arrays()
if observed is None or y_rep is None:
return bundle
ci = float(self.ci_prob)
coverage = []
for p in self._PPC_NOMINAL_LEVELS:
lo = np.percentile(y_rep, (1 - p) / 2 * 100, axis=0)
hi = np.percentile(y_rep, (1 + p) / 2 * 100, axis=0)
emp = float(np.mean((observed >= lo) & (observed <= hi)))
coverage.append({"nominal": float(p), "empirical": emp})
ss_res = float(np.sum((observed - pred_mean) ** 2))
ss_tot = float(np.sum((observed - observed.mean()) ** 2))
r2 = (1.0 - ss_res / ss_tot) if ss_tot > 0 else None
bundle.posterior_predictive = {
"observed": observed,
"pred_mean": pred_mean,
"pred_lower": pred_lower,
"pred_upper": pred_upper,
"samples": self._downsample_rows(y_rep, self._PPC_OVERLAY_CURVES),
"coverage": coverage,
"bayes_p": self._ppc_bayes_p(observed, y_rep),
"ci_level": ci,
"r2": r2,
}
except Exception: # noqa: BLE001 — reporting must never hard-fail
logger.debug("posterior-predictive extraction skipped", exc_info=True)
return bundle
@staticmethod
def _ppc_bayes_p(observed: np.ndarray, y_rep: np.ndarray) -> dict[str, float]:
"""Posterior-predictive p-values: P(T(y_rep) >= T(y_obs)) per statistic."""
stat_fns = {
"mean": lambda a, axis=None: np.mean(a, axis=axis),
"std": lambda a, axis=None: np.std(a, axis=axis),
"min": lambda a, axis=None: np.min(a, axis=axis),
"max": lambda a, axis=None: np.max(a, axis=axis),
}
out: dict[str, float] = {}
for name, fn in stat_fns.items():
obs_stat = float(fn(observed))
rep_stat = fn(y_rep, axis=1)
out[name] = float(np.mean(rep_stat >= obs_stat))
return out
@staticmethod
def _downsample_rows(arr: np.ndarray, n: int) -> np.ndarray:
"""Deterministically thin rows to at most ``n`` (evenly spaced)."""
if arr.shape[0] <= n:
return arr
idx = np.linspace(0, arr.shape[0] - 1, n).astype(int)
return arr[idx]
__all__ = [
"AggregationMixin",
"GeoExtractionMixin",
"ProductExtractionMixin",
"EstimandPPCMixin",
"_finite",
]