Spaces:
Paused
Paused
File size: 3,446 Bytes
5e0b58b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | from __future__ import annotations
import random
from dataclasses import dataclass
from typing import Any, cast
from src.core.spatial_index import DIM_NAMES, pack_coords, unpack_coords
@dataclass(slots=True)
class StimulusResult:
tick: int
mode: str
target_ids: tuple[int, ...]
amplitudes: tuple[float, ...]
total_injected: float
class StimulusEngine:
"""Configurable stimulus injection engine.
The ``config`` dict must contain a ``"diagnostics"`` key with the
stimulus parameters (mode, amplitude, timing, target coordinates, etc.).
"""
def __init__(self, config: dict[str, Any], rng: random.Random) -> None:
self.config = config
self.rng = rng
diag_raw = config.get("diagnostics", {})
if not isinstance(diag_raw, dict):
raise ValueError("config['diagnostics'] must be a dictionary")
self.diag: dict[str, Any] = cast("dict[str, Any]", diag_raw)
def apply(self, network: Any, tick: int) -> StimulusResult:
mode: str | None = cast("str | None", self.diag.get("mode"))
start = int(cast("int | float", self.diag.get("start_tick", 0)))
duration = int(cast("int | float", self.diag.get("duration_ticks", 0)))
amp = float(cast("int | float", self.diag.get("amplitude", 0.0)))
targets: list[int] = []
amps: list[float] = []
def add(nid: int, current: float) -> None:
if nid in network.neurons:
network.inject_current(nid, current)
targets.append(nid)
amps.append(current)
if mode == "single_pulse":
if tick == start:
target_coord = cast(
"list[Any] | tuple[Any, ...]", self.diag.get("target_coord", ())
)
add(pack_coords(*tuple(target_coord)), amp)
elif mode == "single_neuron_drive":
if start <= tick < start + duration:
target_coord = cast(
"list[Any] | tuple[Any, ...]", self.diag.get("target_coord", ())
)
add(pack_coords(*tuple(target_coord)), amp)
elif mode in ("input_plane_pulse", "input_plane_drive"):
active = (
tick == start
if mode == "input_plane_pulse"
else start <= tick < start + duration
)
if active:
dim_name = cast("str", self.diag.get("input_plane_dim", "x"))
dim = DIM_NAMES.get(dim_name, 0)
for nid in network.neurons:
if unpack_coords(nid)[dim] == 0:
add(nid, amp)
elif mode == "poisson_noise":
if tick >= start:
rate_hz = float(
cast("int | float", self.diag.get("poisson_rate_hz", 0.0))
)
p = rate_hz * network.dt_ms / 1000.0
p = min(max(p, 0.0), 1.0)
noise_amp = float(
cast("int | float", self.diag.get("poisson_amplitude", amp))
)
for nid in network.neurons:
if self.rng.random() < p:
add(nid, noise_amp)
elif mode in ("none", None):
pass
else:
raise ValueError(f"Unknown diagnostic mode: {mode}")
return StimulusResult(tick, str(mode), tuple(targets), tuple(amps), sum(amps))
|