import numpy as np, matplotlib.pyplot as plt
rng = np.random.default_rng(0)
P, n, s = 200, 400, 20 # P features; true support = s
beta = np.zeros(P); supp = rng.choice(P, s, replace=False); beta[supp] = rng.normal(0,1,s)
def data(m): X = rng.normal(0,1,(m,P)); return X, X @ beta + 0.3*rng.normal(size=m)
Xtr, ytr = data(n); Xte, yte = data(4000)
w = np.linalg.lstsq(Xtr, ytr, rcond=None)[0] # dense fit
order = np.argsort(np.abs(w))[::-1] # rank by magnitude
ks = [5,10,20,40,80,120,160,200]; e_tic, e_rnd = [], []
for k in ks:
t = order[:k]; e_tic.append(np.mean((Xte[:,t] @ np.linalg.lstsq(Xtr[:,t],ytr,rcond=None)[0] - yte)**2))
r = rng.choice(P,k,replace=False); e_rnd.append(np.mean((Xte[:,r] @ np.linalg.lstsq(Xtr[:,r],ytr,rcond=None)[0] - yte)**2))
plt.figure(figsize=(5.6,2.9))
plt.semilogy([k/P*100 for k in ks], e_tic, "o-", color="#5d2c80", label="winning ticket (magnitude)")
plt.semilogy([k/P*100 for k in ks], e_rnd, "s--", color="crimson", label="random subnet")
plt.axhline(np.mean((Xte@w-yte)**2), ls=":", c="gray", label="dense")
plt.xlabel("% of weights kept"); plt.ylabel("test error (log)"); plt.legend(fontsize=7)
plt.title("A sparse ticket matches the dense net; a random one doesn't"); plt.tight_layout(); plt.show()