processing.imputation

Data imputation for missing values in travel diary survey data.

This module provides configurable imputation using established statistical methods. It is designed as a pipeline step that operates on any combination of canonical survey tables (households, persons, days, trips, tours).

Supported methods

KNN -- single-column imputation via K-Nearest Neighbors similarity matching. Best for isolated missing fields where similar records exist.
Random Forest -- single-column imputation using a supervised RF model that auto-selects classifier vs. regressor based on the column type. Best for complex non-linear relationships or mixed feature types.
MICE -- multi-column imputation via Multiple Imputation by Chained Equations. Best for correlated variables (e.g. depart/arrive/duration).

Additional capabilities

Pre-imputation stash -- optional {column}_preimputed columns that preserve the original value (null, PNTA, MISSING, etc.) before imputation, enabling downstream equity analysis and imputation auditing.
Quality validation -- optional k-fold cross-validation that masks known values, re-imputes them, and reports accuracy / RMSE metrics.
Method comparison -- head-to-head benchmark of KNN, RF, and MICE on every imputed column using the same folds and feature sets.
Cross-table features -- join_tables pulls parent-table columns and auto-generates within-household mode features; aggregate_from pivots child rows up to a parent.

Typical pipeline position

steps:
  - name: load_data
  - name: clean_<project>
  - name: imputation        # ← after cleaning, before linking
  - name: link_trips
  - name: detect_habitual_locations
  - name: detect_joint_trips
  - name: extract_tours

Supported relationships for joining tables

These parent → child relationships are supported for join_tables

Child table	Parent table	Join key
persons	households	hh_id
days	persons / households	person_id / hh_id
unlinked_trips	days / persons / households	day_id / person_id / hh_id
linked_trips	days / persons / households	day_id / person_id / hh_id
tours	persons / households	person_id / hh_id

Missing-data assumptions

KNN assumes similar records (by feature distance) share similar values.
MICE assumes Missing At Random (MAR): missingness may depend on observed values but not on the missing value itself.
If data is Missing Not At Random (MNAR), results may be biased.

Current limitations

No stratified imputation (no group_by option for within-group models).
No support for exogenous data sources (PUMS, land use data).
High-cardinality one-hot encoding can slow MICE convergence -- move ordinal/count variables to numeric_features to mitigate.

References

processing.imputation.comparison

Head-to-head comparison of imputation methods via k-fold cross-validation.

For each imputed column, every supported method (KNN, RF, MICE) is evaluated using the same k-fold splits and the same feature set. The result is a summary DataFrame that makes it easy to pick the best method for each field.

compare_imputation_methods

compare_imputation_methods(
    impute_columns: dict[str, list[dict[str, Any]]],
    tables: dict[str, pl.DataFrame],
    n_folds: int = 5,
    sample_pct: float = 5.0,
    random_state: int | None = None,
    output_path: str | None = None,
) -> pl.DataFrame

Run k-fold validation for every method on every imputed column.

For each unique (table, column) found in impute_columns, KNN, RF, and MICE are each evaluated using the same enrichment and feature set. This produces a comparison table that helps choose the best method per field.

Parameters:

Name	Type	Description	Default
`impute_columns`	`dict[str, list[dict[str, Any]]]`	The same config dict passed to `imputation()`.	required
`tables`	`dict[str, pl.DataFrame]`	Dict of canonical DataFrames (already cleaned).	required
`n_folds`	`int`	Number of cross-validation folds.	`5`
`sample_pct`	`float`	Percentage of non-missing values to test (0-100).	`5.0`
`random_state`	`int \| None`	Random seed for reproducibility.	`None`
`output_path`	`str \| None`	Optional path to save the comparison CSV.	`None`

Returns:

Type	Description
`pl.DataFrame`	Polars DataFrame with columns: table, variable, method, type,
`pl.DataFrame`	n_samples, n_folds, accuracy, precision, recall, f1, rmse, mae, r2.

processing.imputation.flags

Stash pre-imputation values for auditing and downstream analysis.

When stash_preimputed=True (the default), the imputation step copies the original column value before any transforms (missing-value replacement, imputation) into a companion column::

{column}_preimputed   - original value (null, PNTA 999, MISSING 995, etc.)

Examples: mode_preimputed, income_broad_preimputed, gender_preimputed.

This preserves the full original value so downstream consumers can distinguish between genuinely missing data (null), "Prefer Not To Answer" (e.g. 999), "Missing Response" (e.g. 995), or any other substantive value that was overwritten during imputation.

To recover the old boolean "was this imputed?" flag::

was_imputed = df[col + "_preimputed"].is_null() & df[col].is_not_null()

stash_preimputed_column

stash_preimputed_column(
    df: pl.DataFrame, original_df: pl.DataFrame, column: str
) -> pl.DataFrame

Copy the pre-imputation value of column from original_df onto df.

The stashed column is named {column}_preimputed and has the same dtype as the original. It captures the value before prepare_column_for_imputation converted sentinel codes to null, so PNTA / MISSING codes are preserved.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame with imputed values.	required
`original_df`	`pl.DataFrame`	Original DataFrame before imputation.	required
`column`	`str`	Name of the column that was imputed.	required

Returns:

Type	Description
`pl.DataFrame`	DataFrame with added `{column}_preimputed` column.

stash_preimputed_columns

stash_preimputed_columns(
    df: pl.DataFrame, original_df: pl.DataFrame, columns: list[str]
) -> pl.DataFrame

Stash pre-imputation values for multiple columns.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame with imputed values.	required
`original_df`	`pl.DataFrame`	Original DataFrame before imputation.	required
`columns`	`list[str]`	List of column names that were imputed.	required

Returns:

Type	Description
`pl.DataFrame`	DataFrame with added `{column}_preimputed` columns.

processing.imputation.generic_impute

Internal helpers for the imputation pipeline step.

The public entry point lives in :mod:processing.imputation.impute. This module provides the shared scaffolding used by all three methods (KNN, Random Forest, MICE): config grouping, per-config dispatch, and k-fold validation.

processing.imputation.impute

Pipeline step entry point for imputation.

imputation

imputation(
    households: pl.DataFrame | None = None,
    persons: pl.DataFrame | None = None,
    days: pl.DataFrame | None = None,
    unlinked_trips: pl.DataFrame | None = None,
    linked_trips: pl.DataFrame | None = None,
    tours: pl.DataFrame | None = None,
    impute_columns: dict[str, list[dict[str, Any]]] | None = None,
    stash_preimputed: bool = True,
    random_state: int | None = None,
    validate_imputation: dict[str, Any] | None = None,
    canonical_data: CanonicalData | None = None,
) -> dict[str, pl.DataFrame]

Impute missing values using KNN, Random Forest, and/or MICE methods.

Each config block specifies its method (knn, rf, or mice) along with the method-specific parameters. Configs are grouped by method and executed in a fixed order (KNN → RF → MICE) across all tables so that later phases can benefit from values filled in earlier phases.

Handling Missing Values with Enum Labels

Survey data often uses special codes for missing values (e.g. 995 for "Missing Response", 999 for "Prefer not to answer"). Use enum member names (labels) rather than raw numeric values in the config::

missing_values: [MISSING, PNTA]   # enum labels, not 995/999

The module automatically:

Maps the table name to the appropriate codebook module (e.g. households → data_canon.codebook.households).
Finds the enum class whose canonical_field_name matches the target column (e.g. income_broad → IncomeBroad).
Resolves enum member names to their values (e.g. MISSING → 995).
Replaces those values with null before imputation.

For MICE with multiple columns, missing_values can be a dict mapping each column to its own labels, or a single list applied to all columns::

# Per-column
missing_values:
  race: [MISSING]
  ethnicity: [MISSING, PNTA]

# Shared
missing_values: [MISSING, PNTA]   # applied to all columns

Cross-Table Features (join_tables): By default only features from the same table are used. Adding join_tables to a config block pulls columns from parent tables via left-join on known foreign keys, which can significantly improve quality.

Behaviour:

1. Columns from the specified parent table(s) are joined onto the
   child table (e.g. ``persons`` ← ``households`` via ``hh_id``).
2. For each target column a ``hh_mode_{column}`` feature is
   auto-generated — the mode of that column among *other* household
   members (exclude-self).  This captures within-household
   correlation (e.g. siblings sharing race/ethnicity).
3. Auto-generated ``hh_mode_*`` columns are appended to
   ``categorical_features`` automatically.
4. After imputation all joined/aggregated columns are stripped;
   the output schema is unchanged.

Example::

    impute_columns:
      persons:
        - method: knn
          column: gender
          n_neighbors: 5
          join_tables: [households]
          categorical_features: [age, employment, income_bin, residence_type]
          #                                       ^^^^^^^^^^  ^^^^^^^^^^^^^^
          #                              columns from the households table

Child-to-Parent Aggregation (aggregate_from): The reverse of join_tables: aggregate child rows up to a parent table. Useful when imputing parent-level fields that depend on household composition (e.g. predicting household income from the employment/education mix of its members).

For each child table and each field listed under ``pivot_count``, the
module groups child rows by the parent's FK and creates one column per
unique value, counting occurrences.  Generated columns are named
``{child_table}_count_{field}={value}`` and are automatically added to
``numeric_features``.  After imputation, all generated columns are
stripped.

Example::

    impute_columns:
      households:
        - method: mice
          columns: [income_bin]
          aggregate_from:
            persons:
              pivot_count: [employment, education, student]
          categorical_features: [residence_type, residence_rent_own]
          max_iter: 10

Parameters:

Name	Type	Description	Default
`households`	`pl.DataFrame \| None`	Households table (optional).	`None`
`persons`	`pl.DataFrame \| None`	Persons table (optional).	`None`
`days`	`pl.DataFrame \| None`	Days table (optional).	`None`
`unlinked_trips`	`pl.DataFrame \| None`	Unlinked trips table (optional).	`None`
`linked_trips`	`pl.DataFrame \| None`	Linked trips table (optional).	`None`
`tours`	`pl.DataFrame \| None`	Tours table (optional).	`None`
`impute_columns`	`dict[str, list[dict[str, Any]]] \| None`	Dict mapping table names to list of imputation configs. Every config dict must include a `method` key (`knn`, `rf`, or `mice`). The remaining keys are method-specific: KNN (`method: knn`): column: Column name to impute. missing_values: Enum labels to treat as missing. n_neighbors: Number of neighbors (default: 5). neighbor_weights: `'uniform'` or `'distance'` (default: `'distance'`). numeric_features: Numeric feature columns. categorical_features: Categorical feature columns. join_tables: Parent tables to left-join for extra features. aggregate_from: Child-to-parent pivot-count config. Random Forest (`method: rf`): column: Column name to impute. missing_values: Enum labels to treat as missing. n_estimators: Number of trees (default: 100). max_depth: Maximum tree depth (default: None, unlimited). numeric_features: Numeric feature columns. categorical_features: Categorical feature columns. join_tables: Parent tables to left-join for extra features. aggregate_from: Child-to-parent pivot-count config. MICE (`method: mice`): columns: Column names to impute together. missing_values: Dict mapping column → enum labels, or a single list applied to all columns. max_iter: Maximum iterations (default: 10). numeric_features: Numeric feature columns. categorical_features: Categorical feature columns. join_tables: Parent tables to left-join for extra features. aggregate_from: Child-to-parent pivot-count config. At least one of `numeric_features` or `categorical_features` is required in every config block.	`None`
`canonical_data`	`CanonicalData \| None`	Canonical data container. When present, each `{column}_preimputed` stash is registered on it so the writer can tell a column the pipeline created from vendor passthrough.	`None`
`stash_preimputed`	`bool`	Whether to create `{column}_preimputed` columns that preserve the original value before imputation (default: True). The stashed column retains sentinel codes (e.g. PNTA 999, MISSING 995) and nulls so downstream consumers can distinguish the reason a value was imputed. To recover the old boolean flag: `col_preimputed.is_null() & col.is_not_null()`.	`True`
`random_state`	`int \| None`	Random seed for reproducibility across all imputation.	`None`
`validate_imputation`	`dict[str, Any] \| None`	Optional k-fold check of imputation accuracy against known values. Its keys: enabled: Whether to run validation (default: False). n_folds: Number of k-folds (default: 5). sample_pct: Percentage of non-missing values to test (default: 5.0). output_path: Path to save validation or comparison CSV. compare_methods: When True, run all three methods (KNN, RF, MICE) against every column instead of validating only the configured method (default: False).	`None`

Returns:

Type	Description
`dict[str, pl.DataFrame]`	Dictionary of imputed tables. When validation is enabled, an extra
`dict[str, pl.DataFrame]`	key `_validation_summary` contains a Polars DataFrame with columns:
`dict[str, pl.DataFrame]`	table, variable, method, type, n_samples, n_folds, accuracy,
`dict[str, pl.DataFrame]`	precision, recall, f1, rmse, mae, r2.
`dict[str, pl.DataFrame]`	When `compare_methods` is True, an extra key
`dict[str, pl.DataFrame]`	`_method_comparison` contains a Polars DataFrame comparing
`dict[str, pl.DataFrame]`	KNN, RF, and MICE for every imputed column.

Example config

.. code-block:: yaml

impute_columns:
  households:
    - method: knn
      column: income_broad
      missing_values: [MISSING, PNTA]
      n_neighbors: 5
      neighbor_weights: distance
      numeric_features: [num_persons, num_vehicles, num_workers]
  persons:
    - method: knn
      column: gender
      missing_values: [MISSING]
      n_neighbors: 5
      join_tables: [households]
      numeric_features: [age]
      categorical_features: [relationship, employment, income_bin]
    - method: rf
      column: education
      missing_values: [MISSING]
      n_estimators: 200
      max_depth: 15
      numeric_features: [age]
      categorical_features: [employment, occupation]
    - method: mice
      columns: [race, ethnicity]
      missing_values:
        race: [MISSING]
        ethnicity: [MISSING, PNTA]
      join_tables: [households]
      max_iter: 10
      numeric_features: [age]
random_state: 42
stash_preimputed: true
validate_imputation:
  enabled: true
  n_folds: 5
  sample_pct: 5.0

processing.imputation.impute_utils

Shared utilities for imputation methods.

Data type handling:

Numeric columns: imputed directly by KNN / RF / MICE.
Categorical integer columns (e.g. enum codes 1-6): automatically encoded to dense 0..N codes before imputation and decoded back afterwards. This prevents non-contiguous codes (e.g. 1, 2, 3, 995, 999) from distorting distance calculations.
Categorical string columns (e.g. "Hispanic", "White"): automatically encoded to integer codes for MICE, then decoded to original labels after imputation. No manual pre-processing required.

Feature selection guidance:

numeric_features: used as-is (continuous values — e.g. num_trips, num_vehicles).
categorical_features: one-hot encoded into binary columns for distance / regression calculations.

Tip: use numeric_features for ordinal or count variables and categorical_features for unordered enums. Putting high-cardinality integers (e.g. raw age) in categorical_features causes feature explosion and slow performance.

FK_RELATIONSHIPS `module-attribute`

FK_RELATIONSHIPS: dict[tuple[str, str], str] = {
    ("persons", "households"): "hh_id",
    ("days", "persons"): "person_id",
    ("days", "households"): "hh_id",
    ("unlinked_trips", "days"): "day_id",
    ("unlinked_trips", "persons"): "person_id",
    ("unlinked_trips", "households"): "hh_id",
    ("linked_trips", "days"): "day_id",
    ("linked_trips", "persons"): "person_id",
    ("linked_trips", "households"): "hh_id",
    ("tours", "persons"): "person_id",
    ("tours", "households"): "hh_id",
}

log_imputation_stats

log_imputation_stats(
    method: str,
    columns: list[str],
    stats: dict[str, Any] | dict[str, dict[str, Any]],
    n_total: int,
) -> None

Log a summary of imputation results for one config block.

Works with both single-column stats (KNN: flat dict) and multi-column stats (MICE: {col: {...}, ...}).

Parameters:

Name	Type	Description	Default
`method`	`str`	Label for the imputation method (e.g. `"KNN"`, `"MICE"`).	required
`columns`	`list[str]`	Target column names that were imputed.	required
`stats`	`dict[str, Any] \| dict[str, dict[str, Any]]`	Stats dict returned by `impute_knn` or `impute_mice`.	required
`n_total`	`int`	Total number of rows in the DataFrame.	required

validate_features_exist

validate_features_exist(
    df: pl.DataFrame,
    numeric_features: list[str] | None = None,
    categorical_features: list[str] | None = None,
) -> None

Validate that at least one feature type is specified and all features exist.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame to validate against	required
`numeric_features`	`list[str] \| None`	Optional list of numeric feature column names	`None`
`categorical_features`	`list[str] \| None`	Optional list of categorical feature column names	`None`

Raises:

Type	Description
`ValueError`	If no features are specified or if any features are missing

is_categorical

is_categorical(df: pl.DataFrame, column: str) -> bool

Determine if a column should be treated as categorical.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame containing the column	required
`column`	`str`	Column name to check	required

Returns:

Type	Description
`bool`	True if column is categorical (non-float numeric or string), False otherwise

prepare_column_for_imputation

prepare_column_for_imputation(
    df: pl.DataFrame,
    table_name: str,
    column: str,
    missing_value_labels: list[str] | None = None,
) -> tuple[pl.DataFrame, list[Any]]

Prepare a column for imputation by replacing missing values with null.

This function resolves enum labels to their numeric values and replaces them with null, making the column ready for imputation algorithms.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame containing the column	required
`table_name`	`str`	Name of the table (for enum resolution)	required
`column`	`str`	Column name to prepare	required
`missing_value_labels`	`list[str] \| None`	Optional list of enum labels to treat as missing (e.g., ['MISSING', 'PNTA'])	`None`

Returns:

Type	Description
`pl.DataFrame`	Tuple of (prepared_df, resolved_values) where resolved_values is the list
`list[Any]`	of numeric/string values that were replaced with null

Example

df = pl.DataFrame({'income_broad': [1, 2, 995, 999, 3]}) df_prep, values = prepare_column_for_imputation( ... df, 'households', 'income_broad', ['MISSING', 'PNTA'] ... ) values [995, 999]

decode_integer_to_string

decode_integer_to_string(
    imputed_values: np.ndarray, int_to_label: dict[int, str]
) -> list[str]

Decode imputed float values back to string labels.

Rounds each value to the nearest integer and clamps to the valid range of encoded labels.

decode_dense_to_integer

decode_dense_to_integer(
    imputed_values: np.ndarray, dense_to_original: dict[int, int]
) -> list[int]

Decode dense 0..N predictions back to original integer codes.

Rounds each value to the nearest integer and clamps to the valid range of dense keys before looking up the original code.

encode_string_columns

encode_string_columns(
    df: pl.DataFrame, columns: list[str], verbose: bool = False
) -> tuple[pl.DataFrame, dict[str, dict[int, str]]]

Encode string columns to integer values for use in numeric imputation.

Maps each unique non-null string value to an integer, preserving nulls. Returns the modified DataFrame and a mapping to decode back to strings.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame containing the columns to encode	required
`columns`	`list[str]`	List of column names to check and encode	required
`verbose`	`bool`	Whether to log encoding details	`False`

Returns:

Type	Description
`pl.DataFrame`	Tuple of (encoded_df, encodings) where encodings is a dict mapping
`dict[str, dict[int, str]]`	column name -> {int: label} for each encoded column

encode_integer_categoricals

encode_integer_categoricals(
    df: pl.DataFrame, columns: list[str], verbose: bool = False
) -> tuple[pl.DataFrame, dict[str, dict[int, int]]]

Encode integer categorical columns to dense 0..N codes.

Integer-coded enums (e.g. income_bin with codes 1-6) are already close to dense, but other fields may have non-contiguous codes. Mapping to dense 0-based integers keeps predictions in a valid range so that rounding always lands on a real category.

Columns that are already coded 0..N-1 are skipped.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame containing the columns to encode	required
`columns`	`list[str]`	Target column names to check	required
`verbose`	`bool`	Whether to log encoding details	`False`

Returns:

Type	Description
`pl.DataFrame`	Tuple of (encoded_df, encodings) where encodings maps
`dict[str, dict[int, int]]`	`column_name -> {dense_int: original_int}` for every column
`tuple[pl.DataFrame, dict[str, dict[int, int]]]`	that was re-coded.

build_feature_matrix

build_feature_matrix(
    df: pl.DataFrame,
    target_columns: list[str],
    numeric_features: list[str],
    categorical_features: list[str],
) -> tuple[np.ndarray, dict[str, int], list[str]]

Build a numpy feature matrix for imputation algorithms.

Constructs a matrix from:

Continuous columns — all numeric_features present in df plus any target_columns that are numeric. Order-preserving deduplication ensures each column appears at most once.
One-hot encoded categoricals — each unique value in a categorical feature becomes its own 0/1 column. Target columns are excluded from one-hot expansion.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame to convert	required
`target_columns`	`list[str]`	Columns being imputed (included as continuous, excluded from one-hot expansion)	required
`numeric_features`	`list[str]`	Numeric / continuous feature column names	required
`categorical_features`	`list[str]`	Categorical feature column names (one-hot encoded)	required

Returns:

Type	Description
`np.ndarray`	`(matrix, column_indices, feature_names)` where column_indices
`dict[str, int]`	maps each target column name to its column offset in matrix, and
`list[str]`	feature_names lists all column names in matrix order. One-hot
`tuple[np.ndarray, dict[str, int], list[str]]`	encoded columns are named `{cat_col}={value}`.

join_parent_tables

join_parent_tables(
    df: pl.DataFrame,
    table_name: str,
    tables: dict[str, pl.DataFrame],
    join_tables: list[str],
) -> tuple[pl.DataFrame, list[str]]

Join parent table columns onto a child table for use as imputation features.

Looks up the foreign key relationship between the target table and each requested parent table, then left-joins all parent columns (excluding the join key itself, which already exists on the child).

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	The child DataFrame to enrich (e.g., persons)	required
`table_name`	`str`	Name of the child table (e.g., "persons")	required
`tables`	`dict[str, pl.DataFrame]`	Dict of all available canonical tables	required
`join_tables`	`list[str]`	List of parent table names to join (e.g., ["households"])	required

Returns:

Type	Description
`tuple[pl.DataFrame, list[str]]`	Tuple of (enriched_df, list of added column names)

Raises:

Type	Description
`ValueError`	If a requested parent table is not available or if no foreign key relationship is defined.

add_household_agg_features

add_household_agg_features(
    df: pl.DataFrame,
    target_columns: list[str],
    hh_id_col: str = "hh_id",
    person_id_col: str = "person_id",
) -> tuple[pl.DataFrame, list[str]]

Add household-aggregated features for target columns.

For each target column, computes the mode of that column among other household members (excluding the current person). This captures within-household correlation (e.g., siblings likely share race/ethnicity).

The new columns are named hh_mode_{column}.

For single-person households or when all other members have null values, the aggregated value will be null.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame with person-level records (must contain hh_id and person_id)	required
`target_columns`	`list[str]`	List of columns to compute household aggregates for	required
`hh_id_col`	`str`	Name of the household ID column	`'hh_id'`
`person_id_col`	`str`	Name of the person ID column	`'person_id'`

Returns:

Type	Description
`tuple[pl.DataFrame, list[str]]`	Tuple of (enriched_df, list of added column names)

aggregate_from_children

aggregate_from_children(
    df: pl.DataFrame,
    table_name: str,
    tables: dict[str, pl.DataFrame],
    aggregate_config: dict[str, dict[str, list[str]]],
    verbose: bool = True,
) -> tuple[pl.DataFrame, list[str]]

Aggregate child table columns onto a parent table via pivot counts.

For each child table and each field listed under pivot_count, creates one column per unique value counting occurrences within each parent group.

For example, pivot_count: [employment] on persons grouped by hh_id creates persons_count_employment=1, persons_count_employment=2, etc. The sum across a field's pivoted columns equals the household size, so a separate count option is unnecessary.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	The parent DataFrame to enrich (e.g., households)	required
`table_name`	`str`	Name of the parent table (e.g., "households")	required
`tables`	`dict[str, pl.DataFrame]`	Dict of all available canonical tables	required
`aggregate_config`	`dict[str, dict[str, list[str]]]`	Mapping of child table name to aggregation spec. Each spec supports: - `pivot_count`: list of child column names to pivot-count	required
`verbose`	`bool`	Whether to log aggregation details	`True`

Returns:

Type	Description
`tuple[pl.DataFrame, list[str]]`	Tuple of (enriched_df, list of added column names)

Raises:

Type	Description
`ValueError`	If no FK relationship is defined for the child→parent pair.

strip_joined_columns

strip_joined_columns(
    df: pl.DataFrame, added_columns: list[str]
) -> pl.DataFrame

Remove columns that were temporarily added for imputation.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame with extra columns	required
`added_columns`	`list[str]`	List of column names to remove	required

Returns:

Type	Description
`pl.DataFrame`	DataFrame with added columns removed

enrich_dataframe

enrich_dataframe(
    df: pl.DataFrame,
    table_name: str,
    tables: dict[str, pl.DataFrame] | None,
    join_tables_list: list[str],
    target_columns: list[str],
    categorical_features: list[str] | None,
    aggregate_from_config: (
        dict[str, dict[str, list[str]]] | None
    ) = None,
) -> tuple[pl.DataFrame, list[str], list[str] | None]

Enrich a DataFrame with parent joins and child aggregations.

Handles both directions

Parent→child joins via join_tables config
Child→parent aggregations via aggregate_from config

Returns:

Type	Description
`tuple[pl.DataFrame, list[str], list[str] \| None]`	Tuple of (enriched_df, added_column_names, updated_categorical_features)

processing.imputation.knn

KNN-based imputation for missing values.

impute_knn

impute_knn(
    df: pl.DataFrame,
    column: str,
    n_neighbors: int = 5,
    neighbor_weights: Literal["uniform", "distance"] = "distance",
    numeric_features: list[str] | None = None,
    categorical_features: list[str] | None = None,
) -> tuple[pl.DataFrame, dict[str, Any]]

Impute missing values in a single column using K-Nearest Neighbors.

Best for: single columns with isolated missing values where similar records exist in the dataset.

How it works:

Build a feature matrix from numeric_features (used as-is) and categorical_features (one-hot encoded for distance calculation).
Non-contiguous integer codes (e.g. enum values 1, 2, 3, 995, 999) are automatically encoded to dense 0..N codes so they don't distort distance calculations, then decoded back after imputation.
For each row with a missing value, find the K most similar records based on Euclidean distance across all features.
Impute the missing value using the weighted average (or mode for categoricals) of the K neighbours.

neighbor_weights='distance' weights closer neighbours more heavily; neighbor_weights='uniform' treats all K neighbours equally.

Example use cases:

Missing trip mode when other trip attributes are known.
Missing person age when household/demographic info is available.
Missing trip distance when other spatial/temporal features exist.

Performance: O(n log n) complexity; scales well to medium-large datasets.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame containing the column to impute.	required
`column`	`str`	Name of the column to impute.	required
`n_neighbors`	`int`	Number of similar records to use (default: 5).	`5`
`neighbor_weights`	`Literal['uniform', 'distance']`	`'distance'` or `'uniform'` (default: `'distance'`).	`'distance'`
`numeric_features`	`list[str] \| None`	Numeric/continuous feature columns. Used as-is.	`None`
`categorical_features`	`list[str] \| None`	Categorical feature columns. One-hot encoded into binary columns for distance calculation.	`None`

Returns:

Type	Description
`pl.DataFrame`	Tuple of (imputed DataFrame, stats dict). The stats dict contains
`dict[str, Any]`	`n_missing`, `n_imputed`, and `pct_imputed`.

processing.imputation.mice

MICE-based imputation for missing values.

impute_mice

impute_mice(
    df: pl.DataFrame,
    columns: list[str],
    max_iter: int = 10,
    random_state: int | None = None,
    numeric_features: list[str] | None = None,
    categorical_features: list[str] | None = None,
    verbose: bool = True,
) -> tuple[pl.DataFrame, dict[str, Any]]

Impute missing values in multiple correlated columns using MICE.

Best for: multiple correlated columns with missing values (e.g. depart_hour / arrive_hour / duration, or race / ethnicity).

MICE (Multiple Imputation by Chained Equations) imputes several variables together, preserving their joint distribution.

How it works:

Initialise missing values with simple imputation (mean/mode).
For each column with missing values:
1. Treat it as the target variable.
2. Use the other columns as predictors in a regression model.
3. Predict and update missing values.
Repeat iteratively until convergence (max_iter rounds).

Categorical integer columns (e.g. enum codes 1-6) are automatically encoded to dense 0..N codes before imputation and decoded back afterwards. String columns are auto-encoded to integers for the MICE model and decoded to original labels after imputation.

Assumes Missing At Random (MAR): missingness may depend on observed values but not on the missing value itself. If data is Missing Not At Random (MNAR), results may be biased.

Example use cases:

Time fields (depart_hour, arrive_hour, duration) — highly correlated.
Spatial coordinates (origin_lat, origin_lon) — spatially correlated.
Socio-demographic variables (income, education, employment) — often correlated.

Performance: iterative, can be slow for many columns or large datasets.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame containing the columns to impute.	required
`columns`	`list[str]`	Column names to impute together.	required
`max_iter`	`int`	Maximum number of imputation rounds (default: 10).	`10`
`random_state`	`int \| None`	Random seed for reproducibility.	`None`
`numeric_features`	`list[str] \| None`	Numeric/continuous feature columns.	`None`
`categorical_features`	`list[str] \| None`	Categorical feature columns (one-hot encoded).	`None`
`verbose`	`bool`	Whether to log progress during imputation.	`True`

Returns:

Type	Description
`pl.DataFrame`	Tuple of (imputed DataFrame, stats dict). The stats dict is keyed
`dict[str, Any]`	by column name, each containing `n_missing`, `n_imputed`, and
`tuple[pl.DataFrame, dict[str, Any]]`	`pct_imputed`.

processing.imputation.random_forest

Random Forest imputation for missing values.

impute_random_forest

impute_random_forest(
    df: pl.DataFrame,
    column: str,
    n_estimators: int = 100,
    max_depth: int | None = None,
    random_state: int | None = None,
    numeric_features: list[str] | None = None,
    categorical_features: list[str] | None = None,
    verbose: bool = True,
) -> tuple[pl.DataFrame, dict[str, Any]]

Impute missing values in a single column using Random Forest.

Best for: single columns with complex non-linear relationships or mixed feature types where KNN may struggle with decision boundaries.

How it works:

Split rows into known (have a value) and missing (need imputation).
Train a Random Forest model on the known rows using all features.
Automatically select RandomForestClassifier for categorical targets (integer / string dtypes) or RandomForestRegressor for continuous targets (float dtypes).
Predict missing values using the trained model.
NaN values in features are filled with column medians before training.

Non-contiguous integer codes (e.g. enum values 1, 2, 3, 995, 999) are automatically encoded to dense 0..N codes so they don't distort the model, then decoded back after prediction.

Example use cases:

Missing education level when employment, occupation, and age are available.
Missing income category with many mixed-type predictors.
Cases where KNN struggles with non-linear decision boundaries.

Performance: trains on known values only; handles mixed types well but can be memory-intensive with many trees.

Parameters:

Name	Type	Description	Default
`df`	`pl.DataFrame`	DataFrame containing the column to impute.	required
`column`	`str`	Name of the column to impute.	required
`n_estimators`	`int`	Number of trees in the forest (default: 100).	`100`
`max_depth`	`int \| None`	Maximum tree depth (default: None = unlimited).	`None`
`random_state`	`int \| None`	Random seed for reproducibility.	`None`
`numeric_features`	`list[str] \| None`	Numeric/continuous feature columns.	`None`
`categorical_features`	`list[str] \| None`	Categorical feature columns (one-hot encoded).	`None`
`verbose`	`bool`	If True, logs top features by importance.	`True`

Returns:

Type	Description
`pl.DataFrame`	Tuple of (imputed DataFrame, stats dict). The stats dict contains
`dict[str, Any]`	`n_missing`, `n_imputed`, `pct_imputed`, and
`tuple[pl.DataFrame, dict[str, Any]]`	`feature_importance` (a `dict[str, float]` mapping each input
`tuple[pl.DataFrame, dict[str, Any]]`	feature to its Mean Decrease in Impurity score, sorted descending).
`tuple[pl.DataFrame, dict[str, Any]]`	See the Feature Importance section in `docs/pipeline_steps/imputation.md`
`tuple[pl.DataFrame, dict[str, Any]]`	for interpretation guidance.

processing.imputation.validation

K-fold cross-validation for imputation quality assessment.

Optional validation that assesses how accurate the imputation is by:

Sampling a percentage of non-missing values (user-configurable, e.g. 5%).
Artificially masking those values (setting them to null).
Imputing them using k-fold cross-validation.
Comparing imputed vs. actual values.
Computing and logging quality metrics.

Metrics by data type:

Categorical columns (e.g. mode, purpose): Accuracy, Precision, Recall, F1-Score.
Continuous columns (e.g. distance, duration): RMSE, MAE, R².

Configuration::

random_state: 42
validate_imputation:
  enabled: true
  n_folds: 5          # number of CV folds (default: 5)
  sample_pct: 5.0     # % of complete values to test (default: 5%)

Example output::

============================================================
Imputation Validation Results
============================================================

Column: mode (categorical, n=250 test samples)
  Accuracy:  0.876
  Precision: 0.883
  Recall:    0.876
  F1-Score:  0.872

Column: distance (continuous, n=250 test samples)
  RMSE: 2.34
  MAE:  1.82
  R²:   0.721
============================================================

Validation uses the same enrichment (joins, aggregations) as the real imputation pipeline, so metrics reflect the full feature set. Note that validation adds computational overhead (k-fold = k x the imputation time); it is recommended for development/testing and optional in production.

validate_knn_imputation

validate_knn_imputation(
    df: pl.DataFrame,
    column: str,
    n_folds: int,
    sample_pct: float,
    n_neighbors: int,
    neighbor_weights: Literal["uniform", "distance"],
    random_state: int | None = None,
    numeric_features: list[str] | None = None,
    categorical_features: list[str] | None = None,
) -> dict[str, Any]

Validate KNN imputation quality using k-fold cross-validation.

validate_mice_imputation

validate_mice_imputation(
    df: pl.DataFrame,
    columns: list[str],
    n_folds: int,
    sample_pct: float,
    max_iter: int,
    random_state: int | None = None,
    numeric_features: list[str] | None = None,
    categorical_features: list[str] | None = None,
) -> dict[str, dict[str, Any]]

Validate MICE imputation quality using k-fold cross-validation.

validate_rf_imputation

validate_rf_imputation(
    df: pl.DataFrame,
    column: str,
    n_folds: int,
    sample_pct: float,
    n_estimators: int = 100,
    max_depth: int | None = None,
    random_state: int | None = None,
    numeric_features: list[str] | None = None,
    categorical_features: list[str] | None = None,
) -> dict[str, Any]

Validate Random Forest imputation quality using k-fold cross-validation.

log_validation_results

log_validation_results(
    metrics: dict[str, Any] | dict[str, dict[str, Any]],
) -> None

Log validation metrics in a readable format.

Parameters:

Name	Type	Description	Default
`metrics`	`dict[str, Any] \| dict[str, dict[str, Any]]`	Dictionary of metrics (single column or per-column dict)	required

processing.imputation

Supported methods

Additional capabilities

Typical pipeline position

Supported relationships for joining tables

Missing-data assumptions

Current limitations

References

processing.imputation.comparison

compare_imputation_methods

processing.imputation.flags

stash_preimputed_column

stash_preimputed_columns

processing.imputation.generic_impute

processing.imputation.impute

imputation

processing.imputation.impute_utils

FK_RELATIONSHIPS module-attribute

log_imputation_stats

validate_features_exist

is_categorical

prepare_column_for_imputation

decode_integer_to_string

decode_dense_to_integer

encode_string_columns

encode_integer_categoricals

build_feature_matrix

join_parent_tables

add_household_agg_features

aggregate_from_children

strip_joined_columns

enrich_dataframe

processing.imputation.knn

impute_knn

processing.imputation.mice

impute_mice

processing.imputation.random_forest

impute_random_forest

processing.imputation.validation

validate_knn_imputation

validate_mice_imputation

validate_rf_imputation

log_validation_results

FK_RELATIONSHIPS `module-attribute`