diff --git a/src/rtichoke/performance_data/probs_distribution.py b/src/rtichoke/performance_data/probs_distribution.py index 42313d2a..46500b27 100644 --- a/src/rtichoke/performance_data/probs_distribution.py +++ b/src/rtichoke/performance_data/probs_distribution.py @@ -12,11 +12,79 @@ _EvaluationMetadata, _build_evaluation_metadata, ) +from rtichoke.processing.transforms import _compute_probability_quantile_bin_indices class _PredictionDistributionData(TypedDict): bins: pl.DataFrame operating_points: pl.DataFrame + rank_bins: pl.DataFrame + + +def _aggregate_rank_bins_for_evaluation( + probabilities: np.ndarray, + outcomes: np.ndarray, + by: float, + evaluation_metadata: _EvaluationMetadata, +) -> pl.DataFrame: + """Aggregate observed positive and negative mass into probability-quantile rank bins.""" + by = float(by) + q = int(round(1 / by)) + grid_rows = [] + for i in range(q): + grid_rows.append( + { + "stratum_id": i, + "evaluation": evaluation_metadata.evaluation, + "model": evaluation_metadata.model, + "population": evaluation_metadata.population, + "rank_lower": float(round(i * by, 10)), + "rank_upper": float(round((i + 1) * by, 10)), + } + ) + + grid_schema = { + "stratum_id": pl.Int64, + "evaluation": pl.String, + "model": pl.String, + "population": pl.String, + "rank_lower": pl.Float64, + "rank_upper": pl.Float64, + } + + complete_grid = pl.DataFrame(grid_rows, schema=grid_schema) + + if len(probabilities) == 0: + return complete_grid.with_columns( + pl.lit(0, dtype=pl.Int64).alias("n_positive"), + pl.lit(0, dtype=pl.Int64).alias("n_negative"), + ).drop("stratum_id") + + bin_indices, _ = _compute_probability_quantile_bin_indices(probabilities, by) + + obs_df = pl.DataFrame( + { + "stratum_id": bin_indices, + "is_pos": (outcomes == 1).astype(int), + "is_neg": (outcomes == 0).astype(int), + } + ) + + counts_df = obs_df.group_by("stratum_id").agg( + pl.col("is_pos").sum().cast(pl.Int64).alias("n_positive"), + pl.col("is_neg").sum().cast(pl.Int64).alias("n_negative"), + ) + + aggregated_rank_bins = ( + complete_grid.join(counts_df, on="stratum_id", how="left") + .with_columns( + pl.col("n_positive").fill_null(0), + pl.col("n_negative").fill_null(0), + ) + .drop("stratum_id") + ) + + return aggregated_rank_bins def _aggregate_bins_for_evaluation( @@ -236,10 +304,29 @@ def _prepare_probs_distribution_data( ) eval_bins_frames.append(eval_bins) + eval_rank_bins_frames = [] + for evaluation_key in evaluation_keys: + evaluation_metadata = evaluation_metadata_by_group[evaluation_key] + probabilities = np.asarray(probs[evaluation_key], dtype=float) + if isinstance(aligned_reals, dict): + outcomes = np.asarray(aligned_reals[evaluation_key], dtype=int) + else: + outcomes = np.asarray(aligned_reals, dtype=int) + + eval_rank_bins = _aggregate_rank_bins_for_evaluation( + probabilities=probabilities, + outcomes=outcomes, + by=by, + evaluation_metadata=evaluation_metadata, + ) + eval_rank_bins_frames.append(eval_rank_bins) + bins = pl.concat(eval_bins_frames, how="vertical") operating_points = pl.DataFrame(operating_point_rows, schema=operating_point_schema) + rank_bins = pl.concat(eval_rank_bins_frames, how="vertical") return _PredictionDistributionData( bins=bins, operating_points=operating_points, + rank_bins=rank_bins, ) diff --git a/src/rtichoke/processing/transforms.py b/src/rtichoke/processing/transforms.py index f5568ccf..594ac996 100644 --- a/src/rtichoke/processing/transforms.py +++ b/src/rtichoke/processing/transforms.py @@ -1,9 +1,33 @@ +from typing import Dict, Sequence, Union import numpy as np import polars as pl -from typing import Dict, Union from rtichoke.processing.combinations import create_breaks_values +def _compute_probability_quantile_bin_indices( + probs: Union[np.ndarray, Sequence[float]], by: float +) -> tuple[np.ndarray, int]: + """Compute observation-level probability-quantile bin indices 0..q-1. + + Reproduces equal-frequency quantile binning using R Type 7 linear quantiles, + accumulated boundary monotonicity, pinned endpoints [0, 1], and right-closed + digitization. + """ + by = float(by) + q = int(round(1 / by)) # e.g. 0.2 -> 5 bins + + probs_arr = np.asarray(probs, float) + + edges = np.quantile(probs_arr, np.linspace(0.0, 1.0, q + 1), method="linear") + edges = np.maximum.accumulate(edges) + + edges[0] = 0.0 + edges[-1] = 1.0 + + bin_idx = np.digitize(probs_arr, bins=edges[1:-1], right=True) + return bin_idx, q + + def add_cutoff_strata(data: pl.DataFrame, by: float, stratified_by) -> pl.DataFrame: def transform_group(group: pl.DataFrame, by: float) -> pl.DataFrame: probs = group["probs"].to_numpy() @@ -32,20 +56,7 @@ def transform_group(group: pl.DataFrame, by: float) -> pl.DataFrame: ) if "ppcr" in stratified_by: - # --- Compute strata_ppcr as equal-frequency quantile bins by rank --- - by = float(by) - q = int(round(1 / by)) # e.g. 0.2 -> 5 bins - - probs = np.asarray(probs, float) - - edges = np.quantile(probs, np.linspace(0.0, 1.0, q + 1), method="linear") - - edges = np.maximum.accumulate(edges) - - edges[0] = 0.0 - edges[-1] = 1.0 - - bin_idx = np.digitize(probs, bins=edges[1:-1], right=True) + bin_idx, q = _compute_probability_quantile_bin_indices(probs, by) s = str(by) decimals = len(s.split(".")[-1]) if "." in s else 0 diff --git a/tests/test_probs_distribution.py b/tests/test_probs_distribution.py index 06443e8c..4d139b79 100644 --- a/tests/test_probs_distribution.py +++ b/tests/test_probs_distribution.py @@ -112,12 +112,14 @@ def test_structure_and_schemas(): res = _prepare_probs_distribution_data(probs, reals, by=0.1) - assert set(res.keys()) == {"bins", "operating_points"} + assert set(res.keys()) == {"bins", "operating_points", "rank_bins"} bins = res["bins"] ops = res["operating_points"] + rank_bins = res["rank_bins"] assert isinstance(bins, pl.DataFrame) assert isinstance(ops, pl.DataFrame) + assert isinstance(rank_bins, pl.DataFrame) assert bins.schema["evaluation"] == pl.String assert bins.schema["model"] == pl.String @@ -137,6 +139,14 @@ def test_structure_and_schemas(): assert ops.schema["cutoff"] == pl.Float64 assert ops.schema["realized_ppcr"] == pl.Float64 + assert rank_bins.schema["evaluation"] == pl.String + assert rank_bins.schema["model"] == pl.String + assert rank_bins.schema["population"] == pl.String + assert rank_bins.schema["rank_lower"] == pl.Float64 + assert rank_bins.schema["rank_upper"] == pl.Float64 + assert rank_bins.schema["n_positive"] == pl.Int64 + assert rank_bins.schema["n_negative"] == pl.Int64 + def test_nullable_model_single_keyed_population(): probs = {"validation_population": np.array([0.2, 0.8])} @@ -386,3 +396,119 @@ def test_invalid_inputs(): ValueError, match="Estimated probabilities must be between 0 and 1" ): _prepare_probs_distribution_data({"m1": np.array([-0.1, 0.5])}, reals) + + +def test_rank_bins_stratification_invariance(): + probs = {"m1": np.array([0.0, 0.1, 0.2, 0.5, 0.5, 0.8, 1.0])} + reals = np.array([0, 0, 1, 1, 0, 1, 0]) + + res_thresh = _prepare_probs_distribution_data( + probs, reals, stratified_by=("probability_threshold",), by=0.2 + ) + res_ppcr = _prepare_probs_distribution_data( + probs, reals, stratified_by=("ppcr",), by=0.2 + ) + + rb_thresh = res_thresh["rank_bins"] + rb_ppcr = res_ppcr["rank_bins"] + + assert rb_thresh.equals(rb_ppcr) + + +def test_rank_bins_properties(): + # 1. Distinct scores mixed outcomes + probs = { + "m1": np.array([0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65, 0.75, 0.85, 0.95]) + } + reals = np.array([0, 0, 1, 0, 1, 0, 1, 1, 0, 1]) + + res = _prepare_probs_distribution_data(probs, reals, by=0.2) + rb = res["rank_bins"] + + # Retains complete requested grid (q = 1 / 0.2 = 5) + assert len(rb) == 5 + assert np.allclose(rb["rank_lower"].to_list(), [0.0, 0.2, 0.4, 0.6, 0.8]) + assert np.allclose(rb["rank_upper"].to_list(), [0.2, 0.4, 0.6, 0.8, 1.0]) + + # Total mass conservation + assert rb["n_positive"].sum() == (reals == 1).sum() + assert rb["n_negative"].sum() == (reals == 0).sum() + + +def test_rank_bins_ties_never_split(): + # All scores tied + probs = {"m1": np.array([0.5, 0.5, 0.5, 0.5, 0.5])} + reals = np.array([1, 0, 1, 0, 1]) + + res = _prepare_probs_distribution_data(probs, reals, by=0.2) + rb = res["rank_bins"] + + # Retains complete q=5 grid + assert len(rb) == 5 + # All mass assigned to a single bin, empty strata retained explicitly + non_zero_strata = rb.filter((pl.col("n_positive") > 0) | (pl.col("n_negative") > 0)) + assert len(non_zero_strata) == 1 + assert non_zero_strata["n_positive"][0] == 3 + assert non_zero_strata["n_negative"][0] == 2 + + assert rb["n_positive"].sum() == 3 + assert rb["n_negative"].sum() == 2 + + +def test_primary_golden_fixture(): + probs = {"m1": np.array([0.00, 0.15, 0.30, 0.50, 0.50, 0.50, 0.65, 0.80, 1.00])} + reals = np.array([0, 1, 0, 1, 0, 1, 1, 0, 1]) + + res = _prepare_probs_distribution_data(probs, reals, by=0.20) + rb = res["rank_bins"] + + expected_df = pl.DataFrame( + { + "evaluation": ["m1"] * 5, + "model": ["m1"] * 5, + "population": ["__shared_population__"] * 5, + "rank_lower": [0.00, 0.20, 0.40, 0.60, 0.80], + "rank_upper": [0.20, 0.40, 0.60, 0.80, 1.00], + "n_positive": [1, 2, 0, 1, 1], + "n_negative": [1, 2, 0, 0, 1], + } + ) + + assert rb.equals(expected_df) + + +def test_secondary_golden_fixture_n_less_than_q(): + probs = {"m1": np.array([0.10, 0.50, 0.90])} + reals = np.array([0, 1, 1]) + + res = _prepare_probs_distribution_data(probs, reals, by=0.20) + rb = res["rank_bins"] + + expected_df = pl.DataFrame( + { + "evaluation": ["m1"] * 5, + "model": ["m1"] * 5, + "population": ["__shared_population__"] * 5, + "rank_lower": [0.00, 0.20, 0.40, 0.60, 0.80], + "rank_upper": [0.20, 0.40, 0.60, 0.80, 1.00], + "n_positive": [0, 0, 1, 0, 1], + "n_negative": [1, 0, 0, 0, 0], + } + ) + + assert rb.equals(expected_df) + + +def test_primary_golden_fixture_order_invariance(): + probs_orig = np.array([0.00, 0.15, 0.30, 0.50, 0.50, 0.50, 0.65, 0.80, 1.00]) + reals_orig = np.array([0, 1, 0, 1, 0, 1, 1, 0, 1]) + + # Permuted order including within the tied 0.50 group + perm_idx = np.array([4, 0, 5, 2, 3, 8, 1, 7, 6]) + probs_perm = probs_orig[perm_idx] + reals_perm = reals_orig[perm_idx] + + res_orig = _prepare_probs_distribution_data({"m1": probs_orig}, reals_orig, by=0.20) + res_perm = _prepare_probs_distribution_data({"m1": probs_perm}, reals_perm, by=0.20) + + assert res_orig["rank_bins"].equals(res_perm["rank_bins"]) diff --git a/tests/test_transforms.py b/tests/test_transforms.py new file mode 100644 index 00000000..00e7ddbd --- /dev/null +++ b/tests/test_transforms.py @@ -0,0 +1,45 @@ +"""Tests for transforms module helper functions and add_cutoff_strata.""" + +import numpy as np +import polars as pl + +from rtichoke.processing.transforms import ( + _compute_probability_quantile_bin_indices, + add_cutoff_strata, +) + + +def test_compute_probability_quantile_bin_indices_basic(): + probs = np.array([0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]) + by = 0.2 + bin_idx, q = _compute_probability_quantile_bin_indices(probs, by) + + assert q == 5 + assert len(bin_idx) == len(probs) + assert bin_idx.min() >= 0 + assert bin_idx.max() < q + + +def test_add_cutoff_strata_unchanged(): + df = pl.DataFrame( + { + "reference_group": ["g1"] * 5, + "probs": [0.0, 0.25, 0.5, 0.75, 1.0], + "reals": [0, 1, 0, 1, 0], + } + ) + + res_prob = add_cutoff_strata(df, by=0.2, stratified_by=("probability_threshold",)) + assert "strata_probability_threshold" in res_prob.columns + assert res_prob["strata_probability_threshold"].to_list() == [ + "[0.00, 0.20)", + "[0.20, 0.40)", + "[0.40, 0.60)", + "[0.60, 0.80)", + "[0.80, 1.00]", + ] + + res_ppcr = add_cutoff_strata(df, by=0.2, stratified_by=("ppcr",)) + assert "strata_ppcr" in res_ppcr.columns + assert isinstance(res_ppcr["strata_ppcr"].dtype, pl.Enum) + assert res_ppcr["strata_ppcr"].to_list() == ["0.2", "0.4", "0.6", "0.8", "1.0"]