analysis/bootstrap.py

"""Paired bootstrap with prompts as sampling units and independent model seeds."""
import numpy as np
from eval.metrics import axis_metrics,corr,matching

def paired_cluster_mean(a,b,prompt_ids,n=10000,seed=191):
    """Bootstrap an example mean while resampling whole prompt clusters.

    Every prompt has equal sampling probability; every observation retains
    the weight it has in the reported metric. Seeds are paired across methods.
    Unlike an equal-prompt mean, this matches the overall leakage mean when
    different prompts have different numbers of accepted interventions.
    """
    a=np.asarray(a,float);b=np.asarray(b,float);pids=np.asarray(prompt_ids)
    if a.shape!=b.shape or a.ndim!=2 or a.shape[1]!=len(pids):
        raise ValueError('Expected aligned [seed, example] arrays and prompt IDs')
    if not np.isfinite(a).all() or not np.isfinite(b).all():
        raise ValueError('Missing predictions must not silently change the estimand')
    if not a.size:raise ValueError('Empty evaluation')
    unique,inverse=np.unique(pids,return_inverse=True);counts=np.bincount(inverse)
    diff=np.zeros((len(unique),len(a)))
    for s in range(len(a)):np.add.at(diff[:,s],inverse,b[s]-a[s])
    g=np.random.default_rng(seed);boot=[]
    for _ in range(n):
        si=g.integers(len(a),size=len(a));pi=g.integers(len(unique),size=len(unique))
        boot.append(float(diff[pi][:,si].sum()/(counts[pi].sum()*len(a))))
    return {'difference':float((b-a).mean()),'ci95':np.quantile(boot,[.025,.975]).tolist() if len(unique)>1 else None,'bootstrap_samples':n,'prompts':len(unique),'examples':a.shape[1],'units':'paired prompt clusters and paired training seeds','prompt_weighting':'uniform cluster sampling; example-weighted metric','random_seed':seed}

def paired_prompt_mean(a,b,prompt_ids,n=10000,seed=190):
    """a,b [seed,example], nan allowed; each prompt gets equal weight."""
    valid=np.isfinite(a).any(0)&np.isfinite(b).any(0);a=a[:,valid];b=b[:,valid];prompt_ids=np.asarray(prompt_ids)[valid]
    unique,inverse=np.unique(prompt_ids,return_inverse=True);diff=np.empty((len(a),len(unique)))
    for j in range(len(unique)):diff[:,j]=np.nanmean(b[:,inverse==j]-a[:,inverse==j],axis=1)
    g=np.random.default_rng(seed);boots=[]
    for _ in range(n):
        si=g.integers(len(a),size=len(a));pi=g.integers(len(unique),size=len(unique));boots.append(np.nanmean(diff[si][:,pi]))
    return {'difference':float(np.nanmean(diff)),'ci95':np.quantile(boots,[.025,.975]).tolist(),'bootstrap_samples':n,'prompts':len(unique),'units':'paired prompt resampling and training-seed resampling','prompt_weighting':'equal'}

def semantic_bootstrap(a,b,attributes,prompt_ids,n=1000,seed=193):
    """Exact rank/pair statistic with shared prompt draws, conditional on fitted seeds.

    Caller must verify response identities and paired run seeds. Axis purity
    rematches the baseline on each draw; named ILR axes keep their identity.
    """
    a=np.asarray(a);b=np.asarray(b);attributes=np.asarray(attributes);prompt_ids=np.asarray(prompt_ids)
    if a.ndim!=3 or b.ndim!=3 or a.shape[:2]!=b.shape[:2] or a.shape[0]<2:
        raise ValueError('Need at least two aligned [seed,response,axis] model arrays')
    if attributes.shape!=(a.shape[1],4) or prompt_ids.shape!=(a.shape[1],) or min(a.shape[2],b.shape[2])<4:
        raise ValueError('Expected four human attributes and matching response/prompt axes')
    if not all(np.isfinite(v).all() for v in [a,b,attributes]):raise ValueError('Nonfinite semantic input')
    pairs=[(i,j) for i in range(len(a)) for j in range(i+1,len(a))]
    def statistic(ix):
        if any(np.any(np.std(v,axis=-2)<1e-12) for v in [a[:,ix],b[:,ix],attributes[ix]]):return np.array([np.nan,np.nan])
        pd=[];asd=[]
        for s in range(len(a)):
            ma,_=axis_metrics(a[s][ix],attributes[ix],rank=True,named=False);mb,_=axis_metrics(b[s][ix],attributes[ix],rank=True,named=True);pd.append(mb['axis_purity']-ma['axis_purity'])
        for i,j in pairs:
            ca=corr(a[i][ix],a[j][ix]);order,_=matching(a[i][ix],a[j][ix]);sa=abs(ca[order,np.arange(a[j].shape[1])]).mean()
            sb=np.diag(corr(b[i][ix,:4],b[j][ix,:4])).mean();asd.append(sb-sa)
        return np.array([np.mean(pd),np.mean(asd)])
    unique=np.unique(prompt_ids);groups=[np.flatnonzero(prompt_ids==p) for p in unique];g=np.random.default_rng(seed)
    point=statistic(np.arange(a.shape[1]));boots=[]
    for _ in range(n):
        ix=np.concatenate([groups[i] for i in g.integers(len(groups),size=len(groups))]);boots.append(statistic(ix))
    boots=np.array(boots);out={}
    for j,name in enumerate(['axis_purity','axis_stability']):
        v=boots[:,j];undefined=int((~np.isfinite(v)).sum());ok=not undefined and np.isfinite(point[j]) and len(unique)>1
        out[name]={'difference':float(point[j]) if np.isfinite(point[j]) else None,'ci95':np.quantile(v,[.025,.975]).tolist() if ok else None,'bootstrap_mean':float(v.mean()) if not undefined else None,'bootstrap_samples':n,'bootstrap_undefined_samples':undefined,'interval_status':'ok' if ok else 'degenerate_or_insufficient_clusters','prompts':len(unique),'training_seeds':len(a),'distinct_seed_pairs':len(pairs),'random_seed':seed,'units':'prompt clusters conditional on supplied trained seeds; exact rank recomputation; seed SD reported separately'}
    return out