Digital Drosophila Archive
219 lines
14 KiB
Python
219 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Phase 4N read-only reconstruction of frozen Phase 4M acquisition and replay."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import yaml
|
|
|
|
from phase4h_audit import sha256, verify_checksum, write_json
|
|
from phase4i import Validator
|
|
from phase4k import seal
|
|
from phase4m import ReplayEngine, compare_A_B, frozen_checks, frozen_hashes, load_odor_inputs, replay_effect
|
|
from phase4l import OdorSimulator, pn_drive
|
|
from phase4j import effect
|
|
|
|
|
|
GROUPS = ("A-specific", "shared", "B-specific", "inactive")
|
|
|
|
|
|
def load_config(path: Path):
|
|
ncfg = yaml.safe_load(path.read_text())
|
|
mpath = Path(ncfg["frozen_phase4m_config"])
|
|
mcfg = yaml.safe_load(mpath.read_text())
|
|
lpath = Path(mcfg["frozen_phase4l_config"])
|
|
cfg = yaml.safe_load(lpath.read_text())
|
|
cfg["paths"]["phase4l_root"] = cfg["paths"]["experiment_root"]
|
|
cfg["paths"]["phase4l_model_root"] = str(lpath.parent)
|
|
return ncfg, mcfg, cfg, lpath
|
|
|
|
|
|
def prepare(ncfg, mcfg, cfg, exp: Path):
|
|
if exp.exists() and (exp / "edge-identity/definition.json").exists():
|
|
raise RuntimeError("Phase 4N identity/definition already frozen; refusing to overwrite")
|
|
previous = frozen_checks(cfg)
|
|
mroot = Path(mcfg["experiment_root"])
|
|
mmodel = Path(mcfg["model_root"])
|
|
for root, rel in ((mroot, "manifests/checksums.sha256"), (mmodel, "checksums.sha256")):
|
|
check = verify_checksum(root, root / rel)
|
|
if not check["passed"]:
|
|
raise RuntimeError(f"Phase 4M frozen checksum failed: {root}")
|
|
for sub in ("edge-identity", "acquisition-trajectory", "specificity-curve", "dan-source", "endogenous-dan", "local-eligibility", "shared-edge", "b-contamination", "isolated", "whole-brain", "controls", "stopping-analysis", "performance", "regression", "audit", "manifests"):
|
|
(exp / sub).mkdir(parents=True, exist_ok=True)
|
|
src = Path(cfg["paths"]["phase4l_root"]) / "eligible-edges/edge-identity.parquet"
|
|
identity = pd.read_parquet(src)
|
|
required = {"edge_id", "source_neuron_idx", "source_body_id", "target_neuron_idx", "target_body_id", "raw_synapse_count", "target_compartment_verbatim", "odor_eligibility_category"}
|
|
if not required.issubset(identity.columns) or len(identity) != 14232:
|
|
raise RuntimeError("Phase 4L identity schema/size mismatch")
|
|
if set(identity.odor_eligibility_category.astype(str)) != set(GROUPS):
|
|
raise RuntimeError("identity groups differ from four frozen groups")
|
|
identity["anatomical_weight_mv"] = identity.raw_synapse_count.astype(np.float32) * np.float32(cfg["model"]["w_syn_mv_per_anatomical_synapse"])
|
|
identity["initial_E"] = np.float32(1.0)
|
|
identity.rename(columns={"odor_eligibility_category": "identity_group", "target_compartment_verbatim": "compartment"}, inplace=True)
|
|
out = exp / "edge-identity/frozen-edge-identity.parquet"
|
|
identity.to_parquet(out, index=False)
|
|
definition = {"source": str(src), "source_sha256": sha256(src), "source_semantics": ncfg["edge_identity_source"],
|
|
"group_counts": identity.identity_group.value_counts().to_dict(), "edge_count": len(identity),
|
|
"checkpoint_rule": ncfg["checkpoints"], "CSI": ncfg["credit_specificity_index"],
|
|
"functional_separation": ncfg["functional_separation"], "peak_rule": ncfg["peak_rule"],
|
|
"generalization_onset_rule": ncfg["generalization_onset_rule"], "frozen_before_learning": True,
|
|
"phase4m_experiment_checksum_sha256": sha256(mroot / "manifests/checksums.sha256"),
|
|
"phase4m_model_checksum_sha256": sha256(mmodel / "checksums.sha256"), "prior_checksum_passed": previous["all_passed"]}
|
|
write_json(exp / "edge-identity/definition.json", definition)
|
|
shutil.copy2(Path(__file__).resolve(), exp / "audit/phase4n.py")
|
|
return identity, definition
|
|
|
|
|
|
def reconstruct(ncfg, mcfg, cfg, exp: Path, identity: pd.DataFrame):
|
|
hist = pd.read_parquet(Path(mcfg["experiment_root"]) / "efficacy-overlay/update-history.parquet")
|
|
required = {"train_seed", "pairing", "edge_id", "E_bounded_before", "E_bounded_after", "E_multiplicative_before", "E_multiplicative_after", "eligibility"}
|
|
if not required.issubset(hist.columns):
|
|
raise RuntimeError("Phase 4M update history schema mismatch")
|
|
edge_ids = identity.edge_id.to_numpy(np.int64)
|
|
lookup = pd.Series(np.arange(len(edge_ids), dtype=np.int32), index=edge_ids)
|
|
checkpoints = {}
|
|
verify = []
|
|
points = list(map(int, ncfg["checkpoints"]))
|
|
all_rows = []
|
|
for seed in ncfg["train_seeds"]:
|
|
by = hist[hist.train_seed == seed]
|
|
E = np.ones(len(edge_ids), np.float32)
|
|
U = np.ones(len(edge_ids), np.float32)
|
|
checkpoints[seed] = {0: {"bounded": E.copy(), "unbounded": U.copy()}}
|
|
np.savez_compressed(exp / "acquisition-trajectory" / f"E-seed-{seed}-pairings-0.npz", E_bounded=E, E_unbounded=U, edge_id=edge_ids)
|
|
for mechanism in ("bounded", "unbounded"):
|
|
for group in GROUPS:
|
|
mask = identity.identity_group.to_numpy() == group
|
|
all_rows.append({"seed": seed, "pairings": 0, "mechanism": mechanism, "group": group,
|
|
"total_edges": int(mask.sum()), "updated_edges": 0, "mean_delta_E": 0.0, "median_delta_E": 0.0,
|
|
"weighted_transmission_reduction": 0.0, "E_change_total": 0.0})
|
|
for pairing in range(1, max(points) + 1):
|
|
item = by[by.pairing == pairing]
|
|
local = lookup.loc[item.edge_id.to_numpy(np.int64)].to_numpy(np.int32)
|
|
before_error = max(float(np.max(np.abs(E[local] - item.E_bounded_before.to_numpy(np.float32)))) if len(local) else 0.0,
|
|
float(np.max(np.abs(U[local] - item.E_multiplicative_before.to_numpy(np.float32)))) if len(local) else 0.0)
|
|
if before_error > 1e-6:
|
|
raise RuntimeError(f"Phase 4M event history not reconstructible seed={seed} pairing={pairing}: {before_error}")
|
|
E[local] = item.E_bounded_after.to_numpy(np.float32)
|
|
U[local] = item.E_multiplicative_after.to_numpy(np.float32)
|
|
if pairing not in points:
|
|
continue
|
|
checkpoints[seed][pairing] = {"bounded": E.copy(), "unbounded": U.copy()}
|
|
np.savez_compressed(exp / "acquisition-trajectory" / f"E-seed-{seed}-pairings-{pairing}.npz", E_bounded=E, E_unbounded=U, edge_id=edge_ids)
|
|
old = Path(mcfg["experiment_root"]) / "efficacy-overlay" / f"E-seed-{seed}-pairings-{pairing}.npz"
|
|
if old.exists():
|
|
with np.load(old) as f:
|
|
err = max(float(np.max(np.abs(E - f["E_bounded_multiplicative"]))), float(np.max(np.abs(U - f["E_multiplicative"]))))
|
|
verify.append({"seed": seed, "pairings": pairing, "max_abs_error": err})
|
|
if err > 1e-6:
|
|
raise RuntimeError("reconstructed E differs from sealed Phase 4M checkpoint")
|
|
for mechanism, value in (("bounded", E), ("unbounded", U)):
|
|
for group in GROUPS:
|
|
mask = identity.identity_group.to_numpy() == group
|
|
w = identity.anatomical_weight_mv.to_numpy(np.float32)[mask]
|
|
reduction = 1 - value[mask]
|
|
all_rows.append({"seed": seed, "pairings": pairing, "mechanism": mechanism, "group": group,
|
|
"total_edges": int(mask.sum()), "updated_edges": int(np.count_nonzero(reduction > 1e-7)),
|
|
"mean_delta_E": float((-reduction).mean()), "median_delta_E": float(np.median(-reduction)),
|
|
"weighted_transmission_reduction": float(np.dot(w.astype(np.float64), reduction.astype(np.float64)) / w.sum()),
|
|
"E_change_total": float(reduction.sum())})
|
|
table = pd.DataFrame(all_rows)
|
|
table.to_parquet(exp / "acquisition-trajectory/group-plasticity.parquet", index=False)
|
|
wide = table.pivot_table(index=["seed", "pairings", "mechanism"], columns="group", values="weighted_transmission_reduction").reset_index()
|
|
a, b, shared = (wide[x].to_numpy(np.float64) for x in ("A-specific", "B-specific", "shared"))
|
|
wide["CSI"] = a - np.maximum(shared, b)
|
|
wide["A_to_shared_ratio"] = np.divide(a, shared, out=np.full_like(a, np.nan), where=shared != 0)
|
|
wide["A_to_B_ratio"] = np.divide(a, b, out=np.full_like(a, np.nan), where=b != 0)
|
|
wide.to_parquet(exp / "specificity-curve/CSI.parquet", index=False)
|
|
write_json(exp / "regression/phase4m-history-reconstruction.json", {"passed": True, "old_checkpoint_comparisons": verify, "max_error": max((v["max_abs_error"] for v in verify), default=0.0)})
|
|
return checkpoints, table, wide, hist
|
|
|
|
|
|
def isolated(ncfg, mcfg, cfg, exp: Path, identity: pd.DataFrame, checkpoints):
|
|
v = Validator(cfg, Path(mcfg["frozen_phase4l_config"]), exp)
|
|
if not np.array_equal(v.edge_ids, identity.edge_id.to_numpy(np.int64)):
|
|
raise RuntimeError("plastic candidate ordering mismatch")
|
|
sim = OdorSimulator(v, cfg, exp)
|
|
comp = identity.compartment.astype(str).to_numpy() == str(ncfg["matched_compartment_verbatim"])
|
|
targets = np.unique(identity.loc[comp, "target_neuron_idx"].to_numpy(np.int32))
|
|
replay = ReplayEngine(sim, targets)
|
|
anatomical = identity.anatomical_weight_mv.to_numpy(np.float32)
|
|
old = Path(mcfg["experiment_root"])
|
|
evaldata = {}
|
|
zero = np.zeros(len(v.edge_ids), np.float32)
|
|
for seed in ncfg["evaluation_seeds"]:
|
|
for odor in ("A", "B"):
|
|
spikes = pd.read_parquet(old / "isolated-replay" / f"KC-spikes-ODOR_{odor}-seed-{seed}.parquet")
|
|
evaldata[(seed, odor)] = (spikes, replay.run(spikes, zero, float(mcfg["protocol"]["duration_ms"])))
|
|
rows = []
|
|
replay_wall = []
|
|
for train_seed, eval_seed in zip(ncfg["train_seeds"], ncfg["evaluation_seeds"]):
|
|
for pair in ncfg["checkpoints"]:
|
|
E = checkpoints[train_seed][pair]["bounded"]
|
|
delta = anatomical * (E - 1)
|
|
rec = {"train_seed": train_seed, "eval_seed": eval_seed, "pairings": pair}
|
|
effects = {}
|
|
for odor in ("A", "B"):
|
|
spikes, baseline = evaldata[(eval_seed, odor)]
|
|
post = replay.run(spikes, delta, float(mcfg["protocol"]["duration_ms"]))
|
|
replay_wall.append(post["wall_seconds"])
|
|
effects[odor] = replay_effect(baseline, post, sim.dt)
|
|
rec.update({f"{odor}_{k}": value for k, value in effects[odor].items() if k != "isolated_input_scope"})
|
|
rec["A_specific_functional"] = compare_A_B(effects["A"], effects["B"])
|
|
rec["functional_separation_spike_L1"] = float(effects["A"]["MBON_spike_L1"] - effects["B"]["MBON_spike_L1"])
|
|
al = effects["A"]["first_latency_delta_ms"] or 0.0
|
|
bl = effects["B"]["first_latency_delta_ms"] or 0.0
|
|
rec["functional_separation_latency_ms"] = float(al - bl)
|
|
rows.append(rec)
|
|
frame = pd.DataFrame(rows)
|
|
frame.to_parquet(exp / "isolated/response-by-checkpoint.parquet", index=False)
|
|
curve = frame.groupby("pairings").agg(A_specific_functional_seeds=("A_specific_functional", "sum"),
|
|
mean_spike_L1_separation=("functional_separation_spike_L1", "mean"),
|
|
mean_latency_separation_ms=("functional_separation_latency_ms", "mean"),
|
|
A_MBON_spike_L1=("A_MBON_spike_L1", "mean"), B_MBON_spike_L1=("B_MBON_spike_L1", "mean"),
|
|
A_transmission_reduction_percent=("A_transmitted_reduction_percent", "mean"),
|
|
B_transmission_reduction_percent=("B_transmitted_reduction_percent", "mean")).reset_index()
|
|
curve.to_parquet(exp / "specificity-curve/memory-specificity.parquet", index=False)
|
|
valid = curve[curve.A_specific_functional_seeds >= 3].copy()
|
|
if valid.empty:
|
|
peak = None
|
|
else:
|
|
best_spikes = float(valid.mean_spike_L1_separation.max())
|
|
valid = valid[np.isclose(valid.mean_spike_L1_separation, best_spikes, atol=1e-9, rtol=0)]
|
|
best_latency = float(valid.mean_latency_separation_ms.max())
|
|
valid = valid[valid.mean_latency_separation_ms >= best_latency - sim.dt / 2]
|
|
peak = int(valid.pairings.min())
|
|
if peak is not None:
|
|
# Equal score/latency checkpoints are resolved to earliest by the predeclared rule.
|
|
peak_csi = float(np.mean(pd.read_parquet(exp / "specificity-curve/CSI.parquet").query("mechanism == 'bounded' and pairings == @peak").CSI))
|
|
else:
|
|
peak_csi = None
|
|
write_json(exp / "specificity-curve/peak.json", {"status": "SPECIFICITY-PEAK-IDENTIFIED" if peak is not None else "UNIDENTIFIED", "pairings": peak,
|
|
"peak_CSI": peak_csi, "selection_rule": ncfg["peak_rule"]})
|
|
write_json(exp / "performance/isolated.json", {"mean_replay_wall_seconds": float(np.mean(replay_wall)), "total_replay_wall_seconds": float(np.sum(replay_wall)), "count": len(replay_wall)})
|
|
return v, sim, frame, curve, peak
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--config", required=True, type=Path)
|
|
args = ap.parse_args()
|
|
ncfg, mcfg, cfg, base_path = load_config(args.config)
|
|
exp = Path(ncfg["experiment_root"])
|
|
identity, definition = prepare(ncfg, mcfg, cfg, exp)
|
|
print(json.dumps({"stage": "identity_frozen", "group_counts": definition["group_counts"], "identity_sha256": sha256(exp / "edge-identity/frozen-edge-identity.parquet")}), flush=True)
|
|
checkpoints, group, csi, history = reconstruct(ncfg, mcfg, cfg, exp, identity)
|
|
print(json.dumps({"stage": "trajectory_reconstructed", "checkpoint_count": len(ncfg["checkpoints"]), "updates": len(history)}), flush=True)
|
|
v, sim, frame, curve, peak = isolated(ncfg, mcfg, cfg, exp, identity, checkpoints)
|
|
print(json.dumps({"stage": "isolated_complete", "peak": peak, "curve": curve.to_dict("records")}), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|