Spiking neural net example

examples/snn-fully-connected.yaml

@@ -0,0 +1,71 @@
+noob_id: snn-fully-connected
+noob_model: noob.TubeSpecification
+noob_version: 1000.0.0
+description: |
+  A fully connected LIF spiking neural network with one neuron with constant input
+
+  Demonstration of:
+  - using a custom runner! this requires the ZMQFreeRunner in nobes.snn
+  - completely ignoring noob's epoch structure and just using its eventing backbone
+  - time-based, queue-driven non-alignment, where all nodes just run as fast as possible
+    but still accept input whenever it exists
+  - cyclic graphs!!!!
+  - tube params!
+
+  Note that this demo does **not make a whole lot of sense** - 
+  in no universe would you want to actually run a neural net where every neuron runs in a separate process
+  and 99.999999% of the compute effort is spent serializing and deserializing messages between neurons
+  whose internal operation takes a microsecond.
+  This is just a demo of how you can use and abuse noob by messing with its internals,
+  taking bits and pieces as needed.
+params:
+  cycles_allowed: true
+
+
+nodes:
+  a:
+    type: nobes.snn.LIFNeuron
+    params:
+      I_static: 1
+      weights: {I_static: 1}
+    depends:
+      - b: b.spike
+      - c: c.spike
+      - d: d.spike
+      - e: e.spike
+      - f: f.spike
+      - g: g.spike
+      - h: h.spike
+      - i: i.spike
+      - j: j.spike
+  b:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, c: c.spike, d: d.spike, e: e.spike, f: f.spike, g: g.spike, h: h.spike, i: i.spike, j: j.spike]
+  c:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, b: b.spike, d: d.spike, e: e.spike, f: f.spike, g: g.spike, h: h.spike, i: i.spike, j: j.spike]
+  d:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, b: b.spike, c: c.spike, e: e.spike, f: f.spike, g: g.spike, h: h.spike, i: i.spike, j: j.spike]
+  e:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, b: b.spike, c: c.spike, d: d.spike, f: f.spike, g: g.spike, h: h.spike, i: i.spike, j: j.spike]
+  f:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, b: b.spike, c: c.spike, d: d.spike, e: e.spike, g: g.spike, h: h.spike, i: i.spike, j: j.spike]
+  g:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, b: b.spike, c: c.spike, d: d.spike, e: e.spike, f: f.spike, h: h.spike, i: i.spike, j: j.spike]
+  h:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, b: b.spike, c: c.spike, d: d.spike, e: e.spike, f: f.spike, g: g.spike, i: i.spike, j: j.spike]
+  i:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, b: b.spike, c: c.spike, d: d.spike, e: e.spike, f: f.spike, g: g.spike, h: h.spike, j: j.spike]
+  j:
+    type: nobes.snn.LIFNeuron
+    depends: [a: a.spike, b: b.spike, c: c.spike, d: d.spike, e: e.spike, f: f.spike, g: g.spike, h: h.spike, i: i.spike]
+
+  monitor:
+    type: nobes.snn.Monitor
+    depends:  [a: a.voltage, b: b.voltage, c: c.voltage, d: d.voltage, e: e.voltage, f: f.voltage, g: g.voltage, h: h.voltage, i: i.voltage, j: j.voltage]

packages/nobes-snn/.copier-answers.yml

@@ -0,0 +1,6 @@
+# Changes here will be overwritten by Copier
+_src_path: templates/nobes
+author_email: sneakers-the-rat@protonmail.com
+author_name: sneakers-the-rat
+description: spiking neural nets
+package_name: snn

packages/nobes-snn/README.md

@@ -0,0 +1,3 @@
+# nobes-snn
+
+spiking neural nets

packages/nobes-snn/nobes/snn/__init__.py

@@ -0,0 +1,10 @@
+from nobes.snn.lif import LIFNeuron
+from nobes.snn.monitor import Monitor
+from nobes.snn.runner import NodeFreeRunner, ZMQFreeRunner
+
+__all__ = [
+    "LIFNeuron",
+    "Monitor",
+    "NodeFreeRunner",
+    "ZMQFreeRunner",
+]

packages/nobes-snn/nobes/snn/hooks.py

@@ -0,0 +1,6 @@
+from nobes.snn.runner import ZMQFreeRunner
+from noob.runner import TubeRunner
+
+
+def add_runners() -> dict[str, type[TubeRunner]]:
+    return {"zmq-freerun": ZMQFreeRunner}

packages/nobes-snn/nobes/snn/lif.py

@@ -0,0 +1,67 @@
+import time
+from random import random
+from typing import Annotated as A
+
+from pydantic import Field
+
+from noob import Name, Node, NodeSpecification
+from noob.event import MetaSignal
+from noob.node import Slot
+from noob.node.base import all_slots_optional
+
+
+class LIFNeuron(Node):
+    """
+    A leaky integrate and fire neuron!!!
+
+    I strongly doubt this implementation is right, for example it doesn't exactly leak,
+    but hey.
+    """
+
+    spec: NodeSpecification  # must have a spec to run LIFNeuron
+    weights: dict[str, float] = Field(default_factory=dict)
+    resistance: float = 10_000_000
+    capacitance: float = 10
+    v_membrane: float = -70e-3
+    v_rest: float = -70e-3
+    v_reset: float = -80e-3
+    I_static: float = 0
+    threshold: float = 10e-3
+    spike_current: float = 10
+    last_step: float = 0
+
+    def init(self) -> None:
+        for dep in self.spec.depends:
+            key, val = next(iter(dep.items()))
+            if key not in self.weights:
+                self.weights[key] = random()
+
+    def process(
+        self, **kwargs: float
+    ) -> tuple[A[float, Name("voltage")], A[float | MetaSignal, Name("spike")]]:
+        # ignore the first step, nothing matters in the first step
+        if self.last_step == 0:
+            self.last_step = time.time()
+            return self.v_membrane, MetaSignal.NoEvent
+
+        current = self.I_static
+        for k, v in kwargs.items():
+            current += self.weights.get(k, 1) * v
+
+        now = time.time()
+        dt = now - self.last_step
+        tau = self.resistance * self.capacitance
+        self.last_step = now
+        self.v_membrane += (
+            ((self.v_rest - self.v_membrane) + (self.resistance * current)) * dt
+        ) / tau
+
+        spike = MetaSignal.NoEvent
+        if self.v_membrane > self.threshold:
+            self.v_membrane = self.v_reset
+            spike = self.spike_current
+
+        return self.v_membrane, spike
+
+    def _collect_slots(self) -> dict[str, Slot]:
+        return all_slots_optional(self.spec)

packages/nobes-snn/nobes/snn/monitor.py

@@ -0,0 +1,55 @@
+import logging
+
+from pydantic import PrivateAttr
+from rich.console import Console
+from rich.progress import BarColumn, Progress, TaskID, TextColumn
+from rich.table import Column
+
+from noob import Node, NodeSpecification
+from noob.logging import init_logger
+from noob.node import Slot
+from noob.node.base import all_slots_optional
+
+
+class Monitor(Node):
+    """A monitor that can display numerical values on the CLI baby"""
+
+    spec: NodeSpecification  # specs not optional for the monitor
+    range: tuple[float, float] = (-70e-3, 10e-3)
+
+    _console: Console | None = None
+    _progress: Progress | None = None
+    _tasks: dict[str, TaskID] = PrivateAttr(default_factory=dict)
+    _logger: logging.Logger = init_logger("monitor")
+
+    def process(self, **kwargs: float) -> None:
+        for key, val in kwargs.items():
+            if key not in self._tasks:
+                continue
+            self._progress.update(
+                self._tasks[key], completed=val - self.range[0], mv=round(val * 1000, ndigits=2)
+            )
+        self._progress.refresh()
+
+    def init(self) -> None:
+        self._console = Console()
+        self._progress = Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(bar_width=None, table_column=Column(ratio=4)),
+            TextColumn("{task.fields[mv]} mV", table_column=Column(ratio=1)),
+            auto_refresh=False,
+            console=self._console,
+        )
+        for dep in self.spec.depends:
+            key = list(dep.keys())[0]
+            self._tasks[key] = self._progress.add_task(
+                key, total=self.range[1] - self.range[0], mv=0
+            )
+
+        self._progress.start()
+
+    def deinit(self) -> None:
+        self._progress.stop()
+
+    def _collect_slots(self) -> dict[str, Slot]:
+        return all_slots_optional(self.spec)

packages/nobes-snn/nobes/snn/runner.py

@@ -0,0 +1,83 @@
+import asyncio
+import contextlib
+import time
+from collections import defaultdict, deque
+from collections.abc import AsyncGenerator, Generator
+from typing import Any, ClassVar
+
+from noob.event import MetaSignal
+from noob.network.message import EventMsg
+from noob.runner.base import TubeRunner
+from noob.runner.zmq import NodeRunner, ZMQRunner
+from noob.types import Epoch, ReturnNodeType
+
+
+class NodeFreeRunner(NodeRunner):
+    def __init__(self, interval: float = 0.01, *args: Any, **kwargs: Any):
+        """
+        Args:
+            interval (float): minimum time (in seconds) that each iteration should take.
+                (i.e., if the iteration takes less time than this, sleep for the difference)
+        """
+        super().__init__(*args, **kwargs)
+        self._interval = interval
+        self._signals = defaultdict(deque)
+        self._event_lock = asyncio.Lock()
+
+    async def await_inputs(self) -> AsyncGenerator[tuple[tuple[Any], dict[str, Any], Epoch]]:
+        """epochs are meaningless to free runner"""
+        current_epoch = 0
+
+        while not self._quitting.is_set():
+            await self._freerun.wait()
+            epoch = Epoch(current_epoch)
+            input_events = []
+            async with self._event_lock:
+                for dep in self.depends:
+                    with contextlib.suppress(IndexError):
+                        input_events.append(self._signals[dep].popleft())
+            inputs = self.store.transform_events(self._node.edges, input_events)
+            args, kwargs = self.store.split_args_kwargs(inputs)
+            yielded_time = time.time()
+
+            yield args, kwargs, epoch
+
+            runtime = time.time() - yielded_time
+            await asyncio.sleep(max(0.0, self._interval - runtime))
+            current_epoch += 1
+
+    async def on_event(self, msg: EventMsg) -> None:
+        """Just stack the events up in queues, we don't care about epochs here"""
+        depended_events = [
+            e
+            for e in msg.value
+            if (e["node_id"], e["signal"]) in self.depends and e["value"] is not MetaSignal.NoEvent
+        ]
+
+        async with self._event_lock:
+            for e in depended_events:
+                self._signals[(e["node_id"], e["signal"])].append(e)
+
+
+class ZMQFreeRunner(ZMQRunner):
+    noderunner_cls: ClassVar[type[TubeRunner]] = NodeFreeRunner
+
+    def run(self, n: int | None = None) -> None | list[ReturnNodeType]:
+        if n is not None:
+            raise NotImplementedError(
+                "What do you mean freerun for only a little bit you are either in or you're out"
+            )
+        super().run()
+
+    def iter(self, n: int | None = None) -> Generator[ReturnNodeType, None, None]:
+        raise NotImplementedError(
+            "iteration mode doesn't really make sense for the free runner! use run!"
+        )
+
+    def process(self, **kwargs: Any) -> ReturnNodeType:
+        raise NotImplementedError(
+            "process mode doesn't really make sense for the free runner! use run!"
+        )
+
+    async def on_event(self, msg: EventMsg) -> None:
+        pass

packages/nobes-snn/pyproject.toml

@@ -0,0 +1,53 @@
+[project]
+name = "nobes-snn"
+description = """spiking neural nets"""
+authors = [
+  {name = "sneakers-the-rat", email = "sneakers-the-rat@protonmail.com"},
+]
+dynamic=["version"]
+requires-python = ">=3.11"
+readme = "README.md"
+license = {text = "EUPL-1.2"}
+keywords = [
+  "noob",
+  "noob-nodes",
+  "nobes",
+]
+dependencies = [
+  "noob[zmq]",
+  "rich>=14.3.3",
+]
+
+[project.optional-dependencies]
+tests = [
+    "pytest",
+]
+dev = [
+    "nobes-snn[tests]",
+]
+
+[project.entry-points."noob.add_runners"]
+runners = "nobes.snn.hooks:add_runners"
+
+[project.urls]
+homepage = "https://noob.readthedocs.io"
+documentation = "https://noob.readthedocs.io"
+repository = "https://github.com/miniscope/noob"
+changelog = "https://noob.readthedocs.io/en/latest/changelog.html"
+
+[build-system]
+requires = ["pdm-backend"]
+build-backend = "pdm.backend"
+
+[tool.pdm]
+distribution = true
+
+[tool.pdm.build]
+includes = ["nobes"]
+
+[tool.pdm.version]
+# versions from tags like "v0.1.2"
+fallback_version = "0.0.0"
+source = "scm"
+tag_filter="nobes-snn-v*"
+tag_regex = '^nobes-snn-v(?:\D*)?(?P<version>([1-9][0-9]*!)?(0|[1-9][0-9]*)(\.(0|[1-9][0-9]*))*((a|b|c|rc)(0|[1-9][0-9]*))?(\.post(0|[1-9][0-9]*))?(\.dev(0|[1-9][0-9]*))?$)$'