Archive MaleCNS research sources, provenance, bibliography and immutable data inventory
This commit is contained in:
Digital Drosophila Archive
@@ -0,0 +1,645 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user