rnba 0.2.1__tar.gz → 0.2.2__tar.gz

{rnba-0.2.1 → rnba-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: rnba
-Version: 0.2.1
+Version: 0.2.2
 Summary: Recurrent Neural Blackboard Architecture - a simulation framework for neuro-cognitive modelling.
 Author-Email: RNBA Contributors <m.dekamps@leeds.ac.uk>
 License: MIT
@@ -163,6 +163,8 @@ python -m rnba.visualize_activity_flow \
 
 To zoom in on the top of the blackboard, add `--first-rows N` (e.g. `--first-rows 5`): the view is cropped to the first N rows and markers are enlarged for legibility, while the layout and trace mapping stay identical to the full 15-row view. This works for both the single-frame preview (`--preview-time-ms`) and full video rendering, and mirrors the `--first-rows` option of the static schematic renderer (see below).
 
+Frame rendering dominates the runtime. Pass `--workers N` to render frames across N processes (each holds its own copy of the base schematic); output is identical to a sequential run. `--dpi` trades resolution for size but has little effect on speed, since most per-frame cost is drawing the schematic rather than rasterising pixels.
+
 See `python -m rnba.visualize_activity_flow --help` for the full option set (codec, styling, frame selection).
 
 ## Quick Start: Replicating the QA Simulation

{rnba-0.2.1 → rnba-0.2.2}/README.md

@@ -126,6 +126,8 @@ python -m rnba.visualize_activity_flow \
 
 To zoom in on the top of the blackboard, add `--first-rows N` (e.g. `--first-rows 5`): the view is cropped to the first N rows and markers are enlarged for legibility, while the layout and trace mapping stay identical to the full 15-row view. This works for both the single-frame preview (`--preview-time-ms`) and full video rendering, and mirrors the `--first-rows` option of the static schematic renderer (see below).
 
+Frame rendering dominates the runtime. Pass `--workers N` to render frames across N processes (each holds its own copy of the base schematic); output is identical to a sequential run. `--dpi` trades resolution for size but has little effect on speed, since most per-frame cost is drawing the schematic rather than rasterising pixels.
+
 See `python -m rnba.visualize_activity_flow --help` for the full option set (codec, styling, frame selection).
 
 ## Quick Start: Replicating the QA Simulation

{rnba-0.2.1 → rnba-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "rnba"
-version = "0.2.1"
+version = "0.2.2"
 description = "Recurrent Neural Blackboard Architecture - a simulation framework for neuro-cognitive modelling."
 readme = "README.md"
 requires-python = ">=3.8"

{rnba-0.2.1 → rnba-0.2.2}/rnba/__init__.py

@@ -3,7 +3,7 @@
 Simple package with a greet() helper used in tests and CLI.
 """
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
 
 def greet(name: str = "RNBA") -> str:

{rnba-0.2.1 → rnba-0.2.2}/rnba/visualize_activity_flow.py

@@ -14,11 +14,12 @@ from __future__ import annotations
 import argparse
 import csv
 import io
+import multiprocessing as mp
 import re
 import subprocess
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, Iterable, List, Optional
+from typing import Dict, Iterable, List, Optional, Tuple
 
 import matplotlib.pyplot as plt
 from matplotlib.collections import LineCollection
@@ -1255,24 +1256,33 @@ def write_dynamics_audit_csv(
             )
 
 
-def render_preview(
+@dataclass
+class _OverlayContext:
+    """Static rendering state shared by every frame of one video.
+
+    The base schematic figure is expensive to build (thousands of artists);
+    it is created once here, and per-frame overlays are added and removed
+    on top of it.
+    """
+
+    fig: "plt.Figure"
+    ax: "plt.Axes"
+    symbols: List[Symbol]
+    edges: List[SymbolEdge]
+    by_symbol: Dict[str, MappingResult]
+    symbol_pos: Dict[str, Tuple[float, float]]
+    hidden_ids: set
+    overlay_marker_scale: float
+    overlay_line_scale: float
+
+
+def _build_overlay_context(
     symbols: List[Symbol],
     mappings: List[MappingResult],
     edges: List[SymbolEdge],
-    trace_values: Dict[int, float],
-    out_png: Path,
-    vmax: float,
-    base_grey: str,
-    base_alpha: float,
-    dpi: int,
-    base_size: float,
-    activation_threshold: float,
-    activation_gamma: float,
-    active_size_max: float,
-    supersample_scale: float,
-    first_rows: Optional[int] = None,
-    crop_edges: bool = False,
-) -> None:
+    first_rows: Optional[int],
+    crop_edges: bool,
+) -> _OverlayContext:
     by_symbol = {m.symbol_id: m for m in mappings}
     symbol_pos = {s.symbol_id: (s.x, s.y) for s in symbols}
     hidden_families = {"WORD_INHIB"}
@@ -1298,7 +1308,41 @@ def render_preview(
         show_control_nodes=True,
         activity_flow_theme=True,
     )
-    ax = fig.axes[0]
+    return _OverlayContext(
+        fig=fig,
+        ax=fig.axes[0],
+        symbols=symbols,
+        edges=edges,
+        by_symbol=by_symbol,
+        symbol_pos=symbol_pos,
+        hidden_ids=hidden_ids,
+        overlay_marker_scale=overlay_marker_scale,
+        overlay_line_scale=overlay_line_scale,
+    )
+
+
+def _draw_activity_overlay(
+    ctx: _OverlayContext,
+    trace_values: Dict[int, float],
+    vmax: float,
+    activation_threshold: float,
+    activation_gamma: float,
+    active_size_max: float,
+) -> list:
+    """Draw one frame's activity overlay onto the base figure.
+
+    Returns the created artists so a video loop can ``remove()`` them
+    before drawing the next frame.
+    """
+    ax = ctx.ax
+    by_symbol = ctx.by_symbol
+    symbol_pos = ctx.symbol_pos
+    hidden_ids = ctx.hidden_ids
+    edges = ctx.edges
+    symbols = ctx.symbols
+    overlay_marker_scale = ctx.overlay_marker_scale
+    overlay_line_scale = ctx.overlay_line_scale
+    artists: list = []
 
     # Source representative activity drives outgoing edge intensity.
     src_activity: Dict[str, float] = {}
@@ -1349,6 +1393,7 @@ def render_preview(
                 zorder=5,
             )
             ax.add_collection(active_edge_collection)
+            artists.append(active_edge_collection)
 
     x_act: List[float] = []
     y_act: List[float] = []
@@ -1398,8 +1443,18 @@ def render_preview(
         s_act.append(size)
 
     if x_act:
-        ax.scatter(x_act, y_act, s=s_act, c=c_act, edgecolors="none", zorder=6)
+        scatter_artist = ax.scatter(x_act, y_act, s=s_act, c=c_act, edgecolors="none", zorder=6)
+        artists.append(scatter_artist)
+
+    return artists
 
+
+def _save_frame(
+    fig: "plt.Figure",
+    out_png: Path,
+    dpi: int,
+    supersample_scale: float,
+) -> None:
     out_png.parent.mkdir(parents=True, exist_ok=True)
     scale = supersample_scale if supersample_scale > 1.0 else 1.0
     target_w = max(1, int(round(fig.get_figwidth() * dpi)))
@@ -1432,7 +1487,39 @@ def render_preview(
             fig.savefig(out_png, dpi=dpi)
     else:
         fig.savefig(out_png, dpi=super_dpi)
-    plt.close(fig)
+
+
+def render_preview(
+    symbols: List[Symbol],
+    mappings: List[MappingResult],
+    edges: List[SymbolEdge],
+    trace_values: Dict[int, float],
+    out_png: Path,
+    vmax: float,
+    base_grey: str,
+    base_alpha: float,
+    dpi: int,
+    base_size: float,
+    activation_threshold: float,
+    activation_gamma: float,
+    active_size_max: float,
+    supersample_scale: float,
+    first_rows: Optional[int] = None,
+    crop_edges: bool = False,
+) -> None:
+    ctx = _build_overlay_context(symbols, mappings, edges, first_rows, crop_edges)
+    try:
+        _draw_activity_overlay(
+            ctx,
+            trace_values,
+            vmax=vmax,
+            activation_threshold=activation_threshold,
+            activation_gamma=activation_gamma,
+            active_size_max=active_size_max,
+        )
+        _save_frame(ctx.fig, out_png, dpi=dpi, supersample_scale=supersample_scale)
+    finally:
+        plt.close(ctx.fig)
 
 
 def _iter_sampled_trace_rows(
@@ -1479,6 +1566,54 @@ def _row_to_values(row: List[str], col_to_idx: Dict[int, int]) -> Dict[int, floa
     return values
 
 
+# ----------------------------------------------------------------------
+# Parallel frame rendering
+# ----------------------------------------------------------------------
+# Each worker process builds its own base schematic once (the figure is a
+# live matplotlib object and cannot cross the process boundary), then
+# renders the frames it is handed. Frames are fully independent, so this
+# scales close to linearly with the worker count.
+
+_FRAME_WORKER_STATE: Dict[str, object] = {}
+
+
+def _frame_worker_init(
+    symbols, mappings, edges, first_rows, crop_edges,
+    frames_dir, dpi, supersample_scale,
+    vmax, activation_threshold, activation_gamma, active_size_max,
+) -> None:
+    ctx = _build_overlay_context(symbols, mappings, edges, first_rows, crop_edges)
+    _FRAME_WORKER_STATE.update(
+        ctx=ctx,
+        frames_dir=frames_dir,
+        dpi=dpi,
+        supersample_scale=supersample_scale,
+        vmax=vmax,
+        activation_threshold=activation_threshold,
+        activation_gamma=activation_gamma,
+        active_size_max=active_size_max,
+    )
+
+
+def _frame_worker_render(task) -> int:
+    idx, values = task
+    st = _FRAME_WORKER_STATE
+    ctx = st["ctx"]
+    overlay_artists = _draw_activity_overlay(
+        ctx,
+        values,
+        vmax=st["vmax"],
+        activation_threshold=st["activation_threshold"],
+        activation_gamma=st["activation_gamma"],
+        active_size_max=st["active_size_max"],
+    )
+    frame_path = st["frames_dir"] / f"frame_{idx:06d}.png"
+    _save_frame(ctx.fig, frame_path, dpi=st["dpi"], supersample_scale=st["supersample_scale"])
+    for artist in overlay_artists:
+        artist.remove()
+    return idx
+
+
 def render_video(
     symbols: List[Symbol],
     mappings: List[MappingResult],
@@ -1510,6 +1645,7 @@ def render_video(
     video_out: Path,
     first_rows: Optional[int] = None,
     crop_edges: bool = False,
+    workers: int = 1,
 ) -> None:
     out_dir.mkdir(parents=True, exist_ok=True)
     frames_dir.mkdir(parents=True, exist_ok=True)
@@ -1522,31 +1658,57 @@ def render_video(
         raise ValueError(f"No sampled rows found in {trace_csv}")
 
     total = len(sampled_rows)
+
+    # Build per-frame activity maps sequentially, up front: the optional
+    # input aligner advances a monotonic cursor, so it must see frames in
+    # time order regardless of how they are later rendered.
+    frame_tasks: List[Tuple[int, Dict[int, float]]] = []
     for idx, (t_ms, row) in enumerate(sampled_rows):
         values = _row_to_values(row, col_to_idx)
         if input_aligner is not None:
             values.update(input_aligner.values_at(t_ms))
-        frame_path = frames_dir / f"frame_{idx:06d}.png"
-        render_preview(
-            symbols=symbols,
-            mappings=mappings,
-            edges=edges,
-            trace_values=values,
-            out_png=frame_path,
-            vmax=vmax,
-            base_grey=base_grey,
-            base_alpha=base_alpha,
-            dpi=dpi,
-            base_size=base_size,
-            activation_threshold=activation_threshold,
-            activation_gamma=activation_gamma,
-            active_size_max=active_size_max,
-            supersample_scale=supersample_scale,
-            first_rows=first_rows,
-            crop_edges=crop_edges,
+        frame_tasks.append((idx, values))
+
+    workers = max(1, int(workers))
+    if workers > 1:
+        workers = min(workers, total)
+        print(f"Rendering {total} frames with {workers} workers ...")
+        init_args = (
+            symbols, mappings, edges, first_rows, crop_edges,
+            frames_dir, dpi, supersample_scale,
+            vmax, activation_threshold, activation_gamma, active_size_max,
         )
-        if (idx + 1) % 25 == 0 or (idx + 1) == total:
-            print(f"Rendered {idx + 1}/{total} frames")
+        done = 0
+        with mp.Pool(
+            processes=workers,
+            initializer=_frame_worker_init,
+            initargs=init_args,
+        ) as pool:
+            for _ in pool.imap_unordered(_frame_worker_render, frame_tasks, chunksize=4):
+                done += 1
+                if done % 25 == 0 or done == total:
+                    print(f"Rendered {done}/{total} frames")
+    else:
+        # Single-process path: reuse one base schematic across all frames.
+        ctx = _build_overlay_context(symbols, mappings, edges, first_rows, crop_edges)
+        try:
+            for idx, values in frame_tasks:
+                overlay_artists = _draw_activity_overlay(
+                    ctx,
+                    values,
+                    vmax=vmax,
+                    activation_threshold=activation_threshold,
+                    activation_gamma=activation_gamma,
+                    active_size_max=active_size_max,
+                )
+                frame_path = frames_dir / f"frame_{idx:06d}.png"
+                _save_frame(ctx.fig, frame_path, dpi=dpi, supersample_scale=supersample_scale)
+                for artist in overlay_artists:
+                    artist.remove()
+                if (idx + 1) % 25 == 0 or (idx + 1) == total:
+                    print(f"Rendered {idx + 1}/{total} frames")
+        finally:
+            plt.close(ctx.fig)
 
     cmd = [
         "ffmpeg",
@@ -1661,6 +1823,13 @@ def main() -> None:
     parser.add_argument("--fps", type=int, default=25, help="Output video frames per second")
     parser.add_argument("--step-ms", type=float, default=10.0, help="Trace sampling stride in ms for video frames")
     parser.add_argument("--max-frames", type=int, default=0, help="Optional cap on rendered frames (0 means no cap)")
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help="Parallel frame-rendering worker processes (1 = sequential). "
+             "Each worker holds its own copy of the base schematic.",
+    )
     parser.add_argument(
         "--frames-dir",
         type=Path,
@@ -1774,6 +1943,7 @@ def main() -> None:
             video_out=args.video_out,
             first_rows=args.first_rows,
             crop_edges=args.crop_edges,
+            workers=args.workers,
         )
         print(f"Wrote: {args.video_out}")