Digital Drosophila Archive
646 lines
30 KiB
Python
646 lines
30 KiB
Python
#!/usr/bin/env python3
|
|
"""Auditable Phase 4A whole-connectome LIF baseline for MaleCNS v1.0."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import platform
|
|
import random
|
|
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
|
|
|
|
|
|
LOG = logging.getLogger("phase4a")
|
|
|
|
|
|
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, obj: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(obj, 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 load_config(path: Path) -> dict[str, Any]:
|
|
cfg = yaml.safe_load(path.read_text(encoding="utf-8"))
|
|
if not isinstance(cfg, dict):
|
|
raise ValueError("configuration root must be a mapping")
|
|
return cfg
|
|
|
|
|
|
def require_exact_inputs(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"connectome_root must be exactly {allowed}, got {root}")
|
|
paths = {
|
|
"indptr": root / "gpu/csr-indptr.int64.npy",
|
|
"indices": root / "gpu/csr-indices.int32.npy",
|
|
"weights": root / "gpu/csr-weights.int64.npy",
|
|
"metadata": root / "neuron-index/unified-metadata.parquet",
|
|
"csr_metadata": root / "gpu/csr-metadata.json",
|
|
}
|
|
missing = [str(p) for p in paths.values() if not p.is_file()]
|
|
if missing:
|
|
raise FileNotFoundError(f"missing Phase 3B inputs: {missing}")
|
|
return paths
|
|
|
|
|
|
def classify_nt(value: Any, cfg: dict[str, Any]) -> tuple[str, int]:
|
|
sign_cfg = cfg["sign_policy"]
|
|
if value is None or (isinstance(value, float) and math.isnan(value)):
|
|
return "UNKNOWN", int(sign_cfg["unknown_sign"])
|
|
key = str(value).strip().lower()
|
|
if not key:
|
|
return "UNKNOWN", int(sign_cfg["unknown_sign"])
|
|
if key == "unclear":
|
|
return "UNCLEAR", int(sign_cfg["unclear_sign"])
|
|
mapping = sign_cfg["neurotransmitter_sign"]
|
|
if key not in mapping:
|
|
if sign_cfg.get("unmapped_policy") == "error":
|
|
raise RuntimeError(f"unmapped neurotransmitter value: {value!r}")
|
|
return f"UNMAPPED:{key}", int(sign_cfg["unmapped_sign"])
|
|
sign = int(mapping[key])
|
|
if sign not in (-1, 0, 1):
|
|
raise ValueError(f"invalid sign for {key}: {sign}")
|
|
return key, sign
|
|
|
|
|
|
def transform_weights(raw: np.ndarray, transform: str, scale: float) -> np.ndarray:
|
|
x = raw.astype(np.float64, copy=False)
|
|
if transform == "linear":
|
|
y = x
|
|
elif transform == "sqrt":
|
|
y = np.sqrt(x)
|
|
elif transform == "log1p":
|
|
y = np.log1p(x)
|
|
else:
|
|
raise ValueError(f"unsupported weight transform: {transform}")
|
|
return (float(scale) * y).astype(np.float32)
|
|
|
|
|
|
def prepare_derived(cfg: dict[str, Any], inputs: dict[str, Path], exp: Path) -> dict[str, Any]:
|
|
expected = cfg["expected_graph"]
|
|
n = int(expected["neurons"])
|
|
nnz = int(expected["directed_edges"])
|
|
expected_sum = int(expected["synaptic_contact_weight_sum"])
|
|
indptr = np.load(inputs["indptr"], mmap_mode="r")
|
|
indices = np.load(inputs["indices"], mmap_mode="r")
|
|
raw = np.load(inputs["weights"], mmap_mode="r")
|
|
if indptr.dtype != np.int64 or indices.dtype != np.int32 or raw.dtype != np.int64:
|
|
raise RuntimeError(f"unexpected CSR dtypes: {indptr.dtype}, {indices.dtype}, {raw.dtype}")
|
|
if indptr.shape != (n + 1,) or indices.shape != (nnz,) or raw.shape != (nnz,):
|
|
raise RuntimeError(f"unexpected CSR shapes: {indptr.shape}, {indices.shape}, {raw.shape}")
|
|
if int(indptr[-1]) != nnz or int(raw.sum(dtype=np.int64)) != expected_sum:
|
|
raise RuntimeError("Phase 3B graph invariant failed")
|
|
if int(indices.min()) < 0 or int(indices.max()) >= n or int(raw.min()) <= 0:
|
|
raise RuntimeError("CSR index or raw-weight range invalid")
|
|
|
|
table = pq.read_table(inputs["metadata"], columns=["neuron_idx", "body_id", "annotation_type", "neurotransmitter"])
|
|
idx = table["neuron_idx"].to_numpy(zero_copy_only=False)
|
|
if len(table) != n or idx.dtype != np.int32 or not np.array_equal(idx, np.arange(n, dtype=np.int32)):
|
|
raise RuntimeError("metadata neuron_idx is not the canonical contiguous int32 mapping")
|
|
nt_values = table["neurotransmitter"].to_pylist()
|
|
classes: list[str] = []
|
|
signs = np.empty(n, dtype=np.int8)
|
|
for i, value in enumerate(nt_values):
|
|
cls, sign = classify_nt(value, cfg)
|
|
classes.append(cls)
|
|
signs[i] = sign
|
|
|
|
derived = exp / "outputs/derived"
|
|
derived.mkdir(parents=True, exist_ok=True)
|
|
sign_path = derived / "synaptic-sign.int8.npy"
|
|
np.save(sign_path, signs, allow_pickle=False)
|
|
|
|
transform = str(cfg["weights"]["transform"])
|
|
scale = float(cfg["weights"]["scale"])
|
|
effective_path = derived / "effective-weights.float32.npy"
|
|
effective = np.lib.format.open_memmap(effective_path, mode="w+", dtype=np.float32, shape=(nnz,))
|
|
chunk = int(cfg["weights"].get("chunk_edges", 2_000_000))
|
|
for start in range(0, nnz, chunk):
|
|
stop = min(nnz, start + chunk)
|
|
values = transform_weights(raw[start:stop], transform, scale)
|
|
src = np.searchsorted(indptr, np.arange(start, stop, dtype=np.int64), side="right") - 1
|
|
effective[start:stop] = values * signs[src].astype(np.float32, copy=False)
|
|
effective.flush()
|
|
del effective
|
|
|
|
nt_counts = Counter(classes)
|
|
raw_nt_counts = Counter("UNKNOWN" if v is None else str(v).strip().lower() or "UNKNOWN" for v in nt_values)
|
|
sign_counts = {str(k): int(v) for k, v in sorted(Counter(signs.tolist()).items())}
|
|
result = {
|
|
"neurons": n,
|
|
"directed_edges": nnz,
|
|
"raw_synaptic_contact_weight_sum": expected_sum,
|
|
"raw_weight_path": str(inputs["weights"]),
|
|
"raw_weight_sha256": sha256(inputs["weights"]),
|
|
"raw_weight_modified": False,
|
|
"effective_weight_path": str(effective_path),
|
|
"effective_weight_sha256": sha256(effective_path),
|
|
"effective_weight_transform": transform,
|
|
"effective_weight_scale": scale,
|
|
"neurotransmitter_source_field": "neurotransmitter",
|
|
"neurotransmitter_counts": dict(sorted(raw_nt_counts.items())),
|
|
"classification_counts": dict(sorted(nt_counts.items())),
|
|
"sign_counts": sign_counts,
|
|
"unknown_metadata_count": int(nt_counts.get("UNKNOWN", 0)),
|
|
"unclear_metadata_count": int(nt_counts.get("UNCLEAR", 0)),
|
|
"histamine_count": int(nt_counts.get("histamine", 0)),
|
|
"histamine_configured_sign": int(cfg["sign_policy"]["neurotransmitter_sign"]["histamine"]),
|
|
}
|
|
write_json(derived / "sign-and-weight-audit.json", result)
|
|
return result
|
|
|
|
|
|
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_pos = torch.searchsorted(cumulative, seq, right=True)
|
|
previous = cumulative - counts
|
|
return starts[row_pos] + seq - previous[row_pos]
|
|
|
|
|
|
def stimulus_descriptor(cfg: dict[str, Any], table: pa.Table, test: dict[str, Any]) -> dict[str, Any]:
|
|
stim_name = test["stimulation"]
|
|
spec = cfg["stimulations"][stim_name]
|
|
n = len(table)
|
|
if spec["kind"] == "single_neuron_pulse":
|
|
targets = np.array([int(spec["neuron_idx"])], dtype=np.int64)
|
|
elif spec["kind"] == "known_type_pulse":
|
|
wanted = str(spec["type_name"])
|
|
values = table["annotation_type"].to_pylist()
|
|
targets = np.array([i for i, value in enumerate(values) if value == wanted], dtype=np.int64)
|
|
max_neurons = int(spec.get("max_neurons", len(targets)))
|
|
targets = targets[:max_neurons]
|
|
if targets.size == 0:
|
|
raise RuntimeError(f"known type not found: {wanted}")
|
|
elif spec["kind"] == "sparse_poisson":
|
|
rng = np.random.default_rng(int(test["seed"]))
|
|
count = max(1, round(n * float(spec["target_fraction"])))
|
|
targets = np.sort(rng.choice(n, size=count, replace=False)).astype(np.int64)
|
|
else:
|
|
raise ValueError(f"unsupported stimulation kind: {spec['kind']}")
|
|
return {"name": stim_name, "spec": spec, "targets": targets}
|
|
|
|
|
|
def save_events(exp: Path, run_name: str, times: np.ndarray, neurons: np.ndarray, duration_ms: float, window_ms: float) -> None:
|
|
out = exp / "outputs/spikes"
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
event_table = pa.table({
|
|
"timestamp_ms": pa.array(times, type=pa.float64()),
|
|
"neuron_idx": pa.array(neurons, type=pa.int32()),
|
|
})
|
|
pq.write_table(event_table, out / f"{run_name}.events.parquet", compression="zstd")
|
|
starts = np.arange(0.0, duration_ms, window_ms, dtype=np.float64)
|
|
counts, _ = np.histogram(times, bins=np.append(starts, duration_ms))
|
|
rates = counts.astype(np.float64) / ((np.minimum(starts + window_ms, duration_ms) - starts) / 1000.0)
|
|
rate_table = pa.table({
|
|
"window_start_ms": starts,
|
|
"window_end_ms": np.minimum(starts + window_ms, duration_ms),
|
|
"spike_count": counts.astype(np.int64),
|
|
"population_spikes_per_second": rates,
|
|
})
|
|
pq.write_table(rate_table, out / f"{run_name}.window-rates.parquet", compression="zstd")
|
|
|
|
|
|
def run_gpu_smoke(
|
|
cfg: dict[str, Any], inputs: dict[str, Path], exp: Path, test: dict[str, Any],
|
|
graph: dict[str, torch.Tensor], metadata: pa.Table, self_loop_weight: torch.Tensor,
|
|
) -> dict[str, Any]:
|
|
device = graph["indptr"].device
|
|
n = int(cfg["expected_graph"]["neurons"])
|
|
dt = float(cfg["model"]["dt_ms"])
|
|
duration = float(test["duration_ms"])
|
|
steps = int(round(duration / dt))
|
|
if not math.isclose(steps * dt, duration, abs_tol=1e-9):
|
|
raise ValueError("duration must be divisible by dt")
|
|
delay_steps = int(round(float(cfg["model"]["synaptic_delay_ms"]) / dt))
|
|
queue_len = delay_steps + 1
|
|
rest = float(cfg["model"]["resting_potential_mv"])
|
|
reset = float(cfg["model"]["reset_potential_mv"])
|
|
threshold = float(cfg["model"]["threshold_mv"])
|
|
alpha = dt / float(cfg["model"]["membrane_tau_ms"])
|
|
syn_decay = math.exp(-dt / float(cfg["model"]["synaptic_tau_ms"]))
|
|
refractory_steps = int(round(float(cfg["model"]["refractory_ms"]) / dt))
|
|
stim = stimulus_descriptor(cfg, metadata, test)
|
|
spec = stim["spec"]
|
|
targets_cpu = stim["targets"]
|
|
targets = torch.as_tensor(targets_cpu, dtype=torch.int64, device=device)
|
|
generator = torch.Generator(device=device)
|
|
generator.manual_seed(int(test["seed"]))
|
|
|
|
v = torch.full((n,), rest, dtype=torch.float32, device=device)
|
|
syn = torch.zeros(n, dtype=torch.float32, device=device)
|
|
refractory = torch.zeros(n, dtype=torch.int32, device=device)
|
|
queue = torch.zeros((queue_len, n), dtype=torch.float32, device=device)
|
|
event_t: list[np.ndarray] = []
|
|
event_n: list[np.ndarray] = []
|
|
peak_step_spikes = 0
|
|
self_loop_events = 0
|
|
self_loop_effective_sum = 0.0
|
|
numerical_ok = True
|
|
memory_samples: list[int] = []
|
|
torch.cuda.synchronize(device)
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
persistent_bytes = torch.cuda.memory_allocated(device)
|
|
started = time.perf_counter()
|
|
|
|
for step in range(steps):
|
|
slot = step % queue_len
|
|
syn.mul_(syn_decay).add_(queue[slot])
|
|
queue[slot].zero_()
|
|
ext = torch.zeros(n, dtype=torch.float32, device=device)
|
|
kind = spec["kind"]
|
|
t_ms = step * dt
|
|
if kind in ("single_neuron_pulse", "known_type_pulse"):
|
|
if float(spec["start_ms"]) <= t_ms < float(spec["start_ms"]) + float(spec["duration_ms"]):
|
|
ext[targets] = float(spec["amplitude_mv_equivalent"])
|
|
elif kind == "sparse_poisson":
|
|
probability = float(spec["rate_hz"]) * dt / 1000.0
|
|
chosen = torch.rand(targets.numel(), generator=generator, device=device) < probability
|
|
if bool(chosen.any().item()):
|
|
ext[targets[chosen]] = float(spec["event_amplitude_mv_equivalent"])
|
|
|
|
available = refractory == 0
|
|
v[available] += alpha * (rest - v[available] + syn[available] + ext[available])
|
|
refractory.clamp_min_(1).sub_(1)
|
|
spikes = available & (v >= threshold)
|
|
sources = torch.nonzero(spikes, as_tuple=False).flatten()
|
|
count = int(sources.numel())
|
|
if count:
|
|
peak_step_spikes = max(peak_step_spikes, count)
|
|
src_np = sources.detach().cpu().numpy().astype(np.int32, copy=False)
|
|
event_t.append(np.full(count, (step + 1) * dt, dtype=np.float64))
|
|
event_n.append(src_np.copy())
|
|
v[sources] = reset
|
|
refractory[sources] = refractory_steps
|
|
loop_values = self_loop_weight[sources]
|
|
loop_active = loop_values != 0
|
|
self_loop_events += int(loop_active.sum().item())
|
|
self_loop_effective_sum += float(loop_values.sum().item())
|
|
edge_ids = expand_rows(graph["indptr"], sources)
|
|
if edge_ids.numel():
|
|
posts = graph["indices"][edge_ids].to(torch.int64)
|
|
arrivals = graph["effective"][edge_ids]
|
|
future = (step + delay_steps) % queue_len
|
|
queue[future].scatter_add_(0, posts, arrivals)
|
|
|
|
if step % int(cfg["validation"]["finite_check_interval_steps"]) == 0 or step == steps - 1:
|
|
finite = torch.isfinite(v).all() & torch.isfinite(syn).all() & torch.isfinite(queue).all()
|
|
if not bool(finite.item()):
|
|
numerical_ok = False
|
|
break
|
|
memory_samples.append(torch.cuda.memory_allocated(device))
|
|
|
|
torch.cuda.synchronize(device)
|
|
wall = time.perf_counter() - started
|
|
peak_vram = torch.cuda.max_memory_allocated(device)
|
|
peak_reserved_vram = torch.cuda.max_memory_reserved(device)
|
|
final_allocated = torch.cuda.memory_allocated(device)
|
|
times = np.concatenate(event_t) if event_t else np.empty(0, dtype=np.float64)
|
|
neurons = np.concatenate(event_n) if event_n else np.empty(0, dtype=np.int32)
|
|
save_events(exp, test["name"], times, neurons, duration, float(cfg["output"]["rate_window_ms"]))
|
|
total_spikes = int(len(times))
|
|
duration_s = duration / 1000.0
|
|
active = int(np.unique(neurons).size)
|
|
synchronous_fraction = peak_step_spikes / n
|
|
observed_growth = int(max(memory_samples, default=final_allocated) - min(memory_samples, default=final_allocated))
|
|
result = {
|
|
"name": test["name"],
|
|
"stimulation": stim["name"],
|
|
"stimulation_kind": spec["kind"],
|
|
"stimulated_neuron_count": int(len(targets_cpu)),
|
|
"stimulated_neuron_indices": targets_cpu.tolist() if len(targets_cpu) <= 64 else targets_cpu[:64].tolist(),
|
|
"stimulated_indices_truncated": bool(len(targets_cpu) > 64),
|
|
"duration_ms": duration,
|
|
"steps": steps,
|
|
"wall_clock_seconds": wall,
|
|
"realtime_factor_wall_over_biological": wall / duration_s,
|
|
"steps_per_second": steps / wall,
|
|
"total_spikes": total_spikes,
|
|
"spikes_per_biological_second": total_spikes / duration_s,
|
|
"mean_firing_rate_hz": total_spikes / (n * duration_s),
|
|
"active_neuron_count": active,
|
|
"peak_spikes_in_one_step": peak_step_spikes,
|
|
"peak_synchronous_fraction": synchronous_fraction,
|
|
"all_network_silent": total_spikes == 0,
|
|
"explosive_synchrony": synchronous_fraction >= float(cfg["validation"]["explosive_step_fraction"]),
|
|
"nan_or_inf": not numerical_ok,
|
|
"peak_vram_bytes": int(peak_vram),
|
|
"peak_reserved_vram_bytes": int(peak_reserved_vram),
|
|
"persistent_model_vram_bytes": int(persistent_bytes),
|
|
"final_allocated_vram_bytes": int(final_allocated),
|
|
"observed_allocated_growth_bytes": observed_growth,
|
|
"memory_leak_suspected": observed_growth > int(cfg["validation"]["memory_growth_tolerance_bytes"]),
|
|
"self_loop_transmission_events": self_loop_events,
|
|
"self_loop_effective_input_sum": self_loop_effective_sum,
|
|
"event_file": str(exp / f"outputs/spikes/{test['name']}.events.parquet"),
|
|
}
|
|
del v, syn, refractory, queue
|
|
torch.cuda.empty_cache()
|
|
return result
|
|
|
|
|
|
def induced_subgraph(indptr: np.ndarray, indices: np.ndarray, effective: np.ndarray, n: int, size: int, seed: int) -> dict[str, np.ndarray]:
|
|
rng = np.random.default_rng(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] = []
|
|
vals: list[float] = []
|
|
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
|
|
if keep.any():
|
|
count = int(keep.sum())
|
|
rows.extend([local_src] * count)
|
|
cols.extend(local_targets[keep].tolist())
|
|
vals.extend(effective[start:stop][keep].astype(np.float64).tolist())
|
|
return {
|
|
"nodes": nodes,
|
|
"rows": np.asarray(rows, dtype=np.int64),
|
|
"cols": np.asarray(cols, dtype=np.int64),
|
|
"weights": np.asarray(vals, dtype=np.float64),
|
|
}
|
|
|
|
|
|
def reference_numpy(cfg: dict[str, Any], sub: dict[str, np.ndarray], ext: np.ndarray) -> tuple[np.ndarray, np.ndarray, list[tuple[int, int]]]:
|
|
size = len(sub["nodes"])
|
|
steps = ext.shape[0]
|
|
m = cfg["model"]
|
|
dt = float(m["dt_ms"])
|
|
rest, reset, threshold = float(m["resting_potential_mv"]), float(m["reset_potential_mv"]), float(m["threshold_mv"])
|
|
alpha = dt / float(m["membrane_tau_ms"])
|
|
decay = math.exp(-dt / float(m["synaptic_tau_ms"]))
|
|
refractory_steps = int(round(float(m["refractory_ms"]) / dt))
|
|
delay_steps = int(round(float(m["synaptic_delay_ms"]) / dt))
|
|
queue = np.zeros((delay_steps + 1, size), dtype=np.float64)
|
|
v = np.full(size, rest, dtype=np.float64)
|
|
syn = np.zeros(size, dtype=np.float64)
|
|
ref = np.zeros(size, dtype=np.int32)
|
|
events: list[tuple[int, int]] = []
|
|
for step in range(steps):
|
|
slot = step % len(queue)
|
|
syn *= decay
|
|
syn += queue[slot]
|
|
queue[slot].fill(0)
|
|
available = ref == 0
|
|
v[available] += alpha * (rest - v[available] + syn[available] + ext[step, available])
|
|
ref = np.maximum(ref - 1, 0)
|
|
sources = np.flatnonzero(available & (v >= threshold))
|
|
if len(sources):
|
|
for source in sources:
|
|
events.append((step, int(source)))
|
|
v[sources] = reset
|
|
ref[sources] = refractory_steps
|
|
active_edges = np.isin(sub["rows"], sources)
|
|
np.add.at(queue[(step + delay_steps) % len(queue)], sub["cols"][active_edges], sub["weights"][active_edges])
|
|
return v, syn, events
|
|
|
|
|
|
def reference_torch(cfg: dict[str, Any], sub: dict[str, np.ndarray], ext: np.ndarray, device: torch.device) -> tuple[np.ndarray, np.ndarray, list[tuple[int, int]]]:
|
|
size = len(sub["nodes"])
|
|
steps = ext.shape[0]
|
|
m = cfg["model"]
|
|
dt = float(m["dt_ms"])
|
|
rest, reset, threshold = float(m["resting_potential_mv"]), float(m["reset_potential_mv"]), float(m["threshold_mv"])
|
|
alpha = dt / float(m["membrane_tau_ms"])
|
|
decay = math.exp(-dt / float(m["synaptic_tau_ms"]))
|
|
refractory_steps = int(round(float(m["refractory_ms"]) / dt))
|
|
delay_steps = int(round(float(m["synaptic_delay_ms"]) / dt))
|
|
queue = torch.zeros((delay_steps + 1, size), dtype=torch.float64, device=device)
|
|
v = torch.full((size,), rest, dtype=torch.float64, device=device)
|
|
syn = torch.zeros(size, dtype=torch.float64, device=device)
|
|
ref = torch.zeros(size, dtype=torch.int32, device=device)
|
|
rows = torch.as_tensor(sub["rows"], dtype=torch.int64, device=device)
|
|
cols = torch.as_tensor(sub["cols"], dtype=torch.int64, device=device)
|
|
weights = torch.as_tensor(sub["weights"], dtype=torch.float64, device=device)
|
|
ext_t = torch.as_tensor(ext, dtype=torch.float64, device=device)
|
|
events: list[tuple[int, int]] = []
|
|
for step in range(steps):
|
|
slot = step % len(queue)
|
|
syn.mul_(decay).add_(queue[slot])
|
|
queue[slot].zero_()
|
|
available = ref == 0
|
|
v[available] += alpha * (rest - v[available] + syn[available] + ext_t[step, available])
|
|
ref.clamp_min_(1).sub_(1)
|
|
sources = torch.nonzero(available & (v >= threshold), as_tuple=False).flatten()
|
|
if sources.numel():
|
|
for source in sources.cpu().tolist():
|
|
events.append((step, int(source)))
|
|
v[sources] = reset
|
|
ref[sources] = refractory_steps
|
|
active_edges = torch.isin(rows, sources)
|
|
queue[(step + delay_steps) % len(queue)].scatter_add_(0, cols[active_edges], weights[active_edges])
|
|
torch.cuda.synchronize(device)
|
|
return v.cpu().numpy(), syn.cpu().numpy(), events
|
|
|
|
|
|
def run_reference(cfg: dict[str, Any], inputs: dict[str, Path], exp: Path, effective_path: Path, device: torch.device) -> dict[str, Any]:
|
|
spec = cfg["cpu_gpu_reference"]
|
|
n = int(cfg["expected_graph"]["neurons"])
|
|
indptr = np.load(inputs["indptr"], mmap_mode="r")
|
|
indices = np.load(inputs["indices"], mmap_mode="r")
|
|
effective = np.load(effective_path, mmap_mode="r")
|
|
sub = induced_subgraph(indptr, indices, effective, n, int(spec["subgraph_neurons"]), int(spec["seed"]))
|
|
steps = int(spec["steps"])
|
|
size = len(sub["nodes"])
|
|
ext = np.zeros((steps, size), dtype=np.float64)
|
|
start = int(spec["pulse_start_step"])
|
|
stop = min(steps, start + int(spec["pulse_steps"]))
|
|
ext[start:stop, : int(spec["pulse_neurons"])] = float(spec["pulse_amplitude_mv_equivalent"])
|
|
v_cpu, syn_cpu, events_cpu = reference_numpy(cfg, sub, ext)
|
|
v_gpu, syn_gpu, events_gpu = reference_torch(cfg, sub, ext, device)
|
|
v_error = float(np.max(np.abs(v_cpu - v_gpu)))
|
|
syn_error = float(np.max(np.abs(syn_cpu - syn_gpu)))
|
|
tolerance = float(spec["absolute_tolerance"])
|
|
result = {
|
|
"seed": int(spec["seed"]),
|
|
"selection": "uniform random without replacement, then induced edges",
|
|
"subgraph_neurons": size,
|
|
"induced_edges": int(len(sub["rows"])),
|
|
"steps": steps,
|
|
"cpu_event_count": len(events_cpu),
|
|
"gpu_event_count": len(events_gpu),
|
|
"events_exact_match": events_cpu == events_gpu,
|
|
"max_abs_voltage_difference": v_error,
|
|
"max_abs_synaptic_state_difference": syn_error,
|
|
"absolute_tolerance": tolerance,
|
|
"numerical_match": v_error <= tolerance and syn_error <= tolerance,
|
|
"passed": events_cpu == events_gpu and v_error <= tolerance and syn_error <= tolerance,
|
|
"selected_global_neuron_indices": sub["nodes"].tolist(),
|
|
}
|
|
write_json(exp / "benchmark/cpu-gpu-reference.json", result)
|
|
return result
|
|
|
|
|
|
def build_self_loop_vector(inputs: dict[str, Path], effective_path: Path, n: int, device: torch.device) -> tuple[torch.Tensor, int]:
|
|
indptr = np.load(inputs["indptr"], mmap_mode="r")
|
|
indices = np.load(inputs["indices"], mmap_mode="r")
|
|
effective = np.load(effective_path, mmap_mode="r")
|
|
loops = np.zeros(n, dtype=np.float32)
|
|
count = 0
|
|
for src in range(n):
|
|
start, stop = int(indptr[src]), int(indptr[src + 1])
|
|
hit = np.flatnonzero(indices[start:stop] == src)
|
|
if len(hit):
|
|
count += len(hit)
|
|
loops[src] = float(effective[start + int(hit[0])])
|
|
return torch.as_tensor(loops, dtype=torch.float32, device=device), count
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", required=True, type=Path)
|
|
args = parser.parse_args()
|
|
cfg = load_config(args.config)
|
|
exp = Path(cfg["paths"]["experiment_root"]).resolve()
|
|
for sub in ("config", "benchmark", "logs", "outputs", "audit"):
|
|
(exp / sub).mkdir(parents=True, exist_ok=True)
|
|
setup_logging(exp / "logs/phase4a.log")
|
|
LOG.info("Phase 4A start; configuration=%s", args.config)
|
|
random.seed(int(cfg["reproducibility"]["seed"]))
|
|
np.random.seed(int(cfg["reproducibility"]["seed"]))
|
|
torch.manual_seed(int(cfg["reproducibility"]["seed"]))
|
|
inputs = require_exact_inputs(cfg)
|
|
input_hashes_before = {name: sha256(path) for name, path in inputs.items()}
|
|
derived_audit = prepare_derived(cfg, inputs, exp)
|
|
|
|
expected_gpu_name = str(cfg["device"]["required_name_substring"])
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA is unavailable")
|
|
device = torch.device("cuda:0")
|
|
props = torch.cuda.get_device_properties(device)
|
|
if expected_gpu_name not in props.name:
|
|
raise RuntimeError(f"required GPU containing {expected_gpu_name!r}, got {props.name!r}")
|
|
if torch.cuda.device_count() != 1:
|
|
raise RuntimeError("launcher must expose exactly one GPU; mixed multi-GPU is forbidden")
|
|
|
|
n = int(cfg["expected_graph"]["neurons"])
|
|
effective_path = Path(derived_audit["effective_weight_path"])
|
|
graph = {
|
|
"indptr": torch.as_tensor(np.load(inputs["indptr"]), dtype=torch.int64, device=device),
|
|
"indices": torch.as_tensor(np.load(inputs["indices"]), dtype=torch.int32, device=device),
|
|
"effective": torch.as_tensor(np.load(effective_path), dtype=torch.float32, device=device),
|
|
}
|
|
metadata = pq.read_table(inputs["metadata"], columns=["neuron_idx", "annotation_type", "neurotransmitter"])
|
|
self_loop_weight, self_loop_count = build_self_loop_vector(inputs, effective_path, n, device)
|
|
if self_loop_count != int(cfg["expected_graph"]["self_loops"]):
|
|
raise RuntimeError(f"self-loop invariant failed: {self_loop_count}")
|
|
|
|
reference = run_reference(cfg, inputs, exp, effective_path, device)
|
|
if not reference["passed"]:
|
|
raise RuntimeError("CPU/GPU reference comparison failed")
|
|
|
|
smoke_results = []
|
|
for test in cfg["smoke_tests"]:
|
|
LOG.info("running %s", test["name"])
|
|
result = run_gpu_smoke(cfg, inputs, exp, test, graph, metadata, self_loop_weight)
|
|
write_json(exp / f"benchmark/{test['name']}.json", result)
|
|
smoke_results.append(result)
|
|
LOG.info("completed %s: wall=%.6fs spikes=%d", test["name"], result["wall_clock_seconds"], result["total_spikes"])
|
|
|
|
input_hashes_after = {name: sha256(path) for name, path in inputs.items()}
|
|
if input_hashes_before != input_hashes_after:
|
|
raise RuntimeError("an input file changed during Phase 4A")
|
|
peak = max(r["peak_vram_bytes"] for r in smoke_results)
|
|
peak_reserved = max(r["peak_reserved_vram_bytes"] for r in smoke_results)
|
|
summary = {
|
|
"status": "PASS",
|
|
"model_loaded_neurons": n,
|
|
"model_loaded_edges": int(cfg["expected_graph"]["directed_edges"]),
|
|
"raw_synaptic_contact_weight_sum": int(cfg["expected_graph"]["synaptic_contact_weight_sum"]),
|
|
"device": {
|
|
"name": props.name,
|
|
"total_vram_bytes": int(props.total_memory),
|
|
"torch_visible_device_count": torch.cuda.device_count(),
|
|
"mixed_multi_gpu_used": False,
|
|
},
|
|
"peak_vram_bytes_across_smoke_tests": int(peak),
|
|
"peak_reserved_vram_bytes_across_smoke_tests": int(peak_reserved),
|
|
"derived_audit": derived_audit,
|
|
"cpu_gpu_reference": reference,
|
|
"smoke_tests": smoke_results,
|
|
"numerical_instability_detected": any(r["nan_or_inf"] for r in smoke_results),
|
|
"any_explosive_synchrony": any(r["explosive_synchrony"] for r in smoke_results),
|
|
"any_silent_run": any(r["all_network_silent"] for r in smoke_results),
|
|
"any_memory_leak_suspected": any(r["memory_leak_suspected"] for r in smoke_results),
|
|
"self_loops_loaded": self_loop_count,
|
|
"isolated_neurons_retained": int(cfg["expected_graph"]["isolated_neurons"]),
|
|
"input_hashes_unchanged": True,
|
|
"limitations": [
|
|
"Dynamics parameters are a configurable FlyWire-inspired baseline, not MaleCNS single-cell measurements.",
|
|
"Synthetic benchmark drive amplitudes are test inputs, not biological calibration.",
|
|
"No training, plasticity, RL, game, web service, visual input, or 3D frontend is included.",
|
|
],
|
|
}
|
|
write_json(exp / "benchmark/summary.json", summary)
|
|
|
|
# Seal the log before hashing outputs so the manifest remains verifiable.
|
|
LOG.info("PHASE4A_PASS")
|
|
for handler in logging.getLogger().handlers:
|
|
handler.flush()
|
|
audit_controls = {exp / "audit/manifest.json", exp / "audit/checksums.sha256"}
|
|
output_files = sorted(p for p in exp.rglob("*") if p.is_file() and p not in audit_controls)
|
|
manifest = {
|
|
"phase": "4A",
|
|
"created_unix_time": time.time(),
|
|
"configuration": str(args.config.resolve()),
|
|
"configuration_sha256": sha256(args.config),
|
|
"code": str(Path(__file__).resolve()),
|
|
"code_sha256": sha256(Path(__file__).resolve()),
|
|
"python": platform.python_version(),
|
|
"torch": torch.__version__,
|
|
"numpy": np.__version__,
|
|
"pyarrow": pa.__version__,
|
|
"input_files": {name: {"path": str(inputs[name]), "sha256": digest} for name, digest in input_hashes_after.items()},
|
|
"output_files": [{"path": str(p), "bytes": p.stat().st_size, "sha256": sha256(p)} for p in output_files],
|
|
}
|
|
manifest_path = exp / "audit/manifest.json"
|
|
write_json(manifest_path, manifest)
|
|
checksum_targets = output_files + [manifest_path]
|
|
checksum_lines = [f"{sha256(path)} {path.relative_to(exp)}" for path in checksum_targets]
|
|
(exp / "audit/checksums.sha256").write_text("\n".join(checksum_lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|