diff --git a/src/rtichoke/processing/time_reference_lines.py b/src/rtichoke/processing/time_reference_lines.py index 8774aa78..e5a629a0 100644 --- a/src/rtichoke/processing/time_reference_lines.py +++ b/src/rtichoke/processing/time_reference_lines.py @@ -9,6 +9,7 @@ from rtichoke.performance_data.performance_data_times import ( prepare_performance_data_times, ) +from rtichoke.processing.evaluation_semantics import _build_evaluation_metadata from rtichoke.processing.plotly_helper_functions import ( _check_if_multiple_populations_are_being_validated_times, _create_plotly_curve_times, @@ -62,6 +63,7 @@ def _replace_reference_data_times( curve: str, min_p_threshold: float = 0.0, max_p_threshold: float = 1.0, + multiple_populations: bool | None = None, ) -> dict: """Rebuild prevalence-dependent references from cutoff-0 event risk.""" aj_estimates = _get_reference_aj_estimates_times(performance_data) @@ -69,13 +71,16 @@ def _replace_reference_data_times( for horizon in curve_list["fixed_time_horizons"]: aj_horizon = aj_estimates.filter(pl.col("fixed_time_horizon") == horizon) + horizon_has_multiple_populations = ( + multiple_populations + if multiple_populations is not None + else _check_if_multiple_populations_are_being_validated_times(aj_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) - ), + multiple_populations=horizon_has_multiple_populations, min_p_threshold=min_p_threshold, max_p_threshold=max_p_threshold, ).with_columns(pl.lit(horizon).alias("fixed_time_horizon")) @@ -99,7 +104,7 @@ def _create_rtichoke_plotly_curve_times_reference_safe( color_values=None, curve: str = "precision recall", ) -> Figure: - """Create a time-dependent curve with corrected reference prevalence.""" + """Create a time-dependent curve with population-scoped references.""" performance_data = prepare_performance_data_times( probs, reals, @@ -118,6 +123,12 @@ def _create_rtichoke_plotly_curve_times_reference_safe( max_p_threshold=max_p_threshold, ) curve_list = _apply_color_values_times(curve_list, color_values) + + evaluation_metadata = _build_evaluation_metadata(probs, reals, times) + multiple_populations = ( + len({metadata.population for metadata in evaluation_metadata.values()}) > 1 + ) + return _create_plotly_curve_times( _replace_reference_data_times( curve_list, @@ -125,5 +136,6 @@ def _create_rtichoke_plotly_curve_times_reference_safe( curve=curve, min_p_threshold=min_p_threshold, max_p_threshold=max_p_threshold, + multiple_populations=multiple_populations, ) ) diff --git a/tests/test_time_reference_ownership.py b/tests/test_time_reference_ownership.py new file mode 100644 index 00000000..8076dfc2 --- /dev/null +++ b/tests/test_time_reference_ownership.py @@ -0,0 +1,140 @@ +import numpy as np +import polars as pl +import pytest + +from rtichoke import create_precision_recall_curve_times +from rtichoke.processing.time_reference_lines import _replace_reference_data_times + +HORIZONS = [5.0, 10.0] +HEURISTICS = [ + { + "censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative", + } +] + + +def _equal_risk_performance_data() -> pl.DataFrame: + return pl.DataFrame( + { + "reference_group": [ + "Population A", + "Population B", + "Population A", + "Population B", + ], + "fixed_time_horizon": [5.0, 5.0, 10.0, 10.0], + "chosen_cutoff": [0.0, 0.0, 0.0, 0.0], + "real_positives": [2.0, 2.0, 3.0, 3.0], + "n": [6.0, 6.0, 6.0, 6.0], + } + ) + + +def _nonempty_random_reference_names(fig) -> set[str]: + return { + trace.name + for trace in fig.data + if "random_guess" in trace.name and len(trace.x) > 0 + } + + +@pytest.mark.parametrize( + ("curve", "expected"), + [ + ("roc", {"random_guess"}), + ( + "precision recall", + {"random_guess_Population A", "random_guess_Population B"}, + ), + ( + "gains", + { + "random_guess", + "perfect_model_Population A", + "perfect_model_Population B", + }, + ), + ( + "lift", + { + "random_guess", + "perfect_model_Population A", + "perfect_model_Population B", + }, + ), + ( + "decision", + {"treat_none", "treat_all_Population A", "treat_all_Population B"}, + ), + ( + "interventions avoided", + {"treat_all", "treat_none_Population A", "treat_none_Population B"}, + ), + ], +) +def test_equal_risk_populations_keep_population_scoped_references(curve, expected): + curve_list = {"fixed_time_horizons": HORIZONS, "reference_data": pl.DataFrame()} + reference_data = _replace_reference_data_times( + curve_list, + _equal_risk_performance_data(), + curve=curve, + multiple_populations=True, + )["reference_data"] + + for horizon in HORIZONS: + actual = set( + reference_data.filter(pl.col("fixed_time_horizon") == horizon)[ + "reference_group" + ].unique() + ) + assert actual == expected + + +def test_public_equal_risk_populations_keep_distinct_precision_recall_references(): + probs = { + "Population A": np.array([0.05, 0.15, 0.35, 0.55, 0.75, 0.95]), + "Population B": np.array([0.10, 0.25, 0.45, 0.65, 0.80, 0.90]), + } + reals = { + "Population A": np.array([1, 1, 0, 0, 0, 0]), + "Population B": np.array([1, 1, 0, 0, 0, 0]), + } + times = { + "Population A": np.array([3.0, 8.0, 12.0, 13.0, 14.0, 15.0]), + "Population B": np.array([4.0, 9.0, 12.0, 13.0, 14.0, 15.0]), + } + + fig = create_precision_recall_curve_times( + probs, + reals, + times, + fixed_time_horizons=HORIZONS, + heuristics_sets=HEURISTICS, + by=0.25, + ) + + assert _nonempty_random_reference_names(fig) == { + "random_guess_Population A", + "random_guess_Population B", + } + + +def test_public_multiple_models_still_share_population_reference(): + probs = { + "Model A": np.array([0.05, 0.15, 0.35, 0.55, 0.75, 0.95]), + "Model B": np.array([0.10, 0.25, 0.45, 0.65, 0.80, 0.90]), + } + reals = np.array([1, 1, 0, 0, 0, 0]) + times = np.array([3.0, 8.0, 12.0, 13.0, 14.0, 15.0]) + + fig = create_precision_recall_curve_times( + probs, + reals, + times, + fixed_time_horizons=HORIZONS, + heuristics_sets=HEURISTICS, + by=0.25, + ) + + assert _nonempty_random_reference_names(fig) == {"random_guess"}