AugmentedBalancing

Composable outcome-model and balancing corrections for panel ATT

Code 
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.

For an end-to-end staggered-adoption workflow, including a MatrixCompletion nuisance surface, specification comparisons, event-time plots, and weight diagnostics, see Augmented Balancing for Panel Data.

2 Data and treatment contract

Input Shape Requirements
y (n_units, n_periods) Finite, nonempty, two-dimensional balanced outcome matrix
w Same as y Finite binary values; treatment must be absorbing within each row
outcome_model Same as y Optional finite untreated-outcome prediction surface; omitted means an all-zero surface

The treatment parser infers each treated unit’s first-treatment period, the distinct adoption cohorts, and the never-treated donor rows. Every fitted panel must contain at least one treated unit and one never-treated donor. Every cohort must start after period zero so that a pre-treatment history exists. These requirements still apply to balance="none", because ATT and cohort/event-time summaries are defined from the same panel treatment contract.

The estimator does not mask or refit outcome_model. Callers are responsible for constructing it without leakage from treated post-treatment outcomes. Passing a same-shaped finite matrix verifies the array contract, not its causal validity.

3 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\).

4 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"

5 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.

6 Computational scaling

With cohort unit targets, the estimator solves one unit-weight problem per adoption cohort. With individual targets, it solves one per treated unit. time_target="all" solves one time-weight problem per cohort, while time_target="period" solves one for every cohort-specific post-treatment period. Unit problems have one simplex coordinate per never-treated donor; time problems have one coordinate per cohort-specific pre-period.

The implementation stores dense weight matrices in original panel coordinates. Non-donor columns in unit_weights and non-pre-period columns in time_weights are zero. counterfactual and treatment_effect have the same shape as y, with finite rows for ever-treated units and NaN rows for never-treated donors. A successful fit reports converged=True; an invalid input or exhausted/failed weight solve raises ValueError instead of returning a partial fit.

7 Python API

Constructor: cm.AugmentedBalancing

Code 
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.

Method Return value Meaning
fit(y, w, outcome_model=None) None Validate the panel, fit cohort/target-specific weights, and store counterfactual state
predict() (n_units, n_periods) array Estimated untreated counterfactual surface for treated rows; donor rows are NaN
treatment_effect() (n_units, n_periods) array Cellwise y - predict() for treated rows; donor rows are NaN
summary() dict ATT, diagnostics, weights, mappings, configuration, and event/cohort summaries

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.

Code 
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, /)

8 Minimal example

Code 
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.8001119872078418
(10, 14)

9 summary() contract

Key Contract
att Scalar mean of treatment_effect over cells with w == 1
unit_weights Dense (n_unit_targets, n_units) matrix; only never-treated donor columns can be nonzero
time_weights Dense (n_time_targets, n_periods) matrix; only the target cohort’s pre-period columns can be nonzero
counterfactual Same-shaped untreated prediction matrix returned by predict()
treatment_effect Same-shaped observed-minus-counterfactual matrix returned by treatment_effect()
outcome_model The supplied nuisance surface, or the all-zero surface used when it was omitted
event_study Nested unweighted and treated-count-weighted event-time dictionaries with event_time, estimate, and n arrays
group_means Cohort/event-time rows plus nested unweighted and weighted aggregates of treated means, counterfactual means, and effects
pre_rmse Root-mean-square pre-treatment gap over treated-unit counterfactual rows
zeta_omega, zeta_lambda Fitted regularization scales in unit-target and time-target row order; unused or penalized-SCM unit scales are zero
target_units Unit-weight row targets; -1 denotes a cohort-average target
target_cohorts Zero-based adoption period for every unit-weight row
time_target_periods Time-weight row targets; -1 denotes one all-post target
time_target_cohorts Zero-based adoption period for every time-weight row
control_units, treated_units, cohorts Inferred zero-based panel row indices and distinct adoption periods
balance, unit_target, time_target, balance_on, unit_loss, unit_penalty Echo of the fitted constructor configuration
converged True for a completed fit; failed solves raise instead of returning a false flag

The live example below verifies the exact keys and runtime shapes against the installed extension.

Code 
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 ()