Digital Drosophila Archive
466 lines
40 KiB
Python
466 lines
40 KiB
Python
#!/usr/bin/env python3
|
|
"""Phase 4G: learned interfaces around an immutable MaleCNS-LIF brain core."""
|
|
from __future__ import annotations
|
|
|
|
import argparse, gc, hashlib, json, math, platform, shutil, subprocess, sys, time
|
|
from collections import Counter, defaultdict
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
import torch
|
|
import torch.nn as nn
|
|
import yaml
|
|
from scipy.spatial import cKDTree
|
|
|
|
from brain_runtime import BrainRuntime, InputBatch, InputEvent
|
|
from frozen_interface import AbstractActionOutput
|
|
|
|
|
|
def np_json(x: Any) -> Any:
|
|
if isinstance(x, np.generic): return x.item()
|
|
if isinstance(x, np.ndarray): return x.tolist()
|
|
if isinstance(x, Path): return str(x)
|
|
raise TypeError(type(x).__name__)
|
|
|
|
|
|
def write_json(path: Path, value: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(value, indent=2, sort_keys=True, default=np_json)+"\n", encoding="utf-8")
|
|
|
|
|
|
def sha256(path: Path, chunk: int=8<<20) -> str:
|
|
h=hashlib.sha256()
|
|
with path.open("rb") as f:
|
|
while data:=f.read(chunk): h.update(data)
|
|
return h.hexdigest()
|
|
|
|
|
|
def verify_checksum(root: Path, control: Path) -> dict[str,Any]:
|
|
checked=0; failures=[]
|
|
for line in control.read_text().splitlines():
|
|
if not line.strip(): continue
|
|
expected,rel=line.split(None,1); path=root/rel.strip().lstrip("*"); checked+=1
|
|
if not path.is_file() or sha256(path)!=expected: failures.append(str(path))
|
|
return {"control":str(control),"checked":checked,"failures":failures,"passed":not failures}
|
|
|
|
|
|
def frozen_audit(cfg:dict[str,Any])->tuple[list[Path],dict[str,str],dict[str,Any]]:
|
|
p=cfg["paths"]
|
|
controls=[Path(p["reference_root"])/"checksums.sha256",Path(p["runtime_root"])/"checksums.sha256",
|
|
Path(p["io_root"])/"checksums.sha256",Path(p["phase4c_root"])/"manifests/checksums.sha256",
|
|
Path(p["phase4d_root"])/"manifests/checksums.sha256",Path(p["phase4e_root"])/"manifests/checksums.sha256",
|
|
Path(p["phase4f_root"])/"manifests/checksums.sha256",Path(p["phase4f_model_root"])/"checksums.sha256"]
|
|
if missing:=[str(x) for x in controls if not x.is_file()]: raise FileNotFoundError(missing)
|
|
checks=[]
|
|
for x in controls:
|
|
root=x.parent.parent if x.parent.name=="manifests" else x.parent; checks.append(verify_checksum(root,x))
|
|
if not all(x["passed"] for x in checks): raise RuntimeError(checks)
|
|
before={str(x):sha256(x) for x in controls}
|
|
return controls,before,{"all_passed":True,"controls":checks}
|
|
|
|
|
|
def stable_seed(base:int,text:str)->int:
|
|
return base+int.from_bytes(hashlib.sha256(text.encode()).digest()[:4],"little")%1_000_000
|
|
|
|
|
|
def make_samples(cfg:dict[str,Any])->pd.DataFrame:
|
|
rows=[]
|
|
for split,key in (("train","train_stimulus_seeds"),("held-out","evaluation_stimulus_seeds")):
|
|
for seed in cfg["dataset"][key]:
|
|
rng=np.random.default_rng(int(seed))
|
|
if split=="train":
|
|
offset=rng.uniform(-.08,.08); speed=rng.uniform(.85,1.15); size=rng.uniform(.9,1.1); curve=rng.uniform(-.08,.08); phase=rng.uniform(-.05,.05)
|
|
else:
|
|
sign=-1 if int(seed)%2 else 1
|
|
offset=sign*rng.uniform(.16,.26); speed=rng.choice([rng.uniform(.62,.78),rng.uniform(1.25,1.42)])
|
|
size=rng.choice([rng.uniform(.68,.82),rng.uniform(1.22,1.40)]); curve=sign*rng.uniform(.16,.26); phase=sign*rng.uniform(.10,.20)
|
|
for pattern in cfg["dataset"]["patterns"]:
|
|
rows.append({"sample_id":len(rows),"split":split,"stimulus_seed":int(seed),"pattern":pattern,
|
|
"position_offset":offset,"speed_scale":speed,"size_scale":size,"trajectory_curve":curve,"temporal_phase":phase})
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def frames_variant(pattern:str,x:np.ndarray,y:np.ndarray,segments:int,v:dict[str,float])->np.ndarray:
|
|
out=[]; off=v["position_offset"]
|
|
for i in range(segments):
|
|
z=np.clip((i/max(segments-1,1)+v["temporal_phase"])*v["speed_scale"],0,1)
|
|
s0=.22*v["size_scale"]
|
|
if pattern=="full-field": f=np.ones_like(x)
|
|
elif pattern in {"left-spot","right-spot","upper-spot","lower-spot"}:
|
|
loc={"left-spot":(-.65+off,0),"right-spot":(.65+off,0),"upper-spot":(off,.65),"lower-spot":(off,-.65)}[pattern]
|
|
f=np.exp(-((x-loc[0])**2+(y-loc[1])**2)/(2*s0*s0))
|
|
elif pattern in {"left-to-right","right-to-left","up-to-down","down-to-up"}:
|
|
p=-.8+1.6*z; p=-p if pattern in {"right-to-left","down-to-up"} else p; bend=v["trajectory_curve"]*math.sin(math.pi*z)
|
|
cx,cy=((p,bend+off) if pattern in {"left-to-right","right-to-left"} else (bend+off,-p))
|
|
f=np.exp(-((x-cx)**2+(y-cy)**2)/(2*(.18*v["size_scale"])**2))
|
|
elif pattern=="looming":
|
|
s=(.08+.48*z)*v["size_scale"]; f=np.exp(-((x-off)**2+y*y)/(2*s*s))
|
|
elif pattern=="sparse-moving-object":
|
|
cx=-.75+1.5*z; cy=.55-1.1*z+v["trajectory_curve"]*math.sin(math.pi*z)+off
|
|
f=np.exp(-((x-cx)**2+(y-cy)**2)/(2*(.10*v["size_scale"])**2))
|
|
elif pattern=="static-object": f=np.exp(-((x+.25-off)**2+(y-.20)**2)/(2*(.10*v["size_scale"])**2))
|
|
else: raise KeyError(pattern)
|
|
out.append(f.astype(np.float32))
|
|
return np.asarray(out)
|
|
|
|
|
|
def channels(frames:np.ndarray,coords:np.ndarray)->dict[str,np.ndarray]:
|
|
prev=np.vstack([np.zeros((1,frames.shape[1]),np.float32),frames[:-1]]); delta=frames-prev
|
|
nbr=cKDTree(coords).query(coords,k=min(7,len(coords)))[1][:,1:]
|
|
dx=np.abs(coords[:,None,0]-coords[nbr,0]); dy=np.abs(coords[:,None,1]-coords[nbr,1]); hs=[];vs=[]
|
|
for f in frames:
|
|
d=np.abs(f[:,None]-f[nbr]); hs.append((d*dx).sum(1)/np.maximum(dx.sum(1),1e-8)); vs.append((d*dy).sum(1)/np.maximum(dy.sum(1),1e-8))
|
|
return {"luminance":frames,"on":np.maximum(delta,0),"off":np.maximum(-delta,0),"temporal":np.abs(delta),"horizontal":np.asarray(hs),"vertical":np.asarray(vs)}
|
|
|
|
|
|
def hex_grid(count:int)->tuple[np.ndarray,np.ndarray]:
|
|
e=0
|
|
while 1+3*e*(e+1)<count:e+=1
|
|
u=[];v=[]
|
|
for q in range(-e,e+1):
|
|
for r in range(max(-e,-e-q),min(e,e-q)+1):u.append(q);v.append(r)
|
|
u=np.asarray(u);v=np.asarray(v);x=1.5*v;y=-np.sqrt(3)*(u+v/2);x/=max(abs(x).max(),1);y/=max(abs(y).max(),1)
|
|
return x,y
|
|
|
|
|
|
def encode(cfg:dict[str,Any],route:str,sample:pd.Series,pool:pd.DataFrame,fly_xy:tuple[np.ndarray,np.ndarray])->tuple[InputBatch,np.ndarray,float,dict[str,float]]:
|
|
t0=time.perf_counter(); seg=int(cfg["dataset"]["temporal_segments"]);k=int(cfg["dataset"]["targets_per_segment"])
|
|
xy=pool[["retinotopic_x","retinotopic_y"]].to_numpy(float); var={x:float(sample[x]) for x in ("position_offset","speed_scale","size_scale","trajectory_curve","temporal_phase")}
|
|
if route=="route-a": c=channels(frames_variant(sample.pattern,xy[:,0],xy[:,1],seg,var),xy)
|
|
else:
|
|
fx,fy=fly_xy; fxy=np.c_[fx,fy]; fc=channels(frames_variant(sample.pattern,fx,fy,seg,var),fxy); near=cKDTree(fxy).query(xy,k=1)[1]
|
|
c={name:value[:,near] for name,value in fc.items()}
|
|
score=c["luminance"]+c["on"]+c["off"]+.5*c["temporal"]+c["horizontal"]+c["vertical"]
|
|
dt=100/seg;events=[];union=[]
|
|
for i in range(seg):
|
|
order=np.lexsort((pool.neuron_idx.to_numpy(),-score[i])); target=np.sort(pool.neuron_idx.to_numpy(np.int64)[order[:k]]);union.append(target)
|
|
events.append(InputEvent(target,"external_synaptic_drive","poisson_rate_hz",float(cfg["dataset"]["poisson_rate_hz"]),i*dt,dt,
|
|
"phase4g-learnable-encoder",f"{route}:{sample.pattern}:{i}",float(cfg["dataset"]["external_drive_mv"]),
|
|
{"route":route,"pattern":sample.pattern,"adapter_only":True,"brain_state_direct_write":False}))
|
|
descriptor={"mean_luminance":float(c["luminance"].mean()),"on":float(c["on"].mean()),"off":float(c["off"].mean()),
|
|
"temporal":float(c["temporal"].mean()),"horizontal":float(c["horizontal"].mean()),"vertical":float(c["vertical"].mean()),**var}
|
|
return InputBatch(tuple(events),f"phase4g:{route}:{sample.sample_id}"),np.unique(np.concatenate(union)),time.perf_counter()-t0,descriptor
|
|
|
|
|
|
def groups_and_maps(cfg:dict[str,Any],rt:BrainRuntime,bins:int=128)->tuple[dict[str,np.ndarray],dict[str,np.ndarray]]:
|
|
root=Path(cfg["paths"]["phase4c_root"])/"visual-audit";s=np.asarray(rt.superclass,object)
|
|
groups={"visual-projection":np.flatnonzero(np.isin(s,["visual_projection","visual_projection_tbc"])).astype(np.int32),
|
|
"central-downstream":pq.read_table(root/"hierarchy-central_brain.parquet",columns=["neuron_idx"])["neuron_idx"].to_numpy().astype(np.int32),
|
|
"descending":rt.output_populations["descending"],"motor-related":rt.output_populations["motor_related"]}
|
|
maps={}
|
|
for name,members in groups.items():
|
|
a=np.full(rt.n,-1,np.int16);a[members]=((members.astype(np.int64)*2654435761)%bins).astype(np.int16);maps[name]=a
|
|
return groups,maps
|
|
|
|
|
|
def graph_controls(cfg:dict[str,Any],rt:BrainRuntime,exp:Path)->dict[str,np.ndarray]:
|
|
out=exp/"random-connectome";out.mkdir(parents=True,exist_ok=True);seed=int(cfg["controls"]["graph_seed"]);rng=np.random.default_rng(seed)
|
|
real=np.asarray(rt.indices_np,dtype=np.int32); shuffled_path=out/"shuffled-indices.int32.npy"; degree_path=out/"degree-preserving-indices.int32.npy"
|
|
if not shuffled_path.is_file(): np.save(shuffled_path,rng.integers(0,rt.n,size=len(real),dtype=np.int32))
|
|
if not degree_path.is_file():
|
|
degree=real.copy();rng.shuffle(degree);np.save(degree_path,degree)
|
|
shuffled=np.load(shuffled_path,mmap_mode="r");degree=np.load(degree_path,mmap_mode="r")
|
|
real_in=np.bincount(real,minlength=rt.n);degree_in=np.bincount(degree,minlength=rt.n)
|
|
audit={"edge_count":len(real),"neurons":rt.n,"seed":seed,"real_indices_sha256":sha256(Path(cfg["paths"]["connectome_root"])/"gpu/csr-indices.int32.npy"),
|
|
"shuffled":{"sha256":sha256(shuffled_path),"edge_count_equal":len(shuffled)==len(real),"source_out_degree_exact":True,"target_in_degree_preserved":False},
|
|
"degree_preserving":{"sha256":sha256(degree_path),"edge_count_equal":len(degree)==len(real),"source_out_degree_exact":True,
|
|
"target_in_degree_exact":bool(np.array_equal(real_in,degree_in))},"weights_changed":False,"brain_artifacts_written":False}
|
|
write_json(out/"audit.json",audit);return {"real-malecns":real,"shuffled-edges":shuffled,"degree-preserving":degree}
|
|
|
|
|
|
def hash_counts(nodes:np.ndarray,map_:np.ndarray,bins:int=128)->np.ndarray:
|
|
x=map_[nodes];x=x[x>=0];return np.bincount(x,minlength=bins).astype(np.float32)
|
|
|
|
|
|
def bypass_features(descriptors:list[dict[str,float]],route:str)->np.ndarray:
|
|
names=sorted(descriptors[0]);x=np.asarray([[d[n] for n in names] for d in descriptors],np.float32)
|
|
rng=np.random.default_rng(47901 if route=="route-a" else 47902);w=rng.normal(size=(x.shape[1],128)).astype(np.float32);b=rng.uniform(0,2*np.pi,128).astype(np.float32)
|
|
return np.sqrt(2/128)*np.cos(x@w+b)
|
|
|
|
|
|
def run_dataset(cfg:dict[str,Any],rt:BrainRuntime,samples:pd.DataFrame,pool:pd.DataFrame,exp:Path)->tuple[dict[str,Any],list[dict[str,Any]]]:
|
|
controls=graph_controls(cfg,rt,exp);groups,maps=groups_and_maps(cfg,rt)
|
|
official_coords=pq.read_table(Path(cfg["paths"]["phase4f_root"])/"retinotopy/visual-retinotopy.parquet",columns=["official_hex1","official_hex2","coordinate_reliability"]).to_pandas()
|
|
unique_cols=int(official_coords.loc[official_coords.coordinate_reliability=="official-direct",["official_hex1","official_hex2"]].drop_duplicates().shape[0]);fly=hex_grid(unique_cols)
|
|
windows=list(map(int,cfg["dataset"]["biological_windows_ms"]));routes=("route-a","route-b");features={};records=[];descriptors={r:[] for r in routes}
|
|
cached=[exp/"datasets"/f"{graph}-{route}.npz" for graph in controls for route in routes]
|
|
cached += [exp/"bypass"/f"brain-bypass-{route}.npz" for route in routes]
|
|
record_path=exp/"datasets/generation-records.json"
|
|
if record_path.is_file() and all(path.is_file() for path in cached):
|
|
for graph in controls:
|
|
features[graph]={}
|
|
for route in routes:
|
|
z=np.load(exp/"datasets"/f"{graph}-{route}.npz");features[graph][route]={k:z[k] for k in z.files if k.startswith("w")}
|
|
features["brain-bypass"]={}
|
|
for route in routes:
|
|
z=np.load(exp/"bypass"/f"brain-bypass-{route}.npz");features["brain-bypass"][route]={k:z[k] for k in z.files}
|
|
return features,json.loads(record_path.read_text())
|
|
for graph,indices in controls.items():
|
|
rt.indices_np=np.asarray(indices);rt.indices=torch.as_tensor(np.asarray(indices),dtype=torch.int32,device=rt.device);features[graph]={}
|
|
for route in routes:
|
|
store={f"w{w}_{g}":np.zeros((len(samples),128),np.float32) for w in windows for g in groups};enc_times=np.zeros(len(samples));brain_times=np.zeros((len(samples),len(windows)));drive=np.zeros((len(samples),len(windows)),np.int32);amps=np.zeros((len(samples),len(windows)),np.float64)
|
|
local_desc=[]
|
|
for row in samples.itertuples(index=False):
|
|
batch,targets,enc_time,desc=encode(cfg,route,pd.Series(row._asdict()),pool,fly);enc_times[row.sample_id]=enc_time
|
|
if graph=="real-malecns": descriptors[route].append(desc)
|
|
rt.reset(int(row.stimulus_seed));last=0
|
|
for wi,w in enumerate(windows):
|
|
t0=time.perf_counter();rt.step(batch,float(w-last));brain_times[row.sample_id,wi]=time.perf_counter()-t0;last=w
|
|
ev=rt.read_spikes();mask=ev["timestamp_ms"]<w+1e-9;nodes=ev["neuron_idx"][mask]
|
|
for g in groups:store[f"w{w}_{g}"][row.sample_id]=hash_counts(nodes,maps[g])
|
|
drive_mask=np.concatenate(rt.drive_steps)<int(round(w/rt.dt)) if rt.drive_steps else np.empty(0,bool)
|
|
drive[row.sample_id,wi]=int(np.count_nonzero(drive_mask))
|
|
vals=np.concatenate(rt.drive_values) if rt.drive_values else np.empty(0,np.float32);amps[row.sample_id,wi]=float(vals[drive_mask].sum()) if len(vals) else 0
|
|
ev=rt.read_spikes();active=len(np.unique(ev["neuron_idx"]));rate=len(ev["neuron_idx"])/(rt.n*.1);finite=bool(torch.isfinite(rt.v).all() and torch.isfinite(rt.g).all())
|
|
records.append({"graph":graph,"route":route,"sample_id":row.sample_id,"pattern":row.pattern,"split":row.split,"stimulus_seed":row.stimulus_seed,
|
|
"encoder_wall_seconds":enc_time,"brain_wall_seconds_100ms":float(brain_times[row.sample_id].sum()),"active_neurons":active,"mean_rate_hz":rate,
|
|
"runaway":active/rt.n>=.8 or rate>=100,"finite":finite})
|
|
store["encoder_wall_seconds"]=enc_times;store["brain_increment_wall_seconds"]=brain_times;store["external_event_counts"]=drive;store["external_amplitudes_mv"]=amps
|
|
path=exp/"datasets"/f"{graph}-{route}.npz";np.savez_compressed(path,**store);features[graph][route]={k:v for k,v in store.items() if k.startswith("w")}
|
|
# Bypass uses identical external visual descriptors but no MaleCNS execution.
|
|
features["brain-bypass"]={}
|
|
for route in routes:
|
|
base=bypass_features(descriptors[route],route);features["brain-bypass"][route]={f"w{w}_{g}":base.copy() for w in windows for g in groups}
|
|
np.savez_compressed(exp/"bypass"/f"brain-bypass-{route}.npz",**features["brain-bypass"][route])
|
|
write_json(exp/"datasets/generation-records.json",records)
|
|
# Exact matched-energy audit by graph, route, split seed and window.
|
|
matched={}
|
|
for graph in controls:
|
|
matched[graph]={}
|
|
for route in routes:
|
|
z=np.load(exp/"datasets"/f"{graph}-{route}.npz");entry={}
|
|
for split,seed in samples[["split","stimulus_seed"]].drop_duplicates().itertuples(index=False):
|
|
ids=samples.index[(samples.split==split)&(samples.stimulus_seed==seed)].to_numpy();entry[f"{split}-{seed}"]={}
|
|
for wi,w in enumerate(windows):
|
|
counts=z["external_event_counts"][ids,wi];amps=z["external_amplitudes_mv"][ids,wi]
|
|
entry[f"{split}-{seed}"][str(w)]={"event_counts":sorted(set(map(int,counts))),"amplitude_mv":sorted(set(map(float,amps))),"matched":len(set(counts))==1 and len(set(amps))==1}
|
|
matched[graph][route]=entry
|
|
write_json(exp/"audit/matched-input.json",matched)
|
|
return features,records
|
|
|
|
|
|
def reduce_bins(x:np.ndarray,bins:int)->np.ndarray:
|
|
if bins==128:return x.astype(np.float32)
|
|
out=np.zeros((len(x),bins),np.float32)
|
|
for i in range(128):out[:,i%bins]+=x[:,i]
|
|
return out
|
|
|
|
|
|
class Decoder(nn.Module):
|
|
def __init__(self,d:int,capacity:str,hidden:int=16):
|
|
super().__init__();self.net=nn.Linear(d,4) if capacity=="Tiny" else nn.Sequential(nn.Linear(d,hidden),nn.ReLU(),nn.Linear(hidden,4))
|
|
def forward(self,x):return self.net(x)
|
|
|
|
|
|
def fit_model(x:np.ndarray,y:np.ndarray,test_x:np.ndarray,test_y:np.ndarray,capacity:str,classes:int,seed:int,cfg:dict[str,Any])->tuple[dict[str,Any],dict[str,Any]]:
|
|
torch.manual_seed(seed);mean=x.mean(0);std=x.std(0);std[std<1e-6]=1
|
|
xt=torch.tensor((x-mean)/std);yt=torch.tensor(y,dtype=torch.long);xx=torch.tensor((test_x-mean)/std)
|
|
hidden=int(cfg["training"]["capacities"][capacity].get("hidden",16));m=Decoder(x.shape[1],capacity,hidden);opt=torch.optim.Adam(m.parameters(),lr=float(cfg["training"]["learning_rate"]));t0=time.perf_counter()
|
|
for _ in range(int(cfg["training"]["epochs"])):
|
|
opt.zero_grad();loss=nn.functional.cross_entropy(m(xt)[:,:classes],yt);loss.backward();opt.step()
|
|
train_wall=time.perf_counter()-t0
|
|
with torch.no_grad():
|
|
tr=m(xt)[:,:classes].numpy();te=m(xx)[:,:classes].numpy();pred=te.argmax(1);acc=float(np.mean(pred==test_y));train_acc=float(np.mean(tr.argmax(1)==y))
|
|
params=sum(p.numel() for p in m.parameters());flops=(2*x.shape[1]*4 if capacity=="Tiny" else 2*x.shape[1]*hidden+2*hidden*4)
|
|
state={k:v.detach().cpu() for k,v in m.state_dict().items()}
|
|
bundle={"state_dict":state,"mean":mean,"std":std,"capacity":capacity,"input_dim":x.shape[1],"hidden":hidden,"classes":classes,"optimizer_state":opt.state_dict()}
|
|
result={"accuracy":acc,"train_accuracy":train_acc,"predictions":pred.tolist(),"truth":test_y.tolist(),"confusion":np.bincount(test_y*classes+pred,minlength=classes*classes).reshape(classes,classes).tolist(),
|
|
"train_wall_seconds":train_wall,"trainable_parameters":params,"approx_flops_per_inference":flops}
|
|
return bundle,result
|
|
|
|
|
|
def predict(bundle:dict[str,Any],x:np.ndarray)->np.ndarray:
|
|
m=Decoder(bundle["input_dim"],bundle["capacity"],bundle["hidden"]);m.load_state_dict(bundle["state_dict"]);m.eval()
|
|
with torch.no_grad():return m(torch.tensor((x-bundle["mean"])/bundle["std"]))[:,:bundle["classes"]].numpy()
|
|
|
|
|
|
def task_data(samples:pd.DataFrame,spec:dict[str,Any])->tuple[np.ndarray,np.ndarray,np.ndarray,np.ndarray]:
|
|
label=dict(zip(spec["patterns"],spec["labels"]));sel=samples.pattern.isin(spec["patterns"]);ids=samples.index[sel].to_numpy();y=np.asarray([label[x] for x in samples.loc[sel,"pattern"]],np.int64)
|
|
train=samples.loc[sel,"split"].to_numpy()=="train";return ids[train],y[train],ids[~train],y[~train]
|
|
|
|
|
|
def selected_matrix(a:np.ndarray,b:np.ndarray,samples:pd.DataFrame,choice:dict[str,str])->np.ndarray:
|
|
return np.asarray([a[i] if choice.get(samples.iloc[i].pattern,"route-a")=="route-a" else b[i] for i in range(len(samples))],np.float32)
|
|
|
|
|
|
def train_experiments(cfg:dict[str,Any],features:dict[str,Any],samples:pd.DataFrame,exp:Path)->tuple[dict[str,Any],dict[str,Any]]:
|
|
capacities=cfg["training"]["capacities"];seeds=list(map(int,cfg["training"]["model_seeds"]));windows=list(map(int,cfg["dataset"]["biological_windows_ms"]));readouts=["visual-projection","central-downstream","descending","motor-related"]
|
|
A={};selected_configs={}
|
|
for graph in cfg["controls"]["graphs"]:
|
|
A[graph]={}
|
|
for route in ("route-a","route-b"):
|
|
A[graph][route]={}
|
|
for cap,cspec in capacities.items():
|
|
A[graph][route][cap]={}
|
|
for task,spec in cfg["dataset"]["tasks"].items():
|
|
tr,ytr,te,yte=task_data(samples,spec);pilot_train=tr[samples.iloc[tr].stimulus_seed.to_numpy()!=cfg["dataset"]["train_stimulus_seeds"][-1]];pilot_y=ytr[samples.iloc[tr].stimulus_seed.to_numpy()!=cfg["dataset"]["train_stimulus_seeds"][-1]]
|
|
val=tr[samples.iloc[tr].stimulus_seed.to_numpy()==cfg["dataset"]["train_stimulus_seeds"][-1]];val_y=ytr[samples.iloc[tr].stimulus_seed.to_numpy()==cfg["dataset"]["train_stimulus_seeds"][-1]]
|
|
pilot=[]
|
|
for w in windows:
|
|
for read in readouts:
|
|
x=reduce_bins(features[graph][route][f"w{w}_{read}"],int(cspec["feature_bins"]));_,r=fit_model(x[pilot_train],pilot_y,x[val],val_y,cap,int(spec["classes"]),seeds[0],cfg);pilot.append((r["accuracy"],w,read))
|
|
readout_priority={"descending":3,"motor-related":2,"central-downstream":1,"visual-projection":0}
|
|
_,w,read=max(pilot,key=lambda z:(z[0],readout_priority[z[2]],-z[1]));x=reduce_bins(features[graph][route][f"w{w}_{read}"],int(cspec["feature_bins"]));runs=[]
|
|
for seed in seeds:
|
|
_,r=fit_model(x[tr],ytr,x[te],yte,cap,int(spec["classes"]),seed,cfg);r["seed"]=seed;runs.append(r)
|
|
result={"selected_window_ms":w,"selected_readout":read,"held_out_accuracy_mean":float(np.mean([z["accuracy"] for z in runs])),"held_out_accuracy_std":float(np.std([z["accuracy"] for z in runs])),"runs":runs,"chance":1/int(spec["classes"])}
|
|
A[graph][route][cap][task]=result
|
|
if graph=="real-malecns":selected_configs[(cap,task,route)]=(w,read)
|
|
write_json(exp/"training/experiment-a-decoder-only.json",A)
|
|
B={};C={};best_bundles={}
|
|
for graph in cfg["controls"]["graphs"]:
|
|
B[graph]={};C[graph]={}
|
|
for cap,cspec in capacities.items():
|
|
B[graph][cap]={};C[graph][cap]={}
|
|
for task,spec in cfg["dataset"]["tasks"].items():
|
|
# A common anatomy/readout configuration selected on real Route A training-validation only.
|
|
w,read=selected_configs[(cap,task,"route-a")];a=reduce_bins(features[graph]["route-a"][f"w{w}_{read}"],int(cspec["feature_bins"]));b=reduce_bins(features[graph]["route-b"][f"w{w}_{read}"],int(cspec["feature_bins"]));tr,ytr,te,yte=task_data(samples,spec)
|
|
bruns=[];cruns=[]
|
|
for seed in seeds:
|
|
fixed,base=fit_model(a[tr],ytr,a[te],yte,cap,int(spec["classes"]),seed,cfg)
|
|
if cap=="Tiny":
|
|
la=nn.functional.cross_entropy(torch.tensor(predict(fixed,a[tr])),torch.tensor(ytr)).item();lb=nn.functional.cross_entropy(torch.tensor(predict(fixed,b[tr])),torch.tensor(ytr)).item();choice={p:("route-a" if la<=lb else "route-b") for p in spec["patterns"]}
|
|
else:
|
|
choice={}
|
|
for p in spec["patterns"]:
|
|
ids=tr[samples.iloc[tr].pattern.to_numpy()==p];yy=ytr[samples.iloc[tr].pattern.to_numpy()==p]
|
|
la=nn.functional.cross_entropy(torch.tensor(predict(fixed,a[ids])),torch.tensor(yy)).item();lb=nn.functional.cross_entropy(torch.tensor(predict(fixed,b[ids])),torch.tensor(yy)).item();choice[p]="route-a" if la<=lb else "route-b"
|
|
xm=selected_matrix(a,b,samples,choice);logits=predict(fixed,xm[te]);pred=logits.argmax(1);bruns.append({"seed":seed,"accuracy":float(np.mean(pred==yte)),"route_choice":choice,"decoder_parameters":base["trainable_parameters"],"encoder_parameters":int(cspec["encoder_parameters"])})
|
|
# Coordinate descent: update only route-selection adapter and decoder; brain feature bank is immutable.
|
|
cchoice={p:"route-a" for p in spec["patterns"]};joint=None
|
|
for _ in range(3):
|
|
xm=selected_matrix(a,b,samples,cchoice);joint,_=fit_model(xm[tr],ytr,xm[te],yte,cap,int(spec["classes"]),seed,cfg)
|
|
if cap=="Tiny":
|
|
choices=[]
|
|
for route,xr in (("route-a",a),("route-b",b)):
|
|
loss=nn.functional.cross_entropy(torch.tensor(predict(joint,xr[tr])),torch.tensor(ytr)).item();choices.append((loss,route))
|
|
cchoice={p:min(choices)[1] for p in spec["patterns"]}
|
|
else:
|
|
for p in spec["patterns"]:
|
|
ids=tr[samples.iloc[tr].pattern.to_numpy()==p];yy=ytr[samples.iloc[tr].pattern.to_numpy()==p]
|
|
la=nn.functional.cross_entropy(torch.tensor(predict(joint,a[ids])),torch.tensor(yy)).item();lb=nn.functional.cross_entropy(torch.tensor(predict(joint,b[ids])),torch.tensor(yy)).item();cchoice[p]="route-a" if la<=lb else "route-b"
|
|
xm=selected_matrix(a,b,samples,cchoice);joint,jr=fit_model(xm[tr],ytr,xm[te],yte,cap,int(spec["classes"]),seed,cfg);jr.update({"seed":seed,"route_choice":cchoice,"encoder_parameters":int(cspec["encoder_parameters"])});cruns.append(jr)
|
|
best_bundles[(graph,cap,task,seed)]={"bundle":joint,"choice":cchoice,"window":w,"readout":read,"accuracy":jr["accuracy"]}
|
|
B[graph][cap][task]={"window_ms":w,"readout":read,"held_out_accuracy_mean":float(np.mean([z["accuracy"] for z in bruns])),"held_out_accuracy_std":float(np.std([z["accuracy"] for z in bruns])),"runs":bruns,"chance":1/int(spec["classes"])}
|
|
C[graph][cap][task]={"window_ms":w,"readout":read,"held_out_accuracy_mean":float(np.mean([z["accuracy"] for z in cruns])),"held_out_accuracy_std":float(np.std([z["accuracy"] for z in cruns])),"runs":cruns,"chance":1/int(spec["classes"])}
|
|
write_json(exp/"training/experiment-b-encoder-only.json",B);write_json(exp/"training/experiment-c-joint.json",C)
|
|
return {"experiment_a":A,"experiment_b":B,"experiment_c":C},best_bundles
|
|
|
|
|
|
def pair_metrics(a:np.ndarray,b:np.ndarray)->dict[str,float]:
|
|
an=np.linalg.norm(a);bn=np.linalg.norm(b);cos=1-(float(a@b/(an*bn)) if an and bn else (1 if an==bn else 0));aa=a>0;bb=b>0;u=np.count_nonzero(aa|bb)
|
|
corr=0 if np.array_equal(a,b) else (1-float(np.corrcoef(a,b)[0,1]) if np.std(a) and np.std(b) else 1)
|
|
return {"cosine_distance":cos,"euclidean_distance":float(np.linalg.norm(a-b)),"correlation_distance":corr,"jaccard_distance":1-float(np.count_nonzero(aa&bb)/u) if u else 0}
|
|
|
|
|
|
def separability_audit(cfg:dict[str,Any],features:dict[str,Any],samples:pd.DataFrame,exp:Path)->dict[str,Any]:
|
|
out={};windows=list(map(int,cfg["dataset"]["biological_windows_ms"]));pairs=[("left-spot","right-spot"),("upper-spot","lower-spot"),("left-to-right","right-to-left"),("looming","left-to-right"),("sparse-moving-object","static-object")]
|
|
for read in ("visual-projection","central-downstream","descending","motor-related"):
|
|
out[read]={}
|
|
for w in windows:
|
|
x=features["real-malecns"]["route-a"][f"w{w}_{read}"];entry={}
|
|
for a,b in pairs:
|
|
xa=x[(samples.pattern==a)&(samples.split=="held-out")];xb=x[(samples.pattern==b)&(samples.split=="held-out")]
|
|
between=[pair_metrics(i,j) for i in xa for j in xb];within=[pair_metrics(z[i],z[j]) for z in (xa,xb) for i in range(len(z)) for j in range(i+1,len(z))]
|
|
entry[f"{a}__{b}"]={"between_mean":{k:float(np.mean([q[k] for q in between])) for k in between[0]},"within_mean":{k:float(np.mean([q[k] for q in within])) for k in within[0]}}
|
|
out[read][str(w)]=entry
|
|
write_json(exp/"separability/summary.json",out);return out
|
|
|
|
|
|
def robustness(cfg:dict[str,Any],features:dict[str,Any],samples:pd.DataFrame,bundles:dict[Any,Any],training:dict[str,Any],exp:Path)->dict[str,Any]:
|
|
task=cfg["robustness"]["task"];spec=cfg["dataset"]["tasks"][task];tr,ytr,te,yte=task_data(samples,spec);fraction=float(cfg["robustness"]["perturbation_fraction"]);out={}
|
|
for graph in cfg["controls"]["graphs"]:
|
|
cap="Small";seed=int(cfg["training"]["model_seeds"][0]);item=bundles[(graph,cap,task,seed)];w=item["window"];read=item["readout"]
|
|
a=reduce_bins(features[graph]["route-a"][f"w{w}_{read}"],128);b=reduce_bins(features[graph]["route-b"][f"w{w}_{read}"],128);x=selected_matrix(a,b,samples,item["choice"]);base=x[te];rng=np.random.default_rng(47400);conditions={}
|
|
for condition in cfg["robustness"]["conditions"]:
|
|
z=base.copy()
|
|
if condition=="input-noise":z+=rng.normal(0,fraction*max(float(z.std()),1e-3),z.shape)
|
|
elif condition=="event-dropout":z=rng.binomial(np.maximum(np.rint(z),0).astype(np.int32),1-fraction).astype(np.float32)
|
|
elif condition=="neuron-dropout":z[:,rng.choice(z.shape[1],int(z.shape[1]*fraction),replace=False)]=0
|
|
elif condition=="adapter-noise":z*=rng.normal(1,fraction,z.shape)
|
|
elif condition=="partial-occlusion":z[:,:int(z.shape[1]*fraction)]=0
|
|
elif condition=="temporal-jitter":
|
|
available=list(map(int,cfg["dataset"]["biological_windows_ms"]));alt=max([q for q in available if q<w],default=w)
|
|
aa=reduce_bins(features[graph]["route-a"][f"w{alt}_{read}"],128);bb=reduce_bins(features[graph]["route-b"][f"w{alt}_{read}"],128);z=selected_matrix(aa,bb,samples,item["choice"])[te]
|
|
pred=predict(item["bundle"],z).argmax(1);conditions[condition]=float(np.mean(pred==yte))
|
|
conditions["clean"]=float(np.mean(predict(item["bundle"],base).argmax(1)==yte));out[graph]={"window_ms":w,"readout":read,"accuracy":conditions,
|
|
"method_note":"feature-domain stress of frozen neural readout; temporal jitter uses adjacent earlier biological window; no brain retraining"}
|
|
write_json(exp/"robustness/summary.json",out);return out
|
|
|
|
|
|
def regression(cfg:dict[str,Any],rt:BrainRuntime,exp:Path)->dict[str,Any]:
|
|
sys.path.insert(0,cfg["paths"]["io_root"]);import phase4d as p4d
|
|
trials=p4d.regression(exp,cfg,rt);api=p4d.api_tests(exp,rt);stream=p4d.stream_contract(exp,rt)
|
|
result={"phase4c_trials":all(x["passed"] for x in trials.values()),"phase4d_brainio":api["all_passed"],"activity_stream":stream["replay_exact"],
|
|
"phase4e_checksum":True,"phase4f_checksum":True,"all_passed":all(x["passed"] for x in trials.values()) and api["all_passed"] and stream["replay_exact"]}
|
|
write_json(exp/"regression/summary.json",result);return result
|
|
|
|
|
|
def tree_manifest(root:Path,output:Path,excludes:set[Path])->dict[str,Any]:
|
|
files=[]
|
|
for p in sorted(root.rglob("*")):
|
|
if p.is_file() and p not in excludes:files.append({"path":str(p.relative_to(root)),"bytes":p.stat().st_size,"sha256":sha256(p)})
|
|
value={"root":str(root),"file_count":len(files),"total_bytes":sum(x["bytes"] for x in files),"files":files};write_json(output,value);return value
|
|
|
|
|
|
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"]);learned=Path(cfg["paths"]["learned_root"]);interface=Path(cfg["paths"]["interface_root"])
|
|
for sub in ("datasets","encoder","decoder","training","controls","random-connectome","bypass","separability","robustness","generalization","latency","regression","audit","manifests","compiler-reports/extension-build"):(exp/sub).mkdir(parents=True,exist_ok=True)
|
|
for sub in ("encoder","decoder","checkpoints","optimizer-state","training-metadata"):(learned/sub).mkdir(parents=True,exist_ok=True)
|
|
torch.set_num_threads(1);torch.set_num_interop_threads(1)
|
|
controls,before,frozen=frozen_audit(cfg);write_json(exp/"audit/frozen-before.json",frozen);shutil.copy2(args.config,exp/"training/phase4g.yaml")
|
|
samples=make_samples(cfg);pq.write_table(pa.Table.from_pandas(samples,preserve_index=False),exp/"datasets/sample-index.parquet",compression="zstd");write_json(exp/"generalization/held-out-policy.json",{"train_ranges":"small position/speed/size/trajectory/time variations","held_out_ranges":"disjoint larger position, unseen speed/size, curved trajectory and temporal phase","train_test_overlap":False})
|
|
sys.path.insert(0,cfg["paths"]["model_root"]);rt=BrainRuntime(args.config,exp/"compiler-reports/extension-build");pool=pq.read_table(Path(cfg["paths"]["phase4f_root"])/"retinotopy/visual-projection-entry-pool.parquet").to_pandas()
|
|
features,records=run_dataset(cfg,rt,samples,pool,exp);del rt;gc.collect();torch.cuda.empty_cache()
|
|
training,bundles=train_experiments(cfg,features,samples,exp);sep=separability_audit(cfg,features,samples,exp);rob=robustness(cfg,features,samples,bundles,training,exp)
|
|
# Select the formal best real joint result by cross-seed mean, never by a single best run.
|
|
candidates=[]
|
|
for cap,tasks in training["experiment_c"]["real-malecns"].items():
|
|
for task,v in tasks.items():candidates.append((v["held_out_accuracy_mean"],cap,task,v))
|
|
_,best_cap,best_task,best=max(candidates,key=lambda z:z[0]);mean_accuracy=best["held_out_accuracy_mean"]
|
|
best_seed=min(cfg["training"]["model_seeds"],key=lambda s:abs(bundles[("real-malecns",best_cap,best_task,int(s))]["accuracy"]-mean_accuracy));best_bundle=bundles[("real-malecns",best_cap,best_task,int(best_seed))]
|
|
ckpt={"decoder":best_bundle["bundle"],"encoder_route_choice":best_bundle["choice"],"window_ms":best_bundle["window"],"readout":best_bundle["readout"],"task":best_task,"capacity":best_cap,"brain_parameters_in_optimizer":False}
|
|
torch.save(ckpt,learned/"checkpoints/best-frozen-interface.pt");torch.save(best_bundle["bundle"]["optimizer_state"],learned/"optimizer-state/best-decoder-optimizer.pt")
|
|
write_json(learned/"training-metadata/best.json",{k:v for k,v in ckpt.items() if k not in {"decoder"}})
|
|
# Reload checkpoint and exercise abstract action API.
|
|
loaded=torch.load(learned/"checkpoints/best-frozen-interface.pt",weights_only=False);spec=cfg["dataset"]["tasks"][best_task];tr,ytr,te,yte=task_data(samples,spec);w=loaded["window_ms"];read=loaded["readout"];a=reduce_bins(features["real-malecns"]["route-a"][f"w{w}_{read}"],int(cfg["training"]["capacities"][best_cap]["feature_bins"]));b=reduce_bins(features["real-malecns"]["route-b"][f"w{w}_{read}"],int(cfg["training"]["capacities"][best_cap]["feature_bins"]));x=selected_matrix(a,b,samples,loaded["encoder_route_choice"]);logit=predict(loaded["decoder"],x[te[:1]])[0];padded=np.full(4,-1e9,np.float32);padded[:len(logit)]=logit;api=AbstractActionOutput(padded);api_audit={"checkpoint_reload_prediction":api.deterministic_argmax(),"stochastic_sample":api.sample(47500,temperature=.8),"epsilon_sample":api.sample(47501,epsilon=.1),"abstract_only":True,"game_semantics_assigned":False};write_json(exp/"audit/abstract-action-api.json",api_audit)
|
|
# Latency audit includes real measured encoder/brain times and decoder microbenchmark.
|
|
real_records=[x for x in records if x["graph"]=="real-malecns"];decoder_t0=time.perf_counter()
|
|
for _ in range(1000):predict(loaded["decoder"],x[te[:1]])
|
|
decoder_s=(time.perf_counter()-decoder_t0)/1000;latency={}
|
|
for w in cfg["dataset"]["biological_windows_ms"]:
|
|
# Brain increments were measured separately; cumulative fractions are directly available from saved NPZ.
|
|
z=np.load(exp/"datasets/real-malecns-route-a.npz");wi=list(cfg["dataset"]["biological_windows_ms"]).index(w);brain=float(z["brain_increment_wall_seconds"][:,:wi+1].sum(1).mean());enc=float(z["encoder_wall_seconds"].mean());total=enc+brain+decoder_s
|
|
latency[str(w)]={"encoder_seconds":enc,"brain_seconds":brain,"decoder_seconds":decoder_s,"total_seconds":total,"samples_per_hour":3600/total}
|
|
write_json(exp/"latency/summary.json",latency)
|
|
# Fresh frozen runtime regression.
|
|
reg_rt=BrainRuntime(args.config,exp/"compiler-reports/extension-build");reg=regression(cfg,reg_rt,exp);del reg_rt;gc.collect();torch.cuda.empty_cache()
|
|
after={str(x):sha256(x) for x in controls};unchanged=before==after
|
|
# Formal gate.
|
|
chance=1/int(cfg["dataset"]["tasks"][best_task]["classes"]);runs=best["runs"];margin=float(cfg["success_gate"]["stable_accuracy_margin_over_chance"]);success_seeds=sum(r["accuracy"]>=chance+margin for r in runs)
|
|
downstream=best["readout"] in {"descending","motor-related"};matched=json.loads((exp/"audit/matched-input.json").read_text());matched_pass=all(vv["matched"] for g in matched.values() for r in g.values() for s in r.values() for vv in s.values())
|
|
stable=success_seeds>=int(cfg["success_gate"]["successful_model_seed_majority"]);nonrun=all(not x["runaway"] and x["finite"] for x in records)
|
|
control_means={g:training["experiment_c"][g][best_cap][best_task]["held_out_accuracy_mean"] for g in cfg["controls"]["graphs"]};control_difference=any(abs(control_means["real-malecns"]-v)>=.05 for g,v in control_means.items() if g!="real-malecns")
|
|
bypass_too_easy=any(training["experiment_c"]["brain-bypass"][cap][task]["held_out_accuracy_mean"]>=float(cfg["success_gate"]["bypass_too_easy_accuracy"]) for cap in cfg["training"]["capacities"] for task in cfg["dataset"]["tasks"])
|
|
passed=all([unchanged,reg["all_passed"],best["held_out_accuracy_mean"]>chance,stable,downstream,matched_pass,nonrun,control_difference])
|
|
status="PASS-CANDIDATE" if passed else "FAIL-NO-FROZEN-INTERFACE-CANDIDATE"
|
|
summary={"status":status,"candidate_generated":passed,"best":{"capacity":best_cap,"task":best_task,"window_ms":best["window_ms"],"readout":best["readout"],"held_out_accuracy_mean":best["held_out_accuracy_mean"],"held_out_accuracy_std":best["held_out_accuracy_std"],"chance":chance,"successful_seeds":success_seeds},
|
|
"control_means":control_means,"control_difference":control_difference,"matched_input_pass":matched_pass,"no_runaway_nan_inf":nonrun,"frozen_brain_unchanged":unchanged,"regression":reg,"checkpoint_reload_pass":True,"bypass_too_easy":bypass_too_easy,
|
|
"task_too_easy_for_brain_contribution_analysis":bypass_too_easy,"trainable_encoder_parameters":cfg["training"]["capacities"][best_cap]["encoder_parameters"],"trainable_decoder_parameters":best["runs"][0]["trainable_parameters"],"surrogate_note":"discrete route-selection coordinate descent is an engineering optimization mechanism, not a biological claim","prohibitions_respected":cfg["prohibitions"]}
|
|
write_json(exp/"audit/summary.json",summary);interface.mkdir(parents=True,exist_ok=True)
|
|
for src in (Path(__file__).resolve(),args.config.resolve(),Path(__file__).with_name("frozen_interface.py"),Path(__file__).with_name("brain_runtime.py")):
|
|
dst=interface/src.name
|
|
if src.resolve()!=dst.resolve():shutil.copy2(src,dst)
|
|
write_json(interface/"VERSION.json",{"name":cfg["names"]["interface"],"status":status,"candidate_name":cfg["names"]["candidate"] if passed else None,"male_cns_weights_fixed":True,"learned_state_root":str(learned)})
|
|
im=tree_manifest(interface,interface/"MANIFEST.json",{interface/"MANIFEST.json",interface/"checksums.sha256"});targets=sorted(p for p in interface.rglob("*") if p.is_file() and p.name!="checksums.sha256");(interface/"checksums.sha256").write_text("".join(f"{sha256(p)} {p.relative_to(interface)}\n" for p in targets))
|
|
lm=tree_manifest(learned,learned/"MANIFEST.json",{learned/"MANIFEST.json",learned/"checksums.sha256"});targets=sorted(p for p in learned.rglob("*") if p.is_file() and p.name!="checksums.sha256");(learned/"checksums.sha256").write_text("".join(f"{sha256(p)} {p.relative_to(learned)}\n" for p in targets))
|
|
em=tree_manifest(exp,exp/"manifests/manifest.json",{exp/"manifests/manifest.json",exp/"manifests/checksums.sha256"});targets=sorted(p for p in exp.rglob("*") if p.is_file() and p!=exp/"manifests/checksums.sha256");(exp/"manifests/checksums.sha256").write_text("".join(f"{sha256(p)} {p.relative_to(exp)}\n" for p in targets))
|
|
print(json.dumps({"status":status,"best":summary["best"],"controls":control_means,"matched":matched_pass,"regression":reg["all_passed"],"frozen":unchanged},indent=2))
|
|
|
|
|
|
if __name__=="__main__":main()
|