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)

Plot all metrics directly from an xarray results Dataset.

See :meth:simpl.SIMPL.plot_all_metrics for the user-facing API.

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:
    """Plot all metrics directly from an xarray results Dataset.

    See :meth:`simpl.SIMPL.plot_all_metrics` for the user-facing API.
    """

    # 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]
        plot_metric(results, var_name, show_neurons=show_neurons, ax=ax, **plot_kwargs)

    # 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)

Plot a fitting summary directly from an xarray results Dataset.

See :meth:simpl.SIMPL.plot_fitting_summary for the user-facing API.

Source code in src/simpl/plotting.py

def plot_fitting_summary(
    results: xr.Dataset,
    show_neurons: bool = True,
    **plot_kwargs,
) -> np.ndarray:
    """Plot a fitting summary directly from an xarray results Dataset.

    See :meth:`simpl.SIMPL.plot_fitting_summary` for the user-facing API.
    """
    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 latent trajectories directly from an xarray results Dataset.

See :meth:simpl.SIMPL.plot_latent_trajectory for the user-facing API.

Source code in src/simpl/plotting.py

def plot_latent_trajectory(
    results: xr.Dataset,
    time_range: float | tuple[float, float] | None = None,
    iterations: int | tuple[int, ...] | None = None,
    include_ground_truth: bool = True,
    **plot_kwargs,
) -> np.ndarray:
    """Plot latent trajectories directly from an xarray results Dataset.

    See :meth:`simpl.SIMPL.plot_latent_trajectory` for the user-facing API.
    """
    iterations_to_plot = _resolve_iterations(iterations, results)

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

    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 predictions directly from an xarray prediction-results Dataset.

See :meth:simpl.SIMPL.plot_prediction for the user-facing API.

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 predictions directly from an xarray prediction-results Dataset.

    See :meth:`simpl.SIMPL.plot_prediction` for the user-facing API.
    """
    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:
        if len(time_range) != 2:
            raise ValueError("time_range must contain exactly two values: (t_start, t_end).")
        t_start, t_end = time_range
        available_start = float(prediction_results.time.values[0])
        available_end = float(prediction_results.time.values[-1])
        if t_start > t_end:
            raise ValueError(f"time_range start ({t_start}) must not exceed end ({t_end}).")
        if t_start < available_start or t_end > available_end:
            raise ValueError(
                f"time_range ({t_start}, {t_end}) is outside the available prediction range "
                f"({available_start}, {available_end})."
            )
        tslice = slice(t_start, t_end)
        mask = (prediction_results.time.values >= t_start) & (prediction_results.time.values <= t_end)
    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, max_neurons=None, ncols=4, threshold=0, **plot_kwargs)

Plot receptive fields directly from an xarray results Dataset.

See :meth:simpl.SIMPL.plot_receptive_fields for the user-facing API.

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,
    max_neurons: int | None = None,
    ncols: int = 4,
    threshold: float = 0,
    **plot_kwargs,
) -> np.ndarray:
    """Plot receptive fields directly from an xarray results Dataset.

    See :meth:`simpl.SIMPL.plot_receptive_fields` for the user-facing API.
    """
    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 max_neurons is not None:
        neurons = neurons[:max_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)

Plot spike counts directly from an xarray results Dataset.

See :meth:simpl.SIMPL.plot_spikes for the user-facing API.

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:
    """Plot spike counts directly from an xarray results Dataset.

    See :meth:`simpl.SIMPL.plot_spikes` for the user-facing API.
    """
    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 (index in displayed order)")
    outset_axes(ax)

    return ax