#!/usr/bin/env python3 """Stable game-agnostic Brain I/O API for MaleCNS-LIF-runtime-v1.""" from __future__ import annotations import copy import json import math import sys from dataclasses import dataclass, field from pathlib import Path from typing import Any, Iterable import numpy as np import pyarrow.parquet as pq import torch import yaml API_VERSION = "1.0.0" INPUT_BATCH_VERSION = "1.0.0" ACTIVITY_STREAM_VERSION = "1.0.0" @dataclass(frozen=True) class InputEvent: target_neuron_idx: np.ndarray stimulation_mode: str value_kind: str value: float | np.ndarray start_ms: float duration_ms: float source_id: str channel_id: str amplitude_mv: float = 0.0 metadata: dict[str, Any] = field(default_factory=dict) def validate(self, neuron_count: int, dt_ms: float) -> None: target = np.asarray(self.target_neuron_idx) if target.ndim != 1 or target.dtype.kind not in "iu": raise ValueError("target_neuron_idx must be a one-dimensional integer array") if len(target) and (int(target.min()) < 0 or int(target.max()) >= neuron_count): raise ValueError("target neuron_idx outside canonical range") if len(np.unique(target)) != len(target): raise ValueError("target neuron_idx values must be unique") if self.stimulation_mode not in {"spike_injection", "external_synaptic_drive"}: raise ValueError(f"unsupported stimulation_mode {self.stimulation_mode!r}") if self.value_kind not in {"binary", "scalar_intensity", "poisson_rate_hz", "signed_modulation"}: raise ValueError(f"unsupported value_kind {self.value_kind!r}") if self.start_ms < 0 or self.duration_ms <= 0: raise ValueError("timestamps/duration must be non-negative/positive") for x in (self.start_ms, self.duration_ms): if not math.isclose(x / dt_ms, round(x / dt_ms), abs_tol=1e-8): raise ValueError(f"{x} ms is not aligned to dt={dt_ms} ms") value = np.asarray(self.value) if value.ndim > 1 or (value.ndim == 1 and len(value) not in (1, len(target))): raise ValueError("value must be scalar or match target count") if self.value_kind == "poisson_rate_hz" and np.any(value < 0): raise ValueError("Poisson rate cannot be negative") @dataclass(frozen=True) class InputBatch: events: tuple[InputEvent, ...] batch_id: str version: str = INPUT_BATCH_VERSION metadata: dict[str, Any] = field(default_factory=dict) def validate(self, neuron_count: int, dt_ms: float) -> None: if self.version != INPUT_BATCH_VERSION: raise ValueError(f"InputBatch version must be {INPUT_BATCH_VERSION}") for event in self.events: event.validate(neuron_count, dt_ms) class VisualAdapterV1: """Maps abstract channels to audited visual hierarchy populations.""" def __init__(self, config: dict[str, Any], phase4c_root: Path): self.config = config self.phase4c_root = phase4c_root self.modes: dict[str, np.ndarray] = {} for mode, spec in config["visual_adapter"]["modes"].items(): population = spec["hierarchy_population"] path = phase4c_root / "visual-audit" / f"hierarchy-{population}.parquet" if not path.is_file(): raise FileNotFoundError(path) self.modes[mode] = pq.read_table(path, columns=["neuron_idx"])["neuron_idx"].to_numpy().astype(np.int64) def population(self, mode: str) -> np.ndarray: if mode not in self.modes: raise KeyError(mode) return self.modes[mode].copy() def encode( self, mode: str, channel_values: np.ndarray, start_ms: float, duration_ms: float, source_id: str, channel_id: str, value_kind: str = "poisson_rate_hz", ) -> InputBatch: if mode not in self.modes: raise KeyError(mode) spec = self.config["visual_adapter"]["modes"][mode] targets = self.modes[mode] values = np.asarray(channel_values, dtype=np.float32) if values.shape != (len(targets),): raise ValueError(f"visual channel must have shape {(len(targets),)}, got {values.shape}") selected = np.flatnonzero(values != 0).astype(np.int64) target = targets[selected] selected_values = values[selected] event = InputEvent( target_neuron_idx=target, stimulation_mode=spec["input_mode"], value_kind=value_kind, value=selected_values, amplitude_mv=float(self.config["visual_adapter"]["external_drive_mv"]), start_ms=start_ms, duration_ms=duration_ms, source_id=source_id, channel_id=channel_id, metadata={ "adapter": self.config["names"]["visual_adapter"], "adapter_mode": mode, "biological_fidelity": spec["biological_fidelity"], "connectome_weights_modified": False, }, ) return InputBatch(events=(event,), batch_id=f"{source_id}:{channel_id}:{start_ms}", metadata={"adapter_mode": mode}) class BrainRuntime: """Persistent exact-state BrainRuntime API. It has no game-specific dependency.""" def __init__(self, config_path: Path, extension_build_dir: Path): self.config_path = config_path.resolve() self.cfg = yaml.safe_load(self.config_path.read_text(encoding="utf-8")) self.dt = float(self.cfg["model"]["dt_ms"]) self.n = 166700 self.model_root = Path(self.cfg["paths"]["model_root"]) self.runtime_root = Path(self.cfg["paths"]["runtime_root"]) self.input_root = Path(self.cfg["paths"]["connectome_root"]) runtime_module_dir = str(self.runtime_root) if runtime_module_dir not in sys.path: sys.path.insert(0, runtime_module_dir) import phase4c as runtime_module self.runtime_module = runtime_module p4cfg = yaml.safe_load((self.runtime_root / "phase4c.yaml").read_text(encoding="utf-8")) self.paths = runtime_module.exact_paths(p4cfg) self.indptr_np, self.indices_np, self.raw_np = runtime_module.p4b.validate_graph(p4cfg, self.paths) self.table = runtime_module.full_metadata(self.paths["metadata"]) self.population_masks = runtime_module.populations(self.table) if not torch.cuda.is_available() or torch.cuda.device_count() != 1: raise RuntimeError("BrainRuntime requires exactly one visible CUDA GPU") self.device = torch.device("cuda:0") props = torch.cuda.get_device_properties(self.device) if self.cfg["device"]["required_name_substring"] not in props.name: raise RuntimeError(f"required RTX 4090, got {props.name}") extension_build_dir.mkdir(parents=True, exist_ok=True) build_experiment_root = extension_build_dir.parent.parent (build_experiment_root / "compiler-reports/extension-build").mkdir(parents=True, exist_ok=True) self.extension, report = runtime_module.build_extension(self.runtime_root, build_experiment_root) if self.extension is None: raise RuntimeError(f"runtime-v1 extension unavailable: {report}") base = runtime_module.Runtime(p4cfg, self.table, self.indptr_np, self.indices_np, self.raw_np, self.device, self.extension) self.indptr, self.indices, self.weights = base.indptr, base.indices, base.weights self.signs, self.superclass = base.signs, base.superclass self.neuron_type = np.asarray(self.table["annotation_type"].to_pylist(), dtype=object) self.neuron_class = np.asarray(self.table["annotation_class"].to_pylist(), dtype=object) self.body_id = self.table["body_id"].to_numpy().astype(np.uint64) m = self.cfg["model"] self.delay_steps = int(round(float(m["synaptic_delay_ms"]) / self.dt)) self.queue_len = self.delay_steps + 1 self.refractory_steps = int(round(float(m["refractory_ms"]) / self.dt)) self.alpha = self.dt / float(m["membrane_tau_ms"]) self.g_decay = 1.0 - self.dt / float(m["synaptic_tau_ms"]) self.rest = float(m["resting_potential_mv"]) self.reset_value = float(m["reset_potential_mv"]) self.threshold = float(m["threshold_mv"]) self.output_populations = self._output_populations() self.reset(0) def _output_populations(self) -> dict[str, np.ndarray]: s = self.superclass return { "descending": np.flatnonzero(np.isin(s, ["descending_neuron", "descending_neuron_tbc"])).astype(np.int32), "motor_related": np.flatnonzero(np.isin(s, ["vnc_motor", "cb_motor"])).astype(np.int32), "ascending": np.flatnonzero(np.isin(s, ["ascending_neuron", "sensory_ascending", "sensory_ascending_tbc", "efferent_ascending"])).astype(np.int32), } def reset(self, seed: int, initial_state: dict[str, Any] | None = None) -> dict[str, Any]: self.seed = int(seed) self.generator = torch.Generator(device=self.device); self.generator.manual_seed(self.seed) self.v = torch.full((self.n,), self.rest, dtype=torch.float32, device=self.device) self.g = torch.zeros(self.n, dtype=torch.float32, device=self.device) self.refractory = torch.zeros(self.n, dtype=torch.int32, device=self.device) self.delay_queue = torch.zeros((self.queue_len, self.n), dtype=torch.float32, device=self.device) self.spike_counts = torch.zeros(self.n, dtype=torch.int32, device=self.device) self.simulation_step = 0; self.sequence_id = 0 self.spike_steps: list[np.ndarray] = []; self.spike_nodes: list[np.ndarray] = [] self.spike_origins: list[np.ndarray] = [] self.drive_steps: list[np.ndarray] = []; self.drive_nodes: list[np.ndarray] = [] self.drive_values: list[np.ndarray] = []; self.input_audit: list[dict[str, Any]] = [] if initial_state is not None: self.restore_state(initial_state) return self.get_state_summary() @property def simulation_time_ms(self) -> float: return self.simulation_step * self.dt def _event_values(self, event: InputEvent, targets: torch.Tensor) -> torch.Tensor: value = np.asarray(event.value, dtype=np.float32) if value.ndim == 0: return torch.full((len(targets),), float(value), dtype=torch.float32, device=self.device) return torch.as_tensor(value, dtype=torch.float32, device=self.device) def _active(self, event: InputEvent, time_ms: float) -> bool: return event.start_ms - 1e-9 <= time_ms < event.start_ms + event.duration_ms - 1e-9 def _hybrid_edges(self, sources: torch.Tensor) -> torch.Tensor: starts = self.indptr[sources] degrees = self.indptr[sources + 1] - starts cumulative = torch.cumsum(degrees, 0) total = int(cumulative[-1].item()) if cumulative.numel() else 0 output = torch.empty(total, dtype=torch.int64, device=self.device) if total: self.extension.fill_edges(self.indptr, sources, cumulative - degrees, output) return output def step(self, input_batch: InputBatch | None, biological_ms: float) -> dict[str, Any]: steps = int(round(float(biological_ms) / self.dt)) if not math.isclose(steps * self.dt, float(biological_ms), abs_tol=1e-8) or steps <= 0: raise ValueError("biological_ms must be a positive multiple of dt") batch = input_batch or InputBatch(events=(), batch_id="zero-input") batch.validate(self.n, self.dt) batch_start = self.simulation_time_ms prepared = [(e, torch.as_tensor(e.target_neuron_idx, dtype=torch.int64, device=self.device)) for e in batch.events] for event in batch.events: self.input_audit.append({ "batch_id": batch.batch_id, "source_id": event.source_id, "channel_id": event.channel_id, "mode": event.stimulation_mode, "value_kind": event.value_kind, "start_ms": event.start_ms, "duration_ms": event.duration_ms, "target_count": int(len(event.target_neuron_idx)), "metadata": event.metadata, "external_drive_is_connectome_transmission": False, }) start_event_count = sum(len(x) for x in self.spike_nodes) start_drive_count = sum(len(x) for x in self.drive_nodes) for _ in range(steps): global_step = self.simulation_step; time_ms = self.simulation_time_ms slot = global_step % self.queue_len self.g.mul_(self.g_decay).add_(self.delay_queue[slot]); self.delay_queue[slot].zero_() spike_requests: list[tuple[InputEvent, torch.Tensor, torch.Tensor]] = [] for event, targets in prepared: if not self._active(event, time_ms) or targets.numel() == 0: continue values = self._event_values(event, targets) if event.value_kind == "poisson_rate_hz": selected = torch.rand(len(targets), generator=self.generator, device=self.device) < (values * self.dt / 1000.0) elif event.value_kind == "binary": selected = values != 0 elif event.value_kind == "scalar_intensity": selected = values != 0 else: # signed_modulation selected = torch.rand(len(targets), generator=self.generator, device=self.device) < (torch.abs(values) * self.dt / 1000.0) if event.stimulation_mode == "external_synaptic_drive": chosen = targets[selected] if chosen.numel(): signed_value = values[selected] * float(event.amplitude_mv) if event.value_kind == "poisson_rate_hz": signed_value = torch.full_like(signed_value, float(event.amplitude_mv)) self.g.index_add_(0, chosen, signed_value) node_np = chosen.cpu().numpy().astype(np.int32, copy=True) val_np = signed_value.cpu().numpy().astype(np.float32, copy=True) self.drive_steps.append(np.full(len(node_np), global_step, dtype=np.int32)) self.drive_nodes.append(node_np); self.drive_values.append(val_np) else: spike_requests.append((event, targets, selected)) self.refractory.clamp_min_(1).sub_(1) available = self.refractory == 0 self.v[available] += self.alpha * (self.g[available] - (self.v[available] - self.rest)) self.v[~available] = self.reset_value natural = available & (self.v >= self.threshold) spikes = natural.clone() injected = torch.zeros(self.n, dtype=torch.bool, device=self.device) for _event, targets, selected in spike_requests: accepted = selected & available[targets] chosen = targets[accepted] if chosen.numel(): spikes[chosen] = True; injected[chosen] = ~natural[chosen] sources = torch.nonzero(spikes, as_tuple=False).flatten() self.spike_counts.add_(spikes.to(torch.int32)) if sources.numel(): self.v[sources] = self.reset_value; self.refractory[sources] = self.refractory_steps edge_ids = self._hybrid_edges(sources) if edge_ids.numel(): self.delay_queue[(global_step + self.delay_steps) % self.queue_len].scatter_add_(0, self.indices[edge_ids].long(), self.weights[edge_ids]) src_np = sources.cpu().numpy().astype(np.int32, copy=True) origin_np = np.where(injected[sources].cpu().numpy(), "external_spike_injection", "network_spike") self.spike_steps.append(np.full(len(src_np), global_step, dtype=np.int32)) self.spike_nodes.append(src_np); self.spike_origins.append(origin_np) self.simulation_step += 1 torch.cuda.synchronize() self.sequence_id += 1 return { "api_version": API_VERSION, "sequence_id": self.sequence_id, "batch_id": batch.batch_id, "start_time_ms": batch_start, "end_time_ms": self.simulation_time_ms, "biological_ms": biological_ms, "new_spikes": sum(len(x) for x in self.spike_nodes) - start_event_count, "new_external_drive_events": sum(len(x) for x in self.drive_nodes) - start_drive_count, } def _events(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: steps = np.concatenate(self.spike_steps) if self.spike_steps else np.empty(0, np.int32) nodes = np.concatenate(self.spike_nodes) if self.spike_nodes else np.empty(0, np.int32) origins = np.concatenate(self.spike_origins) if self.spike_origins else np.empty(0, object) return steps, nodes, origins def read_spikes(self, window_ms: float | None = None) -> dict[str, Any]: steps, nodes, origins = self._events() if window_ms is not None: threshold = self.simulation_step - int(round(window_ms / self.dt)) mask = steps >= threshold; steps, nodes, origins = steps[mask], nodes[mask], origins[mask] return { "simulation_time_ms": self.simulation_time_ms, "step": steps, "timestamp_ms": steps.astype(np.float64) * self.dt, "neuron_idx": nodes, "origin": origins, } def read_population(self, population_id: str, window_ms: float) -> dict[str, Any]: if population_id not in self.output_populations: raise KeyError(population_id) members = self.output_populations[population_id] events = self.read_spikes(window_ms) mask = np.isin(events["neuron_idx"], members) nodes = events["neuron_idx"][mask] counts = np.zeros(len(members), dtype=np.int32) if len(nodes): lookup = np.full(self.n, -1, dtype=np.int32); lookup[members] = np.arange(len(members), dtype=np.int32) np.add.at(counts, lookup[nodes], 1) seconds = window_ms / 1000.0 return { "simulation_time_ms": self.simulation_time_ms, "population_id": population_id, "neuron_idx": members.copy(), "spike_counts": counts, "firing_rate_hz": counts.astype(np.float32) / seconds, "active_mask": counts > 0, "aggregate_rate_hz": float(counts.sum() / (len(members) * seconds)), "total_spikes": int(counts.sum()), "active_neurons": int(np.count_nonzero(counts)), } def read_output(self, output_spec: dict[str, Any], window_ms: float) -> dict[str, Any]: population = output_spec.get("population", "descending") aggregation = output_spec.get("aggregation", "per-neuron") base = self.read_population(population, window_ms) members = base["neuron_idx"]; counts = base["spike_counts"] if aggregation == "per-neuron": return base if aggregation not in {"per-type", "per-class"}: raise ValueError("aggregation must be per-neuron, per-type, or per-class") labels = self.neuron_type if aggregation == "per-type" else self.neuron_class result: dict[str, dict[str, float | int]] = {} for label in sorted({"" if labels[i] is None else str(labels[i]) for i in members}): mask = np.asarray([("" if labels[i] is None else str(labels[i])) == label for i in members]) result[label] = {"neuron_count": int(mask.sum()), "spike_count": int(counts[mask].sum()), "mean_rate_hz": float(base["firing_rate_hz"][mask].mean())} return {"simulation_time_ms": self.simulation_time_ms, "population_id": population, "aggregation": aggregation, "features": result} def fixed_temporal_bins(self, population_id: str, bin_ms: float, window_ms: float) -> dict[str, Any]: if bin_ms not in self.cfg["readout"]["bins_ms"]: raise ValueError("unsupported fixed bin") ev = self.read_spikes(window_ms); members = self.output_populations[population_id] mask = np.isin(ev["neuron_idx"], members); times = ev["timestamp_ms"][mask] start = self.simulation_time_ms - window_ms edges = np.arange(start, self.simulation_time_ms + bin_ms + 1e-9, bin_ms) counts, _ = np.histogram(times, bins=edges) return {"simulation_time_ms": self.simulation_time_ms, "population_id": population_id, "bin_ms": bin_ms, "bin_start_ms": edges[:-1], "spike_counts": counts.astype(np.int32)} def get_state_summary(self) -> dict[str, Any]: return { "api_version": API_VERSION, "runtime": self.cfg["names"]["runtime"], "simulation_step": self.simulation_step, "simulation_time_ms": self.simulation_time_ms, "sequence_id": self.sequence_id, "seed": self.seed, "neurons": self.n, "edges": int(self.indices.numel()), "active_neurons_total": int(torch.count_nonzero(self.spike_counts).item()), "spikes_total": int(self.spike_counts.sum().item()), "external_drive_events_total": sum(len(x) for x in self.drive_nodes), } def checkpoint_state(self) -> dict[str, Any]: return { "api_version": API_VERSION, "seed": self.seed, "simulation_step": self.simulation_step, "sequence_id": self.sequence_id, "generator_state": self.generator.get_state().cpu(), "v": self.v.clone(), "g": self.g.clone(), "refractory": self.refractory.clone(), "delay_queue": self.delay_queue.clone(), "spike_counts": self.spike_counts.clone(), "spike_steps": copy.deepcopy(self.spike_steps), "spike_nodes": copy.deepcopy(self.spike_nodes), "spike_origins": copy.deepcopy(self.spike_origins), "drive_steps": copy.deepcopy(self.drive_steps), "drive_nodes": copy.deepcopy(self.drive_nodes), "drive_values": copy.deepcopy(self.drive_values), "input_audit": copy.deepcopy(self.input_audit), } def restore_state(self, state: dict[str, Any]) -> dict[str, Any]: if state.get("api_version") != API_VERSION: raise ValueError("checkpoint API version mismatch") self.seed = int(state["seed"]); self.simulation_step = int(state["simulation_step"]); self.sequence_id = int(state["sequence_id"]) self.generator.set_state(state["generator_state"].cpu()) for name in ("v", "g", "refractory", "delay_queue", "spike_counts"): setattr(self, name, state[name].clone().to(self.device)) for name in ("spike_steps", "spike_nodes", "spike_origins", "drive_steps", "drive_nodes", "drive_values", "input_audit"): setattr(self, name, copy.deepcopy(state[name])) return self.get_state_summary() def activity_packet(self, window_ms: float, input_metadata: dict[str, Any] | None = None) -> dict[str, Any]: ev = self.read_spikes(window_ms) aggregates = {name: {k: v for k, v in self.read_population(name, window_ms).items() if k in ("active_neurons", "total_spikes", "aggregate_rate_hz")} for name in self.output_populations} return { "contract_version": ACTIVITY_STREAM_VERSION, "simulation_time_ms": self.simulation_time_ms, "sequence_id": self.sequence_id, "window_ms": window_ms, "spike_neuron_idx": ev["neuron_idx"].tolist(), "spike_timestamp_ms": ev["timestamp_ms"].tolist(), "population_aggregates": aggregates, "active_neuron_count": int(len(np.unique(ev["neuron_idx"]))), "input_channel_metadata": input_metadata or {}, "representation": "event-only", }