Fix time-dependent gains perfect reference

src/rtichoke/discrimination/gains.py

@@ -8,11 +8,60 @@
     _create_rtichoke_plotly_curve_times,
     _create_rtichoke_plotly_curve_binary,
     _plot_rtichoke_curve_binary,
+    _create_rtichoke_curve_list_times,
+    _create_plotly_curve_times,
+    _create_reference_lines_data,
+    _check_if_multiple_populations_are_being_validated_times,
 )
+from rtichoke.performance_data.performance_data_times import prepare_performance_data_times
 import numpy as np
 import polars as pl
 
 
+def _get_gains_aj_estimates_times(performance_data: pl.DataFrame) -> pl.DataFrame:
+    """Return one horizon-specific event estimate per reference group.
+
+    For a gains reference curve the required event probability is the overall
+    event probability at the horizon. At probability threshold 0 everyone is
+    classified positive, so ``real_positives / n`` gives that quantity without
+    mixing in the cutoff-specific estimate at threshold 1.
+    """
+    return (
+        performance_data.filter(pl.col("chosen_cutoff") == 0)
+        .select("reference_group", "fixed_time_horizon", "real_positives", "n")
+        .unique()
+        .with_columns((pl.col("real_positives") / pl.col("n")).alias("aj_estimate"))
+        .select("reference_group", "fixed_time_horizon", "aj_estimate")
+        .sort(["reference_group", "fixed_time_horizon"])
+    )
+
+
+def _replace_gains_reference_data_times(
+    curve_list: dict, performance_data: pl.DataFrame
+) -> dict:
+    """Replace time-dependent gains references with one AJ estimate per horizon."""
+    aj_estimates = _get_gains_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)
+        multiple_populations = _check_if_multiple_populations_are_being_validated_times(
+            aj_horizon
+        )
+        references.append(
+            _create_reference_lines_data(
+                curve="gains",
+                aj_estimates_from_performance_data=aj_horizon,
+                multiple_populations=multiple_populations,
+            ).with_columns(pl.lit(horizon).alias("fixed_time_horizon"))
+        )
+
+    curve_list["reference_data"] = (
+        pl.concat(references, how="vertical") if references else pl.DataFrame()
+    )
+    return curve_list
+
+
 def create_gains_curve(
     probs: Dict[str, np.ndarray],
     reals: Union[np.ndarray, Dict[str, np.ndarray]],
@@ -185,18 +234,23 @@ def create_gains_curve_times(
     Figure
         A Plotly ``Figure`` object for the time-dependent Gains curve.
     """
-
-    fig = _create_rtichoke_plotly_curve_times(
+    performance_data = prepare_performance_data_times(
         probs,
         reals,
         times,
         fixed_time_horizons=fixed_time_horizons,
         heuristics_sets=heuristics_sets,
         by=by,
         stratified_by=stratified_by,
+    )
+
+    curve_list = _create_rtichoke_curve_list_times(
+        performance_data,
+        stratified_by=stratified_by[0],
         size=size,
-        color_values=color_values,
+        color_value=color_values,
         curve="gains",
     )
+    curve_list = _replace_gains_reference_data_times(curve_list, performance_data)
 
-    return fig
+    return _create_plotly_curve_times(curve_list)

tests/test_gains_times.py

@@ -0,0 +1,96 @@
+import polars as pl
+import pytest
+
+from rtichoke.discrimination.gains import (
+    _get_gains_aj_estimates_times,
+    _replace_gains_reference_data_times,
+)
+
+
+def test_gains_reference_uses_cutoff_zero_event_probability():
+    performance_data = pl.DataFrame(
+        {
+            "reference_group": ["model", "model"],
+            "fixed_time_horizon": [5.0, 5.0],
+            "chosen_cutoff": [0.0, 1.0],
+            "real_positives": [20.0, 5.0],
+            "n": [100.0, 100.0],
+        }
+    )
+
+    aj = _get_gains_aj_estimates_times(performance_data)
+
+    assert aj.height == 1
+    assert aj["aj_estimate"].item() == 0.2
+
+    curve_list = {
+        "fixed_time_horizons": [5.0],
+        "reference_data": pl.DataFrame(),
+    }
+    fixed = _replace_gains_reference_data_times(curve_list, performance_data)
+    perfect = fixed["reference_data"].filter(
+        pl.col("reference_group") == "perfect_model"
+    )
+
+    # For p=0.2, the perfect gains curve has y=x/p, so at x=0.1 y=0.5.
+    y_at_point_one = perfect.filter(pl.col("x") == 0.1)["y"].item()
+    assert y_at_point_one == 0.5
+
+
+def test_gains_reference_is_population_and_horizon_specific():
+    performance_data = pl.DataFrame(
+        {
+            "reference_group": [
+                "population_a",
+                "population_a",
+                "population_b",
+                "population_b",
+                "population_a",
+                "population_a",
+                "population_b",
+                "population_b",
+            ],
+            "fixed_time_horizon": [5.0, 5.0, 5.0, 5.0, 10.0, 10.0, 10.0, 10.0],
+            "chosen_cutoff": [0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0],
+            "real_positives": [20.0, 5.0, 40.0, 10.0, 30.0, 7.0, 60.0, 15.0],
+            "n": [100.0] * 8,
+        }
+    )
+
+    aj = _get_gains_aj_estimates_times(performance_data)
+    assert aj.height == 4
+
+    expected_aj = {
+        ("population_a", 5.0): 0.2,
+        ("population_b", 5.0): 0.4,
+        ("population_a", 10.0): 0.3,
+        ("population_b", 10.0): 0.6,
+    }
+    for (population, horizon), expected in expected_aj.items():
+        value = aj.filter(
+            (pl.col("reference_group") == population)
+            & (pl.col("fixed_time_horizon") == horizon)
+        )["aj_estimate"].item()
+        assert value == pytest.approx(expected)
+
+    curve_list = {
+        "fixed_time_horizons": [5.0, 10.0],
+        "reference_data": pl.DataFrame(),
+    }
+    reference_data = _replace_gains_reference_data_times(
+        curve_list, performance_data
+    )["reference_data"]
+
+    expected_y_at_point_one = {
+        ("perfect_model_population_a", 5.0): 0.1 / 0.2,
+        ("perfect_model_population_b", 5.0): 0.1 / 0.4,
+        ("perfect_model_population_a", 10.0): 0.1 / 0.3,
+        ("perfect_model_population_b", 10.0): 0.1 / 0.6,
+    }
+    for (reference_group, horizon), expected in expected_y_at_point_one.items():
+        y = reference_data.filter(
+            (pl.col("reference_group") == reference_group)
+            & (pl.col("fixed_time_horizon") == horizon)
+            & (pl.col("x") == 0.1)
+        )["y"].item()
+        assert y == pytest.approx(expected)