{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9d11760f",
   "metadata": {},
   "source": [
    "# Ridge regression for a marketing mix model, from scratch\n",
    "\n",
    "From the Vexpower MMM library: [vexpower.com/mmm/ridge-regression-mmm](https://vexpower.com/mmm/ridge-regression-mmm)\n",
    "\n",
    "This notebook fits a small MMM to the 104-week practice dataset with ordinary least squares and with ridge regression, using only numpy, pandas and matplotlib. It shows:\n",
    "\n",
    "1. why the channels are standardised before the penalty is applied,\n",
    "2. the ridge trace (every coefficient across a range of penalties),\n",
    "3. choosing λ with time-series cross-validation (train on the past, test on the next weeks), and\n",
    "4. how the channel contributions change between OLS and the chosen ridge model.\n",
    "\n",
    "The objective is glmnet's ridge objective (`alpha = 0`), the one Robyn uses:\n",
    "(1/2n)·Σ(y − ŷ)² + (λ/2)·Σβ², on standardised channels and a standardised outcome.\n",
    "\n",
    "Put `mmm-practice-synthetic.csv` in the same folder (download it from vexpower.com/mmm), or change `CSV` below to your own file.\n",
    "\n",
    "A note on the practice data: it was generated with no noise, no carryover and no saturation (revenue = 20,000 + 1.4 × search + 0.8 × social + 0.6 × video + seasonality + 1,800 × promotion), and the three channels are almost uncorrelated. That makes it a good place to see the mechanics, but not a test of ridge. The fixed adstock retentions below are illustrative and actually wrong for this data (the true carryover is zero), which is why the fitted returns land near, not on, 1.4, 0.8 and 0.6, and why ridge barely changes them. On real data with correlated channels and noise, expect a much bigger gap between OLS and ridge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e181cdaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "CSV = \"mmm-practice-synthetic.csv\"\n",
    "OUTCOME = \"revenue\"\n",
    "CHANNELS = [\"search_spend\", \"social_spend\", \"video_spend\"]\n",
    "CONTROLS = [\"promotion\"]\n",
    "# Fixed geometric adstock retention per channel (illustrative: the practice data has no carryover; in a real model these are tuned too)\n",
    "RETENTION = {\"search_spend\": 0.1, \"social_spend\": 0.3, \"video_spend\": 0.5}\n",
    "\n",
    "df = pd.read_csv(CSV, parse_dates=[\"week\"])\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f942d155",
   "metadata": {},
   "outputs": [],
   "source": [
    "def geometric_adstock(x, retention):\n",
    "    out, carry = np.zeros(len(x)), 0.0\n",
    "    for t, v in enumerate(x):\n",
    "        carry = v + retention * carry\n",
    "        out[t] = carry\n",
    "    return out\n",
    "\n",
    "X_raw = pd.DataFrame({c: geometric_adstock(df[c].to_numpy(float), RETENTION[c]) for c in CHANNELS})\n",
    "for c in CONTROLS:\n",
    "    X_raw[c] = df[c].astype(float)\n",
    "# simple yearly seasonality so the channels don't absorb it\n",
    "t = np.arange(len(df))\n",
    "X_raw[\"sin52\"] = np.sin(2 * np.pi * t / 52)\n",
    "X_raw[\"cos52\"] = np.cos(2 * np.pi * t / 52)\n",
    "y = df[OUTCOME].to_numpy(float)\n",
    "X_raw.corr().round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c411000",
   "metadata": {},
   "source": [
    "## Ridge with standardisation\n",
    "\n",
    "The penalty treats every coefficient the same, so the columns must be on the same scale first; otherwise a channel measured in dollars and a 0/1 promotion flag would be penalised very differently. Like glmnet (`standardize = TRUE`), we standardise, fit, and convert the coefficients back to the original units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60cc8e29",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ridge_fit(X, y, lam):\n",
    "    X = np.asarray(X, float); y = np.asarray(y, float)\n",
    "    n, p = X.shape\n",
    "    mx, sx = X.mean(0), X.std(0)\n",
    "    my, sy = y.mean(), y.std()\n",
    "    Z = (X - mx) / sx\n",
    "    yz = (y - my) / sy\n",
    "    # (1/2n)RSS + (lam/2)|b|^2  <=>  (Z'Z + n*lam*I) b = Z'y\n",
    "    b = np.linalg.solve(Z.T @ Z + n * lam * np.eye(p), Z.T @ yz)\n",
    "    beta = b * sy / sx\n",
    "    intercept = my - mx @ beta\n",
    "    return intercept, beta\n",
    "\n",
    "def predict(model, X):\n",
    "    a, b = model\n",
    "    return a + np.asarray(X, float) @ b\n",
    "\n",
    "ols = ridge_fit(X_raw, y, 0.0)\n",
    "pd.Series(ols[1], index=X_raw.columns).round(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97256c9b",
   "metadata": {},
   "source": [
    "## The ridge trace\n",
    "\n",
    "Each line is one coefficient as λ grows. Coefficients that are well identified hold their value for longer; correlated or weakly identified ones shrink first and towards each other. Every coefficient reaches zero eventually, which is a reminder that ridge shrinks towards *no effect*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "801586de",
   "metadata": {},
   "outputs": [],
   "source": [
    "lams = np.logspace(-4, 1.5, 80)\n",
    "path = np.array([ridge_fit(X_raw, y, l)[1] for l in lams])\n",
    "fig, ax = plt.subplots(figsize=(8, 4.5))\n",
    "for j, c in enumerate(CHANNELS):\n",
    "    ax.plot(np.log10(lams), path[:, j], label=c)\n",
    "ax.axhline(0, color=\"grey\", lw=0.8)\n",
    "ax.set_xlabel(\"log10 lambda\")\n",
    "ax.set_ylabel(\"revenue per $1 of adstocked spend\")\n",
    "ax.set_title(\"Ridge trace for the media channels\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81bbc95a",
   "metadata": {},
   "source": [
    "## Choosing λ with time-series cross-validation\n",
    "\n",
    "Weeks aren't independent, so don't shuffle them into random folds. Train on an expanding window of the past and score the next block of weeks, then average. We also report the \"one standard error\" choice: the largest λ whose error is within one standard error of the best, which gives a more conservative model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b655c0cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ts_cv_error(X, y, lam, first_train=52, horizon=8):\n",
    "    X = np.asarray(X, float); errs = []\n",
    "    for start in range(first_train, len(y) - horizon + 1, horizon):\n",
    "        m = ridge_fit(X[:start], y[:start], lam)\n",
    "        pred = predict(m, X[start:start + horizon])\n",
    "        errs.append(np.sqrt(np.mean((y[start:start + horizon] - pred) ** 2)))\n",
    "    return np.mean(errs), np.std(errs) / np.sqrt(len(errs))\n",
    "\n",
    "cv = np.array([ts_cv_error(X_raw, y, l) for l in lams])\n",
    "best = int(np.argmin(cv[:, 0]))\n",
    "one_se = max(i for i in range(len(lams)) if cv[i, 0] <= cv[best, 0] + cv[best, 1])\n",
    "print(f\"OLS (lambda ~ 0) CV RMSE: {cv[0, 0]:,.0f}\")\n",
    "print(f\"Best lambda: {lams[best]:.4f}  CV RMSE: {cv[best, 0]:,.0f}\")\n",
    "print(f\"One-SE lambda: {lams[one_se]:.4f}  CV RMSE: {cv[one_se, 0]:,.0f}\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.plot(np.log10(lams), cv[:, 0])\n",
    "ax.fill_between(np.log10(lams), cv[:, 0] - cv[:, 1], cv[:, 0] + cv[:, 1], alpha=0.2)\n",
    "ax.axvline(np.log10(lams[best]), ls=\"--\", label=\"best\")\n",
    "ax.axvline(np.log10(lams[one_se]), ls=\":\", label=\"one-SE\")\n",
    "ax.set_xlabel(\"log10 lambda\"); ax.set_ylabel(\"CV RMSE\"); ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e39df6c6",
   "metadata": {},
   "source": [
    "## How the channel story changes\n",
    "\n",
    "Compare channel contributions (coefficient × adstocked spend, summed over the period) under OLS and the chosen ridge model. If the ranking changes a lot, the data alone doesn't settle it, and that's worth reporting rather than hiding behind the penalty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a64e7557",
   "metadata": {},
   "outputs": [],
   "source": [
    "chosen = ridge_fit(X_raw, y, lams[one_se])\n",
    "rows = []\n",
    "for j, c in enumerate(X_raw.columns):\n",
    "    if c in CHANNELS:\n",
    "        spend = df[c].sum()\n",
    "        for name, m in [(\"OLS\", ols), (\"Ridge (one-SE)\", chosen)]:\n",
    "            contrib = m[1][j] * X_raw[c].sum()\n",
    "            rows.append({\"channel\": c, \"model\": name, \"coef\": m[1][j], \"contribution\": contrib, \"roi\": contrib / spend})\n",
    "summary = pd.DataFrame(rows).pivot(index=\"channel\", columns=\"model\", values=[\"coef\", \"roi\"]).round(3)\n",
    "summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "595fe443",
   "metadata": {},
   "source": [
    "Ridge makes these estimates more stable, not more causal. It can't fix a missing confounder, and it pulls every channel towards zero and towards its correlated neighbours. If two channels still trade places between OLS and ridge, the next step is new information: an experiment, a prior from one, or a deliberate change to the media plan."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}