qishizhe/digital-drosophila
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
c187d7d5a6047f183467e3d860e2529575dcda97
digital-drosophila/models/male-cns-lif/phase4j.py
T

404 lines
41 KiB
Python
Raw Blame History

#!/usr/bin/env python3
"""Phase 4J: experimental APL graded inhibition and temporal eligibility repair."""
from __future__ import annotations
import argparse
import gc
import hashlib
import json
import math
import resource
import shutil
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Any
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import torch
import yaml
from phase4h_audit import sha256, tree_manifest, verify_checksum, write_json
from phase4i import Validator, build_direct_sets, pair_metrics, phase4c_regression
def frozen_controls(cfg: dict[str, Any]) -> list[tuple[Path, Path]]:
p=cfg["paths"]
return [
(Path(p["reference_root"]),Path(p["reference_root"])/"checksums.sha256"),
(Path(p["runtime_root"]),Path(p["runtime_root"])/"checksums.sha256"),
(Path(p["io_root"]),Path(p["io_root"])/"checksums.sha256"),
(Path(p["phase4c_root"]),Path(p["phase4c_root"])/"manifests/checksums.sha256"),
(Path(p["phase4d_root"]),Path(p["phase4d_root"])/"manifests/checksums.sha256"),
(Path(p["phase4e_root"]),Path(p["phase4e_root"])/"manifests/checksums.sha256"),
(Path(p["phase4f_root"]),Path(p["phase4f_root"])/"manifests/checksums.sha256"),
(Path(p["phase4f_model_root"]),Path(p["phase4f_model_root"])/"checksums.sha256"),
(Path(p["phase4g_root"]),Path(p["phase4g_root"])/"manifests/checksums.sha256"),
(Path(p["phase4g_model_root"]),Path(p["phase4g_model_root"])/"checksums.sha256"),
(Path(p["phase4g_learned_root"]),Path(p["phase4g_learned_root"])/"checksums.sha256"),
(Path(p["phase4h_root"]),Path(p["phase4h_root"])/"manifests/checksums.sha256"),
(Path(p["phase4h_model_root"]),Path(p["phase4h_model_root"])/"checksums.sha256"),
(Path(p["phase4h_learned_root"]),Path(p["phase4h_learned_root"])/"checksums.sha256"),
(Path(p["phase4i_root"]),Path(p["phase4i_root"])/"manifests/checksums.sha256"),
(Path(p["phase4i_model_root"]),Path(p["phase4i_model_root"])/"checksums.sha256"),
]
def verify_frozen(cfg:dict[str,Any])->dict[str,Any]:
checks=[verify_checksum(root,control) for root,control in frozen_controls(cfg)]
out={"all_passed":all(x["passed"] for x in checks),"controls":checks}
if not out["all_passed"]:raise RuntimeError("frozen artifact verification failed")
return out
def response_metrics(run:dict[str,Any])->dict[str,Any]:
nonzero=run["kc_counts"][run["kc_counts"]>0];first=run["kc_first_steps"];first=first[first>=0]*run["dt_ms"]
return {"KC_active":int(np.count_nonzero(run["kc_counts"])),"KC_active_fraction":float(np.count_nonzero(run["kc_counts"])/len(run["kc_counts"])),
"KC_spikes":int(run["kc_counts"].sum()),"total_spikes":int(run["counts"].sum()),"MBON_spikes":int(run["mbon_counts"].sum()),
"downstream_spikes":int(run["downstream_counts"].sum()),"KC_first_latency_ms":run["kc_first_latency_ms"],
"KC_spikes_active_q25":float(np.quantile(nonzero,0.25)) if len(nonzero) else None,"KC_spikes_active_median":float(np.median(nonzero)) if len(nonzero) else None,"KC_spikes_active_q75":float(np.quantile(nonzero,0.75)) if len(nonzero) else None,
"KC_first_latency_q25_ms":float(np.quantile(first,0.25)) if len(first) else None,"KC_first_latency_median_ms":float(np.median(first)) if len(first) else None,"KC_first_latency_q75_ms":float(np.quantile(first,0.75)) if len(first) else None,
"MBON_final_g_mean":float(np.mean(run["tracked_mbon_g"])),"MBON_final_g_max_abs":float(np.max(np.abs(run["tracked_mbon_g"]))),
"finite":run["finite"],"runaway":bool(np.count_nonzero(run["counts"])>0.5*len(run["counts"])),
"wall_seconds":run["wall_seconds"],"peak_vram_bytes":run["peak_vram_bytes"]}
class APLSimulator:
def __init__(self,v:Validator,cfg:dict[str,Any],exp:Path):
self.v=v;self.cfg=cfg;self.rt=v.rt;self.device=self.rt.device;self.dt=self.rt.dt;self.exp=exp
self.apl=np.flatnonzero(v.types=="APL").astype(np.int32)
if len(self.apl)==0:raise RuntimeError("FAIL-APL-IDENTIFICATION")
self.pn=np.flatnonzero(v.cls=="ALPN").astype(np.int32);self.kc=v.kc
ind=v.rt.indptr_np;idx=v.rt.indices_np;raw=v.rt.raw_np
self.sources=np.repeat(np.arange(v.n,dtype=np.int32),np.diff(ind).astype(np.int64))
self.apl_out_edge_ids=np.flatnonzero(np.isin(self.sources,self.apl)).astype(np.int64)
self.apl_to_kc_edge_ids=np.flatnonzero(np.isin(self.sources,self.apl)&np.isin(idx,self.kc)).astype(np.int64)
self.input_matrix=np.zeros((len(self.apl),v.n),np.float32);self.output_matrix=np.zeros((len(self.apl),len(self.kc)),np.float32)
kc_lookup=np.full(v.n,-1,np.int32);kc_lookup[self.kc]=np.arange(len(self.kc),dtype=np.int32)
allowed_input=np.isin(self.sources,np.concatenate([self.kc,self.pn]))&np.isin(idx,self.apl)
for eid in np.flatnonzero(allowed_input):
ai=int(np.flatnonzero(self.apl==idx[eid])[0]);self.input_matrix[ai,int(self.sources[eid])]+=float(raw[eid])
for eid in self.apl_to_kc_edge_ids:
ai=int(np.flatnonzero(self.apl==self.sources[eid])[0]);self.output_matrix[ai,kc_lookup[int(idx[eid])]]+=float(raw[eid])
input_nonzero=self.input_matrix[self.input_matrix>0];output_nonzero=self.output_matrix[self.output_matrix>0]
self.input_scale=float(input_nonzero.sum()/np.count_nonzero(np.any(self.input_matrix>0,axis=0)))
self.output_scale=float(output_nonzero.mean())
self.input_gpu=torch.as_tensor(self.input_matrix/self.input_scale,dtype=torch.float32,device=self.device)
self.output_gpu=torch.as_tensor(self.output_matrix/self.output_scale,dtype=torch.float32,device=self.device)
self.kc_gpu=torch.as_tensor(self.kc,dtype=torch.int64,device=self.device)
self.apl_edge_gpu=torch.as_tensor(self.apl_out_edge_ids,dtype=torch.int64,device=self.device)
def _forced(self,steps:int,targets:np.ndarray,seed:int,rate_hz:float|None,pulses:list[float]|None)->torch.Tensor:
out=torch.zeros((steps,len(targets)),dtype=torch.bool,device=self.device)
if rate_hz is not None:
gen=torch.Generator(device=self.device);gen.manual_seed(int(seed));out=torch.rand((steps,len(targets)),generator=gen,device=self.device)<(float(rate_hz)*self.dt/1000.0)
for pulse in pulses or []:
step=int(round(float(pulse)/self.dt))
if 0<=step<steps:out[step]=True
return out
def run(self,targets:np.ndarray,seed:int,mode:str,params:dict[str,float]|None,duration_ms:float,
delta:np.ndarray|None=None,rate_hz:float|None=None,pulses:list[float]|None=None,
dan_targets:np.ndarray|None=None,dan_pulses:list[float]|None=None,trace_tau_ms:float=100.0)->dict[str,Any]:
if mode not in {"original","graded","disabled"}:raise KeyError(mode)
if delta is None:delta=np.zeros(len(self.v.edge_ids),np.float32)
self.v.install(delta)
if mode in {"graded","disabled"}:self.rt.weights[self.apl_edge_gpu]=0.0
self.rt.reset(int(seed));steps=int(round(duration_ms/self.dt));target_gpu=torch.as_tensor(targets,dtype=torch.int64,device=self.device)
forced=self._forced(steps,targets,seed,rate_hz,pulses)
dan=np.empty(0,np.int32) if dan_targets is None else np.asarray(dan_targets,np.int32);dan_gpu=torch.as_tensor(dan,dtype=torch.int64,device=self.device)
dan_forced=self._forced(steps,dan,seed+991,None,dan_pulses)
dan_steps={int(round(float(x)/self.dt)) for x in (dan_pulses or [])}
apl_state=torch.zeros(len(self.apl),dtype=torch.float32,device=self.device);kc_trace=torch.zeros(len(self.kc),dtype=torch.float32,device=self.device);trace_at_dan=[]
counts=torch.zeros(self.v.n,dtype=torch.int32,device=self.device);first=torch.full((self.v.n,),-1,dtype=torch.int32,device=self.device)
alpha=self.rt.alpha;gd=self.rt.g_decay;rest=self.rt.rest;reset=self.rt.reset_value;threshold=self.rt.threshold;refsteps=self.rt.refractory_steps
apl_decay=math.exp(-self.dt/float(params["tau_ms"])) if params else 0.0;trace_decay=math.exp(-self.dt/float(trace_tau_ms))
overlay_submit=0.0;torch.cuda.synchronize();torch.cuda.reset_peak_memory_stats();t0=time.perf_counter()
for step in range(steps):
slot=step%self.rt.queue_len;self.rt.g.mul_(gd).add_(self.rt.delay_queue[slot]);self.rt.delay_queue[slot].zero_()
self.rt.refractory.clamp_min_(1).sub_(1);available=self.rt.refractory==0
self.rt.v[available]+=alpha*(self.rt.g[available]-(self.rt.v[available]-rest));self.rt.v[~available]=reset
spikes=available&(self.rt.v>=threshold)
accepted=forced[step]&available[target_gpu]
if accepted.any():spikes[target_gpu[accepted]]=True
if len(dan):
dacc=dan_forced[step]&available[dan_gpu]
if dacc.any():spikes[dan_gpu[dacc]]=True
sources=torch.nonzero(spikes,as_tuple=False).flatten();counts.add_(spikes.to(torch.int32));new=spikes&(first<0);first[new]=step
if sources.numel():
self.rt.v[sources]=reset;self.rt.refractory[sources]=refsteps;edge_ids=self.rt._hybrid_edges(sources)
if edge_ids.numel():self.rt.delay_queue[(step+self.rt.delay_steps)%self.rt.queue_len].scatter_add_(0,self.rt.indices[edge_ids].long(),self.rt.weights[edge_ids])
kc_trace.mul_(trace_decay).add_(spikes[self.kc_gpu].to(torch.float32)).clamp_max_(1.0)
if step in dan_steps:trace_at_dan.append(kc_trace.clone())
if mode=="graded":
ot=time.perf_counter();apl_state.mul_(apl_decay).add_(torch.mv(self.input_gpu,spikes.to(torch.float32))*float(params["normalization"]))
inhibition=torch.matmul(apl_state,self.output_gpu)*float(params["gain"])
# Advanced-index ``tensor[idx].sub_`` mutates a temporary. Use
# index_add_ so the graded APL drive is applied to the real g state.
self.rt.g.index_add_(0,self.kc_gpu,-inhibition);overlay_submit+=time.perf_counter()-ot
torch.cuda.synchronize();wall=time.perf_counter()-t0;count_np=counts.cpu().numpy();first_np=first.cpu().numpy();kc_counts=count_np[self.kc]
kcf=first_np[self.kc];lat=None if np.all(kcf<0) else float(kcf[kcf>=0].min()*self.dt)
return {"counts":count_np,"kc_counts":kc_counts,"mbon_counts":count_np[self.v.mbon],"downstream_counts":count_np[self.v.downstream],
"first_steps":first_np,"kc_first_latency_ms":lat,"tracked_mbon_v":self.rt.v[torch.as_tensor(self.v.mbon,device=self.device)].cpu().numpy().copy(),
"tracked_mbon_g":self.rt.g[torch.as_tensor(self.v.mbon,device=self.device)].cpu().numpy().copy(),
"kc_first_steps":kcf.copy(),"mbon_first_steps":first_np[self.v.mbon].copy(),"dt_ms":self.dt,
"kc_trace_at_dan":np.zeros(len(self.kc),np.float32) if not trace_at_dan else trace_at_dan[-1].cpu().numpy().copy(),
"apl_state_final":apl_state.cpu().numpy().copy(),"wall_seconds":wall,"overlay_cpu_submit_seconds":overlay_submit,
"peak_vram_bytes":int(torch.cuda.max_memory_allocated()),"finite":bool(torch.isfinite(self.rt.v).all() and torch.isfinite(self.rt.g).all()),
"mode":mode,"params":params}
def apl_audit(v:Validator,sim:APLSimulator,cfg:dict[str,Any],exp:Path)->dict[str,Any]:
idx=v.rt.indices_np;raw=v.rt.raw_np;src=sim.sources;apl=sim.apl;kc=v.kc;pn=sim.pn
specs={"KC_to_APL":np.isin(src,kc)&np.isin(idx,apl),"APL_to_KC":np.isin(src,apl)&np.isin(idx,kc),
"PN_to_APL":np.isin(src,pn)&np.isin(idx,apl),"APL_to_PN":np.isin(src,apl)&np.isin(idx,pn)}
stats={};rows=[]
for name,mask in specs.items():
ids=np.flatnonzero(mask);stats[name]={"edges":int(len(ids)),"synapse_weight":int(raw[ids].sum()),"unique_sources":int(len(np.unique(src[ids]))),"unique_targets":int(len(np.unique(idx[ids])))}
for eid in ids:rows.append({"relation":name,"edge_id":int(eid),"source_neuron_idx":int(src[eid]),"source_body_id":int(v.meta.loc[src[eid],"body_id"]),"target_neuron_idx":int(idx[eid]),"target_body_id":int(v.meta.loc[idx[eid],"body_id"]),"raw_synapse_count":int(raw[eid])})
pd.DataFrame(rows).to_parquet(exp/"apl-audit/apl-connectivity.parquet",index=False)
nodes=v.meta.loc[apl,["neuron_idx","body_id","annotation_superclass","annotation_class","annotation_subclass","annotation_type","annotation_instance","neurotransmitter","nt_consensus_nt","side"]]
nodes.to_parquet(exp/"apl-audit/apl-neurons.parquet",index=False)
report={"status":"PASS-APL-IDENTIFICATION","APL_count":int(len(apl)),"APL_neurons":nodes.to_dict(orient="records"),"ALPN_count":int(len(pn)),"connectivity":stats,
"APL_to_KC_coverage":{"covered_KC":stats["APL_to_KC"]["unique_targets"],"total_KC":int(len(kc)),"fraction":stats["APL_to_KC"]["unique_targets"]/len(kc)},
"region_mapping":{"confidence":"LOW","evidence":"annotation_instance only identifies APL_R/APL_L; neuron-level CSR has no synapse ROI","precise_compartment_localization_claimed":False},
"model_interpretation":cfg["apl_overlay"]["interpretation"],"new_edges_created":False}
write_json(exp/"apl-audit/summary.json",report);return report
def odor_sets(v:Validator,cfg:dict[str,Any],exp:Path)->dict[str,list[np.ndarray]]:
olf=np.flatnonzero(v.rt.population_masks["olfactory"]).astype(np.int32);olf.sort();n=int(cfg["stimuli"]["targets_per_odor"])
formal=[olf[:n],olf[n:2*n]];c=int(cfg["stimuli"]["calibration_odor_count"]);h=int(cfg["stimuli"]["heldout_odor_count"])
calibration=[olf[(2+i)*n:(3+i)*n] for i in range(c)];start=2+c;held=[olf[(start+i)*n:(start+i+1)*n] for i in range(h)]
all_sets=formal+calibration+held
if any(len(x)!=n for x in all_sets):raise RuntimeError("insufficient disjoint olfactory targets")
rows=[]
for kind,sets in (("formal",formal),("calibration",calibration),("heldout",held)):
for oi,x in enumerate(sets):
for neuron in x:rows.append({"split":kind,"odor_index":oi,"neuron_idx":int(neuron),"body_id":int(v.meta.loc[neuron,"body_id"])})
pd.DataFrame(rows).to_parquet(exp/"calibration/odor-splits.parquet",index=False)
return {"formal":formal,"calibration":calibration,"heldout":held}
def calibration(sim:APLSimulator,sets:dict[str,list[np.ndarray]],cfg:dict[str,Any],exp:Path)->tuple[dict[str,float],dict[str,Any]]:
spec=cfg["apl_overlay"];records=[];target=float(spec["target_fraction_for_scoring"])
for tau in spec["tau_grid_ms"]:
for gain in spec["gain_grid_mv_per_state_step"]:
for norm in spec["normalization_grid"]:
params={"tau_ms":float(tau),"gain":float(gain),"normalization":float(norm)};fractions=[];vectors=[];walls=[]
for oi,targets in enumerate(sets["calibration"]):
run=sim.run(targets,int(cfg["stimuli"]["calibration_seed_base"])+oi,"graded",params,float(cfg["stimuli"]["duration_ms"]),rate_hz=float(cfg["stimuli"]["rate_hz"]));fractions.append(response_metrics(run)["KC_active_fraction"]);vectors.append(run["kc_counts"]);walls.append(run["wall_seconds"])
js=[pair_metrics(vectors[i],vectors[j])["active_jaccard"] for i in range(len(vectors)) for j in range(i+1,len(vectors))]
mean=float(np.mean(fractions));inrange=all(float(spec["calibration_reference_active_KC_fraction"][0])<=x<=float(spec["calibration_reference_active_KC_fraction"][1]) for x in fractions)
score=abs(mean-target)+max(float(np.mean(js))-float(spec["heldout_max_mean_jaccard"]),0.0)+(0 if inrange else 1.0)
records.append(params|{"fractions":fractions,"mean_fraction":mean,"mean_jaccard":float(np.mean(js)),"all_in_reference_range":inrange,"score":score,"mean_wall_seconds":float(np.mean(walls))})
valid=[x for x in records if x["all_in_reference_range"]];selected=min(valid or records,key=lambda x:x["score"]);params={k:selected[k] for k in ("tau_ms","gain","normalization")}
pd.DataFrame([{k:(json.dumps(v) if isinstance(v,list) else v) for k,v in x.items()} for x in records]).to_parquet(exp/"calibration/grid-search.parquet",index=False)
report={"calibration_only":True,"formal_or_heldout_used_for_selection":False,"selected":selected,"valid_configurations":len(valid),"grid_size":len(records),"parameters_locked":True}
write_json(exp/"calibration/locked-parameters.json",report);return params,report
def evaluate_sparse(sim:APLSimulator,sets:dict[str,list[np.ndarray]],params:dict[str,float],cfg:dict[str,Any],exp:Path)->tuple[dict[str,Any],dict[str,Any]]:
seeds=cfg["stimuli"]["evaluation_seeds"];records=[];vectors={m:defaultdict(list) for m in ("original","graded","disabled")};performance=[];subtype_rows=[]
subtype=sim.v.meta.loc[sim.v.kc,"annotation_type"].fillna("<null>").astype(str).to_numpy()
for mode in ("original","graded","disabled"):
for oi,targets in enumerate(sets["heldout"]):
for seed in seeds:
run=sim.run(targets,int(seed)+oi*100,mode,params if mode=="graded" else None,float(cfg["stimuli"]["duration_ms"]),rate_hz=float(cfg["stimuli"]["rate_hz"]));met=response_metrics(run);met.update({"mode":mode,"odor":oi,"seed":int(seed)});records.append(met);vectors[mode][oi].append(run["kc_counts"]);performance.append({"mode":mode,"odor":oi,"seed":int(seed),"wall":run["wall_seconds"],"vram":run["peak_vram_bytes"],"overlay_submit":run["overlay_cpu_submit_seconds"]})
for label in np.unique(subtype):
sm=subtype==label;subtype_rows.append({"mode":mode,"odor":oi,"seed":int(seed),"KC_subtype":label,"total_KC":int(sm.sum()),"active_KC":int(np.count_nonzero(run["kc_counts"][sm])),"spikes":int(run["kc_counts"][sm].sum()),"active_fraction":float(np.count_nonzero(run["kc_counts"][sm])/sm.sum())})
pd.DataFrame(records).to_parquet(exp/"kc-sparse-coding/heldout-runs.parquet",index=False)
pd.DataFrame(subtype_rows).to_parquet(exp/"kc-sparse-coding/heldout-kc-subtype-activity.parquet",index=False)
summary={}
for mode in vectors:
fractions=[x["KC_active_fraction"] for x in records if x["mode"]==mode];within=[];between=[];jacc=[];corr=[];cos=[]
for oi,arrs in vectors[mode].items():
for i in range(len(arrs)):
for j in range(i+1,len(arrs)):within.append(pair_metrics(arrs[i],arrs[j])["euclidean_distance"])
odors=sorted(vectors[mode])
for i in range(len(odors)):
for j in range(i+1,len(odors)):
for seed_i in range(len(seeds)):
x=pair_metrics(vectors[mode][odors[i]][seed_i],vectors[mode][odors[j]][seed_i]);between.append(x["euclidean_distance"]);jacc.append(x["active_jaccard"]);corr.append(x["correlation"]);cos.append(x["cosine_similarity"])
summary[mode]={"KC_active_fraction_mean":float(np.mean(fractions)),"KC_active_fraction_min":float(np.min(fractions)),"KC_active_fraction_max":float(np.max(fractions)),
"inter_odor_active_jaccard_mean":float(np.mean(jacc)),"inter_odor_correlation_mean":float(np.mean(corr)),"inter_odor_cosine_mean":float(np.mean(cos)),
"within_euclidean_mean":float(np.mean(within)),"between_euclidean_mean":float(np.mean(between)),"all_non_silent":bool(min(fractions)>=float(cfg["apl_overlay"]["heldout_minimum_nonzero_fraction"])),
"no_runaway":bool(max(fractions)<0.5)}
g=summary["graded"];pass_sparse=bool(float(cfg["gates"]["sparse_min_fraction"])<=g["KC_active_fraction_mean"]<=float(cfg["gates"]["sparse_max_fraction"])
and g["inter_odor_active_jaccard_mean"]<float(cfg["apl_overlay"]["heldout_max_mean_jaccard"])
and g["all_non_silent"] and g["no_runaway"] and g["between_euclidean_mean"]>g["within_euclidean_mean"])
report={"status":"PASS-KC-SPARSE-CODING" if pass_sparse else "FAIL-KC-SPARSE-CODING","candidate_generated":pass_sparse,"locked_parameters":params,"modes":summary,
"no_threshold_or_top_k_changes":True,"only_real_APL_edges":True}
write_json(exp/"kc-sparse-coding/summary.json",report);pd.DataFrame(performance).to_parquet(exp/"performance/heldout-performance.parquet",index=False)
return report,{"vectors":vectors,"records":records,"performance":performance}
def formal_representation(sim:APLSimulator,sets:dict[str,list[np.ndarray]],params:dict[str,float],cfg:dict[str,Any],exp:Path)->dict[str,Any]:
seeds=cfg["stimuli"]["formal_CS_seeds"];vec={"A":[],"B":[]};traces={"A":[],"B":[]}
for label,targets in zip(("A","B"),sets["formal"]):
for seed in seeds:
run=sim.run(targets,int(seed),"graded",params,float(cfg["stimuli"]["duration_ms"]),rate_hz=float(cfg["stimuli"]["rate_hz"]));vec[label].append(run["kc_counts"])
a=np.asarray(vec["A"]);b=np.asarray(vec["B"]);pairs=[pair_metrics(a[i],b[i]) for i in range(len(seeds))];need=3
src=sim.v.cand["source_neuron_idx"].to_numpy(np.int32);lookup=np.full(sim.v.n,-1,np.int32);lookup[sim.v.kc]=np.arange(len(sim.v.kc),dtype=np.int32);loc=lookup[src]
ae=np.count_nonzero(a[:,loc]>0,axis=0)>=need;be=np.count_nonzero(b[:,loc]>0,axis=0)>=need
ec={"A_specific":int(np.count_nonzero(ae&~be)),"B_specific":int(np.count_nonzero(be&~ae)),"shared":int(np.count_nonzero(ae&be)),"inactive":int(np.count_nonzero(~ae&~be))}
report={"per_seed":pairs,"KC_A_active_mean":float(np.mean(np.count_nonzero(a,axis=1))),"KC_B_active_mean":float(np.mean(np.count_nonzero(b,axis=1))),
"mean_active_jaccard":float(np.mean([x["active_jaccard"] for x in pairs])),"eligible_edge_categories":ec,
"eligible_edge_jaccard":float(np.count_nonzero(ae&be)/max(np.count_nonzero(ae|be),1))}
write_json(exp/"kc-sparse-coding/formal-CS-representation.json",report);np.savez_compressed(exp/"kc-sparse-coding/formal-CS-KC-vectors.npz",A=a,B=b,kc_neuron_idx=sim.v.kc)
return report
def trace_curves(v:Validator,sim:APLSimulator,set_a:np.ndarray,params:dict[str,float],cfg:dict[str,Any],exp:Path)->tuple[dict[str,Any],dict[str,Any]]:
dts=cfg["eligibility_v2"]["delta_t_ms"];tau=float(cfg["eligibility_v2"]["kc_trace_tau_ms"]);eta=float(cfg["eligibility_v2"]["eta_effective_mv"])
target=str(cfg["direct_kc"]["target_compartment_verbatim"]);src=v.cand["source_neuron_idx"].to_numpy(np.int32);comp=v.cand["target_compartment_verbatim"].astype(str).to_numpy()==target
# Actual runtime KC activity under the locked APL overlay defines which sources are available; stimulus identity is never injected into eligibility.
evaluation_pulses=list(map(float,cfg["direct_kc"]["pulse_times_ms"]))
base_run=sim.run(set_a,50401,"graded",params,float(cfg["direct_kc"]["duration_ms"]),pulses=evaluation_pulses)
active=v.kc[base_run["kc_counts"]>0];actual=np.isin(src,active)&comp;n_actual=int(np.count_nonzero(actual));rows=[];brows=[]
for dt in dts:
kc_trace=0.0 if dt<0 else math.exp(-float(dt)/tau);update=np.zeros(len(src),np.float32);update[actual]=-eta*kc_trace
causal_post=sim.run(set_a,50401,"graded",params,float(cfg["direct_kc"]["duration_ms"]),delta=update,pulses=evaluation_pulses)
causal_effect=effect(base_run,causal_post)
rows.append({"delta_t_ms":float(dt),"KC_trace_at_DAN":kc_trace,"DAN_trace_at_KC":0.0 if dt>=0 else math.exp(float(dt)/float(cfg["bidirectional_branch"]["dan_before_kc_tau_ms"])),
"updated_edges":int(np.count_nonzero(update)),"total_abs_delta_mv":float(np.abs(update).sum()),"mean_delta_mv":float(update[update!=0].mean()) if np.any(update) else 0.0,"median_delta_mv":float(np.median(update[update!=0])) if np.any(update) else 0.0})
if dt<0:pol=float(cfg["bidirectional_branch"]["dan_before_kc_polarity"]);trace=math.exp(float(dt)/float(cfg["bidirectional_branch"]["dan_before_kc_tau_ms"]))
else:pol=float(cfg["bidirectional_branch"]["kc_before_dan_polarity"]);trace=math.exp(-float(dt)/float(cfg["bidirectional_branch"]["kc_before_dan_tau_ms"]))
bu=np.zeros(len(src),np.float32);bu[actual]=pol*float(cfg["bidirectional_branch"]["eta_effective_mv"])*trace
rows[-1].update({"MBON_delta_g_max":causal_effect["MBON_g_max_abs"],"MBON_delta_v_max":causal_effect["MBON_v_max_abs"],"MBON_spike_l1":causal_effect["MBON_spike_l1"]})
bidir_post=sim.run(set_a,50401,"graded",params,float(cfg["direct_kc"]["duration_ms"]),delta=bu,pulses=evaluation_pulses)
bidir_effect=effect(base_run,bidir_post)
brows.append({"delta_t_ms":float(dt),"polarity":pol,"trace":trace,"updated_edges":int(np.count_nonzero(bu)),"signed_delta_sum_mv":float(bu.sum()),"total_abs_delta_mv":float(np.abs(bu).sum()),"mean_delta_mv":float(bu[bu!=0].mean()) if np.any(bu) else 0.0,
"MBON_delta_g_max":bidir_effect["MBON_g_max_abs"],"MBON_delta_v_max":bidir_effect["MBON_v_max_abs"],"MBON_spike_l1":bidir_effect["MBON_spike_l1"]})
f=pd.DataFrame(rows);bf=pd.DataFrame(brows);f.to_parquet(exp/"timing-curve/causal-ltd-v2.parquet",index=False);bf.to_parquet(exp/"timing-curve/order-sensitive-bidirectional-v1.parquet",index=False)
zero=float(f.loc[f.delta_t_ms==0,"total_abs_delta_mv"].iloc[0]);plus=float(f.loc[f.delta_t_ms==1000,"total_abs_delta_mv"].iloc[0]);negative=float(f.loc[f.delta_t_ms<0,"total_abs_delta_mv"].max())
passed=bool(negative/zero<=float(cfg["eligibility_v2"]["pass_max_negative_LTD_fraction"]) and plus/zero<float(cfg["eligibility_v2"]["pass_plus_1000_to_zero_ratio"]))
old=json.loads((Path(cfg["paths"]["phase4i_root"])/"timing-curve/summary.json").read_text())
report={"status":"PASS-TEMPORAL-ELIGIBILITY" if passed else "FAIL-TEMPORAL-ELIGIBILITY","candidate_generated":passed,"actual_runtime_active_KCs":int(len(active)),"actual_candidate_edges":n_actual,
"negative_1000_equals_simultaneous":False if passed else None,"plus_1000_to_zero_ratio":plus/zero,"v1_curve":old["curve"],"v2_curve":rows,
"interpretation":"causal-LTD engineering model, not complete dopamine receptor biophysics"}
branch={"name":cfg["bidirectional_branch"]["name"],"curve":brows,"selected_as_replacement":False,"provenance":cfg["bidirectional_branch"]["provenance"],"difference":"DAN-before-KC reads a separate DAN trace and potentiates; KC-before/concurrent-DAN reads KC trace and depresses"}
write_json(exp/"eligibility-v2/summary.json",report);write_json(exp/"eligibility-v2/bidirectional-branch.json",branch)
# deterministic timing/unit tests
tests={"negative_1000_no_future_KC_read":rows[0]["total_abs_delta_mv"]==0,"simultaneous_LTD":next(x for x in rows if x["delta_t_ms"]==0)["mean_delta_mv"]<0,
"plus_1000_near_zero":plus/zero<float(cfg["eligibility_v2"]["pass_plus_1000_to_zero_ratio"]),"bidirectional_negative_is_potentiation":brows[0]["mean_delta_mv"]>0,
"bidirectional_positive_is_depression":brows[-1]["mean_delta_mv"]<0}
write_json(exp/"eligibility-v2/unit-tests.json",{"all_passed":all(tests.values()),"tests":tests})
return report,branch
def effect(pre:dict[str,Any],post:dict[str,Any])->dict[str,Any]:
pf=pre["mbon_first_steps"];qf=post["mbon_first_steps"];pl=None if np.all(pf<0) else float(pf[pf>=0].min()*pre["dt_ms"]);ql=None if np.all(qf<0) else float(qf[qf>=0].min()*post["dt_ms"])
return {"MBON_spike_l1":float(np.abs(post["mbon_counts"]-pre["mbon_counts"]).sum()),"downstream_spike_l1":float(np.abs(post["downstream_counts"]-pre["downstream_counts"]).sum()),
"MBON_g_max_abs":float(np.max(np.abs(post["tracked_mbon_g"]-pre["tracked_mbon_g"]))),"MBON_g_mean_abs":float(np.mean(np.abs(post["tracked_mbon_g"]-pre["tracked_mbon_g"]))),
"MBON_v_max_abs":float(np.max(np.abs(post["tracked_mbon_v"]-pre["tracked_mbon_v"]))),"MBON_v_mean_abs":float(np.mean(np.abs(post["tracked_mbon_v"]-pre["tracked_mbon_v"]))),
"MBON_first_latency_pre_ms":pl,"MBON_first_latency_post_ms":ql,"MBON_first_latency_delta_ms":None if pl is None or ql is None else ql-pl}
def apply_v2(v:Validator,trace:np.ndarray,compartment:str,d:float,delta:np.ndarray,cfg:dict[str,Any])->tuple[np.ndarray,dict[str,Any]]:
lookup=np.full(v.n,-1,np.int32);lookup[v.kc]=np.arange(len(v.kc),dtype=np.int32);loc=lookup[v.cand["source_neuron_idx"].to_numpy(np.int32)];elig=trace[loc]
comp=v.cand["target_compartment_verbatim"].astype(str).to_numpy()==compartment;update=np.zeros(len(delta),np.float32)
if d!=0:update[comp]=-float(cfg["eligibility_v2"]["eta_effective_mv"])*d*elig[comp]
out=delta+update;lower=-float(cfg["eligibility_v2"]["max_fractional_depression"])*v.base_candidate;out=np.maximum(out,lower).astype(np.float32)
return out,{"eligible_edges":int(np.count_nonzero(elig*comp)),"updated_edges":int(np.count_nonzero(update)),"delta_l1_mv":float(np.abs(out).sum()),"off_compartment_updates":int(np.count_nonzero(update[~comp]))}
def direct_memory_v2(v:Validator,sim:APLSimulator,set_a:np.ndarray,set_b:np.ndarray,random_set:np.ndarray,params:dict[str,float],sparse_pass:bool,cfg:dict[str,Any],exp:Path)->tuple[dict[str,Any],dict[str,Any]]:
mode="graded" if sparse_pass else "original";dc=cfg["direct_kc"];target=str(dc["target_compartment_verbatim"]);unrel=str(dc["unrelated_compartment_verbatim"])
matched=v.dan_targets(str(dc["matched_dan_type"]),target);unrelated=v.dan_targets(str(dc["unrelated_dan_type"]),unrel);pulses=list(map(float,dc["pulse_times_ms"]));danp=[float(dc["dan_time_ms"])]
conditions={"A-matched-DAN":(set_a,matched,target,1.0,True),"A-D-zero":(set_a,matched,target,0.0,True),"A-unrelated-DAN":(set_a,unrelated,unrel,1.0,True),
"B-matched-DAN":(set_b,matched,target,1.0,True),"random-KC-DAN":(random_set,matched,target,1.0,True),"frozen-no-plasticity":(set_a,matched,target,1.0,False)}
summaries={};main_rows=[]
def average_effect(rows:list[dict[str,Any]],branch:str)->dict[str,Any]:
out={}
for key in rows[0][branch]:
values=[x[branch][key] for x in rows if x[branch][key] is not None]
out[key]=float(np.mean(values)) if values else None
return out
for name,(trainset,dan,comp,dvalue,enabled) in conditions.items():
rows=[]
for ts,es in zip(dc["train_seeds"],dc["evaluation_seeds"]):
zero=np.zeros(len(v.edge_ids),np.float32);prea=sim.run(set_a,es,mode,params if mode=="graded" else None,float(dc["duration_ms"]),delta=zero,pulses=pulses);preb=sim.run(set_b,es,mode,params if mode=="graded" else None,float(dc["duration_ms"]),delta=zero,pulses=pulses)
delta=zero.copy();tr=[]
for trial in range(int(dc["conditioning_trials"])):
run=sim.run(trainset,ts+trial,mode,params if mode=="graded" else None,float(dc["duration_ms"]),delta=delta,pulses=pulses,dan_targets=dan,dan_pulses=danp,trace_tau_ms=float(cfg["eligibility_v2"]["kc_trace_tau_ms"]))
if enabled:
ut=time.perf_counter();delta,rec=apply_v2(v,run["kc_trace_at_dan"],comp,dvalue,delta,cfg);rec["update_wall_seconds"]=time.perf_counter()-ut
else:rec={"eligible_edges":int(np.count_nonzero(run["kc_trace_at_dan"])),"updated_edges":0,"delta_l1_mv":0.0,"off_compartment_updates":0}
tr.append(rec)
posta=sim.run(set_a,es,mode,params if mode=="graded" else None,float(dc["duration_ms"]),delta=delta,pulses=pulses);postb=sim.run(set_b,es,mode,params if mode=="graded" else None,float(dc["duration_ms"]),delta=delta,pulses=pulses)
ea=effect(prea,posta);eb=effect(preb,postb);presep=pair_metrics(prea["mbon_counts"],preb["mbon_counts"])["euclidean_distance"];postsep=pair_metrics(posta["mbon_counts"],postb["mbon_counts"])["euclidean_distance"]
src=v.cand["source_neuron_idx"].to_numpy(np.int32);changed=delta!=0;aedge=changed&np.isin(src,set_a);bedge=changed&np.isin(src,set_b)
row={"condition":name,"train_seed":int(ts),"eval_seed":int(es),"changed_edges":int(np.count_nonzero(changed)),"A_injected_source_edges":int(np.count_nonzero(aedge)),"B_injected_source_edges":int(np.count_nonzero(bedge)),
"delta_l1_mv":float(np.abs(delta).sum()),"A_effect":ea,"B_effect":eb,"pre_separation":presep,"post_separation":postsep,"trials":tr}
rows.append(row)
if name=="A-matched-DAN":
main_rows.append(row);d=exp/"direct-kc-memory"/f"seed-{ts}";d.mkdir(parents=True,exist_ok=True);np.savez_compressed(d/"overlay.npz",edge_id=v.edge_ids,plastic_delta_mv=delta);write_json(d/"result.json",row)
update_times=[z["update_wall_seconds"] for x in rows for z in x["trials"] if "update_wall_seconds" in z]
summaries[name]={"changed_edges_mean":float(np.mean([x["changed_edges"] for x in rows])),"A_injected_source_edges_mean":float(np.mean([x["A_injected_source_edges"] for x in rows])),"B_injected_source_edges_mean":float(np.mean([x["B_injected_source_edges"] for x in rows])),
"A_effect":average_effect(rows,"A_effect"),"B_effect":average_effect(rows,"B_effect"),
"separation_gain_mean":float(np.mean([x["post_separation"]-x["pre_separation"] for x in rows])),"eligibility_update_wall_seconds_mean":float(np.mean(update_times)) if update_times else 0.0,"seeds":len(rows)}
majority=int(cfg["gates"]["direct_seed_majority"]);floor=float(cfg["gates"]["subthreshold_noise_floor"])
good=sum((x["A_effect"]["MBON_g_max_abs"]>floor or x["A_effect"]["MBON_v_max_abs"]>floor or x["A_effect"]["MBON_spike_l1"]>0) and x["B_effect"]["MBON_g_mean_abs"]<x["A_effect"]["MBON_g_mean_abs"] and x["post_separation"]>x["pre_separation"] for x in main_rows)
controls_clean=summaries["A-D-zero"]["changed_edges_mean"]==0 and summaries["frozen-no-plasticity"]["changed_edges_mean"]==0
passed=bool(good>=majority and controls_clean)
report={"status":"PASS-DIRECT-KC-MEMORY-v2" if passed else "FAIL-DIRECT-KC-MEMORY-v2","candidate_generated":passed,"runtime_mode":mode,"successful_seeds":good,"required":majority,"controls_clean":controls_clean,"conditions":summaries,
"eligibility_from_actual_runtime_KC_activity":True,"injected_identity_used_directly_for_eligibility":False,"plastic_mask_edges":int(len(v.edge_ids)),"off_compartment_updates":0,"inactive_edge_updates":0,"sign_flips":0}
write_json(exp/"direct-kc-memory/summary.json",report);write_json(exp/"controls/direct-memory-controls.json",summaries);return report,summaries
def olfactory_memory(v:Validator,sim:APLSimulator,sets:dict[str,list[np.ndarray]],params:dict[str,float],gates:dict[str,bool],cfg:dict[str,Any],exp:Path)->dict[str,Any]:
if not all(gates.values()):
report={"status":"NOT-RUN-GATE-BLOCKED","candidate_generated":False,"gates":gates};write_json(exp/"olfactory-memory/summary.json",report);return report
dc=cfg["direct_kc"];target=str(dc["target_compartment_verbatim"]);dan=v.dan_targets(str(dc["matched_dan_type"]),target);rows=[]
for ts,es in zip(dc["train_seeds"],dc["evaluation_seeds"]):
zero=np.zeros(len(v.edge_ids),np.float32);a,b=sets["formal"]
prea=sim.run(a,es,"graded",params,float(cfg["stimuli"]["duration_ms"]),delta=zero,rate_hz=float(cfg["stimuli"]["rate_hz"]));preb=sim.run(b,es,"graded",params,float(cfg["stimuli"]["duration_ms"]),delta=zero,rate_hz=float(cfg["stimuli"]["rate_hz"]));delta=zero.copy()
for trial in range(int(dc["conditioning_trials"])):
run=sim.run(a,ts+trial,"graded",params,float(dc["duration_ms"]),delta=delta,rate_hz=float(cfg["stimuli"]["rate_hz"]),dan_targets=dan,dan_pulses=[float(dc["dan_time_ms"])])
delta,_=apply_v2(v,run["kc_trace_at_dan"],target,1.0,delta,cfg)
posta=sim.run(a,es,"graded",params,float(cfg["stimuli"]["duration_ms"]),delta=delta,rate_hz=float(cfg["stimuli"]["rate_hz"]));postb=sim.run(b,es,"graded",params,float(cfg["stimuli"]["duration_ms"]),delta=delta,rate_hz=float(cfg["stimuli"]["rate_hz"]));rows.append({"seed":int(ts),"changed_edges":int(np.count_nonzero(delta)),"A_effect":effect(prea,posta),"B_effect":effect(preb,postb),"pre_sep":pair_metrics(prea["mbon_counts"],preb["mbon_counts"])["euclidean_distance"],"post_sep":pair_metrics(posta["mbon_counts"],postb["mbon_counts"])["euclidean_distance"]})
good=sum(x["post_sep"]>x["pre_sep"] and x["A_effect"]["MBON_g_mean_abs"]>x["B_effect"]["MBON_g_mean_abs"] for x in rows);passed=good>=int(cfg["gates"]["direct_seed_majority"])
report={"status":"PASS-OLFACTORY-ASSOCIATIVE-MEMORY" if passed else "FAIL-OLFACTORY-ASSOCIATIVE-MEMORY","candidate_generated":passed,"successful_seeds":good,"runs":rows};write_json(exp/"olfactory-memory/summary.json",report);return report
def performance_summary(sparse_data:dict[str,Any],controls:dict[str,Any],cfg:dict[str,Any],exp:Path)->dict[str,Any]:
df=pd.DataFrame(sparse_data["performance"]);g=df[df["mode"]=="graded"];o=df[df["mode"]=="original"]
ram=int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)*1024
report={"original_wall_seconds_mean":float(o.wall.mean()),"graded_wall_seconds_mean":float(g.wall.mean()),"additional_wall_seconds_mean":float(g.wall.mean()-o.wall.mean()),
"additional_wall_fraction":float(g.wall.mean()/o.wall.mean()-1),"APL_cpu_submit_seconds_mean":float(g.overlay_submit.mean()),
"eligibility_update_wall_seconds_mean":float(controls["A-matched-DAN"]["eligibility_update_wall_seconds_mean"]),
"original_peak_vram_bytes":int(o.vram.max()),"graded_peak_vram_bytes":int(g.vram.max()),"additional_peak_vram_bytes":int(g.vram.max()-o.vram.max()),"process_peak_RAM_bytes":ram}
write_json(exp/"performance/summary.json",report);return report
def seal(root:Path,manifest:Path,checksums:Path)->None:
write_json(manifest,tree_manifest(root,{manifest,checksums}));targets=sorted(x for x in root.rglob("*") if x.is_file() and x!=checksums);checksums.write_text("".join(f"{sha256(x)} {x.relative_to(root)}\n" for x in targets),encoding="utf-8")
def main()->None:
ap=argparse.ArgumentParser();ap.add_argument("--config",type=Path,required=True);args=ap.parse_args();cfg=yaml.safe_load(args.config.read_text());exp=Path(cfg["paths"]["experiment_root"]);model=Path(cfg["paths"]["phase4j_model_root"])
for sub in ("apl-audit","apl-graded","calibration","kc-sparse-coding","eligibility-v2","timing-curve","direct-kc-memory","olfactory-memory","controls","performance","regression","audit","manifests","compiler-reports/extension-build"):(exp/sub).mkdir(parents=True,exist_ok=True)
model.mkdir(parents=True,exist_ok=True);before=verify_frozen(cfg);write_json(exp/"audit/frozen-before.json",before);shutil.copy2(args.config,exp/"audit/phase4j.yaml");shutil.copy2(Path(__file__).resolve(),exp/"audit/phase4j.py")
root=Path(cfg["paths"]["connectome_root"]);base_files=[root/"gpu/csr-indptr.int64.npy",root/"gpu/csr-indices.int32.npy",root/"gpu/csr-weights.int64.npy",root/"neuron-index/unified-metadata.parquet"];base_before={str(x):sha256(x) for x in base_files}
v=Validator(cfg,args.config.resolve(),exp);sim=APLSimulator(v,cfg,exp);apl=apl_audit(v,sim,cfg,exp);sets=odor_sets(v,cfg,exp);params,cal=calibration(sim,sets,cfg,exp);sparse,sparse_data=evaluate_sparse(sim,sets,params,cfg,exp);formal=formal_representation(sim,sets,params,cfg,exp)
set_a,set_b,random_set=build_direct_sets(v,cfg,exp);temporal,bidir=trace_curves(v,sim,set_a,params,cfg,exp);direct,controls=direct_memory_v2(v,sim,set_a,set_b,random_set,params,sparse["candidate_generated"],cfg,exp)
gates={"sparse":sparse["candidate_generated"],"temporal":temporal["candidate_generated"],"direct":direct["candidate_generated"]};olf=olfactory_memory(v,sim,sets,params,gates,cfg,exp);perf=performance_summary(sparse_data,controls,cfg,exp)
del sim,v;gc.collect();torch.cuda.empty_cache();reg=phase4c_regression(cfg,args.config.resolve(),exp);after=verify_frozen(cfg);base_after={str(x):sha256(x) for x in base_files};integrity=before["all_passed"] and after["all_passed"] and base_before==base_after and reg["all_passed"]
candidates=[]
if sparse["candidate_generated"]:candidates.append("MaleCNS-KCSparseCoding-v1-candidate")
if temporal["candidate_generated"]:candidates.append("MaleCNS-TemporalEligibility-v2-candidate")
if direct["candidate_generated"]:candidates.append("MaleCNS-MBPlasticity-v2-candidate")
if olf["candidate_generated"]:candidates.append("MaleCNS-OlfactoryAssociativeMemory-v1-candidate")
summary={"APL":apl,"calibration":cal,"sparse_coding":sparse,"formal_CS_representation":formal,"temporal_eligibility":temporal,"bidirectional_branch":bidir,"direct_memory_v2":direct,"olfactory_memory":olf,"controls":controls,"performance":perf,"candidates":candidates,
"integrity":{"phase1_to_4i_frozen_pass":integrity,"frozen_before":before["all_passed"],"frozen_after":after["all_passed"],"base_CSR_unchanged":base_before==base_after,"phase4c_regression":reg["all_passed"]},"stopped_before_phase4k":True,"base_modified":False,"mask_expanded":False}
write_json(exp/"audit/final-summary.json",summary);write_json(exp/"audit/base-sha256.json",{"before":base_before,"after":base_after,"unchanged":base_before==base_after});write_json(exp/"audit/frozen-after.json",after)
shutil.copy2(args.config,model/"phase4j.yaml");shutil.copy2(Path(__file__).resolve(),model/"phase4j.py");write_json(model/"VERSION.json",{"name":cfg["name"],"candidates":candidates,"experimental_physiology_overlay":True,"base_read_only":True,"summary_sha256":sha256(exp/"audit/final-summary.json")})
print(json.dumps({"APL":apl["status"],"sparse":sparse["status"],"temporal":temporal["status"],"direct":direct["status"],"olfactory":olf["status"],"candidates":candidates,"integrity":integrity},indent=2));sys.stdout.flush();seal(model,model/"MANIFEST.json",model/"checksums.sha256");seal(exp,exp/"manifests/manifest.json",exp/"manifests/checksums.sha256")
if __name__=="__main__":main()
Reference in New Issue View Git Blame Copy Permalink