Digital Drosophila Archive
103 lines
6.0 KiB
Python
103 lines
6.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Mechanosensory runtime mode matrix and optional state-window forensics."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pyarrow.parquet as pq
|
|
import torch
|
|
import yaml
|
|
|
|
from phase4o_a import EXP,FROZEN,ROOT,sha,first_event_diff,save
|
|
|
|
NAME='mechanosensory-50hz'
|
|
SEED=42003
|
|
|
|
|
|
def run(mode: str, focus_step: int, focus_neuron: int):
|
|
from brain_runtime import BrainRuntime
|
|
from phase4i import poisson_batch
|
|
cfgpath=ROOT/'models/male-cns-lif/mb-mechanistic-v1/phase4i.yaml'
|
|
if mode=='deterministic-flags':
|
|
torch.use_deterministic_algorithms(True)
|
|
torch.backends.cudnn.deterministic=True
|
|
torch.backends.cudnn.benchmark=False
|
|
torch.backends.cudnn.allow_tf32=False
|
|
torch.backends.cuda.matmul.allow_tf32=False
|
|
torch.set_float32_matmul_precision('highest')
|
|
rt=BrainRuntime(cfgpath,EXP/'compiler-reports/extension-build')
|
|
target=np.flatnonzero(rt.population_masks['mechanosensory']).astype(np.int64)
|
|
refroot=FROZEN/NAME
|
|
ref=pq.read_table(refroot/'spike-events.parquet')
|
|
rstep=ref['step'].to_numpy().astype(np.int32);rnode=ref['neuron_idx'].to_numpy().astype(np.int32)
|
|
rcount=np.load(refroot/'per-neuron-spike-counts.int32.npy')
|
|
rv=np.load(refroot/'final-voltage.float32.npy');rg=np.load(refroot/'final-g.float32.npy')
|
|
rt.reset(SEED)
|
|
batch=poisson_batch(target,50.,1000.,NAME)
|
|
backend={'python_module':str(Path(rt.runtime_module.__file__)), 'python_sha256':sha(Path(rt.runtime_module.__file__)),
|
|
'extension':str(Path(rt.extension.__file__)), 'extension_sha256':sha(Path(rt.extension.__file__))}
|
|
window=[]
|
|
if mode in ('logging-off','logging-on'):
|
|
p4cfg=yaml.safe_load((ROOT/'models/male-cns-lif/runtime-v1/phase4c.yaml').read_text())
|
|
direct=rt.runtime_module.Runtime(p4cfg,rt.table,rt.indptr_np,rt.indices_np,rt.raw_np,rt.device,rt.extension)
|
|
out=direct.run('hybrid-exact',target,50.,SEED,1000.,capture_events=(mode=='logging-on'))
|
|
step=out['event_steps'];node=out['event_nodes'];counts=out['counts'];v=out['v'];g=out['g']
|
|
elif mode in ('instrumentation-on','single-stream-sync'):
|
|
for t in range(10000):
|
|
oldlen=len(rt.spike_steps)
|
|
if mode=='instrumentation-on' and focus_step-6<=t<=focus_step+5 and focus_neuron>=0:
|
|
slot=t%rt.queue_len
|
|
window.append({'step':t,'when':'before','neuron_idx':focus_neuron,
|
|
'voltage':float(rt.v[focus_neuron].item()),'g':float(rt.g[focus_neuron].item()),
|
|
'refractory':int(rt.refractory[focus_neuron].item()),
|
|
'arrival_buffer':float(rt.delay_queue[slot,focus_neuron].item()),
|
|
'ring_buffer_sha256':hashlib.sha256(rt.delay_queue.cpu().numpy().tobytes()).hexdigest()})
|
|
rt.step(batch,0.1)
|
|
torch.cuda.synchronize()
|
|
if mode=='instrumentation-on' and focus_step-6<=t<=focus_step+5 and focus_neuron>=0:
|
|
new_nodes=np.concatenate(rt.spike_nodes[oldlen:]) if len(rt.spike_nodes)>oldlen else np.empty(0,np.int32)
|
|
window.append({'step':t,'when':'after','neuron_idx':focus_neuron,
|
|
'voltage':float(rt.v[focus_neuron].item()),'g':float(rt.g[focus_neuron].item()),
|
|
'refractory':int(rt.refractory[focus_neuron].item()),'focus_spiked':int(focus_neuron in new_nodes),
|
|
'active_presynaptic_count':len(new_nodes),
|
|
'fanout_edge_count':int(np.sum(rt.indptr_np[new_nodes+1]-rt.indptr_np[new_nodes])) if len(new_nodes) else 0,
|
|
'active_presynaptic_neuron_idx':new_nodes.tolist()[:500]})
|
|
events=rt.read_spikes();step=events['step'].astype(np.int32);node=events['neuron_idx'].astype(np.int32)
|
|
counts=rt.spike_counts.cpu().numpy().astype(np.int32);v=rt.v.cpu().numpy();g=rt.g.cpu().numpy()
|
|
else:
|
|
rt.step(batch,1000.)
|
|
events=rt.read_spikes();step=events['step'].astype(np.int32);node=events['neuron_idx'].astype(np.int32)
|
|
counts=rt.spike_counts.cpu().numpy().astype(np.int32);v=rt.v.cpu().numpy();g=rt.g.cpu().numpy()
|
|
event_known=mode!='logging-off'
|
|
result={'mode':mode,'backend':backend,'logging_events_captured':event_known,
|
|
'spike_events_exact':bool(np.array_equal(step,rstep) and np.array_equal(node,rnode)) if event_known else None,
|
|
'per_neuron_counts_exact':bool(np.array_equal(counts,rcount)),
|
|
'final_voltage_max_abs_error':float(np.max(np.abs(v-rv))),
|
|
'final_g_max_abs_error':float(np.max(np.abs(g-rg))),
|
|
'first_event_difference':first_event_diff(rstep,rnode,step,node) if event_known else None,
|
|
'event_sha256':hashlib.sha256(step.tobytes()+node.tobytes()).hexdigest() if event_known else None,
|
|
'count_sha256':hashlib.sha256(counts.tobytes()).hexdigest(),
|
|
'torch_deterministic_algorithms':torch.are_deterministic_algorithms_enabled(),
|
|
'CUBLAS_WORKSPACE_CONFIG':os.getenv('CUBLAS_WORKSPACE_CONFIG'),
|
|
'state_window':window}
|
|
return result
|
|
|
|
|
|
def main():
|
|
ap=argparse.ArgumentParser();ap.add_argument('mode',choices=['default','logging-off','logging-on','instrumentation-off','instrumentation-on','single-stream-sync','deterministic-flags'])
|
|
ap.add_argument('--focus-step',type=int,default=-1000);ap.add_argument('--focus-neuron',type=int,default=-1)
|
|
ap.add_argument('--repeat',type=int,default=0);a=ap.parse_args()
|
|
try: result=run(a.mode,a.focus_step,a.focus_neuron)
|
|
except Exception as exc: result={'mode':a.mode,'error':repr(exc),'unsupported_or_failed':True}
|
|
result['repeat']=a.repeat
|
|
save(EXP/'divergence'/f'mode-{a.mode}-repeat-{a.repeat}.json',result)
|
|
print(json.dumps({k:result.get(k) for k in ('mode','repeat','spike_events_exact','per_neuron_counts_exact','first_event_difference','error')}),flush=True)
|
|
|
|
|
|
if __name__=='__main__':main()
|