Augmented Balancing Across Panel DGPs

A Python reproduction of the simulation plots in the slide deck

This report reproduces the data-generating processes and plot structure in ../drafts/slides.pdf. It uses the standalone panel_dgps package in this folder and the AugmentedBalancing estimator in the linked crabbymetrics checkout.

The deck compares five DGP families: strong factors, sparse synthetic-control weights, stationary and integrated time series, weak factors, and mixed factors. Each design has 160 control units, 40 treated units, 40 pre-treatment periods, and 10 post-treatment periods. Treatment effects increase by 0.2 in each post-treatment period.

Setup

from __future__ import annotations

import os
import time
from dataclasses import dataclass
from typing import Callable

import crabbymetrics as cm
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

from panel_dgps import (
    PanelConfig,
    PanelData,
    classic_factor,
    mixed_factor,
    synthetic_control,
    time_series,
    weak_factor,
)

mpl.rcParams.update(
    {
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.titleweight": "bold",
        "figure.dpi": 120,
        "font.size": 10,
    }
)

CONTROL_COLOR = "#8a8f98"
TREATED_COLOR = "#2f6db3"
EFFECT_COLOR = "#c94435"
BACKGROUND = "#fbfaf7"


@dataclass(frozen=True)
class Scenario:
    key: str
    title: str
    generator: Callable[..., PanelData]
    kwargs: dict

    def draw(self, config: PanelConfig, seed: int) -> PanelData:
        return self.generator(config=config, seed=seed, **self.kwargs)


scenario_groups = {
    "Strong factor model": [
        Scenario("classic_good", "Good overlap: a = 0", classic_factor, {"overlap": 0.0}),
        Scenario("classic_poor", "Poor overlap: a = 1", classic_factor, {"overlap": 1.0}),
    ],
    "Sparse synthetic-control model": [
        Scenario("sc_sparse", "Active controls: 10%", synthetic_control, {"active_share": 0.1}),
        Scenario("sc_dense", "Active controls: 50%", synthetic_control, {"active_share": 0.5}),
    ],
    "Time-series model": [
        Scenario("ar02", "Stationary: phi = 0.2", time_series, {"coefficient": 0.2, "integrated": False}),
        Scenario("ar09", "Stationary: phi = 0.9", time_series, {"coefficient": 0.9, "integrated": False}),
        Scenario("i102", "Integrated: phi = 0.2", time_series, {"coefficient": 0.2, "integrated": True}),
        Scenario("i109", "Integrated: phi = 0.9", time_series, {"coefficient": 0.9, "integrated": True}),
    ],
    "Weak factor model": [
        Scenario("weak_good", "Good overlap: a = 0", weak_factor, {"overlap": 0.0}),
        Scenario("weak_poor", "Poor overlap: a = 1", weak_factor, {"overlap": 1.0}),
    ],
    "Mixed factor model": [
        Scenario("mixed_good", "Good overlap: a = 0", mixed_factor, {"overlap": 0.0}),
        Scenario("mixed_poor", "Poor overlap: a = 1", mixed_factor, {"overlap": 1.0}),
    ],
}

all_scenarios = [scenario for group in scenario_groups.values() for scenario in group]

Data-generating processes

Strong factor model

The untreated outcome is

$$Y_{it}(0)=\lambda_i'f_t+\varepsilon_{it},$$

where $f_t$ contains two scaled trending AR(1) factors. Treated loadings have mean $a$, control loadings have mean $-a$, and each loading has unit variance. Thus, $a=0$ gives overlap and $a=1$ gives poor overlap.

Sparse synthetic-control model

Control outcomes follow the same two-factor model. Each treated untreated outcome is an equal-weighted mean of the latent signals from the controls with the largest loading sums, plus independent noise. The active donor share is 10% or 50%.

Time-series model

The stationary design follows

$$Y_{it}(0)=c_i+\phi Y_{i,t-1}(0)+\varepsilon_{it}.$$

The integrated design applies the same recursion to $\Delta Y_{it}(0)$. The deck sets $c_i=0$ for controls and $c_i=0.25/(1-\phi)$ for treated units. We use $\phi\in\{0.2,0.9\}$.

Weak and mixed factor models

The weak design uses ten factors and scales the factor signal by 0.2. Five factors trend, and five follow a sine path. The mixed design assigns the trending factors to all treated units and 60 controls. The remaining 100 controls follow the periodic factors.

Raw panel paths

The next plots use the deck dimensions and one fixed draw per design. Gray lines are controls. Blue lines are treated units. The shaded region begins with treatment, and the red line shows the treated-group mean effect.

def plot_paths(ax, panel: PanelData, title: str) -> None:
    periods = np.arange(panel.outcome.shape[1]) + 1
    for unit in panel.control_units:
        ax.plot(periods, panel.outcome[unit], color=CONTROL_COLOR, alpha=0.18, lw=0.7)
    for unit in panel.treated_units:
        ax.plot(periods, panel.outcome[unit], color=TREATED_COLOR, alpha=0.24, lw=0.8)
    mean_effect = panel.treatment_effect[panel.treated_units].mean(axis=0)
    ax.plot(periods, mean_effect, color=EFFECT_COLOR, lw=2.0, label="Mean treatment effect")
    ax.axvspan(panel.n_pre + 0.5, periods[-1] + 0.5, color="#d8d8d8", alpha=0.45)
    ax.axvline(panel.n_pre + 0.5, color="#555555", lw=1.0, ls="--")
    ax.set_title(title)
    ax.set_xlabel("Period")
    ax.set_ylabel("Outcome")
    ax.grid(color="#dedad2", alpha=0.55, lw=0.6)
    ax.set_facecolor(BACKGROUND)


raw_config = PanelConfig()
raw_scenarios = [
    ("Strong factor - good overlap", scenario_groups["Strong factor model"][0]),
    ("Strong factor - poor overlap", scenario_groups["Strong factor model"][1]),
    ("Synthetic control - 10% active", scenario_groups["Sparse synthetic-control model"][0]),
    ("Synthetic control - 50% active", scenario_groups["Sparse synthetic-control model"][1]),
    ("Time series - stationary, phi = 0.2", scenario_groups["Time-series model"][0]),
    ("Time series - integrated, phi = 0.2", scenario_groups["Time-series model"][2]),
    ("Weak factor - good overlap", scenario_groups["Weak factor model"][0]),
    ("Weak factor - poor overlap", scenario_groups["Weak factor model"][1]),
    ("Mixed factor - good overlap", scenario_groups["Mixed factor model"][0]),
    ("Mixed factor - poor overlap", scenario_groups["Mixed factor model"][1]),
]

fig, axes = plt.subplots(5, 2, figsize=(14, 19), constrained_layout=True)
fig.patch.set_facecolor(BACKGROUND)
for index, (ax, (title, scenario)) in enumerate(zip(axes.flat, raw_scenarios)):
    panel = scenario.draw(raw_config, seed=8100 + index)
    plot_paths(ax, panel, title)
fig.suptitle("Raw panel draws from the slide-deck DGPs", fontsize=17, weight="bold")
plt.show()

Outcome-model surfaces

The augmented estimators need a same-shaped untreated-outcome prediction surface $\hat m$. The report uses four package-compatible nuisance fits:

• FE: an additive unit-and-time model fit only on untreated cells;
• IFE: iterative rank-two factor imputation that calls InteractiveFixedEffects;
• MC: MatrixCompletion with nuclear-norm shrinkage;
• H-Ridge: HorizontalPanelRidge, with observed never-treated outcomes retained as the donor surface.

The nuisance fits do not use treated post-treatment outcomes. This rule is essential. The report estimates each surface once per replication, then sends it to each augmented balancing variant.

def additive_surface(y: np.ndarray, w: np.ndarray, iterations: int = 80) -> np.ndarray:
    observed = w < 0.5
    grand_mean = float(y[observed].mean())
    unit_effect = np.zeros(y.shape[0])
    time_effect = np.zeros(y.shape[1])
    for _ in range(iterations):
        for unit in range(y.shape[0]):
            mask = observed[unit]
            unit_effect[unit] = np.mean(y[unit, mask] - grand_mean - time_effect[mask])
        unit_effect -= unit_effect.mean()
        for period in range(y.shape[1]):
            mask = observed[:, period]
            time_effect[period] = np.mean(y[mask, period] - grand_mean - unit_effect[mask])
    return grand_mean + unit_effect[:, None] + time_effect[None, :]


def interactive_surface(
    y: np.ndarray,
    w: np.ndarray,
    rank: int = 2,
    iterations: int = 30,
    tolerance: float = 1e-7,
) -> np.ndarray:
    observed = w < 0.5
    surface = additive_surface(y, w)
    filled = np.where(observed, y, surface)
    for _ in range(iterations):
        model = cm.InteractiveFixedEffects(rank=rank, force=3)
        model.fit(filled)
        candidate = np.asarray(model.predict())
        change = np.max(np.abs(candidate[~observed] - surface[~observed]))
        surface = candidate
        filled = np.where(observed, y, surface)
        if change < tolerance:
            break
    return surface


def matrix_completion_surface(y: np.ndarray, w: np.ndarray) -> np.ndarray:
    model = cm.MatrixCompletion(
        lambda_fraction=0.08,
        max_iterations=120,
        effect_iterations=3,
        tolerance=1e-5,
    )
    model.fit(y, w)
    return np.asarray(model.predict())


def horizontal_ridge_surface(y: np.ndarray, w: np.ndarray) -> np.ndarray:
    model = cm.HorizontalPanelRidge(penalty=1.0)
    model.fit(y, w)
    predicted = np.asarray(model.predict())
    surface = y.copy()
    treated_units = np.flatnonzero(w.sum(axis=1) > 0.0)
    surface[treated_units] = predicted[treated_units]
    return surface


outcome_fitters = {
    "FE": additive_surface,
    "IFE": interactive_surface,
    "MC": matrix_completion_surface,
    "H-Ridge": horizontal_ridge_surface,
}

Balancing estimators

For residuals $R_{it}=Y_{it}-\hat m_{it}$, the augmented double-balanced counterfactual is

$$\widehat Y_{it}(0) =\hat m_{it} +\sum_j\hat\omega_jR_{jt} +\sum_s\hat\lambda_sR_{is} -\sum_j\sum_s\hat\omega_j\hat\lambda_sR_{js}.$$

The heatmap rows correspond to the deck:

• unit balancing;
• outcome-model imputation;
• double balancing;
• augmented double balancing with residual-first cohort weights;
• augmented double balancing with residual-first individual weights;
• augmented double balancing with raw-outcome individual weights;
• augmented double balancing with raw-outcome cohort weights.

The simplex losses use scale-aware ridge penalties. If $\hat\sigma$ is the control first-difference standard deviation, the report normally sets $\zeta_\omega=100\hat\sigma$ and $\zeta_\lambda=3\hat\sigma$. The integrated process with $\phi=0.9$ uses $3000\hat\sigma$ and $100\hat\sigma$ because its high-dimensional individual-target problems otherwise reach the optimizer iteration limit. These fixed rules make the Monte Carlo comparison deterministic. The values are simulation settings, not recommended defaults for applied work.

ROW_LABELS = [
    "Unit balance",
    "Outcome",
    "Double balance",
    "Aug. double\nresidual, cohort",
    "Aug. double\nresidual, individual",
    "Aug. double\nraw, individual",
    "Aug. double\nraw, cohort",
]
COLUMN_LABELS = ["None", *outcome_fitters]


def estimator_specs(outcome_models: dict[str, np.ndarray]):
    yield 0, 0, {"balance": "unit"}, None
    yield 2, 0, {"balance": "double"}, None
    for column, name in enumerate(outcome_fitters, start=1):
        model_surface = outcome_models[name]
        yield 1, column, {"balance": "none"}, model_surface
        yield 3, column, {
            "balance": "double",
            "target": "cohort",
            "balance_on": "residual",
        }, model_surface
        yield 4, column, {
            "balance": "double",
            "target": "individual",
            "balance_on": "residual",
        }, model_surface
        yield 5, column, {
            "balance": "double",
            "target": "individual",
            "balance_on": "raw",
        }, model_surface
        yield 6, column, {
            "balance": "double",
            "target": "cohort",
            "balance_on": "raw",
        }, model_surface


def estimate_att(
    y,
    w,
    options,
    outcome_model,
    zeta_omega_scale=100.0,
    zeta_lambda_scale=3.0,
):
    control_units = np.flatnonzero(w.sum(axis=1) == 0.0)
    treated_units = np.flatnonzero(w.sum(axis=1) > 0.0)
    first_treatment = min(np.flatnonzero(w[unit] > 0.5)[0] for unit in treated_units)
    differences = np.diff(y[np.ix_(control_units, np.arange(first_treatment))], axis=1)
    noise_scale = max(float(differences.std(ddof=1)), 1e-4)
    estimator = cm.AugmentedBalancing(
        zeta_omega=zeta_omega_scale * noise_scale,
        zeta_lambda=zeta_lambda_scale * noise_scale,
        max_iterations=5000,
        **options,
    )
    estimator.fit(y, w, outcome_model)
    return float(estimator.summary()["att"])

Monte Carlo profile

The document renders with a verification profile by default. It uses the deck dimensions - 160 controls, 40 treated units, 40 pre-treatment periods, and 10 post-treatment periods - with 12 replications. This profile checks every DGP, nuisance model, and estimator combination in a practical render time.

Set AUGBAL_PROFILE=deck before rendering to use the original dimensions and 100 replications:

AUGBAL_PROFILE=deck quarto render balancing-reproduction.qmd

Both profiles use fixed seeds and the same code path.

profile = os.environ.get("AUGBAL_PROFILE", "verification")
if profile == "verify":
    profile = "verification"
if profile == "deck":
    mc_config = PanelConfig()
    replications = 100
elif profile == "verification":
    mc_config = PanelConfig()
    replications = 12
else:
    raise ValueError("AUGBAL_PROFILE must be 'verify', 'verification', or 'deck'")

print(
    f"Simulation profile: {profile}; "
    f"N={mc_config.n_units}, T={mc_config.n_periods}, replications={replications}"
)

Simulation profile: verification; N=200, T=50, replications=12

RMSE simulation

Each replication sends one panel draw through every outcome model and estimator. RMSE uses the known average treatment effect for that draw. A failed fit remains missing and increments the cell-specific failure count.

def run_scenario(scenario: Scenario, config: PanelConfig, reps: int, base_seed: int):
    estimates = [[[] for _ in COLUMN_LABELS] for _ in ROW_LABELS]
    failures = np.zeros((len(ROW_LABELS), len(COLUMN_LABELS)), dtype=int)
    truth = None
    start = time.perf_counter()
    penalty_scales = (3000.0, 100.0) if scenario.key == "i109" else (100.0, 3.0)

    for replication in range(reps):
        panel = scenario.draw(config, seed=base_seed + replication)
        truth = panel.true_att
        outcome_models = {}
        for name, fitter in outcome_fitters.items():
            try:
                outcome_models[name] = fitter(panel.outcome, panel.treatment)
            except Exception:
                outcome_models[name] = np.full_like(panel.outcome, np.nan)

        for row, column, options, outcome_model in estimator_specs(outcome_models):
            if outcome_model is not None and not np.all(np.isfinite(outcome_model)):
                failures[row, column] += 1
                continue
            try:
                estimate = estimate_att(
                    panel.outcome,
                    panel.treatment,
                    options,
                    outcome_model,
                    zeta_omega_scale=penalty_scales[0],
                    zeta_lambda_scale=penalty_scales[1],
                )
            except Exception:
                failures[row, column] += 1
                continue
            if np.isfinite(estimate):
                estimates[row][column].append(estimate)
            else:
                failures[row, column] += 1

    rmse = np.full((len(ROW_LABELS), len(COLUMN_LABELS)), np.nan)
    for row in range(len(ROW_LABELS)):
        for column in range(len(COLUMN_LABELS)):
            values = np.asarray(estimates[row][column], dtype=float)
            if values.size:
                rmse[row, column] = np.sqrt(np.mean((values - truth) ** 2))

    elapsed = time.perf_counter() - start
    print(
        f"Completed {scenario.title}: {reps} replications in {elapsed:.1f} seconds; "
        f"failed fits={int(failures.sum())}"
    )
    return {"rmse": rmse, "failures": failures, "truth": truth, "elapsed": elapsed}


simulation_results = {}
for scenario_index, scenario in enumerate(all_scenarios):
    simulation_results[scenario.key] = run_scenario(
        scenario,
        mc_config,
        replications,
        base_seed=920_000 + 10_000 * scenario_index,
    )

Completed Good overlap: a = 0: 12 replications in 3.6 seconds; failed fits=0
Completed Poor overlap: a = 1: 12 replications in 4.3 seconds; failed fits=0
Completed Active controls: 10%: 12 replications in 3.7 seconds; failed fits=0
Completed Active controls: 50%: 12 replications in 3.6 seconds; failed fits=0
Completed Stationary: phi = 0.2: 12 replications in 3.3 seconds; failed fits=0
Completed Stationary: phi = 0.9: 12 replications in 4.8 seconds; failed fits=0
Completed Integrated: phi = 0.2: 12 replications in 5.9 seconds; failed fits=0
Completed Integrated: phi = 0.9: 12 replications in 4.0 seconds; failed fits=0
Completed Good overlap: a = 0: 12 replications in 3.6 seconds; failed fits=0
Completed Poor overlap: a = 1: 12 replications in 3.7 seconds; failed fits=0
Completed Good overlap: a = 0: 12 replications in 3.5 seconds; failed fits=0
Completed Poor overlap: a = 1: 12 replications in 4.0 seconds; failed fits=0

RMSE heatmaps

The cells show RMSE. Blank cells are combinations that the deck does not define. A small F suffix marks the number of failed fits when it is positive. Each family uses one color scale across its panels.

def plot_heatmap_family(family: str, scenarios: list[Scenario]) -> None:
    n_panels = len(scenarios)
    n_columns = 2 if n_panels > 1 else 1
    n_rows = int(np.ceil(n_panels / n_columns))
    fig, axes = plt.subplots(
        n_rows,
        n_columns,
        figsize=(7.2 * n_columns, 5.2 * n_rows),
        squeeze=False,
        constrained_layout=True,
    )
    matrices = [simulation_results[scenario.key]["rmse"] for scenario in scenarios]
    finite = np.concatenate([matrix[np.isfinite(matrix)] for matrix in matrices])
    color_max = max(float(np.quantile(finite, 0.95)), 1e-6)
    cmap = mpl.colormaps["YlGnBu_r"].copy()
    cmap.set_bad("#eeeeec")

    image = None
    for ax, scenario, matrix in zip(axes.flat, scenarios, matrices):
        image = ax.imshow(matrix, cmap=cmap, vmin=0.0, vmax=color_max, aspect="auto")
        failed = simulation_results[scenario.key]["failures"]
        for row in range(matrix.shape[0]):
            for column in range(matrix.shape[1]):
                value = matrix[row, column]
                if not np.isfinite(value):
                    continue
                suffix = f"\n{failed[row, column]}F" if failed[row, column] else ""
                text_color = "white" if value < 0.45 * color_max else "#202020"
                ax.text(
                    column,
                    row,
                    f"{value:.3f}{suffix}",
                    ha="center",
                    va="center",
                    fontsize=8.5,
                    color=text_color,
                    weight="bold",
                )
        ax.set_xticks(np.arange(len(COLUMN_LABELS)), COLUMN_LABELS, rotation=25, ha="right")
        ax.set_yticks(np.arange(len(ROW_LABELS)), ROW_LABELS)
        ax.set_title(f"{scenario.title}; true ATT = {simulation_results[scenario.key]['truth']:.2f}")
        ax.set_xlabel("Outcome model")
        ax.set_ylabel("Balancing method")
        ax.tick_params(length=0)

    for ax in axes.flat[n_panels:]:
        ax.set_visible(False)
    fig.colorbar(image, ax=axes.ravel().tolist(), label="RMSE", shrink=0.82)
    fig.suptitle(family, fontsize=16, weight="bold")
    plt.show()

Strong factor model

plot_heatmap_family("Strong factor model", scenario_groups["Strong factor model"])

Sparse synthetic-control model

plot_heatmap_family(
    "Sparse synthetic-control model",
    scenario_groups["Sparse synthetic-control model"],
)

Time-series model

plot_heatmap_family("Time-series model", scenario_groups["Time-series model"])

Weak factor model

plot_heatmap_family("Weak factor model", scenario_groups["Weak factor model"])

Mixed factor model

plot_heatmap_family("Mixed factor model", scenario_groups["Mixed factor model"])

Reading the reproduction

The report is a reproducible implementation of the deck design, not a claim that 12 replications give stable Monte Carlo rankings. Use the deck profile before making numerical comparisons. The verification profile has a narrower purpose: it proves that every DGP, nuisance fit, balancing mode, and plot executes at the original panel dimensions through the public APIs.

The heatmaps also separate two questions that are easy to conflate. balance_on selects the data used to fit weights. The augmentation formula always applies the fitted weights to outcome-model residuals. target selects whether one unit-weight vector matches a cohort mean or whether each treated unit receives its own vector.