Digital Drosophila Archive
757 lines
43 KiB
Python
757 lines
43 KiB
Python
#!/usr/bin/env python3
|
|
"""Phase 4C: freeze reference, audit visual semantics, and optimize exact runtime."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
import torch
|
|
import yaml
|
|
from torch.utils.cpp_extension import load as load_extension
|
|
|
|
import phase4b as p4b
|
|
|
|
LOG = logging.getLogger("phase4c")
|
|
|
|
|
|
def sha256(path: Path, chunk: int = 8 << 20) -> str:
|
|
h = hashlib.sha256()
|
|
with path.open("rb") as f:
|
|
while data := f.read(chunk):
|
|
h.update(data)
|
|
return h.hexdigest()
|
|
|
|
|
|
def write_json(path: Path, value: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
|
|
|
|
|
def copy_immutable(source: Path, destination: Path) -> None:
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
if destination.exists():
|
|
if sha256(source) != sha256(destination):
|
|
raise RuntimeError(f"immutable destination differs: {destination}")
|
|
return
|
|
shutil.copy2(source, destination)
|
|
destination.chmod(0o444)
|
|
|
|
|
|
def exact_paths(cfg: dict[str, Any]) -> dict[str, Path]:
|
|
root = Path(cfg["paths"]["connectome_root"]).resolve()
|
|
allowed = Path("/data/digital-drosophila/datasets/male-cns-v1.0/processed/neuron-only").resolve()
|
|
if root != allowed:
|
|
raise RuntimeError(f"only allowed input is {allowed}")
|
|
paths = {
|
|
"indptr": root / "gpu/csr-indptr.int64.npy",
|
|
"indices": root / "gpu/csr-indices.int32.npy",
|
|
"raw_weights": root / "gpu/csr-weights.int64.npy",
|
|
"metadata": root / "neuron-index/unified-metadata.parquet",
|
|
"mapping": root / "neuron-index/body-id-neuron-idx.parquet",
|
|
"csr_metadata": root / "gpu/csr-metadata.json",
|
|
}
|
|
for path in paths.values():
|
|
if not path.is_file():
|
|
raise FileNotFoundError(path)
|
|
return paths
|
|
|
|
|
|
def freeze_reference(cfg: dict[str, Any], config_path: Path, paths: dict[str, Path]) -> dict[str, Any]:
|
|
model_root = Path(cfg["paths"]["model_root"])
|
|
ref = model_root / "reference-v1"
|
|
ref.mkdir(parents=True, exist_ok=True)
|
|
source_dir = Path(__file__).resolve().parent
|
|
phase4b_root = Path(cfg["paths"]["phase4b_root"])
|
|
source_files = {
|
|
"model.py": source_dir / "phase4b.py",
|
|
"config.yaml": source_dir / "phase4b.yaml",
|
|
}
|
|
for name, source in source_files.items():
|
|
copy_immutable(source, ref / name)
|
|
phase4b_seal = phase4b_root / "manifests/checksums.sha256"
|
|
phase4b_manifest = phase4b_root / "manifests/manifest.json"
|
|
if not phase4b_seal.is_file() or not phase4b_manifest.is_file():
|
|
raise RuntimeError("Phase 4B manifest/seal missing")
|
|
seal_check = subprocess.run(
|
|
["sha256sum", "-c", "--quiet", str(phase4b_seal)], cwd=phase4b_root,
|
|
text=True, capture_output=True,
|
|
)
|
|
if seal_check.returncode:
|
|
raise RuntimeError(f"Phase 4B seal failed: {seal_check.stderr}")
|
|
frozen = {
|
|
"name": cfg["names"]["reference"],
|
|
"description": cfg["names"]["description"],
|
|
"not_claimed": "strict Shiu 2024 original-model reproduction",
|
|
"reason": "MaleCNS v1.0 transmitter inventory includes histamine and differs from the original dataset/inventory",
|
|
"parameters": cfg["model"],
|
|
"sign_policy": cfg["sign_policy"],
|
|
"graph": cfg["expected_graph"],
|
|
"source_phase": "Phase 4B PASS",
|
|
"phase4b_manifest_sha256": sha256(phase4b_manifest),
|
|
"phase4b_checksums_sha256": sha256(phase4b_seal),
|
|
"phase4b_checksum_verification": "PASS",
|
|
"local_repository_commit": cfg["reproducibility"]["local_repository_commit"],
|
|
"remote_model_directory_is_git_repository": False,
|
|
"code": {name: {"path": str(ref / name), "sha256": sha256(ref / name)} for name in source_files},
|
|
"inputs": {name: {"path": str(path), "sha256": sha256(path)} for name, path in paths.items()},
|
|
}
|
|
version_path = ref / "VERSION.json"
|
|
if version_path.exists():
|
|
previous = json.loads(version_path.read_text(encoding="utf-8"))
|
|
if previous != frozen:
|
|
raise RuntimeError("existing immutable reference VERSION differs")
|
|
else:
|
|
write_json(version_path, frozen)
|
|
version_path.chmod(0o444)
|
|
return frozen
|
|
|
|
|
|
def full_metadata(path: Path) -> pa.Table:
|
|
table = pq.read_table(path)
|
|
required = [
|
|
"neuron_idx", "body_id", "annotation_superclass", "annotation_class",
|
|
"annotation_subclass", "annotation_type", "annotation_instance", "side",
|
|
"neurotransmitter",
|
|
]
|
|
missing = [x for x in required if x not in table.column_names]
|
|
if missing:
|
|
raise RuntimeError(f"metadata columns missing: {missing}")
|
|
return table
|
|
|
|
|
|
def list_mask(values: list[Any], accepted: set[str]) -> np.ndarray:
|
|
return np.asarray([x in accepted for x in values], dtype=bool)
|
|
|
|
|
|
def populations(table: pa.Table) -> dict[str, np.ndarray]:
|
|
cls = table["annotation_class"].to_pylist()
|
|
sub = table["annotation_subclass"].to_pylist()
|
|
mech = {"mechanosensory", "mechanosensory_tactile", "mechanosensory_proprioceptive"}
|
|
return {
|
|
"visual": np.asarray([x == "visual" for x in cls]),
|
|
"mechanosensory": np.asarray([(x in mech) and y != "auditory" for x, y in zip(cls, sub)]),
|
|
"olfactory": np.asarray([x == "olfactory" for x in cls]),
|
|
"gustatory": np.asarray([x == "gustatory" for x in cls]),
|
|
"auditory": np.asarray([x == "auditory" for x in sub]),
|
|
}
|
|
|
|
|
|
def layer_edge_stats(indptr: np.ndarray, indices: np.ndarray, raw: np.ndarray,
|
|
layers: dict[str, np.ndarray]) -> dict[str, Any]:
|
|
n = len(indptr) - 1
|
|
label = np.full(n, -1, dtype=np.int16)
|
|
names = list(layers)
|
|
for j, name in enumerate(names):
|
|
label[layers[name]] = j
|
|
counts = np.zeros((len(names), len(names)), dtype=np.int64)
|
|
weights = np.zeros_like(counts)
|
|
for src in range(n):
|
|
a, b = int(indptr[src]), int(indptr[src + 1])
|
|
x = label[src]
|
|
if x < 0 or a == b:
|
|
continue
|
|
y = label[indices[a:b]]
|
|
valid = y >= 0
|
|
if np.any(valid):
|
|
counts[x] += np.bincount(y[valid], minlength=len(names))
|
|
weights[x] += np.bincount(y[valid], weights=raw[a:b][valid], minlength=len(names)).astype(np.int64)
|
|
return {
|
|
"layer_order": names,
|
|
"edge_counts": counts.tolist(),
|
|
"synapse_weight_sums": weights.tolist(),
|
|
}
|
|
|
|
|
|
def visual_audit(exp: Path, table: pa.Table, indptr: np.ndarray,
|
|
indices: np.ndarray, raw: np.ndarray) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
|
|
out = exp / "visual-audit"
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
pops = populations(table)
|
|
visual = pops["visual"]
|
|
nodes = np.flatnonzero(visual).astype(np.int32)
|
|
fields = [
|
|
"neuron_idx", "body_id", "annotation_superclass", "annotation_class",
|
|
"annotation_subclass", "annotation_type", "annotation_instance", "side", "neurotransmitter",
|
|
]
|
|
visual_table = table.select(fields).filter(pa.array(visual))
|
|
pq.write_table(visual_table, out / "phase4b-visual-6091.parquet", compression="zstd")
|
|
arrays = {name: table[name].to_pylist() for name in fields}
|
|
breakdown: dict[str, Any] = {}
|
|
for field in fields[2:]:
|
|
c = Counter("<null>" if arrays[field][i] is None else str(arrays[field][i]) for i in nodes)
|
|
breakdown[field] = dict(sorted(c.items(), key=lambda x: (-x[1], x[0])))
|
|
visual_types = sorted({str(arrays["annotation_type"][i]) for i in nodes if arrays["annotation_type"][i] is not None})
|
|
exact_photoreceptor = visual & list_mask(arrays["annotation_superclass"], {"ol_sensory"})
|
|
|
|
# These anatomical memberships are derived only from explicit metadata fields.
|
|
super_values = arrays["annotation_superclass"]
|
|
type_values = ["" if x is None else str(x) for x in arrays["annotation_type"]]
|
|
ol_intrinsic = list_mask(super_values, {"ol_intrinsic"})
|
|
visual_projection = list_mask(super_values, {"visual_projection", "visual_projection_tbc"})
|
|
lamina_types = {"L1", "L2", "L3", "L4", "L5", "C2", "C3", "T1", "Lawf1", "Lawf2"}
|
|
lamina = ol_intrinsic & np.asarray([x in lamina_types for x in type_values])
|
|
lobula_plate = ol_intrinsic & np.asarray([
|
|
x.startswith(("T4", "T5", "LPi", "LPLC", "LLPC")) for x in type_values
|
|
])
|
|
lobula = ol_intrinsic & ~lobula_plate & np.asarray([
|
|
x.startswith(("LC", "LT", "Li", "Lobula")) for x in type_values
|
|
])
|
|
medulla = ol_intrinsic & ~lamina & ~lobula & ~lobula_plate & np.asarray([
|
|
x.startswith(("Mi", "Dm", "Tm", "TmY", "Pm", "Cm", "Me")) for x in type_values
|
|
])
|
|
optic_other = ol_intrinsic & ~(lamina | medulla | lobula | lobula_plate)
|
|
anatomy = {
|
|
"photoreceptor": exact_photoreceptor,
|
|
"retinal_sensory": exact_photoreceptor,
|
|
"lamina": lamina,
|
|
"medulla": medulla,
|
|
"lobula": lobula,
|
|
"lobula_plate": lobula_plate,
|
|
"optic_lobe_interneurons": ol_intrinsic,
|
|
"visual_projection_neurons": visual_projection,
|
|
"optic_lobe_interneurons_unassigned_to_named_region": optic_other,
|
|
}
|
|
membership_rows: list[dict[str, Any]] = []
|
|
for role, mask in anatomy.items():
|
|
for i in np.flatnonzero(mask):
|
|
membership_rows.append({"anatomical_role": role, "neuron_idx": int(i), "body_id": int(arrays["body_id"][i])})
|
|
pq.write_table(pa.Table.from_pylist(membership_rows), out / "anatomical-memberships.parquet", compression="zstd")
|
|
|
|
# Exact graph hierarchy uses shortest directed distance, not inferred body IDs.
|
|
distance = p4b.directed_distances(indptr, indices, nodes.astype(np.int64), len(visual))
|
|
optic_super = {"ol_intrinsic", "visual_projection", "visual_projection_tbc", "visual_centrifugal"}
|
|
optic = list_mask(super_values, optic_super)
|
|
layers = {
|
|
"visual_sensory": distance == 0,
|
|
"first_order": distance == 1,
|
|
"second_order": distance == 2,
|
|
"higher_visual": (distance >= 3) & optic,
|
|
"central_brain": (distance >= 3) & ~optic,
|
|
}
|
|
for name, mask in layers.items():
|
|
ids = np.flatnonzero(mask).astype(np.int32)
|
|
pq.write_table(pa.table({
|
|
"neuron_idx": pa.array(ids, type=pa.int32()),
|
|
"body_id": pa.array(np.asarray(arrays["body_id"], dtype=np.uint64)[ids], type=pa.uint64()),
|
|
"shortest_visual_path_depth": pa.array(distance[ids], type=pa.int16()),
|
|
}), out / f"hierarchy-{name}.parquet", compression="zstd")
|
|
edge_stats = layer_edge_stats(indptr, indices, raw, layers)
|
|
write_json(out / "hierarchy-edge-statistics.json", edge_stats)
|
|
summary = {
|
|
"phase4b_visual_count": int(visual.sum()),
|
|
"all_histaminergic": all(arrays["neurotransmitter"][i] == "histamine" for i in nodes),
|
|
"all_ol_sensory": all(arrays["annotation_superclass"][i] == "ol_sensory" for i in nodes),
|
|
"raw_visual_types": visual_types,
|
|
"conclusion": "The 6091 annotation_class=visual bodies are photoreceptor/retinal sensory input neurons, not the whole visual processing network.",
|
|
"field_breakdown": breakdown,
|
|
"anatomical_membership_counts": {k: int(v.sum()) for k, v in anatomy.items()},
|
|
"anatomical_rule_provenance": {
|
|
"photoreceptor_retinal_sensory": "annotation_class=visual AND annotation_superclass=ol_sensory; observed types recorded verbatim",
|
|
"optic_lobe_interneurons": "annotation_superclass=ol_intrinsic",
|
|
"visual_projection_neurons": "annotation_superclass in {visual_projection, visual_projection_tbc}",
|
|
"named_optic_regions": "explicit annotation_type prefix/set rules, preserved in this audit; no body_id was guessed",
|
|
},
|
|
"hierarchy_definition": "exclusive shortest directed connectome distance from the 6091 sensory neurons: depth 0 sensory, depth 1 first-order, depth 2 second-order, depth >=3 split into optic/visual-superclass higher_visual versus non-optic downstream central_brain; unreachable bodies are excluded",
|
|
"hierarchy_counts": {k: int(v.sum()) for k, v in layers.items()},
|
|
"edge_statistics_file": str(out / "hierarchy-edge-statistics.json"),
|
|
}
|
|
write_json(out / "summary.json", summary)
|
|
return summary, layers
|
|
|
|
|
|
def generate_forced(steps: int, target_count: int, rate_hz: float, dt: float,
|
|
seed: int, device: torch.device) -> torch.Tensor:
|
|
probability = rate_hz * dt / 1000.0
|
|
generator = torch.Generator(device=device)
|
|
generator.manual_seed(seed)
|
|
forced = torch.empty((steps, target_count), dtype=torch.uint8, device=device)
|
|
# Deliberately preserve Phase 4B's one rand() call per timestep.
|
|
for step in range(steps):
|
|
forced[step].copy_((torch.rand(target_count, generator=generator, device=device) < probability).to(torch.uint8))
|
|
return forced
|
|
|
|
|
|
def expand_rows(indptr: torch.Tensor, sources: torch.Tensor) -> torch.Tensor:
|
|
if sources.numel() == 0:
|
|
return torch.empty(0, dtype=torch.int64, device=indptr.device)
|
|
starts = indptr[sources]
|
|
counts = indptr[sources + 1] - starts
|
|
cumulative = torch.cumsum(counts, 0)
|
|
total = int(cumulative[-1].item()) if cumulative.numel() else 0
|
|
if total == 0:
|
|
return torch.empty(0, dtype=torch.int64, device=indptr.device)
|
|
seq = torch.arange(total, dtype=torch.int64, device=indptr.device)
|
|
row = torch.searchsorted(cumulative, seq, right=True)
|
|
return starts[row] + seq - (cumulative - counts)[row]
|
|
|
|
|
|
class Runtime:
|
|
def __init__(self, cfg: dict[str, Any], table: pa.Table, indptr_np: np.ndarray,
|
|
indices_np: np.ndarray, raw_np: np.ndarray, device: torch.device,
|
|
extension: Any | None):
|
|
self.cfg, self.table, self.device, self.extension = cfg, table, device, extension
|
|
self.n = int(cfg["expected_graph"]["neurons"])
|
|
self.indptr = torch.as_tensor(np.asarray(indptr_np), dtype=torch.int64, device=device)
|
|
self.indices = torch.as_tensor(np.asarray(indices_np), dtype=torch.int32, device=device)
|
|
nt = table["neurotransmitter"].to_pylist()
|
|
policy = cfg["sign_policy"]
|
|
signs = np.empty(self.n, dtype=np.int8)
|
|
for i, value in enumerate(nt):
|
|
key = "" if value is None else str(value).strip().lower()
|
|
if not key:
|
|
signs[i] = int(policy["unknown_sign"])
|
|
elif key == "unclear":
|
|
signs[i] = int(policy["unclear_sign"])
|
|
elif key in policy["neurotransmitter_sign"]:
|
|
signs[i] = int(policy["neurotransmitter_sign"][key])
|
|
else:
|
|
raise RuntimeError(f"unmapped neurotransmitter: {value!r}")
|
|
effective = np.empty(len(raw_np), dtype=np.float32)
|
|
for src in range(self.n):
|
|
a, b = int(indptr_np[src]), int(indptr_np[src + 1])
|
|
effective[a:b] = raw_np[a:b].astype(np.float32) * signs[src] * float(cfg["model"]["w_syn_mv_per_anatomical_synapse"])
|
|
self.weights = torch.as_tensor(effective, dtype=torch.float32, device=device)
|
|
self.signs = signs
|
|
self.superclass = np.asarray(table["annotation_superclass"].to_pylist(), dtype=object)
|
|
|
|
def compiled_state(self, mode: str):
|
|
m = self.cfg["model"]
|
|
gd = 1.0 - float(m["dt_ms"]) / float(m["synaptic_tau_ms"])
|
|
alpha = float(m["dt_ms"]) / float(m["membrane_tau_ms"])
|
|
rest, reset, threshold = (float(m[x]) for x in ("resting_potential_mv", "reset_potential_mv", "threshold_mv"))
|
|
|
|
def fn(v, g, refractory, arrival):
|
|
g2 = g * gd
|
|
g2 = g2 + arrival
|
|
r2 = torch.clamp_min(refractory, 1) - 1
|
|
available = r2 == 0
|
|
leak = v - rest
|
|
drive = g2 - leak
|
|
integrated = v + alpha * drive
|
|
v2 = torch.where(available, integrated, torch.full_like(v, reset))
|
|
natural = available & (v2 >= threshold)
|
|
return v2, g2, r2, natural, available
|
|
|
|
kwargs: dict[str, Any] = {"backend": "inductor", "fullgraph": True}
|
|
if mode == "compile-reduce-overhead":
|
|
kwargs["mode"] = "reduce-overhead"
|
|
return torch.compile(fn, **kwargs), fn
|
|
|
|
def run(self, backend: str, targets_np: np.ndarray, rate_hz: float, seed: int,
|
|
duration_ms: float, capture_events: bool, monitor: bool = False,
|
|
forced_override: torch.Tensor | None = None) -> dict[str, Any]:
|
|
m = self.cfg["model"]
|
|
dt = float(m["dt_ms"]); steps = int(round(duration_ms / dt))
|
|
delay_steps = int(round(float(m["synaptic_delay_ms"]) / dt)); qlen = delay_steps + 1
|
|
ref_steps = int(round(float(m["refractory_ms"]) / dt))
|
|
alpha = dt / float(m["membrane_tau_ms"]); gd = 1.0 - dt / float(m["synaptic_tau_ms"])
|
|
rest = float(m["resting_potential_mv"]); reset = float(m["reset_potential_mv"]); threshold = float(m["threshold_mv"])
|
|
targets = torch.as_tensor(targets_np, dtype=torch.int64, device=self.device)
|
|
forced = forced_override if forced_override is not None else generate_forced(steps, len(targets_np), rate_hz, dt, seed, self.device)
|
|
if forced.shape != (steps, len(targets_np)):
|
|
raise RuntimeError("forced input shape mismatch")
|
|
v = torch.full((self.n,), rest, dtype=torch.float32, device=self.device)
|
|
g = torch.zeros(self.n, dtype=torch.float32, device=self.device)
|
|
refractory = torch.zeros(self.n, dtype=torch.int32, device=self.device)
|
|
queue = torch.zeros((qlen, self.n), dtype=torch.float32, device=self.device)
|
|
counts = torch.zeros(self.n, dtype=torch.int32, device=self.device)
|
|
event_steps: list[np.ndarray] = []; event_nodes: list[np.ndarray] = []
|
|
compiled = eager_fn = None
|
|
if backend.startswith("compile-"):
|
|
compiled, eager_fn = self.compiled_state(backend)
|
|
# Exclude compile latency and record a graph-break diagnostic separately.
|
|
tv, tg, tr, ta = v.clone(), g.clone(), refractory.clone(), queue[0].clone()
|
|
compiled(tv, tg, tr, ta); torch.cuda.synchronize()
|
|
del tv, tg, tr, ta
|
|
if backend.startswith("custom") or backend in ("cudagraph-custom-state", "hybrid-exact"):
|
|
if self.extension is None:
|
|
raise RuntimeError("custom extension unavailable")
|
|
spike_u8 = torch.zeros(self.n, dtype=torch.uint8, device=self.device)
|
|
graph = staging = None
|
|
if backend == "cudagraph-custom-state":
|
|
staging = torch.zeros(self.n, dtype=torch.float32, device=self.device)
|
|
graph = torch.cuda.CUDAGraph()
|
|
# Warm-up on disposable state, then capture only the graph-safe fixed-shape state kernel.
|
|
sv, sg = v.clone(), g.clone(); sr = refractory.clone(); ss = spike_u8.clone(); sc = counts.clone(); sa = staging.clone()
|
|
self.extension.state_step(sv, sg, sr, sa, ss, sc, gd, alpha, rest, reset, threshold, ref_steps)
|
|
torch.cuda.synchronize()
|
|
del sv, sg, sr, ss, sc, sa
|
|
with torch.cuda.graph(graph):
|
|
self.extension.state_step(v, g, refractory, staging, spike_u8, counts, gd, alpha, rest, reset, threshold, ref_steps)
|
|
sampler = p4b.GPUSampler(self.cfg["device"]["launcher_visible_gpu_uuid"], float(self.cfg["benchmark"]["gpu_sample_interval_seconds"])) if monitor else None
|
|
if sampler: sampler.start()
|
|
torch.cuda.synchronize(); torch.cuda.reset_peak_memory_stats(); start = time.perf_counter()
|
|
cpu_start = time.process_time()
|
|
for step in range(steps):
|
|
slot = step % qlen
|
|
if backend in ("eager", "hybrid-exact"):
|
|
g.mul_(gd).add_(queue[slot]); queue[slot].zero_()
|
|
refractory.clamp_min_(1).sub_(1)
|
|
available = refractory == 0
|
|
v[available] += alpha * (g[available] - (v[available] - rest))
|
|
v[~available] = reset
|
|
spikes = available & (v >= threshold)
|
|
accepted = forced[step].bool() & available[targets]
|
|
forced_sources = targets[accepted]
|
|
if forced_sources.numel(): spikes[forced_sources] = True
|
|
sources = torch.nonzero(spikes, as_tuple=False).flatten()
|
|
counts.add_(spikes.to(torch.int32))
|
|
if sources.numel():
|
|
v[sources] = reset; refractory[sources] = ref_steps
|
|
if backend == "hybrid-exact":
|
|
starts = self.indptr[sources]
|
|
degrees = self.indptr[sources + 1] - starts
|
|
cumulative = torch.cumsum(degrees, 0)
|
|
total = int(cumulative[-1].item()) if cumulative.numel() else 0
|
|
edges = torch.empty(total, dtype=torch.int64, device=self.device)
|
|
if total:
|
|
self.extension.fill_edges(self.indptr, sources, cumulative - degrees, edges)
|
|
else:
|
|
edges = expand_rows(self.indptr, sources)
|
|
if edges.numel():
|
|
queue[(step + delay_steps) % qlen].scatter_add_(0, self.indices[edges].long(), self.weights[edges])
|
|
elif backend.startswith("compile-"):
|
|
assert compiled is not None
|
|
arrival = queue[slot].clone(); queue[slot].zero_()
|
|
v, g, refractory, spikes, available = compiled(v, g, refractory, arrival)
|
|
accepted = forced[step].bool() & available[targets]
|
|
forced_sources = targets[accepted]
|
|
if forced_sources.numel(): spikes[forced_sources] = True
|
|
sources = torch.nonzero(spikes, as_tuple=False).flatten()
|
|
counts.add_(spikes.to(torch.int32))
|
|
if sources.numel():
|
|
v[sources] = reset; refractory[sources] = ref_steps
|
|
edges = expand_rows(self.indptr, sources)
|
|
if edges.numel(): queue[(step + delay_steps) % qlen].scatter_add_(0, self.indices[edges].long(), self.weights[edges])
|
|
else:
|
|
assert self.extension is not None
|
|
if backend == "cudagraph-custom-state":
|
|
assert staging is not None and graph is not None
|
|
staging.copy_(queue[slot]); queue[slot].zero_(); graph.replay()
|
|
else:
|
|
self.extension.state_step(v, g, refractory, queue[slot], spike_u8, counts, gd, alpha, rest, reset, threshold, ref_steps)
|
|
self.extension.force_step(v, refractory, spike_u8, counts, targets, forced[step], reset, ref_steps)
|
|
self.extension.fanout(spike_u8, self.indptr, self.indices, self.weights, queue[(step + delay_steps) % qlen])
|
|
sources = torch.nonzero(spike_u8, as_tuple=False).flatten() if capture_events else None
|
|
if capture_events:
|
|
if not (backend.startswith("custom") or backend == "cudagraph-custom-state"):
|
|
pass
|
|
assert sources is not None
|
|
src = sources.detach().cpu().numpy().astype(np.int32, copy=True)
|
|
event_nodes.append(src); event_steps.append(np.full(len(src), step, dtype=np.int32))
|
|
torch.cuda.synchronize(); wall = time.perf_counter() - start; cpu = time.process_time() - cpu_start
|
|
resources = sampler.stop() if sampler else {}
|
|
v_np = v.cpu().numpy().copy(); g_np = g.cpu().numpy().copy(); counts_np = counts.cpu().numpy().copy()
|
|
es = np.concatenate(event_steps) if event_steps else np.empty(0, dtype=np.int32)
|
|
en = np.concatenate(event_nodes) if event_nodes else np.empty(0, dtype=np.int32)
|
|
active = counts_np > 0
|
|
super_values = self.superclass
|
|
desc = np.isin(super_values, ["descending_neuron", "descending_neuron_tbc"])
|
|
motor = np.isin(super_values, ["vnc_motor", "cb_motor"])
|
|
bio_s = duration_ms / 1000.0
|
|
result = {
|
|
"backend": backend, "duration_ms": duration_ms, "steps": steps,
|
|
"wall_clock_seconds": wall, "steps_per_second": steps / wall,
|
|
"biological_realtime_speed": bio_s / wall, "realtime_factor_wall_over_biological": wall / bio_s,
|
|
"total_spikes": int(counts_np.sum(dtype=np.int64)), "active_neurons": int(active.sum()),
|
|
"mean_firing_rate_hz": float(counts_np.sum(dtype=np.int64) / (self.n * bio_s)),
|
|
"descending_active_neurons": int(np.count_nonzero(active & desc)),
|
|
"descending_spikes": int(counts_np[desc].sum(dtype=np.int64)),
|
|
"motor_active_neurons": int(np.count_nonzero(active & motor)),
|
|
"motor_spikes": int(counts_np[motor].sum(dtype=np.int64)),
|
|
"peak_vram_bytes": int(torch.cuda.max_memory_allocated()),
|
|
"cpu_process_utilization_percent_of_one_core": 100.0 * cpu / wall,
|
|
"resource_sampling": resources,
|
|
"event_count": int(len(en)),
|
|
"event_sha256": hashlib.sha256(es.tobytes() + en.tobytes()).hexdigest() if capture_events else None,
|
|
"spike_counts_sha256": hashlib.sha256(counts_np.tobytes()).hexdigest(),
|
|
"final_voltage_sha256": hashlib.sha256(v_np.tobytes()).hexdigest(),
|
|
"final_g_sha256": hashlib.sha256(g_np.tobytes()).hexdigest(),
|
|
}
|
|
return {"metrics": result, "event_steps": es, "event_nodes": en, "counts": counts_np, "v": v_np, "g": g_np}
|
|
|
|
|
|
def compare(reference: dict[str, Any], candidate: dict[str, Any], tolerance: float) -> dict[str, Any]:
|
|
events = np.array_equal(reference["event_steps"], candidate["event_steps"]) and np.array_equal(reference["event_nodes"], candidate["event_nodes"])
|
|
counts = np.array_equal(reference["counts"], candidate["counts"])
|
|
vdiff = np.abs(reference["v"] - candidate["v"])
|
|
gdiff = np.abs(reference["g"] - candidate["g"])
|
|
return {
|
|
"spike_events_exact": bool(events), "per_neuron_spike_counts_exact": bool(counts),
|
|
"final_voltage_max_abs_error": float(vdiff.max(initial=0)),
|
|
"final_voltage_mean_abs_error": float(vdiff.mean()),
|
|
"final_g_max_abs_error": float(gdiff.max(initial=0)), "final_g_mean_abs_error": float(gdiff.mean()),
|
|
"state_within_tolerance": bool(vdiff.max(initial=0) <= tolerance and gdiff.max(initial=0) <= tolerance),
|
|
"eligible_for_runtime_v1": bool(events and counts),
|
|
}
|
|
|
|
|
|
def save_reference_trial(exp: Path, spec: dict[str, Any], run: dict[str, Any]) -> None:
|
|
base = exp / "regression/reference" / spec["name"]
|
|
base.mkdir(parents=True, exist_ok=True)
|
|
pq.write_table(pa.table({
|
|
"step": pa.array(run["event_steps"], type=pa.int32()),
|
|
"timestamp_ms": pa.array(run["event_steps"].astype(np.float64) * 0.1, type=pa.float64()),
|
|
"neuron_idx": pa.array(run["event_nodes"], type=pa.int32()),
|
|
}), base / "spike-events.parquet", compression="zstd")
|
|
np.save(base / "per-neuron-spike-counts.int32.npy", run["counts"], allow_pickle=False)
|
|
np.save(base / "final-voltage.float32.npy", run["v"], allow_pickle=False)
|
|
np.save(base / "final-g.float32.npy", run["g"], allow_pickle=False)
|
|
write_json(base / "statistics.json", run["metrics"] | {"trial": spec})
|
|
|
|
|
|
def propagation_stats(run: dict[str, Any], targets: np.ndarray, distance: np.ndarray,
|
|
superclass: np.ndarray) -> dict[str, Any]:
|
|
active = run["counts"] > 0
|
|
nonnegative = distance[active & (distance >= 0)]
|
|
return {
|
|
"active_neurons": int(active.sum()),
|
|
"downstream_active_neurons": int(np.count_nonzero(active & ~np.isin(np.arange(len(active)), targets))),
|
|
"propagation_depth": int(nonnegative.max(initial=0)),
|
|
"descending_active_neurons": int(np.count_nonzero(active & np.isin(superclass, ["descending_neuron", "descending_neuron_tbc"]))),
|
|
"motor_active_neurons": int(np.count_nonzero(active & np.isin(superclass, ["vnc_motor", "cb_motor"]))),
|
|
}
|
|
|
|
|
|
def visual_dynamics(exp: Path, cfg: dict[str, Any], runtime: Runtime,
|
|
layers: dict[str, np.ndarray], indptr: np.ndarray, indices: np.ndarray) -> dict[str, Any]:
|
|
spec = cfg["visual_dynamics"]; duration = float(spec["duration_ms"]); rate = float(spec["rate_hz"]); seed = int(spec["seed"])
|
|
results: dict[str, Any] = {}
|
|
tests: list[tuple[str, np.ndarray, float, int, torch.Tensor | None, str]] = []
|
|
sensory = np.flatnonzero(layers["visual_sensory"]).astype(np.int64)
|
|
tests.append(("A-forced-histaminergic-sensory", sensory, rate, seed, None, "engineering forcing of sensory output spikes"))
|
|
steps = int(round(duration / float(cfg["model"]["dt_ms"])))
|
|
tonic = generate_forced(steps, len(sensory), float(spec["tonic_histamine_baseline_hz"]), float(cfg["model"]["dt_ms"]), seed + 1, runtime.device)
|
|
mod_start = int(round(float(spec["modulation_start_ms"]) / float(cfg["model"]["dt_ms"])))
|
|
mod = generate_forced(steps - mod_start, len(sensory), float(spec["modulation_rate_hz"]), float(cfg["model"]["dt_ms"]), seed + 2, runtime.device)
|
|
tonic[mod_start:] = mod
|
|
tests.append(("B-tonic-histamine-rate-modulation", sensory, 0.0, seed + 1, tonic, "exploratory tonic output plus rate modulation; not baseline"))
|
|
first = np.flatnonzero(layers["first_order"]).astype(np.int64)
|
|
second = np.flatnonzero(layers["second_order"]).astype(np.int64)
|
|
tests.append(("C-first-order-external-input", first, rate, seed + 3, None, "engineering input interface moved to real first-order targets"))
|
|
tests.append(("C-second-order-external-input", second, rate, seed + 4, None, "engineering input interface moved to real second-order targets"))
|
|
for j, name in enumerate(("visual_sensory", "first_order", "second_order", "higher_visual", "central_brain")):
|
|
targets = np.flatnonzero(layers[name]).astype(np.int64)
|
|
# Keep the audit bounded without inventing body IDs: deterministic full layer up to 25k, otherwise fixed index stride.
|
|
if len(targets) > 25000:
|
|
targets = targets[np.linspace(0, len(targets) - 1, 25000, dtype=np.int64)]
|
|
tests.append((f"D-layer-{name}", targets, rate, seed + 10 + j, None, f"engineering external input at hierarchy layer {name}"))
|
|
for name, targets, hz, trial_seed, override, interpretation in tests:
|
|
LOG.info("visual dynamics %s targets=%d", name, len(targets))
|
|
run = runtime.run("eager", targets, hz, trial_seed, duration, False, forced_override=override)
|
|
distance = p4b.directed_distances(indptr, indices, targets, runtime.n)
|
|
record = run["metrics"] | propagation_stats(run, targets, distance, runtime.superclass) | {
|
|
"test": name, "stimulated_population_size": int(len(targets)), "interpretation": interpretation,
|
|
"biological_interpretation": "histamine remains inhibitory; forced sensory spikes are an engineering interface, not phototransduction/receptor-potential biology",
|
|
"baseline_status": "exploratory only" if name.startswith("B-") else "semantic test; no automatic selection",
|
|
}
|
|
results[name] = record
|
|
write_json(exp / "visual-audit/dynamics" / f"{name}.json", record)
|
|
write_json(exp / "visual-audit/dynamics-summary.json", results)
|
|
return results
|
|
|
|
|
|
def compiler_report(exp: Path, runtime: Runtime) -> dict[str, Any]:
|
|
report: dict[str, Any] = {"torch_version": torch.__version__, "modes": {}}
|
|
n = runtime.n
|
|
m = runtime.cfg["model"]
|
|
for mode in ("compile-default", "compile-reduce-overhead"):
|
|
try:
|
|
compiled, fn = runtime.compiled_state(mode)
|
|
v = torch.full((n,), float(m["resting_potential_mv"]), device=runtime.device)
|
|
g = torch.zeros(n, device=runtime.device); r = torch.zeros(n, dtype=torch.int32, device=runtime.device); a = torch.zeros(n, device=runtime.device)
|
|
compiled(v, g, r, a); torch.cuda.synchronize()
|
|
report["modes"][mode] = {"status": "compiled", "fullgraph": True, "graph_breaks_in_fused_static_state": 0}
|
|
except Exception as exc:
|
|
report["modes"][mode] = {"status": "failed", "error": repr(exc)}
|
|
report["fullgraph_dynamic_csr"] = {
|
|
"status": "not eligible",
|
|
"graph_break_reasons": [
|
|
"torch.nonzero produces data-dependent dynamic shape",
|
|
"Python branching on sources.numel() and edge_ids.numel()",
|
|
"expand_rows requires cumulative[-1].item() host synchronization",
|
|
"dynamic torch.arange allocation and event-dependent CSR fan-out",
|
|
],
|
|
"action": "kept eager reference; custom CUDA prototype keeps spike detection/fan-out/accumulation on GPU",
|
|
}
|
|
write_json(exp / "compiler-reports/torch-compile.json", report)
|
|
return report
|
|
|
|
|
|
def build_extension(source_dir: Path, exp: Path) -> tuple[Any | None, dict[str, Any]]:
|
|
cpp, cu = source_dir / "phase4c_ext.cpp", source_dir / "phase4c_ext.cu"
|
|
started = time.perf_counter()
|
|
try:
|
|
ext = load_extension(
|
|
name="male_cns_lif_runtime_v1_ext", sources=[str(cpp), str(cu)],
|
|
extra_cflags=["-O3"], extra_cuda_cflags=["-O3"], verbose=True,
|
|
build_directory=str(exp / "compiler-reports/extension-build"),
|
|
)
|
|
result = {"status": "PASS", "build_seconds": time.perf_counter() - started, "sources": {str(cpp): sha256(cpp), str(cu): sha256(cu)}}
|
|
return ext, result
|
|
except Exception as exc:
|
|
result = {"status": "FAILED", "build_seconds": time.perf_counter() - started, "error": repr(exc), "sources": {str(cpp): sha256(cpp), str(cu): sha256(cu)}}
|
|
return None, result
|
|
|
|
|
|
def seal_outputs(exp: Path, cfg: dict[str, Any], paths: dict[str, Path], code_paths: list[Path], summary: dict[str, Any]) -> None:
|
|
controls = {exp / "manifests/manifest.json", exp / "manifests/checksums.sha256"}
|
|
outputs = sorted(x for x in exp.rglob("*") if x.is_file() and x not in controls)
|
|
manifest = {
|
|
"phase": "4C", "created_unix_time": time.time(), "status": summary["status"],
|
|
"software": {"python": platform.python_version(), "torch": torch.__version__, "numpy": np.__version__, "pyarrow": pa.__version__},
|
|
"code": [{"path": str(x), "bytes": x.stat().st_size, "sha256": sha256(x)} for x in code_paths],
|
|
"inputs": {name: {"path": str(path), "bytes": path.stat().st_size, "sha256": sha256(path)} for name, path in paths.items()},
|
|
"outputs": [{"path": str(x), "bytes": x.stat().st_size, "sha256": sha256(x)} for x in outputs],
|
|
"prohibitions_respected": cfg["prohibitions"],
|
|
}
|
|
write_json(exp / "manifests/manifest.json", manifest)
|
|
targets = outputs + [exp / "manifests/manifest.json"]
|
|
(exp / "manifests/checksums.sha256").write_text("".join(f"{sha256(x)} {x.relative_to(exp)}\n" for x in targets), encoding="utf-8")
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(); ap.add_argument("--config", required=True, type=Path); args = ap.parse_args()
|
|
cfg = yaml.safe_load(args.config.read_text(encoding="utf-8"))
|
|
exp = Path(cfg["paths"]["experiment_root"])
|
|
for sub in ("configs", "visual-audit", "benchmarks", "compiler-reports", "profiler", "regression", "manifests", "audit"):
|
|
(exp / sub).mkdir(parents=True, exist_ok=True)
|
|
p4b.setup_logging(exp / "audit/phase4c.log")
|
|
shutil.copy2(args.config, exp / "configs/phase4c.yaml")
|
|
paths = exact_paths(cfg); hashes_before = {k: sha256(v) for k, v in paths.items()}
|
|
indptr, indices, raw = p4b.validate_graph(cfg, paths)
|
|
table = full_metadata(paths["metadata"]); pops = populations(table)
|
|
expected_pop = {"visual": 6091, "auditory": 114}
|
|
for name, count in expected_pop.items():
|
|
if int(pops[name].sum()) != count: raise RuntimeError(f"{name} count mismatch")
|
|
frozen = freeze_reference(cfg, args.config, paths)
|
|
visual_summary, layers = visual_audit(exp, table, indptr, indices, raw)
|
|
if not torch.cuda.is_available() or torch.cuda.device_count() != 1: raise RuntimeError("exactly one visible CUDA GPU required")
|
|
device = torch.device("cuda:0"); props = torch.cuda.get_device_properties(device)
|
|
if cfg["device"]["required_name_substring"] not in props.name: raise RuntimeError(f"wrong GPU: {props.name}")
|
|
source_dir = Path(__file__).resolve().parent
|
|
extension, extension_report = build_extension(source_dir, exp)
|
|
write_json(exp / "compiler-reports/custom-extension.json", extension_report)
|
|
runtime = Runtime(cfg, table, indptr, indices, raw, device, extension)
|
|
compiler = compiler_report(exp, runtime)
|
|
|
|
reference_runs: dict[str, dict[str, Any]] = {}
|
|
regression: dict[str, Any] = {}
|
|
tol = float(cfg["regression"]["float_absolute_tolerance"])
|
|
for spec in cfg["reference_trials"]:
|
|
targets = np.flatnonzero(pops[spec["population"]]).astype(np.int64)
|
|
LOG.info("reference trial %s", spec["name"])
|
|
ref = runtime.run("eager", targets, float(spec["rate_hz"]), int(spec["seed"]), float(spec["duration_ms"]), True)
|
|
dist = p4b.directed_distances(indptr, indices, targets, runtime.n)
|
|
ref["metrics"].update(propagation_stats(ref, targets, dist, runtime.superclass))
|
|
save_reference_trial(exp, spec, ref); reference_runs[spec["name"]] = ref
|
|
candidates: dict[str, Any] = {}
|
|
for backend in (["hybrid-exact", "custom-cuda"] if extension is not None else []):
|
|
cand = runtime.run(backend, targets, float(spec["rate_hz"]), int(spec["seed"]), float(spec["duration_ms"]), True)
|
|
comp = compare(ref, cand, tol)
|
|
comp.update({
|
|
"propagation_statistics_exact": propagation_stats(ref, targets, dist, runtime.superclass) == propagation_stats(cand, targets, dist, runtime.superclass),
|
|
"reference_event_sha256": ref["metrics"]["event_sha256"], "candidate_event_sha256": cand["metrics"]["event_sha256"],
|
|
})
|
|
candidates[backend] = comp
|
|
regression[spec["name"]] = candidates
|
|
write_json(exp / "regression" / f"{spec['name']}.json", candidates)
|
|
visual_results = visual_dynamics(exp, cfg, runtime, layers, indptr, indices)
|
|
|
|
bench: dict[str, Any] = {}
|
|
auditory = np.flatnonzero(pops["auditory"]).astype(np.int64)
|
|
baseline_by_duration: dict[float, dict[str, Any]] = {}
|
|
for duration in cfg["benchmark"]["durations_ms"]:
|
|
duration = float(duration); bench[str(duration)] = {}
|
|
for backend in cfg["benchmark"]["backends"]:
|
|
if (backend.startswith("custom") or backend in ("cudagraph-custom-state", "hybrid-exact")) and extension is None:
|
|
bench[str(duration)][backend] = {"status": "unavailable", "reason": extension_report.get("error")}; continue
|
|
LOG.info("benchmark backend=%s duration_ms=%s", backend, duration)
|
|
try:
|
|
run = runtime.run(backend, auditory, float(cfg["benchmark"]["rate_hz"]), int(cfg["benchmark"]["seed"]), duration, False, monitor=True)
|
|
record = run["metrics"] | {"status": "PASS"}
|
|
if backend == "eager": baseline_by_duration[duration] = run
|
|
else:
|
|
ref = baseline_by_duration.get(duration)
|
|
if ref is not None:
|
|
state_cmp = compare({**ref, "event_steps": np.empty(0,np.int32), "event_nodes": np.empty(0,np.int32)}, {**run, "event_steps": np.empty(0,np.int32), "event_nodes": np.empty(0,np.int32)}, tol)
|
|
# Performance runs deliberately disable event capture. Do not
|
|
# mislabel two empty arrays as an event-sequence regression.
|
|
state_cmp["spike_events_exact"] = None
|
|
state_cmp["event_sequence_checked"] = False
|
|
record["benchmark_state_regression"] = state_cmp
|
|
bench[str(duration)][backend] = record
|
|
except Exception as exc:
|
|
bench[str(duration)][backend] = {"status": "FAILED", "error": repr(exc)}
|
|
write_json(exp / "benchmarks/runtime-matrix.json", bench)
|
|
|
|
eligible = ["eager"]
|
|
if extension is not None and all(regression[x].get("hybrid-exact", {}).get("eligible_for_runtime_v1") for x in regression):
|
|
eligible.append("hybrid-exact")
|
|
if extension is not None and all(regression[x].get("custom-cuda", {}).get("eligible_for_runtime_v1") for x in regression):
|
|
eligible.append("custom-cuda")
|
|
# Compiled/CUDA-graph paths need an exact 1 s event regression before eligibility.
|
|
audit_spec = cfg["reference_trials"][0]; audit_ref = reference_runs[audit_spec["name"]]
|
|
targets = np.flatnonzero(pops[audit_spec["population"]]).astype(np.int64)
|
|
backend_regression: dict[str, Any] = {}
|
|
for backend in ("compile-default", "compile-reduce-overhead", "cudagraph-custom-state"):
|
|
if backend == "cudagraph-custom-state" and extension is None: continue
|
|
try:
|
|
candidate = runtime.run(backend, targets, float(audit_spec["rate_hz"]), int(audit_spec["seed"]), 1000.0, True)
|
|
backend_regression[backend] = compare(audit_ref, candidate, tol)
|
|
if backend_regression[backend]["eligible_for_runtime_v1"]: eligible.append(backend)
|
|
except Exception as exc:
|
|
backend_regression[backend] = {"eligible_for_runtime_v1": False, "error": repr(exc)}
|
|
write_json(exp / "regression/backend-eligibility.json", backend_regression)
|
|
|
|
def speed(backend: str) -> float:
|
|
x = bench.get("10000.0", {}).get(backend, {})
|
|
return float(x.get("biological_realtime_speed", 0.0)) if x.get("status") == "PASS" else 0.0
|
|
fastest = max(eligible, key=speed)
|
|
runtime_root = Path(cfg["paths"]["model_root"]) / "runtime-v1"
|
|
runtime_root.mkdir(parents=True, exist_ok=True)
|
|
runtime_sources = (Path(__file__).resolve(), source_dir / "phase4c_ext.cpp", source_dir / "phase4c_ext.cu", source_dir / "requirements.txt", args.config.resolve())
|
|
for source in runtime_sources:
|
|
copy_immutable(source, runtime_root / source.name)
|
|
runtime_version = {
|
|
"name": cfg["names"]["runtime"], "backend": fastest,
|
|
"reference": cfg["names"]["reference"], "eligibility": backend_regression | regression,
|
|
"ten_second_biological_realtime_speed": speed(fastest),
|
|
"spike_behavior_requirement": "exact on all five 1 s frozen reference trials for custom runtime; exact auditory trial for compiler variants",
|
|
"algorithm": "exact sparse event-driven CSR fan-out; no dense full-edge scan; exact 18-step GPU-resident ring buffer",
|
|
"fallback_backend": "eager",
|
|
}
|
|
write_json(runtime_root / "VERSION.json", runtime_version); (runtime_root / "VERSION.json").chmod(0o444)
|
|
hashes_after = {k: sha256(v) for k, v in paths.items()}
|
|
if hashes_before != hashes_after: raise RuntimeError("Phase 3B inputs changed")
|
|
summary = {
|
|
"status": "PASS" if fastest else "FAILED", "reference_frozen": True,
|
|
"reference": frozen, "visual_audit": visual_summary, "visual_dynamics": visual_results,
|
|
"compiler_report": compiler, "custom_extension": extension_report,
|
|
"regression": regression, "backend_regression": backend_regression,
|
|
"runtime_v1_backend": fastest, "benchmarks": bench,
|
|
"device": {"name": props.name, "total_memory_bytes": int(props.total_memory)},
|
|
"input_hashes_unchanged": True, "entered_phase4d": False, "game_rl_web_3d_training": False,
|
|
}
|
|
write_json(exp / "benchmarks/summary.json", summary)
|
|
seal_outputs(exp, cfg, paths, list(runtime_sources), summary)
|
|
LOG.info("PHASE4C_COMPLETE runtime=%s", fastest)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|