140 lines
7.6 KiB
Python
140 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Five read-only repeats of sealed Phase 4M's 500-trial eligibility trajectory."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
|
|
from phase4n import load_config
|
|
from phase4i import Validator
|
|
from phase4j import apply_v2
|
|
from phase4l import OdorSimulator, pn_drive
|
|
from phase4m import load_odor_inputs
|
|
|
|
ROOT = Path("/data/digital-drosophila")
|
|
EXP = ROOT / "experiments/phase4o-runtime-credit"
|
|
|
|
|
|
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 arrhash(x: np.ndarray) -> str:
|
|
return hashlib.sha256(x.tobytes()).hexdigest()
|
|
|
|
|
|
def first_diff(a: np.ndarray, b: np.ndarray):
|
|
n = min(len(a), len(b))
|
|
mismatched = np.flatnonzero(a[:n] != b[:n])
|
|
j = int(mismatched[0]) if len(mismatched) else (n if len(a) != len(b) else None)
|
|
if j is None:
|
|
return None
|
|
return {"event_position": j, "reference": None if j >= len(a) else int(a[j]), "observed": None if j >= len(b) else int(b[j])}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--repeats", type=int, default=5)
|
|
args = ap.parse_args()
|
|
ncfg, mcfg, cfg, cfgpath = load_config(ROOT / "models/male-cns-lif/credit-assignment-v1/phase4n.yaml")
|
|
history = pd.read_parquet(Path(mcfg["experiment_root"]) / "efficacy-overlay/update-history.parquet")
|
|
frozen = {(int(s), int(p)): f for (s, p), f in history.groupby(["train_seed", "pairing"])}
|
|
old = pd.read_parquet(ROOT / "experiments/phase4n-credit-assignment/local-eligibility/KC-spike-history.parquet")
|
|
old_events = {(int(s), int(p)): f[["step", "neuron_idx"]].to_numpy(np.int32)
|
|
for (s, p), f in old.groupby(["seed", "pairing"])}
|
|
v = Validator(cfg, cfgpath, EXP)
|
|
sim = OdorSimulator(v, cfg, EXP)
|
|
panel, vectors, selection, mapping = load_odor_inputs(cfg)
|
|
oid = selection["ODOR_A"]["odor_id"]
|
|
pn = pn_drive(oid, vectors[oid], mapping, selection["glomeruli_order"], str(mcfg["protocol"]["input_mode"]),
|
|
float(np.median(panel.mapped_response_sum)), cfg)[:2]
|
|
target = str(mcfg["protocol"]["target_compartment_verbatim"])
|
|
dan = v.dan_targets(str(mcfg["protocol"]["matched_dan_type"]), target)
|
|
src = v.cand.source_neuron_idx.to_numpy(np.int32)
|
|
lookup = np.full(v.n, -1, np.int32)
|
|
lookup[v.kc] = np.arange(len(v.kc), dtype=np.int32)
|
|
loc = lookup[src]
|
|
comp = v.cand.target_compartment_verbatim.astype(str).to_numpy() == target
|
|
ids = v.edge_ids.astype(np.int64)
|
|
id_to_local = pd.Series(np.arange(len(ids), dtype=np.int32), index=ids)
|
|
seen = {}
|
|
rows = []
|
|
first_failures = []
|
|
for repeat in range(args.repeats):
|
|
for seed in ncfg["train_seeds"]:
|
|
delta = np.zeros(len(ids), np.float32)
|
|
for pairing in range(1, 101):
|
|
start = time.perf_counter()
|
|
run = sim.run_rates(*pn, seed+pairing, float(mcfg["protocol"]["duration_ms"]),
|
|
delta=delta, dan_targets=dan, dan_pulses=[float(mcfg["protocol"]["dan_time_ms"])],
|
|
collect_kc_spikes=True)
|
|
trace = run["kc_trace_at_dan"]
|
|
events = run["kc_spikes"][["step", "neuron_idx"]].to_numpy(np.int32)
|
|
old_ev = old_events.get((seed, pairing), np.empty((0, 2), np.int32))
|
|
rec = frozen[(seed, pairing)]
|
|
local = id_to_local.loc[rec.edge_id.to_numpy(np.int64)].to_numpy(np.int32)
|
|
value = trace[loc[local]]
|
|
recorded = rec.eligibility.to_numpy(np.float32)
|
|
max_err = float(np.max(np.abs(value-recorded))) if len(value) else 0.0
|
|
nonzero_now = np.flatnonzero(comp & (trace[loc] > 0))
|
|
frozen_active = set(map(int, local))
|
|
active_symmetric_diff = len(set(map(int, nonzero_now)) ^ frozen_active)
|
|
key = (seed,pairing)
|
|
h = arrhash(events)
|
|
if key not in seen:
|
|
seen[key] = h
|
|
row = {"repeat": repeat, "seed": seed, "pairing": pairing,
|
|
"KC_event_count": len(events), "KC_event_sha256": h,
|
|
"KC_events_same_as_first_repeat": h == seen[key],
|
|
"KC_events_same_as_phase4n_audit": bool(np.array_equal(events, old_ev)),
|
|
"first_KC_event_difference_from_phase4n": first_diff(old_ev[:,0],events[:,0]) if not np.array_equal(events,old_ev) else None,
|
|
"eligibility_max_abs_error_vs_frozen_phase4m": max_err,
|
|
"eligibility_active_edge_symmetric_difference": active_symmetric_diff,
|
|
"exact_eligibility": bool(max_err <= 2e-5 and active_symmetric_diff == 0),
|
|
"trace_sha256": arrhash(trace), "runtime_wall_seconds": run["wall_seconds"],
|
|
"audit_wall_seconds": time.perf_counter()-start}
|
|
if not row["exact_eligibility"]:
|
|
err = np.abs(value-recorded)
|
|
j = int(np.argmax(err)) if len(err) else None
|
|
row["first_edge_mismatch"] = None if j is None else {
|
|
"edge_id": int(rec.edge_id.iloc[j]), "KC_neuron_idx": int(src[local[j]]),
|
|
"DAN_pulse_ms": float(mcfg["protocol"]["dan_time_ms"]),
|
|
"observed_trace": float(value[j]), "frozen_trace": float(recorded[j]),
|
|
"delta_trace": float(value[j]-recorded[j])}
|
|
if len(first_failures) < 20:
|
|
first_failures.append(row)
|
|
path = EXP / "eligibility-replay" / f"mismatch-repeat-{repeat}-seed-{seed}-pairing-{pairing}-KC-events.parquet"
|
|
path.parent.mkdir(parents=True,exist_ok=True)
|
|
run["kc_spikes"].to_parquet(path,index=False)
|
|
rows.append(row)
|
|
# Keep subsequent brain input equal to the sealed historical additive trajectory,
|
|
# even when a diagnostic trial diverged; no frozen state is ever overwritten.
|
|
sealed_trace = np.zeros(len(v.kc), np.float32)
|
|
sealed_trace[loc[local]] = recorded
|
|
delta, _ = apply_v2(v, sealed_trace, target, 1.0, delta, cfg)
|
|
if pairing % 20 == 0:
|
|
save(EXP / "eligibility-replay/reproducibility-matrix.json", rows)
|
|
print(json.dumps({"repeat":repeat,"seed":seed,"pairing":pairing,
|
|
"exact_in_block":sum(r["exact_eligibility"] for r in rows[-20:]),
|
|
"event_repeat_exact_in_block":sum(r["KC_events_same_as_first_repeat"] for r in rows[-20:])}),flush=True)
|
|
save(EXP / "eligibility-replay/reproducibility-matrix.json", rows)
|
|
save(EXP / "eligibility-replay/summary.json", {"complete_repeats": args.repeats, "trials_per_repeat":500,
|
|
"exact_frozen_eligibility_count":sum(r["exact_eligibility"] for r in rows),
|
|
"exact_repeat_event_count":sum(r["KC_events_same_as_first_repeat"] for r in rows),
|
|
"all_exact":all(r["exact_eligibility"] and r["KC_events_same_as_first_repeat"] for r in rows),
|
|
"first_failures":first_failures,
|
|
"historical_training_KC_event_sequence_available":False,
|
|
"phase4n_event_sequence_used_as_comparator":True,
|
|
"note":"Phase4M sealed only eligibility/update history, not all training KC events; Phase4N KC events are a later audit comparator."})
|
|
|
|
|
|
if __name__ == "__main__": main()
|