BalancingWeights

Calibration weights for covariate balance

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, entropy-like, or Cressie-Read power-divergence-like, with optional lower/upper bounds, approximate balance, and dual 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 Cressie-Read calibration, missing baseline weights default to uniform source weights \(q_i=1/n\). The dual map is

\[ w_i(\beta)=q_i h_\lambda(z_i'\beta), \qquad h_\lambda(u)= \begin{cases} \exp(u), & |\lambda|\approx 0,\\ (1+\lambda u)_+^{1/\lambda}, & \text{otherwise}, \end{cases} \]

where \(\lambda\) is divergence_power. Thus divergence_power=0 is the entropy/KL limit, while divergence_power=1 is the Pearson chi-square member of the Cressie-Read family.

Rényi divergence for two normalized weight vectors \(p\) and \(q\) is

\[ D_\alpha(p\|q)=\frac{1}{\alpha-1}\log\sum_i p_i^\alpha q_i^{1-\alpha}. \]

It shares the same power moment as Cressie-Read after setting \(\lambda=\alpha-1\). In the usual Cressie-Read statistic,

\[ CR_\lambda(p,q)=\frac{2}{\lambda(\lambda+1)} \left\{\sum_i p_i^{\lambda+1}q_i^{-\lambda}-1\right\}, \]

so \(D_\alpha\) and \(CR_{\alpha-1}\) are monotone transformations of the same scalar moment, with the usual KL limit at \(\alpha=1,\lambda=0\). BalancingWeights currently implements the Cressie-Read calibration map; it does not yet expose a separate objective="renyi" 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+\frac{\eta}{2}\|\beta_{-0}\|_2^2, \qquad r(\beta)=Z'w(\beta)-e_1+(0,s(\beta_{-0})')'. \]

Here \(\eta\) is dual_ridge, which defaults to zero. Gauss-Newton uses the analytic residual Jacobian with a small \(10^{-8}\) normal-equation ridge and backtracking. solver="lbfgs" uses argmin’s L-BFGS with the same analytic gradient. Automatic mode tries Gauss-Newton first, then L-BFGS, and falls back to BFGS only when needed. These are generic smooth optimizers for the residual criterion above. They are not the Bregman proximal-gradient algorithm from the Rényi paper, where the update itself is defined by a Bregman geometry and the gradient has a moment-gap form. A future Bregman solver would be a new solver path rather than a relabeling of BFGS/L-BFGS. 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, quadratic, or Cressie-Read 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 L-BFGS and BFGS alternatives expose \(\|r(\beta)\|^2/2\) and \(J'r\) to Argmin and use More-Thuente search; BFGS initializes the inverse Hessian at identity, while L-BFGS retains a ten-update correction history. Either method 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 L-BFGS. If L-BFGS does not improve the criterion, BFGS is tried from the Gauss-Newton iterate. A fallback replaces Gauss-Newton only when its objective is no larger; 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)\). L-BFGS/BFGS add repeated residual and gradient evaluations. Tight box constraints, a Cressie-Read dual index outside its natural domain, 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.

Code 
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, divergence_power=0.0, dual_ridge=0.0)
Code 
cls = cm.BalancingWeights
display(HTML(html_table(["Public method"], public_methods(cls))))
Public method
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, divergence_power=0.0, dual_ridge=0.0)
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.

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