Order time-dependent performance data for cross-group comparison

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]

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