Static mixture explorer build script: rank-INT swarm targets, DSP per-task fits, Tristan's gamma, Pareto-improver, embedded JS gradient-slider mechanic

build_static.py

"""Build a single-file static HTML mixture explorer from the swarm CSVs.

Output: docs/index.html (data inlined, Plotly via CDN).
Open it directly with `file://` or serve with `python -m http.server -d docs/`.
The docs/ folder is what GitHub Pages serves.

Per-task fits use Calvin's DSP canonical form:
  e_p,d = c_p,d * w_p,d
  z_d   = e_0,d + gamma * e_1,d
  y_hat(w) = -b0 + sum_d a_d (1 - exp(-rho_d z_d))
                  - sum_d p_d softplus(log(1+z_d) - tau_d)^2
Fit with L-BFGS-B against per-task rank-INT z's; params dumped to JSON
so the JS can evaluate the predictor live on the by-mixture tab.
"""
import json
import os

import jax
import jax.numpy as jnp
import numpy as np
import polars as pl
from jax import jit, value_and_grad
from jax.flatten_util import ravel_pytree
from scipy.optimize import minimize
from scipy.stats import norm, rankdata

raw = pl.read_csv("raw_metric_matrix_300m.csv", infer_schema_length=500).filter(
    pl.col("status") == "completed"
)

_MMLU_KEEP = {"lm_eval/mmlu_sl_verb_5shot/bpb"}
# Collapse all paloma sub-benchmarks into the single macro-average task.
_PALOMA_KEEP = {"eval/paloma/macro_bpb"}
_AGG_DROP = {
    "eval/bpb", "eval/macro_bpb",
    "eval/paloma/bpb",
    "eval/uncheatable_eval/bpb", "eval/uncheatable_eval/macro_bpb",
}
_TASK_DROP = {
    "teacher_forced/gsm8k_5shot_answer_hash/bpb",
    "lm_eval/arc_easy/bpb",
    "lm_eval/piqa/bpb",
    "mcq_smooth/swag_0shot/bpb",
}
_EXTRA = [("lm_eval/socialiqa_5shot/choice_logprob", +1.0)]


def keep(c):
    # Override: paloma/macro_bpb doesn't end in /bpb, force-keep it.
    if c in _PALOMA_KEEP:
        return True
    if not c.endswith("/bpb"):
        return False
    if c in _AGG_DROP or c in _TASK_DROP:
        return False
    if not c.startswith(("eval/paloma/", "eval/uncheatable_eval/",
                         "lm_eval/", "mcq_smooth/", "teacher_forced/")):
        return False
    if c.startswith("lm_eval/mmlu_") and c not in _MMLU_KEEP:
        return False
    # Drop every individual paloma subdomain — only the macro stays.
    if c.startswith("eval/paloma/"):
        return False
    return True


cols_raw = set(raw.columns)
task_cols, signs = [], []
for c in [x for x in raw.columns if keep(x)]:
    base = c.removesuffix("/bpb")
    cl = base + "/choice_logprob"
    if base.startswith(("lm_eval/", "mcq_smooth/")) and cl in cols_raw:
        task_cols.append(cl); signs.append(+1.0)
    else:
        task_cols.append(c); signs.append(-1.0)
for c, s in _EXTRA:
    if c in cols_raw and c not in task_cols:
        task_cols.append(c); signs.append(s)
signs = np.array(signs)

X = raw.select(task_cols).to_numpy().astype(np.float64) * signs[None, :]
n = X.shape[0]
Z = norm.ppf((np.apply_along_axis(rankdata, 0, X) - 0.5) / n)
# Integer rank 0..N-1 per (run, task) for the slider. argsort twice = ranks.
# runByRank[t][r] gives the run with rank r on task t (O(1) lookup on drag).
R = np.argsort(np.argsort(Z, axis=0), axis=0)
runByRank = np.argsort(Z, axis=0).T

domains = sorted(c.removeprefix("phase_0_") for c in raw.columns
                 if c.startswith("phase_0_"))
w0 = raw.select([f"phase_0_{d}" for d in domains]).to_numpy().astype(np.float64)
w1 = raw.select([f"phase_1_{d}" for d in domains]).to_numpy().astype(np.float64)
w0 = w0 / w0.sum(axis=1, keepdims=True)
w1 = w1 / w1.sum(axis=1, keepdims=True)

meta = pl.read_csv("two_phase_many_epoch_metadata.csv")
meta_row = {r["domain_name"]: r for r in meta.iter_rows(named=True)}
c0 = np.array([meta_row[d]["phase_0_epoch_multiplier"] for d in domains])
c1 = np.array([meta_row[d]["phase_1_epoch_multiplier"] for d in domains])
w0_prop = ((1.0 / c0) / (1.0 / c0).sum()).tolist()
w1_prop = ((1.0 / c1) / (1.0 / c1).sum()).tolist()

short = [c.removesuffix("/bpb").removesuffix("/choice_logprob") for c in task_cols]

# Fixed y-axis order for the bar chart: largest natural-size domains on top
# (w_d_prop ∝ 1/c_d). Stable across runs, no reshuffle on slider drag.
domain_order = np.argsort(-np.asarray(w0_prop)).tolist()

# Stable x-axis upper bound for absolute-weight view: round the global swarm
# max up to the nearest 0.05 so the axis label is tidy.
wmax = float(np.ceil(max(w0.max(), w1.max()) * 20) / 20)

# Tristan's perplexity-correlations estimator γ_j (Thrush et al. 2024):
#   γ_j = Σ_{k,l} sign(y_k - y_l) (rank_j(x_k) - rank_j(x_l))
# computed for each (task, phase, domain). Normalised to [-1, 1] by
# dividing by its maximum n(n²-1)/3 (the value at perfect agreement)
# so the heatmap is read on the same scale as a correlation. The signed
# pairwise sum is mathematically equivalent to Spearman's ρ.
print("computing Tristan's γ_j…")
_T_tasks = Z.shape[1]
_D_dom = w0.shape[1]
corr0 = np.zeros((_T_tasks, _D_dom))
corr1 = np.zeros((_T_tasks, _D_dom))
# Combined-phase weight per (run, domain) = w0 + w1. Lets the user see
# "does *total* weight on this domain across both phases associate with
# this task" without per-phase confounding.
_w_sum = w0 + w1
corr_sum = np.zeros((_T_tasks, _D_dom))


def _tristan_gamma(x, y):
    # γ = Σ_{k,l} sign(y_k - y_l)(rank(x_k) - rank(x_l))
    # = 2 Σ_k rank(x_k) · [2 rank(y_k) - n - 1]  (by k↔l swap symmetry,
    # since Σ_l sign(y_k - y_l) = 2 rank(y_k) - n - 1 when y has no ties).
    _n = len(x)
    _rx = rankdata(x)
    _ry = rankdata(y)
    _gamma = 2.0 * float(np.sum(_rx * (2 * _ry - _n - 1)))
    _gamma_max = _n * (_n * _n - 1) / 3.0  # value at perfect agreement
    return _gamma / _gamma_max


for _t in range(_T_tasks):
    for _d in range(_D_dom):
        corr0[_t, _d]    = _tristan_gamma(w0[:, _d], Z[:, _t])
        corr1[_t, _d]    = _tristan_gamma(w1[:, _d], Z[:, _t])
        corr_sum[_t, _d] = _tristan_gamma(_w_sum[:, _d], Z[:, _t])


# Fit Calvin's DSP form per task on the rank-INT'd targets. ~30-60s for 41
# tasks; runs once at build time, params dumped to JSON for live JS eval.
_c0_jax = jnp.asarray(c0)
_c1_jax = jnp.asarray(c1)
_D = len(domains)


def _spinv(x):
    return float(np.log(np.exp(x) - 1.0))


def _predict_dsp(theta, w0_b, w1_b):
    b0 = theta["b0"]
    rho = jax.nn.softplus(theta["log_rho"])
    tau = theta["tau"]
    gamma = jax.nn.softplus(theta["log_gamma"])
    a = jax.nn.softplus(theta["log_a"])
    p = jax.nn.softplus(theta["log_p"])
    e0 = w0_b * _c0_jax[None, :]
    e1 = w1_b * _c1_jax[None, :]
    z = e0 + gamma * e1
    signal = a[None, :] * (1.0 - jnp.exp(-rho[None, :] * z))
    u = jnp.log1p(z) - tau[None, :]
    penalty = p[None, :] * jax.nn.softplus(u) ** 2
    return -(b0 - signal.sum(axis=1) + penalty.sum(axis=1))


_theta_init = {
    "b0":        jnp.array(0.0),
    "log_rho":   jnp.full(_D, _spinv(0.3)),
    "tau":       jnp.full(_D, 2.0),
    "log_gamma": jnp.array(_spinv(1.0)),
    "log_a":     jnp.full(_D, _spinv(0.1)),
    "log_p":     jnp.full(_D, _spinv(0.01)),
}
_flat_init, _unravel = ravel_pytree(_theta_init)


@jit
def _loss_vg(theta_flat, w0_b, w1_b, y_b):
    def loss(t):
        return jnp.mean((_predict_dsp(_unravel(t), w0_b, w1_b) - y_b) ** 2)
    return value_and_grad(loss)(theta_flat)


def _fit_one(y):
    def f(t):
        v, g = _loss_vg(jnp.asarray(t),
                        jnp.asarray(w0), jnp.asarray(w1), jnp.asarray(y))
        return float(v), np.asarray(g, dtype=np.float64)
    res = minimize(f, np.asarray(_flat_init, dtype=np.float64),
                   method="L-BFGS-B", jac=True,
                   options={"maxiter": 300, "ftol": 1e-9})
    return _unravel(jnp.asarray(res.x))


print(f"fitting DSP form for {Z.shape[1]} tasks…")
dsp_params = []
for _t in range(Z.shape[1]):
    th = _fit_one(Z[:, _t])
    # Residual RMSE on the swarm = per-task prediction-uncertainty std,
    # in rank-INT z units. Used in the UI to draw error bars around the
    # predicted percentile.
    _pred = np.asarray(_predict_dsp(th,
                                    jnp.asarray(w0), jnp.asarray(w1)))
    _sigma = float(np.sqrt(np.mean((_pred - Z[:, _t]) ** 2)))
    dsp_params.append({
        "b0":         round(float(th["b0"]),         5),
        "log_gamma":  round(float(th["log_gamma"]),  5),
        "log_rho":   [round(float(x), 5) for x in th["log_rho"]],
        "tau":       [round(float(x), 5) for x in th["tau"]],
        "log_a":     [round(float(x), 5) for x in th["log_a"]],
        "log_p":     [round(float(x), 5) for x in th["log_p"]],
        "sigma":      round(_sigma, 4),
    })
    print(f"  [{_t+1:>2d}/{Z.shape[1]}] {short[_t]:50s} sigma={_sigma:.3f}")


# Predicted-Pareto-improver mixture vs token-proportional. We optimise
#   L(w) = sum_t max(0, z_prop_t + eps - z_pred_t(w))^2
# i.e. squared-shortfall against the per-task predictions at the
# proportional baseline, with a small margin eps so the optimiser pushes
# for strict improvement. Adam on softmax(theta) keeps w on the simplex
# without projection.
print("optimising Pareto-improver mixture vs proportional…")
_w0_prop_jax = jnp.asarray((1.0 / c0) / (1.0 / c0).sum())
_w1_prop_jax = jnp.asarray((1.0 / c1) / (1.0 / c1).sum())

_b0_arr        = jnp.asarray([p["b0"]        for p in dsp_params])
_log_gamma_arr = jnp.asarray([p["log_gamma"] for p in dsp_params])
_log_rho_arr   = jnp.asarray([p["log_rho"]   for p in dsp_params])
_tau_arr       = jnp.asarray([p["tau"]       for p in dsp_params])
_log_a_arr     = jnp.asarray([p["log_a"]     for p in dsp_params])
_log_p_arr     = jnp.asarray([p["log_p"]     for p in dsp_params])


@jit
def _predict_all(w0v, w1v):
    gamma = jax.nn.softplus(_log_gamma_arr)[:, None]
    rho   = jax.nn.softplus(_log_rho_arr)
    a_t   = jax.nn.softplus(_log_a_arr)
    p_t   = jax.nn.softplus(_log_p_arr)
    e0 = (w0v * _c0_jax)[None, :]
    e1 = (w1v * _c1_jax)[None, :]
    z = e0 + gamma * e1
    signal = a_t * (1.0 - jnp.exp(-rho * z))
    u = jnp.log1p(z) - _tau_arr
    penalty = p_t * jax.nn.softplus(u) ** 2
    return -(_b0_arr - signal.sum(axis=1) + penalty.sum(axis=1))


_z_baseline = _predict_all(_w0_prop_jax, _w1_prop_jax)
_eps_margin = 0.05


@jit
def _pareto_loss(theta0, theta1):
    w0v = jax.nn.softmax(theta0)
    w1v = jax.nn.softmax(theta1)
    zs = _predict_all(w0v, w1v)
    shortfall = jnp.maximum(0.0, _z_baseline + _eps_margin - zs)
    return jnp.sum(shortfall ** 2)


@jit
def _adam_step(theta0, theta1, m0, m1, v0, v1, t):
    g0, g1 = jax.grad(_pareto_loss, argnums=(0, 1))(theta0, theta1)
    lr, b1, b2, eps = 0.05, 0.9, 0.999, 1e-8
    m0 = b1 * m0 + (1 - b1) * g0
    m1 = b1 * m1 + (1 - b1) * g1
    v0 = b2 * v0 + (1 - b2) * g0 ** 2
    v1 = b2 * v1 + (1 - b2) * g1 ** 2
    bc = 1 - b1 ** (t + 1)
    bv = 1 - b2 ** (t + 1)
    theta0 = theta0 - lr * (m0 / bc) / (jnp.sqrt(v0 / bv) + eps)
    theta1 = theta1 - lr * (m1 / bc) / (jnp.sqrt(v1 / bv) + eps)
    return theta0, theta1, m0, m1, v0, v1


_theta0 = jnp.log(_w0_prop_jax + 1e-12)
_theta1 = jnp.log(_w1_prop_jax + 1e-12)
_m0 = jnp.zeros_like(_theta0); _m1 = jnp.zeros_like(_theta1)
_v0 = jnp.zeros_like(_theta0); _v1 = jnp.zeros_like(_theta1)
for _step in range(1000):
    _theta0, _theta1, _m0, _m1, _v0, _v1 = _adam_step(
        _theta0, _theta1, _m0, _m1, _v0, _v1, _step
    )

w_pareto0 = np.asarray(jax.nn.softmax(_theta0))
w_pareto1 = np.asarray(jax.nn.softmax(_theta1))
_z_pareto = np.asarray(_predict_all(jax.nn.softmax(_theta0),
                                    jax.nn.softmax(_theta1)))
_delta = _z_pareto - np.asarray(_z_baseline)
_n_improved = int((_delta > 0).sum())
_n_strict   = int((_delta > _eps_margin).sum())
print(f"  improved on {_n_improved}/{_T_tasks} tasks "
      f"(strict >+{_eps_margin}: {_n_strict})")
print(f"  min Δz = {_delta.min():+.3f}  median Δz = "
      f"{float(np.median(_delta)):+.3f}  max Δz = {_delta.max():+.3f}")

data = {
    "task_names": short,
    "domain_names": domains,
    "domain_order": domain_order,
    "wmax": wmax,
    "c0": [round(float(x), 4) for x in c0],
    "c1": [round(float(x), 4) for x in c1],
    "R": R.tolist(),                    # rank per (run, task), 0..N-1
    "runByRank": runByRank.tolist(),    # T × N: run idx with rank r on task t
    "w0": np.round(w0, 4).tolist(),
    "w1": np.round(w1, 4).tolist(),
    "w0_prop": [round(x, 4) for x in w0_prop],
    "w1_prop": [round(x, 4) for x in w1_prop],
    "w0_pareto": [round(float(x), 4) for x in w_pareto0],
    "w1_pareto": [round(float(x), 4) for x in w_pareto1],
    "dsp_params": dsp_params,
    "corr0":    np.round(corr0,    3).tolist(),
    "corr1":    np.round(corr1,    3).tolist(),
    "corr_sum": np.round(corr_sum, 3).tolist(),
}

HTML = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<title>Swarm Mixture Explorer</title>
<style>
  *, *::before, *::after { box-sizing: border-box; }
  body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
         margin: 0; padding: clamp(12px, 3vw, 28px); color: #1a1a1a;
         max-width: 1400px; margin: 0 auto; }
  h1 { font-size: clamp(18px, 3vw, 22px); margin: 0 0 6px; }
  .intro { color: #555; font-size: clamp(12px, 1.5vw, 13px); line-height: 1.45;
           max-width: 900px; margin-bottom: 18px; }
  .run-info { font-weight: 600; margin: 18px 0 6px;
              font-size: clamp(13px, 1.6vw, 14px); }
  .tabs { display: flex; gap: 0; border-bottom: 1px solid #ddd;
          margin: 8px 0 14px; }
  .tab { padding: 7px 14px; border: 1px solid transparent; border-bottom: none;
         background: transparent; cursor: pointer; font: inherit;
         font-size: 13px; color: #666; border-radius: 4px 4px 0 0; }
  .tab.active { border-color: #ddd; background: white; color: #1a1a1a;
                font-weight: 600; margin-bottom: -1px; }
  .tab-content { display: none; }
  .tab-content.active { display: block; }
  .controls { display: flex; gap: 12px; align-items: center; margin: 12px 0 0;
              font-size: 13px; flex-wrap: wrap; }
  .controls select, .controls button { font-size: 13px; padding: 3px 8px;
                                       font-family: inherit; }
  .controls button { cursor: pointer; background: #f5f5f5;
                     border: 1px solid #ccc; border-radius: 3px; }
  .controls button:hover { background: #ebebeb; }
  /* The mixture-view dropdown only applies to the mixture chart, which we
     only show on the "By task results" tab. */
  body[data-tab="mixture"] .view-mode-control { display: none; }
  /* The reset button only makes sense on the mix tab. */
  body:not([data-tab="mixture"]) .mix-only { display: none; }
  /* Correlations tab is a read-only viewer. */
  body[data-tab="corr"] .controls { display: none; }
  .section { margin-bottom: 14px; }
  .section-title { font-weight: 600; font-size: 12px; margin: 0 0 4px;
                   color: #333; text-transform: uppercase; letter-spacing: 0.04em; }
  /* auto-fit makes columns wrap when there's no room (1 col on phones, up to
     4 cols on a 1400px desktop). */
  .slider-grid { display: grid;
                 grid-template-columns: repeat(auto-fit, minmax(260px, 1fr));
                 gap: 4px 14px; }
  .slider-row { display: flex; align-items: center; gap: 6px; font-size: 11px;
                line-height: 1.3; min-width: 0; }
  .slider-label { flex: 0 0 clamp(110px, 28%, 170px); white-space: normal;
                  line-height: 1.2; overflow: hidden; color: #333;
                  overflow-wrap: anywhere; }
  .slider-label .quality { font-size: 9px; color: #888; font-style: italic; }
  .slider-row input[type=range] { flex: 1; min-width: 0; accent-color: #2c6aa0;
                                  touch-action: manipulation; }
  .slider-val { flex: 0 0 36px; text-align: right;
                font-variant-numeric: tabular-nums; color: #666; }
  #chart { margin-top: 14px; width: 100%; }
  .model-info { margin: 18px 0 0; padding: 10px 14px; background: #fafafa;
                border: 1px solid #e0e0e0; border-radius: 5px;
                font-size: 12.5px; line-height: 1.5; color: #333; }
  .model-info summary { cursor: pointer; font-weight: 600; padding: 4px 0;
                        user-select: none; color: #1a1a1a; }
  .model-info[open] summary { margin-bottom: 8px; }
  .model-info p { margin: 8px 0; }
  .model-info .formula { overflow-x: auto; margin: 8px 0; }
</style>
<script src="https://cdn.plot.ly/plotly-2.35.0.min.js"></script>
<link rel="stylesheet"
      href="https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/katex.min.css">
<script defer
        src="https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/katex.min.js"></script>
<script defer
        src="https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/contrib/auto-render.min.js"
        onload="renderMathInElement(document.body)"></script>
</head>
<body>
<h1>Swarm Mixture Explorer</h1>
<p class="intro">
Two views on the 241-run swarm.
<b>By mixture</b> &mdash; predictive: per-(phase, domain) weight sliders.
Drag any slider to set the demanded mixture (other sliders in the same phase
auto-adjust proportionally to keep &Sigma;w = 1); we evaluate Calvin's DSP fit
on that mixture and the chart below shows the <i>predicted</i> per-task
percentiles. Predictions near the swarm hull are reliable; far away they're
the form extrapolating.
<b>By task results</b> &mdash; empirical: each task has a slider whose stops
are the 241 runs ordered by outcome on that task. Drag any slider to pick the
run at that rank; all other sliders snap to that run's per-task percentiles,
and the chart below shows its actual mixture.
</p>
<div class="tabs">
  <button class="tab active" data-tab="mixture">By mixture</button>
  <button class="tab" data-tab="results">By task results</button>
  <button class="tab" data-tab="corr">Rank correlations</button>
</div>
<div id="tab-mixture" class="tab-content active">
  <div id="run-info-mix" class="run-info"></div>
  <div id="sliders-mix"></div>
  <details class="model-info">
    <summary>About the model (Calvin's DSP form)</summary>
    <p>Per-task fit of the rank-INT'd swarm-percentile z to the mixture
       \((w_0, w_1)\). For each domain \(d\), effective exposure stacks
       training time across the two phases:</p>
    <div class="formula">
      \[ e_{0,d} = c_{0,d}\, w_{0,d}, \qquad
         e_{1,d} = c_{1,d}\, w_{1,d}, \qquad
         z_d = e_{0,d} + \gamma\, e_{1,d}. \]
    </div>
    <p>Predicted rank-INT z is a saturating per-domain benefit minus an
       overexposure penalty:</p>
    <div class="formula">
      \[ \hat y(w) = -b_0
         + \sum_{d=1}^{D} a_d \bigl(1 - e^{-\rho_d z_d}\bigr)
         - \sum_{d=1}^{D} p_d\,\operatorname{softplus}\!\bigl(\log(1+z_d) - \tau_d\bigr)^{2}. \]
    </div>
    <p>Predicted swarm percentile is \(\Phi(\hat y) \times 100\) (targets
       are rank-INT'd to \(\mathcal N(0,1)\) per task by construction).
       Per task: \(4D + 2 = 158\) parameters (\(D = 39\) domains), fit by
       L-BFGS-B against the 241-run swarm. The \(\pm 1\sigma\) band on each
       bar of the chart below is the fit's residual RMSE for that task, run
       through \(\Phi\) to get an asymmetric percentile interval.</p>
  </details>
</div>
<div id="tab-results" class="tab-content">
  <div id="run-info-task" class="run-info"></div>
  <div id="sliders"></div>
</div>
<div id="tab-corr" class="tab-content">
  <p class="intro" style="margin-top: 12px;">
    <b>Tristan's γ<sub>j</sub></b> (Thrush et al. 2024,
    perplexity-correlations): \(\gamma_j = \sum_{k,l} \mathrm{sign}(y_k - y_l)
    (\mathrm{rank}_j(x_k) - \mathrm{rank}_j(x_l))\), normalised by its
    max so cells lie in [-1, 1]. <b>Blue = more weight on this (phase,
    domain) associates with higher z on this task; red = opposite</b>;
    white ≈ no monotone relationship. Equivalent to Spearman ρ up to
    scaling. Monotonic by construction — can't see the inverted-U
    overexposure dynamics that DSP's \(-p_d\,\mathrm{softplus}(\log(1+z_d)
    - \tau_d)^2\) penalty term models explicitly. So a domain that's
    helpful up to ~1 epoch and harmful past it will look white here even
    when DSP gives it a sharp τ. <b>Third panel</b> aggregates the two
    phases: γ between \(w_{0,d} + w_{1,d}\) and z, useful when you care
    about total presence of a domain rather than which phase it's in.
  </p>
</div>
<div class="controls">
  <label class="view-mode-control">Mixture view:
    <select id="view-mode">
      <option value="delta" selected>&Delta; from token-proportional</option>
      <option value="abs">Absolute mixture weight</option>
    </select>
  </label>
  <button id="download-mix">Download mixture (CSV)</button>
  <button id="download-results">Download task results (CSV)</button>
  <button id="reset-mix">Reset to proportional</button>
  <button id="reset-david">Reset to David's favorite mix</button>
  <button id="reset-pareto" title="DSP-optimised mixture predicted to beat proportional on every task by ≥ +0.05 z">Reset to predicted Pareto-improver</button>
</div>
<div id="chart"></div>
<div id="chart-mix-detail" class="mix-only" style="width:100%;margin-top:18px;"></div>

<script>
const DATA = __DATA__;
const N = DATA.runByRank[0].length;   // 241 swarm runs
const T = DATA.task_names.length;

const SECTIONS = [
  {key: "lm_eval",          prefixes: ["lm_eval/", "mcq_smooth/", "teacher_forced/"]},
  {key: "paloma",           prefixes: ["eval/paloma/"]},
  {key: "uncheatable_eval", prefixes: ["eval/uncheatable_eval/"]},
];

const STRIP = ["eval/paloma/", "eval/uncheatable_eval/",
               "lm_eval/", "mcq_smooth/", "teacher_forced/"];
// Common dataset-name prefixes. Strip so the slider/y-axis label of
// "dolma3_cc/art_and_design_high" becomes "art_and_design_high" instead of
// getting truncated by the label column width.
const DOMAIN_STRIP = ["dolma3_cc/", "dolma3_", "dolmino_"];

function sectionOf(name) {
  for (const s of SECTIONS)
    if (s.prefixes.some(p => name.startsWith(p))) return s.key;
  return "other";
}

function shortLabel(name) {
  for (const p of STRIP) if (name.startsWith(p)) return name.slice(p.length);
  return name;
}

// [stripped name, "High Quality"/"Low Quality" or null]
function domainParts(name) {
  let s = name;
  for (const p of DOMAIN_STRIP)
    if (s.startsWith(p)) { s = s.slice(p.length); break; }
  if (s.endsWith("_high")) return [s.slice(0, -5), "High Quality"];
  if (s.endsWith("_low"))  return [s.slice(0, -4), "Low Quality"];
  return [s, null];
}
function domainLabelHTML(name) {
  const [m, q] = domainParts(name);
  return q ? `${m}<br><span class="quality">${q}</span>` : m;
}
function domainLabelChart(name) {
  // Plotly y-axis: strip the common prefix but keep the _high/_low suffix
  // inline. Multi-line labels look bad in the chart; the suffix alone is
  // enough to distinguish a pair.
  let s = name;
  for (const p of DOMAIN_STRIP)
    if (s.startsWith(p)) { s = s.slice(p.length); break; }
  return s;
}

// Rank r ∈ [0, N-1] → percentile in (0, 100). N+1 spacing so neither endpoint
// hits exactly 0 or 100.
const pctOf = r => ((r + 1) / (N + 1)) * 100;

function buildSliders() {
  const container = document.getElementById("sliders");
  const groups = {};
  DATA.task_names.forEach((name, idx) => {
    const sec = sectionOf(name);
    (groups[sec] = groups[sec] || []).push({name, idx});
  });
  SECTIONS.forEach(s => {
    const items = groups[s.key];
    if (!items) return;
    const sec = document.createElement("div");
    sec.className = "section";
    const title = document.createElement("div");
    title.className = "section-title";
    title.textContent = `${s.key} · ${items.length}`;
    sec.appendChild(title);
    const grid = document.createElement("div");
    grid.className = "slider-grid";
    items.forEach(({name, idx}) => {
      const z0 = predictDSP(DATA.dsp_params[idx], MIX0, MIX1);
      const init = normCdf(z0) * 100;
      const row = document.createElement("div");
      row.className = "slider-row";
      row.innerHTML = `
        <span class="slider-label" title="${name}">${shortLabel(name)}</span>
        <input type="range" min="0" max="100" step="0.5" value="${init}" id="s_${idx}">
        <span class="slider-val" id="v_${idx}">${init.toFixed(1)}</span>`;
      grid.appendChild(row);
      const slider = row.querySelector("input");
      // Drag = small gradient step on MIX toward higher (or lower)
      // predicted percentile on this task. Live, continuous, simplex-safe.
      slider.addEventListener("input", () => {
        const curZ = predictDSP(DATA.dsp_params[idx], MIX0, MIX1);
        const curPct = normCdf(curZ) * 100;
        const newPct = parseFloat(slider.value);
        stepMixTowardTask(idx, newPct - curPct);
        syncTaskSlidersFromMix();
        syncMixSliders(0);
        syncMixSliders(1);
        updateChart();
        scheduleURL();
      });
    });
    sec.appendChild(grid);
    container.appendChild(sec);
  });
}

// "By mixture" tab state: free-floating demanded mixture; not tied to any
// swarm run. Starts at the token-proportional baseline so initial predictions
// reflect "what plain proportional sampling would get."
const MIX0 = DATA.w0_prop.slice();
const MIX1 = DATA.w1_prop.slice();

// Keep the mixture on the simplex: when the user pushes one slider by Δ,
// take it proportionally from every other domain in the same phase (or give
// it back proportionally if Δ is negative). Preserves Σw = 1 exactly.
function redistribute(phase, idx, newVal) {
  const mix = phase === 0 ? MIX0 : MIX1;
  const D = mix.length;
  const oldVal = mix[idx];
  newVal = Math.min(1, Math.max(0, newVal));
  if (newVal >= 1) {
    for (let d = 0; d < D; d++) mix[d] = (d === idx) ? 1 : 0;
    return;
  }
  const otherSum = 1 - oldVal;
  if (otherSum <= 1e-9) {
    // All other weights were ~0; share the freed mass evenly.
    const share = (1 - newVal) / (D - 1);
    for (let d = 0; d < D; d++) mix[d] = (d === idx) ? newVal : share;
    return;
  }
  const scale = (1 - newVal) / otherSum;
  for (let d = 0; d < D; d++) {
    mix[d] = (d === idx) ? newVal : mix[d] * scale;
  }
}

function syncMixSliders(phase) {
  const mix = phase === 0 ? MIX0 : MIX1;
  const D = mix.length;
  for (let d = 0; d < D; d++) {
    const v = mix[d];
    document.getElementById(`m_${phase}_${d}`).value = v;
    document.getElementById(`mv_${phase}_${d}`).textContent = v.toFixed(3);
  }
}

function buildMixSliders() {
  const container = document.getElementById("sliders-mix");
  const D = DATA.domain_names.length;
  const mixArr = [MIX0, MIX1];
  ["phase 0", "phase 1"].forEach((phaseName, p) => {
    const sec = document.createElement("div");
    sec.className = "section";
    const title = document.createElement("div");
    title.className = "section-title";
    title.textContent = `${phaseName} · ${D}`;
    sec.appendChild(title);
    const grid = document.createElement("div");
    grid.className = "slider-grid";
    DATA.domain_names.forEach((d, idx) => {
      // Sliders span the full simplex (0..1) since redistribution can push
      // any single domain anywhere in that range. Step 0.001.
      const init = mixArr[p][idx];
      const row = document.createElement("div");
      row.className = "slider-row";
      row.innerHTML = `
        <span class="slider-label" title="${d}">${domainLabelHTML(d)}</span>
        <input type="range" min="0" max="1" step="0.001"
               value="${init}" id="m_${p}_${idx}">
        <span class="slider-val" id="mv_${p}_${idx}">${init.toFixed(3)}</span>`;
      grid.appendChild(row);
      const slider = row.querySelector("input");
      slider.addEventListener("input", () => {
        redistribute(p, idx, parseFloat(slider.value));
        syncMixSliders(p);
        updateChart();
        scheduleURL();
      });
    });
    sec.appendChild(grid);
    container.appendChild(sec);
  });
}

// JS-side DSP evaluator. Mirrors _predict_dsp() in build_static.py.
function softplus(x) {
  return x > 0 ? x + Math.log1p(Math.exp(-x))
               : Math.log1p(Math.exp(x));
}
function erf(x) {
  const a1 = 0.254829592, a2 = -0.284496736, a3 = 1.421413741;
  const a4 = -1.453152027, a5 = 1.061405429, p = 0.3275911;
  const s = x < 0 ? -1 : 1; x = Math.abs(x);
  const t = 1 / (1 + p * x);
  return s * (1 - ((((a5*t + a4)*t + a3)*t + a2)*t + a1)*t * Math.exp(-x*x));
}
function normCdf(x) { return 0.5 * (1 + erf(x / Math.SQRT2)); }

function predictDSP(p, w0, w1) {
  const D = DATA.c0.length;
  const gamma = softplus(p.log_gamma);
  let signalSum = 0, penaltySum = 0;
  for (let i = 0; i < D; i++) {
    const z = w0[i] * DATA.c0[i] + gamma * w1[i] * DATA.c1[i];
    const rho = softplus(p.log_rho[i]);
    const a = softplus(p.log_a[i]);
    const pen = softplus(p.log_p[i]);
    signalSum += a * (1 - Math.exp(-rho * z));
    const u = Math.log1p(z) - p.tau[i];
    const sp = softplus(u);
    penaltySum += pen * sp * sp;
  }
  return -(p.b0 - signalSum + penaltySum);
}

// One prediction call per task; returns the central estimate + a 1-sigma band
// in percentile units. Sigma is the swarm-residual RMSE from the fit, so
// noisier (less predictable) tasks get visibly wider error bars.
function predictAllWithUncertainty(w0, w1) {
  return DATA.dsp_params.map(p => {
    const z = predictDSP(p, w0, w1);
    return {
      pct:    normCdf(z) * 100,
      pctLow: normCdf(z - p.sigma) * 100,
      pctHi:  normCdf(z + p.sigma) * 100,
    };
  });
}

// Analytical gradient of predictDSP w.r.t. (w0, w1). DSP form:
//   z_d = c0[d]*w0[d] + gamma * c1[d]*w1[d]
//   pred = -(b0 - Σ a_d (1 - e^{-rho_d z_d}) + Σ p_d softplus(log(1+z_d) - tau_d)^2)
// dpred/dz_d = a_d * rho_d * e^{-rho_d z_d}
//              - p_d * 2 * softplus(u_d) * sigmoid(u_d) / (1+z_d)
// with u_d = log(1+z_d) - tau_d.
function gradPredict(p, w0, w1) {
  const D = w0.length;
  const gamma = softplus(p.log_gamma);
  const g0 = new Array(D);
  const g1 = new Array(D);
  for (let d = 0; d < D; d++) {
    const z = w0[d] * DATA.c0[d] + gamma * w1[d] * DATA.c1[d];
    const rho = softplus(p.log_rho[d]);
    const a   = softplus(p.log_a[d]);
    const pen = softplus(p.log_p[d]);
    const u = Math.log1p(z) - p.tau[d];
    const sp = softplus(u);
    const sig = 1.0 / (1.0 + Math.exp(-u));
    const dpred_dz = a * rho * Math.exp(-rho * z)
                   - pen * 2 * sp * sig / (1 + z);
    g0[d] = dpred_dz * DATA.c0[d];
    g1[d] = dpred_dz * gamma * DATA.c1[d];
  }
  return [g0, g1];
}

// Take a gradient step on MIX0/MIX1 toward Δpct on a target task.
// Converts Δpct → Δz via the local norm.cdf slope so the step has the
// right scale, then projects back to the simplex.
function stepMixTowardTask(taskIdx, deltaPct) {
  const p = DATA.dsp_params[taskIdx];
  const curZ = predictDSP(p, MIX0, MIX1);
  // d(pct)/dz at curZ = phi(curZ) * 100; invert to get Δz from Δpct.
  const phi = Math.exp(-curZ * curZ / 2) / Math.sqrt(2 * Math.PI);
  const deltaZ = (deltaPct / 100) / Math.max(phi, 1e-4);
  const [g0, g1] = gradPredict(p, MIX0, MIX1);
  let gNormSq = 0;
  for (let d = 0; d < g0.length; d++) gNormSq += g0[d]*g0[d] + g1[d]*g1[d];
  if (gNormSq < 1e-12) return;
  const alpha = deltaZ / gNormSq;
  for (let d = 0; d < MIX0.length; d++) {
    MIX0[d] = Math.max(0, MIX0[d] + alpha * g0[d]);
    MIX1[d] = Math.max(0, MIX1[d] + alpha * g1[d]);
  }
  // Renormalise each phase back onto the simplex.
  let s0 = 0, s1 = 0;
  for (let d = 0; d < MIX0.length; d++) { s0 += MIX0[d]; s1 += MIX1[d]; }
  if (s0 > 0) for (let d = 0; d < MIX0.length; d++) MIX0[d] /= s0;
  if (s1 > 0) for (let d = 0; d < MIX0.length; d++) MIX1[d] /= s1;
}

// Re-sync task sliders to current MIX predictions.
function syncTaskSlidersFromMix() {
  for (let t = 0; t < T; t++) {
    const z = predictDSP(DATA.dsp_params[t], MIX0, MIX1);
    const pct = normCdf(z) * 100;
    const sld = document.getElementById("s_" + t);
    if (sld) sld.value = pct;
    const lbl = document.getElementById("v_" + t);
    if (lbl) lbl.textContent = pct.toFixed(1);
  }
  const ri = document.getElementById("run-info-task");
  if (ri) updateTaskInfo();
}

// ---------- URL state ----------
// Active tab + shared MIX. Both mix and task tabs steer the same MIX.
//   #tab=mixture&w0=...&w1=...
//   #tab=results&w0=...&w1=...&view=delta
//   #tab=corr
function updateURL() {
  const params = new URLSearchParams();
  const tab = activeTab();
  params.set("tab", tab);
  if (tab !== "corr") {
    params.set("w0", MIX0.map(x => x.toFixed(4)).join(","));
    params.set("w1", MIX1.map(x => x.toFixed(4)).join(","));
    if (tab === "results") {
      params.set("view", document.getElementById("view-mode").value);
    }
  }
  history.replaceState(null, "", "#" + params.toString());
}
let _urlT;
function scheduleURL() {
  clearTimeout(_urlT);
  _urlT = setTimeout(updateURL, 200);
}
function activateTab(tab) {
  document.querySelectorAll(".tab").forEach(b =>
    b.classList.toggle("active", b.dataset.tab === tab));
  document.querySelectorAll(".tab-content").forEach(c =>
    c.classList.toggle("active", c.id === "tab-" + tab));
  document.body.dataset.tab = tab;
}
function loadFromURL() {
  if (!location.hash) return false;
  const params = new URLSearchParams(location.hash.slice(1));
  const tab = params.get("tab");
  if (!tab) return false;
  activateTab(tab);
  const w0s = params.get("w0"), w1s = params.get("w1");
  if (w0s && w1s) {
    const v0 = w0s.split(",").map(parseFloat);
    const v1 = w1s.split(",").map(parseFloat);
    if (v0.length === MIX0.length && v1.length === MIX1.length
        && v0.every(Number.isFinite) && v1.every(Number.isFinite)) {
      for (let d = 0; d < MIX0.length; d++) {
        MIX0[d] = v0[d];
        MIX1[d] = v1[d];
      }
      syncMixSliders(0);
      syncMixSliders(1);
      syncTaskSlidersFromMix();
    }
  }
  const view = params.get("view");
  if (view) document.getElementById("view-mode").value = view;
  updateChart();
  return true;
}

// Tab toggle. Also flip the chart between mixture and task-results.
document.querySelectorAll(".tab").forEach(btn => {
  btn.addEventListener("click", () => {
    activateTab(btn.dataset.tab);
    updateChart();
    updateURL();
  });
});
document.body.dataset.tab = "mixture";

function updateMixInfo(pcts) {
  const mean = pcts.reduce((a,b)=>a+b, 0) / pcts.length;
  document.getElementById("run-info-mix").textContent =
    `Predicted mean percentile ${mean.toFixed(1)} for the current demanded mixture`;
}

function updateTaskInfo() {
  let sumZ = 0;
  for (let t = 0; t < T; t++) sumZ += predictDSP(DATA.dsp_params[t], MIX0, MIX1);
  const meanPct = normCdf(sumZ / T) * 100;
  document.getElementById("run-info-task").textContent =
    `Predicted mean percentile ${meanPct.toFixed(1)} for the current demanded mixture`;
}

function activeTab() {
  return document.querySelector(".tab.active").dataset.tab;
}

function updateChart() {
  const t = activeTab();
  if (t === "mixture") {
    renderPredictedTaskChart();
    renderMixDetail();
  } else if (t === "corr") {
    renderCorrelations();
  } else {
    renderMixChart(MIX0, MIX1);
  }
}

function renderMixDetail() {
  // Demanded-mixture bar chart at the bottom of the mix tab. Always shows
  // delta from token-proportional so it's a direct "how far am I asking
  // the optimiser to push each domain?" view.
  const d0 = MIX0.map((w, i) => w - DATA.w0_prop[i]);
  const d1 = MIX1.map((w, i) => w - DATA.w1_prop[i]);
  const order = DATA.domain_order;
  const names = order.map(i => domainLabelChart(DATA.domain_names[i]));
  const d0o = order.map(i => d0[i]);
  const d1o = order.map(i => d1[i]);
  const color = arr => arr.map(x =>
    x > 0 ? "rgba(44,160,44,0.85)" : "rgba(214,39,40,0.85)");

  const narrow = window.innerWidth < 700;
  const traces = [
    {type:"bar", orientation:"h", x:d0o, y:names, xaxis:"x",  yaxis:"y",
     marker:{color: color(d0o)}, showlegend:false,
     hovertemplate:"%{y}: %{x:+.3f}<extra></extra>"},
    {type:"bar", orientation:"h", x:d1o, y:names, xaxis:"x2", yaxis:"y2",
     marker:{color: color(d1o)}, showlegend:false,
     hovertemplate:"%{y}: %{x:+.3f}<extra></extra>"},
  ];
  const layout = {
    grid: {rows: narrow ? 2 : 1, columns: narrow ? 1 : 2, pattern:"independent"},
    height: narrow ? 1100 : 620,
    margin: {l: narrow ? 140 : 240, r: 20, t: 60, b: 50},
    title: {text: "Demanded mixture — Δ from token-proportional",
            font:{size: 14}, x: 0.5, xanchor: "center"},
    xaxis:  {title:{text:"Δ phase 0 weight"}, range:[-0.25, 0.25],
             zeroline:true, zerolinecolor:"rgba(0,0,0,0.4)"},
    xaxis2: {title:{text:"Δ phase 1 weight"}, range:[-0.25, 0.25],
             zeroline:true, zerolinecolor:"rgba(0,0,0,0.4)"},
    yaxis:  {autorange:"reversed",
             tickfont:{size: narrow ? 8 : 9}, automargin:true},
    yaxis2: {autorange:"reversed",
             tickfont:{size: narrow ? 8 : 9}, automargin:true,
             matches: narrow ? undefined : "y"},
    plot_bgcolor:"white", paper_bgcolor:"white",
  };
  Plotly.react("chart-mix-detail", traces, layout, {responsive: true});
}

function renderCorrelations() {
  // Three γ heatmaps: phase 0, phase 1, and (w0 + w1) summed across phases.
  // Side by side on wide screens, stacked on narrow.
  const order = DATA.domain_order;
  const xLabels = order.map(i => domainLabelChart(DATA.domain_names[i]));
  const yLabels = DATA.task_names.map(shortLabel);
  const reorder = mat => mat.map(row => order.map(i => row[i]));
  const z0 = reorder(DATA.corr0);
  const z1 = reorder(DATA.corr1);
  const zs = reorder(DATA.corr_sum);
  const narrow = window.innerWidth < 700;

  const baseTrace = {
    type: "heatmap", x: xLabels, y: yLabels,
    colorscale: "RdBu", reversescale: true,
    zmid: 0, zmin: -1, zmax: 1,
  };
  const traces = [
    {...baseTrace, z: z0, xaxis: "x",  yaxis: "y",  showscale: false,
     hovertemplate: "task: %{y}<br>domain: %{x}<br>γ = %{z:+.3f}<extra>phase 0</extra>"},
    {...baseTrace, z: z1, xaxis: "x2", yaxis: "y2", showscale: false,
     hovertemplate: "task: %{y}<br>domain: %{x}<br>γ = %{z:+.3f}<extra>phase 1</extra>"},
    {...baseTrace, z: zs, xaxis: "x3", yaxis: "y3", showscale: true,
     colorbar: {title: "γ", thickness: 12, len: 0.7},
     hovertemplate: "task: %{y}<br>domain: %{x}<br>γ = %{z:+.3f}<extra>w0+w1</extra>"},
  ];

  const layout = {
    grid: {rows: narrow ? 3 : 1, columns: narrow ? 1 : 3,
           pattern: "independent"},
    height: narrow ? 1800 : 700,
    margin: {l: narrow ? 140 : 200, r: 60, t: 60, b: 110},
    title: {text: "Tristan's γⱼ — phase 0 / phase 1 / w0+w1 summed",
            font: {size: 14}, x: 0.5, xanchor: "center"},
    xaxis:  {title: {text: "phase 0 domain"}, tickfont: {size: 7},
             tickangle: -60, automargin: true},
    xaxis2: {title: {text: "phase 1 domain"}, tickfont: {size: 7},
             tickangle: -60, automargin: true},
    xaxis3: {title: {text: "w0 + w1 (summed) domain"}, tickfont: {size: 7},
             tickangle: -60, automargin: true},
    yaxis:  {tickfont: {size: 9}, automargin: true,
             autorange: "reversed"},
    yaxis2: {tickfont: {size: 9}, automargin: true,
             autorange: "reversed", matches: narrow ? undefined : "y"},
    yaxis3: {tickfont: {size: 9}, automargin: true,
             autorange: "reversed", matches: narrow ? undefined : "y"},
    plot_bgcolor: "white", paper_bgcolor: "white",
  };
  Plotly.react("chart", traces, layout, {responsive: true});
}

function renderPredictedTaskChart() {
  // Predicted per-task percentile from the DSP fits at the demanded mixture,
  // with ±1σ uncertainty bands from the fit's residual RMSE.
  const preds = predictAllWithUncertainty(MIX0, MIX1);
  const pcts = preds.map(p => p.pct);
  updateMixInfo(pcts);
  const order = [...pcts.keys()].sort((a, b) => pcts[b] - pcts[a]);
  const names = order.map(i => shortLabel(DATA.task_names[i]));
  const vals = order.map(i => pcts[i]);
  const errPlus  = order.map(i => preds[i].pctHi  - preds[i].pct);
  const errMinus = order.map(i => preds[i].pct    - preds[i].pctLow);
  const colors = vals.map(v =>
    v >= 50 ? "rgba(44,160,44,0.85)" : "rgba(214,39,40,0.85)");
  const trace = {
    type: "bar", orientation: "h", x: vals, y: names,
    marker: {color: colors}, showlegend: false,
    error_x: {
      type: "data", symmetric: false,
      array: errPlus, arrayminus: errMinus,
      color: "rgba(0,0,0,0.45)", thickness: 1, width: 4,
    },
    hovertemplate:
      "%{y}: %{x:.1f} pct (±1σ band %{customdata[0]:.1f}–%{customdata[1]:.1f})<extra></extra>",
    customdata: order.map(i => [preds[i].pctLow, preds[i].pctHi]),
  };
  const narrow = window.innerWidth < 700;
  const layout = {
    height: narrow ? 900 : 700,
    margin: {l: narrow ? 140 : 220, r: 20, t: 60, b: 50},
    title: {text: "Predicted task percentiles for current demanded mixture",
            font: {size: 14}, x: 0.5, xanchor: "center"},
    xaxis: {title: {text: "Predicted swarm percentile"}, range: [0, 100],
            zeroline: false},
    yaxis: {autorange: "reversed", tickfont: {size: narrow ? 8 : 9},
            automargin: true},
    plot_bgcolor: "white", paper_bgcolor: "white",
    shapes: [
      {type: "line", xref: "x", yref: "paper", x0: 50, x1: 50, y0: 0, y1: 1,
       line: {color: "rgba(0,0,0,0.4)", width: 1, dash: "dot"}},
    ],
  };
  Plotly.react("chart", [trace], layout, {responsive: true});
}

function renderMixChart(w0r, w1r) {
  const mode = document.getElementById("view-mode").value;
  const x0 = mode === "delta"
    ? w0r.map((w, i) => w - DATA.w0_prop[i]) : w0r;
  const x1 = mode === "delta"
    ? w1r.map((w, i) => w - DATA.w1_prop[i]) : w1r;

  // Stable order across runs (precomputed by natural size).
  const order = DATA.domain_order;
  const names = order.map(i => domainLabelChart(DATA.domain_names[i]));
  const x0o = order.map(i => x0[i]);
  const x1o = order.map(i => x1[i]);

  // Delta view uses signed colors; absolute uses one neutral blue.
  const signColor = arr => arr.map(x => x > 0 ? "rgba(44,160,44,0.85)"
                                              : "rgba(214,39,40,0.85)");
  const c0 = mode === "delta" ? signColor(x0o) : "rgba(31,119,180,0.85)";
  const c1 = mode === "delta" ? signColor(x1o) : "rgba(31,119,180,0.85)";

  const xrange = mode === "delta" ? [-0.25, 0.25] : [0, DATA.wmax];
  const xt0 = mode === "delta" ? "Δ phase 0 weight" : "phase 0 weight";
  const xt1 = mode === "delta" ? "Δ phase 1 weight" : "phase 1 weight";
  const title = mode === "delta"
    ? "Current demanded mixture — Δ from token-proportional"
    : "Current demanded mixture — absolute weights";
  const hover = mode === "delta"
    ? "%{y}: %{x:+.3f}<extra></extra>"
    : "%{y}: %{x:.3f}<extra></extra>";

  // Stack subplots vertically on narrow screens.
  const narrow = window.innerWidth < 700;
  const traces = [
    {type:"bar", orientation:"h", x:x0o, y:names, xaxis:"x",  yaxis:"y",
     marker:{color: c0}, showlegend:false, hovertemplate: hover},
    {type:"bar", orientation:"h", x:x1o, y:names, xaxis:"x2", yaxis:"y2",
     marker:{color: c1}, showlegend:false, hovertemplate: hover},
  ];

  const layout = {
    grid: {rows: narrow ? 2 : 1, columns: narrow ? 1 : 2, pattern: "independent"},
    height: narrow ? 1100 : 620,
    margin: {l: narrow ? 140 : 240, r: 20, t: 60, b: 50},
    title: {text: title, font: {size: 14}, x: 0.5, xanchor: "center"},
    xaxis:  {title:{text: xt0}, range: xrange,
             zeroline:true, zerolinecolor:"rgba(0,0,0,0.4)"},
    xaxis2: {title:{text: xt1}, range: xrange,
             zeroline:true, zerolinecolor:"rgba(0,0,0,0.4)"},
    yaxis:  {autorange:"reversed",
             tickfont:{size: narrow ? 8 : 9}, automargin:true},
    yaxis2: {autorange:"reversed",
             tickfont:{size: narrow ? 8 : 9}, automargin:true,
             matches: narrow ? undefined : "y"},
    plot_bgcolor:"white", paper_bgcolor:"white",
  };

  Plotly.react("chart", traces, layout, {responsive: true});
}

// Re-render when the user toggles the view dropdown.
document.getElementById("view-mode").addEventListener("change",
  () => { updateChart(); updateURL(); });

function downloadCSV(rows, filename) {
  const csv = rows.map(r => r.map(c =>
    (typeof c === "string" && /[",\n]/.test(c)) ? `"${c.replace(/"/g, '""')}"` : c
  ).join(",")).join("\n");
  const url = URL.createObjectURL(new Blob([csv], {type: "text/csv"}));
  const a = document.createElement("a");
  a.href = url; a.download = filename;
  document.body.appendChild(a); a.click(); a.remove();
  setTimeout(() => URL.revokeObjectURL(url), 0);
}

// Download the demanded mixture + the DSP-predicted task results.
document.getElementById("download-mix").addEventListener("click", () => {
  const rows = [["domain", "w0", "w1", "w0_prop", "w1_prop",
                 "delta_w0", "delta_w1"]];
  DATA.domain_names.forEach((d, j) => {
    const wp0 = DATA.w0_prop[j], wp1 = DATA.w1_prop[j];
    rows.push([d, MIX0[j].toFixed(4), MIX1[j].toFixed(4),
               wp0, wp1, (MIX0[j]-wp0).toFixed(4),
               (MIX1[j]-wp1).toFixed(4)]);
  });
  downloadCSV(rows, "demanded_mixture.csv");
});

// On the mix tab, reset edits the demanded mixture directly. On the task
// tab, "proportional" maps to the swarm run nearest token-proportional and
// "David's mix" is just swarm run #157.
function resetMixTo(w0Vec, w1Vec) {
  for (let d = 0; d < MIX0.length; d++) {
    MIX0[d] = w0Vec[d];
    MIX1[d] = w1Vec[d];
  }
  syncMixSliders(0);
  syncMixSliders(1);
  syncTaskSlidersFromMix();
  updateChart();
  updateURL();
}

document.getElementById("reset-mix").addEventListener("click",
  () => resetMixTo(DATA.w0_prop, DATA.w1_prop));
document.getElementById("reset-david").addEventListener("click",
  () => resetMixTo(DATA.w0[157], DATA.w1[157]));
document.getElementById("reset-pareto").addEventListener("click",
  () => resetMixTo(DATA.w0_pareto, DATA.w1_pareto));

document.getElementById("download-results").addEventListener("click", () => {
  const rows = [["task", "predicted_percentile",
                 "pct_1sigma_lower", "pct_1sigma_upper", "sigma_z"]];
  const preds = predictAllWithUncertainty(MIX0, MIX1);
  DATA.task_names.forEach((t, j) => rows.push([
    t, preds[j].pct.toFixed(2),
    preds[j].pctLow.toFixed(2), preds[j].pctHi.toFixed(2),
    DATA.dsp_params[j].sigma.toFixed(4),
  ]));
  downloadCSV(rows, "predicted_task_results.csv");
});

buildSliders();
buildMixSliders();
updateChart();
// If the URL has shareable state, apply it; otherwise show the default view.
if (!loadFromURL()) updateURL();

// Re-render on resize (debounced) so the layout flips between
// stacked and side-by-side when the viewport crosses ~700px.
let _resizeT;
window.addEventListener("resize", () => {
  clearTimeout(_resizeT);
  _resizeT = setTimeout(() => updateChart(), 120);
});
</script>
</body>
</html>
"""

out = HTML.replace("__DATA__", json.dumps(data))

os.makedirs("docs", exist_ok=True)
with open("docs/index.html", "w") as f:
    f.write(out)
print(f"wrote docs/index.html ({len(out)/1024:.1f} KB, "
      f"{len(data['R'])} runs × {len(data['task_names'])} tasks "
      f"× {len(data['domain_names'])} domains)")