diff --git a/src/rtichoke/discrimination/gains.py b/src/rtichoke/discrimination/gains.py index 06fea386..51d506f4 100644 --- a/src/rtichoke/discrimination/gains.py +++ b/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,8 +234,7 @@ 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, @@ -194,9 +242,15 @@ def create_gains_curve_times( 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) diff --git a/tests/test_gains_times.py b/tests/test_gains_times.py new file mode 100644 index 00000000..8cb697c6 --- /dev/null +++ b/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)