diff --git a/src/rtichoke/utility/decision.py b/src/rtichoke/utility/decision.py index 7c3c0615..b47f1296 100644 --- a/src/rtichoke/utility/decision.py +++ b/src/rtichoke/utility/decision.py @@ -5,13 +5,11 @@ from typing import Dict, List, Sequence, Union from plotly.graph_objs._figure import Figure from rtichoke.processing.binary_color_values import _apply_color_values_binary -from rtichoke.processing.plotly_helper_functions import ( - _create_rtichoke_plotly_curve_binary, - _plot_rtichoke_curve_binary, -) +from rtichoke.processing.plotly_helper_functions import _plot_rtichoke_curve_binary from rtichoke.processing.time_reference_lines import ( _create_rtichoke_plotly_curve_times_reference_safe, ) +from rtichoke.performance_data.performance_data import prepare_performance_data import numpy as np import polars as pl @@ -90,17 +88,21 @@ def create_decision_curve( else: curve = "interventions avoided" - fig = _create_rtichoke_plotly_curve_binary( - probs, - reals, - by=by, + performance_data = prepare_performance_data( + probs=probs, + reals=reals, stratified_by=stratified_by, - size=size, - color_values=color_values, + by=by, + ) + fig = _plot_rtichoke_curve_binary( + performance_data=performance_data, + stratified_by=stratified_by[0], curve=curve, + size=size, min_p_threshold=min_p_threshold, max_p_threshold=max_p_threshold, ) + fig.update_xaxes(range=[min_p_threshold, max_p_threshold]) return _apply_color_values_binary(fig, color_values) @@ -151,6 +153,7 @@ def plot_decision_curve( min_p_threshold=min_p_threshold, max_p_threshold=max_p_threshold, ) + fig.update_xaxes(range=[min_p_threshold, max_p_threshold]) return fig diff --git a/tests/test_reference_line_audit.py b/tests/test_reference_line_audit.py new file mode 100644 index 00000000..012eec28 --- /dev/null +++ b/tests/test_reference_line_audit.py @@ -0,0 +1,122 @@ +import numpy as np +import pytest + +from rtichoke import ( + create_decision_curve, + create_gains_curve, + create_lift_curve, + create_precision_recall_curve, + create_gains_curve_times, +) + + +def _binary_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]), + } + return probs, reals + + +def _visible_trace(fig, name): + matches = [trace for trace in fig.data if trace.name == name and trace.visible is not False] + assert len(matches) == 1 + return matches[0] + + +def _y_at_x(trace, x): + matches = [float(y) for tx, y in zip(trace.x, trace.y) if float(tx) == pytest.approx(x)] + assert len(matches) == 1 + return matches[0] + + +def test_binary_precision_recall_reference_is_population_specific(): + probs, reals = _binary_inputs() + fig = create_precision_recall_curve(probs, reals, by=0.1) + + a = _visible_trace(fig, "random_guess_population_a") + b = _visible_trace(fig, "random_guess_population_b") + + assert float(a.y[0]) == pytest.approx(3 / 6) + assert float(b.y[0]) == pytest.approx(3 / 4) + + +def test_binary_lift_reference_is_population_specific(): + probs, reals = _binary_inputs() + fig = create_lift_curve(probs, reals, by=0.1) + + a = _visible_trace(fig, "perfect_model_population_a") + b = _visible_trace(fig, "perfect_model_population_b") + + assert float(a.y[0]) == pytest.approx(2.0) + assert float(b.y[0]) == pytest.approx(4 / 3) + + +def test_binary_decision_reference_is_population_specific(): + probs, reals = _binary_inputs() + fig = create_decision_curve( + probs, reals, by=0.1, min_p_threshold=0.1, max_p_threshold=0.9 + ) + + a = _visible_trace(fig, "treat_all_population_a") + b = _visible_trace(fig, "treat_all_population_b") + + x = float(a.x[0]) + assert x == pytest.approx(0.1) + assert float(a.x[-1]) == pytest.approx(0.9) + assert float(b.x[0]) == pytest.approx(x) + assert float(a.y[0]) == pytest.approx(0.5 - 0.5 * x / (1 - x)) + assert float(b.y[0]) == pytest.approx(0.75 - 0.25 * x / (1 - x)) + + +def test_binary_interventions_avoided_honors_displayed_threshold_range(): + probs, reals = _binary_inputs() + fig = create_decision_curve( + probs, + reals, + decision_type="interventions avoided", + by=0.1, + min_p_threshold=0.1, + max_p_threshold=0.9, + ) + + assert tuple(float(x) for x in fig.layout.xaxis.range) == pytest.approx((0.1, 0.9)) + + +def test_binary_gains_reference_is_population_specific(): + probs, reals = _binary_inputs() + fig = create_gains_curve(probs, reals, by=0.1) + + a = _visible_trace(fig, "perfect_model_population_a") + b = _visible_trace(fig, "perfect_model_population_b") + + assert _y_at_x(a, 0.01) == pytest.approx(0.01 / 0.5) + assert _y_at_x(b, 0.01) == pytest.approx(0.01 / 0.75) + + +def test_time_gains_reference_is_population_and_horizon_specific(): + 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]), + } + fig = create_gains_curve_times( + probs, reals, times, fixed_time_horizons=[5.0, 10.0], by=0.1 + ) + + traces = {(t.name, t.visible): t for t in fig.data} + assert _y_at_x(traces[("perfect_model_population_a", True)], 0.01) == pytest.approx(0.06) + assert _y_at_x(traces[("perfect_model_population_b", True)], 0.01) == pytest.approx(0.04) + assert _y_at_x(traces[("perfect_model_population_a", False)], 0.01) == pytest.approx(0.03) + assert _y_at_x(traces[("perfect_model_population_b", False)], 0.01) == pytest.approx(0.02)