Digital Drosophila Archive
1262 lines
63 KiB
Python
1262 lines
63 KiB
Python
#!/usr/bin/env python3
|
|
"""MaleCNS Phase 4B: strict LIF dynamics, propagation and performance audit."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import bisect
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import platform
|
|
import random
|
|
import resource
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from collections import Counter, defaultdict
|
|
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
|
|
|
|
|
|
LOG = logging.getLogger("phase4b")
|
|
|
|
|
|
def sha256(path: Path, chunk: int = 8 << 20) -> str:
|
|
h = hashlib.sha256()
|
|
with path.open("rb") as f:
|
|
while b := f.read(chunk):
|
|
h.update(b)
|
|
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 setup_logging(path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|
handlers=[logging.FileHandler(path, encoding="utf-8"), logging.StreamHandler()],
|
|
force=True,
|
|
)
|
|
|
|
|
|
def classify_nt(value: Any, cfg: dict[str, Any]) -> tuple[str, int]:
|
|
policy = cfg["sign_policy"]
|
|
if value is None or not str(value).strip():
|
|
return "UNKNOWN", int(policy["unknown_sign"])
|
|
key = str(value).strip().lower()
|
|
if key == "unclear":
|
|
return "UNCLEAR", int(policy["unclear_sign"])
|
|
if key not in policy["neurotransmitter_sign"]:
|
|
if policy["unmapped_policy"] == "error":
|
|
raise RuntimeError(f"unmapped neurotransmitter {value!r}")
|
|
return f"UNMAPPED:{key}", 0
|
|
return key, int(policy["neurotransmitter_sign"][key])
|
|
|
|
|
|
def exact_input_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 connectome root is {allowed}; got {root}")
|
|
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 validate_graph(cfg: dict[str, Any], paths: dict[str, Path]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
expected = cfg["expected_graph"]
|
|
n, nnz, weight_sum = int(expected["neurons"]), int(expected["directed_edges"]), int(expected["synaptic_contact_weight_sum"])
|
|
indptr = np.load(paths["indptr"], mmap_mode="r")
|
|
indices = np.load(paths["indices"], mmap_mode="r")
|
|
raw = np.load(paths["raw_weights"], mmap_mode="r")
|
|
if indptr.dtype != np.int64 or indices.dtype != np.int32 or raw.dtype != np.int64:
|
|
raise RuntimeError(f"CSR dtype mismatch: {indptr.dtype}, {indices.dtype}, {raw.dtype}")
|
|
if indptr.shape != (n + 1,) or indices.shape != (nnz,) or raw.shape != (nnz,):
|
|
raise RuntimeError(f"CSR shape mismatch: {indptr.shape}, {indices.shape}, {raw.shape}")
|
|
if int(indptr[-1]) != nnz or int(raw.sum(dtype=np.int64)) != weight_sum:
|
|
raise RuntimeError("fixed graph invariant failed")
|
|
if int(raw.min()) <= 0 or int(indices.min()) < 0 or int(indices.max()) >= n:
|
|
raise RuntimeError("invalid CSR range")
|
|
return indptr, indices, raw
|
|
|
|
|
|
def metadata_arrays(cfg: dict[str, Any], path: Path) -> dict[str, Any]:
|
|
columns = [
|
|
"neuron_idx", "body_id", "neurotransmitter", "annotation_type", "annotation_superclass",
|
|
"annotation_class", "annotation_subclass",
|
|
]
|
|
table = pq.read_table(path, columns=columns)
|
|
n = int(cfg["expected_graph"]["neurons"])
|
|
neuron_idx = table["neuron_idx"].to_numpy(zero_copy_only=False)
|
|
if len(table) != n or neuron_idx.dtype != np.int32 or not np.array_equal(neuron_idx, np.arange(n, dtype=np.int32)):
|
|
raise RuntimeError("metadata mapping invariant failed")
|
|
nt = table["neurotransmitter"].to_pylist()
|
|
classes: list[str] = []
|
|
signs = np.empty(n, dtype=np.int8)
|
|
for i, value in enumerate(nt):
|
|
cls, sign = classify_nt(value, cfg)
|
|
classes.append(cls)
|
|
signs[i] = sign
|
|
return {
|
|
"table": table,
|
|
"body_id": table["body_id"].to_numpy(zero_copy_only=False),
|
|
"nt": nt,
|
|
"nt_class": np.asarray(classes, dtype=object),
|
|
"sign": signs,
|
|
"type": np.asarray(table["annotation_type"].to_pylist(), dtype=object),
|
|
"superclass": np.asarray(table["annotation_superclass"].to_pylist(), dtype=object),
|
|
"class": np.asarray(table["annotation_class"].to_pylist(), dtype=object),
|
|
"subclass": np.asarray(table["annotation_subclass"].to_pylist(), dtype=object),
|
|
}
|
|
|
|
|
|
def population_masks(cfg: dict[str, Any], meta: dict[str, Any]) -> dict[str, np.ndarray]:
|
|
cls, sub = meta["class"], meta["subclass"]
|
|
mech = {"mechanosensory", "mechanosensory_tactile", "mechanosensory_proprioceptive"}
|
|
masks = {
|
|
"visual": cls == "visual",
|
|
"mechanosensory": np.asarray([(c in mech) and s != "auditory" for c, s in zip(cls, sub)], dtype=bool),
|
|
"olfactory": cls == "olfactory",
|
|
"gustatory": cls == "gustatory",
|
|
"auditory": sub == "auditory",
|
|
}
|
|
if set(masks) != set(cfg["sensory_populations"]):
|
|
raise RuntimeError("population implementation/config mismatch")
|
|
if any(not mask.any() for mask in masks.values()):
|
|
raise RuntimeError("empty sensory population")
|
|
return masks
|
|
|
|
|
|
def directed_distances(indptr: np.ndarray, indices: np.ndarray, sources: np.ndarray, n: int) -> np.ndarray:
|
|
distances = np.full(n, -1, dtype=np.int16)
|
|
distances[sources] = 0
|
|
frontier = sources.astype(np.int64, copy=False)
|
|
depth = 0
|
|
while len(frontier):
|
|
depth += 1
|
|
next_mask = np.zeros(n, dtype=bool)
|
|
for src in frontier:
|
|
next_mask[indices[int(indptr[src]):int(indptr[src + 1])]] = True
|
|
next_mask[distances >= 0] = False
|
|
frontier = np.flatnonzero(next_mask).astype(np.int64)
|
|
if not len(frontier):
|
|
break
|
|
if depth >= np.iinfo(np.int16).max:
|
|
raise RuntimeError("graph distance overflow")
|
|
distances[frontier] = depth
|
|
return distances
|
|
|
|
|
|
def audit_populations(
|
|
cfg: dict[str, Any], exp: Path, meta: dict[str, Any], masks: dict[str, np.ndarray],
|
|
indptr: np.ndarray, indices: np.ndarray, raw: np.ndarray,
|
|
) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
|
|
output: dict[str, Any] = {}
|
|
distances: dict[str, np.ndarray] = {}
|
|
pop_dir = exp / "propagation/populations"
|
|
pop_dir.mkdir(parents=True, exist_ok=True)
|
|
n = int(cfg["expected_graph"]["neurons"])
|
|
for name, mask in masks.items():
|
|
nodes = np.flatnonzero(mask).astype(np.int32)
|
|
starts, stops = indptr[nodes], indptr[nodes + 1]
|
|
outgoing_edges = int(np.sum(stops - starts, dtype=np.int64))
|
|
outgoing_weight = 0
|
|
first = np.zeros(n, dtype=bool)
|
|
for start, stop in zip(starts, stops):
|
|
outgoing_weight += int(raw[int(start):int(stop)].sum(dtype=np.int64))
|
|
first[indices[int(start):int(stop)]] = True
|
|
first[mask] = False
|
|
d = directed_distances(indptr, indices, nodes, n)
|
|
distances[name] = d
|
|
membership = pa.table({
|
|
"neuron_idx": pa.array(nodes, type=pa.int32()),
|
|
"body_id": pa.array(meta["body_id"][nodes], type=pa.uint64()),
|
|
"neurotransmitter": pa.array([meta["nt"][i] for i in nodes], type=pa.string()),
|
|
})
|
|
membership_path = pop_dir / f"{name}.parquet"
|
|
pq.write_table(membership, membership_path, compression="zstd")
|
|
nt_counts = Counter(meta["nt_class"][nodes].tolist())
|
|
record = {
|
|
"name": name,
|
|
"filter_expression": cfg["sensory_populations"][name]["expression"],
|
|
"neuron_count": int(len(nodes)),
|
|
"neuron_indices": nodes.tolist(),
|
|
"body_ids": [int(x) for x in meta["body_id"][nodes]],
|
|
"neurotransmitter_composition": dict(sorted(nt_counts.items())),
|
|
"outgoing_edge_count": outgoing_edges,
|
|
"outgoing_synapse_weight": outgoing_weight,
|
|
"downstream_first_hop_neuron_count": int(first.sum()),
|
|
"max_reachable_depth": int(d.max()),
|
|
"reachable_neuron_count": int((d >= 0).sum()),
|
|
"membership_parquet": str(membership_path),
|
|
}
|
|
output[name] = record
|
|
write_json(pop_dir / f"{name}.json", record)
|
|
np.save(pop_dir / f"{name}.distances.int16.npy", d, allow_pickle=False)
|
|
LOG.info("population %s neurons=%d outgoing_edges=%d first_hop=%d", name, len(nodes), outgoing_edges, first.sum())
|
|
write_json(exp / "propagation/sensory-populations.json", output)
|
|
return output, distances
|
|
|
|
|
|
def tiny_numpy(
|
|
signs: np.ndarray, raw: np.ndarray, rows: np.ndarray, cols: np.ndarray,
|
|
forced: dict[int, list[int]], cfg: dict[str, Any], steps: int,
|
|
) -> dict[str, Any]:
|
|
n = int(len(signs))
|
|
m = cfg["model"]
|
|
dt = float(m["dt_ms"])
|
|
alpha = dt / float(m["membrane_tau_ms"])
|
|
g_decay = 1.0 - dt / float(m["synaptic_tau_ms"])
|
|
delay = int(round(float(m["synaptic_delay_ms"]) / dt))
|
|
ref_steps = int(round(float(m["refractory_ms"]) / dt))
|
|
rest, reset, threshold = float(m["resting_potential_mv"]), float(m["reset_potential_mv"]), float(m["threshold_mv"])
|
|
weight = signs[rows].astype(np.float64) * raw.astype(np.float64) * float(m["w_syn_mv_per_anatomical_synapse"])
|
|
v = np.full(n, rest, dtype=np.float64)
|
|
g = np.zeros(n, dtype=np.float64)
|
|
ref = np.zeros(n, dtype=np.int32)
|
|
queue = np.zeros((delay + 1, n), dtype=np.float64)
|
|
events: list[tuple[int, int]] = []
|
|
g_history = np.zeros((steps, n), dtype=np.float64)
|
|
rejected = 0
|
|
for step in range(steps):
|
|
slot = step % len(queue)
|
|
g *= g_decay
|
|
g += queue[slot]
|
|
queue[slot].fill(0)
|
|
ref = np.maximum(ref - 1, 0)
|
|
available = ref == 0
|
|
v[available] += alpha * (g[available] - (v[available] - rest))
|
|
v[~available] = reset
|
|
spikes = available & (v >= threshold)
|
|
for src in forced.get(step, []):
|
|
if available[src]:
|
|
spikes[src] = True
|
|
else:
|
|
rejected += 1
|
|
sources = np.flatnonzero(spikes)
|
|
if len(sources):
|
|
v[sources] = reset
|
|
ref[sources] = ref_steps
|
|
for src in sources:
|
|
events.append((step, int(src)))
|
|
hit = rows == src
|
|
np.add.at(queue[(step + delay) % len(queue)], cols[hit], weight[hit])
|
|
g_history[step] = g
|
|
return {"v": v, "g": g, "events": events, "g_history": g_history, "rejected": rejected}
|
|
|
|
|
|
def tiny_torch(
|
|
signs: np.ndarray, raw: np.ndarray, rows: np.ndarray, cols: np.ndarray,
|
|
forced: dict[int, list[int]], cfg: dict[str, Any], steps: int, device: torch.device,
|
|
) -> dict[str, Any]:
|
|
n = int(len(signs))
|
|
m = cfg["model"]
|
|
dt = float(m["dt_ms"])
|
|
alpha = dt / float(m["membrane_tau_ms"])
|
|
g_decay = 1.0 - dt / float(m["synaptic_tau_ms"])
|
|
delay = int(round(float(m["synaptic_delay_ms"]) / dt))
|
|
ref_steps = int(round(float(m["refractory_ms"]) / dt))
|
|
rest, reset, threshold = float(m["resting_potential_mv"]), float(m["reset_potential_mv"]), float(m["threshold_mv"])
|
|
rows_t = torch.as_tensor(rows, dtype=torch.int64, device=device)
|
|
cols_t = torch.as_tensor(cols, dtype=torch.int64, device=device)
|
|
weight = torch.as_tensor(signs[rows].astype(np.float64) * raw.astype(np.float64) * float(m["w_syn_mv_per_anatomical_synapse"]), dtype=torch.float64, device=device)
|
|
v = torch.full((n,), rest, dtype=torch.float64, device=device)
|
|
g = torch.zeros(n, dtype=torch.float64, device=device)
|
|
ref = torch.zeros(n, dtype=torch.int32, device=device)
|
|
queue = torch.zeros((delay + 1, n), dtype=torch.float64, device=device)
|
|
events: list[tuple[int, int]] = []
|
|
history: list[np.ndarray] = []
|
|
rejected = 0
|
|
for step in range(steps):
|
|
slot = step % len(queue)
|
|
g.mul_(g_decay).add_(queue[slot])
|
|
queue[slot].zero_()
|
|
ref.clamp_min_(1).sub_(1)
|
|
available = ref == 0
|
|
v[available] += alpha * (g[available] - (v[available] - rest))
|
|
v[~available] = reset
|
|
spikes = available & (v >= threshold)
|
|
for src in forced.get(step, []):
|
|
if bool(available[src].item()):
|
|
spikes[src] = True
|
|
else:
|
|
rejected += 1
|
|
sources = torch.nonzero(spikes, as_tuple=False).flatten()
|
|
if sources.numel():
|
|
v[sources] = reset
|
|
ref[sources] = ref_steps
|
|
for src in sources.cpu().tolist():
|
|
events.append((step, int(src)))
|
|
hit = rows_t == src
|
|
queue[(step + delay) % len(queue)].scatter_add_(0, cols_t[hit], weight[hit])
|
|
history.append(g.cpu().numpy().copy())
|
|
torch.cuda.synchronize(device)
|
|
return {"v": v.cpu().numpy(), "g": g.cpu().numpy(), "events": events, "g_history": np.stack(history), "rejected": rejected}
|
|
|
|
|
|
def run_dynamics_unit_tests(cfg: dict[str, Any], exp: Path, device: torch.device) -> dict[str, Any]:
|
|
# 0->1 excitatory; 2->3 inhibitory; 4->5 ten contacts; 6->6 self-loop;
|
|
# 7->9 and 8->9 synchronous convergent input.
|
|
rows = np.array([0, 2, 4, 6, 7, 8], dtype=np.int64)
|
|
cols = np.array([1, 3, 5, 6, 9, 9], dtype=np.int64)
|
|
raw = np.array([1, 1, 10, 1, 1, 1], dtype=np.int64)
|
|
signs = np.array([1, 0, -1, 0, 1, 0, 1, 1, 1, 0], dtype=np.int8)
|
|
forced = {0: [0, 2, 4, 6, 7, 8], 1: [0], 22: [0]}
|
|
steps = 30
|
|
cpu = tiny_numpy(signs, raw, rows, cols, forced, cfg, steps)
|
|
gpu = tiny_torch(signs, raw, rows, cols, forced, cfg, steps, device)
|
|
delay = int(round(float(cfg["model"]["synaptic_delay_ms"]) / float(cfg["model"]["dt_ms"])))
|
|
w = float(cfg["model"]["w_syn_mv_per_anatomical_synapse"])
|
|
checks = {
|
|
"single_excitatory_edge": math.isclose(cpu["g_history"][delay, 1], w, abs_tol=1e-12),
|
|
"single_inhibitory_edge": math.isclose(cpu["g_history"][delay, 3], -w, abs_tol=1e-12),
|
|
"multi_synapse_high_weight_edge": math.isclose(cpu["g_history"][delay, 5], 10 * w, abs_tol=1e-12),
|
|
"delay_exact_1_8_ms": bool(np.all(cpu["g_history"][:delay, [1, 3, 5, 6, 9]] == 0)),
|
|
"refractory_rejects_early_forced_spike": cpu["rejected"] == 1 and (22, 0) in cpu["events"],
|
|
"self_loop": math.isclose(cpu["g_history"][delay, 6], w, abs_tol=1e-12),
|
|
"synchronous_presynaptic_sum": math.isclose(cpu["g_history"][delay, 9], 2 * w, abs_tol=1e-12),
|
|
"g_decays_during_refractory": math.isclose(cpu["g_history"][delay + 1, 6], w * (1 - 0.1 / 5.0), abs_tol=1e-12),
|
|
"cpu_gpu_events_equal": cpu["events"] == gpu["events"],
|
|
"cpu_gpu_voltage_tolerance": float(np.max(np.abs(cpu["v"] - gpu["v"]))) <= 1e-10,
|
|
"cpu_gpu_g_tolerance": float(np.max(np.abs(cpu["g"] - gpu["g"]))) <= 1e-10,
|
|
}
|
|
result = {
|
|
"passed": all(checks.values()),
|
|
"checks": checks,
|
|
"delay_steps": delay,
|
|
"refractory_steps": int(round(float(cfg["model"]["refractory_ms"]) / float(cfg["model"]["dt_ms"]))),
|
|
"cpu_events": cpu["events"],
|
|
"gpu_events": gpu["events"],
|
|
"max_abs_voltage_difference": float(np.max(np.abs(cpu["v"] - gpu["v"]))),
|
|
"max_abs_synaptic_state_difference": float(np.max(np.abs(cpu["g"] - gpu["g"]))),
|
|
"refractory_g_policy": cfg["model"]["refractory_g_policy"],
|
|
}
|
|
write_json(exp / "audit/dynamics-unit-tests.json", result)
|
|
if not result["passed"]:
|
|
raise RuntimeError(f"dynamics unit tests failed: {result}")
|
|
return result
|
|
|
|
|
|
class GPUSampler:
|
|
def __init__(self, uuid: str, interval: float):
|
|
self.uuid = uuid
|
|
self.interval = interval
|
|
self.samples: list[dict[str, float]] = []
|
|
self.stop_event = threading.Event()
|
|
self.thread: threading.Thread | None = None
|
|
|
|
def _run(self) -> None:
|
|
while not self.stop_event.is_set():
|
|
try:
|
|
out = subprocess.check_output(
|
|
["nvidia-smi", "--query-gpu=uuid,utilization.gpu,memory.used", "--format=csv,noheader,nounits"],
|
|
text=True, timeout=2,
|
|
)
|
|
for line in out.splitlines():
|
|
parts = [x.strip() for x in line.split(",")]
|
|
if len(parts) == 3 and parts[0] == self.uuid:
|
|
self.samples.append({"gpu_util_percent": float(parts[1]), "memory_used_mib": float(parts[2])})
|
|
break
|
|
except Exception:
|
|
pass
|
|
self.stop_event.wait(self.interval)
|
|
|
|
def start(self) -> None:
|
|
self.thread = threading.Thread(target=self._run, daemon=True)
|
|
self.thread.start()
|
|
|
|
def stop(self) -> dict[str, Any]:
|
|
self.stop_event.set()
|
|
if self.thread:
|
|
self.thread.join(timeout=3)
|
|
gpu = [x["gpu_util_percent"] for x in self.samples]
|
|
mem = [x["memory_used_mib"] for x in self.samples]
|
|
return {
|
|
"samples": len(self.samples),
|
|
"gpu_util_mean_percent": float(np.mean(gpu)) if gpu else None,
|
|
"gpu_util_max_percent": float(np.max(gpu)) if gpu else None,
|
|
"nvidia_smi_peak_memory_used_mib": float(np.max(mem)) if mem else None,
|
|
}
|
|
|
|
|
|
def tensor_record(name: str, tensor: torch.Tensor) -> dict[str, Any]:
|
|
return {
|
|
"name": name,
|
|
"shape": list(tensor.shape),
|
|
"dtype": str(tensor.dtype),
|
|
"device": str(tensor.device),
|
|
"bytes": int(tensor.numel() * tensor.element_size()),
|
|
}
|
|
|
|
|
|
def array_record(name: str, array: np.ndarray) -> dict[str, Any]:
|
|
return {
|
|
"name": name,
|
|
"shape": list(array.shape),
|
|
"dtype": str(array.dtype),
|
|
"device": "cpu",
|
|
"bytes": int(array.nbytes),
|
|
}
|
|
|
|
|
|
def expand_rows_profiled(indptr: torch.Tensor, sources: torch.Tensor, profile: dict[str, float] | None) -> 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)
|
|
sync_start = time.perf_counter()
|
|
total = int(cumulative[-1].item()) if cumulative.numel() else 0
|
|
if profile is not None:
|
|
profile["synchronization"] += time.perf_counter() - sync_start
|
|
if total == 0:
|
|
return torch.empty(0, dtype=torch.int64, device=indptr.device)
|
|
seq = torch.arange(total, dtype=torch.int64, device=indptr.device)
|
|
row_position = torch.searchsorted(cumulative, seq, right=True)
|
|
previous = cumulative - counts
|
|
return starts[row_position] + seq - previous[row_position]
|
|
|
|
|
|
def category_activity(values: np.ndarray, active: np.ndarray, spike_counts: np.ndarray) -> dict[str, Any]:
|
|
result: dict[str, dict[str, int]] = {}
|
|
for value in sorted({str(x) for x in values if x is not None and str(x).strip()}):
|
|
mask = values == value
|
|
result[value] = {
|
|
"active_neurons": int(np.count_nonzero(active & mask)),
|
|
"spikes": int(spike_counts[mask].sum(dtype=np.int64)),
|
|
}
|
|
null_mask = np.asarray([x is None or not str(x).strip() for x in values], dtype=bool)
|
|
result["<null>"] = {
|
|
"active_neurons": int(np.count_nonzero(active & null_mask)),
|
|
"spikes": int(spike_counts[null_mask].sum(dtype=np.int64)),
|
|
}
|
|
return result
|
|
|
|
|
|
class FullGraphRuntime:
|
|
def __init__(
|
|
self, cfg: dict[str, Any], paths: dict[str, Path], meta: dict[str, Any],
|
|
indptr_np: np.ndarray, indices_np: np.ndarray, raw_np: np.ndarray, device: torch.device,
|
|
):
|
|
self.cfg = cfg
|
|
self.paths = paths
|
|
self.meta = meta
|
|
self.device = device
|
|
self.n = int(cfg["expected_graph"]["neurons"])
|
|
self.indptr_np, self.indices_np, self.raw_np = indptr_np, indices_np, raw_np
|
|
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)
|
|
self.sign = torch.as_tensor(meta["sign"], dtype=torch.int8, device=device)
|
|
# Exact raw counts fit losslessly in float32 because max is 2591.
|
|
signed_raw_np = np.empty(len(raw_np), dtype=np.float32)
|
|
for src in range(self.n):
|
|
start, stop = int(indptr_np[src]), int(indptr_np[src + 1])
|
|
signed_raw_np[start:stop] = raw_np[start:stop].astype(np.float32) * float(meta["sign"][src])
|
|
self.signed_raw_np = signed_raw_np
|
|
self.effective: torch.Tensor | None = None
|
|
|
|
def set_weight(self, w_syn: float) -> None:
|
|
# Keep exactly one edge-weight tensor on GPU; the rescalable base stays on CPU.
|
|
self.effective = None
|
|
torch.cuda.empty_cache()
|
|
self.effective = torch.as_tensor(self.signed_raw_np, dtype=torch.float32, device=self.device)
|
|
self.effective.mul_(float(w_syn))
|
|
|
|
def persistent_tensor_audit(self) -> list[dict[str, Any]]:
|
|
if self.effective is None:
|
|
raise RuntimeError("effective weights not initialized")
|
|
return [
|
|
array_record("source_csr_indptr_readonly_memmap", self.indptr_np),
|
|
array_record("source_csr_indices_readonly_memmap", self.indices_np),
|
|
array_record("source_raw_synapse_count_int64_readonly_memmap", self.raw_np),
|
|
tensor_record("csr_indptr", self.indptr),
|
|
tensor_record("csr_indices", self.indices),
|
|
tensor_record("presynaptic_sign", self.sign),
|
|
array_record("signed_raw_synapse_count_rescale_source", self.signed_raw_np),
|
|
tensor_record("runtime_effective_weights", self.effective),
|
|
]
|
|
|
|
def run(
|
|
self, population: str, targets_np: np.ndarray, distances: np.ndarray, rate_hz: float,
|
|
seed: int, duration_ms: float, w_syn: float, log_events: bool,
|
|
event_path: Path | None = None, profile_categories: bool = False,
|
|
monitor_resources: bool = False, abort_supercritical: bool | None = None,
|
|
) -> tuple[dict[str, Any], np.ndarray]:
|
|
self.set_weight(w_syn)
|
|
assert self.effective is not None
|
|
cfg, m = self.cfg, self.cfg["model"]
|
|
dt = float(m["dt_ms"])
|
|
steps = int(round(duration_ms / dt))
|
|
delay_steps = int(round(float(m["synaptic_delay_ms"]) / dt))
|
|
queue_len = delay_steps + 1
|
|
refractory_steps = int(round(float(m["refractory_ms"]) / dt))
|
|
alpha = dt / float(m["membrane_tau_ms"])
|
|
g_decay = 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)
|
|
target_mask_np = np.zeros(self.n, dtype=bool)
|
|
target_mask_np[targets_np] = True
|
|
probability = float(rate_hz) * dt / 1000.0
|
|
generator = torch.Generator(device=self.device)
|
|
generator.manual_seed(int(seed))
|
|
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)
|
|
delay_queue = torch.zeros((queue_len, self.n), dtype=torch.float32, device=self.device)
|
|
spike_counts = torch.zeros(self.n, dtype=torch.int32, device=self.device)
|
|
event_steps: list[np.ndarray] = []
|
|
event_neurons: list[np.ndarray] = []
|
|
stimulated_attempts = 0
|
|
stimulated_spikes = 0
|
|
peak_step_spikes = 0
|
|
total_spikes_running = 0
|
|
self_loop_transmissions = 0
|
|
self_loop_sources = np.zeros(self.n, dtype=bool)
|
|
for src in range(self.n):
|
|
if np.any(self.indices_np[int(self.indptr_np[src]):int(self.indptr_np[src + 1])] == src):
|
|
self_loop_sources[src] = True
|
|
self_loop_sources_t = torch.as_tensor(self_loop_sources, dtype=torch.bool, device=self.device)
|
|
numerical_ok = True
|
|
termination_reason = None
|
|
timing = defaultdict(float) if profile_categories else None
|
|
memory_samples: list[int] = []
|
|
runtime_tensors = self.persistent_tensor_audit() + [
|
|
tensor_record("membrane_voltage", v), tensor_record("synaptic_state_g", g),
|
|
tensor_record("refractory_state", refractory), tensor_record("delay_event_buffer", delay_queue),
|
|
tensor_record("spike_counts", spike_counts), tensor_record("stimulated_targets", targets),
|
|
]
|
|
sampler = None
|
|
if monitor_resources:
|
|
sampler = GPUSampler(cfg["device"]["launcher_visible_gpu_uuid"], float(cfg["performance"]["gpu_sample_interval_seconds"]))
|
|
sampler.start()
|
|
cpu_before = resource.getrusage(resource.RUSAGE_SELF)
|
|
torch.cuda.synchronize(self.device)
|
|
torch.cuda.reset_peak_memory_stats(self.device)
|
|
started = time.perf_counter()
|
|
completed_steps = 0
|
|
abort_flag = cfg["classification"]["abort_supercritical"] if abort_supercritical is None else abort_supercritical
|
|
|
|
def timed_sync_start() -> float:
|
|
if profile_categories:
|
|
torch.cuda.synchronize(self.device)
|
|
return time.perf_counter()
|
|
|
|
def timed_sync_end(name: str, started_at: float) -> None:
|
|
if profile_categories:
|
|
torch.cuda.synchronize(self.device)
|
|
assert timing is not None
|
|
timing[name] += time.perf_counter() - started_at
|
|
|
|
for step in range(steps):
|
|
t0 = timed_sync_start()
|
|
slot = step % queue_len
|
|
g.mul_(g_decay).add_(delay_queue[slot])
|
|
delay_queue[slot].zero_()
|
|
timed_sync_end("delay_queue_and_g", t0)
|
|
|
|
t0 = timed_sync_start()
|
|
refractory.clamp_min_(1).sub_(1)
|
|
available = refractory == 0
|
|
v[available] += alpha * (g[available] - (v[available] - rest))
|
|
v[~available] = reset
|
|
timed_sync_end("neuron_state_integration", t0)
|
|
|
|
t0 = timed_sync_start()
|
|
natural = available & (v >= threshold)
|
|
poisson = torch.rand(targets.numel(), generator=generator, device=self.device) < probability
|
|
stimulated_attempts += int(poisson.sum().item())
|
|
accepted = poisson & available[targets]
|
|
forced_sources = targets[accepted]
|
|
stimulated_spikes += int(forced_sources.numel())
|
|
spikes = natural
|
|
if forced_sources.numel():
|
|
spikes[forced_sources] = True
|
|
sources = torch.nonzero(spikes, as_tuple=False).flatten()
|
|
count = int(sources.numel())
|
|
total_spikes_running += count
|
|
peak_step_spikes = max(peak_step_spikes, count)
|
|
spike_counts.add_(spikes.to(torch.int32))
|
|
timed_sync_end("spike_detection_and_poisson", t0)
|
|
|
|
if count:
|
|
v[sources] = reset
|
|
refractory[sources] = refractory_steps
|
|
self_loop_transmissions += int(self_loop_sources_t[sources].sum().item())
|
|
|
|
t0 = timed_sync_start()
|
|
edge_ids = expand_rows_profiled(self.indptr, sources, timing)
|
|
if edge_ids.numel():
|
|
posts = self.indices[edge_ids].to(torch.int64)
|
|
arrivals = self.effective[edge_ids]
|
|
delay_queue[(step + delay_steps) % queue_len].scatter_add_(0, posts, arrivals)
|
|
timed_sync_end("csr_fanout_and_accumulation", t0)
|
|
|
|
if log_events:
|
|
t0 = timed_sync_start()
|
|
src_np = sources.detach().cpu().numpy().astype(np.int32, copy=False)
|
|
event_steps.append(np.full(count, step, dtype=np.int32))
|
|
event_neurons.append(src_np.copy())
|
|
timed_sync_end("event_memory_transfer_and_logging", t0)
|
|
|
|
if step % int(cfg["audit"]["finite_check_interval_steps"]) == 0 or step == steps - 1:
|
|
finite = torch.isfinite(v).all() & torch.isfinite(g).all() & torch.isfinite(delay_queue).all()
|
|
if not bool(finite.item()):
|
|
numerical_ok = False
|
|
termination_reason = "nan_or_inf"
|
|
completed_steps = step + 1
|
|
break
|
|
memory_samples.append(torch.cuda.memory_allocated(self.device))
|
|
|
|
completed_steps = step + 1
|
|
if abort_flag and (step + 1) % int(cfg["classification"]["check_interval_steps"]) == 0:
|
|
elapsed_s = (step + 1) * dt / 1000.0
|
|
active_now = int(torch.count_nonzero(spike_counts).item())
|
|
mean_rate_now = total_spikes_running / (self.n * elapsed_s)
|
|
if peak_step_spikes / self.n >= float(cfg["classification"]["supercritical_peak_step_fraction"]):
|
|
termination_reason = "supercritical_peak_synchrony"
|
|
elif active_now / self.n >= float(cfg["classification"]["supercritical_active_fraction"]):
|
|
termination_reason = "supercritical_active_fraction"
|
|
elif mean_rate_now >= float(cfg["classification"]["supercritical_mean_rate_hz"]):
|
|
termination_reason = "supercritical_mean_rate"
|
|
if termination_reason:
|
|
break
|
|
|
|
torch.cuda.synchronize(self.device)
|
|
wall = time.perf_counter() - started
|
|
cpu_after = resource.getrusage(resource.RUSAGE_SELF)
|
|
resource_result = sampler.stop() if sampler else {"samples": 0, "gpu_util_mean_percent": None, "gpu_util_max_percent": None, "nvidia_smi_peak_memory_used_mib": None}
|
|
cpu_seconds = (cpu_after.ru_utime + cpu_after.ru_stime) - (cpu_before.ru_utime + cpu_before.ru_stime)
|
|
spike_counts_np = spike_counts.cpu().numpy().astype(np.int32, copy=True)
|
|
active = spike_counts_np > 0
|
|
simulated_ms = completed_steps * dt
|
|
simulated_s = simulated_ms / 1000.0
|
|
total_spikes = int(spike_counts_np.sum(dtype=np.int64))
|
|
active_outside = active & ~target_mask_np
|
|
sign_np = self.meta["sign"]
|
|
descending = np.isin(self.meta["superclass"], ["descending_neuron", "descending_neuron_tbc"])
|
|
motor = np.isin(self.meta["superclass"], ["vnc_motor", "cb_motor"])
|
|
active_distances = distances[active]
|
|
max_depth = int(active_distances[active_distances >= 0].max()) if np.any(active_distances >= 0) else 0
|
|
unreachable_active = int(np.count_nonzero(active & (distances < 0)))
|
|
active_fraction = int(active.sum()) / self.n
|
|
mean_rate = total_spikes / (self.n * simulated_s) if simulated_s else 0.0
|
|
supercritical = (
|
|
peak_step_spikes / self.n >= float(cfg["classification"]["supercritical_peak_step_fraction"])
|
|
or active_fraction >= float(cfg["classification"]["supercritical_active_fraction"])
|
|
or mean_rate >= float(cfg["classification"]["supercritical_mean_rate_hz"])
|
|
)
|
|
if supercritical:
|
|
region = "supercritical"
|
|
elif max_depth >= int(cfg["classification"]["propagating_min_depth"]) and int(active_outside.sum()) > 0:
|
|
region = "propagating"
|
|
else:
|
|
region = "subcritical"
|
|
|
|
if log_events and event_path is not None:
|
|
steps_np = np.concatenate(event_steps) if event_steps else np.empty(0, dtype=np.int32)
|
|
neurons_np = np.concatenate(event_neurons) if event_neurons else np.empty(0, dtype=np.int32)
|
|
timestamps_np = steps_np.astype(np.float64) * dt
|
|
runtime_tensors.extend([
|
|
array_record("event_steps", steps_np),
|
|
array_record("event_neuron_indices", neurons_np),
|
|
array_record("event_timestamps_ms", timestamps_np),
|
|
])
|
|
event_path.parent.mkdir(parents=True, exist_ok=True)
|
|
pq.write_table(pa.table({
|
|
"step": pa.array(steps_np, type=pa.int32()),
|
|
"timestamp_ms": pa.array(timestamps_np, type=pa.float64()),
|
|
"neuron_idx": pa.array(neurons_np, type=pa.int32()),
|
|
}), event_path, compression="zstd")
|
|
|
|
category_seconds = dict(timing) if timing is not None else {}
|
|
measured = sum(category_seconds.values())
|
|
category_seconds["python_and_unattributed_overhead"] = max(0.0, wall - measured)
|
|
category_percent = {k: (100.0 * value / wall if wall else 0.0) for k, value in category_seconds.items()}
|
|
peak_allocated = int(torch.cuda.max_memory_allocated(self.device))
|
|
peak_reserved = int(torch.cuda.max_memory_reserved(self.device))
|
|
result = {
|
|
"population": population,
|
|
"population_neurons": int(len(targets_np)),
|
|
"rate_hz": float(rate_hz),
|
|
"seed": int(seed),
|
|
"w_syn_mv": float(w_syn),
|
|
"requested_duration_ms": float(duration_ms),
|
|
"simulated_duration_ms": simulated_ms,
|
|
"requested_steps": steps,
|
|
"completed_steps": completed_steps,
|
|
"terminated_early": completed_steps < steps,
|
|
"termination_reason": termination_reason,
|
|
"stimulation_mode": cfg["stimulation"]["mode"],
|
|
"stimulated_spike_attempts": stimulated_attempts,
|
|
"stimulated_spikes": stimulated_spikes,
|
|
"total_spikes": total_spikes,
|
|
"active_neurons": int(active.sum()),
|
|
"active_neurons_outside_stimulated_population": int(active_outside.sum()),
|
|
"mean_firing_rate_hz": mean_rate,
|
|
"max_firing_rate_hz": float(spike_counts_np.max(initial=0) / simulated_s) if simulated_s else 0.0,
|
|
"spikes_per_biological_second": total_spikes / simulated_s if simulated_s else 0.0,
|
|
"first_hop_active_neurons": int(np.count_nonzero(active & (distances == 1))),
|
|
"second_hop_active_neurons": int(np.count_nonzero(active & (distances == 2))),
|
|
"three_or_more_hop_active_neurons": int(np.count_nonzero(active & (distances >= 3))),
|
|
"propagation_depth": max_depth,
|
|
"unreachable_active_neurons": unreachable_active,
|
|
"excitatory_active_neurons": int(np.count_nonzero(active & (sign_np > 0))),
|
|
"inhibitory_active_neurons": int(np.count_nonzero(active & (sign_np < 0))),
|
|
"zero_sign_active_neurons": int(np.count_nonzero(active & (sign_np == 0))),
|
|
"descending_active_neurons": int(np.count_nonzero(active & descending)),
|
|
"descending_spikes": int(spike_counts_np[descending].sum(dtype=np.int64)),
|
|
"motor_active_neurons": int(np.count_nonzero(active & motor)),
|
|
"motor_spikes": int(spike_counts_np[motor].sum(dtype=np.int64)),
|
|
"activity_by_superclass": category_activity(self.meta["superclass"], active, spike_counts_np),
|
|
"activity_by_class": category_activity(self.meta["class"], active, spike_counts_np),
|
|
"peak_spikes_one_step": peak_step_spikes,
|
|
"peak_synchronous_fraction": peak_step_spikes / self.n,
|
|
"dynamics_region": region,
|
|
"nan_or_inf": not numerical_ok,
|
|
"self_loop_transmission_events": self_loop_transmissions,
|
|
"wall_clock_seconds": wall,
|
|
"realtime_factor_wall_over_biological": wall / simulated_s if simulated_s else None,
|
|
"biological_realtime_speed": simulated_s / wall if wall else None,
|
|
"steps_per_second": completed_steps / wall if wall else None,
|
|
"peak_allocated_vram_bytes": peak_allocated,
|
|
"peak_reserved_vram_bytes": peak_reserved,
|
|
"observed_allocated_growth_bytes": int(max(memory_samples, default=peak_allocated) - min(memory_samples, default=peak_allocated)),
|
|
"memory_leak_suspected": int(max(memory_samples, default=peak_allocated) - min(memory_samples, default=peak_allocated)) > int(cfg["audit"]["memory_growth_tolerance_bytes"]),
|
|
"cpu_process_seconds": cpu_seconds,
|
|
"cpu_process_utilization_percent_of_one_core": 100.0 * cpu_seconds / wall if wall else None,
|
|
"resource_sampling": resource_result,
|
|
"event_logging": log_events,
|
|
"event_file": str(event_path) if event_path else None,
|
|
"runtime_tensors": runtime_tensors,
|
|
"profile_category_seconds": category_seconds,
|
|
"profile_category_percent": category_percent,
|
|
}
|
|
del v, g, refractory, delay_queue, spike_counts, targets, self_loop_sources_t
|
|
torch.cuda.empty_cache()
|
|
return result, spike_counts_np
|
|
|
|
|
|
def run_induced_reference(
|
|
cfg: dict[str, Any], exp: Path, meta: dict[str, Any], indptr: np.ndarray,
|
|
indices: np.ndarray, raw: np.ndarray, device: torch.device,
|
|
) -> dict[str, Any]:
|
|
spec = cfg["cpu_gpu_reference"]
|
|
n = int(cfg["expected_graph"]["neurons"])
|
|
size = int(spec["subgraph_neurons"])
|
|
rng = np.random.default_rng(int(spec["seed"]))
|
|
nodes = np.sort(rng.choice(n, size=size, replace=False)).astype(np.int64)
|
|
inverse = np.full(n, -1, dtype=np.int32)
|
|
inverse[nodes] = np.arange(size, dtype=np.int32)
|
|
rows: list[int] = []
|
|
cols: list[int] = []
|
|
weights: list[int] = []
|
|
for local_src, global_src in enumerate(nodes):
|
|
start, stop = int(indptr[global_src]), int(indptr[global_src + 1])
|
|
local_targets = inverse[indices[start:stop]]
|
|
keep = local_targets >= 0
|
|
count = int(keep.sum())
|
|
rows.extend([local_src] * count)
|
|
cols.extend(local_targets[keep].tolist())
|
|
weights.extend(raw[start:stop][keep].tolist())
|
|
rows_np = np.asarray(rows, dtype=np.int64)
|
|
cols_np = np.asarray(cols, dtype=np.int64)
|
|
raw_np = np.asarray(weights, dtype=np.int64)
|
|
local_signs = meta["sign"][nodes]
|
|
forced = {0: [0, 1, 2, 3], 22: [0, 1, 2, 3], 44: [0, 1, 2, 3]}
|
|
steps = int(spec["steps"])
|
|
cpu = tiny_numpy(local_signs, raw_np, rows_np, cols_np, forced, cfg, steps)
|
|
gpu = tiny_torch(local_signs, raw_np, rows_np, cols_np, forced, cfg, steps, device)
|
|
tolerance = float(spec["absolute_tolerance"])
|
|
v_diff = float(np.max(np.abs(cpu["v"] - gpu["v"])))
|
|
g_diff = float(np.max(np.abs(cpu["g"] - gpu["g"])))
|
|
result = {
|
|
"selection": "uniform random without replacement using fixed seed; induced directed edges",
|
|
"seed": int(spec["seed"]),
|
|
"subgraph_neurons": size,
|
|
"induced_edges": int(len(rows_np)),
|
|
"steps": steps,
|
|
"cpu_events": len(cpu["events"]),
|
|
"gpu_events": len(gpu["events"]),
|
|
"events_exact_match": cpu["events"] == gpu["events"],
|
|
"max_abs_voltage_difference": v_diff,
|
|
"max_abs_synaptic_state_difference": g_diff,
|
|
"absolute_tolerance": tolerance,
|
|
"selected_global_neuron_indices": nodes.tolist(),
|
|
"passed": cpu["events"] == gpu["events"] and v_diff <= tolerance and g_diff <= tolerance,
|
|
}
|
|
write_json(exp / "benchmarks/cpu-gpu-induced-reference.json", result)
|
|
if not result["passed"]:
|
|
raise RuntimeError("CPU/GPU induced reference failed")
|
|
return result
|
|
|
|
|
|
def save_trial(exp: Path, category: str, result: dict[str, Any]) -> Path:
|
|
path = trial_path(exp, category, result["population"], result["rate_hz"], result["seed"], result["w_syn_mv"])
|
|
write_json(path, result)
|
|
return path
|
|
|
|
|
|
def trial_path(exp: Path, category: str, population: str, rate: float, seed: int, weight: float) -> Path:
|
|
safe = f"{population}-rate{float(rate):g}-seed{int(seed)}-w{float(weight):g}.json"
|
|
return exp / f"benchmarks/{category}/{safe}"
|
|
|
|
|
|
def causal_validation(
|
|
cfg: dict[str, Any], exp: Path, event_path: Path, targets: np.ndarray,
|
|
indptr: np.ndarray, indices: np.ndarray, raw: np.ndarray, signs: np.ndarray, w_syn: float,
|
|
) -> dict[str, Any]:
|
|
table = pq.read_table(event_path, columns=["step", "neuron_idx"])
|
|
steps = table["step"].to_numpy(zero_copy_only=False).astype(np.int32, copy=False)
|
|
neurons = table["neuron_idx"].to_numpy(zero_copy_only=False).astype(np.int32, copy=False)
|
|
by_neuron: dict[int, np.ndarray] = {}
|
|
for neuron in np.unique(neurons):
|
|
by_neuron[int(neuron)] = np.sort(steps[neurons == neuron])
|
|
active_set = set(by_neuron)
|
|
target_set = {int(x) for x in targets}
|
|
delay = int(round(float(cfg["model"]["synaptic_delay_ms"]) / float(cfg["model"]["dt_ms"])))
|
|
max_edges = int(cfg["causal_validation"]["max_edges_examined"])
|
|
parent: dict[int, dict[str, Any]] = {}
|
|
inhibitory: list[dict[str, Any]] = []
|
|
edges_examined = 0
|
|
edge_integrity_failures = 0
|
|
sign_failures = 0
|
|
active_mask = np.zeros(len(indptr) - 1, dtype=bool)
|
|
active_mask[np.fromiter(active_set, dtype=np.int64)] = True
|
|
last_spike = np.full(len(indptr) - 1, -1, dtype=np.int32)
|
|
for neuron, neuron_steps in by_neuron.items():
|
|
last_spike[neuron] = int(neuron_steps[-1])
|
|
for pre, pre_steps in by_neuron.items():
|
|
start, stop = int(indptr[pre]), int(indptr[pre + 1])
|
|
pre_sign = int(signs[pre])
|
|
if pre_sign == 0:
|
|
continue
|
|
row_posts = indices[start:stop]
|
|
candidate_offsets = np.flatnonzero(active_mask[row_posts])
|
|
if len(candidate_offsets) == 0:
|
|
continue
|
|
remaining = max_edges - edges_examined
|
|
if remaining <= 0:
|
|
break
|
|
candidate_offsets = candidate_offsets[:remaining]
|
|
candidate_posts = row_posts[candidate_offsets].astype(np.int64, copy=False)
|
|
valid_temporal = last_spike[candidate_posts] >= int(pre_steps[0]) + delay
|
|
candidate_offsets = candidate_offsets[valid_temporal]
|
|
candidate_posts = candidate_posts[valid_temporal]
|
|
edges_examined += len(candidate_posts)
|
|
for offset, post_np in zip(candidate_offsets, candidate_posts):
|
|
edge_id = start + int(offset)
|
|
post = int(post_np)
|
|
if not (start <= edge_id < stop and int(indices[edge_id]) == post):
|
|
edge_integrity_failures += 1
|
|
effective = pre_sign * int(raw[edge_id]) * float(w_syn)
|
|
if (pre_sign > 0 and effective <= 0) or (pre_sign < 0 and effective >= 0):
|
|
sign_failures += 1
|
|
post_steps = by_neuron[post]
|
|
post_pos = int(np.searchsorted(post_steps, int(pre_steps[0]) + delay, side="left"))
|
|
if post_pos >= len(post_steps):
|
|
continue
|
|
post_step = int(post_steps[post_pos])
|
|
pre_pos = int(np.searchsorted(pre_steps, post_step - delay, side="right")) - 1
|
|
if pre_pos < 0:
|
|
continue
|
|
pre_step = int(pre_steps[pre_pos])
|
|
record = {
|
|
"pre_neuron_idx": pre,
|
|
"post_neuron_idx": post,
|
|
"pre_spike_step": pre_step,
|
|
"arrival_step": pre_step + delay,
|
|
"post_spike_step": post_step,
|
|
"edge_id": edge_id,
|
|
"raw_synapse_count": int(raw[edge_id]),
|
|
"presynaptic_sign": pre_sign,
|
|
"effective_weight_mv": effective,
|
|
"delay_steps": delay,
|
|
}
|
|
if pre_sign < 0 and len(inhibitory) < int(cfg["causal_validation"]["requested_inhibitory_transmissions"]):
|
|
inhibitory.append(record)
|
|
if pre_sign > 0 and post not in parent:
|
|
parent[post] = record
|
|
if edges_examined >= max_edges:
|
|
break
|
|
|
|
supported_neurons = set(parent)
|
|
outside_active = active_set - target_set
|
|
unsupported = sorted(outside_active - supported_neurons)
|
|
candidate_keys = [post for post in parent if post not in target_set]
|
|
rng = random.Random(int(cfg["causal_validation"]["seed"]))
|
|
rng.shuffle(candidate_keys)
|
|
requested = int(cfg["causal_validation"]["requested_chains"])
|
|
chains: list[dict[str, Any]] = []
|
|
for key in candidate_keys:
|
|
edges: list[dict[str, Any]] = []
|
|
seen: set[int] = set()
|
|
current = key
|
|
while current in parent and current not in seen and len(edges) < 64:
|
|
seen.add(current)
|
|
edge = parent[current]
|
|
edges.append(edge)
|
|
next_neuron = int(edge["pre_neuron_idx"])
|
|
previous = parent.get(next_neuron)
|
|
if previous is not None and int(previous["post_spike_step"]) > int(edge["pre_spike_step"]):
|
|
break
|
|
current = next_neuron
|
|
if current in target_set:
|
|
break
|
|
if edges:
|
|
edges.reverse()
|
|
valid = all(
|
|
int(indices[e["edge_id"]]) == e["post_neuron_idx"]
|
|
and e["arrival_step"] - e["pre_spike_step"] == delay
|
|
and e["presynaptic_sign"] == int(signs[e["pre_neuron_idx"]])
|
|
for e in edges
|
|
)
|
|
chains.append({"edges": edges, "length": len(edges), "valid": valid})
|
|
if len(chains) >= requested:
|
|
break
|
|
inhibitory = inhibitory[: int(cfg["causal_validation"]["requested_inhibitory_transmissions"])]
|
|
result = {
|
|
"event_file": str(event_path),
|
|
"active_neurons": len(active_set),
|
|
"active_outside_stimulated": len(outside_active),
|
|
"edges_examined": edges_examined,
|
|
"edge_scan_truncated": edges_examined >= max_edges,
|
|
"requested_chains": requested,
|
|
"validated_chains": len(chains),
|
|
"all_chain_edges_valid": all(x["valid"] for x in chains),
|
|
"maximum_validated_chain_length": max((x["length"] for x in chains), default=0),
|
|
"unsupported_active_neuron_count": len(unsupported),
|
|
"unsupported_active_neuron_indices": unsupported,
|
|
"edge_integrity_failures": edge_integrity_failures,
|
|
"sign_failures": sign_failures,
|
|
"requested_inhibitory_transmissions": int(cfg["causal_validation"]["requested_inhibitory_transmissions"]),
|
|
"validated_inhibitory_transmissions": len(inhibitory),
|
|
"all_inhibitory_weights_negative": all(x["effective_weight_mv"] < 0 for x in inhibitory),
|
|
"delay_steps": delay,
|
|
"delay_ms": delay * float(cfg["model"]["dt_ms"]),
|
|
"chains": chains,
|
|
"inhibitory_transmissions": inhibitory,
|
|
"no_connectionless_transmission": edge_integrity_failures == 0,
|
|
"passed": (
|
|
len(chains) >= requested and all(x["valid"] for x in chains)
|
|
and edge_integrity_failures == 0 and sign_failures == 0
|
|
and len(unsupported) == 0
|
|
and len(inhibitory) >= int(cfg["causal_validation"]["requested_inhibitory_transmissions"])
|
|
and all(x["effective_weight_mv"] < 0 for x in inhibitory)
|
|
),
|
|
}
|
|
write_json(exp / "propagation/causal-validation.json", result)
|
|
return result
|
|
|
|
|
|
def compact_trial(result: dict[str, Any]) -> dict[str, Any]:
|
|
keys = [
|
|
"population", "population_neurons", "rate_hz", "seed", "w_syn_mv",
|
|
"requested_duration_ms", "simulated_duration_ms", "terminated_early", "termination_reason",
|
|
"stimulated_spikes", "total_spikes", "active_neurons",
|
|
"active_neurons_outside_stimulated_population", "mean_firing_rate_hz", "max_firing_rate_hz",
|
|
"first_hop_active_neurons", "second_hop_active_neurons", "three_or_more_hop_active_neurons",
|
|
"propagation_depth", "unreachable_active_neurons", "excitatory_active_neurons",
|
|
"inhibitory_active_neurons", "zero_sign_active_neurons", "descending_active_neurons",
|
|
"descending_spikes", "motor_active_neurons", "motor_spikes", "dynamics_region", "nan_or_inf",
|
|
"wall_clock_seconds", "realtime_factor_wall_over_biological", "steps_per_second",
|
|
"peak_allocated_vram_bytes", "peak_reserved_vram_bytes", "memory_leak_suspected",
|
|
]
|
|
return {key: result[key] for key in keys}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", required=True, type=Path)
|
|
args = parser.parse_args()
|
|
cfg = yaml.safe_load(args.config.read_text(encoding="utf-8"))
|
|
exp = Path(cfg["paths"]["experiment_root"]).resolve()
|
|
for sub in ("configs", "benchmarks", "propagation", "profiler", "events", "audit", "manifests"):
|
|
(exp / sub).mkdir(parents=True, exist_ok=True)
|
|
setup_logging(exp / "audit/phase4b.log")
|
|
LOG.info("Phase 4B start config=%s", args.config)
|
|
random.seed(int(cfg["reproducibility"]["global_seed"]))
|
|
np.random.seed(int(cfg["reproducibility"]["global_seed"]))
|
|
torch.manual_seed(int(cfg["reproducibility"]["global_seed"]))
|
|
paths = exact_input_paths(cfg)
|
|
hashes_before = {name: sha256(path) for name, path in paths.items()}
|
|
indptr, indices, raw = validate_graph(cfg, paths)
|
|
meta = metadata_arrays(cfg, paths["metadata"])
|
|
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA unavailable")
|
|
if torch.cuda.device_count() != 1:
|
|
raise RuntimeError("exactly one visible GPU is required; mixed multi-GPU is forbidden")
|
|
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"required RTX 4090, got {props.name}")
|
|
|
|
masks = population_masks(cfg, meta)
|
|
population_audit, distances = audit_populations(cfg, exp, meta, masks, indptr, indices, raw)
|
|
unit_tests = run_dynamics_unit_tests(cfg, exp, device)
|
|
induced_reference = run_induced_reference(cfg, exp, meta, indptr, indices, raw, device)
|
|
runtime = FullGraphRuntime(cfg, paths, meta, indptr, indices, raw, device)
|
|
baseline_w = float(cfg["model"]["w_syn_mv_per_anatomical_synapse"])
|
|
runtime.set_weight(baseline_w)
|
|
assert runtime.effective is not None
|
|
sign_path = exp / "audit/synaptic-sign.int8.npy"
|
|
np.save(sign_path, meta["sign"], allow_pickle=False)
|
|
effective_path = exp / "audit/effective-weights-0.275.float32.npy"
|
|
np.save(effective_path, runtime.effective.cpu().numpy(), allow_pickle=False)
|
|
sign_classes = Counter(meta["nt_class"].tolist())
|
|
sign_audit = {
|
|
"source_field": "neurotransmitter",
|
|
"classification_counts": dict(sorted(sign_classes.items())),
|
|
"sign_counts": {str(k): int(v) for k, v in sorted(Counter(meta["sign"].tolist()).items())},
|
|
"unknown_count": int(sign_classes.get("UNKNOWN", 0)),
|
|
"unclear_count": int(sign_classes.get("UNCLEAR", 0)),
|
|
"formula": "effective_weight(j->i) = sign(j) * raw_synapse_count(j->i) * W_syn",
|
|
"baseline_w_syn_mv": baseline_w,
|
|
"transform": "strict linear; no sqrt/log1p/normalization/degree normalization/threshold",
|
|
"raw_weights_path": str(paths["raw_weights"]),
|
|
"raw_weights_sha256": hashes_before["raw_weights"],
|
|
"raw_weights_read_only": True,
|
|
"effective_weights_path": str(effective_path),
|
|
"effective_weights_sha256": sha256(effective_path),
|
|
}
|
|
write_json(exp / "audit/sign-and-effective-weight.json", sign_audit)
|
|
|
|
baseline_results: list[dict[str, Any]] = []
|
|
total_baseline = len(masks) * len(cfg["baseline_sweep"]["rates_hz"]) * len(cfg["baseline_sweep"]["seeds"])
|
|
trial_number = 0
|
|
for population, mask in masks.items():
|
|
targets = np.flatnonzero(mask).astype(np.int64)
|
|
for rate in cfg["baseline_sweep"]["rates_hz"]:
|
|
for seed in cfg["baseline_sweep"]["seeds"]:
|
|
trial_number += 1
|
|
LOG.info("baseline trial %d/%d population=%s rate=%s seed=%s", trial_number, total_baseline, population, rate, seed)
|
|
checkpoint = trial_path(exp, "baseline-trials", population, float(rate), int(seed), baseline_w)
|
|
if checkpoint.is_file():
|
|
result = json.loads(checkpoint.read_text(encoding="utf-8"))
|
|
LOG.info("checkpoint reused %s", checkpoint.name)
|
|
else:
|
|
result, _ = runtime.run(
|
|
population, targets, distances[population], float(rate), int(seed),
|
|
float(cfg["baseline_sweep"]["duration_ms"]), baseline_w, False,
|
|
)
|
|
save_trial(exp, "baseline-trials", result)
|
|
baseline_results.append(result)
|
|
write_json(exp / "benchmarks/baseline-summary.json", {
|
|
"trial_count": len(baseline_results),
|
|
"expected_trial_count": total_baseline,
|
|
"trials": [compact_trial(x) for x in baseline_results],
|
|
"region_counts": dict(Counter(x["dynamics_region"] for x in baseline_results)),
|
|
"early_termination_count": sum(bool(x["terminated_early"]) for x in baseline_results),
|
|
"maximum_propagation_depth": max(x["propagation_depth"] for x in baseline_results),
|
|
})
|
|
|
|
weight_results: list[dict[str, Any]] = []
|
|
weight_pop = str(cfg["weight_sweep"]["population"])
|
|
weight_targets = np.flatnonzero(masks[weight_pop]).astype(np.int64)
|
|
for weight in cfg["weight_sweep"]["values_mv"]:
|
|
LOG.info("weight sweep W_syn=%s", weight)
|
|
checkpoint = trial_path(
|
|
exp, "weight-sweep", weight_pop, float(cfg["weight_sweep"]["rate_hz"]),
|
|
int(cfg["weight_sweep"]["seed"]), float(weight),
|
|
)
|
|
if checkpoint.is_file():
|
|
result = json.loads(checkpoint.read_text(encoding="utf-8"))
|
|
LOG.info("checkpoint reused %s", checkpoint.name)
|
|
else:
|
|
result, _ = runtime.run(
|
|
weight_pop, weight_targets, distances[weight_pop], float(cfg["weight_sweep"]["rate_hz"]),
|
|
int(cfg["weight_sweep"]["seed"]), float(cfg["weight_sweep"]["duration_ms"]), float(weight), False,
|
|
)
|
|
save_trial(exp, "weight-sweep", result)
|
|
weight_results.append(result)
|
|
write_json(exp / "benchmarks/weight-sweep-summary.json", {
|
|
"controlled_population": weight_pop,
|
|
"controlled_rate_hz": cfg["weight_sweep"]["rate_hz"],
|
|
"controlled_seed": cfg["weight_sweep"]["seed"],
|
|
"classification_thresholds_fixed_before_run": cfg["classification"],
|
|
"trials": [compact_trial(x) for x in weight_results],
|
|
"regions": {str(x["w_syn_mv"]): x["dynamics_region"] for x in weight_results},
|
|
})
|
|
|
|
perf_pop = str(cfg["performance"]["population"])
|
|
perf_targets = np.flatnonzero(masks[perf_pop]).astype(np.int64)
|
|
perf_common = dict(
|
|
population=perf_pop, targets_np=perf_targets, distances=distances[perf_pop],
|
|
rate_hz=float(cfg["performance"]["rate_hz"]), seed=int(cfg["performance"]["seed"]),
|
|
duration_ms=float(cfg["performance"]["duration_ms"]), w_syn=baseline_w,
|
|
monitor_resources=True,
|
|
)
|
|
LOG.info("performance benchmark logging disabled")
|
|
perf_disabled, counts_disabled = runtime.run(**perf_common, log_events=False)
|
|
LOG.info("performance benchmark full event logging")
|
|
perf_event_path = exp / "events/performance-full.events.parquet"
|
|
perf_full, counts_full = runtime.run(**perf_common, log_events=True, event_path=perf_event_path)
|
|
event_consistency = bool(np.array_equal(counts_disabled, counts_full))
|
|
performance = {
|
|
"logging_disabled": compact_trial(perf_disabled) | {
|
|
"resource_sampling": perf_disabled["resource_sampling"],
|
|
"cpu_process_utilization_percent_of_one_core": perf_disabled["cpu_process_utilization_percent_of_one_core"],
|
|
},
|
|
"full_event_logging": compact_trial(perf_full) | {
|
|
"resource_sampling": perf_full["resource_sampling"],
|
|
"cpu_process_utilization_percent_of_one_core": perf_full["cpu_process_utilization_percent_of_one_core"],
|
|
"event_file": perf_full["event_file"],
|
|
},
|
|
"spike_count_arrays_exact_match": event_consistency,
|
|
"logging_disabled_spike_count_sha256": hashlib.sha256(counts_disabled.tobytes()).hexdigest(),
|
|
"full_logging_spike_count_sha256": hashlib.sha256(counts_full.tobytes()).hexdigest(),
|
|
}
|
|
write_json(exp / "benchmarks/performance-logging-comparison.json", performance)
|
|
if not event_consistency:
|
|
raise RuntimeError("logging changed spike events")
|
|
|
|
profile_pop = str(cfg["profiler"]["population"])
|
|
profile_targets = np.flatnonzero(masks[profile_pop]).astype(np.int64)
|
|
LOG.info("synchronized category profiler")
|
|
profile_result, profile_counts_disabled = runtime.run(
|
|
profile_pop, profile_targets, distances[profile_pop], float(cfg["profiler"]["rate_hz"]),
|
|
int(cfg["profiler"]["seed"]), float(cfg["profiler"]["duration_ms"]), baseline_w,
|
|
False, profile_categories=True,
|
|
)
|
|
write_json(exp / "profiler/category-profile.json", profile_result)
|
|
LOG.info("synchronized category profiler with full event logging")
|
|
profiler_event_path = exp / "events/profiler-full.events.parquet"
|
|
profile_full_result, profile_counts_full = runtime.run(
|
|
profile_pop, profile_targets, distances[profile_pop], float(cfg["profiler"]["rate_hz"]),
|
|
int(cfg["profiler"]["seed"]), float(cfg["profiler"]["duration_ms"]), baseline_w,
|
|
True, event_path=profiler_event_path, profile_categories=True,
|
|
)
|
|
profile_event_consistency = bool(np.array_equal(profile_counts_disabled, profile_counts_full))
|
|
write_json(exp / "profiler/category-profile-full-logging.json", profile_full_result)
|
|
if not profile_event_consistency:
|
|
raise RuntimeError("profile logging changed spike events")
|
|
|
|
eligible = [x for x in baseline_results if not x["terminated_early"] and x["active_neurons_outside_stimulated_population"] > 0]
|
|
candidate_pool = eligible if eligible else baseline_results
|
|
selected = max(candidate_pool, key=lambda x: (x["propagation_depth"], x["active_neurons_outside_stimulated_population"], -x["rate_hz"]))
|
|
causal_pop = selected["population"]
|
|
causal_targets = np.flatnonzero(masks[causal_pop]).astype(np.int64)
|
|
causal_event_path = exp / "events/causal-trial.events.parquet"
|
|
causal_trial_path = exp / "propagation/causal-trial.json"
|
|
causal_audit_path = exp / "propagation/causal-validation.json"
|
|
if causal_event_path.is_file() and causal_trial_path.is_file() and causal_audit_path.is_file():
|
|
LOG.info("causal trial and validation checkpoints reused")
|
|
causal_trial = json.loads(causal_trial_path.read_text(encoding="utf-8"))
|
|
causal = json.loads(causal_audit_path.read_text(encoding="utf-8"))
|
|
else:
|
|
LOG.info("causal trial population=%s rate=%s seed=%s", causal_pop, selected["rate_hz"], selected["seed"])
|
|
causal_trial, _ = runtime.run(
|
|
causal_pop, causal_targets, distances[causal_pop], float(selected["rate_hz"]), int(selected["seed"]),
|
|
float(cfg["baseline_sweep"]["duration_ms"]), baseline_w, True, event_path=causal_event_path,
|
|
)
|
|
write_json(causal_trial_path, causal_trial)
|
|
causal = causal_validation(cfg, exp, causal_event_path, causal_targets, indptr, indices, raw, meta["sign"], baseline_w)
|
|
|
|
tensor_records = perf_disabled["runtime_tensors"]
|
|
gpu_tensor_bytes = sum(x["bytes"] for x in tensor_records if x["device"].startswith("cuda"))
|
|
peak_allocated = perf_disabled["peak_allocated_vram_bytes"]
|
|
runtime_memory = {
|
|
"tensors_and_arrays": tensor_records,
|
|
"full_logging_event_arrays": [x for x in perf_full["runtime_tensors"] if x["name"].startswith("event_")],
|
|
"listed_gpu_tensor_bytes": gpu_tensor_bytes,
|
|
"peak_allocated_vram_bytes": peak_allocated,
|
|
"peak_reserved_vram_bytes": perf_disabled["peak_reserved_vram_bytes"],
|
|
"peak_minus_listed_tensor_bytes": peak_allocated - gpu_tensor_bytes,
|
|
"explanation": "peak minus listed persistent/state tensors is transient fan-out indexing, spike detection, Poisson, and scatter workspace; no padding copies were created",
|
|
"phase4a_234_7_mib_explained": abs(peak_allocated / (1024 ** 2) - 234.7) < 32.0,
|
|
}
|
|
write_json(exp / "audit/runtime-tensor-memory.json", runtime_memory)
|
|
|
|
hashes_after = {name: sha256(path) for name, path in paths.items()}
|
|
if hashes_before != hashes_after:
|
|
raise RuntimeError("Phase 3B input changed during Phase 4B")
|
|
summary = {
|
|
"status": "PASS_WITH_REPORTED_ANOMALIES" if not causal["passed"] else "PASS",
|
|
"graph": cfg["expected_graph"],
|
|
"device": {"name": props.name, "total_vram_bytes": int(props.total_memory), "visible_gpu_count": 1},
|
|
"strict_shiu_baseline": {
|
|
"parameters": cfg["model"],
|
|
"formula": sign_audit["formula"],
|
|
"forbidden_transforms_used": False,
|
|
"implementation_semantics": "dv/dt=(g-(v-V_rest))/tau_m; dg/dt=-g/tau_syn; forward Euler dt=0.1 ms; delayed g increment; v reset; g continues during refractory",
|
|
},
|
|
"population_audit": {name: {k: v for k, v in rec.items() if k not in ("body_ids", "neuron_indices")} for name, rec in population_audit.items()},
|
|
"baseline": {
|
|
"trials": len(baseline_results),
|
|
"region_counts": dict(Counter(x["dynamics_region"] for x in baseline_results)),
|
|
"early_termination_count": sum(bool(x["terminated_early"]) for x in baseline_results),
|
|
"maximum_propagation_depth": max(x["propagation_depth"] for x in baseline_results),
|
|
"maximum_descending_active": max(x["descending_active_neurons"] for x in baseline_results),
|
|
"maximum_motor_active": max(x["motor_active_neurons"] for x in baseline_results),
|
|
},
|
|
"weight_sweep": {str(x["w_syn_mv"]): compact_trial(x) for x in weight_results},
|
|
"unit_tests": unit_tests,
|
|
"cpu_gpu_induced_reference": induced_reference,
|
|
"causal_validation": {k: v for k, v in causal.items() if k not in ("chains", "inhibitory_transmissions", "unsupported_active_neuron_indices")},
|
|
"performance": performance,
|
|
"profiler": {
|
|
"logging_disabled_wall_clock_seconds": profile_result["wall_clock_seconds"],
|
|
"full_logging_wall_clock_seconds": profile_full_result["wall_clock_seconds"],
|
|
"simulated_duration_ms": profile_result["simulated_duration_ms"],
|
|
"logging_disabled_category_seconds": profile_result["profile_category_seconds"],
|
|
"logging_disabled_category_percent": profile_result["profile_category_percent"],
|
|
"full_logging_category_seconds": profile_full_result["profile_category_seconds"],
|
|
"full_logging_category_percent": profile_full_result["profile_category_percent"],
|
|
"spike_count_arrays_exact_match": profile_event_consistency,
|
|
},
|
|
"runtime_memory": runtime_memory,
|
|
"input_hashes_unchanged": True,
|
|
"no_training_game_web_or_frontend": True,
|
|
}
|
|
write_json(exp / "benchmarks/summary.json", summary)
|
|
|
|
LOG.info("PHASE4B_COMPLETE status=%s", summary["status"])
|
|
for handler in logging.getLogger().handlers:
|
|
handler.flush()
|
|
controls = {exp / "manifests/manifest.json", exp / "manifests/checksums.sha256"}
|
|
outputs = sorted(path for path in exp.rglob("*") if path.is_file() and path not in controls)
|
|
manifest = {
|
|
"phase": "4B",
|
|
"created_unix_time": time.time(),
|
|
"code": {"path": str(Path(__file__).resolve()), "sha256": sha256(Path(__file__).resolve())},
|
|
"config": {"path": str(args.config.resolve()), "sha256": sha256(args.config)},
|
|
"software": {"python": platform.python_version(), "torch": torch.__version__, "numpy": np.__version__, "pyarrow": pa.__version__},
|
|
"inputs": {name: {"path": str(paths[name]), "sha256": digest} for name, digest in hashes_after.items()},
|
|
"outputs": [{"path": str(path), "bytes": path.stat().st_size, "sha256": sha256(path)} for path in outputs],
|
|
}
|
|
manifest_path = exp / "manifests/manifest.json"
|
|
write_json(manifest_path, manifest)
|
|
targets = outputs + [manifest_path]
|
|
(exp / "manifests/checksums.sha256").write_text(
|
|
"\n".join(f"{sha256(path)} {path.relative_to(exp)}" for path in targets) + "\n", encoding="utf-8"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|