Fix time-dependent reference prevalence by horizon

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,

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,

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,
+        )
+    )

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,

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))