rpx_benchmark.jedi

  1"""JEDI: Joint Empirical Desirability Index.
  2
  3Implements the composite desirability score from paper §3.3. JEDI is a
  4direction-aware geometric mean over per-metric individual desirabilities
  5:math:`d_k \\in [0, 1]`. Failing any one metric drives the composite toward
  6zero unless the caller opts into an :math:`\\varepsilon`-floor.
  7
  8Pipeline (paper Eqs. 4–6)
  9-------------------------
 10
 111. Look up each metric's :class:`MetricSpec` (direction, best, worst).
 122. Map raw score :math:`x_k` to individual desirability::
 13
 14       d_k = clip( (x_k - x_k^worst) / (x_k^best - x_k^worst), epsilon, 1 )
 15
 16   The numerator and denominator share the *same sign* by construction
 17   (paper §3.3), so :math:`d_k \\in [0, 1]` without ad-hoc negation. The
 18   ``clip(·, eps, 1)`` guards out-of-bounds values.
 19
 203. Take the geometric mean::
 21
 22       J = ( prod_k d_k ) ** (1/K)
 23
 244. For per-(scene, phase) aggregation, average :math:`J_{s,p}` across the
 25   :math:`3N` cells to get :math:`J_{model}`.
 26
 27Default :math:`\\varepsilon = 0` keeps the strict "any failure tanks J"
 28semantics the paper argues for; practitioners who want softer behaviour
 29can pass a small positive ``epsilon`` (e.g. 0.01).
 30"""
 31
 32from __future__ import annotations
 33
 34from dataclasses import dataclass
 35from typing import Dict, List, Mapping, Optional, Sequence
 36
 37import numpy as np
 38
 39from .exceptions import MetricError
 40from .logging_utils import get_logger
 41from .metrics.specs import MetricSpec, get_spec
 42
 43log = get_logger(__name__)
 44
 45
 46# --------------------------------------------------------------------------- #
 47# Per-cell JEDI
 48# --------------------------------------------------------------------------- #
 49
 50
 51@dataclass
 52class JEDIResult:
 53    """Outcome of a single (scene, phase) JEDI computation.
 54
 55    Attributes
 56    ----------
 57    score : float
 58        The composite J = ( ∏ d_k )^(1/K) ∈ [eps^1, 1].
 59    individual : Dict[str, float]
 60        Per-metric individual desirabilities d_k.
 61    out_of_bounds : List[str]
 62        Metric names whose raw value fell outside the registered
 63        (worst, best) window and were clipped. Useful for flagging
 64        models that exceed the frozen bounds (paper §3.3).
 65    epsilon : float
 66        The ε floor used; 0.0 by default (strict).
 67    """
 68
 69    score: float
 70    individual: Dict[str, float]
 71    out_of_bounds: List[str]
 72    epsilon: float = 0.0
 73
 74    def to_dict(self) -> Dict[str, object]:
 75        return {
 76            "jedi": self.score,
 77            "individual": dict(self.individual),
 78            "out_of_bounds": list(self.out_of_bounds),
 79            "epsilon": self.epsilon,
 80        }
 81
 82
 83def _desirability(raw: float, spec: MetricSpec, epsilon: float) -> tuple[float, bool]:
 84    """Compute one individual desirability d_k and the out-of-bounds flag.
 85
 86    Returns
 87    -------
 88    d_k : float
 89        Clipped into ``[max(epsilon, 0), 1]``.
 90    out_of_bounds : bool
 91        True iff the raw value fell outside ``(worst, best)`` before
 92        clipping. Helps surface models that exceed the frozen window.
 93    """
 94    span = spec.best - spec.worst
 95    # span is positive for higher-is-better, negative for lower-is-better;
 96    # numerator (raw - worst) carries the same sign by construction, so the
 97    # ratio is in [0, 1] when raw ∈ [worst, best] regardless of direction.
 98    ratio = (raw - spec.worst) / span
 99
100    floor = max(epsilon, 0.0)
101    clipped = float(np.clip(ratio, floor, 1.0))
102    oob = bool(ratio < 0.0 or ratio > 1.0)
103    return clipped, oob
104
105
106def compute_jedi(
107    raw_metrics: Mapping[str, float],
108    *,
109    epsilon: float = 0.0,
110    specs: Optional[Mapping[str, MetricSpec]] = None,
111) -> JEDIResult:
112    """Compute JEDI for one (scene, phase) cell.
113
114    Parameters
115    ----------
116    raw_metrics : Mapping[str, float]
117        ``{metric_name: raw_score}`` for the K metrics of one task on
118        one (scene, phase). Keys must have specs registered (either
119        via the built-in registry in :mod:`rpx_benchmark.metrics.specs`
120        or via explicit ``specs=`` override).
121    epsilon : float
122        Lower floor on individual desirabilities. Default 0.0 (strict
123        Derringer–Suich: a single failure zeros J). Pass e.g. 0.01 for
124        a softer penalty (paper §3.3, "Practitioners who prefer a softer
125        penalty …").
126    specs : Mapping[str, MetricSpec], optional
127        Override the global spec registry. When None, each name is
128        resolved via :func:`get_spec`.
129
130    Returns
131    -------
132    JEDIResult
133        The composite score, per-metric d_k breakdown, and any
134        out-of-bounds flags.
135
136    Raises
137    ------
138    MetricError
139        If ``raw_metrics`` is empty or any metric has no registered spec.
140    """
141    if not raw_metrics:
142        raise MetricError(
143            "compute_jedi requires at least one metric",
144            hint="Pass {metric_name: raw_score, ...} with K >= 1 entries.",
145        )
146    if epsilon < 0.0 or epsilon >= 1.0:
147        raise MetricError(
148            f"epsilon must be in [0, 1), got {epsilon}",
149            hint="Use epsilon=0 for strict, or a small positive value (e.g. 0.01).",
150        )
151
152    individual: Dict[str, float] = {}
153    oob: List[str] = []
154    for name, raw in raw_metrics.items():
155        if not np.isfinite(raw):
156            raise MetricError(
157                f"Metric {name!r} has non-finite raw value {raw!r}",
158                hint="Filter NaN/Inf upstream or impute before computing JEDI.",
159            )
160        spec = specs[name] if specs is not None and name in specs else get_spec(name)
161        d_k, was_oob = _desirability(float(raw), spec, epsilon)
162        individual[name] = d_k
163        if was_oob:
164            oob.append(name)
165
166    # Geometric mean. Numerically safer via exp(mean(log(·))) when no d_k=0.
167    values = np.array(list(individual.values()), dtype=np.float64)
168    if (values == 0.0).any():
169        score = 0.0
170    else:
171        score = float(np.exp(np.log(values).mean()))
172
173    return JEDIResult(
174        score=score,
175        individual=individual,
176        out_of_bounds=oob,
177        epsilon=epsilon,
178    )
179
180
181# --------------------------------------------------------------------------- #
182# Per-model aggregation across (scene, phase)
183# --------------------------------------------------------------------------- #
184
185
186@dataclass
187class JEDIModelResult:
188    """Per-model JEDI aggregated over all (scene, phase) cells.
189
190    Attributes
191    ----------
192    jedi : float
193        ``mean_{s,p} J_{s,p}`` — paper Eq. 6.
194    per_phase : Dict[str, float]
195        Per-phase means: ``{"clu": J_clu, "int": J_int, "cln": J_cln}``.
196        Phase names are passed through from input keys.
197    n_cells : int
198        Number of (scene, phase) cells contributing to the mean.
199    n_out_of_bounds : int
200        Total count of out-of-bounds metric flags across all cells.
201    """
202
203    jedi: float
204    per_phase: Dict[str, float]
205    n_cells: int
206    n_out_of_bounds: int
207
208    def to_dict(self) -> Dict[str, object]:
209        return {
210            "jedi": self.jedi,
211            "per_phase": dict(self.per_phase),
212            "n_cells": self.n_cells,
213            "n_out_of_bounds": self.n_out_of_bounds,
214        }
215
216
217def aggregate_jedi(
218    per_cell: Sequence[tuple[str, JEDIResult]],
219) -> JEDIModelResult:
220    """Aggregate per-cell JEDI scores into a per-model summary.
221
222    Parameters
223    ----------
224    per_cell : sequence of (phase_label, JEDIResult)
225        One tuple per (scene, phase) cell. ``phase_label`` is the
226        column key (e.g. ``"clu"`` / ``"int"`` / ``"cln"``); rows
227        within the same phase are averaged together.
228
229    Returns
230    -------
231    JEDIModelResult
232    """
233    if not per_cell:
234        raise MetricError(
235            "aggregate_jedi requires at least one per-cell result",
236            hint="Pass [(phase_label, JEDIResult), ...] with len >= 1.",
237        )
238
239    by_phase: Dict[str, List[float]] = {}
240    total_oob = 0
241    for phase_label, r in per_cell:
242        by_phase.setdefault(phase_label, []).append(r.score)
243        total_oob += len(r.out_of_bounds)
244
245    per_phase = {p: float(np.mean(vs)) for p, vs in by_phase.items()}
246    overall = float(np.mean([r.score for _, r in per_cell]))
247
248    return JEDIModelResult(
249        jedi=overall,
250        per_phase=per_phase,
251        n_cells=len(per_cell),
252        n_out_of_bounds=total_oob,
253    )
254
255
256__all__ = [
257    "JEDIResult",
258    "JEDIModelResult",
259    "compute_jedi",
260    "aggregate_jedi",
261]