| """Preregister, acquire, then evaluate every eligible external cohort.""" |
| from __future__ import annotations |
|
|
| import argparse |
| from concurrent.futures import ThreadPoolExecutor |
| import csv |
| import gzip |
| import json |
| from pathlib import Path |
| import time |
|
|
| import torch |
|
|
| from unionmind.decision_sufficiency.core import ( |
| calibration, count_samples, dispersion, encode, estimate, plan, push, realized, |
| ) |
| from .data import ROOT, PROTOCOL, acquire, prepare, sha, write |
| from .runtime import Timer, replay, select_samples |
|
|
|
|
| def lock(): |
| ROOT.mkdir(parents=True,exist_ok=True) |
| assert not (ROOT/'preconfirmation_lock.json').exists() |
| old = Path('artifacts/decision_sufficiency_01') |
| files = json.loads((old/'preservation_lock.json').read_text())['files'] |
| for p,digest in json.loads((old/'final_snapshot.json').read_text())['files'].items(): |
| assert p not in files or files[p] == digest |
| files[p] = digest |
| for name in ('final_snapshot.json','finalize.log'): |
| path = old/name |
| if path.exists(): files[str(path)] = sha(path) |
| def changed(item): |
| p,d = item |
| return p if not Path(p).exists() or sha(p) != d else None |
| with ThreadPoolExecutor(max_workers=8) as pool: |
| errors = [p for p in pool.map(changed,files.items()) if p] |
| assert not errors, errors |
| write(ROOT/'preservation_lock.json',dict(files=files)) |
| code = list(Path('unionmind/external_sufficiency').glob('*.py'))+[ |
| PROTOCOL,Path('tests/test_external_sufficiency.py'), |
| Path('unionmind/decision_sufficiency/core.py')] |
| write(ROOT/'preconfirmation_lock.json',dict(files={str(p):sha(p) for p in code}, |
| locked_before_any_source_download=True,registered_threshold_changes_permitted=False, |
| archived_results_protected=len(files),created_unix=time.time())) |
| print(json.dumps({'locked_files':len(code),'prior_files':len(files)}),flush=True) |
|
|
|
|
| def verify_lock(): |
| locked = json.loads((ROOT/'preconfirmation_lock.json').read_text()) |
| changed = [p for p,d in locked['files'].items() if sha(p) != d] |
| assert not changed, changed |
|
|
|
|
| def predictive_scores(model, truth): |
| scores = calibration(model,truth) |
| mass = model.sum((2,3)) |
| conditional = model/mass[...,None,None] |
| scores['joint_conditional_log_loss'] = -(truth*conditional.clamp_min(1e-300).log()).sum((1,2,3)) |
| |
| given_h = model/model.sum(3,keepdim=True).clamp_min(1e-300) |
| scores['observation_given_h_log_loss'] = -(truth*given_h.clamp_min(1e-300).log()).sum((1,2,3)) |
| return {k:float(v[0]) for k,v in scores.items()} |
|
|
|
|
| def evaluate(protocol): |
| verify_lock() |
| manifest = json.loads((ROOT/'prepared/manifest.json').read_text()) |
| out = ROOT/'evaluation' |
| out.mkdir(exist_ok=True) |
| assert not (out/'metrics.csv.gz').exists() |
| device = protocol['device'] |
| fieldnames = ['task_index','dataset','environment','task','regime','sample_seed','sample_count', |
| 'model','representation','budget','cost','utility','queries','abstention','estimated_utility', |
| 'fitted_compression_loss','fitted_dispersion','empirical_hindsight_utility', |
| 'empirical_compression_loss','empirical_model_loss','posterior_brier','posterior_log_loss', |
| 'posterior_absolute_error','joint_conditional_log_loss','observation_given_h_log_loss', |
| 'raw_posterior_brier','raw_posterior_log_loss','raw_joint_conditional_log_loss', |
| 'construction_seconds','evaluation_records'] |
| rows = 0 |
| maximum_replay_error = 0. |
| maximum_decomposition_error = 0. |
| primary_trajectory_records = 0 |
| skipped = [] |
| start = time.perf_counter() |
| torch.cuda.reset_peak_memory_stats(device) |
| with gzip.open(out/'metrics.csv.gz','wt',newline='') as f: |
| writer = csv.DictWriter(f,fieldnames=fieldnames) |
| writer.writeheader() |
| for task in manifest['tasks']: |
| archive = torch.load(ROOT/'prepared'/task['file'],map_location=device,weights_only=True) |
| held = archive['evaluation'] |
| truth_counts = count_samples(held[None],6) |
| truth = truth_counts/truth_counts.sum() |
| hindsight = plan(truth)['expected'] |
| saved = {} |
| for regime in protocol['regimes']: |
| for seed in protocol['adapter']['sample_seeds']: |
| samples = select_samples(archive,regime,seed) |
| if samples is None: |
| skipped.append(dict(task_index=task['index'],regime=regime,seed=seed, |
| reason='fewer_than_256_target_adaptation_records')) |
| continue |
| for model_name in protocol['estimation_models']: |
| with Timer(device) as fitted: |
| model = estimate(count_samples(samples[None],6),model_name) |
| raw_scores = predictive_scores(model,truth) |
| full_fitted = plan(model)['expected'] |
| for method in protocol['representations']: |
| with Timer(device) as constructed: |
| code,k = encode(model,method) |
| compact = push(model,code,k) |
| policy = plan(compact) |
| truth_code = push(truth,code,k) |
| scores = predictive_scores(compact,truth_code) |
| actual = realized(truth_code,policy) |
| empirical_code_hindsight = plan(truth_code)['expected'] |
| fitted_loss = full_fitted-policy['expected'] |
| empirical_compression = hindsight-empirical_code_hindsight |
| empirical_model = empirical_code_hindsight-actual['utility'] |
| maxerr = (hindsight-actual['utility']-empirical_compression-empirical_model).abs().max() |
| maximum_decomposition_error = max(maximum_decomposition_error,float(maxerr)) |
| d = float(dispersion(model,code,compact)[0]) |
| primary = regime == protocol['primary']['regime'] and ( |
| model_name == protocol['primary']['model'] or method == 'identity') |
| if primary: |
| trace = [] |
| for budget in protocol['planner']['query_budgets']: |
| for ci,cost in enumerate(protocol['planner']['query_costs']): |
| action,queries,seen = replay(policy,code.to(torch.uint8), |
| (held//4)%64,held%4,budget,ci) |
| utility = torch.where(action == 2,0., |
| torch.where(action == held//256,1.,-2.)).double()-cost*queries.double() |
| error = abs(float(utility.mean())-float(actual['utility'][0,budget,ci])) |
| maximum_replay_error = max(maximum_replay_error,error) |
| assert error < 1e-12 |
| assert bool((queries <= budget).all()) |
| assert torch.equal(((seen&1)>0).long()+((seen&2)>0).long(),queries) |
| trace.append(torch.stack((action,queries,seen),1).to(torch.uint8)) |
| primary_trajectory_records += len(held) |
| saved[f'{seed}:{model_name}:{method}'] = dict( |
| samples=samples,model=model,code=code.to(torch.uint8), |
| compact=compact,choices=policy['choices'],values=policy['values'], |
| expected=policy['expected'],traces=torch.stack(trace)) |
| for budget in protocol['planner']['query_budgets']: |
| for ci,cost in enumerate(protocol['planner']['query_costs']): |
| writer.writerow(dict(task_index=task['index'],dataset=task['dataset'], |
| environment=task['environment'],task=task['task'],regime=regime, |
| sample_seed=seed,sample_count=len(samples),model=model_name, |
| representation=method,budget=budget,cost=cost, |
| utility=float(actual['utility'][0,budget,ci]), |
| queries=float(actual['queries'][0,budget,ci]), |
| abstention=float(actual['abstention'][0,budget,ci]), |
| estimated_utility=float(policy['expected'][0,budget,ci]), |
| fitted_compression_loss=float(fitted_loss[0,budget,ci]), |
| fitted_dispersion=d,empirical_hindsight_utility=float(hindsight[0,budget,ci]), |
| empirical_compression_loss=float(empirical_compression[0,budget,ci]), |
| empirical_model_loss=float(empirical_model[0,budget,ci]), |
| **scores,raw_posterior_brier=raw_scores['posterior_brier'], |
| raw_posterior_log_loss=raw_scores['posterior_log_loss'], |
| raw_joint_conditional_log_loss=raw_scores['joint_conditional_log_loss'], |
| construction_seconds=fitted.seconds+constructed.seconds, |
| evaluation_records=len(held))) |
| rows += 1 |
| torch.save(saved,out/f'primary_task_{task["index"]:03}.pt') |
| print(json.dumps({'evaluated_task':task['index'],'dataset':task['dataset'], |
| 'environment':task['environment'],'rows_so_far':rows, |
| 'elapsed_seconds':time.perf_counter()-start}),flush=True) |
| verify_lock() |
| write(out/'summary.json',dict(metric_rows=rows,primary_trajectory_records=primary_trajectory_records, |
| maximum_replay_error=maximum_replay_error,maximum_decomposition_error=maximum_decomposition_error, |
| skipped_secondary_regimes=skipped,evaluation_seconds=time.perf_counter()-start, |
| peak_gpu_allocated_bytes=torch.cuda.max_memory_allocated(device), |
| peak_gpu_reserved_bytes=torch.cuda.max_memory_reserved(device), |
| empirical_hindsight_oracle_is_not_a_population_oracle=True, |
| primary_selection_changed=False)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('command',choices=['lock','acquire','prepare','evaluate']) |
| args = parser.parse_args() |
| if args.command == 'lock': return lock() |
| verify_lock() |
| protocol = json.loads(PROTOCOL.read_text()) |
| if args.command == 'acquire': return acquire(protocol) |
| device = protocol['device'] |
| torch.cuda.set_device(device) |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
| if args.command == 'prepare': return prepare(protocol) |
| return evaluate(protocol) |
|
|
|
|
| if __name__ == '__main__': main() |
|
|
