Implement producer-owned prediction distribution rank_bins

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

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

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

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"]