Digital Drosophila Archive
120 lines
7.7 KiB
Python
120 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Exact-order replay of all frozen Phase 4N KC events through Phase 4M eligibility."""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow.parquet as pq
|
|
import yaml
|
|
|
|
ROOT=Path('/data/digital-drosophila')
|
|
EXP=ROOT/'experiments/phase4o-runtime-credit'
|
|
MODEL=ROOT/'models/male-cns-lif'
|
|
|
|
|
|
def save(p,x):
|
|
p.parent.mkdir(parents=True,exist_ok=True)
|
|
p.write_text(json.dumps(x,indent=2,sort_keys=True,default=str)+'\n')
|
|
|
|
|
|
def h(x): return hashlib.sha256(x.tobytes()).hexdigest()
|
|
|
|
|
|
def main():
|
|
ncfg=yaml.safe_load((MODEL/'credit-assignment-v1/phase4n.yaml').read_text())
|
|
mcfg=yaml.safe_load((MODEL/'kcmbon-efficacy-v1/phase4m.yaml').read_text())
|
|
lcfg=yaml.safe_load(Path(mcfg['frozen_phase4l_config']).read_text())
|
|
cfg=lcfg
|
|
identity=pd.read_parquet(ROOT/'experiments/phase4n-credit-assignment/edge-identity/frozen-edge-identity.parquet')
|
|
historical=pd.read_parquet(Path(mcfg['experiment_root'])/'efficacy-overlay/update-history.parquet')
|
|
kc_events=pd.read_parquet(ROOT/'experiments/phase4n-credit-assignment/local-eligibility/KC-spike-history.parquet')
|
|
graph=Path(cfg['paths']['connectome_root'])
|
|
meta=pq.read_table(graph/'neuron-index/unified-metadata.parquet',columns=['neuron_idx','annotation_class']).to_pandas()
|
|
kc=meta.loc[meta.annotation_class=='Kenyon_Cell','neuron_idx'].to_numpy(np.int32)
|
|
kc.sort()
|
|
lookup=np.full(len(meta),-1,np.int32);lookup[kc]=np.arange(len(kc),dtype=np.int32)
|
|
src=identity.source_neuron_idx.to_numpy(np.int32);local=lookup[src]
|
|
if (local<0).any(): raise RuntimeError('non-KC edge source')
|
|
comp=identity.compartment.astype(str).to_numpy()==str(ncfg['matched_compartment_verbatim'])
|
|
edgeid=identity.edge_id.to_numpy(np.int64)
|
|
id_to_local=pd.Series(np.arange(len(identity),dtype=np.int32),index=edgeid)
|
|
events={(int(seed),int(pair)):group[['step','neuron_idx']].to_numpy(np.int32)
|
|
for (seed,pair),group in kc_events.groupby(['seed','pairing'])}
|
|
histories={(int(seed),int(pair)):group for (seed,pair),group in historical.groupby(['train_seed','pairing'])}
|
|
dt=float(cfg['model']['dt_ms']);tau=float(cfg['eligibility_v2']['kc_trace_tau_ms'])
|
|
decay=np.float32(math.exp(-dt/tau));step_gate=int(round(float(mcfg['protocol']['dan_time_ms'])/dt))
|
|
k=np.float32(cfg['eligibility_v2']['eta_effective_mv']/cfg['model']['w_syn_mv_per_anatomical_synapse'])
|
|
bound=np.float32(1.0-cfg['eligibility_v2']['max_fractional_depression'])
|
|
results=[];first_signatures={};mismatch_examples=[]
|
|
for repeat in range(5):
|
|
for seed in ncfg['train_seeds']:
|
|
E=np.ones(len(identity),np.float32)
|
|
for pairing in range(1,101):
|
|
arr=events.get((seed,pairing),np.empty((0,2),np.int32))
|
|
# Phase4N emitted events in simulation-time, then neuron_idx order.
|
|
if len(arr) and (np.any(arr[:,0]<0) or np.any(arr[:,0]>=int(round(mcfg['protocol']['duration_ms']/dt)))): raise RuntimeError('invalid event timestamp')
|
|
if len(arr)>1 and (np.any(arr[1:,0]<arr[:-1,0]) or np.any((arr[1:,0]==arr[:-1,0])&(arr[1:,1]<=arr[:-1,1]))): raise RuntimeError('event tie/duplicate/order ambiguity')
|
|
relevant=arr[arr[:,0]<=step_gate]
|
|
if len(relevant) and (lookup[relevant[:,1]]<0).any(): raise RuntimeError('non-KC event')
|
|
by_step={int(step):lookup[relevant[relevant[:,0]==step,1]] for step in np.unique(relevant[:,0])}
|
|
trace=np.zeros(len(kc),np.float32)
|
|
for step in range(step_gate+1):
|
|
trace*=decay
|
|
firing=by_step.get(step)
|
|
if firing is not None: trace[firing]+=np.float32(1)
|
|
np.minimum(trace,np.float32(1),out=trace)
|
|
rec=histories[(seed,pairing)]
|
|
rid=id_to_local.loc[rec.edge_id.to_numpy(np.int64)].to_numpy(np.int32)
|
|
observed=trace[local[rid]]
|
|
expected=rec.eligibility.to_numpy(np.float32)
|
|
err=np.abs(observed-expected)
|
|
actual=np.flatnonzero(comp&(trace[local]>0))
|
|
missing=np.setdiff1d(rid,actual,assume_unique=True)
|
|
extra=np.setdiff1d(actual,rid,assume_unique=True)
|
|
predicted=np.maximum(bound,E[rid]*(np.float32(1)-k*observed))
|
|
frozen_after=rec.E_bounded_after.to_numpy(np.float32)
|
|
update_error=float(np.max(np.abs(predicted-frozen_after))) if len(predicted) else 0.0
|
|
trace_error=float(err.max()) if len(err) else 0.0
|
|
signature=h(np.concatenate((arr.flatten(),observed.view(np.int32),predicted.view(np.int32))))
|
|
key=(seed,pairing)
|
|
first_signatures.setdefault(key,signature)
|
|
row={'repeat':repeat,'seed':seed,'pairing':pairing,'event_count':len(arr),'event_sha256':h(arr),
|
|
'gate_step':step_gate,'eligibility_sha256':h(observed),'update_sha256':h(predicted),
|
|
'repeat_exact':signature==first_signatures[key],
|
|
'frozen_trace_max_abs_error':trace_error,'frozen_update_max_abs_error':update_error,
|
|
'missing_frozen_active_edges':len(missing),'extra_active_edges':len(extra),
|
|
'exact_to_phase4m':trace_error<=2e-5 and update_error<=2e-5 and len(missing)==0 and len(extra)==0}
|
|
if not row['exact_to_phase4m'] and len(mismatch_examples)<30:
|
|
j=int(np.argmax(err)) if len(err) else None
|
|
row['example']={'edge_id':int(rec.edge_id.iloc[j]) if j is not None else None,
|
|
'KC_neuron_idx':int(src[rid[j]]) if j is not None else None,
|
|
'timestamp_ms':float(mcfg['protocol']['dan_time_ms']),
|
|
'event_trace':float(observed[j]) if j is not None else None,
|
|
'frozen_trace':float(expected[j]) if j is not None else None,
|
|
'frozen_delta_E':float(frozen_after[j]-rec.E_bounded_before.iloc[j]) if j is not None else None,
|
|
'missing_edge_ids':edgeid[missing[:20]].tolist(),'extra_edge_ids':edgeid[extra[:20]].tolist()}
|
|
mismatch_examples.append(row.copy())
|
|
results.append(row)
|
|
# Each trial is run against the sealed Phase4M before/after trajectory.
|
|
# A divergent diagnostic event trace never feeds the next trial.
|
|
E[rid]=frozen_after
|
|
print(json.dumps({'repeat':repeat,'seed':seed,'total_so_far':len(results),
|
|
'exact_to_phase4m_so_far':sum(r['exact_to_phase4m'] for r in results)}),flush=True)
|
|
save(EXP/'eligibility-replay/event-reproducibility-matrix.json',results)
|
|
save(EXP/'eligibility-replay/event-reproducibility-matrix.json',results)
|
|
save(EXP/'eligibility-replay/event-replay-summary.json',{'complete_repeats':5,'events_per_repeat':500,
|
|
'same_event_and_update_sequence_5_of_5':all(r['repeat_exact'] for r in results),
|
|
'all_exact_to_frozen_phase4m':all(r['exact_to_phase4m'] for r in results),
|
|
'frozen_exact_count_by_repeat':{str(rep):sum(r['exact_to_phase4m'] for r in results if r['repeat']==rep) for rep in range(5)},
|
|
'mismatch_examples':mismatch_examples,
|
|
'event_order_contract_observed':'simulation_step ascending, neuron_idx ascending, no duplicate KC/step; trace updated after spike; external gate sampled after step 900 trace update',
|
|
'limitation':'Phase4M did not seal its training KC event sequence. Phase4N sealed the later event replay; exact repeat of those events does not prove the original GPU trajectory.'})
|
|
|
|
|
|
if __name__=='__main__':main()
|