plotting

25@dataclass(frozen=True)
26class RasterGroup:
27    """
28    Definition of a neuron group for spike raster plotting.
29
30    Parameters
31    ----------
32    name:
33        Unique identifier for the group. Used to derive default labels.
34    ids:
35        Specification of neuron membership. May be a slice, range, iterable of ids,
36        NumPy array, or callable returning a boolean mask when applied to neuron ids.
37    color:
38        Matplotlib-compatible color specification for the group's spikes.
39    marker:
40        Marker symbol passed to ``Axes.scatter``.
41    size:
42        Marker size passed to ``Axes.scatter``.
43    label:
44        Optional display label overriding automatic label resolution.
45
46    Examples
47    --------
48    ```python
49    groups = [
50        RasterGroup("exc_a", ids=range(0, 5), color="#1f77b4", label="Exc A"),
51        RasterGroup("inh", ids=range(5, 7), color="#8B0000", label="Inh"),
52    ]
53    ```
54
55    Expected output
56    ---------------
57    Passing `groups` into `plot_spike_raster(...)` draws each group with its own
58    color, marker, and label.
59    """
60
61    name: str
62    ids: GroupIndexer
63    color: str = "black"
64    marker: str = "."
65    size: float = 4.0
66    label: Optional[str] = None

 69@dataclass
 70class RasterLabels:
 71    """
 72    Configuration for annotating neuron groups within a raster plot.
 73
 74    Parameters
 75    ----------
 76    show:
 77        Whether to annotate groups.
 78    mapping:
 79        Explicit mapping of group name -> label text.
 80    excitatory:
 81        Fallback text for groups whose name starts with ``"exc"``.
 82    inhibitory:
 83        Fallback text for groups whose name starts with ``"inh"``.
 84    location:
 85        ``\"right\"`` (default) or ``\"left\"`` indicating where labels should be placed
 86        relative to the plot area.
 87    offset:
 88        Fraction of the x-range used to offset labels from the axis boundary.
 89    kwargs:
 90        Additional keyword arguments forwarded to ``Axes.text``.
 91
 92    Examples
 93    --------
 94    ```python
 95    labels = RasterLabels(
 96        mapping={"exc_a": "Exc A", "inh": "Inh"},
 97        location="right",
 98        kwargs={"fontsize": 9},
 99    )
100    ```
101
102    Expected output
103    ---------------
104    The raster receives one text label per group at the chosen side of the
105    axes.
106    """
107
108    show: bool = True
109    mapping: Mapping[str, str] = field(default_factory=dict)
110    excitatory: Optional[str] = None
111    inhibitory: Optional[str] = None
112    location: str = "right"
113    offset: float = 0.02
114    kwargs: Mapping[str, Any] = field(default_factory=dict)
115
116    def resolve_label(self, group: RasterGroup) -> Optional[str]:
117        if not self.show:
118            return None
119        if group.name in self.mapping:
120            return self.mapping[group.name]
121        if group.label is not None:
122            return group.label
123        gname = group.name.lower()
124        if self.excitatory and gname.startswith("exc"):
125            return self.excitatory
126        if self.inhibitory and gname.startswith("inh"):
127            return self.inhibitory
128        return None

131def plot_spike_raster(
132    ax: plt.Axes,
133    spike_times_ms: Sequence[float],
134    spike_ids: Sequence[int],
135    *,
136    n_exc: Optional[int] = None,
137    n_inh: Optional[int] = None,
138    groups: Optional[Sequence[RasterGroup]] = None,
139    stride: int = 1,
140    t_start: Optional[float] = None,
141    t_end: Optional[float] = None,
142    align_time: Optional[float] = None,
143    time_reference: str = "absolute",
144    reference_time: Optional[float] = None,
145    marker: str = ".",
146    marker_size: float = 4.0,
147    exc_color: str = "black",
148    inh_color: str = "#8B0000",
149    labels: Optional[RasterLabels] = None,
150) -> plt.Axes:
151    """
152    Plot a configurable spike raster on *ax*.
153
154    Parameters
155    ----------
156    ax:
157        Target Matplotlib axes.
158    spike_times_ms, spike_ids:
159        1-D sequences of spike times (ms) and corresponding neuron ids.
160    n_exc, n_inh:
161        Sizes of excitatory and inhibitory populations used when *groups* is not provided.
162    groups:
163        Optional explicit group definitions overriding *n_exc* / *n_inh*.
164    stride:
165        Keep every ``stride``-th neuron id (e.g., ``stride=10`` shows every 10th neuron).
166    t_start, t_end:
167        Optional temporal window (after alignment) to display.
168    align_time:
169        Time (ms) subtracted from all spike times prior to plotting.
170    time_reference:
171        Either ``\"absolute\"`` (default) or ``\"relative\"``. When ``\"relative\"`` and
172        *reference_time* is ``None``, the minimum time after alignment is used as reference.
173    reference_time:
174        Time (ms) used as reference when ``time_reference=\"relative\"``.
175    marker, marker_size:
176        Defaults for marker appearance when *groups* is not provided.
177    exc_color, inh_color:
178        Default colors for excitatory/inhibitory groups.
179    labels:
180        Optional :class:`RasterLabels` controlling group annotations.
181
182    Examples
183    --------
184    ```python
185    fig, ax = plt.subplots(figsize=(4, 2))
186    groups = [
187        RasterGroup("exc_a", ids=range(0, 3), color="#1f77b4", label="Exc A"),
188        RasterGroup("exc_b", ids=range(3, 5), color="#2ca02c", label="Exc B"),
189        RasterGroup("inh", ids=range(5, 7), color="#8B0000", label="Inh"),
190    ]
191    plot_spike_raster(
192        ax,
193        spike_times_ms=[5, 8, 11, 13, 21, 23, 29],
194        spike_ids=[0, 1, 2, 3, 4, 5, 6],
195        groups=groups,
196        labels=RasterLabels(location="right", kwargs={"fontsize": 8}),
197    )
198    ```
199
200    Expected output
201    ---------------
202    The axes contains one grouped raster with group-specific colors and three
203    group labels at the right margin.
204
205    ![Spike raster example](plotting_assets/grouped_spike_raster_example.png)
206    """
207    times = np.asarray(spike_times_ms, dtype=float)
208    neuron_ids = np.asarray(spike_ids, dtype=int)
209    if times.shape != neuron_ids.shape:
210        raise ValueError("spike_times_ms and spike_ids must have matching shapes.")
211
212    if align_time is not None:
213        times = times - float(align_time)
214
215    time_reference = time_reference.lower()
216    if time_reference not in {"absolute", "relative"}:
217        raise ValueError("time_reference must be 'absolute' or 'relative'.")
218    if time_reference == "relative":
219        ref = reference_time
220        if ref is None:
221            ref = times.min() if times.size else 0.0
222        times = times - float(ref)
223
224    # Temporal selection
225    mask = np.ones(times.shape, dtype=bool)
226    if t_start is not None:
227        mask &= times >= float(t_start)
228    if t_end is not None:
229        mask &= times <= float(t_end)
230
231    # Stride selection
232    stride = max(int(stride), 1)
233    if stride > 1:
234        mask &= (neuron_ids % stride) == 0
235
236    times = times[mask]
237    neuron_ids = neuron_ids[mask]
238
239    resolved_groups = _resolve_groups(
240        groups=groups,
241        n_exc=n_exc,
242        n_inh=n_inh,
243        marker=marker,
244        marker_size=marker_size,
245        exc_color=exc_color,
246        inh_color=inh_color,
247    )
248
249    for group in resolved_groups:
250        group_mask = _evaluate_group_mask(group.ids, neuron_ids)
251        if not np.any(group_mask):
252            continue
253        ax.scatter(
254            times[group_mask],
255            neuron_ids[group_mask],
256            s=float(group.size) ** 2,
257            c=group.color,
258            marker=group.marker,
259            linewidths=0,
260            edgecolors="none",
261        )
262
263    # Limits
264    if times.size:
265        xmin = 0.0 if t_start is None else float(t_start)
266        xmax = float(times.max() if t_end is None else t_end)
267        if t_start is not None:
268            xmin = float(t_start)
269        if t_end is not None:
270            xmax = float(t_end)
271        if xmin == xmax:
272            xmax = xmin + 1.0
273        ax.set_xlim(xmin, xmax)
274    else:
275        xmin = float(t_start if t_start is not None else 0.0)
276        xmax = float(t_end if t_end is not None else xmin + 1.0)
277        if xmin == xmax:
278            xmax = xmin + 1.0
279        ax.set_xlim(xmin, xmax)
280
281    if neuron_ids.size:
282        ymin = neuron_ids.min()
283        ymax = neuron_ids.max()
284    else:
285        ymin, ymax = _group_global_bounds(resolved_groups)
286    if ymin == ymax:
287        ymax = ymin + 1
288    padding = max(1, int((ymax - ymin) * 0.02))
289    ax.set_ylim(ymin - padding, ymax + padding)
290
291    if labels is not None and labels.show:
292        _apply_group_labels(ax, resolved_groups, labels)
293
294    return ax

 27@dataclass
 28class BinaryStateSource:
 29    """Describe how binary neuron states are provided for raster plotting.
 30
 31    Parameters
 32    ----------
 33    inline_states:
 34        Optional in-memory state matrix with shape `(steps, neurons)`.
 35    chunk_files:
 36        Optional paths to `.npy` or `.npz` chunks storing state matrices.
 37    neuron_count:
 38        Total number of neurons represented by the source.
 39    update_log, delta_log:
 40        Optional diff-log representation where each column corresponds to a
 41        simulation step.
 42    initial_state:
 43        Optional initial state for diff-log based sources.
 44
 45    Examples
 46    --------
 47    Wrap an in-memory state matrix:
 48
 49    >>> states = np.array([[0, 1], [1, 1], [1, 0]], dtype=np.uint8)
 50    >>> source = BinaryStateSource.from_array(states)
 51    >>> source.neuron_count
 52    2
 53
 54    Wrap diff logs directly:
 55
 56    >>> updates = np.array([[0, 1], [1, 0]], dtype=np.uint16)
 57    >>> deltas = np.array([[1, -1], [0, 1]], dtype=np.int8)
 58    >>> source = BinaryStateSource.from_diff_logs(updates, deltas, neuron_count=2)
 59    >>> source.update_log.shape
 60    (2, 2)
 61    """
 62
 63    inline_states: Optional[np.ndarray] = None
 64    chunk_files: Sequence[Path] = field(default_factory=tuple)
 65    neuron_count: int = 0
 66    update_log: Optional[np.ndarray] = None
 67    delta_log: Optional[np.ndarray] = None
 68    initial_state: Optional[np.ndarray] = None
 69
 70    def iter_chunks(self) -> Iterator[np.ndarray]:
 71        """Yield state chunks with shape `(steps, neurons)`."""
 72        if self.inline_states is not None and self.inline_states.size:
 73            yield np.asarray(self.inline_states, dtype=np.uint8)
 74            return
 75        for entry in self.chunk_files:
 76            path = Path(entry)
 77            if not path.exists():
 78                continue
 79            data = np.load(path, allow_pickle=False, mmap_mode="r")
 80            yield np.asarray(data, dtype=np.uint8)
 81
 82    @classmethod
 83    def from_array(cls, states: Union[np.ndarray, Iterable[Sequence[int]]]) -> "BinaryStateSource":
 84        """Convenience helper for wrapping an in-memory state matrix.
 85
 86        Examples
 87        --------
 88        >>> source = BinaryStateSource.from_array([[0, 1], [1, 1]])
 89        >>> source.neuron_count
 90        2
 91        """
 92        array = np.asarray(states)
 93        neuron_count = int(array.shape[1]) if array.ndim == 2 else 0
 94        return cls(inline_states=array, chunk_files=tuple(), neuron_count=neuron_count)
 95
 96    @classmethod
 97    def from_diff_logs(
 98        cls,
 99        updates: np.ndarray,
100        deltas: np.ndarray,
101        *,
102        neuron_count: int,
103        initial_state: Optional[np.ndarray] = None,
104    ) -> "BinaryStateSource":
105        """Construct a source from diff-log traces.
106
107        Examples
108        --------
109        >>> source = BinaryStateSource.from_diff_logs(
110        ...     np.array([[0, 1]], dtype=np.uint16),
111        ...     np.array([[1, -1]], dtype=np.int8),
112        ...     neuron_count=2,
113        ... )
114        >>> source.neuron_count
115        2
116        """
117        update_arr = np.asarray(updates, dtype=np.uint16)
118        delta_arr = np.asarray(deltas, dtype=np.int8)
119        init_state = None
120        if initial_state is not None:
121            init_state = np.asarray(initial_state, dtype=np.uint8)
122        return cls(
123            inline_states=None,
124            chunk_files=tuple(),
125            neuron_count=int(neuron_count),
126            update_log=update_arr,
127            delta_log=delta_arr,
128            initial_state=init_state,
129        )

132def collect_binary_onset_events(
133    state_source: BinaryStateSource,
134    sample_interval: int,
135    *,
136    window: Optional[Tuple[float, float]] = None,
137) -> tuple[np.ndarray, np.ndarray]:
138    """Collect onset events from a binary state stream.
139
140    Parameters
141    ----------
142    state_source:
143        Binary state source, either as full states or diff logs.
144    sample_interval:
145        Time step between successive recorded samples.
146    window:
147        Optional `(start, end)` filter in the same units as `sample_interval`.
148
149    Returns
150    -------
151    tuple[np.ndarray, np.ndarray]
152        Event times and neuron ids suitable for `plot_spike_raster(...)`.
153
154    Examples
155    --------
156    >>> source = BinaryStateSource.from_array(np.array([[0, 0], [1, 0], [1, 1]], dtype=np.uint8))
157    >>> times, ids = collect_binary_onset_events(source, sample_interval=5)
158    >>> times.tolist()
159    [5.0, 10.0]
160    >>> ids.tolist()
161    [0, 1]
162    """
163    sample_interval = max(1, int(sample_interval))
164    window_start = float(window[0]) if window else None
165    window_end = float(window[1]) if window else None
166    if state_source.update_log is not None and state_source.delta_log is not None:
167        return _collect_from_diff_log(state_source, sample_interval, window=window)
168    times: list[np.ndarray] = []
169    neurons: list[np.ndarray] = []
170    prev_state: Optional[np.ndarray] = None
171    sample_index = 0
172    for chunk in state_source.iter_chunks():
173        block = np.asarray(chunk, dtype=np.uint8)
174        if block.ndim != 2 or block.shape[0] == 0:
175            continue
176        for row in block:
177            sample_index += 1
178            if prev_state is None:
179                prev_state = row.copy()
180                continue
181            transitions = (prev_state == 0) & (row == 1)
182            transition_time = (sample_index - 1) * sample_interval
183            prev_state = row.copy()
184            if window_end is not None and transition_time > window_end:
185                return _finalize(times, neurons)
186            if not transitions.any():
187                continue
188            if window_start is not None and transition_time < window_start:
189                continue
190            idx = np.flatnonzero(transitions)
191            if idx.size == 0:
192                continue
193            times.append(np.full(idx.size, transition_time, dtype=np.float64))
194            neurons.append(idx.astype(np.int64))
195    return _finalize(times, neurons)

198def plot_binary_raster(
199    ax: Axes,
200    *,
201    state_source: BinaryStateSource,
202    sample_interval: int,
203    n_exc: int,
204    n_inh: Optional[int] = None,
205    total_neurons: Optional[int] = None,
206    window: Optional[Tuple[float, float]] = None,
207    time_scale: float = 1.0,
208    stride: int = 1,
209    labels: Optional[RasterLabels] = None,
210    groups: Optional[Sequence[RasterGroup]] = None,
211    marker: str = ".",
212    marker_size: float = 4.0,
213    empty_text: str = "No neuron onset events",
214    **raster_kwargs,
215) -> tuple[np.ndarray, np.ndarray]:
216    """
217    Plot a binary-network onset raster.
218
219    The function first extracts onset events from the provided binary state
220    source and then forwards those events to `plot_spike_raster(...)`.
221
222    Examples
223    --------
224    ```python
225    fig, ax = plt.subplots(figsize=(4, 2))
226    source = BinaryStateSource.from_array(
227        np.array([[0, 0], [1, 0], [1, 1], [0, 1]], dtype=np.uint8)
228    )
229    times, ids = plot_binary_raster(
230        ax,
231        state_source=source,
232        sample_interval=10,
233        n_exc=1,
234        n_inh=1,
235    )
236    ```
237
238    Expected output
239    ---------------
240    `times` is `array([10., 20.])` and `ids` is `array([0, 1])`. The axes
241    shows two onset markers at those coordinates.
242
243    ![Binary raster example](plotting_assets/binary_raster_example.png)
244    """
245    events = collect_binary_onset_events(state_source, sample_interval, window=window)
246    spike_times, spike_ids = events
247    if spike_times.size == 0 or spike_ids.size == 0:
248        ax.text(0.5, 0.5, empty_text, ha="center", va="center", transform=ax.transAxes)
249        ax.set_axis_off()
250        return events
251    safe_scale = time_scale if time_scale > 0 else 1.0
252    scaled_times = spike_times / safe_scale
253    t_start = (window[0] / safe_scale) if window else None
254    t_end = (window[1] / safe_scale) if window else None
255
256    total = total_neurons or state_source.neuron_count or 0
257    n_exc = max(0, int(n_exc))
258    if n_inh is None:
259        n_inh = max(0, total - n_exc)
260    n_inh = max(0, int(n_inh))
261
262    plot_spike_raster(
263        ax=ax,
264        spike_times_ms=scaled_times,
265        spike_ids=spike_ids,
266        n_exc=n_exc,
267        n_inh=n_inh,
268        groups=groups,
269        stride=max(1, int(stride)),
270        t_start=t_start,
271        t_end=t_end,
272        marker=marker,
273        marker_size=marker_size,
274        labels=labels,
275        **raster_kwargs,
276    )
277    if t_start is not None or t_end is not None:
278        xmin = t_start if t_start is not None else ax.get_xlim()[0]
279        xmax = t_end if t_end is not None else ax.get_xlim()[1]
280        if xmin == xmax:
281            xmax = xmin + 1.0
282        ax.set_xlim(xmin, xmax)
283    return events

34def add_image_ax(ax, path, label=None, fc: Optional[FontCfg]=None, *, pdf_page: int = 0, pdf_zoom: float = 2.0):
35    """Render an image into an axes and optionally add a panel label.
36
37    Parameters
38    ----------
39    ax:
40        Target Matplotlib axes.
41    path:
42        Image or PDF path.
43    label:
44        Optional panel label placed near the upper-left corner.
45    fc:
46        Optional `FontCfg` used for the label size.
47    pdf_page, pdf_zoom:
48        PDF rendering options used when `path` points to a PDF.
49
50    Examples
51    --------
52    ```python
53    fig, ax = plt.subplots(figsize=(3, 2))
54    add_image_ax(ax, "docs/plotting_assets/spike_raster_example.png", label="A")
55    ```
56
57    Expected output
58    ---------------
59    The axes displays the image and, when `label` is provided, a bold panel
60    label is drawn near the upper-left corner.
61
62    ![add_image_ax example](plotting_assets/add_image_ax_example.png)
63    """
64    ax.axis('off')
65    try:
66        img = _imread_any(path, pdf_page=pdf_page, pdf_zoom=pdf_zoom)
67        h, w = img.shape[0], img.shape[1]
68        ax.imshow(img)
69        ax.set_box_aspect(h / w if w else 1)
70    except FileNotFoundError:
71        ax.text(0.5, 0.5, f"Missing image:\n{path}", ha='center', va='center', fontsize=10)
72        ax.set_box_aspect(1)
73    except Exception as e:
74        ax.text(0.5, 0.5, f"Error loading:\n{os.path.basename(path)}\n{e}", ha='center', va='center', fontsize=10)
75        ax.set_box_aspect(1)
76    if label:
77        label_cfg = fc if fc is not None else FontCfg().resolve()
78        add_panel_label(ax, label, label_cfg)
79    return ax

20@dataclass
21class FontCfg:
22    """Font-size configuration shared across multi-panel figures.
23
24    Call `resolve()` once after construction to derive unset sizes from the
25    `base` and `scale` values.
26
27    Examples
28    --------
29    >>> cfg = FontCfg(base=10.0, scale=1.2).resolve()
30    >>> round(cfg.label, 1), round(cfg.letter, 1)
31    (12.0, 13.2)
32    """
33    base: float = 12.0
34    scale: float = 1.4
35    title: float = None
36    label: float = None
37    tick: float = None
38    legend: float = None
39    panel: float = None
40    labelpad: float = 6.0
41    letter: float = None          # NEW: subplot letter size
42
43    def resolve(self):
44        """Fill unset size fields from the base configuration."""
45        if self.title  is None: self.title  = self.base * self.scale * 1.20
46        if self.label  is None: self.label  = self.base * self.scale * 1.00
47        if self.tick   is None: self.tick   = self.base * self.scale * 0.95
48        if self.legend is None: self.legend = self.base * self.scale * 0.95
49        if self.panel  is None: self.panel  = self.base * self.scale * 0.95
50        if self.letter is None: self.letter = self.base * self.scale * 1.10
51        return self

54def add_corner_tag(ax, text, color, fc: FontCfg, *, x=0.985, y=0.985):
55    """Add a boxed annotation tag in the upper-right corner of an axes.
56
57    Expected output
58    ---------------
59    The axes receives one bold text box near its upper-right corner.
60
61    ![add_corner_tag example](plotting_assets/add_corner_tag_example.png)
62    """
63    ax.text(
64        x, y, text,
65        transform=ax.transAxes,
66        ha="right", va="top",
67        fontsize=fc.label, fontweight="bold", color=color,
68        bbox=dict(
69            facecolor="white",
70            edgecolor="darkgrey",
71            linewidth=0.9,
72            boxstyle="square,pad=0.25"
73        ),
74        zorder=10,
75        clip_on=False
76    )

78def add_panel_label(ax, text, fc: FontCfg, *, x=-0.12, y=1.03):
79    """Add a bold panel label in axes coordinates.
80
81    Expected output
82    ---------------
83    One bold label appears slightly above and left of the axes, matching the
84    figure panel style used across this repository.
85
86    ![add_panel_label example](plotting_assets/add_panel_label_example.png)
87    """
88    ax.text(x, y, text, transform=ax.transAxes,
89            ha="left", va="top", fontsize=fc.letter, fontweight="bold", clip_on=False)