File size: 6,328 Bytes
fce6c09
 
 
 
a58bdd0
fce6c09
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
baf1286
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fce6c09
 
a58bdd0
 
fce6c09
 
 
 
 
 
 
 
 
 
 
 
a58bdd0
 
 
 
 
fce6c09
 
a58bdd0
 
 
baf1286
fce6c09
a58bdd0
baf1286
a288069
 
 
 
baf1286
 
 
 
 
 
 
a288069
 
a58bdd0
baf1286
 
 
 
 
 
 
 
 
 
 
a288069
 
baf1286
a58bdd0
baf1286
a58bdd0
 
fce6c09
a288069
fce6c09
 
a288069
fce6c09
 
 
 
 
 
a288069
fce6c09
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
from __future__ import annotations

import importlib
from pathlib import Path
import warnings 
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,
    )


DATASET_FILES = {
    "L63": [
        "l63_ensemble_results.nc",
        "UKI_results/uki_l63_ensemble_results.nc",
        "bayesian/l63_abc.nc",
        "bayesian/l63_hm.nc"
    ],
    "L96": [
        "l96_ensemble_results.nc",
        "UKI_results/uki_l96_ensemble_results.nc",
        "bayesian/l96_abc.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"
    ],
}



EXPECTED_DIMS = ("algorithm_type", "rmse_target", "ensemble_size", "random_seed")


@st.cache_data(show_spinner=False)
def load_metric_store() -> pd.DataFrame:
    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"]
                
                # Track failures (metric == -1) as a percentage
                failures_count = (metric == -1).sum(dim="random_seed")

                failure_rate = (failures_count / dataset.sizes["random_seed"]) * 100
                
                # Filter out failed calibrations (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
                        
                missing_cols = [c for c in ["benchmark", "algorithm_type", "rmse_target", "ensemble_size", "metric", "failure_rate"] 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

                records.append(df[["benchmark", "algorithm_type", "rmse_target", "ensemble_size", "metric", "failure_rate"]])

    if not records:
        return pd.DataFrame(columns=["benchmark", "algorithm_type", "rmse_target", "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", "metric"])
    merged["ensemble_size"] = merged["ensemble_size"].astype(int)
    merged["forward_model_runs"] = merged["metric"]

    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


def build_leaderboard(metric_store: pd.DataFrame) -> pd.DataFrame:
    if metric_store.empty:
        return pd.DataFrame(
            columns=[
                "family",
                "Method",
                "abbreviation",
                "benchmark",
                "Forward Model Runs",
                "Optimal Ensemble Size",
                "Target Level",
            ]
        )

    best_idx = metric_store.groupby(["benchmark", "algorithm_type"])["metric"].idxmin()
    best = metric_store.loc[best_idx].copy()
    best = best.rename(columns={"metric": "Forward Model Runs", "rmse_target": "Target Level", "ensemble_size": "Optimal Ensemble Size"})
    return best[["family", "Method", "abbreviation", "benchmark", "Forward Model Runs", "Optimal Ensemble Size", "Target Level"]].sort_values(
        ["benchmark", "Forward Model Runs"], ascending=[True, True]
    )