digital-drosophila/models/male-cns-lif/deterministic-runtime-v2/phase4p.py

#!/usr/bin/env python3
"""Canonical deterministic MaleCNS LIF reference backend and validation tasks."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import time
from pathlib import Path
from typing import Any

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import torch
import yaml
from torch.utils.cpp_extension import load as load_extension

import phase4c


def sha_bytes(*arrays: np.ndarray) -> str:
    h=hashlib.sha256()
    for x in arrays:h.update(np.ascontiguousarray(x).tobytes())
    return h.hexdigest()


def save(path: Path,x: object):
    path.parent.mkdir(parents=True,exist_ok=True)
    path.write_text(json.dumps(x,indent=2,sort_keys=True,default=str)+'\n')


def build_extension(model:Path,exp:Path):
    build=exp/'gpu-reference/extension-build';build.mkdir(parents=True,exist_ok=True)
    start=time.perf_counter()
    ext=load_extension(name='male_cns_lif_deterministic_v2_ext',
        sources=[str(model/'phase4p_ext.cpp'),str(model/'phase4p_ext.cu')],
        build_directory=str(build),extra_cflags=['-O3','-ffp-contract=off'],
        extra_cuda_cflags=['-O3','--fmad=false','-std=c++20'],verbose=True)
    report={'status':'PASS','build_seconds':time.perf_counter()-start,'binary':str(Path(ext.__file__)),
            'binary_sha256':phase4c.sha256(Path(ext.__file__)),'flags':{'cxx':['-O3','-ffp-contract=off'],'cuda':['-O3','--fmad=false','-std=c++20']}}
    save(exp/'gpu-reference/extension.json',report);return ext


class CanonicalRuntime:
    def __init__(self,cfg:dict[str,Any],model:Path,exp:Path):
        self.cfg=cfg;self.exp=exp;self.device=torch.device('cuda:0')
        if torch.cuda.device_count()!=1 or cfg['device']['required_name_substring'] not in torch.cuda.get_device_name(0):
            raise RuntimeError('exactly one visible RTX 4090 required')
        p4cfg=yaml.safe_load((Path(cfg['paths']['historical_runtime_root'])/'phase4c.yaml').read_text())
        paths=phase4c.exact_paths(p4cfg)
        self.indptr_np,self.indices_np,self.raw_np=phase4c.p4b.validate_graph(p4cfg,paths)
        self.table=phase4c.full_metadata(paths['metadata']);self.pops=phase4c.populations(self.table)
        base=phase4c.Runtime(p4cfg,self.table,self.indptr_np,self.indices_np,self.raw_np,self.device,None)
        self.indptr,self.indices,self.weights=base.indptr,base.indices,base.weights
        self.ext=build_extension(model,exp)
        m=cfg['model'];self.n=int(cfg['expected_graph']['neurons']);self.dt=float(m['dt_ms'])
        self.delay_steps=int(round(m['synaptic_delay_ms']/self.dt));self.qlen=self.delay_steps+1
        self.ref_steps=int(round(m['refractory_ms']/self.dt));self.gd=np.float32(1-self.dt/m['synaptic_tau_ms'])
        self.alpha=np.float32(self.dt/m['membrane_tau_ms']);self.rest=np.float32(m['resting_potential_mv'])
        self.resetv=np.float32(m['reset_potential_mv']);self.threshold=np.float32(m['threshold_mv'])
        self.superclass=np.asarray(self.table['annotation_superclass'].to_pylist(),object)

    def reset(self):
        d=self.device;self.v=torch.full((self.n,),float(self.rest),dtype=torch.float32,device=d)
        self.g=torch.zeros(self.n,dtype=torch.float32,device=d);self.ref=torch.zeros(self.n,dtype=torch.int32,device=d)
        self.queue=torch.zeros((self.qlen,self.n),dtype=torch.float32,device=d)
        self.spike_u8=torch.zeros(self.n,dtype=torch.uint8,device=d);self.counts=torch.zeros(self.n,dtype=torch.int32,device=d)
        self.step_index=0;self.event_steps=[];self.event_nodes=[]

    def canonical_fanout(self,sources:torch.Tensor,destination:torch.Tensor)->dict[str,float|int]:
        if not sources.numel():return {'contributions':0,'segments':0,'sort_seconds':0.,'reduce_seconds':0.}
        starts=self.indptr[sources];degree=self.indptr[sources+1]-starts;cum=torch.cumsum(degree,0);total=int(cum[-1].item())
        edges=torch.empty(total,dtype=torch.int64,device=self.device)
        self.ext.fill_edges(self.indptr,sources,cum-degree,edges)
        t0=time.perf_counter();targets=self.indices[edges]
        order=torch.argsort(targets,stable=True);sorted_edges=edges[order];sorted_targets=targets[order]
        boundary=torch.ones(total,dtype=torch.bool,device=self.device)
        if total>1:boundary[1:]=sorted_targets[1:]!=sorted_targets[:-1]
        seg_starts=torch.nonzero(boundary,as_tuple=False).flatten().to(torch.int64)
        unique_targets=sorted_targets[seg_starts].to(torch.int32);sort_sec=time.perf_counter()-t0
        t1=time.perf_counter();self.ext.segmented_sum(sorted_edges,seg_starts,unique_targets,self.weights,destination)
        return {'contributions':total,'segments':int(len(unique_targets)),'sort_seconds':sort_sec,'reduce_seconds':time.perf_counter()-t1}

    def run(self,targets_np:np.ndarray,rate:float,seed:int,duration_ms:float,capture=True,forced_override=None):
        self.reset();steps=int(round(duration_ms/self.dt));targets=torch.as_tensor(targets_np,dtype=torch.int64,device=self.device)
        forced=forced_override if forced_override is not None else phase4c.generate_forced(steps,len(targets_np),rate,self.dt,seed,self.device)
        torch.cuda.synchronize();t0=time.perf_counter();sort_t=reduce_t=0.;contrib=0
        for step in range(steps):
            slot=step%self.qlen
            self.ext.state_step(self.v,self.g,self.ref,self.queue[slot],self.spike_u8,self.counts,
                float(self.gd),float(self.alpha),float(self.rest),float(self.resetv),float(self.threshold),self.ref_steps)
            self.ext.force_step(self.v,self.ref,self.spike_u8,self.counts,targets,forced[step],float(self.resetv),self.ref_steps)
            sources=torch.nonzero(self.spike_u8,as_tuple=False).flatten()
            stats=self.canonical_fanout(sources,self.queue[(step+self.delay_steps)%self.qlen])
            sort_t+=stats['sort_seconds'];reduce_t+=stats['reduce_seconds'];contrib+=stats['contributions']
            if capture and sources.numel():
                x=sources.cpu().numpy().astype(np.int32,copy=True);self.event_nodes.append(x);self.event_steps.append(np.full(len(x),step,np.int32))
            self.step_index+=1
        torch.cuda.synchronize();wall=time.perf_counter()-t0
        es=np.concatenate(self.event_steps) if self.event_steps else np.empty(0,np.int32)
        en=np.concatenate(self.event_nodes) if self.event_nodes else np.empty(0,np.int32)
        return {'step':es,'neuron_idx':en,'counts':self.counts.cpu().numpy().copy(),'v':self.v.cpu().numpy().copy(),'g':self.g.cpu().numpy().copy(),
                'metrics':{'wall_seconds':wall,'biological_ms':duration_ms,'realtime_factor':wall/(duration_ms/1000),
                           'peak_vram_bytes':int(torch.cuda.max_memory_allocated()),'contributions':int(contrib),
                           'sort_submit_seconds':sort_t,'reduce_submit_seconds':reduce_t,'events':len(en),'active':int((self.counts>0).sum().item())}}


def cpu_ordered_sum(edge_ids:np.ndarray,indices:np.ndarray,weights:np.ndarray):
    edge_ids=np.asarray(edge_ids,np.int64);targets=np.asarray(indices[edge_ids],np.int32)
    order=np.argsort(targets,kind='stable');ordered=edge_ids[order];ordered_target=targets[order]
    unique=[];sums=[];begin=0
    while begin<len(ordered):
        target=int(ordered_target[begin]);end=begin+1;acc=np.float32(0)
        while end<len(ordered) and ordered_target[end]==target:end+=1
        for e in ordered[begin:end]:acc=np.float32(acc+np.float32(weights[e]))
        unique.append(target);sums.append(acc);begin=end
    return np.asarray(unique,np.int32),np.asarray(sums,np.float32),ordered


def save_trial(root:Path,run:dict):
    root.mkdir(parents=True,exist_ok=True)
    pq.write_table(pa.table({'step':run['step'],'timestamp_ms':run['step'].astype(np.float64)*.1,'neuron_idx':run['neuron_idx']}),root/'spike-events.parquet',compression='zstd')
    for name,key,dtype in [('per-neuron-spike-counts.int32.npy','counts',np.int32),('final-voltage.float32.npy','v',np.float32),('final-g.float32.npy','g',np.float32)]:np.save(root/name,run[key].astype(dtype),allow_pickle=False)
    save(root/'statistics.json',run['metrics']|{'event_sha256':sha_bytes(run['step'],run['neuron_idx']),
         'counts_sha256':sha_bytes(run['counts']),'v_sha256':sha_bytes(run['v']),'g_sha256':sha_bytes(run['g'])})


def phase4c_trials(cfg,rt,repeats):
    exp=Path(cfg['paths']['experiment_root']);rows=[]
    for spec in cfg['reference_trials']:
        target=np.flatnonzero(rt.pops[spec['population']]).astype(np.int64)
        baseline=None
        for repeat in range(repeats):
            run=rt.run(target,float(spec['rate_hz']),int(spec['seed']),float(spec['duration_ms']),True)
            if baseline is None:
                baseline={k:run[k].copy() for k in ('step','neuron_idx','counts','v','g')};save_trial(exp/'phase4c-v2/reference'/spec['name'],run)
            exact={k:bool(np.array_equal(run[k],baseline[k])) for k in ('step','neuron_idx','counts','v','g')}
            row={'trial':spec['name'],'repeat':repeat,**exact,**run['metrics']};rows.append(row)
            save(exp/'phase4c-v2/reproducibility-matrix.json',rows)
            print(json.dumps({'trial':spec['name'],'repeat':repeat,'exact':all(exact.values()),'wall':run['metrics']['wall_seconds']}),flush=True)
    summary={'repeats':repeats,'matrix':{s['name']:{k:sum(r[k] for r in rows if r['trial']==s['name']) for k in ('step','neuron_idx','counts','v','g')} for s in cfg['reference_trials']},
             'all_exact':all(all(r[k] for k in ('step','neuron_idx','counts','v','g')) for r in rows),
             'performance':{s['name']:{'mean_wall_seconds':float(np.mean([r['wall_seconds'] for r in rows if r['trial']==s['name']])),
                'mean_realtime_factor':float(np.mean([r['realtime_factor'] for r in rows if r['trial']==s['name']])),
                'peak_vram_bytes':max(r['peak_vram_bytes'] for r in rows if r['trial']==s['name']),
                'mean_contributions':float(np.mean([r['contributions'] for r in rows if r['trial']==s['name']]))} for s in cfg['reference_trials']}}
    save(exp/'phase4c-v2/summary.json',summary)


def main():
    ap=argparse.ArgumentParser();ap.add_argument('--config',type=Path,required=True);ap.add_argument('task',choices=['phase4c']);ap.add_argument('--repeats',type=int)
    a=ap.parse_args();cfg=yaml.safe_load(a.config.read_text());exp=Path(cfg['paths']['experiment_root']);model=Path(cfg['paths']['model_root'])
    for sub in ('numerical-contract','cpu-reference','gpu-reference','threshold-3942','fanout-tests','phase4c-v2','historical-comparison','brainio','activity-stream','eligibility-replay','performance','revalidation','regression','audit','manifests'):(exp/sub).mkdir(parents=True,exist_ok=True)
    rt=CanonicalRuntime(cfg,model,exp)
    phase4c_trials(cfg,rt,a.repeats or int(cfg['repeats_per_trial']))


if __name__=='__main__':main()