From 1397deb7c12fa2dbf12cf11e248aa2d52048ce3c Mon Sep 17 00:00:00 2001 From: Robin Date: Wed, 17 Sep 2025 16:15:40 +0200 Subject: [PATCH 1/6] add plotting and eval scripts --- codes/benchmark/bench_fcts.py | 9 +- codes/benchmark/bench_plots.py | 27 +- config.yaml | 20 +- paper_eval.py | 100 +++++ paper_plots.py | 647 +++++++++++++++++++++++++++++++++ 5 files changed, 787 insertions(+), 16 deletions(-) create mode 100644 paper_eval.py create mode 100644 paper_plots.py diff --git a/codes/benchmark/bench_fcts.py b/codes/benchmark/bench_fcts.py index 407f42e..30c5235 100644 --- a/codes/benchmark/bench_fcts.py +++ b/codes/benchmark/bench_fcts.py @@ -1,3 +1,4 @@ +import os from contextlib import redirect_stdout from typing import Any @@ -95,6 +96,8 @@ def run_benchmark(surr_name: str, surrogate_class, conf: dict) -> dict[str, Any] tolerance=conf["dataset"]["tolerance"], per_species=conf["dataset"].get("normalise_per_species", False), ) + # TEMP + print(conf["dataset"]["name"], train_data.shape, val_data.shape, test_data.shape) model_config = get_model_config(surr_name, conf) n_timesteps = train_data.shape[1] @@ -1272,6 +1275,10 @@ def compare_errors(metrics: dict[str, dict], config: dict) -> None: if log_errors: plot_errors_over_time(mean_log, median_log, timesteps, config, mode="deltadex") plot_error_distribution_comparative(log_errors, config, mode="deltadex") + # TEMP + dataset = config["dataset"]["name"] + os.makedirs(f"scripts/pp/{dataset}", exist_ok=True) + np.savez(f"scripts/pp/{dataset}/all_log_errors.npz", log_errors) def compare_inference_time( @@ -1540,7 +1547,7 @@ def compare_UQ(all_metrics: dict, config: dict) -> None: ensemble_errors, ensemble_std, config, - flag_fractions=(0.01, 0.05, 0.10, 0.20, 0.30, 0.40, 0.50), + flag_fractions=(0, 0.025, 0.05, 0.10, 0.20, 0.30, 0.40, 0.50), save=True, show_title=True, ) diff --git a/codes/benchmark/bench_plots.py b/codes/benchmark/bench_plots.py index 8acf26c..40f0ad1 100644 --- a/codes/benchmark/bench_plots.py +++ b/codes/benchmark/bench_plots.py @@ -2458,6 +2458,10 @@ def plot_catastrophic_detection_curves( names = list(errors_log.keys()) colors = plt.cm.viridis(np.linspace(0, 0.95, len(names))) summary: dict[str, dict[float, dict[str, float]]] = {} + # TEMP + recall_99 = np.zeros((len(names), len(flag_fractions))) + recall_90 = np.zeros((len(names), len(flag_fractions))) + dataset_name = conf["dataset"]["name"] # --- Recall vs fraction flagged (per catastrophic percentile) --- for ax, perc in zip(axes[:-1], percentiles): @@ -2476,11 +2480,21 @@ def plot_catastrophic_detection_curves( xs, ys = [], [] for f in flag_fractions: - unc_thr = np.percentile(u, 100.0 * (1.0 - float(f))) - flagged = u >= unc_thr - recall = (flagged & is_cat).sum() / n_cat if n_cat > 0 else 0.0 - xs.append(100.0 * flagged.mean()) - ys.append(100.0 * recall) + if f <= 0.0: + xs.append(0.0) + ys.append(0.0) + recall = 0.0 + else: + unc_thr = np.percentile(u, 100.0 * (1.0 - float(f))) + flagged = u >= unc_thr + recall = (flagged & is_cat).sum() / n_cat if n_cat > 0 else 0.0 + xs.append(100.0 * flagged.mean()) + ys.append(100.0 * recall) + # TEMP + if perc == 99.0: + recall_99[i, flag_fractions.index(f)] = recall + if perc == 90.0: + recall_90[i, flag_fractions.index(f)] = recall ax.plot( xs, @@ -2507,6 +2521,9 @@ def plot_catastrophic_detection_curves( f"Detection @ {perc}th percentile (Top {100 - perc:.0f}% Δdex)" ) + np.savez(f"scripts/pp/{dataset_name}/catastrophic_recall_99.npz", recall_99) + np.savez(f"scripts/pp/{dataset_name}/catastrophic_recall_90.npz", recall_90) + # MAE improvement plot ax_mae = axes[-1] for i, name in enumerate(names): diff --git a/config.yaml b/config.yaml index 1642023..e1df0a6 100644 --- a/config.yaml +++ b/config.yaml @@ -11,32 +11,32 @@ dataset: tolerance: 1e-25 normalise_per_species: True log_timesteps: True -devices: ["cuda:3", "cuda:2", "cuda:3", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"] +devices: ["cuda:0", "cuda:2", "cuda:3", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"] seed: 42 verbose: False checkpoint: True # Models to train interpolation: - enabled: False + enabled: True intervals: [2, 3, 4, 5, 6, 7, 8, 9, 10] extrapolation: - enabled: False + enabled: True cutoffs: [50, 60, 70, 80, 90] sparse: - enabled: False + enabled: True factors: [2, 4, 8, 16, 32] batch_scaling: - enabled: False + enabled: True sizes: [1/16, 1/8, 1/4, 1/2] uncertainty: - enabled: False + enabled: True ensemble_size: 5 # Number of models for deep ensemble # Evaluations during benchmark -iterative: False -losses: False -gradients: False +iterative: True +losses: True +gradients: True timing: True -compute: False +compute: True compare: True # Whether to compare the surrogates diff --git a/paper_eval.py b/paper_eval.py new file mode 100644 index 0000000..b89be1d --- /dev/null +++ b/paper_eval.py @@ -0,0 +1,100 @@ +""" +Convenience script to re-run paper evaluations across multiple training runs. + +Behavior: +- For each specified training_id, load its saved config from trained//config.yaml +- Override the devices in that config with the hardcoded DEVICE below +- Run the same evaluation flow as run_eval.py using that modified config + +Note: This script intentionally does NOT read the top-level config.yaml. +""" + +import os +from typing import Dict + +from codes.benchmark import ( + check_benchmark, + check_surrogate, + compare_models, + get_surrogate, + run_benchmark, +) +from codes.utils import download_data, nice_print, read_yaml_config + +# Hardcoded device override for all evaluations +DEVICE = "cuda:0" + +# Training IDs to evaluate +TRAINING_IDS = [ + "_cloud_finetuned", + "_cloud_parametric_finetuned", + "_primordial_finetuned", + "_primordial_parametric_finetuned", +] + + +def evaluate_with_config(config: Dict) -> None: + """Run the evaluation loop for a single configuration dict.""" + # Basic checks and data + check_benchmark(config) + download_data(config["dataset"]["name"], verbose=config.get("verbose", False)) + + surrogates = config["surrogates"] + all_metrics = {surrogate: {} for surrogate in surrogates} + + # Evaluate each surrogate + for surrogate_name in surrogates: + surrogate_class = get_surrogate(surrogate_name) + if surrogate_class is None: + print(f"Surrogate {surrogate_name} not recognized. Skipping.") + continue + + nice_print(f"Running benchmark for {surrogate_name}") + check_surrogate(surrogate_name, config) + metrics = run_benchmark(surrogate_name, surrogate_class, config) + all_metrics[surrogate_name] = metrics + + # Compare models if requested + if config.get("compare", False): + if len(surrogates) < 2: + nice_print("At least two surrogate models are required to compare.") + else: + nice_print("Comparing models") + compare_models(all_metrics, config) + + +def load_trained_config(training_id: str) -> Dict | None: + """Load the saved config for a given training_id from trained//config.yaml.""" + cfg_path = os.path.join("trained", training_id, "config.yaml") + if not os.path.exists(cfg_path): + print( + f"Config not found for training_id '{training_id}': {cfg_path}. Skipping." + ) + return None + config = read_yaml_config(cfg_path) + return config + + +def main(): + for tid in TRAINING_IDS: + nice_print(f"Evaluating {tid}") + + config = load_trained_config(tid) + if config is None: + continue + + # Override devices with the hardcoded DEVICE + config["devices"] = [DEVICE] + + try: + evaluate_with_config(config) + except Exception as e: + print(f"Evaluation failed for {tid}: {e}") + # Continue with the next training_id + continue + + nice_print("All requested evaluations processed") + + +if __name__ == "__main__": + main() diff --git a/paper_plots.py b/paper_plots.py new file mode 100644 index 0000000..e64c011 --- /dev/null +++ b/paper_plots.py @@ -0,0 +1,647 @@ +#!/usr/bin/env python3 +""" +Paper plots: comparative Δdex error distributions across datasets. + +This script loads the per-dataset error dictionaries saved by the benchmark +(compare_errors -> scripts/pp//all_log_errors.npz), and creates a 2x2 grid +with one subplot per dataset, each showing the same comparative plot as +plot_error_distribution_comparative(..., mode="deltadex"). + +Usage: + python paper_plots.py --root scripts/pp \ + [--output plots/paper/error_dist_deltadex_by_dataset.png] \ + [--cols 2] + +Notes: +- Each dataset directory must contain an NPZ file named 'all_log_errors.npz'. + For robustness, we also try 'all_errors_log.npz' as a fallback. +- The NPZ file typically contains a single object array 'arr_0' which is a + dict mapping surrogate_name -> numpy array of Δdex errors with shape [N, T, Q]. +""" +from __future__ import annotations + +import argparse +import os +from typing import Dict, List, Tuple + +import matplotlib.pyplot as plt +import numpy as np +from scipy.ndimage import gaussian_filter1d + +# Reuse palette from project for consistent styling, but keep a safe fallback +try: + from codes.benchmark.bench_plots import get_custom_palette +except Exception: # pragma: no cover - fallback if import fails + + def get_custom_palette(n: int): + return plt.cm.viridis(np.linspace(0, 0.95, n)) + + +def _format_dataset_title(name: str) -> str: + """Map dataset folder names to display titles with proper capitalization.""" + mapping = { + "primordial": "Primordial", + "primordial_parametric": "Primordial Parametric", + } + if name in mapping: + return mapping[name] + # Fallback: replace underscores with space and title-case + return name.replace("_", " ").title() + + +def _load_errors_npz(path_npz: str) -> Dict[str, np.ndarray]: + """ + Load a dict[str, np.ndarray] from an NPZ file created by np.savez. + + Supports the common patterns: + - arr_0 (object array) holding a Python dict + - Named arrays per surrogate (if saved with kwargs) + - A single key 'log_errors' holding the dict + """ + if not os.path.exists(path_npz): + raise FileNotFoundError(path_npz) + + data = np.load(path_npz, allow_pickle=True) + try: + # Preferred: saved as a single object array containing the dict + if "arr_0" in data.files and isinstance(data["arr_0"], np.ndarray): + obj = data["arr_0"] + # Could be 0-d object array with dict + if obj.dtype == object: + d = obj.item() + if isinstance(d, dict): + return d + # Alternative: explicit key name + if "log_errors" in data.files: + d = data["log_errors"].item() + if isinstance(d, dict): + return d + # Fallback: construct dict from per-surrogate arrays + out: Dict[str, np.ndarray] = {} + for k in data.files: + arr = data[k] + # Only accept ND arrays + if isinstance(arr, np.ndarray) and arr.ndim >= 1: + out[k] = arr + if out: + return out + finally: + data.close() + + raise ValueError(f"Could not interpret NPZ structure in {path_npz}") + + +def load_dataset_errors(root: str, dataset: str) -> Dict[str, np.ndarray]: + """Try both file names and return the errors dict for a dataset.""" + cand1 = os.path.join(root, dataset, "all_log_errors.npz") + cand2 = os.path.join(root, dataset, "all_errors_log.npz") # user-mentioned alt + last_err: Exception | None = None + for p in (cand1, cand2): + try: + return _load_errors_npz(p) + except Exception as e: + last_err = e + continue + raise FileNotFoundError( + f"No error file found for dataset '{dataset}'. Tried: {cand1}, {cand2}. Last error: {last_err}" + ) + + +def compute_global_range( + datasets_errors: Dict[str, Dict[str, np.ndarray]], + low_pct: float = 2.0, + high_pct: float = 98.0, +) -> Tuple[float, float]: + """ + Compute global x-range in log10 space across all datasets and surrogates, + following the same logic as plot_error_distribution_comparative. + """ + log_arrays: List[np.ndarray] = [] + for ds, err_dict in datasets_errors.items(): + for _, arr in err_dict.items(): + flat = arr.astype(float).ravel() + # Filter finite and strictly positive (avoid log10(0) and NaN) + mask = np.isfinite(flat) & (flat > 0) + if not np.any(mask): + continue + log_arrays.append(np.log10(flat[mask])) + if not log_arrays: + # Default safe range if everything is empty + return -8.0, 0.0 + + mins = [np.percentile(x, low_pct) for x in log_arrays if x.size > 0] + maxs = [np.percentile(x, high_pct) for x in log_arrays if x.size > 0] + global_min = float(np.min(mins)) + global_max = float(np.max(maxs)) + # Expand to nice boundaries + x_min = float(np.floor(global_min)) + x_max = float(np.ceil(global_max)) + return x_min, x_max + + +def build_color_map(datasets_errors: Dict[str, Dict[str, np.ndarray]]): + """Build a consistent surrogate->color map across all datasets using viridis. + + Ensures deterministic ordering by sorting surrogate names alphabetically. + Returns a dict preserving this order and the ordered list of names. + """ + name_set = set() + for err_dict in datasets_errors.values(): + name_set.update(list(err_dict.keys())) + names = sorted(name_set) + # Permute to specific legend ordering + if len(names) == 4: + names = [names[i] for i in [3, 0, 2, 1]] + colors = plt.cm.viridis(np.linspace(0, 0.95, len(names))) + # Dict preserves insertion order, matching `names` sequence + color_map = {name: colors[i] for i, name in enumerate(names)} + return color_map, names + + +def reorder_legend_entries_rowwise( + handles: List, labels: List[str], max_ncols: int +) -> Tuple[List, List[str], int]: + """ + Reorder legend entries so they display row-wise when Matplotlib fills columns first. + + Provide `handles`/`labels` in the desired row-wise reading order; this + function returns a permutation that, when passed to Matplotlib legend with + ncol=legend_ncol, yields that row-wise order. + + Returns (final_handles, final_labels, legend_ncol). + """ + N = len(handles) + if N == 0: + return [], [], 1 + + legend_ncol = max(1, min(max_ncols, N)) + rows = int(np.ceil(N / legend_ncol)) + + final_handles: List = [] + final_labels: List[str] = [] + # Convert a row-wise ordered list to the column-first order expected by Matplotlib + # Example (N=6, ncol=2): input [a,b,c,d,e,f] -> output [a,c,e,b,d,f] + for c in range(legend_ncol): + for r in range(rows): + idx = r * legend_ncol + c + if idx < N: + final_handles.append(handles[idx]) + final_labels.append(labels[idx]) + + return final_handles, final_labels, legend_ncol + + +def plot_grid_deltadex( + datasets: List[str], + datasets_errors: Dict[str, Dict[str, np.ndarray]], + x_log_min: float, + x_log_max: float, + color_map: Dict[str, Tuple[float, float, float, float]], + dpi: int = 300, + n_cols: int = 2, +): + """ + Render a 2x2 grid: one subplot per dataset, reproducing the comparative + error distribution plot for Δdex with consistent axes and colors. + """ + # Prepare x bin edges in log10 space and transform to linear for plotting + x_vals = np.linspace(x_log_min, x_log_max + 0.1, 100) + + n = max(1, len(datasets)) + n_cols = max(1, n_cols) + n_rows = int(np.ceil(n / n_cols)) + fig, axes = plt.subplots( + n_rows, + n_cols, + figsize=(5 * n_cols, 3 * n_rows), + sharex=True, # sharey=True + ) + if isinstance(axes, np.ndarray): + axes = axes.flatten() + else: + axes = [axes] + + for idx, (ax, dataset) in enumerate(zip(axes, datasets)): + err_dict = datasets_errors.get(dataset, {}) + if not err_dict: + ax.text(0.5, 0.5, f"No data for {dataset}", ha="center", va="center") + ax.set_axis_off() + continue + + # For legend ordering, use color_map order + for model_name, color in color_map.items(): + if model_name not in err_dict: + continue + arr = err_dict[model_name] + flat = arr.astype(float).ravel() + mask = np.isfinite(flat) & (flat > 0) + if not np.any(mask): + continue + vals = flat[mask] + logs = np.log10(vals) + + hist, bin_edges = np.histogram(logs, bins=x_vals, density=True) + smoothed = gaussian_filter1d(hist, sigma=2) + + ax.plot(10 ** bin_edges[:-1], smoothed, label=model_name, color=color) + + # Mean and median markers (on linear scale) + mean_val = float(np.mean(vals)) + median_val = float(np.median(vals)) + ax.axvline( + x=mean_val, color=color, linestyle="--", linewidth=1.0, alpha=0.9 + ) + ax.axvline( + x=median_val, color=color, linestyle="-.", linewidth=1.0, alpha=0.9 + ) + + ax.set_xscale("log") + ax.set_xlim(left=1e-4, right=10) + # Y label only on first column + if (idx % n_cols) == 0: + ax.set_ylabel("Smoothed Histogram Count") + ax.set_ylim(0, None) + ax.set_title(_format_dataset_title(dataset)) + + # Common X label and legend + for ax in axes[-n_cols:]: + ax.set_xlabel(r"Log-MAE ($\Delta dex$)") + + # Build a single legend using first axis handles for present models + handles, labels = [], [] + for model_name, color in color_map.items(): + # Proxy lines for legend + line = plt.Line2D([0], [0], color=color, label=model_name) + handles.append(line) + labels.append(model_name) + # Mean/median style proxies + handles.append(plt.Line2D([0], [0], color="black", linestyle="--", label="Mean")) + labels.append("Mean") + handles.append(plt.Line2D([0], [0], color="black", linestyle="-.", label="Median")) + labels.append("Median") + + final_handles, final_labels, legend_ncol = reorder_legend_entries_rowwise( + handles, labels, max_ncols=2 + ) + + # Place legend below plots, arranged row-wise + fig.legend( + final_handles, + final_labels, + loc="lower center", + bbox_to_anchor=(0.52, 0.05), + fontsize="small", + frameon=True, + ncol=legend_ncol, + ) + + # No overall title; leave space at the bottom for legend + plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) + + out_path = "scripts/pp/error_dist_deltadex_by_dataset.png" + os.makedirs(os.path.dirname(out_path), exist_ok=True) + fig.savefig(out_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + + +def plot_grid_deltadex_percentiles( + datasets: List[str], + datasets_errors: Dict[str, Dict[str, np.ndarray]], + timesteps: np.ndarray, + color_map: Dict[str, Tuple[float, float, float, float]], + dpi: int = 300, + n_cols: int = 2, +): + """ + Create a grid of subplots (one per dataset) showing Δdex error percentiles over time. + + For each dataset: + - Draw neutral grey one-sided percentile bands (50, 90, 99) aggregated across surrogates. + - Overlay each surrogate's mean and median Δdex over time using the shared viridis color map. + """ + # Prepare layout + n = max(1, len(datasets)) + n_cols = max(1, n_cols) + n_rows = int(np.ceil(n / n_cols)) + fig, axes = plt.subplots( + n_rows, n_cols, figsize=(5 * n_cols, 3 * n_rows), sharex=False, sharey=True + ) + if isinstance(axes, np.ndarray): + axes = axes.flatten() + else: + axes = [axes] + + # Legend proxies for mean and 99th percentile styles + mean_proxy = plt.Line2D([0], [0], color="black", linestyle="-", label="Mean") + p99_proxy = plt.Line2D( + [0], [0], color="black", linestyle="--", label="99th Percentile" + ) + + surrogate_proxies = [] + surrogate_labels = [] + for name, color in color_map.items(): + surrogate_proxies.append(plt.Line2D([0], [0], color=color, label=name)) + surrogate_labels.append(name) + + for idx, (ax, dataset) in enumerate(zip(axes, datasets)): + err_dict = datasets_errors.get(dataset, {}) + if not err_dict: + ax.text(0.5, 0.5, f"No data for {dataset}", ha="center", va="center") + ax.set_axis_off() + continue + + # Assume same T across surrogates within a dataset + any_arr = next(iter(err_dict.values())) + T = any_arr.shape[1] + + # Aggregate across surrogates for percentile bands + pooled = [] + for arr in err_dict.values(): + if arr.shape[1] != T: + continue + pooled.append(arr) + + # Plot mean and 99th percentile per surrogate + for model_name, color in color_map.items(): + if model_name not in err_dict: + continue + arr = err_dict[model_name] + if arr.shape[1] != T: + continue + mean_ts = np.mean(arr, axis=(0, 2)) + p99_ts = np.percentile(arr, 99, axis=(0, 2)) + ax.plot(timesteps, mean_ts, color=color, linestyle="-", linewidth=1.2) + ax.plot(timesteps, p99_ts, color=color, linestyle="--", linewidth=1.0) + + ax.set_xscale("log") + ax.set_xlim(left=max(1, timesteps[0]), right=timesteps[-1]) + # Y label only on first column + if (idx % n_cols) == 0: + ax.set_ylabel(r"$\Delta dex$") + ax.set_ylim(2 * 1e-2, 20) + ax.set_title(_format_dataset_title(dataset)) + ax.set_yscale("log") + ax.grid(False) + + # Label bottom row x-axis + for ax in axes[-n_cols:]: + ax.set_xlabel("Time (y)") + + # Turn off tick labels on upper rows if multiple rows + if n_rows > 1: + for ax in axes[:-n_cols]: + ax.set_xticklabels([]) + + # Build combined legend below plots: surrogates + mean/99th style proxies + handles = surrogate_proxies + [mean_proxy, p99_proxy] + labels = surrogate_labels + ["Mean", "99th Percentile"] + + handles, labels, legend_ncol = reorder_legend_entries_rowwise(handles, labels, 2) + + fig.legend( + handles, + labels, + loc="lower center", + bbox_to_anchor=(0.52, 0.04), + fontsize="small", + frameon=True, + ncol=legend_ncol, + ) + + plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) + + out_path = "scripts/pp/error_percentiles_deltadex_by_dataset.png" + os.makedirs(os.path.dirname(out_path), exist_ok=True) + fig.savefig(out_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + + +def _load_catastrophic_recall( + root: str, dataset: str, percentile: int +) -> np.ndarray | None: + """Load catastrophic recall matrix saved as npz for a dataset. + + Expected path: scripts/pp//catastrophic_recall_.npz + Returns array of shape [S, F] (surrogates x flag fractions), or None if missing. + """ + path = os.path.join(root, dataset, f"catastrophic_recall_{percentile}.npz") + if not os.path.exists(path): + return None + data = np.load(path, allow_pickle=True) + try: + if "arr_0" in data.files: + arr = data["arr_0"] + else: + # Fallback to the first entry + arr = data[data.files[0]] + if isinstance(arr, np.ndarray) and arr.ndim == 2: + return arr + finally: + data.close() + return None + + +def plot_grid_catastrophic_detection( + datasets: List[str], + datasets_errors: Dict[str, Dict[str, np.ndarray]], + color_map: Dict[str, Tuple[float, float, float, float]], + root: str, + recall_percentile: int = 99, + flag_fractions: Tuple[float, ...] = ( + 0.0, + 0.025, + 0.05, + 0.10, + 0.20, + 0.30, + 0.40, + 0.50, + ), + dpi: int = 300, + n_cols: int = 2, +): + """Create a grid of catastrophic error detection curves across datasets. + + Plots recall (%) vs fraction flagged (%) for each surrogate using precomputed + recall matrices (90 or 99). Uses consistent color mapping and a single legend below. + """ + # Layout + n = max(1, len(datasets)) + n_cols = max(1, n_cols) + n_rows = int(np.ceil(n / n_cols)) + fig, axes = plt.subplots( + n_rows, n_cols, figsize=(5 * n_cols, 3 * n_rows), sharex=False, sharey=True + ) + if isinstance(axes, np.ndarray): + axes = axes.flatten() + else: + axes = [axes] + + # Precompute legend proxies for surrogates in desired order + leg_handles: List = [] + leg_labels: List[str] = [] + for name, color in color_map.items(): + leg_handles.append(plt.Line2D([0], [0], color=color, marker="o", label=name)) + leg_labels.append(name) + + for idx, (ax, dataset) in enumerate(zip(axes, datasets)): + err_dict = datasets_errors.get(dataset, {}) + recall_mat = _load_catastrophic_recall(root, dataset, recall_percentile) + if not err_dict or recall_mat is None: + ax.text(0.5, 0.5, f"No recall data for {dataset}", ha="center", va="center") + ax.set_axis_off() + continue + + # Surrogate names order used when recall was saved + saved_names = list(err_dict.keys()) + F = recall_mat.shape[1] + # Build X values. Prefer the canonical fractions if lengths match; fallback to linear spacing + if F == len(flag_fractions): + xs = np.array(flag_fractions, dtype=float) * 100.0 + else: + xs = np.linspace(0.0, 100.0 * max(flag_fractions), F) + + # Plot curves in our global color order, mapping into the saved row index + for name, color in color_map.items(): + if name not in saved_names: + continue + row = saved_names.index(name) + ys = recall_mat[row, : len(xs)] * 100.0 + ax.plot(xs, ys, marker="o", color=color, linewidth=1.2) + + ax.set_ylabel("Catastrophic error recall (%)" if (idx % n_cols) == 0 else "") + ax.set_xlim(0, float(xs.max())) + ax.set_ylim(0, 100) + if (idx % n_cols) == 0: + ax.set_ylabel("Catastrophic error recall (%)") + ax.grid(True, alpha=0.3) + ax.set_title(_format_dataset_title(dataset)) + + # Label bottom row x-axis + for ax in axes[-n_cols:]: + ax.set_xlabel("Flagged fraction (%)") + + # Turn off tick labels on upper rows if multiple rows + if n_rows > 1: + for ax in axes[:-n_cols]: + ax.set_xticklabels([]) + + # Figure-level legend + handles, labels, legend_ncol = reorder_legend_entries_rowwise( + leg_handles, leg_labels, 2 + ) + fig.legend( + handles, + labels, + loc="lower center", + bbox_to_anchor=(0.52, 0.05), + fontsize="small", + frameon=True, + ncol=legend_ncol, + ) + + plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) + + out_path = f"scripts/pp/catastrophic_detection_{recall_percentile}_by_dataset.png" + os.makedirs(os.path.dirname(out_path), exist_ok=True) + fig.savefig(out_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + + +def main(): + parser = argparse.ArgumentParser( + description="2x2 comparative Δdex error distributions across datasets" + ) + parser.add_argument( + "--root", + type=str, + default="scripts/pp", + help="Root directory containing per-dataset subdirectories", + ) + parser.add_argument( + "--cols", + type=int, + default=2, + help="Number of columns in the subplot grid", + ) + parser.add_argument("--dpi", type=int, default=300, help="Figure DPI") + + args = parser.parse_args() + + # Auto-discover datasets: list subdirectories under root that contain an NPZ + if not os.path.isdir(args.root): + raise NotADirectoryError(args.root) + + all_subdirs = sorted( + d for d in os.listdir(args.root) if os.path.isdir(os.path.join(args.root, d)) + ) + datasets: List[str] = [] + datasets_errors: Dict[str, Dict[str, np.ndarray]] = {} + for ds in all_subdirs: + # Only include if an NPZ exists + npz1 = os.path.join(args.root, ds, "all_log_errors.npz") + npz2 = os.path.join(args.root, ds, "all_errors_log.npz") + if not (os.path.exists(npz1) or os.path.exists(npz2)): + continue + try: + datasets_errors[ds] = load_dataset_errors(args.root, ds) + datasets.append(ds) + except Exception: + # Skip subdirs with unreadable NPZs + continue + timesteps_path = os.path.join(args.root, "timesteps.npz") + timesteps = np.load(timesteps_path)["arr_0"] + + if not datasets: + raise RuntimeError( + f"No datasets with error NPZs found under root '{args.root}'." + ) + + # Global x-range and consistent colors across all surrogates + x_min, x_max = compute_global_range(datasets_errors) + color_map, _ = build_color_map(datasets_errors) + + plot_grid_deltadex( + datasets=datasets, + datasets_errors=datasets_errors, + x_log_min=x_min, + x_log_max=x_max, + color_map=color_map, + dpi=args.dpi, + n_cols=args.cols, + ) + + # Percentiles-over-time grid (deltadex mode) + plot_grid_deltadex_percentiles( + datasets=datasets, + datasets_errors=datasets_errors, + timesteps=timesteps, + color_map=color_map, + dpi=args.dpi, + n_cols=args.cols, + ) + + # Catastrophic detection grid + plot_grid_catastrophic_detection( + datasets=datasets, + datasets_errors=datasets_errors, + color_map=color_map, + root=args.root, + recall_percentile=99, + dpi=args.dpi, + n_cols=args.cols, + ) + + plot_grid_catastrophic_detection( + datasets=datasets, + datasets_errors=datasets_errors, + color_map=color_map, + root=args.root, + recall_percentile=90, + dpi=args.dpi, + n_cols=args.cols, + ) + + +if __name__ == "__main__": + main() From c1b90d4eeb8743865e94d4ee4c98fa50a733436b Mon Sep 17 00:00:00 2001 From: Robin Date: Mon, 22 Dec 2025 15:11:59 +0100 Subject: [PATCH 2/6] additional plots --- codes/benchmark/bench_fcts.py | 19 +- codes/benchmark/bench_plots.py | 1 + .../AbstractSurrogate/abstract_surrogate.py | 9 + codes/tune/evaluate_tuning.py | 583 +++++++++++++-- config.yaml | 2 +- config_full.yaml | 2 +- paper_eval.py | 2 +- paper_plots.py | 663 +++++++++++++++++- 8 files changed, 1202 insertions(+), 79 deletions(-) diff --git a/codes/benchmark/bench_fcts.py b/codes/benchmark/bench_fcts.py index f591619..73a587a 100644 --- a/codes/benchmark/bench_fcts.py +++ b/codes/benchmark/bench_fcts.py @@ -381,7 +381,6 @@ def evaluate_iterative_predictions( # container for the piecewise predictions; seed t=0 with ground truth so errors # are computed only on actual predictions for t>=1 while keeping shape intact iterative_preds = np.zeros_like(targets) - iterative_preds[:, 0, :] = targets[:, 0, :] # number of chunks n_chunks = (n_timesteps + iter_interval - 1) // iter_interval @@ -436,14 +435,17 @@ def evaluate_iterative_predictions( ) # We predict steps 1..(chunk_len-1) relative to the provided init state (index 0). # Map these to global indices [start+1 .. end] inclusively. + if i == 0: + iterative_preds[:, start : end + 1, :] = preds_chunk[:, : model.n_timesteps, :].detach().cpu().numpy() iterative_preds[:, start + 1 : end + 1, :] = ( preds_chunk[:, 1 : model.n_timesteps, :].detach().cpu().numpy() ) iterative_preds_log = model.denormalize(iterative_preds, leave_log=True) + full_preds_log = model.denormalize(full_preds, leave_log=True) targets_log = model.denormalize(targets, leave_log=True) iterative_preds = model.denormalize(iterative_preds) - full_preds = model.denormalize(full_preds.detach().cpu().numpy()) + full_preds_real = model.denormalize(full_preds.detach().cpu().numpy()) targets = model.denormalize(targets) # compute error metrics @@ -469,7 +471,7 @@ def evaluate_iterative_predictions( surr_name, conf, iterative_preds, - full_preds, + full_preds_real, targets, timesteps, iter_interval=iter_interval, @@ -1313,6 +1315,11 @@ def compare_iterative(metrics: dict[str, dict], config: dict) -> None: iterative_errors[surrogate], axis=(0, 2) ) + # TEMP + dataset = config["dataset"]["name"] + os.makedirs(f"scripts/pp/{dataset}", exist_ok=True) + np.savez(f"scripts/pp/{dataset}/all_iterative_errors.npz", iterative_errors) + plot_errors_over_time( mean_iterative_errors, median_iterative_errors, @@ -1595,6 +1602,12 @@ def compare_UQ(all_metrics: dict, config: dict) -> None: show_title=True, ) + # TEMP + dataset = config["dataset"]["name"] + os.makedirs(f"scripts/pp/{dataset}", exist_ok=True) + np.savez(f"scripts/pp/{dataset}/all_uq_errors.npz", ensemble_errors) + np.savez(f"scripts/pp/{dataset}/all_uq_std.npz", ensemble_std) + def tabular_comparison(all_metrics: dict, config: dict) -> None: """ diff --git a/codes/benchmark/bench_plots.py b/codes/benchmark/bench_plots.py index 97cd8ac..26091c9 100644 --- a/codes/benchmark/bench_plots.py +++ b/codes/benchmark/bench_plots.py @@ -1661,6 +1661,7 @@ def plot_errors_over_time( elif mode == "iterative": # Single backslash inside raw string to render the LaTeX Delta properly plt.ylabel(r"Log-MAE ($\Delta dex$)") + plt.ylim(bottom=0, top=min(np.max(list(mean_errors.values())) * 1.1, 5)) fname = "iterative_delta_dex_time.png" title = "Comparison of Δdex Errors Over Time for Iterative Predictions" # Add subtle dashed vertical lines at every n-th timestep if provided and valid diff --git a/codes/surrogates/AbstractSurrogate/abstract_surrogate.py b/codes/surrogates/AbstractSurrogate/abstract_surrogate.py index 8c62d7b..4b2a735 100644 --- a/codes/surrogates/AbstractSurrogate/abstract_surrogate.py +++ b/codes/surrogates/AbstractSurrogate/abstract_surrogate.py @@ -449,8 +449,10 @@ def denormalize( Returns: Tensor | np.ndarray: The denormalized data. """ + data_type = None if self.normalisation is not None: if not leave_norm: + data_type = data.dtype if self.normalisation["mode"] == "disabled": ... elif self.normalisation["mode"] == "minmax": @@ -475,6 +477,13 @@ def denormalize( if self.normalisation["log10_transform"] and not leave_log: data = 10**data + # Conserve dtype + if data_type is not None: + if isinstance(data, Tensor): + return data.to(dtype=data_type) + if isinstance(data, np.ndarray): + return data.astype(data_type) + return data def denormalize_old(self, data: Tensor) -> Tensor: diff --git a/codes/tune/evaluate_tuning.py b/codes/tune/evaluate_tuning.py index a421835..d10b4df 100644 --- a/codes/tune/evaluate_tuning.py +++ b/codes/tune/evaluate_tuning.py @@ -11,6 +11,8 @@ import optuna import psycopg2 import torch +from matplotlib.lines import Line2D +from matplotlib.ticker import FormatStrFormatter from optuna.trial import TrialState from psycopg2 import sql @@ -22,9 +24,35 @@ from codes.tune import load_yaml_config from codes.utils import nice_print - -def pareto_front(points: np.ndarray) -> np.ndarray: - # lower-is-better for both objectives +os.environ.setdefault("PGCONNECT_TIMEOUT", "3") + +TRIAL_COLOR = "#c7c7c7" +PARETO_COLOR = "#1f77b4" +BEST_F1_COLOR = "#d62728" +CHOSEN_COLOR = "#2ca02c" +EDGE_COLOR = "#000000" +PARETO_X_RANGE = (0, 6) +SURROGATE_ORDER = [ + "multionet", + "fullyconnected", + "latentpoly", + "latentneuralode", +] + + +def format_surrogate_title(name: str) -> str: + mapping = { + "multionet": "MultiONet", + "fullyconnected": "FullyConnected", + "latentneuralode": "LatentNeuralODE", + "latentpoly": "LatentPoly", + } + return mapping.get(name.lower(), name) + + +def pareto_front_mask(points: np.ndarray) -> np.ndarray: + if points.size == 0: + return np.array([], dtype=bool) is_efficient = np.ones(points.shape[0], dtype=bool) for i, p in enumerate(points): if not is_efficient[i]: @@ -34,7 +62,12 @@ def pareto_front(points: np.ndarray) -> np.ndarray: dominated = better & (np.arange(points.shape[0]) != i) if np.any(dominated): is_efficient[i] = False - return points[is_efficient] + return is_efficient + + +def pareto_front(points: np.ndarray) -> np.ndarray: + mask = pareto_front_mask(points) + return points[mask] def hypervolume_2d(pareto_points: np.ndarray, reference: np.ndarray) -> float: @@ -99,13 +132,391 @@ def compute_hypervolume_over_time( return hypervolumes, reference +def _pareto_legend_handles() -> list[Line2D]: + return [ + Line2D( + [0], + [0], + marker="o", + color="none", + markerfacecolor=TRIAL_COLOR, + markeredgecolor=TRIAL_COLOR, + label="Trial Outcome", + markersize=6, + ), + Line2D( + [0], + [0], + marker="o", + color="none", + markerfacecolor=PARETO_COLOR, + markeredgecolor=EDGE_COLOR, + label="Pareto front", + markersize=6, + ), + Line2D( + [0], + [0], + marker="o", + color="none", + markerfacecolor=BEST_F1_COLOR, + markeredgecolor=EDGE_COLOR, + label="Lowest Δdex99", + markersize=6, + ), + Line2D( + [0], + [0], + marker="o", + color="none", + markerfacecolor=CHOSEN_COLOR, + markeredgecolor=EDGE_COLOR, + label="Chosen Trial", + markersize=6, + ), + ] + + +def compute_pareto_plot_data( + study: optuna.Study, + suffix: str, + ignore_last_n: int = 10, + chosen_trial_number: int | None = None, +): + if len(study.directions) != 2: + return None + completed = [t for t in study.trials if t.state == TrialState.COMPLETE] + if not completed: + print(f"Skipping Pareto plot for {suffix}: no completed trials.") + return None + completed.sort(key=lambda t: t.datetime_complete or t.datetime_start) + cutoff = max(0, len(completed) - max(0, ignore_last_n)) + eligible_trials = completed[:cutoff] + filtered = [] + for t in eligible_trials: + if t.values is None: + continue + vals = np.array(t.values, dtype=float) + if vals.shape[0] != 2 or np.any(~np.isfinite(vals)): + continue + filtered.append((vals, t.number)) + if not filtered: + print( + f"Skipping Pareto plot for {suffix}: no eligible trials after ignoring last {ignore_last_n}." + ) + return None + points = np.vstack([vals for vals, _ in filtered]) + trial_numbers = np.array([num for _, num in filtered]) + mask = pareto_front_mask(points) + best_idx = int(np.argmin(points[:, 0])) + best_point = points[best_idx] + chosen_point = None + if chosen_trial_number is not None: + matches = np.where(trial_numbers == chosen_trial_number)[0] + if matches.size: + chosen_point = points[matches[0]] + else: + print( + f"Chosen trial {chosen_trial_number} not in eligible set for {suffix}; " + "skipping highlight." + ) + + y = points[:, 1] + y_min = 0 + y_max = best_point[1] * 1.3 + + return { + "suffix": suffix, + "points": points, + "mask": mask, + "best_point": best_point, + "chosen_point": chosen_point, + "xlim": PARETO_X_RANGE, + "ylim": (y_min, y_max), + "ignore_last_n": ignore_last_n, + } + + +def summarize_pareto_tradeoff(data: dict): + suffix = data["suffix"] + best_point = data["best_point"] + chosen_point = data["chosen_point"] + pts = data["points"][data["mask"]] + x_span = float(pts[:, 0].max() - pts[:, 0].min()) + y_span = float(pts[:, 1].max() - pts[:, 1].min()) + min_x = float(pts[:, 0].min()) + min_y = float(pts[:, 1].min()) + x_rel = x_span / max(abs(min_x), 1e-12) + y_rel = y_span / max(abs(min_y), 1e-12) + print( + f"{suffix}: Pareto span Δdex={x_span:.4f} " + f"(x{ x_rel:.2f} relative), inference time span={y_span:.4f}s " + f"(x{ y_rel:.2f} relative)" + ) + best_err, best_time = float(best_point[0]), float(best_point[1]) + print( + f"{suffix}: lowest-error trial Δdex={best_err:.4f}, inference time={best_time:.4f}s" + ) + if chosen_point is None: + print( + f"{suffix}: chosen trial not available after filtering; skipping summary." + ) + return + chosen_err, chosen_time = float(chosen_point[0]), float(chosen_point[1]) + print( + f"{suffix}: chosen trial Δdex={chosen_err:.4f}, inference time={chosen_time:.4f}s" + ) + err_reduction = 1.0 - best_err / chosen_err if chosen_err > 0 else float("nan") + time_ratio = best_time / chosen_time if chosen_time > 0 else float("nan") + if math.isnan(err_reduction) or math.isnan(time_ratio): + print(f"{suffix}: insufficient data to compute tradeoff summary.") + return + err_phrase = ( + f"{abs(err_reduction) * 100:.1f}% lower error" + if err_reduction >= 0 + else f"{abs(err_reduction) * 100:.1f}% higher error" + ) + ratio_phrase = ( + f"but {time_ratio:.2f}x higher inference time" + if time_ratio >= 1 + else f"and {time_ratio:.2f}x lower inference time" + ) + print(f"{suffix}: {err_phrase}, {ratio_phrase}.") + + +def _render_pareto_scatter( + ax, + data, + show_xlabel: bool, + show_ylabel: bool, + title: str | None, + hide_xticklabels: bool = False, +): + pts = data["points"] + mask = data["mask"] + best_point = data["best_point"] + chosen_point = data["chosen_point"] + + ax.scatter(pts[:, 0], pts[:, 1], color=TRIAL_COLOR, alpha=0.7, label=None) + ax.scatter( + pts[mask, 0], + pts[mask, 1], + color=PARETO_COLOR, + edgecolor=EDGE_COLOR, + linewidth=0.5, + label=None, + ) + ax.scatter( + best_point[0], + best_point[1], + color=BEST_F1_COLOR, + edgecolor=EDGE_COLOR, + linewidth=0.6, + label=None, + zorder=3, + ) + if chosen_point is not None: + ax.scatter( + chosen_point[0], + chosen_point[1], + color=CHOSEN_COLOR, + edgecolor=EDGE_COLOR, + linewidth=0.6, + label=None, + zorder=4, + ) + + ax.set_xlim(*data["xlim"]) + ax.set_ylim(*data["ylim"]) + if show_xlabel: + ax.set_xlabel("99th-percentile Δdex") + if show_ylabel: + ax.set_ylabel("inference time (s)") + if title: + ax.set_title(title) + ax.tick_params(labelbottom=not hide_xticklabels) + ax.grid(True, linestyle="--", alpha=0.3) + + ax.yaxis.set_major_formatter(FormatStrFormatter("%.3f")) + + +def save_individual_pareto_plot(data, out_dir: str): + fig, ax = plt.subplots(figsize=(5, 3)) + _render_pareto_scatter( + ax, data, show_xlabel=True, show_ylabel=True, title=data["suffix"] + ) + ax.legend(handles=_pareto_legend_handles(), loc="best", ncol=2) + fig.tight_layout() + os.makedirs(out_dir, exist_ok=True) + fig.savefig(os.path.join(out_dir, f"pareto_front_{data['suffix']}.png"), dpi=300) + plt.close(fig) + print( + f"Saved Pareto front plot for {data['suffix']}; " + f"ignored last {data['ignore_last_n']} completed trials." + ) + + +def save_pareto_front_grid(datasets: list[dict], out_dir: str): + if not datasets: + return + n_cols = 2 + n_rows = max(1, math.ceil(len(datasets) / n_cols)) + fig, axes = plt.subplots( + n_rows, n_cols, figsize=(5 * n_cols, 3 * n_rows), squeeze=False + ) + + for idx, data in enumerate(datasets): + row, col = divmod(idx, n_cols) + ax = axes[row][col] + show_xlabel = row == n_rows - 1 + show_ylabel = col == 0 + proper_title = format_surrogate_title(data["suffix"]) + _render_pareto_scatter( + ax, + data, + show_xlabel=show_xlabel, + show_ylabel=show_ylabel, + title=proper_title, + hide_xticklabels=not show_xlabel, + ) + for idx in range(len(datasets), n_rows * n_cols): + row, col = divmod(idx, n_cols) + axes[row][col].axis("off") + + fig.tight_layout(rect=(0.01, 0.035, 0.99, 0.99)) + fig.legend( + handles=_pareto_legend_handles(), + loc="lower center", + ncol=2, + bbox_to_anchor=(0.53, 0.0), + ) + os.makedirs(out_dir, exist_ok=True) + fig.savefig(os.path.join(out_dir, "pareto_front_grid.png"), dpi=300) + plt.close(fig) + print( + f"Saved combined Pareto front grid with {len(datasets)} subplot(s) " + f"to {os.path.join(out_dir, 'pareto_front_grid.png')}." + ) + + +def save_relative_hv_grid(datasets: list[dict], out_dir: str): + if not datasets: + return + n_cols = 2 + n_rows = max(1, math.ceil(len(datasets) / n_cols)) + fig, axes = plt.subplots( + n_rows, n_cols, figsize=(5 * n_cols, 3 * n_rows), squeeze=False + ) + + for idx, data in enumerate(datasets): + row, col = divmod(idx, n_cols) + ax = axes[row][col] + show_xlabel = row == n_rows - 1 + show_ylabel = col == 0 + ax.plot(data["x"], data["y"], color="#1f77b4") + ax.set_xlim(0, max(data["x"]) if len(data["x"]) else 1) + ax.set_ylim(0.6, 1.03) + ax.set_title(format_surrogate_title(data["suffix"])) + if show_xlabel: + ax.set_xlabel("Completed Trials") + if show_ylabel: + ax.set_ylabel("Fraction of Final HV") + else: + ax.tick_params(labelleft=False) + ax.grid(True, linestyle="--", alpha=0.3) + + for idx in range(len(datasets), n_rows * n_cols): + row, col = divmod(idx, n_cols) + axes[row][col].axis("off") + + fig.tight_layout(rect=(0.01, 0.01, 0.99, 0.99)) + os.makedirs(out_dir, exist_ok=True) + fig.savefig(os.path.join(out_dir, "hypervolume_relative_grid.png"), dpi=300) + plt.close(fig) + print( + f"Saved combined relative-hypervolume grid with {len(datasets)} subplot(s) " + f"to {os.path.join(out_dir, 'hypervolume_relative_grid.png')}." + ) + + +def save_relative_hv_combined(datasets: list[dict], out_dir: str): + if len(datasets) < 2: + return + fig, ax = plt.subplots(figsize=(6, 4)) + color_palette = plt.cm.viridis(np.linspace(0, 0.95, len(SURROGATE_ORDER))) + order_index = {name: idx for idx, name in enumerate(SURROGATE_ORDER)} + + def color_for(name: str): + key = name.lower() + idx = order_index.get(key) + if idx is not None: + return color_palette[idx] + return plt.cm.tab10(0) + + ordered = sorted( + datasets, + key=lambda d: order_index.get(d["suffix"].lower(), len(SURROGATE_ORDER)), + ) + + for data in ordered: + color = color_for(data["suffix"]) + ax.plot( + data["x"], + data["y"], + label=format_surrogate_title(data["suffix"]), + color=color, + ) + if len(data["x"]): + last_x = data["x"][-1] + last_y = data["y"][-1] + ax.vlines( + last_x, + max(0, last_y - 0.01), + min(1.05, last_y + 0.01), + colors=color, + linewidth=2, + ) + ax.set_xlim(0, max(max(d["x"]) for d in datasets if len(d["x"])) if datasets else 1) + ax.set_ylim(0.6, 1.02) + ax.set_xlabel("Completed Trials") + ax.set_ylabel("Fraction of Final HV") + ax.set_title("Relative Hypervolume") + ax.grid(True, linestyle="--", alpha=0.3) + ax.legend(ncol=2, loc="lower right") + fig.tight_layout() + os.makedirs(out_dir, exist_ok=True) + fig.savefig(os.path.join(out_dir, "hypervolume_relative_combined.png"), dpi=300) + plt.close(fig) + print( + "Saved combined relative-hypervolume line plot to " + f"{os.path.join(out_dir, 'hypervolume_relative_combined.png')}." + ) + + def load_loss_history(model_path: str) -> tuple[np.ndarray, np.ndarray, int]: """ - Load loss histories from a saved model file (.pth). - Returns (train_loss, test_loss, n_epochs). + Load loss histories from a saved .pth. + If the checkpoint contains pickled Optuna storage (which may try to connect + to a remote DB during unpickle), bound the timeout and fall back to skipping. """ - model_dict = torch.load(model_path, map_location="cpu", weights_only=False) - attributes = model_dict.get("attributes", {}) + # Prefer the Tensor-only safe path if your PyTorch supports it (PyTorch >= 2.0) + try: + obj = torch.load(model_path, map_location="cpu", weights_only=True) + # weights_only returns just state_dict; no custom 'attributes' available + # We can’t recover losses from state_dict → skip gracefully. + return None, None, 0 + except TypeError: + # weights_only not supported → careful unpickle with timeout + catch + pass + + try: + obj = torch.load(model_path, map_location="cpu") + except Exception as e: + # Any error (including timeout from unpickling Optuna storage) → skip losses + print(f"[warn] Could not safely load {model_path}: {e}. Skipping loss curves.") + return None, None, 0 + + attributes = obj.get("attributes", {}) if isinstance(obj, dict) else {} train_loss = ( np.array(attributes.get("train_loss")) if attributes.get("train_loss") is not None @@ -219,6 +630,7 @@ def evaluate_tuning( top_n: int = 10, storage_name: str = "optuna_db", ignore_last_n: int = 10, + chosen_indices: list[int] | None = None, ) -> None: """ For all surrogate studies named '_', @@ -243,6 +655,8 @@ def evaluate_tuning( connect_timeout=5, ).close() except Exception: + if str(pg.get("host", "localhost")) not in ("localhost", "127.0.0.1", "::1"): + raise # don’t try to start a local server if host isn't local pg_data = pg.get("data_dir", os.path.expanduser("~/postgres/data")) pg_ctl = os.path.join( pg.get("database_folder", os.path.expanduser("~/postgres")), "bin", "pg_ctl" @@ -276,10 +690,10 @@ def evaluate_tuning( from optuna.study import get_all_study_summaries summaries = get_all_study_summaries(storage=storage_url) - study_names = [ + study_names_all = [ s.study_name for s in summaries if s.study_name.startswith(f"{study_prefix}_") ] - if not study_names: + if not study_names_all: print(f"No studies found with prefix '{study_prefix}_' in {storage_url}") return @@ -287,8 +701,41 @@ def evaluate_tuning( save_dir = os.path.join("tuned", study_prefix) os.makedirs(save_dir, exist_ok=True) + # Derive deterministic ordering from config surrogate list if available + surrogate_entries = config.get("surrogates", []) + surrogate_names = [ + str(entry.get("name")) + for entry in surrogate_entries + if isinstance(entry, dict) and entry.get("name") + ] + study_names_ordered = [] + used: set[str] = set() + study_lookup = {name.lower(): name for name in study_names_all} + if surrogate_names: + for surrogate in surrogate_names: + candidate = f"{study_prefix}_{surrogate}".lower() + match = study_lookup.get(candidate) + if match and match not in used: + study_names_ordered.append(match) + used.add(match) + for name in sorted(study_names_all): + if name not in used: + study_names_ordered.append(name) + used.add(name) + study_names = study_names_ordered + + if chosen_indices is not None and len(chosen_indices) != len(study_names): + print( + "Provided chosen_indices length does not match number of studies; " + "skipping chosen-trial highlighting." + ) + chosen_indices = None + + pareto_datasets: list[dict] = [] + hypervolume_datasets: list[dict] = [] + # Loop over each surrogate study - for full_name in study_names: + for idx, full_name in enumerate(study_names): suffix = full_name[len(study_prefix) + 1 :] print(f"--- Evaluating study {full_name} -> surrogate '{suffix}' ---") try: @@ -336,12 +783,33 @@ def evaluate_tuning( dpi=300, ) plt.close() + max_x = max(1, len(rel_hvs) - max(0, ignore_last_n)) + hypervolume_datasets.append( + { + "suffix": suffix, + "x": np.arange(1, len(rel_hvs) + 1)[:max_x], + "y": rel_hvs[:max_x], + } + ) print( f"Saved hypervolume plots for {suffix} (final HV={final_hv:.3e}); " f"ignored last {ignore_last_n} trials in Pareto/HV." ) else: print("No hypervolume computed (no complete trials).") + chosen_trial = None + if chosen_indices is not None: + chosen_trial = chosen_indices[idx] + plot_data = compute_pareto_plot_data( + study, + suffix, + ignore_last_n=ignore_last_n, + chosen_trial_number=chosen_trial, + ) + if plot_data: + save_individual_pareto_plot(plot_data, save_dir) + pareto_datasets.append(plot_data) + summarize_pareto_tradeoff(plot_data) else: print("Skipping hypervolume: study is not two-objective.") @@ -369,40 +837,44 @@ def evaluate_tuning( continue models_dir = os.path.join(models_root, match) - # Collect top trials' loss histories - records = [] - epochs = None - for fname in os.listdir(models_dir): - if not fname.endswith(".pth"): - continue - m = re.search(r"_(\d+)\.pth$", fname) - if not m: - continue - tnum = int(m.group(1)) - if tnum not in best: - continue - path = os.path.join(models_dir, fname) - train, test, ep = load_loss_history(path) - if test is None: - continue - records.append((tnum, test)) - epochs = ep - if not records: - print(f"No top-{top_n} trials for surrogate '{suffix}'") - continue - records.sort(key=lambda x: best.index(x[0])) - tnums, losses = zip(*records) - labels = [f"Trial {n}" for n in tnums] - out_dir = save_dir - plot_losses( - list(losses), - epochs, - labels, - title=f"{suffix} Top-{top_n}", - save=True, - out_dir=out_dir, - mode=suffix, - ) + # # Collect top trials' loss histories + # records = [] + # epochs = None + # for fname in os.listdir(models_dir): + # if not fname.endswith(".pth"): + # continue + # m = re.search(r"_(\d+)\.pth$", fname) + # if not m: + # continue + # tnum = int(m.group(1)) + # if tnum not in best: + # continue + # path = os.path.join(models_dir, fname) + # train, test, ep = load_loss_history(path) + # if test is None: + # continue + # records.append((tnum, test)) + # epochs = ep + # if not records: + # print(f"No top-{top_n} trials for surrogate '{suffix}'") + # continue + # records.sort(key=lambda x: best.index(x[0])) + # tnums, losses = zip(*records) + # labels = [f"Trial {n}" for n in tnums] + # out_dir = save_dir + # plot_losses( + # list(losses), + # epochs, + # labels, + # title=f"{suffix} Top-{top_n}", + # save=True, + # out_dir=out_dir, + # mode=suffix, + # ) + + save_pareto_front_grid(pareto_datasets, save_dir) + save_relative_hv_grid(hypervolume_datasets, save_dir) + save_relative_hv_combined(hypervolume_datasets, save_dir) def parse_args(): @@ -435,16 +907,37 @@ def parse_args(): "Number of most-recent completed trials to exclude from Pareto/hypervolume" ), ) + p.add_argument( + "--chosen_indices", + type=str, + # default="171,114,135,237", # cloud_final + # default="27,61,13,299", # cloud_parametric_final + default="18,1,16,234", # primordial_parametric_final + # default="196,107,31,243", # primordial_final + help=( + "Comma-separated list of Optuna trial numbers chosen per study (order " + "matches sorted study names)." + ), + ) return p.parse_args() def main(): args = parse_args() + chosen_list = None + if args.chosen_indices: + chosen_list = [] + for item in args.chosen_indices.split(","): + item = item.strip() + if not item: + continue + chosen_list.append(int(item)) evaluate_tuning( args.study_name, args.top_n, args.storage_name, ignore_last_n=args.ignore_last_n, + chosen_indices=chosen_list, ) diff --git a/config.yaml b/config.yaml index e1df0a6..4d10e24 100644 --- a/config.yaml +++ b/config.yaml @@ -11,7 +11,7 @@ dataset: tolerance: 1e-25 normalise_per_species: True log_timesteps: True -devices: ["cuda:0", "cuda:2", "cuda:3", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"] +devices: ["cuda:4", "cuda:2", "cuda:3", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"] seed: 42 verbose: False checkpoint: True diff --git a/config_full.yaml b/config_full.yaml index e20ee71..b578f98 100644 --- a/config_full.yaml +++ b/config_full.yaml @@ -13,7 +13,7 @@ dataset: tolerance: 1e-25 subset_factor: 1 log_timesteps: True -devices: ["cuda:0", "cuda:1", "cuda:2", "cuda:3"] # ["cuda:0", "cuda:1", "cuda:2", "cuda:3", "cuda:4", "cuda:5", "cuda:6", "cuda:7", "cuda:8"] +devices: ["cuda:1", "cuda:1", "cuda:2", "cuda:3"] # ["cuda:0", "cuda:1", "cuda:2", "cuda:3", "cuda:4", "cuda:5", "cuda:6", "cuda:7", "cuda:8"] seed: 42 verbose: False relative_error_threshold: 1e-10 diff --git a/paper_eval.py b/paper_eval.py index b89be1d..72d3565 100644 --- a/paper_eval.py +++ b/paper_eval.py @@ -22,7 +22,7 @@ from codes.utils import download_data, nice_print, read_yaml_config # Hardcoded device override for all evaluations -DEVICE = "cuda:0" +DEVICE = "cuda:2" # Training IDs to evaluate TRAINING_IDS = [ diff --git a/paper_plots.py b/paper_plots.py index e64c011..68c7419 100644 --- a/paper_plots.py +++ b/paper_plots.py @@ -21,12 +21,17 @@ from __future__ import annotations import argparse +import csv import os +import re +import warnings from typing import Dict, List, Tuple import matplotlib.pyplot as plt import numpy as np +import yaml from scipy.ndimage import gaussian_filter1d +from sklearn.metrics import roc_auc_score # Reuse palette from project for consistent styling, but keep a safe fallback try: @@ -158,6 +163,153 @@ def build_color_map(datasets_errors: Dict[str, Dict[str, np.ndarray]]): return color_map, names +SURROGATE_YAML_FILENAMES: Dict[str, str] = { + "LatentNeuralODE": "latentneuralode_metrics.yaml", + "MultiONet": "multionet_metrics.yaml", + "FullyConnected": "fullyconnected_metrics.yaml", + "LatentPoly": "latentpoly_metrics.yaml", +} + + +def _extract_first_number(token: str) -> float | None: + """Return the first numeric value found in a label like 'interval 4'.""" + matches = re.findall(r"[-+]?\d*\.?\d+", token) + if not matches: + return None + try: + return float(matches[0]) + except ValueError: + return None + + +def load_all_surrogate_metrics( + root: str, datasets: List[str], surrogates: List[str] +) -> Dict[str, Dict[str, dict]]: + """Load per-surrogate YAML metric dictionaries for each dataset.""" + results: Dict[str, Dict[str, dict]] = {} + for dataset in datasets: + dataset_metrics: Dict[str, dict] = {} + for surrogate in surrogates: + yaml_name = SURROGATE_YAML_FILENAMES.get(surrogate) + if not yaml_name: + continue + path = os.path.join(root, dataset, yaml_name) + if not os.path.exists(path): + continue + try: + with open(path, "r", encoding="utf-8") as fh: + data = yaml.safe_load(fh) or {} + except Exception: + continue + if isinstance(data, dict): + dataset_metrics[surrogate] = data + results[dataset] = dataset_metrics + return results + + +def _extract_interpolation_series(metrics: dict) -> Tuple[np.ndarray, np.ndarray]: + section = metrics.get("interpolation") + if not isinstance(section, dict): + return np.array([]), np.array([]) + + entries: List[Tuple[float, float]] = [] + for label, payload in section.items(): + if not isinstance(payload, dict): + continue + position = _extract_first_number(label) + mae = payload.get("MAE_log") + if position is None or mae is None: + continue + try: + entries.append((float(position), float(mae))) + except (TypeError, ValueError): + continue + + if not entries: + return np.array([]), np.array([]) + + entries.sort(key=lambda item: item[0]) + xs, ys = zip(*entries) + return np.asarray(xs, dtype=float), np.asarray(ys, dtype=float) + + +def _extract_extrapolation_series(metrics: dict) -> Tuple[np.ndarray, np.ndarray]: + section = metrics.get("extrapolation") + if not isinstance(section, dict): + return np.array([]), np.array([]) + + entries: List[Tuple[float, float]] = [] + for label, payload in section.items(): + if not isinstance(payload, dict): + continue + cutoff = _extract_first_number(label) + mae = payload.get("MAE_log") + if cutoff is None or mae is None: + continue + try: + entries.append((float(cutoff), float(mae))) + except (TypeError, ValueError): + continue + + if not entries: + return np.array([]), np.array([]) + + entries.sort(key=lambda item: item[0]) + xs, ys = zip(*entries) + return np.asarray(xs, dtype=float), np.asarray(ys, dtype=float) + + +def _extract_sparse_series(metrics: dict) -> Tuple[np.ndarray, np.ndarray]: + section = metrics.get("sparse") + if not isinstance(section, dict): + return np.array([]), np.array([]) + + entries: List[Tuple[float, float]] = [] + for payload in section.values(): + if not isinstance(payload, dict): + continue + mae = payload.get("MAE_log") + samples = payload.get("n_train_samples") + if mae is None or samples is None: + continue + try: + entries.append((float(samples), float(mae))) + except (TypeError, ValueError): + continue + + if not entries: + return np.array([]), np.array([]) + + entries.sort(key=lambda item: item[0]) + xs, ys = zip(*entries) + return np.asarray(xs, dtype=float), np.asarray(ys, dtype=float) + + +def collect_modality_series( + dataset_metrics: Dict[str, Dict[str, dict]], modality: str +) -> Dict[str, Dict[str, Tuple[np.ndarray, np.ndarray]]]: + """Convert raw YAML metric dictionaries into plottable series per modality.""" + extractor_map = { + "interpolation": _extract_interpolation_series, + "extrapolation": _extract_extrapolation_series, + "sparse": _extract_sparse_series, + } + extractor = extractor_map.get(modality) + if extractor is None: + raise ValueError(f"Unknown modality '{modality}'") + + modality_data: Dict[str, Dict[str, Tuple[np.ndarray, np.ndarray]]] = {} + for dataset, surrogate_metrics in dataset_metrics.items(): + per_surrogate: Dict[str, Tuple[np.ndarray, np.ndarray]] = {} + for surrogate, metrics in surrogate_metrics.items(): + xs, ys = extractor(metrics) + if xs.size == 0 or ys.size == 0: + continue + per_surrogate[surrogate] = (xs, ys) + modality_data[dataset] = per_surrogate + return modality_data + + def reorder_legend_entries_rowwise( handles: List, labels: List[str], max_ncols: int ) -> Tuple[List, List[str], int]: @@ -191,6 +343,132 @@ def reorder_legend_entries_rowwise( return final_handles, final_labels, legend_ncol +def plot_modality_comparison_grid( + datasets: List[str], + modality_data: Dict[str, Dict[str, Tuple[np.ndarray, np.ndarray]]], + color_map: Dict[str, Tuple[float, float, float, float]], + xlabel: str, + filename: str, + dpi: int = 300, + n_cols: int = 2, + xlog: bool = False, + linear_fit: bool = True, +) -> None: + """Render a grid comparing Log-MAE trends across surrogates for one modality. + + When `linear_fit` is True, draw a least-squares fit line per surrogate. + """ + + n = max(1, len(datasets)) + n_cols = max(1, n_cols) + n_rows = int(np.ceil(n / n_cols)) + fig, axes = plt.subplots( + n_rows, + n_cols, + figsize=(5 * n_cols, 3.5 * n_rows), + sharey=True, + squeeze=False, + ) + axes_flat = axes.flatten() + + for idx, dataset in enumerate(datasets): + ax = axes_flat[idx] + per_surrogate = modality_data.get(dataset, {}) + if not per_surrogate: + ax.text(0.5, 0.5, "No metrics", ha="center", va="center") + ax.set_axis_off() + continue + + for surrogate, color in color_map.items(): + series = per_surrogate.get(surrogate) + if series is None: + continue + xs, ys = series + if xs.size == 0 or ys.size == 0: + continue + ax.scatter(xs, ys, color=color, s=25) + ax.plot(xs, ys, color=color, alpha=0.6, linestyle="-") + + if linear_fit: + mask = np.isfinite(xs) & np.isfinite(ys) + if xlog: + mask &= xs > 0 + if np.count_nonzero(mask) >= 2: + fit_x = xs[mask] + fit_y = ys[mask] + if xlog: + fit_x = np.log10(fit_x) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", np.RankWarning) + coeffs = np.polyfit(fit_x, fit_y, deg=1) + # print(f"Dataset '{dataset}', Surrogate '{surrogate}': fit coeffs {coeffs}") + fit_domain = np.linspace(fit_x.min(), fit_x.max(), 100) + fit_values = np.polyval(coeffs, fit_domain) + if xlog: + fit_domain = np.power(10.0, fit_domain) + ax.plot( + fit_domain, + fit_values, + color=color, + linestyle="--", + linewidth=1.1, + alpha=0.9, + ) + + if xlog: + ax.set_xscale("log") + # ax.set_yscale("log") + ax.grid(True, which="major", linestyle="--", linewidth=0.5, alpha=0.4) + if (idx % n_cols) == 0: + ax.set_ylabel(r"Log-MAE($\Delta dex$)") + ax.set_title(_format_dataset_title(dataset)) + + # Hide unused axes (when datasets < grid slots) + for ax in axes_flat[len(datasets) :]: + ax.set_axis_off() + + # Label bottom row x-axes + for ax in axes_flat[-n_cols:]: + if ax.has_data(): + ax.set_xlabel(xlabel) + + legend_handles: List = [] + legend_labels: List[str] = [] + for surrogate, color in color_map.items(): + handle = plt.Line2D( + [0], + [0], + color=color, + marker="o", + linestyle="-", + linewidth=1.2, + markersize=5, + label=surrogate, + ) + legend_handles.append(handle) + legend_labels.append(surrogate) + + legend_handles, legend_labels, legend_ncol = reorder_legend_entries_rowwise( + legend_handles, legend_labels, max_ncols=2 + ) + + fig.legend( + legend_handles, + legend_labels, + loc="lower center", + bbox_to_anchor=(0.52, 0.03), + fontsize="small", + frameon=True, + ncol=legend_ncol, + ) + + plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.96]) + + os.makedirs(os.path.dirname(filename), exist_ok=True) + fig.savefig(filename, dpi=dpi, bbox_inches="tight") + plt.close(fig) + + def plot_grid_deltadex( datasets: List[str], datasets_errors: Dict[str, Dict[str, np.ndarray]], @@ -441,6 +719,89 @@ def _load_catastrophic_recall( return None +def _load_npz_dict(path: str) -> Dict[str, np.ndarray] | None: + if not os.path.exists(path): + return None + data = np.load(path, allow_pickle=True) + try: + if "arr_0" in data.files: + obj = data["arr_0"] + if isinstance(obj, np.ndarray) and obj.dtype == object and obj.shape == (): + d = obj.item() + return {k: np.asarray(v) for k, v in d.items()} + # fallthrough if it wasn't a pickled dict + # handle the case np.savez(key=array, ...) was used + return {k: np.asarray(data[k]) for k in data.files} + finally: + data.close() + + +def compute_and_report_auroc_for_datasets( + datasets: List[str], + root: str, + percentile: int = 99, + save_csv: bool = True, +): + p = percentile / 100.0 + for dataset in datasets: + err_path = os.path.join(root, dataset, "all_uq_errors.npz") + std_path = os.path.join(root, dataset, "all_uq_std.npz") + + errors_dict = _load_npz_dict(err_path) + std_dict = _load_npz_dict(std_path) + + if errors_dict is None or std_dict is None: + print(f"[WARN] Missing NPZ for {dataset}; skipping.") + continue + + keys = [k for k in errors_dict.keys() if k in std_dict] + if not keys: + print(f"[WARN] No overlapping surrogates for {dataset}; skipping.") + continue + + print( + f"\n=== AUROC — {dataset} (per-model catastrophic = top {100*(1-p):.1f}% errors) ===" + ) + out_rows = [("surrogate", "auroc", "n_used", "k_cat")] + + for k in keys: + e = np.ravel(errors_dict[k]) + u = np.ravel(std_dict[k]) + m = np.isfinite(e) & np.isfinite(u) + e, u = e[m], u[m] + + N = e.size + if N <= 1: + print(f" {k:<20s}: insufficient samples (N={N}).") + out_rows.append((k, "", N, 0)) + continue + + # top-k by error as catastrophics (model-specific) + k_cat = int(np.ceil((1.0 - p) * N)) + k_cat = max(1, min(N - 1, k_cat)) + + idx = np.argsort(e) # ascending + cat_idx = idx[-k_cat:] # largest errors + y = np.zeros(N, dtype=int) + y[cat_idx] = 1 + + try: + auc = roc_auc_score(y, u) + print(f" {k:<20s}: AUROC = {auc:.3f} (N={N}, k={k_cat})") + out_rows.append((k, f"{auc:.6f}", N, k_cat)) + except Exception as ex: + print(f" {k:<20s}: AUROC error: {ex}") + out_rows.append((k, "", N, k_cat)) + + if save_csv: + out_dir = os.path.join(root, dataset) + os.makedirs(out_dir, exist_ok=True) + csv_path = os.path.join(out_dir, f"auroc_per_model_topk_{percentile}.csv") + with open(csv_path, "w", newline="") as f: + csv.writer(f).writerows(out_rows) + print(f" → Saved: {csv_path}") + + def plot_grid_catastrophic_detection( datasets: List[str], datasets_errors: Dict[str, Dict[str, np.ndarray]], @@ -548,6 +909,206 @@ def plot_grid_catastrophic_detection( plt.close(fig) +def _load_iterative_errors(root: str, dataset: str) -> Dict[str, np.ndarray] | None: + """Load dict[str, np.ndarray] of iterative Δdex errors for a dataset. + + Expected file: scripts/pp//all_iterative_errors.npz + Returns None if missing or unreadable. + """ + path = os.path.join(root, dataset, "all_iterative_errors_3.npz") + if not os.path.exists(path): + return None + try: + return _load_errors_npz(path) + except Exception: + return None + + +def plot_grid_iterative_deltadex_percentiles( + datasets: List[str], + root: str, + timesteps: np.ndarray, + color_map: Dict[str, Tuple[float, float, float, float]], + dpi: int = 300, + n_cols: int = 2, + iter_interval: int = 10, +): + """ + Create a grid (one per dataset) showing iterative Δdex percentiles over time. + + For each dataset, load scripts/pp//all_iterative_errors.npz (dict surrogate-> [N,T,Q]). + Plot each surrogate's mean and 99th percentile Δdex over time, with subtle dashed vertical + lines at every `iter_interval`-th timestep. + """ + n = max(1, len(datasets)) + n_cols = max(1, n_cols) + n_rows = int(np.ceil(n / n_cols)) + fig, axes = plt.subplots( + n_rows, n_cols, figsize=(5 * n_cols, 3 * n_rows), sharex=False, sharey=True + ) + if isinstance(axes, np.ndarray): + axes = axes.flatten() + else: + axes = [axes] + + # Legend proxies for line styles + non_iter_proxy = plt.Line2D( + [0], + [0], + color="black", + linestyle="-", + label="Non-iterative mean", + ) + iter_proxy = plt.Line2D( + [0], + [0], + color="black", + linestyle="--", + label="Iterative mean", + ) + # iter_p99_proxy = plt.Line2D( + # [0], + # [0], + # color="black", + # linestyle=":", + # label="Iterative 99th percentile", + # ) + + surrogate_proxies = [] + surrogate_labels = [] + for name, color in color_map.items(): + surrogate_proxies.append(plt.Line2D([0], [0], color=color, label=name)) + surrogate_labels.append(name) + + for idx, (ax, dataset) in enumerate(zip(axes, datasets)): + iter_err_dict = _load_iterative_errors(root, dataset) + try: + non_iter_err_dict = load_dataset_errors(root, dataset) + except Exception: + non_iter_err_dict = {} + + if not iter_err_dict and not non_iter_err_dict: + ax.text( + 0.5, + 0.5, + f"No error data for {dataset}", + ha="center", + va="center", + ) + ax.set_axis_off() + continue + + T = len(timesteps) + + # Plot per-surrogate series + for model_name, color in color_map.items(): + non_iter_arr = ( + non_iter_err_dict.get(model_name) + if model_name in non_iter_err_dict + else None + ) + iter_arr = ( + iter_err_dict.get(model_name) + if iter_err_dict and model_name in iter_err_dict + else None + ) + + if non_iter_arr is not None and non_iter_arr.shape[1] == T: + axes_to_reduce = tuple(i for i in range(non_iter_arr.ndim) if i != 1) + mean_ts = ( + np.mean(non_iter_arr, axis=axes_to_reduce) + if axes_to_reduce + else non_iter_arr + ) + + ax.plot( + timesteps, + mean_ts, + color=color, + linestyle="-", + linewidth=1.3, + ) + + if iter_arr is not None and iter_arr.shape[1] == T: + axes_to_reduce = tuple(i for i in range(iter_arr.ndim) if i != 1) + iter_mean = ( + np.mean(iter_arr, axis=axes_to_reduce) + if axes_to_reduce + else iter_arr + ) + # iter_p99 = ( + # np.percentile(iter_arr, 99, axis=axes_to_reduce) + # if axes_to_reduce + # else iter_arr + # ) + ax.plot( + timesteps, + iter_mean, + color=color, + linestyle="--", + linewidth=1.2, + ) + # ax.plot( + # timesteps, + # iter_p99, + # color=color, + # linestyle=":", + # linewidth=1.0, + # ) + + # Vertical dashed lines every iter_interval steps (skip initial boundary) + if isinstance(iter_interval, int) and iter_interval > 0: + for i in range(iter_interval, T, iter_interval): + if i < len(timesteps): + x = timesteps[i] + ax.axvline( + x=x, linestyle="--", color="gray", alpha=0.3, linewidth=0.8 + ) + + ax.set_xscale("log") + ax.set_xlim(left=timesteps[0], right=timesteps[-1]) + if (idx % n_cols) == 0: + ax.set_ylabel(r"$\Delta dex$") + ax.set_ylim(0, 3.2) + ax.set_title(_format_dataset_title(dataset)) + # ax.set_yscale("log") + ax.grid(False) + + # Label bottom row x-axis + for ax in axes[-n_cols:]: + ax.set_xlabel("Time (y)") + + if n_rows > 1: + for ax in axes[:-n_cols]: + ax.set_xticklabels([]) + + # Combined legend + handles = surrogate_proxies + [non_iter_proxy, iter_proxy] # iter_p99_proxy + labels = surrogate_labels + [ + "One-shot mean", + "Iterative mean", + # "Iterative 99th percentile", + ] + handles, labels, legend_ncol = reorder_legend_entries_rowwise(handles, labels, 2) + + fig.legend( + handles, + labels, + loc="lower center", + bbox_to_anchor=(0.52, 0.025), + fontsize="small", + frameon=True, + ncol=legend_ncol, + ) + + plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) + + out_path = "scripts/pp/iterative_error_percentiles_deltadex_by_dataset.png" + os.makedirs(os.path.dirname(out_path), exist_ok=True) + fig.savefig(out_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + + def main(): parser = argparse.ArgumentParser( description="2x2 comparative Δdex error distributions across datasets" @@ -601,45 +1162,91 @@ def main(): x_min, x_max = compute_global_range(datasets_errors) color_map, _ = build_color_map(datasets_errors) - plot_grid_deltadex( - datasets=datasets, - datasets_errors=datasets_errors, - x_log_min=x_min, - x_log_max=x_max, - color_map=color_map, - dpi=args.dpi, - n_cols=args.cols, + surrogate_names = list(color_map.keys()) + dataset_metrics = load_all_surrogate_metrics( + root=args.root, datasets=datasets, surrogates=surrogate_names ) - # Percentiles-over-time grid (deltadex mode) - plot_grid_deltadex_percentiles( + modality_specs = [ + ("interpolation", "Interpolation Interval", False), + ("extrapolation", "Extrapolation Cutoff (%)", False), + ("sparse", "Training Samples", True), + ] + + # for modality, xlabel, xlog in modality_specs: + # modality_data = collect_modality_series(dataset_metrics, modality) + # out_path = os.path.join( + # args.root, f"{modality}_logmae_comparison_by_dataset.png" + # ) + # plot_modality_comparison_grid( + # datasets=datasets, + # modality_data=modality_data, + # color_map=color_map, + # xlabel=xlabel, + # filename=out_path, + # dpi=args.dpi, + # n_cols=args.cols, + # xlog=xlog, + # ) + + # plot_grid_deltadex( + # datasets=datasets, + # datasets_errors=datasets_errors, + # x_log_min=x_min, + # x_log_max=x_max, + # color_map=color_map, + # dpi=args.dpi, + # n_cols=args.cols, + # ) + + # # Percentiles-over-time grid (deltadex mode) + # plot_grid_deltadex_percentiles( + # datasets=datasets, + # datasets_errors=datasets_errors, + # timesteps=timesteps, + # color_map=color_map, + # dpi=args.dpi, + # n_cols=args.cols, + # ) + + # Iterative percentiles-over-time grid (deltadex with vertical guide lines) + plot_grid_iterative_deltadex_percentiles( datasets=datasets, - datasets_errors=datasets_errors, + root=args.root, timesteps=timesteps, color_map=color_map, dpi=args.dpi, n_cols=args.cols, + iter_interval=3, ) - # Catastrophic detection grid - plot_grid_catastrophic_detection( - datasets=datasets, - datasets_errors=datasets_errors, - color_map=color_map, - root=args.root, - recall_percentile=99, - dpi=args.dpi, - n_cols=args.cols, - ) - - plot_grid_catastrophic_detection( + # # Catastrophic detection grid + # plot_grid_catastrophic_detection( + # datasets=datasets, + # datasets_errors=datasets_errors, + # color_map=color_map, + # root=args.root, + # recall_percentile=99, + # dpi=args.dpi, + # n_cols=args.cols, + # ) + + # plot_grid_catastrophic_detection( + # datasets=datasets, + # datasets_errors=datasets_errors, + # color_map=color_map, + # root=args.root, + # recall_percentile=90, + # dpi=args.dpi, + # n_cols=args.cols, + # ) + + # AUROC reports + compute_and_report_auroc_for_datasets( datasets=datasets, - datasets_errors=datasets_errors, - color_map=color_map, root=args.root, - recall_percentile=90, - dpi=args.dpi, - n_cols=args.cols, + percentile=99, + save_csv=True, ) From afbb151ff5718f851e89d759bc06c1e504d03aef Mon Sep 17 00:00:00 2001 From: Robin Date: Mon, 22 Dec 2025 16:58:00 +0100 Subject: [PATCH 3/6] minor plot updates --- codes/tune/evaluate_tuning.py | 37 ++++++++++++++++++----------------- paper_plots.py | 31 +++++++++++++++-------------- 2 files changed, 35 insertions(+), 33 deletions(-) diff --git a/codes/tune/evaluate_tuning.py b/codes/tune/evaluate_tuning.py index d10b4df..2fa3e4a 100644 --- a/codes/tune/evaluate_tuning.py +++ b/codes/tune/evaluate_tuning.py @@ -329,9 +329,9 @@ def _render_pareto_scatter( ax.set_xlim(*data["xlim"]) ax.set_ylim(*data["ylim"]) if show_xlabel: - ax.set_xlabel("99th-percentile Δdex") + ax.set_xlabel(r"LAE$_{99}$ [dex]") if show_ylabel: - ax.set_ylabel("inference time (s)") + ax.set_ylabel("Inference Time [s]") if title: ax.set_title(title) ax.tick_params(labelbottom=not hide_xticklabels) @@ -389,6 +389,7 @@ def save_pareto_front_grid(datasets: list[dict], out_dir: str): loc="lower center", ncol=2, bbox_to_anchor=(0.53, 0.0), + fontsize=10, ) os.makedirs(out_dir, exist_ok=True) fig.savefig(os.path.join(out_dir, "pareto_front_grid.png"), dpi=300) @@ -420,7 +421,7 @@ def save_relative_hv_grid(datasets: list[dict], out_dir: str): if show_xlabel: ax.set_xlabel("Completed Trials") if show_ylabel: - ax.set_ylabel("Fraction of Final HV") + ax.set_ylabel("Normalized Hypervolume") else: ax.tick_params(labelleft=False) ax.grid(True, linestyle="--", alpha=0.3) @@ -431,11 +432,11 @@ def save_relative_hv_grid(datasets: list[dict], out_dir: str): fig.tight_layout(rect=(0.01, 0.01, 0.99, 0.99)) os.makedirs(out_dir, exist_ok=True) - fig.savefig(os.path.join(out_dir, "hypervolume_relative_grid.png"), dpi=300) + fig.savefig(os.path.join(out_dir, "hypervolume_normalized_grid.png"), dpi=300) plt.close(fig) print( - f"Saved combined relative-hypervolume grid with {len(datasets)} subplot(s) " - f"to {os.path.join(out_dir, 'hypervolume_relative_grid.png')}." + f"Saved combined normalized-hypervolume grid with {len(datasets)} subplot(s) " + f"to {os.path.join(out_dir, 'hypervolume_normalized_grid.png')}." ) @@ -479,17 +480,17 @@ def color_for(name: str): ax.set_xlim(0, max(max(d["x"]) for d in datasets if len(d["x"])) if datasets else 1) ax.set_ylim(0.6, 1.02) ax.set_xlabel("Completed Trials") - ax.set_ylabel("Fraction of Final HV") - ax.set_title("Relative Hypervolume") + ax.set_ylabel("Normalized Hypervolume") + # ax.set_title("Normalized Hypervolume") ax.grid(True, linestyle="--", alpha=0.3) ax.legend(ncol=2, loc="lower right") fig.tight_layout() os.makedirs(out_dir, exist_ok=True) - fig.savefig(os.path.join(out_dir, "hypervolume_relative_combined.png"), dpi=300) + fig.savefig(os.path.join(out_dir, "hypervolume_normalized_combined.png"), dpi=300) plt.close(fig) print( - "Saved combined relative-hypervolume line plot to " - f"{os.path.join(out_dir, 'hypervolume_relative_combined.png')}." + "Saved combined normalized-hypervolume line plot to " + f"{os.path.join(out_dir, 'hypervolume_normalized_combined.png')}." ) @@ -771,11 +772,11 @@ def evaluate_tuning( plt.plot( np.arange(1, len(rel_hvs) + 1), rel_hvs, - label="Relative Hypervolume", + label="Normalized Hypervolume", ) plt.xlabel("Completed Trials") - plt.ylabel("Fraction of Final HV") - plt.title(f"{suffix} Relative Hypervolume") + plt.ylabel("Normalized Hypervolume") + plt.title(f"{suffix} Normalized Hypervolume") plt.grid(True) plt.tight_layout() plt.savefig( @@ -884,13 +885,13 @@ def parse_args(): p.add_argument( "--study_name", type=str, - default="primordial_parametric_final", + default="primordial_final", help="Main study prefix (e.g. lvparams5)", ) p.add_argument( "--storage_name", type=str, - default="primordial_parametric_final", + default="primordial_final", help="Main study prefix (e.g. lvparams5)", ) p.add_argument( @@ -912,8 +913,8 @@ def parse_args(): type=str, # default="171,114,135,237", # cloud_final # default="27,61,13,299", # cloud_parametric_final - default="18,1,16,234", # primordial_parametric_final - # default="196,107,31,243", # primordial_final + # default="18,1,16,234", # primordial_parametric_final + default="196,107,31,243", # primordial_final help=( "Comma-separated list of Optuna trial numbers chosen per study (order " "matches sorted study names)." diff --git a/paper_plots.py b/paper_plots.py index 68c7419..6b837d2 100644 --- a/paper_plots.py +++ b/paper_plots.py @@ -420,7 +420,7 @@ def plot_modality_comparison_grid( # ax.set_yscale("log") ax.grid(True, which="major", linestyle="--", linewidth=0.5, alpha=0.4) if (idx % n_cols) == 0: - ax.set_ylabel(r"Log-MAE($\Delta dex$)") + ax.set_ylabel(r"LAE [dex]") ax.set_title(_format_dataset_title(dataset)) # Hide unused axes (when datasets < grid slots) @@ -537,13 +537,13 @@ def plot_grid_deltadex( ax.set_xlim(left=1e-4, right=10) # Y label only on first column if (idx % n_cols) == 0: - ax.set_ylabel("Smoothed Histogram Count") + ax.set_ylabel("Smoothed normalized\nlog-space frequency") ax.set_ylim(0, None) ax.set_title(_format_dataset_title(dataset)) # Common X label and legend for ax in axes[-n_cols:]: - ax.set_xlabel(r"Log-MAE ($\Delta dex$)") + ax.set_xlabel(r"LAE [dex]") # Build a single legend using first axis handles for present models handles, labels = [], [] @@ -567,7 +567,7 @@ def plot_grid_deltadex( final_handles, final_labels, loc="lower center", - bbox_to_anchor=(0.52, 0.05), + bbox_to_anchor=(0.52, 0.045), fontsize="small", frameon=True, ncol=legend_ncol, @@ -575,6 +575,7 @@ def plot_grid_deltadex( # No overall title; leave space at the bottom for legend plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) + fig.align_ylabels(axes) out_path = "scripts/pp/error_dist_deltadex_by_dataset.png" os.makedirs(os.path.dirname(out_path), exist_ok=True) @@ -655,7 +656,7 @@ def plot_grid_deltadex_percentiles( ax.set_xlim(left=max(1, timesteps[0]), right=timesteps[-1]) # Y label only on first column if (idx % n_cols) == 0: - ax.set_ylabel(r"$\Delta dex$") + ax.set_ylabel(r"LAE [dex]") ax.set_ylim(2 * 1e-2, 20) ax.set_title(_format_dataset_title(dataset)) ax.set_yscale("log") @@ -1068,8 +1069,8 @@ def plot_grid_iterative_deltadex_percentiles( ax.set_xscale("log") ax.set_xlim(left=timesteps[0], right=timesteps[-1]) if (idx % n_cols) == 0: - ax.set_ylabel(r"$\Delta dex$") - ax.set_ylim(0, 3.2) + ax.set_ylabel("LAE [dex]") + ax.set_ylim(0, 4) ax.set_title(_format_dataset_title(dataset)) # ax.set_yscale("log") ax.grid(False) @@ -1199,7 +1200,7 @@ def main(): # n_cols=args.cols, # ) - # # Percentiles-over-time grid (deltadex mode) + # Percentiles-over-time grid (deltadex mode) # plot_grid_deltadex_percentiles( # datasets=datasets, # datasets_errors=datasets_errors, @@ -1241,13 +1242,13 @@ def main(): # n_cols=args.cols, # ) - # AUROC reports - compute_and_report_auroc_for_datasets( - datasets=datasets, - root=args.root, - percentile=99, - save_csv=True, - ) + # # AUROC reports + # compute_and_report_auroc_for_datasets( + # datasets=datasets, + # root=args.root, + # percentile=99, + # save_csv=True, + # ) if __name__ == "__main__": From 4c069daccfb8630d501aa9a76cdb23c78d734107 Mon Sep 17 00:00:00 2001 From: Robin Date: Mon, 19 Jan 2026 16:01:02 +0100 Subject: [PATCH 4/6] Refine tuning eval --- codes/tune/evaluate_tuning.py | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/codes/tune/evaluate_tuning.py b/codes/tune/evaluate_tuning.py index 2fa3e4a..c1ae48d 100644 --- a/codes/tune/evaluate_tuning.py +++ b/codes/tune/evaluate_tuning.py @@ -157,7 +157,7 @@ def _pareto_legend_handles() -> list[Line2D]: Line2D( [0], [0], - marker="o", + marker="v", color="none", markerfacecolor=BEST_F1_COLOR, markeredgecolor=EDGE_COLOR, @@ -167,7 +167,7 @@ def _pareto_legend_handles() -> list[Line2D]: Line2D( [0], [0], - marker="o", + marker="^", color="none", markerfacecolor=CHOSEN_COLOR, markeredgecolor=EDGE_COLOR, @@ -297,10 +297,29 @@ def _render_pareto_scatter( best_point = data["best_point"] chosen_point = data["chosen_point"] + # Remove best point and chosen point from pareto points to avoid double-plotting + if mask.sum() > 0: + pareto_pts = pts[mask] + pareto_pts = pareto_pts[ + ~np.all(pareto_pts == best_point, axis=1) + ] # remove best point + if chosen_point is not None: + pareto_pts = pareto_pts[ + ~np.all(pareto_pts == chosen_point, axis=1) + ] # remove chosen point + # Recompute mask + new_mask = np.array( + [any(np.all(p == pp) for pp in pareto_pts) for p in pts], dtype=bool + ) + mask = new_mask + else: + mask = np.array([False] * pts.shape[0], dtype=bool) + ax.scatter(pts[:, 0], pts[:, 1], color=TRIAL_COLOR, alpha=0.7, label=None) ax.scatter( pts[mask, 0], pts[mask, 1], + marker="o", color=PARETO_COLOR, edgecolor=EDGE_COLOR, linewidth=0.5, @@ -309,6 +328,7 @@ def _render_pareto_scatter( ax.scatter( best_point[0], best_point[1], + marker="v", color=BEST_F1_COLOR, edgecolor=EDGE_COLOR, linewidth=0.6, @@ -319,6 +339,7 @@ def _render_pareto_scatter( ax.scatter( chosen_point[0], chosen_point[1], + marker="^", color=CHOSEN_COLOR, edgecolor=EDGE_COLOR, linewidth=0.6, From 1e8ccdfaefe34657cd06962f28fc585920c48ee3 Mon Sep 17 00:00:00 2001 From: Robin Date: Tue, 20 Jan 2026 13:21:04 +0100 Subject: [PATCH 5/6] Remove unneccesary files --- paper_eval.py | 100 ---- paper_plots.py | 1255 ------------------------------------------------ 2 files changed, 1355 deletions(-) delete mode 100644 paper_eval.py delete mode 100644 paper_plots.py diff --git a/paper_eval.py b/paper_eval.py deleted file mode 100644 index 72d3565..0000000 --- a/paper_eval.py +++ /dev/null @@ -1,100 +0,0 @@ -""" -Convenience script to re-run paper evaluations across multiple training runs. - -Behavior: -- For each specified training_id, load its saved config from trained//config.yaml -- Override the devices in that config with the hardcoded DEVICE below -- Run the same evaluation flow as run_eval.py using that modified config - -Note: This script intentionally does NOT read the top-level config.yaml. -""" - -import os -from typing import Dict - -from codes.benchmark import ( - check_benchmark, - check_surrogate, - compare_models, - get_surrogate, - run_benchmark, -) -from codes.utils import download_data, nice_print, read_yaml_config - -# Hardcoded device override for all evaluations -DEVICE = "cuda:2" - -# Training IDs to evaluate -TRAINING_IDS = [ - "_cloud_finetuned", - "_cloud_parametric_finetuned", - "_primordial_finetuned", - "_primordial_parametric_finetuned", -] - - -def evaluate_with_config(config: Dict) -> None: - """Run the evaluation loop for a single configuration dict.""" - # Basic checks and data - check_benchmark(config) - download_data(config["dataset"]["name"], verbose=config.get("verbose", False)) - - surrogates = config["surrogates"] - all_metrics = {surrogate: {} for surrogate in surrogates} - - # Evaluate each surrogate - for surrogate_name in surrogates: - surrogate_class = get_surrogate(surrogate_name) - if surrogate_class is None: - print(f"Surrogate {surrogate_name} not recognized. Skipping.") - continue - - nice_print(f"Running benchmark for {surrogate_name}") - check_surrogate(surrogate_name, config) - metrics = run_benchmark(surrogate_name, surrogate_class, config) - all_metrics[surrogate_name] = metrics - - # Compare models if requested - if config.get("compare", False): - if len(surrogates) < 2: - nice_print("At least two surrogate models are required to compare.") - else: - nice_print("Comparing models") - compare_models(all_metrics, config) - - -def load_trained_config(training_id: str) -> Dict | None: - """Load the saved config for a given training_id from trained//config.yaml.""" - cfg_path = os.path.join("trained", training_id, "config.yaml") - if not os.path.exists(cfg_path): - print( - f"Config not found for training_id '{training_id}': {cfg_path}. Skipping." - ) - return None - config = read_yaml_config(cfg_path) - return config - - -def main(): - for tid in TRAINING_IDS: - nice_print(f"Evaluating {tid}") - - config = load_trained_config(tid) - if config is None: - continue - - # Override devices with the hardcoded DEVICE - config["devices"] = [DEVICE] - - try: - evaluate_with_config(config) - except Exception as e: - print(f"Evaluation failed for {tid}: {e}") - # Continue with the next training_id - continue - - nice_print("All requested evaluations processed") - - -if __name__ == "__main__": - main() diff --git a/paper_plots.py b/paper_plots.py deleted file mode 100644 index 6b837d2..0000000 --- a/paper_plots.py +++ /dev/null @@ -1,1255 +0,0 @@ -#!/usr/bin/env python3 -""" -Paper plots: comparative Δdex error distributions across datasets. - -This script loads the per-dataset error dictionaries saved by the benchmark -(compare_errors -> scripts/pp//all_log_errors.npz), and creates a 2x2 grid -with one subplot per dataset, each showing the same comparative plot as -plot_error_distribution_comparative(..., mode="deltadex"). - -Usage: - python paper_plots.py --root scripts/pp \ - [--output plots/paper/error_dist_deltadex_by_dataset.png] \ - [--cols 2] - -Notes: -- Each dataset directory must contain an NPZ file named 'all_log_errors.npz'. - For robustness, we also try 'all_errors_log.npz' as a fallback. -- The NPZ file typically contains a single object array 'arr_0' which is a - dict mapping surrogate_name -> numpy array of Δdex errors with shape [N, T, Q]. -""" -from __future__ import annotations - -import argparse -import csv -import os -import re -import warnings -from typing import Dict, List, Tuple - -import matplotlib.pyplot as plt -import numpy as np -import yaml -from scipy.ndimage import gaussian_filter1d -from sklearn.metrics import roc_auc_score - -# Reuse palette from project for consistent styling, but keep a safe fallback -try: - from codes.benchmark.bench_plots import get_custom_palette -except Exception: # pragma: no cover - fallback if import fails - - def get_custom_palette(n: int): - return plt.cm.viridis(np.linspace(0, 0.95, n)) - - -def _format_dataset_title(name: str) -> str: - """Map dataset folder names to display titles with proper capitalization.""" - mapping = { - "primordial": "Primordial", - "primordial_parametric": "Primordial Parametric", - } - if name in mapping: - return mapping[name] - # Fallback: replace underscores with space and title-case - return name.replace("_", " ").title() - - -def _load_errors_npz(path_npz: str) -> Dict[str, np.ndarray]: - """ - Load a dict[str, np.ndarray] from an NPZ file created by np.savez. - - Supports the common patterns: - - arr_0 (object array) holding a Python dict - - Named arrays per surrogate (if saved with kwargs) - - A single key 'log_errors' holding the dict - """ - if not os.path.exists(path_npz): - raise FileNotFoundError(path_npz) - - data = np.load(path_npz, allow_pickle=True) - try: - # Preferred: saved as a single object array containing the dict - if "arr_0" in data.files and isinstance(data["arr_0"], np.ndarray): - obj = data["arr_0"] - # Could be 0-d object array with dict - if obj.dtype == object: - d = obj.item() - if isinstance(d, dict): - return d - # Alternative: explicit key name - if "log_errors" in data.files: - d = data["log_errors"].item() - if isinstance(d, dict): - return d - # Fallback: construct dict from per-surrogate arrays - out: Dict[str, np.ndarray] = {} - for k in data.files: - arr = data[k] - # Only accept ND arrays - if isinstance(arr, np.ndarray) and arr.ndim >= 1: - out[k] = arr - if out: - return out - finally: - data.close() - - raise ValueError(f"Could not interpret NPZ structure in {path_npz}") - - -def load_dataset_errors(root: str, dataset: str) -> Dict[str, np.ndarray]: - """Try both file names and return the errors dict for a dataset.""" - cand1 = os.path.join(root, dataset, "all_log_errors.npz") - cand2 = os.path.join(root, dataset, "all_errors_log.npz") # user-mentioned alt - last_err: Exception | None = None - for p in (cand1, cand2): - try: - return _load_errors_npz(p) - except Exception as e: - last_err = e - continue - raise FileNotFoundError( - f"No error file found for dataset '{dataset}'. Tried: {cand1}, {cand2}. Last error: {last_err}" - ) - - -def compute_global_range( - datasets_errors: Dict[str, Dict[str, np.ndarray]], - low_pct: float = 2.0, - high_pct: float = 98.0, -) -> Tuple[float, float]: - """ - Compute global x-range in log10 space across all datasets and surrogates, - following the same logic as plot_error_distribution_comparative. - """ - log_arrays: List[np.ndarray] = [] - for ds, err_dict in datasets_errors.items(): - for _, arr in err_dict.items(): - flat = arr.astype(float).ravel() - # Filter finite and strictly positive (avoid log10(0) and NaN) - mask = np.isfinite(flat) & (flat > 0) - if not np.any(mask): - continue - log_arrays.append(np.log10(flat[mask])) - if not log_arrays: - # Default safe range if everything is empty - return -8.0, 0.0 - - mins = [np.percentile(x, low_pct) for x in log_arrays if x.size > 0] - maxs = [np.percentile(x, high_pct) for x in log_arrays if x.size > 0] - global_min = float(np.min(mins)) - global_max = float(np.max(maxs)) - # Expand to nice boundaries - x_min = float(np.floor(global_min)) - x_max = float(np.ceil(global_max)) - return x_min, x_max - - -def build_color_map(datasets_errors: Dict[str, Dict[str, np.ndarray]]): - """Build a consistent surrogate->color map across all datasets using viridis. - - Ensures deterministic ordering by sorting surrogate names alphabetically. - Returns a dict preserving this order and the ordered list of names. - """ - name_set = set() - for err_dict in datasets_errors.values(): - name_set.update(list(err_dict.keys())) - names = sorted(name_set) - # Permute to specific legend ordering - if len(names) == 4: - names = [names[i] for i in [3, 0, 2, 1]] - colors = plt.cm.viridis(np.linspace(0, 0.95, len(names))) - # Dict preserves insertion order, matching `names` sequence - color_map = {name: colors[i] for i, name in enumerate(names)} - return color_map, names - - -SURROGATE_YAML_FILENAMES: Dict[str, str] = { - "LatentNeuralODE": "latentneuralode_metrics.yaml", - "MultiONet": "multionet_metrics.yaml", - "FullyConnected": "fullyconnected_metrics.yaml", - "LatentPoly": "latentpoly_metrics.yaml", -} - - -def _extract_first_number(token: str) -> float | None: - """Return the first numeric value found in a label like 'interval 4'.""" - matches = re.findall(r"[-+]?\d*\.?\d+", token) - if not matches: - return None - try: - return float(matches[0]) - except ValueError: - return None - - -def load_all_surrogate_metrics( - root: str, datasets: List[str], surrogates: List[str] -) -> Dict[str, Dict[str, dict]]: - """Load per-surrogate YAML metric dictionaries for each dataset.""" - results: Dict[str, Dict[str, dict]] = {} - for dataset in datasets: - dataset_metrics: Dict[str, dict] = {} - for surrogate in surrogates: - yaml_name = SURROGATE_YAML_FILENAMES.get(surrogate) - if not yaml_name: - continue - path = os.path.join(root, dataset, yaml_name) - if not os.path.exists(path): - continue - try: - with open(path, "r", encoding="utf-8") as fh: - data = yaml.safe_load(fh) or {} - except Exception: - continue - if isinstance(data, dict): - dataset_metrics[surrogate] = data - results[dataset] = dataset_metrics - return results - - -def _extract_interpolation_series(metrics: dict) -> Tuple[np.ndarray, np.ndarray]: - section = metrics.get("interpolation") - if not isinstance(section, dict): - return np.array([]), np.array([]) - - entries: List[Tuple[float, float]] = [] - for label, payload in section.items(): - if not isinstance(payload, dict): - continue - position = _extract_first_number(label) - mae = payload.get("MAE_log") - if position is None or mae is None: - continue - try: - entries.append((float(position), float(mae))) - except (TypeError, ValueError): - continue - - if not entries: - return np.array([]), np.array([]) - - entries.sort(key=lambda item: item[0]) - xs, ys = zip(*entries) - return np.asarray(xs, dtype=float), np.asarray(ys, dtype=float) - - -def _extract_extrapolation_series(metrics: dict) -> Tuple[np.ndarray, np.ndarray]: - section = metrics.get("extrapolation") - if not isinstance(section, dict): - return np.array([]), np.array([]) - - entries: List[Tuple[float, float]] = [] - for label, payload in section.items(): - if not isinstance(payload, dict): - continue - cutoff = _extract_first_number(label) - mae = payload.get("MAE_log") - if cutoff is None or mae is None: - continue - try: - entries.append((float(cutoff), float(mae))) - except (TypeError, ValueError): - continue - - if not entries: - return np.array([]), np.array([]) - - entries.sort(key=lambda item: item[0]) - xs, ys = zip(*entries) - return np.asarray(xs, dtype=float), np.asarray(ys, dtype=float) - - -def _extract_sparse_series(metrics: dict) -> Tuple[np.ndarray, np.ndarray]: - section = metrics.get("sparse") - if not isinstance(section, dict): - return np.array([]), np.array([]) - - entries: List[Tuple[float, float]] = [] - for payload in section.values(): - if not isinstance(payload, dict): - continue - mae = payload.get("MAE_log") - samples = payload.get("n_train_samples") - if mae is None or samples is None: - continue - try: - entries.append((float(samples), float(mae))) - except (TypeError, ValueError): - continue - - if not entries: - return np.array([]), np.array([]) - - entries.sort(key=lambda item: item[0]) - xs, ys = zip(*entries) - return np.asarray(xs, dtype=float), np.asarray(ys, dtype=float) - - -def collect_modality_series( - dataset_metrics: Dict[str, Dict[str, dict]], modality: str -) -> Dict[str, Dict[str, Tuple[np.ndarray, np.ndarray]]]: - """Convert raw YAML metric dictionaries into plottable series per modality.""" - extractor_map = { - "interpolation": _extract_interpolation_series, - "extrapolation": _extract_extrapolation_series, - "sparse": _extract_sparse_series, - } - extractor = extractor_map.get(modality) - if extractor is None: - raise ValueError(f"Unknown modality '{modality}'") - - modality_data: Dict[str, Dict[str, Tuple[np.ndarray, np.ndarray]]] = {} - for dataset, surrogate_metrics in dataset_metrics.items(): - per_surrogate: Dict[str, Tuple[np.ndarray, np.ndarray]] = {} - for surrogate, metrics in surrogate_metrics.items(): - xs, ys = extractor(metrics) - if xs.size == 0 or ys.size == 0: - continue - per_surrogate[surrogate] = (xs, ys) - modality_data[dataset] = per_surrogate - return modality_data - - -def reorder_legend_entries_rowwise( - handles: List, labels: List[str], max_ncols: int -) -> Tuple[List, List[str], int]: - """ - Reorder legend entries so they display row-wise when Matplotlib fills columns first. - - Provide `handles`/`labels` in the desired row-wise reading order; this - function returns a permutation that, when passed to Matplotlib legend with - ncol=legend_ncol, yields that row-wise order. - - Returns (final_handles, final_labels, legend_ncol). - """ - N = len(handles) - if N == 0: - return [], [], 1 - - legend_ncol = max(1, min(max_ncols, N)) - rows = int(np.ceil(N / legend_ncol)) - - final_handles: List = [] - final_labels: List[str] = [] - # Convert a row-wise ordered list to the column-first order expected by Matplotlib - # Example (N=6, ncol=2): input [a,b,c,d,e,f] -> output [a,c,e,b,d,f] - for c in range(legend_ncol): - for r in range(rows): - idx = r * legend_ncol + c - if idx < N: - final_handles.append(handles[idx]) - final_labels.append(labels[idx]) - - return final_handles, final_labels, legend_ncol - - -def plot_modality_comparison_grid( - datasets: List[str], - modality_data: Dict[str, Dict[str, Tuple[np.ndarray, np.ndarray]]], - color_map: Dict[str, Tuple[float, float, float, float]], - xlabel: str, - filename: str, - dpi: int = 300, - n_cols: int = 2, - xlog: bool = False, - linear_fit: bool = True, -) -> None: - """Render a grid comparing Log-MAE trends across surrogates for one modality. - - When `linear_fit` is True, draw a least-squares fit line per surrogate. - """ - - n = max(1, len(datasets)) - n_cols = max(1, n_cols) - n_rows = int(np.ceil(n / n_cols)) - fig, axes = plt.subplots( - n_rows, - n_cols, - figsize=(5 * n_cols, 3.5 * n_rows), - sharey=True, - squeeze=False, - ) - axes_flat = axes.flatten() - - for idx, dataset in enumerate(datasets): - ax = axes_flat[idx] - per_surrogate = modality_data.get(dataset, {}) - if not per_surrogate: - ax.text(0.5, 0.5, "No metrics", ha="center", va="center") - ax.set_axis_off() - continue - - for surrogate, color in color_map.items(): - series = per_surrogate.get(surrogate) - if series is None: - continue - xs, ys = series - if xs.size == 0 or ys.size == 0: - continue - ax.scatter(xs, ys, color=color, s=25) - ax.plot(xs, ys, color=color, alpha=0.6, linestyle="-") - - if linear_fit: - mask = np.isfinite(xs) & np.isfinite(ys) - if xlog: - mask &= xs > 0 - if np.count_nonzero(mask) >= 2: - fit_x = xs[mask] - fit_y = ys[mask] - if xlog: - fit_x = np.log10(fit_x) - with warnings.catch_warnings(): - warnings.simplefilter("ignore", np.RankWarning) - coeffs = np.polyfit(fit_x, fit_y, deg=1) - # print(f"Dataset '{dataset}', Surrogate '{surrogate}': fit coeffs {coeffs}") - fit_domain = np.linspace(fit_x.min(), fit_x.max(), 100) - fit_values = np.polyval(coeffs, fit_domain) - if xlog: - fit_domain = np.power(10.0, fit_domain) - ax.plot( - fit_domain, - fit_values, - color=color, - linestyle="--", - linewidth=1.1, - alpha=0.9, - ) - - if xlog: - ax.set_xscale("log") - # ax.set_yscale("log") - ax.grid(True, which="major", linestyle="--", linewidth=0.5, alpha=0.4) - if (idx % n_cols) == 0: - ax.set_ylabel(r"LAE [dex]") - ax.set_title(_format_dataset_title(dataset)) - - # Hide unused axes (when datasets < grid slots) - for ax in axes_flat[len(datasets) :]: - ax.set_axis_off() - - # Label bottom row x-axes - for ax in axes_flat[-n_cols:]: - if ax.has_data(): - ax.set_xlabel(xlabel) - - legend_handles: List = [] - legend_labels: List[str] = [] - for surrogate, color in color_map.items(): - handle = plt.Line2D( - [0], - [0], - color=color, - marker="o", - linestyle="-", - linewidth=1.2, - markersize=5, - label=surrogate, - ) - legend_handles.append(handle) - legend_labels.append(surrogate) - - legend_handles, legend_labels, legend_ncol = reorder_legend_entries_rowwise( - legend_handles, legend_labels, max_ncols=2 - ) - - fig.legend( - legend_handles, - legend_labels, - loc="lower center", - bbox_to_anchor=(0.52, 0.03), - fontsize="small", - frameon=True, - ncol=legend_ncol, - ) - - plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.96]) - - os.makedirs(os.path.dirname(filename), exist_ok=True) - fig.savefig(filename, dpi=dpi, bbox_inches="tight") - plt.close(fig) - - -def plot_grid_deltadex( - datasets: List[str], - datasets_errors: Dict[str, Dict[str, np.ndarray]], - x_log_min: float, - x_log_max: float, - color_map: Dict[str, Tuple[float, float, float, float]], - dpi: int = 300, - n_cols: int = 2, -): - """ - Render a 2x2 grid: one subplot per dataset, reproducing the comparative - error distribution plot for Δdex with consistent axes and colors. - """ - # Prepare x bin edges in log10 space and transform to linear for plotting - x_vals = np.linspace(x_log_min, x_log_max + 0.1, 100) - - n = max(1, len(datasets)) - n_cols = max(1, n_cols) - n_rows = int(np.ceil(n / n_cols)) - fig, axes = plt.subplots( - n_rows, - n_cols, - figsize=(5 * n_cols, 3 * n_rows), - sharex=True, # sharey=True - ) - if isinstance(axes, np.ndarray): - axes = axes.flatten() - else: - axes = [axes] - - for idx, (ax, dataset) in enumerate(zip(axes, datasets)): - err_dict = datasets_errors.get(dataset, {}) - if not err_dict: - ax.text(0.5, 0.5, f"No data for {dataset}", ha="center", va="center") - ax.set_axis_off() - continue - - # For legend ordering, use color_map order - for model_name, color in color_map.items(): - if model_name not in err_dict: - continue - arr = err_dict[model_name] - flat = arr.astype(float).ravel() - mask = np.isfinite(flat) & (flat > 0) - if not np.any(mask): - continue - vals = flat[mask] - logs = np.log10(vals) - - hist, bin_edges = np.histogram(logs, bins=x_vals, density=True) - smoothed = gaussian_filter1d(hist, sigma=2) - - ax.plot(10 ** bin_edges[:-1], smoothed, label=model_name, color=color) - - # Mean and median markers (on linear scale) - mean_val = float(np.mean(vals)) - median_val = float(np.median(vals)) - ax.axvline( - x=mean_val, color=color, linestyle="--", linewidth=1.0, alpha=0.9 - ) - ax.axvline( - x=median_val, color=color, linestyle="-.", linewidth=1.0, alpha=0.9 - ) - - ax.set_xscale("log") - ax.set_xlim(left=1e-4, right=10) - # Y label only on first column - if (idx % n_cols) == 0: - ax.set_ylabel("Smoothed normalized\nlog-space frequency") - ax.set_ylim(0, None) - ax.set_title(_format_dataset_title(dataset)) - - # Common X label and legend - for ax in axes[-n_cols:]: - ax.set_xlabel(r"LAE [dex]") - - # Build a single legend using first axis handles for present models - handles, labels = [], [] - for model_name, color in color_map.items(): - # Proxy lines for legend - line = plt.Line2D([0], [0], color=color, label=model_name) - handles.append(line) - labels.append(model_name) - # Mean/median style proxies - handles.append(plt.Line2D([0], [0], color="black", linestyle="--", label="Mean")) - labels.append("Mean") - handles.append(plt.Line2D([0], [0], color="black", linestyle="-.", label="Median")) - labels.append("Median") - - final_handles, final_labels, legend_ncol = reorder_legend_entries_rowwise( - handles, labels, max_ncols=2 - ) - - # Place legend below plots, arranged row-wise - fig.legend( - final_handles, - final_labels, - loc="lower center", - bbox_to_anchor=(0.52, 0.045), - fontsize="small", - frameon=True, - ncol=legend_ncol, - ) - - # No overall title; leave space at the bottom for legend - plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) - fig.align_ylabels(axes) - - out_path = "scripts/pp/error_dist_deltadex_by_dataset.png" - os.makedirs(os.path.dirname(out_path), exist_ok=True) - fig.savefig(out_path, dpi=dpi, bbox_inches="tight") - plt.close(fig) - - -def plot_grid_deltadex_percentiles( - datasets: List[str], - datasets_errors: Dict[str, Dict[str, np.ndarray]], - timesteps: np.ndarray, - color_map: Dict[str, Tuple[float, float, float, float]], - dpi: int = 300, - n_cols: int = 2, -): - """ - Create a grid of subplots (one per dataset) showing Δdex error percentiles over time. - - For each dataset: - - Draw neutral grey one-sided percentile bands (50, 90, 99) aggregated across surrogates. - - Overlay each surrogate's mean and median Δdex over time using the shared viridis color map. - """ - # Prepare layout - n = max(1, len(datasets)) - n_cols = max(1, n_cols) - n_rows = int(np.ceil(n / n_cols)) - fig, axes = plt.subplots( - n_rows, n_cols, figsize=(5 * n_cols, 3 * n_rows), sharex=False, sharey=True - ) - if isinstance(axes, np.ndarray): - axes = axes.flatten() - else: - axes = [axes] - - # Legend proxies for mean and 99th percentile styles - mean_proxy = plt.Line2D([0], [0], color="black", linestyle="-", label="Mean") - p99_proxy = plt.Line2D( - [0], [0], color="black", linestyle="--", label="99th Percentile" - ) - - surrogate_proxies = [] - surrogate_labels = [] - for name, color in color_map.items(): - surrogate_proxies.append(plt.Line2D([0], [0], color=color, label=name)) - surrogate_labels.append(name) - - for idx, (ax, dataset) in enumerate(zip(axes, datasets)): - err_dict = datasets_errors.get(dataset, {}) - if not err_dict: - ax.text(0.5, 0.5, f"No data for {dataset}", ha="center", va="center") - ax.set_axis_off() - continue - - # Assume same T across surrogates within a dataset - any_arr = next(iter(err_dict.values())) - T = any_arr.shape[1] - - # Aggregate across surrogates for percentile bands - pooled = [] - for arr in err_dict.values(): - if arr.shape[1] != T: - continue - pooled.append(arr) - - # Plot mean and 99th percentile per surrogate - for model_name, color in color_map.items(): - if model_name not in err_dict: - continue - arr = err_dict[model_name] - if arr.shape[1] != T: - continue - mean_ts = np.mean(arr, axis=(0, 2)) - p99_ts = np.percentile(arr, 99, axis=(0, 2)) - ax.plot(timesteps, mean_ts, color=color, linestyle="-", linewidth=1.2) - ax.plot(timesteps, p99_ts, color=color, linestyle="--", linewidth=1.0) - - ax.set_xscale("log") - ax.set_xlim(left=max(1, timesteps[0]), right=timesteps[-1]) - # Y label only on first column - if (idx % n_cols) == 0: - ax.set_ylabel(r"LAE [dex]") - ax.set_ylim(2 * 1e-2, 20) - ax.set_title(_format_dataset_title(dataset)) - ax.set_yscale("log") - ax.grid(False) - - # Label bottom row x-axis - for ax in axes[-n_cols:]: - ax.set_xlabel("Time (y)") - - # Turn off tick labels on upper rows if multiple rows - if n_rows > 1: - for ax in axes[:-n_cols]: - ax.set_xticklabels([]) - - # Build combined legend below plots: surrogates + mean/99th style proxies - handles = surrogate_proxies + [mean_proxy, p99_proxy] - labels = surrogate_labels + ["Mean", "99th Percentile"] - - handles, labels, legend_ncol = reorder_legend_entries_rowwise(handles, labels, 2) - - fig.legend( - handles, - labels, - loc="lower center", - bbox_to_anchor=(0.52, 0.04), - fontsize="small", - frameon=True, - ncol=legend_ncol, - ) - - plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) - - out_path = "scripts/pp/error_percentiles_deltadex_by_dataset.png" - os.makedirs(os.path.dirname(out_path), exist_ok=True) - fig.savefig(out_path, dpi=dpi, bbox_inches="tight") - plt.close(fig) - - -def _load_catastrophic_recall( - root: str, dataset: str, percentile: int -) -> np.ndarray | None: - """Load catastrophic recall matrix saved as npz for a dataset. - - Expected path: scripts/pp//catastrophic_recall_.npz - Returns array of shape [S, F] (surrogates x flag fractions), or None if missing. - """ - path = os.path.join(root, dataset, f"catastrophic_recall_{percentile}.npz") - if not os.path.exists(path): - return None - data = np.load(path, allow_pickle=True) - try: - if "arr_0" in data.files: - arr = data["arr_0"] - else: - # Fallback to the first entry - arr = data[data.files[0]] - if isinstance(arr, np.ndarray) and arr.ndim == 2: - return arr - finally: - data.close() - return None - - -def _load_npz_dict(path: str) -> Dict[str, np.ndarray] | None: - if not os.path.exists(path): - return None - data = np.load(path, allow_pickle=True) - try: - if "arr_0" in data.files: - obj = data["arr_0"] - if isinstance(obj, np.ndarray) and obj.dtype == object and obj.shape == (): - d = obj.item() - return {k: np.asarray(v) for k, v in d.items()} - # fallthrough if it wasn't a pickled dict - # handle the case np.savez(key=array, ...) was used - return {k: np.asarray(data[k]) for k in data.files} - finally: - data.close() - - -def compute_and_report_auroc_for_datasets( - datasets: List[str], - root: str, - percentile: int = 99, - save_csv: bool = True, -): - p = percentile / 100.0 - for dataset in datasets: - err_path = os.path.join(root, dataset, "all_uq_errors.npz") - std_path = os.path.join(root, dataset, "all_uq_std.npz") - - errors_dict = _load_npz_dict(err_path) - std_dict = _load_npz_dict(std_path) - - if errors_dict is None or std_dict is None: - print(f"[WARN] Missing NPZ for {dataset}; skipping.") - continue - - keys = [k for k in errors_dict.keys() if k in std_dict] - if not keys: - print(f"[WARN] No overlapping surrogates for {dataset}; skipping.") - continue - - print( - f"\n=== AUROC — {dataset} (per-model catastrophic = top {100*(1-p):.1f}% errors) ===" - ) - out_rows = [("surrogate", "auroc", "n_used", "k_cat")] - - for k in keys: - e = np.ravel(errors_dict[k]) - u = np.ravel(std_dict[k]) - m = np.isfinite(e) & np.isfinite(u) - e, u = e[m], u[m] - - N = e.size - if N <= 1: - print(f" {k:<20s}: insufficient samples (N={N}).") - out_rows.append((k, "", N, 0)) - continue - - # top-k by error as catastrophics (model-specific) - k_cat = int(np.ceil((1.0 - p) * N)) - k_cat = max(1, min(N - 1, k_cat)) - - idx = np.argsort(e) # ascending - cat_idx = idx[-k_cat:] # largest errors - y = np.zeros(N, dtype=int) - y[cat_idx] = 1 - - try: - auc = roc_auc_score(y, u) - print(f" {k:<20s}: AUROC = {auc:.3f} (N={N}, k={k_cat})") - out_rows.append((k, f"{auc:.6f}", N, k_cat)) - except Exception as ex: - print(f" {k:<20s}: AUROC error: {ex}") - out_rows.append((k, "", N, k_cat)) - - if save_csv: - out_dir = os.path.join(root, dataset) - os.makedirs(out_dir, exist_ok=True) - csv_path = os.path.join(out_dir, f"auroc_per_model_topk_{percentile}.csv") - with open(csv_path, "w", newline="") as f: - csv.writer(f).writerows(out_rows) - print(f" → Saved: {csv_path}") - - -def plot_grid_catastrophic_detection( - datasets: List[str], - datasets_errors: Dict[str, Dict[str, np.ndarray]], - color_map: Dict[str, Tuple[float, float, float, float]], - root: str, - recall_percentile: int = 99, - flag_fractions: Tuple[float, ...] = ( - 0.0, - 0.025, - 0.05, - 0.10, - 0.20, - 0.30, - 0.40, - 0.50, - ), - dpi: int = 300, - n_cols: int = 2, -): - """Create a grid of catastrophic error detection curves across datasets. - - Plots recall (%) vs fraction flagged (%) for each surrogate using precomputed - recall matrices (90 or 99). Uses consistent color mapping and a single legend below. - """ - # Layout - n = max(1, len(datasets)) - n_cols = max(1, n_cols) - n_rows = int(np.ceil(n / n_cols)) - fig, axes = plt.subplots( - n_rows, n_cols, figsize=(5 * n_cols, 3 * n_rows), sharex=False, sharey=True - ) - if isinstance(axes, np.ndarray): - axes = axes.flatten() - else: - axes = [axes] - - # Precompute legend proxies for surrogates in desired order - leg_handles: List = [] - leg_labels: List[str] = [] - for name, color in color_map.items(): - leg_handles.append(plt.Line2D([0], [0], color=color, marker="o", label=name)) - leg_labels.append(name) - - for idx, (ax, dataset) in enumerate(zip(axes, datasets)): - err_dict = datasets_errors.get(dataset, {}) - recall_mat = _load_catastrophic_recall(root, dataset, recall_percentile) - if not err_dict or recall_mat is None: - ax.text(0.5, 0.5, f"No recall data for {dataset}", ha="center", va="center") - ax.set_axis_off() - continue - - # Surrogate names order used when recall was saved - saved_names = list(err_dict.keys()) - F = recall_mat.shape[1] - # Build X values. Prefer the canonical fractions if lengths match; fallback to linear spacing - if F == len(flag_fractions): - xs = np.array(flag_fractions, dtype=float) * 100.0 - else: - xs = np.linspace(0.0, 100.0 * max(flag_fractions), F) - - # Plot curves in our global color order, mapping into the saved row index - for name, color in color_map.items(): - if name not in saved_names: - continue - row = saved_names.index(name) - ys = recall_mat[row, : len(xs)] * 100.0 - ax.plot(xs, ys, marker="o", color=color, linewidth=1.2) - - ax.set_ylabel("Catastrophic error recall (%)" if (idx % n_cols) == 0 else "") - ax.set_xlim(0, float(xs.max())) - ax.set_ylim(0, 100) - if (idx % n_cols) == 0: - ax.set_ylabel("Catastrophic error recall (%)") - ax.grid(True, alpha=0.3) - ax.set_title(_format_dataset_title(dataset)) - - # Label bottom row x-axis - for ax in axes[-n_cols:]: - ax.set_xlabel("Flagged fraction (%)") - - # Turn off tick labels on upper rows if multiple rows - if n_rows > 1: - for ax in axes[:-n_cols]: - ax.set_xticklabels([]) - - # Figure-level legend - handles, labels, legend_ncol = reorder_legend_entries_rowwise( - leg_handles, leg_labels, 2 - ) - fig.legend( - handles, - labels, - loc="lower center", - bbox_to_anchor=(0.52, 0.05), - fontsize="small", - frameon=True, - ncol=legend_ncol, - ) - - plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) - - out_path = f"scripts/pp/catastrophic_detection_{recall_percentile}_by_dataset.png" - os.makedirs(os.path.dirname(out_path), exist_ok=True) - fig.savefig(out_path, dpi=dpi, bbox_inches="tight") - plt.close(fig) - - -def _load_iterative_errors(root: str, dataset: str) -> Dict[str, np.ndarray] | None: - """Load dict[str, np.ndarray] of iterative Δdex errors for a dataset. - - Expected file: scripts/pp//all_iterative_errors.npz - Returns None if missing or unreadable. - """ - path = os.path.join(root, dataset, "all_iterative_errors_3.npz") - if not os.path.exists(path): - return None - try: - return _load_errors_npz(path) - except Exception: - return None - - -def plot_grid_iterative_deltadex_percentiles( - datasets: List[str], - root: str, - timesteps: np.ndarray, - color_map: Dict[str, Tuple[float, float, float, float]], - dpi: int = 300, - n_cols: int = 2, - iter_interval: int = 10, -): - """ - Create a grid (one per dataset) showing iterative Δdex percentiles over time. - - For each dataset, load scripts/pp//all_iterative_errors.npz (dict surrogate-> [N,T,Q]). - Plot each surrogate's mean and 99th percentile Δdex over time, with subtle dashed vertical - lines at every `iter_interval`-th timestep. - """ - n = max(1, len(datasets)) - n_cols = max(1, n_cols) - n_rows = int(np.ceil(n / n_cols)) - fig, axes = plt.subplots( - n_rows, n_cols, figsize=(5 * n_cols, 3 * n_rows), sharex=False, sharey=True - ) - if isinstance(axes, np.ndarray): - axes = axes.flatten() - else: - axes = [axes] - - # Legend proxies for line styles - non_iter_proxy = plt.Line2D( - [0], - [0], - color="black", - linestyle="-", - label="Non-iterative mean", - ) - iter_proxy = plt.Line2D( - [0], - [0], - color="black", - linestyle="--", - label="Iterative mean", - ) - # iter_p99_proxy = plt.Line2D( - # [0], - # [0], - # color="black", - # linestyle=":", - # label="Iterative 99th percentile", - # ) - - surrogate_proxies = [] - surrogate_labels = [] - for name, color in color_map.items(): - surrogate_proxies.append(plt.Line2D([0], [0], color=color, label=name)) - surrogate_labels.append(name) - - for idx, (ax, dataset) in enumerate(zip(axes, datasets)): - iter_err_dict = _load_iterative_errors(root, dataset) - try: - non_iter_err_dict = load_dataset_errors(root, dataset) - except Exception: - non_iter_err_dict = {} - - if not iter_err_dict and not non_iter_err_dict: - ax.text( - 0.5, - 0.5, - f"No error data for {dataset}", - ha="center", - va="center", - ) - ax.set_axis_off() - continue - - T = len(timesteps) - - # Plot per-surrogate series - for model_name, color in color_map.items(): - non_iter_arr = ( - non_iter_err_dict.get(model_name) - if model_name in non_iter_err_dict - else None - ) - iter_arr = ( - iter_err_dict.get(model_name) - if iter_err_dict and model_name in iter_err_dict - else None - ) - - if non_iter_arr is not None and non_iter_arr.shape[1] == T: - axes_to_reduce = tuple(i for i in range(non_iter_arr.ndim) if i != 1) - mean_ts = ( - np.mean(non_iter_arr, axis=axes_to_reduce) - if axes_to_reduce - else non_iter_arr - ) - - ax.plot( - timesteps, - mean_ts, - color=color, - linestyle="-", - linewidth=1.3, - ) - - if iter_arr is not None and iter_arr.shape[1] == T: - axes_to_reduce = tuple(i for i in range(iter_arr.ndim) if i != 1) - iter_mean = ( - np.mean(iter_arr, axis=axes_to_reduce) - if axes_to_reduce - else iter_arr - ) - # iter_p99 = ( - # np.percentile(iter_arr, 99, axis=axes_to_reduce) - # if axes_to_reduce - # else iter_arr - # ) - ax.plot( - timesteps, - iter_mean, - color=color, - linestyle="--", - linewidth=1.2, - ) - # ax.plot( - # timesteps, - # iter_p99, - # color=color, - # linestyle=":", - # linewidth=1.0, - # ) - - # Vertical dashed lines every iter_interval steps (skip initial boundary) - if isinstance(iter_interval, int) and iter_interval > 0: - for i in range(iter_interval, T, iter_interval): - if i < len(timesteps): - x = timesteps[i] - ax.axvline( - x=x, linestyle="--", color="gray", alpha=0.3, linewidth=0.8 - ) - - ax.set_xscale("log") - ax.set_xlim(left=timesteps[0], right=timesteps[-1]) - if (idx % n_cols) == 0: - ax.set_ylabel("LAE [dex]") - ax.set_ylim(0, 4) - ax.set_title(_format_dataset_title(dataset)) - # ax.set_yscale("log") - ax.grid(False) - - # Label bottom row x-axis - for ax in axes[-n_cols:]: - ax.set_xlabel("Time (y)") - - if n_rows > 1: - for ax in axes[:-n_cols]: - ax.set_xticklabels([]) - - # Combined legend - handles = surrogate_proxies + [non_iter_proxy, iter_proxy] # iter_p99_proxy - labels = surrogate_labels + [ - "One-shot mean", - "Iterative mean", - # "Iterative 99th percentile", - ] - handles, labels, legend_ncol = reorder_legend_entries_rowwise(handles, labels, 2) - - fig.legend( - handles, - labels, - loc="lower center", - bbox_to_anchor=(0.52, 0.025), - fontsize="small", - frameon=True, - ncol=legend_ncol, - ) - - plt.tight_layout(rect=[0.03, 0.08, 0.97, 0.98]) - - out_path = "scripts/pp/iterative_error_percentiles_deltadex_by_dataset.png" - os.makedirs(os.path.dirname(out_path), exist_ok=True) - fig.savefig(out_path, dpi=dpi, bbox_inches="tight") - plt.close(fig) - - -def main(): - parser = argparse.ArgumentParser( - description="2x2 comparative Δdex error distributions across datasets" - ) - parser.add_argument( - "--root", - type=str, - default="scripts/pp", - help="Root directory containing per-dataset subdirectories", - ) - parser.add_argument( - "--cols", - type=int, - default=2, - help="Number of columns in the subplot grid", - ) - parser.add_argument("--dpi", type=int, default=300, help="Figure DPI") - - args = parser.parse_args() - - # Auto-discover datasets: list subdirectories under root that contain an NPZ - if not os.path.isdir(args.root): - raise NotADirectoryError(args.root) - - all_subdirs = sorted( - d for d in os.listdir(args.root) if os.path.isdir(os.path.join(args.root, d)) - ) - datasets: List[str] = [] - datasets_errors: Dict[str, Dict[str, np.ndarray]] = {} - for ds in all_subdirs: - # Only include if an NPZ exists - npz1 = os.path.join(args.root, ds, "all_log_errors.npz") - npz2 = os.path.join(args.root, ds, "all_errors_log.npz") - if not (os.path.exists(npz1) or os.path.exists(npz2)): - continue - try: - datasets_errors[ds] = load_dataset_errors(args.root, ds) - datasets.append(ds) - except Exception: - # Skip subdirs with unreadable NPZs - continue - timesteps_path = os.path.join(args.root, "timesteps.npz") - timesteps = np.load(timesteps_path)["arr_0"] - - if not datasets: - raise RuntimeError( - f"No datasets with error NPZs found under root '{args.root}'." - ) - - # Global x-range and consistent colors across all surrogates - x_min, x_max = compute_global_range(datasets_errors) - color_map, _ = build_color_map(datasets_errors) - - surrogate_names = list(color_map.keys()) - dataset_metrics = load_all_surrogate_metrics( - root=args.root, datasets=datasets, surrogates=surrogate_names - ) - - modality_specs = [ - ("interpolation", "Interpolation Interval", False), - ("extrapolation", "Extrapolation Cutoff (%)", False), - ("sparse", "Training Samples", True), - ] - - # for modality, xlabel, xlog in modality_specs: - # modality_data = collect_modality_series(dataset_metrics, modality) - # out_path = os.path.join( - # args.root, f"{modality}_logmae_comparison_by_dataset.png" - # ) - # plot_modality_comparison_grid( - # datasets=datasets, - # modality_data=modality_data, - # color_map=color_map, - # xlabel=xlabel, - # filename=out_path, - # dpi=args.dpi, - # n_cols=args.cols, - # xlog=xlog, - # ) - - # plot_grid_deltadex( - # datasets=datasets, - # datasets_errors=datasets_errors, - # x_log_min=x_min, - # x_log_max=x_max, - # color_map=color_map, - # dpi=args.dpi, - # n_cols=args.cols, - # ) - - # Percentiles-over-time grid (deltadex mode) - # plot_grid_deltadex_percentiles( - # datasets=datasets, - # datasets_errors=datasets_errors, - # timesteps=timesteps, - # color_map=color_map, - # dpi=args.dpi, - # n_cols=args.cols, - # ) - - # Iterative percentiles-over-time grid (deltadex with vertical guide lines) - plot_grid_iterative_deltadex_percentiles( - datasets=datasets, - root=args.root, - timesteps=timesteps, - color_map=color_map, - dpi=args.dpi, - n_cols=args.cols, - iter_interval=3, - ) - - # # Catastrophic detection grid - # plot_grid_catastrophic_detection( - # datasets=datasets, - # datasets_errors=datasets_errors, - # color_map=color_map, - # root=args.root, - # recall_percentile=99, - # dpi=args.dpi, - # n_cols=args.cols, - # ) - - # plot_grid_catastrophic_detection( - # datasets=datasets, - # datasets_errors=datasets_errors, - # color_map=color_map, - # root=args.root, - # recall_percentile=90, - # dpi=args.dpi, - # n_cols=args.cols, - # ) - - # # AUROC reports - # compute_and_report_auroc_for_datasets( - # datasets=datasets, - # root=args.root, - # percentile=99, - # save_csv=True, - # ) - - -if __name__ == "__main__": - main() From f1df561b6047e39d23f2e7367982ac12d728a009 Mon Sep 17 00:00:00 2001 From: Robin Date: Tue, 20 Jan 2026 13:38:34 +0100 Subject: [PATCH 6/6] Remove temporary code snippets --- codes/benchmark/bench_fcts.py | 17 ----------------- codes/benchmark/bench_plots.py | 20 -------------------- 2 files changed, 37 deletions(-) diff --git a/codes/benchmark/bench_fcts.py b/codes/benchmark/bench_fcts.py index 73a587a..265f7a3 100644 --- a/codes/benchmark/bench_fcts.py +++ b/codes/benchmark/bench_fcts.py @@ -96,8 +96,6 @@ def run_benchmark(surr_name: str, surrogate_class, conf: dict) -> dict[str, Any] tolerance=conf["dataset"]["tolerance"], per_species=conf["dataset"].get("normalise_per_species", False), ) - # TEMP - print(conf["dataset"]["name"], train_data.shape, val_data.shape, test_data.shape) model_config = get_model_config(surr_name, conf) n_timesteps = train_data.shape[1] @@ -1282,10 +1280,6 @@ def compare_errors(metrics: dict[str, dict], config: dict) -> None: if log_errors: plot_errors_over_time(mean_log, median_log, timesteps, config, mode="deltadex") plot_error_distribution_comparative(log_errors, config, mode="deltadex") - # TEMP - dataset = config["dataset"]["name"] - os.makedirs(f"scripts/pp/{dataset}", exist_ok=True) - np.savez(f"scripts/pp/{dataset}/all_log_errors.npz", log_errors) def compare_iterative(metrics: dict[str, dict], config: dict) -> None: @@ -1315,11 +1309,6 @@ def compare_iterative(metrics: dict[str, dict], config: dict) -> None: iterative_errors[surrogate], axis=(0, 2) ) - # TEMP - dataset = config["dataset"]["name"] - os.makedirs(f"scripts/pp/{dataset}", exist_ok=True) - np.savez(f"scripts/pp/{dataset}/all_iterative_errors.npz", iterative_errors) - plot_errors_over_time( mean_iterative_errors, median_iterative_errors, @@ -1602,12 +1591,6 @@ def compare_UQ(all_metrics: dict, config: dict) -> None: show_title=True, ) - # TEMP - dataset = config["dataset"]["name"] - os.makedirs(f"scripts/pp/{dataset}", exist_ok=True) - np.savez(f"scripts/pp/{dataset}/all_uq_errors.npz", ensemble_errors) - np.savez(f"scripts/pp/{dataset}/all_uq_std.npz", ensemble_std) - def tabular_comparison(all_metrics: dict, config: dict) -> None: """ diff --git a/codes/benchmark/bench_plots.py b/codes/benchmark/bench_plots.py index 26091c9..f2804c4 100644 --- a/codes/benchmark/bench_plots.py +++ b/codes/benchmark/bench_plots.py @@ -2036,8 +2036,6 @@ def inference_time_bar_plot( # Calculate the upper y-limit to provide space for text max_bar = max(means[i] + stds[i] for i in range(len(means))) # min_bar = min(means[i] - stds[i] for i in range(len(means))) - # Temp! - # ax.set_ylim(min_bar * 0.3, max_bar * 2) # Set limits with some padding ax.set_ylim(0, max_bar * 1.2) # Set limits with some padding # Add inference time as text to the bars using the format_time function @@ -2054,8 +2052,6 @@ def inference_time_bar_plot( ax.set_xlabel("Surrogate Model") ax.set_ylabel("Mean Inference Time per Run") - # Temp! - # ax.set_yscale("log") if show_title: ax.set_title("Surrogate Mean Inference Time Comparison") @@ -2490,10 +2486,6 @@ def plot_catastrophic_detection_curves( names = list(errors_log.keys()) colors = plt.cm.viridis(np.linspace(0, 0.95, len(names))) summary: dict[str, dict[float, dict[str, float]]] = {} - # TEMP - recall_99 = np.zeros((len(names), len(flag_fractions))) - recall_90 = np.zeros((len(names), len(flag_fractions))) - dataset_name = conf["dataset"]["name"] # --- Recall vs fraction flagged (per catastrophic percentile) --- for ax, perc in zip(axes[:-1], percentiles): @@ -2522,11 +2514,6 @@ def plot_catastrophic_detection_curves( recall = (flagged & is_cat).sum() / n_cat if n_cat > 0 else 0.0 xs.append(100.0 * flagged.mean()) ys.append(100.0 * recall) - # TEMP - if perc == 99.0: - recall_99[i, flag_fractions.index(f)] = recall - if perc == 90.0: - recall_90[i, flag_fractions.index(f)] = recall ax.plot( xs, @@ -2553,9 +2540,6 @@ def plot_catastrophic_detection_curves( f"Detection @ {perc}th percentile (Top {100 - perc:.0f}% Δdex)" ) - np.savez(f"scripts/pp/{dataset_name}/catastrophic_recall_99.npz", recall_99) - np.savez(f"scripts/pp/{dataset_name}/catastrophic_recall_90.npz", recall_90) - # MAE improvement plot ax_mae = axes[-1] for i, name in enumerate(names): @@ -3006,8 +2990,6 @@ def rel_errors_and_uq( ax1.set_xlabel("Time") ax1.set_xlim(timesteps[0], timesteps[-1]) - # Temp! - # ax1.set_ylim(3e-4, 1) ax1.set_ylabel("Relative Error") ax1.set_yscale("log") ax1.set_title("Comparison of Relative Errors Over Time") @@ -3038,8 +3020,6 @@ def rel_errors_and_uq( ax2.set_xlabel("Time") ax2.set_xlim(timesteps[0], timesteps[-1]) - # Temp! - # ax2.set_ylim(0, 0.04) ax2.set_ylabel("Uncertainty/Absolute Error") if show_title: ax2.set_title("Comparison of Predictive Uncertainty Over Time")