cursus.steps.scripts.stratified_sampling

Stratified Sampling Script

Applies stratified sampling to input data with four allocation strategies: 1. Balanced — equal samples per stratum (class imbalance correction) 2. Proportional with minimum — proportional allocation with floor constraints (causal analysis) 3. Optimal (Neyman) — variance-weighted allocation (minimizes sampling error) 4. External proportional — sample to match an external reference distribution with multiplier

Features: - Sampling with replacement (allow_replacement) for oversampling when target > available - NaN guard: warns and excludes NaN strata values - Empty DataFrame guard: returns empty result gracefully - Per-split diagnostics JSON output (requested vs achieved per stratum) - Format preservation: reads and writes CSV/TSV/Parquet maintaining input format - Split-aware: processes train/val splits for training job type, copies test unchanged - Reference counts loaded from sidecar file (reference_counts.json) or env var fallback

Input: /opt/ml/processing/input/data/{split}/{split}_processed_data.{csv|tsv|parquet} Output: /opt/ml/processing/output/{split}/{split}_processed_data.{csv|tsv|parquet} Diagnostics: /opt/ml/processing/output/{split}/sampling_diagnostics.json

class StratifiedSampler(random_state=42)[source]

Bases: object

Stratified sampling implementation with four allocation strategies: 1. Balanced allocation - for class imbalance 2. Proportional with minimum constraints - for causal analysis 3. Optimal allocation (Neyman) - for variance optimization 4. External proportional - sample to match an external reference distribution

sample(df, strata_column, target_size, strategy='balanced', min_samples_per_stratum=10, variance_column=None, reference_counts=None, multiplier=1.0, allow_replacement=False)[source]

Perform stratified sampling on a DataFrame.

Parameters:

  • df (DataFrame) – Input DataFrame

  • strata_column (str) – Column name to stratify by

  • target_size (int) – Total desired sample size

  • strategy (str) – Sampling strategy (‘balanced’, ‘proportional_min’, ‘optimal’, ‘external_proportional’)

  • min_samples_per_stratum (int) – Minimum samples per stratum

  • variance_column (str | None) – Column for variance calculation (needed for optimal strategy)

  • reference_counts (Dict[str, int] | None) – External reference distribution {stratum: count} (for external_proportional)

  • multiplier (float) – Multiplier for reference counts (e.g., 5.0 for 5× oversampling)

  • allow_replacement (bool) – Allow sampling with replacement when target > available

Returns:

Sampled DataFrame

Return type:

DataFrame

main(input_paths, output_paths, environ_vars, job_args, logger=None)[source]

Main logic for stratified sampling, following tabular_preprocessing format.

Parameters:

  • input_paths (Dict[str, str]) – Dictionary of input paths with logical names

  • output_paths (Dict[str, str]) – Dictionary of output paths with logical names

  • environ_vars (Dict[str, str]) – Dictionary of environment variables

  • job_args (Namespace) – Command line arguments

  • logger (Callable[[str], None] | None) – Optional logger object (defaults to print if None)

Returns:

Dictionary of sampled DataFrames by split name

Return type:

Dict[str, DataFrame]