| """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() |
|
|
