#!/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("" 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()