Digital Drosophila Archive
228 lines
17 KiB
Python
228 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""Phase 4N event-level KC/DAN side-channel on unmodified frozen LIF dynamics."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import resource
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
|
|
from phase4h_audit import write_json
|
|
from phase4i import Validator
|
|
from phase4l import OdorSimulator, pn_drive
|
|
from phase4m import load_odor_inputs
|
|
from phase4n import load_config
|
|
|
|
|
|
class EventSimulator(OdorSimulator):
|
|
def run_events(self, targets, rates_hz, seed, duration_ms, delta, dan_targets, dan_pulse_ms, all_dan,
|
|
save_all_events=False):
|
|
targets = np.asarray(targets, np.int32); rates_hz = np.asarray(rates_hz, np.float32)
|
|
self.v.install(delta); self.rt.weights[self.apl_edge_gpu] = 0.0; self.rt.reset(int(seed))
|
|
steps = int(round(duration_ms / self.dt)); tg = torch.as_tensor(targets, dtype=torch.int64, device=self.device)
|
|
rg = torch.as_tensor(rates_hz * self.dt / 1000.0, dtype=torch.float32, device=self.device)
|
|
gen = torch.Generator(device=self.device); gen.manual_seed(int(seed))
|
|
forced = torch.rand((steps, len(targets)), generator=gen, device=self.device) < rg[None, :]
|
|
dan_targets = np.asarray(dan_targets, np.int32); dg = torch.as_tensor(dan_targets, dtype=torch.int64, device=self.device)
|
|
dan_forced = self._forced(steps, dan_targets, seed + 991, None, [dan_pulse_ms])
|
|
dan_all_gpu = torch.as_tensor(all_dan, dtype=torch.int64, device=self.device)
|
|
apl = self.cfg["apl_overlay"]; apl_decay = math.exp(-self.dt / float(apl["tau_ms"]))
|
|
trace_decay = math.exp(-self.dt / float(self.cfg["eligibility_v2"]["kc_trace_tau_ms"]))
|
|
apl_state = torch.zeros(len(self.apl), dtype=torch.float32, device=self.device)
|
|
trace = torch.zeros(len(self.kc), dtype=torch.float32, device=self.device)
|
|
kc_history = torch.zeros((steps, len(self.kc)), dtype=torch.bool, device=self.device)
|
|
dan_history = torch.zeros((steps, len(all_dan)), dtype=torch.bool, device=self.device)
|
|
natural_dan_history = torch.zeros_like(dan_history)
|
|
trace_at_gate = None; all_nodes = []; all_steps = []
|
|
alpha, gd, rest, reset, threshold, refsteps = self.rt.alpha, self.rt.g_decay, self.rt.rest, self.rt.reset_value, self.rt.threshold, self.rt.refractory_steps
|
|
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[tg]
|
|
if accepted.any(): spikes[tg[accepted]] = True
|
|
natural_dan_history[step] = spikes[dan_all_gpu]
|
|
if len(dan_targets):
|
|
da = dan_forced[step] & available[dg]
|
|
if da.any(): spikes[dg[da]] = True
|
|
sources = torch.nonzero(spikes, as_tuple=False).flatten()
|
|
kc_history[step] = spikes[self.kc_gpu]
|
|
dan_history[step] = spikes[dan_all_gpu]
|
|
if save_all_events and sources.numel():
|
|
all_nodes.append(sources.to(torch.int32)); all_steps.append(torch.full((sources.numel(),), step, dtype=torch.int32, device=self.device))
|
|
if sources.numel():
|
|
self.rt.v[sources] = reset; self.rt.refractory[sources] = refsteps; eids = self.rt._hybrid_edges(sources)
|
|
if eids.numel(): self.rt.delay_queue[(step + self.rt.delay_steps) % self.rt.queue_len].scatter_add_(0, self.rt.indices[eids].long(), self.rt.weights[eids])
|
|
trace.mul_(trace_decay).add_(spikes[self.kc_gpu].to(torch.float32)).clamp_max_(1.0)
|
|
if step == int(round(dan_pulse_ms / self.dt)): trace_at_gate = trace.clone()
|
|
apl_state.mul_(apl_decay).add_(torch.mv(self.input_gpu, spikes.to(torch.float32)) * float(apl["normalization"]))
|
|
inhibition = torch.matmul(apl_state, self.output_gpu) * float(apl["gain"])
|
|
self.rt.g.index_add_(0, self.kc_gpu, -inhibition)
|
|
torch.cuda.synchronize()
|
|
event_table = None
|
|
if save_all_events:
|
|
event_table = pd.DataFrame({"step": torch.cat(all_steps).cpu().numpy() if all_steps else np.empty(0,np.int32),
|
|
"neuron_idx": torch.cat(all_nodes).cpu().numpy() if all_nodes else np.empty(0,np.int32)})
|
|
return {"kc": kc_history.cpu().numpy(), "dan": dan_history.cpu().numpy(), "natural_dan": natural_dan_history.cpu().numpy(), "all_events": event_table,
|
|
"trace_at_external": trace_at_gate.cpu().numpy(), "wall_seconds": time.perf_counter()-t0,
|
|
"peak_vram_bytes": int(torch.cuda.max_memory_allocated())}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(); ap.add_argument("--config", type=Path, required=True); ap.add_argument("--max-pairings", type=int, default=100); ap.add_argument("--train-seed", type=int)
|
|
args = ap.parse_args(); ncfg, mcfg, cfg, basepath = load_config(args.config)
|
|
exp = Path(ncfg["experiment_root"])
|
|
identity = pd.read_parquet(exp / "edge-identity/frozen-edge-identity.parquet")
|
|
history = pd.read_parquet(Path(mcfg["experiment_root"]) / "efficacy-overlay/update-history.parquet")
|
|
v = Validator(cfg, basepath, exp); sim = EventSimulator(v, cfg, exp)
|
|
if not np.array_equal(v.edge_ids, identity.edge_id.to_numpy(np.int64)):
|
|
raise RuntimeError("edge identity order mismatch")
|
|
all_dan = np.unique(v.confirmed.source_neuron_idx.to_numpy(np.int32))
|
|
matched = v.dan_targets(str(ncfg["matched_dan_type"]), str(ncfg["matched_compartment_verbatim"]))
|
|
matched_set = set(map(int, matched)); all_dan_to_local = {int(n): i for i,n in enumerate(all_dan)}
|
|
matched_loc = np.array([all_dan_to_local[int(n)] for n in matched], np.int32)
|
|
panel, vectors, selection, mapping = load_odor_inputs(cfg)
|
|
gloms = selection["glomeruli_order"]; matched_sum = float(np.median(panel.mapped_response_sum))
|
|
oid = selection["ODOR_A"]["odor_id"]
|
|
pn = pn_drive(oid, vectors[oid], mapping, gloms, str(mcfg["protocol"]["input_mode"]), matched_sum, cfg)[:2]
|
|
src = identity.source_neuron_idx.to_numpy(np.int32)
|
|
kclook = np.full(v.n, -1, np.int32); kclook[v.kc] = np.arange(len(v.kc), dtype=np.int32)
|
|
loc = kclook[src]
|
|
comp = identity.compartment.astype(str).to_numpy() == str(ncfg["matched_compartment_verbatim"])
|
|
bound = np.float32(1.0 - cfg["eligibility_v2"]["max_fractional_depression"])
|
|
k = np.float32(cfg["eligibility_v2"]["eta_effective_mv"] / cfg["model"]["w_syn_mv_per_anatomical_synapse"])
|
|
tau = float(cfg["eligibility_v2"]["kc_trace_tau_ms"])
|
|
short_tau = float(ncfg["shadow_trace_tau_ms"])
|
|
pulse_step = int(round(float(mcfg["protocol"]["dan_time_ms"]) / sim.dt))
|
|
hist_by = {(int(s),int(p)): frame for (s,p),frame in history.groupby(["train_seed","pairing"])}
|
|
eid_to_local = pd.Series(np.arange(len(identity), dtype=np.int32), index=identity.edge_id.to_numpy(np.int64))
|
|
audit_rows = []; gate_rows = []; shadow_rows = []; perf = []; kc_event_frames = []; dan_event_frames = []; mismatches = []
|
|
key = {10,20,50,60,75,100}
|
|
for seed in ([args.train_seed] if args.train_seed is not None else ncfg["train_seeds"]):
|
|
additive = np.zeros(len(identity), np.float32)
|
|
E_ext = np.ones(len(identity), np.float32)
|
|
E_all = np.ones(len(identity), np.float32)
|
|
rng = np.random.default_rng(seed + int(ncfg["random_us_seed_offset"]))
|
|
for pairing in range(1, args.max_pairings+1):
|
|
audit_start = len(audit_rows); gate_start = len(gate_rows)
|
|
capture = pairing in key
|
|
run = sim.run_events(*pn, seed+pairing, float(mcfg["protocol"]["duration_ms"]), additive, matched,
|
|
float(mcfg["protocol"]["dan_time_ms"]), all_dan, save_all_events=capture)
|
|
perf.append({"seed":seed,"pairing":pairing,"wall_seconds":run["wall_seconds"],"peak_vram_bytes":run["peak_vram_bytes"]})
|
|
kc = run["kc"]; dan = run["natural_dan"]
|
|
ks, kl = np.nonzero(kc)
|
|
if len(ks):
|
|
kc_event_frames.append(pd.DataFrame({"seed":seed,"pairing":pairing,"step":ks.astype(np.int32),
|
|
"neuron_idx":v.kc[kl].astype(np.int32)}))
|
|
ds, dl = np.nonzero(dan)
|
|
if len(ds):
|
|
dan_event_frames.append(pd.DataFrame({"seed":seed,"pairing":pairing,"step":ds.astype(np.int32),
|
|
"neuron_idx":all_dan[dl].astype(np.int32),"source":"endogenous-spike"}))
|
|
trace = np.zeros(len(v.kc), np.float32); short = np.zeros_like(trace)
|
|
last = np.full(len(v.kc), -1, np.int32); first = np.full(len(v.kc), -1, np.int32)
|
|
external_at_gate = None
|
|
endogenous_steps = 0; endogenous_events = 0
|
|
random_step = int(round(float(rng.uniform(*ncfg["random_us_window_ms"])) / sim.dt))
|
|
random_trace = None
|
|
for step in range(len(kc)):
|
|
spike = kc[step]
|
|
trace *= np.float32(math.exp(-sim.dt / tau)); trace += spike.astype(np.float32); np.minimum(trace, 1.0, out=trace)
|
|
short *= np.float32(math.exp(-sim.dt / short_tau)); short += spike.astype(np.float32); np.minimum(short, 1.0, out=short)
|
|
last[spike] = step; first[spike & (first < 0)] = step
|
|
if step == random_step: random_trace = trace[loc].copy()
|
|
endogenous = bool(np.any(dan[step, matched_loc]))
|
|
external = step == pulse_step
|
|
if endogenous:
|
|
endogenous_steps += 1; endogenous_events += int(np.count_nonzero(dan[step, matched_loc]))
|
|
if external:
|
|
external_at_gate = trace[loc].copy()
|
|
max_error = float(np.max(np.abs(trace - run["trace_at_external"])))
|
|
if max_error > 2e-5: raise RuntimeError(f"KC trace reconstruction mismatch {max_error}")
|
|
if external or endogenous:
|
|
elig = trace[loc]; active = comp & (elig > 0)
|
|
before = E_all[active].copy()
|
|
E_all[active] = np.maximum(bound, E_all[active] * (np.float32(1.0)-k*elig[active]))
|
|
gate_rows.append({"seed":seed,"pairing":pairing,"step":step,"time_ms":step*sim.dt,
|
|
"source": "both" if external and endogenous else ("external-US-associated" if external else "endogenous-feedback-associated"),
|
|
"matched_endogenous_DAN_spikes":int(np.count_nonzero(dan[step,matched_loc])),
|
|
"eligible_edges":int(active.sum()),"E_change_sum":float(np.sum(before-E_all[active]))})
|
|
if external:
|
|
last_age_ms = np.where(last[loc] >= 0, (step-last[loc])*sim.dt, np.inf)
|
|
a = comp & (trace[loc] > 0)
|
|
for local in np.flatnonzero(a):
|
|
source = int(src[local]); age = float(last_age_ms[local])
|
|
audit_rows.append({"seed":seed,"pairing":pairing,"edge_id":int(identity.edge_id.iat[local]),
|
|
"source_neuron_idx":source,"source_body_id":int(identity.source_body_id.iat[local]),
|
|
"target_neuron_idx":int(identity.target_neuron_idx.iat[local]),"compartment":str(identity.compartment.iat[local]),
|
|
"identity_group":str(identity.identity_group.iat[local]),"update_time_ms":step*sim.dt,
|
|
"eligibility_v2":float(trace[loc[local]]),"eligibility_short_shadow":float(short[loc[local]]),
|
|
"event_coincidence_shadow":bool(age <= float(ncfg["event_coincidence_window_ms"])),
|
|
"first_activation_only_shadow":bool(first[loc[local]] == last[loc[local]]),
|
|
"last_local_KC_spike_age_ms":age,"first_local_KC_spike_ms":float(first[loc[local]]*sim.dt),
|
|
"matched_external_DAN_arrival_ms":step*sim.dt,
|
|
"endogenous_matched_DAN_simultaneous":bool(endogenous),
|
|
"source_class": "external-US-associated", "delta_E_bounded":float(-k*trace[loc[local]]*E_ext[local])})
|
|
if external_at_gate is None: raise RuntimeError("external gate missing")
|
|
rec = hist_by[(int(seed),int(pairing))]
|
|
active_loc = eid_to_local.loc[rec.edge_id.to_numpy(np.int64)].to_numpy(np.int32)
|
|
err = float(np.max(np.abs(external_at_gate[active_loc] - rec.eligibility.to_numpy(np.float32)))) if len(active_loc) else 0.0
|
|
exact = err <= 2e-5
|
|
if not exact:
|
|
mismatches.append({"seed":seed,"pairing":pairing,"max_abs_eligibility_error":err})
|
|
for row in audit_rows[audit_start:]: row["frozen_trace_match"] = exact
|
|
for row in gate_rows[gate_start:]: row["frozen_trace_match"] = exact
|
|
# Advance the next trial from the sealed Phase 4M overlay, never from a divergent audit replay.
|
|
old_E = rec.E_bounded_before.to_numpy(np.float32); new_E = rec.E_bounded_after.to_numpy(np.float32)
|
|
if len(active_loc) and float(np.max(np.abs(E_ext[active_loc]-old_E))) > 1e-6:
|
|
raise RuntimeError(f"sealed E history before-state mismatch {seed}/{pairing}")
|
|
E_ext[active_loc] = new_E
|
|
true_elig = rec.eligibility.to_numpy(np.float32)
|
|
additive[active_loc] = np.maximum(-np.float32(cfg["eligibility_v2"]["max_fractional_depression"])*v.base_candidate[active_loc],
|
|
additive[active_loc] - np.float32(cfg["eligibility_v2"]["eta_effective_mv"])*true_elig)
|
|
if pairing in ncfg["checkpoints"]:
|
|
with np.load(exp / "acquisition-trajectory" / f"E-seed-{seed}-pairings-{pairing}.npz") as f:
|
|
chk_err = float(np.max(np.abs(E_ext-f["E_bounded"])))
|
|
if chk_err > 1e-6: raise RuntimeError(f"external E reconstruction mismatch {seed}/{pairing}: {chk_err}")
|
|
np.savez_compressed(exp / "endogenous-dan" / f"E-all-DAN-shadow-seed-{seed}-pairings-{pairing}.npz", E_all_DAN=E_all,
|
|
E_external_only=E_ext, edge_id=identity.edge_id.to_numpy(np.int64))
|
|
shadow_rows.append({"seed":seed,"pairing":pairing,"random_US_step":random_step,
|
|
"external_matched_E_reduction_sum":float(np.sum(1-E_ext)),"all_DAN_shadow_E_reduction_sum":float(np.sum(1-E_all)),
|
|
"matched_endogenous_DAN_steps":endogenous_steps,"matched_endogenous_DAN_spikes":endogenous_events,
|
|
"random_US_trace_positive_edges":int(np.count_nonzero(comp & (random_trace>0))) if random_trace is not None else 0,
|
|
"frozen_trace_match":exact,"max_abs_eligibility_error":err})
|
|
if capture and run["all_events"] is not None:
|
|
run["all_events"].to_parquet(exp / "whole-brain" / f"spike-events-seed-{seed}-pairings-{pairing}.parquet", index=False)
|
|
if pairing % 20 == 0:
|
|
print(json.dumps({"seed":seed,"pairing":pairing,"endogenous_steps":endogenous_steps,"E_external_sum":float(np.sum(1-E_ext)),"E_all_sum":float(np.sum(1-E_all))}),flush=True)
|
|
pd.DataFrame(audit_rows).to_parquet(exp / "local-eligibility/event-level-updates.parquet", index=False)
|
|
pd.DataFrame(gate_rows).to_parquet(exp / "dan-source/event-level-gates.parquet", index=False)
|
|
pd.concat(kc_event_frames,ignore_index=True).to_parquet(exp / "local-eligibility/KC-spike-history.parquet",index=False)
|
|
if dan_event_frames:
|
|
pd.concat(dan_event_frames,ignore_index=True).to_parquet(exp / "dan-source/endogenous-DAN-spike-history.parquet",index=False)
|
|
pd.DataFrame(shadow_rows).to_parquet(exp / "endogenous-dan/shadow-trajectory.parquet", index=False)
|
|
write_json(exp / "regression/event-trace-reproduction.json", {"all_passed":len(mismatches)==0,"mismatched_trials":mismatches,
|
|
"mismatched_count":len(mismatches),"audited_trial_count":len(shadow_rows),
|
|
"next_trial_overlay_restored_from_sealed_Phase4M_history":True})
|
|
pd.DataFrame(perf).to_parquet(exp / "performance/event-audit-calls.parquet", index=False)
|
|
write_json(exp / "endogenous-dan/summary.json", {"mode_A": "observational frozen Phase 4M LIF with endogenous DAN spikes recorded",
|
|
"mode_B": "external US-only plasticity gate, exact Phase 4M E reproduction",
|
|
"all_DAN_mode": "counterfactual open-loop shadow E; same frozen Phase 4M additive brain trajectory; not a new brain model",
|
|
"all_DAN_scope": "only metadata-confirmed DAN mapped to matched compartment; other DAN unresolved",
|
|
"frozen_brain_dynamics_changed": False, "training_seed_count":len(ncfg["train_seeds"]),"pairings_per_seed":args.max_pairings})
|
|
write_json(exp / "performance/event-audit.json", {"calls":len(perf),"total_wall_seconds":float(sum(x["wall_seconds"] for x in perf)),
|
|
"mean_wall_seconds":float(np.mean([x["wall_seconds"] for x in perf])),"peak_VRAM_bytes":max(x["peak_vram_bytes"] for x in perf),
|
|
"peak_RAM_KiB":int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss),
|
|
"KC_boolean_history_bytes_per_trial":int(len(run["kc"])*len(v.kc)),"DAN_boolean_history_bytes_per_trial":int(len(run["dan"])*len(all_dan)),
|
|
"no_dense_time_by_edge_tensor":True})
|
|
|
|
|
|
if __name__ == "__main__": main()
|