diff --git a/scripts/analyze_certainty.py b/scripts/analyze_certainty.py index 0dfc514..32be6db 100644 --- a/scripts/analyze_certainty.py +++ b/scripts/analyze_certainty.py @@ -3627,19 +3627,19 @@ OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL") INPUT_CSV = "/home/shahin/Lab/Doktorarbeit/Barcelona/data/processed/MS_Briefe_400_with_unique_id_SHA3_explore_cleaned_unique.csv" -EDSS_INSTRUCTIONS_PATH = "/home/shahin/Lab/Doktorarbeit/Barcelona/prompts/Komplett.txt" +EDSS_INSTRUCTIONS_PATH = "/home/shahin/Lab/Doktorarbeit/Barcelona/prompts/edss_guideline_DRAFT.txt" RESULTS_ROOT = "/home/shahin/Lab/Doktorarbeit/Barcelona/results/benchmark_runs" MODEL_CONFIGS = [ { - "model_name": "gpt-oss-120b", + "model_name": "openai/gpt-oss-120b", "use_response_format": True, "temperature": 0.0, "max_tokens": 4096, "extra_body": None, }, { - "model_name": "qwen3.6-27b", + "model_name": "qwen3.8-27b", "use_response_format": False, "temperature": 0.0, "max_tokens": 4096, @@ -3650,7 +3650,7 @@ MODEL_CONFIGS = [ }, }, { - "model_name": "gemma-4-31B-it", + "model_name": "google/gemma-4-31B-it", "use_response_format": False, "temperature": 0.0, "max_tokens": 4096, diff --git a/scripts/show_plots.py b/scripts/show_plots.py index 96b550d..46e79a5 100644 --- a/scripts/show_plots.py +++ b/scripts/show_plots.py @@ -12904,7 +12904,8 @@ print(OUTPUT_WIDE_XLSX) ## -# %% name +# %% ALTERNATIVE WIDE TABLE FOR PAPER + # ========================= # ALTERNATIVE WIDE TABLE FOR PAPER # Rows: Predicted EDSS range + Confidence bracket @@ -12963,4 +12964,1810 @@ print(OUTPUT_MODEL_COLUMNS_MD) print(OUTPUT_MODEL_COLUMNS_XLSX) ## +# %% Accuracy-time trade-off plot directly from model result CSVs + +from pathlib import Path + +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt + + +# ========================= +# CONFIGURATION +# ========================= + +GT_PATH = Path( + "/home/shahin/Lab/Doktorarbeit/Barcelona/data/processed/" + "MS_Briefe_400_with_unique_id_SHA3_explore_cleaned_unique.csv" +) + +RUN_DIR = Path( + "/home/shahin/Lab/Doktorarbeit/Barcelona/results/benchmark_runs/run_20260528_103942" +) + +TARGET_ITERATION = 1 + +OUTPUT_DIR = RUN_DIR / f"accuracy_time_tradeoff_iter_{TARGET_ITERATION}" +OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + +OUTPUT_TABLE = OUTPUT_DIR / f"accuracy_time_tradeoff_table_iter_{TARGET_ITERATION}.csv" + +OUTPUT_MAE_SVG = OUTPUT_DIR / f"accuracy_time_tradeoff_mae_iter_{TARGET_ITERATION}.svg" +OUTPUT_MAE_PNG = OUTPUT_DIR / f"accuracy_time_tradeoff_mae_iter_{TARGET_ITERATION}.png" + +OUTPUT_ACC_SVG = OUTPUT_DIR / f"accuracy_time_tradeoff_accuracy_within_0_5_iter_{TARGET_ITERATION}.svg" +OUTPUT_ACC_PNG = OUTPUT_DIR / f"accuracy_time_tradeoff_accuracy_within_0_5_iter_{TARGET_ITERATION}.png" + +GT_EDSS_COL = "EDSS" +PRED_EDSS_COL = "EDSS_numeric" +PRED_EDSS_FALLBACK_COL = "EDSS" + +ANALYSIS_MODE = "primary_valid_only" +# ANALYSIS_MODE = "sensitivity_all_numeric" + +plt.rcParams["font.family"] = "Arial" + + +# ========================= +# HELPERS +# ========================= + +def to_num(series): + return pd.to_numeric( + series.astype(str).str.replace(",", ".", regex=False), + errors="coerce" + ) + + +def to_bool(series): + return series.astype(str).str.lower().isin(["true", "1", "yes", "ja"]) + + +def clean_model_name(name): + replacements = { + "gpt-oss-120b": "GPT-OSS-120B", + "qwen3.6-27b": "Qwen3.6-27B", + "gemma-4-31B-it": "Gemma-4-31B-it", + "GPT-OSS-120B": "GPT-OSS-120B", + "Qwen3.6-27B": "Qwen3.6-27B", + "Gemma-4-31B-it": "Gemma-4-31B-it", + } + return replacements.get(str(name), str(name)) + + +def find_iter_file(model_dir, iteration): + files = sorted(model_dir.glob(f"*results_iter_{iteration}_*.csv")) + + files = [ + f for f in files + if "incremental" not in f.name.lower() + and "summary" not in f.name.lower() + and "all_results" not in f.name.lower() + and "primary_valid_only" not in f.name.lower() + and "sensitivity_all_numeric" not in f.name.lower() + ] + + return files[0] if files else None + + +def get_model_name(pred, model_dir): + if "model" in pred.columns and pred["model"].notna().any(): + return str(pred["model"].dropna().iloc[0]) + return model_dir.name + + +def evaluate_model(gt_numeric, model_dir, result_file): + pred_raw = pd.read_csv(result_file) + + if "row_index" not in pred_raw.columns: + print(f"Skipping {model_dir.name}: missing row_index") + return None + + model_name = get_model_name(pred_raw, model_dir) + model_display = clean_model_name(model_name) + + pred = pred_raw.copy() + + pred["row_index"] = pd.to_numeric(pred["row_index"], errors="coerce") + pred = pred.dropna(subset=["row_index"]).copy() + pred["row_index"] = pred["row_index"].astype(int) + + raw_prediction_rows = len(pred) + + if "success" in pred.columns: + pred = pred[to_bool(pred["success"])].copy() + + if ANALYSIS_MODE == "primary_valid_only": + if "EDSS_is_numeric" in pred.columns: + pred = pred[to_bool(pred["EDSS_is_numeric"])].copy() + + if "EDSS_in_valid_range" in pred.columns: + pred = pred[to_bool(pred["EDSS_in_valid_range"])].copy() + + elif ANALYSIS_MODE == "sensitivity_all_numeric": + if "EDSS_is_numeric" in pred.columns: + pred = pred[to_bool(pred["EDSS_is_numeric"])].copy() + else: + raise ValueError(f"Unknown ANALYSIS_MODE: {ANALYSIS_MODE}") + + if PRED_EDSS_COL in pred.columns: + pred_col = PRED_EDSS_COL + elif PRED_EDSS_FALLBACK_COL in pred.columns: + pred_col = PRED_EDSS_FALLBACK_COL + else: + print(f"Skipping {model_dir.name}: missing EDSS prediction column") + return None + + pred["PRED_EDSS_numeric"] = to_num(pred[pred_col]) + pred = pred.dropna(subset=["PRED_EDSS_numeric"]).copy() + + pred = pred.drop_duplicates("row_index", keep="first").copy() + + merged = gt_numeric.merge( + pred, + on="row_index", + how="inner", + suffixes=("_gt", "_pred") + ) + + if merged.empty: + print(f"Skipping {model_dir.name}: no evaluable rows") + return None + + merged["error"] = merged["PRED_EDSS_numeric"] - merged["GT_EDSS_numeric"] + merged["abs_error"] = merged["error"].abs() + + n_evaluable = len(merged) + + mae = merged["abs_error"].mean() + median_absolute_error = merged["abs_error"].median() + rmse = np.sqrt((merged["error"] ** 2).mean()) + mean_signed_error = merged["error"].mean() + + exact_accuracy = (merged["abs_error"] == 0).mean() + accuracy_within_0_5 = (merged["abs_error"] <= 0.5).mean() + accuracy_within_1_0 = (merged["abs_error"] <= 1.0).mean() + + if "inference_time_sec" in merged.columns: + mean_inference_time_sec = pd.to_numeric( + merged["inference_time_sec"], + errors="coerce" + ).mean() + else: + mean_inference_time_sec = np.nan + + return { + "model": model_name, + "model_display": model_display, + "result_file": str(result_file), + "iteration": TARGET_ITERATION, + "analysis_mode": ANALYSIS_MODE, + + "raw_prediction_rows": raw_prediction_rows, + "n_after_filtering": len(pred), + "n_evaluable": n_evaluable, + + "mean_inference_time_sec": mean_inference_time_sec, + + "MAE": mae, + "median_absolute_error": median_absolute_error, + "RMSE": rmse, + "mean_signed_error": mean_signed_error, + + "exact_accuracy": exact_accuracy, + "accuracy_within_0_5": accuracy_within_0_5, + "accuracy_within_1_0": accuracy_within_1_0, + + "exact_accuracy_percent": exact_accuracy * 100, + "accuracy_within_0_5_percent": accuracy_within_0_5 * 100, + "accuracy_within_1_0_percent": accuracy_within_1_0 * 100, + } + + +# ========================= +# LOAD GROUND TRUTH +# ========================= + +gt = pd.read_csv(GT_PATH, sep=";") +gt["row_index"] = gt.index +gt["GT_EDSS_numeric"] = to_num(gt[GT_EDSS_COL]) +gt_numeric = gt.dropna(subset=["GT_EDSS_numeric"]).copy() + +print(f"GT rows: {len(gt)}") +print(f"GT rows with numeric EDSS: {len(gt_numeric)}") + + +# ========================= +# EVALUATE MODELS +# ========================= + +model_dirs = [ + p for p in sorted(RUN_DIR.iterdir()) + if p.is_dir() + and not p.name.startswith("confusion") + and not p.name.startswith("functional_system") + and not p.name.startswith("repeated_run") + and not p.name.startswith("edss_error_distribution") + and not p.name.startswith("edss_threshold_metrics") + and not p.name.startswith("edss_severity_group_metrics") + and not p.name.startswith("structured_output_validity") + and not p.name.startswith("confidence") + and not p.name.startswith("accuracy_time_tradeoff") +] + +rows = [] + +for model_dir in model_dirs: + result_file = find_iter_file(model_dir, TARGET_ITERATION) + + if result_file is None: + print(f"No iter_{TARGET_ITERATION} result file found for {model_dir.name}") + continue + + print("\n" + "=" * 100) + print(f"Model folder: {model_dir.name}") + print(f"Result file: {result_file}") + + row = evaluate_model( + gt_numeric=gt_numeric, + model_dir=model_dir, + result_file=result_file + ) + + if row is not None: + rows.append(row) + +metrics_df = pd.DataFrame(rows) + +if metrics_df.empty: + raise ValueError("No model metrics could be calculated.") + +model_order = [ + "GPT-OSS-120B", + "Qwen3.6-27B", + "Gemma-4-31B-it", +] + +metrics_df["model_display"] = pd.Categorical( + metrics_df["model_display"], + categories=model_order, + ordered=True +) + +metrics_df = metrics_df.sort_values("model_display").reset_index(drop=True) + +metrics_df.to_csv(OUTPUT_TABLE, index=False) + +print("\nAccuracy-time trade-off table:") +print( + metrics_df[ + [ + "model_display", + "n_evaluable", + "mean_inference_time_sec", + "MAE", + "accuracy_within_0_5_percent", + ] + ] +) + + +# ========================= +# PLOT SETTINGS +# ========================= + +colors = { + "GPT-OSS-120B": "#1F77B4", + "Qwen3.6-27B": "#FF7F0E", + "Gemma-4-31B-it": "#2CA02C", +} + +markers = { + "GPT-OSS-120B": "o", + "Qwen3.6-27B": "s", + "Gemma-4-31B-it": "^", +} + + +# ========================= +# FIGURE 1: TIME VS MAE +# ========================= + +fig, ax = plt.subplots(figsize=(7.8, 5.6)) + +for _, row in metrics_df.iterrows(): + model = str(row["model_display"]) + + ax.scatter( + row["mean_inference_time_sec"], + row["MAE"], + s=200, + color=colors.get(model, "gray"), + marker=markers.get(model, "o"), + edgecolor="black", + linewidth=0.8, + alpha=0.9, + ) + + ax.text( + row["mean_inference_time_sec"], + row["MAE"], + f" {model}", + ha="left", + va="center", + fontsize=10, + fontweight="bold", + ) + +ax.set_xlabel("Mean inference time per record (seconds)", fontsize=11, fontweight="bold") +ax.set_ylabel("Mean absolute error (EDSS points)", fontsize=11, fontweight="bold") + +ax.set_title( + "Accuracy–time trade-off: EDSS MAE", + fontsize=13, + fontweight="bold", + pad=15, +) + +ax.grid(True, linestyle="--", alpha=0.3) +ax.set_axisbelow(True) + +for spine in ["top", "right"]: + ax.spines[spine].set_visible(False) + +ax.text( + 0.02, + 0.02, + "Better region: lower-left\nLower MAE and faster inference", + transform=ax.transAxes, + ha="left", + va="bottom", + fontsize=9, + bbox=dict( + boxstyle="round,pad=0.35", + facecolor="white", + edgecolor="#BBBBBB", + alpha=0.9, + ), +) + +plt.tight_layout() + +plt.savefig(OUTPUT_MAE_SVG, format="svg", bbox_inches="tight") +plt.savefig(OUTPUT_MAE_PNG, dpi=300, bbox_inches="tight") + +plt.show() + + +# ========================= +# FIGURE 2: TIME VS ±0.5 ACCURACY +# ========================= + +fig, ax = plt.subplots(figsize=(7.8, 5.6)) + +for _, row in metrics_df.iterrows(): + model = str(row["model_display"]) + + ax.scatter( + row["mean_inference_time_sec"], + row["accuracy_within_0_5_percent"], + s=200, + color=colors.get(model, "gray"), + marker=markers.get(model, "o"), + edgecolor="black", + linewidth=0.8, + alpha=0.9, + ) + + ax.text( + row["mean_inference_time_sec"], + row["accuracy_within_0_5_percent"], + f" {model}", + ha="left", + va="center", + fontsize=10, + fontweight="bold", + ) + +ax.set_xlabel("Mean inference time per record (seconds)", fontsize=11, fontweight="bold") +ax.set_ylabel("Accuracy within ±0.5 EDSS (%)", fontsize=11, fontweight="bold") + +ax.set_title( + "Accuracy–time trade-off: EDSS ±0.5 accuracy", + fontsize=13, + fontweight="bold", + pad=15, +) + +y_min = max(0, metrics_df["accuracy_within_0_5_percent"].min() - 5) +y_max = min(100, metrics_df["accuracy_within_0_5_percent"].max() + 5) +ax.set_ylim(y_min, y_max) + +ax.grid(True, linestyle="--", alpha=0.3) +ax.set_axisbelow(True) + +for spine in ["top", "right"]: + ax.spines[spine].set_visible(False) + +ax.text( + 0.02, + 0.98, + "Better region: upper-left\nHigher accuracy and faster inference", + transform=ax.transAxes, + ha="left", + va="top", + fontsize=9, + bbox=dict( + boxstyle="round,pad=0.35", + facecolor="white", + edgecolor="#BBBBBB", + alpha=0.9, + ), +) + +plt.tight_layout() + +plt.savefig(OUTPUT_ACC_SVG, format="svg", bbox_inches="tight") +plt.savefig(OUTPUT_ACC_PNG, dpi=300, bbox_inches="tight") + +plt.show() + + +print("\nSaved:") +print(OUTPUT_TABLE) +print(OUTPUT_MAE_SVG) +print(OUTPUT_MAE_PNG) +print(OUTPUT_ACC_SVG) +print(OUTPUT_ACC_PNG) +## + +# %% Accuracy-time trade-off bubble plot directly from model result CSVs +# Bubble size = estimated FP16 model size in GB + +from pathlib import Path + +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt + + +# ========================= +# CONFIGURATION +# ========================= + +GT_PATH = Path( + "/home/shahin/Lab/Doktorarbeit/Barcelona/data/processed/" + "MS_Briefe_400_with_unique_id_SHA3_explore_cleaned_unique.csv" +) + +RUN_DIR = Path( + "/home/shahin/Lab/Doktorarbeit/Barcelona/results/benchmark_runs/run_20260528_103942" +) + +TARGET_ITERATION = 1 + +OUTPUT_DIR = RUN_DIR / f"accuracy_time_tradeoff_iter_{TARGET_ITERATION}" +OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + +OUTPUT_TABLE = OUTPUT_DIR / f"accuracy_time_tradeoff_bubble_table_iter_{TARGET_ITERATION}.csv" + +OUTPUT_MAE_SVG = OUTPUT_DIR / f"accuracy_time_tradeoff_mae_bubble_iter_{TARGET_ITERATION}.svg" +OUTPUT_MAE_PNG = OUTPUT_DIR / f"accuracy_time_tradeoff_mae_bubble_iter_{TARGET_ITERATION}.png" + +OUTPUT_ACC_SVG = OUTPUT_DIR / f"accuracy_time_tradeoff_accuracy_within_0_5_bubble_iter_{TARGET_ITERATION}.svg" +OUTPUT_ACC_PNG = OUTPUT_DIR / f"accuracy_time_tradeoff_accuracy_within_0_5_bubble_iter_{TARGET_ITERATION}.png" + +GT_EDSS_COL = "EDSS" +PRED_EDSS_COL = "EDSS_numeric" +PRED_EDSS_FALLBACK_COL = "EDSS" + +ANALYSIS_MODE = "primary_valid_only" +# ANALYSIS_MODE = "sensitivity_all_numeric" + +# Public parameter counts. +# Estimated FP16 size: parameters * 2 bytes. + +MODEL_SIZE_INFO = { + "GPT-OSS-120B": { + "total_size_bytes": 195764040609, + "size_gb": 195.76, + "size_gib": 182.32, + "note": "Hugging Face repository total_size", + }, + "Qwen3.6-27B": { + "total_size_bytes": 55586107940, + "size_gb": 55.59, + "size_gib": 51.77, + "note": "Hugging Face repository total_size", + }, + "Gemma-4-31B-it": { + "total_size_bytes": 62578663749, + "size_gb": 62.58, + "size_gib": 58.28, + "note": "Hugging Face repository total_size", + }, +} +# Increase/decrease this to make bubbles visually larger/smaller. +BUBBLE_SIZE_SCALE = 10 + +plt.rcParams["font.family"] = "Arial" + + +# ========================= +# HELPERS +# ========================= + +def to_num(series): + return pd.to_numeric( + series.astype(str).str.replace(",", ".", regex=False), + errors="coerce" + ) + + +def to_bool(series): + return series.astype(str).str.lower().isin(["true", "1", "yes", "ja"]) + + +def clean_model_name(name): + replacements = { + "gpt-oss-120b": "GPT-OSS-120B", + "qwen3.6-27b": "Qwen3.6-27B", + "gemma-4-31B-it": "Gemma-4-31B-it", + "GPT-OSS-120B": "GPT-OSS-120B", + "Qwen3.6-27B": "Qwen3.6-27B", + "Gemma-4-31B-it": "Gemma-4-31B-it", + } + return replacements.get(str(name), str(name)) + + +def find_iter_file(model_dir, iteration): + files = sorted(model_dir.glob(f"*results_iter_{iteration}_*.csv")) + + files = [ + f for f in files + if "incremental" not in f.name.lower() + and "summary" not in f.name.lower() + and "all_results" not in f.name.lower() + and "primary_valid_only" not in f.name.lower() + and "sensitivity_all_numeric" not in f.name.lower() + ] + + return files[0] if files else None + + +def get_model_name(pred, model_dir): + if "model" in pred.columns and pred["model"].notna().any(): + return str(pred["model"].dropna().iloc[0]) + return model_dir.name + + +def bubble_area_from_gb(size_gb): + """ + Matplotlib scatter uses area, not radius. + """ + return float(size_gb) * BUBBLE_SIZE_SCALE + + +def evaluate_model(gt_numeric, model_dir, result_file): + pred_raw = pd.read_csv(result_file) + + if "row_index" not in pred_raw.columns: + print(f"Skipping {model_dir.name}: missing row_index") + return None + + model_name = get_model_name(pred_raw, model_dir) + model_display = clean_model_name(model_name) + + pred = pred_raw.copy() + + pred["row_index"] = pd.to_numeric(pred["row_index"], errors="coerce") + pred = pred.dropna(subset=["row_index"]).copy() + pred["row_index"] = pred["row_index"].astype(int) + + raw_prediction_rows = len(pred) + + if "success" in pred.columns: + pred = pred[to_bool(pred["success"])].copy() + + if ANALYSIS_MODE == "primary_valid_only": + if "EDSS_is_numeric" in pred.columns: + pred = pred[to_bool(pred["EDSS_is_numeric"])].copy() + + if "EDSS_in_valid_range" in pred.columns: + pred = pred[to_bool(pred["EDSS_in_valid_range"])].copy() + + elif ANALYSIS_MODE == "sensitivity_all_numeric": + if "EDSS_is_numeric" in pred.columns: + pred = pred[to_bool(pred["EDSS_is_numeric"])].copy() + else: + raise ValueError(f"Unknown ANALYSIS_MODE: {ANALYSIS_MODE}") + + if PRED_EDSS_COL in pred.columns: + pred_col = PRED_EDSS_COL + elif PRED_EDSS_FALLBACK_COL in pred.columns: + pred_col = PRED_EDSS_FALLBACK_COL + else: + print(f"Skipping {model_dir.name}: missing EDSS prediction column") + return None + + pred["PRED_EDSS_numeric"] = to_num(pred[pred_col]) + pred = pred.dropna(subset=["PRED_EDSS_numeric"]).copy() + + pred = pred.drop_duplicates("row_index", keep="first").copy() + + merged = gt_numeric.merge( + pred, + on="row_index", + how="inner", + suffixes=("_gt", "_pred") + ) + + if merged.empty: + print(f"Skipping {model_dir.name}: no evaluable rows") + return None + + merged["error"] = merged["PRED_EDSS_numeric"] - merged["GT_EDSS_numeric"] + merged["abs_error"] = merged["error"].abs() + + n_evaluable = len(merged) + + mae = merged["abs_error"].mean() + median_absolute_error = merged["abs_error"].median() + rmse = np.sqrt((merged["error"] ** 2).mean()) + mean_signed_error = merged["error"].mean() + + exact_accuracy = (merged["abs_error"] == 0).mean() + accuracy_within_0_5 = (merged["abs_error"] <= 0.5).mean() + accuracy_within_1_0 = (merged["abs_error"] <= 1.0).mean() + + if "inference_time_sec" in merged.columns: + mean_inference_time_sec = pd.to_numeric( + merged["inference_time_sec"], + errors="coerce" + ).mean() + else: + mean_inference_time_sec = np.nan + + size_info = MODEL_SIZE_INFO.get(model_display, {}) + params_billion = size_info.get("params_billion", np.nan) + estimated_fp16_gb = size_info.get("estimated_fp16_gb", np.nan) + size_note = size_info.get("note", "") + + return { + "model": model_name, + "model_display": model_display, + "result_file": str(result_file), + "iteration": TARGET_ITERATION, + "analysis_mode": ANALYSIS_MODE, + + "raw_prediction_rows": raw_prediction_rows, + "n_after_filtering": len(pred), + "n_evaluable": n_evaluable, + + "mean_inference_time_sec": mean_inference_time_sec, + + "MAE": mae, + "median_absolute_error": median_absolute_error, + "RMSE": rmse, + "mean_signed_error": mean_signed_error, + + "exact_accuracy": exact_accuracy, + "accuracy_within_0_5": accuracy_within_0_5, + "accuracy_within_1_0": accuracy_within_1_0, + + "exact_accuracy_percent": exact_accuracy * 100, + "accuracy_within_0_5_percent": accuracy_within_0_5 * 100, + "accuracy_within_1_0_percent": accuracy_within_1_0 * 100, + + "params_billion": params_billion, + "estimated_fp16_gb": estimated_fp16_gb, + "model_size_note": size_note, + "bubble_area": bubble_area_from_gb(estimated_fp16_gb), + } + + +# ========================= +# LOAD GROUND TRUTH +# ========================= + +gt = pd.read_csv(GT_PATH, sep=";") +gt["row_index"] = gt.index +gt["GT_EDSS_numeric"] = to_num(gt[GT_EDSS_COL]) +gt_numeric = gt.dropna(subset=["GT_EDSS_numeric"]).copy() + +print(f"GT rows: {len(gt)}") +print(f"GT rows with numeric EDSS: {len(gt_numeric)}") + + +# ========================= +# EVALUATE MODELS +# ========================= + +model_dirs = [ + p for p in sorted(RUN_DIR.iterdir()) + if p.is_dir() + and not p.name.startswith("confusion") + and not p.name.startswith("functional_system") + and not p.name.startswith("repeated_run") + and not p.name.startswith("edss_error_distribution") + and not p.name.startswith("edss_threshold_metrics") + and not p.name.startswith("edss_severity_group_metrics") + and not p.name.startswith("structured_output_validity") + and not p.name.startswith("confidence") + and not p.name.startswith("accuracy_time_tradeoff") +] + +rows = [] + +for model_dir in model_dirs: + result_file = find_iter_file(model_dir, TARGET_ITERATION) + + if result_file is None: + print(f"No iter_{TARGET_ITERATION} result file found for {model_dir.name}") + continue + + print("\n" + "=" * 100) + print(f"Model folder: {model_dir.name}") + print(f"Result file: {result_file}") + + row = evaluate_model( + gt_numeric=gt_numeric, + model_dir=model_dir, + result_file=result_file + ) + + if row is not None: + rows.append(row) + +metrics_df = pd.DataFrame(rows) + +if metrics_df.empty: + raise ValueError("No model metrics could be calculated.") + +model_order = [ + "GPT-OSS-120B", + "Qwen3.6-27B", + "Gemma-4-31B-it", +] + +metrics_df["model_display"] = pd.Categorical( + metrics_df["model_display"], + categories=model_order, + ordered=True +) + +metrics_df = metrics_df.sort_values("model_display").reset_index(drop=True) + +metrics_df.to_csv(OUTPUT_TABLE, index=False) + +print("\nAccuracy-time trade-off table:") +print( + metrics_df[ + [ + "model_display", + "n_evaluable", + "mean_inference_time_sec", + "MAE", + "accuracy_within_0_5_percent", + "params_billion", + "estimated_fp16_gb", + "model_size_note", + ] + ] +) + + +# ========================= +# PLOT SETTINGS +# ========================= + +colors = { + "GPT-OSS-120B": "#1F77B4", + "Qwen3.6-27B": "#FF7F0E", + "Gemma-4-31B-it": "#2CA02C", +} + +markers = { + "GPT-OSS-120B": "o", + "Qwen3.6-27B": "s", + "Gemma-4-31B-it": "^", +} + + +# ========================= +# FIGURE 1: TIME VS MAE +# ========================= + +fig, ax = plt.subplots(figsize=(8.2, 5.8)) + +for _, row in metrics_df.iterrows(): + model = str(row["model_display"]) + + ax.scatter( + row["mean_inference_time_sec"], + row["MAE"], + s=row["bubble_area"], + color=colors.get(model, "gray"), + marker=markers.get(model, "o"), + edgecolor="black", + linewidth=0.9, + alpha=0.75, + ) + + ax.text( + row["mean_inference_time_sec"], + row["MAE"], + f" {model}\n ~{row['estimated_fp16_gb']:.0f} GB", + ha="left", + va="center", + fontsize=9, + fontweight="bold", + ) + +ax.set_xlabel("Mean inference time per record (seconds)", fontsize=11, fontweight="bold") +ax.set_ylabel("Mean absolute error (EDSS points)", fontsize=11, fontweight="bold") + +ax.set_title( + "Accuracy–time trade-off: EDSS MAE", + fontsize=13, + fontweight="bold", + pad=15, +) + +ax.grid(True, linestyle="--", alpha=0.3) +ax.set_axisbelow(True) + +for spine in ["top", "right"]: + ax.spines[spine].set_visible(False) + +ax.text( + 0.02, + 0.02, + "Better region: lower-left\nBubble size: estimated FP16 model size", + transform=ax.transAxes, + ha="left", + va="bottom", + fontsize=9, + bbox=dict( + boxstyle="round,pad=0.35", + facecolor="white", + edgecolor="#BBBBBB", + alpha=0.9, + ), +) + +# Bubble-size legend +size_legend_values = [50, 100, 200] +size_handles = [ + ax.scatter( + [], + [], + s=bubble_area_from_gb(v), + color="lightgray", + edgecolor="black", + alpha=0.75, + label=f"~{v} GB", + ) + for v in size_legend_values +] + +ax.legend( + handles=size_handles, + title="Estimated FP16 size", + loc="upper right", + frameon=True, +) + +plt.tight_layout() + +plt.savefig(OUTPUT_MAE_SVG, format="svg", bbox_inches="tight") +plt.savefig(OUTPUT_MAE_PNG, dpi=300, bbox_inches="tight") + +plt.show() + + +# ========================= +# FIGURE 2: TIME VS ±0.5 ACCURACY +# ========================= + +fig, ax = plt.subplots(figsize=(8.2, 5.8)) + +for _, row in metrics_df.iterrows(): + model = str(row["model_display"]) + + ax.scatter( + row["mean_inference_time_sec"], + row["accuracy_within_0_5_percent"], + s=row["bubble_area"], + color=colors.get(model, "gray"), + marker=markers.get(model, "o"), + edgecolor="black", + linewidth=0.9, + alpha=0.75, + ) + + ax.text( + row["mean_inference_time_sec"], + row["accuracy_within_0_5_percent"], + f" {model}\n ~{row['estimated_fp16_gb']:.0f} GB", + ha="left", + va="center", + fontsize=9, + fontweight="bold", + ) + +ax.set_xlabel("Mean inference time per record (seconds)", fontsize=11, fontweight="bold") +ax.set_ylabel("Accuracy within ±0.5 EDSS (%)", fontsize=11, fontweight="bold") + +ax.set_title( + "Accuracy–time trade-off: EDSS ±0.5 accuracy", + fontsize=13, + fontweight="bold", + pad=15, +) + +y_min = max(0, metrics_df["accuracy_within_0_5_percent"].min() - 5) +y_max = min(100, metrics_df["accuracy_within_0_5_percent"].max() + 5) +ax.set_ylim(y_min, y_max) + +ax.grid(True, linestyle="--", alpha=0.3) +ax.set_axisbelow(True) + +for spine in ["top", "right"]: + ax.spines[spine].set_visible(False) + +ax.text( + 0.02, + 0.98, + "Better region: upper-left\nBubble size: estimated FP16 model size", + transform=ax.transAxes, + ha="left", + va="top", + fontsize=9, + bbox=dict( + boxstyle="round,pad=0.35", + facecolor="white", + edgecolor="#BBBBBB", + alpha=0.9, + ), +) + +size_handles = [ + ax.scatter( + [], + [], + s=bubble_area_from_gb(v), + color="lightgray", + edgecolor="black", + alpha=0.75, + label=f"~{v} GB", + ) + for v in size_legend_values +] + +ax.legend( + handles=size_handles, + title="Estimated FP16 size", + loc="lower right", + frameon=True, +) + +plt.tight_layout() + +plt.savefig(OUTPUT_ACC_SVG, format="svg", bbox_inches="tight") +plt.savefig(OUTPUT_ACC_PNG, dpi=300, bbox_inches="tight") + +plt.show() + + +print("\nSaved:") +print(OUTPUT_TABLE) +print(OUTPUT_MAE_SVG) +print(OUTPUT_MAE_PNG) +print(OUTPUT_ACC_SVG) +print(OUTPUT_ACC_PNG) +## + + +# %% Compact accuracy-time bubble plot + +from pathlib import Path +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt + +GT_PATH = Path("/home/shahin/Lab/Doktorarbeit/Barcelona/data/processed/MS_Briefe_400_with_unique_id_SHA3_explore_cleaned_unique.csv") +RUN_DIR = Path("/home/shahin/Lab/Doktorarbeit/Barcelona/results/benchmark_runs/run_20260528_103942") +ITER = 1 + +OUT_DIR = RUN_DIR / f"accuracy_time_tradeoff_iter_{ITER}" +OUT_DIR.mkdir(parents=True, exist_ok=True) + +OUT_CSV = OUT_DIR / f"accuracy_time_tradeoff_compact_iter_{ITER}.csv" +OUT_SVG = OUT_DIR / f"accuracy_time_tradeoff_compact_iter_{ITER}.svg" +OUT_PNG = OUT_DIR / f"accuracy_time_tradeoff_compact_iter_{ITER}.png" + +MODEL_SIZE_GB = { + "GPT-OSS-120B": 195.76, + "Qwen3.6-27B": 55.59, + "Gemma-4-31B-it": 62.58, +} + +COLORS = { + "GPT-OSS-120B": "#1F77B4", + "Qwen3.6-27B": "#FF7F0E", + "Gemma-4-31B-it": "#2CA02C", +} + +MARKERS = { + "GPT-OSS-120B": "o", + "Qwen3.6-27B": "s", + "Gemma-4-31B-it": "^", +} + +MODEL_NAME_MAP = { + "gpt-oss-120b": "GPT-OSS-120B", + "qwen3.6-27b": "Qwen3.6-27B", + "gemma-4-31B-it": "Gemma-4-31B-it", +} + + +def num(s): + return pd.to_numeric(s.astype(str).str.replace(",", ".", regex=False), errors="coerce") + + +def bools(s): + return s.astype(str).str.lower().isin(["true", "1", "yes", "ja"]) + + +def clean_model(x): + return MODEL_NAME_MAP.get(str(x), str(x)) + + +def result_file(model_dir): + files = sorted(model_dir.glob(f"*results_iter_{ITER}_*.csv")) + files = [f for f in files if "incremental" not in f.name.lower()] + return files[0] if files else None + + +gt = pd.read_csv(GT_PATH, sep=";") +gt["row_index"] = gt.index +gt["GT_EDSS"] = num(gt["EDSS"]) +gt = gt.dropna(subset=["GT_EDSS"])[["row_index", "GT_EDSS"]] + +rows = [] + +for model_dir in sorted(RUN_DIR.iterdir()): + if not model_dir.is_dir(): + continue + + if model_dir.name.startswith(("confusion", "functional", "repeated", "edss_", "structured", "confidence", "accuracy_time")): + continue + + f = result_file(model_dir) + if f is None: + continue + + pred = pd.read_csv(f) + + if "row_index" not in pred.columns: + continue + + model_raw = pred["model"].dropna().iloc[0] if "model" in pred.columns and pred["model"].notna().any() else model_dir.name + model = clean_model(model_raw) + + pred["row_index"] = pd.to_numeric(pred["row_index"], errors="coerce") + pred = pred.dropna(subset=["row_index"]).copy() + pred["row_index"] = pred["row_index"].astype(int) + + if "success" in pred.columns: + pred = pred[bools(pred["success"])].copy() + + if "EDSS_is_numeric" in pred.columns: + pred = pred[bools(pred["EDSS_is_numeric"])].copy() + + if "EDSS_in_valid_range" in pred.columns: + pred = pred[bools(pred["EDSS_in_valid_range"])].copy() + + pred_col = "EDSS_numeric" if "EDSS_numeric" in pred.columns else "EDSS" + pred["PRED_EDSS"] = num(pred[pred_col]) + pred = pred.dropna(subset=["PRED_EDSS"]).drop_duplicates("row_index") + + merged = gt.merge(pred, on="row_index", how="inner") + + if merged.empty: + continue + + err = merged["PRED_EDSS"] - merged["GT_EDSS"] + abs_err = err.abs() + + rows.append({ + "model": model, + "result_file": str(f), + "n_evaluable": len(merged), + "mean_inference_time_sec": pd.to_numeric(merged.get("inference_time_sec", np.nan), errors="coerce").mean(), + "MAE": abs_err.mean(), + "accuracy_within_0_5_percent": (abs_err <= 0.5).mean() * 100, + "accuracy_within_1_0_percent": (abs_err <= 1.0).mean() * 100, + "HF_repo_size_GB": MODEL_SIZE_GB.get(model, np.nan), + }) + +df = pd.DataFrame(rows) +df = df.sort_values("model") +df.to_csv(OUT_CSV, index=False) + +print(df[["model", "n_evaluable", "mean_inference_time_sec", "MAE", "accuracy_within_0_5_percent", "HF_repo_size_GB"]]) + +plt.rcParams["font.family"] = "Arial" + +fig, ax = plt.subplots(figsize=(8, 5.8)) + +for _, r in df.iterrows(): + model = r["model"] + size = r["HF_repo_size_GB"] + + ax.scatter( + r["mean_inference_time_sec"], + r["accuracy_within_0_5_percent"], + s=size * 12, + color=COLORS.get(model, "gray"), + marker=MARKERS.get(model, "o"), + edgecolor="black", + linewidth=0.8, + alpha=0.75, + ) + + ax.text( + r["mean_inference_time_sec"], + r["accuracy_within_0_5_percent"], + f" {model}\n {size:.1f} GB", + ha="left", + va="center", + fontsize=9, + fontweight="bold", + ) + +ax.set_xlabel("Mean inference time per record (seconds)", fontsize=11, fontweight="bold") +ax.set_ylabel("Accuracy within ±0.5 EDSS (%)", fontsize=11, fontweight="bold") +ax.set_title("Accuracy–time trade-off by model size", fontsize=13, fontweight="bold", pad=15) + +ax.grid(True, linestyle="--", alpha=0.3) +ax.set_axisbelow(True) + +for spine in ["top", "right"]: + ax.spines[spine].set_visible(False) + +ax.text( + 0.02, + 0.98, + "Better region: upper-left\nBubble size: Hugging Face repository size", + transform=ax.transAxes, + ha="left", + va="top", + fontsize=9, + bbox=dict(boxstyle="round,pad=0.35", facecolor="white", edgecolor="#BBBBBB", alpha=0.9), +) + +legend_sizes = [50, 100, 200] +handles = [ + ax.scatter([], [], s=s * 12, color="lightgray", edgecolor="black", alpha=0.75, label=f"{s} GB") + for s in legend_sizes +] + +ax.legend(handles=handles, title="HF repo size", loc="lower right", frameon=True) + +plt.tight_layout() +plt.savefig(OUT_SVG, format="svg", bbox_inches="tight") +plt.savefig(OUT_PNG, dpi=300, bbox_inches="tight") +plt.show() + +print("\nSaved:") +print(OUT_CSV) +print(OUT_SVG) +print(OUT_PNG) +## + +# %% Improved EDSS absolute error distribution plot + +from pathlib import Path +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt + + +# ========================= +# CONFIGURATION +# ========================= + +GT_PATH = Path( + "/home/shahin/Lab/Doktorarbeit/Barcelona/data/processed/" + "MS_Briefe_400_with_unique_id_SHA3_explore_cleaned_unique.csv" +) + +RUN_DIR = Path( + "/home/shahin/Lab/Doktorarbeit/Barcelona/results/benchmark_runs/run_20260528_103942" +) + +ITER = 1 + +OUTPUT_DIR = RUN_DIR / f"edss_absolute_error_boxplot_iter_{ITER}" +OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + +OUTPUT_LONG = OUTPUT_DIR / f"edss_absolute_errors_long_iter_{ITER}.csv" +OUTPUT_SUMMARY = OUTPUT_DIR / f"edss_absolute_error_summary_iter_{ITER}.csv" +OUTPUT_SVG = OUTPUT_DIR / f"edss_absolute_error_raincloud_iter_{ITER}.svg" +OUTPUT_PNG = OUTPUT_DIR / f"edss_absolute_error_raincloud_iter_{ITER}.png" + +GT_EDSS_COL = "EDSS" +PRED_EDSS_COL = "EDSS_numeric" +PRED_EDSS_FALLBACK_COL = "EDSS" + +ANALYSIS_MODE = "primary_valid_only" + +plt.rcParams["font.family"] = "Arial" + + +# ========================= +# HELPERS +# ========================= + +def to_num(s): + return pd.to_numeric( + s.astype(str).str.replace(",", ".", regex=False), + errors="coerce" + ) + + +def to_bool(s): + return s.astype(str).str.lower().isin(["true", "1", "yes", "ja"]) + + +def clean_model_name(name): + mapping = { + "gpt-oss-120b": "GPT-OSS-120B", + "qwen3.6-27b": "Qwen3.6-27B", + "gemma-4-31B-it": "Gemma-4-31B-it", + } + return mapping.get(str(name), str(name)) + + +def find_result_file(model_dir): + files = sorted(model_dir.glob(f"*results_iter_{ITER}_*.csv")) + files = [ + f for f in files + if "incremental" not in f.name.lower() + and "summary" not in f.name.lower() + and "all_results" not in f.name.lower() + ] + return files[0] if files else None + + +# ========================= +# LOAD GT +# ========================= + +gt = pd.read_csv(GT_PATH, sep=";") +gt["row_index"] = gt.index +gt["GT_EDSS_numeric"] = to_num(gt[GT_EDSS_COL]) +gt = gt.dropna(subset=["GT_EDSS_numeric"])[["row_index", "GT_EDSS_numeric"]].copy() + + +# ========================= +# BUILD LONG ERROR DATA +# ========================= + +long_rows = [] + +skip_prefixes = ( + "confusion", + "functional", + "repeated", + "edss_error", + "edss_threshold", + "edss_severity", + "structured", + "confidence", + "accuracy_time", +) + +for model_dir in sorted(RUN_DIR.iterdir()): + if not model_dir.is_dir() or model_dir.name.startswith(skip_prefixes): + continue + + result_file = find_result_file(model_dir) + if result_file is None: + continue + + pred = pd.read_csv(result_file) + + if "row_index" not in pred.columns: + continue + + model_raw = ( + pred["model"].dropna().iloc[0] + if "model" in pred.columns and pred["model"].notna().any() + else model_dir.name + ) + model = clean_model_name(model_raw) + + pred["row_index"] = pd.to_numeric(pred["row_index"], errors="coerce") + pred = pred.dropna(subset=["row_index"]).copy() + pred["row_index"] = pred["row_index"].astype(int) + + if "success" in pred.columns: + pred = pred[to_bool(pred["success"])].copy() + + if ANALYSIS_MODE == "primary_valid_only": + if "EDSS_is_numeric" in pred.columns: + pred = pred[to_bool(pred["EDSS_is_numeric"])].copy() + if "EDSS_in_valid_range" in pred.columns: + pred = pred[to_bool(pred["EDSS_in_valid_range"])].copy() + + elif ANALYSIS_MODE == "sensitivity_all_numeric": + if "EDSS_is_numeric" in pred.columns: + pred = pred[to_bool(pred["EDSS_is_numeric"])].copy() + + pred_col = PRED_EDSS_COL if PRED_EDSS_COL in pred.columns else PRED_EDSS_FALLBACK_COL + pred["PRED_EDSS_numeric"] = to_num(pred[pred_col]) + + pred = pred.dropna(subset=["PRED_EDSS_numeric"]) + pred = pred.drop_duplicates("row_index", keep="first") + + merged = gt.merge(pred, on="row_index", how="inner") + + if merged.empty: + continue + + merged["error"] = merged["PRED_EDSS_numeric"] - merged["GT_EDSS_numeric"] + merged["abs_error"] = merged["error"].abs() + + for _, row in merged.iterrows(): + long_rows.append({ + "model": model, + "row_index": row["row_index"], + "GT_EDSS_numeric": row["GT_EDSS_numeric"], + "PRED_EDSS_numeric": row["PRED_EDSS_numeric"], + "error": row["error"], + "abs_error": row["abs_error"], + }) + +long_df = pd.DataFrame(long_rows) + +if long_df.empty: + raise ValueError("No evaluable rows found.") + +long_df.to_csv(OUTPUT_LONG, index=False) + + +# ========================= +# SUMMARY +# ========================= + +summary = ( + long_df + .groupby("model") + .agg( + n=("abs_error", "count"), + MAE=("abs_error", "mean"), + median=("abs_error", "median"), + q1=("abs_error", lambda x: x.quantile(0.25)), + q3=("abs_error", lambda x: x.quantile(0.75)), + max_error=("abs_error", "max"), + exact_percent=("abs_error", lambda x: (x == 0).mean() * 100), + within_0_5_percent=("abs_error", lambda x: (x <= 0.5).mean() * 100), + within_1_0_percent=("abs_error", lambda x: (x <= 1.0).mean() * 100), + ) + .reset_index() +) + +summary.to_csv(OUTPUT_SUMMARY, index=False) + +print(summary) + + +# ========================= +# PLOT +# ========================= + +model_order = [ + "GPT-OSS-120B", + "Qwen3.6-27B", + "Gemma-4-31B-it", +] + +model_order = [m for m in model_order if m in long_df["model"].unique()] + +colors = { + "GPT-OSS-120B": "#1F77B4", + "Qwen3.6-27B": "#FF7F0E", + "Gemma-4-31B-it": "#2CA02C", +} + +fig, ax = plt.subplots(figsize=(9.2, 5.8)) + +rng = np.random.default_rng(42) + +for i, model in enumerate(model_order, start=1): + vals = long_df.loc[long_df["model"] == model, "abs_error"].values + color = colors[model] + + # Jittered points + x = rng.normal(i, 0.055, size=len(vals)) + + ax.scatter( + x, + vals, + s=14, + color=color, + alpha=0.18, + edgecolor="none", + zorder=1, + ) + + # Boxplot + bp = ax.boxplot( + vals, + positions=[i], + widths=0.42, + patch_artist=True, + showfliers=False, + showmeans=False, + zorder=2, + ) + + bp["boxes"][0].set_facecolor(color) + bp["boxes"][0].set_alpha(0.35) + bp["boxes"][0].set_edgecolor("black") + bp["boxes"][0].set_linewidth(1.2) + + for median in bp["medians"]: + median.set_color("black") + median.set_linewidth(2.0) + + for whisker in bp["whiskers"]: + whisker.set_color("black") + whisker.set_linewidth(1.1) + + for cap in bp["caps"]: + cap.set_color("black") + cap.set_linewidth(1.1) + + # Mean marker + mean_val = np.mean(vals) + ax.scatter( + i, + mean_val, + s=90, + color="white", + edgecolor="black", + linewidth=1.2, + marker="D", + zorder=4, + ) + + # Summary label + row = summary[summary["model"] == model].iloc[0] + + label = ( + f"n={int(row['n'])}\n" + f"MAE={row['MAE']:.2f}\n" + f"Med={row['median']:.2f}\n" + f"Max={row['max_error']:.2f}" + ) + + ax.text( + i, + 4.25, + label, + ha="center", + va="top", + fontsize=9, + fontweight="bold", + bbox=dict( + boxstyle="round,pad=0.3", + facecolor="white", + edgecolor=color, + alpha=0.95, + ), + ) + +# Reference lines +for y, label in [(0.5, "±0.5"), (1.0, "±1.0")]: + ax.axhline( + y, + color="gray", + linestyle="--", + linewidth=1, + alpha=0.6, + ) + ax.text( + len(model_order) + 0.47, + y, + label, + va="center", + ha="left", + fontsize=9, + color="gray", + ) + +ax.set_xticks(range(1, len(model_order) + 1)) +ax.set_xticklabels(model_order, fontsize=11, fontweight="bold") + +ax.set_ylabel("Absolute EDSS error", fontsize=12, fontweight="bold") +ax.set_xlabel("Model", fontsize=12, fontweight="bold") + +ax.set_title( + "Absolute EDSS prediction error by model", + fontsize=15, + fontweight="bold", + pad=15, +) + +ax.set_ylim(-0.08, 4.45) + +ax.yaxis.grid(True, linestyle="--", alpha=0.25) +ax.set_axisbelow(True) + +for spine in ["top", "right"]: + ax.spines[spine].set_visible(False) + +ax.text( + 0.02, + 0.97, + "Diamond = mean; black line = median\nDashed lines mark clinically relevant error thresholds", + transform=ax.transAxes, + ha="left", + va="top", + fontsize=9, + bbox=dict( + boxstyle="round,pad=0.35", + facecolor="white", + edgecolor="#BBBBBB", + alpha=0.95, + ), +) + +plt.tight_layout() + +plt.savefig(OUTPUT_SVG, format="svg", bbox_inches="tight") +plt.savefig(OUTPUT_PNG, dpi=300, bbox_inches="tight") + +plt.show() + +print("\nSaved:") +print(OUTPUT_LONG) +print(OUTPUT_SUMMARY) +print(OUTPUT_SVG) +print(OUTPUT_PNG) +## + +# %% EDSS error category count plot + +from pathlib import Path +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt + + +# ========================= +# CONFIGURATION +# ========================= + +INPUT_PATH = Path( + "/home/shahin/Lab/Doktorarbeit/Barcelona/results/benchmark_runs/" + "run_20260528_103942/edss_error_distribution_iter_1/" + "edss_error_distribution_iter_1.csv" +) + +OUTPUT_DIR = INPUT_PATH.parent / "figures" +OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + +OUTPUT_SVG = OUTPUT_DIR / "edss_error_category_counts_stacked_bar.svg" +OUTPUT_PNG = OUTPUT_DIR / "edss_error_category_counts_stacked_bar.png" +OUTPUT_TABLE = OUTPUT_DIR / "edss_error_category_counts_plot_table.csv" + +plt.rcParams["font.family"] = "Arial" + + +# ========================= +# HELPERS +# ========================= + +def clean_model_name(name): + mapping = { + "gpt-oss-120b": "GPT-OSS-120B", + "qwen3.6-27b": "Qwen3.6-27B", + "gemma-4-31B-it": "Gemma-4-31B-it", + } + return mapping.get(str(name), str(name)) + + +# ========================= +# LOAD DATA +# ========================= + +df = pd.read_csv(INPUT_PATH) +df["model_display"] = df["model"].apply(clean_model_name) + +model_order = [ + "GPT-OSS-120B", + "Qwen3.6-27B", + "Gemma-4-31B-it", +] + +df["model_display"] = pd.Categorical( + df["model_display"], + categories=model_order, + ordered=True, +) + +df = df.sort_values("model_display").reset_index(drop=True) + + +# ========================= +# ERROR CATEGORIES +# ========================= + +categories = [ + { + "label": "Exact", + "n_col": "exact_match_n", + "color": "#1F77B4", + }, + { + "label": "0.5 error", + "n_col": "error_equal_0_5_n", + "color": "#9ECAE1", + }, + { + "label": ">0.5–1.0", + "n_col": "error_gt_0_5_le_1_0_n", + "color": "#FDDC7A", + }, + { + "label": ">1.0", + "n_col": "error_gt_1_0_n", + "color": "#F28E2B", + }, + { + "label": "Missing/invalid", + "n_col": "n_missing_or_invalid_against_gt_numeric", + "color": "#D62728", + }, +] + +plot_rows = [] + +for _, row in df.iterrows(): + total = int(row["n_gt_numeric"]) + + for cat in categories: + n = int(row[cat["n_col"]]) + percent = n / total * 100 if total > 0 else np.nan + + plot_rows.append({ + "model": row["model_display"], + "category": cat["label"], + "n": n, + "percent": percent, + "total_numeric_gt": total, + }) + +plot_df = pd.DataFrame(plot_rows) +plot_df.to_csv(OUTPUT_TABLE, index=False) + +print("\nPlot table:") +print(plot_df) + + +# ========================= +# PLOT +# ========================= + +fig, ax = plt.subplots(figsize=(11, 5.8)) + +y = np.arange(len(df)) +left = np.zeros(len(df)) + +for cat in categories: + values = [] + percents = [] + + for model in df["model_display"]: + sub = plot_df[ + (plot_df["model"] == model) + & (plot_df["category"] == cat["label"]) + ].iloc[0] + + values.append(sub["n"]) + percents.append(sub["percent"]) + + values = np.array(values) + percents = np.array(percents) + + bars = ax.barh( + y, + values, + left=left, + color=cat["color"], + edgecolor="white", + linewidth=0.8, + label=cat["label"], + ) + + for i, (n, pct) in enumerate(zip(values, percents)): + if n == 0: + continue + + # Label inside if enough space, otherwise just skip tiny labels + if n >= 5: + ax.text( + left[i] + n / 2, + i, + f"{int(n)}\n({pct:.1f}%)", + ha="center", + va="center", + fontsize=8, + fontweight="bold", + color="white" if cat["label"] in ["Exact", "Missing/invalid"] else "black", + ) + + left += values + + +# Add total n to the right +for i, row in df.iterrows(): + ax.text( + row["n_gt_numeric"] + 4, + i, + f"total n={int(row['n_gt_numeric'])}", + va="center", + ha="left", + fontsize=9, + fontweight="bold", + ) + +ax.set_yticks(y) +ax.set_yticklabels(df["model_display"], fontsize=11, fontweight="bold") + +ax.set_xlabel("Number of numeric reference EDSS records", fontsize=11, fontweight="bold") +ax.set_ylabel("Model", fontsize=11, fontweight="bold") + +ax.set_title( + "EDSS prediction error categories by model", + fontsize=14, + fontweight="bold", + pad=15, +) + +ax.set_xlim(0, int(df["n_gt_numeric"].max()) + 45) + +ax.xaxis.grid(True, linestyle="--", alpha=0.25) +ax.set_axisbelow(True) + +for spine in ["top", "right", "left"]: + ax.spines[spine].set_visible(False) + +ax.tick_params(axis="y", length=0) + +ax.legend( + loc="lower center", + bbox_to_anchor=(0.5, 1.02), + ncol=5, + frameon=False, +) + +plt.tight_layout(rect=[0, 0, 1, 0.90]) + +plt.savefig(OUTPUT_SVG, format="svg", bbox_inches="tight") +plt.savefig(OUTPUT_PNG, dpi=300, bbox_inches="tight") + +plt.show() + +print("\nSaved:") +print(OUTPUT_SVG) +print(OUTPUT_PNG) +print(OUTPUT_TABLE) +##