Fit an MMM in Python

Start with a small synthetic model to learn the lifecycle. This self-contained example builds the graph, checks priors, fits, retains predictive draws and saves the result. Its sampler settings are starting values, not an acceptance criterion.

from pathlib import Path

import arviz as az
import numpy as np
import pandas as pd
from ammm.mmm import MMM, GeometricAdstock, LogisticSaturation

rng = np.random.default_rng(42)
X = pd.DataFrame({
    "date": pd.date_range("2025-01-06", periods=48, freq="W-MON"),
    "tv": rng.uniform(0, 100, 48),
    "social": rng.uniform(0, 60, 48),
})
y = pd.Series(
    40 + 30 * np.tanh(X["tv"] / 60)
    + 20 * np.tanh(X["social"] / 40)
    + rng.normal(0, 2, len(X)),
    index=X.index,
)
model = MMM(
    date_column="date",
    channel_columns=["tv", "social"],
    adstock=GeometricAdstock(l_max=1),
    saturation=LogisticSaturation(),
)
model.build_model(X, y)
model.add_original_scale_contribution_variable(["y", "channel_contribution"])
prior = model.sample_prior_predictive(X, y, samples=100, random_seed=42)
model.fit(
    X, y, nuts_sampler="pymc", chains=4, cores=1,
    tune=1000, draws=1000, target_accept=0.95, random_seed=42,
)
print(az.summary(model.idata, var_names=["saturation_beta", "y_sigma"]))
print("Divergences:", int(model.idata.sample_stats["diverging"].sum()))
predictions = model.sample_posterior_predictive(
    X, var_names=["y", "y_original_scale", "channel_contribution_original_scale"],
    combined=False, random_seed=43,
)
output = Path("sandbox/python-quickstart.nc")
output.parent.mkdir(parents=True, exist_ok=True)
if output.exists():
    raise FileExistsError(output)
model.save(output)
restored = MMM.load(output)

Prior predictive draws show the range implied before conditioning on outcomes; inspect model.idata["prior"]["y_original_scale"] against plausible outcome levels. The returned prior dataset contains the observed variable y, because the sampler extracts the prior_predictive group (src/ammm/model_builder.py:673). Posterior predictive draws include observation variation. The fitted channel contribution is a model component, so its uncertainty differs from predictive uncertainty.

Inspect sampling diagnostics, then evaluate a held-out period before interpreting real data. This example uses all synthetic rows for demonstration and does not provide an independent predictive assessment.

Check both data interfaces

The existing executable example also checks persistence by comparing canonical training data, fixed-parameter log density and seeded predictions after loading:

uv run --no-sync python scripts/examples/mmm_fit_save_load.py --output sandbox/mmm-dataframe.nc
uv run --no-sync python scripts/examples/mmm_fit_save_load.py --layout xarray --output sandbox/mmm-xarray.nc

Use fresh output paths. Add --draws 20 --tune 20 for an execution smoke check; its diagnostics are not evidence of adequate posterior inference.

Continue with data and lifecycle rules, model specification, exporting results and budget optimisation.

Implementation reference at 7cb7f20: src/ammm/mmm/mmm.py:1731, scripts/examples/mmm_fit_save_load.py:15.