ZD-717: Share multiclass proba preamble; align classes on estimator.classes_ (follow-up to #535/#539)

[hunner]

Pull Request Description

What and why?

Implements the three follow-up items from the review of #539 (ZD-717 multiclass work, continuing #535):

1. Extract the shared preamble. ROCCurve, PrecisionRecallCurve, PopulationStabilityIndex ([SC-17128] ZD-717: Add multiclass (one-vs-rest) support to VM tests #535) and GINITable ([SC-17187] ZD-717: GINITable fails on multiclass models (follow-up to PR #535) #539) each carried an identical ~40-line block: reach the underlying estimator's predict_proba (the VMModel wrapper only exposes the positive-class column), guard unusable models with SkipTestError, and label_binarize the true labels. That now lives in validmind/tests/model_validation/sklearn/_multiclass_proba.py (same private-module pattern as _diagnosis_metrics.py), and all four tests call it. ConfusionMatrix never had the block (it uses argmax predictions) and is untouched.

2. Align classes on estimator.classes_. The preamble assumed np.unique(dataset.y) matches predict_proba column order, but sklearn orders columns by the classes seen in training. The helper aligns on classes_ (falling back to np.unique), binarizes against the full training class list, and emits per-class output only for classes present in the evaluated y. Net behavior change: an evaluation slice missing a training class — e.g. a test split without a rare class, exactly the ZD-717 customer's input_grid over [train, test] — now computes metrics for the present classes plus micro-average instead of raising SkipTestError.

3. Close the coverage gap on the shape/missing-class branch (see tests below).

Two guards the old width check provided implicitly are kept explicit, with clear skip messages: a binary-trained model evaluated on a >2-class dataset skips (the new alignment would otherwise IndexError), and evaluation labels the model was never trained on skip rather than being silently absorbed as all-negative rows.

Binary paths of all four tests are unchanged, as is each test's len(np.unique(dataset.y)) > 2 branch decision and output structure (RawData keys, table columns, trace layout).

How to test

python -m pytest tests/unit_tests/model_validation/sklearn/test_multiclass_proba.py \
  tests/unit_tests/model_validation/statsmodels/test_GINITable.py \
  tests/unit_tests/model_validation/sklearn/test_ROCCurve.py \
  tests/unit_tests/model_validation/sklearn/test_PrecisionRecallCurve.py \
  tests/unit_tests/model_validation/sklearn/test_PopulationStabilityIndex.py -q

Expected: 22 passed, 5 skipped (the pre-existing xgboost-extra gates). Full tests/unit_tests/model_validation: 114 passed, 5 skipped (21 collection errors from missing optional extras are identical on main — environmental).

New coverage:

• test_multiclass_proba.py — 10 unit tests for the helper: happy path, classes_ alignment, missing-class, no-predict_proba / raising-predict_proba / wrong-width skips, binary-trained-model skip, unseen-label skip.
• Missing-class end-to-end tests for all four call sites (model fit on 4 classes, dataset slice omitting class 3 → per-present-class output + micro, no skip).
• Un-gated sklearn (LogisticRegression) multiclass tests for ROC/PR — their existing multiclass tests are entirely xgboost-gated, so the refactored paths previously had zero runtime coverage without the extra. PSI gains its first unit test file.

What needs special review?

• The two new skip guards' judgment calls: a binary-trained model on a multiclass dataset and unseen evaluation labels both skip with explanatory messages rather than crash/silently compute. If we'd rather compute OvR treating unseen labels as all-negative, that's a one-line change in the helper — but silent felt wrong for a model-risk tool.
• The micro-average now pools only the present classes' aligned columns (previously all np.unique(dataset.y) columns — identical when nothing is missing, which was the only case that previously computed).

Dependencies, breaking changes, and deployment notes

Follow-up to #535 / #539. Library-only, no migrations, no env vars, no version bump (left to the release process). No breaking changes: every previously-computing case produces identical output; the only behavior change is previously-skipped missing-class runs now computing.

Release notes

ROC Curve, Precision-Recall Curve, Population Stability Index, and GINITable multiclass metrics now align class columns with the model's training classes. Evaluating on a dataset slice that is missing one of the model's training classes now produces metrics for the classes present (plus the micro-average) instead of skipping the test.

Checklist

• What and why
• Screenshots or videos (Frontend) — N/A
• How to test
• What needs special review
• Dependencies, breaking changes, and deployment notes
• Labels applied — enhancement
• PR linked to Shortcut — no story yet; items originate from the review comments on [SC-17187] ZD-717: GINITable fails on multiclass models (follow-up to PR #535) #539
• Unit tests added (Backend)
• Tested locally — 114 passed / 5 skipped across tests/unit_tests/model_validation
• Documentation updated (if required) — N/A
• Environment variable additions/changes documented (if required) — N/A

🤖 Generated with Claude Code

[juanmleng]

Really nice refactor — pulling the shared "reach the estimator for the full per-class probability matrix" preamble into one helper cleans up four tests at once, and switching the column alignment to estimator.classes_ is the right call. I traced it against the previously-shipped work and the all-classes-present cases come out identical (same sorted class order), while the case where a training class is missing from the evaluated slice now computes instead of skipping. The per-class indexing and the after-loop micro-average slicing both look correct. A few small, non-blocking thoughts:

• Present-class detection (np.any(y_true == cls)) and label_binarize(..., classes=class_list) both lean on estimator.classes_ and dataset.y sharing a dtype. In the normal flow they always will, but if a model trained on integer labels ever meets a dataset whose target loads as strings, classes could be quietly marked absent or the binarization could misalign. A small normalization or an assertion would harden it, and a targeted test would lock it in.

• In the missing-class scenario the micro-average now pools only the present classes, so the absent class's column is dropped from the micro curve. That seems like the right behavior, but since it makes the micro number differ from a naive full-matrix computation, a one-line note in the module docstring would save a future reader some head-scratching.

• GINITable now imports the helper from the sklearn package's private _multiclass_proba module. The dependency direction already existed, so this is minor, but if the helper is meant to be shared across both test families a neutral (non-underscore) home might read a little better.

• One test gap worth a follow-up: in the PSI path the set of present classes is derived from the initial dataset only — the second dataset just gets the column-width check, not the present-set or unknown-label checks. It's probably fine since PSI compares score distributions rather than positives, but a case where the two slices differ in class membership would confirm it.

[hunner]

Thanks for the careful trace — addressed in 59cbf8d:

1. dtype mismatch (int-trained model, string-labeled target): good instinct, but this case can't quietly misalign — the unknown-label guard fires first, since string labels compare unequal to the integer class_list, and it skips loudly naming every label ("found label(s) ['0', '1', '2'] … not trained on"). Rather than add normalization (which would make a genuinely suspicious setup compute silently), I locked the behavior in with test_string_labels_with_int_trained_model_skip, which asserts the skip and the message contents.

2. Micro-average pooling: agreed — added the sentence to the module docstring: micro-averages pool only the present classes' aligned columns, so an absent training class is deliberately excluded from the micro number.

3. Private-module home for the helper: leaving as-is for now. _diagnosis_metrics.py set the private-module precedent in this directory, and promoting the helper to a neutral non-underscore home would create public API surface for an internal detail. If a second cross-family consumer shows up, that's the moment to relocate it — happy to do so then.

4. PSI second-dataset checks: added test_new_dataset_missing_class_uses_initial_present_set — model on 4 classes, initial slice with all 4, new slice missing class 3. Confirms your reading: the present set is driven by the initial dataset, all 4 per-class tables compute without skipping, and the new dataset only needs matching columns (PSI compares score distributions, not positives), so no present-set/unknown-label checks are needed on it.

Suite is now 24 passed / 5 xgboost-gated skips on the affected files.

🤖 Automated reply

[github-actions]

PR Summary

This pull request introduces a set of enhancements and additions to the model evaluation framework for multiclass classification. The key changes include:

1. New and updated unit tests:

   • Tests for PopulationStabilityIndex now validate one-vs-rest tables, figures, and raw outputs in scenarios where the evaluated dataset is missing one or more classes, or when the class membership between initial and new slices differs.
   • PrecisionRecallCurve tests have been refactored to cover both standard and missing-class cases, ensuring that per-class curves and micro-averaging are correctly computed.
   • ROCCurve tests have been modified to assess proper alignment of the probability matrix and to verify micro-average calculations, specifically when a training class is absent from evaluation.
   • Additional tests for the multiclass probability module (_multiclass_proba.py) verify that per-class probabilities are aligned with the model's training classes. These tests also cover error handling for situations such as missing or mismatched probability matrices, cases with unseen labels, and conflicts between integer and string label types.
   • Updates to GINITable tests to handle missing classes by computing metrics for present classes plus a micro-average.

2. Code refactoring and enhancements:

   • Introduction of a unified helper, multiclass_proba, that centralizes the logic for retrieving and aligning per-class probability matrices. This helper ensures that the probability columns are aligned with the model's training classes and that the one-hot encoding for the true labels remains consistent even when some classes are absent in the evaluation data.
   • Modifications to PopulationStabilityIndex, PrecisionRecallCurve, ROCCurve, and GINITable functions to utilize the new helper, thereby reducing code duplication and enhancing consistency across multiclass metric computations.
   • Improvements in error handling via SkipTestError for models that do not provide a suitable probability matrix or when the dataset contains labels the model was not trained on.

Overall, these changes strengthen the robustness and flexibility of the model validation process in handling multiclass scenarios, ensuring that metrics are computed accurately even when the evaluation dataset does not completely mirror the training data.

Test Suggestions

• Run all unit tests to validate that the new multiclass logic works across various scenarios (all classes present, missing classes, and mismatched class labels).
• Simulate cases where the model's predict_proba function is missing or fails, and verify that a SkipTestError is raised with the appropriate message.
• Check that the generated figures (plots) contain the correct number of traces and proper legends for both per-class and micro-average metrics.
• Validate that the probability matrices are correctly aligned, and that the per-instance probabilities sum to 1 within a small tolerance.
• Test edge cases such as datasets with unseen labels or labels in string format when the model is trained with integers.