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