diff --git a/src/rtichoke/performance_data/performance_data_times.py b/src/rtichoke/performance_data/performance_data_times.py index 53678155..d3177406 100644 --- a/src/rtichoke/performance_data/performance_data_times.py +++ b/src/rtichoke/performance_data/performance_data_times.py @@ -79,7 +79,6 @@ def prepare_performance_data_times( thresholds and time horizons. It includes columns for cutoffs, time points, heuristics, and performance measures. """ - # 1. Get the underlying binned time-dependent classification data final_adjusted_data = prepare_binned_classification_data_times( probs=probs, reals=reals, @@ -91,13 +90,47 @@ def prepare_performance_data_times( risk_set_scope=["pooled_by_cutoff"], ) - # 2. Apply AJ cumulative machinery cumulative_aj_data = _calculate_cumulative_aj_data(final_adjusted_data) - - # 3. Turn AJ output into performance metrics performance_data = _turn_cumulative_aj_to_performance_data(cumulative_aj_data) - return performance_data + group_order = {group: index for index, group in enumerate(probs)} + horizon_order = { + float(horizon): index for index, horizon in enumerate(fixed_time_horizons) + } + heuristic_order = { + f'{heuristics["censoring_heuristic"]}\x1f{heuristics["competing_heuristic"]}': index + for index, heuristics in enumerate(heuristics_sets) + } + + return ( + performance_data.with_columns( + pl.col("reference_group") + .replace_strict(group_order, default=len(group_order)) + .alias("_reference_group_order"), + pl.col("fixed_time_horizon") + .replace_strict(horizon_order, default=len(horizon_order)) + .alias("_fixed_time_horizon_order"), + pl.concat_str( + ["censoring_heuristic", "competing_heuristic"], separator="\x1f" + ) + .replace_strict(heuristic_order, default=len(heuristic_order)) + .alias("_heuristic_order"), + ) + .sort( + [ + "_fixed_time_horizon_order", + "_heuristic_order", + "stratified_by", + "chosen_cutoff", + "_reference_group_order", + ] + ) + .drop( + "_reference_group_order", + "_fixed_time_horizon_order", + "_heuristic_order", + ) + ) def prepare_binned_classification_data_times( @@ -131,7 +164,7 @@ def prepare_binned_classification_data_times( A dictionary mapping model or dataset names (str) to their predicted probabilities. reals : Union[np.ndarray, Dict[str, np.ndarray]] - The true event statuses (e.g., 0=censored, 1=event, 2=competing). + The true event statuses (e.g., 0=censored, 1=event, 2=competing event). times : Union[np.ndarray, Dict[str, np.ndarray]] The event or censoring times. fixed_time_horizons : list[float] diff --git a/tests/test_time_performance_order.py b/tests/test_time_performance_order.py new file mode 100644 index 00000000..1e877731 --- /dev/null +++ b/tests/test_time_performance_order.py @@ -0,0 +1,57 @@ +import numpy as np + +from rtichoke import prepare_performance_data_times + + +def test_prepare_performance_data_times_groups_reference_groups_for_comparison(): + probs = { + "population_b": np.array([0.8, 0.2, 0.6, 0.4]), + "population_a": np.array([0.7, 0.1, 0.9, 0.3]), + } + reals = { + "population_b": np.array([1, 0, 1, 0]), + "population_a": np.array([0, 1, 1, 0]), + } + times = { + "population_b": np.array([2.0, 7.0, 4.0, 9.0]), + "population_a": np.array([8.0, 3.0, 5.0, 10.0]), + } + horizons = [8.0, 5.0] + heuristics_sets = [ + { + "censoring_heuristic": "excluded", + "competing_heuristic": "excluded", + }, + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + }, + ] + + result = prepare_performance_data_times( + probs, + reals, + times, + fixed_time_horizons=horizons, + heuristics_sets=heuristics_sets, + by=0.2, + ) + + comparison_columns = [ + "fixed_time_horizon", + "censoring_heuristic", + "competing_heuristic", + "stratified_by", + "chosen_cutoff", + ] + expected_group_order = list(probs) + + for block in result.partition_by(comparison_columns, maintain_order=True): + assert block["reference_group"].to_list() == expected_group_order + + observed_horizons = ( + result.select("fixed_time_horizon") + .unique(maintain_order=True)["fixed_time_horizon"] + .to_list() + ) + assert observed_horizons == horizons