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