diff --git a/src/rtichoke/performance_data/performance_data.py b/src/rtichoke/performance_data/performance_data.py index d734ac99..27bac46f 100644 --- a/src/rtichoke/performance_data/performance_data.py +++ b/src/rtichoke/performance_data/performance_data.py @@ -19,6 +19,31 @@ import numpy as np +_PERFORMANCE_DATA_COLUMNS = [ + "reference_group", + "stratified_by", + "chosen_cutoff", + "true_positives", + "true_negatives", + "false_positives", + "false_negatives", + "predicted_positives", + "predicted_negatives", + "real_positives", + "real_negatives", + "n", + "sensitivity", + "specificity", + "ppv", + "npv", + "false_positive_rate", + "lift", + "net_benefit", + "net_benefit_interventions_avoided", + "ppcr", +] + + def _probs_with_r_binary_cutoff_semantics( probs: Dict[str, np.ndarray], by: float ) -> Dict[str, np.ndarray]: @@ -192,6 +217,6 @@ def prepare_performance_data( performance_data = _turn_cumulative_aj_to_performance_data(cumulative_aj_data) - return performance_data.sort( + return performance_data.select(_PERFORMANCE_DATA_COLUMNS).sort( ["reference_group", "stratified_by", "chosen_cutoff"] ) diff --git a/tests/test_performance_schema_order.py b/tests/test_performance_schema_order.py new file mode 100644 index 00000000..5a21bc9d --- /dev/null +++ b/tests/test_performance_schema_order.py @@ -0,0 +1,72 @@ +import numpy as np +import polars as pl +from polars.testing import assert_frame_equal + +from rtichoke import prepare_performance_data + + +PROBS = {"model": np.array([0.1, 0.2, 0.5, 0.5, 0.8, 0.9])} +REALS = np.array([0, 0, 1, 0, 1, 1]) + + +EXPECTED_COLUMNS = [ + "reference_group", + "stratified_by", + "chosen_cutoff", + "true_positives", + "true_negatives", + "false_positives", + "false_negatives", + "predicted_positives", + "predicted_negatives", + "real_positives", + "real_negatives", + "n", + "sensitivity", + "specificity", + "ppv", + "npv", + "false_positive_rate", + "lift", + "net_benefit", + "net_benefit_interventions_avoided", + "ppcr", +] + + +def test_binary_performance_schema_order_is_stable_across_stratification(): + for stratified_by in [ + ("probability_threshold",), + ("ppcr",), + ("probability_threshold", "ppcr"), + ]: + result = prepare_performance_data( + probs=PROBS, + reals=REALS, + stratified_by=stratified_by, + by=0.5, + ) + + assert result.columns == EXPECTED_COLUMNS + + +def test_combined_ppcr_values_match_ppcr_only(): + ppcr_only = prepare_performance_data( + probs=PROBS, + reals=REALS, + stratified_by=("ppcr",), + by=0.5, + ).sort(["reference_group", "chosen_cutoff"]) + + combined_ppcr = ( + prepare_performance_data( + probs=PROBS, + reals=REALS, + stratified_by=("probability_threshold", "ppcr"), + by=0.5, + ) + .filter(pl.col("stratified_by") == "ppcr") + .sort(["reference_group", "chosen_cutoff"]) + ) + + assert_frame_equal(combined_ppcr, ppcr_only)