Digital Drosophila Archive
88 lines
5.2 KiB
Python
88 lines
5.2 KiB
Python
"""Fixed Brian/GPU stress suite, with source-checked immutable inputs."""
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import time
|
|
import numpy as np
|
|
from context import ROOT,MODEL,EXP,CFG,save,check
|
|
|
|
def load(path):
|
|
with np.load(path,allow_pickle=False) as f:return {k:f[k].copy() for k in f.files}
|
|
|
|
def compare(a,b,case,precision):
|
|
checks={k:bool(np.array_equal(a[k],b[k])) for k in ('step','neuron_idx','counts','ref')}
|
|
ah=np.asarray(a['history'],np.float64);bh=np.asarray(b['history'],np.float64)
|
|
assert ah.shape==bh.shape
|
|
delta=np.abs(ah[:,:2]-bh[:,:2]);limit=CFG['max_abs_volt'][precision]
|
|
checks['history_ref_exact']=bool(np.array_equal(ah[:,2],bh[:,2]))
|
|
checks['finite_states']=bool(np.isfinite(ah).all() and np.isfinite(bh).all())
|
|
checks['state_tolerance']=bool(delta.max()<=limit)
|
|
checks['weights_exact']=bool(np.array_equal(a['weights'],b['weights']))
|
|
# Distinguish logical event requests from realised requests after refractory.
|
|
for key in ('request_attempted','request_accepted','request_rejected'):
|
|
checks[key+'_exact']=bool(key in a and key in b and np.array_equal(a[key],b[key]))
|
|
data=load(case/'arrays.npz');targets=data['dst']
|
|
def arrival_sequence(result):
|
|
tick=result['arrival_step'];eid=result['arrival_edge_index'];guard=result['arrival_accepted']
|
|
assert tick.shape==eid.shape==guard.shape
|
|
# Different targets can interleave; order of writes to one target cannot.
|
|
order=np.lexsort((np.arange(len(eid)),targets[eid],tick))
|
|
return np.stack((tick[order],eid[order],guard[order]),axis=1)
|
|
aa=arrival_sequence(a);bb=arrival_sequence(b)
|
|
checks['arrival_order_and_guard_exact']=bool(np.array_equal(aa,bb))
|
|
first=None
|
|
for t in range(len(ah)):
|
|
if not np.array_equal(a['neuron_idx'][a['step']==t],b['neuron_idx'][b['step']==t]):first=t;break
|
|
bits_equal=a['history'].dtype==b['history'].dtype and np.array_equal(np.ascontiguousarray(a['history']).view(np.uint8),np.ascontiguousarray(b['history']).view(np.uint8))
|
|
return dict(pass_all=all(checks.values()),checks=checks,bitwise_history=bool(bits_equal),
|
|
max_abs_volt=float(delta.max()),mean_abs_volt=float(delta.mean()),fixed_limit_volt=limit,
|
|
first_event_divergence_step=first,oracle_events=int(len(a['step'])),gpu_events=int(len(b['step'])),
|
|
due_edges=int(len(aa)),accepted_edges=int(a['arrival_accepted'].sum()),
|
|
attempted_requests=int(a['request_attempted'].sum()),accepted_requests=int(a['request_accepted'].sum()))
|
|
|
|
def oracle():
|
|
from oracle_stress import run_case
|
|
t0=time.perf_counter();rows=[]
|
|
for case in sorted((EXP/'cases').iterdir()):
|
|
for precision in CFG['precisions']:
|
|
output=EXP/'oracle'/(case.name+'-'+precision)
|
|
run_case(case,output,precision)
|
|
r=load(output/'state.npz');assert np.isfinite(r['history']).all()
|
|
rows.append(dict(case=case.name,precision=precision,events=len(r['step'])))
|
|
print(json.dumps(rows[-1]),flush=True)
|
|
save(EXP/'regression/oracle-complete.json',dict(cases=rows,wall_seconds=time.perf_counter()-t0))
|
|
|
|
def gpu():
|
|
assert (EXP/'regression/oracle-complete.json').exists(),'No complete oracle suite'
|
|
os.environ['CUDA_VISIBLE_DEVICES']=CFG['gpu_uuid'];os.environ['CUDA_HOME']='/usr/local/cuda'
|
|
os.environ['PATH']=str(ROOT/'models/male-cns-lif/.venv/bin')+':/usr/local/cuda/bin:'+os.environ['PATH']
|
|
os.environ['MAX_JOBS']='2';os.environ['TORCH_CUDA_ARCH_LIST']='8.9'
|
|
sys.path.insert(0,str(ROOT/'models/male-cns-lif/terminal-scope-diagnostic-v1'))
|
|
from diagnostic import queues
|
|
queues()
|
|
from gpu_backend import run_gpu
|
|
t0=time.perf_counter();rows=[]
|
|
for case in sorted((EXP/'cases').iterdir()):
|
|
for precision in CFG['precisions']:
|
|
queues();refdir=EXP/'oracle'/(case.name+'-'+precision)
|
|
reference=load(refdir/'state.npz');coeff=json.loads((refdir/'coefficients.json').read_text());first=None
|
|
for repeat in range(CFG['gpu_repeats']):
|
|
out=EXP/'gpu'/(case.name+'-'+precision+f'-r{repeat}')
|
|
run_gpu(case,out,precision,coeff,EXP/'extension-build',weights=reference['weights'])
|
|
result=load(out/'state.npz');cmp=compare(reference,result,case,precision)
|
|
if first is None:first=result
|
|
rep={k:bool(np.array_equal(v,result[k])) for k,v in first.items()}
|
|
row=dict(case=case.name,precision=precision,repeat=repeat,comparison=cmp,repeat_arrays=rep,
|
|
pass_all=cmp['pass_all'] and all(rep.values()))
|
|
save(out/'comparison.json',row);rows.append(row)
|
|
print(json.dumps(dict(case=case.name,precision=precision,repeat=repeat,pass_all=row['pass_all'],
|
|
error_volt=cmp['max_abs_volt'],events=cmp['oracle_events'],bitwise_history=cmp['bitwise_history'])),flush=True)
|
|
save(EXP/'regression/gpu-complete.json',dict(pass_all=all(x['pass_all'] for x in rows),trials=rows,
|
|
wall_seconds=time.perf_counter()-t0,whole_brain_tested=False,learning_achieved=False))
|
|
|
|
if __name__=='__main__':
|
|
p=argparse.ArgumentParser();p.add_argument('mode',choices=['oracle','gpu']);args=p.parse_args()
|
|
check();globals()[args.mode]()
|