processing.weighting.data_prep

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

Census geography (census_geo) -- download and cache TIGER PUMA and block shapefiles via pygris, including block-level population inputs.
Geography crosswalk (crosswalk) -- construct a population-weighted PUMA-to-target-zone allocation table and assign/allocate records across project geographies.
PUMS data (pums_data) -- fetch ACS 1-year PUMS microdata from the Census API or load local extracts, with chunking and caching helpers.
Control data (control_data) -- recode PUMS variables into weighting control categories and aggregate them into zone-level control totals.
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.

processing.weighting.data_prep.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.

puma_vintage_for_pums_year

puma_vintage_for_pums_year(pums_year: int) -> int

Return the PUMA geography vintage for a given PUMS year.

ACS PUMS 2022+ uses 2020 PUMAs; 2012-2021 uses 2010 PUMAs.

get_puma_gdf

get_puma_gdf(
    state_fips: str,
    pums_year: int,
    *,
    cache_dir: Path | None = None
) -> gpd.GeoDataFrame

Download PUMA polygons for state_fips via pygris.

Returns a GeoDataFrame with puma_id (str) and geometry.

get_block_gdf

get_block_gdf(
    state_fips: str,
    pums_year: int,
    *,
    cache_dir: Path | None = None
) -> gpd.GeoDataFrame

Download block geometries and join decennial population.

Uses the same vintage as PUMAs so block GEOIDs match the sample plan. Population is fetched from the Census decennial API and joined by GEOID.

The population is used for dasymetric block-level weighting of cross PUMA-to-custom zone allocations.

Returns a GeoDataFrame with block_id (str), block_pop (int), and geometry.

processing.weighting.data_prep.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:

Load PUMS data; join persons to households to carry PUMA geography.
Join crosswalk; multiply PUMS weight by allocation_weight to distribute into custom zones.
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

recode_pums_households

recode_pums_households(
    hh_df: pl.DataFrame,
    person_df: pl.DataFrame,
    targets: list[str] | None = None,
) -> pl.DataFrame

Recode PUMS household records into control categories.

Derives person-level aggregates (workers, children) then loops over household-level controls calling ctrl.from_pums_row.

Parameters:

Name	Type	Description	Default
`hh_df`	`pl.DataFrame`	PUMS household microdata.	required
`person_df`	`pl.DataFrame`	PUMS person microdata (used to derive hh-level aggregates).	required
`targets`	`list[str] \| None`	Registry keys to recode. `None` → all household controls.	`None`

recode_pums_persons

recode_pums_persons(
    person_df: pl.DataFrame, targets: list[str] | None = None
) -> pl.DataFrame

Recode PUMS person records into control categories.

Parameters:

Name	Type	Description	Default
`person_df`	`pl.DataFrame`	PUMS person microdata.	required
`targets`	`list[str] \| None`	Registry keys to recode. `None` → all person controls.	`None`

build_control_totals

build_control_totals(
    hh_df: pl.DataFrame,
    person_df: pl.DataFrame,
    controls: list[ControlSpec],
    geo_col: str = "PUMA",
) -> ControlTotals

Build weighted control totals from recoded PUMS data.

Parameters:

Name	Type	Description	Default
`hh_df`	`pl.DataFrame`	Recoded PUMS household data (from `recode_pums_households`).	required
`person_df`	`pl.DataFrame`	Recoded PUMS person data (from `recode_pums_persons`).	required
`controls`	`list[ControlSpec]`	Control specifications (which variables to include).	required
`geo_col`	`str`	Column name for the geography identifier. Defaults to `"PUMA"`.	`'PUMA'`

Returns:

Type	Description
`ControlTotals`	Tidy totals frame and metadata.

apply_zone_groups

apply_zone_groups(
    control_totals: ControlTotals,
    seed: pl.DataFrame,
    zone_groups: dict[str, list[str]],
    geo_col: str = "ctrl_geoid",
) -> tuple[ControlTotals, pl.DataFrame]

Merge zones for balancing while preserving crosswalk granularity.

Zones listed under a group name are remapped to that group in both the control totals and the seed table. Unmapped zones pass through unchanged.

processing.weighting.data_prep.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.

Load target zone polygons; auto-discover overlapping PUMAs.
Download/cache TIGER PUMA and block shapefiles (via census_geo).
Rasterize block population into a density grid.
Rasterize PUMA IDs into a categorical label grid.
exactextract: compute sum(population) per target zone, grouped by PUMA label.
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)

TargetZoneConfig `pydantic-model`

Target zone polygon source.

Config:

extra: forbid

Fields:

file (str)
id_field (str | None)

model_config `class-attribute` `instance-attribute`

model_config = ConfigDict(extra='forbid')

file `pydantic-field`

file: str

id_field `pydantic-field`

id_field: str | None

Zone-ID column; null for one study-area boundary.

GeographyConfig `pydantic-model`

Geography block of the weighting YAML config.

Attributes:

Name	Type	Description
`target_zones`	`TargetZoneConfig`	Polygon file and optional zone-ID column.
`resolution`	`int`	Raster cell size in metres.
`min_allocation`	`float`	PUMA slivers allocated less than this share are dropped.
`zone_groups`	`dict[str, list[str]] \| None`	Named groups of target zones balanced as one; null for none.

Config:

extra: forbid

Fields:

target_zones (TargetZoneConfig)
resolution (int)
min_allocation (float)
zone_groups (dict[str, list[str]] | None)

Validators:

_positive_resolution → resolution
_valid_min_allocation → min_allocation

model_config `class-attribute` `instance-attribute`

model_config = ConfigDict(extra='forbid')

target_zones `pydantic-field`

target_zones: TargetZoneConfig

resolution `pydantic-field`

resolution: int

min_allocation `pydantic-field`

min_allocation: float

zone_groups `pydantic-field`

zone_groups: dict[str, list[str]] | None

PumaCrosswalk

Population-weighted PUMA-to-target-zone crosswalk.

Loads target zones and Census geographies on construction, rasterizes block population, and cross-tabulates to produce allocation weights. Exposes crosswalk_df, puma_ids, and target_gdf for downstream use. The crosswalk produces both study_geoid (raw target zone from the shapefile) and ctrl_geoid (balancing geography — equal to study_geoid unless zone groups are configured).

zone_groups `instance-attribute`

zone_groups = zone_groups or {}

target_gdf `instance-attribute`

target_gdf = _load_target_zones(
    config.target_zones.file, config.target_zones.id_field
)

crosswalk_df `instance-attribute`

crosswalk_df = pl.read_parquet(cache_dir_xw / 'crosswalk.parquet')

puma_gdf `instance-attribute`

puma_gdf = gpd.read_parquet(cache_dir_xw / 'puma_clipped.parquet')

block_gdf `instance-attribute`

block_gdf = gpd.read_parquet(cache_dir_xw / "block_clipped.parquet")

puma_ids `instance-attribute`

puma_ids = sorted(self.puma_gdf['puma_id'].unique().tolist())

zone_populations `property`

zone_populations: pl.DataFrame

Census block population totals per balancing zone.

Returns a DataFrame with columns [geo_id, target_population] where target_population is the sum of Census block population allocated to each ctrl_geoid.

block_group_populations `property`

block_group_populations: pl.DataFrame

Census block population summed to block-group level.

Block group GEOID = first 12 characters of the 15-character Census block GEOID (FIPS state + county + tract + block group).

Returns a DataFrame with columns [bg_geo_id, bg_population].

assign_block_groups

assign_block_groups(
    households: pl.DataFrame,
    lon_col: str = "home_lon",
    lat_col: str = "home_lat",
) -> pl.DataFrame

Point-in-polygon assignment of households to Census block groups.

Returns households with a bg_geo_id column added (null for points outside all block groups). A block group is the first 12 characters of its blocks' GEOIDs, so each point is matched to the block it lies in rather than to block groups dissolved from blocks, which is far slower and gives the same answer.

A point on the edge between blocks lies within neither, though it does lie within their block group when they share one. Such a point takes the block group of the blocks it touches when they all agree, as it would against dissolved block groups.

assign_households

assign_households(
    households: pl.DataFrame,
    lon_col: str = "home_lon",
    lat_col: str = "home_lat",
) -> pl.DataFrame

Point-in-polygon assignment of households to target zones.

Returns households with study_geoid and ctrl_geoid columns added (null for points outside all zones).

allocate_pums_weights

allocate_pums_weights(
    hh_df: pl.DataFrame,
    person_df: pl.DataFrame,
    geo_col: str = "PUMA",
) -> tuple[pl.DataFrame, pl.DataFrame]

Join crosswalk to PUMS and scale weights by allocation factor.

Produces _xw_WGTP / _xw_PWGTP columns (original weight times allocation_weight) and adds study_geoid and ctrl_geoid to both frames.

allocate_pums_incidence

allocate_pums_incidence(
    pums_incidence: pl.DataFrame,
    *,
    puma_col: str = "PUMA",
    weight_col: str = "WGTP"
) -> pl.DataFrame

Assign target-zone IDs and scale PUMS weights via the crosswalk.

Each PUMS household is duplicated per overlapping target zone, and its weight is scaled by the zone's allocation_weight.

Parameters:

Name	Type	Description	Default
`pums_incidence`	`pl.DataFrame`	Pivoted PUMS incidence table. Must contain puma_col and weight_col.	required
`puma_col`	`str`	Column identifying the PUMA (default `"PUMA"`).	`'PUMA'`
`weight_col`	`str`	Original PUMS weight column to scale (default `"WGTP"`).	`'WGTP'`

Returns:

Type	Description
`pl.DataFrame`	Incidence table with `study_geoid` and `ctrl_geoid` added,
`pl.DataFrame`	weight_col scaled by `allocation_weight`, and rows
`pl.DataFrame`	multiplied per zone overlap.

processing.weighting.data_prep.fractional_impute

Fractional imputation for null-deficient incidence rows.

When a survey record has a null control value (e.g. missing gender), the incidence pivot produces all-zero member columns for that control — the record effectively vanishes from that constraint. This module replaces those zeros with probability distributions predicted by a Random Forest trained on complete PUMS data, implementing fractional imputation (Kim & Fuller, 2004).

The max-entropy balancer accepts non-negative floats in the incidence matrix, so fractional entries are valid without modification.

This is not the survey imputation step (:mod:processing.imputation), and the two should not be read as versions of each other. That step decides a value for a missing field and writes it into canonical data, where every downstream consumer sees it. This one decides nothing: it spreads a record across a control's categories, only inside the seed incidence matrix, only for this balancing run, and the matrix is discarded once weights are fitted. Control totals are never touched -- they come from PUMS and stay fixed. It is a last resort that keeps a record in a constraint it would otherwise drop out of, not a statement about what that record is.

Public API

:func:fill_null_incidence Top-level orchestrator called from the weighting pipeline.

fill_null_incidence

fill_null_incidence(
    survey_bundle: IncidenceBundle,
    pums_bundle: IncidenceBundle,
    targets: list[str],
    *,
    cache_dir: Path | None = None
) -> tuple[pl.DataFrame, list[ImputationSummary]]

Replace null-induced zeros in survey incidence with RF-predicted fractions.

Parameters

survey_bundle : IncidenceBundle Survey incidence bundle (may have null-induced zeros). pums_bundle : IncidenceBundle PUMS incidence bundle (complete, no nulls — used as training data). targets : list[str] Control registry names. cache_dir : Path | None When set, fitted RF models and CV metrics are cached here so that repeated runs with the same PUMS data skip training entirely.

Returns:

tuple[pl.DataFrame, list[ImputationSummary]] Modified survey incidence table with fractional values replacing zeros for null-deficient rows, plus per-control imputation metadata.

processing.weighting.data_prep.incidence

Incidence-table construction and control-total aggregation.

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

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.
Incidence table — :func:build_incidence_table orchestrates the helpers into a single household-level table for the balancer.
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.

pivot_hh_controls

pivot_hh_controls(
    hh_recoded: pl.DataFrame,
    targets: list[str],
    hh_id_col: str = "hh_id",
) -> pl.DataFrame

Pivot recoded household controls into {ctrl}__{member} indicator columns.

Each non-structural household control becomes a set of 0/1 columns (e.g. h_size__size_1, h_income__low). Structural controls (like h_total) are carried through as-is.

Parameters

hh_recoded : pl.DataFrame Recoded households with encoded control columns. targets : list[str] Control registry names. hh_id_col : str Household identifier column.

Returns:

pl.DataFrame [hh_id_col, h_ctrl__member_1, h_ctrl__member_2, ...]

pivot_person_controls

pivot_person_controls(
    per_recoded: pl.DataFrame,
    targets: list[str],
    hh_id_col: str = "hh_id",
) -> tuple[pl.DataFrame, pl.DataFrame]

Pivot recoded person controls into per-person indicators and per-HH counts.

Each non-structural person control becomes a set of 0/1 columns per person (e.g. p_gender__male). Structural controls (p_total) become a literal 1 per person.

Returns both the per-person pivot (useful for fractional imputation) and the per-household aggregation (used in the incidence table).

Parameters

per_recoded : pl.DataFrame Recoded persons with encoded control columns and hh_id_col. targets : list[str] Control registry names. hh_id_col : str Household identifier column.

Returns:

tuple[pl.DataFrame, pl.DataFrame] (person_pivot, person_agg)

*person_pivot* — one row per person with 0/1 indicator columns.
*person_agg* — one row per household with integer count columns
(sum of person_pivot grouped by household).

build_incidence_table

build_incidence_table(
    hh_recoded: pl.DataFrame,
    per_recoded: pl.DataFrame,
    targets: list[str],
    *,
    hh_id_col: str = "hh_id",
    extra_cols: list[str] | None = None
) -> IncidenceBundle

Build household-level incidence table with fully-pivoted control columns.

This is the unified pivoter for both survey and PUMS data. After recoding, both datasets share the same control-column schema (h_size, h_income, p_gender, …). This function pivots those encoded columns into the incidence layout expected by the balancer:

HH controls as 0/1 indicators: h_size__size_1, …
Person controls as per-household counts: p_gender__male, …
Structural controls remain unpivoted: h_total, p_total.

Non-structural controls use the naming {control_name}__{member_name}.

Parameters

hh_recoded : pl.DataFrame Recoded households (survey or PUMS). per_recoded : pl.DataFrame Recoded persons (survey or PUMS). targets : list[str] Control registry names used for recoding. hh_id_col : str Household identifier column ("hh_id" for survey, "SERIALNO" for PUMS). extra_cols : list[str] | None Additional columns from hh_recoded to carry forward (e.g. ["WGTP", "PUMA"] for PUMS).

Returns:

IncidenceBundle Bundle containing the combined incidence table, per-person pivot, and household-only pivot.

aggregate_control_totals

aggregate_control_totals(
    pums_incidence: pl.DataFrame,
    targets: list[str],
    *,
    weight_col: str = "WGTP",
    geo_col: str = "ctrl_geoid"
) -> ControlTotals

Aggregate PUMS incidence into per-zone control totals.

Reads the pivoted {control}__{member} columns and structural columns directly from the incidence table, multiplies each by the household weight, and sums by geography.

Parameters

pums_incidence : pl.DataFrame Zone-allocated PUMS incidence table. Must have geo_col, weight_col, and the {control}__{member} / structural columns. targets : list[str] Control registry names. weight_col : str Weight column (default "WGTP"). geo_col : str Geography column (default "ctrl_geoid").

processing.weighting.data_prep.merges

Category merging for incidence tables and control totals.

Three concerns:

Crosstab merges — :func:apply_crosstab_merges collapses composite {xtab}__{dim0}_{dim1} columns along one or more dimensions, dropping originals. Applied before 1-D merges.
1-D merges — :func:apply_1d_merges handles both global (drop originals) and zone-specific (keep originals) merge specs based on spec.zones.
Control-total merges — :func:merge_control_totals adjusts aggregated target rows for zone-specific merges.

apply_1d_merges

apply_1d_merges(
    df: pl.DataFrame, merges: list[MergeSpec]
) -> pl.DataFrame

Collapse or add incidence columns for 1-D merge specs.

Global (zones=None): sums constituent {ctrl}__{member} columns into {ctrl}__{merged_label} and drops the originals.
Zone-specific (zones set): adds the merged column but keeps the originals — non-merged zones still need them. The per-zone selection happens in control totals, not here.

merge_control_totals

merge_control_totals(
    control_totals: ControlTotals, merges: list[MergeSpec]
) -> ControlTotals

Collapse control-total rows for zone-specific merges.

For each zone merge the constituent category rows are summed into a single merged row and the originals removed — but only for the specified zones. Other zones keep their original rows.

apply_crosstab_merges

apply_crosstab_merges(
    df: pl.DataFrame, merges: list[MergeSpec]
) -> pl.DataFrame

Collapse composite incidence columns along crosstab dimensions.

Each spec's groups is {dim_name: {merged_label: [sources]}}. For each dimension merge, enumerates combinations of the other dimensions' current members, sums constituent columns, and drops originals. Dimensions within one spec are processed sequentially so that a merge in one dimension sees already-merged columns from another.

processing.weighting.data_prep.pums_api

The Census API client for PUMS microdata.

All PUMAs batched in a single API request.
Column chunking when >48 columns (API limit ~50), parallel via ThreadPoolExecutor; the chunks are joined on their record keys rather than stacked.
A response that fails its checks is retried.
JSON → Polars directly (no pandas intermediate).
Streaming download with tqdm progress bars.

CENSUS_BASE `module-attribute`

CENSUS_BASE = 'https://api.census.gov/data'

MAX_COLS_PER_REQUEST `module-attribute`

MAX_COLS_PER_REQUEST = 48

MAX_FETCH_ATTEMPTS `module-attribute`

MAX_FETCH_ATTEMPTS = 4

available_variables

available_variables(base_url: str) -> set[str]

Fetch the set of valid variable names from the Census API.

Hits {base_url}/variables.json (typically <1 MB) and returns the set of variable names. Falls back to an empty set on error (skipping validation rather than blocking the pipeline).

fetch_table

fetch_table(
    base_url: str,
    all_cols: list[str],
    state_fips: str,
    puma_geo: str,
    *,
    label: str = "table",
    keys: tuple[str, ...] = ("SERIALNO",)
) -> pl.DataFrame

Fetch a full PUMS table, chunking columns if needed.

The Census API limits ~50 variables per request. When all_cols exceeds that, the request is split into chunks and fetched in parallel, then joined back together on keys.

The join is what makes this safe, and it has to be a join. The API does not return a stable row order across requests: two 48-column requests for the same 167,075 California households were measured returning the same rows in orders differing at 110,255 positions. Stacking the chunks side by side by position -- which this did, while discarding each chunk's own copy of the key -- therefore attached most values to the wrong record, and left every row carrying exactly one key inherited from the first chunk, so the result looked well formed. It silently scrambled 65% of the replicate weights on every fetch that chunked.

keys must identify a row uniquely and is added to every chunk request: SERIALNO for households, (SERIALNO, SPORDER) for persons, where SERIALNO alone repeats once per household member.

Raises:

Type	Description
`ValueError`	If a chunk does not carry exactly the rows of the first, or if keys does not uniquely identify a row -- either would mean the join silently dropped or multiplied records.

processing.weighting.data_prep.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.

check_pums_integrity

check_pums_integrity(
    hh_df: pl.DataFrame, person_df: pl.DataFrame
) -> None

Verify PUMS records are internally consistent, or raise.

Fetching assembles each table from several API requests, so a record can end up carrying another record's values while still looking well formed -- right keys, plausible numbers, correct row count. These checks look for the disagreements that misattribution produces and that nothing downstream can see:

NP is the household's own count of its members, so it must equal the number of person records sharing that SERIALNO. This is the sharpest check available: the two facts come from different tables, so they only agree when both are attached to the right household. It found 965 disagreements in a corrupted fetch and none in a sound one.
TYPEHUGQ distinguishes housing units from the two kinds of group quarters and admits nothing else. A foreign value there is proof of misattribution -- it is how this class of corruption was first noticed.
Keys identify rows, or the tables cannot be joined at all.

Parameters:

Name	Type	Description	Default
`hh_df`	`pl.DataFrame`	PUMS household records, before group-quarters filtering.	required
`person_df`	`pl.DataFrame`	PUMS person records.	required

Raises:

Type	Description
`ValueError`	If any check fails, naming the count and what it means.

fetch_pums_data

fetch_pums_data(
    source: PUMSSource,
    extra_hh_vars: set[str] | None = None,
    extra_person_vars: set[str] | None = None,
    load_replicate_weights: bool = False,
    cache_dir: Path | None = None,
) -> tuple[pl.DataFrame, pl.DataFrame]

Download PUMS household and person microdata from the Census API.

Parameters:

Name	Type	Description	Default
`source`	`PUMSSource`	State, year, and optional PUMA filter.	required
`extra_hh_vars`	`set[str] \| None`	Additional household PUMS variable names to fetch beyond the defaults.	`None`
`extra_person_vars`	`set[str] \| None`	Additional person PUMS variable names to fetch beyond the defaults.	`None`
`load_replicate_weights`	`bool`	If `True`, also fetch the 80 replicate weight columns per table (`WGTP1` to `WGTP80` and `PWGTP1` to `PWGTP80`). Required for MOE-based importance calculation.	`False`
`cache_dir`	`Path \| None`	If set, raw PUMS data is cached as Parquet files under `cache_dir/pums/`. Subsequent calls with the same state/year load from cache instead of hitting the API.	`None`

Returns:

Type	Description
`pl.DataFrame`	Tuple of `(households, persons)` Polars DataFrames with PUMS
`pl.DataFrame`	data typed to appropriate dtypes.

load_pums_from_files

load_pums_from_files(
    hh_path: str,
    person_path: str,
    state_fips: str | None = None,
    puma_ids: list[str] | None = None,
    load_replicate_weights: bool = False,
) -> tuple[pl.DataFrame, pl.DataFrame]

Load PUMS data from local CSV/Parquet files.

Parameters:

Name	Type	Description	Default
`hh_path`	`str`	Path to household PUMS file (CSV or Parquet).	required
`person_path`	`str`	Path to person PUMS file (CSV or Parquet).	required
`state_fips`	`str \| None`	Optional filter to a specific state FIPS code.	`None`
`puma_ids`	`list[str] \| None`	Optional filter to specific PUMAs.	`None`
`load_replicate_weights`	`bool`	If `True`, retain `WGTP1` to `WGTP80` and `PWGTP1` to `PWGTP80` replicate weight columns.	`False`

Returns:

Type	Description
`tuple[pl.DataFrame, pl.DataFrame]`	Tuple of `(households, persons)` Polars DataFrames.

processing.weighting.data_prep.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.

recode_survey_households

recode_survey_households(
    households: pl.DataFrame,
    persons: pl.DataFrame,
    targets: list[str],
    strict_nulls: bool = False,
) -> pl.DataFrame

Recode canonical survey households into control categories.

Parameters:

Name	Type	Description	Default
`households`	`pl.DataFrame`	Canonical households table (must have `hh_id`).	required
`persons`	`pl.DataFrame`	Canonical persons table (must have `hh_id`).	required
`targets`	`list[str]`	Registry keys to recode. Person-level keys are ignored.	required
`strict_nulls`	`bool`	If `True`, raise on any null recode output.	`False`

Raises:

Type	Description
`ValueError`	Unknown target.
`KeyError`	Missing required column.
`ValueError`	A required column is null on every row.

recode_survey_persons

recode_survey_persons(
    persons: pl.DataFrame,
    targets: list[str],
    strict_nulls: bool = False,
) -> pl.DataFrame

Recode canonical survey persons into control categories.

Parameters:

Name	Type	Description	Default
`persons`	`pl.DataFrame`	Canonical persons table.	required
`targets`	`list[str]`	Registry keys to recode. Household-level keys are ignored.	required
`strict_nulls`	`bool`	If `True`, raise on any null recode output.	`False`

Raises:

Type	Description
`ValueError`	Unknown target.
`KeyError`	Missing required column.
`ValueError`	A required column is null on every row.

processing.weighting.data_prep

Modules

processing.weighting.data_prep.census_geo

puma_vintage_for_pums_year

get_puma_gdf

get_block_gdf

processing.weighting.data_prep.control_data

recode_pums_households

recode_pums_persons

build_control_totals

apply_zone_groups

processing.weighting.data_prep.crosswalk

TargetZoneConfig pydantic-model

model_config class-attribute instance-attribute

file pydantic-field

id_field pydantic-field

GeographyConfig pydantic-model

model_config class-attribute instance-attribute

target_zones pydantic-field

resolution pydantic-field

min_allocation pydantic-field

zone_groups pydantic-field

PumaCrosswalk

zone_groups instance-attribute

target_gdf instance-attribute

crosswalk_df instance-attribute

puma_gdf instance-attribute

block_gdf instance-attribute

puma_ids instance-attribute

zone_populations property

block_group_populations property

assign_block_groups

assign_households

allocate_pums_weights

allocate_pums_incidence

processing.weighting.data_prep.fractional_impute

Public API

fill_null_incidence

Parameters

Returns:

processing.weighting.data_prep.incidence

pivot_hh_controls

Parameters

Returns:

pivot_person_controls

Parameters

Returns:

build_incidence_table

Parameters

Returns:

aggregate_control_totals

Parameters

processing.weighting.data_prep.merges

apply_1d_merges

merge_control_totals

apply_crosstab_merges

processing.weighting.data_prep.pums_api

CENSUS_BASE module-attribute

MAX_COLS_PER_REQUEST module-attribute

MAX_FETCH_ATTEMPTS module-attribute

available_variables

fetch_table

processing.weighting.data_prep.pums_data

check_pums_integrity

fetch_pums_data

load_pums_from_files

processing.weighting.data_prep.seed_data

recode_survey_households

recode_survey_persons

TargetZoneConfig `pydantic-model`

model_config `class-attribute` `instance-attribute`

file `pydantic-field`

id_field `pydantic-field`

GeographyConfig `pydantic-model`

model_config `class-attribute` `instance-attribute`

target_zones `pydantic-field`

resolution `pydantic-field`

min_allocation `pydantic-field`

zone_groups `pydantic-field`

zone_groups `instance-attribute`

target_gdf `instance-attribute`

crosswalk_df `instance-attribute`

puma_gdf `instance-attribute`

block_gdf `instance-attribute`

puma_ids `instance-attribute`

zone_populations `property`

block_group_populations `property`

CENSUS_BASE `module-attribute`

MAX_COLS_PER_REQUEST `module-attribute`

MAX_FETCH_ATTEMPTS `module-attribute`