BREAKING: use population metadata for time references

src/rtichoke/processing/time_reference_lines.py

@@ -9,6 +9,7 @@
 from rtichoke.performance_data.performance_data_times import (
     prepare_performance_data_times,
 )
+from rtichoke.processing.evaluation_semantics import _build_evaluation_metadata
 from rtichoke.processing.plotly_helper_functions import (
     _check_if_multiple_populations_are_being_validated_times,
     _create_plotly_curve_times,
@@ -62,20 +63,24 @@ def _replace_reference_data_times(
     curve: str,
     min_p_threshold: float = 0.0,
     max_p_threshold: float = 1.0,
+    multiple_populations: bool | None = None,
 ) -> dict:
     """Rebuild prevalence-dependent references from cutoff-0 event risk."""
     aj_estimates = _get_reference_aj_estimates_times(performance_data)
     references = []
 
     for horizon in curve_list["fixed_time_horizons"]:
         aj_horizon = aj_estimates.filter(pl.col("fixed_time_horizon") == horizon)
+        horizon_has_multiple_populations = (
+            multiple_populations
+            if multiple_populations is not None
+            else _check_if_multiple_populations_are_being_validated_times(aj_horizon)
+        )
         references.append(
             _create_reference_lines_data(
                 curve=curve,
                 aj_estimates_from_performance_data=aj_horizon,
-                multiple_populations=(
-                    _check_if_multiple_populations_are_being_validated_times(aj_horizon)
-                ),
+                multiple_populations=horizon_has_multiple_populations,
                 min_p_threshold=min_p_threshold,
                 max_p_threshold=max_p_threshold,
             ).with_columns(pl.lit(horizon).alias("fixed_time_horizon"))
@@ -99,7 +104,7 @@ def _create_rtichoke_plotly_curve_times_reference_safe(
     color_values=None,
     curve: str = "precision recall",
 ) -> Figure:
-    """Create a time-dependent curve with corrected reference prevalence."""
+    """Create a time-dependent curve with population-scoped references."""
     performance_data = prepare_performance_data_times(
         probs,
         reals,
@@ -118,12 +123,19 @@ def _create_rtichoke_plotly_curve_times_reference_safe(
         max_p_threshold=max_p_threshold,
     )
     curve_list = _apply_color_values_times(curve_list, color_values)
+
+    evaluation_metadata = _build_evaluation_metadata(probs, reals, times)
+    multiple_populations = (
+        len({metadata.population for metadata in evaluation_metadata.values()}) > 1
+    )
+
     return _create_plotly_curve_times(
         _replace_reference_data_times(
             curve_list,
             performance_data,
             curve=curve,
             min_p_threshold=min_p_threshold,
             max_p_threshold=max_p_threshold,
+            multiple_populations=multiple_populations,
         )
     )

tests/test_time_reference_ownership.py

@@ -0,0 +1,140 @@
+import numpy as np
+import polars as pl
+import pytest
+
+from rtichoke import create_precision_recall_curve_times
+from rtichoke.processing.time_reference_lines import _replace_reference_data_times
+
+HORIZONS = [5.0, 10.0]
+HEURISTICS = [
+    {
+        "censoring_heuristic": "adjusted",
+        "competing_heuristic": "adjusted_as_negative",
+    }
+]
+
+
+def _equal_risk_performance_data() -> pl.DataFrame:
+    return pl.DataFrame(
+        {
+            "reference_group": [
+                "Population A",
+                "Population B",
+                "Population A",
+                "Population B",
+            ],
+            "fixed_time_horizon": [5.0, 5.0, 10.0, 10.0],
+            "chosen_cutoff": [0.0, 0.0, 0.0, 0.0],
+            "real_positives": [2.0, 2.0, 3.0, 3.0],
+            "n": [6.0, 6.0, 6.0, 6.0],
+        }
+    )
+
+
+def _nonempty_random_reference_names(fig) -> set[str]:
+    return {
+        trace.name
+        for trace in fig.data
+        if "random_guess" in trace.name and len(trace.x) > 0
+    }
+
+
+@pytest.mark.parametrize(
+    ("curve", "expected"),
+    [
+        ("roc", {"random_guess"}),
+        (
+            "precision recall",
+            {"random_guess_Population A", "random_guess_Population B"},
+        ),
+        (
+            "gains",
+            {
+                "random_guess",
+                "perfect_model_Population A",
+                "perfect_model_Population B",
+            },
+        ),
+        (
+            "lift",
+            {
+                "random_guess",
+                "perfect_model_Population A",
+                "perfect_model_Population B",
+            },
+        ),
+        (
+            "decision",
+            {"treat_none", "treat_all_Population A", "treat_all_Population B"},
+        ),
+        (
+            "interventions avoided",
+            {"treat_all", "treat_none_Population A", "treat_none_Population B"},
+        ),
+    ],
+)
+def test_equal_risk_populations_keep_population_scoped_references(curve, expected):
+    curve_list = {"fixed_time_horizons": HORIZONS, "reference_data": pl.DataFrame()}
+    reference_data = _replace_reference_data_times(
+        curve_list,
+        _equal_risk_performance_data(),
+        curve=curve,
+        multiple_populations=True,
+    )["reference_data"]
+
+    for horizon in HORIZONS:
+        actual = set(
+            reference_data.filter(pl.col("fixed_time_horizon") == horizon)[
+                "reference_group"
+            ].unique()
+        )
+        assert actual == expected
+
+
+def test_public_equal_risk_populations_keep_distinct_precision_recall_references():
+    probs = {
+        "Population A": np.array([0.05, 0.15, 0.35, 0.55, 0.75, 0.95]),
+        "Population B": np.array([0.10, 0.25, 0.45, 0.65, 0.80, 0.90]),
+    }
+    reals = {
+        "Population A": np.array([1, 1, 0, 0, 0, 0]),
+        "Population B": np.array([1, 1, 0, 0, 0, 0]),
+    }
+    times = {
+        "Population A": np.array([3.0, 8.0, 12.0, 13.0, 14.0, 15.0]),
+        "Population B": np.array([4.0, 9.0, 12.0, 13.0, 14.0, 15.0]),
+    }
+
+    fig = create_precision_recall_curve_times(
+        probs,
+        reals,
+        times,
+        fixed_time_horizons=HORIZONS,
+        heuristics_sets=HEURISTICS,
+        by=0.25,
+    )
+
+    assert _nonempty_random_reference_names(fig) == {
+        "random_guess_Population A",
+        "random_guess_Population B",
+    }
+
+
+def test_public_multiple_models_still_share_population_reference():
+    probs = {
+        "Model A": np.array([0.05, 0.15, 0.35, 0.55, 0.75, 0.95]),
+        "Model B": np.array([0.10, 0.25, 0.45, 0.65, 0.80, 0.90]),
+    }
+    reals = np.array([1, 1, 0, 0, 0, 0])
+    times = np.array([3.0, 8.0, 12.0, 13.0, 14.0, 15.0])
+
+    fig = create_precision_recall_curve_times(
+        probs,
+        reals,
+        times,
+        fixed_time_horizons=HORIZONS,
+        heuristics_sets=HEURISTICS,
+        by=0.25,
+    )
+
+    assert _nonempty_random_reference_names(fig) == {"random_guess"}