Skip to content

Data Preparation

Data-preparation toolbox for the weighting workflow.

This sub-package does not orchestrate the weighting pipeline itself. Instead, it provides the reusable building blocks that the weighting pipeline uses to prepare geography, PUMS inputs, control totals, and survey seed data.

Modules

  1. Census geography (census_geo) -- download and cache TIGER PUMA and block shapefiles via pygris, including block-level population inputs.
  2. Geography crosswalk (crosswalk) -- construct a population-weighted PUMA-to-target-zone allocation table and assign/allocate records across project geographies.
  3. PUMS data (pums_data) -- fetch ACS 1-year PUMS microdata from the Census API or load local extracts, with chunking and caching helpers.
  4. Control data (control_data) -- recode PUMS variables into weighting control categories and aggregate them into zone-level control totals.
  5. Seed data (seed_data) -- recode canonical survey variables into the same control categories used by the PUMS-based controls.

Together these modules provide the shared data-preparation utilities used by WeightingPipeline before balancing begins.

crosswalk

PUMA-specific crosswalk wrapper.

PumaCrosswalk fetches Census PUMA and block geographies and delegates the heavy lifting to build_crosswalk.

See docs/pipeline_steps/weighting/crosswalk.md for the detailed crosswalk explanation, diagrams, and worked example.

  1. Load target zone polygons; auto-discover overlapping PUMAs.
  2. Download/cache TIGER PUMA and block shapefiles (via census_geo).
  3. Rasterize block population into a density grid.
  4. Rasterize PUMA IDs into a categorical label grid.
  5. exactextract: compute sum(population) per target zone, grouped by PUMA label.
  6. Normalise: allocation_weight = pop(puma, target) / pop(puma) per PUMA.

Resolution only affects within-block population distribution granularity — boundary accuracy is exact at any resolution due to exactextract's analytical sub-cell coverage.

Outputs:

  • crosswalk_df: puma_id, study_geoid, population, allocation_weight.
  • puma_ids: list of PUMAs overlapping the study area.
Example 
xw = PumaCrosswalk(geo_cfg, state_fips="06", pums_year=2023)
households = xw.assign_households(households)
pums_hh, pums_per = xw.allocate_pums_weights(pums_hh, pums_per)

census_geo

Census TIGER geography loading.

PUMA and block geometries share the same decennial vintage (2010 for pre-2022 PUMS, 2020 for 2022+) so that block GEOIDs match the sample plan. Block population is fetched from the Census decennial API (PL 94-171 for 2020, SF1 for 2010) and joined by GEOID — the TIGER shapefiles themselves do not always carry population columns.

Processed GeoDataFrames are cached as GeoParquet under <pipeline-cache-dir>/weighting/census_geo/ for fast repeat loads.

pums_data

PUMS microdata I/O.

Downloads ACS PUMS 1-year microdata from the Census Bureau API (through pums_api) or loads it from local CSV / Parquet files, checks it, and casts the API's string responses to proper numeric dtypes. Downloads are cached as Parquet at <cache_dir>/pums/{state}_{year}_{hh|person}.parquet.

Transformation (recoding, aggregation) lives in control_data.

control_data

PUMS control-data transformation.

Recodes raw PUMS variables into the shared control categories defined in controls.py, then aggregates weighted totals by geography (PUMA by default).

All mapping logic lives in controls.py. This module only orchestrates recoding (via ctrl.pums_expr()) and aggregation.

Approach:

  1. Load PUMS data; join persons to households to carry PUMA geography.
  2. Join crosswalk; multiply PUMS weight by allocation_weight to distribute into custom zones.
  3. For each control: apply filter, recode variable into bins/groups, aggregate weighted sum by (ctrl_geoid, category).

Control Variable YAML Configuration Example::

controls:
  # Simple marginal — household size
  - name: h_size
    table: households
    variable: NP
    bins:
      "1":  [1, 1]
      "2":  [2, 2]
      "3":  [3, 3]
      "4+": [4, 99]

  # Grouped marginal — commute mode
  - name: commute_mode
    table: persons
    variable: JWTRNS
    groups:
      drove_alone: [1]
      carpool:     [2, 3]
      transit:     [4, 5, 6, 7, 8, 9]
      other:       [10, 11, 12]
    filter: "ESR in [1,2,4,5]"   # employed persons only

seed_data

Survey seed-data preparation for weighting.

Uses ControlTarget.survey_expr() to recode canonical survey data into control-category ints via native Polars expressions (vectorised, no map_elements). Every control is handled uniformly.

Survey field mapping is hardcoded in each ControlTarget subclass — the survey_fields class attribute declares which canonical survey columns the control reads, and survey_expr() returns the Polars expression that maps those values to control-category ints. For example:

  • HHIncomeControl.survey_fields = ("income_bin",) — reads the already-binned income_bin column via identity_expr.
  • GenderControl.survey_fields = ("gender",) — maps Gender enum values to GenderCategory ints via replace_strict.
  • HHSizeControl.survey_fields = ("_n_persons",) — clips the person-count aggregate column to [1, 10].

No YAML config is involved; the mapping lives entirely in the control class definitions in processing.weighting.controls.

incidence

Incidence-table construction and control-total aggregation.

Unified operations that apply identically to both survey and PUMS datasets after recoding:

  1. Pivot helpers — :func:pivot_hh_controls and :func:pivot_person_controls convert encoded control columns into the {ctrl}__{member} incidence layout at HH and person levels.
  2. Incidence table — :func:build_incidence_table orchestrates the helpers into a single household-level table for the balancer.
  3. Aggregate — :func:aggregate_control_totals sums the pivoted PUMS incidence into per-zone target totals.

Category merging lives in :mod:processing.weighting.data_prep.merges.

Code: processing.weighting.data_prep