From 618b7345f6ba7295207e6b45aa37c7e68784ec9c Mon Sep 17 00:00:00 2001 From: Uriah Finkel Date: Thu, 20 Aug 2026 12:27:33 +0300 Subject: [PATCH 1/4] Add time input alignment validation helper --- .../processing/time_input_validation.py | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 src/rtichoke/processing/time_input_validation.py diff --git a/src/rtichoke/processing/time_input_validation.py b/src/rtichoke/processing/time_input_validation.py new file mode 100644 index 00000000..32fd4c4c --- /dev/null +++ b/src/rtichoke/processing/time_input_validation.py @@ -0,0 +1,82 @@ +"""Validation helpers for time-dependent performance inputs.""" + +from typing import Dict, Union + +import numpy as np + + +def _validate_time_input_alignment( + probs: Dict[str, np.ndarray], + reals: Union[np.ndarray, Dict[str, np.ndarray]], + times: Union[np.ndarray, Dict[str, np.ndarray]], +) -> None: + """Validate supported array/dict layouts before time-dependent processing.""" + if not isinstance(probs, dict) or not probs: + raise ValueError("`probs` must be a non-empty dictionary of probability arrays.") + + groups = list(probs) + multiple_groups = len(groups) > 1 + reals_is_dict = isinstance(reals, dict) + times_is_dict = isinstance(times, dict) + + if multiple_groups and reals_is_dict != times_is_dict: + raise ValueError( + "For multiple groups, `reals` and `times` must both be arrays or both be dictionaries." + ) + + if multiple_groups and reals_is_dict: + expected_keys = set(groups) + if set(reals) != expected_keys or set(times) != expected_keys: + raise ValueError( + "For multiple populations, `reals` and `times` dictionary keys must exactly match `probs`." + ) + for group in groups: + n_probs = len(np.asarray(probs[group])) + n_reals = len(np.asarray(reals[group])) + n_times = len(np.asarray(times[group])) + if n_probs != n_reals or n_probs != n_times: + raise ValueError( + f"Input lengths must match within group {group!r}: " + f"len(probs)={n_probs}, len(reals)={n_reals}, len(times)={n_times}." + ) + return + + if multiple_groups: + n_reals = len(np.asarray(reals)) + n_times = len(np.asarray(times)) + if n_reals != n_times: + raise ValueError( + "For multiple models sharing outcomes, `reals` and `times` must have the same length." + ) + for group in groups: + n_probs = len(np.asarray(probs[group])) + if n_probs != n_reals: + raise ValueError( + f"Shared outcome length must match probabilities for group {group!r}: " + f"len(probs)={n_probs}, len(reals)={n_reals}, len(times)={n_times}." + ) + return + + group = groups[0] + if reals_is_dict: + if group not in reals: + raise ValueError(f"`reals` is missing the key {group!r} required by `probs`.") + reals_values = reals[group] + else: + reals_values = reals + + if times_is_dict: + if group not in times: + raise ValueError(f"`times` is missing the key {group!r} required by `probs`.") + times_values = times[group] + else: + times_values = times + + n_probs = len(np.asarray(probs[group])) + n_reals = len(np.asarray(reals_values)) + n_times = len(np.asarray(times_values)) + if n_probs != n_reals or n_probs != n_times: + raise ValueError( + f"Input lengths must match for group {group!r}: " + f"len(probs)={n_probs}, len(reals)={n_reals}, len(times)={n_times}." + ) From 6eab24782e86bf2f76a22bc46359370f5f35339d Mon Sep 17 00:00:00 2001 From: Uriah Finkel Date: Thu, 20 Aug 2026 12:27:59 +0300 Subject: [PATCH 2/4] Validate time-dependent input alignment early --- .../performance_data_times.py | 99 ++----------------- 1 file changed, 9 insertions(+), 90 deletions(-) diff --git a/src/rtichoke/performance_data/performance_data_times.py b/src/rtichoke/performance_data/performance_data_times.py index d3177406..a0170047 100644 --- a/src/rtichoke/performance_data/performance_data_times.py +++ b/src/rtichoke/performance_data/performance_data_times.py @@ -3,13 +3,17 @@ """ from typing import Dict, Union -import polars as pl from collections.abc import Sequence + +import numpy as np +import polars as pl + from rtichoke.processing.adjustments import create_adjusted_data from rtichoke.processing.combinations import ( create_aj_data_combinations, create_breaks_values, ) +from rtichoke.processing.time_input_validation import _validate_time_input_alignment from rtichoke.processing.transforms import ( _calculate_cumulative_aj_data, _create_list_data_to_adjust, @@ -17,8 +21,6 @@ cast_and_join_adjusted_data, ) -import numpy as np - def prepare_performance_data_times( probs: Dict[str, np.ndarray], @@ -34,51 +36,7 @@ def prepare_performance_data_times( stratified_by: Sequence[str] = ("probability_threshold",), by: float = 0.01, ) -> pl.DataFrame: - """Prepare performance data for models with time-to-event outcomes. - - This function calculates a comprehensive set of performance metrics for - models predicting time-to-event outcomes. It handles censored data and - competing events by applying specified heuristics at different time - horizons. The function first bins the data using - `prepare_binned_classification_data_times` and then computes cumulative, - Aalen-Johansen-based performance metrics. - - The resulting dataframe is the primary input for time-dependent plotting - functions. - - Parameters - ---------- - probs : Dict[str, np.ndarray] - A dictionary mapping model or dataset names (str) to their predicted - probabilities of an event occurring by a given time. - reals : Union[np.ndarray, Dict[str, np.ndarray]] - The true event statuses. Can be a single array or a dictionary. - Labels should be integers indicating the outcome (e.g., 0=censored, - 1=event of interest, 2=competing event). - times : Union[np.ndarray, Dict[str, np.ndarray]] - The event or censoring times corresponding to the `reals`. Can be a - single array or a dictionary. - fixed_time_horizons : list[float] - A list of numeric time points at which to evaluate the model's - performance. Integer inputs are accepted and normalized to floats. - heuristics_sets : list[Dict], optional - A list of dictionaries, each specifying how to handle censored data - and competing events. The default is - ``[{"censoring_heuristic": "adjusted", - "competing_heuristic": "adjusted_as_negative"}]``. - stratified_by : Sequence[str], optional - Variables by which to stratify the analysis. Defaults to - ``("probability_threshold",)``. - by : float, optional - The step size for probability thresholds. Defaults to ``0.01``. - - Returns - ------- - pl.DataFrame - A Polars DataFrame with performance metrics computed across probability - thresholds and time horizons. It includes columns for cutoffs, time - points, heuristics, and performance measures. - """ + """Prepare performance data for models with time-to-event outcomes.""" final_adjusted_data = prepare_binned_classification_data_times( probs=probs, reals=reals, @@ -148,45 +106,8 @@ def prepare_binned_classification_data_times( by: float = 0.01, risk_set_scope: Sequence[str] = ["pooled_by_cutoff", "within_stratum"], ) -> pl.DataFrame: - """ - Prepare binned, time-dependent classification data. - - This function constructs the foundational binned data needed for - time-to-event performance analysis. It bins predictions by probability - thresholds, applies censoring and competing event heuristics, and stratifies - the data across specified time horizons. The output is a detailed breakdown - of outcomes within each bin, which can be used for calibration or passed to - `prepare_performance_data_times` for full performance metric calculation. - - Parameters - ---------- - probs : Dict[str, np.ndarray] - A dictionary mapping model or dataset names (str) to their predicted - probabilities. - reals : Union[np.ndarray, Dict[str, np.ndarray]] - The true event statuses (e.g., 0=censored, 1=event, 2=competing event). - times : Union[np.ndarray, Dict[str, np.ndarray]] - The event or censoring times. - fixed_time_horizons : list[float] - A list of numeric time points for performance evaluation. Integer - inputs are accepted and normalized to floats. - heuristics_sets : list[Dict], optional - Specifies how to handle censored data and competing events. - stratified_by : Sequence[str], optional - Variables for stratification. Defaults to ``("probability_threshold",)``. - by : float, optional - The step size for probability thresholds. Defaults to ``0.01``. - risk_set_scope : Sequence[str], optional - Defines the scope for risk set calculations. Defaults to - ``["pooled_by_cutoff", "within_stratum"]``. - - Returns - ------- - pl.DataFrame - A Polars DataFrame with binned, time-dependent data. Each row - represents a unique combination of dataset, bin, time horizon, - heuristic, and other strata. - """ + """Prepare binned, time-dependent classification data.""" + _validate_time_input_alignment(probs=probs, reals=reals, times=times) fixed_time_horizons = [float(horizon) for horizon in fixed_time_horizons] breaks = create_breaks_values(None, "probability_threshold", by) @@ -219,9 +140,7 @@ def prepare_binned_classification_data_times( risk_set_scope=risk_set_scope, ) - final_adjusted_data = cast_and_join_adjusted_data( + return cast_and_join_adjusted_data( aj_data_combinations, adjusted_data, ).with_columns(pl.col("reals_estimate").fill_null(0.0)) - - return final_adjusted_data From 18b2049b68c541bbbee27b3d45748dfbed854a0e Mon Sep 17 00:00:00 2001 From: Uriah Finkel Date: Thu, 20 Aug 2026 12:28:16 +0300 Subject: [PATCH 3/4] Add regression tests for time input validation --- tests/test_time_input_validation.py | 65 +++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 tests/test_time_input_validation.py diff --git a/tests/test_time_input_validation.py b/tests/test_time_input_validation.py new file mode 100644 index 00000000..3b34bf57 --- /dev/null +++ b/tests/test_time_input_validation.py @@ -0,0 +1,65 @@ +import numpy as np +import pytest + +from rtichoke import prepare_performance_data_times + + +PROBS = { + "train": np.array([0.1, 0.4, 0.7, 0.9]), + "test": np.array([0.2, 0.3, 0.6, 0.8]), +} +REALS = { + "train": np.array([0, 1, 0, 1]), + "test": np.array([0, 0, 1, 1]), +} +TIMES = { + "train": np.array([8.0, 2.0, 7.0, 3.0]), + "test": np.array([9.0, 8.0, 4.0, 2.0]), +} + + +def _call(probs=PROBS, reals=REALS, times=TIMES): + return prepare_performance_data_times( + probs=probs, + reals=reals, + times=times, + fixed_time_horizons=[5.0], + by=0.5, + ) + + +def test_multiple_population_keys_must_match_probs(): + bad_reals = {"train": REALS["train"], "validation": REALS["test"]} + + with pytest.raises(ValueError, match="keys must exactly match `probs`"): + _call(reals=bad_reals) + + +def test_multiple_groups_reject_mixed_dict_and_array_outcomes(): + with pytest.raises(ValueError, match="must both be arrays or both be dictionaries"): + _call(reals=REALS, times=TIMES["train"]) + + +def test_multiple_population_lengths_must_match_within_group(): + bad_times = {**TIMES, "test": TIMES["test"][:-1]} + + with pytest.raises(ValueError, match="Input lengths must match within group 'test'"): + _call(times=bad_times) + + +def test_multiple_models_can_share_array_outcomes(): + shared_reals = np.array([0, 1, 0, 1]) + shared_times = np.array([8.0, 2.0, 7.0, 3.0]) + + result = _call(reals=shared_reals, times=shared_times) + + assert set(result["reference_group"].to_list()) == {"train", "test"} + + +def test_single_group_accepts_keyed_reals_with_array_times(): + probs = {"train": PROBS["train"]} + reals = {"train": REALS["train"]} + + result = _call(probs=probs, reals=reals, times=TIMES["train"]) + + assert result.height > 0 From 1bb35d300c39b080d3e7d04c5e478d9696658b39 Mon Sep 17 00:00:00 2001 From: Uriah Finkel Date: Thu, 20 Aug 2026 12:28:59 +0300 Subject: [PATCH 4/4] Restore docs while keeping early time input validation --- .../performance_data_times.py | 97 +++++++++++++++++-- 1 file changed, 90 insertions(+), 7 deletions(-) diff --git a/src/rtichoke/performance_data/performance_data_times.py b/src/rtichoke/performance_data/performance_data_times.py index a0170047..7a00175a 100644 --- a/src/rtichoke/performance_data/performance_data_times.py +++ b/src/rtichoke/performance_data/performance_data_times.py @@ -3,11 +3,8 @@ """ from typing import Dict, Union -from collections.abc import Sequence - -import numpy as np import polars as pl - +from collections.abc import Sequence from rtichoke.processing.adjustments import create_adjusted_data from rtichoke.processing.combinations import ( create_aj_data_combinations, @@ -21,6 +18,8 @@ cast_and_join_adjusted_data, ) +import numpy as np + def prepare_performance_data_times( probs: Dict[str, np.ndarray], @@ -36,7 +35,51 @@ def prepare_performance_data_times( stratified_by: Sequence[str] = ("probability_threshold",), by: float = 0.01, ) -> pl.DataFrame: - """Prepare performance data for models with time-to-event outcomes.""" + """Prepare performance data for models with time-to-event outcomes. + + This function calculates a comprehensive set of performance metrics for + models predicting time-to-event outcomes. It handles censored data and + competing events by applying specified heuristics at different time + horizons. The function first bins the data using + `prepare_binned_classification_data_times` and then computes cumulative, + Aalen-Johansen-based performance metrics. + + The resulting dataframe is the primary input for time-dependent plotting + functions. + + Parameters + ---------- + probs : Dict[str, np.ndarray] + A dictionary mapping model or dataset names (str) to their predicted + probabilities of an event occurring by a given time. + reals : Union[np.ndarray, Dict[str, np.ndarray]] + The true event statuses. Can be a single array or a dictionary. + Labels should be integers indicating the outcome (e.g., 0=censored, + 1=event of interest, 2=competing event). + times : Union[np.ndarray, Dict[str, np.ndarray]] + The event or censoring times corresponding to the `reals`. Can be a + single array or a dictionary. + fixed_time_horizons : list[float] + A list of numeric time points at which to evaluate the model's + performance. Integer inputs are accepted and normalized to floats. + heuristics_sets : list[Dict], optional + A list of dictionaries, each specifying how to handle censored data + and competing events. The default is + ``[{"censoring_heuristic": "adjusted", + "competing_heuristic": "adjusted_as_negative"}]``. + stratified_by : Sequence[str], optional + Variables by which to stratify the analysis. Defaults to + ``("probability_threshold",)``. + by : float, optional + The step size for probability thresholds. Defaults to ``0.01``. + + Returns + ------- + pl.DataFrame + A Polars DataFrame with performance metrics computed across probability + thresholds and time horizons. It includes columns for cutoffs, time + points, heuristics, and performance measures. + """ final_adjusted_data = prepare_binned_classification_data_times( probs=probs, reals=reals, @@ -106,7 +149,45 @@ def prepare_binned_classification_data_times( by: float = 0.01, risk_set_scope: Sequence[str] = ["pooled_by_cutoff", "within_stratum"], ) -> pl.DataFrame: - """Prepare binned, time-dependent classification data.""" + """ + Prepare binned, time-dependent classification data. + + This function constructs the foundational binned data needed for + time-to-event performance analysis. It bins predictions by probability + thresholds, applies censoring and competing event heuristics, and stratifies + the data across specified time horizons. The output is a detailed breakdown + of outcomes within each bin, which can be used for calibration or passed to + `prepare_performance_data_times` for full performance metric calculation. + + Parameters + ---------- + probs : Dict[str, np.ndarray] + A dictionary mapping model or dataset names (str) to their predicted + probabilities. + reals : Union[np.ndarray, Dict[str, np.ndarray]] + The true event statuses (e.g., 0=censored, 1=event, 2=competing event). + times : Union[np.ndarray, Dict[str, np.ndarray]] + The event or censoring times. + fixed_time_horizons : list[float] + A list of numeric time points for performance evaluation. Integer + inputs are accepted and normalized to floats. + heuristics_sets : list[Dict], optional + Specifies how to handle censored data and competing events. + stratified_by : Sequence[str], optional + Variables for stratification. Defaults to ``("probability_threshold",)``. + by : float, optional + The step size for probability thresholds. Defaults to ``0.01``. + risk_set_scope : Sequence[str], optional + Defines the scope for risk set calculations. Defaults to + ``["pooled_by_cutoff", "within_stratum"]``. + + Returns + ------- + pl.DataFrame + A Polars DataFrame with binned, time-dependent data. Each row + represents a unique combination of dataset, bin, time horizon, + heuristic, and other strata. + """ _validate_time_input_alignment(probs=probs, reals=reals, times=times) fixed_time_horizons = [float(horizon) for horizon in fixed_time_horizons] @@ -140,7 +221,9 @@ def prepare_binned_classification_data_times( risk_set_scope=risk_set_scope, ) - return cast_and_join_adjusted_data( + final_adjusted_data = cast_and_join_adjusted_data( aj_data_combinations, adjusted_data, ).with_columns(pl.col("reals_estimate").fill_null(0.0)) + + return final_adjusted_data