same basic structure for each of the modelling sections

Author: drbenvincent

docs/source/notebooks/graded_intervention_time_series_single_channel_ols.ipynb

@@ -212,8 +212,6 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 1. Generate Simulated Data\n",
-    "\n",
     "We'll simulate weekly water consumption data for a catchment area in a **dry climate** over 3 years with:\n",
     "- **Baseline drivers**: temperature (seasonal) and rainfall (very low, with drought periods)\n",
     "- **Responsive policy**: public communications intensity that activates only during sustained drought\n",
@@ -517,12 +515,14 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 2. Model Fitting\n",
+    "## Modelling with HAC\n",
     "\n",
     "Fitting a transfer function model involves finding both the optimal transform parameters and the regression coefficients. This is accomplished through a nested optimization procedure. In the outer loop, the algorithm searches for the best saturation and adstock parameters—either by exhaustively evaluating all combinations on a discrete grid, or by using continuous optimization to search more efficiently through the parameter space. For each candidate set of transform parameters, the inner loop applies these transformations to the raw treatment variable and fits a regression model (OLS or ARIMAX) to the data. The root mean squared error (RMSE) of each fitted model is computed, and the parameter combination that minimizes this error is selected.\n",
     "\n",
     "This nested approach is computationally tractable because ordinary least squares has a closed-form solution based on matrix operations, making each individual model fit very fast. When using grid search with, say, 10 values for each of 3 parameters, the algorithm evaluates 1,000 model fits, which typically completes in under a second. Continuous optimization via gradient-based methods can be even faster, though it may settle on local optima rather than finding the global best. For models with ARIMAX error structures, each fit requires numerical optimization and takes longer, but the overall approach remains the same.\n",
     "\n",
+    "### Fit Model\n",
+    "\n",
     "Let's fit a model using grid search to estimate the transform parameters. We'll use a coarse grid for speed in this demonstration, though you can make it finer for production use to achieve more precise parameter estimates.\n"
    ]
   },
@@ -578,7 +578,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 3. Visualize Estimated vs True Transform Parameters\n",
+    "### Visualize Estimated vs True Transform Parameters\n",
     "\n",
     "Since we know the true parameters used to generate the data, we can compare the estimated transforms to the true transforms. This helps us assess **parameter recovery** - how well the estimation procedure identifies the true data-generating process.\n",
     "\n",
@@ -665,10 +665,6 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 4. Model Methods and Diagnostics\n",
-    "\n",
-    "Now that we have a fitted model with estimated transforms, let's explore the available methods for analysis and diagnostics.\n",
-    "\n",
     "### Model Summary\n",
     "\n",
     "View the fitted model coefficients and their **HAC standard errors** (robust to autocorrelation and heteroskedasticity):\n"
@@ -1080,7 +1076,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1102,20 +1098,22 @@
     "    scale=0.0,  # Zero out all communications (no policy counterfactual)\n",
     ")\n",
     "\n",
+    "# Visualize the counterfactual analysis\n",
+    "fig, ax = result_estimated.plot_effect(effect_result)\n",
+    "plt.show()\n",
+    "\n",
     "print(\n",
-    "    f\"Total water saved by communications policy: {-effect_result['total_effect']:.0f} ML\"\n",
-    ")\n",
-    "print(f\"Average weekly savings: {-effect_result['mean_effect']:.0f} ML/week\")\n",
-    "print(\n",
-    "    f\"Percentage reduction: {-100 * effect_result['mean_effect'] / df['water_consumption'].mean():.1f}%\"\n",
+    "    f\"\\nThe communications policy saved approximately {-effect_result['total_effect']:.0f} ML of water \"\n",
+    "    f\"over the 3-year period, representing a {-100 * effect_result['mean_effect'] / df['water_consumption'].mean():.1f}% \"\n",
+    "    f\"reduction in average consumption.\"\n",
     ")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 5. Alternative Error Model: ARIMAX\n",
+    "## Modelling with ARIMAX\n",
     "\n",
     "So far we've used **HAC (Newey-West) standard errors**, which provide robust inference without requiring us to specify the autocorrelation structure. This is the recommended default approach.\n",
     "\n",
@@ -1143,7 +1141,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Fit Model with ARIMAX Errors\n",
+    "### Fit Model\n",
     "\n",
     "Since we generated the data with AR(2) errors (`rho1=0.5`, `rho2=0.2`), the true error structure is ARIMA(2,0,0). For demonstration purposes, we'll fit an ARIMA(1,0,0) model, which is a slight misspecification. This shows how ARIMAX still performs reasonably well even when the order is not perfectly matched. In practice, you would use ACF/PACF plots to guide ARIMA order selection:\n"
    ]
@@ -1192,9 +1190,73 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Visualize Estimated vs True Transform Parameters\n",
+    "\n",
+    "Since we know the true parameters used to generate the data, we can compare the estimated transforms to the true transforms. This helps us assess **parameter recovery** - how well the estimation procedure identifies the true data-generating process.\n",
+    "\n",
+    "We'll visualize:\n",
+    "1. **Saturation curves**: How raw communication intensity gets transformed by saturation\n",
+    "2. **Adstock weights**: How effects carry over across weeks\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create true transform objects (parameters used for data generation)\n",
+    "true_saturation = cp.LogisticSaturation(lam=0.5)\n",
+    "true_adstock = cp.GeometricAdstock(half_life=1.5, l_max=8, normalize=True)\n",
+    "\n",
+    "# Plot estimated transforms with comparison to true transforms\n",
+    "fig, ax = result_arimax.plot_transforms(\n",
+    "    true_saturation=true_saturation, true_adstock=true_adstock, x_range=(0, 10)\n",
+    ")\n",
+    "plt.show()\n",
+    "\n",
+    "# Parameter Recovery Assessment\n",
+    "true_params = true_saturation.get_params()\n",
+    "est_params = result_arimax.treatments[0].saturation.get_params()\n",
+    "true_adstock_params = true_adstock.get_params()\n",
+    "est_adstock_params = result_arimax.treatments[0].adstock.get_params()\n",
+    "\n",
+    "print(\"\\nParameter Recovery Assessment:\")\n",
+    "print(f\"Saturation - lam error: {abs(est_params['lam'] - true_params['lam']):.2f}\")\n",
+    "print(\n",
+    "    f\"Adstock - half_life error: {abs(est_adstock_params['half_life'] - true_adstock_params['half_life']):.2f} weeks\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Interpretation:**\n",
+    "\n",
+    "- **Saturation curve** (left): Shows how raw communication intensity (0-10) gets transformed by diminishing returns. The curve flattens at higher intensities, meaning the 10th message has much less impact than the 1st.\n",
+    "\n",
+    "- **Adstock weights** (right): Shows how a communication \"impulse\" at week 0 affects water consumption over the following weeks. The bars show the relative contribution of each lag.\n",
+    "\n",
+    "- **Parameter recovery**: In this simulated example with known ground truth, we can assess how well the estimation recovered the true parameters. The ARIMAX model should recover similar transform parameters as the HAC model, since both use the same estimation procedure for transforms.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Model Summary\n",
+    "\n",
+    "View the fitted model coefficients and their standard errors. Note the ARIMA order is displayed:\n"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1224,7 +1286,6 @@
     }
    ],
    "source": [
-    "# View summary - note the ARIMA order is displayed\n",
     "result_arimax.summary(round_to=2)"
    ]
   },
@@ -1448,7 +1509,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1476,15 +1537,9 @@
     }
    ],
    "source": [
-    "# Visualize the counterfactual analysis\n",
-    "fig, ax = result_estimated.plot_effect(effect_result)\n",
-    "plt.show()\n",
-    "\n",
-    "print(\n",
-    "    f\"\\nThe communications policy saved approximately {-effect_result['total_effect']:.0f} ML of water \"\n",
-    "    f\"over the 3-year period, representing a {-100 * effect_result['mean_effect'] / df['water_consumption'].mean():.1f}% \"\n",
-    "    f\"reduction in average consumption.\"\n",
-    ")"
+    "# Note: This cell was removed - the HAC counterfactual visualization\n",
+    "# is now properly located in the HAC section above\n",
+    "pass"
    ]
   },
   {