simjeb-bracket-fea-mass-prediction-real

cd /app && ls -la && echo "---- CONFIG ----" && cat task_config.json && echo "---- SCHEMA ----" && cat feature_schema.json

cd /app && echo "==== candidate_model.py ====" && cat candidate_model.py && echo "==== dataset.py ====" && cat dataset.py && echo "==== train.py ====" && cat train.py

cd /app && ls -la data && echo "---" && for s in train val public_test; do echo "== $s =="; ls -la data/$s; done

cd /app && head -n 3 data/train/manifest.jsonl && echo "..." && wc -l data/train/manifest.jsonl data/val/manifest.jsonl data/public_test/manifest.jsonl

cd /app && python -c "import torch, sys; print('torch', torch.__version__); print('cuda', torch.cuda.is_available()); print('py', sys.version)" && nproc && free -h 2>/dev/null | head -2

cd /app && nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv 2>/dev/null || echo "no nvidia-smi"

cd /app && python -c "
import torch
for s in ['train','val','public_test']:
    d = torch.load(f'data/{s}/samples.pt', weights_only=True)
    f = d['features']; l = d['labels']
    print(f'== {s} ==')
    print('  features', f.shape, f.dtype, 'labels', l.shape, l.dtype)
    ff = f.float(); ll = l.float()
    print('  coord min/max', ff.min().item(), ff.max().item())
    print('  per-axis min', ff.amin(dim=(0,1)).tolist())
    print('  per-axis max', ff.amax(dim=(0,1)).tolist())
    print('  label min/max/mean/std', ll.min().item(), ll.max().item(), ll.mean().item(), ll.std().item())
"

cd /app && python -c "
import json
def ids(s):
    out=[]
    with open(f'data/{s}/manifest.jsonl') as fh:
        for line in fh:
            out.append(json.loads(line)['source_bracket_id'])
    return out
tr=ids('train'); va=ids('val'); te=ids('public_test')
print('train n=',len(tr),'unique=',len(set(tr)))
print('val   n=',len(va),'unique=',len(set(va)))
print('test  n=',len(te),'unique=',len(set(te)))
print('train∩val', sorted(set(tr)&set(va)))
print('train∩test', sorted(set(tr)&set(te)))
print('val∩test', sorted(set(va)&set(te)))
print('all unique source ids:', len(set(tr)|set(va)|set(te)))
print('train ids sorted:', sorted(set(tr)))
"

cd /app && python -c "
import torch
for s in ['train','val','public_test']:
    d = torch.load(f'data/{s}/samples.pt', weights_only=True)
    f = d['features'].float()
    # per-sample bbox extents
    mn = f.amin(dim=1); mx = f.amax(dim=1)  # [N,3]
    ext = mx - mn  # [N,3]
    ctr = (mx+mn)/2
    print(f'== {s} ==  N={f.shape[0]}')
    print('  extent mean', ext.mean(0).tolist())
    print('  extent min ', ext.amin(0).tolist())
    print('  extent max ', ext.amax(0).tolist())
    print('  max-extent over axes per sample: min/max =', ext.max(1).values.min().item(), ext.max(1).values.max().item())
    print('  center mean', ctr.mean(0).tolist(), 'center absmax', ctr.abs().amax(0).tolist())
"

cd /app && python -c "
import torch, numpy as np
d = torch.load('data/train/samples.pt', weights_only=True)
f = d['features'].float(); y = d['labels'].float().squeeze(1)
N = f.shape[0]
mn = f.amin(dim=1); mx = f.amax(dim=1); ext = mx-mn
feats = {}
feats['ext_x']=ext[:,0]; feats['ext_z']=ext[:,2]
feats['bbox_vol']=ext[:,0]*ext[:,1]*ext[:,2]
feats['xz_area']=ext[:,0]*ext[:,2]
# radial spread
r = f.norm(dim=2)  # [N,2048]
feats['mean_r']=r.mean(1); feats['std_r']=r.std(1)
# per-axis std of points
feats['std_x']=f[:,:,0].std(1); feats['std_y']=f[:,:,1].std(1); feats['std_z']=f[:,:,2].std(1)
# mean abs coord
feats['mabs_x']=f[:,:,0].abs().mean(1); feats['mabs_z']=f[:,:,2].abs().mean(1)
# proxy 'compactness': average nearest-neighbor distance (subsample for speed)
for k,v in feats.items():
    v=v.numpy(); yy=y.numpy()
    c=np.corrcoef(v,yy)[0,1]
    print(f'{k:10s} corr={c:+.3f}')
print('label mean/std', y.mean().item(), y.std().item())
"

cd /app && python -c "
import torch
d = torch.load('data/train/samples.pt', weights_only=True)
f = d['features'].float()  # [N,2048,3]
N = f.shape[0]
def chamfer(a,b):
    # a,b: [P,3]  -> mean nearest neighbor both directions, subsample 512 for speed
    a=a[:512]; b=b[:512]
    d2 = (a[:,None,:]-b[None,:,:]).pow(2).sum(-1)  # [P,P]
    return d2.min(1).values.mean().sqrt().item()/2 + d2.min(0).values.mean().sqrt().item()/2
import random
idxs=list(range(N))
for ax,name in [(0,'X'),(1,'Y'),(2,'Z')]:
    vals=[]
    for i in idxs[:30]:
        c=f[i]
        m=c.clone(); m[:,ax]=-m[:,ax]
        vals.append(chamfer(c,m))
    t=torch.tensor(vals)
    print(f'reflect {name}: chamfer mean={t.mean():.4f} median={t.median():.4f} min={t.min():.4f} max={t.max():.4f}')
# compare to chamfer between DIFFERENT brackets (scale reference)
import itertools
vals=[]
for i in range(0,20,2):
    vals.append(chamfer(f[i],f[i+1]))
print('between different brackets chamfer mean=',torch.tensor(vals).mean().item())
"

mkdir -p /app/work && cat > /app/work/lib.py << 'PY'
from __future__ import annotations
import os, sys, json, math, copy
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F

APP = Path('/app')
sys.path.insert(0, str(APP))
from candidate_model import build_model

def load_cfg():
    with open(APP/'task_config.json') as fh:
        return json.load(fh)

def load_split(s):
    d = torch.load(APP/'data'/s/'samples.pt', weights_only=True)
    return d['features'].float(), d['labels'].float()

def compute_norm(feats, labels, isotropic=False):
    fm = feats.mean(dim=(0,1))
    if isotropic:
        fs = feats.std().repeat(3)
    else:
        fs = feats.std(dim=(0,1)).clamp_min(1e-6)
    lm = labels.mean(dim=0)
    ls = labels.std(dim=0).clamp_min(1e-6)
    return fm, fs, lm, ls

@torch.no_grad()
def predict(model, X, fm, fs, lm, ls, device, bs=256):
    # exact inference transform: ((x-fm)/fs) -> model -> *ls + lm
    model.eval()
    out=[]
    for i in range(0, X.shape[0], bs):
        xb = X[i:i+bs].to(device)
        xb = (xb - fm.to(device))/fs.to(device)
        p = model(xb)
        out.append(p.cpu())
    p = torch.cat(out,0)
    return p*ls + lm

def metrics(pred, true):
    pred=pred.reshape(-1).double(); true=true.reshape(-1).double()
    err = pred-true
    rel_l2 = (err.pow(2).sum().sqrt()/true.pow(2).sum().sqrt()).item()
    mae = err.abs().mean().item()
    maxabs = err.abs().max().item()
    return dict(rel_l2=rel_l2, mae=mae, maxabs=maxabs)

def augment(xb, jitter=0.01, reflect=(True,True,True), rot_deg=0.0, scale_jit=0.0, clip=0.03):
    # xb: [B,2048,3] on device
    B = xb.shape[0]; dev=xb.device
    if any(reflect):
        signs = torch.ones(B,1,3, device=dev)
        for a,on in enumerate(reflect):
            if on:
                flip = (torch.rand(B,1,device=dev) < 0.5).float()*(-2)+1  # +1 or -1
                signs[:,:,a] = flip.squeeze(1)
        xb = xb*signs
    if rot_deg>0:
        # small rotation about Y axis (vertical long axis) only, to keep envelope-ish
        ang = (torch.rand(B,device=dev)*2-1)*math.radians(rot_deg)
        c=torch.cos(ang); s=torch.sin(ang)
        R = torch.zeros(B,3,3, device=dev)
        R[:,0,0]=c; R[:,0,2]=s; R[:,2,0]=-s; R[:,2,2]=c; R[:,1,1]=1
        xb = torch.bmm(xb, R.transpose(1,2))
    if scale_jit>0:
        sc = 1+ (torch.rand(B,1,3,device=dev)*2-1)*scale_jit
        xb = xb*sc
    if jitter>0:
        n = torch.randn_like(xb)*jitter
        if clip>0: n = n.clamp(-clip, clip)
        xb = xb + n
    return xb

class EMA:
    def __init__(self, model, decay=0.999):
        self.decay=decay
        self.shadow={k:v.detach().clone() for k,v in model.state_dict().items()}
    def update(self, model):
        for k,v in model.state_dict().items():
            s=self.shadow[k]
            if v.dtype.is_floating_point:
                s.mul_(self.decay).add_(v.detach(), alpha=1-self.decay)
            else:
                s.copy_(v)
    def copy_to(self, model):
        model.load_state_dict(self.shadow, strict=True)
PY
echo "wrote lib.py"

cat > /app/work/trainer.py << 'PY'
from __future__ import annotations
import math, copy, time
import torch, torch.nn.functional as F
from lib import build_model, predict, metrics, augment, EMA

def train_model(cfg, Xtr, Ytr, fm, fs, lm, ls, device,
                Xval=None, Yval=None, epochs=300, bs=32, lr=1e-3, wd=1e-4,
                jitter=0.01, reflect=(True,True,True), rot_deg=0.0, scale_jit=0.0,
                ema_decay=0.999, warmup=10, seed=0, loss='mse', huber_beta=0.1,
                select='ema', verbose=False, sched='cosine', min_lr_ratio=0.01):
    torch.manual_seed(seed); 
    model = build_model(cfg).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
    Xtr_d = Xtr.to(device); Ytr_d = ((Ytr-lm)/ls).to(device)
    N = Xtr_d.shape[0]
    ema = EMA(model, ema_decay)
    steps_per = max(1, math.ceil(N/bs))
    total = epochs*steps_per
    def lr_at(step):
        if step < warmup*steps_per:
            return lr*(step+1)/(warmup*steps_per)
        if sched=='cosine':
            prog=(step-warmup*steps_per)/max(1,(total-warmup*steps_per))
            return lr*(min_lr_ratio + (1-min_lr_ratio)*0.5*(1+math.cos(math.pi*prog)))
        return lr
    best=None; best_state=None; step=0
    for ep in range(epochs):
        model.train()
        perm = torch.randperm(N, device=device)
        for i in range(0, N, bs):
            idx = perm[i:i+bs]
            if idx.numel()<2:  # avoid BN on batch size 1
                continue
            xb = Xtr_d[idx]; yb = Ytr_d[idx]
            xb = augment(xb, jitter=jitter, reflect=reflect, rot_deg=rot_deg, scale_jit=scale_jit)
            xb = (xb - fm.to(device))/fs.to(device)
            for g in opt.param_groups: g['lr']=lr_at(step)
            opt.zero_grad(set_to_none=True)
            pred = model(xb)
            if loss=='mse':
                l = F.mse_loss(pred, yb)
            elif loss=='huber':
                l = F.smooth_l1_loss(pred, yb, beta=huber_beta)
            l.backward()
            opt.step(); ema.update(model); step+=1
        if Xval is not None and (ep+1)%5==0:
            # evaluate EMA
            cur = copy.deepcopy(model.state_dict())
            if select=='ema': ema.copy_to(model)
            m = metrics(predict(model, Xval, fm,fs,lm,ls, device), Yval)
            model.load_state_dict(cur)
            score = m['rel_l2']
            if best is None or score<best:
                best=score
                best_state = {k:v.detach().cpu().clone() for k,v in (ema.shadow if select=='ema' else model.state_dict()).items()}
            if verbose and ((ep+1)%25==0):
                print(f'  ep{ep+1}: val rel_l2={m["rel_l2"]:.4f} mae={m["mae"]:.4f} max={m["maxabs"]:.4f} (best={best:.4f})')
    # final: if no val, use ema end
    if best_state is None:
        if select=='ema': ema.copy_to(model)
        best_state={k:v.detach().cpu().clone() for k,v in model.state_dict().items()}
    return model, best_state, best
PY
echo done

cd /app/work && python -c "
import torch, time
from lib import *
from trainer import train_model
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val'); Xte,Yte=load_split('public_test')
fm,fs,lm,ls=compute_norm(Xtr,Ytr)
print('feat_mean',fm.tolist(),'feat_std',fs.tolist(),'label_mean',lm.tolist(),'label_std',ls.tolist())
# baseline: predict train mean
import copy
pm = torch.full_like(Yval, lm.item())
print('predict-train-mean val:', metrics(pm,Yval))
t=time.time()
model,state,best=train_model(cfg,Xtr,Ytr,fm,fs,lm,ls,device,Xval,Yval,
   epochs=200,bs=32,lr=1e-3,wd=1e-4,jitter=0.01,reflect=(True,True,True),
   ema_decay=0.999,warmup=10,seed=0,select='ema',verbose=True)
print('time',time.time()-t,'best val rel_l2',best)
model.load_state_dict(state)
print('VAL ', metrics(predict(model,Xval,fm,fs,lm,ls,device),Yval))
print('TEST', metrics(predict(model,Xte,fm,fs,lm,ls,device),Yte))
"

cd /app/work && sed -i 's/                signs\[:,:,a\] = flip.squeeze(1)/                signs[:,:,a] = flip/' lib.py && sed -n '52,62p' lib.py

cd /app/work && python -c "
import torch, time
from lib import *
from trainer import train_model
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val'); Xte,Yte=load_split('public_test')
fm,fs,lm,ls=compute_norm(Xtr,Ytr)
t=time.time()
model,state,best=train_model(cfg,Xtr,Ytr,fm,fs,lm,ls,device,Xval,Yval,
   epochs=200,bs=32,lr=1e-3,wd=1e-4,jitter=0.01,reflect=(True,True,True),
   ema_decay=0.999,warmup=10,seed=0,select='ema',verbose=True)
print('time %.1fs'%(time.time()-t),'best val rel_l2',round(best,4))
model.load_state_dict(state)
print('VAL ', {k:round(v,4) for k,v in metrics(predict(model,Xval,fm,fs,lm,ls,device),Yval).items()})
print('TEST', {k:round(v,4) for k,v in metrics(predict(model,Xte,fm,fs,lm,ls,device),Yte).items()})
"

cd /app/work && python -c "
import torch, time, math, copy
import torch.nn.functional as F
from lib import *
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
fm,fs,lm,ls=compute_norm(Xtr,Ytr)
torch.manual_seed(0)
model=build_model(cfg).to(device)
opt=torch.optim.AdamW(model.parameters(),lr=1e-3,weight_decay=1e-4)
Xtr_d=Xtr.to(device); Ytr_d=((Ytr-lm)/ls).to(device); N=Xtr_d.shape[0]
fmd=fm.to(device); fsd=fs.to(device)
for ep in range(150):
    model.train(); perm=torch.randperm(N,device=device); tot=0;nb=0
    for i in range(0,N,32):
        idx=perm[i:i+32]
        xb=Xtr_d[idx]; yb=Ytr_d[idx]
        xb=augment(xb,jitter=0.01,reflect=(True,True,True))
        xb=(xb-fmd)/fsd
        opt.zero_grad(set_to_none=True)
        p=model(xb); l=F.mse_loss(p,yb); l.backward(); opt.step()
        tot+=l.item(); nb+=1
    if (ep+1)%25==0:
        model.eval()
        with torch.no_grad():
            vp=predict(model,Xval,fm,fs,lm,ls,device)
        print(f'ep{ep+1} trainloss={tot/nb:.4f} raw-val={metrics(vp,Yval)[\"rel_l2\"]:.4f}')
"

cd /app/work && python -c "
import torch, math
import torch.nn.functional as F
from lib import *
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train')
fm,fs,lm,ls=compute_norm(Xtr,Ytr)
torch.manual_seed(0)
model=build_model(cfg).to(device)
opt=torch.optim.AdamW(model.parameters(),lr=1e-3,weight_decay=0.0)
Xtr_d=Xtr.to(device); Ytr_d=((Ytr-lm)/ls).to(device); N=Xtr_d.shape[0]
fmd=fm.to(device); fsd=fs.to(device)
Xn=(Xtr_d-fmd)/fsd
for ep in range(400):
    model.train(); perm=torch.randperm(N,device=device)
    for i in range(0,N,32):
        idx=perm[i:i+32]
        opt.zero_grad(set_to_none=True)
        p=model(Xn[idx]); l=F.mse_loss(p,Ytr_d[idx]); l.backward(); opt.step()
    if (ep+1)%50==0:
        model.eval()
        with torch.no_grad():
            tp=predict(model,Xtr,fm,fs,lm,ls,device)
        print(f'ep{ep+1} train rel_l2={metrics(tp,Ytr)[\"rel_l2\"]:.4f} mae={metrics(tp,Ytr)[\"mae\"]:.4f}')
"

cd /app/work && cat > sweep.py << 'PY'
import torch, math, copy
import torch.nn.functional as F
from lib import *

def run(cfg, Xtr,Ytr,Xval,Yval, fm,fs,lm,ls, device,
        epochs=500, bs=32, lr=1e-3, wd=1e-4, jitter=0.0, reflect=(False,False,False),
        rot_deg=0.0, scale_jit=0.0, seed=0, evalevery=10, warmup=15, min_lr_ratio=0.01,
        track_ema=0.0):
    torch.manual_seed(seed)
    model=build_model(cfg).to(device)
    opt=torch.optim.AdamW(model.parameters(),lr=lr,weight_decay=wd)
    Xtr_d=Xtr.to(device); Ytr_d=((Ytr-lm)/ls).to(device); N=Xtr_d.shape[0]
    fmd=fm.to(device); fsd=fs.to(device)
    steps_per=max(1,math.ceil(N/bs)); total=epochs*steps_per
    ema=EMA(model,track_ema) if track_ema>0 else None
    def lr_at(s):
        if s<warmup*steps_per: return lr*(s+1)/(warmup*steps_per)
        prog=(s-warmup*steps_per)/max(1,total-warmup*steps_per)
        return lr*(min_lr_ratio+(1-min_lr_ratio)*0.5*(1+math.cos(math.pi*prog)))
    best=(9,9,9); best_state=None; s=0
    for ep in range(epochs):
        model.train(); perm=torch.randperm(N,device=device)
        for i in range(0,N,bs):
            idx=perm[i:i+bs]
            if idx.numel()<2: continue
            xb=augment(Xtr_d[idx],jitter=jitter,reflect=reflect,rot_deg=rot_deg,scale_jit=scale_jit)
            xb=(xb-fmd)/fsd
            for g in opt.param_groups: g['lr']=lr_at(s)
            opt.zero_grad(set_to_none=True)
            l=F.mse_loss(model(xb),Ytr_d[idx]); l.backward(); opt.step(); s+=1
            if ema: ema.update(model)
        if (ep+1)%evalevery==0:
            if ema:
                cur=copy.deepcopy(model.state_dict()); ema.copy_to(model)
            m=metrics(predict(model,Xval,fm,fs,lm,ls,device),Yval)
            if ema: model.load_state_dict(cur)
            if m['rel_l2']<best[0]:
                best=(m['rel_l2'],m['mae'],m['maxabs'])
                src = ema.shadow if ema else model.state_dict()
                best_state={k:v.detach().cpu().clone() for k,v in src.items()}
    return best, best_state
PY
python -c "
import torch
from lib import *
from sweep import run
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
fm,fs,lm,ls=compute_norm(Xtr,Ytr)
configs={
 'no-aug':dict(jitter=0.0,reflect=(False,False,False)),
 'jit0.01':dict(jitter=0.01,reflect=(False,False,False)),
 'jit0.02':dict(jitter=0.02,reflect=(False,False,False)),
 'refl-all':dict(jitter=0.0,reflect=(True,True,True)),
 'jit0.01+refl':dict(jitter=0.01,reflect=(True,True,True)),
 'jit0.02+refl':dict(jitter=0.02,reflect=(True,True,True)),
}
for name,kw in configs.items():
    rs=[]
    for seed in [0,1,2]:
        best,_=run(cfg,Xtr,Ytr,Xval,Yval,fm,fs,lm,ls,device,epochs=400,seed=seed,**kw)
        rs.append(best)
    import numpy as np
    a=np.array(rs)
    print(f'{name:16s} rel_l2={a[:,0].mean():.4f}±{a[:,0].std():.3f}  mae={a[:,1].mean():.4f}  max={a[:,2].mean():.4f}')
"

cd /app/work && python -c "import scipy, sklearn; print('scipy',scipy.__version__,'sklearn',sklearn.__version__)" 2>&1 | head

cd /app/work && pip install scipy scikit-learn 2>&1 | tail -2

cd /app/work && cat > volfeat.py << 'PY'
import torch, torch.nn.functional as F

def voxel_volume(pc, res=20, close_iters=2):
    # pc: [P,3] in [-0.5,0.5]; returns estimated solid volume fraction + surface count
    dev=pc.device
    idx = ((pc+0.5)*res).long().clamp(0,res-1)  # [P,3]
    occ = torch.zeros(res,res,res, device=dev)
    occ[idx[:,0],idx[:,1],idx[:,2]]=1.0
    surf = occ.clone()
    # morphological closing: dilate then erode, to seal gaps
    def dilate(v,k=3):
        return (F.max_pool3d(v[None,None],k,1,k//2)[0,0]>0).float()
    def erode(v,k=3):
        return (F.max_pool3d((1-v)[None,None],k,1,k//2)[0,0]==0).float()
    v=occ
    for _ in range(close_iters): v=dilate(v)
    for _ in range(close_iters): v=erode(v)
    closed=v
    # flood fill outside: BFS from all border voxels through non-closed cells
    free=(closed==0).float()
    outside=torch.zeros_like(free)
    # seed border
    outside[0,:,:]=free[0,:,:]; outside[-1,:,:]=free[-1,:,:]
    outside[:,0,:]=free[:,0,:]; outside[:,-1,:]=free[:,-1,:]
    outside[:,:,0]=free[:,:,0]; outside[:,:,-1]=free[:,:,-1]
    for _ in range(res*3):
        new=(F.max_pool3d(outside[None,None],3,1,1)[0,0]>0).float()*free
        if torch.equal(new,outside): break
        outside=new
    inside = (1-outside)  # includes closed solid + trapped interior
    vol = inside.sum()/ (res**3)
    return vol.item(), surf.sum().item(), closed.sum().item()/(res**3)

if __name__=='__main__':
    from lib import *
    import numpy as np
    device='cuda'
    feats={}; labels={}
    for s in ['train','val','public_test']:
        X,Y=load_split(s)
        vs=[]
        for i in range(X.shape[0]):
            pc=X[i].to(device)
            vol,surf,closed=voxel_volume(pc,res=20,close_iters=2)
            vol16,_,_=voxel_volume(pc,res=16,close_iters=2)
            vol28,_,_=voxel_volume(pc,res=28,close_iters=3)
            vs.append([vol,surf,closed,vol16,vol28])
        feats[s]=np.array(vs); labels[s]=Y.squeeze(1).numpy()
    names=['vol20','surf','closed20','vol16','vol28']
    tr=feats['train']; ytr=labels['train']
    for j,nm in enumerate(names):
        c=np.corrcoef(tr[:,j],ytr)[0,1]
        print(f'{nm:10s} corr={c:+.3f}')
    np.savez('/app/work/volfeats.npz', 
             tr=feats['train'],va=feats['val'],te=feats['public_test'],
             ytr=labels['train'],yva=labels['val'],yte=labels['public_test'])
    print('saved volfeats.npz')
PY
python volfeat.py

cd /app/work && python -c "
import torch, numpy as np
from lib import *
device='cuda'
from volfeat import voxel_volume
def feat_matrix(X):
    F=[]
    for i in range(X.shape[0]):
        pc=X[i]
        mn=pc.amin(0); mx=pc.amax(0); ext=(mx-mn)
        g=pc.to(device)
        v20,surf,cl=voxel_volume(g,20,2); v16,_,_=voxel_volume(g,16,2); v24,_,_=voxel_volume(g,24,3)
        row=[v20,v16,v24,cl,surf/8000.0,
             ext[0].item(),ext[2].item(),
             pc[:,0].std().item(),pc[:,1].std().item(),pc[:,2].std().item(),
             pc.norm(dim=1).mean().item(), pc[:,0].abs().mean().item()]
        F.append(row)
    return np.array(F)
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0).squeeze(1).numpy()
FM=feat_matrix(X)
np.savez('/app/work/handfeats.npz',F=FM,Y=Y)
print('feat matrix',FM.shape)
# 5-fold CV ridge regression
def cv_ridge(F,Y,alpha=1.0,folds=5,poly=False):
    N=len(Y); idx=np.arange(N); rng=np.random.RandomState(0); rng.shuffle(idx)
    mu=F.mean(0); sd=F.std(0)+1e-8; Fn=(F-mu)/sd
    if poly:
        Fn=np.concatenate([Fn,Fn**2],1)
    preds=np.zeros(N)
    for k in range(folds):
        te=idx[k::folds]; tr=np.setdiff1d(idx,te)
        A=np.concatenate([Fn[tr],np.ones((len(tr),1))],1)
        w=np.linalg.solve(A.T@A+alpha*np.eye(A.shape[1]), A.T@Y[tr])
        At=np.concatenate([Fn[te],np.ones((len(te),1))],1)
        preds[te]=At@w
    err=preds-Y
    rel=np.sqrt((err**2).sum()/(Y**2).sum())
    return rel, np.abs(err).mean(), np.abs(err).max()
for alpha in [0.1,1.0,5.0,20.0]:
    print('ridge alpha',alpha, [round(x,4) for x in cv_ridge(FM,Y,alpha)])
for alpha in [1.0,5.0,20.0]:
    print('ridge+poly alpha',alpha, [round(x,4) for x in cv_ridge(FM,Y,alpha,poly=True)])
"

cd /app/work && cat > cv.py << 'PY'
import torch, math, copy, numpy as np
import torch.nn.functional as F
from lib import *

def train_fold(cfg, Xtr,Ytr, fm,fs,lm,ls, device, epochs,bs,lr,wd,
               jitter,reflect,rot_deg,scale_jit,seed,warmup,min_lr_ratio,ema_decay,
               record=None):
    torch.manual_seed(seed)
    model=build_model(cfg).to(device)
    opt=torch.optim.AdamW(model.parameters(),lr=lr,weight_decay=wd)
    Xtr_d=Xtr.to(device); Ytr_d=((Ytr-lm)/ls).to(device); N=Xtr_d.shape[0]
    fmd=fm.to(device); fsd=fs.to(device)
    steps_per=max(1,math.ceil(N/bs)); total=epochs*steps_per
    ema=EMA(model,ema_decay) if ema_decay>0 else None
    def lr_at(s):
        if s<warmup*steps_per: return lr*(s+1)/(warmup*steps_per)
        prog=(s-warmup*steps_per)/max(1,total-warmup*steps_per)
        return lr*(min_lr_ratio+(1-min_lr_ratio)*0.5*(1+math.cos(math.pi*prog)))
    s=0
    snaps={}
    for ep in range(epochs):
        model.train(); perm=torch.randperm(N,device=device)
        for i in range(0,N,bs):
            idx=perm[i:i+bs]
            if idx.numel()<2: continue
            xb=augment(Xtr_d[idx],jitter=jitter,reflect=reflect,rot_deg=rot_deg,scale_jit=scale_jit)
            xb=(xb-fmd)/fsd
            for g in opt.param_groups: g['lr']=lr_at(s)
            opt.zero_grad(set_to_none=True)
            F.mse_loss(model(xb),Ytr_d[idx]).backward(); opt.step(); s+=1
            if ema: ema.update(model)
        if record and (ep+1) in record:
            if ema:
                cur=copy.deepcopy(model.state_dict()); ema.copy_to(model)
                snaps[ep+1]={k:v.detach().cpu().clone() for k,v in model.state_dict().items()}
                model.load_state_dict(cur)
            else:
                snaps[ep+1]={k:v.detach().cpu().clone() for k,v in model.state_dict().items()}
    if ema: ema.copy_to(model)
    final={k:v.detach().cpu().clone() for k,v in model.state_dict().items()}
    return model, final, snaps

def cv_eval(cfg, X, Y, device, folds=5, epochs=400, bs=32, lr=1e-3, wd=1e-4,
            jitter=0.01, reflect=(False,False,False), rot_deg=0.0, scale_jit=0.0,
            seeds=(0,), warmup=15, min_lr_ratio=0.01, ema_decay=0.0, record_epochs=None,
            isotropic=False):
    N=X.shape[0]; idx=np.arange(N); rng=np.random.RandomState(123); rng.shuffle(idx)
    # pooled predictions per (epoch-snapshot or final), averaged over seeds
    rec = record_epochs or [epochs]
    pred_acc={e:torch.zeros(N) for e in rec}
    for k in range(folds):
        te=idx[k::folds]; tr=np.setdiff1d(idx,te)
        Xtr,Ytr=X[tr],Y[tr]; Xte,Yte=X[te],Y[te]
        fm,fs,lm,ls=compute_norm(Xtr,Ytr,isotropic=isotropic)
        for seed in seeds:
            model,final,snaps=train_fold(cfg,Xtr,Ytr,fm,fs,lm,ls,device,epochs,bs,lr,wd,
                jitter,reflect,rot_deg,scale_jit,seed,warmup,min_lr_ratio,ema_decay,
                record=set(rec))
            for e in rec:
                st = snaps.get(e, final)
                model.load_state_dict(st)
                p=predict(model,Xte,fm,fs,lm,ls,device).squeeze(1)
                pred_acc[e][torch.tensor(te)] += p/len(seeds)
    Yf=Y.squeeze(1)
    out={}
    for e in rec:
        out[e]=metrics(pred_acc[e],Yf)
    return out
PY
echo "cv.py ready"

cd /app/work && python -c "
import torch, time
from lib import *
from cv import cv_eval
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0)
rec=[150,250,400,600]
t=time.time()
configs={
 'no-aug':dict(jitter=0.0,reflect=(False,False,False)),
 'jit0.005':dict(jitter=0.005,reflect=(False,False,False)),
 'jit0.01':dict(jitter=0.01,reflect=(False,False,False)),
}
for name,kw in configs.items():
    out=cv_eval(cfg,X,Y,device,folds=5,epochs=600,bs=32,lr=1e-3,wd=1e-4,
                seeds=(0,1),ema_decay=0.0,record_epochs=rec,**kw)
    line=' | '.join([f'e{e}:rl2={out[e][\"rel_l2\"]:.4f},mae={out[e][\"mae\"]:.3f},mx={out[e][\"maxabs\"]:.3f}' for e in rec])
    print(f'{name:10s} {line}')
print('elapsed %.0fs'%(time.time()-t))
"

restart: true

cd /app/work && cat > exp1.py << 'PY'
import torch, time, sys
from lib import *
from cv import cv_eval
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0)
rec=[150,250,400]
configs={
 'no-aug':dict(jitter=0.0,reflect=(False,False,False)),
 'jit0.005':dict(jitter=0.005,reflect=(False,False,False)),
 'jit0.01':dict(jitter=0.01,reflect=(False,False,False)),
 'jit0.01+Xrefl':dict(jitter=0.01,reflect=(True,False,False)),
}
for name,kw in configs.items():
    t=time.time()
    out=cv_eval(cfg,X,Y,device,folds=5,epochs=400,bs=32,lr=1e-3,wd=1e-4,
                seeds=(0,1),ema_decay=0.0,record_epochs=rec,**kw)
    line=' | '.join([f'e{e}:rl2={out[e]["rel_l2"]:.4f},mae={out[e]["mae"]:.3f},mx={out[e]["maxabs"]:.3f}' for e in rec])
    print(f'{name:14s} {line}  ({time.time()-t:.0f}s)',flush=True)
PY
nohup python exp1.py > exp1.log 2>&1 &
echo "started pid $!"

cd /app/work && sleep 120 && cat exp1.log

cd /app/work && sleep 90 && cat exp1.log; echo "--- still running? ---"; ps aux | grep exp1 | grep -v grep | head

cd /app/work && sleep 280 && cat exp1.log; echo "---"; ps aux|grep exp1|grep -v grep|head

cd /app/work && cat > exp2.py << 'PY'
import torch, time
from lib import *
from cv import cv_eval
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0)
rec=[200,300]
base=dict(folds=5,epochs=300,lr=1e-3,jitter=0.0,reflect=(False,False,False),
          seeds=(0,),record_epochs=rec)
configs={
 'bs32 wd1e-4 ema0':   dict(bs=32,wd=1e-4,ema_decay=0.0),
 'bs32 wd1e-4 ema.99': dict(bs=32,wd=1e-4,ema_decay=0.99),
 'bs32 wd1e-3 ema0':   dict(bs=32,wd=1e-3,ema_decay=0.0),
 'bs32 wd1e-2 ema0':   dict(bs=32,wd=1e-2,ema_decay=0.0),
 'bs16 wd1e-4 ema0':   dict(bs=16,wd=1e-4,ema_decay=0.0),
 'bs64 wd1e-4 ema0':   dict(bs=64,wd=1e-4,ema_decay=0.0),
 'bs32 wd1e-3 iso':    dict(bs=32,wd=1e-3,ema_decay=0.0,isotropic=True),
}
for name,kw in configs.items():
    t=time.time()
    k2={**base,**kw}
    out=cv_eval(cfg,X,Y,device,**k2)
    line=' | '.join([f'e{e}:rl2={out[e]["rel_l2"]:.4f},mae={out[e]["mae"]:.3f},mx={out[e]["maxabs"]:.3f}' for e in rec])
    print(f'{name:20s} {line}  ({time.time()-t:.0f}s)',flush=True)
PY
nohup python exp2.py > exp2.log 2>&1 &
echo "started $!"

cd /app/work && sleep 240 && cat exp2.log

cd /app/work && sleep 180 && cat exp2.log; echo "---"; ps aux|grep exp2|grep -v grep|head

cd /app/work && cat > exp3.py << 'PY'
import torch, time
from lib import *
from cv import cv_eval
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0)
rec=[250,400]
base=dict(folds=5,epochs=400,reflect=(False,False,False),seeds=(0,),record_epochs=rec)
configs={
 'bs32 lr1e-3 wd2e-2':   dict(bs=32,lr=1e-3,wd=2e-2,jitter=0.0),
 'bs32 lr1e-3 wd5e-2':   dict(bs=32,lr=1e-3,wd=5e-2,jitter=0.0),
 'bs32 lr1e-3 wd1e-1':   dict(bs=32,lr=1e-3,wd=1e-1,jitter=0.0),
 'bs16 lr1e-3 wd2e-2':   dict(bs=16,lr=1e-3,wd=2e-2,jitter=0.0),
 'bs16 lr1e-3 wd5e-2':   dict(bs=16,lr=1e-3,wd=5e-2,jitter=0.0),
 'bs32 lr2e-3 wd2e-2':   dict(bs=32,lr=2e-3,wd=2e-2,jitter=0.0),
 'bs32 lr5e-4 wd2e-2':   dict(bs=32,lr=5e-4,wd=2e-2,jitter=0.0),
 'bs32 lr1e-3 wd3e-2 jit.005': dict(bs=32,lr=1e-3,wd=3e-2,jitter=0.005),
}
for name,kw in configs.items():
    t=time.time()
    out=cv_eval(cfg,X,Y,device,**{**base,**kw})
    line=' | '.join([f'e{e}:rl2={out[e]["rel_l2"]:.4f},mae={out[e]["mae"]:.3f},mx={out[e]["maxabs"]:.3f}' for e in rec])
    print(f'{name:28s} {line}  ({time.time()-t:.0f}s)',flush=True)
PY
nohup python exp3.py > exp3.log 2>&1 &
echo "started $!"

cd /app/work && sleep 290 && cat exp3.log

cd /app/work && sleep 220 && cat exp3.log | tail -5; echo "---"; ps aux|grep exp3|grep -v grep|head

cd /app/work && sleep 60 && tail -1 exp3.log

cd /app/work && cat >> cv.py << 'PY'

def cv_soup(cfg, X, Y, device, folds=5, epochs=400, bs=32, lr=1.5e-3, wd=2e-2,
            jitter=0.0, reflect=(False,False,False), seeds=(0,1,2), warmup=15,
            min_lr_ratio=0.01, ema_decay=0.0):
    # Compare: single(seed0), pred-ensemble(seeds), weight-soup(seeds)
    N=X.shape[0]; idx=np.arange(N); rng=np.random.RandomState(123); rng.shuffle(idx)
    pred_single=torch.zeros(N); pred_ens=torch.zeros(N); pred_soup=torch.zeros(N)
    for k in range(folds):
        te=idx[k::folds]; tr=np.setdiff1d(idx,te)
        Xtr,Ytr=X[tr],Y[tr]; Xte,Yte=X[te],Y[te]
        fm,fs,lm,ls=compute_norm(Xtr,Ytr)
        states=[]
        for si,seed in enumerate(seeds):
            model,final,_=train_fold(cfg,Xtr,Ytr,fm,fs,lm,ls,device,epochs,bs,lr,wd,
                jitter,reflect,0.0,0.0,seed,warmup,min_lr_ratio,ema_decay,record=None)
            states.append(final)
            p=predict(model,Xte,fm,fs,lm,ls,device).squeeze(1)
            pred_ens[torch.tensor(te)] += p/len(seeds)
            if si==0: pred_single[torch.tensor(te)] = p
        # soup
        soup={k2: sum(s[k2].float() for s in states)/len(states) for k2 in states[0]}
        model.load_state_dict(soup)
        p=predict(model,Xte,fm,fs,lm,ls,device).squeeze(1)
        pred_soup[torch.tensor(te)] = p
    Yf=Y.squeeze(1)
    return dict(single=metrics(pred_single,Yf), ensemble=metrics(pred_ens,Yf), soup=metrics(pred_soup,Yf))
PY
cat > exp4.py << 'PY'
import torch, time
from lib import *
from cv import cv_soup
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0)
for tag,kw in {
  'lr1.5e-3 wd2e-2 e500 ema0':dict(lr=1.5e-3,wd=2e-2,epochs=500,ema_decay=0.0),
  'lr1.5e-3 wd2e-2 e500 ema.99':dict(lr=1.5e-3,wd=2e-2,epochs=500,ema_decay=0.99),
}.items():
    t=time.time()
    r=cv_soup(cfg,X,Y,device,folds=5,bs=32,seeds=(0,1,2,3),**kw)
    print(tag,flush=True)
    for k,v in r.items():
        print(f'   {k:9s} rl2={v["rel_l2"]:.4f} mae={v["mae"]:.4f} max={v["maxabs"]:.4f}',flush=True)
    print(f'   ({time.time()-t:.0f}s)',flush=True)
PY
nohup python exp4.py > exp4.log 2>&1 &
echo "started $!"

cd /app/work && sleep 280 && cat exp4.log

cd /app/work && sleep 200 && cat exp4.log; echo ---; ps aux|grep exp4|grep -v grep|head -1

cd /app/work && sleep 180 && cat exp4.log; echo ---; ps aux|grep exp4|grep -v grep|head -1

cd /app/work && cat > exp5.py << 'PY'
import torch, time, numpy as np
from lib import *
from cv import train_fold
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0)
N=X.shape[0]; idx=np.arange(N); rng=np.random.RandomState(123); rng.shuffle(idx)
def cv_single(seed, epochs=400, lr=1.5e-3, wd=2e-2, ema_decay=0.99):
    pred=torch.zeros(N)
    for k in range(5):
        te=idx[k::5]; tr=np.setdiff1d(idx,te)
        Xtr2,Ytr2=X[tr],Y[tr]
        fm,fs,lm,ls=compute_norm(Xtr2,Ytr2)
        model,final,_=train_fold(cfg,Xtr2,Ytr2,fm,fs,lm,ls,device,epochs,32,lr,wd,
            0.0,(False,False,False),0.0,0.0,seed,15,0.01,ema_decay,record=None)
        model.load_state_dict(final)
        pred[torch.tensor(te)]=predict(model,X[te],fm,fs,lm,ls,device).squeeze(1)
    return metrics(pred,Y.squeeze(1)), pred
Yf=Y.squeeze(1)
allp=[]
for seed in range(6):
    t=time.time()
    m,pred=cv_single(seed)
    allp.append(pred)
    print(f'seed{seed}: rl2={m["rel_l2"]:.4f} mae={m["mae"]:.4f} max={m["maxabs"]:.4f} ({time.time()-t:.0f}s)',flush=True)
ens=torch.stack(allp).mean(0)
me=metrics(ens,Yf)
print(f'6-seed ens: rl2={me["rel_l2"]:.4f} mae={me["mae"]:.4f} max={me["maxabs"]:.4f}',flush=True)
torch.save(torch.stack(allp),'seedpreds.pt')
PY
nohup python exp5.py > exp5.log 2>&1 &
echo "started $!"

cd /app/work && sleep 250 && cat exp5.log

cd /app/work && sleep 220 && cat exp5.log

cd /app/work && cat > exp6.py << 'PY'
import torch, time, numpy as np, math, copy
import torch.nn.functional as F
from lib import *
from cv import train_fold
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0)
N=X.shape[0]; idx=np.arange(N); rng=np.random.RandomState(123); rng.shuffle(idx)
Yf=Y.squeeze(1)
K=6; EP=400; LR=1.5e-3; WD=2e-2

def distill_student(Xtr2, soft_tr, fm,fs,lm,ls, seed, epochs=400):
    torch.manual_seed(1000+seed)
    model=build_model(cfg).to(device)
    opt=torch.optim.AdamW(model.parameters(),lr=LR,weight_decay=WD)
    Xd=Xtr2.to(device); yd=((soft_tr-lm)/ls).to(device); n=Xd.shape[0]
    fmd=fm.to(device); fsd=fs.to(device); spp=max(1,math.ceil(n/32)); tot=epochs*spp; s=0
    for ep in range(epochs):
        model.train(); perm=torch.randperm(n,device=device)
        for i in range(0,n,32):
            ii=perm[i:i+32]
            if ii.numel()<2: continue
            xb=(Xd[ii]-fmd)/fsd
            for g in opt.param_groups:
                s2=s
                g['lr']=LR*((s2+1)/(15*spp) if s2<15*spp else 0.01+0.99*0.5*(1+math.cos(math.pi*(s2-15*spp)/max(1,tot-15*spp))))
            opt.zero_grad(set_to_none=True)
            F.mse_loss(model(xb),yd[ii]).backward(); opt.step(); s+=1
    return model

pred_single=torch.zeros(N); pred_ens=torch.zeros(N); pred_med=torch.zeros(N); pred_dist=torch.zeros(N)
t0=time.time()
for k in range(5):
    te=idx[k::5]; tr=np.setdiff1d(idx,te)
    Xtr2,Ytr2=X[tr],Y[tr]; Xte=X[te]
    fm,fs,lm,ls=compute_norm(Xtr2,Ytr2)
    teach_te=[]; teach_tr=[]
    for seed in range(K):
        m,final,_=train_fold(cfg,Xtr2,Ytr2,fm,fs,lm,ls,device,EP,32,LR,WD,
            0.0,(False,False,False),0.0,0.0,seed,15,0.01,0.99,record=None)
        m.load_state_dict(final)
        teach_te.append(predict(m,Xte,fm,fs,lm,ls,device).squeeze(1))
        teach_tr.append(predict(m,Xtr2,fm,fs,lm,ls,device).squeeze(1))
    T_te=torch.stack(teach_te); T_tr=torch.stack(teach_tr)
    ens_te=T_te.mean(0); ens_tr=T_tr.mean(0)
    pred_single[torch.tensor(te)]=teach_te[0]
    pred_ens[torch.tensor(te)]=ens_te
    # medoid: teacher closest to ens on TRAIN preds
    d=((T_tr-ens_tr[None])**2).mean(1); j=int(d.argmin())
    pred_med[torch.tensor(te)]=teach_te[j]
    # distill student on ensemble train preds
    stu=distill_student(Xtr2, ens_tr.unsqueeze(1), fm,fs,lm,ls, 0, epochs=EP)
    pred_dist[torch.tensor(te)]=predict(stu,Xte,fm,fs,lm,ls,device).squeeze(1)
    print(f'fold{k} done medoid=seed{j} ({time.time()-t0:.0f}s)',flush=True)
for nm,p in [('single',pred_single),('ensemble',pred_ens),('medoid',pred_med),('distill',pred_dist)]:
    m=metrics(p,Yf); print(f'{nm:9s} rl2={m["rel_l2"]:.4f} mae={m["mae"]:.4f} max={m["maxabs"]:.4f}',flush=True)
PY
nohup python exp6.py > exp6.log 2>&1 &
echo "started $!"

cd /app/work && python -c "
import torch, numpy as np
from lib import *
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
Y=torch.cat([Ytr,Yval],0).squeeze(1)
P=torch.load('seedpreds.pt')  # [6,N]
N=P.shape[1]; idx=np.arange(N); rng=np.random.RandomState(123); rng.shuffle(idx)
ens=P.mean(0)
err=(ens-Y)
# sort by abs err
order=err.abs().argsort(descending=True)
print('worst 12 samples (ensemble): idx  true  pred  err')
for i in order[:12].tolist():
    print(f'  {i:3d}  true={Y[i]:.3f} pred={ens[i]:.3f} err={err[i]:+.3f}')
print('corr(true,err):', np.corrcoef(Y.numpy(),err.numpy())[0,1])
print('mean true of worst-10:', Y[order[:10]].mean().item(), 'overall mean:', Y.mean().item())
# how much of rel_l2^2 numerator from top-5?
num=err.pow(2).sum().item(); top5=err[order[:5]].pow(2).sum().item()
print(f'top-5 samples account for {100*top5/num:.0f}% of squared error')
print('fraction true>1.3:', (Y>1.3).float().mean().item(), 'count', (Y>1.3).sum().item())
"

cd /app/work && python -c "
import torch, numpy as np
from lib import *
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
Y=torch.cat([Ytr,Yval],0).squeeze(1).numpy()
P=torch.load('seedpreds.pt',weights_only=True).numpy()  # [6,N] OOF preds in mass units
def relL2(p): 
    e=p-Y; return np.sqrt((e**2).sum()/(Y**2).sum()), np.abs(e).mean(), np.abs(e).max()
# per-seed and ensemble, with/without linear calibration fit on OOF (proper would be nested, approx here)
for name,p in [('seed0',P[0]),('ens',P.mean(0))]:
    a,b=np.polyfit(p,Y,1)
    pc=a*p+b
    print(f'{name}: raw rl2/mae/max={tuple(round(x,4) for x in relL2(p))}  slope={a:.3f} b={b:.3f}')
    print(f'       calibrated   ={tuple(round(x,4) for x in relL2(pc))}')
"

cd /app/work && sleep 120 && cat exp6.log

cd /app/work && cat > final_train.py << 'PY'
import torch, math, copy, numpy as np, argparse, time
import torch.nn.functional as F
from lib import *
from cv import train_fold

cfg=load_cfg(); device='cuda'
LR=1.5e-3; WD=2e-2; BS=32; EP=500; WARM=15; EMAD=0.99

def get_all_data(include_test=True):
    Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
    parts=[(Xtr,Ytr),(Xval,Yval)]
    if include_test:
        Xte,Yte=load_split('public_test'); parts.append((Xte,Yte))
    X=torch.cat([p[0] for p in parts],0); Y=torch.cat([p[1] for p in parts],0)
    return X,Y

@torch.no_grad()
def recalc_bn(model, Xn, device, passes=3, bs=64):
    # reset and recompute BN running stats from normalized inputs Xn (already standardized)
    for m in model.modules():
        if isinstance(m, torch.nn.BatchNorm1d):
            m.reset_running_stats(); m.momentum=None  # cumulative average
    model.train()
    n=Xn.shape[0]
    for _ in range(passes):
        perm=torch.randperm(n)
        for i in range(0,n,bs):
            xb=Xn[perm[i:i+bs]].to(device)
            if xb.shape[0]<2: continue
            model(xb)
    model.eval()

def train_one(X,Y,fm,fs,lm,ls,seed,epochs=EP):
    model,final,_=train_fold(cfg,X,Y,fm,fs,lm,ls,device,epochs,BS,LR,WD,
        0.0,(False,False,False),0.0,0.0,seed,WARM,0.01,EMAD,record=None)
    model.load_state_dict(final)
    return model, final

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--method',default='medoid',choices=['single','medoid'])
    ap.add_argument('--kseeds',type=int,default=9)
    ap.add_argument('--epochs',type=int,default=EP)
    ap.add_argument('--include_test',type=int,default=1)
    ap.add_argument('--out',default='/app/outputs/model.pt')
    ap.add_argument('--bn_recalc',type=int,default=1)
    args=ap.parse_args()
    X,Y=get_all_data(bool(args.include_test))
    print('final train data:',X.shape[0],'samples',flush=True)
    fm,fs,lm,ls=compute_norm(X,Y)
    Xn=((X.to(device)-fm.to(device))/fs.to(device)).cpu()
    states=[]; preds_tr=[]
    for seed in range(args.kseeds):
        t=time.time()
        model,final=train_one(X,Y,fm,fs,lm,ls,seed,args.epochs)
        if args.bn_recalc:
            model.load_state_dict(final); recalc_bn(model,Xn,device)
            final={k:v.detach().cpu().clone() for k,v in model.state_dict().items()}
        states.append(final)
        model.load_state_dict(final)
        p=predict(model,X,fm,fs,lm,ls,device).squeeze(1)
        preds_tr.append(p)
        print(f'  seed{seed} trained ({time.time()-t:.0f}s) train-rl2={metrics(p,Y.squeeze(1))["rel_l2"]:.4f}',flush=True)
    P=torch.stack(preds_tr); ens=P.mean(0)
    if args.method=='single' or args.kseeds==1:
        chosen=states[0]; ci=0
    else:
        d=((P-ens[None])**2).mean(1); ci=int(d.argmin()); chosen=states[ci]
    print(f'chosen method={args.method} idx={ci}',flush=True)
    payload={
        'state_dict':{k:v.cpu() for k,v in chosen.items()},
        'feat_mean':fm.squeeze().cpu().float(),
        'feat_std':fs.squeeze().cpu().float(),
        'label_mean':lm.squeeze().cpu().float().reshape(1),
        'label_std':ls.squeeze().cpu().float().reshape(1),
    }
    import os; os.makedirs('/app/outputs',exist_ok=True)
    torch.save(payload,args.out)
    print('saved',args.out,flush=True)

if __name__=='__main__':
    main()
PY
echo "final_train.py ready"

cd /app/work && sleep 200 && cat exp6.log

cd /app/work && cat > exp7.py << 'PY'
import torch, time, numpy as np, math
import torch.nn.functional as F
from lib import *
from cv import train_fold
cfg=load_cfg(); device='cuda'
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val')
X=torch.cat([Xtr,Xval],0); Y=torch.cat([Ytr,Yval],0)
N=X.shape[0]; idx=np.arange(N); rng=np.random.RandomState(123); rng.shuffle(idx)
Yf=Y.squeeze(1); K=8; EP=400; LR=1.5e-3; WD=2e-2

def make_aug(Xt, copies, jitter, reflect, seedbase=0):
    outs=[Xt]
    for c in range(copies):
        g=torch.Generator(device=device).manual_seed(seedbase*97+c+1)
        xb=Xt.clone()
        if reflect:
            for a in range(3):
                fl=(torch.rand(xb.shape[0],1,generator=g,device=device)<0.5).float()*(-2)+1
                xb[:,:,a]=xb[:,:,a]*fl
        if jitter>0:
            xb=xb+ (torch.randn(xb.shape,generator=g,device=device)*jitter).clamp(-0.03,0.03)
        outs.append(xb)
    return torch.cat(outs,0)

def train_student(Xin, soft, fm,fs,lm,ls, epochs=EP, seed=0):
    torch.manual_seed(1000+seed)
    model=build_model(cfg).to(device)
    opt=torch.optim.AdamW(model.parameters(),lr=LR,weight_decay=WD)
    yd=((soft-lm.to(device))/ls.to(device)); n=Xin.shape[0]
    fmd=fm.to(device); fsd=fs.to(device); spp=max(1,math.ceil(n/BS)); 
    BS2=64; spp=max(1,math.ceil(n/BS2)); tot=epochs*spp; s=0
    for ep in range(epochs):
        model.train(); perm=torch.randperm(n,device=device)
        for i in range(0,n,BS2):
            ii=perm[i:i+BS2]
            if ii.numel()<2: continue
            xb=(Xin[ii]-fmd)/fsd
            for g in opt.param_groups:
                g['lr']=LR*((s+1)/(15*spp) if s<15*spp else 0.01+0.99*0.5*(1+math.cos(math.pi*(s-15*spp)/max(1,tot-15*spp))))
            opt.zero_grad(set_to_none=True); F.mse_loss(model(xb),yd[ii]).backward(); opt.step(); s+=1
    return model
BS=32
pe={n:torch.zeros(N) for n in ['pure','aug_j','aug_jr','blend']}
t0=time.time()
for k in range(5):
    te=idx[k::5]; tr=np.setdiff1d(idx,te)
    Xtr2,Ytr2=X[tr],Y[tr]; Xte=X[te]
    fm,fs,lm,ls=compute_norm(Xtr2,Ytr2)
    Tt=[]
    for seed in range(K):
        m,final,_=train_fold(cfg,Xtr2,Ytr2,fm,fs,lm,ls,device,EP,BS,LR,WD,0.0,(False,False,False),0.0,0.0,seed,15,0.01,0.99,None)
        m.load_state_dict(final); Tt.append(m)
    Xtr2d=Xtr2.to(device)
    def ens_pred(Xq):
        with torch.no_grad():
            ps=[predict(m,Xq,fm,fs,lm,ls,device).squeeze(1) for m in Tt]
        return torch.stack(ps).mean(0)
    soft_tr=ens_pred(Xtr2)  # cpu
    # variants
    # pure
    stu=train_student(Xtr2d, soft_tr.to(device), fm,fs,lm,ls); pe['pure'][torch.tensor(te)]=predict(stu,Xte,fm,fs,lm,ls,device).squeeze(1)
    # blend with true
    blend=0.5*soft_tr+0.5*Ytr2.squeeze(1)
    stu=train_student(Xtr2d, blend.to(device), fm,fs,lm,ls); pe['blend'][torch.tensor(te)]=predict(stu,Xte,fm,fs,lm,ls,device).squeeze(1)
    # aug jitter only x3
    Xa=make_aug(Xtr2d,3,0.01,False,seedbase=k); soft_a=ens_pred(Xa.cpu())
    stu=train_student(Xa, soft_a.to(device), fm,fs,lm,ls); pe['aug_j'][torch.tensor(te)]=predict(stu,Xte,fm,fs,lm,ls,device).squeeze(1)
    # aug jitter+reflect x3
    Xb=make_aug(Xtr2d,3,0.01,True,seedbase=k+50); soft_b=ens_pred(Xb.cpu())
    stu=train_student(Xb, soft_b.to(device), fm,fs,lm,ls); pe['aug_jr'][torch.tensor(te)]=predict(stu,Xte,fm,fs,lm,ls,device).squeeze(1)
    print(f'fold{k} done ({time.time()-t0:.0f}s)',flush=True)
for nm,p in pe.items():
    m=metrics(p,Yf); print(f'{nm:8s} rl2={m["rel_l2"]:.4f} mae={m["mae"]:.4f} max={m["maxabs"]:.4f}',flush=True)
PY
nohup python exp7.py > exp7.log 2>&1 &
echo "started $!"

cd /app/work && sleep 290 && cat exp7.log

cd /app/work && kill 618 2>/dev/null; sleep 2; pkill -f exp7.py 2>/dev/null; sleep 1
# fix: make all soft targets [n,1]
sed -i 's/    soft_tr=ens_pred(Xtr2)  # cpu/    soft_tr=ens_pred(Xtr2).unsqueeze(1)  # [n,1]/' exp7.py
sed -i 's/    blend=0.5\*soft_tr+0.5\*Ytr2.squeeze(1)/    blend=0.5*soft_tr+0.5*Ytr2/' exp7.py
sed -i 's/    Xa=make_aug(Xtr2d,3,0.01,False,seedbase=k); soft_a=ens_pred(Xa.cpu())/    Xa=make_aug(Xtr2d,3,0.01,False,seedbase=k); soft_a=ens_pred(Xa.cpu()).unsqueeze(1)/' exp7.py
sed -i 's/    Xb=make_aug(Xtr2d,3,0.01,True,seedbase=k+50); soft_b=ens_pred(Xb.cpu())/    Xb=make_aug(Xtr2d,3,0.01,True,seedbase=k+50); soft_b=ens_pred(Xb.cpu()).unsqueeze(1)/' exp7.py
grep -n "soft_tr=ens_pred\|blend=0.5\|soft_a=ens_pred\|soft_b=ens_pred\|def train_student" exp7.py

cd /app/work && nohup python exp7.py > exp7.log 2>&1 &
echo "restarted $!"

cd /app/work && sleep 295 && grep -v Warning exp7.log | grep -v mse_loss

cd /app/work && sleep 295 && grep "fold\|rl2" exp7.log

cd /app/work && sleep 295 && grep "fold\|rl2" exp7.log

cd /app/work && sleep 220 && grep "fold\|rl2" exp7.log; echo ---; ps aux|grep exp7|grep -v grep|head -1

cd /app/work && cat > pipeline.py << 'PY'
import torch, math, numpy as np, time, argparse, os
import torch.nn.functional as F
from lib import *
from cv import train_fold
cfg=load_cfg(); device='cuda'
LR=1.5e-3; WD=2e-2; BS=32; WARM=15; EMAD=0.99

def train_base(X,Y,fm,fs,lm,ls,seed,epochs):
    m,final,_=train_fold(cfg,X,Y,fm,fs,lm,ls,device,epochs,BS,LR,WD,0.0,(False,False,False),0.0,0.0,seed,WARM,0.01,EMAD,None)
    m.load_state_dict(final); return m, final

def train_student(Xin_d, soft_d, fm,fs,lm,ls, epochs, seed):
    # Xin_d,[n,2048,3] on device; soft_d [n,1] on device (mass units)
    torch.manual_seed(2000+seed)
    model=build_model(cfg).to(device)
    opt=torch.optim.AdamW(model.parameters(),lr=LR,weight_decay=WD)
    yd=(soft_d-lm.to(device))/ls.to(device); n=Xin_d.shape[0]
    fmd=fm.to(device); fsd=fs.to(device); spp=max(1,math.ceil(n/BS)); tot=epochs*spp; s=0
    ema=EMA(model,EMAD)
    for ep in range(epochs):
        model.train(); perm=torch.randperm(n,device=device)
        for i in range(0,n,BS):
            ii=perm[i:i+BS]
            if ii.numel()<2: continue
            xb=(Xin_d[ii]-fmd)/fsd
            for g in opt.param_groups:
                g['lr']=LR*((s+1)/(WARM*spp) if s<WARM*spp else 0.01+0.99*0.5*(1+math.cos(math.pi*(s-WARM*spp)/max(1,tot-WARM*spp))))
            opt.zero_grad(set_to_none=True); F.mse_loss(model(xb),yd[ii]).backward(); opt.step(); ema.update(model); s+=1
    ema.copy_to(model)
    return model

def oof_soft(X,Y,K,epochs,folds=5):
    N=X.shape[0]; idx=np.arange(N); rng=np.random.RandomState(7); rng.shuffle(idx)
    oof=torch.zeros(N)
    for k in range(folds):
        te=idx[k::folds]; tr=np.setdiff1d(idx,te)
        fm,fs,lm,ls=compute_norm(X[tr],Y[tr])
        ps=[]
        for seed in range(K):
            m,_=train_base(X[tr],Y[tr],fm,fs,lm,ls,seed,epochs)
            ps.append(predict(m,X[te],fm,fs,lm,ls,device).squeeze(1))
        oof[torch.tensor(te)]=torch.stack(ps).mean(0)
    return oof  # mass units, [N]

def run_pipeline(Xpool,Ypool,K=6,t_ep=400,s_ep=500,alpha=0.5,s_seed=0):
    fm,fs,lm,ls=compute_norm(Xpool,Ypool)
    oof=oof_soft(Xpool,Ypool,K,t_ep)            # OOF ensemble soft targets
    soft=(alpha*Ypool.squeeze(1)+(1-alpha)*oof).unsqueeze(1)
    stu=train_student(Xpool.to(device),soft.to(device),fm,fs,lm,ls,s_ep,s_seed)
    state={k:v.detach().cpu().clone() for k,v in stu.state_dict().items()}
    return stu,state,(fm,fs,lm,ls),oof

if __name__=='__main__':
    ap=argparse.ArgumentParser(); ap.add_argument('--mode',default='validate')
    ap.add_argument('--K',type=int,default=6); ap.add_argument('--tep',type=int,default=400)
    ap.add_argument('--sep',type=int,default=500); ap.add_argument('--alpha',type=float,default=0.5)
    a=ap.parse_args()
    Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val'); Xte,Yte=load_split('public_test')
    if a.mode=='validate':
        Xpool,Ypool=Xtr,Ytr
        Xhold=torch.cat([Xval,Xte],0); Yhold=torch.cat([Yval,Yte],0)
        t=time.time()
        # baseline singles + ensemble on held-out
        fm,fs,lm,ls=compute_norm(Xpool,Ypool); ps=[]
        for seed in range(a.K):
            m,_=train_base(Xpool,Ypool,fm,fs,lm,ls,seed,a.tep)
            ps.append(predict(m,Xhold,fm,fs,lm,ls,device).squeeze(1))
        singles=[metrics(p,Yhold.squeeze(1))['rel_l2'] for p in ps]
        ens=torch.stack(ps).mean(0)
        print(f'[holdout] single mean={np.mean(singles):.4f} best={min(singles):.4f} worst={max(singles):.4f}')
        print(f'[holdout] ensemble: {metrics(ens,Yhold.squeeze(1))}')
        # distilled student (cross-fit OOF on pool)
        for sseed in range(3):
            stu,state,(fm,fs,lm,ls),oof=run_pipeline(Xpool,Ypool,a.K,a.tep,a.sep,a.alpha,sseed)
            ph=predict(stu,Xhold,fm,fs,lm,ls,device).squeeze(1)
            m=metrics(ph,Yhold.squeeze(1))
            print(f'[holdout] distill sseed{sseed}: rl2={m["rel_l2"]:.4f} mae={m["mae"]:.4f} max={m["maxabs"]:.4f}  (oof-rl2={metrics(oof,Ypool.squeeze
…[truncated]

cd /app/work && kill 707 2>/dev/null; pkill -f pipeline.py 2>/dev/null; sleep 2
python - << 'PY'
import re
s=open('pipeline.py').read()
old='''        # distilled student (cross-fit OOF on pool)
        for sseed in range(3):
            stu,state,(fm,fs,lm,ls),oof=run_pipeline(Xpool,Ypool,a.K,a.tep,a.sep,a.alpha,sseed)
            ph=predict(stu,Xhold,fm,fs,lm,ls,device).squeeze(1)
            m=metrics(ph,Yhold.squeeze(1))
            print(f'[holdout] distill sseed{sseed}: rl2={m["rel_l2"]:.4f} mae={m["mae"]:.4f} max={m["maxabs"]:.4f}  (oof-rl2={metrics(oof,Ypool.squeeze(1))["rel_l2"]:.4f})',flush=True)'''
new='''        # distilled student (cross-fit OOF computed ONCE, reused)
        oof=oof_soft(Xpool,Ypool,a.K,a.tep)
        print(f'[pool] oof-ensemble rl2={metrics(oof,Ypool.squeeze(1))["rel_l2"]:.4f}',flush=True)
        for alpha in [0.5,0.0,1.0]:
            soft=(alpha*Ypool.squeeze(1)+(1-alpha)*oof).unsqueeze(1).to(device)
            hs=[]
            for sseed in range(3):
                stu=train_student(Xpool.to(device),soft,fm,fs,lm,ls,a.sep,sseed)
                ph=predict(stu,Xhold,fm,fs,lm,ls,device).squeeze(1)
                hs.append(metrics(ph,Yhold.squeeze(1)))
            rl=[h["rel_l2"] for h in hs]
            print(f'[holdout] distill alpha={alpha}: rl2 mean={np.mean(rl):.4f} (per-seed {[round(x,4) for x in rl]}) mae={np.mean([h["mae"] for h in hs]):.4f} max={np.mean([h["maxabs"] for h in hs]):.4f}',flush=True)'''
s=s.replace(old,new)
open('pipeline.py','w').write(s)
print('patched' if new in s else 'FAILED')
PY
nohup python pipeline.py --mode validate --K 6 --tep 400 --sep 500 > pipe_val.log 2>&1 &
echo "started $!"

cd /app/work && sleep 290 && cat pipe_val.log

cd /app/work && sleep 120 && cat pipe_val.log

cd /app/work && sleep 200 && cat pipe_val.log; echo ---; ps aux|grep pipeline|grep -v grep|head -1

cd /app/work && cat > exp8.py << 'PY'
import torch, math, numpy as np, time
import torch.nn.functional as F
from lib import *
from cv import train_fold
cfg=load_cfg(); device='cuda'
LR=1.5e-3; WD=2e-2; BS=32; WARM=15; EMAD=0.99
Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val'); Xpt,Ypt=load_split('public_test')
Xpool,Ypool=Xtr,Ytr
Xhold=torch.cat([Xval,Xpt],0); Yhold=torch.cat([Yval,Ypt],0).squeeze(1)
fm,fs,lm,ls=compute_norm(Xpool,Ypool)

def train_base(X,Y,seed,epochs=400):
    m,final,_=train_fold(cfg,X,Y,fm,fs,lm,ls,device,epochs,BS,LR,WD,0.0,(False,False,False),0.0,0.0,seed,WARM,0.01,EMAD,None)
    m.load_state_dict(final); return m

def aug_copies(Xd, R, jitter=0.008, seedbase=0):
    outs=[]
    for c in range(R):
        g=torch.Generator(device=device).manual_seed(seedbase*131+c+1)
        xb=Xd.clone()
        for a in range(3):
            fl=(torch.rand(xb.shape[0],1,generator=g,device=device)<0.5).float()*(-2)+1
            xb[:,:,a]=xb[:,:,a]*fl
        xb=xb+(torch.randn(xb.shape,generator=g,device=device)*jitter).clamp(-0.025,0.025)
        outs.append(xb)
    return torch.cat(outs,0)

def ens_label(teachers, Xq_d, bs=256):
    ps=[]
    with torch.no_grad():
        for m in teachers:
            m.eval(); out=[]
            for i in range(0,Xq_d.shape[0],bs):
                xb=(Xq_d[i:i+bs]-fm.to(device))/fs.to(device)
                out.append((m(xb).squeeze(1)*ls.to(device)+lm.to(device)))
            ps.append(torch.cat(out))
    return torch.stack(ps).mean(0)  # [n] mass units on device

def train_student(Xin_d, soft_d, epochs, seed):
    torch.manual_seed(3000+seed)
    model=build_model(cfg).to(device)
    opt=torch.optim.AdamW(model.parameters(),lr=LR,weight_decay=WD)
    yd=((soft_d-lm.to(device))/ls.to(device)).unsqueeze(1); n=Xin_d.shape[0]
    fmd=fm.to(device); fsd=fs.to(device); spp=max(1,math.ceil(n/BS)); tot=epochs*spp; s=0
    ema=EMA(model,EMAD)
    for ep in range(epochs):
        model.train(); perm=torch.randperm(n,device=device)
        for i in range(0,n,BS):
            ii=perm[i:i+BS]
            if ii.numel()<2: continue
            xb=(Xin_d[ii]-fmd)/fsd
            for g in opt.param_groups:
                g['lr']=LR*((s+1)/(WARM*spp) if s<WARM*spp else 0.01+0.99*0.5*(1+math.cos(math.pi*(s-WARM*spp)/max(1,tot-WARM*spp))))
            opt.zero_grad(set_to_none=True); F.mse_loss(model(xb),yd[ii]).backward(); opt.step(); ema.update(model); s+=1
    ema.copy_to(model); return model

t0=time.time()
K=8
teachers=[train_base(Xpool,Ypool,seed) for seed in range(K)]
ens_h=torch.stack([predict(m,Xhold,fm,fs,lm,ls,device).squeeze(1) for m in teachers]).mean(0)
print(f'[holdout] ensemble K={K}: {metrics(ens_h,Yhold)}',flush=True)
Xpool_d=Xpool.to(device)
R=8
Xaug=aug_copies(Xpool_d,R,seedbase=1)
soft_orig=ens_label(teachers,Xpool_d)   # on originals
soft_aug=ens_label(teachers,Xaug)
# variant A: mimic all (orig+aug) with ensemble soft
XA=torch.cat([Xpool_d,Xaug],0); SA=torch.cat([soft_orig,soft_aug],0)
# variant B: orig true + aug soft
SB=torch.cat([Ypool.squeeze(1).to(device),soft_aug],0)
for name,(Xin,S) in {'mimic':(XA,SA),'hybrid':(XA,SB)}.items():
    rs=[]
    for sd in range(3):
        stu=train_student(Xin,S,epochs=300,seed=sd)
        m=metrics(predict(stu,Xhold,fm,fs,lm,ls,device).squeeze(1),Yhold); rs.append(m)
    rl=[r['rel_l2'] for r in rs]
    print(f'[holdout] distill-{name}: rl2 mean={np.mean(rl):.4f} per-seed={[round(x,4) for x in rl]} mae={np.mean([r["mae"] for r in rs]):.4f} max={np.mean([r["maxabs"] for r in rs]):.4f}',flush=True)
print(f'elapsed {time.time()-t0:.0f}s')
PY
nohup python exp8.py > exp8.log 2>&1 &
echo "started $!"

cd /app/work && sleep 290 && cat exp8.log

cd /app/work && sleep 150 && cat exp8.log; echo ---; ps aux|grep exp8|grep -v grep|head -1

cd /app/work && cat > final_build.py << 'PY'
import torch, math, numpy as np, time, argparse, os
import torch.nn.functional as F
from lib import *
from cv import train_fold
cfg=load_cfg(); device='cuda'
LR=1.5e-3; WD=2e-2; BS=32; WARM=15; EMAD=0.99

def train_base(X,Y,fm,fs,lm,ls,seed,epochs):
    m,final,_=train_fold(cfg,X,Y,fm,fs,lm,ls,device,epochs,BS,LR,WD,0.0,(False,False,False),0.0,0.0,seed,WARM,0.01,EMAD,None)
    m.load_state_dict(final); return m

def aug_copies(Xd,R,jitter,seedbase):
    outs=[]
    for c in range(R):
        g=torch.Generator(device=device).manual_seed(seedbase*131+c+1)
        xb=Xd.clone()
        for a in range(3):
            fl=(torch.rand(xb.shape[0],1,generator=g,device=device)<0.5).float()*(-2)+1
            xb[:,:,a]=xb[:,:,a]*fl
        xb=xb+(torch.randn(xb.shape,generator=g,device=device)*jitter).clamp(-0.025,0.025)
        outs.append(xb)
    return torch.cat(outs,0)

def ens_label(teachers,Xq_d,fm,fs,lm,ls,bs=256):
    ps=[]
    with torch.no_grad():
        for m in teachers:
            m.eval(); out=[]
            for i in range(0,Xq_d.shape[0],bs):
                xb=(Xq_d[i:i+bs]-fm.to(device))/fs.to(device)
                out.append(m(xb).squeeze(1)*ls.to(device)+lm.to(device))
            ps.append(torch.cat(out))
    return torch.stack(ps).mean(0)

def train_student(Xin_d,soft_d,fm,fs,lm,ls,epochs,seed):
    torch.manual_seed(3000+seed)
    model=build_model(cfg).to(device)
    opt=torch.optim.AdamW(model.parameters(),lr=LR,weight_decay=WD)
    yd=((soft_d-lm.to(device))/ls.to(device)).unsqueeze(1); n=Xin_d.shape[0]
    fmd=fm.to(device); fsd=fs.to(device); spp=max(1,math.ceil(n/BS)); tot=epochs*spp; s=0
    ema=EMA(model,EMAD)
    for ep in range(epochs):
        model.train(); perm=torch.randperm(n,device=device)
        for i in range(0,n,BS):
            ii=perm[i:i+BS]
            if ii.numel()<2: continue
            xb=(Xin_d[ii]-fmd)/fsd
            for g in opt.param_groups:
                g['lr']=LR*((s+1)/(WARM*spp) if s<WARM*spp else 0.01+0.99*0.5*(1+math.cos(math.pi*(s-WARM*spp)/max(1,tot-WARM*spp))))
            opt.zero_grad(set_to_none=True); F.mse_loss(model(xb),yd[ii]).backward(); opt.step(); ema.update(model); s+=1
    ema.copy_to(model)
    return model

@torch.no_grad()
def bn_recalc(model,Xorig_d,fm,fs,passes=4,bs=64):
    for m in model.modules():
        if isinstance(m,torch.nn.BatchNorm1d):
            m.reset_running_stats(); m.momentum=None
    model.train(); n=Xorig_d.shape[0]; fmd=fm.to(device); fsd=fs.to(device)
    for _ in range(passes):
        perm=torch.randperm(n)
        for i in range(0,n,bs):
            xb=(Xorig_d[perm[i:i+bs]]-fmd)/fsd
            if xb.shape[0]<2: continue
            model(xb)
    model.eval()

def build(Xpool,Ypool,K,R,t_ep,s_ep,jitter,s_seed,do_bn):
    fm,fs,lm,ls=compute_norm(Xpool,Ypool)
    teachers=[train_base(Xpool,Ypool,fm,fs,lm,ls,seed,t_ep) for seed in range(K)]
    Xpool_d=Xpool.to(device)
    Xaug=aug_copies(Xpool_d,R,jitter,seedbase=1)
    soft=torch.cat([ens_label(teachers,Xpool_d,fm,fs,lm,ls), ens_label(teachers,Xaug,fm,fs,lm,ls)],0)
    Xin=torch.cat([Xpool_d,Xaug],0)
    stu=train_student(Xin,soft,fm,fs,lm,ls,s_ep,s_seed)
    if do_bn: bn_recalc(stu,Xpool_d,fm,fs)
    return stu,(fm,fs,lm,ls),teachers

if __name__=='__main__':
    ap=argparse.ArgumentParser()
    ap.add_argument('--mode',default='validate')
    ap.add_argument('--K',type=int,default=10); ap.add_argument('--R',type=int,default=10)
    ap.add_argument('--tep',type=int,default=450); ap.add_argument('--sep',type=int,default=320)
    ap.add_argument('--jitter',type=float,default=0.008); ap.add_argument('--bn',type=int,default=1)
    ap.add_argument('--sseed',type=int,default=0)
    a=ap.parse_args()
    Xtr,Ytr=load_split('train'); Xval,Yval=load_split('val'); Xpt,Ypt=load_split('public_test')
    t0=time.time()
    if a.mode=='validate':
        Xpool,Ypool=Xtr,Ytr
        Xhold=torch.cat([Xval,Xpt],0); Yhold=torch.cat([Yval,Ypt],0).s
…[truncated]

cd /app/work && sleep 290 && cat fb_val.log; echo ---; ps aux|grep final_build|grep -v grep|head -1

cd /app/work && nohup python final_build.py --mode validate --K 8 --R 8 --tep 450 --sep 300 --bn 0 > fb_val_nobn.log 2>&1 &
echo "started $!"; sleep 200; cat fb_val_nobn.log

command: str_replace
path: /app/work/final_build.py
old_str:     else:
        Xpool=torch.cat([Xtr,Xval,Xpt],0); Ypool=torch.cat([Ytr,Yval,Ypt],0)
        stu,(fm,fs,lm,ls),teachers=build(Xpool,Ypool,a.K,a.R,a.tep,a.sep,a.jitter,a.sseed,a.bn)
        os.makedirs('/app/outputs',exist_ok=True)
new_str:     else:
        Xpool=torch.cat([Xtr,Xval,Xpt],0); Ypool=torch.cat([Ytr,Yval,Ypt],0)
        # Build teachers + augmented ensemble-labeled transfer set ONCE
        fm,fs,lm,ls=compute_norm(Xpool,Ypool)
        teachers=[train_base(Xpool,Ypool,fm,fs,lm,ls,seed,a.tep) for seed in range(a.K)]
        Xpool_d=Xpool.to(device)
        Xaug=aug_copies(Xpool_d,a.R,a.jitter,seedbase=1)
        soft_orig=ens_label(teachers,Xpool_d,fm,fs,lm,ls)
        soft=torch.cat([soft_orig, ens_label(teachers,Xaug,fm,fs,lm,ls)],0)
        Xin=torch.cat([Xpool_d,Xaug],0)
        print(f'[SUBMIT] teachers={a.K} transfer-set={Xin.shape[0]}',flush=True)
        # Train several students; pick the one most faithful to ensemble on canonical originals
        best=None; best_stu=None
        for sd in range(3):
            stu=train_student(Xin,soft,fm,fs,lm,ls,a.sep,sd)
            p=predict(stu,Xpool,fm,fs,lm,ls,device).squeeze(1).to(device)
            faith=F.mse_loss(p,soft_orig).item()
            print(f'  student seed{sd} faithfulness(MSE to ens)={faith:.5f}',flush=True)
            if best is None or faith<best:
                best=faith; best_stu=stu
        stu=best_stu
        os.makedirs('/app/outputs',exist_ok=True)

cd /app/work && nohup python final_build.py --mode submit --K 12 --R 10 --tep 450 --sep 320 --jitter 0.008 --bn 0 > fb_submit.log 2>&1 &
echo "started $!"

cd /app/work && cat > verify.py << 'PY'
import torch, json, sys
from pathlib import Path
sys.path.insert(0,'/app')
from candidate_model import build_model
from lib import load_cfg, load_split, metrics

cfg=load_cfg()
ckpt=torch.load('/app/outputs/model.pt', weights_only=True)
print('keys:', sorted(ckpt.keys()))
for k in ['feat_mean','feat_std','label_mean','label_std']:
    v=ckpt[k]; print(f'  {k}: shape={tuple(v.shape)} dtype={v.dtype} val={v.tolist()}')

# Reconstruct architecture EXACTLY and load
model=build_model(cfg)
missing,unexpected=model.load_state_dict(ckpt['state_dict'], strict=False)
print('missing keys:',missing)
print('unexpected keys:',unexpected)
model.eval()

fm=ckpt['feat_mean'].float(); fs=ckpt['feat_std'].float()
lm=ckpt['label_mean'].float(); ls=ckpt['label_std'].float()

@torch.no_grad()
def infer(X, bs=256):
    outs=[]
    for i in range(0,X.shape[0],bs):
        xb=(X[i:i+bs].float()-fm)/fs          # documented input standardization
        p=model(xb)                            # [B,1]
        outs.append(p*ls+lm)                   # documented de-standardization
    return torch.cat(outs,0)

# contract check
dummy=torch.randn(5,2048,3)
out=infer(dummy)
assert out.shape==(5,1), out.shape
print('contract OK: [5,2048,3] ->', tuple(out.shape))

for s in ['train','val','public_test']:
    X,Y=load_split(s)
    p=infer(X).squeeze(1)
    m=metrics(p,Y.squeeze(1))
    print(f'[{s:11s}] rel_l2={m["rel_l2"]:.4f} mae={m["mae"]:.4f} maxabs={m["maxabs"]:.4f}')
print('NOTE: train/val/public_test are IN-SAMPLE for the submitted model (trained on all 148).')
PY
echo "verify.py ready"; sleep 200; echo "=== submit log so far ==="; cat fb_submit.log

cd /app/work && sleep 180 && cat fb_submit.log; echo ---; ps aux|grep final_build|grep -v grep|head -1

cd /app/work && sleep 220 && cat fb_submit.log; echo ---; ls -la /app/outputs/ 2>/dev/null

cd /app/work && sleep 160 && cat fb_submit.log; echo ---; ls -la /app/outputs/ 2>/dev/null

cd /app/work && python verify.py

cd /app/work && nohup python final_build.py --mode validate --K 12 --R 10 --tep 450 --sep 320 --jitter 0.008 --bn 0 > fb_val_matched.log 2>&1 &
echo "started $!"; sleep 230; cat fb_val_matched.log

cd /app/work && nohup python final_build.py --mode validate --K 16 --R 12 --tep 450 --sep 340 --jitter 0.008 --bn 0 > fb_val_k16.log 2>&1 &
echo "started $!"; sleep 260; cat fb_val_k16.log

cd /app/work && sleep 60 && cat fb_val_k16.log

cd /app/work && cp /app/outputs/model.pt /app/work/model_k12_backup.pt && nohup python final_build.py --mode submit --K 16 --R 12 --tep 450 --sep 340 --jitter 0.008 --bn 0 > fb_submit2.log 2>&1 &
echo "started $! (backup of K12 model saved)"

cd /app/work && sleep 280 && cat fb_submit2.log

cd /app/work && sleep 220 && cat fb_submit2.log; echo ---; ls -la /app/outputs/

cd /app/work && sleep 120 && cat fb_submit2.log; echo ---; ls -la /app/outputs/model.pt

cd /app/work && python verify.py

cd /app/work && python - << 'PY'
import torch, sys
sys.path.insert(0,'/app')
from candidate_model import build_model
from lib import load_cfg, load_split, metrics
cfg=load_cfg()
ck=torch.load('/app/outputs/model.pt', weights_only=True)
model=build_model(cfg); model.load_state_dict(ck['state_dict']); model.eval()
fm,fs,lm,ls=[ck[k].float() for k in ['feat_mean','feat_std','label_mean','label_std']]
X,Y=load_split('val')
@torch.no_grad()
def infer(X):
    return (model((X.float()-fm)/fs)*ls+lm).squeeze(1)
base=infer(X)
# (a) float16 input path
Xh=X.half()
with torch.no_grad():
    ph=(model((Xh.float()-fm)/fs)*ls+lm).squeeze(1)
print('float16-input max abs diff vs float32:', (ph-base).abs().max().item())
# (b) permutation invariance
perm=torch.randperm(2048)
pp=infer(X[:,perm,:])
print('permutation max abs diff:', (pp-base).abs().max().item())
# (c) single-sample batch
s1=infer(X[:1]); print('single-sample vs batched diff:', (s1-base[:1]).abs().max().item())
# (d) model params dtype
print('param dtype:', next(model.parameters()).dtype)
print('val rel_l2 (in-sample sanity):', round(metrics(base,Y.squeeze(1))['rel_l2'],4))
print('ALL ROBUSTNESS CHECKS PASSED')
PY

cd /app/work && pkill -f final_build.py 2>/dev/null; pkill -f "exp[0-9]" 2>/dev/null; pkill -f pipeline.py 2>/dev/null; sleep 2
echo "=== running python jobs (should be none) ==="; ps aux|grep -E "python (exp|final|pipeline|sweep)"|grep -v grep|head
echo "=== GPU ==="; nvidia-smi --query-gpu=memory.used --format=csv,noheader
echo "=== deliverable ==="; ls -la /app/outputs/model.pt
python -c "import torch; d=torch.load('/app/outputs/model.pt',weights_only=True); print('keys:',sorted(d)); print({k:tuple(d[k].shape) for k in ['feat_mean','feat_std','label_mean','label_std']}); print('n state_dict tensors:',len(d['state_dict']))"

ps aux | grep python | grep -v grep | head; echo "--- my python procs ---"; pgrep -af python | grep -iE "exp|final|pipeline|sweep|train" || echo "none of mine running"

cd /app/work && python - << 'PY'
import torch, sys
sys.path.insert(0,'/app')
from candidate_model import build_model
from lib import load_cfg, load_split, metrics
cfg=load_cfg()
ck=torch.load('/app/outputs/model.pt', weights_only=True)
m=build_model(cfg); m.load_state_dict(ck['state_dict']); m.eval()
fm,fs,lm,ls=[ck[k].float() for k in ['feat_mean','feat_std','label_mean','label_std']]
# Exact grader-style inference: standardize coords -> forward -> de-standardize output
X,Y=load_split('public_test')
with torch.no_grad():
    z = m((X.float()-fm)/fs)          # standardized-space prediction, [B,1]
    mass = z*ls + lm                  # de-standardize to kg
print('sample preds (kg):', [round(v,3) for v in mass.squeeze(1)[:6].tolist()])
print('sample truth (kg):', [round(v,3) for v in Y.squeeze(1)[:6].tolist()])
print('output shape:', tuple(mass.shape), '| positivity: min pred =', round(mass.min().item(),3))
print('\nFINAL DELIVERABLE: /app/outputs/model.pt  (ready)')
print('Honest out-of-sample estimate (train on 100, test on unseen 48):')
print('  rel_L2 ~ 0.239 | MAE ~ 0.151 kg | worst-case abs err ~ 0.65 kg')
PY