n_samples

Author: michielkallenberg

.gitignore

@@ -1,6 +1,7 @@
 data
 output
 outputs/
+_build/
 .venv/
 venv/
 __pycache__

cybench/runs/analysis/benchmark_run_catalog.py

@@ -196,7 +196,7 @@ def _metrics_from_report_yaml(path: Path) -> dict[str, Any]:
     raw = OmegaConf.to_container(OmegaConf.load(path))
     if not isinstance(raw, dict):
         return {}
-    return flatten_report_metrics(
+    flat = flatten_report_metrics(
         {
             "n_regions": raw.get("n_regions"),
             "n_years": raw.get("n_years"),
@@ -207,6 +207,9 @@ def _metrics_from_report_yaml(path: Path) -> dict[str, Any]:
             "anomaly": raw.get("anomaly") or {},
         }
     )
+    if raw.get("n_train") is not None:
+        flat["n_train"] = raw.get("n_train")
+    return flat
 
 
 def load_run_metrics(run: BenchmarkRun, *, seed: int = 42) -> dict[str, Any] | None:
@@ -221,7 +224,7 @@ def load_run_metrics(run: BenchmarkRun, *, seed: int = 42) -> dict[str, Any] | N
             return None
         return {
             k: summary.get(k)
-            for k in ("n_regions", "n_years", "n_samples") + METRIC_KEYS
+            for k in ("n_regions", "n_years", "n_samples", "n_train") + METRIC_KEYS
         }
 
     metrics_path = _screening_metrics_path(run.path, seed=seed)

cybench/runs/analysis/collect_walk_forward_results.py

@@ -37,6 +37,7 @@
 from pathlib import Path
 from typing import Any, cast
 
+import numpy as np
 import pandas as pd
 import yaml
 from omegaconf import OmegaConf
@@ -93,6 +94,54 @@ def resolve_model_column(run_dir: Path, model_slug: str) -> str:
 
 COUNT_METRIC_KEYS: tuple[str, ...] = ("n_regions", "n_years", "n_samples")
 
+_TRAIN_LOG_RE = re.compile(r"Train final model on (\d+) datapoints", re.IGNORECASE)
+
+
+def _n_train_from_split_metrics(run_dir: Path, seed: int) -> list[int]:
+    """Training row counts recorded per walk-forward split (``report_metrics.yaml``)."""
+    sizes: list[int] = []
+    for split_dir in sorted(run_dir.iterdir()):
+        if not split_dir.is_dir() or not re.fullmatch(r"\d{4}", split_dir.name):
+            continue
+        metrics_path = split_dir / str(seed) / "report_metrics.yaml"
+        if not metrics_path.is_file():
+            continue
+        try:
+            raw = yaml.safe_load(metrics_path.read_text(encoding="utf-8")) or {}
+        except (OSError, yaml.YAMLError):
+            continue
+        n_train = raw.get("n_train")
+        if n_train is not None and not pd.isna(n_train):
+            sizes.append(int(n_train))
+    return sizes
+
+
+def _n_train_from_run_log(run_dir: Path) -> list[int]:
+    """Fallback for runs logged before ``n_train`` was written to report_metrics."""
+    log_path = run_dir / "run_experiments.log"
+    if not log_path.is_file():
+        return []
+    try:
+        text = log_path.read_text(encoding="utf-8", errors="replace")
+    except OSError:
+        return []
+    return [int(m.group(1)) for m in _TRAIN_LOG_RE.finditer(text)]
+
+
+def summarize_walk_forward_n_train(run_dir: Path, *, seed: int) -> dict[str, Any]:
+    """Mean training rows across walk-forward splits (train window grows per year)."""
+    sizes = _n_train_from_split_metrics(run_dir, seed)
+    if not sizes:
+        sizes = _n_train_from_run_log(run_dir)
+    if not sizes:
+        return {"n_train": None, "n_train_splits": 0}
+    return {
+        "n_train": int(round(float(np.mean(sizes)))),
+        "n_train_min": int(min(sizes)),
+        "n_train_max": int(max(sizes)),
+        "n_train_splits": len(sizes),
+    }
+
 
 def discover_run_seeds(run_dir: Path) -> list[int]:
     """Seed subfolders under walk-forward year splits (e.g. 2016/42/test_preds.csv)."""
@@ -278,6 +327,7 @@ def collect_walk_forward_run(
 
     assert plot_df is not None
     summary = aggregate_flat_metrics_across_seeds(per_seed_flat, seeds)
+    summary.update(summarize_walk_forward_n_train(run.path, seed=plot_seed))
     return summary, per_seed_rows, plot_df, model_col, plot_seed
 
 

cybench/runs/analysis/global_insights_lib.py

@@ -63,7 +63,7 @@ def load_summary_frame(summary_paths: list[Path]) -> pd.DataFrame:
     if not frames:
         return pd.DataFrame()
     out = pd.concat(frames, ignore_index=True)
-    for col in ("nrmse", "r2", "n_samples", "n_regions", "n_years"):
+    for col in ("nrmse", "r2", "n_samples", "n_train", "n_regions", "n_years"):
         if col in out.columns:
             out[col] = pd.to_numeric(out[col], errors="coerce")
     return out

cybench/runs/analysis/model_family_radar_lib.py

@@ -195,12 +195,13 @@ def family_for_model(model: str) -> str | None:
 
 
 SAMPLE_SCATTER_METRICS: tuple[dict[str, str], ...] = (
-    {"key": "r2", "label": "Overall R²"},
-    {"key": "r2_spatial", "label": "Spatial R² (med/yr)"},
-    {"key": "r2_spatial_agg", "label": "Spatial R² (agg)"},
-    {"key": "r2_temporal", "label": "Temporal R² (med/reg)"},
-    {"key": "r2_temporal_agg", "label": "Temporal R² (agg)"},
-    {"key": "r2_anomaly", "label": "Anomaly R² (med/reg)"},
+    {"key": "nrmse", "label": "Overall NRMSE", "lower_is_better": True},
+    {"key": "r2", "label": "Overall R²", "lower_is_better": False},
+    {"key": "r2_spatial", "label": "Spatial R² (med/yr)", "lower_is_better": False},
+    {"key": "r2_spatial_agg", "label": "Spatial R² (agg)", "lower_is_better": False},
+    {"key": "r2_temporal", "label": "Temporal R² (med/reg)", "lower_is_better": False},
+    {"key": "r2_temporal_agg", "label": "Temporal R² (agg)", "lower_is_better": False},
+    {"key": "r2_anomaly", "label": "Anomaly R² (med/reg)", "lower_is_better": False},
 )
 
 
@@ -210,15 +211,15 @@ def build_sample_scatter_slice(
     batch_horizon: str,
     crop: str | None = None,
 ) -> list[dict[str, Any]]:
-    """Per-family points for performance vs pooled test sample count."""
+    """Per-family points for performance vs mean walk-forward training set size."""
     work = df[df["batch_horizon"] == batch_horizon].copy() if "batch_horizon" in df.columns else df
     if crop:
         work = work[work["crop"] == crop]
     if work.empty:
         return []
 
     metric_keys = [m["key"] for m in SAMPLE_SCATTER_METRICS]
-    for key in metric_keys + ["n_samples"]:
+    for key in metric_keys + ["n_train"]:
         if key in work.columns:
             work[key] = pd.to_numeric(work[key], errors="coerce")
 
@@ -229,8 +230,8 @@ def build_sample_scatter_slice(
             continue
         points: list[dict[str, Any]] = []
         for _, row in sub.iterrows():
-            n_samples = row.get("n_samples")
-            if pd.isna(n_samples) or int(n_samples) <= 0:
+            n_train = row.get("n_train")
+            if pd.isna(n_train) or int(n_train) <= 0:
                 continue
             model = str(row["model"])
             point: dict[str, Any] = {
@@ -239,7 +240,7 @@ def build_sample_scatter_slice(
                 "country": str(row.get("country", "")),
                 "crop": str(row.get("crop", "")),
                 "dataset": f"{row.get('crop', '')}_{row.get('country', '')}",
-                "n_samples": int(n_samples),
+                "n_train": int(n_train),
                 "metrics": {},
             }
             for metric in metric_keys:

cybench/runs/run_experiments.py

@@ -364,6 +364,7 @@ def main(cfg):
                 target_col=KEY_TARGET,
                 model_col=model_col,
             )
+            report_metrics["n_train"] = len(train_dataset)
             metric_ls.append(eval_metric)
             log.info(f"Split {train_test_split[-1]} (seed {seed}) finished with metrics: {eval_metric}")
             log.info(

cybench/runs/viz/model_family_radar_template.html

@@ -188,8 +188,8 @@ <h2>Raw metrics (median across datasets)</h2>
     </div>
 
     <div class="card">
-      <h2>Performance vs test samples</h2>
-      <p class="muted">Each point is one crop×country dataset. X-axis: pooled walk-forward test rows (region-years).</p>
+      <h2>Performance vs training set size</h2>
+      <p class="muted">Each point is one crop×country dataset. X-axis: mean training rows across walk-forward splits (region-years used to fit).</p>
       <div class="controls">
         <label>Metric
           <select id="scatter-metric-select"></select>
@@ -228,11 +228,19 @@ <h2>Performance vs test samples</h2>
     });
     if (!scatterMetricSelect.options.length) {
       const opt = document.createElement("option");
-      opt.value = "r2";
-      opt.textContent = "Overall R²";
+      opt.value = "nrmse";
+      opt.textContent = "Overall NRMSE";
       scatterMetricSelect.appendChild(opt);
     }
 
+    const scatterMetricByKey = Object.fromEntries(
+      (DATA.sample_scatter_metrics || []).map(m => [m.key, m])
+    );
+
+    function scatterMetricLabel(key) {
+      return (scatterMetricByKey[key] || {}).label || key;
+    }
+
     const horizonLabels = { eos: "End of season", mid: "Mid-season" };
     Object.keys(DATA.by_horizon).forEach(hz => {
       const opt = document.createElement("option");
@@ -336,10 +344,12 @@ <h2>Performance vs test samples</h2>
     }
 
     function renderScatterFamilies() {
-      const metricKey = scatterMetricSelect.value || "r2";
+      const metricKey = scatterMetricSelect.value || "nrmse";
+      const metricMeta = scatterMetricByKey[metricKey] || {};
+      const lowerIsBetter = !!metricMeta.lower_is_better;
       const families = currentScatterFamilies();
       if (!families.length) {
-        scatterGrid.innerHTML = '<p class="muted">No scatter data for this selection.</p>';
+        scatterGrid.innerHTML = '<p class="muted">No scatter data for this selection (re-collect summaries to populate n_train).</p>';
         return;
       }
 
@@ -357,12 +367,17 @@ <h2>Performance vs test samples</h2>
         const innerW = width - pad.l - pad.r;
         const innerH = height - pad.t - pad.b;
 
-        const xs = points.map(p => p.n_samples);
+        const xs = points.map(p => p.n_train);
         const ys = points.map(p => p.metrics[metricKey]);
         const xMin = Math.min(...xs);
         const xMax = Math.max(...xs);
-        const yMin = Math.min(-0.2, ...ys);
-        const yMax = Math.max(1.0, ...ys);
+        const yPad = lowerIsBetter ? 0.05 : 0.2;
+        const yMin = lowerIsBetter
+          ? Math.max(0, Math.min(...ys) - yPad)
+          : Math.min(-0.2, ...ys);
+        const yMax = lowerIsBetter
+          ? Math.max(...ys) + yPad
+          : Math.max(1.0, ...ys);
         const xScale = v => pad.l + ((v - xMin) / (xMax - xMin || 1)) * innerW;
         const yScale = v => pad.t + innerH - ((v - yMin) / (yMax - yMin || 1)) * innerH;
 
@@ -371,7 +386,7 @@ <h2>Performance vs test samples</h2>
 
         let svg = `<svg viewBox="0 0 ${width} ${height}" role="img" aria-label="${fam.family} scatter">`;
         svg += `<rect x="${pad.l}" y="${pad.t}" width="${innerW}" height="${innerH}" fill="#fff" stroke="#e6eaef"/>`;
-        if (yMin < 0 && yMax > 0) {
+        if (!lowerIsBetter && yMin < 0 && yMax > 0) {
           const y0 = yScale(0);
           svg += `<line x1="${pad.l}" y1="${y0}" x2="${pad.l + innerW}" y2="${y0}" stroke="#bbb" stroke-dasharray="4 3"/>`;
         }
@@ -381,15 +396,15 @@ <h2>Performance vs test samples</h2>
           svg += `<text x="${xScale(tx)}" y="${height - 8}" text-anchor="middle" font-size="10" fill="#656d76">${Math.round(tx)}</text>`;
           svg += `<text x="8" y="${yScale(ty) + 3}" text-anchor="start" font-size="10" fill="#656d76">${ty.toFixed(2)}</text>`;
         }
-        svg += `<text x="${pad.l + innerW / 2}" y="${height - 2}" text-anchor="middle" font-size="11" fill="#1f2328">Test samples</text>`;
-        svg += `<text x="14" y="${pad.t + innerH / 2}" text-anchor="middle" font-size="11" fill="#1f2328" transform="rotate(-90 14 ${pad.t + innerH / 2})">R²</text>`;
+        svg += `<text x="${pad.l + innerW / 2}" y="${height - 2}" text-anchor="middle" font-size="11" fill="#1f2328">Training rows (mean)</text>`;
+        svg += `<text x="14" y="${pad.t + innerH / 2}" text-anchor="middle" font-size="11" fill="#1f2328" transform="rotate(-90 14 ${pad.t + innerH / 2})">${scatterMetricLabel(metricKey)}</text>`;
 
         points.forEach(p => {
-          const cx = xScale(p.n_samples);
+          const cx = xScale(p.n_train);
           const cy = yScale(p.metrics[metricKey]);
           const col = modelColor(p.model, modelIdx[p.model]);
           svg += `<circle cx="${cx}" cy="${cy}" r="5" fill="${col}" fill-opacity="0.85" stroke="#fff" stroke-width="1">
-            <title>${p.display_name} · ${p.dataset}\nN=${p.n_samples}, ${metricKey}=${p.metrics[metricKey]}</title>
+            <title>${p.display_name} · ${p.dataset}\nTrain=${p.n_train}, ${metricKey}=${p.metrics[metricKey]}</title>
           </circle>`;
         });
         svg += "</svg>";

tests/runs/test_collect_walk_forward_results.py

@@ -237,3 +237,18 @@ def test_summary_rows_to_dashboard_records(tmp_path: Path):
     scatter_recs = [r for r in records if r.get("images", {}).get("scatter")]
     assert scatter_recs
     assert "maize_NL" in scatter_recs[0]["images"]["scatter"]
+
+
+def test_summarize_walk_forward_n_train_from_log(tmp_path: Path):
+    from cybench.runs.analysis.collect_walk_forward_results import summarize_walk_forward_n_train
+
+    run_dir = tmp_path / "maize_DE_tabpfn_walk_forward_eos_20260616_135557"
+    run_dir.mkdir()
+    (run_dir / "run_experiments.log").write_text(
+        "Train final model on 2783 datapoints (years [...])\n"
+        "Train final model on 3225 datapoints (years [...])\n",
+        encoding="utf-8",
+    )
+    out = summarize_walk_forward_n_train(run_dir, seed=42)
+    assert out["n_train_splits"] == 2
+    assert out["n_train"] == 3004