"""
DAG Model Builder
Main builder class for constructing MMM models from DAG specifications.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Self
import numpy as np
import pandas as pd
from mmm_framework.config import InferenceMethod, MFFConfig, ModelConfig
from .config_translator import (
dag_to_combined_config,
dag_to_mff_config,
dag_to_multivariate_config,
dag_to_nested_config,
dag_to_structural_config,
)
from .dag_spec import DAGEdge, DAGNode, DAGSpec
from .model_type_resolver import ModelType, get_model_class, resolve_model_type
from .validation import ValidationResult, validate_complete, validate_dag
if TYPE_CHECKING:
from mmm_framework.data_loader import PanelDataset
from mmm_framework.model import BayesianMMM, TrendConfig
from mmm_framework.mmm_extensions.models import (
CombinedMMM,
MultivariateMMM,
NestedMMM,
)
[docs]
class DAGBuildError(Exception):
"""Exception raised when model building fails."""
pass
[docs]
class DAGModelBuilder:
"""
Builder for constructing MMM models from DAG specifications.
This is the primary interface for the DAG-based model building workflow.
It supports fluent API for configuration and automatic model type selection.
Examples
--------
Basic usage:
>>> from mmm_framework.dag_model_builder import (
... DAGModelBuilder, DAGSpec, DAGNode, DAGEdge, NodeType
... )
>>> dag = DAGSpec(
... nodes=[
... DAGNode(id="sales", variable_name="Sales", node_type=NodeType.KPI),
... DAGNode(id="tv", variable_name="TV", node_type=NodeType.MEDIA),
... ],
... edges=[DAGEdge(source="tv", target="sales")]
... )
>>> model = (
... DAGModelBuilder()
... .with_dag(dag)
... .with_mff_data("data.csv")
... .bayesian_numpyro()
... .build()
... )
With mediation:
>>> dag = DAGSpec(
... nodes=[
... DAGNode(id="sales", variable_name="Sales", node_type=NodeType.KPI),
... DAGNode(id="tv", variable_name="TV", node_type=NodeType.MEDIA),
... DAGNode(id="awareness", variable_name="Awareness", node_type=NodeType.MEDIATOR),
... ],
... edges=[
... DAGEdge(source="tv", target="awareness"),
... DAGEdge(source="awareness", target="sales"),
... DAGEdge(source="tv", target="sales"),
... ]
... )
>>> # Automatically uses NestedMMM
>>> model = DAGModelBuilder().with_dag(dag).with_mff_data(df).build()
"""
[docs]
def __init__(self) -> None:
self._dag: DAGSpec | None = None
self._panel: PanelDataset | None = None
# The raw MFF long table (kept when data arrives via with_mff_data):
# mediator / additional-outcome series live here, NOT in the panel —
# dag_to_mff_config only emits kpi/media/controls.
self._mff_raw: pd.DataFrame | None = None
self._mff_config: MFFConfig | None = None
self._model_config: ModelConfig | None = None
self._trend_config: TrendConfig | None = None
self._node_config_overrides: dict[str, dict[str, Any]] = {}
self._date_format: str = "%Y-%m-%d"
self._frequency: str = "W"
# =========================================================================
# DAG Specification
# =========================================================================
[docs]
def with_dag(self, dag: DAGSpec) -> Self:
"""
Set the DAG specification.
Parameters
----------
dag : DAGSpec
The DAG specification.
Returns
-------
Self
The builder instance for chaining.
"""
self._dag = dag
return self
[docs]
def with_dag_dict(self, dag_dict: dict) -> Self:
"""
Set DAG from a dictionary.
Parameters
----------
dag_dict : dict
Dictionary with "nodes" and "edges" keys.
Returns
-------
Self
The builder instance for chaining.
"""
nodes = [DAGNode(**n) for n in dag_dict.get("nodes", [])]
edges = [DAGEdge(**e) for e in dag_dict.get("edges", [])]
metadata = dag_dict.get("metadata", {})
self._dag = DAGSpec(nodes=nodes, edges=edges, metadata=metadata)
return self
[docs]
@classmethod
def from_dag(cls, dag: DAGSpec) -> Self:
"""
Create a builder with a DAG specification.
Parameters
----------
dag : DAGSpec
The DAG specification.
Returns
-------
DAGModelBuilder
A new builder instance with the DAG set.
"""
builder = cls()
builder._dag = dag
return builder
[docs]
@classmethod
def from_frontend_json(
cls,
json_data: dict | str,
panel: "PanelDataset | None" = None,
) -> Self:
"""
Create builder from React Flow frontend JSON format.
Parameters
----------
json_data : dict | str
React Flow JSON data or JSON string.
panel : PanelDataset | None
Optional panel dataset.
Returns
-------
DAGModelBuilder
A new builder instance.
"""
from .frontend_adapter import react_flow_to_dag_spec
import json
if isinstance(json_data, str):
json_data = json.loads(json_data)
nodes = json_data.get("nodes", [])
edges = json_data.get("edges", [])
dag = react_flow_to_dag_spec(nodes, edges)
builder = cls()
builder._dag = dag
if panel is not None:
builder._panel = panel
return builder
# =========================================================================
# Data
# =========================================================================
[docs]
def with_panel(self, panel: "PanelDataset") -> Self:
"""
Set the panel dataset.
Parameters
----------
panel : PanelDataset
The panel dataset.
Returns
-------
Self
The builder instance for chaining.
"""
self._panel = panel
return self
[docs]
def with_mff_data(
self,
data: pd.DataFrame | str,
mff_config: MFFConfig | None = None,
) -> Self:
"""
Load data from MFF format.
If mff_config is not provided, it will be auto-generated from the DAG.
Parameters
----------
data : pd.DataFrame | str
MFF data or path to CSV file.
mff_config : MFFConfig | None
Optional MFF configuration. If None, generated from DAG.
Returns
-------
Self
The builder instance for chaining.
"""
from mmm_framework.data_loader import MFFLoader
if mff_config is not None:
self._mff_config = mff_config
elif self._dag is not None:
self._mff_config = dag_to_mff_config(
self._dag,
date_format=self._date_format,
frequency=self._frequency,
)
else:
raise DAGBuildError(
"Cannot load MFF data without DAG or MFF config. "
"Call with_dag() first or provide mff_config."
)
# Load data
if isinstance(data, str):
data = pd.read_csv(data)
self._mff_raw = data
loader = MFFLoader(self._mff_config)
self._panel = loader.load(data).build_panel()
return self
[docs]
def with_frequency(self, frequency: str) -> Self:
"""
Set the data frequency.
Parameters
----------
frequency : str
Data frequency ("W" for weekly, "D" for daily, "M" for monthly).
Returns
-------
Self
The builder instance for chaining.
"""
self._frequency = frequency
return self
# =========================================================================
# Model Configuration
# =========================================================================
[docs]
def with_model_config(self, config: ModelConfig) -> Self:
"""
Set model-level configuration.
Parameters
----------
config : ModelConfig
The model configuration.
Returns
-------
Self
The builder instance for chaining.
"""
self._model_config = config
return self
[docs]
def with_trend_config(self, config: "TrendConfig") -> Self:
"""
Set trend configuration.
Parameters
----------
config : TrendConfig
The trend configuration.
Returns
-------
Self
The builder instance for chaining.
"""
self._trend_config = config
return self
[docs]
def bayesian_pymc(self) -> Self:
"""
Use PyMC backend for inference.
Returns
-------
Self
The builder instance for chaining.
"""
if self._model_config is None:
self._model_config = ModelConfig()
self._model_config = self._model_config.model_copy(
update={"inference_method": InferenceMethod.BAYESIAN_PYMC}
)
return self
[docs]
def bayesian_numpyro(self) -> Self:
"""
Use NumPyro backend for inference (faster).
Returns
-------
Self
The builder instance for chaining.
"""
if self._model_config is None:
self._model_config = ModelConfig()
self._model_config = self._model_config.model_copy(
update={"inference_method": InferenceMethod.BAYESIAN_NUMPYRO}
)
return self
[docs]
def bayesian_nutpie(self) -> Self:
"""
Use the nutpie NUTS sampler (Rust backend).
Returns
-------
Self
The builder instance for chaining.
"""
if self._model_config is None:
self._model_config = ModelConfig()
self._model_config = self._model_config.model_copy(
update={"inference_method": InferenceMethod.BAYESIAN_NUTPIE}
)
return self
[docs]
def with_draws(self, n_draws: int) -> Self:
"""
Set number of posterior draws.
Parameters
----------
n_draws : int
Number of draws per chain.
Returns
-------
Self
The builder instance for chaining.
"""
if self._model_config is None:
self._model_config = ModelConfig()
self._model_config = self._model_config.model_copy(update={"n_draws": n_draws})
return self
[docs]
def with_tune(self, n_tune: int) -> Self:
"""
Set number of tuning samples.
Parameters
----------
n_tune : int
Number of tuning samples per chain.
Returns
-------
Self
The builder instance for chaining.
"""
if self._model_config is None:
self._model_config = ModelConfig()
self._model_config = self._model_config.model_copy(update={"n_tune": n_tune})
return self
[docs]
def with_chains(self, n_chains: int) -> Self:
"""
Set number of chains.
Parameters
----------
n_chains : int
Number of MCMC chains.
Returns
-------
Self
The builder instance for chaining.
"""
if self._model_config is None:
self._model_config = ModelConfig()
self._model_config = self._model_config.model_copy(
update={"n_chains": n_chains}
)
return self
# =========================================================================
# Node-Level Configuration Overrides
# =========================================================================
# =========================================================================
# Validation
# =========================================================================
[docs]
def validate(self) -> ValidationResult:
"""
Validate DAG structure and data compatibility.
Returns
-------
ValidationResult
Validation result with errors and warnings.
Raises
------
DAGBuildError
If no DAG has been set.
"""
if self._dag is None:
raise DAGBuildError("No DAG set. Call with_dag() first.")
return validate_complete(self._dag, self._panel)
# =========================================================================
# Build
# =========================================================================
[docs]
def get_model_type(self) -> ModelType:
"""
Get the resolved model type based on DAG structure.
Returns
-------
ModelType
The model type that will be built.
Raises
------
DAGBuildError
If no DAG has been set.
"""
if self._dag is None:
raise DAGBuildError("No DAG set. Call with_dag() first.")
return resolve_model_type(self._dag)
[docs]
def build_mff_config(self) -> MFFConfig:
"""
Build MFFConfig from DAG without building the full model.
Returns
-------
MFFConfig
The generated MFF configuration.
Raises
------
DAGBuildError
If no DAG has been set.
"""
if self._dag is None:
raise DAGBuildError("No DAG set. Call with_dag() first.")
# Apply overrides to DAG nodes
dag = self._apply_overrides()
return dag_to_mff_config(
dag,
date_format=self._date_format,
frequency=self._frequency,
)
def _apply_overrides(self) -> DAGSpec:
"""Apply node config overrides to DAG."""
if self._dag is None:
raise DAGBuildError("No DAG set.")
if not self._node_config_overrides:
return self._dag
# Create new nodes with overrides applied
new_nodes = []
for node in self._dag.nodes:
if node.variable_name in self._node_config_overrides:
# Merge overrides into node config
new_config = {
**node.config,
**self._node_config_overrides[node.variable_name],
}
new_node = node.model_copy(update={"config": new_config})
new_nodes.append(new_node)
else:
new_nodes.append(node)
return DAGSpec(
nodes=new_nodes,
edges=self._dag.edges,
metadata=self._dag.metadata,
)
[docs]
def build(
self,
) -> "BayesianMMM | NestedMMM | MultivariateMMM | CombinedMMM":
"""
Build the appropriate model based on DAG structure.
Returns
-------
BayesianMMM | NestedMMM | MultivariateMMM | CombinedMMM
The constructed model.
Raises
------
DAGBuildError
If DAG or data is not set, or validation fails.
"""
# Validate DAG is set
if self._dag is None:
raise DAGBuildError("No DAG set. Call with_dag() first.")
# Apply overrides
dag = self._apply_overrides()
# Validate DAG structure
validation = validate_dag(dag)
if not validation.valid:
raise DAGBuildError(
"DAG validation failed:\n"
+ "\n".join(f" - {e}" for e in validation.errors)
)
# Ensure panel data is available
if self._panel is None:
raise DAGBuildError(
"No data set. Call with_panel() or with_mff_data() first."
)
# Get model type and class
model_type = resolve_model_type(dag)
model_class = get_model_class(model_type)
# Get or create model config
model_config = self._model_config or ModelConfig()
# Build the appropriate model
if model_type == ModelType.BAYESIAN_MMM:
return model_class(
panel=self._panel,
model_config=model_config,
trend_config=self._trend_config,
)
# Extension models take ARRAYS (X_media, y/outcome_data, mediator_data),
# not a panel: the MFF panel carries kpi/media/controls only, so
# mediator / additional-outcome series are pulled from the raw MFF
# table and aligned to the panel's period index.
y, X_media, X_controls = self._panel.to_numpy()
channel_names = list(self._panel.X_media.columns)
index = self._panel.X_media.index
if model_type == ModelType.STRUCTURAL_NESTED_MMM:
structural_config, data_requirements = dag_to_structural_config(dag)
mediator_data: dict = {}
mediator_trials: dict = {}
for req in data_requirements:
if req["likelihood"] == "ordered":
# (n_obs, K) count matrix from the K category columns;
# NaN weeks = unobserved (the model masks them).
cols = self._national_series(
req["category_variables"], allow_missing=True
)
mediator_data[req["name"]] = np.column_stack(
[cols[c] for c in req["category_variables"]]
)
else:
mediator_data[req["name"]] = self._national_series(
[req["variable_name"]], allow_missing=True
)[req["variable_name"]]
if req["likelihood"] == "binomial":
mediator_trials[req["name"]] = self._national_series(
[req["trials_variable"]], allow_missing=True
)[req["trials_variable"]]
control_names = (
list(self._panel.X_controls.columns) if X_controls is not None else None
)
return model_class(
X_media=X_media,
y=y,
channel_names=channel_names,
config=structural_config,
mediator_data=mediator_data,
mediator_trials=mediator_trials,
X_controls=X_controls,
control_names=control_names,
index=index,
model_config=model_config,
trend_config=self._trend_config,
)
if model_type == ModelType.NESTED_MMM:
nested_config = dag_to_nested_config(dag)
mediator_data = self._national_series(
[n.variable_name for n in dag.mediator_nodes]
)
return model_class(
X_media=X_media,
y=y,
channel_names=channel_names,
config=nested_config,
mediator_data=mediator_data,
index=index,
# Extensions now honor the spec's baseline dynamics + outcome
# likelihood (seasonality via model_config, trend via
# trend_config) — previously discarded on this branch.
model_config=model_config,
trend_config=self._trend_config,
)
# Multi-outcome: the FIRST outcome is the panel's KPI series; every
# additional outcome node's series comes from the raw MFF table.
outcome_nodes = dag.outcome_nodes
outcome_data: dict[str, Any] = {outcome_nodes[0].variable_name: y}
outcome_data.update(
self._national_series([n.variable_name for n in outcome_nodes[1:]])
)
if model_type == ModelType.MULTIVARIATE_MMM:
multivariate_config = dag_to_multivariate_config(dag)
return model_class(
X_media=X_media,
outcome_data=outcome_data,
channel_names=channel_names,
config=multivariate_config,
index=index,
model_config=model_config,
trend_config=self._trend_config,
)
elif model_type == ModelType.COMBINED_MMM:
combined_config = dag_to_combined_config(dag)
mediator_data = self._national_series(
[n.variable_name for n in dag.mediator_nodes]
)
return model_class(
X_media=X_media,
outcome_data=outcome_data,
channel_names=channel_names,
config=combined_config,
mediator_data=mediator_data,
index=index,
model_config=model_config,
trend_config=self._trend_config,
)
else:
raise DAGBuildError(f"Unknown model type: {model_type}")
def _national_series(
self, names: list[str], *, allow_missing: bool = False
) -> dict[str, Any]:
"""National per-period series for DAG variables that are NOT in the
MFF panel (mediators / additional outcomes), aligned to the panel's
period index. Values are summed across any geo/product rows — the
same national collapse the MFF loader applies.
``allow_missing=True`` keeps NaN for periods with no rows instead of
raising — structural mediator surveys treat NaN as "no survey that
week" (the plain nested/multivariate branches keep strict coverage)."""
if not names:
return {}
if self._mff_raw is None:
raise DAGBuildError(
"Mediator/outcome variables need the raw MFF table — load data "
"with with_mff_data() (a with_panel() panel doesn't carry them)."
)
df = self._mff_raw
for col in ("VariableName", "VariableValue", "Period"):
if col not in df.columns:
raise DAGBuildError(
f"MFF data is missing the '{col}' column needed to extract "
"mediator/outcome series."
)
panel_index = self._panel.X_media.index
out: dict[str, Any] = {}
for name in names:
sub = df[df["VariableName"] == name]
if sub.empty:
raise DAGBuildError(
f"DAG variable '{name}' is not in the MFF data "
"(no VariableName rows) — mediators/outcomes must be "
"observed columns of the dataset."
)
periods = pd.to_datetime(
sub["Period"], format=self._date_format, errors="coerce"
)
if periods.isna().any():
periods = pd.to_datetime(sub["Period"], errors="coerce")
series = (
pd.to_numeric(sub["VariableValue"], errors="coerce")
.groupby(periods.values)
.sum()
)
aligned = series.reindex(pd.to_datetime(panel_index))
if aligned.isna().any() and not allow_missing:
missing = int(aligned.isna().sum())
raise DAGBuildError(
f"DAG variable '{name}' does not cover the panel's periods "
f"({missing} of {len(aligned)} missing after alignment)."
)
out[name] = aligned.to_numpy(dtype=float)
return out