1import json, math, time
  2import numpy as np
  3from sklearn.neural_network import MLPRegressor
  4
  5# Adversarially calibrated residualization MVP.
  6# The implementation follows the displayed constraint literally.
  7
  8def edit_weights(t, F, w0=None, tau=0.0, kappa=0.0):
  9    """Solve the stated local edit problem for the important w0=1 case.
 10    Since w=1 has objective zero and exactly zero (w-1) moments, it is the
 11    unique minimizer whenever w0=1 (the objective is strictly convex).
 12    """
 13    n = len(t)
 14    if w0 is None: w0 = np.ones(n)
 15    w0 = np.asarray(w0)
 16    if np.allclose(w0, 1.0):
 17        w = np.ones(n)
 18    else:
 19        # Small projected-gradient fallback for non-unit initial weights.
 20        w = np.maximum(w0.copy(), 0.0)
 21        A = t[:, None] * F
 22        lr = 0.2 / (np.linalg.norm(A, 2)**2 / n + 1e-8)
 23        for _ in range(2000):
 24            mom = ((w-1)[:, None] * A).mean(0)
 25            viol = np.maximum(np.abs(mom)-tau, 0)
 26            grad = (w-w0)/n
 27            if np.any(viol):
 28                grad += ((A @ (np.sign(mom)*viol)) / n)
 29            w = np.maximum(w-lr*grad, 0)
 30            w *= len(w)/w.sum()
 31    energy = np.mean(w*t*t)
 32    # The stated rejection rule: retain a feasible-energy solution if possible.
 33    if energy < kappa:
 34        return np.ones(n) if np.mean(t*t) >= kappa else w
 35    return w
 36
 37def moments(w,t,F):
 38    return np.mean((w-1)[:,None]*t[:,None]*F, axis=0)
 39
 40def toy_verification(seed=123):
 41    rng=np.random.default_rng(seed); n=600
 42    x=rng.normal(size=(n,8)); t=0.5*x[:,0]+rng.normal(size=n)
 43    F=np.tanh(x[:,:4])
 44    out={"tau_sweep":[], "critic_scale_sweep":[], "dimension_sweep":[]}
 45    for tau in [0,1e-4,1e-2,0.1,1.0]:
 46        w=edit_weights(t,F,tau=tau,kappa=.01)
 47        out["tau_sweep"].append({"tau":tau,"edit_l2":float(np.linalg.norm(w-1)),"max_moment":float(np.max(np.abs(moments(w,t,F))))})
 48    for scale in [0.1,1,10]:
 49        Fs=scale*F
 50        w=edit_weights(t,Fs,tau=.01,kappa=.01)
 51        out["critic_scale_sweep"].append({"scale":scale,"edit_l2":float(np.linalg.norm(w-1)),"max_moment":float(np.max(np.abs(moments(w,t,Fs))))})
 52    for d in [1,2,4,8]:
 53        Fs=np.tanh(x[:,:d])
 54        w=edit_weights(t,Fs,tau=.01,kappa=.01)
 55        out["dimension_sweep"].append({"dimension":d,"edit_l2":float(np.linalg.norm(w-1)),"max_moment":float(np.max(np.abs(moments(w,t,Fs))))})
 56    return out
 57
 58def dml_once(seed, n=500, reverse=False):
 59    rng=np.random.default_rng(seed); p=20; beta=1.0
 60    X=rng.normal(size=(n,p));
 61    # Imbalanced nuisance difficulty: high-frequency outcome versus smooth treatment,
 62    # then the reverse. A small MLP intentionally underfits the high-frequency part.
 63    smooth_pi=0.8*np.sin(X[:,0])+0.3*X[:,1]
 64    hard_mu=1.5*np.sin(5*X[:,0])+0.7*np.cos(4*X[:,1])+0.3*X[:,2]
 65    if not reverse: mu, pi=hard_mu, smooth_pi
 66    else: mu, pi=smooth_pi, hard_mu/1.5
 67    T=pi+rng.normal(size=n)
 68    Y=mu+beta*T+0.8*rng.normal(size=n)
 69    idx=np.arange(n); rng.shuffle(idx); folds=np.array_split(idx,2)
 70    ry=np.zeros(n); rt=np.zeros(n)
 71    for te in folds:
 72        tr=np.setdiff1d(idx,te,assume_unique=False)
 73        my=MLPRegressor(hidden_layer_sizes=(32,16),early_stopping=False,max_iter=100,
 74                        random_state=seed, solver='adam', learning_rate_init=.003)
 75        mt=MLPRegressor(hidden_layer_sizes=(32,16),early_stopping=False,max_iter=100,
 76                        random_state=seed+17, solver='adam', learning_rate_init=.003)
 77        my.fit(X[tr],Y[tr]); mt.fit(X[tr],T[tr])
 78        ry[te]=Y[te]-my.predict(X[te]); rt[te]=T[te]-mt.predict(X[te])
 79    beta_hat=float(np.sum(rt*ry)/np.sum(rt*rt))
 80    F=np.tanh(X[:,:4]); w=edit_weights(rt,F,tau=.02,kappa=.1*np.mean(rt*rt))
 81    cal=float(np.sum(w*rt*ry)/np.sum(w*rt*rt))
 82    ess=float(w.sum()**2/np.sum(w*w)); viol=float(np.max(np.abs(moments(w,rt,F))))
 83    return beta_hat,cal,ess,viol
 84
 85def mini_experiment():
 86    rows=[]
 87    for reverse in [False,True]:
 88        vals=[dml_once(100+i,reverse=reverse) for i in range(8)]
 89        a=np.array(vals)
 90        rows.append({"regime":"hard_mu" if not reverse else "hard_pi",
 91          "baseline_abs_bias":float(np.mean(np.abs(a[:,0]-1))),
 92          "idea_abs_bias":float(np.mean(np.abs(a[:,1]-1))),
 93          "baseline_rmse":float(np.sqrt(np.mean((a[:,0]-1)**2))),
 94          "idea_rmse":float(np.sqrt(np.mean((a[:,1]-1)**2))),
 95          "mean_ess":float(np.mean(a[:,2])),"max_violation":float(np.max(a[:,3])),
 96          "paired_max_difference":float(np.max(np.abs(a[:,0]-a[:,1])))})
 97    return rows
 98
 99if __name__=='__main__':
100    t=time.time(); result={"toy":toy_verification(),"mini":mini_experiment(),"runtime_sec":time.time()-t}
101    print(json.dumps(result,indent=2))