AugmentedBalancing

Composable outcome-model and balancing corrections for panel ATT

from _api_doc_utils import *

1 Where it fits

Group: Causal inference

AugmentedBalancing combines a supplied untreated-outcome surface with unit and time balancing. It implements the panel residual-balancing framework that contains outcome-only imputation, unit balancing, time balancing, double balancing, and their outcome-model-augmented versions in one API.

The class uses the common fit(Y, W, outcome_model=None) panel contract. Y is a finite balanced outcome matrix. W is a same-shaped binary absorbing treatment matrix. outcome_model, when supplied, is a finite same-shaped matrix of untreated-outcome predictions. The class treats this matrix as a fitted nuisance input; it does not fit the outcome model itself.

2 Estimator

For one treated unit \(i\), one target period \(t\), never-treated controls \(C\), and pre-treatment periods \(P\), define residuals

\[ R_{js}=Y_{js}-\hat m_{js}. \]

The double-balanced counterfactual is

\[ \widehat Y_{it}(0) = \hat m_{it} +\sum_{j\in C}\hat\omega_j R_{jt} +\sum_{s\in P}\hat\lambda_s R_{is} -\sum_{j\in C}\sum_{s\in P}\hat\omega_j\hat\lambda_s R_{js}. \]

balance="none" returns \(\hat m_{it}\) without a balancing correction. Every other mode uses the full formula. balance="unit" optimizes \(\omega\) and uses uniform \(\lambda\); balance="time" uses uniform \(\omega\) and optimizes \(\lambda\); balance="double" optimizes both. If no outcome-model matrix is supplied, \(\hat m=0\).

For each treatment cohort, simplex unit weights match the cohort-average pre-period path by default. unit_target="individual" fits one unit-weight vector for each treated unit. time_target="all" fits one time-weight vector against the mean control post path; time_target="period" fits one vector for every post period. balance_on="raw" fits weights with \(Y\) and then applies them to residuals \(R\). balance_on="residual" fits the same problems with \(R\).

unit_loss="ridge" uses the SDID ridge loss. unit_loss="penalized_scm" standardizes the donor matrix and adds the linear donor-distance penalty from the R penSCM implementation. unit_penalty=1e-4 is its default. Time weights always use the SDID ridge loss.

The unit and time weight problems use the same profiled-intercept simplex least-squares solver as SyntheticDID. Omitted penalties use the cohort-specific control first-difference scale:

\[ \zeta_\omega=(N_{\mathrm{target}}T_{\mathrm{post}})^{1/4}\hat\sigma, \qquad \zeta_\lambda=10^{-6}\hat\sigma. \]

For individual targeting, \(N_{\mathrm{target}}=1\). The ATT is the simple mean of \(Y_{it}-\widehat Y_{it}(0)\) over cells where \(W_{it}=1\).

3 Configuration map

Estimator	Constructor and fit choice
Outcome model	balance="none", supply outcome_model
Unit balancing	balance="unit", omit outcome_model; time weights are uniform
Time balancing	balance="time", omit outcome_model; unit weights are uniform
Double balancing	balance="double", omit outcome_model
Augmented unit balancing	balance="unit", supply outcome_model
Augmented double balancing	balance="double", supply outcome_model
Residual-first augmented double balancing	add balance_on="residual"
Individual-target augmented double balancing	add unit_target="individual", time_target="period"
Penalized-SCM unit loss	add unit_loss="penalized_scm"

4 Scope and diagnostics

Every treated cohort must have a pre-period and at least one never-treated donor. The class supports staggered absorbing adoption by fitting each cohort separately. It does not fit nuisance outcome models, select their tuning parameters, cross-fit, or provide analytic or resampling inference. Treat the supplied surface as a causal nuisance estimate only when it was fit without treated post-treatment outcomes.

summary() returns ATT, panel counterfactuals and effects, event-study and group summaries, fitted unit and time weights, weight-target mappings, penalties, pre-period RMSE, and the supplied outcome-model surface. target_units and target_cohorts map unit-weight rows. time_target_periods uses -1 for an all-post target and the zero-based post-period index for period-specific targets. time_target_cohorts maps time-weight rows to cohorts.

5 Python API

Constructor: cm.AugmentedBalancing

print(inspect.signature(cm.AugmentedBalancing))

(balance='double', unit_target='cohort', time_target='all', balance_on='raw', unit_loss='ridge', unit_penalty=0.0001, zeta_omega=None, zeta_lambda=None, max_iterations=1000)

Parameter	Default	Contract
balance	"double"	"none", "unit", "time", or "double"
unit_target	"cohort"	"cohort" or "individual"
time_target	"all"	"all" or "period"
balance_on	"raw"	"raw" or "residual"
unit_loss	"ridge"	"ridge" or "penalized_scm"
unit_penalty	0.0001	Nonnegative penalized-SCM donor-distance coefficient
zeta_omega	None	Optional nonnegative unit-ridge scale; None uses the data-dependent rule
zeta_lambda	None	Optional nonnegative time-ridge scale; None uses the data-dependent rule
max_iterations	1000	Positive active-set iteration budget

fit(y, w, outcome_model=None) fits in place and returns None. predict() and treatment_effect() return arrays with the same shape as y. summary() returns the full fitted-state dictionary described below. These three accessors raise ValueError before a successful fit.

Constructor choices are validated when the class is created. During fit, y and the optional outcome model must be finite, two-dimensional, and same-shaped. w must be a same-shaped binary absorbing-treatment matrix. Each treated cohort must have a pre-period and the panel must contain a never-treated donor. Invalid arrays, unsupported treatment paths, and solver failures raise ValueError.

cls = cm.AugmentedBalancing
display(HTML(html_table(["Public method"], public_methods(cls))))

Public method
fit(self, /, y, w, outcome_model=None)
predict(self, /)
summary(self, /)
treatment_effect(self, /)

6 Minimal example

rng = np.random.default_rng(24)
n_control, n_treated, n_pre, n_post = 8, 2, 10, 4
controls = rng.normal(size=(n_control, n_pre + n_post))
weights = rng.dirichlet(np.ones(n_control))
untreated = weights @ controls
treated = untreated + np.r_[np.zeros(n_pre), np.full(n_post, 0.8)]
y = np.vstack([controls, treated - 0.1, treated + 0.1])
w = np.zeros_like(y)
w[n_control:, n_pre:] = 1.0

# A simple supplied nuisance surface. A production analysis should estimate it
# without using treated post-treatment outcomes.
outcome_model = np.zeros_like(y)
model = cm.AugmentedBalancing(
    balance="double",
    unit_target="cohort",
    time_target="all",
    balance_on="residual",
    zeta_omega=0.01,
    zeta_lambda=0.01,
)
model.fit(y, w, outcome_model)
print(model.summary()["att"])
print(model.predict().shape)

0.8001119872078419
(10, 14)

7 summary() contract

summary = model.summary()
display(HTML(html_table(["summary() key", "shape"], summary_shape_rows(summary))))

summary() key	shape
att	()
unit_weights	(1, 10)
time_weights	(1, 14)
counterfactual	(10, 14)
treatment_effect	(10, 14)
outcome_model	(10, 14)
event_study	()
group_means	()
pre_rmse	()
zeta_omega	(1,)
zeta_lambda	(1,)
target_units	(1,)
target_cohorts	(1,)
time_target_periods	(1,)
time_target_cohorts	(1,)
control_units	(8,)
treated_units	(2,)
cohorts	(1,)
balance	()
unit_target	()
time_target	()
balance_on	()
unit_loss	()
unit_penalty	()
converged	()