1import json, math, random
  2import numpy as np
  3import torch
  4
  5SEED=2543
  6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_num_threads(4)
  8DT=0.1; K=30; N=3000
  9
 10def features(x):
 11    # Smooth characteristic features plus moments; sufficient to expose law mismatch.
 12    fs=[x, x*x, x**3]
 13    for w in [0.35,0.7,1.2,2.0,3.0]:
 14        fs += [torch.cos(w*x), torch.sin(w*x)]
 15    return torch.stack(fs,-1)
 16
 17def np_features(x):
 18    x=torch.as_tensor(x,dtype=torch.float32)
 19    return features(x)
 20
 21def mixture_stats(paths, hazards):
 22    # States x_0,...,x_K: stop at k<K; all residual mass emits x_K.
 23    q=1-torch.exp(-torch.nn.functional.softplus(hazards)*DT)
 24    surv=torch.cumprod(torch.cat([torch.ones(1),1-q[:-1]]),0)
 25    w=torch.zeros_like(q)
 26    w[:-1]=surv[:-1]*q[:-1]
 27    w[-1]=surv[-1]
 28    z=(features(paths)*w[None,:,None]).sum(1)
 29    expected_steps=(surv[:-1]*DT).sum()
 30    return z.mean(0), w, expected_steps
 31
 32def mmd_feature(a,b):
 33    return ((a-b)**2).mean()
 34
 35def exact_checks():
 36    rows=[]
 37    for h in [0.2,0.5,1.0,2.0]:
 38        q=1-math.exp(-h*DT)
 39        s=np.array([(1-q)**k for k in range(K+1)])
 40        # fit log survival slope, excluding endpoint numerical issues
 41        slope=np.polyfit(np.arange(K+1)*DT,np.log(s),1)[0]
 42        # discrete expected number of active intervals, matching sum survival
 43        expected=float(s[:-1].sum()*DT)
 44        pred=(1-math.exp(-h*K*DT))/h
 45        rows.append({'h':h,'predicted_log_slope':-h,'observed_log_slope':float(slope),
 46                     'slope_abs_err':float(abs(slope+h)),
 47                     'predicted_expected_time':pred,'observed_expected_time':expected,
 48                     'time_abs_err':abs(expected-pred)})
 49    # survival-mixture mass for arbitrary positive hazards
 50    hz=np.array([.1,.4,1.1,2.3,.7]*7)[:K+1]
 51    q=1-np.exp(-hz*DT); surv=np.cumprod(np.r_[1.,1-q[:-1]])
 52    mass=float(np.sum(surv*q)+surv[-1]*(1-q[-1]))
 53    # Prediction 3: discrete expected stopping time converges to the
 54    # continuous integral as dt decreases (first-order Riemann error).
 55    conv=[]
 56    h=.8; horizon=3.0
 57    for dt in [.2,.1,.05,.025]:
 58        n=round(horizon/dt); r=math.exp(-h*dt)
 59        discrete=dt*(1-r**n)/(1-r)
 60        continuous=(1-math.exp(-h*horizon))/h
 61        conv.append({'dt':dt,'discrete_expected_time':discrete,
 62                     'continuous_prediction':continuous,
 63                     'abs_error':abs(discrete-continuous)})
 64    return rows, mass, conv
 65
 66def train_policy(paths, target_feat, steps=350):
 67    p=torch.nn.Parameter(torch.full((K+1,),-1.0))
 68    opt=torch.optim.Adam([p],lr=.08)
 69    history=[]
 70    for it in range(steps):
 71        opt.zero_grad()
 72        pred,w,etime=mixture_stats(paths,p)
 73        loss=mmd_feature(pred,target_feat)+1e-4*etime
 74        loss.backward(); opt.step()
 75        if it in [0,49,149,349]: history.append(float(loss.detach()))
 76    with torch.no_grad():
 77        pred,w,etime=mixture_stats(paths,p)
 78        q=1-torch.exp(-torch.nn.functional.softplus(p)*DT)
 79        # empirical stopped samples use weighted mixture CDF proxy; law metric is feature MSE
 80        return float(mmd_feature(pred,target_feat)),float(etime),q.numpy(),history
 81
 82def run():
 83    # Common Brownian paths; X_0=0 and X_k=sqrt(t_k) Gaussian marginals.
 84    inc=np.random.randn(N,K).astype('float32')*math.sqrt(DT)
 85    paths=np.concatenate([np.zeros((N,1),dtype='float32'),np.cumsum(inc,axis=1)],axis=1)
 86    pt=torch.from_numpy(paths)
 87    # Target law is a known constant-hazard stopped Brownian law (independent samples).
 88    htrue=.8
 89    u=np.random.rand(N); tau=-np.log(u)/htrue
 90    kt=np.minimum((tau/DT).astype(int),K)
 91    target=np.zeros(N,dtype='float32')
 92    for k in range(K):
 93        mask=kt==k
 94        target[mask]=paths[mask,k] # couple only for reproducibility; marginal is Brownian
 95    mask=kt==K; target[mask]=paths[mask,K]
 96    target_feat=np_features(target).mean(0)
 97    # Standard fixed-budget baseline: emit X_K.
 98    baseline=float(mmd_feature(np_features(paths[:,:][:,-1]).mean(0),target_feat))
 99    idea,etime,q,history=train_policy(pt,target_feat)
100    # Also report learned hazard averaged over its (time-only) head and fit survival.
101    surv=np.cumprod(np.r_[1.,1-q[:-1]])
102    slope=float(np.polyfit(np.arange(K+1)*DT,np.log(np.maximum(surv,1e-12)),1)[0])
103    return {'checks':{'constant_hazard_sweep':exact_checks()[0],
104                      'mixture_mass_error':abs(exact_checks()[1]-1.0),
105                      'dt_convergence_sweep':exact_checks()[2]},
106            'mini_experiment':{'target_true_hazard':htrue,'baseline_feature_mse':baseline,
107                               'idea_feature_mse':idea,'idea_expected_compute_time':etime,
108                               'fixed_budget_compute_time':K*DT,
109                               'learned_survival_log_slope':slope,
110                               'training_loss_checkpoints':history,
111                               'learned_q_first_last':[float(q[0]),float(q[-1])]}}
112
113if __name__=='__main__':
114    try:
115        # CUDA is allowed, but this tiny verification is deliberately CPU-safe.
116        out=run()
117    except Exception as e:
118        out={'error':repr(e)}
119    print(json.dumps(out,indent=2))