Files
digital-drosophila/models/male-cns-lif/visual-encoder-v1/brain_runtime.py
T

416 lines
23 KiB
Python

#!/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({"<null>" if labels[i] is None else str(labels[i]) for i in members}):
mask = np.asarray([("<null>" 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",
}