Hugging Face's logo

Spaces:
calibrationcomp
/
calibration_benchmark
Running

App Files Community
20
calibration_benchmark / src /common /method_registry.py
odunbar's picture
odunbar
Score-based Levenberg Marquardt (#12)
1ad9fc2
Raw History Blame
11.2 kB
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).
# "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" β€” follows the loss gradient or a Gauss-Newton approximation of it.
# "kalman" β€” a (possibly linearized/unscented) Kalman-style ensemble update.
# "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": "kalman",
"method_goal": "optimization",
"emulator_use": "none",
"aliases": ["teki"],
},
"etki": {
"abbreviation": "ETKI",
"Method": "Ensemble Transform Kalman Inversion",
"parallelism": "parallel-interacting",
"update_type": "kalman",
"method_goal": "optimization",
"emulator_use": "none",
"aliases": ["etki"],
},
"iekf": {
"abbreviation": "IEKF",
"Method": "Iterative Ensemble Kalman Filter",
"parallelism": "parallel-interacting",
"update_type": "kalman",
"method_goal": "uq",
"emulator_use": "none",
"aliases": ["iekf", "gnsl", "gnki"],
},
"uki": {
"abbreviation": "UKI",
"Method": "Unscented Kalman Inversion",
"parallelism": "parallel-interacting",
"update_type": "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": "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": "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": "kalman",
"method_goal": "uq",
"emulator_use": "after-optimize",
"aliases": ["ces-iekf-const"]
},
"adam": {
"abbreviation": "ADAM",
"Method": "Adaptive Moment Estimation",
"parallelism": "serial",
"update_type": "gradient",
"method_goal": "optimization",
"emulator_use": "none",
"aliases": ["adam"],
},
"lm": {
"abbreviation": "LM",
"Method": "Levenberg-Marquardt",
"parallelism": "serial",
"update_type": "gradient",
"method_goal": "optimization",
"emulator_use": "none",
"aliases": ["lm", "levenberg_marquardt", "levenberg-marquardt", "gradient_descent"],
},
"gboed": {
"abbreviation": "GBOED",
"Method": "Goal-oriented Bayesian Optimal Experimental Design",
"parallelism": "parallel-interacting",
"update_type": "general",
"method_goal": "uq",
"emulator_use": "within-optimize",
"aliases": ["gboed", "boed"],
},
"lm-kgmm": {
"abbreviation": "LM-KGMM",
"Method": "Score-based Levenberg-Marquardt (KGMM)",
"parallelism": "serial",
"update_type": "gradient",
"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",
"method_goal": "optimization",
"emulator_use": "none",
"aliases": [
"lm-dsm",
"score-based lm (gfdt, dsm) (extended window)",
"score-based lm (gfdt, dsm) (parallel branches)",
],
},
}
# 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())
}
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")