|
|
|
@@ -0,0 +1,468 @@
|
|
|
|
|
#!/usr/bin/env python3
|
|
|
|
|
"""Phase 4H stage 2: conservative dopamine-gated KC-to-MBON plasticity overlay."""
|
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
|
|
|
|
import argparse
|
|
|
|
|
import gc
|
|
|
|
|
import hashlib
|
|
|
|
|
import json
|
|
|
|
|
import math
|
|
|
|
|
import shutil
|
|
|
|
|
import sys
|
|
|
|
|
import time
|
|
|
|
|
from collections import Counter
|
|
|
|
|
from datetime import datetime, timezone
|
|
|
|
|
from pathlib import Path
|
|
|
|
|
from typing import Any
|
|
|
|
|
|
|
|
|
|
import numpy as np
|
|
|
|
|
import pandas as pd
|
|
|
|
|
import pyarrow.parquet as pq
|
|
|
|
|
import torch
|
|
|
|
|
import yaml
|
|
|
|
|
|
|
|
|
|
from brain_runtime import BrainRuntime, InputBatch, InputEvent
|
|
|
|
|
from phase4h_audit import run_frozen_audit, sha256, tree_manifest, write_json
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def vector_counts(nodes: np.ndarray, members: np.ndarray, n: int) -> np.ndarray:
|
|
|
|
|
lookup = np.full(n, -1, dtype=np.int32)
|
|
|
|
|
lookup[members] = np.arange(len(members), dtype=np.int32)
|
|
|
|
|
local = lookup[nodes]
|
|
|
|
|
local = local[local >= 0]
|
|
|
|
|
return np.bincount(local, minlength=len(members)).astype(np.int32)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def cosine_distance(a: np.ndarray, b: np.ndarray) -> float:
|
|
|
|
|
af = a.astype(np.float64); bf = b.astype(np.float64)
|
|
|
|
|
denom = float(np.linalg.norm(af) * np.linalg.norm(bf))
|
|
|
|
|
return 0.0 if denom == 0 else float(1.0 - np.dot(af, bf) / denom)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def separation(a: dict[str, Any], b: dict[str, Any], key: str) -> dict[str, float]:
|
|
|
|
|
x = a[key].astype(np.float64); y = b[key].astype(np.float64)
|
|
|
|
|
return {
|
|
|
|
|
"euclidean": float(np.linalg.norm(x - y)),
|
|
|
|
|
"cosine_distance": cosine_distance(x, y),
|
|
|
|
|
"l1": float(np.abs(x - y).sum()),
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def make_batch(targets: np.ndarray, rate_hz: float, duration_ms: float, name: str,
|
|
|
|
|
dan_targets: np.ndarray | None = None, dan_start_ms: float = 80.0,
|
|
|
|
|
dan_duration_ms: float = 20.0, dan_rate_hz: float = 100.0) -> tuple[InputBatch, float]:
|
|
|
|
|
events = [InputEvent(
|
|
|
|
|
target_neuron_idx=targets.astype(np.int64), stimulation_mode="spike_injection",
|
|
|
|
|
value_kind="poisson_rate_hz", value=float(rate_hz), start_ms=0.0,
|
|
|
|
|
duration_ms=float(duration_ms), source_id="phase4h-engineering-olfactory-CS",
|
|
|
|
|
channel_id=name, metadata={"interpretation": "engineering olfactory CS interface"},
|
|
|
|
|
)]
|
|
|
|
|
total_ms = duration_ms
|
|
|
|
|
if dan_targets is not None and len(dan_targets):
|
|
|
|
|
events.append(InputEvent(
|
|
|
|
|
target_neuron_idx=dan_targets.astype(np.int64), stimulation_mode="spike_injection",
|
|
|
|
|
value_kind="poisson_rate_hz", value=float(dan_rate_hz), start_ms=float(dan_start_ms),
|
|
|
|
|
duration_ms=float(dan_duration_ms), source_id="phase4h-artificial-US-interface",
|
|
|
|
|
channel_id="DAN-like-reinforcement",
|
|
|
|
|
metadata={
|
|
|
|
|
"interpretation": "artificial unconditioned-signal interface",
|
|
|
|
|
"not_equal_to": "measured dopamine concentration",
|
|
|
|
|
},
|
|
|
|
|
))
|
|
|
|
|
total_ms = max(total_ms, dan_start_ms + dan_duration_ms)
|
|
|
|
|
return InputBatch(events=tuple(events), batch_id=name, metadata={"phase": "4H"}), float(total_ms)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Experiment:
|
|
|
|
|
def __init__(self, cfg: dict[str, Any], exp: Path):
|
|
|
|
|
self.cfg = cfg; self.exp = exp
|
|
|
|
|
self.rt = BrainRuntime(Path(__file__).with_name("phase4h.yaml"), exp / "compiler-reports/extension-build")
|
|
|
|
|
self.base_weights = self.rt.weights.detach().clone()
|
|
|
|
|
self.n = self.rt.n
|
|
|
|
|
self.metadata = self.rt.table.to_pandas().sort_values("neuron_idx").reset_index(drop=True)
|
|
|
|
|
self.classes = self.metadata["annotation_class"].fillna("").astype(str).to_numpy()
|
|
|
|
|
self.superclass = self.metadata["annotation_superclass"].fillna("").astype(str).to_numpy()
|
|
|
|
|
self.types = self.metadata["annotation_type"].fillna("").astype(str).to_numpy()
|
|
|
|
|
self.nt = self.metadata["nt_consensus_nt"].fillna("unclear").astype(str).str.lower().to_numpy()
|
|
|
|
|
self.kc = np.flatnonzero(self.classes == "Kenyon_Cell").astype(np.int32)
|
|
|
|
|
self.mbon = np.flatnonzero(self.classes == "MBON").astype(np.int32)
|
|
|
|
|
self.desc = self.rt.output_populations["descending"]
|
|
|
|
|
self.motor = self.rt.output_populations["motor_related"]
|
|
|
|
|
self.central = np.flatnonzero(np.char.startswith(self.superclass.astype(str), "cb_")).astype(np.int32)
|
|
|
|
|
downstream = pq.read_table(exp / "mappings/mbon-to-downstream-actual-edges.parquet",
|
|
|
|
|
columns=["target_neuron_idx"])["target_neuron_idx"].to_numpy()
|
|
|
|
|
self.downstream = np.unique(downstream).astype(np.int32)
|
|
|
|
|
self.candidates = pq.read_table(exp / "plastic-mask/plastic-edge-candidates.parquet").to_pandas()
|
|
|
|
|
self.edge_ids = self.candidates["edge_id"].to_numpy(dtype=np.int64)
|
|
|
|
|
self.base_candidate = self.base_weights[torch.as_tensor(self.edge_ids, device=self.rt.device)].cpu().numpy()
|
|
|
|
|
self.sign_policy = cfg["sign_policy"]["neurotransmitter_sign"]
|
|
|
|
|
olfactory = np.flatnonzero(self.rt.population_masks["olfactory"]).astype(np.int32)
|
|
|
|
|
take = int(cfg["conditioning"]["cs_targets_per_identity"])
|
|
|
|
|
if len(olfactory) < 2 * take:
|
|
|
|
|
take = len(olfactory) // 2
|
|
|
|
|
if take == 0:
|
|
|
|
|
raise RuntimeError("olfactory population is too small for two disjoint CS identities")
|
|
|
|
|
olfactory.sort(); self.cs_a = olfactory[:take]; self.cs_b = olfactory[take:2 * take]
|
|
|
|
|
if np.intersect1d(self.cs_a, self.cs_b).size:
|
|
|
|
|
raise RuntimeError("CS identities overlap")
|
|
|
|
|
self.dan_targets = {}
|
|
|
|
|
for label in ("appetitive", "aversive"):
|
|
|
|
|
dtype = str(cfg["conditioning"][f"{label}_dan_type_required"])
|
|
|
|
|
comp = str(cfg["conditioning"][f"{label}_compartment_verbatim"])
|
|
|
|
|
rows = self.candidates[self.candidates["target_compartment_verbatim"] == comp]
|
|
|
|
|
confirmed = pq.read_table(exp / "mappings/dan-to-mbon-confirmed.parquet").to_pandas()
|
|
|
|
|
sources = confirmed[(confirmed["source_annotation_type"] == dtype)
|
|
|
|
|
& (confirmed["source_compartment_verbatim"] == comp)]["source_neuron_idx"].unique()
|
|
|
|
|
if rows.empty or len(sources) == 0:
|
|
|
|
|
raise RuntimeError(f"required exact DAN/compartment mapping absent: {label} {dtype} {comp}")
|
|
|
|
|
self.dan_targets[label] = np.sort(sources.astype(np.int32))
|
|
|
|
|
|
|
|
|
|
def install_delta(self, delta: np.ndarray) -> None:
|
|
|
|
|
if delta.shape != (len(self.edge_ids),):
|
|
|
|
|
raise ValueError("plastic delta shape mismatch")
|
|
|
|
|
self.rt.weights.copy_(self.base_weights)
|
|
|
|
|
ids = torch.as_tensor(self.edge_ids, dtype=torch.int64, device=self.rt.device)
|
|
|
|
|
values = torch.as_tensor(self.base_candidate + delta, dtype=torch.float32, device=self.rt.device)
|
|
|
|
|
self.rt.weights[ids] = values
|
|
|
|
|
|
|
|
|
|
def evaluate(self, cs_name: str, seed: int, delta: np.ndarray) -> dict[str, Any]:
|
|
|
|
|
self.install_delta(delta); self.rt.reset(seed)
|
|
|
|
|
targets = self.cs_a if cs_name == "CS-A" else self.cs_b
|
|
|
|
|
batch, duration = make_batch(targets, self.cfg["conditioning"]["cs_rate_hz"],
|
|
|
|
|
self.cfg["conditioning"]["cs_duration_ms"], f"eval-{cs_name}")
|
|
|
|
|
t0 = time.perf_counter(); self.rt.step(batch, duration); wall = time.perf_counter() - t0
|
|
|
|
|
events = self.rt.read_spikes(); nodes = events["neuron_idx"].astype(np.int32)
|
|
|
|
|
origins = events["origin"]
|
|
|
|
|
vectors = {
|
|
|
|
|
"kc": vector_counts(nodes, self.kc, self.n),
|
|
|
|
|
"mbon": vector_counts(nodes, self.mbon, self.n),
|
|
|
|
|
"downstream": vector_counts(nodes, self.downstream, self.n),
|
|
|
|
|
"central": vector_counts(nodes, self.central, self.n),
|
|
|
|
|
"descending": vector_counts(nodes, self.desc, self.n),
|
|
|
|
|
"motor": vector_counts(nodes, self.motor, self.n),
|
|
|
|
|
}
|
|
|
|
|
mbon_nt = self.nt[self.mbon]
|
|
|
|
|
positive_nt = {key for key, value in self.sign_policy.items() if int(value) > 0}
|
|
|
|
|
negative_nt = {key for key, value in self.sign_policy.items() if int(value) < 0}
|
|
|
|
|
mvec = vectors["mbon"]
|
|
|
|
|
pos = int(mvec[np.isin(mbon_nt, list(positive_nt))].sum())
|
|
|
|
|
neg = int(mvec[np.isin(mbon_nt, list(negative_nt))].sum())
|
|
|
|
|
finite = bool(torch.isfinite(self.rt.v).all().item() and torch.isfinite(self.rt.g).all().item())
|
|
|
|
|
active = int(np.unique(nodes).size)
|
|
|
|
|
return {
|
|
|
|
|
"seed": int(seed), "CS": cs_name, "wall_seconds": wall,
|
|
|
|
|
"total_spikes": int(len(nodes)), "network_spikes": int(np.count_nonzero(origins == "network_spike")),
|
|
|
|
|
"active_neurons": active, "active_fraction": float(active / self.n),
|
|
|
|
|
"finite": finite, "runaway": bool(active > 0.5 * self.n),
|
|
|
|
|
"population_spikes": {key: int(value.sum()) for key, value in vectors.items()},
|
|
|
|
|
"mbon_balance": {
|
|
|
|
|
"excitatory_transmitter_spikes": pos, "inhibitory_transmitter_spikes": neg,
|
|
|
|
|
"signed_transmitter_balance": pos - neg,
|
|
|
|
|
"interpretation": "transmitter-stratified MBON activity; not an assigned approach/avoidance valence",
|
|
|
|
|
},
|
|
|
|
|
**{f"_{key}_vector": value for key, value in vectors.items()},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
def update(self, cs_name: str, seed: int, delta: np.ndarray, compartment: str | None,
|
|
|
|
|
dan_label: str | None, delay_ms: float, enabled: bool, d_value: float) -> tuple[np.ndarray, dict[str, Any]]:
|
|
|
|
|
self.install_delta(delta); self.rt.reset(seed)
|
|
|
|
|
targets = self.cs_a if cs_name == "CS-A" else self.cs_b
|
|
|
|
|
dan = self.dan_targets[dan_label] if dan_label else None
|
|
|
|
|
dan_start = self.cfg["conditioning"]["cs_duration_ms"] + delay_ms if delay_ms > 0 else 80.0
|
|
|
|
|
batch, duration = make_batch(targets, self.cfg["conditioning"]["cs_rate_hz"],
|
|
|
|
|
self.cfg["conditioning"]["cs_duration_ms"], f"train-{cs_name}",
|
|
|
|
|
dan_targets=dan, dan_start_ms=dan_start)
|
|
|
|
|
self.rt.step(batch, duration)
|
|
|
|
|
nodes = self.rt.read_spikes()["neuron_idx"].astype(np.int32)
|
|
|
|
|
kc_counts = np.bincount(nodes[np.isin(nodes, self.kc)], minlength=self.n)
|
|
|
|
|
pre = self.candidates["source_neuron_idx"].to_numpy(dtype=np.int32)
|
|
|
|
|
eligible = np.minimum(kc_counts[pre].astype(np.float32), 1.0)
|
|
|
|
|
if delay_ms > 0:
|
|
|
|
|
eligible *= math.exp(-delay_ms / float(self.cfg["plasticity"]["eligibility_tau_ms"]))
|
|
|
|
|
comp_mask = (self.candidates["target_compartment_verbatim"].astype(str).to_numpy() == str(compartment)) if compartment else np.zeros(len(delta), bool)
|
|
|
|
|
update = np.zeros_like(delta)
|
|
|
|
|
if enabled and d_value != 0.0 and compartment is not None:
|
|
|
|
|
polarity = float(self.cfg["plasticity"]["default_confirmed_compartment_polarity"])
|
|
|
|
|
update[comp_mask] = polarity * float(self.cfg["plasticity"]["eta_effective_mv"]) * float(d_value) * eligible[comp_mask]
|
|
|
|
|
out = delta + update
|
|
|
|
|
lower = -float(self.cfg["plasticity"]["max_fractional_depression_per_edge"]) * self.base_candidate
|
|
|
|
|
out = np.maximum(out, lower).astype(np.float32)
|
|
|
|
|
if np.any(self.base_candidate + out < -1e-7):
|
|
|
|
|
raise RuntimeError("plastic overlay made an excitatory KC edge negative")
|
|
|
|
|
record = {
|
|
|
|
|
"CS": cs_name, "seed": int(seed), "compartment_verbatim": compartment,
|
|
|
|
|
"DAN_label": dan_label, "DAN_neuron_count": 0 if dan is None else int(len(dan)),
|
|
|
|
|
"D_value": float(d_value), "D_interpretation": "artificial US interface scalar, not dopamine concentration",
|
|
|
|
|
"delay_ms": float(delay_ms), "plasticity_enabled": bool(enabled),
|
|
|
|
|
"eligible_candidate_edges": int(np.count_nonzero(eligible * comp_mask)),
|
|
|
|
|
"updated_edges_this_trial": int(np.count_nonzero(update)),
|
|
|
|
|
"delta_l1_effective_mv": float(np.abs(out).sum()),
|
|
|
|
|
"delta_min_effective_mv": float(out.min(initial=0.0)), "delta_max_effective_mv": float(out.max(initial=0.0)),
|
|
|
|
|
}
|
|
|
|
|
return out, record
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def public_eval(value: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
return {key: item for key, item in value.items() if not key.startswith("_")}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def eval_pair(experiment: Experiment, seed: int, delta: np.ndarray) -> dict[str, Any]:
|
|
|
|
|
a = experiment.evaluate("CS-A", seed, delta)
|
|
|
|
|
b = experiment.evaluate("CS-B", seed, delta)
|
|
|
|
|
return {
|
|
|
|
|
"CS-A": public_eval(a), "CS-B": public_eval(b),
|
|
|
|
|
"mbon_separation": separation(a, b, "_mbon_vector"),
|
|
|
|
|
"downstream_separation": separation(a, b, "_downstream_vector"),
|
|
|
|
|
"descending_separation": separation(a, b, "_descending_vector"),
|
|
|
|
|
"motor_separation": separation(a, b, "_motor_vector"),
|
|
|
|
|
"balance_difference_A_minus_B": int(a["mbon_balance"]["signed_transmitter_balance"] - b["mbon_balance"]["signed_transmitter_balance"]),
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def protocol_plan(name: str, cfg: dict[str, Any], reliable: list[str], rng: np.random.Generator) -> list[dict[str, Any]]:
|
|
|
|
|
app = str(cfg["conditioning"]["appetitive_compartment_verbatim"])
|
|
|
|
|
av = str(cfg["conditioning"]["aversive_compartment_verbatim"])
|
|
|
|
|
trials = int(cfg["plasticity"]["training_trials"])
|
|
|
|
|
base = {"delay_ms": 0.0, "enabled": True, "d_value": float(cfg["plasticity"]["d_magnitude"])}
|
|
|
|
|
if name == "appetitive-conditioning":
|
|
|
|
|
return [base | {"CS": "CS-A", "compartment": app, "dan_label": "appetitive"} for _ in range(trials)]
|
|
|
|
|
if name == "aversive-conditioning":
|
|
|
|
|
return [base | {"CS": "CS-A", "compartment": av, "dan_label": "aversive"} for _ in range(trials)]
|
|
|
|
|
if name == "delayed-reward":
|
|
|
|
|
return [base | {"CS": "CS-A", "compartment": app, "dan_label": "appetitive",
|
|
|
|
|
"delay_ms": float(cfg["conditioning"]["delayed_reward_ms"])} for _ in range(trials)]
|
|
|
|
|
if name == "retention":
|
|
|
|
|
return [base | {"CS": "CS-A", "compartment": app, "dan_label": "appetitive"} for _ in range(trials)]
|
|
|
|
|
if name == "d-zero-control":
|
|
|
|
|
return [base | {"CS": "CS-A", "compartment": app, "dan_label": None, "d_value": 0.0} for _ in range(trials)]
|
|
|
|
|
if name == "frozen-brain-control":
|
|
|
|
|
return [base | {"CS": "CS-A", "compartment": app, "dan_label": "appetitive", "enabled": False} for _ in range(trials)]
|
|
|
|
|
if name == "random-reward-control":
|
|
|
|
|
plan = []
|
|
|
|
|
for _ in range(trials):
|
|
|
|
|
comp = str(rng.choice(reliable))
|
|
|
|
|
label = "appetitive" if comp == app else ("aversive" if comp == av else None)
|
|
|
|
|
plan.append(base | {"CS": str(rng.choice(["CS-A", "CS-B"])), "compartment": comp, "dan_label": label})
|
|
|
|
|
return plan
|
|
|
|
|
if name == "reversal-learning":
|
|
|
|
|
first = [base | {"CS": "CS-A", "compartment": app, "dan_label": "appetitive", "phase": "acquisition"} for _ in range(trials)]
|
|
|
|
|
second = [base | {"CS": "CS-B", "compartment": app, "dan_label": "appetitive", "phase": "reversal"} for _ in range(trials)]
|
|
|
|
|
return first + second
|
|
|
|
|
raise KeyError(name)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def phase4c_regression(cfg: dict[str, Any], exp: Path) -> dict[str, Any]:
|
|
|
|
|
rt = BrainRuntime(Path(__file__).with_name("phase4h.yaml"), exp / "compiler-reports/extension-build")
|
|
|
|
|
specs = [
|
|
|
|
|
("auditory-10hz", "auditory", 10.0, 42001),
|
|
|
|
|
("auditory-100hz", "auditory", 100.0, 42001),
|
|
|
|
|
("olfactory-50hz", "olfactory", 50.0, 42003),
|
|
|
|
|
("mechanosensory-50hz", "mechanosensory", 50.0, 42003),
|
|
|
|
|
("gustatory-50hz", "gustatory", 50.0, 42003),
|
|
|
|
|
]
|
|
|
|
|
results = {}
|
|
|
|
|
root = Path(cfg["paths"]["phase4c_root"]) / "regression/reference"
|
|
|
|
|
for name, population, rate, seed in specs:
|
|
|
|
|
targets = np.flatnonzero(rt.population_masks[population]).astype(np.int64)
|
|
|
|
|
batch, duration = make_batch(targets, rate, 1000.0, f"phase4c-regression-{name}")
|
|
|
|
|
rt.reset(seed); rt.step(batch, duration)
|
|
|
|
|
events = rt.read_spikes(); steps = events["step"].astype(np.int32); nodes = events["neuron_idx"].astype(np.int32)
|
|
|
|
|
ref = pq.read_table(root / name / "spike-events.parquet")
|
|
|
|
|
ref_steps = ref["step"].to_numpy().astype(np.int32); ref_nodes = ref["neuron_idx"].to_numpy().astype(np.int32)
|
|
|
|
|
counts = rt.spike_counts.cpu().numpy().astype(np.int32)
|
|
|
|
|
ref_counts = np.load(root / name / "per-neuron-spike-counts.int32.npy")
|
|
|
|
|
v = rt.v.cpu().numpy(); g = rt.g.cpu().numpy()
|
|
|
|
|
ref_v = np.load(root / name / "final-voltage.float32.npy"); ref_g = np.load(root / name / "final-g.float32.npy")
|
|
|
|
|
results[name] = {
|
|
|
|
|
"spike_events_exact": bool(np.array_equal(steps, ref_steps) and np.array_equal(nodes, ref_nodes)),
|
|
|
|
|
"per_neuron_counts_exact": bool(np.array_equal(counts, ref_counts)),
|
|
|
|
|
"max_abs_voltage_error": float(np.max(np.abs(v - ref_v))),
|
|
|
|
|
"max_abs_g_error": float(np.max(np.abs(g - ref_g))),
|
|
|
|
|
"event_count": int(len(nodes)),
|
|
|
|
|
}
|
|
|
|
|
results["all_passed"] = bool(all(v["spike_events_exact"] and v["per_neuron_counts_exact"]
|
|
|
|
|
and v["max_abs_voltage_error"] <= 0.00390625
|
|
|
|
|
and v["max_abs_g_error"] <= 0.00390625
|
|
|
|
|
for k, v in results.items() if k != "all_passed"))
|
|
|
|
|
write_json(exp / "regression/phase4c-reference-replay.json", results)
|
|
|
|
|
del rt; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
|
return results
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def seal(root: Path, manifest_path: Path, checksum_path: Path) -> None:
|
|
|
|
|
manifest = tree_manifest(root, {manifest_path, checksum_path})
|
|
|
|
|
write_json(manifest_path, manifest)
|
|
|
|
|
targets = sorted(path for path in root.rglob("*") if path.is_file() and path != checksum_path)
|
|
|
|
|
checksum_path.write_text("".join(f"{sha256(path)} {path.relative_to(root)}\n" for path in targets), encoding="utf-8")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def main() -> None:
|
|
|
|
|
ap = argparse.ArgumentParser(); ap.add_argument("--config", type=Path, required=True); args = ap.parse_args()
|
|
|
|
|
cfg = yaml.safe_load(args.config.read_text(encoding="utf-8"))
|
|
|
|
|
exp = Path(cfg["paths"]["experiment_root"]); learned = Path(cfg["paths"]["learned_root"]); model = Path(cfg["paths"]["model_root"])
|
|
|
|
|
for sub in ("conditioning", "controls", "plasticity", "metrics", "regression", "audit", "manifests", "compiler-reports/extension-build"):
|
|
|
|
|
(exp / sub).mkdir(parents=True, exist_ok=True)
|
|
|
|
|
for sub in ("plastic-delta", "training-metadata"):
|
|
|
|
|
(learned / sub).mkdir(parents=True, exist_ok=True)
|
|
|
|
|
model.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
gate_path = model / "STAGE1_AUDIT_GATE.json"
|
|
|
|
|
if not gate_path.is_file():
|
|
|
|
|
raise RuntimeError("stage-1 audit gate missing")
|
|
|
|
|
gate = json.loads(gate_path.read_text(encoding="utf-8"))
|
|
|
|
|
if gate.get("status") != "PASS-AUDIT" or not gate.get("stage2_authorized"):
|
|
|
|
|
raise RuntimeError(f"stage-1 audit did not authorize conditioning: {gate}")
|
|
|
|
|
report_path = Path(gate["audit_report"])
|
|
|
|
|
if sha256(report_path) != gate["audit_report_sha256"]:
|
|
|
|
|
raise RuntimeError("stage-1 audit report changed after gate creation")
|
|
|
|
|
stage1 = json.loads(report_path.read_text(encoding="utf-8"))
|
|
|
|
|
reliable = stage1["reliable_mb_compartments"]["verbatim_values"]
|
|
|
|
|
|
|
|
|
|
frozen_before = run_frozen_audit(cfg)
|
|
|
|
|
write_json(exp / "audit/frozen-checksums-before-stage2.json", frozen_before)
|
|
|
|
|
shutil.copy2(args.config, exp / "audit/phase4h-stage2-config.yaml")
|
|
|
|
|
shutil.copy2(Path(__file__).resolve(), exp / "audit/phase4h.py")
|
|
|
|
|
experiment = Experiment(cfg, exp)
|
|
|
|
|
write_json(exp / "audit/conditioning-interface.json", {
|
|
|
|
|
"CS_A_neuron_idx": experiment.cs_a.tolist(), "CS_B_neuron_idx": experiment.cs_b.tolist(),
|
|
|
|
|
"CS_populations_are_disjoint": True, "CS_identity": "engineering split of sorted olfactory annotation_class population",
|
|
|
|
|
"DAN_targets": {key: value.tolist() for key, value in experiment.dan_targets.items()},
|
|
|
|
|
"reward_interface": cfg["conditioning"]["reward_interface_label"],
|
|
|
|
|
"reward_value_is_dopamine_concentration": False,
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
all_runs = []
|
|
|
|
|
finite_all = True; no_runaway = True
|
|
|
|
|
for protocol in cfg["conditioning"]["protocols"]:
|
|
|
|
|
protocol_dir = exp / ("controls" if "control" in protocol else "conditioning") / protocol
|
|
|
|
|
protocol_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
for train_seed, eval_seed in zip(cfg["conditioning"]["train_seeds"], cfg["conditioning"]["evaluation_seeds"]):
|
|
|
|
|
rng = np.random.default_rng(int(train_seed))
|
|
|
|
|
delta = np.zeros(len(experiment.edge_ids), dtype=np.float32)
|
|
|
|
|
pre = eval_pair(experiment, int(eval_seed), delta)
|
|
|
|
|
plan = protocol_plan(protocol, cfg, reliable, rng)
|
|
|
|
|
updates = []; trajectory = []
|
|
|
|
|
acquisition_snapshot = None
|
|
|
|
|
for trial_idx, item in enumerate(plan):
|
|
|
|
|
delta, update = experiment.update(item["CS"], int(train_seed) + trial_idx, delta,
|
|
|
|
|
item["compartment"], item["dan_label"], item["delay_ms"],
|
|
|
|
|
item["enabled"], item["d_value"])
|
|
|
|
|
update["trial"] = trial_idx; update["phase"] = item.get("phase", "acquisition")
|
|
|
|
|
updates.append(update)
|
|
|
|
|
pair = eval_pair(experiment, int(eval_seed), delta)
|
|
|
|
|
trajectory.append({"trial": trial_idx, "phase": update["phase"],
|
|
|
|
|
"mbon_separation": pair["mbon_separation"],
|
|
|
|
|
"downstream_separation": pair["downstream_separation"],
|
|
|
|
|
"balance_difference_A_minus_B": pair["balance_difference_A_minus_B"]})
|
|
|
|
|
if protocol == "reversal-learning" and update["phase"] == "acquisition":
|
|
|
|
|
acquisition_snapshot = pair
|
|
|
|
|
post = eval_pair(experiment, int(eval_seed), delta)
|
|
|
|
|
retention = []
|
|
|
|
|
if protocol == "retention":
|
|
|
|
|
for nominal_ms, extra_seed in zip((0, 1000, 10000), cfg["conditioning"]["evaluation_seeds"][:3]):
|
|
|
|
|
pair = eval_pair(experiment, int(extra_seed), delta)
|
|
|
|
|
retention.append({"nominal_retention_interval_ms": nominal_ms,
|
|
|
|
|
"state_policy": "fresh neural state; persistent plastic_delta; no idle dynamics approximation",
|
|
|
|
|
"mbon_separation": pair["mbon_separation"], "downstream_separation": pair["downstream_separation"]})
|
|
|
|
|
run = {
|
|
|
|
|
"protocol": protocol, "train_seed": int(train_seed), "evaluation_seed": int(eval_seed),
|
|
|
|
|
"pre": pre, "post": post, "updates": updates, "trajectory": trajectory,
|
|
|
|
|
"retention": retention,
|
|
|
|
|
"mbon_euclidean_gain": post["mbon_separation"]["euclidean"] - pre["mbon_separation"]["euclidean"],
|
|
|
|
|
"downstream_euclidean_gain": post["downstream_separation"]["euclidean"] - pre["downstream_separation"]["euclidean"],
|
|
|
|
|
"reversal_balance_sign_changed": (
|
|
|
|
|
False if acquisition_snapshot is None else
|
|
|
|
|
bool(np.sign(acquisition_snapshot["balance_difference_A_minus_B"]) != np.sign(post["balance_difference_A_minus_B"]))
|
|
|
|
|
),
|
|
|
|
|
"changed_edges": int(np.count_nonzero(delta)), "delta_l1_effective_mv": float(np.abs(delta).sum()),
|
|
|
|
|
"base_CSR_modified": False,
|
|
|
|
|
}
|
|
|
|
|
run_path = protocol_dir / f"train-{train_seed}-eval-{eval_seed}.json"
|
|
|
|
|
write_json(run_path, run)
|
|
|
|
|
np.savez_compressed(learned / "plastic-delta" / f"{protocol}-seed-{train_seed}.npz",
|
|
|
|
|
edge_id=experiment.edge_ids, plastic_delta_effective_mv=delta)
|
|
|
|
|
finite_all = finite_all and all(x["finite"] for x in (post["CS-A"], post["CS-B"]))
|
|
|
|
|
no_runaway = no_runaway and not any(x["runaway"] for x in (post["CS-A"], post["CS-B"]))
|
|
|
|
|
all_runs.append(run)
|
|
|
|
|
|
|
|
|
|
del experiment; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
|
regression = phase4c_regression(cfg, exp)
|
|
|
|
|
frozen_after = run_frozen_audit(cfg)
|
|
|
|
|
write_json(exp / "audit/frozen-checksums-after-stage2.json", frozen_after)
|
|
|
|
|
|
|
|
|
|
by_protocol = {}
|
|
|
|
|
for protocol in cfg["conditioning"]["protocols"]:
|
|
|
|
|
rows = [row for row in all_runs if row["protocol"] == protocol]
|
|
|
|
|
by_protocol[protocol] = {
|
|
|
|
|
"seeds": len(rows),
|
|
|
|
|
"mbon_gain_mean": float(np.mean([row["mbon_euclidean_gain"] for row in rows])),
|
|
|
|
|
"mbon_gain_positive_seeds": int(sum(row["mbon_euclidean_gain"] > 0 for row in rows)),
|
|
|
|
|
"downstream_gain_mean": float(np.mean([row["downstream_euclidean_gain"] for row in rows])),
|
|
|
|
|
"downstream_gain_positive_seeds": int(sum(row["downstream_euclidean_gain"] > 0 for row in rows)),
|
|
|
|
|
"changed_edges_mean": float(np.mean([row["changed_edges"] for row in rows])),
|
|
|
|
|
"delta_l1_effective_mv_mean": float(np.mean([row["delta_l1_effective_mv"] for row in rows])),
|
|
|
|
|
"reversal_signature_seeds": int(sum(row["reversal_balance_sign_changed"] for row in rows)),
|
|
|
|
|
}
|
|
|
|
|
majority = int(cfg["success_gate"]["require_minimum_seed_majority"])
|
|
|
|
|
app = by_protocol["appetitive-conditioning"]; av = by_protocol["aversive-conditioning"]
|
|
|
|
|
real_mbon = app["mbon_gain_positive_seeds"] >= majority and av["mbon_gain_positive_seeds"] >= majority
|
|
|
|
|
real_downstream = app["downstream_gain_positive_seeds"] >= majority and av["downstream_gain_positive_seeds"] >= majority
|
|
|
|
|
reversal = by_protocol["reversal-learning"]["reversal_signature_seeds"] >= majority
|
|
|
|
|
retention_rows = [row for row in all_runs if row["protocol"] == "retention"]
|
|
|
|
|
retention = sum(bool(row["retention"]) and row["retention"][-1]["mbon_separation"]["euclidean"] > row["pre"]["mbon_separation"]["euclidean"] for row in retention_rows) >= majority
|
|
|
|
|
control_ceiling = max(by_protocol[name]["mbon_gain_mean"] for name in ("d-zero-control", "random-reward-control", "frozen-brain-control"))
|
|
|
|
|
condition_floor = min(app["mbon_gain_mean"], av["mbon_gain_mean"])
|
|
|
|
|
controls_distinct = condition_floor > control_ceiling
|
|
|
|
|
passed = all([frozen_before["all_passed"], frozen_after["all_passed"], regression["all_passed"],
|
|
|
|
|
finite_all, no_runaway, real_mbon, real_downstream, reversal, retention, controls_distinct])
|
|
|
|
|
status = "MaleCNS-DopaminePlasticity-v1-candidate" if passed else "FAIL-NO-DOPAMINE-PLASTICITY-CANDIDATE"
|
|
|
|
|
summary = {
|
|
|
|
|
"status": status, "candidate_generated": passed, "stage1_status": stage1["status"],
|
|
|
|
|
"mapping": {
|
|
|
|
|
"KC": stage1["population_counts"]["KC"], "MBON": stage1["population_counts"]["MBON"],
|
|
|
|
|
"PAM": stage1["population_counts"]["PAM"], "PPL1": stage1["population_counts"]["PPL1"],
|
|
|
|
|
"other_DAN": stage1["population_counts"]["other_DAN"],
|
|
|
|
|
"reliable_compartments": reliable, "plastic_candidate_edges": stage1["plastic_edge_candidate_mask"]["edges"],
|
|
|
|
|
},
|
|
|
|
|
"protocols": by_protocol,
|
|
|
|
|
"gate": {
|
|
|
|
|
"all_frozen_checksums_pass": frozen_before["all_passed"] and frozen_after["all_passed"],
|
|
|
|
|
"phase4c_reference_regression_pass": regression["all_passed"],
|
|
|
|
|
"finite": finite_all, "no_runaway": no_runaway,
|
|
|
|
|
"appetitive_and_aversive_MBON_separation_majority": real_mbon,
|
|
|
|
|
"appetitive_and_aversive_downstream_separation_majority": real_downstream,
|
|
|
|
|
"reversal_signature_majority": reversal, "retention_signature_majority": retention,
|
|
|
|
|
"controls_distinct": controls_distinct,
|
|
|
|
|
"weight_change_alone_used_as_success": False,
|
|
|
|
|
},
|
|
|
|
|
"base_brain": {"modified_on_disk": False, "runtime_base_restored_between_runs": True,
|
|
|
|
|
"equation": cfg["plasticity"]["overlay_equation"]},
|
|
|
|
|
"reward_interface": {"artificial_unconditioned_signal": True, "dopamine_concentration_claimed": False},
|
|
|
|
|
"limitations": [
|
|
|
|
|
"CS-A and CS-B are reproducible engineering subsets of the annotated olfactory population, not identified odor receptor channels.",
|
|
|
|
|
"Compartment membership is based on exact verbatim annotation_instance tokens and direct neuron-level CSR edges; synapse-level ROI is unavailable.",
|
|
|
|
|
"MBON balance is transmitter-stratified and is not relabeled as approach/avoidance without a reliable functional annotation.",
|
|
|
|
|
"Retention uses a fresh neural state with persistent plastic_delta and does not simulate a long idle biological interval.",
|
|
|
|
|
],
|
|
|
|
|
}
|
|
|
|
|
write_json(exp / "audit/stage2-summary.json", summary)
|
|
|
|
|
write_json(model / "VERSION.json", {
|
|
|
|
|
"status": status, "candidate_name": status if passed else None,
|
|
|
|
|
"base_model": cfg["names"]["reference"], "runtime": cfg["names"]["runtime"],
|
|
|
|
|
"plasticity_overlay_only": True, "stage1_gate_sha256": sha256(gate_path),
|
|
|
|
|
"summary_sha256": sha256(exp / "audit/stage2-summary.json"),
|
|
|
|
|
})
|
|
|
|
|
shutil.copy2(args.config, model / "phase4h.yaml")
|
|
|
|
|
shutil.copy2(Path(__file__).resolve(), model / "phase4h.py")
|
|
|
|
|
write_json(learned / "training-metadata/summary.json", {"status": status, "protocols": by_protocol, "candidate": passed})
|
|
|
|
|
|
|
|
|
|
print(json.dumps(summary, indent=2, sort_keys=True))
|
|
|
|
|
sys.stdout.flush()
|
|
|
|
|
# Flush redirected stdout before hashing so a log inside the experiment tree
|
|
|
|
|
# cannot change after its checksum has been recorded.
|
|
|
|
|
seal(model, model / "MANIFEST.json", model / "checksums.sha256")
|
|
|
|
|
seal(learned, learned / "MANIFEST.json", learned / "checksums.sha256")
|
|
|
|
|
seal(exp, exp / "manifests/manifest.json", exp / "manifests/checksums.sha256")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
|
main()
|
Digital Drosophila Archive