utils

Utility functions for pipeline operations.

utils.create_ids

Helper functions for creating globally unique IDs in the pipeline.

This module provides utilities for creating composite IDs from hierarchical identifier columns. IDs follow a hierarchical concatenation structure: hh_id (8 digits) -> person_id (+ 2 digits) -> day_id (+ 2 digits) -> linked_trip_id (+ 2 digits) -> etc.

Each level is created by appending zero-padded digits to the parent level, forming a character string that uniquely identifies entities at each level of the hierarchy.

ID vs Num Columns

ID columns (e.g., day_id, person_id): Pre-concatenated hierarchical identifiers that contain the full path from root (hh_id) to current level. Multiple rows may share the same parent ID (e.g., multiple trips per day).
Num columns (e.g., day_num, person_num): Sequence numbers within a parent scope. These reset for each new parent (e.g., each person has day_num 1, 2, 3).

create_concatenated_id

create_concatenated_id(
    df: pl.DataFrame,
    output_col: str,
    parent_id_col: str,
    sequence_col: str,
    sequence_padding: int = 2,
) -> pl.DataFrame

Create hierarchical ID by concatenating parent ID with padded sequence.

This function creates IDs following a hierarchical structure where each level is formed by appending zero-padded sequence numbers to the parent: - hh_id: "23000075" [Can variable length, typically 8 digits] - person_id: "2300007501" [10-digits] (hh_id + 2-digit person_num) - day_id: "230000750101" [12-digits] (person_id + 2-digit day_num) - unlinked_trip_id: "23000075010101" [14-digits] (day_id + 2-digit trip_num) - linked_trip_id: "23000075010102" [14-digits] (day_id + 2-digit trip_num) - tour_id : "23000075010101" [16-digits] (day_id + 2-digit parent tour_num + 2-digit sub-tour_num)

The parent_id_col may already contain the full hierarchical path (pre-concatenated), making this a simple append operation.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	Input DataFrame	required
`output_col`	`str`	Name for the new ID column to create	required
`parent_id_col`	`str`	Column containing the parent ID (may be pre-concatenated)	required
`sequence_col`	`str`	Column containing sequence number at this level	required
`sequence_padding`	`int`	Number of digits for zero-padding sequence (default: 2)	`2`

Returns:

Type	Description
`pl.DataFrame`	DataFrame with new ID column added

Examples:

Create person_id from hh_id + person_num

>>> df = create_concatenated_id(
...     df,
...     output_col="person_id",
...     parent_id_col="hh_id",
...     sequence_col="person_num",
...     sequence_padding=2
... )

Create day_id from person_id + day_num

>>> df = create_concatenated_id(
...     df,
...     output_col="day_id",
...     parent_id_col="person_id",
...     sequence_col="day_num",
...     sequence_padding=2
... )

Create linked_trip_id from day_id + linked_trip_num

>>> df = create_concatenated_id(
...     df,
...     output_col="linked_trip_id",
...     parent_id_col="day_id",
...     sequence_col="linked_trip_num",
...     sequence_padding=2
... )

create_linked_trip_id

create_linked_trip_id(
    linked_trips: pl.DataFrame,
    day_id_col: str = "day_id",
    sequence_col: str = "linked_trip_num",
    sequence_padding: int = 2,
) -> pl.DataFrame

Create linked_trip_id by appending sequence to day_id.

Convenience wrapper for create_concatenated_id() specifically for linked trip IDs. Assumes day_id is already pre-concatenated with the full hierarchy (hh_id + person_num + day_num).

Parameters:

Name	Type	Description	Default
`linked_trips`	`pl.DataFrame`	Input DataFrame containing day_id and sequence columns	required
`day_id_col`	`str`	Column name for day ID (default: "day_id")	`'day_id'`
`sequence_col`	`str`	Column name for sequence number (default: "linked_trip_num")	`'linked_trip_num'`
`sequence_padding`	`int`	Number of digits for padding (default: 2)	`2`

Returns:

Type	Description
`pl.DataFrame`	DataFrame with "linked_trip_id" column added

Example

trips = trips.with_columns( ... pl.col("new_trip_flag") ... .cum_sum() ... .over("person_id") ... .alias("linked_trip_num") ... ) trips = create_linked_trip_id(trips)

create_tour_ids

create_tour_ids(
    linked_trips: pl.DataFrame,
    day_id_col: str = "day_id",
    tour_num_col: str = "tour_num",
    subtour_num_col: str = "subtour_num",
) -> pl.DataFrame

Create tour_id and parent_tour_id from day_id and tour numbers.

Creates hierarchical tour identifiers by combining day_id with tour_num and subtour_num. The tour_id includes both tour and subtour information, while parent_tour_id only includes the tour number (for linking subtours to their parent tour).

Parameters:

Name	Type	Description	Default
`linked_trips`	`pl.DataFrame`	Input DataFrame containing day_id and tour number columns	required
`day_id_col`	`str`	Column name for day ID (default: "day_id")	`'day_id'`
`tour_num_col`	`str`	Column name for tour number (default: "tour_num")	`'tour_num'`
`subtour_num_col`	`str`	Column name for subtour number (default: "subtour_num")	`'subtour_num'`

Returns:

Type	Description
`pl.DataFrame`	DataFrame with "tour_id" and "parent_tour_id" columns added

Example

linked_trips = create_tour_ids(linked_trips)

utils.crosswalk

Generic population-weighted geography crosswalk utilities.

Provides build_crosswalk, a generic function that maps source polygons to target polygons via a weight polygon layer (e.g. Census blocks with population counts). The three layers are rasterized onto a common grid and exactextract performs the zonal cross-tabulation with sub-pixel coverage fractions.

See src/utils/CROSSWALK.md for the mathematical formulation.

Example

from utils.crosswalk import build_crosswalk

xw = build_crosswalk(
    source_gdf=pumas, target_gdf=counties, weight_gdf=blocks,
    source_id_col="PUMACE20", target_id_col="COUNTYFP",
    weight_col="block_pop", resolution=250,
)

build_crosswalk

build_crosswalk(
    source_gdf: gpd.GeoDataFrame,
    target_gdf: gpd.GeoDataFrame,
    weight_gdf: gpd.GeoDataFrame,
    *,
    source_id_col: str = "source_id",
    target_id_col: str = "target_id",
    weight_col: str = "block_pop",
    resolution: int = 100,
    min_allocation: float = 0.0
) -> pl.DataFrame

Build a population-weighted crosswalk between two polygon layers.

Parameters:

Name	Type	Description	Default
`source_gdf`	`gpd.GeoDataFrame`	Source zone polygons (e.g. PUMAs). Must contain source_id_col.	required
`target_gdf`	`gpd.GeoDataFrame`	Target zone polygons (e.g. counties, TAZs). Must contain target_id_col.	required
`weight_gdf`	`gpd.GeoDataFrame`	Weight polygons carrying a numeric attribute (e.g. Census blocks with `block_pop`). Must contain weight_col and geometry.	required
`source_id_col`	`str`	Column in source_gdf identifying source zones.	`'source_id'`
`target_id_col`	`str`	Column in target_gdf identifying target zones.	`'target_id'`
`weight_col`	`str`	Numeric column in weight_gdf used as the allocation weight (typically population).	`'block_pop'`
`resolution`	`int`	Raster cell size in metres (EPSG:5070).	`100`
`min_allocation`	`float`	Drop source-target pairs whose allocation weight is below this fraction (e.g. 0.02 = 2%). Remaining weights are not re-normalised; the dropped slivers simply become part of the out-of-region remainder. Default 0.0 (keep all).	`0.0`

Returns:

Type	Description
`pl.DataFrame`	DataFrame with columns
`pl.DataFrame`	`[source_id, target_id, population, allocation_weight]`.
`pl.DataFrame`	`allocation_weight` is the fraction of each source zone's
`pl.DataFrame`	total population that falls in the target zone. Weights sum
`pl.DataFrame`	to <= 1.0 per `source_id`; the remainder is population outside
`pl.DataFrame`	all target zones.

utils.enum_helpers

Helper functions for resolving enum labels to values.

This module provides utilities for working with codebook enumerations, particularly for resolving enum member names to their numeric values and preparing data for imputation by replacing sentinel/missing values.

MODEL_MAPPING `module-attribute`

MODEL_MAPPING = {
    "households": HouseholdModel,
    "persons": PersonModel,
    "unlinked_trips": UnlinkedTripModel,
    "linked_trips": LinkedTripModel,
    "tours": TourModel,
    "joint_trips": JointTripModel,
}

get_enum_class_for_field

get_enum_class_for_field(
    table_name: str, field_name: str
) -> type[LabeledEnum] | None

Find the enum class for a given field in a table.

Uses the Pydantic model's type annotations to determine the enum class for a field. For example, if HouseholdModel has income_bin: IncomeBroad, this function will return the IncomeBroad enum class.

Parameters:

Name	Type	Description	Default
`table_name`	`str`	Name of the table (e.g., 'households', 'persons', 'unlinked_trips')	required
`field_name`	`str`	Name of the field/column to find enum for	required

Returns:

Type	Description
`type[LabeledEnum] \| None`	The enum class if found, None otherwise

Example

enum_class = get_enum_class_for_field('households', 'income_bin') enum_class

resolve_enum_labels

resolve_enum_labels(
    table_name: str, field_name: str, enum_labels: list[str | int]
) -> list[int | str]

Resolve enum labels to their values for a given field.

Parameters:

Name	Type	Description	Default
`table_name`	`str`	Name of the table containing the field	required
`field_name`	`str`	Name of the field/column	required
`enum_labels`	`list[str \| int]`	List of enum member names (e.g., ['MISSING', 'PNTA'])	required

Returns:

Type	Description
`list[int \| str]`	List of resolved values (integers or strings)

Example

resolve_enum_labels('households', 'income_bin', ['MISSING', 'PNTA']) [995, 999]

utils.helpers

Utility functions for trip linking and data processing.

get_income_midpoint

get_income_midpoint(income_enum: LabeledEnum) -> int

Calculate the midpoint dollar value for an income category enum.

Parses the income range from the enum label and returns the midpoint. For "Under" categories, uses $0 as the lower bound. For "or more" categories, uses 1.25x multiplier to estimate upper bound.

Parameters:

Name	Type	Description	Default
`income_enum`	`LabeledEnum`	Income category enum (IncomeBroad)	required

Returns:

Type	Description
`int`	Midpoint dollar value for the income bracket

Raises:

Type	Description
`ValueError`	If the label format cannot be parsed or is PNTA/Missing

Example

from data_canon.codebook.households import IncomeBroad midpoint = get_income_midpoint(IncomeBroad.INCOME_50TO75) midpoint 62500

datetime_from_parts

datetime_from_parts(
    date: pl.Expr, hour: pl.Expr, minute: pl.Expr, second: pl.Expr
) -> pl.Expr

Construct datetime from date and time parts.

add_time_columns

add_time_columns(
    trips: pl.DataFrame, datetime_format: str = "%Y-%m-%d %H:%M:%S"
) -> pl.DataFrame

Add datetime columns for departure and arrival times if missing.

If datetime columns exist as strings, parse them to datetime type. Otherwise, construct them from component columns.

expr_haversine

expr_haversine(
    lat1: pl.Expr,
    lon1: pl.Expr,
    lat2: pl.Expr,
    lon2: pl.Expr,
    units: str = "meters",
) -> pl.Expr

Return a Polars expression for Haversine distance.

Returns null if any coordinate is null (e.g., missing work/school locations for non-workers/non-students).

get_age_midpoint

get_age_midpoint(age_enum: LabeledEnum) -> int

Calculate the midpoint age value for an age category enum.

Parses the age range from the enum label and returns the midpoint. For "Under" categories, uses 0 as the lower bound. For "and up" categories, estimates upper bound using the category span.

Parameters:

Name	Type	Description	Default
`age_enum`	`LabeledEnum`	Age category enum (e.g., AgeCategory)	required

Returns:

Type	Description
`int`	Midpoint age value for the age bracket

Raises:

Type	Description
`ValueError`	If the label format cannot be parsed

Example

from data_canon.codebook.persons import AgeCategory midpoint = get_age_midpoint(AgeCategory.AGE_35_TO_44) midpoint 39

utils

utils.create_ids

create_concatenated_id

Create person_id from hh_id + person_num

Create day_id from person_id + day_num

Create linked_trip_id from day_id + linked_trip_num

create_linked_trip_id

create_tour_ids

utils.crosswalk

build_crosswalk

utils.enum_helpers

MODEL_MAPPING module-attribute

get_enum_class_for_field

resolve_enum_labels

utils.helpers

get_income_midpoint

datetime_from_parts

add_time_columns

expr_haversine

get_age_midpoint

MODEL_MAPPING `module-attribute`