BalancingWeights

Calibration weights for covariate balance

from _api_doc_utils import *

1 Where it fits

Group: Causal inference

BalancingWeights chooses weights for a source sample so weighted source covariate means match a target sample:

\[ \sum_i w_i x_i / \sum_i w_i \approx \sum_j q_j x_j / \sum_j q_j. \]

The objective can be quadratic or entropy-like, with optional lower/upper bounds and ridge stabilization.

2 Calibration system and numerical objective

Let \(\bar x_T\) be the optionally weighted target mean and define calibration rows

\[ z_i=(1,\ x_i-\bar x_T)'. \]

Weights are normalized implicitly by solving \(Z'w=e_1\). For entropy calibration, the dual map is

\[ w_i(\beta)= \begin{cases} \exp(z_i'\beta), & \text{without baseline weights},\\ q_i\exp(z_i'\beta-1), & \text{with normalized baseline }q_i, \end{cases} \]

clipped to the configured box when the bounded phase is needed. For quadratic calibration it is

\[ w_i(\beta)= \begin{cases} z_i'\beta, & \text{without baseline weights},\\ q_i(1+z_i'\beta), & \text{with baseline }q_i, \end{cases} \]

again clipped to the weight box. Thus the objective name selects the primal calibration geometry and corresponding dual weight map; it is not the literal function passed to the optimizer.

For allowed L2 imbalance \(\delta\), define

\[ s(\beta_{-0}) = \begin{cases} \delta\,\beta_{-0}/\|\beta_{-0}\|_2,&\|\beta_{-0}\|_2>0,\\ 0,&\text{otherwise}. \end{cases} \]

The numerical solver minimizes

\[ \frac12\|r(\beta)\|_2^2, \qquad r(\beta)=Z'w(\beta)-e_1+(0,s(\beta_{-0})')'. \]

Gauss-Newton uses the analytic residual Jacobian with a small \(10^{-8}\) normal-equation ridge and backtracking. Automatic mode tries Gauss-Newton first and falls back to BFGS only when needed. Entropy fits begin with relaxed bounds and rerun with clipping if the first solution violates the requested box.

3 Implementation walkthrough

The class solves a package-owned nonlinear calibration system; the named objective determines the primal-to-dual weight map, while both solver choices minimize squared calibration residuals.

1. Optional baseline and target weights are validated and normalized to sum one. With autoscaling, each source column is mapped by its source minimum and range and the identical affine map is applied to target rows. A source range at or below \(10^{-12}\) maps both samples to zero in that column; target values are allowed outside \([0,1]\).
2. The target mean is computed in fitting coordinates and the dense design \(Z=[\mathbf1,X-\bar x_T]\) is built. Entropy starts its dual vector at zero. Quadratic calibration instead solves a linearized calibration system for its starting vector, with a baseline-weight adjustment when supplied.
3. Given \(\beta\), one pass computes \(Z\beta\), the entropy or quadratic weight map, and the derivative of each weight with respect to its scalar index. Clipped weights receive derivative zero. The residual is \(Z'w-e_1\) plus the radial L2 slack, and the analytic Jacobian is \(Z'\operatorname{diag}(w')Z\) plus the slack Jacobian.
4. Native Gauss-Newton forms \(J'J+10^{-8}I\), solves its least-squares system for the step, and halves the step until the residual sum of squares strictly decreases or the scale falls below \(10^{-8}\). Only the scaled residual norm reaching tolerance counts as convergence; small steps or small objective changes alone are explicitly labeled nonconverged.
5. The BFGS alternative exposes \(\|r(\beta)\|^2/2\) and \(J'r\) to Argmin, initializes the inverse Hessian at identity, and uses More-Thuente search. It is accepted as converged only when both Argmin’s status and the package’s residual-norm check pass.
6. In solver='auto', a nonconverged Gauss-Newton result seeds BFGS. BFGS replaces it only if the BFGS objective is no larger; otherwise the better Gauss-Newton iterate is retained. A finite nonconverged fit is therefore reportable for diagnostics.
7. Entropy first permits weights up to \(10^8\) and no lower clipping to preserve a smooth solve. If that result violates the requested box, the same solver sequence is rerun from its dual vector with bounded exponential clipping. Quadratic calibration uses the requested box from the start.
8. Final success separately requires solver convergence, unit weight sum within \(10^{-6}\), finite weights, and box feasibility. Original-scale weighted means are recomputed from the unscaled covariates, so diagnostic balance is not inferred from the scaled residual.

This design makes infeasibility observable rather than exceptional: the object can retain the best finite weights with solver_converged=False or success=False. The zero derivative at active clipping bounds also explains why tight boxes can stall either local solver.

4 Diagnostics and inference

There is no outcome model, treatment-effect estimand, standard error, or covariance calculation. The fitted object returns normalized weights, original-scale mean differences, effective sample size

\[ \mathrm{ESS}=\frac{1}{\sum_i w_i^2}, \]

the dual coefficient, and solver diagnostics. solver_converged means that the scaled calibration residual satisfies the requested tolerance and, for BFGS, that the optimizer also reports a converged status. The generic converged, iterations, and termination_reason keys describe that same scaled-coordinate solve. success additionally requires finite normalized weights whose sum and box constraints are feasible.

With autoscaling, covariates are min-max scaled using source-sample minima and ranges before solving. Solver convergence is intentionally not reclassified by applying that scaled tolerance to unscaled covariates. scaled_residual_norm and solver_objective report the numerical problem; original_balance_l2 and original_balance_max_abs report balance in the units supplied by the user. The legacy aliases residual_norm, criterion, l2_diff, and max_abs_diff remain available. Large original-unit imbalance can therefore coexist with solver convergence when feature scales are large, and should be judged substantively rather than against a dimensionless solver tolerance.

5 Performance and numerical behavior

For \(n\) source rows and \(p\) covariates, each residual/Jacobian construction forms dense weighted cross-products with about \(O(np^2)\) work and \(O(np+p^2)\) storage; the Gauss-Newton system costs up to \(O(p^3)\). BFGS fallback adds repeated residual and Jacobian evaluations. Tight box constraints or an infeasible balance tolerance can leave a finite set of weights with success=False rather than raising. Constant source columns are mapped to zero under autoscaling, and target values outside the source range are not clipped. Very disparate original feature scales favor autoscale=True, but original-unit diagnostics must still be interpreted feature by feature.

6 Python API

Constructor: cm.BalancingWeights

Use fit(covariates, target_covariates, baseline_weights=None, target_weights=None). get_weights() returns the fitted source weights. The read-only solver_converged property reports scaled solver convergence, while success also checks weight feasibility. summary() reports both statuses, scaled numerical diagnostics, original-scale balance, the dual coefficients, and effective sample size.

print(inspect.signature(cm.BalancingWeights))

(objective='quadratic', solver='auto', autoscale=False, min_weight=0.0, max_weight=1.0, l2_norm=0.0, max_iterations=200, tolerance=1e-08)

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

Public method
fit(self, /, covariates, target_covariates, baseline_weights=None, target_weights=None)
get_weights(self, /)
summary(self, /)

7 Minimal example

rng = np.random.default_rng(10)
x0 = rng.normal(size=(180, 3))
x1 = rng.normal(loc=np.array([0.4, -0.2, 0.1]), size=(80, 3))
model = cm.BalancingWeights(objective='quadratic', max_iterations=500)
model.fit(x0, x1)
fit = model.summary()
print({key: fit[key] for key in ['solver_converged', 'success', 'scaled_residual_norm', 'original_balance_l2']})
print(fit['mean_diff'])
print(fit['effective_sample_size'])
print(model.get_weights()[:5])

{'solver_converged': True, 'success': True, 'scaled_residual_norm': 2.4007203514854626e-16, 'original_balance_l2': 1.712375058312158e-16}
[1.11022302e-16 5.55111512e-17 1.17961196e-16]
155.1166258737379
[0.00415802 0.00646452 0.00614945 0.00537598 0.00450626]

8 summary() contract

The table below is generated by fitting the live class in this repository and then inspecting summary(). Shapes are shown because most values are plain NumPy arrays or scalars.

rng = np.random.default_rng(110)
x0 = rng.normal(size=(80, 3))
x1 = rng.normal(loc=np.array([0.4, -0.2, 0.1]), size=(40, 3))
model = cm.BalancingWeights()
model.fit(x0, x1)
summary = model.summary()
display(HTML(html_table(["summary() key", "shape"], summary_shape_rows(summary))))

summary() key	shape
weights	(80,)
beta	(4,)
objective	()
solver	()
success	()
nit	()
criterion	()
residual_norm	()
solver_converged	()
scaled_residual_norm	()
solver_objective	()
converged	()
iterations	()
termination_reason	()
weight_sum	()
weighted_mean	(3,)
target_mean	(3,)
mean_diff	(3,)
l2_diff	()
original_balance_l2	()
max_abs_diff	()
original_balance_max_abs	()
effective_sample_size	()
min_weight	()
max_weight	()
l2_norm	()