diff --git a/src/rtichoke/performance_data/performance_data_times.py b/src/rtichoke/performance_data/performance_data_times.py index 7a00175a..57b77b48 100644 --- a/src/rtichoke/performance_data/performance_data_times.py +++ b/src/rtichoke/performance_data/performance_data_times.py @@ -21,6 +21,35 @@ import numpy as np +_PERFORMANCE_DATA_TIMES_COLUMNS = [ + "reference_group", + "fixed_time_horizon", + "censoring_heuristic", + "competing_heuristic", + "stratified_by", + "chosen_cutoff", + "excluded", + "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 prepare_performance_data_times( probs: Dict[str, np.ndarray], reals: Union[np.ndarray, Dict[str, np.ndarray]], @@ -131,6 +160,7 @@ def prepare_performance_data_times( "_fixed_time_horizon_order", "_heuristic_order", ) + .select(_PERFORMANCE_DATA_TIMES_COLUMNS) ) diff --git a/tests/test_time_performance_schema_order.py b/tests/test_time_performance_schema_order.py new file mode 100644 index 00000000..7a253afd --- /dev/null +++ b/tests/test_time_performance_schema_order.py @@ -0,0 +1,44 @@ +import numpy as np +import polars as pl +from polars.testing import assert_frame_equal + +from rtichoke import prepare_performance_data_times + + +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]) +TIMES = np.array([8.0, 7.0, 3.0, 6.0, 2.0, 1.0]) + + +def _prepare(stratified_by): + return prepare_performance_data_times( + probs=PROBS, + reals=REALS, + times=TIMES, + fixed_time_horizons=[5.0], + stratified_by=stratified_by, + by=0.5, + ) + + +def test_time_performance_schema_order_is_stable_across_stratification(): + threshold_only = _prepare(("probability_threshold",)) + ppcr_only = _prepare(("ppcr",)) + combined = _prepare(("probability_threshold", "ppcr")) + + assert ppcr_only.columns == threshold_only.columns + assert combined.columns == threshold_only.columns + + +def test_combined_time_ppcr_values_match_ppcr_only(): + ppcr_only = _prepare(("ppcr",)).sort( + ["reference_group", "fixed_time_horizon", "chosen_cutoff"] + ) + combined_ppcr = ( + _prepare(("probability_threshold", "ppcr")) + .filter(pl.col("stratified_by") == "ppcr") + .select(ppcr_only.columns) + .sort(["reference_group", "fixed_time_horizon", "chosen_cutoff"]) + ) + + assert_frame_equal(combined_ppcr, ppcr_only)