from __future__ import annotations import importlib from pathlib import Path import warnings import numpy as np import pandas as pd import streamlit as st try: from common.method_registry import ( dump_method_registry_snapshot, canonicalize_method_name, get_method_meta, normalize_method_name, ) except ModuleNotFoundError: from src.common.method_registry import ( dump_method_registry_snapshot, canonicalize_method_name, get_method_meta, normalize_method_name, ) # --------------------------------------------------------------------------- # Optimization (RMSE) datasets # --------------------------------------------------------------------------- DATASET_FILES = { "L63": [ "l63_ensemble_results.nc", "UKI_results/uki_l63_ensemble_results.nc", "bayesian/l63_abc.nc", "bayesian/l63_hm.nc", "adam_results/leaderboard_adam_l63_2026-06-26.nc", "levenberg_marquardt_results/leaderboard_lm_l63_2026-06-29.nc", ], "L96": [ "l96_ensemble_results.nc", "UKI_results/uki_l96_ensemble_results.nc", "bayesian/l96_abc.nc", "adam_results/leaderboard_adam_l96_const-force_2026-06-26.nc", "levenberg_marquardt_results/leaderboard_lm_l96_const-force_2026-06-29.nc", ], "L96_NN_FORCING": [ "l96_nn_forcing_ensemble_results.nc", "UKI_results/uki_l96_nn_forcing_ensemble_results.nc", ], "L96_SPATIAL_FORCING": [ "l96_spatial_forcing_ensemble_results.nc", "UKI_results/uki_l96_spatial_forcing_ensemble_results.nc", "bayesian/l96_varying_abc.nc", "adam_results/leaderboard_adam_l96_vec-force_2026-06-26.nc", "levenberg_marquardt_results/leaderboard_lm_l96_vec-force_2026-06-29.nc", ], } EXPECTED_DIMS_OPT = ("algorithm_type", "rmse_target", "ensemble_size", "random_seed") EXPECTED_DIMS_UQ = ("random_seed", "ensemble_size", "k_iter", "coverage_quantile") # --------------------------------------------------------------------------- # Uncertainty Quantification datasets # --------------------------------------------------------------------------- # Add entries here when UQ result NetCDF files are available. Each file must # carry a ``metric`` variable and a ``uq_target`` coordinate (in place of the # ``rmse_target`` used by the optimization files). The rest of the schema # (algorithm_type, ensemble_size, random_seed) is identical. # # Example: # UQ_DATASET_FILES = { # "L63": ["uq/l63_uq_results.nc"], # "L96": ["uq/l96_uq_results.nc"], # } UQ_DATASET_FILES: dict[str, list[str]] = {} # Canonical UQ target-scaling levels always offered in the UQ leaderboard selector. UQ_TARGET_LEVELS: list[float] = [1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5] # Budget-for-coverage source files: maps benchmark → list of (algorithm_type, nc_path) pairs. # Each NC must carry output_coverage (coverage_quantile, k_iter, ensemble_size, random_seed), # target_scaling, and output_dim. UQ_BUDGET_FILES: dict[str, list[tuple[str, str]]] = { "L63": [ ("ces-eki-dmc", "ces-eki-dmc_results/ces-eki-dmc_l63_ensemble_results_2026-06-15_minimal.nc"), ("iekf", "gnki-uq_results/leaderboard_gnki_l63_2026-07-07_minimal.nc"), ], "L96": [ ("ces-eki-dmc", "ces-eki-dmc_results/ces-eki-dmc_l96_ensemble_results_2026-06-15_minimal.nc"), ("iekf", "gnki-uq_results/leaderboard_gnki_l96_const-force_2026-07-07_minimal.nc"), ], "L96_NN_FORCING": [ ("ces-eki-dmc", "ces-eki-dmc_results/ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-06-15_minimal.nc"), ("iekf", "gnki-uq_results/leaderboard_gnki_l96_flux-force_2026-07-07_minimal.nc"), ], "L96_SPATIAL_FORCING": [ ("ces-eki-dmc", "ces-eki-dmc_results/ces-eki-dmc_l96_spatial_forcing_ensemble_results_2026-06-15_minimal.nc"), ("iekf", "gnki-uq_results/leaderboard_gnki_l96_vec-force_2026-07-07_minimal.nc"), ], } # Quantile levels used for the all-quantiles-satisfied budget-for-coverage condition. # Budget = N_ens · k_iter where first k s.t. |S(q)−q| ≤ c·√(q(1−q)/N_y) for ALL q below. UQ_COVERAGE_QUANTILES: list[float] = [0.15, 0.5, 0.85] # Physical dimensions of each benchmark: (param_dim, state_dim, output_dim). # Used only for column-header annotations in the suitability table. BENCHMARK_DIMS: dict[str, tuple[int, int, int]] = { "L63": (2, 3, 9), "L96": (1, 40, 80), "L96_NN_FORCING": (61, 100, 200), "L96_SPATIAL_FORCING": (40, 40, 80), } # --------------------------------------------------------------------------- # Private loader — shared by both public load_* functions # --------------------------------------------------------------------------- def _load_store(dataset_files: dict[str, list[str]], target_col: str, *, drop_nan_metric: bool = True) -> pd.DataFrame: """Load and merge NetCDF metric files into a long-format DataFrame. Parameters ---------- dataset_files: Mapping of benchmark name → list of relative NetCDF paths under ``data/``. target_col: Name of the target-coordinate dimension in the NetCDF files, e.g. ``"rmse_target"`` or ``"uq_target"``. """ if not dataset_files: return pd.DataFrame( columns=["benchmark", "algorithm_type", target_col, "ensemble_size", "metric", "failure_rate"] ) xr = importlib.import_module("xarray") project_root = Path(__file__).resolve().parent.parent data_dir = project_root / "data" records: list[pd.DataFrame] = [] for benchmark_name, filenames in dataset_files.items(): for filename in filenames: file_path = data_dir / filename if not file_path.exists(): warnings.warn(f"Dataset {file_path} does not exist") continue with xr.open_dataset(file_path) as dataset: if "metric" not in dataset: warnings.warn(f"Dataset {file_path} does not contain 'metric' variable") continue metric = dataset["metric"] if not set(metric.dims).issubset(EXPECTED_DIMS_OPT): warnings.warn( f"{file_path.name}: metric dims {metric.dims} contain unexpected " f"dimensions not in EXPECTED_DIMS_OPT {EXPECTED_DIMS_OPT} — skipping" ) continue # Track failures: -1 sentinel OR NaN (some methods use NaN instead of -1) failures_count = ((metric == -1) | metric.isnull()).sum(dim="random_seed") failure_rate = (failures_count / dataset.sizes["random_seed"]) * 100 # Filter out failed runs (metric == -1) metric = metric.where(metric != -1) metric_mean = metric.mean(dim="random_seed", skipna=True) df = metric_mean.to_dataframe(name="metric").reset_index() df_failures = failure_rate.to_dataframe(name="failure_rate").reset_index() df["failure_rate"] = df_failures["failure_rate"] df["benchmark"] = benchmark_name if "algorithm_type" not in df.columns: if "abc" in file_path.name: df["algorithm_type"] = "abc" elif "hm" in file_path.name: df["algorithm_type"] = "hm" if "ensemble_size" not in df.columns: if "abc" in file_path.name: df["ensemble_size"] = 1 required_cols = ["benchmark", "algorithm_type", target_col, "ensemble_size", "metric", "failure_rate"] missing_cols = [c for c in required_cols if c not in df.columns] if missing_cols: st.warning( f"Skipping {file_path.name}: metric dataframe is missing expected columns {missing_cols}" ) continue # Tag rows from files with exactly one ensemble size. Only those rows # are retained when metric is NaN — they represent a method that was # genuinely attempted at that size but every seed failed. NaN rows from # multi-size files are placeholder entries for sizes never actually run. df["_single_ens"] = dataset.sizes.get("ensemble_size", 1) == 1 records.append(df[required_cols + ["_single_ens"]]) if not records: return pd.DataFrame( columns=["benchmark", "algorithm_type", target_col, "ensemble_size", "metric", "failure_rate"] ) merged = pd.concat(records, ignore_index=True) merged["algorithm_alias"] = merged["algorithm_type"].map(normalize_method_name) merged["algorithm_type"] = merged["algorithm_alias"].map(canonicalize_method_name) merged["ensemble_size"] = pd.to_numeric(merged["ensemble_size"], errors="coerce") merged["metric"] = pd.to_numeric(merged["metric"], errors="coerce") merged["failure_rate"] = pd.to_numeric(merged["failure_rate"], errors="coerce").fillna(0.0) merged = merged.dropna(subset=["ensemble_size"]) if drop_nan_metric: # Keep NaN-metric rows only when they came from a single-ensemble-size file # AND every seed failed (failure_rate ≈ 100 %). Those represent a method that # was genuinely attempted but never converged. NaN rows from multi-size files # are placeholders for ensemble sizes that were never actually run. single_ens = merged.pop("_single_ens").fillna(False) merged = merged[merged["metric"].notna() | (single_ens & (merged["failure_rate"] >= 99.9))] else: merged = merged.drop(columns=["_single_ens"], errors="ignore") merged["ensemble_size"] = merged["ensemble_size"].astype(int) merged["abbreviation"] = merged["algorithm_type"].map(lambda method: get_method_meta(method).get("abbreviation")) merged["Method"] = merged["algorithm_type"].map(lambda method: get_method_meta(method).get("Method")) merged["family"] = merged["algorithm_type"].map(lambda method: get_method_meta(method).get("family")) merged["abbreviation"] = merged["abbreviation"].fillna(merged["algorithm_type"].str.upper()) merged["Method"] = merged["Method"].fillna(merged["abbreviation"]) merged["family"] = merged["family"].fillna("Kalman") try: dump_method_registry_snapshot(project_root=project_root, observed_methods=set(merged["algorithm_type"].unique())) except OSError: st.warning("Unable to write method registry snapshot to .cache/known_methods_snapshot.json") return merged # --------------------------------------------------------------------------- # Public cached loaders # --------------------------------------------------------------------------- @st.cache_data(show_spinner=False) def load_metric_store() -> pd.DataFrame: """Load the optimization (RMSE) metric store from ``DATASET_FILES``.""" return _load_store(DATASET_FILES, "rmse_target") @st.cache_data(show_spinner=False) def load_uq_store() -> pd.DataFrame: """Load the UQ metric store derived from budget-for-coverage computations. Returns a DataFrame compatible with ``render_leaderboard`` where ``metric`` is the mean budget (N_ens·k_iter) to satisfy the coverage condition across all three quantiles in ``UQ_COVERAGE_QUANTILES``. ``uq_target`` is the target-scaling level c. Failure means the coverage target was never reached within the k_iter range. """ budget_df = load_uq_budget_store() if budget_df.empty: return pd.DataFrame( columns=[ "benchmark", "algorithm_type", "algorithm_alias", "abbreviation", "Method", "family", "uq_target", "ensemble_size", "metric", "failure_rate", ] ) df = budget_df.copy() df["metric"] = df["mean_budget"] df["algorithm_alias"] = df["algorithm_type"].map(normalize_method_name) df["Method"] = df["algorithm_type"].map( lambda m: get_method_meta(canonicalize_method_name(normalize_method_name(m))).get( "Method", m.upper() ) ) return df[ [ "benchmark", "algorithm_type", "algorithm_alias", "abbreviation", "Method", "family", "uq_target", "ensemble_size", "metric", "failure_rate", ] ] @st.cache_data(show_spinner=False) def load_uq_budget_store() -> pd.DataFrame: """Compute budget-for-coverage and iterations-for-coverage from ensemble result NC files. For each (target_scaling c, ensemble_size N, random_seed): Find the minimum k_iter s.t. |S(q) − q| ≤ c·√(q(1−q)/N_y) for ALL q in UQ_COVERAGE_QUANTILES. Budget = N · k_iter; NaN when target never reached. Returns a DataFrame with columns: benchmark, algorithm_type, abbreviation, family, uq_target, ensemble_size, mean_budget, mean_iters, failure_count, failure_rate, n_seeds """ xr = importlib.import_module("xarray") project_root = Path(__file__).resolve().parent.parent data_dir = project_root / "data" bq_vals = np.array(UQ_COVERAGE_QUANTILES) records: list[dict] = [] for benchmark_name, method_files in UQ_BUDGET_FILES.items(): for algorithm_type, filename in method_files: file_path = data_dir / filename if not file_path.exists(): warnings.warn(f"UQ budget file {file_path} does not exist") continue with xr.open_dataset(file_path) as ds: if "output_coverage" not in ds: warnings.warn(f"{file_path.name}: missing output_coverage") continue if set(ds["output_coverage"].dims) != set(EXPECTED_DIMS_UQ): warnings.warn( f"{file_path.name}: output_coverage dims {ds['output_coverage'].dims} " f"do not match EXPECTED_DIMS_UQ {EXPECTED_DIMS_UQ} — skipping" ) continue cov_q = ds["coverage_quantile"].values # (n_cov_q,) k_vals = ds["k_iter"].values # (n_k,) 1-indexed ens_vals = ds["ensemble_size"].values # (n_ens,) ts_vals = ds["target_scaling"].values # (n_ts,) # xarray drops dimensions that no variable uses; read output_dim # directly from the underlying netCDF4 file to handle minimal files. nc4 = importlib.import_module("netCDF4") with nc4.Dataset(file_path) as _nc: if "output_dim" not in _nc.dimensions: warnings.warn(f"{file_path.name}: missing output_dim dimension") continue n_y = len(_nc.dimensions["output_dim"]) # dims: (coverage_quantile, k_iter, ensemble_size, random_seed) cov_np = ds["output_coverage"].values n_rng = cov_np.shape[3] # Indices into coverage_quantile for [0.15, 0.5, 0.85] bq_idx = [int(np.argmin(np.abs(cov_q - q))) for q in bq_vals] for si, c in enumerate(ts_vals): tol = c * np.sqrt(bq_vals * (1.0 - bq_vals) / n_y) for ei, ens_f in enumerate(ens_vals): N_ens = int(ens_f) budgets = np.full(n_rng, np.nan) kiters = np.full(n_rng, np.nan) for ri in range(n_rng): for ki, kv in enumerate(k_vals): s_q = cov_np[bq_idx, ki, ei, ri] if np.any(np.isnan(s_q)): continue if np.all(np.abs(s_q - bq_vals) <= tol): budgets[ri] = N_ens * float(kv) kiters[ri] = float(kv) break valid_b = budgets[~np.isnan(budgets)] valid_k = kiters[~np.isnan(kiters)] n_fail = n_rng - len(valid_b) records.append( { "benchmark": benchmark_name, "algorithm_type": algorithm_type, "uq_target": float(c), "ensemble_size": N_ens, "mean_budget": float(np.mean(valid_b)) if len(valid_b) > 0 else np.nan, "mean_iters": float(np.mean(valid_k)) if len(valid_k) > 0 else np.nan, "failure_count": n_fail, "failure_rate": 100.0 * n_fail / n_rng, "n_seeds": n_rng, } ) if not records: return pd.DataFrame( columns=[ "benchmark", "algorithm_type", "abbreviation", "family", "uq_target", "ensemble_size", "mean_budget", "mean_iters", "failure_count", "failure_rate", "n_seeds", ] ) df = pd.DataFrame(records) def _meta(m: str, key: str, fallback: str) -> str: return get_method_meta(canonicalize_method_name(normalize_method_name(m))).get(key, fallback) df["abbreviation"] = df["algorithm_type"].map(lambda m: _meta(m, "abbreviation", m.upper())) df["family"] = df["algorithm_type"].map(lambda m: _meta(m, "family", "Kalman")) df["Method"] = df["algorithm_type"].map(lambda m: _meta(m, "Method", m.upper())) return df