app.py

"""
Interactive demo: the least-squares estimator for *shuffled* linear regression is
inconsistent in every dimension.

Companion to the note "Inconsistency and Bias of the Least-Squares Estimator in
Higher-Dimensional Shuffled Linear Regression" (a generalization of Theorem 1 of
Abid, Poon & Zou, 2017, "Linear Regression with Shuffled Labels").

Model:  y = pi_0 (x w0) + e,   x_i ~ N(mu_X, Sigma_X),   e ~ N(0, sigma_E^2),
        pi_0 an unknown permutation.

Theory (see paper):  the shuffled-LS loss converges to
    ell(w) = (w^T mu_X - w0^T mu_X)^2 + (sqrt(w^T Sigma_X w) - sqrt(w0^T Sigma_X w0 + sigma_E^2))^2,
so w_hat converges to the moment-matching set
    { w : w^T mu_X = w0^T mu_X,  w^T Sigma_X w = w0^T Sigma_X w0 + sigma_E^2 }.
=> inconsistent for sigma_E != 0, with the Sigma_X-norm inflated by exactly sigma_E^2,
   independent of the dimension d.  For d=1 this is the paper's closed form, Eq. (4).
"""

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import gradio as gr

import shuffled_ls as sl

PAPER_URL = "https://huggingface.co/spaces/abidlabs/shuffled-linear-regression"

# Enable inline ($...$) as well as display ($$...$$) math in all Markdown.
LATEX = [
    {"left": "$$", "right": "$$", "display": True},
    {"left": "$", "right": "$", "display": False},
]

ACCENT, TRUTH, PRED = "#4f46e5", "#9ca3af", "#dc2626"
plt.rcParams.update({"figure.facecolor": "white", "axes.grid": True,
                     "grid.color": "#e5e7eb", "grid.linewidth": 0.6,
                     "axes.edgecolor": "#d1d5db", "font.size": 11})


def _build_covariance(d, kind, rng):
    if kind == "Identity":
        return np.eye(d)
    if kind == "Diagonal (random variances)":
        return np.diag(rng.uniform(0.5, 2.5, size=d))
    return sl.random_spd(d, rng)            # "General (correlated)"


def run_experiment(d, sigma_E, mu_scale, cov_kind, max_log_n, n_trials, seed):
    d = int(d); n_trials = int(n_trials); seed = int(seed)
    rng = np.random.default_rng(seed)

    # Ground truth and design parameters.
    w0 = rng.standard_normal(d)
    mu_X = mu_scale * rng.standard_normal(d)
    Sigma_X = _build_covariance(d, cov_kind, rng)

    mean_target, var_target = sl.target_invariants(w0, mu_X, Sigma_X, sigma_E)
    w0_energy = var_target - sigma_E ** 2          # w0^T Sigma_X w0

    # Sample sizes from 10^2 up to 10^max_log_n.
    ns = np.unique(np.round(np.logspace(2, max_log_n, 7)).astype(int))

    infl_mean, infl_lo, infl_hi = [], [], []
    mean_err, var_err = [], []
    d1_emp = []                                    # only used when d == 1
    for n in ns:
        infl_t, me_t, ve_t, w1_t = [], [], [], []
        for _ in range(n_trials):
            x, y = sl.make_data(n, w0, mu_X, Sigma_X, sigma_E, rng)
            # Honest estimator: multi-start, NO knowledge of w0.
            w_hat, _ = sl.fit_ls(x, y, n_starts=6, rng=rng)
            mS, vS = sl.invariants(w_hat, mu_X, Sigma_X)
            infl_t.append(vS - w0_energy)
            me_t.append(abs(mS - mean_target))
            ve_t.append(abs(vS - var_target))
            if d == 1:
                w1_t.append(float(w_hat[0]))
        infl_mean.append(np.mean(infl_t))
        infl_lo.append(np.percentile(infl_t, 10)); infl_hi.append(np.percentile(infl_t, 90))
        mean_err.append(np.mean(me_t)); var_err.append(np.mean(ve_t))
        if d == 1:
            d1_emp.append(np.mean(w1_t))
    ns = np.array(ns)
    infl_mean = np.array(infl_mean)

    # ---------- Figure 1: inconsistency (norm amplification does NOT vanish) ----------
    fig1, ax = plt.subplots(figsize=(6.6, 4.4))
    ax.fill_between(ns, infl_lo, infl_hi, color=ACCENT, alpha=0.15,
                    label="10–90% over trials")
    ax.semilogx(ns, infl_mean, "o-", color=ACCENT, lw=2, label="empirical")
    ax.axhline(sigma_E ** 2, ls="--", color=PRED, lw=2,
               label=fr"theory  $\sigma_E^2={sigma_E**2:.3g}$")
    ax.axhline(0.0, ls=":", color=TRUTH, lw=1.5, label="consistent would be 0")
    ax.set_xlabel("sample size  n")
    ax.set_ylabel(r"$\|\hat w\|^2_{\Sigma_X}-\|w_0\|^2_{\Sigma_X}$")
    ax.set_title(f"Inconsistency in d={d}: signal energy is inflated by $\\sigma_E^2$")
    ax.legend(loc="best", fontsize=9)
    fig1.tight_layout()

    # ---------- Figure 2: convergence to the moment-matching set ----------
    fig2, ax2 = plt.subplots(figsize=(6.6, 4.4))
    ax2.loglog(ns, np.maximum(mean_err, 1e-12), "o-", color="#0891b2", lw=2,
               label=r"$|\hat w^\top\mu_X-w_0^\top\mu_X|$")
    ax2.loglog(ns, np.maximum(var_err, 1e-12), "s-", color="#ea580c", lw=2,
               label=r"$|\hat w^\top\Sigma_X\hat w-(w_0^\top\Sigma_X w_0+\sigma_E^2)|$")
    ax2.set_xlabel("sample size  n")
    ax2.set_ylabel("absolute error")
    ax2.set_title("Estimator converges to the moment-matching set")
    ax2.legend(loc="best", fontsize=9)
    fig2.tight_layout()

    # ---------- Summary ----------
    n_big = ns[-1]
    emp_energy = infl_mean[-1] + w0_energy
    lines = [
        f"### Results (d = {d}, {cov_kind} covariance, {n_trials} trials)",
        "",
        f"- **True signal energy** $w_0^\\top\\Sigma_X w_0$ = `{w0_energy:.4f}`",
        f"- **Predicted limit** of $\\|\\hat w\\|^2_{{\\Sigma_X}}$ = `{var_target:.4f}` "
        f"(= signal energy **+ $\\sigma_E^2$ = {sigma_E**2:.4f}**)",
        f"- **Empirical** $\\|\\hat w\\|^2_{{\\Sigma_X}}$ at n={n_big:,} = `{emp_energy:.4f}`",
        "",
        f"The amplification settles near **{infl_mean[-1]:.4f}** (theory: "
        f"**{sigma_E**2:.4f}**) instead of decaying to 0, the estimator is "
        f"**inconsistent**, and the gap is the noise variance, independent of $d$.",
    ]
    if d == 1:
        pred1 = sl.theorem1_limit_1d(w0[0], mu_X[0], np.sqrt(Sigma_X[0, 0]), sigma_E)
        lines += [
            "",
            "#### One-dimensional check vs. the paper's closed form (Eq. 4)",
            f"- true weight $w_0$ = `{w0[0]:.4f}`",
            f"- **Eq. (4) limit** = `{pred1:.4f}`",
            f"- **empirical** $\\hat w$ at n={n_big:,} = `{d1_emp[-1]:.4f}` "
            f"(amplification ≈ {d1_emp[-1]/w0[0]:.3f}×)",
        ]
    summary = "\n".join(lines)
    return fig1, fig2, summary


INTRO = f"""
# 🔀 Shuffled Linear Regression, the LS estimator is inconsistent in every dimension

In **shuffled linear regression** you observe features $x$ and labels $y$, but an
unknown permutation $\\pi_0$ scrambles which label goes with which row:
$$y = \\pi_0\\,(x\\,w_0) + e,\\qquad x_i\\sim\\mathcal N(\\mu_X,\\Sigma_X),\\quad e\\sim\\mathcal N(0,\\sigma_E^2).$$

The natural least-squares estimator $\\hat w_{{\\mathrm{{LS}}}}=\\arg\\min_w\\min_\\pi\\|\\pi x w-y\\|^2$
is **not consistent**. As $n\\to\\infty$ it converges to the *moment-matching set*
$$\\{{\\,w:\\ w^\\top\\mu_X=w_0^\\top\\mu_X,\\quad w^\\top\\Sigma_X w=w_0^\\top\\Sigma_X w_0+\\sigma_E^2\\,\\}},$$
so the feature-covariance norm of the estimate is **inflated by exactly $\\sigma_E^2$,
independent of the dimension $d$**. (For $d=1$ this is the closed form of Abid–Poon–Zou 2017, Theorem 1, Eq. 4.)

Generate fresh data below and watch the amplification **fail to vanish** as $n$ grows.
"""


with gr.Blocks(title="Shuffled Linear Regression") as demo:
    gr.Markdown(INTRO, latex_delimiters=LATEX)
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("#### Experiment settings")
            d = gr.Slider(1, 8, value=3, step=1, label="Dimension  d")
            sigma_E = gr.Slider(0.0, 3.0, value=1.0, step=0.1, label="Noise std  σ_E")
            mu_scale = gr.Slider(0.0, 3.0, value=1.0, step=0.1,
                                 label="Feature-mean scale  ‖μ_X‖")
            cov_kind = gr.Dropdown(
                ["Identity", "Diagonal (random variances)", "General (correlated)"],
                value="General (correlated)", label="Feature covariance  Σ_X")
            max_log_n = gr.Slider(3.0, 5.0, value=4.0, step=0.5,
                                  label="Max sample size  (10^x)")
            n_trials = gr.Slider(1, 10, value=4, step=1, label="Trials per n")
            seed = gr.Number(value=0, precision=0, label="Random seed")
            run = gr.Button("Run experiment", variant="primary")
        with gr.Column(scale=2):
            with gr.Row():
                plot1 = gr.Plot(label="Inconsistency: norm amplification")
                plot2 = gr.Plot(label="Convergence to the moment-matching set")
            summary = gr.Markdown(latex_delimiters=LATEX)

    gr.Markdown(
        "The left panel is the key result: a *consistent* estimator would drive the "
        "curve to 0 (grey dotted), but the shuffled-LS estimate settles on "
        "$\\sigma_E^2$ (red dashed) in every dimension. The right panel shows the "
        "estimate instead converges to the two-moment-matching set the theory predicts.",
        latex_delimiters=LATEX,
    )

    inputs = [d, sigma_E, mu_scale, cov_kind, max_log_n, n_trials, seed]
    run.click(run_experiment, inputs=inputs, outputs=[plot1, plot2, summary])
    demo.load(run_experiment, inputs=inputs, outputs=[plot1, plot2, summary])


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="indigo"))