example04_place_cells_continuous_stimulus.py

#!/usr/bin/env python3
"""Example 04 — Place-Cell Receptive Fields (Gaussian vs Zernike).

This example demonstrates:
  1) Loading hippocampal place-cell data from two animals.
  2) Visualising spike locations overlaid on the animal's path.
  3) Loading precomputed Gaussian and Zernike polynomial receptive-field fits.
  4) Comparing model families via KS, AIC, and BIC statistics using FitSummary.
  5) Generating 2-D heatmaps of place fields for all neurons.
  6) Generating 3-D mesh comparison for selected example cells.

Data provenance:
  Uses ``data/PlaceCellDataAnimal{1,2}.mat`` (trajectories + spike times)
  and ``PlaceCellAnimal{1,2}Results.mat`` (precomputed FitResult structures).

Expected outputs:
  - Figure 1: Example cells — spike locations over path (4 cells per animal).
  - Figure 2: Population model-comparison statistics (dKS, dAIC, dBIC).
  - Figure 3: Gaussian receptive-field heatmaps (Animal 1).
  - Figure 4: Zernike receptive-field heatmaps (Animal 1).
  - Figure 5: Gaussian receptive-field heatmaps (Animal 2).
  - Figure 6: Zernike receptive-field heatmaps (Animal 2).
  - Figure 7: 3-D mesh comparison for selected example cells.

Paper mapping:
  Section 2.3.5 (place-cell continuous-stimulus analysis).
"""
from __future__ import annotations

import argparse
import math
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy.io import loadmat

THIS_DIR = Path(__file__).resolve().parent
REPO_ROOT = THIS_DIR.parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from nstat import (  # noqa: E402
    Covariate,
    FitResult,
    FitSummary,
    TrialConfig,
    ConfigCollection,
)
from nstat.core import nspikeTrain  # noqa: E402
from nstat.data_manager import ensure_example_data  # noqa: E402
from nstat.zernike import zernike_basis_from_cartesian  # noqa: E402


# =====================================================================
# Helpers
# =====================================================================
def _load_animal_data(path):
    """Load place-cell trajectory and spike data from a .mat file."""
    d = loadmat(str(path), squeeze_me=True)
    x = np.asarray(d["x"], dtype=float).ravel()
    y = np.asarray(d["y"], dtype=float).ravel()
    time = np.asarray(d["time"], dtype=float).ravel()
    neurons = np.asarray(d["neuron"], dtype=object).ravel()
    return x, y, time, neurons


def _load_animal_results(path, x, y, time, neurons):
    """Load precomputed FitResult structures and reconstruct Python FitResults."""
    d = loadmat(str(path), squeeze_me=True)
    res_structs = np.asarray(d["resStruct"], dtype=object).ravel()
    fit_results = []

    for i, rs in enumerate(res_structs):
        # Extract lambda signal
        lam = rs["lambda"].item()
        lam_time = np.asarray(lam["time"].item(), dtype=float).ravel()
        lam_data = np.asarray(lam["data"].item(), dtype=float)
        lam_name = str(lam["name"].item()) if lam["name"].size else "\\lambda"

        lambda_cov = Covariate(
            lam_time, lam_data, lam_name, "time", "s", "spikes/sec",
        )

        # Extract coefficients
        b_raw = rs["b"].item()
        if b_raw.dtype == object:
            b_list = [np.asarray(b_raw[j], dtype=float).ravel() for j in range(b_raw.size)]
        else:
            b_list = [np.asarray(b_raw, dtype=float).ravel()]

        numResults = int(np.asarray(rs["numResults"].item()).ravel()[0])

        # Extract AIC/BIC/logLL
        AIC = np.asarray(rs["AIC"].item(), dtype=float).ravel()
        BIC = np.asarray(rs["BIC"].item(), dtype=float).ravel()
        logLL = np.asarray(rs["logLL"].item(), dtype=float).ravel()

        # Config names
        cn_raw = rs["configNames"].item()
        if isinstance(cn_raw, np.ndarray):
            config_names = [str(c) for c in cn_raw.ravel()]
        else:
            config_names = [str(cn_raw)]

        # Covariate labels
        if "covLabels" in rs.dtype.names:
            cl_raw = rs["covLabels"].item()
            cl = cl_raw
        else:
            cl = []
        if isinstance(cl, np.ndarray) and cl.dtype == object:
            cov_labels = []
            for j in range(cl.size):
                item = cl[j]
                if isinstance(item, np.ndarray):
                    cov_labels.append([str(x) for x in item.ravel()])
                else:
                    cov_labels.append([str(item)])
        elif isinstance(cl, str):
            cov_labels = [[cl]] * numResults
        else:
            cov_labels = [config_names] * numResults

        # Create spike train for this neuron
        st = np.asarray(neurons[i]["spikeTimes"].item(), dtype=float).ravel()
        nst = nspikeTrain(st, name=str(i + 1),
                          minTime=float(time[0]), maxTime=float(time[-1]),
                          makePlots=-1)

        # numHist
        if "numHist" in rs.dtype.names:
            nh = rs["numHist"].item()
            num_hist = list(np.asarray(nh, dtype=int).ravel())
        else:
            num_hist = [0] * numResults

        cfgs = ConfigCollection([TrialConfig(name=n) for n in config_names])

        fr = FitResult(
            nst,
            cov_labels,
            num_hist,
            [],              # histObjects
            [],              # ensHistObjects
            lambda_cov,
            b_list,
            [0.0] * numResults,  # dev
            [None] * numResults, # stats
            AIC,
            BIC,
            logLL,
            cfgs,
            [],              # XvalData
            [],              # XvalTime
            "poisson",
        )

        # Load KS statistics if available
        if "KSStats" in rs.dtype.names:
            ks_struct = rs["KSStats"].item()
            if hasattr(ks_struct, "dtype") and ks_struct.dtype.names:
                ks_stat = np.asarray(ks_struct["ks_stat"].item(), dtype=float).ravel()
                pval = np.asarray(ks_struct["pValue"].item(), dtype=float).ravel()
                within = np.asarray(ks_struct["withinConfInt"].item(), dtype=float).ravel()
                if ks_stat.size >= numResults:
                    fr.KSStats = ks_stat[:numResults].reshape(numResults, 1)
                    fr.KSPvalues = pval[:numResults]
                    fr.withinConfInt = within[:numResults]

        fit_results.append(fr)

    return fit_results


def _compute_place_field(coeffs, grid_design, grid_shape, sample_rate=1.0):
    """Compute predicted firing rate on a spatial grid.

    Matches Matlab ``FitResult.evalLambda`` which computes
    ``exp(X * b) * sampleRate`` to convert from conditional intensity
    (per bin) to firing rate (Hz).
    """
    eta = grid_design @ coeffs
    rate = np.exp(eta) * sample_rate
    return rate.reshape(grid_shape)


# =====================================================================
# Main example
# =====================================================================
def run_example04(*, export_figures: bool = False, export_dir: Path | None = None):
    """Run Example 04: Place-cell receptive fields."""
    print("=== Example 04: Place-Cell Receptive Fields ===")

    data_dir = ensure_example_data(download=True)
    if export_dir is None:
        export_dir = THIS_DIR / "figures" / "example04"

    # ==================================================================
    # 1. Load data for both animals
    # ==================================================================
    x1, y1, t1, neurons1 = _load_animal_data(
        data_dir / "Place Cells" / "PlaceCellDataAnimal1.mat")
    x2, y2, t2, neurons2 = _load_animal_data(
        data_dir / "Place Cells" / "PlaceCellDataAnimal2.mat")
    nCells1 = len(neurons1)
    nCells2 = len(neurons2)
    print(f"  Animal 1: {nCells1} cells, {len(t1)} time points")
    print(f"  Animal 2: {nCells2} cells, {len(t2)} time points")

    # ==================================================================
    # 2. Load precomputed FitResults
    # ==================================================================
    fitResults1 = _load_animal_results(
        data_dir / "PlaceCellAnimal1Results.mat", x1, y1, t1, neurons1)
    fitResults2 = _load_animal_results(
        data_dir / "PlaceCellAnimal2Results.mat", x2, y2, t2, neurons2)
    print(f"  Loaded {len(fitResults1)} + {len(fitResults2)} FitResult objects")

    # ==================================================================
    # 3. Build FitSummary for each animal
    # ==================================================================
    summary1 = FitSummary(fitResults1)
    summary2 = FitSummary(fitResults2)

    # Delta statistics
    # dKS: direct subtraction Gaussian - Zernike (matches Matlab line 81-83)
    dKS1 = summary1.KSStats[:, 0] - summary1.KSStats[:, 1]
    dKS2 = summary2.KSStats[:, 0] - summary2.KSStats[:, 1]

    # dAIC/dBIC: Matlab uses getDiffAIC(1) / getDiffBIC(1) which computes
    # Zernike - Gaussian (other columns minus reference column).
    dAIC1 = summary1.AIC[:, 1] - summary1.AIC[:, 0]
    dBIC1 = summary1.BIC[:, 1] - summary1.BIC[:, 0]

    dAIC2 = summary2.AIC[:, 1] - summary2.AIC[:, 0]
    dBIC2 = summary2.BIC[:, 1] - summary2.BIC[:, 0]

    dAIC_all = np.concatenate([dAIC1, dAIC2])
    dBIC_all = np.concatenate([dBIC1, dBIC2])
    dKS_all = np.concatenate([dKS1, dKS2])

    print(f"  Mean dKS  (Gauss-Zern): {np.nanmean(dKS_all):.4f}")
    print(f"  Mean dAIC (Zern-Gauss): {np.nanmean(dAIC_all):.2f}")
    print(f"  Mean dBIC (Zern-Gauss): {np.nanmean(dBIC_all):.2f}")

    # ==================================================================
    # Figure 1: Example cells — spike locations over path (2x2)
    # ==================================================================
    exampleCells = [1, 20, 24, 48]  # 0-indexed (MATLAB: [2 21 25 49])
    fig1, axes1 = plt.subplots(2, 2, figsize=(14, 9))  # MATLAB: 1400x900
    for i, cidx in enumerate(exampleCells):
        ax = axes1.flat[i]
        h1, = ax.plot(x1, y1, "b-", linewidth=0.5)
        n = neurons1[min(cidx, nCells1 - 1)]
        xn = np.asarray(n["xN"].item(), dtype=float).ravel()
        yn = np.asarray(n["yN"].item(), dtype=float).ravel()
        h2, = ax.plot(xn, yn, "r.", markersize=7)
        ax.set_title(f"Cell#{cidx + 1}", fontweight="bold", fontsize=12)
        ax.set_xticks(np.arange(-1, 1.5, 0.5))
        ax.set_yticks(np.arange(-1, 1.5, 0.5))
        ax.set_aspect("equal")
        ax.axis("square")
        if i == 3:
            ax.legend([h1, h2], ["Animal Path", "Location at time of spike"])
    fig1.tight_layout()

    # ==================================================================
    # Figure 2: Population statistics (1x3 box plots)
    # ==================================================================
    fig2, axes2 = plt.subplots(1, 3, figsize=(14, 9))  # MATLAB: 1400x900

    axes2[0].boxplot([dKS1[np.isfinite(dKS1)], dKS2[np.isfinite(dKS2)]],
                     tick_labels=["Animal 1", "Animal 2"],
                     vert=True)
    axes2[0].set_title(r"$\Delta$ KS Statistic", fontsize=14, fontweight="bold",
                       fontfamily="Arial")

    axes2[1].boxplot([dAIC1[np.isfinite(dAIC1)], dAIC2[np.isfinite(dAIC2)]],
                     tick_labels=["Animal 1", "Animal 2"],
                     vert=True)
    axes2[1].set_title(r"$\Delta$ AIC", fontsize=14, fontweight="bold",
                       fontfamily="Arial")

    axes2[2].boxplot([dBIC1[np.isfinite(dBIC1)], dBIC2[np.isfinite(dBIC2)]],
                     tick_labels=["Animal 1", "Animal 2"],
                     vert=True)
    axes2[2].set_title(r"$\Delta$ BIC", fontsize=14, fontweight="bold",
                       fontfamily="Arial")

    fig2.tight_layout()

    # ==================================================================
    # 4. Build spatial grids and design matrices for heatmaps
    # ==================================================================
    grid_res = 201  # Matlab: meshgrid(-1:0.01:1) → 201 points
    xGrid = np.linspace(-1, 1, grid_res)
    yGrid = np.linspace(-1, 1, grid_res)
    xx, yy = np.meshgrid(xGrid, yGrid)
    yy = np.flipud(yy)   # Matlab: y increases bottom-to-top
    xx = np.fliplr(xx)    # Matlab: x increases right-to-left
    xf, yf = xx.ravel(), yy.ravel()

    # Gaussian design: [1, x, y, x^2, y^2, xy] (intercept prepended)
    gridDesignGauss = np.column_stack([
        np.ones(xf.size), xf, yf, xf**2, yf**2, xf * yf
    ])

    # Zernike design: [1, z1, z2, ..., z9] (intercept prepended)
    zBasis = zernike_basis_from_cartesian(xf, yf, fill_value=0.0)
    gridDesignZern = np.column_stack([np.ones(xf.size), zBasis])

    # ==================================================================
    # Figures 3-6: Place field heatmaps
    # ==================================================================
    def _plot_heatmaps(fit_results, nCells, title_prefix, design_gauss,
                       design_zern, grid_shape):
        nRows = math.ceil(nCells / 7)
        nCols = 7

        figG, axesG = plt.subplots(nRows, nCols, figsize=(14, 9))
        figZ, axesZ = plt.subplots(nRows, nCols, figsize=(14, 9))
        if nRows == 1:
            axesG = axesG[np.newaxis, :]
            axesZ = axesZ[np.newaxis, :]

        for i in range(nCells):
            row, col = divmod(i, nCols)
            fr = fit_results[i]
            sr = float(fr.lambda_signal.sampleRate)
            coeffs_g = np.asarray(fr.b[0], dtype=float).ravel()
            coeffs_z = np.asarray(fr.b[1], dtype=float).ravel() if fr.numResults > 1 else coeffs_g

            # Gaussian field
            ax = axesG[row, col]
            try:
                field_g = _compute_place_field(coeffs_g, design_gauss[:, :coeffs_g.size], grid_shape, sr)
                ax.pcolormesh(xx, yy, field_g, shading="gouraud", cmap="jet")
            except Exception:
                pass
            ax.set_aspect("equal")
            ax.set_xticks([])
            ax.set_yticks([])
            ax.spines["top"].set_visible(False)
            ax.spines["right"].set_visible(False)
            ax.spines["left"].set_visible(False)
            ax.spines["bottom"].set_visible(False)

            # Zernike field
            ax = axesZ[row, col]
            try:
                field_z = _compute_place_field(coeffs_z, design_zern[:, :coeffs_z.size], grid_shape, sr)
                ax.pcolormesh(xx, yy, field_z, shading="gouraud", cmap="jet")
            except Exception:
                pass
            ax.set_aspect("equal")
            ax.set_xticks([])
            ax.set_yticks([])
            ax.spines["top"].set_visible(False)
            ax.spines["right"].set_visible(False)
            ax.spines["left"].set_visible(False)
            ax.spines["bottom"].set_visible(False)

        # Hide unused subplots
        for i in range(nCells, nRows * nCols):
            row, col = divmod(i, nCols)
            axesG[row, col].set_visible(False)
            axesZ[row, col].set_visible(False)

        # MATLAB: sgtitle('Gaussian Place Fields - Animal#N')
        animal_num = title_prefix.replace("Animal ", "")
        figG.suptitle(f"Gaussian Place Fields - Animal#{animal_num}",
                      fontweight="bold", fontsize=12)
        figZ.suptitle(f"Zernike Place Fields - Animal#{animal_num}",
                      fontweight="bold", fontsize=12)
        figG.tight_layout()
        figZ.tight_layout()
        return figG, figZ

    figG1, figZ1 = _plot_heatmaps(
        fitResults1, nCells1, "Animal 1",
        gridDesignGauss, gridDesignZern, xx.shape,
    )
    figG2, figZ2 = _plot_heatmaps(
        fitResults2, nCells2, "Animal 2",
        gridDesignGauss, gridDesignZern, xx.shape,
    )
    print("  Figures 3-6: Place field heatmaps")

    # ==================================================================
    # Figure 7: 3-D mesh comparison for an example cell
    # ==================================================================
    exampleCell = min(24, nCells1 - 1)  # 0-indexed → cell 25 in Matlab
    fr_ex = fitResults1[exampleCell]
    sr_ex = float(fr_ex.lambda_signal.sampleRate)
    coeffs_g = np.asarray(fr_ex.b[0], dtype=float).ravel()
    coeffs_z = np.asarray(fr_ex.b[1], dtype=float).ravel()

    field_g = _compute_place_field(
        coeffs_g, gridDesignGauss[:, :coeffs_g.size], xx.shape, sr_ex)
    field_z = _compute_place_field(
        coeffs_z, gridDesignZern[:, :coeffs_z.size], xx.shape, sr_ex)

    fig7 = plt.figure(figsize=(14, 9))
    ax3d = fig7.add_subplot(111, projection="3d")
    # MATLAB: mesh(xGrid, yGrid, lambda, 'FaceAlpha',0.2, 'EdgeAlpha',0.2, 'EdgeColor','b'/'g')
    # plot_wireframe matches MATLAB's mesh() (wireframe only, no filled faces)
    ax3d.plot_wireframe(xx, yy, field_g, color="b", alpha=0.2,
                        rstride=5, cstride=5, linewidth=0.3)
    ax3d.plot_wireframe(xx, yy, field_z, color="g", alpha=0.2,
                        rstride=5, cstride=5, linewidth=0.3)
    # Overlay animal path at z=0 (MATLAB: 'k')
    ax3d.plot(x1, y1, np.zeros_like(x1), "k-", linewidth=0.3)
    # Overlay spike locations (MATLAB: 'r.')
    n_ex = neurons1[exampleCell]
    xn_ex = np.asarray(n_ex["xN"].item(), dtype=float).ravel()
    yn_ex = np.asarray(n_ex["yN"].item(), dtype=float).ravel()
    ax3d.scatter(xn_ex, yn_ex, np.zeros_like(xn_ex), c="r", s=5)
    ax3d.set_xlabel("x position")
    ax3d.set_ylabel("y position")
    ax3d.set_title(f"Animal#1, Cell#{exampleCell + 1}",
                   fontweight="bold", fontsize=12)
    # MATLAB legend (wireframe lines, not filled patches)
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], color="b", alpha=0.5, linewidth=1.5,
               label=r"$\lambda_{Gaussian}$"),
        Line2D([0], [0], color="g", alpha=0.5, linewidth=1.5,
               label=r"$\lambda_{Zernike}$"),
        Line2D([0], [0], color="k", linewidth=0.5, label="Animal Path"),
        Line2D([0], [0], marker=".", color="r", linestyle="",
               label="Spike Locations"),
    ]
    ax3d.legend(handles=legend_elements, loc="best")

    print(f"  Figure 7: 3D mesh for cell {exampleCell + 1}")

    # ==================================================================
    # Save figures
    # ==================================================================
    all_figs = {
        "fig01_example_cells_path_overlay": fig1,
        "fig02_model_summary_statistics": fig2,
        "fig03_gaussian_place_fields_animal1": figG1,
        "fig04_zernike_place_fields_animal1": figZ1,
        "fig05_gaussian_place_fields_animal2": figG2,
        "fig06_zernike_place_fields_animal2": figZ2,
        "fig07_example_cell_mesh_comparison": fig7,
    }

    if export_figures:
        export_dir.mkdir(parents=True, exist_ok=True)
        for name, fig in all_figs.items():
            path = export_dir / f"{name}.png"
            fig.savefig(str(path), dpi=250, facecolor="w", edgecolor="none")
            print(f"  Saved {path}")

    plt.show()
    print(f"\nExample 04 complete. Generated {len(all_figs)} figure(s).")
    return all_figs


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Example 04: Place-Cell Receptive Fields"
    )
    parser.add_argument("--repo-root", type=Path, default=REPO_ROOT)
    parser.add_argument("--export-figures", action="store_true")
    parser.add_argument("--export-dir", type=Path, default=None)
    args = parser.parse_args()

    run_example04(
        export_figures=args.export_figures,
        export_dir=args.export_dir,
    )