diff --git a/CHANGELOG.md b/CHANGELOG.md index 6d2d840..e81cf86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,14 +22,28 @@ - Multiclass flipper finds where the pairwise margin `p[b] − p[a]` crosses 0; multilabel flipper samples pairs per label (new `pair_label` grid column) and crosses each label's own threshold. +- **Single-target regression models.** `task_type: "regression"` in + `metadata.json` (default output `variable`, then `predictions`), and + `--task regression` with true values from `--targets + --target-col `, joined by ID as text. `baselines` are `rmse`, `mae` + and `r2`; resiliency is blackout R² ÷ clean R² (the share of skill kept) + under `metric_based.r2`; jitter variance is scaled by the target's variance. + Flipper pairs patients across a clinical cutoff (`--flip-threshold`, in + target units, passed to both `simulate` and `evaluate`) and finds where the + prediction crosses it. A new `regression` section reports these in target + units; classifier JSON is unchanged. ### Changed - **Breaking (library):** `Model::predict_proba` returns a `DataFrame` with one column per output, `SimulationVariant.y` is a `DataFrame` label matrix, `EvaluationConfig` gains `task`, and `OnnxModel::from_file` takes - `positive_class_index: Option`. CLI usage for binary models is - unchanged. + `positive_class_index: Option`. `TaskType` gains `Regression` and + `EvaluationResult` gains an optional `regression` field. CLI usage for + binary models is unchanged. +- `evaluate`/`compare` `--labels` is no longer required at parse time; it is + required for classifiers and refused for regression. `--flip-threshold` + still defaults to 0.5 for classifiers. - **ROC AUC handles tied scores.** Ties count half (Mann–Whitney) instead of depending on row order. Results without ties agree with the previous implementation to floating-point rounding. diff --git a/cli/README.md b/cli/README.md index ed68aa9..1bda6fd 100644 --- a/cli/README.md +++ b/cli/README.md @@ -48,7 +48,9 @@ pasteur-cli simulate \ | Flag | Default | Description | |------|---------|-------------| | `--input` | *(required)* | Local clean parquet | -| `--labels` | — | Group-membership parquet ([schema](../docs/source/data-contracts.rst)) for flipper pair sampling (omit to skip flipper) | +| `--labels` | — | Group-membership parquet ([schema](../docs/source/data-contracts.rst)) for flipper pair sampling (omit to skip flipper). Classifiers only | +| `--targets`, `--target-col` | — | Regression: targets parquet and its target column ([schema](../docs/source/data-contracts.rst)) | +| `--targets-id-col` | `node_id` | Regression: ID column of `--targets`, matched as text | | `--output` | `output` | Root for sim-type folders | | `--dataset-name` | — | Provenance label for logs | | `--id-col` | `node_id` | Row id column (becomes `source_row_id`) | @@ -58,7 +60,8 @@ pasteur-cli simulate \ | `--jitter-scale` | `1.0` | Jitter noise scale | | `--jitter-iters` | `3` | Number of `jitter_` variants | | `--positive-group-id` | `103` | `group_id` for positive cohort (flipper). Repeat once per class/label for `--task multiclass`/`multilabel` | -| `--task` | `binary` | `binary`, `multiclass`, or `multilabel`: how flipper pairs are sampled | +| `--task` | `binary` | `binary`, `multiclass`, `multilabel`, or `regression`: how flipper pairs are sampled | +| `--flip-threshold` | — | Regression only: clinical cutoff (target units) that flipper pairs straddle; omit to skip flipper | | `--flipper-pairs` | `200` | Cross-label pairs to sample | | `--flipper-steps` | `20` | Interpolation steps per pair | | `--random-state` | `42` | RNG seed | @@ -114,7 +117,9 @@ pasteur-cli evaluate blackout \ |------|---------|-------------| | `sim_type` | *(positional)* | e.g. `blackout`, `jitter`, `flipper` | | `--sim-root` | *(required)* | Local bundle with `clean/` and `/` | -| `--labels` | *(required)* | Local `groups.parquet` for label derivation | +| `--labels` | *(classifiers)* | Local `groups.parquet` for label derivation | +| `--targets`, `--target-col` | *(regression)* | Local targets parquet and its target column | +| `--targets-id-col` | `node_id` | ID column of `--targets`, matched as text | | `--model` | *(required)* | Local `.onnx` path | | `--contract` | — | Explicit `metadata.json` path (default: next to `--model`) | | `--input-name` | `input` | ONNX input tensor name | @@ -123,7 +128,7 @@ pasteur-cli evaluate blackout \ | `--positive-group-id` | `103` | `group_id` for positive cohort; repeat once per model output (in `output.classes` order) for multiclass/multilabel | | `--task` | `binary` | Must match the model's `metadata.json` `task_type` | | `--dataset-name` | `simulation` | Internal dataset key in JSON | -| `--flip-threshold` | `0.5` | Decision threshold for flipper stability (multilabel: labels without a contract threshold; unused for multiclass) | +| `--flip-threshold` | `0.5` | Decision threshold for flipper stability (multilabel: labels without a contract threshold; unused for multiclass). Regression: clinical cutoff in target units, no default, required for `flipper` | | `-o, --output` | stdout | Write `EvaluationResult` JSON to file | --- @@ -148,7 +153,7 @@ Same flags as `evaluate`, except: - `--predictions-out` writes a local per-patient predictions parquet - A single `--contract` applies to every model -**Predictions schema:** `source_row_id`, `variant_name`, `label`, one f64 column per model (stem of model path), `all_agree` (all models on the same side of 0.5). +**Predictions schema:** `source_row_id`, `variant_name`, `label`, one f64 column per model (stem of model path), `all_agree` (all models on the same side of 0.5). Regression writes `target` instead of `label`, and `all_agree` compares sides of `--flip-threshold` (null without one). The `--output` directory must already exist. @@ -181,7 +186,7 @@ pasteur-cli cache clear --kind simulations Defined in [`core/src/schema.rs`](../core/src/schema.rs): -- **`EvaluationResult`**: `baselines`, `evaluations.metric_invariant`, `evaluations.metric_based`, `created_at` +- **`EvaluationResult`**: `baselines`, `evaluations.metric_invariant`, `evaluations.metric_based`, `created_at`, plus optional `multi` (multiclass/multilabel) or `regression` sections - **`ComparisonResult`**: `evaluations: [{ model_label, evaluation }]` --- diff --git a/cli/src/args.rs b/cli/src/args.rs index d2d7edf..daa619e 100644 --- a/cli/src/args.rs +++ b/cli/src/args.rs @@ -36,6 +36,8 @@ pub enum TaskArg { Multiclass, /// One cohort per label; rows may belong to any number. Multilabel, + /// One continuous target per row, from --targets; the model predicts it. + Regression, } impl From for TaskType { @@ -44,10 +46,28 @@ impl From for TaskType { TaskArg::Binary => TaskType::Binary, TaskArg::Multiclass => TaskType::Multiclass, TaskArg::Multilabel => TaskType::Multilabel, + TaskArg::Regression => TaskType::Regression, } } } +/// Where a regression model's true values come from. Classifiers take +/// `--labels` instead. +#[derive(clap::Args, Clone)] +pub struct TargetArgs { + /// Regression only: local parquet with one row per patient, holding an ID + /// column and the true target value. + #[arg(long)] + pub targets: Option, + /// Regression only: the target column in --targets (e.g. `hba1c`). + #[arg(long)] + pub target_col: Option, + /// Regression only: the ID column in --targets, matched as text to the + /// clean input's --id-col values. + #[arg(long, default_value = "node_id")] + pub targets_id_col: String, +} + #[derive(clap::Args)] pub struct CacheArgs { #[command(subcommand)] @@ -122,18 +142,26 @@ pub struct SimulateArgs { pub jitter_iters: u32, #[arg(long, default_value_t = 42)] pub random_state: u64, - /// Group/cohort membership table used to derive labels for flipper pair sampling. + /// Group/cohort membership table used to derive labels for flipper pair + /// sampling. Classification only; regression takes --targets. #[arg(long)] pub labels: Option, + #[command(flatten)] + pub targets: TargetArgs, /// Cohort `group_id` defining the labels. Binary takes one (the positive /// cohort); multiclass and multilabel take one per model output, repeated - /// in the model's `output.classes` order. + /// in the model's `output.classes` order. Unused for regression. #[arg(long = "positive-group-id", default_values_t = [103u32])] pub positive_group_ids: Vec, - /// `binary`, `multiclass`, or `multilabel`; must match each model's - /// `metadata.json` `task_type`. + /// `binary`, `multiclass`, `multilabel`, or `regression`; must match each + /// model's `metadata.json` `task_type`. #[arg(long, value_enum, default_value_t = TaskArg::Binary)] pub task: TaskArg, + /// Regression only: the clinical cutoff, in the target's units, that + /// flipper pairs are sampled across (one patient below it, one at or + /// above). Without it the flipper grid is skipped for regression. + #[arg(long)] + pub flip_threshold: Option, #[arg(long, default_value_t = 200)] pub flipper_pairs: usize, #[arg(long, default_value_t = 20)] @@ -171,8 +199,12 @@ pub struct EvaluateArgs { /// Local root with `clean/` and `/` subfolders. #[arg(long)] pub sim_root: PathBuf, + /// Group/cohort membership table. Required for classification; regression + /// takes --targets. #[arg(long)] - pub labels: PathBuf, + pub labels: Option, + #[command(flatten)] + pub targets: TargetArgs, #[arg(long)] pub model: PathBuf, /// Model contract sidecar (`metadata.json`). Defaults to @@ -193,11 +225,11 @@ pub struct EvaluateArgs { pub positive_class_index: Option, /// Cohort `group_id` defining the labels. Binary takes one (the positive /// cohort); multiclass and multilabel take one per model output, repeated - /// in the model's `output.classes` order. + /// in the model's `output.classes` order. Unused for regression. #[arg(long = "positive-group-id", default_values_t = [103u32])] pub positive_group_ids: Vec, - /// `binary`, `multiclass`, or `multilabel`; must match each model's - /// `metadata.json` `task_type`. + /// `binary`, `multiclass`, `multilabel`, or `regression`; must match each + /// model's `metadata.json` `task_type`. #[arg(long, value_enum, default_value_t = TaskArg::Binary)] pub task: TaskArg, #[arg(long, default_value = "simulation")] @@ -205,8 +237,10 @@ pub struct EvaluateArgs { /// Decision threshold for flipper crossings: the positive-class /// probability for binary models, every label without a contract /// `output.thresholds` for multilabel. Unused for multiclass (argmax). - #[arg(long, default_value_t = 0.5)] - pub flip_threshold: f64, + /// Defaults to 0.5. For regression it is the clinical cutoff in the + /// target's units, has no default, and must match `simulate`'s. + #[arg(long)] + pub flip_threshold: Option, #[arg(short, long)] pub output: Option, } @@ -216,8 +250,12 @@ pub struct CompareArgs { pub sim_type: String, #[arg(long)] pub sim_root: PathBuf, + /// Group/cohort membership table. Required for classification; regression + /// takes --targets. #[arg(long)] - pub labels: PathBuf, + pub labels: Option, + #[command(flatten)] + pub targets: TargetArgs, #[arg(long = "model", required = true)] pub models: Vec, /// Model contract sidecar (`metadata.json`). Defaults to @@ -238,11 +276,11 @@ pub struct CompareArgs { pub positive_class_index: Option, /// Cohort `group_id` defining the labels. Binary takes one (the positive /// cohort); multiclass and multilabel take one per model output, repeated - /// in the model's `output.classes` order. + /// in the model's `output.classes` order. Unused for regression. #[arg(long = "positive-group-id", default_values_t = [103u32])] pub positive_group_ids: Vec, - /// `binary`, `multiclass`, or `multilabel`; must match each model's - /// `metadata.json` `task_type`. + /// `binary`, `multiclass`, `multilabel`, or `regression`; must match each + /// model's `metadata.json` `task_type`. #[arg(long, value_enum, default_value_t = TaskArg::Binary)] pub task: TaskArg, #[arg(long, default_value = "simulation")] @@ -250,8 +288,10 @@ pub struct CompareArgs { /// Decision threshold for flipper crossings: the positive-class /// probability for binary models, every label without a contract /// `output.thresholds` for multilabel. Unused for multiclass (argmax). - #[arg(long, default_value_t = 0.5)] - pub flip_threshold: f64, + /// Defaults to 0.5. For regression it is the clinical cutoff in the + /// target's units, has no default, and must match `simulate`'s. + #[arg(long)] + pub flip_threshold: Option, #[arg(short, long)] pub output: Option, /// Local path for per-patient predictions parquet (agent uploads with `hf`). diff --git a/cli/src/commands/compare.rs b/cli/src/commands/compare.rs index 7a0f434..bd7555f 100644 --- a/cli/src/commands/compare.rs +++ b/cli/src/commands/compare.rs @@ -5,19 +5,26 @@ use pasteur_model::OnnxModel; use crate::args::CompareArgs; use crate::data::{ build_evaluation_inputs, build_predictions_table, flipper_feature_order_from_inputs, + resolve_flip_threshold, LabelSource, }; use crate::io::{print_or_write, write_predictions_parquet}; use crate::model::{load_model, model_label}; pub fn run(args: CompareArgs) -> Result<()> { + let task = args.task.into(); + let source = LabelSource::required( + task, + args.labels.as_deref(), + &args.positive_group_ids, + &args.targets, + )?; let inputs = build_evaluation_inputs( &args.sim_root, &args.sim_type, - &args.labels, - &args.positive_group_ids, - args.task.into(), + &source, + task, &args.dataset_name, - args.flip_threshold, + resolve_flip_threshold(task, &args.sim_type, args.flip_threshold)?, )?; let feature_order = flipper_feature_order_from_inputs(&inputs, &args.sim_type); diff --git a/cli/src/commands/evaluate.rs b/cli/src/commands/evaluate.rs index 130acde..4d08f5a 100644 --- a/cli/src/commands/evaluate.rs +++ b/cli/src/commands/evaluate.rs @@ -2,19 +2,27 @@ use anyhow::Result; use pasteur_core::Evaluator; use crate::args::EvaluateArgs; -use crate::data::{build_evaluation_inputs, flipper_feature_order_from_inputs}; +use crate::data::{ + build_evaluation_inputs, flipper_feature_order_from_inputs, resolve_flip_threshold, LabelSource, +}; use crate::io::print_or_write; use crate::model::load_model; pub fn run(args: EvaluateArgs) -> Result<()> { + let task = args.task.into(); + let source = LabelSource::required( + task, + args.labels.as_deref(), + &args.positive_group_ids, + &args.targets, + )?; let inputs = build_evaluation_inputs( &args.sim_root, &args.sim_type, - &args.labels, - &args.positive_group_ids, - args.task.into(), + &source, + task, &args.dataset_name, - args.flip_threshold, + resolve_flip_threshold(task, &args.sim_type, args.flip_threshold)?, )?; let feature_order = flipper_feature_order_from_inputs(&inputs, &args.sim_type); @@ -25,7 +33,7 @@ pub fn run(args: EvaluateArgs) -> Result<()> { args.positive_class_index, args.null_fill, args.contract.as_deref(), - args.task.into(), + task, )?; let evaluator = Evaluator::new(); diff --git a/cli/src/commands/simulate.rs b/cli/src/commands/simulate.rs index 11ec556..f0ff218 100644 --- a/cli/src/commands/simulate.rs +++ b/cli/src/commands/simulate.rs @@ -8,7 +8,7 @@ use pasteur_hf::layout; use polars::prelude::*; use crate::args::SimulateArgs; -use crate::data::{class_indices, derive_label_matrix}; +use crate::data::{class_indices, LabelSource}; use crate::io::{read_parquet, with_ids, write_variant}; pub fn run(args: SimulateArgs) -> Result<()> { @@ -30,16 +30,7 @@ pub fn run(args: SimulateArgs) -> Result<()> { run_blackout(&args, &features, &source_row_id, &row_ordinal)?; run_jitter_iters(&args, &features, &source_row_id, &row_ordinal)?; - if let Some(labels_path) = &args.labels { - run_flipper( - &args, - &features, - labels_path, - source_row_id.as_materialized_series(), - )?; - } else { - println!(" skipped flipper/ (pass --labels to generate interpolation grid)"); - } + run_flipper(&args, &features, source_row_id.as_materialized_series())?; layout::validate_sim_bundle(&args.output) .context("simulation output failed layout validation")?; @@ -131,29 +122,65 @@ fn run_jitter_iters( Ok(()) } -fn run_flipper( - args: &SimulateArgs, - features: &DataFrame, - labels_path: &std::path::Path, - source_row_id: &Series, -) -> Result<()> { +/// Writes the flipper grid when the flags give it labels to pair across: +/// cohorts for a classifier, targets plus a cutoff for regression. +fn run_flipper(args: &SimulateArgs, features: &DataFrame, source_row_id: &Series) -> Result<()> { let task = TaskType::from(args.task); - let labels = derive_label_matrix(labels_path, &args.positive_group_ids, task, source_row_id)?; + let source = LabelSource::from_args( + task, + args.labels.as_deref(), + &args.positive_group_ids, + &args.targets, + )?; + let cutoff = match (task, args.flip_threshold) { + (TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel, Some(_)) => { + anyhow::bail!( + "simulate reads --flip-threshold only for --task regression; a {task} grid \ + pairs patients across cohorts, and the threshold is set at evaluate time" + ) + } + (TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel, None) => None, + (TaskType::Regression, cutoff) => cutoff, + }; + let Some(source) = source else { + let needs = match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => "--labels", + TaskType::Regression => "--targets, --target-col and --flip-threshold", + }; + println!(" skipped flipper/ (pass {needs} to generate interpolation grid)"); + return Ok(()); + }; + if task == TaskType::Regression && cutoff.is_none() { + println!( + " skipped flipper/ (pass --flip-threshold, the clinical cutoff in target units, \ + to generate interpolation grid)" + ); + return Ok(()); + } + let labels = source.derive(task, source_row_id)?; + // Set by now for regression; classifier grids do not read it. + let flip_threshold = cutoff.unwrap_or(0.5); let config = FlipperConfig { n_pairs: args.flipper_pairs, n_steps: args.flipper_steps, - flip_threshold: 0.5, + flip_threshold, positive_class_label: None, random_state: Some(args.random_state), }; let simulator = FlipperSimulator::new(config); // Binary pairs across the positive cohort; multiclass across two classes // (label_a/label_b hold class indices in --positive-group-id order); - // multilabel across one label per pair (recorded in `pair_label`). + // multilabel across one label per pair (recorded in `pair_label`); + // regression across the cutoff (label_a/label_b hold true targets). let grid = match task { TaskType::Binary => simulator.generate(features, labels[0].as_materialized_series())?, TaskType::Multiclass => simulator.generate(features, &class_indices(&labels)?)?, TaskType::Multilabel => simulator.generate_multilabel(features, &labels)?, + TaskType::Regression => simulator.generate_regression( + features, + labels[0].as_materialized_series(), + flip_threshold, + )?, }; let out = with_flipper_ids(&grid)?; write_variant(&args.output, "flipper", "flipper", &out) diff --git a/cli/src/data/all_agree.rs b/cli/src/data/all_agree.rs index ea39255..83e4418 100644 --- a/cli/src/data/all_agree.rs +++ b/cli/src/data/all_agree.rs @@ -2,8 +2,9 @@ use pasteur_core::TaskType; /// `per_model_preds[model][column][row]`. Each model's decision per row is /// the positive class at a fixed 0.5 for binary (so `--flip-threshold` does -/// not move it), the argmax for multiclass, and the set of labels at or over -/// that model's threshold for multilabel. +/// not move it), the argmax for multiclass, the set of labels at or over +/// that model's threshold for multilabel, and the side of the cutoff +/// (`per_model_thresholds[model][0]`) for regression. pub fn build_all_agree_column( per_model_preds: &[Vec>], per_model_thresholds: &[Vec], @@ -31,6 +32,7 @@ fn decisions( (0..n_rows) .map(|i| match task { TaskType::Binary => vec![preds[0][i] >= 0.5], + TaskType::Regression => vec![preds[0][i] >= thresholds[0]], TaskType::Multiclass => { let best = (0..preds.len()) .max_by(|&a, &b| preds[a][i].total_cmp(&preds[b][i])) @@ -66,6 +68,16 @@ mod tests { assert_eq!(agree, vec![true, false]); } + #[test] + fn regression_agreement_is_the_side_of_the_cutoff() { + // Row 0: 6.0 and 6.4 both under 6.5. Row 1: 6.4 vs 6.6 straddle it. + let a = vec![vec![6.0, 6.4]]; + let b = vec![vec![6.4, 6.6]]; + let t = vec![vec![6.5], vec![6.5]]; + let agree = build_all_agree_column(&[a, b], &t, TaskType::Regression, 2); + assert_eq!(agree, vec![true, false]); + } + #[test] fn multilabel_agreement_uses_each_models_thresholds() { let a = vec![vec![0.6], vec![0.2]]; diff --git a/cli/src/data/eval_inputs.rs b/cli/src/data/eval_inputs.rs index 3fb97e0..495feeb 100644 --- a/cli/src/data/eval_inputs.rs +++ b/cli/src/data/eval_inputs.rs @@ -6,7 +6,7 @@ use pasteur_core::{is_flipper_meta_column, EvaluationConfig, SimulationVariant, use pasteur_hf::layout; use polars::prelude::*; -use crate::data::labels::derive_label_matrix; +use crate::data::source::LabelSource; use crate::io::read_parquet; pub struct EvaluationInputs { @@ -21,14 +21,13 @@ pub struct EvaluationInputs { pub fn build_evaluation_inputs( sim_root: &Path, sim_type: &str, - labels_path: &Path, - group_ids: &[u32], + source: &LabelSource, task: TaskType, dataset_name: &str, flip_threshold: f64, ) -> Result { let (clean_df, source_row_id, feature_order) = load_clean_baseline(sim_root)?; - let labels = derive_label_matrix(labels_path, group_ids, task, &source_row_id)?; + let labels = source.derive(task, &source_row_id)?; let clean = SimulationVariant { x: feature_columns(&clean_df)?, y: labels.clone(), @@ -36,7 +35,7 @@ pub fn build_evaluation_inputs( let variant_files = list_variant_files(sim_root, sim_type)?; let simulation_variants = - load_simulation_variants(&variant_files, sim_type, dataset_name, &labels)?; + load_simulation_variants(&variant_files, sim_type, dataset_name, &labels, task)?; let feature_order = if sim_type == "flipper" { simulation_variants @@ -54,7 +53,11 @@ pub fn build_evaluation_inputs( let eval_config = EvaluationConfig { n_jitter_iters: variant_files.len(), allow_partial: true, - metrics: vec!["roc_auc".to_string()], + metrics: vec![match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => "roc_auc", + TaskType::Regression => "r2", + } + .to_string()], flip_threshold, task, }; @@ -124,21 +127,19 @@ fn load_simulation_variants( sim_type: &str, dataset_name: &str, labels: &DataFrame, + task: TaskType, ) -> Result>> { let mut simulation_variants = HashMap::new(); for path in variant_files { let name = variant_name(path, sim_type); let df = read_parquet(path)?; let (x, y) = if sim_type == "flipper" { - // Grid rows are synthetic patients with no label of their own. + // Grid rows are synthetic patients with no label of their own: + // 0 for classifiers, as before, and no target for regression. let placeholder: Vec = labels .get_column_names() .iter() - .map(|name| { - Int64Chunked::from_vec((*name).clone(), vec![0i64; df.height()]) - .into_series() - .into() - }) + .map(|name| placeholder_label((*name).clone(), task, df.height()).into()) .collect(); let height = df.height(); (df, DataFrame::new(height, placeholder)?) @@ -152,6 +153,15 @@ fn load_simulation_variants( Ok(simulation_variants) } +fn placeholder_label(name: PlSmallStr, task: TaskType, n: usize) -> Series { + match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => { + Int64Chunked::from_vec(name, vec![0i64; n]).into_series() + } + TaskType::Regression => Series::full_null(name, n, &DataType::Float64), + } +} + pub fn model_feature_frame( variant: &SimulationVariant, feature_order: &[String], diff --git a/cli/src/data/labels.rs b/cli/src/data/labels.rs index 9b60471..2522735 100644 --- a/cli/src/data/labels.rs +++ b/cli/src/data/labels.rs @@ -28,7 +28,7 @@ pub fn derive_label_matrix( let row_idx = find_positive_group_row(ids, gid, labels_path)?; let member_ids = load_member_ids(&groups, row_idx)?; let name = match task { - TaskType::Binary => "label".to_string(), + TaskType::Binary | TaskType::Regression => "label".to_string(), TaskType::Multiclass | TaskType::Multilabel => format!("group_{gid}"), }; Ok(labels_for_rows(source_row_ids, &member_ids) @@ -57,6 +57,9 @@ fn check_group_ids(task: TaskType, group_ids: &[u32]) -> Result<()> { "multiclass needs one --positive-group-id per class (at least two), got {n}" ), (TaskType::Multiclass, _) => Ok(()), + (TaskType::Regression, _) => { + anyhow::bail!("regression is scored against --targets, not --positive-group-id cohorts") + } } } diff --git a/cli/src/data/mod.rs b/cli/src/data/mod.rs index c884f9d..2efac00 100644 --- a/cli/src/data/mod.rs +++ b/cli/src/data/mod.rs @@ -2,7 +2,10 @@ mod all_agree; mod eval_inputs; mod labels; mod predictions; +mod source; +mod targets; pub use eval_inputs::{build_evaluation_inputs, flipper_feature_order_from_inputs}; -pub use labels::{class_indices, derive_label_matrix}; +pub use labels::class_indices; pub use predictions::build_predictions_table; +pub use source::{resolve_flip_threshold, LabelSource}; diff --git a/cli/src/data/predictions.rs b/cli/src/data/predictions.rs index 171fcca..68d8c55 100644 --- a/cli/src/data/predictions.rs +++ b/cli/src/data/predictions.rs @@ -12,11 +12,13 @@ pub fn build_predictions_table( ) -> Result { let dataset_key = &inputs.dataset_name; let mut rows = PredictionRows::new(models, inputs.eval_config.task); + // NaN when a regression run was given no cutoff. + let cutoff = Some(inputs.eval_config.flip_threshold).filter(|t| t.is_finite()); rows.append_variant("clean", &inputs.clean_datasets[dataset_key], inputs, models)?; for (name, per_dataset) in &inputs.simulation_variants { rows.append_variant(name, &per_dataset[dataset_key], inputs, models)?; } - Ok(assemble_dataframe(rows, models)) + Ok(assemble_dataframe(rows, models, cutoff)) } struct PredictionRows { @@ -25,8 +27,9 @@ struct PredictionRows { classes: Vec, source_row_id: Vec, variant_name: Vec, - /// `labels[column][row]`. - labels: Vec>, + /// `labels[column][row]`: 0/1 for classifiers, the true target for + /// regression (`None` on synthetic flipper rows). + labels: Vec>>, /// `per_model_preds[model][column][row]`. per_model_preds: Vec>>, } @@ -72,9 +75,9 @@ impl PredictionRows { self.variant_name.push(name.to_string()); } for (out, col) in self.labels.iter_mut().zip(variant.y.columns()) { - let col = col.cast(&DataType::Int64)?; - let col = col.i64()?; - out.extend((0..n).map(|i| col.get(i).unwrap_or(0))); + let col = col.cast(&DataType::Float64)?; + let col = col.f64()?; + out.extend((0..n).map(|i| col.get(i))); } for (preds, (_, model)) in self.per_model_preds.iter_mut().zip(models) { let frame = model.predict_proba(&x)?; @@ -93,11 +96,19 @@ impl PredictionRows { /// Binary keeps its original columns (`label`, one column per model). /// Multi-output models get `label__` and `__` per class. -fn assemble_dataframe(rows: PredictionRows, models: &[(String, OnnxModel)]) -> DataFrame { +/// Regression gets `target` and one prediction column per model. +fn assemble_dataframe( + rows: PredictionRows, + models: &[(String, OnnxModel)], + cutoff: Option, +) -> DataFrame { let n_rows = rows.source_row_id.len(); - let binary = rows.task.is_binary(); + let single_column = match rows.task { + TaskType::Binary | TaskType::Regression => true, + TaskType::Multiclass | TaskType::Multilabel => false, + }; let suffix = |prefix: &str, class: &str| { - if binary { + if single_column { prefix.to_string() } else { format!("{prefix}__{class}") @@ -108,16 +119,25 @@ fn assemble_dataframe(rows: PredictionRows, models: &[(String, OnnxModel)]) -> D Series::new("variant_name".into(), rows.variant_name).into(), ]; for (class, labels) in rows.classes.iter().zip(rows.labels) { - columns.push(Series::new(suffix("label", class).as_str().into(), labels).into()); + columns.push(label_column(rows.task, &suffix("label", class), labels).into()); } let thresholds: Vec> = models .iter() - .map(|(_, m)| { - m.decision_thresholds() - .unwrap_or_else(|| vec![0.5; rows.classes.len()]) + .map(|(_, m)| match rows.task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => m + .decision_thresholds() + .unwrap_or_else(|| vec![0.5; rows.classes.len()]), + TaskType::Regression => cutoff.into_iter().collect(), }) .collect(); - let all_agree = build_all_agree_column(&rows.per_model_preds, &thresholds, rows.task, n_rows); + let all_agree = match (rows.task, cutoff) { + // Without a cutoff, regression predictions have no decision to agree on. + (TaskType::Regression, None) => vec![None; n_rows], + _ => build_all_agree_column(&rows.per_model_preds, &thresholds, rows.task, n_rows) + .into_iter() + .map(Some) + .collect(), + }; for ((model_label, _), preds) in models.iter().zip(rows.per_model_preds) { for (class, col) in rows.classes.iter().zip(preds) { columns.push(Series::new(suffix(model_label, class).as_str().into(), col).into()); @@ -126,3 +146,18 @@ fn assemble_dataframe(rows: PredictionRows, models: &[(String, OnnxModel)]) -> D columns.push(Series::new("all_agree".into(), all_agree).into()); DataFrame::new(n_rows, columns).expect("prediction columns aligned") } + +/// Classifier labels stay `i64` 0/1, as before; regression writes the true +/// target as `target`. +fn label_column(task: TaskType, name: &str, labels: Vec>) -> Series { + match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => { + let labels: Vec = labels + .into_iter() + .map(|v| v.unwrap_or(0.0) as i64) + .collect(); + Series::new(name.into(), labels) + } + TaskType::Regression => Series::new("target".into(), labels), + } +} diff --git a/cli/src/data/source.rs b/cli/src/data/source.rs new file mode 100644 index 0000000..079a60d --- /dev/null +++ b/cli/src/data/source.rs @@ -0,0 +1,120 @@ +use std::path::Path; + +use anyhow::Result; +use pasteur_core::TaskType; +use polars::prelude::*; + +use crate::args::TargetArgs; +use crate::data::labels::derive_label_matrix; +use crate::data::targets::derive_target_column; + +/// Where the true outcome for each row comes from: cohort membership for a +/// classifier, a target table for a regression model. +pub enum LabelSource<'a> { + Cohorts { + path: &'a Path, + group_ids: &'a [u32], + }, + Targets { + path: &'a Path, + id_col: &'a str, + target_col: &'a str, + }, +} + +impl<'a> LabelSource<'a> { + /// The source the flags select for `task`, or `None` when none was given. + /// Flags for the other kind of task are an error rather than ignored, so a + /// regression run never silently scores against cohort labels. + pub fn from_args( + task: TaskType, + labels: Option<&'a Path>, + group_ids: &'a [u32], + targets: &'a TargetArgs, + ) -> Result> { + match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => { + if targets.targets.is_some() || targets.target_col.is_some() { + anyhow::bail!( + "--targets and --target-col are for --task regression; a {task} model \ + is scored against cohorts from --labels" + ); + } + Ok(labels.map(|path| LabelSource::Cohorts { path, group_ids })) + } + TaskType::Regression => { + if labels.is_some() { + anyhow::bail!( + "--labels holds cohorts for classifiers; a regression model is scored \ + against --targets with --target-col" + ); + } + match (&targets.targets, &targets.target_col) { + (Some(path), Some(target_col)) => Ok(Some(LabelSource::Targets { + path, + id_col: &targets.targets_id_col, + target_col, + })), + (Some(_), None) => { + anyhow::bail!("--targets needs --target-col to name the target column") + } + (None, Some(_)) => { + anyhow::bail!("--target-col names a column of --targets; pass --targets") + } + (None, None) => Ok(None), + } + } + } + } + + /// `from_args` for commands that cannot run without labels. + pub fn required( + task: TaskType, + labels: Option<&'a Path>, + group_ids: &'a [u32], + targets: &'a TargetArgs, + ) -> Result { + Self::from_args(task, labels, group_ids, targets)?.ok_or_else(|| match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => { + anyhow::anyhow!("a {task} evaluation needs --labels") + } + TaskType::Regression => { + anyhow::anyhow!("a regression evaluation needs --targets and --target-col") + } + }) + } + + /// The label matrix (classification) or target column (regression) for + /// `source_row_ids`, in row order. + pub fn derive(&self, task: TaskType, source_row_ids: &Series) -> Result { + match self { + LabelSource::Cohorts { path, group_ids } => { + derive_label_matrix(path, group_ids, task, source_row_ids) + } + LabelSource::Targets { + path, + id_col, + target_col, + } => derive_target_column(path, id_col, target_col, source_row_ids), + } + } +} + +/// The `--flip-threshold` to score with: 0.5 by default for a classifier, +/// and for regression the user's cutoff, required when flipper is scored. +/// A regression run that does not score flipper never reads it. +pub fn resolve_flip_threshold(task: TaskType, sim_type: &str, given: Option) -> Result { + match (task, given) { + (TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel, t) => Ok(t.unwrap_or(0.5)), + (TaskType::Regression, Some(t)) if t.is_finite() => Ok(t), + (TaskType::Regression, Some(t)) => { + anyhow::bail!("--flip-threshold must be finite, got {t}") + } + (TaskType::Regression, None) if sim_type == "flipper" => anyhow::bail!( + "regression flipper needs --flip-threshold: the clinical cutoff, in the target's \ + units, that `simulate` sampled pairs across" + ), + // Blackout and jitter scoring never read the cutoff. + (TaskType::Regression, None) => Ok(f64::NAN), + } +} diff --git a/cli/src/data/targets.rs b/cli/src/data/targets.rs new file mode 100644 index 0000000..4ec6a10 --- /dev/null +++ b/cli/src/data/targets.rs @@ -0,0 +1,178 @@ +use std::collections::HashMap; +use std::path::Path; + +use anyhow::{Context, Result}; +use polars::prelude::*; + +use crate::io::read_parquet; + +/// The regression target for `source_row_ids`: one `Float64` column named +/// `target_col`, joined from `targets_path` by ID. IDs are compared as text, +/// so `A00123` and `007` match exactly rather than being integer-coerced. +/// +/// Every cohort row needs exactly one finite target. A missing, duplicated or +/// non-finite target stops the run: scoring a model on whichever rows happen +/// to have a value would quietly change the cohort being evaluated. +pub fn derive_target_column( + targets_path: &Path, + id_col: &str, + target_col: &str, + source_row_ids: &Series, +) -> Result { + let table = read_parquet(targets_path)?; + let column = |name: &str, role: &str| { + table.column(name).with_context(|| { + format!( + "{} has no {role} column {name:?}; its columns are {:?}", + targets_path.display(), + table.get_column_names() + ) + }) + }; + let ids = column(id_col, "ID (--targets-id-col)")?.cast(&DataType::String)?; + let values = column(target_col, "target (--target-col)")? + .cast(&DataType::Float64) + .with_context(|| format!("target column {target_col:?} is not numeric"))?; + let by_id = index_targets(ids.str()?, values.f64()?)?; + + let row_ids = source_row_ids.str()?; + let mut target = Vec::with_capacity(row_ids.len()); + let (mut unmatched, mut not_finite) = (Vec::new(), Vec::new()); + for id in row_ids.iter() { + let id = id.unwrap_or(""); + match by_id.get(id) { + None => unmatched.push(id), + Some(Some(v)) if v.is_finite() => target.push(*v), + Some(_) => not_finite.push(id), + } + } + + let problems: Vec = [ + (unmatched, "have no row in --targets"), + (not_finite, "have a missing or non-finite target"), + ] + .into_iter() + .filter(|(ids, _)| !ids.is_empty()) + .map(|(ids, what)| format!("{} row(s) {what} (e.g. {})", ids.len(), sample(&ids))) + .collect(); + if !problems.is_empty() { + anyhow::bail!( + "targets from {} column {target_col:?} do not cover the cohort: {}", + targets_path.display(), + problems.join("; ") + ); + } + Ok(DataFrame::new( + target.len(), + vec![Series::new(target_col.into(), target).into()], + )?) +} + +/// Target value by ID. An ID listed twice is an error rather than first-wins, +/// since the two rows may disagree. +fn index_targets<'a>( + ids: &'a StringChunked, + values: &Float64Chunked, +) -> Result>> { + let mut by_id = HashMap::with_capacity(ids.len()); + let mut duplicates = Vec::new(); + for (id, value) in ids.iter().zip(values.iter()) { + let Some(id) = id else { continue }; + if by_id.insert(id, value).is_some() { + duplicates.push(id); + } + } + if !duplicates.is_empty() { + duplicates.sort_unstable(); + duplicates.dedup(); + anyhow::bail!( + "--targets lists {} ID(s) more than once (e.g. {}); give each patient one target", + duplicates.len(), + sample(&duplicates) + ); + } + Ok(by_id) +} + +fn sample(ids: &[&str]) -> String { + ids[..ids.len().min(5)].join(", ") +} + +#[cfg(test)] +mod tests { + use super::*; + + fn write_targets(dir: &Path, ids: &[&str], values: &[Option]) -> std::path::PathBuf { + std::fs::create_dir_all(dir).unwrap(); + let path = dir.join("targets.parquet"); + let mut df = df!["node_id" => ids, "hba1c" => values].unwrap(); + let mut file = std::fs::File::create(&path).unwrap(); + ParquetWriter::new(&mut file).finish(&mut df).unwrap(); + path + } + + fn dir(name: &str) -> std::path::PathBuf { + std::env::temp_dir().join(format!("pasteur-targets-{name}-{}", std::process::id())) + } + + fn rows(ids: &[&str]) -> Series { + Series::new("source_row_id".into(), ids) + } + + #[test] + fn targets_are_joined_by_id_in_cohort_order() { + let d = dir("join"); + let path = write_targets(&d, &["b", "A007", "c"], &[Some(7.0), Some(5.5), Some(9.0)]); + let y = derive_target_column(&path, "node_id", "hba1c", &rows(&["A007", "b"])).unwrap(); + assert_eq!(y.get_column_names(), ["hba1c"]); + let v: Vec = y + .column("hba1c") + .unwrap() + .f64() + .unwrap() + .into_no_null_iter() + .collect(); + assert_eq!(v, vec![5.5, 7.0]); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn missing_and_non_finite_targets_are_reported_together() { + let d = dir("missing"); + let path = write_targets(&d, &["1", "2", "3"], &[Some(1.0), None, Some(f64::NAN)]); + let err = derive_target_column(&path, "node_id", "hba1c", &rows(&["1", "2", "3", "4"])) + .unwrap_err() + .to_string(); + assert!( + err.contains("1 row(s) have no row in --targets (e.g. 4)"), + "{err}" + ); + assert!( + err.contains("2 row(s) have a missing or non-finite target (e.g. 2, 3)"), + "{err}" + ); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn duplicate_ids_are_rejected() { + let d = dir("dup"); + let path = write_targets(&d, &["1", "1"], &[Some(1.0), Some(2.0)]); + let err = derive_target_column(&path, "node_id", "hba1c", &rows(&["1"])) + .unwrap_err() + .to_string(); + assert!(err.contains("1 ID(s) more than once (e.g. 1)"), "{err}"); + let _ = std::fs::remove_dir_all(&d); + } + + #[test] + fn a_wrong_column_name_lists_the_real_ones() { + let d = dir("col"); + let path = write_targets(&d, &["1"], &[Some(1.0)]); + let err = derive_target_column(&path, "node_id", "a1c", &rows(&["1"])) + .unwrap_err() + .to_string(); + assert!(err.contains("\"a1c\"") && err.contains("hba1c"), "{err}"); + let _ = std::fs::remove_dir_all(&d); + } +} diff --git a/cli/tests/cli_local.rs b/cli/tests/cli_local.rs index 9d9e0dc..6a2baef 100644 --- a/cli/tests/cli_local.rs +++ b/cli/tests/cli_local.rs @@ -575,3 +575,325 @@ fn binary_takes_exactly_one_group() { assert!(err.contains("--task multiclass"), "got: {err}"); let _ = fs::remove_dir_all(&root); } + +/// Targets for `write_two_feature_input`'s rows, close to the fixtures' +/// `5.5 + f1 − 0.5 f2` so the models are scored on something they fit. +fn write_regression_targets(path: &Path, skip_id: Option<&str>) { + let (ids, values): (Vec, Vec) = (0..12) + .map(|i| { + let f1 = -1.5 + 0.27 * i as f64; + let f2 = ((i * 7) % 12) as f64 / 6.0 - 1.0; + ( + (i + 1).to_string(), + 5.5 + f1 - 0.5 * f2 + 0.05 * (i % 3) as f64, + ) + }) + .filter(|(id, _)| Some(id.as_str()) != skip_id) + .unzip(); + let mut df = df!["node_id" => ids, "hba1c" => values].unwrap(); + let mut file = fs::File::create(path).unwrap(); + ParquetWriter::new(&mut file).finish(&mut df).unwrap(); +} + +/// Simulates the two-feature cohort for regression, with a flipper grid +/// across `cutoff` when one is given. +fn simulate_regression(root: &Path, cutoff: Option<&str>) -> (PathBuf, PathBuf) { + let input = root.join("clean.parquet"); + let targets = root.join("targets.parquet"); + let sim_root = root.join("sim"); + write_two_feature_input(&input); + write_regression_targets(&targets, None); + let mut args = vec![ + "simulate", + "--input", + input.to_str().unwrap(), + "--output", + sim_root.to_str().unwrap(), + "--feature", + "f1", + "--task", + "regression", + "--targets", + targets.to_str().unwrap(), + "--target-col", + "hba1c", + "--jitter-iters", + "2", + "--flipper-pairs", + "6", + "--flipper-steps", + "5", + ]; + if let Some(cutoff) = cutoff { + args.extend(["--flip-threshold", cutoff]); + } + let out = run_cli(&args); + assert!( + out.status.success(), + "simulate failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + (sim_root, targets) +} + +fn regression_args<'a>( + command: &'a str, + sim_type: &'a str, + sim_root: &'a Path, + targets: &'a Path, +) -> Vec<&'a str> { + vec![ + command, + sim_type, + "--sim-root", + sim_root.to_str().unwrap(), + "--targets", + targets.to_str().unwrap(), + "--target-col", + "hba1c", + "--task", + "regression", + ] +} + +#[test] +fn regression_simulate_evaluate_and_compare() { + let root = temp_dir("regression"); + let (sim_root, targets) = simulate_regression(&root, Some("5.5")); + let flipper = + ParquetReader::new(fs::File::open(sim_root.join("flipper/flipper.parquet")).unwrap()) + .finish() + .unwrap(); + let label_a = flipper.column("label_a").unwrap().f64().unwrap(); + let label_b = flipper.column("label_b").unwrap().f64().unwrap(); + assert!((0..flipper.height()) + .all(|i| label_a.get(i).unwrap() < 5.5 && label_b.get(i).unwrap() >= 5.5)); + + let linear = fixture_model( + &root.join("linear"), + "regression_linear", + "regression", + &["hba1c"], + ); + let forest = fixture_model( + &root.join("forest"), + "regression_forest", + "regression", + &["hba1c"], + ); + + for sim_type in ["blackout", "jitter", "flipper"] { + let mut args = regression_args("evaluate", sim_type, &sim_root, &targets); + args.extend([ + "--model", + linear.to_str().unwrap(), + "--flip-threshold", + "5.5", + ]); + let out = run_cli(&args); + assert!( + out.status.success(), + "evaluate {sim_type} failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + let result: serde_json::Value = serde_json::from_slice(&out.stdout).unwrap(); + assert!( + result["baselines"]["rmse"].as_f64().unwrap() < 0.2, + "{result}" + ); + assert!( + result["baselines"]["r2"].as_f64().unwrap() > 0.9, + "{result}" + ); + assert!(result["evaluations"]["metric_based"]["r2"]["resiliency"].is_number()); + assert!(result.get("multi").is_none()); + let detail = &result["regression"]; + assert_eq!(detail["target"], "hba1c"); + assert_eq!(detail["n_rows"], 12); + match sim_type { + "blackout" => { + assert!(detail["blackout_rmse"].is_number()); + assert!(detail["blackout_r2"].is_number()); + } + "jitter" => assert!(detail["jitter_prediction_sd"].is_number()), + _ => { + assert_eq!(detail["flip_threshold"], 5.5); + assert!(result["evaluations"]["metric_invariant"]["flipper_stability"].is_number()); + } + } + } + + let predictions = root.join("predictions.parquet"); + let comparison = root.join("compare.json"); + let mut args = regression_args("compare", "jitter", &sim_root, &targets); + args.extend([ + "--model", + linear.to_str().unwrap(), + "--model", + forest.to_str().unwrap(), + "--flip-threshold", + "5.5", + "--predictions-out", + predictions.to_str().unwrap(), + "--output", + comparison.to_str().unwrap(), + ]); + let out = run_cli(&args); + assert!( + out.status.success(), + "{}", + String::from_utf8_lossy(&out.stderr) + ); + let result: serde_json::Value = + serde_json::from_slice(&fs::read(&comparison).unwrap()).unwrap(); + let df = ParquetReader::new(fs::File::open(&predictions).unwrap()) + .finish() + .unwrap(); + assert_eq!( + df.get_column_names(), + [ + "source_row_id", + "variant_name", + "target", + "regression_linear", + "regression_forest", + "all_agree" + ] + ); + // The JSON's clean RMSE is the RMSE of the clean rows in the parquet. + let clean = df + .clone() + .lazy() + .filter(col("variant_name").eq(lit("clean"))) + .collect() + .unwrap(); + let target = clean.column("target").unwrap().f64().unwrap(); + let pred = clean.column("regression_linear").unwrap().f64().unwrap(); + let mse: f64 = (0..clean.height()) + .map(|i| (pred.get(i).unwrap() - target.get(i).unwrap()).powi(2)) + .sum::() + / clean.height() as f64; + let reported = result["evaluations"][0]["evaluation"]["baselines"]["rmse"] + .as_f64() + .unwrap(); + assert!( + (mse.sqrt() - reported).abs() < 1e-9, + "{} vs {reported}", + mse.sqrt() + ); + assert_eq!(df.column("all_agree").unwrap().null_count(), 0); + let _ = fs::remove_dir_all(&root); +} + +#[test] +fn regression_flipper_needs_the_cutoff_it_was_simulated_with() { + let root = temp_dir("regression-cutoff"); + let (sim_root, targets) = simulate_regression(&root, Some("5.5")); + let model = fixture_model( + &root.join("m"), + "regression_linear", + "regression", + &["hba1c"], + ); + + let mut args = regression_args("evaluate", "flipper", &sim_root, &targets); + args.extend(["--model", model.to_str().unwrap()]); + let out = run_cli(&args); + assert!(!out.status.success()); + let err = String::from_utf8_lossy(&out.stderr); + assert!(err.contains("needs --flip-threshold"), "got: {err}"); + + args.extend(["--flip-threshold", "9"]); + let out = run_cli(&args); + assert!(!out.status.success()); + let err = String::from_utf8_lossy(&out.stderr); + assert!(err.contains("another cutoff or task"), "got: {err}"); + let _ = fs::remove_dir_all(&root); +} + +#[test] +fn regression_without_a_cutoff_skips_flipper() { + let root = temp_dir("regression-no-cutoff"); + let (sim_root, _) = simulate_regression(&root, None); + assert!(sim_root.join("jitter").is_dir()); + assert!(!sim_root.join("flipper").exists()); + let _ = fs::remove_dir_all(&root); +} + +#[test] +fn regression_rejects_cohort_labels_and_classifier_models() { + let root = temp_dir("regression-flags"); + let (sim_root, targets) = simulate_regression(&root, None); + let labels = root.join("groups.parquet"); + write_groups(&labels, &[(103, &[1, 2])]); + let model = fixture_model( + &root.join("m"), + "regression_linear", + "regression", + &["hba1c"], + ); + + let mut args = regression_args("evaluate", "jitter", &sim_root, &targets); + args.extend([ + "--model", + model.to_str().unwrap(), + "--labels", + labels.to_str().unwrap(), + ]); + let err = String::from_utf8_lossy(&run_cli(&args).stderr).to_string(); + assert!( + err.contains("--labels holds cohorts for classifiers"), + "got: {err}" + ); + + // A binary model scored as regression is refused at load. + let binary = fixture_model(&root.join("b"), "binary_zipmap", "binary", &["0", "1"]); + let mut args = regression_args("evaluate", "jitter", &sim_root, &targets); + args.extend(["--model", binary.to_str().unwrap()]); + let err = String::from_utf8_lossy(&run_cli(&args).stderr).to_string(); + assert!(err.contains("is a binary model"), "got: {err}"); + + // And --targets is refused for a classifier. + let out = run_cli(&[ + "evaluate", + "jitter", + "--sim-root", + sim_root.to_str().unwrap(), + "--labels", + labels.to_str().unwrap(), + "--targets", + targets.to_str().unwrap(), + "--model", + binary.to_str().unwrap(), + ]); + let err = String::from_utf8_lossy(&out.stderr); + assert!( + err.contains("--targets and --target-col are for --task regression"), + "got: {err}" + ); + let _ = fs::remove_dir_all(&root); +} + +#[test] +fn regression_targets_must_cover_the_cohort() { + let root = temp_dir("regression-missing"); + let (sim_root, _) = simulate_regression(&root, None); + let partial = root.join("partial.parquet"); + write_regression_targets(&partial, Some("7")); + let model = fixture_model( + &root.join("m"), + "regression_linear", + "regression", + &["hba1c"], + ); + let mut args = regression_args("evaluate", "jitter", &sim_root, &partial); + args.extend(["--model", model.to_str().unwrap()]); + let out = run_cli(&args); + assert!(!out.status.success()); + let err = String::from_utf8_lossy(&out.stderr); + assert!( + err.contains("1 row(s) have no row in --targets (e.g. 7)"), + "got: {err}" + ); + let _ = fs::remove_dir_all(&root); +} diff --git a/core/src/evaluate/flipper.rs b/core/src/evaluate/flipper.rs index 8766fad..8369b07 100644 --- a/core/src/evaluate/flipper.rs +++ b/core/src/evaluate/flipper.rs @@ -162,10 +162,13 @@ pub struct FlipperScores { /// - multiclass: the pairwise margin `p[class_b] − p[class_a]`, since the /// decision is an argmax and no single threshold applies; /// - multilabel: `p[k] − threshold[k]` for the label `k` the pair was sampled -/// for, read from the grid's `pair_label` column. +/// for, read from the grid's `pair_label` column; +/// - regression: `prediction − cutoff`, where the cutoff is `thresholds[0]` +/// in the target's units. Every pair must have one endpoint below the +/// cutoff and one at or above it. /// -/// `thresholds` holds one entry per output column for binary and multilabel -/// and is ignored for multiclass. +/// `thresholds` holds one entry per output column for binary, multilabel and +/// regression, and is ignored for multiclass. pub fn score_flipper( model: &dyn Model, grid: &DataFrame, @@ -191,7 +194,10 @@ pub fn score_flipper( // Each pair's labels and, per grid row, a score that crosses 0 exactly // where the decision changes. let (pair_labels, crossing): (Vec>, Vec) = match task { - TaskType::Binary => { + TaskType::Binary | TaskType::Regression => { + if task == TaskType::Regression { + check_pairs_straddle(grid, thresholds[0])?; + } let preds = Series::new("pred".into(), probs.swap_remove(0)); let results = find_flip_points(grid, &preds, thresholds[0])?; let stability = calculate_flipper_stability(&results); @@ -250,6 +256,27 @@ pub fn score_flipper( }) } +/// A regression grid pairs a patient below the cutoff with one at or above +/// it. A pair that does not straddle `cutoff` came from a grid generated for +/// another cutoff or another task, and its crossing would mean nothing. +fn check_pairs_straddle(grid: &DataFrame, cutoff: f64) -> Result<(), CoreError> { + let label_a = series_to_f64(grid.column(LABEL_A_COL)?.as_materialized_series())?; + let label_b = series_to_f64(grid.column(LABEL_B_COL)?.as_materialized_series())?; + let bad = label_a + .iter() + .zip(&label_b) + .filter(|(a, b)| !(**a < cutoff && **b >= cutoff)) + .count(); + if bad == 0 { + return Ok(()); + } + Err(CoreError::InvalidConfig(format!( + "{bad} flipper grid row(s) do not pair a target below {cutoff} with one at or above \ + it; the grid was generated for another cutoff or task. Regenerate it with \ + `simulate --task regression --flip-threshold {cutoff}`" + ))) +} + fn never_flipped(results: &[FlipperPairResult]) -> f64 { if results.is_empty() { return 0.0; diff --git a/core/src/evaluate/mod.rs b/core/src/evaluate/mod.rs index 1c2ee62..8291a03 100644 --- a/core/src/evaluate/mod.rs +++ b/core/src/evaluate/mod.rs @@ -1,6 +1,7 @@ pub mod flipper; mod metrics; mod multi; +mod regression; pub use flipper::{ calculate_flipper_stability, find_flip_points, flipper_model_features, score_flipper, @@ -14,7 +15,7 @@ pub use metrics::{ use crate::error::CoreError; use crate::schema::{ EvaluationConfig, EvaluationResult, EvaluationScores, MetricBasedScores, MetricInvariantScores, - MultiOutputScores, SimulationVariant, + MultiOutputScores, RegressionScores, SimulationVariant, TaskType, }; use chrono::Utc; use metrics::{jitter_stability_from_variance, series_to_f64}; @@ -51,7 +52,7 @@ pub(crate) fn predict_matrix( } /// `predict_matrix`, plus the model's output column names. -fn predict_named( +pub(crate) fn predict_named( model: &dyn Model, x: &DataFrame, n_outputs: usize, @@ -124,29 +125,36 @@ impl Evaluator { } // Each headline number is the mean over datasets where it is defined. - let mean = |f: fn(&SingleDatasetScores) -> Option| { + let mean = |f: &dyn Fn(&SingleDatasetScores) -> Option| { mean_some(per_dataset_scores.iter().map(f)) }; - let jitter_stability = mean(|s| Some(s.jitter_stability)).unwrap_or(1.0); - let flipper_stability = mean(|s| s.flipper_stability); - let resiliency = mean(|s| Some(s.resiliency)).unwrap_or(1.0); - let baseline = mean(|s| Some(s.baseline)).unwrap_or(0.0); - let multi = average_multi( - per_dataset_scores - .into_iter() - .filter_map(|s| s.multi) - .collect(), - ); - - // Only ROC AUC is computed so far; generalizability is a placeholder. + let jitter_stability = mean(&|s| Some(s.jitter_stability)).unwrap_or(1.0); + let flipper_stability = mean(&|s| s.flipper_stability); + let resiliency = mean(&|s| Some(s.resiliency)).unwrap_or(1.0); + let (metric, keys) = baseline_keys(config.task); + let baselines = keys + .iter() + .enumerate() + .map(|(j, key)| { + let value = mean(&|s| s.baselines.get(j).copied()).unwrap_or(0.0); + (key.to_string(), value) + }) + .collect(); + let (multi, regression): (Vec<_>, Vec<_>) = per_dataset_scores + .into_iter() + .map(|s| (s.multi, s.regression)) + .unzip(); + let multi = average_multi(multi.into_iter().flatten().collect()); + let regression = regression::average_regression(regression.into_iter().flatten().collect()); + + // Generalizability is a placeholder. let metric_based = HashMap::from([( - "roc_auc".to_string(), + metric.to_string(), MetricBasedScores { resiliency, generalizability: 1.0, }, )]); - let baselines = HashMap::from([("roc_auc".to_string(), baseline)]); Ok(EvaluationResult { baselines, @@ -159,6 +167,7 @@ impl Evaluator { }, created_at: Utc::now(), multi, + regression, }) } @@ -168,6 +177,36 @@ impl Evaluator { clean_dataset: &SimulationVariant, simulated_variants: &HashMap, config: &EvaluationConfig, + ) -> Result { + match config.task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => { + self.score_classifier(model, clean_dataset, simulated_variants, config) + } + TaskType::Regression => { + let scored = regression::score_dataset( + model, + clean_dataset, + simulated_variants, + config.flip_threshold, + )?; + Ok(SingleDatasetScores { + baselines: scored.baselines.to_vec(), + resiliency: scored.resiliency.unwrap_or(1.0), + jitter_stability: jitter_stability_from_variance(scored.jitter_variance), + flipper_stability: scored.flipper.map(|f| f.stability), + multi: None, + regression: Some(scored.detail), + }) + } + } + } + + fn score_classifier( + &self, + model: &dyn Model, + clean_dataset: &SimulationVariant, + simulated_variants: &HashMap, + config: &EvaluationConfig, ) -> Result { let task = config.task; let y = label_matrix(&clean_dataset.y, task)?; @@ -215,7 +254,7 @@ impl Evaluator { }); Ok(SingleDatasetScores { - baseline, + baselines: vec![baseline], resiliency: mean_some(resiliency.iter().copied()).unwrap_or(1.0), jitter_stability: jitter_stability_from_variance(headline_jitter_variance( task, @@ -223,14 +262,26 @@ impl Evaluator { )), flipper_stability: flipper.map(|f| f.stability), multi, + regression: None, }) } } +/// The metric that `metric_based` is keyed by, and the `baselines` keys, in +/// the order `SingleDatasetScores::baselines` holds them. +fn baseline_keys(task: TaskType) -> (&'static str, &'static [&'static str]) { + match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => ("roc_auc", &["roc_auc"]), + TaskType::Regression => ("r2", ®ression::BASELINE_KEYS), + } +} + struct SingleDatasetScores { - baseline: f64, + /// One value per `baseline_keys(task)` entry. + baselines: Vec, resiliency: f64, jitter_stability: f64, flipper_stability: Option, multi: Option, + regression: Option, } diff --git a/core/src/evaluate/multi.rs b/core/src/evaluate/multi.rs index da1ad9f..46285c8 100644 --- a/core/src/evaluate/multi.rs +++ b/core/src/evaluate/multi.rs @@ -33,6 +33,9 @@ pub(crate) fn label_matrix(y: &DataFrame, task: TaskType) -> Result {} + (TaskType::Regression, _) => { + return invalid("regression targets are scored by `regression::target_vector`".into()) + } } for i in 0..y.height() { let row: Vec = cols.iter().map(|c| c[i]).collect(); @@ -66,7 +69,8 @@ pub(crate) fn decision_thresholds( t.len() ))), (TaskType::Multilabel, Some(t)) => Ok(t), - (TaskType::Multilabel, None) | (TaskType::Binary | TaskType::Multiclass, _) => { + (TaskType::Multilabel, None) + | (TaskType::Binary | TaskType::Multiclass | TaskType::Regression, _) => { Ok(vec![flip_threshold; k]) } } @@ -98,6 +102,8 @@ pub(crate) fn label_resiliency( _ => None, }) .collect(), + // Regression resiliency is an RMSE ratio, computed in `regression`. + TaskType::Regression => vec![None; clean.len()], } } @@ -113,7 +119,10 @@ pub(crate) fn label_resiliency( /// times more stable for every extra class. pub(crate) fn headline_jitter_variance(task: TaskType, per_column: &[Option]) -> Option { match task { - TaskType::Binary | TaskType::Multilabel => mean_some(per_column.iter().copied()), + // Regression scores its one column in `regression`; it never gets here. + TaskType::Binary | TaskType::Multilabel | TaskType::Regression => { + mean_some(per_column.iter().copied()) + } TaskType::Multiclass => per_column .iter() .copied() @@ -141,7 +150,7 @@ pub(crate) fn jitter_variances(draws: &[Columns], k: usize) -> Result Option { let decide = match task { - TaskType::Binary => return None, + TaskType::Binary | TaskType::Regression => return None, TaskType::Multiclass => argmax_decision, TaskType::Multilabel => threshold_decision, }; diff --git a/core/src/evaluate/regression.rs b/core/src/evaluate/regression.rs new file mode 100644 index 0000000..c536551 --- /dev/null +++ b/core/src/evaluate/regression.rs @@ -0,0 +1,280 @@ +//! Scoring for regression models: one predicted value per patient, compared +//! against a continuous target. The headline fields keep their classifier +//! meaning where one carries over (resiliency is "how much of the model's +//! skill is kept under blackout"), and the `regression` section reports the +//! same numbers in the target's units. + +use super::flipper::{score_flipper, FlipperScores}; +use super::metrics::series_to_f64; +use super::multi::mean_some; +use super::{predict_matrix, predict_named, Model}; +use crate::error::CoreError; +use crate::schema::{RegressionScores, SimulationVariant, TaskType}; +use polars::prelude::*; +use std::collections::HashMap; + +/// The keys of `baselines` for a regression model, in report order. +pub(crate) const BASELINE_KEYS: [&str; 3] = ["rmse", "mae", "r2"]; + +/// Everything one dataset contributes to a regression evaluation. +pub(crate) struct RegressionDataset { + /// `(rmse, mae, r2)` on the clean cohort. + pub baselines: [f64; 3], + pub resiliency: Option, + /// Mean per-patient prediction variance across jitter draws, divided by + /// the target's variance. + pub jitter_variance: Option, + pub flipper: Option, + pub detail: RegressionScores, +} + +pub(crate) fn score_dataset( + model: &dyn Model, + clean: &SimulationVariant, + variants: &HashMap, + flip_threshold: f64, +) -> Result { + let (target_name, y) = target_vector(&clean.y)?; + let target_variance = variance(&y); + if target_variance == 0.0 { + return Err(CoreError::InvalidConfig(format!( + "target {target_name:?} has the same value for every clean row; a constant target \ + cannot tell a good regression model from a bad one" + ))); + } + let (_, mut clean_pred) = predict_named(model, &clean.x, 1)?; + let clean_pred = clean_pred.swap_remove(0); + let clean_rmse = rmse(&y, &clean_pred); + let clean_r2 = r2(&y, &clean_pred, target_variance); + + let blackout_rmse = variants + .get("blackout") + .map(|blackout| -> Result { + Ok(rmse(&y, &predict_column(model, &blackout.x)?)) + }) + .transpose()?; + let blackout_r2 = blackout_rmse.map(|stressed| 1.0 - stressed.powi(2) / target_variance); + + let jitter_draws: Vec> = variants + .iter() + .filter(|(name, _)| name.starts_with("jitter")) + .map(|(_, variant)| predict_column(model, &variant.x)) + .collect::>()?; + let jitter = jitter_spread(&jitter_draws); + + let flipper = variants + .get("flipper") + .map(|flipper| { + score_flipper( + model, + &flipper.x, + TaskType::Regression, + 1, + &[flip_threshold], + ) + }) + .transpose()?; + + Ok(RegressionDataset { + baselines: [clean_rmse, mae(&y, &clean_pred), clean_r2], + resiliency: blackout_r2.and_then(|stressed| r2_resiliency(clean_r2, stressed)), + jitter_variance: jitter.map(|j| j.variance / target_variance), + detail: RegressionScores { + target: target_name, + n_rows: y.len(), + target_sd: target_variance.sqrt(), + blackout_rmse, + blackout_r2, + jitter_prediction_sd: jitter.map(|j| j.sd), + flip_threshold: flipper.as_ref().map(|_| flip_threshold), + flipper_never_flipped: flipper.as_ref().map(|f| f.never_flipped), + }, + flipper, + }) +} + +/// The single target column, which must be finite on every row: a missing +/// target has no error to measure, and dropping the row silently would score +/// the model on a different cohort than the one asked for. +pub(crate) fn target_vector(y: &DataFrame) -> Result<(String, Vec), CoreError> { + let [column] = y.columns() else { + return Err(CoreError::InvalidConfig(format!( + "regression evaluation takes one target column, got {}", + y.width() + ))); + }; + let values = series_to_f64(column.as_materialized_series())?; + let missing = values.iter().filter(|v| !v.is_finite()).count(); + if missing > 0 { + return Err(CoreError::InvalidConfig(format!( + "target {:?} is missing or non-finite on {missing} of {} rows", + column.name(), + values.len() + ))); + } + Ok((column.name().to_string(), values)) +} + +fn predict_column(model: &dyn Model, x: &DataFrame) -> Result, CoreError> { + Ok(predict_matrix(model, x, 1)?.swap_remove(0)) +} + +/// Root-mean-square error. A non-finite prediction makes it NaN rather than +/// being skipped: the model failed on that patient, and that is an error too. +pub(crate) fn rmse(y: &[f64], pred: &[f64]) -> f64 { + mean(y.iter().zip(pred).map(|(y, p)| (p - y).powi(2))).sqrt() +} + +pub(crate) fn mae(y: &[f64], pred: &[f64]) -> f64 { + mean(y.iter().zip(pred).map(|(y, p)| (p - y).abs())) +} + +/// Coefficient of determination, `1 − SSE / SST`. 0 means no better than +/// predicting the mean; it is negative for a model worse than that. +pub(crate) fn r2(y: &[f64], pred: &[f64], target_variance: f64) -> f64 { + let mse = mean(y.iter().zip(pred).map(|(y, p)| (p - y).powi(2))); + 1.0 - mse / target_variance +} + +/// Blackout R² ÷ clean R²: the share of the model's skill over predicting +/// the cohort mean that survives blackout. 1.0 means blackout cost nothing, +/// 0 means the model fell back to no better than the mean, and it is +/// negative when blackout made the model worse than the mean. It exceeds 1.0 +/// when blackout happens to help. +/// +/// A ratio of errors (clean RMSE ÷ blackout RMSE) would reward a model for +/// having little accuracy to lose: its floor is `sqrt(1 − R²)`, so a weak +/// model that collapses below the mean can outscore a strong one that keeps +/// most of its skill. +/// +/// `None` when the clean R² is not positive: a model with no skill has none +/// to keep, just as a classifier label at chance has no resiliency. +pub(crate) fn r2_resiliency(clean: f64, stressed: f64) -> Option { + (clean > 0.0).then(|| stressed / clean) +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct JitterSpread { + /// Mean per-patient variance of the prediction across draws. + pub variance: f64, + /// Mean per-patient standard deviation, in target units. + pub sd: f64, +} + +/// Spread of each patient's prediction across jitter draws, averaged over +/// patients. `None` with fewer than two draws. +pub(crate) fn jitter_spread(draws: &[Vec]) -> Option { + let n_rows = draws.first()?.len(); + if draws.len() < 2 || n_rows == 0 { + return None; + } + let per_row: Vec = (0..n_rows) + .map(|i| variance(&draws.iter().map(|d| d[i]).collect::>())) + .collect(); + Some(JitterSpread { + variance: mean(per_row.iter().copied()), + sd: mean(per_row.iter().map(|v| v.sqrt())), + }) +} + +/// Population variance. +fn variance(xs: &[f64]) -> f64 { + let m = mean(xs.iter().copied()); + mean(xs.iter().map(|x| (x - m).powi(2))) +} + +fn mean(xs: impl IntoIterator) -> f64 { + let (sum, n) = xs + .into_iter() + .fold((0.0, 0usize), |(sum, n), x| (sum + x, n + 1)); + if n == 0 { + f64::NAN + } else { + sum / n as f64 + } +} + +/// Averages regression detail across datasets field by field, as the +/// headline numbers are averaged. +pub(crate) fn average_regression(scores: Vec) -> Option { + if scores.len() <= 1 { + return scores.into_iter().next(); + } + let avg = |f: &dyn Fn(&RegressionScores) -> Option| mean_some(scores.iter().map(f)); + let first = &scores[0]; + Some(RegressionScores { + target: first.target.clone(), + n_rows: scores.iter().map(|s| s.n_rows).sum(), + target_sd: avg(&|s| Some(s.target_sd)).unwrap_or(0.0), + blackout_rmse: avg(&|s| s.blackout_rmse), + blackout_r2: avg(&|s| s.blackout_r2), + jitter_prediction_sd: avg(&|s| s.jitter_prediction_sd), + flip_threshold: first.flip_threshold, + flipper_never_flipped: avg(&|s| s.flipper_never_flipped), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn close(a: f64, b: f64) -> bool { + (a - b).abs() < 1e-12 + } + + #[test] + fn error_metrics_match_hand_computation() { + let y = [1.0, 2.0, 3.0, 4.0]; + let p = [1.5, 2.0, 2.0, 5.0]; + // Errors 0.5, 0, -1, 1: squared 0.25, 0, 1, 1 → MSE 0.5625. + assert!(close(rmse(&y, &p), 0.5625f64.sqrt())); + assert!(close(mae(&y, &p), 0.625)); + // Var(y) = 1.25, so R² = 1 − 0.5625 / 1.25. + assert!(close(r2(&y, &p, variance(&y)), 0.55)); + } + + #[test] + fn resiliency_is_the_share_of_r2_kept() { + assert!(close(r2_resiliency(0.8, 0.4).unwrap(), 0.5)); + assert!(close(r2_resiliency(0.4, 0.8).unwrap(), 2.0)); + // Worse than predicting the mean under blackout: negative. + assert!(r2_resiliency(0.3, -0.1).unwrap() < 0.0); + // No skill to keep. + assert_eq!(r2_resiliency(0.0, -0.5), None); + assert_eq!(r2_resiliency(-0.2, -0.2), None); + } + + #[test] + fn a_weak_model_that_collapses_does_not_outscore_a_strong_one() { + // Target variance 1. Strong: R² 0.66 → 0.27. Weak: R² 0.26 → −0.09. + // The RMSE ratio ranks the weak model as more resilient (0.82 vs + // 0.69); the R² ratio does not. + let rmse_of = |r2: f64| (1.0 - r2).sqrt(); + let strong = (0.66, 0.27); + let weak = (0.26, -0.09); + assert!(rmse_of(weak.0) / rmse_of(weak.1) > rmse_of(strong.0) / rmse_of(strong.1)); + assert!( + r2_resiliency(weak.0, weak.1).unwrap() < r2_resiliency(strong.0, strong.1).unwrap() + ); + } + + #[test] + fn jitter_spread_needs_two_draws() { + assert_eq!(jitter_spread(&[vec![1.0, 2.0]]), None); + // Row 0 moves by ±1 (variance 1, sd 1); row 1 does not move. + let spread = jitter_spread(&[vec![0.0, 5.0], vec![2.0, 5.0]]).unwrap(); + assert!(close(spread.variance, 0.5)); + assert!(close(spread.sd, 0.5)); + } + + #[test] + fn target_must_be_one_finite_column() { + let two = df!["a" => [1.0], "b" => [2.0]].unwrap(); + assert!(target_vector(&two).is_err()); + let nan = df!["a" => [1.0, f64::NAN]].unwrap(); + let err = target_vector(&nan).unwrap_err().to_string(); + assert!(err.contains("1 of 2 rows"), "{err}"); + let null = df!["a" => [Some(1.0), None]].unwrap(); + assert!(target_vector(&null).is_err()); + } +} diff --git a/core/src/schema.rs b/core/src/schema.rs index 0720c15..93a02a0 100644 --- a/core/src/schema.rs +++ b/core/src/schema.rs @@ -55,10 +55,11 @@ pub struct SimulationConfig { pub flipper: Option, } -/// What a model's probability output means. Binary models emit one -/// positive-class column; multiclass models emit K mutually exclusive class -/// probabilities (one row sums to 1); multilabel models emit K independent -/// per-label probabilities. +/// What a model's output means. Binary models emit one positive-class +/// column; multiclass models emit K mutually exclusive class probabilities +/// (one row sums to 1); multilabel models emit K independent per-label +/// probabilities; regression models emit one predicted value, in the +/// target's units. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum TaskType { @@ -66,6 +67,7 @@ pub enum TaskType { Binary, Multiclass, Multilabel, + Regression, } impl TaskType { @@ -80,13 +82,15 @@ impl std::fmt::Display for TaskType { TaskType::Binary => "binary", TaskType::Multiclass => "multiclass", TaskType::Multilabel => "multilabel", + TaskType::Regression => "regression", }) } } /// `y` is a label matrix with one 0/1 column per model output column, in the /// model's class order: one column for binary, a one-hot row for multiclass, -/// a multi-hot row for multilabel. +/// a multi-hot row for multilabel. For regression it is one column holding +/// the true target value. #[derive(Debug, Clone)] pub struct SimulationVariant { pub x: DataFrame, @@ -103,6 +107,7 @@ pub struct EvaluationConfig { /// Probability threshold for flipper stability (decision-boundary crossing). /// For multilabel models without per-label `output.thresholds`, applies /// to every label. Unused for multiclass, whose decision is the argmax. + /// For regression it is a clinical cutoff in the target's units. pub flip_threshold: f64, #[serde(default)] pub task: TaskType, @@ -136,6 +141,31 @@ pub struct EvaluationResult { /// Absent for binary models, so binary output is unchanged. #[serde(default, skip_serializing_if = "Option::is_none")] pub multi: Option, + /// Regression detail, in the target's units. Absent for classifiers. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub regression: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegressionScores { + /// Name of the target column. + pub target: String, + pub n_rows: usize, + /// Standard deviation of the true target on the clean cohort. Jitter + /// stability is scaled by its square, so it reads the same in any unit. + pub target_sd: f64, + /// RMSE on the blackout variant. + pub blackout_rmse: Option, + /// R² on the blackout variant; `metric_based.r2.resiliency` is this + /// divided by `baselines.r2`. + pub blackout_r2: Option, + /// Mean per-patient standard deviation of the prediction across jitter + /// draws, in target units. `None` with fewer than two draws. + pub jitter_prediction_sd: Option, + /// The cutoff flipper pairs were scored against, when flipper ran. + pub flip_threshold: Option, + /// Share of flipper pairs whose prediction never crossed the cutoff. + pub flipper_never_flipped: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/core/src/simulate/flipper.rs b/core/src/simulate/flipper.rs index 41a1c1c..209614e 100644 --- a/core/src/simulate/flipper.rs +++ b/core/src/simulate/flipper.rs @@ -159,6 +159,68 @@ impl FlipperSimulator { self.interpolate(x, &pairs) } + /// Regression counterpart of `generate`. A continuous target has no + /// classes to pair across, so `cutoff` (a clinical threshold in the + /// target's units) supplies them: each pair is a row `a` whose target is + /// below `cutoff` and a row `b` whose target is at or above it. + /// `label_a`/`label_b` hold the two true target values, so the grid + /// records which cutoff it was built for. Rows with a missing or + /// non-finite target are never sampled. + pub fn generate_regression( + &self, + x: &DataFrame, + y: &Series, + cutoff: f64, + ) -> Result { + if x.height() != y.len() { + return Err(CoreError::InvalidConfig(format!( + "x has {} rows but y has {} — must match", + x.height(), + y.len() + ))); + } + if !cutoff.is_finite() { + return Err(CoreError::InvalidConfig(format!( + "flipper cutoff must be a finite number, got {cutoff}" + ))); + } + let y = y.cast(&DataType::Float64)?; + let targets: Vec<(usize, f64)> = y + .f64()? + .iter() + .enumerate() + .filter_map(|(row, v)| Some((row, v.filter(|v| v.is_finite())?))) + .collect(); + let (below, at_or_above): (Vec<_>, Vec<_>) = + targets.into_iter().partition(|(_, v)| *v < cutoff); + if below.is_empty() || at_or_above.is_empty() { + return Err(CoreError::InvalidConfig(format!( + "FlipperSimulator needs targets on both sides of the cutoff {cutoff}: \ + {} row(s) below it, {} at or above it", + below.len(), + at_or_above.len() + ))); + } + + let mut rng = self.rng(); + let pairs: Vec = (0..self.config.n_pairs) + .map(|_| { + let &(row_a, label_a) = below.choose(&mut rng).expect("below is non-empty"); + let &(row_b, label_b) = at_or_above + .choose(&mut rng) + .expect("at_or_above is non-empty"); + Pair { + row_a, + row_b, + label_a, + label_b, + pair_label: None, + } + }) + .collect(); + self.interpolate(x, &pairs) + } + fn rng(&self) -> StdRng { match self.config.random_state { Some(seed) => StdRng::seed_from_u64(seed), @@ -404,6 +466,53 @@ mod tests { Ok(()) } + fn regression_config(seed: u64) -> FlipperConfig { + FlipperConfig { + n_pairs: 30, + n_steps: 4, + flip_threshold: 6.5, + positive_class_label: None, + random_state: Some(seed), + } + } + + #[test] + fn regression_pairs_straddle_the_cutoff() -> Result<(), CoreError> { + let (x, _) = sample_df(); + let y = Series::new("hba1c".into(), &[5.0, 5.5, 6.0, 6.4, 6.5, 7.0, 8.0, 9.0]); + let sim = FlipperSimulator::new(regression_config(3)); + let grid = sim.generate_regression(&x, &y, 6.5)?; + assert_eq!(grid.height(), 30 * 4); + assert!(grid.column(PAIR_LABEL_COL).is_err()); + let label_a = grid.column(LABEL_A_COL)?.f64()?; + let label_b = grid.column(LABEL_B_COL)?.f64()?; + let source_a = grid.column(SOURCE_A_ROW_COL)?.u32()?; + for i in 0..grid.height() { + let (a, b) = (label_a.get(i).unwrap(), label_b.get(i).unwrap()); + assert!(a < 6.5 && b >= 6.5, "row {i}: {a} / {b}"); + // label_a is the source row's own target. + let src = source_a.get(i).unwrap() as usize; + assert_eq!(y.f64()?.get(src), Some(a)); + } + + let again = sim.generate_regression(&x, &y, 6.5)?; + assert!(grid.equals(&again), "same seed must give the same grid"); + Ok(()) + } + + #[test] + fn regression_needs_both_sides_of_the_cutoff() { + let (x, _) = sample_df(); + let y = Series::new("t".into(), &[1.0f64; 8]); + let sim = FlipperSimulator::new(regression_config(1)); + let err = sim + .generate_regression(&x, &y, 6.5) + .unwrap_err() + .to_string(); + assert!(err.contains("8 row(s) below it, 0 at or above"), "{err}"); + assert!(sim.generate_regression(&x, &y, f64::NAN).is_err()); + } + #[test] fn generate_errors_with_fewer_than_two_label_groups() { let x = df!["f1" => [1.0, 2.0, 3.0]].unwrap(); diff --git a/core/tests/regression.rs b/core/tests/regression.rs new file mode 100644 index 0000000..64ca85d --- /dev/null +++ b/core/tests/regression.rs @@ -0,0 +1,294 @@ +use pasteur_core::evaluate::JITTER_STABILITY_ALPHA; +use pasteur_core::*; +use polars::prelude::*; +use std::collections::HashMap; + +/// Predicts feature `x0` unchanged, so each test spells out the model's +/// prediction row by row. +struct PassThrough; + +impl Model for PassThrough { + fn predict_proba(&self, df: &DataFrame) -> Result { + let mut c = df.column("x0")?.clone(); + c.rename("prediction".into()); + Ok(DataFrame::new(df.height(), vec![c])?) + } +} + +/// A binary model for the "classifier output is unchanged" check. +struct Proba; + +impl Model for Proba { + fn predict_proba(&self, df: &DataFrame) -> Result { + let mut c = df.column("x0")?.clone(); + c.rename("proba".into()); + Ok(DataFrame::new(df.height(), vec![c])?) + } +} + +fn variant(pred: &[f64], target: &[f64]) -> SimulationVariant { + SimulationVariant { + x: df!["x0" => pred.to_vec()].unwrap(), + y: df!["hba1c" => target.to_vec()].unwrap(), + } +} + +fn run( + model: &dyn Model, + clean: SimulationVariant, + variants: Vec<(&str, SimulationVariant)>, + task: TaskType, + flip_threshold: f64, +) -> Result { + let config = EvaluationConfig { + n_jitter_iters: 0, + allow_partial: true, + metrics: vec![], + flip_threshold, + task, + }; + let clean_datasets = HashMap::from([("d".to_string(), clean)]); + let simulation_variants = variants + .into_iter() + .map(|(name, v)| (name.to_string(), HashMap::from([("d".to_string(), v)]))) + .collect(); + Evaluator::new().run(model, &clean_datasets, &simulation_variants, &config) +} + +fn close(a: f64, b: f64) -> bool { + (a - b).abs() < 1e-9 +} + +const TARGET: [f64; 4] = [1.0, 2.0, 3.0, 4.0]; + +#[test] +fn clean_baselines_are_rmse_mae_and_r2() { + // Errors 0.5, 0, -1, 1 → MSE 0.5625; Var(y) = 1.25. + let r = run( + &PassThrough, + variant(&[1.5, 2.0, 2.0, 5.0], &TARGET), + vec![], + TaskType::Regression, + 0.0, + ) + .unwrap(); + assert!(close(r.baselines["rmse"], 0.75)); + assert!(close(r.baselines["mae"], 0.625)); + assert!(close(r.baselines["r2"], 1.0 - 0.5625 / 1.25)); + assert_eq!(r.baselines.len(), 3); + assert!(r.evaluations.metric_based.contains_key("r2")); + assert!(!r.evaluations.metric_based.contains_key("roc_auc")); + // Nothing simulated: neutral values, as for classifiers. + assert_eq!(r.evaluations.metric_based["r2"].resiliency, 1.0); + assert_eq!(r.evaluations.metric_invariant.jitter_stability, 1.0); + assert!(r.evaluations.metric_invariant.flipper_stability.is_none()); + + let detail = r.regression.expect("regression detail"); + assert_eq!(detail.target, "hba1c"); + assert_eq!(detail.n_rows, 4); + assert!(close(detail.target_sd, 1.25f64.sqrt())); + assert!(r.multi.is_none()); +} + +#[test] +fn resiliency_is_blackout_r2_over_clean_r2() { + // Var(target) = 1.25. Clean errors all 0.5: MSE 0.25, R² 0.8. Blackout + // errors all 1: MSE 1, R² 0.2. A quarter of the skill survives. + let r = run( + &PassThrough, + variant(&[1.5, 2.5, 3.5, 4.5], &TARGET), + vec![("blackout", variant(&[2.0, 3.0, 4.0, 5.0], &TARGET))], + TaskType::Regression, + 0.0, + ) + .unwrap(); + assert!(close(r.baselines["r2"], 0.8)); + assert!(close(r.evaluations.metric_based["r2"].resiliency, 0.25)); + let detail = r.regression.unwrap(); + assert!(close(detail.blackout_rmse.unwrap(), 1.0)); + assert!(close(detail.blackout_r2.unwrap(), 0.2)); +} + +#[test] +fn a_model_with_no_skill_has_neutral_resiliency() { + // Predicting the target's mean (2.5) everywhere: R² 0, nothing to lose. + let r = run( + &PassThrough, + variant(&[2.5; 4], &TARGET), + vec![("blackout", variant(&[2.5; 4], &TARGET))], + TaskType::Regression, + 0.0, + ) + .unwrap(); + assert!(close(r.baselines["r2"], 0.0)); + assert_eq!(r.evaluations.metric_based["r2"].resiliency, 1.0); +} + +#[test] +fn jitter_variance_is_scaled_by_the_target_variance() { + // Row 0 moves 1.0 ↔ 2.0 (variance 0.25); the rest do not move. + let r = run( + &PassThrough, + variant(&TARGET, &TARGET), + vec![ + ("jitter_0", variant(&[1.0, 2.0, 3.0, 4.0], &TARGET)), + ("jitter_1", variant(&[2.0, 2.0, 3.0, 4.0], &TARGET)), + ], + TaskType::Regression, + 0.0, + ) + .unwrap(); + let v = 0.25 / 4.0 / 1.25; + assert!(close( + r.evaluations.metric_invariant.jitter_stability, + 1.0 / (1.0 + JITTER_STABILITY_ALPHA * v) + )); + // Mean per-patient SD in target units: 0.5 on one of four rows. + assert!(close( + r.regression.unwrap().jitter_prediction_sd.unwrap(), + 0.125 + )); +} + +#[test] +fn jitter_stability_does_not_depend_on_the_target_unit() { + let scaled = |k: f64| { + let t: Vec = TARGET.iter().map(|v| v * k).collect(); + let draw = |first: f64| { + let mut p = t.clone(); + p[0] = first * k; + variant(&p, &t) + }; + run( + &PassThrough, + variant(&t, &t), + vec![("jitter_0", draw(1.0)), ("jitter_1", draw(2.0))], + TaskType::Regression, + 0.0, + ) + .unwrap() + .evaluations + .metric_invariant + .jitter_stability + }; + // mmol/mol vs %: same model, same score. + assert!(close(scaled(1.0), scaled(10.93))); +} + +/// One pair, `n` steps from a target-5 patient to a target-8 patient. +fn regression_grid(preds: &[f64], label_a: f64, label_b: f64) -> SimulationVariant { + let n = preds.len(); + let x = df![ + "pair_id" => vec![0u32; n], + "step" => (0..n as u32).collect::>(), + "t" => (0..n).map(|i| i as f64 / (n - 1) as f64).collect::>(), + "source_a_row" => vec![0u32; n], + "source_b_row" => vec![1u32; n], + "label_a" => vec![label_a; n], + "label_b" => vec![label_b; n], + "x0" => preds.to_vec(), + ] + .unwrap(); + let y = df!["hba1c" => vec![0.0; n]].unwrap(); + SimulationVariant { x, y } +} + +#[test] +fn flipper_finds_where_the_prediction_crosses_the_cutoff() { + // Predictions 5, 6, 7, 8 at t = 0, 1/3, 2/3, 1 cross 6.5 at t = 0.5. + let r = run( + &PassThrough, + variant(&TARGET, &TARGET), + vec![("flipper", regression_grid(&[5.0, 6.0, 7.0, 8.0], 5.0, 8.0))], + TaskType::Regression, + 6.5, + ) + .unwrap(); + assert!(close( + r.evaluations.metric_invariant.flipper_stability.unwrap(), + 0.5 + )); + let detail = r.regression.unwrap(); + assert_eq!(detail.flip_threshold, Some(6.5)); + assert_eq!(detail.flipper_never_flipped, Some(0.0)); +} + +#[test] +fn flipper_rejects_a_grid_built_for_another_cutoff() { + // Both endpoints below 9.0: this pair was not sampled across it. + let err = run( + &PassThrough, + variant(&TARGET, &TARGET), + vec![("flipper", regression_grid(&[5.0, 6.0, 7.0, 8.0], 5.0, 8.0))], + TaskType::Regression, + 9.0, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("--flip-threshold 9"), "{err}"); + + // A binary grid's 0/1 labels do not straddle a clinical cutoff either. + let binary_grid = regression_grid(&[5.0, 6.0, 7.0, 8.0], 0.0, 1.0); + let err = run( + &PassThrough, + variant(&TARGET, &TARGET), + vec![("flipper", binary_grid)], + TaskType::Regression, + 6.5, + ); + assert!(err.is_err()); +} + +#[test] +fn constant_or_missing_targets_are_rejected() { + let err = run( + &PassThrough, + variant(&[1.0, 2.0], &[3.0, 3.0]), + vec![], + TaskType::Regression, + 0.0, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("same value for every clean row"), "{err}"); + assert!(run( + &PassThrough, + variant(&[1.0, 2.0], &[3.0, f64::NAN]), + vec![], + TaskType::Regression, + 0.0, + ) + .is_err()); +} + +#[test] +fn a_regression_model_must_return_one_column() { + struct TwoColumns; + impl Model for TwoColumns { + fn predict_proba(&self, df: &DataFrame) -> Result { + Ok(df!["a" => vec![0.0; df.height()], "b" => vec![0.0; df.height()]]?) + } + } + assert!(run( + &TwoColumns, + variant(&TARGET, &TARGET), + vec![], + TaskType::Regression, + 0.0 + ) + .is_err()); +} + +#[test] +fn binary_results_have_no_regression_section() { + let clean = SimulationVariant { + x: df!["x0" => [0.1, 0.9, 0.4, 0.6]].unwrap(), + y: df!["label" => [0i64, 1, 0, 1]].unwrap(), + }; + let r = run(&Proba, clean, vec![], TaskType::Binary, 0.5).unwrap(); + assert_eq!(r.baselines.keys().collect::>(), vec!["roc_auc"]); + assert!(r.regression.is_none()); + let json = serde_json::to_value(&r).unwrap(); + assert!(json.get("regression").is_none()); + assert!(json.get("multi").is_none()); +} diff --git a/docs/source/cli.rst b/docs/source/cli.rst index 3e82183..560f643 100644 --- a/docs/source/cli.rst +++ b/docs/source/cli.rst @@ -26,7 +26,13 @@ Important options: - Local clean parquet * - ``--labels`` - none - - Cohort-membership parquet; enables flipper + - Cohort-membership parquet; enables flipper (classifiers) + * - ``--targets``, ``--target-col`` + - none + - Regression: targets parquet and its target column + * - ``--targets-id-col`` + - ``node_id`` + - Regression: ID column of ``--targets``, matched as text * - ``--output`` - ``output`` - Simulation bundle root @@ -54,8 +60,12 @@ Important options: for ``--task multiclass`` / ``multilabel`` * - ``--task`` - ``binary`` - - ``binary``, ``multiclass``, or ``multilabel``: how flipper pairs are - sampled + - ``binary``, ``multiclass``, ``multilabel``, or ``regression``: how + flipper pairs are sampled + * - ``--flip-threshold`` + - none + - Regression only: clinical cutoff, in target units, that flipper pairs + are sampled across. Without it, regression skips flipper * - ``--flipper-pairs`` - ``200`` - Opposite-label patient pairs to sample for flipper @@ -114,8 +124,14 @@ from each. - required - Bundle written by ``simulate`` * - ``--labels`` - - required + - classifiers - Cohort-membership parquet (see :doc:`data-contracts`) + * - ``--targets``, ``--target-col`` + - regression + - Targets parquet and its target column (see :doc:`data-contracts`) + * - ``--targets-id-col`` + - ``node_id`` + - ID column of ``--targets``, matched as text * - ``--model`` - required - Local ``.onnx`` file @@ -138,11 +154,13 @@ from each. * - ``--positive-class-index`` - ``1`` - Binary only: column of the probability output that is the positive - class. Rejected for multiclass and multilabel models + class. Rejected for multiclass, multilabel and regression models * - ``--flip-threshold`` - ``0.5`` - Decision threshold used by flipper stability. Multilabel: applies to - labels without a contract threshold. Unused for multiclass (argmax) + labels without a contract threshold. Unused for multiclass (argmax). + Regression: the clinical cutoff in target units, with no default; + required for ``evaluate flipper`` and must match ``simulate``'s * - ``-o``, ``--output`` - stdout - Write the JSON result to a file. The directory must already exist. @@ -184,6 +202,23 @@ For a multiclass or multilabel model, pass the task and one group per output: The model's ``metadata.json`` must declare ``task_type`` and ``output.classes``; see :doc:`data-contracts`. +For a regression model, pass the targets instead of cohorts, and the clinical +cutoff for flipper to both commands: + +.. code-block:: bash + + pasteur-cli simulate --input ./clean.parquet --output ./sim \ + --task regression --targets ./targets.parquet --target-col hba1c \ + --flip-threshold 6.5 + + pasteur-cli evaluate flipper --sim-root ./sim \ + --task regression --targets ./targets.parquet --target-col hba1c \ + --flip-threshold 6.5 --model ./models/model.onnx + +``--predictions-out`` then writes ``target`` and one prediction column per +model. ``all_agree`` is true when every model puts the patient on the same +side of the cutoff, and null when no ``--flip-threshold`` is given. + cache ----- diff --git a/docs/source/data-contracts.rst b/docs/source/data-contracts.rst index c67e255..0f320bf 100644 --- a/docs/source/data-contracts.rst +++ b/docs/source/data-contracts.rst @@ -36,6 +36,8 @@ The original ID is converted to a string, not integer-coerced. A separate Every column other than the ID is carried into the outputs and sent to the model, in file order. Drop outcome, label, and administrative columns before running Pasteur, and keep the features in the order the model was trained on. +A regression target in particular belongs in a separate targets file (see +below), never in the clean cohort, or the model is handed the answer. The input must be parquet; :doc:`model-selection` shows how to convert a CSV. Labels @@ -79,6 +81,41 @@ it cannot check that group 12 really means the model's second class. :doc:`model-selection` includes a short script that builds this file from an outcome column. +Targets (regression) +-------------------- + +A regression model is scored against a continuous true value for each row, +from a separate parquet passed as ``--targets`` with ``--target-col``: + +.. list-table:: + :header-rows: 1 + :widths: 28 22 50 + + * - Column + - Type + - Meaning + * - ID column + - any + - Selected by ``--targets-id-col`` (default ``node_id``) + * - Target column + - numeric + - Selected by ``--target-col``, e.g. ``hba1c``, in the target's units + +IDs are matched to the clean cohort's ``--id-col`` **as text**, so IDs such as +``A00123`` work and are not coerced to integers. Rows in the targets file that +are not in the cohort are ignored. + +Every cohort row needs exactly one finite target. A row with no entry, an ID +listed twice, or a missing or non-finite value stops the run with a count and +example IDs. Pasteur does not drop those rows, because scoring a model on +whichever rows happen to have a value would quietly change the cohort. A +target with the same value on every row is also an error: it cannot tell a +good model from a bad one. + +``--labels`` and ``--positive-group-id`` are for classifiers; passing +``--labels`` with ``--task regression``, or ``--targets`` with a classifier, +is an error. + Model contract -------------- @@ -118,8 +155,8 @@ Multiclass and multilabel models must also declare what their outputs are: * - Field - Meaning * - ``task_type`` - - ``binary`` (default), ``multiclass``, or ``multilabel``. Must match - ``--task``. + - ``binary`` (default), ``multiclass``, ``multilabel``, or + ``regression``. Must match ``--task``. * - ``output.name`` - Graph output holding the probabilities. Defaults to ``output_probability``, then ``probabilities``. @@ -131,6 +168,24 @@ Multiclass and multilabel models must also declare what their outputs are: - Multilabel only, optional: one decision threshold per label. Defaults to ``--flip-threshold`` for every label. +A single-target regression model declares ``task_type`` and, optionally, a +name for its one output column: + +.. code-block:: json + + { + "task_type": "regression", + "input": {"feature_order": ["glucose", "bmi", "age"]}, + "output": {"classes": ["hba1c"]} + } + +For regression, ``output.name`` defaults to ``variable`` (what scikit-learn +regressors export), then ``predictions``. ``output.classes`` is optional and, +if given, holds exactly one name; the prediction column is otherwise called +``prediction``. The output must be a float tensor of shape ``[n]`` or +``[n, 1]``. ``output.thresholds`` and ``--positive-class-index`` are rejected: +a regression model's cutoff is ``--flip-threshold``, in the target's units. + Pasteur reads every probability format scikit-learn models export to ONNX: a ``[n, K]`` tensor, a ZipMap with integer or string keys, and, for multilabel, ``MultiOutputClassifier``'s one ``[n, 2]`` tensor per label. A model declared diff --git a/docs/source/metrics.rst b/docs/source/metrics.rst index fc5c205..525242f 100644 --- a/docs/source/metrics.rst +++ b/docs/source/metrics.rst @@ -184,6 +184,77 @@ grids store the label index in ``pair_label``. Evaluating a grid built for a different task is an error. Evaluating one built with the same task but different groups is not detected. +Regression models +----------------- + +A single-target regression model (``--task regression``) predicts a value in +the target's units instead of a probability. The JSON keeps the same shape, +with these fields: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Field + - Definition + * - ``baselines.rmse``, ``baselines.mae``, ``baselines.r2`` + - Root-mean-square error, mean absolute error, and R² on the clean + cohort. RMSE and MAE are in the target's units; lower is better. R² is + 1 for a perfect model, 0 for one no better than predicting the mean, + and negative for one worse than that. + * - ``metric_based.r2.resiliency`` + - Blackout R² ÷ clean R²: the share of the model's skill over predicting + the cohort mean that survives blackout. 1.0 means blackout cost + nothing; 0 means the model fell back to no better than the mean; + negative means blackout made it worse than the mean. It can exceed + 1.0 when blackout happens to help, which usually means the + blacked-out feature was hurting the model. A model whose clean R² is + 0 or below has no skill to keep and is reported at the neutral 1.0; + read ``baselines.r2`` first. + * - ``metric_invariant.jitter_stability`` + - ``1 / (1 + 100 × v / var(target))``, where ``v`` is the per-patient + variance of the prediction across jitter variants, averaged over + patients. Dividing by the target's variance makes the score + unit-free: the same model scores the same in mmol/mol and in %. + * - ``metric_invariant.flipper_stability`` + - For each pair of one patient whose true target is below the + ``--flip-threshold`` cutoff and one at or above it, the position ``t`` + (0 → 1) where the *prediction* first crosses the cutoff. Pairs that + never cross count as 1.0. + +**Why an R² ratio and not an RMSE ratio.** An error ratio rewards a model +for having little accuracy to lose. When a model loses a feature entirely and +falls back to the mean, clean RMSE ÷ blackout RMSE is ``sqrt(1 − R²)``: about +0.58 for a model with R² 0.66, but 0.86 for one with R² 0.26. On a synthetic +HbA1c cohort with glucose blacked out for every patient, a 1-nearest-neighbour +model dropped from R² 0.26 to −0.09 (worse than predicting the mean) and still +had the best RMSE ratio, 0.82, of the models that used glucose. Its R² ratio +is −0.35, the worst. + +The ``regression`` section reports the same behaviour in the target's units: +``target_sd``, ``blackout_rmse``, ``blackout_r2``, ``jitter_prediction_sd`` (the mean +per-patient standard deviation of the prediction across jitter draws), +``flip_threshold``, and ``flipper_never_flipped``. It is absent for +classifiers. + +**The flipper cutoff is a clinical choice.** There is no decision boundary in +a regression model, so flipper needs one: the threshold your service acts on, +such as HbA1c 6.5%. Pass the same ``--flip-threshold`` to ``simulate`` and to +``evaluate``. A grid evaluated against a cutoff that some pair does not +straddle is rejected. Without a cutoff, ``simulate`` skips flipper. + +**Read ``flipper_never_flipped`` here too.** Pairs are chosen by their true +targets, but crossings are read from predictions. A model whose predictions +stay on one side of the cutoff for patients just above it never crosses, and +those pairs count as 1.0. Regression models shrink towards the mean, so when +the cutoff sits in the tail of the target this is common: on a synthetic +cohort with 8% of HbA1c values at or above 6.5%, 57–68% of pairs never +crossed for well-fitted models, and those pairs dominate the mean. + +**Jitter depends on the cohort.** Because jitter is scaled by the target's +variance in the clean cohort, compare jitter stability between models on the +same cohort only. + What these numbers do not show ------------------------------ diff --git a/docs/source/outputs.rst b/docs/source/outputs.rst index 5ce4d98..98da45b 100644 --- a/docs/source/outputs.rst +++ b/docs/source/outputs.rst @@ -45,8 +45,10 @@ Flipper output also contains ``pair_id``, ``step``, ``t``, ``source_a_row``, ``source_b_row``, ``label_a``, and ``label_b``. Grids simulated with ``--task multilabel`` also contain ``pair_label``, the index of the label the pair was sampled for. For multiclass grids ``label_a`` and ``label_b`` are -class indices in ``--positive-group-id`` order. Evaluation strips these -metadata columns before model scoring. +class indices in ``--positive-group-id`` order. For regression grids they are +the two patients' true target values, one below the ``--flip-threshold`` +cutoff and one at or above it. Evaluation strips these metadata columns +before model scoring. Evaluation JSON --------------- @@ -95,6 +97,25 @@ above hold macro averages; :doc:`metrics` defines each field. Multiclass results also carry ``flipper_detour_rate``. +Regression results key ``baselines`` by ``rmse``, ``mae`` and ``r2``, key +``metric_based`` by ``r2``, and add a ``regression`` object in the target's +units (see :doc:`metrics`): + +.. code-block:: json + + "regression": { + "target": "hba1c", + "n_rows": 300, + "target_sd": 1.07, + "blackout_rmse": 0.42, + "blackout_r2": 0.85, + "jitter_prediction_sd": 0.05, + "flip_threshold": 6.5, + "flipper_never_flipped": 0.23 + } + +Fields for simulation types that were not loaded are ``null``. + Comparison predictions ---------------------- @@ -102,9 +123,11 @@ When ``compare`` receives ``--predictions-out``, the parquet contains: - ``source_row_id`` - ``variant_name`` -- ``label`` (multi-output: ``label__`` per class) +- ``label`` (multi-output: ``label__`` per class; regression: + ``target``, the true value, null on synthetic flipper rows) - one floating-point probability column per model (multi-output: - ``__`` per class) -- ``all_agree`` + ``__`` per class; regression: the predicted value) +- ``all_agree`` (regression: same side of ``--flip-threshold``, null without + one) The model column name is derived from the model file stem. diff --git a/model/src/contract.rs b/model/src/contract.rs index ca404fd..1aaaaa2 100644 --- a/model/src/contract.rs +++ b/model/src/contract.rs @@ -25,13 +25,16 @@ pub(crate) struct ContractInput { #[derive(Debug, Default, serde::Deserialize)] pub(crate) struct ContractOutput { - /// Graph output holding the probabilities. Defaults to the first of - /// `output_probability`/`probabilities` the graph has. + /// Graph output holding the probabilities (or, for regression, the + /// predicted value). Defaults to the first of + /// `output_probability`/`probabilities` the graph has, or of + /// `variable`/`predictions` for regression. #[serde(default)] pub(crate) name: Option, /// One name per output column, in the model's column order. Required for /// multiclass and multilabel models: it is the only record of what each - /// column means, and its length is checked against the graph. + /// column means, and its length is checked against the graph. Optional + /// for regression, where its one entry names the target. #[serde(default)] pub(crate) classes: Option>, /// Multilabel only: per-label decision thresholds, in `classes` order. @@ -50,8 +53,13 @@ pub(crate) struct ModelContract { } const DEFAULT_PROBABILITY_OUTPUTS: [&str; 2] = ["output_probability", "probabilities"]; +/// sklearn-onnx names a regressor's output `variable`. +const DEFAULT_PREDICTION_OUTPUTS: [&str; 2] = ["variable", "predictions"]; /// Column name of a binary model's single positive-class output. pub const BINARY_OUTPUT_COLUMN: &str = "proba"; +/// Column name of a regression model's output when `output.classes` does not +/// name it. +pub const REGRESSION_OUTPUT_COLUMN: &str = "prediction"; /// What `predict_proba` reads and how it labels it, settled at load time so a /// mismatched contract fails before any patient is scored. @@ -61,7 +69,8 @@ pub(crate) struct OutputSpec { pub(crate) output_name: String, /// Binary only. pub(crate) positive_class_index: usize, - /// Output column names: `["proba"]` for binary, `output.classes` otherwise. + /// Output column names: `["proba"]` for binary, `output.classes` for + /// multi-output, and one name for regression. pub(crate) columns: Vec, pub(crate) thresholds: Option>, } @@ -70,23 +79,30 @@ pub(crate) struct OutputSpec { /// when the graph declares one. fn resolve_output_name( contract_label: &str, + task: TaskType, declared: Option<&ContractOutput>, outputs: &[(String, Option)], ) -> Result { + let (kind, defaults) = match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => { + ("probability", DEFAULT_PROBABILITY_OUTPUTS) + } + TaskType::Regression => ("prediction", DEFAULT_PREDICTION_OUTPUTS), + }; match declared.and_then(|o| o.name.clone()) { Some(name) if outputs.iter().any(|(n, _)| *n == name) => Ok(name), Some(name) => Err(CoreError::InvalidConfig(format!( "{contract_label} names output {name:?} but the model's outputs are {:?}", outputs.iter().map(|(n, _)| n).collect::>() ))), - None => match DEFAULT_PROBABILITY_OUTPUTS + None => match defaults .iter() .find(|d| outputs.iter().any(|(n, _)| n == **d)) { Some(name) => Ok(name.to_string()), None => Err(CoreError::InvalidConfig(format!( - "model has no probability output (looked for {DEFAULT_PROBABILITY_OUTPUTS:?}; \ - its outputs are {:?}). Set `output.name` in {contract_label}.", + "model has no {kind} output (looked for {defaults:?}; its outputs are {:?}). \ + Set `output.name` in {contract_label}.", outputs.iter().map(|(n, _)| n).collect::>() ))), }, @@ -190,6 +206,54 @@ fn resolve_multi_spec( }) } +fn resolve_regression_spec( + contract_label: &str, + output_name: String, + positive_class_index: Option, + width: Option, + classes: Option>, + thresholds: Option>, +) -> Result { + let invalid = |msg: String| Err(CoreError::InvalidConfig(msg)); + if positive_class_index.is_some() { + return invalid( + "--positive-class-index selects one column of a binary model; a regression \ + model has one output, its predicted value" + .to_string(), + ); + } + if thresholds.is_some() { + return invalid(format!( + "{contract_label} sets `output.thresholds`, which applies to multilabel models; \ + a regression model's flipper cutoff is --flip-threshold, in the target's units" + )); + } + if let Some(w) = width.filter(|w| *w != 1) { + return invalid(format!( + "output {output_name:?} has {w} columns; Pasteur scores single-target \ + regression models, which predict one value per row" + )); + } + let column = match classes.as_deref() { + None => REGRESSION_OUTPUT_COLUMN.to_string(), + Some([name]) => name.clone(), + Some(names) => { + return invalid(format!( + "{contract_label} lists {} `output.classes` for a regression model; give one \ + name for the target, or leave it out", + names.len() + )) + } + }; + Ok(OutputSpec { + task: TaskType::Regression, + output_name, + positive_class_index: 0, + columns: vec![column], + thresholds: None, + }) +} + pub(crate) fn resolve_output_spec( contract_label: &str, contract: Option<&ModelContract>, @@ -199,7 +263,7 @@ pub(crate) fn resolve_output_spec( let task = contract.map(|c| c.task_type).unwrap_or_default(); let declared = contract.map(|c| &c.output); - let output_name = resolve_output_name(contract_label, declared, outputs)?; + let output_name = resolve_output_name(contract_label, task, declared, outputs)?; let width = outputs .iter() @@ -217,8 +281,8 @@ pub(crate) fn resolve_output_spec( } } - if task.is_binary() { - resolve_binary_spec( + match task { + TaskType::Binary => resolve_binary_spec( contract_label, task, output_name, @@ -226,16 +290,23 @@ pub(crate) fn resolve_output_spec( width, classes, thresholds, - ) - } else { - resolve_multi_spec( + ), + TaskType::Multiclass | TaskType::Multilabel => resolve_multi_spec( contract_label, task, output_name, positive_class_index, classes, thresholds, - ) + ), + TaskType::Regression => resolve_regression_spec( + contract_label, + output_name, + positive_class_index, + width, + classes, + thresholds, + ), } } @@ -518,6 +589,78 @@ mod tests { assert!(err.contains("[0, 1]"), "got {err}"); } + fn regression_outputs(width: Option) -> Vec<(String, Option)> { + vec![("variable".to_string(), width)] + } + + fn resolve_regression( + json: &str, + pci: Option, + width: Option, + ) -> Result { + resolve_output_spec( + "metadata.json", + Some(&contract(json)), + pci, + ®ression_outputs(width), + ) + .map_err(|e| e.to_string()) + } + + #[test] + fn regression_defaults_to_the_variable_output() { + let spec = resolve_regression(r#"{"task_type": "regression"}"#, None, Some(1)).unwrap(); + assert_eq!(spec.task, TaskType::Regression); + assert_eq!(spec.output_name, "variable"); + assert_eq!(spec.columns, vec![REGRESSION_OUTPUT_COLUMN.to_string()]); + let named = resolve_regression( + r#"{"task_type": "regression", "output": {"classes": ["hba1c"]}}"#, + None, + Some(1), + ) + .unwrap(); + assert_eq!(named.columns, vec!["hba1c".to_string()]); + } + + #[test] + fn regression_without_a_prediction_output_names_what_it_looked_for() { + let err = resolve_output_spec( + "metadata.json", + Some(&contract(r#"{"task_type": "regression"}"#)), + None, + &outputs(Some(2)), + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("no prediction output") && err.contains("variable"), + "{err}" + ); + } + + #[test] + fn regression_is_single_output_without_thresholds_or_positive_class() { + let json = r#"{"task_type": "regression"}"#; + let err = resolve_regression(json, None, Some(3)).unwrap_err(); + assert!(err.contains("3 columns"), "{err}"); + let err = resolve_regression(json, Some(1), Some(1)).unwrap_err(); + assert!(err.contains("--positive-class-index"), "{err}"); + let err = resolve_regression( + r#"{"task_type": "regression", "output": {"thresholds": [6.5]}}"#, + None, + Some(1), + ) + .unwrap_err(); + assert!(err.contains("--flip-threshold"), "{err}"); + let err = resolve_regression( + r#"{"task_type": "regression", "output": {"classes": ["a", "b"]}}"#, + None, + None, + ) + .unwrap_err(); + assert!(err.contains("2 `output.classes`"), "{err}"); + } + #[test] fn multi_output_needs_unique_classes() { let err = resolve_err(r#"{"task_type": "multilabel"}"#, None); diff --git a/model/src/onnx.rs b/model/src/onnx.rs index 97a5717..e0a4a1a 100644 --- a/model/src/onnx.rs +++ b/model/src/onnx.rs @@ -9,7 +9,7 @@ use polars::prelude::*; use std::path::Path; use std::sync::Mutex; -pub use crate::contract::{BINARY_OUTPUT_COLUMN, CONTRACT_FILENAME}; +pub use crate::contract::{BINARY_OUTPUT_COLUMN, CONTRACT_FILENAME, REGRESSION_OUTPUT_COLUMN}; /// In-process ONNX Runtime model, loaded from a locally-cached (pulled) file. /// No HTTP server — this is a direct `Model` impl for scoring, per the @@ -33,7 +33,8 @@ impl OnnxModel { /// blackout push patients toward the positive class. /// /// `positive_class_index` picks the scored column of a binary model - /// (default 1) and must be `None` for multiclass and multilabel models. + /// (default 1) and must be `None` for multiclass, multilabel and + /// regression models. pub fn from_file( path: &Path, input_name: impl Into, @@ -132,6 +133,8 @@ fn static_width(dtype: &ValueType) -> Option { } impl Model for OnnxModel { + /// For a regression model the one column is the predicted value, in the + /// target's units, not a probability. fn predict_proba(&self, df: &DataFrame) -> Result { let n_rows = df.height(); let n_features = self.feature_order.len(); diff --git a/model/src/probabilities.rs b/model/src/probabilities.rs index 8432c3a..8c54ae4 100644 --- a/model/src/probabilities.rs +++ b/model/src/probabilities.rs @@ -1,6 +1,6 @@ //! Reading a classifier's probability output into one column per class, -//! from every shape sklearn-onnx exports, and refusing outputs that cannot -//! be scored honestly. +//! or a regressor's predicted value into one column, from every shape +//! sklearn-onnx exports, and refusing outputs that cannot be scored honestly. use crate::contract::{OutputSpec, CONTRACT_FILENAME}; use ort::value::{DynMapValueType, DynSequenceValueType, DynTensorValueType, DynValue}; @@ -21,6 +21,7 @@ pub(crate) fn read_probabilities( TaskType::Multiclass | TaskType::Multilabel => { extract_all_probabilities(value, spec.task, &spec.columns)? } + TaskType::Regression => vec![extract_prediction(value)?], }; check_columns(&columns, spec.task, n_rows)?; Ok(columns) @@ -29,13 +30,17 @@ pub(crate) fn read_probabilities( fn check_columns(columns: &[Vec], task: TaskType, n_rows: usize) -> Result<(), CoreError> { let total: usize = columns.iter().map(Vec::len).sum(); let nans = columns.iter().flatten().filter(|p| p.is_nan()).count(); + let noun = match task { + TaskType::Binary | TaskType::Multiclass | TaskType::Multilabel => "probabilities", + TaskType::Regression => "predictions", + }; if nans > 0 { // A NaN fill only means "missingness" to a NaN-native model. // Anything exported from sklearn-onnx propagates it straight // through, and a NaN probability silently poisons ROC AUC rather // than erroring. Refuse. return Err(CoreError::Generic(format!( - "{nans}/{total} probabilities came back NaN — this model is not NaN-native. \ + "{nans}/{total} {noun} came back NaN — this model is not NaN-native. \ Pass the value it was trained with (e.g. --null-fill 0), or declare \ `input.absent_sentinel` in its {CONTRACT_FILENAME}." ))); @@ -48,7 +53,24 @@ fn check_columns(columns: &[Vec], task: TaskType, n_rows: usize) -> Result< } match task { TaskType::Multiclass => check_rows_sum_to_one(columns), - TaskType::Binary | TaskType::Multilabel => Ok(()), + TaskType::Binary | TaskType::Multilabel | TaskType::Regression => Ok(()), + } +} + +/// A regressor's output: a `tensor(float)` of shape `[n]` or `[n, 1]`. The +/// value is in the target's units, so nothing is range-checked. +fn extract_prediction(value: &DynValue) -> Result, CoreError> { + let (shape, data) = value.try_extract_tensor::().map_err(|e| { + CoreError::Generic(format!( + "a regression model's output must be a float tensor of shape [n] or [n, 1]: {e}" + )) + })?; + match **shape { + [_] | [_, 1] => Ok(data.to_vec()), + _ => Err(CoreError::Generic(format!( + "prediction tensor has shape {shape:?}; a single-target regression model \ + returns [n] or [n, 1]" + ))), } } diff --git a/model/tests/fixtures/expected.json b/model/tests/fixtures/expected.json index a024cfc..5dbf942 100644 --- a/model/tests/fixtures/expected.json +++ b/model/tests/fixtures/expected.json @@ -134,6 +134,26 @@ ] ] }, + "regression_linear": { + "outputs": [ + "variable" + ], + "predictions": [ + 5.499722957611084, + 7.977728366851807, + 3.031151056289673 + ] + }, + "regression_forest": { + "outputs": [ + "variable" + ], + "predictions": [ + 5.471961975097656, + 7.156322479248047, + 3.8161368370056152 + ] + }, "probe": [ [ 0.0, diff --git a/model/tests/fixtures/generate.py b/model/tests/fixtures/generate.py index 9f71bd0..63f91b5 100644 --- a/model/tests/fixtures/generate.py +++ b/model/tests/fixtures/generate.py @@ -5,22 +5,26 @@ task. `expected.json` records onnxruntime's probabilities for PROBE rows so the Rust tests can check column order, not just shape. Probabilities are stored row-major ([row][class]) for tensors, and as per-row dicts for ZipMap. +Regressors record their one predicted value per row instead. pip install scikit-learn skl2onnx onnxruntime python model/tests/fixtures/generate.py """ import json +import sys from pathlib import Path import numpy as np import onnxruntime as ort from skl2onnx import to_onnx -from sklearn.linear_model import LogisticRegression +from sklearn.ensemble import RandomForestRegressor +from sklearn.linear_model import LinearRegression, LogisticRegression from sklearn.multioutput import MultiOutputClassifier from sklearn.neural_network import MLPClassifier HERE = Path(__file__).parent +OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else HERE rng = np.random.default_rng(0) X = rng.normal(size=(200, 2)).astype(np.float32) PROBE = np.array([[0.0, 0.0], [2.0, -1.0], [-1.5, 2.0]], dtype=np.float32) @@ -29,6 +33,11 @@ y_multiclass = np.digitize(X[:, 0], [-0.5, 0.5]) # 0, 1, 2 y_multilabel = np.stack([X[:, 0] > 0, X[:, 1] > 0, X[:, 0] + X[:, 1] > 1], axis=1).astype(int) names = np.array(["low", "mid", "high"]) +# Drawn after every classifier target, so adding it left their fixtures as +# they were. +y_regression = (5.5 + X[:, 0] - 0.5 * X[:, 1] + rng.normal(scale=0.1, size=200)).astype( + np.float32 +) fixtures = { # Binary, ZipMap seq(map(int64, float)) — the pre-existing default path. @@ -49,6 +58,12 @@ MLPClassifier(hidden_layer_sizes=(4,), max_iter=500, random_state=0).fit(X, y_multilabel), {"zipmap": False}, ), + # Regression: one [n, 1] float tensor named `variable`. + "regression_linear": (LinearRegression().fit(X, y_regression), {}), + "regression_forest": ( + RandomForestRegressor(n_estimators=5, max_depth=3, random_state=0).fit(X, y_regression), + {}, + ), } expected = {} @@ -64,12 +79,18 @@ onx.graph.input[0].name = "input" for node in onx.graph.node: node.input[:] = ["input" if i == "X" else i for i in node.input] - path = HERE / f"{name}.onnx" + path = OUT / f"{name}.onnx" path.write_bytes(onx.SerializeToString()) sess = ort.InferenceSession(str(path)) outputs = sess.run(None, {"input": PROBE}) by_name = dict(zip([o.name for o in sess.get_outputs()], outputs)) + if "variable" in by_name: + expected[name] = { + "outputs": [o.name for o in sess.get_outputs()], + "predictions": np.asarray(by_name["variable"]).reshape(-1).tolist(), + } + continue probs = by_name.get("output_probability", by_name.get("probabilities")) if isinstance(probs, list) and probs and isinstance(probs[0], dict): probs = [{str(k): float(v) for k, v in row.items()} for row in probs] @@ -84,5 +105,5 @@ } expected["probe"] = PROBE.tolist() -(HERE / "expected.json").write_text(json.dumps(expected, indent=2) + "\n") +(OUT / "expected.json").write_text(json.dumps(expected, indent=2) + "\n") print("wrote", ", ".join(fixtures)) diff --git a/model/tests/fixtures/regression_forest.onnx b/model/tests/fixtures/regression_forest.onnx new file mode 100644 index 0000000..f2d5ad5 Binary files /dev/null and b/model/tests/fixtures/regression_forest.onnx differ diff --git a/model/tests/fixtures/regression_linear.onnx b/model/tests/fixtures/regression_linear.onnx new file mode 100644 index 0000000..870295d Binary files /dev/null and b/model/tests/fixtures/regression_linear.onnx differ diff --git a/model/tests/onnx_outputs.rs b/model/tests/onnx_outputs.rs index 319bfb5..6c82cbb 100644 --- a/model/tests/onnx_outputs.rs +++ b/model/tests/onnx_outputs.rs @@ -202,7 +202,60 @@ fn unknown_output_name_lists_the_real_ones() { fn unknown_task_type_names_the_valid_ones() { let err = load_err("multiclass_tensor", r#"{"task_type": "multi-label"}"#, None); assert!( - err.contains("binary") && err.contains("multiclass") && err.contains("multilabel"), + err.contains("binary") + && err.contains("multiclass") + && err.contains("multilabel") + && err.contains("regression"), "got {err}" ); } + +fn assert_predictions(model: &OnnxModel, fixture: &str, column: &str) { + let got = model.predict_proba(&probe()).unwrap(); + assert_eq!(got.get_column_names(), [column]); + let want: Vec = expected()[fixture]["predictions"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap()) + .collect(); + let col = got.column(column).unwrap().f64().unwrap(); + for (i, w) in want.iter().enumerate() { + let g = col.get(i).unwrap(); + assert!((g - w).abs() < 1e-5, "{fixture} row {i}: got {g}, want {w}"); + } +} + +#[test] +fn regression_linear_predicts_in_target_units() { + let model = load("regression_linear", r#"{"task_type": "regression"}"#, None).unwrap(); + assert_eq!(model.task(), TaskType::Regression); + assert_eq!(model.decision_thresholds(), None); + assert_predictions(&model, "regression_linear", "prediction"); +} + +#[test] +fn regression_forest_with_a_named_target() { + let contract = r#"{"task_type": "regression", "output": {"classes": ["hba1c"]}}"#; + let model = load("regression_forest", contract, None).unwrap(); + assert_predictions(&model, "regression_forest", "hba1c"); +} + +#[test] +fn a_classifier_declared_regression_has_no_prediction_output() { + let err = load_err("multiclass_tensor", r#"{"task_type": "regression"}"#, None); + assert!(err.contains("no prediction output"), "got {err}"); + // Pointed at its probabilities by name, the [n, 3] width gives it away. + let err = load_err( + "multiclass_tensor", + r#"{"task_type": "regression", "output": {"name": "probabilities"}}"#, + None, + ); + assert!(err.contains("3 columns"), "got {err}"); +} + +#[test] +fn a_regressor_declared_binary_is_refused() { + let err = load_err("regression_linear", "{}", None); + assert!(err.contains("no probability output"), "got {err}"); +}