Files
digital-drosophila/models/male-cns-lif/author-apl-lab-v2/verify_frozen.py
T
2026-09-30 09:26:21 +08:00

36 lines
2.1 KiB
Python

"""Run old reference unchanged, writing only new experiment output."""
import os
import sys
import time
import numpy as np
import pandas as pd
import yaml
from context import ROOT,EXP,CFG,save,check
def main():
check();assert not (EXP/'regression/frozen-five.json').exists()
rtpath=ROOT/'models/male-cns-lif/deterministic-runtime-v2';pc=yaml.safe_load((rtpath/'phase4p.yaml').read_text())
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();sys.path.insert(0,str(ROOT/'models/male-cns-lif'));sys.path.insert(0,str(rtpath))
import torch,phase4p
assert torch.cuda.mem_get_info()[0]>2*1024**3
torch.cuda.reset_peak_memory_stats();t0=time.perf_counter();rt=phase4p.CanonicalRuntime(pc,rtpath,EXP);rows=[]
for spec in pc['reference_trials']:
queues();base=ROOT/'experiments/phase4p-deterministic-runtime/phase4c-v2/reference'/spec['name']
ev=pd.read_parquet(base/'spike-events.parquet')
ref=dict(step=ev.step.to_numpy(),neuron_idx=ev.neuron_idx.to_numpy(),counts=np.load(base/'per-neuron-spike-counts.int32.npy'),v=np.load(base/'final-voltage.float32.npy'),g=np.load(base/'final-g.float32.npy'))
result=rt.run(np.flatnonzero(rt.pops[spec['population']]),spec['rate_hz'],spec['seed'],spec['duration_ms'])
row=dict(trial=spec['name'],exact={k:bool(np.array_equal(v,result[k])) for k,v in ref.items()},wall_seconds=result['metrics']['wall_seconds'])
rows.append(row);print(row,flush=True)
save(EXP/'regression/frozen-five.json',dict(pass_all=all(all(r['exact'].values()) for r in rows),trials=rows,
wall_seconds=time.perf_counter()-t0,peak_allocated_vram=int(torch.cuda.max_memory_allocated()),oracle='old frozen deterministic-v2 only'))
assert all(all(r['exact'].values()) for r in rows)
if __name__=='__main__':main()