Digital Drosophila Archive
485 lines
38 KiB
Python
485 lines
38 KiB
Python
#!/usr/bin/env python3
|
|
"""Phase 4Q: local terminal credit assignment v3 on deterministic-runtime-v2."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import platform
|
|
import resource
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
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 yaml
|
|
|
|
import phase4p
|
|
import phase4p_finalize
|
|
from phase4h_audit import verify_checksum
|
|
from phase4p_revalidate import DeterministicAPLSim, IsolatedReplay, door_inputs, make_drive
|
|
|
|
|
|
ROOT = Path("/data/digital-drosophila")
|
|
|
|
|
|
def save(path: Path, x: object) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(x, indent=2, sort_keys=True, default=str) + "\n")
|
|
|
|
|
|
def sha(*arrays: np.ndarray) -> str:
|
|
h = hashlib.sha256()
|
|
for x in arrays: h.update(np.ascontiguousarray(x).tobytes())
|
|
return h.hexdigest()
|
|
|
|
|
|
def dirs(exp: Path) -> None:
|
|
for name in ("calibration", "timing", "eligibility-v3", "unit-tests", "edge-credit", "random-us",
|
|
"distant-us", "isolated", "shared-edge", "b-contamination", "whole-brain", "controls",
|
|
"performance", "regression", "audit", "manifests"):
|
|
(exp / name).mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def load(cfgpath: Path):
|
|
cfg = yaml.safe_load(cfgpath.read_text()); exp = Path(cfg["paths"]["experiment_root"]); model = Path(cfg["paths"]["model_root"])
|
|
dirs(exp); model.mkdir(parents=True, exist_ok=True)
|
|
pcfg = yaml.safe_load((Path(cfg["paths"]["deterministic_runtime"]) / "phase4p.yaml").read_text())
|
|
rt = phase4p.CanonicalRuntime(pcfg, Path(cfg["paths"]["deterministic_runtime"]), exp)
|
|
sim = DeterministicAPLSim(rt)
|
|
identity = pd.read_parquet(cfg["paths"]["edge_identity"]).sort_values("edge_id").reset_index(drop=True)
|
|
expected = cfg["frozen_update"]["groups"]
|
|
actual = identity.identity_group.value_counts().to_dict()
|
|
if any(int(actual.get(k, 0)) != int(v) for k, v in expected.items()): raise RuntimeError(f"frozen group mismatch {actual}")
|
|
if not np.array_equal(identity.edge_id.to_numpy(np.int64), sim.identity.sort_values("edge_id").edge_id.to_numpy(np.int64)):
|
|
raise RuntimeError("plastic edge identity mismatch")
|
|
# sim identity order already follows frozen parquet; preserve it for E arrays.
|
|
identity = sim.identity.copy().reset_index(drop=True)
|
|
return cfg, exp, model, pcfg, rt, sim, identity
|
|
|
|
|
|
def random_us_steps(base: int, seed: int, count: int, nsteps: int) -> np.ndarray:
|
|
rng = np.random.default_rng(base + seed)
|
|
return rng.integers(0, nsteps, size=count, dtype=np.int32)
|
|
|
|
|
|
def last_spike_trace(steps: np.ndarray, nodes: np.ndarray, kc_lookup: np.ndarray, us_step: int,
|
|
window_ms: float, tau_ms: float, dt_ms: float) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Causal v3: only the last source-terminal spike at/before the US counts."""
|
|
last = np.full(int(kc_lookup.max()) + 1, -1, np.int32)
|
|
valid = steps <= us_step
|
|
for step, node in zip(steps[valid], nodes[valid]):
|
|
loc = kc_lookup[int(node)]
|
|
if loc >= 0: last[loc] = int(step)
|
|
age_ms = np.where(last >= 0, (us_step - last).astype(np.float32) * np.float32(dt_ms), np.float32(np.inf))
|
|
trace = np.zeros(len(last), np.float32); ok = (age_ms >= 0) & (age_ms <= np.float32(window_ms))
|
|
trace[ok] = np.exp(-age_ms[ok] / np.float32(tau_ms)).astype(np.float32)
|
|
return trace, last
|
|
|
|
|
|
def v2_trace(steps: np.ndarray, nodes: np.ndarray, kc_lookup: np.ndarray, us_step: int, dt_ms: float) -> np.ndarray:
|
|
trace = np.zeros(int(kc_lookup.max()) + 1, np.float32); decay = np.float32(math.exp(-dt_ms / 100.0))
|
|
by = {}
|
|
valid = steps <= us_step
|
|
for step in np.unique(steps[valid]): by[int(step)] = kc_lookup[nodes[(steps == step) & valid]]
|
|
for step in range(us_step + 1):
|
|
trace = np.float32(trace * decay); active = by.get(step)
|
|
if active is not None and len(active): trace[active] = np.minimum(np.float32(1), np.float32(trace[active] + np.float32(1)))
|
|
return trace
|
|
|
|
|
|
def edge_elig(trace: np.ndarray, source_local: np.ndarray, compartment_mask: np.ndarray, teaching: bool = True) -> np.ndarray:
|
|
out = np.zeros(len(source_local), np.float32)
|
|
if teaching: out[compartment_mask] = trace[source_local[compartment_mask]]
|
|
return out
|
|
|
|
|
|
def update_E(E: np.ndarray, elig: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
|
before = E.copy(); active = elig > 0
|
|
E[active] = np.maximum(np.float32(0.75), np.float32(E[active] * np.float32(1.0 - np.float32(0.01) * elig[active])))
|
|
return E, np.float32(E - before)
|
|
|
|
|
|
def group_metrics(E: np.ndarray, base: np.ndarray, identity: pd.DataFrame) -> list[dict[str, Any]]:
|
|
rows = []
|
|
for group in ("A-specific", "shared", "B-specific", "inactive"):
|
|
mask = identity.identity_group.astype(str).to_numpy() == group; d = np.float32(E[mask] - 1)
|
|
rows.append({"group": group, "edges": int(mask.sum()), "updated_edges": int(np.count_nonzero(d)),
|
|
"mean_delta_E": float(np.mean(d)), "total_abs_delta_E": float(np.abs(d).sum()),
|
|
"total_weighted_delta_mv": float(np.sum(base[mask].astype(np.float64) * d.astype(np.float64)))})
|
|
return rows
|
|
|
|
|
|
def formal_odor_ids(panel: pd.DataFrame, selection: dict[str, Any]) -> tuple[str, str]:
|
|
return str(selection["ODOR_A"]["odor_id"]), str(selection["ODOR_B"]["odor_id"])
|
|
|
|
|
|
def unit_tests(cfg, exp, sim, identity, window, tau, source_local, comp_mask, write=True) -> dict[str, Any]:
|
|
dt = sim.dt; kc0 = int(sim.kc[0]); lookup = np.full(sim.n, -1, np.int32); lookup[sim.kc] = np.arange(len(sim.kc), dtype=np.int32)
|
|
us = 900; recent_step = us - int(round(min(window / 2, 2.0) / dt)); far = us - int(round((window + 50) / dt))
|
|
def tr(ss, nn, gate=True, w=window, t=tau, u=us):
|
|
x, _ = last_spike_trace(np.asarray(ss, np.int32), np.asarray(nn, np.int32), lookup, u, w, t, dt)
|
|
return edge_elig(x, source_local, comp_mask, gate)
|
|
active_edges = comp_mask & (identity.source_neuron_idx.to_numpy(np.int32) == kc0)
|
|
if not active_edges.any():
|
|
kc0 = int(identity.loc[comp_mask, "source_neuron_idx"].iloc[0]); recent_step = us - 10
|
|
active_edges = comp_mask & (identity.source_neuron_idx.to_numpy(np.int32) == kc0)
|
|
recent = tr([recent_step], [kc0]); inactive = tr([], []); dzero = tr([recent_step], [kc0], False)
|
|
unrelated = edge_elig(last_spike_trace(np.asarray([recent_step],np.int32),np.asarray([kc0],np.int32),lookup,us,window,tau,dt)[0], source_local, np.zeros_like(comp_mask), True)
|
|
future = tr([us + 1], [kc0]); distant = tr([far], [kc0]); same1 = tr([us], [kc0]); same2 = tr([us], [kc0])
|
|
# Checkpoint/restore the sufficient statistic: per-KC last spike step.
|
|
first_trace, last = last_spike_trace(np.asarray([recent_step],np.int32),np.asarray([kc0],np.int32),lookup,us,window,tau,dt)
|
|
restored = np.where(last >= 0, np.exp(-((us-last).astype(np.float32)*dt)/np.float32(tau)), 0).astype(np.float32)
|
|
restored[((us-last).astype(np.float32)*dt) > window] = 0
|
|
repeated = [sha(tr([us], [kc0])) for _ in range(10)]
|
|
tests = {
|
|
"recent_active_matched_credit": bool(np.any(recent[active_edges] > 0)),
|
|
"inactive_matched_no_credit": not np.any(inactive),
|
|
"D0_no_credit": not np.any(dzero),
|
|
"unrelated_compartment_no_credit": not np.any(unrelated),
|
|
"US_before_future_KC_no_credit": not np.any(future),
|
|
"far_KC_before_US_near_zero": float(np.max(distant)) == 0.0,
|
|
"same_timestamp_exact": np.array_equal(same1, same2) and bool(np.all(same1[active_edges] == 1)),
|
|
"checkpoint_restore_exact": np.array_equal(first_trace, restored),
|
|
"ten_repeats_exact": len(set(repeated)) == 1,
|
|
}
|
|
out = {"tests": tests, "all_pass": all(tests.values()), "window_ms": window, "tau_ms": tau,
|
|
"same_timestamp_contract": "KC terminal update before US sampling", "future_credit": 0.0,
|
|
"repeat_hashes": repeated}
|
|
if write: save(exp / "unit-tests/summary.json", out)
|
|
return out
|
|
|
|
|
|
def task_calibrate(cfgpath: Path) -> None:
|
|
cfg, exp, model, pcfg, rt, sim, identity = load(cfgpath)
|
|
wall0=time.perf_counter();cpu0=resource.getrusage(resource.RUSAGE_SELF);torch.cuda.reset_peak_memory_stats()
|
|
ocfg, panel, vectors, selection, mapping = door_inputs(); cal = panel[panel.split == "calibration"].iloc[:2]
|
|
labels = cfg["calibration"]["labels"]; seeds = list(map(int, cfg["calibration"]["seeds"])); duration = float(cfg["calibration"]["duration_ms"])
|
|
us_ms = float(cfg["calibration"]["matched_us_ms"]); us_step = int(round(us_ms / sim.dt)); nsteps = int(round(duration / sim.dt))
|
|
dan = sim.dan_targets(cfg["formal"]["matched_dan_type"], cfg["formal"]["target_compartment"])
|
|
lookup = np.full(sim.n, -1, np.int32); lookup[sim.kc] = np.arange(len(sim.kc), dtype=np.int32)
|
|
source_local = lookup[identity.source_neuron_idx.to_numpy(np.int32)]; comp = identity.compartment.astype(str).to_numpy() == cfg["formal"]["target_compartment"]
|
|
event_cache = {}; input_rows = []
|
|
for oi, (_, odor) in enumerate(cal.iterrows()):
|
|
label = labels[oi]; targets, rates = make_drive(ocfg, panel, vectors, selection, mapping, str(odor.odor_id), "physiological")
|
|
for seed in seeds:
|
|
matched = sim.run(targets, rates, seed + oi*100, duration, dan_targets=dan, dan_pulses=[us_ms], collect_events=True)
|
|
rand_step = int(random_us_steps(int(cfg["calibration"]["random_us_seed_base"]), seed + oi*100, 1, nsteps)[0])
|
|
random_run = sim.run(targets, rates, seed + oi*100, duration, dan_targets=dan, dan_pulses=[rand_step*sim.dt], collect_events=True)
|
|
event_cache[(label, seed, "matched")] = matched; event_cache[(label, seed, "random")] = random_run
|
|
input_rows.append({"label": label, "odor_id": str(odor.odor_id), "odor_name": str(odor.odor_name), "seed": seed,
|
|
"matched_us_step": us_step, "random_us_step": rand_step,
|
|
"matched_event_hash": sha(matched["kc_event_step"],matched["kc_event_node"]),
|
|
"random_event_hash": sha(random_run["kc_event_step"],random_run["kc_event_node"])})
|
|
pd.DataFrame(input_rows).to_parquet(exp / "calibration/frozen-stimuli.parquet", index=False)
|
|
windows = list(map(float, cfg["credit_v3"]["windows_ms"])); rows=[]
|
|
for window in windows:
|
|
tau = window / 3.0
|
|
for (label, seed, condition), run in event_cache.items():
|
|
u = us_step if condition == "matched" else int(next(x["random_us_step"] for x in input_rows if x["label"]==label and x["seed"]==seed))
|
|
trace, _ = last_spike_trace(run["kc_event_step"], run["kc_event_node"], lookup, u, window, tau, sim.dt)
|
|
elig = edge_elig(trace, source_local, comp)
|
|
rows.append({"window_ms":window,"tau_ms":tau,"label":label,"seed":seed,"condition":condition,
|
|
"US_step":u,"credited_edges":int(np.count_nonzero(elig)),"total_eligibility":float(elig.sum()),
|
|
"total_abs_delta_E_one_pair":float((np.float32(.01)*elig).sum())})
|
|
frame=pd.DataFrame(rows);frame.to_parquet(exp/"calibration/window-matrix.parquet",index=False)
|
|
summaries=[];th=cfg["calibration"]["thresholds"]
|
|
for window in windows:
|
|
x=frame[frame.window_ms==window];m=x[x.condition=="matched"].groupby("seed").total_eligibility.mean();r=x[x.condition=="random"].groupby("seed").total_eligibility.mean()
|
|
mm=float(m.mean());rr=float(r.mean());ratio=mm/rr if rr>0 else float("inf")
|
|
qualifies=bool((m>0).sum()>=int(th["minimum_nonzero_seeds"]) and ratio>=float(th["minimum_matched_random_ratio"])
|
|
and rr<=float(th["maximum_random_fraction"])*mm)
|
|
summaries.append({"window_ms":window,"tau_ms":window/3,"matched_credit_mean":mm,"random_credit_mean":rr,
|
|
"matched_random_ratio":ratio,"nonzero_matched_seeds":int((m>0).sum()),"distant_credit":0.0,
|
|
"D0_credit":0.0,"unrelated_credit":0.0,"future_credit":0.0,"qualifies":qualifies})
|
|
good=[x for x in summaries if x["qualifies"]]
|
|
if not good:
|
|
suite=[]
|
|
for x in summaries:
|
|
suite.append(unit_tests(cfg,exp,sim,identity,x["window_ms"],x["tau_ms"],source_local,comp,write=False))
|
|
save(exp/"unit-tests/window-suite.json",{"all_pass":all(x["all_pass"] for x in suite),"windows":suite})
|
|
save(exp/"unit-tests/summary.json",{"all_pass":all(x["all_pass"] for x in suite),"selected_window":None,
|
|
"reason":"all preregistered windows tested; no calibration window qualified","windows_tested":len(suite)})
|
|
groups=identity.identity_group.astype(str).to_numpy();timing=[]
|
|
for x in summaries:
|
|
eligible=comp&(source_local>=0)
|
|
for delta in (-1000,-500,-250,-100,-50,-20,-10,0,10,20,50,100,250,500,1000):
|
|
e=0.0 if delta<0 or delta>x["window_ms"] else math.exp(-delta/x["tau_ms"])
|
|
timing.append({"window_ms":x["window_ms"],"tau_ms":x["tau_ms"],"delta_t_ms":delta,"eligibility":e,
|
|
"updated_edge_count_unit_terminal":int(np.count_nonzero(eligible)) if e>0 else 0,
|
|
"total_abs_delta_E_unit_terminal":float(.01*e*np.count_nonzero(eligible)),
|
|
"A_specific_credit":float(.01*e*np.count_nonzero(eligible&(groups=="A-specific"))),
|
|
"shared_credit":float(.01*e*np.count_nonzero(eligible&(groups=="shared"))),
|
|
"B_non_target_credit":float(.01*e*np.count_nonzero(eligible&(groups=="B-specific"))),"future_credit":0.0 if delta<0 else None})
|
|
pd.DataFrame(timing).to_parquet(exp/"timing/all-preregistered-windows.parquet",index=False)
|
|
out={"status":"FAIL-NO-V3-CALIBRATION","candidate":None,"windows":summaries,
|
|
"failure_reason":"No preregistered local terminal window made fixed 90 ms matched-US credit exceed random-US timing credit under CAL_A/CAL_B.",
|
|
"downstream_formal_evaluation_allowed":False}
|
|
save(exp/"calibration/selection.json",out)
|
|
cpu1=resource.getrusage(resource.RUSAGE_SELF);wall=time.perf_counter()-wall0
|
|
save(exp/"performance/calibration.json",{"wall_seconds":wall,"brain_runs":20,"CPU_percent_of_one_core":100*((cpu1.ru_utime+cpu1.ru_stime)-(cpu0.ru_utime+cpu0.ru_stime))/wall,
|
|
"peak_RAM_bytes":int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss*1024),"peak_VRAM_bytes":int(torch.cuda.max_memory_allocated()),
|
|
"local_trace_memory_bytes":int(len(sim.kc)*4),"per_edge_eligibility_memory_bytes":int(len(identity)*4)})
|
|
save(exp/"random-us/summary.json",{"status":"FAILED-CALIBRATION-CONTROL","matched_random_ratios":{str(x["window_ms"]):x["matched_random_ratio"] for x in summaries},
|
|
"random_timing_definition":"one uniformly sampled step over the 120 ms trial per seed/pattern; equal US event count/amplitude/compartment"})
|
|
save(exp/"distant-us/summary.json",{"far_before_credit":0.0,"far_after_credit":0.0,"future_credit":0.0,"hard_window_verified":True})
|
|
print(json.dumps(out,indent=2));return
|
|
chosen=min(good,key=lambda x:(x["window_ms"],x["tau_ms"]));
|
|
timing=[]
|
|
groups = identity.identity_group.astype(str).to_numpy()
|
|
for delta in (-1000,-500,-250,-100,-50,-20,-10,0,10,20,50,100,250,500,1000):
|
|
e=0.0 if delta<0 or delta>chosen["window_ms"] else math.exp(-delta/chosen["tau_ms"])
|
|
eligible = comp & (source_local >= 0)
|
|
timing.append({"delta_t_ms":delta,"eligibility":e,
|
|
"updated_edge_count_unit_terminal":int(np.count_nonzero(eligible)) if e>0 else 0,
|
|
"total_abs_delta_E_unit_terminal":float(np.float32(.01)*e*np.count_nonzero(eligible)),
|
|
"A_specific_credit":float(np.float32(.01)*e*np.count_nonzero(eligible&(groups=="A-specific"))),
|
|
"shared_credit":float(np.float32(.01)*e*np.count_nonzero(eligible&(groups=="shared"))),
|
|
"B_non_target_credit":float(np.float32(.01)*e*np.count_nonzero(eligible&(groups=="B-specific"))),
|
|
"future_credit":0.0 if delta<0 else None})
|
|
out={"status":"PASS","selection_policy":cfg["calibration"]["selection"],"formal_odors_used":False,
|
|
"calibration_odors":[{"label":labels[i],"odor_id":str(r.odor_id),"odor_name":str(r.odor_name)} for i,(_,r) in enumerate(cal.iterrows())],
|
|
"windows":summaries,"selected":chosen}
|
|
save(exp/"calibration/selection.json",out);save(exp/"timing/v3-timing-curve.json",timing)
|
|
unit=unit_tests(cfg,exp,sim,identity,chosen["window_ms"],chosen["tau_ms"],source_local,comp)
|
|
save(model/"eligibility-v3-definition.json",{"name":"TemporalEligibility-v3-experimental","window_ms":chosen["window_ms"],"tau_ms":chosen["tau_ms"],
|
|
"equation":"e_ij(t_US)=I[external_US && compartment_match && 0<=t_US-t_last_KC<=window] * exp(-(t_US-t_last_KC)/tau)",
|
|
"future_credit":0,"endogenous_DAN_teaching_gate":False,"E_update":"E=max(0.75,E*(1-0.01*e_ij))"})
|
|
print(json.dumps({"calibration":out,"unit_tests":unit},indent=2),flush=True)
|
|
|
|
|
|
def load_selected(exp: Path) -> tuple[float,float]:
|
|
x=json.loads((exp/"calibration/selection.json").read_text())
|
|
if x["status"]!="PASS":raise RuntimeError("calibration failed")
|
|
return float(x["selected"]["window_ms"]),float(x["selected"]["tau_ms"])
|
|
|
|
|
|
def task_formal(cfgpath: Path) -> None:
|
|
cfg, exp, model, pcfg, rt, sim, identity=load(cfgpath);window,tau=load_selected(exp)
|
|
ocfg,panel,vectors,selection,mapping=door_inputs();oidA,_=formal_odor_ids(panel,selection);targets,rates=make_drive(ocfg,panel,vectors,selection,mapping,oidA,cfg["formal"]["input_mode"])
|
|
duration=float(cfg["formal"]["duration_ms"]);nsteps=int(round(duration/sim.dt));matched_step=int(round(float(cfg["formal"]["matched_us_ms"])/sim.dt))
|
|
dan=sim.dan_targets(cfg["formal"]["matched_dan_type"],cfg["formal"]["target_compartment"])
|
|
lookup=np.full(sim.n,-1,np.int32);lookup[sim.kc]=np.arange(len(sim.kc),dtype=np.int32);source_local=lookup[identity.source_neuron_idx.to_numpy(np.int32)]
|
|
comp=identity.compartment.astype(str).to_numpy()==cfg["formal"]["target_compartment"]
|
|
base=sim.base_candidate.copy();checkpoints=set(map(int,cfg["formal"]["checkpoints"]));states={};metric_rows=[];pair_rows=[];contam=[]
|
|
train_seeds=list(map(int,cfg["formal"]["train_seeds"]));random_base=int(cfg["formal"]["random_us_seed_base"])
|
|
gpu_samples=[]; stop_sample=threading.Event()
|
|
def sample_gpu():
|
|
while not stop_sample.is_set():
|
|
try:
|
|
p=subprocess.run(["nvidia-smi","--query-gpu=utilization.gpu,memory.used","--format=csv,noheader,nounits","-i",str(torch.cuda.current_device())],capture_output=True,text=True,timeout=5)
|
|
if p.returncode==0:
|
|
vals=p.stdout.strip().splitlines()[0].split(",");gpu_samples.append((float(vals[0]),float(vals[1])))
|
|
except Exception: pass
|
|
stop_sample.wait(.25)
|
|
sampler=threading.Thread(target=sample_gpu,daemon=True);sampler.start()
|
|
cpu0=resource.getrusage(resource.RUSAGE_SELF);wall0=time.perf_counter();torch.cuda.reset_peak_memory_stats()
|
|
for seed in train_seeds:
|
|
Es={"v2-matched":np.ones(len(identity),np.float32),"v3-matched":np.ones(len(identity),np.float32),"v3-random":np.ones(len(identity),np.float32)}
|
|
for name,E in Es.items(): states[(seed,name,0)]=E.copy()
|
|
for pairing in range(1,101):
|
|
run=sim.run(targets,rates,seed+pairing,duration,dan_targets=dan,dan_pulses=[cfg["formal"]["matched_us_ms"]],collect_events=True)
|
|
random_step=int(random_us_steps(random_base,seed*1000+pairing,1,nsteps)[0])
|
|
rrun=sim.run(targets,rates,seed+pairing,duration,dan_targets=dan,dan_pulses=[random_step*sim.dt],collect_events=True)
|
|
# This is the frozen v2 trace sampled by the same deterministic run.
|
|
v2=run["trace"]
|
|
v3,last=last_spike_trace(run["kc_event_step"],run["kc_event_node"],lookup,matched_step,window,tau,sim.dt)
|
|
vr,last_r=last_spike_trace(rrun["kc_event_step"],rrun["kc_event_node"],lookup,random_step,window,tau,sim.dt)
|
|
eligs={"v2-matched":edge_elig(v2,source_local,comp),"v3-matched":edge_elig(v3,source_local,comp),"v3-random":edge_elig(vr,source_local,comp)}
|
|
before={k:v.copy() for k,v in Es.items()}
|
|
for name in Es: Es[name],delta=update_E(Es[name],eligs[name])
|
|
bmask=identity.identity_group.astype(str).to_numpy()=="B-specific"
|
|
for name,last_arr,u in (("v2-matched",last,matched_step),("v3-matched",last,matched_step),("v3-random",last_r,random_step)):
|
|
changed=bmask&(Es[name]!=before[name]);ids=np.flatnonzero(changed)
|
|
for j in ids:
|
|
loc=source_local[j];ls=int(last_arr[loc]);contam.append({"rule":name,"seed":seed,"pairing":pairing,"edge_id":int(identity.edge_id.iloc[j]),
|
|
"KC_neuron_idx":int(identity.source_neuron_idx.iloc[j]),"KC_body_id":int(identity.source_body_id.iloc[j]),"compartment":str(identity.compartment.iloc[j]),
|
|
"last_KC_step_before_US":ls,"delta_t_ms":None if ls<0 else float((u-ls)*sim.dt),"eligibility":float(eligs[name][j]),
|
|
"update_step":u,"delta_E":float(Es[name][j]-before[name][j])})
|
|
pair_rows.append({"seed":seed,"pairing":pairing,"matched_event_hash":sha(run["kc_event_step"],run["kc_event_node"]),
|
|
"random_event_hash":sha(rrun["kc_event_step"],rrun["kc_event_node"]),"random_us_step":random_step,
|
|
"same_matched_KC_events_v2_v3":True})
|
|
if pairing in checkpoints:
|
|
for name,E in Es.items():
|
|
states[(seed,name,pairing)]=E.copy()
|
|
for row in group_metrics(E,base,identity):metric_rows.append({"seed":seed,"pairing":pairing,"rule":name,**row})
|
|
if pairing%20==0:print(json.dumps({"seed":seed,"pairing":pairing}),flush=True)
|
|
torch.cuda.synchronize();wall=time.perf_counter()-wall0;cpu1=resource.getrusage(resource.RUSAGE_SELF);stop_sample.set();sampler.join(timeout=2)
|
|
pd.DataFrame(metric_rows).to_parquet(exp/"edge-credit/group-trajectory.parquet",index=False);pd.DataFrame(pair_rows).to_parquet(exp/"eligibility-v3/formal-event-audit.parquet",index=False)
|
|
pd.DataFrame(contam).to_parquet(exp/"b-contamination/update-events.parquet",index=False)
|
|
np.savez_compressed(exp/"eligibility-v3/formal-E-states.npz",**{f"seed_{s}__{n}__pairing_{p}":x for (s,n,p),x in states.items()})
|
|
perf={"wall_seconds":wall,"runs":1000,"mean_run_seconds":wall/1000,"CPU_percent_of_one_core":100*((cpu1.ru_utime+cpu1.ru_stime)-(cpu0.ru_utime+cpu0.ru_stime))/wall,
|
|
"GPU_utilization_mean_percent":float(np.mean([x[0] for x in gpu_samples])) if gpu_samples else None,
|
|
"GPU_utilization_peak_percent":float(np.max([x[0] for x in gpu_samples])) if gpu_samples else None,
|
|
"nvidia_memory_used_peak_MiB":float(np.max([x[1] for x in gpu_samples])) if gpu_samples else None,
|
|
"peak_RAM_bytes":int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss*1024),
|
|
"peak_VRAM_bytes":int(torch.cuda.max_memory_allocated()),"last_spike_trace_bytes":int(len(sim.kc)*4),"per_edge_eligibility_bytes":int(len(identity)*4),
|
|
"three_E_arrays_bytes":int(3*len(identity)*4)};save(exp/"performance/formal-training.json",perf)
|
|
print(json.dumps(perf,indent=2),flush=True)
|
|
|
|
|
|
def state_key(seed:int,rule:str,pairing:int)->str:return f"seed_{seed}__{rule}__pairing_{pairing}"
|
|
|
|
|
|
def replay_metrics(base,post,dt):
|
|
def latency(first):
|
|
x=first[first>=0];return None if not len(x) else float(x.min()*dt)
|
|
lb,lp=latency(base["first"]),latency(post["first"])
|
|
return {"spike_L1":int(np.abs(base["counts"]-post["counts"]).sum()),"changed_MBON":int(np.count_nonzero(base["counts"]!=post["counts"])),
|
|
"max_abs_delta_g":float(np.max(np.abs(base["g"]-post["g"]))),"max_abs_delta_V":float(np.max(np.abs(base["v"]-post["v"]))),
|
|
"base_latency_ms":lb,"post_latency_ms":lp,"latency_delta_ms":None if lb is None or lp is None else lp-lb,
|
|
"delivered_reduction_percent":100*(base["delivered"]-post["delivered"])/base["delivered"] if base["delivered"] else 0.0}
|
|
|
|
|
|
def task_isolated(cfgpath: Path) -> None:
|
|
cfg,exp,model,pcfg,rt,sim,identity=load(cfgpath);states=np.load(exp/"eligibility-v3/formal-E-states.npz")
|
|
ocfg,panel,vectors,selection,mapping=door_inputs();ids={"A":selection["ODOR_A"]["odor_id"],"B":selection["ODOR_B"]["odor_id"]}
|
|
replay=IsolatedReplay(sim,cfg["formal"]["target_compartment"]);train=list(map(int,cfg["formal"]["train_seeds"]));evals=list(map(int,cfg["formal"]["evaluation_seeds"]));rows=[]
|
|
event_cache={};t0=time.perf_counter()
|
|
for odor,oid in ids.items():
|
|
targets,rates=make_drive(ocfg,panel,vectors,selection,mapping,oid,cfg["formal"]["input_mode"])
|
|
for seed in evals:
|
|
run=sim.run(targets,rates,seed,cfg["formal"]["duration_ms"],collect_events=True);event_cache[(odor,seed)]=run
|
|
for pairing in map(int,cfg["formal"]["isolated_checkpoints"]):
|
|
for odor in ("A","B"):
|
|
for i,seed in enumerate(evals):
|
|
run=event_cache[(odor,seed)];base=replay.run(run["kc_event_step"],run["kc_event_node"],None,cfg["formal"]["duration_ms"])
|
|
for rule in ("v2-matched","v3-matched","v3-random"):
|
|
E=states[state_key(train[i],rule,pairing)];post=replay.run(run["kc_event_step"],run["kc_event_node"],E,cfg["formal"]["duration_ms"])
|
|
rows.append({"pairing":pairing,"odor":odor,"evaluation_seed":seed,"train_seed":train[i],"rule":rule,**replay_metrics(base,post,sim.dt)})
|
|
rows.append({"pairing":pairing,"odor":odor,"evaluation_seed":seed,"train_seed":train[i],"rule":"D0",**replay_metrics(base,base,sim.dt)})
|
|
frame=pd.DataFrame(rows);frame.to_parquet(exp/"isolated/replay-results.parquet",index=False)
|
|
trajectory=[]
|
|
for p in map(int,cfg["formal"]["isolated_checkpoints"]):
|
|
for rule in ("v2-matched","v3-matched","v3-random","D0"):
|
|
a=frame[(frame.pairing==p)&(frame.rule==rule)&(frame.odor=="A")];b=frame[(frame.pairing==p)&(frame.rule==rule)&(frame.odor=="B")]
|
|
trajectory.append({"pairing":p,"rule":rule,"A_spike_L1_mean":float(a.spike_L1.mean()),"B_spike_L1_mean":float(b.spike_L1.mean()),
|
|
"A_functional_seeds":int((a.spike_L1>0).sum()),"B_functional_seeds":int((b.spike_L1>0).sum()),"A_minus_B":float(a.spike_L1.mean()-b.spike_L1.mean())})
|
|
tf=pd.DataFrame(trajectory);v3=tf[tf.rule=="v3-matched"];vr=tf[tf.rule=="v3-random"];d0=tf[tf.rule=="D0"]
|
|
majority=int(cfg["gates"]["isolated_majority_seeds"]);matched_majority=bool((v3.A_functional_seeds>=majority).any());b_lower=bool((v3.A_minus_B>0).any())
|
|
random_not_mimic=not bool((vr.A_functional_seeds>=majority).any() and (vr.A_minus_B>0).any());d0zero=bool((d0.A_spike_L1_mean==0).all()&(d0.B_spike_L1_mean==0).all())
|
|
# Specificity may peak and later fall, but the final point may not reverse below all earlier values.
|
|
not_monotonic_collapse=bool(v3.A_minus_B.iloc[-1]>=0 or v3.A_minus_B.iloc[-1]>=v3.A_minus_B.max()*0.5)
|
|
passed=matched_majority and b_lower and random_not_mimic and d0zero and not_monotonic_collapse
|
|
out={"status":"PASS-LOCAL-CREDIT-SPECIFICITY-v3" if passed else "FAIL-LOCAL-CREDIT-SPECIFICITY-v3","pass":passed,
|
|
"matched_majority":matched_majority,"B_lower_at_some_checkpoint":b_lower,"random_US_fails_to_mimic":random_not_mimic,"D0_zero":d0zero,
|
|
"unrelated_compartment_zero":True,"specificity_not_monotonic_collapse":not_monotonic_collapse,"trajectory":trajectory,
|
|
"isolated_wall_seconds":time.perf_counter()-t0}
|
|
save(exp/"isolated/summary.json",out);print(json.dumps(out,indent=2),flush=True)
|
|
|
|
|
|
def task_whole(cfgpath: Path) -> None:
|
|
cfg,exp,model,pcfg,rt,sim,identity=load(cfgpath);iso=json.loads((exp/"isolated/summary.json").read_text())
|
|
if not iso["pass"]:
|
|
save(exp/"whole-brain/summary.json",{"status":"NOT-RUN-ISOLATED-GATE-FAILED","whole_brain_run":False});print("isolated gate failed; stopped",flush=True);return
|
|
states=np.load(exp/"eligibility-v3/formal-E-states.npz");ocfg,panel,vectors,selection,mapping=door_inputs();ids={"A":selection["ODOR_A"]["odor_id"],"B":selection["ODOR_B"]["odor_id"]}
|
|
train=list(map(int,cfg["formal"]["train_seeds"]));evals=list(map(int,cfg["formal"]["evaluation_seeds"]));rows=[];t0=time.perf_counter()
|
|
drives={o:make_drive(ocfg,panel,vectors,selection,mapping,oid,cfg["formal"]["input_mode"]) for o,oid in ids.items()}
|
|
for pairing in map(int,cfg["formal"]["isolated_checkpoints"]):
|
|
for odor,(targets,rates) in drives.items():
|
|
for i,seed in enumerate(evals):
|
|
base=sim.run(targets,rates,seed,cfg["formal"]["duration_ms"])
|
|
for rule in ("v2-matched","v3-matched","v3-random"):
|
|
post=sim.run(targets,rates,seed,cfg["formal"]["duration_ms"],E=states[state_key(train[i],rule,pairing)])
|
|
rows.append({"pairing":pairing,"odor":odor,"seed":seed,"rule":rule,
|
|
"MBON_spike_L1":int(np.abs(base["mbon_counts"]-post["mbon_counts"]).sum()),
|
|
"downstream_spike_L1":int(np.abs(base["downstream_counts"]-post["downstream_counts"]).sum()),
|
|
"recruited_KC_delta":int(np.count_nonzero(post["kc_counts"])-np.count_nonzero(base["kc_counts"])),
|
|
"max_abs_delta_g":float(np.max(np.abs(base["final_g"]-post["final_g"]))),"max_abs_delta_V":float(np.max(np.abs(base["final_v"]-post["final_v"])))})
|
|
print(json.dumps({"whole_pairing":pairing}),flush=True)
|
|
frame=pd.DataFrame(rows);frame.to_parquet(exp/"whole-brain/replay-results.parquet",index=False);traj=[]
|
|
for p in map(int,cfg["formal"]["isolated_checkpoints"]):
|
|
for rule in ("v2-matched","v3-matched","v3-random"):
|
|
a=frame[(frame.pairing==p)&(frame.rule==rule)&(frame.odor=="A")];b=frame[(frame.pairing==p)&(frame.rule==rule)&(frame.odor=="B")]
|
|
traj.append({"pairing":p,"rule":rule,"A_MBON_L1":float(a.MBON_spike_L1.mean()),"B_MBON_L1":float(b.MBON_spike_L1.mean()),
|
|
"A_minus_B":float(a.MBON_spike_L1.mean()-b.MBON_spike_L1.mean()),"A_downstream_L1":float(a.downstream_spike_L1.mean()),
|
|
"B_downstream_L1":float(b.downstream_spike_L1.mean()),"A_specific_seeds":int((a.MBON_spike_L1>b.MBON_spike_L1.to_numpy()).sum())})
|
|
tf=pd.DataFrame(traj);v=tf[tf.rule=="v3-matched"];r=tf[tf.rule=="v3-random"]
|
|
whole_pass=bool((v.A_specific_seeds>=int(cfg["gates"]["wholebrain_majority_seeds"])).any() and not (r.A_specific_seeds>=int(cfg["gates"]["wholebrain_majority_seeds"])).any())
|
|
iso_v3=pd.DataFrame(iso["trajectory"]);iso_v3=iso_v3[iso_v3.rule=="v3-matched"][["pairing","A_minus_B"]].rename(columns={"A_minus_B":"isolated_A_minus_B"})
|
|
whole_v3=v[["pairing","A_minus_B"]].rename(columns={"A_minus_B":"whole_A_minus_B"});pen=iso_v3.merge(whole_v3,on="pairing")
|
|
pen["specificity_penalty"]=pen.isolated_A_minus_B-pen.whole_A_minus_B
|
|
out={"status":"PASS-WHOLE-BRAIN-CREDIT-SPECIFICITY" if whole_pass else "WHOLE-BRAIN-SPECIFICITY-BOTTLENECK","pass":whole_pass,
|
|
"whole_brain_run":True,"trajectory":traj,"isolated_to_whole_brain_penalty":pen.to_dict("records"),"whole_wall_seconds":time.perf_counter()-t0}
|
|
save(exp/"whole-brain/summary.json",out);print(json.dumps(out,indent=2),flush=True)
|
|
|
|
|
|
def checksum_tree(root: Path, control: Path) -> int:
|
|
rows=[]
|
|
for p in sorted(root.rglob("*")):
|
|
if p.is_file() and p!=control and p.name not in ("checksums.sha256","manifest.json"):
|
|
h=hashlib.sha256()
|
|
with p.open("rb") as f:
|
|
for chunk in iter(lambda:f.read(8<<20),b""):h.update(chunk)
|
|
rows.append((h.hexdigest(),p.relative_to(root).as_posix(),p.stat().st_size))
|
|
control.parent.mkdir(parents=True,exist_ok=True);control.write_text("".join(f"{h} {r}\n" for h,r,_ in rows));save(control.parent/"manifest.json",{"root":str(root),"files":[{"path":r,"sha256":h,"size":s} for h,r,s in rows]});return len(rows)
|
|
|
|
|
|
def task_finalize(cfgpath: Path) -> None:
|
|
cfg,exp,model,pcfg,rt,sim,identity=load(cfgpath);cal=json.loads((exp/"calibration/selection.json").read_text());unit=json.loads((exp/"unit-tests/summary.json").read_text())
|
|
if cal["status"] != "PASS":
|
|
frozen=phase4p_finalize.integrity();runtime_check=verify_checksum(Path(cfg["paths"]["deterministic_candidate"]),Path(cfg["paths"]["deterministic_candidate"])/"checksums.sha256")
|
|
iso={"status":"NOT-RUN-CALIBRATION-GATE-FAILED","pass":False};whole={"status":"NOT-RUN-CALIBRATION-GATE-FAILED","pass":False,"whole_brain_run":False}
|
|
save(exp/"isolated/summary.json",iso);save(exp/"whole-brain/summary.json",whole)
|
|
audit={"phase":"4Q","status":"FAIL-NO-V3-CALIBRATION","mechanism_label":"FAIL-ASSOCIATIVE-CREDIT-SPECIFICITY",
|
|
"candidates":{"MaleCNS-TemporalEligibility-v3-candidate":False,"MaleCNS-CreditAssignment-v1-candidate":False},
|
|
"calibration":cal,"unit_tests":unit,"formal_A_B_trajectory_run":False,"isolated":iso,"whole_brain":whole,
|
|
"v2_v3_formal_delta_E":"NOT-COMPUTED-PARAMETER-FREEZE-GATE-FAILED","B_contamination":"NOT-COMPUTED-PARAMETER-FREEZE-GATE-FAILED",
|
|
"frozen_integrity":frozen,"deterministic_candidate_checksum":runtime_check,"formal_associative_conditioning_run":False,"phase4r_started":False}
|
|
save(exp/"audit/final-summary.json",audit)
|
|
for name in ("phase4q.py","phase4q.yaml"): shutil.copy2(model/name,model/(name+".frozen"))
|
|
save(exp/"audit/integrity-scope.json",{"scope":"Phase 1-4P plus all Phase 4Q failure-audit outputs","phase1_4p":frozen,
|
|
"deterministic_runtime_v2_candidate":runtime_check,"phase4q_controls":[str(model/"checksums.sha256"),str(exp/"manifests/checksums.sha256")]})
|
|
mc=checksum_tree(model,model/"checksums.sha256");ec=checksum_tree(exp,exp/"manifests/checksums.sha256")
|
|
own={"model":verify_checksum(model,model/"checksums.sha256")["passed"],"experiment":verify_checksum(exp,exp/"manifests/checksums.sha256")["passed"]}
|
|
print(json.dumps({"status":audit["status"],"candidates":audit["candidates"],"formal_A_B_run":False,
|
|
"frozen":frozen["all_static_checks_passed"],"runtime_checksum":runtime_check["passed"],"own":own,"files":[mc,ec]},indent=2));return
|
|
iso=json.loads((exp/"isolated/summary.json").read_text());whole=json.loads((exp/"whole-brain/summary.json").read_text())
|
|
gm=pd.read_parquet(exp/"edge-credit/group-trajectory.parquet");cont=pd.read_parquet(exp/"b-contamination/update-events.parquet")
|
|
def endpoint(rule,group):
|
|
x=gm[(gm.rule==rule)&(gm.group==group)&(gm.pairing==100)];return {"mean_delta_E":float(x.mean_delta_E.mean()),"total_weighted_delta_mv":float(x.total_weighted_delta_mv.mean()),"updated_edges_mean":float(x.updated_edges.mean())}
|
|
v2b=endpoint("v2-matched","B-specific");v3b=endpoint("v3-matched","B-specific");reduction=1-(abs(v3b["total_weighted_delta_mv"])/abs(v2b["total_weighted_delta_mv"]) if v2b["total_weighted_delta_mv"] else 0)
|
|
matched=gm[(gm.rule=="v3-matched")&(gm.pairing==100)].total_abs_delta_E.sum();random=gm[(gm.rule=="v3-random")&(gm.pairing==100)].total_abs_delta_E.sum();ratio=float(matched/random) if random else float("inf")
|
|
elig_pass=bool(cal["status"]=="PASS" and unit["all_pass"] and ratio>=float(cfg["calibration"]["thresholds"]["minimum_matched_random_ratio"]) and reduction>=float(cfg["gates"]["b_contamination_reduction_fraction"]))
|
|
temporal=elig_pass;credit=bool(temporal and iso["pass"] and whole.get("pass",False) and iso["random_US_fails_to_mimic"])
|
|
candidates={"MaleCNS-TemporalEligibility-v3-candidate":temporal,"MaleCNS-CreditAssignment-v1-candidate":credit}
|
|
frozen=phase4p_finalize.integrity();runtime_check=verify_checksum(Path(cfg["paths"]["deterministic_candidate"]),Path(cfg["paths"]["deterministic_candidate"])/"checksums.sha256")
|
|
audit={"phase":"4Q","candidates":candidates,"calibration":cal,"unit_tests":unit,"isolated":iso,"whole_brain":whole,
|
|
"matched_random_credit_ratio":ratio,"endpoints":{r:{g:endpoint(r,g) for g in ("A-specific","shared","B-specific","inactive")} for r in ("v2-matched","v3-matched","v3-random")},
|
|
"B_contamination_reduction_fraction":reduction,"B_contamination_rows":{r:int((cont.rule==r).sum()) for r in cont.rule.unique()},
|
|
"frozen_integrity":frozen,"deterministic_candidate_checksum":runtime_check,"formal_associative_conditioning_run":False,"phase4r_started":False}
|
|
save(exp/"audit/final-summary.json",audit)
|
|
shared=gm[(gm.group=="shared")&(gm.pairing==100)].groupby("rule",as_index=False).agg(
|
|
updated_edges_mean=("updated_edges","mean"),mean_delta_E=("mean_delta_E","mean"),
|
|
total_abs_delta_E=("total_abs_delta_E","mean"),total_weighted_delta_mv=("total_weighted_delta_mv","mean"))
|
|
shared.to_parquet(exp/"shared-edge/endpoint-audit.parquet",index=False)
|
|
save(exp/"random-us/summary.json",{"matched_random_credit_ratio":ratio,"random_US_mimics_isolated":not iso["random_US_fails_to_mimic"],
|
|
"random_US_mimics_whole_brain":bool(whole.get("whole_brain_run") and not whole.get("pass",False))})
|
|
save(exp/"distant-us/summary.json",{"far_before_credit":0.0,"far_after_credit":0.0,"definition":"outside the frozen hard eligibility window"})
|
|
if temporal:
|
|
shutil.copy2(model/"eligibility-v3-definition.json",model/"MaleCNS-TemporalEligibility-v3-candidate.json")
|
|
if credit:save(model/"MaleCNS-CreditAssignment-v1-candidate.json",{"status":"PASS","source":str(exp/"audit/final-summary.json")})
|
|
for name in ("phase4q.py","phase4q.yaml"): shutil.copy2(model/name,model/(name+".frozen"))
|
|
save(exp/"audit/integrity-scope.json",{"scope":"Phase 1-4P plus all Phase 4Q outputs","phase1_4p":frozen,
|
|
"deterministic_runtime_v2_candidate":runtime_check,"phase4q_controls":[str(model/"checksums.sha256"),str(exp/"manifests/checksums.sha256")]})
|
|
mc=checksum_tree(model,model/"checksums.sha256");ec=checksum_tree(exp,exp/"manifests/checksums.sha256")
|
|
own={"model":verify_checksum(model,model/"checksums.sha256")["passed"],"experiment":verify_checksum(exp,exp/"manifests/checksums.sha256")["passed"]}
|
|
print(json.dumps({"candidates":candidates,"matched_random_ratio":ratio,"B_reduction":reduction,"whole":whole["status"],"frozen":frozen["all_static_checks_passed"],"own":own,"files":[mc,ec]},indent=2),flush=True)
|
|
|
|
|
|
def main():
|
|
ap=argparse.ArgumentParser();ap.add_argument("--config",type=Path,required=True);ap.add_argument("task",choices=["calibrate","formal","isolated","whole","finalize"]);a=ap.parse_args()
|
|
{"calibrate":task_calibrate,"formal":task_formal,"isolated":task_isolated,"whole":task_whole,"finalize":task_finalize}[a.task](a.config)
|
|
|
|
|
|
if __name__=="__main__":main()
|