"""Paid recorded-feature acquisition with exact Gaussian rank-interval value.""" from __future__ import annotations import hashlib,json,time import torch from . import factorial,interval_value from .factor_study import ROOT,PROTOCOL,verify,write class Evidence: def __init__(self,values):self._values=values;self.visible=values.new_full(values.shape,float('nan'));self.observed=[[] for _ in values];self.query_log=[] def reveal(self,rows,q): for row in rows.tolist(): assert q not in self.observed[row];self.observed[row].append(q) self.visible[rows,q]=self._values[rows,q] self.query_log.append(dict(rows=rows.tolist(),feature=q)) @torch.no_grad() def episode(model,values,policy,seed,budgets=(2,4,8),cost=.01): evidence=Evidence(values);n,d=values.shape;device=values.device stopped=torch.zeros(n,device=device,dtype=torch.bool);counts=torch.zeros(n,device=device,dtype=torch.long) snapshots={};cache={};start=time.perf_counter();value_evaluations=0 generator=torch.Generator(device=device).manual_seed(seed) static_means=model['mean'];static=(static_means.var(0,unbiased=False)/model['cov'].diagonal()).argsort(descending=True,stable=True) for step in range(1,max(budgets)+1): groups={} for row in range(n): if not bool(stopped[row]):groups.setdefault(tuple(sorted(evidence.observed[row])),[]).append(row) for key,ids in groups.items(): rows=torch.tensor(ids,device=device);visible=evidence.visible[rows] probability=factorial.posterior(model,visible,key) current=probability.topk(min(3,probability.shape[-1]),-1).values.sum(-1) if key not in cache:cache[key]=factorial.conditional_channels(model,key) means,variance=cache[key];available=[q for q in range(d) if q not in key] if not available:stopped[rows]=True;continue if policy.startswith('value'): expected=[] for q in available: expected.append(interval_value.expected_topk(probability,means[:,q],variance[q],3));value_evaluations+=len(rows) future,which=torch.stack(expected,-1).max(-1) chosen=rows.new_tensor(available)[which] stop=(future-current<=cost) if policy.endswith('stop') else torch.zeros_like(current,dtype=torch.bool) elif policy.startswith('variance'): mu=means[:,available];average=probability@mu score=((probability@mu.square())-average.square())/variance[available] chosen=rows.new_tensor(available)[score.argmax(-1)];stop=torch.zeros_like(current,dtype=torch.bool) if policy.endswith('stop'): for q in chosen.unique().tolist(): selected=chosen==q;future=interval_value.expected_topk(probability[selected],means[:,q],variance[q],3) stop[selected]=future-current[selected]<=cost;value_evaluations+=int(selected.sum()) elif policy=='static_forced': q=next(q for q in static.tolist() if q in available);chosen=rows.new_full((len(rows),),q);stop=torch.zeros_like(current,dtype=torch.bool) else: chosen=rows.new_tensor(available)[torch.randint(len(available),(len(rows),),device=device,generator=generator)];stop=torch.zeros_like(current,dtype=torch.bool) stopped[rows[stop]]=True for q in chosen[~stop].unique().tolist(): selected=rows[(chosen==q)&~stop];evidence.reveal(selected,q);counts[selected]+=1 if step in budgets: probs=values.new_zeros((n,len(model['prior']))) groups={} for row in range(n):groups.setdefault(tuple(sorted(evidence.observed[row])),[]).append(row) for key,ids in groups.items(): rows=torch.tensor(ids,device=device);probs[rows]=factorial.posterior(model,evidence.visible[rows],key) snapshots[step]=dict(probability=probs,queries=counts.clone()) torch.cuda.synchronize();return snapshots,dict(seconds=time.perf_counter()-start,queried_features=evidence.query_log, conditional_masks_materialized=len(cache),exact_value_evaluations=value_evaluations, future_observation_draws=0,withheld_values_never_read_by_inference=True) def main(device='cuda:1'): verify();start=time.perf_counter();torch.cuda.set_device(device);torch.backends.cuda.matmul.allow_tf32=False config=json.loads(PROTOCOL.read_text());manifest=json.loads((ROOT/'prepared/manifest.json').read_text());settings=config['secondary_acquisition'];all_results=[] for series in manifest['series']: if not series['primary_eligible']:continue p=torch.load(ROOT/'prepared'/series['file'],map_location=device,weights_only=True) rows=torch.where(p['masks']['unseen'])[0].tolist() rows=sorted(rows,key=lambda i:hashlib.sha256(('20781:'+series['records'][i]['id']).encode()).hexdigest())[:32] selected=torch.tensor(rows,device=device);sizes=[len(v) for v in series['levels']];true=factorial.joint_index(p['h'][selected],sizes) for name in settings['world_models']: path=ROOT/'models'/series['name']/f'n64_s{config["sample_seeds"][0]}'/f'{name}.pt' checkpoint=torch.load(path,map_location=device,weights_only=True);model=checkpoint['model'];x=(p['x'][selected]-checkpoint['mean'])/checkpoint['scale'] for policy in settings['policies']: outputs,trace=episode(model,x,policy,20801,tuple(settings['budgets']),settings['per_scalar_recorded_query_cost']) for budget,out in outputs.items(): credit=factorial.topk_credit(out['probability'],true,3) all_results.append(dict(series=series['name'],system=series['system'],model=name,policy=policy,budget=budget, selected_rows=selected,groups=p['groups'][selected],top3=credit,queries=out['queries'], net_utility=credit-settings['per_scalar_recorded_query_cost']*out['queries'],trace=trace)) print(json.dumps(dict(series=series['name'],model=name,policy=policy,seconds=trace['seconds'])),flush=True) torch.save(dict(results=all_results),ROOT/'acquisition.pt');write('acquisition_execution.json',dict(seconds=time.perf_counter()-start, peak_gpu_bytes=torch.cuda.max_memory_allocated(device),cells=len(all_results),subset_fixed_before_predictions=True, complete_primary_diagnostic_cohort_retained=True,no_globally_optimal_multistep_claim=True)) verify();print('FULL-HYPOTHESIS ACQUISITION COMPLETE',flush=True) if __name__=='__main__':main()