Skip to content

Support binary classification in conformal prediction set methods - #1238

Open
fbonc wants to merge 6 commits into
sunlabuiuc:masterfrom
fbonc:fix-binary-mode-cp-prediction-sets
Open

fbonc wants to merge 6 commits into
sunlabuiuc:masterfrom
fbonc:fix-binary-mode-cp-prediction-sets

Conversation

@fbonc

@fbonc fbonc commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

Problem
Every SetPredictor subclass only accepted multiclass base models. A binary model emits a single positive-class probability y_prob of shape (N, 1), but a conformal prediction set ranges over both labels {0, 1} and needs one column per class, so binary tasks were out entirely, and the binary metrics function couldn't score prediction sets at all.

Fix
Re-present a binary model as a 2-class problem at the two boundaries where a raw (N, 1) probability enters the pipeline: expand it to [P(y=0), P(y=1)] (N, 2) before the non-conformity scoring in calibrate (for the threshold) and in forward (for the set), so the existing split-conformal machinery runs unchanged with K = 2. This sits on top of the new scores.py scoring layer, so binary composes with any score_type (threshold or APS). Then wire the prediction-set metrics into binary_metrics_fn so results are scorable.

Changes

  • Modified pyhealth/metrics/binary.py: binary_metrics_fn now accepts y_predset and computes the conformal set metrics.
  • Modified pyhealth/metrics/prediction_set.py: added PREDICTION_SET_METRICS and compute_prediction_set_metric, a single dispatch point for set metrics.
  • Modified pyhealth/metrics/multiclass.py: delegates set-metric dispatch to the shared helper (removes duplicated block).
  • Modified pyhealth/calib/utils.py: added binary_to_2col (primitive) plus expand_binary_cal (calibrate side, numpy) and expand_binary_pred (forward side, numpy) which expand the probability before the scores.py score functions.
  • Modified pyhealth/calib/predictionset/{label,base_conformal,cluster/cluster_label,cluster/neighborhood_label,covariate/covariate_label}.py: accept binary in the mode gate, and expand y_prob to 2 columns before the true_class_* / all_class_* score functions in calibrate and forward.
  • Modified pyhealth/calib/predictionset/scrib/__init__.py: accept binary in the mode gate; SCRIB keeps its own quicksearch scoring, so the expansion is an inline torch step (it thresholds the probability directly rather than going through scores.py).
  • Created tests/core/test_binary_label.py, tests/core/test_binary_predictionset.py: binary end-to-end coverage for all six methods and the helpers.
  • Modified tests/core/test_{cluster,covariate,neighborhood}_label.py: the "binary raises" tests now assert binary is accepted and that a genuinely unsupported mode still raises.

Core decisions

  • y_prob stays native (N, 1) - only the set is expanded. roc_auc, etc. are unaffected, and self.mode stays honest ("binary") rather than being faked as multiclass.
  • Expansion happens in both calibrate and forward. They're different data; only the threshold self.t passes between them, so each entry point must expand its own raw probability.
  • Binary composes with the scores.py layer: expanding to (N, 2) before the score functions means score_type (threshold / APS) works for binary with no extra code.
  • Set-metric dispatch is centralized in prediction_set.py, and the binary and multiclass metrics functions delegate to the shared helper (previously duplicated in multiclass.py).
  • FavMac is intentionally excluded. It is multilabel, and a binary task doesn't map to its cost-over-multiple-labels contract.

Extend LABEL, BaseConformal, SCRIB, ClusterLabel, NeighborhoodLabel, and
CovariateLabel to accept binary base models by re-presenting the single
positive-class probability as a 2-class problem at the calibrate and forward
boundaries. FavMac stays multilabel-only.
Reconcile binary-mode support with upstream's new scores.py scoring
refactor (sunlabuiuc#1189 real APS/dynamic scoring, sunlabuiuc#1190 SCRIB fixes): expand binary
y_prob to 2 columns before the true_class_*/all_class_* score functions in
each method's calibrate and forward, on top of the new score_type/rng API.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant