Spaces:

calibrationcomp
/

calibration_benchmark

Running

File size: 17,945 Bytes

from __future__ import annotations

import json
from pathlib import Path


# Taxonomy
# --------
# Every method carries four independent tags, each scoped to its own axis:
#
# ``parallelism`` — how the search explores parameter space:
#   "serial"               — a single point estimate advanced step by step (ADAM, LM).
#   "serial (short window)" — serial, but each step advances a sliding statistics window by
#                             a fraction of its length (RL-*-SHORT) rather than a full window.
#   "parallel-independent" — a population of candidates updated with no coupling
#                             between members (ABC's accepted samples, HM's per-wave
#                             resampling from the non-implausible region).
#   "parallel-interacting" — an ensemble whose members are coupled through a shared
#                             update each iteration (any Kalman-based method).
#
# ``update_type`` — the mechanism driving each update step:
#   "gradient (exact)"  — follows the loss gradient or a Gauss-Newton approximation of it,
#                          differentiated through the forward model (ADAM, LM).
#   "gradient (Kalman)" — a (possibly linearized/unscented) Kalman-style ensemble update,
#                          i.e. an ensemble-estimated gradient/Gauss-Newton step.
#   "gradient (score)"  — a gradient built from a learned/estimated score or critic rather
#                          than a simulator Jacobian (*-KGMM, *-DSM, RL-*).
#   "general"           — anything else (e.g. ABC's rejection sampling, HM's implausibility cuts).
#
# ``method_goal`` — what the method is built to report:
#   "optimization" — a single best-fit parameter estimate.
#   "uq"           — the full posterior / parameter uncertainty. UQ methods can still
#                    be scored on the Optimization leaderboard (usually less
#                    competitive there, since they're not optimizing for speed-to-target).
#
# ``emulator_use`` — when/whether a surrogate model of the forward model is used:
#   "none"            — samples/evaluates the true forward model throughout.
#   "within-optimize" — refits a surrogate at each iteration of the search (HM waves).
#   "after-optimize"  — fits a surrogate once, after calibration finishes (CES).
#
# Note: Kalman methods are Bayesian in spirit too (approximate Gaussian posterior
# updates) — ``update_type`` is about mechanism, not a "Bayesian vs not" philosophy.
KNOWN_METHODS = {
    "teki": {
        "abbreviation": "TEKI",
        "Method": "Tikhonov Regularized Ensemble Kalman Inversion",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["teki"],
    },
    "etki": {
        "abbreviation": "ETKI",
        "Method": "Ensemble Transform Kalman Inversion",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["etki"],
    },
    "iekf": {
        "abbreviation": "IEKF",
        "Method": "Iterative Ensemble Kalman Filter",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "uq",
        "emulator_use": "none",
        "aliases": ["iekf", "gnsl", "gnki"],
    },
    "uki": {
        "abbreviation": "UKI",
        "Method": "Unscented Kalman Inversion",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["uki"],
    },
    "abc": {
        "abbreviation": "ABC",
        "Method": "Approximate Bayesian Calibration",
        "parallelism": "parallel-independent",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "none",
        "aliases": ["abc"],
    },
    "hm": {
        "abbreviation": "HM",
        "Method": "History Matching",
        "parallelism": "parallel-independent",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "within-optimize",
        "aliases": ["hm"],
    },
    "ces-eki-dmc": {
        "abbreviation": "CES-EKI-DMC",
        "Method": "Calibrate Emulate Sample (EKI-DataMisfitController)",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "uq",
        "emulator_use": "after-optimize",
        "aliases": ["ces-eki-dmc"]
    },
    "ces-eki-const": {
        "abbreviation": "CES-EKI-CONST",
        "Method": "Calibrate Emulate Sample (EKI-Constant Scheduler)",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "uq",
        "emulator_use": "after-optimize",
        "aliases": ["ces-eki-const"]
    },
    "ces-iekf-const": {
        "abbreviation": "CES-IEKF-CONST",
        "Method": "Calibrate Emulate Sample (IEKF-Constant Scheduler)",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "uq",
        "emulator_use": "after-optimize",
        "aliases": ["ces-iekf-const"]
    },
    "adam": {
        "abbreviation": "ADAM",
        "Method": "Adaptive Moment Estimation",
        "parallelism": "serial",
        "update_type": "gradient (exact)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["adam"],
    },
    "lm": {
        "abbreviation": "LM",
        "Method": "Levenberg-Marquardt",
        "parallelism": "serial",
        "update_type": "gradient (exact)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["lm", "levenberg_marquardt", "levenberg-marquardt", "gradient_descent"],
    },
    "gboed-tmcmc": {
        "abbreviation": "GBOED-TMCMC",
        "Method": "Goal-oriented Bayesian Optimal Experimental Design (TMCMC)",
        "parallelism": "parallel-interacting",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "within-optimize",
        "aliases": ["gboed", "gboed-tmcmc", "boed", "gboed-tmcmc-eig", "boed-tmcmc-eig", "tmcmc-eig"],
    },
    "lm-kgmm": {
        "abbreviation": "LM-KGMM",
        "Method": "Score-based Levenberg-Marquardt (KGMM)",
        "parallelism": "serial",
        "update_type": "gradient (score)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": [
            "lm-kgmm",
            "score-based lm (gfdt, kgmm) (extended window)",
            "score-based lm (gfdt, kgmm) (parallel branches)",
        ],
    },
    "lm-dsm": {
        "abbreviation": "LM-DSM",
        "Method": "Score-based Levenberg-Marquardt (DSM)",
        "parallelism": "serial",
        "update_type": "gradient (score)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": [
            "lm-dsm",
            "score-based lm (gfdt, dsm) (extended window)",
            "score-based lm (gfdt, dsm) (parallel branches)",
        ],
    },
    "adam-kgmm": {
        "abbreviation": "ADAM-KGMM",
        "Method": "Score-based Adam (KGMM)",
        "parallelism": "serial",
        "update_type": "gradient (score)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": [
            "adam-kgmm",
            "score-based adam (gfdt, kgmm) (extended window)",
            "score-based adam (gfdt, kgmm) (parallel branches)",
        ],
    },
    "adam-dsm": {
        "abbreviation": "ADAM-DSM",
        "Method": "Score-based Adam (DSM)",
        "parallelism": "serial",
        "update_type": "gradient (score)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": [
            "adam-dsm",
            "score-based adam (gfdt, dsm) (extended window)",
            "score-based adam (gfdt, dsm) (parallel branches)",
        ],
    },
    "rl-ddpg-short": {
        "abbreviation": "RL-DDPG-SHORT",
        "Method": "Reinforcement Learning (DDPG, short RL step)",
        "parallelism": "serial (short window)",
        "update_type": "gradient (score)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["rl-ddpg-short"],
    },
    "rl-ddpg": {
        "abbreviation": "RL-DDPG",
        "Method": "Reinforcement Learning (DDPG)",
        "parallelism": "serial",
        "update_type": "gradient (score)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["rl-ddpg", "rl-ddpg-long"],
    },
    "rl-td3-short": {
        "abbreviation": "RL-TD3-SHORT",
        "Method": "Reinforcement Learning (TD3, short RL step)",
        "parallelism": "serial (short window)",
        "update_type": "gradient (score)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["rl-td3-short"],
    },
    "rl-td3": {
        "abbreviation": "RL-TD3",
        "Method": "Reinforcement Learning (TD3)",
        "parallelism": "serial",
        "update_type": "gradient (score)",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["rl-td3", "rl-td3-long"],
    },
    "ies-iekf": {
        "abbreviation": "IES-IEKF",
        "Method": "Iterated Emulate-Sample (IEKF)",
        "parallelism": "parallel-interacting",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "within-optimize",
        "aliases": ["ies-iekf", "boed-iekf"],
    },
    "gboed-iekf": {
        "abbreviation": "GBOED-IEKF",
        "Method": "Goal-oriented Bayesian Optimal Experimental Design (IEKF)",
        "parallelism": "parallel-interacting",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "within-optimize",
        "aliases": ["gboed-iekf", "gboed-iekf-eig", "boed-iekf-eig"],
    },
    "ies-tmcmc": {
        "abbreviation": "IES-TMCMC",
        "Method": "Iterated Emulate-Sample (TMCMC)",
        "parallelism": "parallel-interacting",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "within-optimize",
        "aliases": ["ies-tmcmc", "boed-tmcmc"],
    },
    "eks": {
        "abbreviation": "EKS",
        "Method": "Ensemble Kalman Sampler (EKS a.k.a ALDI)",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "uq",
        "emulator_use": "none",
        "aliases": ["eks", "sampler"],
    },
    "ces-eks": {
        "abbreviation": "CES-EKS",
        "Method": "Calibrate Emulate Sample (EKS)",
        "parallelism": "parallel-interacting",
        "update_type": "gradient (Kalman)",
        "method_goal": "uq",
        "emulator_use": "after-optimize",
        "aliases": ["ces-eks"],
    },
    "ies-eks": {
        "abbreviation": "IES-EKS",
        "Method": "Iterated Emulate-Sample (EKS)",
        "parallelism": "parallel-interacting",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "within-optimize",
        "aliases": ["ies-eks", "boed-eks"],
    },
    "gboed-eks": {
        "abbreviation": "GBOED-EKS",
        "Method": "Goal-oriented Bayesian Optimal Experimental Design (EKS)",
        "parallelism": "parallel-interacting",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "within-optimize",
        "aliases": ["gboed-eks", "gboed-eks-eig", "boed-eks-eig"],
    },
}


# One slot per method in KNOWN_METHODS declaration order. New methods appended
# to KNOWN_METHODS get the next slot.
#
# Generated by a greedy farthest-point search in OKLab: each slot was chosen to
# maximize its minimum distance to every slot before it, simultaneously under
# normal vision and simulated protan/deutan color-vision deficiency (Machado,
# Oliveira & Fernandes 2009 severity-1.0 transforms — the same model used by
# the dataviz skill's validate_palette.js). Because this optimizes the running
# ALL-PAIRS minimum (not just adjacent-slot pairs), the printed ΔE at slot k
# below is a genuine worst-case guarantee for any subset of the first k slots
# shown together — relevant here because any subset of methods can appear
# together via the "Methods to display in charts" multiselect, not just
# neighbors in this list.
#
# Honest limits (OKLab ΔE x100; CVD = min(protan, deutan)):
#   - Slots  0-15 (16 colors): CVD >= 8.0  — clears the full colorblind-safety
#     target for every pair among any subset of these.
#   - Slots 16-27            : CVD in [6.0, 8.0) — the legal "floor" band per
#     the skill (secondary encoding required). This app already provides it:
#     every legend and tooltip in leaderboard.py shows the method abbreviation
#     as text, so identity is never color-alone even where hue alone is weak.
#   - Slots 28-29            : CVD ~5.8, just under the 6.0 floor for their
#     closest pair — weakest two slots, kept only for the requested headroom.
#   - Normal-vision floor (>=15, a hard gate when checked in isolation): holds
#     for the first ~6 slots against ALL other slots; beyond that, some pairs
#     among 30 simultaneously-distinguishable colors are close even for
#     full-color vision. This is a geometry limit of a finite perceptual
#     space, not a search-quality issue — the reference palette in the dataviz
#     skill's own docs states plainly that no ordering clears all-pairs floors
#     past ~8 categories. 30 slots were requested anyway for growth headroom;
#     rely on the legend/tooltip text labels (already present) for identity,
#     not on hue alone, especially past slot ~16.
# Re-run the search (seeded with all prior slots fixed) if extending further.
_METHOD_PALETTE = [
    "#d80065", "#007fff", "#ff9700", "#00cae4", "#6d00b8", "#970f00",
    "#007ca4", "#de8580", "#00a0d7", "#fa0000", "#0100ff", "#90006d",
    "#e500a0", "#007656", "#008200", "#006dd1", "#00b7ad", "#47b3ff",
    "#0087d2", "#e97e00", "#ec9f72", "#4b4589", "#00d0ae", "#c87181",
    "#5194ff", "#9c5683", "#0058ff", "#892348", "#e2003d", "#00b5d5",
]

# Stable abbreviation → hex color mapping.  Import this wherever Altair charts are built
# so every plot in the app assigns the same color to each method.  With 30 palette
# slots and 14 registered methods, every method gets a unique color — the modulo
# below only matters as a defensive fallback if KNOWN_METHODS grows past 30 without
# the palette being extended first (see comment above).
METHOD_COLORS: dict[str, str] = {
    meta["abbreviation"]: _METHOD_PALETTE[i % len(_METHOD_PALETTE)]
    for i, meta in enumerate(KNOWN_METHODS.values())
}


# Method families used to group rows in the suitability tables (the online table in
# common/leaderboard.py and the LaTeX tables from offline-plots/suitability_tables.py).
# The idea is to put methods of similar construction — and so similar performance —
# next to each other.  Dict order = family order; list order = row order within a
# family.  A family only appears on a board if one of its methods has data there.
# Every method that can appear in a suitability table must be listed here.
SUITABILITY_FAMILIES: dict[str, list[str]] = {
    "CES": ["CES-EKI-CONST", "CES-EKI-DMC", "CES-EKS", "CES-IEKF-CONST"],
    "Kalman": ["EKS", "ETKI", "IEKF", "TEKI", "UKI"],
    "Gradient": ["LM", "ADAM"],
    "GFDT": ["LM-DSM", "ADAM-DSM", "LM-KGMM", "ADAM-KGMM"],
    "RL": ["RL-DDPG", "RL-TD3", "RL-DDPG-SHORT", "RL-TD3-SHORT"],
    "IES": ["IES-EKS", "IES-IEKF", "IES-TMCMC"],
    "GBOED": ["GBOED-EKS", "GBOED-IEKF", "GBOED-TMCMC"],
    "General": ["ABC", "HM"],
}

# Spelled-out display names for the SUITABILITY_FAMILIES keys (suitability table,
# home-page key, and the grouped method selector on the leaderboards).
FAMILY_FULL_NAMES: dict[str, str] = {
    "CES": "Calibrate, Emulate, Sample (CES)",
    "Kalman": "Kalman",
    "Gradient": "Gradient",
    "GFDT": "Generalized Fluctuation-Dissipation (GFDT)",
    "RL": "Reinforcement Learning (RL)",
    "IES": "Iterated Emulate-Sample (IES)",
    "GBOED": "Goal-oriented Bayesian Experimental Design (GBOED)",
    "General": "General",
}


def family_full_name(family: str) -> str:
    return FAMILY_FULL_NAMES.get(family, family)


def build_alias_lookup() -> dict[str, str]:
    lookup: dict[str, str] = {}
    for canonical_name, meta in KNOWN_METHODS.items():
        lookup[canonical_name] = canonical_name
        lookup[canonical_name.upper()] = canonical_name
        for alias in meta.get("aliases", []):
            lookup[alias.lower()] = canonical_name
            lookup[alias.upper()] = canonical_name
    return lookup


ALIAS_TO_CANONICAL = build_alias_lookup()


def normalize_method_name(name: object) -> str:
    text = str(name).strip()

    if text.startswith("b'") and text.endswith("'"):
        text = text[2:-1]
    elif text.startswith('b"') and text.endswith('"'):
        text = text[2:-1]

    text = text.strip("\"'").strip()
    return text.lower()


def canonicalize_method_name(name: object) -> str:
    normalized = normalize_method_name(name)
    return ALIAS_TO_CANONICAL.get(normalized, normalized)


def get_method_meta(canonical_name: str) -> dict[str, str]:
    return KNOWN_METHODS.get(canonical_name, {})


def dump_method_registry_snapshot(project_root: Path, observed_methods: set[str]) -> None:
    snapshot = {
        "known_methods": KNOWN_METHODS,
        "observed_methods": sorted(observed_methods),
        "unmapped_observed_methods": sorted([method for method in observed_methods if method not in KNOWN_METHODS]),
    }
    cache_dir = project_root / ".cache"
    cache_dir.mkdir(parents=True, exist_ok=True)
    target_file = cache_dir / "known_methods_snapshot.json"
    target_file.write_text(json.dumps(snapshot, indent=2), encoding="utf-8")