Files
digital-drosophila/models/male-cns-lif/runtime-credit-v1/phase4o_eligibility.py
T

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()