processing.weighting.balancing

Balancing sub-package (balancer, base weights, importance).

Orchestrates the core balancing loop:

Base weights (base_weights) -- compute initial expansion factors per zone: target_hh_pop / n_responses.
Importance (importance) -- derive MOE-based per-control importance from PUMS replicate weights, with explicit YAML overrides.
Balancer (balancer) -- maximum-entropy list balancing via np_balancer_numba, vendored from PopulationSim. Runs independently per geography zone.

Household weights leave here for processing.weighting.core.propagation, which carries them down the canonical hierarchy.

processing.weighting.balancing.balancer

Maximum-entropy list balancer.

Thin Polars→numpy bridge around np_balancer_numba, vendored from PopulationSim. Runs independently per geography zone.

Algorithm

Find weight vector w closest to seed weights w₀ (KL-divergence) subject to marginal constraints:

$$ \min \sum_i w_i \ln(w_i / w_{0i}) \quad \text{s.t.} \quad Aw = t,\; w_i \ge 0 $$

where A is the incidence matrix and t is the target totals vector.

Implementation

Calls np_balancer_numba from _np_balancer, vendored from PopulationSim v0.10.0 -- a pure @njit function taking numpy arrays, so PopulationSim is not a dependency. Zones are independent and run in parallel via ThreadPoolExecutor.

Configuration (YAML)

The balancing: block of compute_weights (see BalancingConfig):

balancing:
  max_expansion_factor: 5.0     # upper bound = initial_weight x factor
  min_expansion_factor: 0.01    # lower bound = initial_weight x factor
  min_weight: 1                 # absolute floor
  # max_weight: null            # absolute ceiling (optional)
  max_iterations: 10000
  n_workers: 8                  # zones balanced in parallel
  default_importance: 100
  moe_based_importance: true

balance_weights

balance_weights(
    seed: pl.DataFrame,
    control_totals: ControlTotals,
    targets: list[str],
    balancing: BalancingConfig,
    importance: ImportanceConfig,
    *,
    verbose: bool = True
) -> tuple[pl.DataFrame, list[ZoneStatus]]

Balance household weights to match control totals per zone.

Parameters:

Name	Type	Description	Default
`seed`	`pl.DataFrame`	Incidence table with `hh_id`, `ctrl_geoid`, `base_weight`, and pivoted control columns (`{ctrl}__{member}` or structural). All merges (global and zone-specific) must already be applied.	required
`control_totals`	`ControlTotals`	Per-zone targets (with merges already applied).	required
`targets`	`list[str]`	Control registry names.	required
`balancing`	`BalancingConfig`	Solver bounds, iteration limits, and parallelism.	required
`importance`	`ImportanceConfig`	Per-control importance weights.	required
`verbose`	`bool`	Log per-zone convergence (default `True`).	`True`

Returns:

Type	Description
`pl.DataFrame`	Tuple of `(weights, statuses)` where weights is a DataFrame
`list[ZoneStatus]`	with columns `hh_id`, `hh_weight`, `geo_id`, and
`tuple[pl.DataFrame, list[ZoneStatus]]`	statuses is one `ZoneStatus`
`tuple[pl.DataFrame, list[ZoneStatus]]`	entry per zone with convergence info.

grid_search_expansion_factor

grid_search_expansion_factor(
    seed: pl.DataFrame,
    control_totals: ControlTotals,
    targets: list[str],
    ef_grid: list[float],
    selected_ef: float,
    balancing: BalancingConfig,
    importance: ImportanceConfig,
) -> list[GridPoint]

Re-run the balancer across a grid of max_expansion_factor values.

Returns one GridPoint per EF value with aggregate fit and weight-quality metrics. The selected_ef is automatically included in the grid if not already present.

This is a diagnostics-only function — it does not modify the production weights.

Grid points are independent, so with balancing.n_workers > 1 they run in parallel, each balancing its zones in turn. That keeps the thread count at n_workers while using it fully: a run has few zones, and zone-level parallelism alone left most workers idle.

processing.weighting.balancing.base_weights

Initial (base) expansion weights for the balancer.

Before balancing, each survey household needs a starting weight that reflects its basic expansion factor — the number of real-world households it represents. Without meaningful initial weights, the Newton-Raphson solver starts from 1.0 and must bridge the gap to expansion factors of hundreds or thousands, quickly slamming into the expansion-factor constraints.

This module provides two paths:

Response inversion (default) — target_hh_pop / n_responses per zone. Works whenever PUMS-derived control totals are available (always, in our pipeline).
Sample plan — a SamplePlan object mapping Census block groups to sampling strata (segments). Block-group populations are sourced from the crosswalk (Census block population summed to BG level). Each household is assigned a block group (via spatial join), mapped to a segment via the plan, and receives a segment-level initial weight: segment_pop / segment_responses.

The public entry point is compute_base_weights, which adds a base_weight column to the seed table.

load_sample_plan

load_sample_plan(path: str | Path) -> SamplePlan

Read a sample-plan CSV into a SamplePlan.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Path to a CSV file. Must contain at minimum `bg_geo_id` and `sample_segment`. Block-group population totals are sourced from the crosswalk, not from the CSV.	required

Returns:

Type	Description
`SamplePlan`	`SamplePlan`.

Raises:

Type	Description
`FileNotFoundError`	If path does not exist.
`ValueError`	If required columns are missing (raised by `SamplePlan`).

compute_base_weights

compute_base_weights(
    seed: pl.DataFrame,
    control_totals: ControlTotals,
    targets: list[str],
    geo_col: str = "ctrl_geoid",
    *,
    sample_plan: str | Path | SamplePlan | None = None,
    bg_populations: pl.DataFrame | None = None
) -> pl.DataFrame

Add base_weight column to the seed table.

Parameters:

Name	Type	Description	Default
`seed`	`pl.DataFrame`	Household seed table from `build_incidence_table`. Must contain `hh_id` and geo_col. When using a sample plan, must also contain `bg_geo_id` (assigned via `PumaCrosswalk.assign_block_groups`).	required
`control_totals`	`ControlTotals`	PUMS-derived targets from `build_control_totals`.	required
`targets`	`list[str]`	Control registry names (used to identify the master HH control).	required
`geo_col`	`str`	Geography column on seed.	`'ctrl_geoid'`
`sample_plan`	`str \| Path \| SamplePlan \| None`	Optional sample plan. Accepts a file path (loaded via `load_sample_plan`) or an already-loaded `SamplePlan`. When `None`, default response inversion is used.	`None`
`bg_populations`	`pl.DataFrame \| None`	Census block-group population totals with columns `[bg_geo_id, bg_population]`. Required when sample_plan is provided; sourced from `PumaCrosswalk.block_group_populations`.	`None`

Returns:

Type	Description
`pl.DataFrame`	seed with an additional `base_weight` column (Float64).

Raises:

Type	Description
`ValueError`	If no household-level control is found in targets, or if a zone has zero survey responses.

processing.weighting.balancing.importance

MOE-based importance weight calculation.

Uses PUMS successive-difference replicate weights (WGTP1-80 / PWGTP1-80) to estimate the Margin of Error (MOE) for each weighted control total, then converts to importance weights inversely proportional to the coefficient of variation (CV).

Controls with higher sampling uncertainty (larger CV) receive lower importance so the balancer doesn't chase noisy targets. Controls absent from the returned dict (e.g. structural totals whose MOE is meaningless) fall back to the balancer's default importance.

DEFAULT_IMPORTANCE `module-attribute`

DEFAULT_IMPORTANCE = 100.0

compute_moe_importance

compute_moe_importance(
    hh_df: pl.DataFrame,
    person_df: pl.DataFrame,
    target_names: list[str],
    *,
    geo_col: str = "ctrl_geoid"
) -> dict[str, float]

Compute per-control importance weights from PUMS replicate-weight MOE.

Parameters:

Name	Type	Description	Default
`hh_df`	`pl.DataFrame`	Crosswalk-allocated PUMS households with `_xw_WGTP` and `_xw_WGTP1` … `_xw_WGTP80` columns (plus recoded controls).	required
`person_df`	`pl.DataFrame`	Crosswalk-allocated PUMS persons with `_xw_PWGTP` and `_xw_PWGTP1` … `_xw_PWGTP80` columns (plus recoded controls).	required
`target_names`	`list[str]`	Control registry names to compute importance for.	required
`geo_col`	`str`	Geography column (default `"ctrl_geoid"`).	`'ctrl_geoid'`

Returns:

Type	Description
`dict[str, float]`	`{control_name: importance}` for controls where MOE could be
`dict[str, float]`	computed. Controls with no PUMS records (data sparsity) are
`dict[str, float]`	omitted — the balancer will apply its default importance.

Raises:

Type	Description
`ValueError`	If a target name is not in the control registry, or if required replicate weight / control columns are missing from the data.

compute_control_moe

compute_control_moe(
    hh_df: pl.DataFrame,
    person_df: pl.DataFrame,
    target_names: list[str],
    *,
    geo_col: str = "ctrl_geoid"
) -> pl.DataFrame

Compute per-cell MOE from PUMS replicate weights for all controls.

Returns a DataFrame with columns [geo_id, control_name, category, target_total, se, moe_pct]. Controls with no PUMS records are silently skipped.

processing.weighting.balancing.merges

Category merge operations for balancing controls.

Handles collapsing control categories to reduce matrix size and simplify sparse cross-tabs. Supports both 1D merges (simple category lists) and N-D merges (cartesian products for cross-tabulated controls).

apply_category_merges

apply_category_merges(
    incidence: np.ndarray,
    targets: np.ndarray,
    row_labels: list[tuple[str, str]],
    master_idx: int,
    merges: list[MergeSpec],
    importance: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, int, np.ndarray]

Collapse incidence rows + target entries for merged categories.

For each MergeSpec, every group maps several base member names to a single merged row (1D merge) or maps a cartesian product of dimension members to a single merged row (N-D merge for cross-tabs).

The incidence rows are summed element-wise and the target values are summed. Original rows are replaced by the single merged row.

Parameters

incidence : np.ndarray Shape (n_controls, n_households). targets : np.ndarray One entry per row. row_labels : list[tuple[str, str]] (control_name, member_name) per row. Modified in place to reflect the merged labels. master_idx : int Row index of the total-HH master control. merges : list[MergeSpec] Merge specifications to apply. importance : np.ndarray Importance weights per row. Merged rows keep the first row's value.

Returns:

incidence, targets, master_idx, importance Reduced arrays and (possibly shifted) master index.

expand_crosstab_dimensions

expand_crosstab_dimensions(
    control_name: str,
    dimensions: dict[str, list[str]],
    row_labels: list[tuple[str, str]],
) -> list[str]

Expand N-D merge dimensions to composite enum member names.

Looks up the cross-tab control to determine the correct dimension order, then builds composite member names matching the enum's naming convention.

Parameters

control_name : str Name of the cross-tab control (for dimension order lookup). dimensions : dict[str, list[str]] Dimension control name -> list of member names to include. E.g. {"h_income": ["income_under25", "income_25to50"], "h_size": ["size_1", "size_2"]} row_labels : list[tuple[str, str]] Current row labels (control_name, member_name) - used to extract dimension order from actual cross-tab structure.

Returns:

list[str] Composite member names from cartesian product in correct dimension order. E.g. for h_size_by_income (dimensions: h_size, h_income): ["size_1_income_under25", "size_1_income_25to50", "size_2_income_under25", "size_2_income_25to50"]

Examples:

expand_crosstab_dimensions( ... "h_size_by_income", ... {"h_size": ["size_1", "size_2"], "h_income": ["income_under25"]}, ... [("h_size_by_income", "size_1_income_under25"), ...] ... ) ['size_1_income_under25', 'size_2_income_under25']

processing.weighting.balancing

processing.weighting.balancing.balancer

Algorithm

Implementation

Configuration (YAML)

balance_weights

grid_search_expansion_factor

processing.weighting.balancing.base_weights

load_sample_plan

compute_base_weights

processing.weighting.balancing.importance

DEFAULT_IMPORTANCE module-attribute

compute_moe_importance

compute_control_moe

processing.weighting.balancing.merges

apply_category_merges

Parameters

Returns:

expand_crosstab_dimensions

Parameters

Returns:

Examples:

DEFAULT_IMPORTANCE `module-attribute`