外观
熵平衡
牛顿法求解矩约束的熵平衡权重。
核心代码
py
def run_ebalance(
df: pd.DataFrame,
covariates: List[str],
*,
mode: str,
treat_var: Optional[str],
target_orders: List[int],
manual_targets: Optional[List[float]],
base_weight_var: Optional[str],
weight_treat: bool,
normconst: float,
maxiter: int,
tolerance: float,
weight_var_name: str,
) -> EbalanceResult:
"""执行与 Stata ebalance 行为保持一致的熵平衡。"""
normalized_mode = str(mode or 'treat').strip().lower()
if normalized_mode not in {'treat', 'manual'}:
raise ValueError(_('熵平衡模式无效'))
covariates = [str(item).strip() for item in covariates if str(item).strip()]
if not covariates:
raise ValueError(_('熵平衡至少需要选择一个配平协变量'))
if normalized_mode == 'manual':
if weight_treat:
raise ValueError(_('单组手动目标模式不支持“处理组权重也使用 base weight”'))
if not manual_targets or len(manual_targets) != len(covariates):
raise ValueError(_('单组手动目标模式下，目标均值个数必须与协变量个数一致'))
target_orders = [1] * len(covariates)
else:
if not treat_var:
raise ValueError(_('请先选择处理组变量'))
if not target_orders or len(target_orders) != len(covariates):
raise ValueError(_('每个协变量都必须指定目标阶数'))
for order in target_orders:
if int(order) not in {1, 2, 3}:
raise ValueError(_('targets 仅支持 1、2、3 阶矩'))
full_work = df.copy()
needed_columns = list(covariates)
if normalized_mode == 'treat':
needed_columns.append(str(treat_var).strip())
if base_weight_var:
needed_columns.append(str(base_weight_var).strip())
work = full_work[needed_columns].copy()
for covariate in covariates:
work[covariate] = _to_numeric_series(work, covariate)
if normalized_mode == 'treat':
treat_name = str(treat_var).strip()
work[treat_name] = _to_numeric_series(work, treat_name)
base_mask = work[covariates + [treat_name]].notna().all(axis=1)
else:
treat_name = '__ebalance_treat__'
base_mask = work[covariates].notna().all(axis=1)
notes: List[str] = []
if base_weight_var:
weight_name = str(base_weight_var).strip()
work[weight_name] = _to_numeric_series(work, weight_name)
if normalized_mode == 'manual':
missing_weight_mask = base_mask & work[weight_name].isna()
if bool(missing_weight_mask.any()):
notes.append(_('已删除 %(count)s 个 base weight 缺失的样本。') % {'count': int(missing_weight_mask.sum())})
touse_mask = base_mask & work[weight_name].notna()
base_weights = work.loc[touse_mask, weight_name].astype(float).to_numpy()
treated_base_weights = np.array([1.0], dtype=float)
else:
treat_full = work[treat_name].copy()
invalid_treat_mask = base_mask & ~treat_full.isin([0, 1])
if bool(invalid_treat_mask.any()):
raise ValueError(_('处理组变量 %(treat_name)s 必须是 0/1 二元变量') % {'treat_name': treat_name})
if weight_treat:
missing_weight_mask = base_mask & work[weight_name].isna()
if bool(missing_weight_mask.any()):
notes.append(_('已删除 %(count)s 个 base weight 缺失的样本。') % {'count': int(missing_weight_mask.sum())})
touse_mask = base_mask & work[weight_name].notna()
base_weights = work.loc[touse_mask, weight_name].astype(float).to_numpy()
treated_base_weights = base_weights[work.loc[touse_mask, treat_name].astype(int).to_numpy() == 1]
else:
missing_control_weight_mask = base_mask & work[weight_name].isna() & (treat_full == 0)
if bool(missing_control_weight_mask.any()):
notes.append(_('已删除 %(count)s 个控制组 base weight 缺失样本。') % {'count': int(missing_control_weight_mask.sum())})
ignored_treated_weight_mask = base_mask & work[weight_name].notna() & (treat_full == 1)
if bool(ignored_treated_weight_mask.any()):
notes.append(_('处理组 base weight 已忽略；若希望处理组也带权，请勾选“处理组也使用 base weight”。'))
touse_mask = base_mask & ~(missing_control_weight_mask)
filtered_treat = work.loc[touse_mask, treat_name].astype(int).to_numpy()
base_weights = np.where(
filtered_treat == 0,
work.loc[touse_mask, weight_name].astype(float).to_numpy(),
1.0,
)
treated_base_weights = np.ones(int((filtered_treat == 1).sum()), dtype=float)
else:
if weight_treat:
raise ValueError(_('勾选“处理组也使用 base weight”时，必须指定 base weight 变量'))
touse_mask = base_mask.copy()
base_weights = np.ones(int(touse_mask.sum()), dtype=float)
treated_base_weights = np.ones(1, dtype=float)
sample_df = full_work.loc[touse_mask, :].copy().reset_index(drop=True)
if sample_df.empty:
raise ValueError(_('熵平衡可用样本为空，请检查变量缺失情况'))
if normalized_mode == 'manual':
sample_df[treat_name] = 0
treat_array = sample_df[treat_name].astype(int).to_numpy()
if len(sample_df) <= 1:
raise ValueError(_('单组手动目标模式下，有效样本量至少需要 2 个'))
else:
sample_df[treat_name] = sample_df[treat_name].astype(int)
treat_array = sample_df[treat_name].to_numpy()
if np.any(~np.isin(treat_array, [0, 1])):
raise ValueError(_('处理组变量 %(treat_name)s 必须是 0/1 二元变量') % {'treat_name': treat_name})
if treat_array.max() == 0 or treat_array.min() == 1:
raise ValueError(_('处理组变量 %(treat_name)s 必须同时包含处理组和控制组样本') % {'treat_name': treat_name})
if normalized_mode == 'treat' and int((treat_array == 0).sum()) <= 1:
raise ValueError(_('控制组有效样本量不足，无法执行熵平衡'))
sample_base_weights = np.asarray(base_weights, dtype=float)
first_order_blocks: List[tuple[str, np.ndarray]] = []
second_order_blocks: List[tuple[str, np.ndarray]] = []
third_order_blocks: List[tuple[str, np.ndarray]] = []
skipped_second_order: List[str] = []
skipped_third_order: List[str] = []
covariate_arrays: Dict[str, np.ndarray] = {
covariate: sample_df[covariate].astype(float).to_numpy()
for covariate in covariates
}
if normalized_mode == 'manual':
control_mask = np.ones(len(sample_df), dtype=bool)
control_count = int(control_mask.sum())
treated_count = 1
for covariate in covariates:
values = covariate_arrays[covariate]
target_value = float(manual_targets[covariates.index(covariate)])
if target_value < np.min(values) or target_value > np.max(values):
raise ValueError(_('手动目标值超出了 %(covariate)s 的样本范围') % {'covariate': covariate})
first_order_blocks.append((covariate, values.copy()))
treated_matrix = np.array([manual_targets], dtype=float)
control_matrix = np.column_stack([values for _, values in first_order_blocks])
treated_constraint_weights = np.array([1.0], dtype=float)
control_base_weights = sample_base_weights.copy()
else:
control_mask = treat_array == 0
treated_mask = treat_array == 1
control_count = int(control_mask.sum())
treated_count = int(treated_mask.sum())
for covariate, target_order in zip(covariates, target_orders):
values = covariate_arrays[covariate]
treated_values = values[treated_mask]
control_values = values[control_mask]
if treated_values.max() < control_values.min() or control_values.max() < treated_values.min():
raise ValueError(_('%(covariate)s 在处理组与控制组之间没有重叠区间') % {'covariate': covariate})
first_order_blocks.append((covariate, values.copy()))
if int(target_order) >= 2:
if np.std(treated_values, ddof=1) == 0 or np.std(control_values, ddof=1) == 0:
skipped_second_order.append(covariate)
if int(target_order) >= 3:
skipped_third_order.append(covariate)
continue
treated_mean = _weighted_mean(treated_values, sample_base_weights[treated_mask])
treated_n = int(treated_mask.sum())
control_n = int(control_mask.sum())
second_order = np.empty_like(values, dtype=float)
second_order[control_mask] = (control_values - treated_mean) ** 2 * control_n / (control_n - 1)
second_order[treated_mask] = (treated_values - treated_mean) ** 2 * treated_n / (treated_n - 1)
second_order_blocks.append((covariate, second_order))
if int(target_order) >= 3:
third_order = np.empty_like(values, dtype=float)
third_order[control_mask] = (control_values - treated_mean) ** 3 * (control_n / (control_n - 1)) ** 1.5
third_order[treated_mask] = (treated_values - treated_mean) ** 3 * (treated_n / (treated_n - 1)) ** 1.5
third_order_blocks.append((covariate, third_order))
if skipped_second_order:
notes.append(_('以下变量因样本内方差不足，自动跳过了高阶矩约束：%(columns)s') % {'columns': ', '.join(skipped_second_order)})
all_constraints = first_order_blocks + second_order_blocks + third_order_blocks
treated_matrix = np.column_stack([values[treated_mask] for _, values in all_constraints])
control_matrix = np.column_stack([values[control_mask] for _, values in all_constraints])
treated_constraint_weights = sample_base_weights[treated_mask].copy()
control_base_weights = sample_base_weights[control_mask].copy()
if normalized_mode == 'manual':
all_constraints = first_order_blocks
else:
all_constraints = first_order_blocks + second_order_blocks + third_order_blocks
constraint_labels = [name for name, _ in all_constraints]
constraint_orders = ([1] * len(first_order_blocks)) + ([2] * len(second_order_blocks)) + ([3] * len(third_order_blocks))
treated_moments = np.sum(treated_constraint_weights[:, None] * treated_matrix, axis=0) / np.sum(treated_constraint_weights)
moments_with_constant = np.concatenate(([1.0], treated_moments))
moment_scale = _safe_power_of_ten(moments_with_constant)
scaled_target_moments = moments_with_constant / moment_scale
scaled_control_matrix = np.column_stack([np.ones(control_matrix.shape[0]), control_matrix]) / moment_scale
coefficients = np.zeros(scaled_control_matrix.shape[1], dtype=float)
raw_control_weights = np.exp(scaled_control_matrix @ coefficients) * control_base_weights
gradient = raw_control_weights @ scaled_control_matrix - scaled_target_moments
maxdiff = float(np.max(np.abs(gradient)))
maxdist = int(np.argmax(np.abs(gradient)))
iteration_logs: List[str] = []
zero_step = False
current_iter = 1
while maxdiff > float(tolerance) and current_iter <= int(maxiter):
raw_control_weights = np.exp(scaled_control_matrix @ coefficients) * control_base_weights
weighted_sums = raw_control_weights @ scaled_control_matrix
gradient = weighted_sums - scaled_target_moments
hessian = scaled_control_matrix.T @ (raw_control_weights[:, None] * scaled_control_matrix)
previous_coefficients = coefficients.copy()
newton_step = _solve_linear_system(hessian, gradient)
coefficients = coefficients - newton_step
loss_new = float(np.max(np.abs((np.exp(scaled_control_matrix @ coefficients) * control_base_weights) @ scaled_control_matrix - scaled_target_moments)))
loss_old = float(np.max(np.abs((np.exp(scaled_control_matrix @ previous_coefficients) * control_base_weights) @ scaled_control_matrix - scaled_target_moments)))
if loss_old <= loss_new:
best_loss = float('inf')
best_coefficients = coefficients.copy()
for divisor in range(1, 6):
candidate = previous_coefficients - (newton_step / divisor)
candidate_loss = float(np.max(np.abs((np.exp(scaled_control_matrix @ candidate) * control_base_weights) @ scaled_control_matrix - scaled_target_moments)))
if candidate_loss < best_loss:
best_loss = candidate_loss
best_coefficients = candidate
coefficients = best_coefficients
zero_step = False
maxdiff = float(np.max(np.abs(gradient)) / scaled_target_moments[0])
maxdist = int(np.argmax(np.abs(gradient)))
iteration_logs.append(f'Iteration {current_iter}: Max Difference = {maxdiff}')
if zero_step:
break
current_iter += 1
converged = bool(maxdiff <= float(tolerance))
if not converged:
notes.append(_format_constraint_note(
maxdist=maxdist,
xnum1=len(first_order_blocks),
xnum2=len(second_order_blocks),
constrt1=[name for name, _ in first_order_blocks],
constrt2=[name for name, _ in second_order_blocks],
))
# 与 Stata 保持一致：输出权重使用“最后一次迭代更新前”的控制组权重。
final_control_weights = raw_control_weights.copy()
if normalized_mode == 'manual':
treated_weight_total = float(len(sample_df))
control_weight_total = float(np.sum(final_control_weights))
normalized_control_weights = final_control_weights * float(normconst) * len(sample_df) / control_weight_total
final_weights_array = normalized_control_weights
treated_total_weight = float(np.sum(normalized_control_weights))
control_total_weight = treated_total_weight
else:
treated_weights_array = sample_base_weights[treated_mask].copy() if weight_treat else np.ones(treated_count, dtype=float)
treated_weight_total = float(np.sum(treated_weights_array))
control_weight_total = float(np.sum(final_control_weights))
normalized_control_weights = final_control_weights * float(normconst) * treated_weight_total / control_weight_total
final_weights_array = np.empty(len(sample_df), dtype=float)
final_weights_array[treated_mask] = treated_weights_array
final_weights_array[control_mask] = normalized_control_weights
treated_total_weight = float(np.sum(treated_weights_array))
control_total_weight = float(np.sum(normalized_control_weights))
final_weights = pd.Series(final_weights_array, name=weight_var_name)
if normalized_mode == 'manual':
pre_rows = []
post_rows = []
for covariate in covariates:
values = sample_df[covariate].astype(float).to_numpy()
pre_stats = _summary_triplet(values, sample_base_weights)
post_stats = _summary_triplet(values, final_weights.to_numpy())
pre_rows.append({
'Xname': covariate,
'mean_Pre': pre_stats['mean'],
'var_Pre': pre_stats['variance'],
'skew_Pre': pre_stats['skewness'],
})
post_rows.append({
'Xname': covariate,
'mean_Post': post_stats['mean'],
'var_Post': post_stats['variance'],
'skew_Post': post_stats['skewness'],
})
pre_balance = pd.DataFrame(pre_rows)
post_balance = pd.DataFrame(post_rows)
balance_table = pre_balance.merge(post_balance, on='Xname', how='left')
weight_summary = pd.DataFrame([
{'项目': _('调整样本量'), '取值': int(len(sample_df))},
{'项目': _('权重总和'), '取值': float(control_total_weight)},
{'项目': _('是否收敛'), '取值': _('是') if converged else _('否')},
{'项目': _('最大偏差'), '取值': float(maxdiff)},
])
else:
pre_rows = []
post_rows = []
for covariate in covariates:
treated_values = sample_df.loc[treat_array == 1, covariate].astype(float).to_numpy()
control_values = sample_df.loc[treat_array == 0, covariate].astype(float).to_numpy()
treated_stats = _summary_triplet(treated_values, sample_base_weights[treat_array == 1] if weight_treat else np.ones(treated_count, dtype=float))
control_pre_stats = _summary_triplet(control_values, sample_base_weights[treat_array == 0])
control_post_stats = _summary_triplet(control_values, final_weights.loc[treat_array == 0].to_numpy())
pre_rows.append({
'Xname': covariate,
'mean_Tr': treated_stats['mean'],
'var_Tr': treated_stats['variance'],
'skew_Tr': treated_stats['skewness'],
'mean_Co_Pre': control_pre_stats['mean'],
'var_Co_Pre': control_pre_stats['variance'],
'skew_Co_Pre': control_pre_stats['skewness'],
})
post_rows.append({
'Xname': covariate,
'mean_Tr': treated_stats['mean'],
'var_Tr': treated_stats['variance'],
'skew_Tr': treated_stats['skewness'],
'mean_Co_Post': control_post_stats['mean'],
'var_Co_Post': control_post_stats['variance'],
'skew_Co_Post': control_post_stats['skewness'],
})
pre_balance = pd.DataFrame(pre_rows)
post_balance = pd.DataFrame(post_rows)
balance_table = pre_balance.merge(post_balance[['Xname', 'mean_Co_Post', 'var_Co_Post', 'skew_Co_Post']], on='Xname', how='left')
balance_table['sdiff_Pre'] = (balance_table['mean_Tr'] - balance_table['mean_Co_Pre']) / np.sqrt(balance_table['var_Tr'])
balance_table['sdiff_Post'] = (balance_table['mean_Tr'] - balance_table['mean_Co_Post']) / np.sqrt(balance_table['var_Tr'])
weight_summary = pd.DataFrame([
{'项目': _('处理组样本量'), '取值': int(treated_count)},
{'项目': _('处理组权重总和'), '取值': float(treated_total_weight)},
{'项目': _('控制组样本量'), '取值': int(control_count)},
{'项目': _('控制组权重总和'), '取值': float(control_total_weight)},
{'项目': _('是否收敛'), '取值': _('是') if converged else _('否')},
{'项目': _('最大偏差'), '取值': float(maxdiff)},
])
lambdas_df = pd.DataFrame({
'Xname': constraint_labels,
'order': constraint_orders,
'lambda': coefficients[1:],
})
moments_df = pd.DataFrame({
'Xname': constraint_labels,
'order': constraint_orders,
'mconstrnt': moments_with_constant[1:],
})
return EbalanceResult(
sample_df=sample_df,
final_weights=final_weights,
weight_var_name=weight_var_name,
mode=normalized_mode,
converged=converged,
maxdiff=float(maxdiff),
maxdist=maxdist,
iteration_logs=iteration_logs,
notes=notes,
covariates=covariates,
constraint_labels=constraint_labels,
constraint_orders=constraint_orders,
lambdas=lambdas_df,
moments=moments_df,
pre_balance=pre_balance,
post_balance=post_balance,
balance_table=balance_table,
weight_summary=weight_summary,
treated_count=treated_count,
control_count=control_count,
treated_total_weight=float(treated_total_weight),
control_total_weight=float(control_total_weight),
)