diff --git a/src/rtichoke/discrimination/lift.py b/src/rtichoke/discrimination/lift.py index a8460b12..4efc3256 100644 --- a/src/rtichoke/discrimination/lift.py +++ b/src/rtichoke/discrimination/lift.py @@ -5,10 +5,12 @@ from typing import Dict, List, Sequence, Union from plotly.graph_objs._figure import Figure from rtichoke.processing.plotly_helper_functions import ( - _create_rtichoke_plotly_curve_times, _create_rtichoke_plotly_curve_binary, _plot_rtichoke_curve_binary, ) +from rtichoke.processing.time_reference_lines import ( + _create_rtichoke_plotly_curve_times_reference_safe, +) import numpy as np import polars as pl @@ -186,7 +188,7 @@ def create_lift_curve_times( A Plotly ``Figure`` object for the time-dependent Lift curve. """ - fig = _create_rtichoke_plotly_curve_times( + fig = _create_rtichoke_plotly_curve_times_reference_safe( probs, reals, times, diff --git a/src/rtichoke/discrimination/precision_recall.py b/src/rtichoke/discrimination/precision_recall.py index 7770cbdc..83ccf695 100644 --- a/src/rtichoke/discrimination/precision_recall.py +++ b/src/rtichoke/discrimination/precision_recall.py @@ -5,10 +5,12 @@ from typing import Dict, List, Sequence, Union from plotly.graph_objs._figure import Figure from rtichoke.processing.plotly_helper_functions import ( - _create_rtichoke_plotly_curve_times, _create_rtichoke_plotly_curve_binary, _plot_rtichoke_curve_binary, ) +from rtichoke.processing.time_reference_lines import ( + _create_rtichoke_plotly_curve_times_reference_safe, +) import numpy as np import polars as pl @@ -187,7 +189,7 @@ def create_precision_recall_curve_times( A Plotly ``Figure`` object for the time-dependent Precision-Recall curve. """ - fig = _create_rtichoke_plotly_curve_times( + fig = _create_rtichoke_plotly_curve_times_reference_safe( probs, reals, times, diff --git a/src/rtichoke/processing/time_reference_lines.py b/src/rtichoke/processing/time_reference_lines.py new file mode 100644 index 00000000..3b56e496 --- /dev/null +++ b/src/rtichoke/processing/time_reference_lines.py @@ -0,0 +1,98 @@ +"""Helpers for horizon-specific time-dependent reference curves.""" + +from typing import Dict, Sequence, Union + +import numpy as np +import polars as pl +from plotly.graph_objs._figure import Figure + +from rtichoke.performance_data.performance_data_times import prepare_performance_data_times +from rtichoke.processing.plotly_helper_functions import ( + _check_if_multiple_populations_are_being_validated_times, + _create_plotly_curve_times, + _create_reference_lines_data, + _create_rtichoke_curve_list_times, +) + + +def _get_reference_aj_estimates_times(performance_data: pl.DataFrame) -> pl.DataFrame: + """Return the cutoff-0 event risk for each group and horizon.""" + 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_reference_data_times( + curve_list: dict, + performance_data: pl.DataFrame, + curve: str, + min_p_threshold: float = 0.0, + max_p_threshold: float = 1.0, +) -> 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) + 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) + ), + min_p_threshold=min_p_threshold, + max_p_threshold=max_p_threshold, + ).with_columns(pl.lit(horizon).alias("fixed_time_horizon")) + ) + + curve_list["reference_data"] = pl.concat(references, how="vertical") + return curve_list + + +def _create_rtichoke_plotly_curve_times_reference_safe( + probs: Dict[str, np.ndarray], + reals: Union[np.ndarray, Dict[str, np.ndarray]], + times: Union[np.ndarray, Dict[str, np.ndarray]], + fixed_time_horizons: list[float], + heuristics_sets: list[Dict], + min_p_threshold: float = 0, + max_p_threshold: float = 1, + by: float = 0.01, + stratified_by: Sequence[str] = ("probability_threshold",), + size: int = 600, + color_values=None, + curve: str = "precision recall", +) -> Figure: + """Create a time-dependent curve with corrected reference prevalence.""" + performance_data = prepare_performance_data_times( + probs, + reals, + times, + by=by, + fixed_time_horizons=fixed_time_horizons, + heuristics_sets=heuristics_sets, + stratified_by=stratified_by, + ) + curve_list = _create_rtichoke_curve_list_times( + performance_data, + stratified_by=stratified_by[0], + curve=curve, + min_p_threshold=min_p_threshold, + max_p_threshold=max_p_threshold, + ) + 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, + ) + ) diff --git a/src/rtichoke/utility/decision.py b/src/rtichoke/utility/decision.py index 12a0592a..ab218b74 100644 --- a/src/rtichoke/utility/decision.py +++ b/src/rtichoke/utility/decision.py @@ -6,9 +6,11 @@ from plotly.graph_objs._figure import Figure from rtichoke.processing.plotly_helper_functions import ( _create_rtichoke_plotly_curve_binary, - _create_rtichoke_plotly_curve_times, _plot_rtichoke_curve_binary, ) +from rtichoke.processing.time_reference_lines import ( + _create_rtichoke_plotly_curve_times_reference_safe, +) import numpy as np import polars as pl @@ -234,7 +236,7 @@ def create_decision_curve_times( else: curve = "interventions avoided" - fig = _create_rtichoke_plotly_curve_times( + fig = _create_rtichoke_plotly_curve_times_reference_safe( probs, reals, times, diff --git a/tests/test_time_reference_lines.py b/tests/test_time_reference_lines.py new file mode 100644 index 00000000..ad5485f5 --- /dev/null +++ b/tests/test_time_reference_lines.py @@ -0,0 +1,199 @@ +import numpy as np +import polars as pl +import pytest + +from rtichoke import ( + create_decision_curve_times, + create_lift_curve_times, + create_precision_recall_curve_times, +) +from rtichoke.processing.time_reference_lines import ( + _get_reference_aj_estimates_times, + _replace_reference_data_times, +) + + +def _performance_data_with_boundary_drift() -> pl.DataFrame: + return 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, + } + ) + + +def _public_curve_inputs(): + probs = { + "population_a": np.array([0.05, 0.15, 0.35, 0.55, 0.75, 0.95]), + "population_b": np.array([0.10, 0.30, 0.50, 0.70]), + } + reals = { + "population_a": np.array([1, 0, 1, 0, 1, 0]), + "population_b": np.array([1, 1, 0, 1]), + } + times = { + "population_a": np.array([3.0, 11.0, 7.0, 12.0, 13.0, 14.0]), + "population_b": np.array([4.0, 8.0, 11.0, 12.0]), + } + return probs, reals, times + + +def _trace(fig, name: str, visible: bool): + matches = [trace for trace in fig.data if trace.name == name and trace.visible is visible] + assert len(matches) == 1 + return matches[0] + + +def test_reference_prevalence_uses_cutoff_zero_per_population_and_horizon(): + aj = _get_reference_aj_estimates_times(_performance_data_with_boundary_drift()) + + assert aj.height == 4 + expected = { + ("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), value in expected.items(): + actual = aj.filter( + (pl.col("reference_group") == population) + & (pl.col("fixed_time_horizon") == horizon) + )["aj_estimate"].item() + assert actual == pytest.approx(value) + + +@pytest.mark.parametrize( + ("curve", "reference_group", "x", "expected"), + [ + ("precision recall", "random_guess_population_a", 0.1, 0.2), + ("precision recall", "random_guess_population_b", 0.1, 0.4), + ("lift", "perfect_model_population_a", 0.1, 5.0), + ("lift", "perfect_model_population_b", 0.1, 2.5), + ("decision", "treat_all_population_a", 0.1, 0.11111111111111112), + ("decision", "treat_all_population_b", 0.1, 0.33333333333333337), + ], +) +def test_prevalence_dependent_references_are_population_specific_at_horizon_five( + curve, reference_group, x, expected +): + curve_list = { + "fixed_time_horizons": [5.0, 10.0], + "reference_data": pl.DataFrame(), + } + reference_data = _replace_reference_data_times( + curve_list, + _performance_data_with_boundary_drift(), + curve=curve, + )["reference_data"] + + y = reference_data.filter( + (pl.col("reference_group") == reference_group) + & (pl.col("fixed_time_horizon") == 5.0) + & (pl.col("x") == x) + )["y"].item() + + assert y == pytest.approx(expected) + + +def test_precision_recall_reference_changes_with_horizon(): + curve_list = { + "fixed_time_horizons": [5.0, 10.0], + "reference_data": pl.DataFrame(), + } + reference_data = _replace_reference_data_times( + curve_list, + _performance_data_with_boundary_drift(), + curve="precision recall", + )["reference_data"] + + p5 = reference_data.filter( + (pl.col("reference_group") == "random_guess_population_a") + & (pl.col("fixed_time_horizon") == 5.0) + & (pl.col("x") == 0.1) + )["y"].item() + p10 = reference_data.filter( + (pl.col("reference_group") == "random_guess_population_a") + & (pl.col("fixed_time_horizon") == 10.0) + & (pl.col("x") == 0.1) + )["y"].item() + + assert p5 == pytest.approx(0.2) + assert p10 == pytest.approx(0.3) + + +def test_public_precision_recall_references_are_population_and_horizon_specific(): + probs, reals, times = _public_curve_inputs() + fig = create_precision_recall_curve_times( + probs, reals, times, fixed_time_horizons=[5.0, 10.0], by=0.1 + ) + + a5 = _trace(fig, "random_guess_population_a", True) + b5 = _trace(fig, "random_guess_population_b", True) + a10 = _trace(fig, "random_guess_population_a", False) + b10 = _trace(fig, "random_guess_population_b", False) + + assert float(a5.y[0]) == pytest.approx(1 / 6) + assert float(b5.y[0]) == pytest.approx(1 / 4) + assert float(a10.y[0]) == pytest.approx(2 / 6) + assert float(b10.y[0]) == pytest.approx(2 / 4) + + +def test_public_lift_references_are_population_and_horizon_specific(): + probs, reals, times = _public_curve_inputs() + fig = create_lift_curve_times( + probs, reals, times, fixed_time_horizons=[5.0, 10.0], by=0.1 + ) + + a5 = _trace(fig, "perfect_model_population_a", True) + b5 = _trace(fig, "perfect_model_population_b", True) + a10 = _trace(fig, "perfect_model_population_a", False) + b10 = _trace(fig, "perfect_model_population_b", False) + + assert float(a5.y[0]) == pytest.approx(6.0) + assert float(b5.y[0]) == pytest.approx(4.0) + assert float(a10.y[0]) == pytest.approx(3.0) + assert float(b10.y[0]) == pytest.approx(2.0) + + +def test_public_decision_references_are_population_and_horizon_specific(): + probs, reals, times = _public_curve_inputs() + fig = create_decision_curve_times( + probs, + reals, + times, + fixed_time_horizons=[5.0, 10.0], + by=0.1, + min_p_threshold=0.1, + max_p_threshold=0.9, + ) + + a5 = _trace(fig, "treat_all_population_a", True) + b5 = _trace(fig, "treat_all_population_b", True) + a10 = _trace(fig, "treat_all_population_a", False) + b10 = _trace(fig, "treat_all_population_b", False) + + assert float(a5.x[0]) == pytest.approx(0.1) + assert float(b5.x[0]) == pytest.approx(0.1) + assert float(a10.x[0]) == pytest.approx(0.1) + assert float(b10.x[0]) == pytest.approx(0.1) + + def treat_all(p): + x = 0.1 + return p - (1 - p) * x / (1 - x) + + assert float(a5.y[0]) == pytest.approx(treat_all(1 / 6)) + assert float(b5.y[0]) == pytest.approx(treat_all(1 / 4)) + assert float(a10.y[0]) == pytest.approx(treat_all(2 / 6)) + assert float(b10.y[0]) == pytest.approx(treat_all(2 / 4))