spikeforge 0.3.0__py3-none-any.whl

main.py

@@ -0,0 +1,26 @@
+"""Entry point: train a rate-coded MNIST subset and export all visuals."""
+
+from spikeforge.exporters.presentation_exporter import (
+    PresentationGifExporter,
+)
+from spikeforge.exporters.raster_exporter import RasterExporter
+from spikeforge.exporters.reconstruction_exporter import (
+    ReconstructionExporter,
+)
+from spikeforge.exporters.spike_gif_exporter import SpikeGifExporter
+from spikeforge.exporters.video_exporter import VideoExporter
+from spikeforge.training.logger import SNNTrainerLogger
+
+
+def main() -> None:
+    """Build the trainer and write every output artifact."""
+    trainer = SNNTrainerLogger(animation_interval=100)
+    VideoExporter(trainer).export()
+    SpikeGifExporter(trainer).export()
+    ReconstructionExporter(trainer).export()
+    RasterExporter(trainer).export()
+    PresentationGifExporter(trainer).export()
+
+
+if __name__ == "__main__":
+    main()

main_encodings.py

@@ -0,0 +1,43 @@
+"""Run the additional tutorial-1 encoding demos (latency/delta/random).
+
+Produces latency curve + rasters + animation, delta plots, and random
+spike visuals under build/.
+"""
+
+from spikeforge.encoding.delta_trainer import DeltaTrainer
+from spikeforge.encoding.latency_trainer import LatencyTrainer
+from spikeforge.encoding.random_spikegen import RandomSpikeGenerator
+from spikeforge.exporters.delta_exporter import DeltaPlotExporter
+from spikeforge.exporters.latency_curve_exporter import (
+    LatencyCurveExporter,
+)
+from spikeforge.exporters.latency_raster_exporter import (
+    LatencyRasterExporter,
+)
+from spikeforge.exporters.latency_video_exporter import (
+    LatencyVideoExporter,
+)
+from spikeforge.exporters.random_spike_raster_exporter import (
+    RandomSpikeRasterExporter,
+)
+from spikeforge.exporters.random_spike_video_exporter import (
+    RandomSpikeVideoExporter,
+)
+
+
+def main() -> None:
+    """Train the latency/delta/random encoders and export their visuals."""
+    latency_trainer = LatencyTrainer(animation_interval=100)
+    LatencyCurveExporter(latency_trainer).export()
+    LatencyRasterExporter(latency_trainer).export()
+    LatencyVideoExporter(latency_trainer).export(key="clip")
+
+    DeltaPlotExporter(DeltaTrainer()).export()
+
+    random_spikegen = RandomSpikeGenerator(num_steps=100)
+    RandomSpikeRasterExporter(random_spikegen).export()
+    RandomSpikeVideoExporter(random_spikegen).export()
+
+
+if __name__ == "__main__":
+    main()

spikeforge/__init__.py

@@ -0,0 +1,19 @@
+"""Public exports for the spikeforge package."""
+
+from spikeforge.encoding.delta_trainer import DeltaTrainer
+from spikeforge.encoding.latency_trainer import (
+    LatencyTrainer,
+    convert_to_time,
+)
+from spikeforge.encoding.random_spikegen import RandomSpikeGenerator
+from spikeforge.training.logger import SNNTrainerLogger
+from spikeforge.training.trainer import SSNTrainer
+
+__all__ = [
+    "SSNTrainer",
+    "SNNTrainerLogger",
+    "LatencyTrainer",
+    "DeltaTrainer",
+    "RandomSpikeGenerator",
+    "convert_to_time",
+]

spikeforge/benchmark/__init__.py

@@ -0,0 +1,44 @@
+"""Performance and memory benchmarks for the interpreter's execution modes.
+
+This package measures wall time and memory of a topology run under
+production and educational modes (and, opt-in, a compiled production path).
+It is deliberately separate from ``introspection``, which describes a
+recorded trajectory rather than the cost of producing one.
+
+Use :func:`run_benchmark` from Python, or ``python -m
+spikeforge.benchmark`` from a shell; both return the same
+JSON-serialisable report.
+"""
+
+from spikeforge.benchmark.compare import (
+    DEFAULT_THRESHOLD,
+    compare_runs,
+    exit_code,
+)
+from spikeforge.benchmark.config import BenchmarkConfig, default_config
+from spikeforge.benchmark.harness import run_benchmark
+from spikeforge.benchmark.serving import (
+    ServingBenchmarkConfig,
+    percentile,
+    run_serving_benchmark,
+    run_serving_suite,
+)
+from spikeforge.benchmark.store import BenchmarkStore, default_directory
+from spikeforge.benchmark.suite import run_suite, with_metadata
+
+__all__ = [
+    "BenchmarkConfig",
+    "BenchmarkStore",
+    "DEFAULT_THRESHOLD",
+    "ServingBenchmarkConfig",
+    "compare_runs",
+    "default_config",
+    "default_directory",
+    "exit_code",
+    "percentile",
+    "run_benchmark",
+    "run_serving_benchmark",
+    "run_serving_suite",
+    "run_suite",
+    "with_metadata",
+]

spikeforge/benchmark/__main__.py

@@ -0,0 +1,14 @@
+"""Module entry point: ``python -m spikeforge.benchmark``.
+
+The parser and handlers live in
+:mod:`spikeforge.benchmark.cli`; this module only exposes ``main`` so
+the module runner and the ``spikeforge-benchmark`` console script share
+one path.
+"""
+
+import sys
+
+from spikeforge.benchmark.cli import main
+
+if __name__ == "__main__":
+    sys.exit(main())

spikeforge/benchmark/cli.py

@@ -0,0 +1,192 @@
+"""Command-line interface for the benchmark suite.
+
+``python -m spikeforge.benchmark`` (or the ``spikeforge-benchmark`` script)
+runs a fixture and prints JSON. ``--save``/``--list`` persist and list runs
+through the :class:`~spikeforge.benchmark.store.BenchmarkStore`, and
+``--compare`` diffs a stored baseline against a fresh or stored run. When
+``--fail-on-regression`` is given the command exits non-zero on a regression,
+so it works as a CI gate.
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional
+
+from spikeforge.benchmark.compare import (
+    DEFAULT_THRESHOLD,
+    compare_runs,
+    exit_code,
+)
+from spikeforge.benchmark.config import BenchmarkConfig
+from spikeforge.benchmark.harness import run_benchmark
+from spikeforge.benchmark.serving import (
+    ServingBenchmarkConfig,
+    run_serving_benchmark,
+    run_serving_suite,
+)
+from spikeforge.benchmark.store import BenchmarkStore
+from spikeforge.benchmark.suite import with_metadata
+
+
+def _add_fixture_args(parser: argparse.ArgumentParser) -> None:
+    """Register the fixture-shaping arguments."""
+    parser.add_argument("--topology", action="append", dest="topologies")
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--steps", type=int, default=8)
+    parser.add_argument("--repeats", type=int, default=3)
+    parser.add_argument("--warmup", type=int, default=1)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--compiled", action="store_true")
+    parser.add_argument("--no-backward", action="store_true")
+    parser.add_argument("--energy", action="store_true")
+    parser.add_argument("--energy-target", dest="energy_target",
+                        default="reference")
+    parser.add_argument("--out", default=None)
+
+
+def _add_serving_args(parser: argparse.ArgumentParser) -> None:
+    """Register the serving-mode arguments."""
+    parser.add_argument(
+        "--serving",
+        action="store_true",
+        help="benchmark an in-process serving bundle instead of a topology",
+    )
+    parser.add_argument(
+        "--bundle",
+        default=None,
+        help="path to the .spkf deployment bundle served",
+    )
+    parser.add_argument("--calls", type=int, default=8)
+    parser.add_argument("--concurrency", type=int, default=1)
+
+
+def _add_store_args(parser: argparse.ArgumentParser) -> None:
+    """Register the store, listing, and comparison arguments."""
+    parser.add_argument("--save", action="store_true")
+    parser.add_argument("--label", default=None)
+    parser.add_argument("--store-dir", dest="store_dir", default=None)
+    parser.add_argument("--list", action="store_true", dest="list_runs")
+    parser.add_argument("--compare", default=None)
+    parser.add_argument("--against", default=None)
+    parser.add_argument("--threshold", type=float, default=DEFAULT_THRESHOLD)
+    parser.add_argument("--fail-on-regression", action="store_true")
+
+
+def _parser() -> argparse.ArgumentParser:
+    """Return the argument parser for the benchmark CLI."""
+    parser = argparse.ArgumentParser(
+        prog="spikeforge-benchmark",
+        description="Benchmark interpreter execution modes.",
+    )
+    _add_fixture_args(parser)
+    _add_serving_args(parser)
+    _add_store_args(parser)
+    return parser
+
+
+def _config(args: argparse.Namespace) -> BenchmarkConfig:
+    """Build a :class:`BenchmarkConfig` from parsed CLI arguments."""
+    return BenchmarkConfig(
+        topologies=tuple(args.topologies or ("fc_small",)),
+        batch_size=args.batch_size,
+        steps=args.steps,
+        repeats=args.repeats,
+        warmup=args.warmup,
+        seed=args.seed,
+        device=args.device,
+        compiled=args.compiled,
+        backward=not args.no_backward,
+        energy=args.energy,
+        energy_target=args.energy_target,
+    )
+
+
+def _serving_config(args: argparse.Namespace) -> ServingBenchmarkConfig:
+    """Build a :class:`ServingBenchmarkConfig` from parsed arguments."""
+    if not args.bundle:
+        raise SystemExit("--serving requires --bundle PATH")
+    return ServingBenchmarkConfig(
+        bundle=args.bundle,
+        calls=args.calls,
+        concurrency=args.concurrency,
+        warmup=args.warmup,
+        seed=args.seed,
+        device=args.device,
+    )
+
+
+def _store(args: argparse.Namespace) -> BenchmarkStore:
+    """Return the store configured by ``--store-dir``."""
+    return BenchmarkStore(args.store_dir)
+
+
+def _emit(payload: Mapping[str, Any], out: Optional[str]) -> int:
+    """Print ``payload`` or write it to ``out``; return success."""
+    text = json.dumps(payload, indent=2)
+    if out:
+        Path(out).write_text(text, encoding="utf-8")
+    else:
+        print(text)
+    return 0
+
+
+def _run_list(args: argparse.Namespace) -> int:
+    """Print the stored run summaries as JSON."""
+    return _emit({"runs": _store(args).list_runs()}, args.out)
+
+
+def _run_fixture(args: argparse.Namespace) -> int:
+    """Run the fixture, optionally save it, and emit the report."""
+    report = run_benchmark(_config(args))
+    if args.save:
+        report = with_metadata(report)
+        report["run_id"] = _store(args).save(report, label=args.label)
+    return _emit(report, args.out)
+
+
+def _run_serving(args: argparse.Namespace) -> int:
+    """Run the serving fixture, optionally save it, and emit the report."""
+    config = _serving_config(args)
+    if args.save:
+        report = run_serving_suite(
+            config, store=_store(args), label=args.label or "serving"
+        )
+    else:
+        report = run_serving_benchmark(config)
+    return _emit(report, args.out)
+
+
+def _candidate(args: argparse.Namespace) -> Dict[str, Any]:
+    """Return a stored run when ``--against`` is set, else a fresh run."""
+    if args.against:
+        return _store(args).load(args.against)
+    if args.serving:
+        return run_serving_benchmark(_serving_config(args))
+    return run_benchmark(_config(args))
+
+
+def _run_compare(args: argparse.Namespace) -> int:
+    """Compare a baseline run and print the deltas plus a status."""
+    baseline = _store(args).load(args.compare)
+    result = compare_runs(baseline, _candidate(args), args.threshold)
+    print(json.dumps(result, indent=2))
+    return exit_code(result) if args.fail_on_regression else 0
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    """Parse ``argv`` and dispatch list, compare, or a fixture run."""
+    args = _parser().parse_args(argv)
+    if args.list_runs:
+        return _run_list(args)
+    if args.compare:
+        return _run_compare(args)
+    if args.serving:
+        return _run_serving(args)
+    return _run_fixture(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

spikeforge/benchmark/compare.py

@@ -0,0 +1,192 @@
+"""Compare two benchmark records and flag regressions past a threshold.
+
+The comparator matches records on ``(topology, mode)`` and reports the
+relative change of the throughput, latency, and memory metrics. A candidate
+regresses when a throughput metric moves the wrong way by more than the
+threshold, or when its peak memory or p99 latency grows by more than it. The
+result is JSON-able, and :func:`exit_code` turns it into a CI gate.
+
+Serving records (:mod:`spikeforge.benchmark.serving`) carry an extra
+``serving`` block; its ``p99_ms`` and ``throughput_per_second`` are compared
+here too, and are simply absent (and therefore skipped) on training records.
+"""
+
+from typing import Any, Dict, List, Mapping, Optional, Tuple
+
+#: Default relative change (fraction) that counts as a regression.
+DEFAULT_THRESHOLD = 0.10
+
+#: (reported metric, direction) where "lower" means a larger value is worse.
+_METRICS: Tuple[Tuple[str, str], ...] = (
+    ("ms_per_step", "lower"),
+    ("steps_per_second", "higher"),
+    ("p99_ms", "lower"),
+    ("throughput_per_second", "higher"),
+    ("memory_bytes", "lower"),
+)
+
+#: Serving metrics read from the record's ``serving`` block, not ``forward``.
+_SERVING_KEYS: Tuple[str, ...] = ("p99_ms", "throughput_per_second")
+
+#: Memory keys probed in order; the first non-zero one is compared.
+_MEMORY_KEYS = (
+    "tracemalloc_peak_bytes",
+    "process_rss_bytes",
+    "cuda_peak_bytes",
+)
+
+
+def _case_key(record: Mapping[str, Any]) -> Tuple[str, str]:
+    """Return the ``(topology, mode)`` key a record is matched on."""
+    return str(record.get("topology")), str(record.get("mode"))
+
+
+def _memory(record: Mapping[str, Any]) -> Optional[float]:
+    """Return the first non-zero memory metric for a record, or None."""
+    block = record.get("memory") or {}
+    for key in _MEMORY_KEYS:
+        value = block.get(key)
+        if value:
+            return float(value)
+    return None
+
+
+def _value(record: Mapping[str, Any], metric: str) -> Optional[float]:
+    """Return the numeric value of ``metric`` for a record, or None."""
+    if metric == "memory_bytes":
+        return _memory(record)
+    if metric in _SERVING_KEYS:
+        value = (record.get("serving") or {}).get(metric)
+        return None if value is None else float(value)
+    key = "mean_ms_per_step" if metric == "ms_per_step" else metric
+    value = (record.get("forward") or {}).get(key)
+    return None if value is None else float(value)
+
+
+def _change(
+    candidate: Optional[float], baseline: Optional[float]
+) -> Optional[float]:
+    """Return the relative change from baseline to candidate, or None."""
+    if not baseline or candidate is None or baseline is None:
+        return None
+    return (candidate - baseline) / baseline
+
+
+def _regressed(
+    direction: str, fraction: Optional[float], threshold: float
+) -> bool:
+    """Return True when ``fraction`` breaches ``threshold`` the wrong way."""
+    if fraction is None:
+        return False
+    if direction == "lower":
+        return fraction > threshold
+    return fraction < -threshold
+
+
+def _entry(
+    metric: str,
+    direction: str,
+    baseline: Mapping[str, Any],
+    candidate: Mapping[str, Any],
+    threshold: float,
+) -> Dict[str, Any]:
+    """Return one metric comparison entry."""
+    base_value = _value(baseline, metric)
+    cand_value = _value(candidate, metric)
+    fraction = _change(cand_value, base_value)
+    return {
+        "metric": metric,
+        "baseline": base_value,
+        "candidate": cand_value,
+        "change_fraction": fraction,
+        "regressed": _regressed(direction, fraction, threshold),
+    }
+
+
+def _case(
+    baseline: Mapping[str, Any],
+    candidate: Mapping[str, Any],
+    threshold: float,
+) -> Dict[str, Any]:
+    """Compare one matched topology/mode pair."""
+    entries = [
+        _entry(metric, direction, baseline, candidate, threshold)
+        for metric, direction in _METRICS
+    ]
+    return {
+        "topology": baseline.get("topology"),
+        "mode": baseline.get("mode"),
+        "metrics": entries,
+        "regressed": any(entry["regressed"] for entry in entries),
+    }
+
+
+def _identity(record: Mapping[str, Any]) -> Dict[str, Any]:
+    """Return the identifying fields of a run."""
+    return {
+        "run_id": record.get("run_id"),
+        "label": record.get("label"),
+        "created_at": record.get("created_at"),
+    }
+
+
+def _name(case: Mapping[str, Any]) -> str:
+    """Return a ``topology:mode`` label for a regressed case."""
+    return f"{case['topology']}:{case['mode']}"
+
+
+def _pairs(
+    baseline: Mapping[str, Any], candidate: Mapping[str, Any]
+) -> List[Any]:
+    """Return matched ``(baseline, candidate)`` records in baseline order."""
+    base_map = {_case_key(r): r for r in baseline.get("results", [])}
+    cand_map = {_case_key(r): r for r in candidate.get("results", [])}
+    return [
+        (base_map[key], cand_map[key])
+        for key in base_map
+        if key in cand_map
+    ]
+
+
+def _result(
+    cases: List[Dict[str, Any]],
+    threshold: float,
+    baseline: Mapping[str, Any],
+    candidate: Mapping[str, Any],
+) -> Dict[str, Any]:
+    """Assemble the JSON-able comparison result."""
+    regressions: List[str] = [
+        _name(case) for case in cases if case["regressed"]
+    ]
+    return {
+        "threshold": float(threshold),
+        "compared": len(cases),
+        "cases": cases,
+        "regressions": regressions,
+        "regressed": bool(regressions),
+        "baseline": _identity(baseline),
+        "candidate": _identity(candidate),
+    }
+
+
+def compare_runs(
+    baseline: Mapping[str, Any],
+    candidate: Mapping[str, Any],
+    threshold: float = DEFAULT_THRESHOLD,
+) -> Dict[str, Any]:
+    """Return per-config deltas and whether ``candidate`` regressed.
+
+    ``baseline`` and ``candidate`` are reports from
+    :func:`~spikeforge.benchmark.harness.run_benchmark` (or a stored
+    run). Only configs present in both are compared.
+    """
+    cases = [
+        _case(base, cand, threshold)
+        for base, cand in _pairs(baseline, candidate)
+    ]
+    return _result(cases, threshold, baseline, candidate)
+
+
+def exit_code(result: Mapping[str, Any]) -> int:
+    """Return the process exit status: 1 on regression, else 0."""
+    return 1 if result.get("regressed") else 0

spikeforge/benchmark/config.py

@@ -0,0 +1,40 @@
+"""Configuration for a topology benchmark run.
+
+The benchmark lives in its own package, not under ``introspection``: it
+measures wall time and memory of an execution run rather than describing a
+recorded trajectory, and it exposes its own tiny, JSON-returning fixture.
+"""
+
+from dataclasses import dataclass
+from typing import Tuple
+
+
+@dataclass(frozen=True)
+class BenchmarkConfig:
+    """A small, explicit benchmark fixture.
+
+    ``topologies`` are registry topology names. ``steps`` and ``batch_size``
+    shape the random spike input, ``seed`` fixes both the module weights and
+    the spikes, and ``warmup`` untimed runs precede ``repeats`` timed samples.
+    ``compiled`` additionally measures a ``torch.compile`` production path and
+    ``backward`` measures a backward pass where one is meaningful. ``energy``
+    (default off) additionally attaches an event-driven energy estimate for
+    ``energy_target`` to every record; the fields are additive.
+    """
+
+    topologies: Tuple[str, ...] = ("fc_small",)
+    batch_size: int = 2
+    steps: int = 4
+    repeats: int = 2
+    warmup: int = 1
+    seed: int = 0
+    device: str = "auto"
+    compiled: bool = False
+    backward: bool = True
+    energy: bool = False
+    energy_target: str = "reference"
+
+
+def default_config() -> BenchmarkConfig:
+    """Return the tiny default fixture so tests and CI stay fast."""
+    return BenchmarkConfig()

spikeforge/benchmark/energy.py

@@ -0,0 +1,31 @@
+"""Opt-in event-driven energy block for the benchmark harness.
+
+Kept out of :mod:`spikeforge.benchmark.harness` so that module stays
+within the project's file-length limit; the harness calls this only when its
+config asks for energy, and the returned block is additive on the record.
+"""
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from spikeforge.benchmark.config import BenchmarkConfig
+from spikeforge.topology.stage_module import StageModule
+from spikeforge_targets.energy.accounting import account
+from spikeforge_targets.event_runtime.sparse_runner import sparse_run
+
+
+def energy_block(
+    module: StageModule, spikes: torch.Tensor, config: BenchmarkConfig
+) -> Optional[Dict[str, Any]]:
+    """Return the opt-in event-driven energy estimate, or None."""
+    if not config.energy:
+        return None
+    with torch.no_grad():
+        result = sparse_run(module, spikes)
+    report = account(result, config.energy_target)
+    return {
+        "report": report.to_dict(),
+        "counts": dict(result.counts),
+        "density": result.density,
+    }