Plotting

Plotting functions for SIMPL results.

All functions accept an xr.Dataset (the results_ attribute of a fitted SIMPL model) and return matplotlib Axes so users can customise further. matplotlib is imported lazily inside each function so the rest of the package stays lightweight.

outset_axes(ax, offset_mm=2)

Outset bottom/left spines by offset_mm mm.

Source code in src/simpl/plotting.py

def outset_axes(ax, offset_mm: float = 2) -> None:
    """Outset bottom/left spines by *offset_mm* mm."""
    ax.spines["bottom"].set_position(("outward", offset_mm * 72 / 25.4))
    ax.spines["left"].set_position(("outward", offset_mm * 72 / 25.4))

plot_all_metrics(results, show_neurons=True, ncols=3, **plot_kwargs)

Auto-discover and plot all per-iteration metrics.

Parameters:

Name	Type	Description	Default
results	Dataset		required
show_neurons	bool	Show individual neuron dots for per-neuron metrics.	True
ncols	int	Number of columns in the grid.	3
**plot_kwargs		Forwarded to line/scatter calls.	{}

Returns:

Name	Type	Description
axes	np.ndarray of Axes

Source code in src/simpl/plotting.py

def plot_all_metrics(
    results: xr.Dataset,
    show_neurons: bool = True,
    ncols: int = 3,
    **plot_kwargs,
) -> np.ndarray:
    """Auto-discover and plot all per-iteration metrics.

    Parameters
    ----------
    results : xr.Dataset
    show_neurons : bool
        Show individual neuron dots for per-neuron metrics.
    ncols : int
        Number of columns in the grid.
    **plot_kwargs
        Forwarded to line/scatter calls.

    Returns
    -------
    axes : np.ndarray of Axes
    """
    iterations = _non_negative_iterations(results)
    last_iteration = int(iterations[-1])
    baselines = _baseline_iterations(results)

    # discover metric variables: anything with iteration dim and only neuron/place_field remaining
    # skip _val variants — they are plotted alongside their train counterpart
    metric_names = []
    for var_name in results.data_vars:
        da = results[var_name]
        if "iteration" not in da.dims:
            continue
        other_dims = set(da.dims) - {"iteration"}
        if other_dims <= {"neuron", "place_field"}:
            if var_name.endswith("_val") and var_name[:-4] in results.data_vars:
                continue  # will be plotted with the train variant
            metric_names.append(var_name)

    n_metrics = len(metric_names)
    if n_metrics == 0:
        warnings.warn("No metrics found to plot.", stacklevel=2)
        return np.array([])

    nrows = int(np.ceil(n_metrics / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=(3.0 * ncols, 2.5 * nrows), squeeze=False, layout="constrained")

    for i, var_name in enumerate(metric_names):
        ax = axes.flat[i]
        da = results[var_name]
        other_dims = [d for d in da.dims if d != "iteration"]
        attrs = da.attrs
        ylabel = attrs.get("axis_title", attrs.get("axis title", var_name))

        is_scalar = len(other_dims) == 0
        has_place_field = "place_field" in other_dims

        # Only plot iterations that have data for this variable (skip all-NaN iterations)
        var_iterations = []
        for e in iterations:
            vals_e = da.sel(iteration=e)
            if not np.all(np.isnan(vals_e.values)):
                var_iterations.append(e)
        var_iterations = np.array(var_iterations) if var_iterations else iterations

        if is_scalar:
            # line plot
            val_name = f"{var_name}_val"
            has_val = val_name in results.data_vars
            vals = [float(da.sel(iteration=e)) for e in var_iterations]
            vals_v = [float(results[val_name].sel(iteration=e)) for e in var_iterations] if has_val else None
            for j in range(len(var_iterations)):
                c = _iteration_color(var_iterations[j], last_iteration)
                ax.scatter(var_iterations[j], vals[j], color=c, zorder=5, **plot_kwargs)
                if has_val:
                    ax.scatter(
                        var_iterations[j],
                        vals_v[j],
                        color=c,
                        marker="o",
                        facecolors="none",
                        linewidth=1.5,
                        zorder=5,
                        **plot_kwargs,
                    )
            for j in range(len(var_iterations) - 1):
                c = _iteration_color(var_iterations[j + 1], last_iteration)
                ax.plot(var_iterations[j : j + 2], vals[j : j + 2], color=c, lw=0.8, zorder=3)
                if has_val:
                    ax.plot(var_iterations[j : j + 2], vals_v[j : j + 2], color=c, lw=0.8, ls="--", zorder=3)
            # baseline: only iteration -1 ("best model")
            if -1 in baselines:
                y_gt = float(da.sel(iteration=-1))
                ax.axhline(y_gt, color="k", ls="--", lw=0.8)
                ax.text(
                    0.0,
                    y_gt,
                    " ground truth",
                    va="bottom",
                    ha="left",
                    fontsize="x-small",
                    color="k",
                    transform=ax.get_yaxis_transform(),
                )
            # ML baseline for LL / bps panels
            # if var_name in ("bits_per_spike", "logPYXF") and "mode_l" in results and 1 in results.iteration.values:
            # from simpl.utils import get_ML_loglikelihoods

            # ml = get_ML_loglikelihoods(results)
            # ml_key = f"{var_name}_val"
            # if ml_key in ml:
            #     y_ml = ml[ml_key]
            #     ax.axhline(y_ml, color="k", ls=":", lw=0.8)
            #     ax.text(
            #         0.0,
            #         y_ml,
            #         " naive ML",
            #         va="bottom",
            #         ha="left",
            #         fontsize="x-small",
            #         color="k",
            #         transform=ax.get_yaxis_transform(),
            #     )
        else:
            # per-neuron (possibly mean over place_field first)
            means = []
            for e in var_iterations:
                c = _iteration_color(e, last_iteration)
                vals_e = da.sel(iteration=e)
                if has_place_field:
                    vals_e = vals_e.mean(dim="place_field", skipna=True)
                v = vals_e.values
                if show_neurons:
                    jitter = np.random.default_rng(int(e)).uniform(-0.15, 0.15, size=len(v))
                    ax.scatter(e + jitter, v, color=c, alpha=0.15, s=5, linewidths=0)
                if np.all(np.isnan(v)):
                    means.append(np.nan)
                else:
                    means.append(float(np.nanmean(v)))
                ax.scatter(e, means[-1], color=c, s=60, zorder=5, linewidths=0)
            for j in range(len(var_iterations) - 1):
                c = _iteration_color(var_iterations[j + 1], last_iteration)
                ax.plot(var_iterations[j : j + 2], means[j : j + 2], color=c, lw=0.8, zorder=3)

        ax.set(xlabel="Iteration", ylabel=ylabel, xlim=(-0.5, int(iterations[-1]) + 0.5))
        outset_axes(ax)
        ax.spines["bottom"].set_bounds(0, int(iterations[-1]))

    # legend on first panel showing train/val distinction
    first_ax = axes.flat[0]
    first_ax.plot([], [], color="gray", lw=0.8, label="train")
    first_ax.plot([], [], color="gray", lw=0.8, ls="--", label="val")
    first_ax.legend(fontsize="small", frameon=False)

    # hide unused axes
    for j in range(n_metrics, len(axes.flat)):
        axes.flat[j].set_visible(False)

    return axes

plot_fitting_summary(results, show_neurons=True, **plot_kwargs)

Two-panel summary: bits per spike (left) and mutual information (right).

Parameters:

Name	Type	Description	Default
results	Dataset	The results_ Dataset from a fitted SIMPL model.	required
show_neurons	bool	Show individual neuron dots for per-neuron metrics (mutual information).	True
**plot_kwargs		Forwarded to the main scatter calls.	{}

Returns:

Name	Type	Description
axes	np.ndarray of Axes, shape (2,)

Source code in src/simpl/plotting.py

def plot_fitting_summary(
    results: xr.Dataset,
    show_neurons: bool = True,
    **plot_kwargs,
) -> np.ndarray:
    """Two-panel summary: bits per spike (left) and mutual information (right).

    Parameters
    ----------
    results : xr.Dataset
        The ``results_`` Dataset from a fitted SIMPL model.
    show_neurons : bool
        Show individual neuron dots for per-neuron metrics (mutual information).
    **plot_kwargs
        Forwarded to the main scatter calls.

    Returns
    -------
    axes : np.ndarray of Axes, shape (2,)
    """
    iterations = _non_negative_iterations(results)
    last_iteration = int(iterations[-1])

    fig, axes = plt.subplots(1, 2, figsize=(0.7 * FIG_WIDTH, 0.7 * FIG_WIDTH * 0.35), layout="constrained")
    ax_bps, ax_mi = axes

    bps_train, bps_val, mi_means = [], [], []
    for e in iterations:
        c = _iteration_color(e, last_iteration)
        bps_train.append(float(results.bits_per_spike.sel(iteration=e)))
        bps_val.append(float(results.bits_per_spike_val.sel(iteration=e)))
        ax_bps.scatter(e, bps_train[-1], color=c, zorder=5, **plot_kwargs)
        ax_bps.scatter(e, bps_val[-1], color=c, marker="o", facecolors="none", linewidth=1.5, zorder=5, **plot_kwargs)

        mi = results.mutual_information.sel(iteration=e).values
        if show_neurons:
            jitter = np.random.default_rng(int(e)).uniform(-0.15, 0.15, size=len(mi))
            ax_mi.scatter(e + jitter, mi, color=c, alpha=0.15, s=5, linewidths=0)
        mi_means.append(float(np.mean(mi)))
        ax_mi.scatter(e, mi_means[-1], color=c, s=60, zorder=5, linewidths=0)

    # connecting lines
    for i in range(len(iterations) - 1):
        c = _iteration_color(iterations[i + 1], last_iteration)
        ax_bps.plot(iterations[i : i + 2], bps_train[i : i + 2], color=c, lw=0.8, zorder=3)
        ax_bps.plot(iterations[i : i + 2], bps_val[i : i + 2], color=c, lw=0.8, ls="--", zorder=3)
        ax_mi.plot(iterations[i : i + 2], mi_means[i : i + 2], color=c, lw=0.8, zorder=3)

    # baseline: only iteration -1 ("best model")
    if -1 in results.iteration.values:
        if "bits_per_spike" in results:
            y_gt = float(results.bits_per_spike.sel(iteration=-1))
            ax_bps.axhline(y_gt, color="k", ls="--", lw=0.8)
            ax_bps.text(
                0.0,
                y_gt,
                " ground truth",
                va="bottom",
                ha="left",
                fontsize="x-small",
                color="k",
                transform=ax_bps.get_yaxis_transform(),
            )
        if "mutual_information" in results:
            y_gt_mi = float(results.mutual_information.sel(iteration=-1).mean())
            ax_mi.axhline(y_gt_mi, color="k", ls="--", lw=0.8)
            ax_mi.text(
                0.0,
                y_gt_mi,
                " ground truth",
                va="bottom",
                ha="left",
                fontsize="x-small",
                color="k",
                transform=ax_mi.get_yaxis_transform(),
            )

    # # ML baseline
    # if "mode_l" in results and 1 in results.iteration.values:
    #     from simpl.utils import get_ML_loglikelihoods

    #     ml = get_ML_loglikelihoods(results)
    #     y_ml = ml["bits_per_spike"]
    #     ax_bps.axhline(y_ml, color="k", ls=":", lw=0.8)
    #     ax_bps.text(
    #         0.0,
    #         y_ml,
    #         " naive ML",
    #         va="bottom",
    #         ha="left",
    #         fontsize="x-small",
    #         color="k",
    #         transform=ax_bps.get_yaxis_transform(),
    #     )

    # legend on first panel
    ax_bps.plot([], [], color="gray", lw=0.8, label="train")
    ax_bps.plot([], [], color="gray", lw=0.8, ls="--", label="val")
    ax_bps.legend(fontsize="small", frameon=False)

    ax_bps.set(xlabel="Iteration", ylabel="Bits per spike")
    ax_mi.set(xlabel="Iteration", ylabel="Mutual information (bits/s)")
    for ax in axes:
        outset_axes(ax)
        ax.spines["bottom"].set_bounds(0, int(iterations[-1]))
    return axes

plot_latent_trajectory(results, time_range=None, iterations=None, include_ground_truth=True, **plot_kwargs)

Plot decoded latent trajectory (one subplot per spatial dimension).

Parameters:

Name	Type	Description	Default
results	Dataset		required
time_range	tuple	(t_start, t_end) in seconds. Default: first 120 s.	None
iterations	int or tuple of ints	Which iteration(s) to show. Negative values index from the end of the non-negative iterations (-1 = last iteration). Default: (0, -1) (behavior and final iteration).	None
include_ground_truth	bool	Show Xt as "k--" if present.	True
**plot_kwargs		Forwarded to ax.plot.	{}

Returns:

Name	Type	Description
axes	np.ndarray of Axes, shape (D,)

Source code in src/simpl/plotting.py

def plot_latent_trajectory(
    results: xr.Dataset,
    time_range: tuple[float, float] | None = None,
    iterations: int | tuple[int, ...] | None = None,
    include_ground_truth: bool = True,
    **plot_kwargs,
) -> np.ndarray:
    """Plot decoded latent trajectory (one subplot per spatial dimension).

    Parameters
    ----------
    results : xr.Dataset
    time_range : tuple, optional
        ``(t_start, t_end)`` in seconds.  Default: first 120 s.
    iterations : int or tuple of ints, optional
        Which iteration(s) to show.  Negative values index from the end of the
        non-negative iterations (``-1`` = last iteration).  Default: ``(0, -1)``
        (behavior and final iteration).
    include_ground_truth : bool
        Show ``Xt`` as ``"k--"`` if present.
    **plot_kwargs
        Forwarded to ``ax.plot``.

    Returns
    -------
    axes : np.ndarray of Axes, shape (D,)
    """
    iterations_to_plot = _resolve_iterations(iterations, results)

    if time_range is None:
        t0 = float(results.time.values[0])
        time_range = (t0, t0 + 120)

    dim_names = list(results.dim.values)
    tslice = slice(*time_range)
    t = results.time.sel(time=tslice).values
    last_iteration = _last_non_negative_iteration(results)

    traces = []
    for ep in iterations_to_plot:
        label = f"Iteration {ep} (behavior)" if ep == 0 else f"Iteration {ep}"
        traces.append(
            (
                results.X.sel(iteration=ep, time=tslice).values,
                dict(color=_iteration_color(ep, last_iteration), alpha=0.8, label=label),
            )
        )

    Xt = results.Xt.sel(time=tslice).values if (include_ground_truth and "Xt" in results) else None

    # Extract trial boundary times (shaded bands between end of one trial and start of next)
    trial_boundary_times = None
    tb_indices = results.attrs.get("trial_boundaries", None)
    if tb_indices is not None and len(tb_indices) > 1:
        all_t = results.time.values
        pairs = []
        for b in tb_indices[1:]:
            t_end = all_t[b - 1]  # last timestep of previous trial
            t_start = all_t[b]  # first timestep of next trial
            if t_start >= time_range[0] and t_end <= time_range[1]:
                pairs.append((t_end, t_start))
        if pairs:
            trial_boundary_times = np.array(pairs)

    is_angular = bool(results.attrs.get("is_1D_angular", 0))
    return _plot_trajectory_panel(
        t, traces, Xt, dim_names, trial_boundary_times=trial_boundary_times, is_1D_angular=is_angular, **plot_kwargs
    )

plot_prediction(prediction_results, Xb=None, Xt=None, time_range=None, **plot_kwargs)

Plot predicted trajectory from predict().

Parameters:

Name	Type	Description	Default
prediction_results	Dataset	The prediction_results_ Dataset from SIMPL.predict.	required
Xb	ndarray	Behavioral positions for the prediction window, shape (T, D).	None
Xt	ndarray	Ground truth positions for the prediction window, shape (T, D).	None
time_range	tuple	(t_start, t_end) in seconds. Default: full prediction range.	None
**plot_kwargs		Forwarded to ax.plot.	{}

Returns:

Name	Type	Description
axes	np.ndarray of Axes, shape (D,)

Source code in src/simpl/plotting.py

def plot_prediction(
    prediction_results: xr.Dataset,
    Xb: np.ndarray | None = None,
    Xt: np.ndarray | None = None,
    time_range: tuple[float, float] | None = None,
    **plot_kwargs,
) -> np.ndarray:
    """Plot predicted trajectory from ``predict()``.

    Parameters
    ----------
    prediction_results : xr.Dataset
        The ``prediction_results_`` Dataset from ``SIMPL.predict``.
    Xb : np.ndarray, optional
        Behavioral positions for the prediction window, shape ``(T, D)``.
    Xt : np.ndarray, optional
        Ground truth positions for the prediction window, shape ``(T, D)``.
    time_range : tuple, optional
        ``(t_start, t_end)`` in seconds.  Default: full prediction range.
    **plot_kwargs
        Forwarded to ``ax.plot``.

    Returns
    -------
    axes : np.ndarray of Axes, shape (D,)
    """
    dim_names = list(prediction_results.dim.values)

    T = len(prediction_results.time)
    if Xb is not None:
        assert Xb.shape[0] == T, f"Xb length {Xb.shape[0]} != prediction_results time length {T}"
    if Xt is not None:
        assert Xt.shape[0] == T, f"Xt length {Xt.shape[0]} != prediction_results time length {T}"

    if time_range is not None:
        tslice = slice(*time_range)
        mask = (prediction_results.time.values >= time_range[0]) & (prediction_results.time.values <= time_range[1])
    else:
        tslice = slice(None)
        mask = slice(None)

    t = prediction_results.time.sel(time=tslice).values
    traces = []
    if Xb is not None:
        traces.append((Xb[mask], dict(color=_iteration_color(0, 1), alpha=0.8, label="Behavior")))
    traces.append(
        (
            prediction_results.mu_s.sel(time=tslice).values,
            dict(color=_iteration_color(1, 1), alpha=0.8, label="Predicted"),
        )
    )

    Xt_sliced = Xt[mask] if Xt is not None else None

    # Extract trial boundary times if available
    trial_boundary_times = None
    tb_indices = prediction_results.attrs.get("trial_boundaries", None)
    if tb_indices is not None and len(tb_indices) > 1:
        all_t = prediction_results.time.values
        t0 = t[0] if len(t) > 0 else -np.inf
        t1 = t[-1] if len(t) > 0 else np.inf
        pairs = []
        for b in tb_indices[1:]:
            t_end = all_t[b - 1]
            t_start = all_t[b]
            if t_start >= t0 and t_end <= t1:
                pairs.append((t_end, t_start))
        if pairs:
            trial_boundary_times = np.array(pairs)

    is_angular = bool(prediction_results.attrs.get("is_1D_angular", 0))
    return _plot_trajectory_panel(
        t,
        traces,
        Xt_sliced,
        dim_names,
        title="Prediction on held-out data",
        trial_boundary_times=trial_boundary_times,
        is_1D_angular=is_angular,
        **plot_kwargs,
    )

plot_receptive_fields(results, extent=None, iterations=None, neurons=None, include_baselines=False, sort_by_spatial_information=False, ncols=4, threshold=0, **plot_kwargs)

Plot receptive fields for selected neurons.

Parameters:

Name	Type	Description	Default
results	Dataset		required
extent	tuple	Matplotlib extent (xmin, xmax, ymin, ymax, ...). Used for 2-D imshow.	None
iterations	int or tuple of int	Which iteration(s) to show. Negative values index from the end of the non-negative iterations (-1 = last iteration). Default: (0, -1) (behavior and final iteration).	None
neurons	array - like	Subset of neuron indices. Default: all neurons.	None
include_baselines	bool	Show ground-truth fields (Ft) if present.	False
sort_by_spatial_information	bool	If True, reorder neurons so that the most spatially informative appear first (uses the last training iteration).	False
ncols	int	Maximum number of neuron-columns in the grid.	4
**plot_kwargs		Forwarded to imshow (2-D) or plot (1-D).	{}

Returns:

Name	Type	Description
axes	np.ndarray of Axes

Source code in src/simpl/plotting.py

def plot_receptive_fields(
    results: xr.Dataset,
    extent: tuple | None = None,
    iterations: int | tuple[int, ...] | None = None,
    neurons: list[int] | np.ndarray | None = None,
    include_baselines: bool = False,
    sort_by_spatial_information: bool = False,
    ncols: int = 4,
    threshold: float = 0,
    **plot_kwargs,
) -> np.ndarray:
    """Plot receptive fields for selected neurons.

    Parameters
    ----------
    results : xr.Dataset
    extent : tuple, optional
        Matplotlib extent ``(xmin, xmax, ymin, ymax, ...)``.  Used for 2-D imshow.
    iterations : int or tuple of int, optional
        Which iteration(s) to show.  Negative values index from the end of the
        non-negative iterations (``-1`` = last iteration).  Default: ``(0, -1)``
        (behavior and final iteration).
    neurons : array-like, optional
        Subset of neuron indices.  Default: all neurons.
    include_baselines : bool
        Show ground-truth fields (``Ft``) if present.
    sort_by_spatial_information : bool
        If ``True``, reorder neurons so that the most spatially informative
        appear first (uses the last training iteration).
    ncols : int
        Maximum number of neuron-columns in the grid.
    **plot_kwargs
        Forwarded to ``imshow`` (2-D) or ``plot`` (1-D).

    Returns
    -------
    axes : np.ndarray of Axes
    """
    dim_names = list(results.dim.values)
    D = len(dim_names)
    if D > 2:
        raise ValueError(f"plot_receptive_fields only supports 1-D and 2-D environments, got {D}-D.")

    iterations = _resolve_iterations(iterations, results)

    if neurons is None:
        neurons = results.neuron.values
    neurons = np.asarray(neurons)

    if sort_by_spatial_information:
        neurons = _sort_neurons_by_si(results, neurons)

    if len(neurons) > 50:
        warnings.warn(f"Plotting {len(neurons)} neurons — this may be slow.", stacklevel=2)

    # Resolve baseline source: prefer Ft, fall back to F at iteration -1
    has_baselines = False
    baseline_label = None
    if include_baselines:
        if "Ft" in results:
            has_baselines = True
            baseline_label = "GT"
        elif -1 in results.iteration.values:
            has_baselines = True
            baseline_label = "Best"

    # Build column labels per neuron
    col_labels = []
    for ep in iterations:
        col_labels.append(f"It {ep}" if ep != 0 else "It 0 (behavior)")
    if has_baselines:
        col_labels.append(baseline_label)
    n_cols_per_neuron = len(col_labels)

    # Layout: neurons along rows, with a spacer column between neuron groups
    n_neurons = len(neurons)
    n_neuron_cols = min(n_neurons, ncols)
    n_neuron_rows = int(np.ceil(n_neurons / n_neuron_cols))

    # Each neuron group gets n_cols_per_neuron + 1 spacer, except the last group
    total_cols = n_neuron_cols * n_cols_per_neuron + (n_neuron_cols - 1)
    total_rows = n_neuron_rows

    is_polar = bool(results.attrs.get("is_1D_angular", 0)) and D == 1

    if is_polar:
        return _plot_receptive_fields_polar(
            results,
            iterations,
            neurons,
            has_baselines,
            baseline_label,
            dim_names,
            n_cols_per_neuron,
            n_neuron_cols,
            n_neuron_rows,
            total_cols,
            total_rows,
            col_labels,
            **plot_kwargs,
        )

    # Convert from spikes-per-bin to firing rate (Hz) if dt is available
    dt = results.attrs.get("dt", None)
    hz_scale = (1.0 / dt) if dt is not None else 1.0

    # Pre-compute per-neuron vmax (shared across iterations) for 2D plots
    neuron_vmax = {}
    if D == 2:
        for n in neurons:
            vals = [float(results.F.sel(iteration=ep, neuron=n).values.max()) for ep in iterations]
            if has_baselines:
                if "Ft" in results:
                    vals.append(float(results.Ft.sel(neuron=n).values.max()))
                else:
                    vals.append(float(results.F.sel(iteration=-1, neuron=n).values.max()))
            neuron_vmax[int(n)] = max(max(vals) * hz_scale, threshold)

    # Width ratios: data columns are 1, cbar columns are 0.05 (2D only), spacer columns are 0.3
    cols_per_group = n_cols_per_neuron + (1 if D == 2 else 0)
    total_cols = n_neuron_cols * cols_per_group + (n_neuron_cols - 1)

    width_ratios = []
    for g in range(n_neuron_cols):
        width_ratios.extend([1] * n_cols_per_neuron)
        if D == 2:
            width_ratios.append(0.05)
        if g < n_neuron_cols - 1:
            width_ratios.append(0.3)

    fig, axes = plt.subplots(
        total_rows,
        total_cols,
        figsize=(FIG_WIDTH / 6 * n_neuron_cols * n_cols_per_neuron, FIG_WIDTH / 6 * total_rows),
        squeeze=False,
        gridspec_kw={"width_ratios": width_ratios, "hspace": 0.1, "wspace": 0.3},
    )

    # Pre-compute the starting column index for each neuron group (skipping spacers + cbars)
    group_col_starts = []
    for g in range(n_neuron_cols):
        group_col_starts.append(g * (cols_per_group + 1))

    # Track which axes are used for plotting
    used_axes = set()
    cbar_axes = set()

    # build extent for imshow
    if D == 2 and extent is not None:
        ext = [extent[0], extent[1], extent[2], extent[3]]
    elif D == 2:
        ext = [
            float(results[dim_names[0]].values[0]),
            float(results[dim_names[0]].values[-1]),
            float(results[dim_names[1]].values[0]),
            float(results[dim_names[1]].values[-1]),
        ]
    else:
        ext = None

    imkw = dict(cmap=FIELD_CMAP, origin="lower", aspect="equal", **plot_kwargs)
    if ext is not None:
        imkw["extent"] = ext

    for idx, n in enumerate(neurons):
        row = idx // n_neuron_cols
        group = idx % n_neuron_cols
        col_base = group_col_starts[group]

        col_offset = 0
        im = None  # track last imshow for colorbar

        # Per-neuron normalization for 2D
        if D == 2:
            imkw_n = {**imkw, "vmin": threshold, "vmax": neuron_vmax[int(n)]}
        else:
            imkw_n = imkw

        # iteration columns
        for ep in iterations:
            ax = axes[row, col_base + col_offset]
            used_axes.add((row, col_base + col_offset))
            F_ep = np.clip(results.F.sel(iteration=ep, neuron=n).values * hz_scale, threshold, None)
            if D == 2:
                im = ax.imshow(F_ep.T, **imkw_n)
            else:
                ax.plot(results[dim_names[0]].values, F_ep, **plot_kwargs)
            if row == 0:
                label = f"It {ep}" if ep != 0 else "It 0 (behavior)"
                ax.set_title(label, fontsize=8)
            col_offset += 1

        # baseline column (Ft if available, else F at iteration -1)
        if has_baselines:
            ax = axes[row, col_base + col_offset]
            used_axes.add((row, col_base + col_offset))
            if "Ft" in results:
                F_base = np.clip(results.Ft.sel(neuron=n).values * hz_scale, threshold, None)
            else:
                F_base = np.clip(results.F.sel(iteration=-1, neuron=n).values * hz_scale, threshold, None)
            if D == 2:
                im = ax.imshow(F_base.T, **imkw_n)
            else:
                ax.plot(results[dim_names[0]].values, F_base, **plot_kwargs)
            if row == 0:
                ax.set_title(baseline_label, fontsize=8)

        # colorbar column for 2D
        if D == 2 and im is not None:
            cbar_col = col_base + n_cols_per_neuron
            cbar_ax = axes[row, cbar_col]
            used_axes.add((row, cbar_col))
            cbar_axes.add((row, cbar_col))
            vmax = neuron_vmax[int(n)]
            cb = fig.colorbar(im, cax=cbar_ax)
            cb.set_ticks([threshold, vmax])
            hz_label = " Hz" if dt is not None else ""
            low_label = f"<{threshold:g}" if threshold > 0 else "0"
            cb.set_ticklabels([low_label, f"{vmax:.1f}{hz_label}"])
            cb.ax.tick_params(labelsize=6)

        # label
        axes[row, col_base].set_ylabel(f"N{n}", fontsize=7, rotation=0, labelpad=15)

    # Clean up: remove ticks from used axes, turn off unused/spacer axes entirely
    for r in range(total_rows):
        for c in range(total_cols):
            ax = axes[r, c]
            if (r, c) in cbar_axes:
                ax.set_xticks([])
            elif (r, c) in used_axes:
                ax.set_xticks([])
                ax.set_yticks([])
                for spine in ax.spines.values():
                    spine.set_visible(False)
            else:
                ax.axis("off")

    # Match colorbar heights to their neighboring data axes (which may be shorter
    # than the grid cell due to aspect="equal").
    if D == 2:
        fig.canvas.draw()
        for idx, n in enumerate(neurons):
            row = idx // n_neuron_cols
            group = idx % n_neuron_cols
            col_base = group_col_starts[group]
            cbar_col = col_base + n_cols_per_neuron
            data_ax = axes[row, col_base + n_cols_per_neuron - 1]
            cbar_ax = axes[row, cbar_col]
            data_pos = data_ax.get_position()
            cbar_pos = cbar_ax.get_position()
            cbar_ax.set_position([cbar_pos.x0, data_pos.y0, cbar_pos.width, data_pos.height])

    return axes

plot_spikes(results, time_range=None, neurons=None, sort_by_spatial_information=False, cmap='Greys', **plot_kwargs)

Visualise spike counts as a heatmap (time × neurons).

Parameters:

Name	Type	Description	Default
results	Dataset	The results_ Dataset from a fitted SIMPL model.	required
time_range	tuple	(t_start, t_end) in seconds. Default: first 120 s.	None
neurons	array - like	Subset of neuron indices to display. Default: all neurons.	None
sort_by_spatial_information	bool	If True, reorder neurons so that the most spatially informative appear at the top of the heatmap (uses the last training iteration).	False
cmap	str	Colormap for imshow. Default: "Greys".	'Greys'
**plot_kwargs		Forwarded to ax.imshow.	{}

Returns:

Name	Type	Description
ax	matplotlib Axes

Source code in src/simpl/plotting.py

def plot_spikes(
    results: xr.Dataset,
    time_range: tuple[float, float] | None = None,
    neurons: list[int] | np.ndarray | None = None,
    sort_by_spatial_information: bool = False,
    cmap: str = "Greys",
    **plot_kwargs,
) -> matplotlib.axes.Axes:
    """Visualise spike counts as a heatmap (time × neurons).

    Parameters
    ----------
    results : xr.Dataset
        The ``results_`` Dataset from a fitted SIMPL model.
    time_range : tuple, optional
        ``(t_start, t_end)`` in seconds.  Default: first 120 s.
    neurons : array-like, optional
        Subset of neuron indices to display.  Default: all neurons.
    sort_by_spatial_information : bool
        If ``True``, reorder neurons so that the most spatially informative
        appear at the top of the heatmap (uses the last training iteration).
    cmap : str
        Colormap for ``imshow``.  Default: ``"Greys"``.
    **plot_kwargs
        Forwarded to ``ax.imshow``.

    Returns
    -------
    ax : matplotlib Axes
    """
    if "Y" not in results:
        raise ValueError("results Dataset does not contain 'Y' (spike counts).")

    Y = results.Y
    t = results.time.values

    # Default to first 120 s (same as plot_latent_trajectory)
    if time_range is None:
        t0 = float(t[0])
        time_range = (t0, t0 + 120)

    tslice = slice(*time_range)
    Y = Y.sel(time=tslice)
    t = Y.time.values

    # Neuron subset
    if neurons is not None:
        neurons = np.asarray(neurons)
    else:
        neurons = results.neuron.values

    if sort_by_spatial_information:
        neurons = _sort_neurons_by_si(results, neurons)

    Y = Y.sel(neuron=neurons)
    data = np.array(Y.values)  # (T, N)

    fig, ax = plt.subplots(
        1,
        1,
        figsize=(FIG_WIDTH, FIG_WIDTH * 0.35),
        layout="constrained",
    )

    extent = [float(t[0]), float(t[-1]), -0.5, data.shape[1] - 0.5]
    imkw = dict(
        cmap=cmap,
        aspect="auto",
        interpolation="none",
        origin="lower",
        extent=extent,
    )
    imkw.update(plot_kwargs)
    im = ax.imshow(data.T, **imkw)
    fig.colorbar(im, ax=ax, label="Spike count", shrink=0.8, pad=0.02)

    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Neuron")
    outset_axes(ax)

    return ax