"""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)) # Predict withheld observations conditional on known H and current evidence. 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()