"""Interactive Brain 5D observatory with optional X-Y heatmap.""" from __future__ import annotations from pathlib import Path from typing import Any import matplotlib.pyplot as plt import numpy as np from src.core.spatial_index import unpack_coords from src.visualization.heatmap import HeatmapKind, HeatmapProjector, HeatmapView def build_raster_points(frames: Any, sample_ids: list[int]) -> list[tuple[int, int]]: """Convert spike-history frames to raster plot points.""" id_to_row = {neuron_id: row for row, neuron_id in enumerate(sample_ids)} points: list[tuple[int, int]] = [] for frame in frames: for neuron_id in frame.neuron_ids: row = id_to_row.get(neuron_id) if row is not None: points.append((frame.tick, row)) return points class Observatory: """Interactive visual observer for sparse Brain 5D simulations.""" def __init__( self, network: Any, config: dict[str, Any], spike_history: Any, history: Any, probes: Any | None = None, ) -> None: self.net = network self.config = config self.spike_history = spike_history self.history = history self.probes = probes self.dims = tuple(config["dimensions"]) visualization = config["visualization"] self.raster_n = int(visualization["spike_raster_neurons"]) self.proj_dim = str(visualization.get("projection_4d", "d4")) self.tau = float(visualization.get("activity_tau_ticks", 50.0)) self.show_heatmap = bool(visualization.get("show_heatmap", True)) self.heatmap_kind: HeatmapKind = self._parse_heatmap_kind( visualization.get("heatmap_type", "activity") ) plt.ion() # type: ignore[reportUnknownMemberType] self.fig = plt.figure(figsize=(16, 10)) # type: ignore[reportUnknownMemberType] grid = self.fig.add_gridspec(2, 3) # type: ignore[reportUnknownMemberType] self.ax1 = self.fig.add_subplot(grid[0, 0], projection="3d") # type: ignore[reportUnknownMemberType] self.ax2 = self.fig.add_subplot(grid[0, 1]) # type: ignore[reportUnknownMemberType] self.ax_heat = self.fig.add_subplot(grid[0, 2]) # type: ignore[reportUnknownMemberType] self.ax3 = self.fig.add_subplot(grid[1, 0]) # type: ignore[reportUnknownMemberType] self.ax4 = self.fig.add_subplot(grid[1, 1:]) # type: ignore[reportUnknownMemberType] self.scatter_xyz = self.ax1.scatter( # type: ignore[reportUnknownMemberType] [], [], [], c=[], cmap="hot", vmin=0, vmax=1, s=10, alpha=0.7 # type: ignore[reportArgumentType] ) self.scatter_xd = self.ax2.scatter( # type: ignore[reportUnknownMemberType] [], [], c=[], cmap="plasma", vmin=0, vmax=1, s=8, alpha=0.7 ) self.raster_scatter = self.ax3.scatter( # type: ignore[reportUnknownMemberType] [], [], s=8, marker="|" ) (self.line_spikes,) = self.ax4.plot([], [], label="Spikes/tick") # type: ignore[reportUnknownMemberType] (self.line_v,) = self.ax4.plot([], [], label="Mean V") # type: ignore[reportUnknownMemberType] (self.line_queue,) = self.ax4.plot([], [], label="Queue") # type: ignore[reportUnknownMemberType] self.ax4.legend(fontsize=8) # type: ignore[reportUnknownMemberType] self.ax1.set_title("XYZ Activity") # type: ignore[reportUnknownMemberType] self.ax2.set_title(f"X vs {self.proj_dim.upper()}") # type: ignore[reportUnknownMemberType] self.ax3.set_title("Real Spike Raster") # type: ignore[reportUnknownMemberType] self.ax4.set_title("Time Series") # type: ignore[reportUnknownMemberType] self.heatmap_projector = HeatmapProjector(self.net, self.tau) self.heatmap_view = HeatmapView(self.ax_heat) if not self.show_heatmap: self.ax_heat.set_visible(False) self.status_text = self.fig.text(0.02, 0.02, "", family="monospace", fontsize=9) # type: ignore[reportUnknownMemberType] self.probe_text = self.fig.text(0.72, 0.02, "", family="monospace", fontsize=9) # type: ignore[reportUnknownMemberType] self.snapshot_dir = Path("artifacts/snapshots") self.snapshot_dir.mkdir(parents=True, exist_ok=True) self.fig.canvas.mpl_connect("key_press_event", self._on_key) plt.show(block=False) # type: ignore[reportUnknownMemberType] @staticmethod def _parse_heatmap_kind(value: object) -> HeatmapKind: kind = str(value).lower() if kind not in {"activity", "weights", "energy"}: raise ValueError( "visualization.heatmap_type must be activity, weights or energy" ) return kind # type: ignore[return-value] def _on_key(self, event: Any) -> None: if event.key and str(event.key).lower() == "s": self.save_snapshot() def _activity(self, neuron: Any) -> float: if neuron.last_spike_tick < 0: return 0.0 age = max(0, self.net.current_tick - neuron.last_spike_tick) return float(np.exp(-age / self.tau)) def draw(self) -> None: """Refresh all Observatory panels from the current network state.""" sample = list(self.net.neurons.items())[:2000] x_values: list[int] = [] y_values: list[int] = [] z_values: list[int] = [] activity: list[float] = [] xd_values: list[int] = [] yd_values: list[int] = [] activity_2d: list[float] = [] for neuron_id, neuron in sample: x_coord, y_coord, z_coord, d4_coord, d5_coord = unpack_coords(neuron_id) value = self._activity(neuron) x_values.append(x_coord) y_values.append(y_coord) z_values.append(z_coord) activity.append(value) xd_values.append(x_coord) yd_values.append(d4_coord if self.proj_dim == "d4" else d5_coord) activity_2d.append(value) self.scatter_xyz._offsets3d = (x_values, y_values, z_values) # type: ignore[reportAttributeAccessIssue] self.scatter_xyz.set_array(np.asarray(activity)) # type: ignore[reportUnknownMemberType] points_2d = ( np.column_stack((xd_values, yd_values)) if xd_values else np.empty((0, 2)) ) self.scatter_xd.set_offsets(points_2d) # type: ignore[reportUnknownMemberType] self.scatter_xd.set_array(np.asarray(activity_2d)) # type: ignore[reportUnknownMemberType] self.ax1.set( # type: ignore[reportUnknownMemberType] xlim=(0, self.dims[0] - 1), ylim=(0, self.dims[1] - 1), zlim=(0, self.dims[2] - 1), ) projection_size = self.dims[3] if self.proj_dim == "d4" else self.dims[4] self.ax2.set( # type: ignore[reportUnknownMemberType] xlim=(0, self.dims[0] - 1), ylim=(0, projection_size - 1), ) frames = self.spike_history.get_frames()[-100:] sample_ids = list(self.net.neurons)[: self.raster_n] raster_points = build_raster_points(frames, sample_ids) self.raster_scatter.set_offsets( # type: ignore[reportUnknownMemberType] np.asarray(raster_points) if raster_points else np.empty((0, 2)) ) if frames: self.ax3.set_xlim( # type: ignore[reportUnknownMemberType] frames[0].tick, max(frames[0].tick + 1, frames[-1].tick + 1), ) self.ax3.set_ylim(0, max(1, len(sample_ids))) # type: ignore[reportUnknownMemberType] history = self.history.get_all() if history: ticks = [item["tick"] for item in history] self.line_spikes.set_data( # type: ignore[reportUnknownMemberType] ticks, [item["spikes_this_tick"] for item in history] ) self.line_v.set_data( # type: ignore[reportUnknownMemberType] ticks, [item["mean_v"] for item in history] ) self.line_queue.set_data( # type: ignore[reportUnknownMemberType] ticks, [item["queued_events"] for item in history] ) self.ax4.relim() # type: ignore[reportUnknownMemberType] self.ax4.autoscale_view() # type: ignore[reportUnknownMemberType] if self.show_heatmap: self.heatmap_view.render(self.heatmap_projector.build(self.heatmap_kind)) self.status_text.set_text( # type: ignore[reportUnknownMemberType] f"Tick {self.net.current_tick} | neurons {len(self.net.neurons)} | " f"synapses {self.net.synapse_count} | queue {self.net.queued_event_count} | " f"spikes {self.net.total_spikes}" ) if self.probes and self.probes.probes: rows: list[str] = [] for probe_id in self.probes.probes[:3]: data = self.probes.get_probe_data(probe_id) rows.append( f"{probe_id}: v={data.get('v', 0):.1f} " f"spk={data.get('spike_counter', 0)}" ) self.probe_text.set_text("Probes\n" + "\n".join(rows)) # type: ignore[reportUnknownMemberType] self.fig.canvas.draw_idle() # type: ignore[reportUnknownMemberType] plt.pause(0.001) # type: ignore[reportUnknownMemberType] def save_snapshot(self) -> Path: """Save the current Observatory figure and return its path.""" path = self.snapshot_dir / f"brain5d_tick_{self.net.current_tick:06d}.png" self.fig.savefig(path, dpi=150, bbox_inches="tight") # type: ignore[reportUnknownMemberType] return path @staticmethod def block_until_closed() -> None: """Switch Matplotlib to blocking mode until the window is closed.""" plt.ioff() # type: ignore[reportUnknownMemberType] plt.show(block=True) # type: ignore[reportUnknownMemberType]