Skip to content
222 changes: 222 additions & 0 deletions src/rtichoke/_performance_table_spec.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,222 @@
"""Internal canonical PerformanceTableSpec builders.

These helpers translate already-calculated production performance data plus
semantic evaluation metadata. They do not calculate or render statistics.
"""

from __future__ import annotations

import math
from collections.abc import Mapping
from typing import Any, TypedDict, cast

import polars as pl

from rtichoke.processing.evaluation_semantics import _EvaluationMetadata

_METRICS: tuple[tuple[str, str], ...] = (
("true_positives", "True Positives"),
("true_negatives", "True Negatives"),
("false_positives", "False Positives"),
("false_negatives", "False Negatives"),
("sensitivity", "Sensitivity"),
("specificity", "Specificity"),
("false_positive_rate", "False Positive Rate"),
("ppv", "PPV"),
("npv", "NPV"),
("lift", "Lift"),
("predicted_positives", "Predicted Positives"),
("ppcr", "PPCR"),
("net_benefit", "Net Benefit"),
("net_benefit_interventions_avoided", "Interventions Avoided"),
)


class _EvaluationSpec(TypedDict, total=False):
id: str
model: str
population: str


class _MetricDefinition(TypedDict):
id: str
label: str


class _OperatingPoint(TypedDict):
type: str
value: float


class _MetricValue(TypedDict):
metricId: str
estimate: float | int | None


class _EvaluationContext(TypedDict):
censoringHeuristic: str
competingEventHeuristic: str


class _PerformanceTableRow(TypedDict, total=False):
evaluationId: str
operatingPoint: _OperatingPoint
values: list[_MetricValue]
horizon: float
context: _EvaluationContext


class _PerformanceTableSpec(TypedDict):
schemaVersion: str
type: str
evaluations: list[_EvaluationSpec]
metrics: list[_MetricDefinition]
rows: list[_PerformanceTableRow]


def _performance_table_spec_from_performance_data(
performance_data: pl.DataFrame,
evaluation_metadata: Mapping[str, _EvaluationMetadata],
) -> _PerformanceTableSpec:
"""Build the canonical static PerformanceTableSpec."""
return _build_performance_table_spec(
performance_data, evaluation_metadata, time_dependent=False
)


def _performance_table_times_spec_from_performance_data(
performance_data: pl.DataFrame,
evaluation_metadata: Mapping[str, _EvaluationMetadata],
) -> _PerformanceTableSpec:
"""Build the canonical time-dependent PerformanceTableSpec."""
return _build_performance_table_spec(
performance_data, evaluation_metadata, time_dependent=True
)


def _build_performance_table_spec(
performance_data: pl.DataFrame,
evaluation_metadata: Mapping[str, _EvaluationMetadata],
*,
time_dependent: bool,
) -> _PerformanceTableSpec:
required = {"reference_group", "stratified_by", "chosen_cutoff"}
if time_dependent:
required |= {
"fixed_time_horizon",
"censoring_heuristic",
"competing_heuristic",
}
missing = required.difference(performance_data.columns)
if missing:
raise ValueError(
"Performance-table data is missing columns: " + ", ".join(sorted(missing))
)

metric_definitions: list[_MetricDefinition] = [
cast(_MetricDefinition, {"id": metric_id, "label": label})
for metric_id, label in _METRICS
if metric_id in performance_data.columns
]
metric_ids = [definition["id"] for definition in metric_definitions]

rows = performance_data.select(
[
column
for column in (
"reference_group",
"fixed_time_horizon",
"censoring_heuristic",
"competing_heuristic",
"stratified_by",
"chosen_cutoff",
*metric_ids,
)
if column in performance_data.columns
]
).to_dicts()
row_groups = {str(row["reference_group"]) for row in rows}
missing_metadata = row_groups.difference(evaluation_metadata)
if missing_metadata:
raise ValueError(
"Performance-table rows are missing evaluation metadata: "
+ ", ".join(sorted(missing_metadata))
)

ordered_groups = [group for group in evaluation_metadata if group in row_groups]
evaluation_ids = {
group: f"evaluation-{index}"
for index, group in enumerate(ordered_groups, start=1)
}
evaluations: list[_EvaluationSpec] = []
for group in ordered_groups:
metadata = evaluation_metadata[group]
evaluation: _EvaluationSpec = {
"id": evaluation_ids[group],
"population": metadata.population,
}
if metadata.model is not None:
evaluation["model"] = metadata.model
evaluations.append(evaluation)

canonical_rows: list[_PerformanceTableRow] = []
for row in rows:
group = str(row["reference_group"])
stratified_by = str(row["stratified_by"])
operating_point: _OperatingPoint
if stratified_by == "probability_threshold":
operating_point = {
"type": "probability_threshold",
"value": _number(row["chosen_cutoff"]),
}
elif stratified_by == "ppcr":
operating_point = {"type": "ppcr", "value": _number(row["ppcr"])}
else:
raise ValueError(
"Canonical PerformanceTableSpec supports probability_threshold "
f"or ppcr operating points, not {stratified_by!r}"
)

canonical_row: _PerformanceTableRow = {
"evaluationId": evaluation_ids[group],
"operatingPoint": operating_point,
"values": cast(
list[_MetricValue],
[
{
"metricId": metric_id,
"estimate": _nullable_number(row[metric_id]),
}
for metric_id in metric_ids
],
),
}
if time_dependent:
canonical_row["horizon"] = _number(row["fixed_time_horizon"])
canonical_row["context"] = {
"censoringHeuristic": str(row["censoring_heuristic"]),
"competingEventHeuristic": str(row["competing_heuristic"]),
}
canonical_rows.append(canonical_row)

return {
"schemaVersion": "2.0",
"type": "performance_table",
"evaluations": evaluations,
"metrics": metric_definitions,
"rows": canonical_rows,
}


def _nullable_number(value: Any) -> float | int | None:
if value is None:
return None
if isinstance(value, float) and math.isnan(value):
return None
return value


def _number(value: Any) -> float:
if value is None:
raise ValueError("Operating point and horizon values must not be null")
return float(value)
17 changes: 12 additions & 5 deletions src/rtichoke/processing/adjustments.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
from collections.abc import Sequence

import pandas as pd
import polars as pl
from polarstate import predict_aj_estimates
from polarstate import prepare_event_table
from collections.abc import Sequence
from polarstate import predict_aj_estimates, prepare_event_table

from rtichoke.processing.transforms import assign_and_explode_polars


Expand Down Expand Up @@ -529,7 +530,7 @@ def _aj_estimates_by_cutoff_per_horizon(
df.filter(pl.col("fixed_time_horizon") == h)
.group_by("strata")
.map_groups(
lambda group: extract_aj_estimate_by_cutoffs(
lambda group, h=h: extract_aj_estimate_by_cutoffs(
group, [h], breaks, stratified_by, full_event_table=False
)
)
Expand All @@ -547,7 +548,7 @@ def _aj_estimates_per_horizon(
df.filter(pl.col("fixed_time_horizon") == h)
.group_by("strata")
.map_groups(
lambda group: extract_aj_estimate_for_strata(
lambda group, h=h: extract_aj_estimate_for_strata(
group, [h], full_event_table
)
)
Expand Down Expand Up @@ -650,6 +651,8 @@ def _aj_adjusted_events(
adjusted = extract_aj_estimate_by_cutoffs(
non_competing, horizons, breaks, stratified_by, full_event_table
)
else:
raise ValueError(f"Unsupported risk-set scope: {risk_set_scope!r}")

adjusted = adjusted.with_columns(
[
Expand Down Expand Up @@ -704,6 +707,8 @@ def _aj_adjusted_events(
adjusted = _aj_estimates_by_cutoff_per_horizon(
base_df, horizons, breaks, stratified_by
)
else:
raise ValueError(f"Unsupported risk-set scope: {risk_set_scope!r}")

adjusted = adjusted.with_columns(
pl.lit(risk_set_scope)
Expand All @@ -730,6 +735,8 @@ def _aj_adjusted_events(
adjusted = extract_aj_estimate_by_cutoffs(
base_df, horizons, breaks, stratified_by, full_event_table
)
else:
raise ValueError(f"Unsupported risk-set scope: {risk_set_scope!r}")

adjusted = adjusted.with_columns(
[
Expand Down
8 changes: 5 additions & 3 deletions tests/test_cutoff_grid_endpoint.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from typing import cast

import numpy as np
import pytest

Expand All @@ -17,7 +19,7 @@ def test_probability_threshold_breaks_match_r_seq_semantics(by, expected):
breaks = create_breaks_values(None, "probability_threshold", by)

assert breaks.tolist() == expected
assert np.max(breaks) <= 1.0
assert float(np.max(breaks)) <= 1.0


@pytest.mark.parametrize(
Expand All @@ -33,7 +35,7 @@ def test_binary_performance_data_matches_r_cutoff_endpoint(by, expected_max):

assert cutoffs.min() == 0.0
assert cutoffs.max() == expected_max
assert cutoffs.max() <= 1.0
assert cast(float, cutoffs.max()) <= 1.0


@pytest.mark.parametrize(
Expand All @@ -56,4 +58,4 @@ def test_time_performance_data_matches_r_cutoff_endpoint(by, expected_max):

assert cutoffs.min() == 0.0
assert cutoffs.max() == expected_max
assert cutoffs.max() <= 1.0
assert cast(float, cutoffs.max()) <= 1.0
10 changes: 7 additions & 3 deletions tests/test_performance_table.py
Original file line number Diff line number Diff line change
@@ -1,18 +1,20 @@
from typing import cast

import numpy as np
import pytest
from great_tables import GT
from reactable import Reactable

import rtichoke.performance_table as performance_table_module
from rtichoke.performance_table_reactable import _bar_style, _net_benefit_style

from rtichoke import (
create_performance_table,
create_performance_table_times,
prepare_performance_data,
prepare_performance_data_times,
render_performance_table,
)
from rtichoke.performance_table import PerformanceTableRenderer
from rtichoke.performance_table_reactable import _bar_style, _net_benefit_style


def _example():
Expand Down Expand Up @@ -215,7 +217,9 @@ def test_invalid_renderer_is_rejected():
probs, reals = _example()
data = prepare_performance_data(probs, reals, by=0.1)
with pytest.raises(ValueError, match="renderer"):
render_performance_table(data, renderer="unknown")
render_performance_table(
data, renderer=cast(PerformanceTableRenderer, "unknown")
)


def test_reactable_metric_bar_matches_r_colors_and_geometry():
Expand Down
Loading
Loading