1import json, math, random
  2import numpy as np
  3from sklearn.linear_model import Lasso
  4from sklearn.metrics import mean_squared_error
  5
  6SEED = 3053
  7rng = np.random.default_rng(SEED)
  8
  9
 10def solve_lasso(A, b, lam, alpha=None, max_iter=20000):
 11    # sklearn objective is (1/(2n))*||Aw-b||^2 + alpha ||w||_1
 12    if alpha is None: alpha = lam / len(b)
 13    model = Lasso(alpha=alpha, fit_intercept=False, max_iter=max_iter,
 14                  tol=1e-11, selection='cyclic', random_state=SEED)
 15    model.fit(A, b)
 16    w = model.coef_.copy()
 17    resid = A @ w - b
 18    primal = .5 * np.dot(resid, resid) + lam * np.abs(w).sum()
 19    # Dual y convention: y = Aw-b; maximize -.5||y||^2 - b^T y,
 20    # constrained by ||A^T y||inf <= lambda.
 21    yraw = resid.copy()
 22    aty = A.T @ yraw
 23    scale = min(1.0, lam / (np.max(np.abs(aty)) + 1e-30))
 24    y = yraw * scale
 25    dual = -.5*np.dot(y,y) - np.dot(b,y)
 26    gap = max(0.0, primal-dual)
 27    return w, y, primal, dual, gap
 28
 29
 30def core_math_check():
 31    n, d = 80, 24
 32    A = rng.normal(size=(n,d))
 33    # Correlated response with deliberately many zero optimum coordinates.
 34    truth = np.zeros(d); truth[[1,5,9,14]] = [2.0,-1.5,.8,1.2]
 35    b = A @ truth + .15*rng.normal(size=n)
 36    lam = 2.0
 37    w,y,pr,du,gap = solve_lasso(A,b,lam)
 38    aty = A.T @ y
 39    R = math.sqrt(2*gap + 1e-18) # dual objective is 1-strongly concave
 40    safe = np.abs(aty) + R*np.linalg.norm(A,axis=0) < lam
 41    zeros = np.abs(w) < 2e-6
 42    no_false_safe = bool(np.all(~safe | zeros))
 43    # At optimum, KKT complementarity: nonzero coordinates have |A'y|=lambda.
 44    kkt_nonzero = float(np.max(np.abs(np.abs(aty[np.abs(w)>2e-6])-lam))) if np.any(np.abs(w)>2e-6) else 0.
 45    # Equivariance: selecting columns then dualizing is exactly the same
 46    # reduced matrix as row/sample selection then dualizing.
 47    fmask = np.abs(aty) >= np.quantile(np.abs(aty), .5)
 48    smask = np.arange(n) % 2 == 0
 49    reduced_both = A[smask][:,fmask]
 50    dualized_after_feature = reduced_both
 51    feature_after_dual = A[smask][:,fmask]
 52    equiv_err = float(np.max(np.abs(dualized_after_feature-feature_after_dual)))
 53    # Show safety becomes useful as gap shrinks using exact-ish solution.
 54    caught = int(np.sum(safe)); zero_count=int(np.sum(zeros))
 55    return dict(n=n,d=d, primal=pr,dual=du,gap=gap,radius=R,
 56                safe_features=caught, optimizer_zero_features=zero_count,
 57                no_false_safe=no_false_safe,kkt_nonzero_error=kkt_nonzero,
 58                mask_equivariance_max_error=equiv_err)
 59
 60
 61def fit_on_masks(Atr, btr, Aval, bval, fmask, smask, lam):
 62    cols=np.flatnonzero(fmask); rows=np.flatnonzero(smask)
 63    if len(cols)==0 or len(rows)==0: return (float('inf'), float('inf'), 0)
 64    w,y,pr,du,gap=solve_lasso(Atr[np.ix_(rows,cols)], btr[rows], lam)
 65    pred=Aval[:,cols]@w
 66    return mean_squared_error(bval,pred), pr, len(cols)
 67
 68
 69def mini_experiment():
 70    # Feature pruning is tested on a held-out set; sample pruning only changes
 71    # training rows. Dual sample score is the separable conjugate excess score.
 72    ntr,nva,d=240,160,60
 73    A=rng.normal(size=(ntr+nva,d)); truth=np.zeros(d)
 74    truth[[2,7,13,31,44]]=[2.2,-1.7,1.1,.9,-1.3]
 75    b=A@truth + .35*rng.normal(size=ntr+nva)
 76    Atr,Ava=A[:ntr],A[ntr:]; btr,bva=b[:ntr],b[ntr:]
 77    lam=1.0
 78    fullmask=np.ones(d,dtype=bool)
 79    # Obtain dual scores from full training solution.
 80    w,y,pr,du,gap=solve_lasso(Atr,btr,lam); aty=Atr.T@y
 81    R=math.sqrt(2*gap+1e-18)
 82    safe=np.abs(aty)+R*np.linalg.norm(Atr,axis=0)<lam
 83    # Explicit heuristic fallback ranks by dual correlation, as specified.
 84    dual_feat_order=np.argsort(np.abs(aty))
 85    mag_order=np.argsort(np.abs(w))
 86    # squared-loss f_i^* excess over its minimum = .5*(A_i w)^2
 87    sample_score=.5*(Atr@w)**2
 88    dual_sample_order=np.argsort(sample_score) # low contribution first
 89    random_order=rng.permutation(ntr)
 90    results=[]
 91    for frac in [0.25,0.50,0.75]:
 92        kf=max(1,int(d*(1-frac))); ks=max(1,int(ntr*(1-frac)))
 93        # Baseline: magnitude features + random samples.
 94        bm=np.ones(d,bool); bm[mag_order[:d-kf]]=False
 95        bs=np.ones(ntr,bool); bs[random_order[:ntr-ks]]=False
 96        # Idea: permanently safe-remove features, then dual-correlation rank
 97        # remaining features if target exceeds safe removals; dual sample scores.
 98        im=np.ones(d,bool); im[safe]=False
 99        if im.sum()>kf:
100            remaining=np.flatnonzero(im)
101            im[remaining[np.argsort(np.abs(aty[remaining]))[:len(remaining)-kf]]]=False
102        elif im.sum()<kf:
103            # conservative rule cannot retain fewer than requested; leave mask safe
104            pass
105        isamp=np.ones(ntr,bool); isamp[dual_sample_order[:ntr-ks]]=False
106        bval_m,_,bd=fit_on_masks(Atr,btr,Ava,bva,bm,bs,lam)
107        ival_m,_,idd=fit_on_masks(Atr,btr,Ava,bva,im,isamp,lam)
108        results.append(dict(reduction=frac,baseline_val_mse=bval_m,idea_val_mse=ival_m,
109                            baseline_features=int(bm.sum()),idea_features=int(im.sum()),
110                            baseline_samples=int(bs.sum()),idea_samples=int(isamp.sum()),
111                            safe_removed=int(safe.sum())))
112    return dict(train=ntr,validation=nva,features=d,full_gap=gap,
113                full_nonzero=int(np.sum(np.abs(w)>1e-6)),results=results)
114
115
116def main():
117    math_result=core_math_check()
118    exp=mini_experiment()
119    out={'seed':SEED,'math_check':math_result,'mini_experiment':exp}
120    print(json.dumps(out,indent=2))
121
122if __name__=='__main__': main()